Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration
Approximating softmax in ten frozen LLMs yields a FlashAttention-4 kernel up to 25.8% faster on NVIDIA B200.
The paper approximates softmax at inference in ten frozen decoder-only models ranging from 0.5B to 72B parameters. It finds probability mass and within-row resolution can be cut, but uniform weighting of the same positions hurts, and resolution near each row maximum is preferred. Rowmax-PoT and its FlashAttention-4 specialization, Rowmax-H15, make the patched FP8 attention forward 12.4% faster at causal 8K and 25.8% faster at non-causal 8K on NVIDIA Blackwell B200, while cutting board energy 8.4% at causal 16K. On a separate BF16 path at 2K, perplexity rises 0.091–0.492% across five models from three families.
- Ten frozen decoder-only models from 0.5B to 72B were tested.
- FP8 attention is 12.4% faster at causal 8K on B200.
- Non-causal 8K FP8 forward is 25.8% faster in host latency.
- BF16 perplexity rises 0.091–0.492% across five models.
- Board energy per causal 16K forward falls 8.4%.
- Vendors
- NVIDIA
- Products
- Blackwell B200FlashAttention-4
From the original article
1 min read at huggingface.co
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.
By Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.