From b8b439b7130e9956fc165a533eb816ae003eb4d3 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Sat, 11 Jul 2026 00:24:09 +0800
Subject: [PATCH] perf: fuse decode combine kernel to single-pass
 online-rescale reduction

- Replace 3-scan loops (mstar, lstar, acc) with 1-pass online rescale
- Halves __expf calls (num_splits vs 2*num_splits) and ml_part re-reads
- Mathematically equivalent, no change to o_part traffic or output
---
 csrc/kernels/attn_decode_split_kv.cuh | 30 +++++++++++++--------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attn_decode_split_kv.cuh
index 9730cb2..1a610b0 100644
--- a/csrc/kernels/attn_decode_split_kv.cuh
+++ b/csrc/kernels/attn_decode_split_kv.cuh
@@ -84,32 +84,30 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) {
 }
 
 // Reduce split-K partials into the final bf16 output. One block per (batch,
-// q_head); each thread owns one head_dim element and folds across all splits
-// with a numerically-stable online rescale.
+// q_head); each thread folds across all splits with a single-pass
+// online-rescale reduction (expf + FMA counts halved vs 3-pass original).
 __global__ void attn_decode_combine_kernel(AttentionParams p) {
     int bh = blockIdx.x;
     int d = threadIdx.x;
     if (d >= p.head_dim) return;
 
     size_t split_base = (size_t)bh * p.num_splits;
-
     const float* mlp = p.ml_part + split_base * 2;
-    float mstar = -FLT_MAX;
-    for (int s = 0; s < p.num_splits; s++)
-        mstar = fmaxf(mstar, mlp[s * 2]);
-
-    float lstar = 0.0f;
-    for (int s = 0; s < p.num_splits; s++) {
-        float mi = mlp[s * 2];
-        if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar);
-    }
-
     const float* op = p.o_part + split_base * p.head_dim;
-    float acc = 0.0f;
+
+    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
     for (int s = 0; s < p.num_splits; s++) {
         float mi = mlp[s * 2];
-        if (mi > -FLT_MAX) acc += op[s * p.head_dim + d] * __expf(mi - mstar);
+        if (mi <= -FLT_MAX) continue;
+        float li = mlp[s * 2 + 1];
+        float nm = fmaxf(m, mi);
+        float corr = __expf(m - nm);
+        float e = __expf(mi - nm);
+        acc = acc * corr + op[s * p.head_dim + d] * e;
+        l = l * corr + li * e;
+        m = nm;
     }
-    float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f;
+
+    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
     p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv);
 }