 template <typename T, int Dh, int Dh_MAX, int THREADS_PER_KEY,
           int THREADS_PER_VALUE, int THREADS_PER_BLOCK>
 __global__ void masked_multihead_attention_kernel(
-    Masked_multihead_attention_params<T> params, int pad_active_groups) {
+    Masked_multihead_attention_params<T> params) {
 #if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
 
-  static_assert(Dh % THREADS_PER_KEY == 0, "");
-  static_assert(Dh % THREADS_PER_VALUE == 0, "");
+  static_assert(Dh_MAX % THREADS_PER_KEY == 0, "");
+  static_assert(Dh_MAX % THREADS_PER_VALUE == 0, "");
 
   constexpr int WARP_SIZE = 32;
   constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE;

    case 128:
      fmha_launch_kernel<T, 128, 128>(params, dev_ctx.stream());
      break;
    default:
      PADDLE_THROW(platform::errors::Unimplemented(
          "dim_head = %d is unsupport, only support "
          "dim_head = 32, 64, 96 or 128 for now.",
          dim_head));
  }
}



  using K_vec = typename K_vec_<T, THREADS_PER_KEY>::Type;
  constexpr int K_VEC_SIZE = sizeof(K_vec) / sizeof(T);
  static_assert(Dh_MAX % K_VEC_SIZE == 0, "");
  constexpr int K_ELTS_PER_THREAD = Dh_MAX / THREADS_PER_KEY;
  constexpr int K_VECS_PER_THREAD = K_ELTS_PER_THREAD / K_VEC_SIZE;

  int ko = tid / THREADS_PER_KEY;
  int ki = (tid % THREADS_PER_KEY) * K_VEC_SIZE;

  static_assert(Dh_MAX == THREADS_PER_KEY * K_VEC_SIZE * K_VECS_PER_THREAD, "");

  K_vec q[K_VECS_PER_THREAD];
#pragma unroll
