diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index c413a5c7..c66d4598 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -32,6 +32,7 @@ class FmhaKernel: self.batch_size = batch_size self.normalize = normalize # D5a: False = emit un-normalized O + lse self.apply_swa_mask = apply_swa_mask # D3: mask logits at positions >= swa_lens + self.is_causal = False # D4: causal mask (k_coord > m_coord) on SWA branch self.acc_dtype = Float32; self.qk_acc_dtype = Float32 self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 @@ -410,23 +411,33 @@ class FmhaKernel: cute.copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS) cute.arch.fence_view_async_tmem_load() - # D3: In-kernel SWA sequence length masking. - # After loading S from TMEM, mask positions >= swa_lens[batch_idx] to -inf. + # D3/D4: In-kernel logit masking. + # After loading S from TMEM, mask invalid positions to -inf. # Uses tTMEM_LOADcS coordinate tensor to map register indices to (row, col). - # col = position in KV sequence. For kt > 0, actual pos = kt*128 + col. - # This is the PROPER approach: post-QK masking in the softmax, - # not pre-masking K with BF16 min (which can't produce true -inf). - if const_expr(self.apply_swa_mask): + # D3: SWA mask — positions >= swa_lens[batch_idx] → -inf + # D4: Causal mask — positions where k_coord > m_coord → -inf + # Both use the same coordinate mapping from tTMEM_LOADcS. + # For kt > 0, absolute KV pos = kt*128 + k_coord. + if const_expr(self.apply_swa_mask or self.is_causal): _bidx, _bidy, _bidz = cute.arch.block_idx() - swa_len = swa_lens[_bidz] + if const_expr(self.apply_swa_mask): + swa_len = swa_lens[_bidz] kt_offset = Int32(kt * 128) # KV position offset for this tile # Iterate using same coordinate indexing as SMEM-P path for j0 in range(32): for j1 in range(4): coord = tTMEM_LOADcS[(j0, 0), j1, 0, 0] + m_coord = coord[0] # query row position k_coord = coord[1] # position within this KV tile kv_pos = kt_offset + k_coord # absolute KV position - if kv_pos >= swa_len: + should_mask = Boolean(0) + if const_expr(self.apply_swa_mask): + if kv_pos >= swa_len: + should_mask = Boolean(1) + if const_expr(self.is_causal): + if k_coord > m_coord: + should_mask = Boolean(1) + if should_mask: tTMEM_LOADrS[(j0, 0), j1, 0, 0] = -Float32.inf old_row_max = row_max diff --git a/tests/unit/test_d4_causal_mask.py b/tests/unit/test_d4_causal_mask.py new file mode 100644 index 00000000..a652fe44 --- /dev/null +++ b/tests/unit/test_d4_causal_mask.py @@ -0,0 +1,211 @@ +""" +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. This is the proper causal attention mask. + +Combined with D3 SWA length masking: both conditions can be active +simultaneously (OR logic). + +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_causal_attention(q, k, v, scale, swa_lens=None): + """FP32 reference with causal mask (and optional SWA length mask). + + Args: + q: (M, hd) BF16 + k: (s_k, hd) BF16 + v: (s_k, hd) BF16 + scale: float + swa_lens: (M,) int32 or None — per-row valid KV count + + Returns: + o: (M, hd) BF16 + """ + 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: positions >= swa_lens + if swa_lens is not None: + for i in range(M): + sl = swa_lens[i].item() + if sl < s_k: + 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) + + +def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False, + apply_swa_mask=False, is_causal=False, swa_lens_tensor=None): + """Run FMHA with masking and return normalized output.""" + 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) + + # swa_lens as CuTe tensor + if swa_lens_tensor is not None: + mSwaLens = ct.from_dlpack(swa_lens_tensor).mark_layout_dynamic( + leading_dim=ct.get_leading_dim(swa_lens_tensor) + ) + else: + mSwaLens = None + + 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, mSwaLens) + + compiled(mQ, mK, mV, mC, stream, mLSE, mSwaLens) + o_unnorm[:, pv * pv_n_tile:(pv + 1) * pv_n_tile] = c_tile[:, :, 0].float() + + # External normalization + q_flat = q_3d[:, :, 0] + k_flat = k_3d[:, :, 0] + swa_lens_for_ref = swa_lens_tensor.cpu().expand(m) if swa_lens_tensor is not None else None + ref = reference_causal_attention(q_flat, k_flat, v, scale, swa_lens=swa_lens_for_ref) + + # Use reference attn_sum for normalization (same pattern as other tests) + scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale + for i in range(m): + scores[i, i + 1:] = float('-inf') + if swa_lens_tensor is not None: + sl = swa_lens_tensor[0].item() + if sl < s_k: + scores[:, sl:] = 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_lens=64, hd=64).""" + print("\n=== Test 2: Causal + SWA swa_lens=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') + swa_lens = torch.tensor([64], dtype=torch.int32, device='cuda') + + cos = _run_fmha(q, k, v, m, s_k, hd, apply_swa_mask=True, is_causal=True, swa_lens_tensor=swa_lens) + 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 very short window (swa_lens=32, hd=64).""" + print("\n=== Test 4: Causal + SWA swa_lens=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') + swa_lens = torch.tensor([32], dtype=torch.int32, device='cuda') + + cos = _run_fmha(q, k, v, m, s_k, hd, apply_swa_mask=True, is_causal=True, swa_lens_tensor=swa_lens) + 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 + scale = 1.0 / math.sqrt(hd) + + 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()