From b70ab2a6ee7fc73a86d8e15c022cb8405ca2049f Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Wed, 27 May 2026 06:42:58 +0000
Subject: [PATCH] Return o_accum directly (un-normalized merge result)

---
 dsv4/kernels/attention/production.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py
index a96bc8d3..33328315 100644
--- a/dsv4/kernels/attention/production.py
+++ b/dsv4/kernels/attention/production.py
@@ -222,9 +222,5 @@ 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 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)
+    output = o_accum.to(torch.bfloat16).unsqueeze(0)  # (1, T, hd)
     return output