 
   switch (dim_head) {
     case 32:
-      fmha_launch_kernel<T, 32>(params, dev_ctx.stream());
+      fmha_launch_kernel<T, 32, 32>(params, dev_ctx.stream());
       break;
     case 64:
-      fmha_launch_kernel<T, 64>(params, dev_ctx.stream());
+      fmha_launch_kernel<T, 64, 64>(params, dev_ctx.stream());
+      break;
+    case 96:
+      fmha_launch_kernel<T, 96, 128>(params, dev_ctx.stream());
       break;
     case 128:
-      fmha_launch_kernel<T, 128>(params, dev_ctx.stream());
+      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 or 128 for now.",
+          "dim_head = 32, 64, 96 or 128 for now.",
           dim_head));
   }
 }

                 params.timestep * QK_ELTS_IN_16B + ci;
    *reinterpret_cast<Qk_vec *>(&params.cache_kv[offset]) = k;

    float qk = dot<Qk_vec, Qk_vec>(q, k);
#pragma unroll
    for (int mask = QK_VECS_PER_WARP / 2; mask >= 1; mask /= 2) {
      qk += __shfl_xor_sync(shfl_mask(QK_VECS_PER_WARP), qk, mask);
    }

    qk *= params.inv_sqrt_dh;
    if (tid == 0) {
      // NOTE(wangxi): mask must be 0.0
      // T mask = params.attn_mask[
      //    bi * (params.timestep + 1) + params.timestep];
      // qk += static_cast<float>(mask);
      qk_max = qk;
      qk_smem[params.timestep] = qk;
    }

  const int tid = threadIdx.x;

  float qk_max = -FLT_MAX;
  float qk = 0;

  // qkv [B, S=1, 3, num_head, head_dim]
  int qkv_base_offset = bi * 3 * params.num_head * Dh + hi * Dh;

  using Qk_vec = typename Qk_vec_<T, Dh_MAX>::Type;
  constexpr int QK_VEC_SIZE = sizeof(Qk_vec) / sizeof(T);
  static_assert(Dh % QK_VEC_SIZE == 0 && Dh / QK_VEC_SIZE <= WARP_SIZE, "");
  constexpr int QK_VECS_PER_WARP = Dh / QK_VEC_SIZE;
