D3: add SWA sequence length mask test (reference oracle + full-window regression)

tests/unit/test_d3_swa_mask.py

@@ -0,0 +1,148 @@
+"""
+FMHA D3: SWA sequence length mask.
+
+Adds swa_lens[b] masking to the softmax: positions >= swa_lens are -inf.
+This handles variable-length SWA windows (early positions have fewer tokens).
+
+Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_swa_mask.py
+"""
+import torch
+import math
+import cutlass.cute as cute
+import cutlass.torch as ct
+import cuda.bindings.driver as cuda
+from dsv4.kernels.attention.fmha import FmhaKernel
+
+
+def reference_swa_attention(q, k, v, swa_lens, scale):
+    """FP32 reference: q (M, hd), k (s_k, hd), v (s_k, hd), swa_lens (M,) → o (M, hd)"""
+    scores = torch.matmul(q.float(), k.float().T) * scale
+    # Apply SWA mask: positions >= swa_lens are -inf
+    for i in range(q.shape[0]):
+        sl = swa_lens[i].item()
+        if sl < k.shape[0]:
+            scores[i, sl:] = float('-inf')
+    max_s = scores.max(dim=-1, keepdim=True).values
+    exp_s = (scores - max_s).exp()
+    sum_s = exp_s.sum(dim=-1, keepdim=True)
+    p = exp_s / sum_s
+    o = torch.matmul(p, v.float())
+    return o.to(torch.bfloat16), sum_s
+
+
+def test_d3_full_window():
+    """Full SWA window (swa_lens=128): no masking, same as dense attention."""
+    print("\n=== Test 1: Full SWA window (swa_lens=128, hd=64) ===")
+    torch.manual_seed(42)
+    m, s_k, hd = 128, 128, 64
+    scale = 1.0 / math.sqrt(hd)
+    
+    q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
+    swa_lens = torch.full((m,), s_k, dtype=torch.int32, device='cuda')
+    
+    # Run FMHA (same as dense, no masking needed)
+    q_3d = q; k_3d = k.unsqueeze(-1)
+    kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False)
+    pv_n_tile = kernel.pv_n_tile; n_pv_tiles = kernel.n_pv_tiles
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    
+    v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
+    c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
+    lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
+    
+    mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d))
+    mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d))
+    mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
+    mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
+    mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
+    
+    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
+    
+    o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda')
+    for pv in range(n_pv_tiles):
+        v_tile = v[:, pv*pv_n_tile:(pv+1)*pv_n_tile].contiguous().unsqueeze(-1)
+        c_tile.zero_(); lse_tensor.zero_()
+        mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
+        mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
+        mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
+        compiled(mQ, mK, mV, mC, stream, mLSE)
+        o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float()
+    
+    # Reference normalization
+    scores = torch.matmul(q[:,:,0].float(), k.float().T) * scale
+    max_s = scores.max(dim=-1, keepdim=True).values
+    attn_sum = (scores - max_s).exp().sum(dim=-1, keepdim=True)
+    o_norm = (o_unnorm / attn_sum).to(torch.bfloat16)
+    
+    ref, _ = reference_swa_attention(q[:,:,0], k, v, swa_lens, scale)
+    cos = torch.nn.functional.cosine_similarity(
+        o_norm.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
+    ).item()
+    print(f"  cos = {cos:.6f}")
+    assert cos >= 0.995, f"cosine too low: {cos}"
+    print("  ✅ PASS")
+
+
+def test_d3_partial_window():
+    """Partial SWA window (swa_lens=64): first 64 tokens valid, rest masked."""
+    print("\n=== Test 2: Partial SWA window (swa_lens=64, hd=64) ===")
+    print("  (Testing reference oracle — kernel SWA mask not yet implemented)")
+    torch.manual_seed(42)
+    m, s_k, hd = 128, 128, 64
+    scale = 1.0 / math.sqrt(hd)
+    
+    q = torch.randn(m, hd, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
+    swa_lens = torch.full((m,), 64, dtype=torch.int32, device='cuda')
+    
+    ref, _ = reference_swa_attention(q, k, v, swa_lens, scale)
+    
+    # Full attention (no masking)
+    scores_full = torch.matmul(q.float(), k.float().T) * scale
+    max_s = scores_full.max(dim=-1, keepdim=True).values
+    o_full = (torch.softmax(scores_full, dim=-1) @ v.float()).to(torch.bfloat16)
+    
+    # Verify reference masking works: full and masked should differ
+    cos_full = torch.nn.functional.cosine_similarity(
+        ref.flatten().float().unsqueeze(0), o_full.flatten().float().unsqueeze(0)
+    ).item()
+    print(f"  cos (masked vs full) = {cos_full:.6f} (should be < 1.0, proving mask works)")
+    assert cos_full < 0.999, f"Masking should change output, got cos={cos_full}"
+    print("  ✅ PASS (reference oracle works)")
+
+
+def test_d3_varying_lens():
+    """Varying SWA lens across rows: simulates batch with different positions."""
+    print("\n=== Test 3: Varying swa_lens (hd=64) ===")
+    print("  (Testing reference oracle — kernel SWA mask not yet implemented)")
+    torch.manual_seed(42)
+    m, s_k, hd = 128, 128, 64
+    scale = 1.0 / math.sqrt(hd)
+    
+    q = torch.randn(m, hd, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
+    # Varying lens: some rows have 128, some 64, some 32
+    swa_lens = torch.full((m,), 128, dtype=torch.int32, device='cuda')
+    swa_lens[0:32] = 32
+    swa_lens[32:64] = 64
+    
+    ref, _ = reference_swa_attention(q, k, v, swa_lens, scale)
+    print(f"  Output shape: {ref.shape}")
+    print(f"  swa_lens: min={swa_lens.min()}, max={swa_lens.max()}")
+    print("  ✅ PASS (reference oracle works)")
+
+
+def test():
+    print("=== D3: SWA Sequence Length Mask ===")
+    test_d3_full_window()
+    test_d3_partial_window()
+    test_d3_varying_lens()
+    print("\n=== ALL TESTS PASSED ===")
+
+
+if __name__ == '__main__':
+    test()