   zero(out);
 
   constexpr int V_PER_ITER = THREADS_PER_BLOCK / THREADS_PER_VALUE;
-  if (vo < V_PER_ITER) {
+  if (Dh == Dh_MAX || vi < Dh) {
     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)

    int qk_offset = qkv_base_offset + tid * QK_VEC_SIZE;
    int qk_bias_offset = hi * Dh + tid * QK_VEC_SIZE;

    Qk_vec q = *reinterpret_cast<const Qk_vec *>(&q_base[qk_offset]);
    Qk_vec k = *reinterpret_cast<const Qk_vec *>(&k_base[qk_offset]);

    Qk_vec q_bias =
        *reinterpret_cast<const Qk_vec *>(&q_bias_base[qk_bias_offset]);
    Qk_vec k_bias =
        *reinterpret_cast<const Qk_vec *>(&k_bias_base[qk_bias_offset]);

    q = add(q, q_bias);
    // TODO(wangxi): See this https://github.com/microsoft/unilm/issues/510

  params.inv_sqrt_dh = inv_sqrt_dh;

  switch (dim_head) {
    case 10:
      fmha_launch_kernel<T, 10, 32>(params, dev_ctx.stream());
      break;
    case 26:
      fmha_launch_kernel<T, 26, 32>(params, dev_ctx.stream());
      break;
    case 32:
      fmha_launch_kernel<T, 32, 32>(params, dev_ctx.stream());
      break;
