Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

On Blackwell B200, tensor cores outrun the exponential unit by ~500× (8192 vs 16 ops/clk/SM), so exp becomes exposed in fused attention. We ask what a frozen pretrained LLM actually needs from softmax — and use the answer to replace exp2 inside 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.
• Where resolution 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!