From 57615029a4e5edc9bc71d56a91821b1d4678430c Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Tue, 19 May 2026 08:00:08 +0000
Subject: [PATCH] =?UTF-8?q?Fix=20KV=20expand=20for=20SDPA:=20(T,HD)=20?=
 =?UTF-8?q?=E2=86=92=20(T*NH,=20T,=20HD)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cutedsl/csa_attention.py | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/cutedsl/csa_attention.py b/cutedsl/csa_attention.py
index 7f30eab9..53bbaa02 100644
--- a/cutedsl/csa_attention.py
+++ b/cutedsl/csa_attention.py
@@ -392,15 +392,22 @@ def full_attention_reference(
     """
     T, NH, HD = q.shape
 
-    # K=V from kv latent (shared across heads, so expand)
-    # kv: (T, HD) → broadcast to all heads and all query positions
-    k = kv.unsqueeze(0).unsqueeze(2).expand(T, NH, T, -1).contiguous()  # (T, NH, T, HD)
-    v = k.clone()
+    # K=V from kv latent (shared across all heads and all query positions)
+    # kv: (T, HD) → each token's KV is seen by all heads at all query positions
+    k = kv.unsqueeze(1).expand(-1, NH, -1).contiguous()  # (T, NH, HD)
+    # For cross-attention where each Q attends to all KV positions:
+    # K needs to be (T_q, NH, T_kv, HD) — repeat for each query position
+    k = k.unsqueeze(0).expand(T, -1, -1, -1).contiguous()  # (T, T, NH, HD) → wrong order
+    # Actually: for self-attention, K/V shape for SDPA is (batch, seq_kv, HD)
+    # where batch = T*NH (each query token is a batch, each head independent)
+    # K/V: (T*NH, T, HD) — each (query, head) pair attends to all T KV positions
+    kv_expanded = kv.unsqueeze(1).expand(-1, NH, -1).contiguous()  # (T, NH, HD)
+    # Repeat KV for each query: (T, NH, HD) → (T*NH, T, HD)
+    k_2d = kv_expanded.permute(1, 0, 2).unsqueeze(1).expand(NH, T, T, -1).contiguous().reshape(T * NH, T, HD)
+    v_2d = k_2d.clone()
 
-    # Reshape for SDPA: Q (T*NH, 1, HD), K (T*NH, T, HD), V (T*NH, T, HD)
+    # Q: (T, NH, HD) → (T*NH, 1, HD)
     q_2d = q.reshape(T * NH, 1, HD)
-    k_2d = k.reshape(T * NH, T, HD)
-    v_2d = v.reshape(T * NH, T, HD)
 
     # Causal mask
     causal_mask = torch.tril(torch.ones(T, T, device=q.device, dtype=torch.bool)).unsqueeze(0)