Fix: unsqueeze k/v when dim==2

2026-05-27 06:52:43 +00:00
parent fc4172937c
commit 36a6f07a7e
1 changed files with 2 additions and 2 deletions
--- 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,