""" FMHA D4: Causal mask on SWA branch. In-kernel causal masking: for each query row m, mask KV positions where k_coord > m_coord to -inf. Combined with D3 SWA length masking. Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d4_causal_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_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 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 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) def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False, apply_swa_mask=False, is_causal=False, swa_len_val=None): """Run FMHA with masking and return cosine similarity vs reference.""" scale = 1.0 / math.sqrt(hd) kernel = FmhaKernel( head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, apply_swa_mask=apply_swa_mask, is_causal=is_causal, normalize=False, ) pv_n_tile = kernel.pv_n_tile n_pv_tiles = kernel.n_pv_tiles stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) 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 = 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)) if pv == 0: compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, swa_len_val) 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() # Reference q_flat = q_3d[:, :, 0] k_flat = k_3d[:, :, 0] 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 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) o_norm = (o_unnorm / attn_sum).to(torch.bfloat16) cos = torch.nn.functional.cosine_similarity( o_norm.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) ).item() return cos def test_d4_causal_hd64(): """Causal mask only (hd=64).""" print("\n=== Test 1: Causal mask only (hd=64) ===") torch.manual_seed(42) m, s_k, hd = 128, 128, 64 q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') cos = _run_fmha(q, k, v, m, s_k, hd, is_causal=True) print(f" cos = {cos:.6f}") assert cos >= 0.99, f"cosine too low: {cos}" print(" ✅ PASS") def test_d4_causal_swa64(): """Causal + SWA mask combined (swa_len=64, hd=64).""" print("\n=== Test 2: Causal + SWA swa_len=64 (hd=64) ===") torch.manual_seed(42) m, s_k, hd = 128, 128, 64 q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') cos = _run_fmha(q, k, v, m, s_k, hd, apply_swa_mask=True, is_causal=True, swa_len_val=64) print(f" cos = {cos:.6f}") assert cos >= 0.99, f"cosine too low: {cos}" print(" ✅ PASS") def test_d4_causal_hd128(): """Causal mask at hd=128 (SMEM-P path).""" print("\n=== Test 3: Causal mask only (hd=128) ===") torch.manual_seed(42) m, s_k, hd = 128, 128, 128 q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') cos = _run_fmha(q, k, v, m, s_k, hd, use_smem_p=True, is_causal=True) print(f" cos = {cos:.6f}") assert cos >= 0.99, f"cosine too low: {cos}" print(" ✅ PASS") def test_d4_causal_swa32(): """Causal + SWA with short window (swa_len=32, hd=64).""" print("\n=== Test 4: Causal + SWA swa_len=32 (hd=64) ===") torch.manual_seed(42) m, s_k, hd = 128, 128, 64 q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') cos = _run_fmha(q, k, v, m, s_k, hd, apply_swa_mask=True, is_causal=True, swa_len_val=32) print(f" cos = {cos:.6f}") assert cos >= 0.99, f"cosine too low: {cos}" print(" ✅ PASS") def test_d4_no_mask_regression(): """No masking (regression — should match D1 results).""" print("\n=== Test 5: No mask regression (hd=64) ===") torch.manual_seed(42) m, s_k, hd = 128, 128, 64 q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda') k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') cos = _run_fmha(q, k, v, m, s_k, hd) print(f" cos = {cos:.6f}") assert cos >= 0.995, f"Regression: cosine too low: {cos}" print(" ✅ PASS") def test(): print("=== D4: Causal Mask on SWA Branch ===") test_d4_no_mask_regression() test_d4_causal_hd64() test_d4_causal_swa64() test_d4_causal_hd128() test_d4_causal_swa32() print("\n=== ALL TESTS PASSED ===") if __name__ == '__main__': test()