Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration
Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski
University of Illinois Urbana-Champaign
Published September 27, 2026 · on the feed September 29, 2026 · ▲ 4 · arXiv 2609.33586
The abstract
On NVIDIA Blackwell B200, tensor-core throughput outpaces special-function exponential throughput by more than two orders of magnitude, exposing exponential evaluation in fused attention kernels. A pretrained Transformer, however, may not need it evaluated accurately at every element. We characterize what a pretrained model does need by approximating softmax at inference in ten frozen decoder-only models (0.5B-72B). The number of positions the softmax map assigns probability to and within-row resolution can be cut substantially, yet uniform weighting of the same positions is damaging. Where a fixed resolution budget is placed matters as much as its size, with resolution near the row maximum consistently favored. Perturbations matched on scalar distortion produce model-dependent responses of opposite sign. These findings motivate Rowmax-PoT, a coarse logarithmic weight representation anchored at each row maximum, and Rowmax-H15, its hardware specialization in FlashAttention-4. On B200, the patched FP8 attention forward is 12.4% faster at causal 8K and 25.8% faster at non-causal 8K in host-side call-latency measurements; board energy per forward falls by 8.4% at causal 16K. Measured separately on the BF16 kernel path at 2K, Rowmax-H15 increases perplexity by 0.091-0.492% across five models from three families.
The Brief
2:29A short podcast summary of the paper. Written and voiced by AI.
Transcript
Approximating Softmax in Pretrained LLMs. Model Sensitivity and Kernel Acceleration. By Shangzhen Zhu and two others, from the University of Illinois Urbana-Champaign. Published on the twenty-seventh of September. No code repository is linked to it.
On NVIDIA's Blackwell B200, tensor cores deliver more than two orders of magnitude more throughput than the special-function units that compute exponentials. Inside a fused attention kernel, that leaves the exponential in softmax exposed as the slow part. The paper's claim is that a pretrained Transformer may not need that exponential evaluated accurately at every element, and it sets out to characterize what such a model does need.
The method is to approximate softmax at inference time in ten frozen decoder-only models, ranging from zero point five billion to seventy-two billion parameters. Two quantities turn out to be cuttable by a substantial margin: the number of positions that softmax assigns any probability to, and the resolution within a row. But weighting those same positions uniformly does real damage. Where a fixed resolution budget is spent matters as much as how big it is, and resolution near the row maximum is consistently the favoured place. One more result cautions against simple summaries: perturbations matched on a scalar measure of distortion draw responses of opposite sign from different models.
Those findings motivate Rowmax-PoT, a coarse logarithmic representation of the weights anchored at each row's maximum, and Rowmax-H15, its hardware specialization inside FlashAttention-4. On B200, the patched FP8 attention forward pass is twelve point four percent faster at causal eight thousand tokens, and twenty-five point eight percent faster at non-causal eight thousand. Board energy per forward pass drops by eight point four percent at causal sixteen thousand.
Several qualifiers travel with those numbers. The speedups are host-side call-latency measurements on B200. Accuracy was measured separately, on the BF16 kernel path at two thousand tokens, where Rowmax-H15 raises perplexity by between zero point zero nine one and zero point four nine two percent, across five models from three families. And the sensitivity findings are model-dependent by the paper's own account, so the favoured placement of resolution is a tendency across these ten models rather than a law.
Where the precision goes, then, turns out to matter more than how much of it there is. Keep detail near the largest score in each row, let the rest go coarse, and on this evidence a pretrained model barely notices, while the attention kernel built on that idea runs measurably faster.
Paper Quest · 7 chapters · about 41 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
A GPU that multiplies fast and exponentiates slowly. Something has to give.
You are running large language models on a new GPU. Its matrix units outpace the unit that computes exponentials by more than two orders of magnitude, so the exponential inside attention has become the slow part.
The road ahead
- 1 · The problem ~7 min On a new GPU the exponential inside attention has become the slow step, yet you cannot retrain the models or know which of its precision they need.
- 2 · What came before ~3 min
- 3 · The key idea ~9 min Probe frozen models to learn what softmax weights must keep, then approximate the exponential coarsely, keeping fine steps near each row's maximum.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min Probe frozen models, learn that resolution near the row maximum matters most, build Rowmax-PoT around that, then specialize it as Rowmax-H15 in FlashAttention-4.
- 6 · The experiment ~6 min You’ll guess: Replacing the exact exponential in attention with the coarse Rowmax-H15 approximation, how much did perplexity rise across five pretrained models?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.