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

Normalize O_accum by exp(lse) before returning

Browse Source
This commit is contained in:
biondizzle
2026-05-27 06:39:36 +00:00
parent ddc701af9b
commit 312ac52d15
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -193,8 +193,8 @@ def _attention_single_head(
row_sums_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
# Prepare CuTe tensors
q_input = q[0].unsqueeze(-1) # (T, hd, 1)
k_input = k_seg[0].unsqueeze(-1) # (s_k_per_seg, hd, 1)
q_input = q[0].contiguous().unsqueeze(-1) # (T, hd, 1)
k_input = k_seg[0].contiguous().unsqueeze(-1) # (s_k_per_seg, hd, 1)
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
mQ = ct.from_dlpack(q_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_input))
@@ -222,6 +222,9 @@ def _attention_single_head(
o_accum = (e_old * o_accum + e_new * seg_o) / e_sum
lse_accum = torch.log(e_sum)
# o_accum is already normalized (in the merge above)
output = o_accum.to(torch.bfloat16).unsqueeze(0) # (1, T, hd)
# o_accum is the LSE-merged un-normalized O. Normalize by the final LSE.
# O_norm = O_unnorm / row_sum, where row_sum = exp(lse)
row_sum = torch.exp(lse_accum).clamp(min=1e-30)
o_norm = o_accum / row_sum
output = o_norm.to(torch.bfloat16).unsqueeze(0) # (1, T, hd)
return output