From e3e01071f483f819eeb7d0d0b4de4c7473dd08b0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 26 May 2026 10:54:41 +0000 Subject: [PATCH] fix: swa_len as Int32 scalar instead of CuTe tensor CuTeDSL @cute.kernel cannot handle dynamic-shape tensors as parameters. Pass swa_len as Int32 scalar instead of a 1D tensor. This works for batch_size=1 (current config). Updated D3 and D4 tests to pass swa_len as int. --- dsv4/kernels/attention/fmha.py | 19 ++--- tests/unit/test_d3_inkernel_mask.py | 118 +++++++++------------------- tests/unit/test_d4_causal_mask.py | 70 +++++------------ 3 files changed, 62 insertions(+), 145 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index d73e1df3..d7a459d7 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -105,7 +105,7 @@ class FmhaKernel: cute.size_in_bytes(self.q_dtype, v_s)) * cta @cute.jit - def __call__(self, q, k, v, c, stream, lse=None, swa_lens=None): + def __call__(self, q, k, v, c, stream, lse=None, swa_len=None): self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() @@ -133,17 +133,17 @@ class FmhaKernel: # CuTeDSL doesn't support None parameters in @cute.kernel. 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 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)) + if const_expr(swa_len is None): + # No SWA masking — pass max int (no positions masked) + swa_len = Int32(2147483647) + else: + swa_len = Int32(swa_len) # 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) + 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_len).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream) @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE, swa_lens): + def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE, swa_len): warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) tidx,_,_ = cute.arch.thread_idx() if warp_idx == self.tma_warp_id: @@ -419,9 +419,6 @@ class FmhaKernel: # 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() - 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): diff --git a/tests/unit/test_d3_inkernel_mask.py b/tests/unit/test_d3_inkernel_mask.py index 390197e9..814b0d20 100644 --- a/tests/unit/test_d3_inkernel_mask.py +++ b/tests/unit/test_d3_inkernel_mask.py @@ -1,13 +1,10 @@ """ FMHA D3: In-kernel SWA sequence length masking. -Proper approach: the kernel receives swa_lens and masks logits to -inf +Proper approach: the kernel receives swa_len (int) 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 @@ -18,24 +15,11 @@ 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 - """ +def reference_swa_attention(q, k, v, swa_len, scale): + """FP32 reference with proper -inf masking.""" 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') + if swa_len < k.shape[0]: + 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) @@ -44,18 +28,8 @@ def reference_swa_attention(q, k, v, swa_lens, scale): 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 - """ +def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_len_val, use_smem_p=False): + """Run FMHA with in-kernel SWA 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, @@ -65,13 +39,7 @@ def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_lens_tensor, use_smem_p=Fals 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) @@ -85,19 +53,15 @@ def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_lens_tensor, use_smem_p=Fals 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 = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, swa_len_val) - compiled(mQ, mK, mV, mC, stream, mLSE, mSwaLens) + 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() - 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). + # External normalization using reference attn_sum 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 @@ -107,19 +71,17 @@ def _run_fmha_masked(q_3d, k_3d, v, m, s_k, hd, swa_lens_tensor, use_smem_p=Fals 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) ===") + """Full window (swa_len=128): no masking, regression test.""" + print("\n=== Test 1: No masking (swa_len=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) + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=s_k) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, s_k, 1.0 / math.sqrt(hd)) cos = torch.nn.functional.cosine_similarity( o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) @@ -130,19 +92,17 @@ def test_d3_no_mask(): def test_d3_swa64(): - """SWA with swa_lens=64: mask positions 64-127 to -inf.""" - print("\n=== Test 2: swa_lens=64 (hd=64) ===") + """SWA with swa_len=64: mask positions 64-127 to -inf.""" + print("\n=== Test 2: swa_len=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') + 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) + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=64) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 64, 1.0 / math.sqrt(hd)) cos = torch.nn.functional.cosine_similarity( o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) @@ -153,19 +113,17 @@ def test_d3_swa64(): def test_d3_swa32(): - """SWA with swa_lens=32: only 32 valid positions.""" - print("\n=== Test 3: swa_lens=32 (hd=64) ===") + """SWA with swa_len=32: only 32 valid positions.""" + print("\n=== Test 3: swa_len=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) + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=32) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 32, 1.0 / math.sqrt(hd)) cos = torch.nn.functional.cosine_similarity( o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) @@ -176,19 +134,17 @@ def test_d3_swa32(): def test_d3_swa1(): - """Edge case: swa_lens=1, only one valid KV position.""" - print("\n=== Test 4: swa_lens=1 (hd=64) ===") + """Edge case: swa_len=1, only one valid KV position.""" + print("\n=== Test 4: swa_len=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) + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=1) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 1, 1.0 / math.sqrt(hd)) cos = torch.nn.functional.cosine_similarity( o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) @@ -200,18 +156,16 @@ def test_d3_swa1(): def test_d3_hd128(): """SWA masking at hd=128 (SMEM-P path).""" - print("\n=== Test 5: swa_lens=64 (hd=128) ===") + print("\n=== Test 5: swa_len=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) + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=64, use_smem_p=True) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, 64, 1.0 / math.sqrt(hd)) cos = torch.nn.functional.cosine_similarity( o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) @@ -222,19 +176,17 @@ def test_d3_hd128(): 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) ===") + """No masking at hd=128: regression test.""" + print("\n=== Test 6: No masking (swa_len=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) + o = _run_fmha_masked(q, k, v, m, s_k, hd, swa_len_val=s_k, use_smem_p=True) + ref = reference_swa_attention(q[:, :, 0], k[:, :, 0], v, s_k, 1.0 / math.sqrt(hd)) cos = torch.nn.functional.cosine_similarity( o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0) @@ -256,4 +208,4 @@ def test(): if __name__ == '__main__': - test() \ No newline at end of file + test() diff --git a/tests/unit/test_d4_causal_mask.py b/tests/unit/test_d4_causal_mask.py index a652fe44..1a5aee04 100644 --- a/tests/unit/test_d4_causal_mask.py +++ b/tests/unit/test_d4_causal_mask.py @@ -2,10 +2,7 @@ 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). +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 """ @@ -17,31 +14,17 @@ 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 - """ +def reference_causal_attention(q, k, v, scale, swa_len=None): + """FP32 reference with causal mask (and optional SWA length mask).""" 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') + # SWA length mask + 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) @@ -51,8 +34,8 @@ def reference_causal_attention(q, k, v, scale, swa_lens=None): 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.""" + 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, @@ -63,14 +46,6 @@ def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False, 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): @@ -85,25 +60,21 @@ def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False, 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 = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, swa_len_val) - compiled(mQ, mK, mV, mC, stream, mLSE, mSwaLens) + 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 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) + ref = reference_causal_attention(q_flat, k_flat, v, scale, swa_len=swa_len_val) - # 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') + if 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) @@ -131,17 +102,16 @@ def test_d4_causal_hd64(): 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) ===") + """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') - 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) + 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") @@ -164,17 +134,16 @@ def test_d4_causal_hd128(): 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) ===") + """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') - 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) + 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") @@ -185,7 +154,6 @@ def test_d4_no_mask_regression(): 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')