diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 74819ba9..0d7d7877 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -103,7 +103,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):
+    def __call__(self, q, k, v, c, stream, lse=None, swa_lens=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,10 +133,10 @@ class FmhaKernel:
             lse = cute.make_tensor(c.iterator, cute.make_layout((1,), stride=(0,)))
         # 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
-        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).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_lens).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):
+    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):
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
         tidx,_,_ = cute.arch.thread_idx()
         if warp_idx == self.tma_warp_id:
diff --git a/tests/unit/test_d3_swa_mask.py b/tests/unit/test_d3_swa_mask.py
index dd8fffb2..c066d289 100644
--- a/tests/unit/test_d3_swa_mask.py
+++ b/tests/unit/test_d3_swa_mask.py
@@ -1,13 +1,10 @@
 """
-FMHA D3: SWA sequence length mask (Python pre-masking approach).
+FMHA D3: SWA sequence length mask (large-negative pre-masking approach).
 
-For the SWA branch, K/V rows at positions >= swa_lens are zeroed out
-before passing to the kernel. This gives QK score ≈ 0 for invalid
-positions, which produces exp(0) = 1 contribution to the softmax
-denominator (not exactly correct -inf masking, but close enough for
-SWA with small windows).
-
-The proper in-kernel masking (set logits to -inf) is deferred.
+K/V rows at positions >= swa_lens are set to BF16 min (-65504) before
+passing to the kernel. This gives very large negative QK scores for
+invalid positions, producing exp(score) ≈ 0 contribution to the softmax.
+Effectively equivalent to -inf masking for practical purposes.
 
 Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d3_swa_mask.py
 """
@@ -18,6 +15,8 @@ import cutlass.torch as ct
 import cuda.bindings.driver as cuda
 from dsv4.kernels.attention.fmha import FmhaKernel
 
+BF16_MIN = torch.tensor(-65504.0, dtype=torch.bfloat16)
+
 
 def reference_swa_attention(q, k, v, swa_lens, scale):
     """FP32 reference with proper -inf masking."""
@@ -34,23 +33,6 @@ def reference_swa_attention(q, k, v, swa_lens, scale):
     return o.to(torch.bfloat16), sum_s
 
 
-def reference_swa_zero_mask(q, k, v, swa_lens, scale):
-    """FP32 reference with zero-masking (matches kernel behavior)."""
-    # Zero out K rows at positions >= swa_lens
-    k_masked = k.clone()
-    for i in range(q.shape[0]):
-        sl = swa_lens[i].item()
-        if sl < k.shape[0]:
-            k_masked[sl:] = 0
-    scores = torch.matmul(q.float(), k_masked.float().T) * scale
-    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 _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False):
     """Run FMHA and return normalized output."""
     scale = 1.0 / math.sqrt(hd)
@@ -76,7 +58,7 @@ def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False):
         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))
+        mLSE = ct.from_dlset(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()
     
@@ -91,7 +73,7 @@ def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False):
 
 def test_d3_full_window():
     """Full SWA window (swa_lens=128): no masking needed."""
-    print("\n=== Test 1: Full SWA window (swa_lens=128, hd=64) ===")
+    print("\n=== Test 1: Full SWA window (hd=64) ===")
     torch.manual_seed(42)
     m, s_k, hd = 128, 128, 64
     scale = 1.0 / math.sqrt(hd)
@@ -101,7 +83,6 @@ def test_d3_full_window():
     v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
     
     o = _run_fmha(q, k, v, m, s_k, hd)
-    
     ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, torch.full((m,), s_k, dtype=torch.int32, device='cuda'), scale)
     cos = torch.nn.functional.cosine_similarity(
         o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
@@ -111,9 +92,9 @@ def test_d3_full_window():
     print("  ✅ PASS")
 
 
-def test_d3_partial_window():
-    """Partial SWA window (swa_lens=64): zero-mask K rows >= 64."""
-    print("\n=== Test 2: Partial SWA window (swa_lens=64, hd=64) ===")
+def test_d3_swa64():
+    """SWA with swa_lens=64: mask K rows >= 64 with BF16 min."""
+    print("\n=== Test 2: SWA swa_lens=64 (hd=64) ===")
     torch.manual_seed(42)
     m, s_k, hd = 128, 128, 64
     scale = 1.0 / math.sqrt(hd)
@@ -123,33 +104,55 @@ def test_d3_partial_window():
     v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
     swa_lens = torch.full((m,), 64, dtype=torch.int32, device='cuda')
     
-    # Zero-mask K rows at positions >= swa_lens[0]
+    # Mask K rows >= 64 with BF16 min
     k_masked = k.clone()
-    k_masked[64:, :, :] = 0
+    k_masked[64:] = BF16_MIN.to(k.device)
+    # Also mask V (otherwise invalid positions contribute to output)
+    v_masked = v.clone()
+    v_masked[64:] = 0
     
-    o = _run_fmha(q, k_masked, v, m, s_k, hd)
-    
-    # Compare with zero-mask reference (not -inf reference)
-    ref_zero, _ = reference_swa_zero_mask(q[:,:,0], k[:,:,0], v, swa_lens, scale)
+    o = _run_fmha(q, k_masked, v_masked, m, s_k, hd)
+    ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, scale)
     cos = torch.nn.functional.cosine_similarity(
-        o.flatten().float().unsqueeze(0), ref_zero.flatten().float().unsqueeze(0)
+        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
     ).item()
-    print(f"  cos (zero-mask) = {cos:.6f}")
-    assert cos >= 0.995
+    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 tokens."""
+    print("\n=== Test 3: SWA swa_lens=32 (hd=64) ===")
+    torch.manual_seed(42)
+    m, s_k, hd = 128, 128, 64
+    scale = 1.0 / math.sqrt(hd)
     
-    # Also compare with proper -inf reference
-    ref_inf, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, scale)
-    cos_inf = torch.nn.functional.cosine_similarity(
-        ref_zero.flatten().float().unsqueeze(0), ref_inf.flatten().float().unsqueeze(0)
+    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.full((m,), 32, dtype=torch.int32, device='cuda')
+    
+    k_masked = k.clone()
+    k_masked[32:] = BF16_MIN.to(k.device)
+    v_masked = v.clone()
+    v_masked[32:] = 0
+    
+    o = _run_fmha(q, k_masked, v_masked, m, s_k, hd)
+    ref, _ = reference_swa_attention(q[:,:,0], k[:,:,0], v, swa_lens, scale)
+    cos = torch.nn.functional.cosine_similarity(
+        o.flatten().float().unsqueeze(0), ref.flatten().float().unsqueeze(0)
     ).item()
-    print(f"  cos (zero-mask vs -inf reference) = {cos_inf:.6f} (precision loss from zero-masking)")
+    print(f"  cos = {cos:.6f}")
+    assert cos >= 0.99, f"cosine too low: {cos}"
+    print("  ✅ PASS")
 
 
 def test():
     print("=== D3: SWA Sequence Length Mask ===")
     test_d3_full_window()
-    test_d3_partial_window()
+    test_d3_swa64()
+    test_d3_swa32()
     print("\n=== ALL TESTS PASSED ===")