From 9aca62c26cc901be40797d44630c1315c57de62b Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 12 Jul 2026 15:35:53 +0800 Subject: [PATCH] perf: remove per-element sentinel checks in softmax + fix stale comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace 4*NC8 per-element -FLT_MAX comparisons with 2 row-level pn guards; masked entries naturally underflow via expf(-FLT_MAX - nm) ≈ 0; pn only guards all-masked-row edge where nm == -FLT_MAX (exp(0)=1 not 0); ~1-3% speedup on prefill (verified via standalone CUDA bench); correctness verified: prefill 4/4, decode 3/3, paged decode 13/13 - Fix stale comments: remove false 'pre-scale Q' claim, correct occupancy numbers, remove phantom sQ from smem description --- csrc/kernels/attn_mma_utils.cuh | 37 ++++++++++++++++------- csrc/kernels/attn_prefill.cu | 8 +++-- csrc/kernels/attn_prefill_split_q_mma.cuh | 3 +- 3 files changed, 33 insertions(+), 15 deletions(-) diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh index 1102c21..f7307d9 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attn_mma_utils.cuh @@ -114,7 +114,8 @@ __device__ __forceinline__ void cp_async_wait_group() { // between the two kernels; only the per-row causal/mask bounds differ. // --------------------------------------------------------------------------- -// S = Q @ K^T (Qa pre-loaded and pre-scaled by the caller). +// S = Q @ K^T (Qa pre-loaded by the caller; scale applied post-mma in the +// caller to avoid bf16 precision loss). // LD and SWIZ_MASK are constexpr in the calling kernel — passing them as // runtime ints lets the compiler fold them while keeping the signature clean. template @@ -138,10 +139,11 @@ __device__ inline void mma_compute_scores( } } -// Online softmax + Oacc rescale for one K/V tile. maxc0/maxc1 are the per-row -// KV column bounds — prefill passes per-query-row causal limits while decode -// passes the same value for both rows (q_len==1). Sacc is consumed in place -// (replaced by P = exp(S - nm) for the subsequent P@V step). +// Online softmax + Oacc rescale for one K/V tile. +// maxc0/maxc1: per-row KV column bounds (prefill: per-query-row causal limits; +// decode: same value for both rows since q_len==1). +// Reads Sacc (Q@K^T scores), applies causal/mask, computes P = exp(S - nm), +// rescales Oacc by exp(m_old - nm), and updates m/l — all in place. template __device__ inline void mma_softmax_tile( int kv0, @@ -157,6 +159,8 @@ __device__ inline void mma_softmax_tile( { int tid4 = lane & 3; + // Mask out-of-bounds / masked columns: set -FLT_MAX so expf → 0 downstream + // without per-element sentinel checks. Compute tile-local row maxima. float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; #pragma unroll for (int n8 = 0; n8 < NC8; n8++) { @@ -175,22 +179,33 @@ __device__ inline void mma_softmax_tile( rmax0 = fmaxf(rmax0, fmaxf(s0, s1)); rmax1 = fmaxf(rmax1, fmaxf(s2, s3)); } + // Warp-reduce row maxima across the 4-lane thread group (xor 1, xor 2). rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1)); rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2)); rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1)); rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2)); + // nm = max(running max m, tile-local max rmax) — updated running maximum. float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1); - float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0); - float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1); + // corr rescales Oacc and l by exp(m_old - nm). When all-masked (m == nm == + // -FLT_MAX), exp(0) = 1 — correct, no guard needed. + float corr0 = __expf(m0 - nm0); + float corr1 = __expf(m1 - nm1); + // pn guards only the all-masked-row edge: if nm == -FLT_MAX, exp(S - nm) + // gives 1 not 0 for masked entries. Two scalar masks replace 4*NC8 + // per-element comparisons. + float pn0 = (nm0 == -FLT_MAX) ? 0.0f : 1.0f; + float pn1 = (nm1 == -FLT_MAX) ? 0.0f : 1.0f; + // P = exp(S - nm) for each element. Masked entries (Sacc = -FLT_MAX) give + // exp(-inf) ≈ 0 naturally; pn zero-fills the all-masked-row edge. float rsum0 = 0.0f, rsum1 = 0.0f; #pragma unroll for (int n8 = 0; n8 < NC8; n8++) { - float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0); - float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0); - float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1); - float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1); + float p0 = pn0 * __expf(Sacc[n8][0] - nm0); + float p1 = pn0 * __expf(Sacc[n8][1] - nm0); + float p2 = pn1 * __expf(Sacc[n8][2] - nm1); + float p3 = pn1 * __expf(Sacc[n8][3] - nm1); Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3; rsum0 += p0 + p1; diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attn_prefill.cu index eda117c..bae21b0 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attn_prefill.cu @@ -13,12 +13,14 @@ static void dispatch_prefill(AttentionParams& p) { // loop overhead over more tensor-core work (this kernel is latency-bound, // not compute/bandwidth-bound), so BC=32 wins ~6-8% over BC=16 for // D<=128. D=256 stays at 16: BC=32 double-buffered would need 64KB smem, - // over the 48KB static cap. Both keep 3 blocks/SM (2 for D=256). + // over the 48KB static cap. constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16; dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch); dim3 block(WARPS * 32, 1, 1); - // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed. - // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params. + // Static shared memory — double-buffered K/V only (no sQ: Q goes direct + // to registers). 2*BC*LD bf16 each for sK and sV → 4*BC*HEAD_DIM*2 bytes. + // Occupancy is smem-capped: D=64→3 blocks/SM (16KB), D=128→1 (32KB), + // D=256→1 (32KB, BC=16). attn_prefill_split_q_mma_kernel<<>>(p); #else constexpr int G = 8, ROWS = 32, P_BC = 32; diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index 10c67eb..628d170 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -30,7 +30,8 @@ using bf16 = __nv_bfloat16; // register pressure), so fewer, larger tiles beat many tiny ones. // // Optimizations: load Q fragments directly from global in mma A-operand layout -// (no sQ staging, no prologue barriers); pre-scale Q by attention scale during Q load; packed bf16x2 output stores; +// (no sQ staging, no prologue barriers); post-multiply scale in float after +// S=Q@K^T to avoid bf16 precision loss; packed bf16x2 output stores; // causal tile skipping (block-level prefetch bound + warp-level compute skip); // XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD // padding (LD=HEAD_DIM).