| title | Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration |
|---|---|
| authors | Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski |
| date | 2026-09-27 |
| venue | arXiv (preprint, not peer reviewed) |
| id | arXiv:2609.33586 |
| affiliation | University of Illinois at Urbana-Champaign |
| upvotes | 4 (Hugging Face Daily Papers, 2026-11-01) |
| links | abs pdf hf |
| 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. |
cite
@misc{zhu2026approximating,
title={Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration},
author={Shangzhen Zhu and Muyan Hu and Tomasz Kozlowski},
year={2026},
eprint={2609.33586},
archivePrefix={arXiv},
url={https://arxiv.org/abs/2609.33586}
}