perf: remove per-element sentinel checks in softmax + fix stale comments

- 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
This commit is contained in:
2026-07-12 15:45:08 +08:00
parent b5cdea98ad
commit 9aca62c26c
3 changed files with 33 additions and 15 deletions
+2 -1
View File
@@ -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).