Fix KV expand for SDPA: (T,HD) → (T*NH, T, HD)

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)