From 60b6f8ee7929a4130c3d0bf5a97ac5cc3420f9e0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 26 May 2026 15:23:21 +0000 Subject: [PATCH] fix: use separate kernels for segments with/without compressed KV (n_comp is compile-time) --- tests/unit/test_d5c_multitile.py | 34 +++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_d5c_multitile.py b/tests/unit/test_d5c_multitile.py index f6aad639..dac270c9 100644 --- a/tests/unit/test_d5c_multitile.py +++ b/tests/unit/test_d5c_multitile.py @@ -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