diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py
index 6dea0070..fb527c40 100644
--- a/dsv4/kernels/attention/production.py
+++ b/dsv4/kernels/attention/production.py
@@ -96,8 +96,8 @@ def dsv4_attention(
 
     for h in range(n_h):
         q_h = q[h:h+1]  # (1, T, hd)
-        k_h = k if k.dim() == 2 else k[h:h+1]  # (1, N, hd)
-        v_h = v if v.dim() == 2 else v[h:h+1]
+        k_h = k.unsqueeze(0) if k.dim() == 2 else k[h:h+1]  # (1, N, hd)
+        v_h = v.unsqueeze(0) if v.dim() == 2 else v[h:h+1]  # (1, N, hd)
 
         o_h = _attention_single_head(
             q_h, k_h, v_h, scale=scale,