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

fix: simplify run_segment to use full hd V tensor (was incorrectly splitting by pv_n_tile)

Browse Source
This commit is contained in:
biondizzle
2026-05-26 15:34:57 +00:00
parent aa2df1a202
commit 2a5f9dc6e3
1 changed files with 22 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
tests/unit/test_d5c_multitile.py
View File

@@ -51,48 +51,40 @@ def run_segment(q, k_seg, v_seg, kernel, compiled, stream,
"""Run one 128-token segment of FMHA, return normalized O and LSE."""
m = q.shape[0]
hd = v_seg.shape[1]
pv_n_tile = kernel.pv_n_tile
# 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))
for nt in range(kernel.n_pv_tiles):
v_start = nt * pv_n_tile
v_end = v_start + pv_n_tile
v_tile = v_seg[:, v_start:v_end].contiguous()
c_tile = torch.zeros(m, pv_n_tile, 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')
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_tile.unsqueeze(-1))
mC = to_cute(c_tile)
mLSE = to_cute(lse_tile)
mRS = to_cute(rs_tile)
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)
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_seg[:, v_start:v_end] = c_tile[:, :, 0]
if nt == 0:
lse_seg = lse_tile[:, 0, 0].clone()
rs_seg = rs_tile[:, 0, 0].clone()
# Note: LSE and row_sum are the same across PV tiles (same softmax)
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()
# Normalize using row_sum
o_norm = o_seg.float() / rs_seg.unsqueeze(1).clamp(min=1e-30)
return o_norm, lse_seg
return o_norm, lse_val
def python_kv_merge(segment_results):