From 36a6f07a7ec449b16b796a9e42c90614faa8a3af Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 06:52:43 +0000 Subject: [PATCH] Fix: unsqueeze k/v when dim==2 --- dsv4/kernels/attention/production.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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,