   zero(out);
 
   constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE;
-  for (int ti = vo; ti < params.timestep; ti += V_PER_ITER) {
-    V_vec v = *reinterpret_cast<const V_vec *>(&v_cache[ti * Dh]);
+  if (vo < V_PER_ITER) {
+    for (int ti = vo; ti < params.timestep; ti += V_PER_ITER) {
+      V_vec v = *reinterpret_cast<const V_vec *>(&v_cache[ti * Dh]);
 #if defined(MMHA_USE_FP32_ACUM_FOR_LOGITS)
-    float logit = logits_smem[ti];
-    out = fma(logit, cast_to_float(v), out);
+      float logit = logits_smem[ti];
+      out = fma(logit, cast_to_float(v), out);
 #else
-    T logit = logits_smem[ti];
-    // Update the partial sums.
-    out = fma(logit, v, out);
-#endif
+      T logit = logits_smem[ti];
+      // Update the partial sums.
+      out = fma(logit, v, out);
+#endif
+    }
   }
 
 #ifdef _DEBUG_FUSED_MULTI_TRANSFORMER

                 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;
    }

  dst = tmp.raw;
}

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) {
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)

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