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

diag: rewrite multi-tile test with explicit per-segment compile and reference

Browse Source
This commit is contained in:
biondizzle
2026-05-26 15:39:39 +00:00
parent 7f983fb855
commit c9eab3c7e0
1 changed files with 97 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files

329
tests/unit/test_d5c_multitile.py
View File

@@ -1,17 +1,6 @@
"""
FMHA D5c: Sink bias + Python KV merge for multi-tile (s_k > 128).
Verifies the full DSV4 attention pipeline:
1. Concatenate KV: [compressed_K; swa_K] (total s_k > 128)
2. Split into 128-token segments
3. Run FMHA per segment (with sink bias on SWA positions, D3/D4 masking)
4. Merge segments using Python KV merge formula:
O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i))
5. Normalize using row_sum
This is the production path for DSV4 Pro (s_k=1152, 9 KV tiles)
until the D1.5 TMEM round-trip fix enables in-kernel O rescale.
Run: ~/.openclaw/workspace/fire_b200_test tests/unit/test_d5c_multitile.py
"""
import torch
@@ -22,15 +11,18 @@ import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel
def to_cute(t):
return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
def reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa,
attn_sink, scale, swa_len, is_causal=False):
"""FP32 reference: single softmax over combined KV with sink bias on SWA."""
m, hd = q.shape
n_comp = k_comp.shape[0]
n_swa = k_swa.shape[0]
k_combined = torch.cat([k_comp, k_swa], dim=0)
v_combined = torch.cat([v_comp, v_swa], dim=0)
scores = torch.matmul(q.float(), k_combined.float().T) * scale
k_comb = torch.cat([k_comp, k_swa], dim=0)
v_comb = torch.cat([v_comp, v_swa], dim=0)
scores = q.float() @ k_comb.float().T * scale
scores[:, n_comp:] += attn_sink
if swa_len < n_swa:
scores[:, n_comp + swa_len:] = float('-inf')
@@ -42,82 +34,51 @@ def reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa,
max_s = scores.max(dim=-1, keepdim=True).values
exp_s = (scores - max_s).exp()
sum_s = exp_s.sum(dim=-1, keepdim=True).clamp(min=1e-30)
o = torch.matmul(exp_s / sum_s, v_combined.float())
return o.to(torch.bfloat16)
def run_segment(q, k_seg, v_seg, kernel, compiled, stream,
sink_bias=None, n_comp_in_seg=0, swa_len=999999):
"""Run one 128-token segment of FMHA, return normalized O and LSE."""
m = q.shape[0]
hd = v_seg.shape[1]
# Allocate per-segment outputs
o_seg = torch.zeros(m, hd, dtype=torch.bfloat16, device='cuda')
lse_seg = torch.zeros(m, dtype=torch.float32, device='cuda')
rs_seg = torch.zeros(m, dtype=torch.float32, device='cuda')
def to_cute(t):
return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
c_tile = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
lse_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
rs_tile = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
mQ = to_cute(q.unsqueeze(-1))
mK = to_cute(k_seg.unsqueeze(-1))
mV = to_cute(v_seg.unsqueeze(-1)) # (n_k, hd, 1)
mC = to_cute(c_tile)
mLSE = to_cute(lse_tile)
mRS = to_cute(rs_tile)
if sink_bias is not None:
mSB = to_cute(sink_bias)
compiled(mQ, mK, mV, mC, stream, mLSE,
swa_len=swa_len, sink_bias=mSB, row_sums=mRS)
else:
compiled(mQ, mK, mV, mC, stream, mLSE,
swa_len=swa_len, row_sums=mRS)
torch.cuda.synchronize()
rs = rs_tile[:, 0, 0].float()
o_norm = c_tile[:, :, 0].float() / rs.unsqueeze(1).clamp(min=1e-30)
lse_val = lse_tile[:, 0, 0].clone()
return o_norm, lse_val
return (exp_s / sum_s @ v_comb.float()).to(torch.bfloat16)
def python_kv_merge(segment_results):
"""Merge multiple FMHA segments using Python KV merge formula.
O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i))
"""
# Compute max LSE for numerical stability
lse_stack = torch.stack([r[1] for r in segment_results], dim=0) # (n_seg, m)
lse_max = lse_stack.max(dim=0).values # (m,)
numerator = torch.zeros_like(segment_results[0][0]) # (m, hd)
"""O = sum_i(exp(lse_i) * O_i_norm) / sum_i(exp(lse_i))"""
lse_stack = torch.stack([r[1] for r in segment_results], dim=0)
lse_max = lse_stack.max(dim=0).values
numerator = torch.zeros_like(segment_results[0][0])
denominator = torch.zeros(lse_max.shape[0], dtype=torch.float32, device=lse_max.device)
for o_norm, lse in segment_results:
exp_lse = (lse - lse_max).exp() # (m,)
exp_lse = (lse - lse_max).exp()
numerator += exp_lse.unsqueeze(1) * o_norm.float()
denominator += exp_lse
return numerator / denominator.unsqueeze(1).clamp(min=1e-30)
o_merged = numerator / denominator.unsqueeze(1).clamp(min=1e-30)
return o_merged
def run_fmha(q, k, v, kernel_obj, compiled, stream, swa_len, sink_bias=None):
"""Run FMHA with single KV tile, return (O_norm, LSE)."""
m = q.shape[0]
hd = v.shape[1]
c = torch.zeros(m, hd, 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.unsqueeze(-1))
mK = to_cute(k.unsqueeze(-1))
mV = to_cute(v.unsqueeze(-1))
mC = to_cute(c); mLSE = to_cute(lse); mRS = to_cute(rs)
if sink_bias is not None:
mSB = to_cute(sink_bias)
compiled(mQ, mK, mV, mC, stream, mLSE, swa_len=swa_len, sink_bias=mSB, row_sums=mRS)
else:
compiled(mQ, mK, mV, mC, stream, mLSE, swa_len=swa_len, row_sums=mRS)
torch.cuda.synchronize()
o_norm = c[:, :, 0].float() / rs[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
return o_norm, lse[:, 0, 0].clone()
def test_d5c_multitile():
print("=== D5c Multi-Tile: Sink Bias + Python KV Merge ===\n")
hd = 64
m = 128
n_comp = 96 # compressed KV tokens
n_swa = 160 # SWA tokens
n_total = n_comp + n_swa # 256, 2 KV tiles
swa_len = 100 # valid SWA fill (within the 160 SWA window)
scale = 1.0 / math.sqrt(hd)
hd = 64; m = 128; n_comp = 96; n_swa = 160; n_total = 256
swa_len = 100; scale = 1.0 / math.sqrt(hd); seg_size = 128
torch.manual_seed(42)
q = torch.randn(m, hd, dtype=torch.bfloat16, device='cuda')
@@ -125,176 +86,80 @@ def test_d5c_multitile():
v_comp = torch.randn(n_comp, hd, dtype=torch.bfloat16, device='cuda')
k_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda')
v_swa = torch.randn(n_swa, hd, dtype=torch.bfloat16, device='cuda')
attn_sink_val = 0.5
attn_sink = torch.tensor([attn_sink_val], dtype=torch.float32, device='cuda')
# Direct test of segment 0 (verify run_segment isn't the issue)
print('--- Direct segment 0 test ---')
k_seg0 = k_combined[0:seg_size]
v_seg0 = v_combined[0:seg_size]
c0 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
lse0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
rs0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
mQ0 = to_cute(q.unsqueeze(-1))
mK0 = to_cute(k_seg0.unsqueeze(-1))
mV0 = to_cute(v_seg0.unsqueeze(-1))
mC0 = to_cute(c0); mLSE0 = to_cute(lse0); mRS0 = to_cute(rs0)
mSB0 = to_cute(attn_sink)
# Compile FRESH for segment 0
k0 = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=n_comp, apply_sink_bias=True)
comp0 = cute.compile(k0, mQ0, mK0, mV0, mC0, stream, mLSE0,
swa_len=swa_len, sink_bias=mSB0, row_sums=mRS0)
comp0(mQ0, mK0, mV0, mC0, stream, mLSE0,
swa_len=swa_len, sink_bias=mSB0, row_sums=mRS0)
torch.cuda.synchronize()
ok0 = c0[:, :, 0].float() / rs0[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
# Per-segment reference for segment 0
scores0_ref = q.float() @ k_seg0.float().T * scale
scores0_ref[:, n_comp:] += attn_sink_val
if swa_len < (seg_size - n_comp): # 100 < 32? No
scores0_ref[:, n_comp + swa_len:] = float('-inf')
ref0 = (torch.softmax(scores0_ref, dim=-1) @ v_seg0.float()).to(torch.bfloat16)
cos0 = torch.nn.functional.cosine_similarity(ok0.flatten().unsqueeze(0), ref0.flatten().unsqueeze(0).float()).item()
print(f'Seg0 direct: cos {cos0:.6f} LSE_kern={lse0[0,0,0].item():.4f} LSE_ref={torch.logsumexp(scores0_ref, dim=-1)[0].item():.4f}')
print()
k_combined = torch.cat([k_comp, k_swa], dim=0)
v_combined = torch.cat([v_comp, v_swa], dim=0)
# Reference
ref = reference_combined_attention(
q, k_comp, v_comp, k_swa, v_swa,
attn_sink_val, scale, swa_len
)
# Full reference
ref = reference_combined_attention(q, k_comp, v_comp, k_swa, v_swa,
attn_sink_val, scale, swa_len)
# Combined KV
k_combined = torch.cat([k_comp, k_swa], dim=0) # (256, hd)
v_combined = torch.cat([v_comp, v_swa], dim=0) # (256, hd)
# Split into 128-token segments and run FMHA per segment
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
seg_size = 128
n_segs = (n_total + seg_size - 1) // seg_size
# 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.
# Per-segment reference and kernel verification
# Segment 0: positions 0-127 (n_comp=96, 32 SWA tokens, all valid, sink bias on SWA)
k0 = k_combined[0:128]; v0 = v_combined[0:128]
scores0 = q.float() @ k0.float().T * scale
scores0[:, n_comp:] += attn_sink_val
# swa_len=100 but only 32 SWA positions → no D3 masking
ref0 = (torch.softmax(scores0, dim=-1) @ v0.float()).to(torch.bfloat16)
lse0_ref = torch.logsumexp(scores0, dim=-1)
# 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_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_comp = cute.compile(kernel_comp, _mQ, _mK, _mV, _mC, stream, _mLSE,
swa_len=swa_len, sink_bias=_mSB, row_sums=_mRS)
k_seg0 = FmhaKernel(head_dim=hd, s_k=128, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=n_comp, apply_sink_bias=True)
c0 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
l0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
r0 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
comp0 = cute.compile(k_seg0, to_cute(q.unsqueeze(-1)), to_cute(k0.unsqueeze(-1)),
to_cute(v0.unsqueeze(-1)), to_cute(c0), stream, to_cute(l0),
swa_len=100, sink_bias=to_cute(attn_sink), row_sums=to_cute(r0))
comp0(to_cute(q.unsqueeze(-1)), to_cute(k0.unsqueeze(-1)),
to_cute(v0.unsqueeze(-1)), to_cute(c0), stream, to_cute(l0),
swa_len=100, sink_bias=to_cute(attn_sink), row_sums=to_cute(r0))
torch.cuda.synchronize()
ok0 = c0[:, :, 0].float() / r0[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
cos0 = torch.nn.functional.cosine_similarity(ok0.flatten().unsqueeze(0), ref0.flatten().unsqueeze(0).float()).item()
lse0_kern = l0[:, 0, 0]
print(f'Seg0: cos {cos0:.6f} LSE_kern[0]={lse0_kern[0].item():.4f} LSE_ref[0]={lse0_ref[0].item():.4f}')
# Kernel 2: SWA only (n_comp=0, but sink bias still applies to all positions)
kernel_swa = FmhaKernel(head_dim=hd, s_k=seg_size, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=0, apply_sink_bias=True)
# 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)
# Segment 1: positions 128-255 (all SWA, n_comp=0, sink bias on all, D3 mask at position 68)
k1 = k_combined[128:256]; v1 = v_combined[128:256]
scores1 = q.float() @ k1.float().T * scale
scores1 += attn_sink_val # all SWA → sink bias on all
# Valid SWA: absolute 96-195. In this segment (128-255): positions 0-67 valid, 68-127 masked
swa_len_seg1 = 68 # n_comp + swa_len - k_start = 96 + 100 - 128 = 68
scores1[:, swa_len_seg1:] = float('-inf')
ref1 = (torch.softmax(scores1, dim=-1) @ v1.float()).to(torch.bfloat16)
lse1_ref = torch.logsumexp(scores1, dim=-1)
# Run each segment
segment_results = []
n_comp_global = n_comp
swa_len_global = swa_len # number of valid SWA tokens (relative to SWA region start at n_comp)
k_seg1 = FmhaKernel(head_dim=hd, s_k=128, normalize=False,
apply_swa_mask=True, is_causal=False, n_comp=0, apply_sink_bias=True)
c1 = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
l1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
r1 = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
comp1 = cute.compile(k_seg1, to_cute(q.unsqueeze(-1)), to_cute(k1.unsqueeze(-1)),
to_cute(v1.unsqueeze(-1)), to_cute(c1), stream, to_cute(l1),
swa_len=swa_len_seg1, sink_bias=to_cute(attn_sink), row_sums=to_cute(r1))
comp1(to_cute(q.unsqueeze(-1)), to_cute(k1.unsqueeze(-1)),
to_cute(v1.unsqueeze(-1)), to_cute(c1), stream, to_cute(l1),
swa_len=swa_len_seg1, sink_bias=to_cute(attn_sink), row_sums=to_cute(r1))
torch.cuda.synchronize()
ok1 = c1[:, :, 0].float() / r1[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30)
cos1 = torch.nn.functional.cosine_similarity(ok1.flatten().unsqueeze(0), ref1.flatten().unsqueeze(0).float()).item()
lse1_kern = l1[:, 0, 0]
print(f'Seg1: cos {cos1:.6f} LSE_kern[0]={lse1_kern[0].item():.4f} LSE_ref[0]={lse1_ref[0].item():.4f}')
for seg_idx in range(n_segs):
k_start = seg_idx * seg_size
k_end = min(k_start + seg_size, n_total)
k_seg = k_combined[k_start:k_end]
v_seg = v_combined[k_start:k_end]
if k_seg.shape[0] == 0:
continue
# n_comp_local: compressed KV tokens in this segment
# If segment starts before n_comp_global, some positions are compressed.
# If segment starts at or after n_comp_global, all positions are SWA.
if k_start < n_comp_global:
n_comp_local = min(n_comp_global - k_start, seg_size)
else:
n_comp_local = 0
# swa_len_local: the kernel masks kv_pos >= n_comp_local + swa_len_local
# We want to mask absolute positions >= n_comp_global + swa_len_global
# So: n_comp_local + swa_len_local = n_comp_global + swa_len_global - k_start
# => swa_len_local = n_comp_global + swa_len_global - k_start - n_comp_local
swa_len_local = n_comp_global + swa_len_global - k_start - n_comp_local
# Clamp: if > remaining segment size, no masking needed
swa_len_local = min(swa_len_local, seg_size)
# If <= 0, all SWA positions are masked (no valid SWA)
if swa_len_local <= 0 and n_comp_local == 0:
continue # Skip empty segment
# 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, kern, comp, stream,
sink_bias=attn_sink if has_swa else None,
n_comp_in_seg=n_comp_local,
swa_len=swa_len_local
)
segment_results.append((o_norm, lse))
# Merge segments
o_merged = python_kv_merge(segment_results)
# Compare
# Merge
o_merged = python_kv_merge([(ok0, lse0_kern), (ok1, lse1_kern)])
cos = torch.nn.functional.cosine_similarity(
o_merged.flatten().unsqueeze(0),
ref.flatten().unsqueeze(0).float()
).item()
max_abs = (o_merged - ref.float()).abs().max().item()
o_merged.flatten().unsqueeze(0), ref.flatten().unsqueeze(0).float()).item()
status = "PASS" if cos >= 0.99 else "FAIL"
print(f'D5c multi-tile: cos {cos:.6f} max_abs {max_abs:.4f} {status}')
print(f'\nMerged: cos {cos:.6f} {status}')
if cos < 0.99:
print(f' kernel[0,:4]={o_merged[0,:4].tolist()}')
print(f' ref[0,:4]={ref[0,:4].tolist()}')
# Debug: check each segment individually against per-segment reference
for i, (o_n, lse) in enumerate(segment_results):
k_start = i * seg_size
k_end = min(k_start + seg_size, n_total)
k_seg = k_combined[k_start:k_end]
v_seg = v_combined[k_start:k_end]
n_comp_local = min(max(n_comp - k_start, 0), seg_size)
# Per-segment reference
scores_ref = q.float() @ k_seg.float().T * scale
if k_start < n_comp:
scores_ref[:, n_comp - k_start:] += attn_sink_val
swa_end = n_comp + swa_len - k_start
if swa_end < seg_size:
scores_ref[:, swa_end:] = float('-inf')
else:
scores_ref += attn_sink_val
swa_len_in_seg = n_comp + swa_len - k_start
if swa_len_in_seg < seg_size:
scores_ref[:, swa_len_in_seg:] = float('-inf')
o_ref_seg = (torch.softmax(scores_ref, dim=-1) @ v_seg.float()).to(torch.bfloat16)
lse_ref = torch.logsumexp(scores_ref, dim=-1)
cos_seg = torch.nn.functional.cosine_similarity(
o_n.flatten().unsqueeze(0), o_ref_seg.flatten().unsqueeze(0).float()
).item()
print(f' Seg {i}: LSE_kern={lse[0].item():.4f} LSE_ref={lse_ref[0].item():.4f} cos={cos_seg:.6f}')
if __name__ == '__main__':