diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 79604bc1..c413a5c7 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -16,7 +16,7 @@ import math class FmhaKernel: - def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True, num_query_heads=1, batch_size=1): + def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True, num_query_heads=1, batch_size=1, apply_swa_mask=False): self.head_dim = head_dim self.s_k = s_k self.n_kv_tiles = s_k // 128 @@ -31,6 +31,7 @@ class FmhaKernel: self.num_query_heads = num_query_heads 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.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 @@ -132,10 +133,12 @@ class FmhaKernel: if const_expr(lse is None): lse = cute.make_tensor(c.iterator, cute.make_layout((1,), stride=(0,))) if const_expr(swa_lens is None): - # No SWA masking — pass a dummy tensor - swa_lens = cute.make_tensor(c.iterator, cute.make_layout((1,), stride=(0,))) + # No SWA masking — pass a dummy tensor with large value (no positions masked) + _swa_dummy = torch.tensor([2147483647], dtype=torch.int32, device='cuda') + swa_lens = ct.from_dlpack(_swa_dummy).mark_layout_dynamic(leading_dim=ct.get_leading_dim(_swa_dummy)) # Grid: (M_tiles, 1, batch) where M = n_h * T packed into M dimension # For single-head (n_h=1): grid=(1,1,1) — backward compatible + # block_idx_z = batch index, used for swa_lens[batch_idx] in D3 masking self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_lens).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream) @cute.kernel @@ -407,6 +410,25 @@ 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. + # 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): + _bidx, _bidy, _bidz = cute.arch.block_idx() + 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] + k_coord = coord[1] # position within this KV tile + kv_pos = kt_offset + k_coord # absolute KV position + if kv_pos >= swa_len: + tTMEM_LOADrS[(j0, 0), j1, 0, 0] = -Float32.inf + old_row_max = row_max frg_cnt = 4 frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt diff --git a/tests/unit/test_d3_inkernel_mask.py b/tests/unit/test_d3_inkernel_mask.py new file mode 100644 index 00000000..390197e9 --- /dev/null +++ b/tests/unit/test_d3_inkernel_mask.py @@ -0,0 +1,259 @@ +""" +FMHA D3: In-kernel SWA sequence length masking. + +Proper approach: the kernel receives swa_lens and masks logits to -inf +inside the softmax, using the tTMEM_LOADcS coordinate tensor to map +register fragment positions to (row, col) in the QK matrix. + +This replaces the pre-masking approach (BF16 min on K) which cannot +produce true -inf QK scores. + +Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_inkernel_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 with proper -inf masking. + + Args: + q: (M, hd) BF16 + k: (s_k, hd) BF16 + v: (s_k, hd) BF16 + swa_lens: (M,) int32 — per-row number of valid KV positions + scale: float + + Returns: + o: (M, hd) BF16 + """ + scores = torch.matmul(q.float(), k.float().T) * scale + 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) + + +def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_lens_tensor, use_smem_p=False): + """Run FMHA with in-kernel SWA masking and return normalized output. + + Args: + q_3d: (M, hd, 1) BF16 + k_3d: (s_k, hd, 1) BF16 + v: (s_k, hd) BF16 + swa_lens_tensor: (1,) int32 — number of valid KV positions + + Returns: + o_norm: (M, hd) BF16 + """ + scale = 1.0 / math.sqrt(hd) + kernel = FmhaKernel( + head_dim=hd, s_k=s_k, use_smem_p=use_smem_p, + apply_swa_mask=True, 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 (1D, int32) + mSwaLens = ct.from_dlpack(swa_lens_tensor).mark_layout_dynamic( + leading_dim=ct.get_leading_dim(swa_lens_tensor) + ) + + o_unnorm = torch.zeros(m, hd, dtype=torch.float32, device='cuda') + lse_all = torch.zeros(m, 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() + lse_all += lse_tensor[:, 0, 0] + + # External normalization using LSE + # O_norm = O_unnorm / exp(LSE) ... but per-row LSE only row 0 is written. + # Use reference attn_sum for normalization (same as head-packed tests). + q_flat = q_3d[:, :, 0] + k_flat = k_3d[:, :, 0] + scores = torch.matmul(q_flat.float(), k_flat.float().T) * scale + swa_len_val = swa_lens_tensor[0].item() + if 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) + return o_norm + + +def test_d3_no_mask(): + """Full window (swa_lens=128): no masking, regression test.""" + print("\n=== Test 1: No masking (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, 1, dtype=torch.bfloat16, device='cuda') + v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') + swa_lens = torch.tensor([s_k], dtype=torch.int32, device='cuda') + + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale) + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) + ).item() + print(f" cos = {cos:.6f}") + assert cos >= 0.995, f"Regression: cosine too low: {cos}" + print(" ✅ PASS") + + +def test_d3_swa64(): + """SWA with swa_lens=64: mask positions 64-127 to -inf.""" + print("\n=== Test 2: swa_lens=64 (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') + swa_lens = torch.tensor([64], dtype=torch.int32, device='cuda') + + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale) + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) + ).item() + print(f" cos = {cos:.6f}") + assert cos >= 0.99, f"cosine too low: {cos}" + print(" ✅ PASS") + + +def test_d3_swa32(): + """SWA with swa_lens=32: only 32 valid positions.""" + print("\n=== Test 3: swa_lens=32 (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') + swa_lens = torch.tensor([32], dtype=torch.int32, device='cuda') + + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale) + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) + ).item() + print(f" cos = {cos:.6f}") + assert cos >= 0.99, f"cosine too low: {cos}" + print(" ✅ PASS") + + +def test_d3_swa1(): + """Edge case: swa_lens=1, only one valid KV position.""" + print("\n=== Test 4: swa_lens=1 (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') + swa_lens = torch.tensor([1], dtype=torch.int32, device='cuda') + + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale) + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) + ).item() + print(f" cos = {cos:.6f}") + assert cos >= 0.99, f"cosine too low: {cos}" + print(" ✅ PASS") + + +def test_d3_hd128(): + """SWA masking at hd=128 (SMEM-P path).""" + print("\n=== Test 5: swa_lens=64 (hd=128) ===") + torch.manual_seed(42) + m, s_k, hd = 128, 128, 128 + 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') + swa_lens = torch.tensor([64], dtype=torch.int32, device='cuda') + + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens, use_smem_p=True) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale) + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) + ).item() + print(f" cos = {cos:.6f}") + assert cos >= 0.99, f"cosine too low: {cos}" + print(" ✅ PASS") + + +def test_d3_swa128_hd128(): + """No masking at hd=128: regression test (should match existing D1 results).""" + print("\n=== Test 6: No masking (swa_lens=128, hd=128) ===") + torch.manual_seed(42) + m, s_k, hd = 128, 128, 128 + 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') + swa_lens = torch.tensor([s_k], dtype=torch.int32, device='cuda') + + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_lens, use_smem_p=True) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, swa_lens.cpu().expand(m), scale) + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) + ).item() + print(f" cos = {cos:.6f}") + assert cos >= 0.995, f"Regression: cosine too low: {cos}" + print(" ✅ PASS") + + +def test(): + print("=== D3: In-Kernel SWA Sequence Length Mask ===") + test_d3_no_mask() + test_d3_swa64() + test_d3_swa32() + test_d3_swa1() + test_d3_hd128() + test_d3_swa128_hd128() + print("\n=== ALL TESTS PASSED ===") + + +if __name__ == '__main__': + test() \ No newline at end of file