fix: D4 test reference computation only applies causal mask when is_causal=True

tests/unit/test_d4_causal_mask.py

@@ -14,15 +14,14 @@ import cuda.bindings.driver as cuda
 from dsv4.kernels.attention.fmha import FmhaKernel
 
 
-def reference_causal_attention(q, k, v, scale, swa_len=None):
-    """FP32 reference with causal mask (and optional SWA length mask)."""
+def reference_attention(q, k, v, scale, is_causal=False, swa_len=None):
+    """FP32 reference attention with optional causal and SWA masking."""
     M, hd = q.shape
     s_k = k.shape[0]
     scores = torch.matmul(q.float(), k.float().T) * scale
-    # Causal mask: row i can only attend to positions 0..i
-    for i in range(M):
-        scores[i, i + 1:] = float('-inf')
-    # SWA length mask
+    if is_causal:
+        for i in range(M):
+            scores[i, i + 1:] = float('-inf')
     if swa_len is not None and swa_len < s_k:
         scores[:, swa_len:] = float('-inf')
     max_s = scores.max(dim=-1, keepdim=True).values
@@ -65,15 +64,17 @@ def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False,
         compiled(mQ, mK, mV, mC, stream, mLSE, swa_len_val)
         o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float()
 
-    # External normalization
+    # Reference
     q_flat = q_3d[:, :, 0]
     k_flat = k_3d[:, :, 0]
-    ref = reference_causal_attention(q_flat, k_flat, v, scale, swa_len=swa_len_val)
+    ref = reference_attention(q_flat, k_flat, v, scale, is_causal=is_causal, swa_len=swa_len_val)
 
+    # Normalize using the same mask as the reference
     scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale
-    for i in range(m):
-        scores[i, i + 1:] = float('-inf')
-    if swa_len_val is not None and swa_len_val < s_k:
+    if is_causal:
+        for i in range(m):
+            scores[i, i + 1:] = float('-inf')
+    if apply_swa_mask and swa_len_val is not None and swa_len_val < s_k:
         scores[:, swa_len_val:] = float('-inf')
     max_s = scores.max(dim=-1, keepdim=True).values
     attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True)