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

fix: use separate kernels for segments with/without compressed KV (n_comp is compile-time)

Browse Source
This commit is contained in:
biondizzle
2026-05-26 15:23:21 +00:00
parent 2efd15c852
commit 60b6f8ee79
1 changed files with 25 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
tests/unit/test_d5c_multitile.py
View File

@@ -152,24 +152,33 @@ def test_d5c_multitile():
seg_size = 128
n_segs = (n_total + seg_size - 1) // seg_size
# Compile kernel for s_k=128 (single KV tile, no O rescale)
kernel = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=n_comp)
# Compile two kernels: one with n_comp (for segments with compressed+SWA),
# one without n_comp (for segments with SWA only)
# n_comp is a compile-time parameter baked into const_expr guards.
# Pre-compile with dummy data
# Kernel 1: has compressed KV + SWA (n_comp > 0)
kernel_comp = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=n_comp)
def to_cute(t):
return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
_q = q.unsqueeze(-1)
_k = k_combined[:seg_size].unsqueeze(-1)
_v = v_combined[:seg_size, :kernel.pv_n_tile].contiguous().unsqueeze(-1)
_c = torch.zeros(m, kernel.pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
_v = v_combined[:seg_size, :kernel_comp.pv_n_tile].contiguous().unsqueeze(-1)
_c = torch.zeros(m, kernel_comp.pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
_lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
_rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
_mQ = to_cute(_q); _mK = to_cute(_k); _mV = to_cute(_v)
_mC = to_cute(_c); _mLSE = to_cute(_lse); _mRS = to_cute(_rs)
_mSB = to_cute(attn_sink)
compiled = cute.compile(kernel, _mQ, _mK, _mV, _mC, stream, _mLSE,
swa_len=swa_len, sink_bias=_mSB, row_sums=_mRS)
compiled_comp = cute.compile(kernel_comp, _mQ, _mK, _mV, _mC, stream, _mLSE,
swa_len=swa_len, sink_bias=_mSB, row_sums=_mRS)
# Kernel 2: SWA only (n_comp=0, all positions get sink bias)
kernel_swa = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=0)
# Re-compile for n_comp=0 (different const_expr path)
compiled_swa = cute.compile(kernel_swa, _mQ, _mK, _mV, _mC, stream, _mLSE,
swa_len=swa_len, sink_bias=_mSB, row_sums=_mRS)
# Run each segment
segment_results = []
@@ -206,9 +215,16 @@ def test_d5c_multitile():
# Whether this segment has SWA positions (needs sink bias)
has_swa = (k_start + seg_size) > n_comp_global
has_comp = n_comp_local > 0
# Pick the right kernel/compiled object
if has_comp:
kern = kernel_comp; comp = compiled_comp
else:
kern = kernel_swa; comp = compiled_swa
o_norm, lse = run_segment(
q, k_seg, v_seg, kernel, compiled, stream,
q, k_seg, v_seg, kern, comp, stream,
sink_bias=attn_sink if has_swa else None,
n_comp_in_seg=n_comp_local,
swa_len=swa_len_local