From 8f87109f869a462de2057beffb19d9d414875473 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Wed, 27 May 2026 06:48:56 +0000
Subject: [PATCH] Single-segment: use normalize=False + per-row normalization
 from row_sums

---
 dsv4/kernels/attention/production.py | 49 ++++++++++++++++++++++++----
 1 file changed, 43 insertions(+), 6 deletions(-)

diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py
index 7a7232d0..81d2a282 100644
--- a/dsv4/kernels/attention/production.py
+++ b/dsv4/kernels/attention/production.py
@@ -134,7 +134,7 @@ def _attention_single_head_normalized(
     sink_bias: torch.Tensor = None,
     use_smem_p: bool = False,
 ) -> torch.Tensor:
-    """Run FMHA for a single head with in-kernel normalization (single KV tile)."""
+    """Run FMHA for a single head with Python normalization (single KV tile)."""
     _, T, hd = q.shape
     N = k.shape[1]
     apply_swa_mask = swa_len is not None
@@ -142,14 +142,15 @@ def _attention_single_head_normalized(
 
     compiled, kernel = _get_or_compile_kernel(
         head_dim=hd, s_k=N, use_smem_p=use_smem_p,
-        normalize=True, apply_swa_mask=apply_swa_mask,
+        normalize=False, apply_swa_mask=apply_swa_mask,
         is_causal=is_causal, n_comp=n_comp if n_comp > 0 else None,
         apply_sink_bias=apply_sink_bias,
     )
     pv_n_tile = kernel.pv_n_tile
     n_pv_tiles = kernel.n_pv_tiles
 
-    output = torch.zeros(T, hd, dtype=torch.bfloat16, device='cuda')
+    output_unnorm = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
+    lse_val = None
 
     for nt in range(n_pv_tiles):
         v_start = nt * pv_n_tile
@@ -157,6 +158,8 @@ def _attention_single_head_normalized(
         v_tile = v[0, :, v_start:v_end].contiguous()
         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')
 
         q_input = q[0].contiguous().unsqueeze(-1)
         k_input = k[0].contiguous().unsqueeze(-1)
@@ -166,13 +169,47 @@ def _attention_single_head_normalized(
         mK = ct.from_dlpack(k_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_input))
         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)
+        compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
         torch.cuda.synchronize()
 
-        output[:, v_start:v_end] = c_tile[:, :, 0]
+        output_unnorm[:, v_start:v_end] = c_tile[:, :, 0].float()
+        if nt == 0:
+            lse_val = lse_tensor[0, 0, 0].item()
 
-    return output.unsqueeze(0)  # (1, T, hd)
+    # Normalize: O_norm = O_unnorm / row_sum
+    # row_sum is computed from lse: row_sum = exp(lse - row_max * ln2) ... complex
+    # But we have row_sums from the kernel!
+    # row_sums_tensor[0,0,0] has the row_sum for row 0 only.
+    # For per-row normalization, we'd need per-row row_sums.
+    # The kernel outputs row_sums[sfw_idx, 0, 0] for each of 128 rows.
+    # But we can also compute from the un-normalized O and the reference.
+    #
+    # Simpler: compute row_sum from the un-normalized O and the LSE.
+    # lse = ln(row_sum) + row_max * ln(2)
+    # exp(lse) = row_sum * exp(row_max * ln(2)) = row_sum * 2^row_max
+    # O_unnorm = P @ V where P = softmax * row_sum
+    # Wait, the kernel's P = exp2(S*scale - row_max) which is NOT softmax.
+    # softmax = P / row_sum
+    # O_unnorm = P @ V = softmax * row_sum @ V = O_norm * row_sum
+    # So: O_norm = O_unnorm / row_sum
+    #
+    # We need row_sum per row. The kernel outputs it in row_sums_tensor.
+    # But row_sums_tensor only has values for 128 rows, and we need all T rows.
+    # Actually T=128, so row_sums_tensor[0:128, 0, 0] has all rows.
+    #
+    # Let me extract per-row row_sums.
+    # But the kernel only writes row_sums[sfw_idx, 0, 0] for sfw_idx 0..127.
+    # And row_sums_tensor is (T, 1, 1) = (128, 1, 1).
+    # So row_sums_tensor[:, 0, 0] should have all 128 rows.
+
+    # Re-run to get row_sums (we already have them from the last call)
+    row_sums_per_row = row_sums_tensor[:, 0, 0].float().unsqueeze(1)  # (T, 1)
+    row_sums_per_row = row_sums_per_row.clamp(min=1e-30)
+    output_norm = output_unnorm / row_sums_per_row
+    return output_norm.to(torch.bfloat16).unsqueeze(0)  # (1, T, hd)
 
 
 def _attention_single_head(