Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 24 of its 30 references that are in the map. First posted on arXiv 27 Sep 2026.
Builds on
24 of 30 references.
References · 24
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 507
- 2021I-BERT: Integer-only BERT Quantizationcited by 250
- 2021PTQ4ViT: Post-training quantization for vision transformers with twin uniform quantizationcited by 173
- 2024SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Accelerationcited by 158
- 2021FQ-ViT: Post-Training Quantization for Fully Quantized Vision Transformercited by 148
- 2022RepQ-ViT: Scale Reparameterization for Post-Training Quantization of Vision Transformerscited by 134
- 2024SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantizationcited by 113
- 2022I-ViT: Integer-only Quantization for Efficient Vision Transformer Inferencecited by 93
- 2023Towards Accurate Post-Training Quantization for Vision Transformercited by 70
- 2025SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Trainingcited by 63
12 more
- 2026FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalingcited by 58
- 2021Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformerscited by 52
- 2024PTQ4SAM: Post-Training Quantization for Segment Anythingcited by 32
- 2025Efficient Attention Mechanisms for Large Language Models: A Surveycited by 21
- 2024ConSmax: Hardware-Friendly Alternative Softmax with Learnable Parameterscited by 14
- 2023ITA: An Energy-Efficient Attention and Softmax Accelerator for Quantized Transformerscited by 12
- 2024INT-FlashAttention: Enabling Flash Attention for INT8 Quantizationcited by 9
- 2024LRP-QViT: Mixed-Precision Vision Transformer Quantization via Layer-wise Relevance Propagationcited by 8
- 2024EXAQ: Exponent Aware Quantization For LLMs Accelerationcited by 3
- 2025IntAttention: A Fully Integer Attention Pipeline for Efficient Edge Inferencecited by 3
- 2026VFA: Relieving Vector Operations in Flash Attention with Global Maximum Pre-computationcited by 2
- 2026MXAttention: Data-Free Optimal Scaling and Pre-Normalization Quantization for MXFP4 Attentioncited by 2
Built on by
Not cited yet.