                  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 = dot<Qk_vec, Qk_vec>(q, k);
+  }
+  if (tid < WARP_SIZE) {
+    for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
+      qk += __shfl_xor_sync(uint32_t(-1), 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 *= params.inv_sqrt_dh;
       qk_max = qk;
       qk_smem[params.timestep] = qk;
     }

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) {
  size_t qk_sz = div_up(params.timestep + 1, 4) * 16;
  size_t logits_sz = 0;


  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;
