From 312ac52d15e23f34941392e271db60a2379b0364 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 06:39:36 +0000 Subject: [PATCH] Normalize O_accum by exp(lse) before returning --- dsv4/kernels/attention/production.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py index f33c7e03..c2a62831 100644 --- a/dsv4/kernels/attention/production.py +++ b/dsv4/kernels/attention/production.py @@ -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