biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: swa_len as Int32 scalar instead of CuTe tensor

Browse Source 
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.
This commit is contained in:
biondizzle
2026-05-26 10:54:41 +00:00
parent df84420414
commit e3e01071f4
3 changed files with 62 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
dsv4/kernels/attention/fmha.py
View File

@@ -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):

118
tests/unit/test_d3_inkernel_mask.py
View File

@@ -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()
test()

70
tests/unit/test_d4_causal_mask.py
View File

@@ -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')