Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

Hugging Face Daily Papers Papers
softmax-approximation llm-inference kernel-acceleration flashattention hardware-optimization fp8-training model-analysis

Summary

This paper explores approximating softmax in pretrained LLMs for kernel acceleration, demonstrating performance gains like up to 25.8% speedup on Blackwell B200 with minimal perplexity impact.

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.
Original Article
View Cached Full Text

Cached at: 09/30/26, 04:23 AM

Paper page - Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

Source: https://huggingface.co/papers/2609.33586 On Blackwell B200, tensor cores outrun the exponential unit by ~500× (8192 vs 16 ops/clk/SM), soexpbecomes exposed in fused attention. We ask what afrozenpretrained LLM actually needs from softmax — and use the answer to replaceexp2inside FlashAttention-4.

What frozen models need(10 models, 0.5B–72B, no retraining):

  • Attention support and within-row resolution can be cut substantially — but uniform weighting on the same support hurts at every tested layer.
  • Whereresolution goes matters as much as how much: finer intervals near the row maximum lower NLL in all 10 models, even when overall approximation error goes up.
  • A scalar distortion budget is not enough: flattening vs. sharpening at equal attention JSD gives opposite-sign losses depending on the model.

**Rowmax-H15 in FA4:**weights snap to {1, 1.5}×2^k anchored at the kernel’s running row max — no calibration, one FP32 add + one bit shift per element.

  • FP8 attention forward on B200:+12.4%(causal 8K),+25.8%(non-causal 8K); 9.3% faster than the best stock FA4 emulation setting we tested
  • **−8.4%**board energy per forward (causal 16K)
  • Perplexity**+0.09–0.49%**across 5 models from 3 families (BF16 kernel, 2K)

Scope: attention forward on B200; Approximate Softmax Characterizing.

Companion training-side paper (pretraining from scratch with quantized softmax — why the backward rule and calibration gradients matter):https://arxiv.org/abs/2609.33591

Kernel patch and code will be released. Questions and feedback welcome!

Similar Articles

MiniMax Sparse Attention

Hugging Face Daily Papers

MiniMax Sparse Attention introduces a blockwise sparse attention mechanism that achieves significant speedups for ultra-long-context LLMs, reducing per-token attention compute by 28.4x at 1M context with wall-clock speedups of 14.2x for prefill and 7.6x for decoding on H800 GPUs. The method is accompanied by an open-source inference kernel and a publicly released multimodal model.