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

Fix KV merge: use NORMALIZED O (O_unnorm/row_sum) with LSE

Browse Source
This commit is contained in:
biondizzle
2026-05-27 07:07:51 +00:00
parent 06e7f7ab48
commit 0736a04d9b
1 changed files with 15 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
dsv4/kernels/attention/production.py
View File

@@ -165,6 +165,7 @@ def _attention_single_head(
seg_o = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
seg_lse = torch.zeros(T, 1, dtype=torch.float32, device='cuda')
seg_row_sums = torch.zeros(T, 1, dtype=torch.float32, device='cuda')
for nt in range(n_pv_tiles):
v_start = nt * pv_n_tile
@@ -173,34 +174,38 @@ def _attention_single_head(
v_kernel = v_tile.unsqueeze(-1)
c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
lse_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
row_sums_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d))
mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg))
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
compiled(mQ, mK, mV, mC, stream, lse=mLSE)
compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
torch.cuda.synchronize()
seg_o[:, v_start:v_end] = c_tile[:, :, 0].float()
if nt == 0:
seg_lse[:, 0] = lse_tensor[:, 0, 0].float()
seg_row_sums[:, 0] = row_sums_tensor[:, 0, 0].float()
# Merge with accumulator using log-sum-exp
# Normalize segment output: O_norm = O_unnorm / row_sum (per row)
seg_row_sums = seg_row_sums.clamp(min=1e-30)
seg_o_norm = seg_o / seg_row_sums # (T, hd) normalized
# Merge with accumulator using CORRECT formula:
# O = sum(exp(lse_i) * O_i_norm) / sum(exp(lse_i))
# where O_i_norm = O_i_unnorm / row_sum_i
#
# o_accum is tracked in NORMALIZED form.
e_old = torch.exp(lse_accum)
e_new = torch.exp(seg_lse)
e_sum = e_old + e_new
o_accum = (e_old * o_accum + e_new * seg_o) / e_sum
o_accum = (e_old * o_accum + e_new * seg_o_norm) / e_sum
lse_accum = torch.log(e_sum)
# Debug: check LSE values
if seg == 0:
print(f' seg 0: lse[0]={seg_lse[0,0].item():.4f}, o[0,0]={seg_o[0,0].item():.4f}')
elif seg == 1:
print(f' seg 1: lse[0]={seg_lse[0,0].item():.4f}, o[0,0]={seg_o[0,0].item():.4f}')
print(f' merged: lse[0]={lse_accum[0,0].item():.4f}, o[0,0]={o_accum[0,0].item():.4f}')
output = o_accum.to(torch.bfloat16).unsqueeze(0)
return output