 #endif
   size_t softmax_sz = qk_sz + logits_sz;
 
-  int rows_per_red = threads_per_block / threads_per_value;
+  int rows_per_red = pad_active_groups;
   size_t red_sz = rows_per_red * dim_head * sizeof(T) / 2;
 
   return max(softmax_sz, red_sz);
 }
 
-#define MMHA_LAUNCH_KERNEL(T, Dh, THDS_PER_KEY, THDS_PER_VALUE,          \
-                           THDS_PER_BLOCK, stream)                       \
-  size_t smem_sz =                                                       \
-      smem_size_in_bytes<T>(params, Dh, THDS_PER_VALUE, THDS_PER_BLOCK); \
-  dim3 grid(params.num_head, params.batch_size);                         \
-  masked_multihead_attention_kernel<                                     \
-      T, Dh, THDS_PER_KEY, THDS_PER_VALUE,                               \
-      THDS_PER_BLOCK><<<grid, THDS_PER_BLOCK, smem_sz, stream>>>(params)
-
-template <typename T, int Dh>
+#define MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE,        \
+                           THDS_PER_BLOCK, stream)                             \
+  int pad_active_groups =                                                      \
+      1 << static_cast<int>(ceil(std::log2(THDS_PER_BLOCK / THDS_PER_VALUE))); \
+  size_t smem_sz = smem_size_in_bytes<T>(params, Dh, THDS_PER_VALUE,           \
+                                         THDS_PER_BLOCK, pad_active_groups);   \
+  dim3 grid(params.num_head, params.batch_size);                               \
+  masked_multihead_attention_kernel<                                           \
+      T, Dh, Dh_MAX, THDS_PER_KEY, THDS_PER_VALUE,                             \
+      THDS_PER_BLOCK><<<grid, THDS_PER_BLOCK, smem_sz, stream>>>(              \
+      params, pad_active_groups)
+
+template <typename T, int Dh, int Dh_MAX>
 void fmha_launch_kernel(const Masked_multihead_attention_params<T> &params,
                         const cudaStream_t &stream) {
   constexpr int THREADS_PER_VALUE = Dh * sizeof(T) / 16;
   if (params.timestep < 32) {
-    MMHA_LAUNCH_KERNEL(T, Dh, 4, THREADS_PER_VALUE, 64, stream);
+    MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 4, THREADS_PER_VALUE, 64, stream);
   } else if (params.timestep < 2048) {
-    MMHA_LAUNCH_KERNEL(T, Dh, 2, THREADS_PER_VALUE, 128, stream);
+    MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 2, THREADS_PER_VALUE, 128, stream);
   } else {
-    MMHA_LAUNCH_KERNEL(T, Dh, 1, THREADS_PER_VALUE, 256, stream);
+    MMHA_LAUNCH_KERNEL(T, Dh, Dh_MAX, 1, THREADS_PER_VALUE, 256, stream);
   }
 }
 

  dst = tmp.raw;
}

template <typename T, int Dh, int THREADS_PER_KEY, int THREADS_PER_VALUE,
          int THREADS_PER_BLOCK>
__global__ void masked_multihead_attention_kernel(
    Masked_multihead_attention_params<T> params) {
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)

  static_assert(Dh % THREADS_PER_KEY == 0, "");

template <typename T>
inline size_t smem_size_in_bytes(
    const Masked_multihead_attention_params<T> &params, int dim_head,
    int threads_per_value, int threads_per_block, int pad_active_groups) {
  size_t qk_sz = div_up(params.timestep + 1, 4) * 16;
  size_t logits_sz = 0;

