From 0736a04d9b7d0f9ae4db27c593f7188259a75e2c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 07:07:51 +0000 Subject: [PATCH] Fix KV merge: use NORMALIZED O (O_unnorm/row_sum) with LSE --- dsv4/kernels/attention/production.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py index 4d931449..3f89f1a8 100644 --- a/dsv4/kernels/attention/production.py +++ b/dsv4/kernels/attention/production.py @@ -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