Fix attention: manual causal mask for batched single-query

cutedsl/csa_attention.py

@@ -409,13 +409,20 @@ def full_attention_reference(
     # Q: (T, NH, HD) → (T*NH, 1, HD)
     q_2d = q.reshape(T * NH, 1, HD)
 
-    # Causal mask: (1, 1, T, T) broadcast over batch dim T*NH
-    causal_mask = torch.tril(torch.ones(T, T, device=q.device, dtype=torch.bool)).unsqueeze(0).unsqueeze(0)
-
-    out = F.scaled_dot_product_attention(
-        q_2d, k_2d, v_2d,
-        attn_mask=causal_mask,
-        scale=scale,
-    )
+    # Manual attention (SDPA mask handling is tricky with batched single-query)
+    # scores: (T*NH, 1, T) = Q @ K^T
+    scores = torch.matmul(q_2d, k_2d.transpose(-1, -2)) * scale
+    # Causal mask: each query at position i can only attend to positions <= i
+    # Since each batch is (query_pos, head), and KV has all T positions,
+    # we need position-aware masking
+    # For single-query batches: batch i corresponds to (pos i // NH, head i % NH)
+    # All positions <= i // NH are valid
+    # Simple approach: use a per-query mask
+    query_positions = torch.arange(T, device=device).unsqueeze(1).repeat(1, NH).reshape(T * NH)  # (T*NH,)
+    kv_positions = torch.arange(T, device=device).unsqueeze(0)  # (1, T)
+    causal = kv_positions <= query_positions.unsqueeze(1)  # (T*NH, T)
+    scores = scores.squeeze(1).masked_fill(~causal, float('-inf'))  # (T*NH, T)
+    weights = F.softmax(scores.float(), dim=-1).to(q.dtype)  # (T*NH, T)
+    out = torch.matmul(weights.unsqueeze(1), v_2d)  # (T*NH, 1, HD)
 
     return out.squeeze(1).reshape(T, NH, HD)