Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

Shangzhen Zhu; Muyan Hu; Tomasz Kozlowski

cs.LGSep 27, 2026

Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration

Authors: Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski

Organizations: University of Illinois Urbana-Champaign

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.

Figures & tables

Degree of freedom	Intervention (others fixed)	Endpoint	Where
Support	top- $k$ , mean threshold	$\Delta\mathrm{NLL}$	§ 4.1 , Fig. 1 A
Within-support weighting	uniform vs. softmax, same support	$D_{W}(\ell)$	§ 4.1 , Fig. 1 B
Intervals $K$ ; weight map	full-range grid; exp. vs. linear map	$\Delta\mathrm{NLL}$	§ 4.1 , Fig. 1 C, Fig. S4
Allocation $R$	$K=21$ , $R\in[0.25,8]$	$D_{R}(21)$ , $I_{R}$	§ 4.2 , Fig. 2
Reconstruction	upper edge, nearest, interpolation	$\Delta\mathrm{NLL}$	§ 4.4
Exponential evaluation	Rowmax-PoT; Rowmax-H15 in FA4	$\Delta\mathrm{NLL}$ , latency	§ 4.4 , § 5

Table 1: Intervention dimensions. Each experiment varies the listed property and fixes the others.

Figure 1: Attention support, weighting, and a 32-interval grid. (A) Support sweep in Qwen2.5-1.5B with softmax (blue) or uniform (orange) weights; both axes are logarithmic. Diamonds mark the rule that retains scores above the row mean, plotted at each run’s own kept fraction (0.493 for softmax; 0.526 for uniform). The interventions run in every layer and change later scores, so panel A does not hold support fixed between weighting rules. (B) Separate single-layer control in Qwen2.5-1.5B with identical mean-threshold supports; the y-axis is logarithmic.

D_{W}(\ell)

is NLL with uniform weights minus NLL with softmax weights at target layer

\ell

, so positive values mean uniform weighting increases NLL. All 28 pointwise 95% paired-bootstrap intervals exclude zero. (C) Perplexity increase across ten models using 32 full-range intervals, narrower near the row maximum (

R=4

). Whiskers are 95% paired-bootstrap intervals; the dashed line marks a 1% increase.

Figure 2: Allocation of a fixed number of intervals.

R>1

makes intervals narrower near the row maximum;

R=1

gives uniform widths. (A)

\Delta\mathrm{NLL}

in Qwen2.5-1.5B as

R

varies at

K=21

; whiskers are 95% paired-bootstrap intervals. (B)

\Delta\mathrm{NLL}

and RMS log-weight errors relative to their values at

R=1

. The errors use scores from the softmax baseline;

\Delta\mathrm{NLL}

evaluates the approximation in every layer. From

R=1

R=4

, unweighted error rises while probability-weighted error and

\Delta\mathrm{NLL}

fall. (C) Across ten models,

D_{R}(21)=\Delta\mathrm{NLL}_{R=1}-\Delta\mathrm{NLL}_{R=4}

. Positive values favor finer intervals near the row maximum; filled markers have 95% intervals excluding zero.

Figure 3: Attention divergence and model loss. (A) NLL increase against mean row-wise attention JSD for 56 interventions on Qwen2.5-1.5B; the line is an in-sample log–log fit. (B) Flattening-minus-sharpening NLL contrast

D_{T}

after the two perturbations are calibrated within each model to the same target mean JSD. Positive values mean flattening increases NLL more than sharpening.

J_{0}

is each model’s JSD under the octave PoT perturbation; circles target

J_{0}

and squares target

10J_{0}

. Horizontal bars are 95% paired-bootstrap intervals; hollow markers have intervals including zero. The signed horizontal axis uses a symlog scale.

Figure 4: Lattice controls, Rowmax-H15 fidelity, and B200 performance. (A) Qwen2.5-1.5B model evaluations, not B200 kernel measurements. Orange compares absolute and rowmax anchors at octave spacing; green compares octave and half-octave spacing at the absolute anchor. Arrow labels give the resulting decreases in

\Delta\mathrm{NLL}

; whiskers are 95% paired-bootstrap intervals. (B) Perplexity increase over stock FA4 at 2K. Filled points are direct BF16 kernel measurements in five models; hollow points to the right of the divider are BF16 and FP8 semantic simulations on Qwen2.5-1.5B. Model prefixes Q, M, and L denote Qwen2.5, Mistral-7B-v0.3, and Llama-3.1, respectively. Intervals are in Table 2 ; FP8 kernel fidelity was not measured directly. (C) Attention-forward speed-up over stock FA4 at 8K by dtype and mask (causal or non-causal). Bars show speed-ups computed from pooled median call latencies; points show measurement rounds and whiskers their standard deviation, not a confidence interval. (D) Separate FP8 causal 8K sweep against FA4’s shipping default. Green crosses vary the share

\beta

of polynomially emulated elements and emulation period

f

; the +2.7% label marks the fastest tested stock setting. The orange star is Rowmax-H15, which has no

\beta

. The horizontal dotted line marks zero speed-up; the vertical dashed line separates Rowmax-H15 from the stock-setting sweep.

Fidelity path	Model	$\Delta$ NLL [95% CI]	PPL change (%)
BF16 kernel, measured	Qwen2.5-1.5B	0.00121 [0.000862, 0.00156]	0.121
BF16 kernel, measured	Qwen2.5-72B	0.00101 [0.000665, 0.00139]	0.101
BF16 kernel, measured	Mistral-7B-v0.3	0.000914 [0.000600, 0.00125]	0.091
BF16 kernel, measured	Llama-3.1-8B	0.00246 [0.00206, 0.00289]	0.246
BF16 kernel, measured	Llama-3.1-70B	0.00491 [0.00402, 0.00582]	0.492
BF16 semantic simulation	Qwen2.5-1.5B	0.00107 [0.000680, 0.00145]	0.107

Table 2: Rowmax-H15 fidelity, B200 call latency and energy. Top: 97 blocks, 95% paired-bootstrap intervals; measured rows evaluate the patched kernel directly on the BF16 path; simulation rows are Qwen2.5-1.5B only, and the FP8 row is a semantic simulation, not a direct FP8 kernel measurement. Bottom: pooled medians of three rounds; energy from one sustained loop per kernel.

Model	Params (B)	$K{=}32,R{=}4$ PPL inc. (%)	$D_{R}(21)$ $[95\%\ \mathrm{CI}]$	Rowmax-PoT $K_{\max}{=}20$ PPL inc. (%)	Rowmax-PoT $K_{\max}{=}20$ top-1 agree. (%)
Qwen2.5-0.5B	0.49	0.134	0.00611 [0.00525, 0.00697]	0.382	95.8
Qwen2.5-1.5B	1.54	0.091	0.00351 [0.00282, 0.00420]	0.310	96.2
Qwen2.5-3B	3.09	0.082	0.00331 [0.00256, 0.00404]	0.289	96.4
Qwen2.5-72B	72.71	0.055	0.00196 [0.00137, 0.00258]	0.252	97.6
Llama-3.2-1B	1.24	0.223	0.0109 [0.00976, 0.0120]	0.950	94.4
Llama-3.2-3B	3.21	0.121	0.00629 [0.00540, 0.00714]	0.704	95.2

Table 3: Cross-model summary. Perplexity increases are point estimates;

D_{R}(21)

has model-specific paired-bootstrap intervals; top-1 is the vocabulary-output agreement with softmax under Rowmax-PoT. No interval is pooled across models.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

change	reduction in $\Delta$ NLL [95% CI]
anchor: absolute $\rightarrow$ rowmax, octave spacing	0.00355 [0.00268, 0.00448]
anchor: absolute $\rightarrow$ rowmax, half-octave spacing	0.00122 [0.000708, 0.00171]
spacing: octave $\rightarrow$ half-octave, absolute anchor	0.00443 [0.00367, 0.00517]
spacing: octave $\rightarrow$ half-octave, rowmax anchor	0.00210 [0.00132, 0.00286]

Appendix

Table 8

Figure S1: Correctness gates for the patched B200 kernel. Element-wise readbacks measured on B200. (A) Single-tile gate: residual of the kernel probabilities against the Rowmax-H15 reference over 8,256 valid elements of 128 rows. (B) Multi-tile gate over 131,328 valid elements: filled bars, the measured normalized ratio $\log_{2}(P_{1}/P_{0})$ against the grey $\pm 0.585$ -octave bound it must satisfy; outlined bars, the recovered unnormalized error $\log_{2}(\hat{w}/w)$ against the yellow analytic envelope $[-0.25,0.335]$ ; green ticks, the per-row normalization shift read at the row maximum. Values in Appendix B.2 .

Measurement (semantics)	$\Delta$ NLL [95% CI]	PPL change	Evidence status
B200 patched FlashAttention-4 kernel (A1 $-$ A0) (BF16 path, rescale_threshold = 8.0 )	+0.00121 [0.000862, 0.00156]	+0.121%	direct kernel measurement
Semantic simulation of the BF16 kernel path (BF16 semantics, threshold 8.0)	+0.00107 [+0.000680, +0.00145]	+0.107%	cross-check of the simulator against the measured kernel
Semantic simulation of the FP8 kernel path (FP8 semantics, per-tile running max (tile 128), threshold 0)	+0.000716 [+0.000370, +0.00106]	+0.072%	semantic simulation — not a direct FP8 kernel measurement
Offline whole-row Rowmax-H15 (control) (rowmax anchor, no tiling)	+0.00105 [+0.000712, +0.00139]	+0.105%	semantic control
Absolute-anchored $\{1,1.5\}\times 2^{k}$ lattice (control; condition ilpot_h15_v1 ) (absolute anchor)	+0.00224 [+0.00179, +0.00266]	+0.224%	semantic control (phase-misaligned)
PoT m = 2 (half-octave reference) (rowmax anchor)	+0.000927 [+0.000543, +0.00132]	+0.093%	reference

Appendix

Table 10

Contrast	$\Delta$ ( $\Delta$ NLL) [95% CI]	Reading
FP8 semantic Rowmax-H15 $-$ PoT m = 2	-0.000211 [-0.000612, +0.000183]	No resolved difference at the reported precision; the paired 95% interval includes zero.
FP8 semantic Rowmax-H15 $-$ ILPoT-H-A	-0.000143 [-0.000590, +0.000304]	No resolved difference at the reported precision; the paired 95% interval includes zero.
FP8 semantic Rowmax-H15 $-$ offline whole-row Rowmax-H15	-0.000333 [-0.000648, -0.0000278]	per-tile running-max semantics differ from whole-row rowmax
BF16 semantic Rowmax-H15 $-$ FP8 semantic Rowmax-H15	+0.000358 [-0.00000769, +0.000736]	Higher point estimate under BF16 conditional-rescale semantics; the paired 95% interval includes zero.
Simulation vs measured kernel (BF16 path)	measured +0.00121 vs simulated +0.00107	relative gap 11%; intervals overlap

Appendix

Table 11

Offline operator	$\Delta$ NLL vs softmax [95% CI]	PPL change	Paired $\Delta$ NLL vs Rowmax-H15 [95% CI]
Rowmax-H15, $H(x)=S(Q_{1/2}(x))$	+0.00105 [+0.000712, +0.00139]	+0.105%	reference
Rowmax-PoT $m=2$ (Eq. ( 18 ))	+0.000927 [+0.000543, +0.00132]	+0.093%	$-0.000122$ [ $-0.000551$ , +0.000301]
Unshifted Schraudolph $S(x)$	$-0.0000231$ [ $-0.000306$ , +0.000270]	$-0.00231\%$	$-0.00107$ [ $-0.00140$ , $-0.000747$ ]
$S(Q_{2^{-2}}(x))$ , $k=2$	+0.000329 [+0.0000270, +0.000624]	+0.033%	$-0.000721$ [ $-0.00106$ , $-0.000393$ ]
$S(Q_{2^{-3}}(x))$ , $k=3$	+0.0000796 [ $-0.000156$ , +0.000316]	+0.00796%	$-0.000970$ [ $-0.00129$ , $-0.000649$ ]

Appendix

Table 12

Run group	Measurements	GPU used	Driver	PyTorch / CUDA	Clock and power record
Characterization, core	Section 4 interventions on the core models; anchor controls and semantic simulations (Qwen2.5-1.5B)	one GPU, model not recorded	not recorded	2.5.1+cu121 / 12.1	not recorded
Characterization, large	Section 4 interventions on Llama-3.1-70B and Qwen2.5-72B	three A100-SXM4-80GB	580.173.02	2.5.1+cu121 / 12.1	not recorded
B200 S1	patch and gates; Qwen2.5-1.5B BF16 fidelity; main matrix; sequence sweep and 1K/2K retest; FP8 non-causal control; head, batch and head-dimension-64 controls; five tuning settings and Rowmax-H15; energy loop; cuDNN and SDPA references; ncu refused	one B200	580.126.09	2.14.0+cu130 / 13.0	lock not permitted; clock, temperature and power after each window; energy loop sampled
B200 S2	gates re-run with element dumps; GQA (three rounds); ncu refused	one B200	580.126.09	2.14.0+cu130 / 13.0	after each window (GQA at 1965 MHz throughout)
B200 S3	Nsight Systems timeline (kernel duration, call latency, launch cost); ncu refused	one B200	595.91.07	2.14.0+cu130 / 13.0 a	none
B200 S4	FP16 causal and FP16, BF16 non-causal (three rounds); three tuning settings	one B200	580.126.09	2.14.0 / 13.0 b	after each window; 1485–1965 MHz and 258–546 W in the three-round cells

Appendix

Table 13

Figure S2: Sequence length: the sweep and the traced run, kept apart. FP8 causal attention-forward speed-up of A1 over A0 (head dimension 128). (A) First-session sweep in call latency: pooled median over three alternating rounds (line), the rounds (small dots, sd annotated at 4K and 8K) and the separate higher-precision 1K/2K retest (triangles; its

-2.2\%

and

-3.7\%

are the values quoted in the text, the pooled sweep gives

-3.6\%

and

-2.6\%

). (B) Third-session traced run: host-timed call latency and kernel-only speed-up from the same run. The panels are different sessions and are not compared across panels.

seq	mask	kernel ( $\mu$ s) A0 / A1		call latency ( $\mu$ s) A0 / A1		kernel-only	call latency	non-kernel % (A0 / A1)
1K	causal	17.840	16.672	45.611	42.529	+7.0%	+7.2%	60.9 / 60.8
2K	causal	30.688	28.000	42.916	45.027	+9.6%	-4.7%	28.5 / 37.8
4K	causal	57.280	52.256	58.495	52.803	+9.6%	+10.8%	2.1 / 1.0
8K	causal	200.640	178.368	201.617	179.274	+12.5%	+12.5%	0.5 / 0.5
16K	causal	725.920	653.296	726.815	655.539	+11.1%	+10.9%	0.1 / 0.3
32K	causal	2747.570	2508.674	2799.928	2548.133	+9.5%	+9.9%	1.9 / 1.5

Appendix

Table 15

dtype	seq	A0 ms (pooled median)	A1 ms	speed-up (pooled)	per-round sd	rounds
FP8 E4M3	4K	0.0621	0.0580	+7.1%	2.13	3
FP8 E4M3	8K	0.2063	0.1835	+12.4%	0.02	3
FP8 E4M3	16K	0.7320	0.6509	+12.5%	0.08	3
BF16	4K	0.0586	0.0564	+3.8%	0.93	3
BF16	8K	0.2020	0.1917	+5.4%	0.17	3
BF16	16K	0.7871	0.7350	+7.1%	0.52	3

Appendix

Table 16

seq	pooled speed-up	per-round min … max	sd	1K/2K retest (n = 9)
1K	-3.6%	-5.5% … +1.1%	2.97	-2.2%
2K	-2.6%	-6.0% … +4.5%	4.53	-3.7%
4K	+7.9%	+7.7% … +8.3%	0.25	—
8K	+12.1%	+12.0% … +12.1%	0.06	—
16K	+12.5%	+12.1% … +12.7%	0.28	—
32K	+9.9%	+9.9% … +10.1%	0.08	—

Appendix

Table 17

dtype	mask	seq	speed-up (pooled)	rounds
FP8	non-causal	8K	+25.8% (sd 0.03)	3
BF16	non-causal	8K	+13.0% (sd 0.50)	3
FP16	non-causal	8K	+11.8% (sd 0.70)	3
FP8	non-causal	4K	+27.2%	1
FP8	non-causal	16K	+25.4%	1
BF16	non-causal	4K	+11.4%	1

Appendix

Table 18

A0 configuration	emulation fraction $\beta$	vs stock
freq 0 (hardware MUFU only)	0%	+1.4%
freq 16, start 1	12.5%	+1.6%
freq 8, start 1 (stock)	25%	+0.0%
freq 8, start 0	37.5%	+2.7%
freq 2, start 1	50%	+0.2%
freq 4, start 1	50%	-1.2%

Appendix

Table 19

configuration	speed-up
8K, heads 4	+12.6%
8K, heads 8	+12.3%
8K, heads 32	+12.1%
8K, batch 4	+11.7%
8K, GQA (32 query / 8 KV heads): 0.38885 ms $\rightarrow$ 0.34583 ms	+12.4% (per-round sd 0.03)
head dim 64, FP8: 4K / 8K / 16K	+2.0% / +2.5% / +3.7%

Appendix

Table 20

kernel	forwards	mean power (W)	time per forward (ms)	energy per forward (mJ)
A0 stock FA4	59550	982.3	0.7559	742.5
A1 FA4 + Rowmax-H15	65050	982.9	0.6922	680.3

Appendix

Table 21

seq	torch SDPA flash	torch SDPA mem-efficient	cuDNN 9.24	FA4 stock (A0)	FA4 + Rowmax-H15 (A1)	Rowmax-H15 vs cuDNN	FA4 vs cuDNN
4K	0.2876	0.5596	0.0536	0.0586	0.0564	-5.1%	-8.6%
8K	0.9117	1.9075	0.2230	0.2020	0.1917	+16.4%	+10.4%
16K	3.1568	6.9877	0.8316	0.7871	0.7350	+13.1%	+5.6%

Appendix

Table 22

Figure S3: Long-context within-context robustness. (A) Allocation contrast

D_{R}

and (B) matched-JSD contrast

D_{T}

, re-matched within each context length, with 95% paired block-bootstrap intervals; filled if the interval excludes zero, hollow if it includes zero. Each context has its own baseline,

J_{0}

and temperatures; no cross-context comparison is made.

Model	Context	Blocks	Predictions	$J_{0}$	$A^{-}$ achieved (rel. err. %)	$A^{+}$ achieved (rel. err. %)
Qwen2.5-72B	2K	97	198,559	0.001228	0.001229 (0.084)	0.001235 (0.591)
Qwen2.5-72B	8K	24	196,584	0.001323	0.001315 (0.614)	0.001333 (0.725)
Qwen2.5-72B	16K	12	196,596	0.001362	0.001371 (0.629)	0.001353 (0.666)
Llama-3.1-70B	2K	97	198,559	0.001778	0.001787 (0.510)	0.001785 (0.393)
Llama-3.1-70B	8K	24	196,584	0.001137	0.001129 (0.687)	0.001138 (0.100)
Llama-3.1-70B	16K	12	196,596	0.001236	0.001244 (0.621)	0.001237 (0.062)

Appendix

Table C1: Per-context matched-JSD reproducibility for the two approximately 70B models. Each context has its own softmax baseline, reference level

J_{0}

(PoT

m=1

), and independently resolved flattening (

A^{-}

) and sharpening (

A^{+}

) temperatures. Parentheses give relative target error in percent; all twelve matches are within 1%. Contexts are not compared by a cross-context significance test.

Qwen2.5-72B
Condition	LAMBADA	HellaSwag	ARC-E	ARC-C	PIQA	WinoGrande	LAMBADA PPL (%)
Softmax	77.72 (+0.00)	86.10 (+0.00)	83.38 (+0.00)	62.71 (+0.00)	83.62 (+0.00)	77.35 (+0.00)	—
Mean-threshold support	76.73 ( $-$ 0.99)	85.72 ( $-$ 0.38)	84.81 (+1.43)	62.71 (+0.00)	81.83 ( $-$ 1.80)	76.16 ( $-$ 1.18)	+6.727
Exp. $K=32,R=4$	77.59 ( $-$ 0.14)	86.13 (+0.03)	83.25 ( $-$ 0.13)	62.54 ( $-$ 0.17)	83.62 (+0.00)	77.35 (+0.00)	+0.037
Exp. $K=64,R=4$	77.64 ( $-$ 0.08)	86.05 ( $-$ 0.05)	83.21 ( $-$ 0.17)	62.63 ( $-$ 0.09)	83.68 (+0.05)	78.30 (+0.95)	+0.033
Rowmax-PoT $K_{\max}=20$	77.53 ( $-$ 0.19)	86.06 ( $-$ 0.04)	83.46 (+0.08)	62.20 ( $-$ 0.51)	83.73 (+0.11)	77.74 (+0.39)	+0.231

Appendix

Table C2: Zero-shot scores under attention approximations. Each cell is the score in percent followed by its descriptive change from the same-model softmax baseline in percentage points. Accuracy for LAMBADA and WinoGrande; length-normalized accuracy ( acc_norm ) for HellaSwag, ARC-Easy, ARC-Challenge and PIQA. The last column is the LAMBADA perplexity change (%) relative to softmax, reported for the two larger models. No pooled score or paired cross-condition interval is defined.

Figure S4: Weighting geometry under full support.

\Delta\mathrm{NLL}

of the exponential and the tested linear level-to-weight maps against

K

(full support,

R=1

, upper-edge reconstruction): Qwen2.5-1.5B (lines) and the other models with this task at

K=16

(crosses); 95% paired-bootstrap intervals.

model	params (B)	$D_{R}$ (21) [95% CI]	resolved?	$D_{R}$ (16)	$D_{R}$ (32)	$I_{R}$ [95% CI]
Qwen2.5-0.5B	0.49	+0.00611 [+0.00525, +0.00697]	yes	+0.0105	+0.00192	+0.00855 [+0.00712, +0.0101]
Llama-3.2-1B	1.24	+0.0109 [+0.00976, +0.0120]	yes	—	—	—
Qwen2.5-1.5B	1.54	+0.00351 [+0.00282, +0.00420]	yes	+0.00612	+0.00133	+0.00480 [+0.00356, +0.00600]
Gemma-2-2B	2.61	+0.00527 [-0.00369, +0.0132]	no (interval crosses 0)	—	—	—
Qwen2.5-3B	3.09	+0.00331 [+0.00256, +0.00404]	yes	+0.00532	+0.00113	+0.00419 [+0.00305, +0.00535]
Llama-3.2-3B	3.21	+0.00629 [+0.00540, +0.00714]	yes	—	—	—

Appendix

Table 27

model	$D_{T}$ (J 0 ) [95% CI]	resolved	$D_{T}$ (10J 0 ) [95% CI]	resolved
Qwen2.5-0.5B	+0.00472 [+0.00331, +0.00607]	+	+0.0300 [+0.0258, +0.0341]	+
Llama-3.2-1B	-0.00313 [-0.00428, -0.00202]	$-$ (resolved negative)	-0.000204 [-0.00339, +0.00285]	unresolved
Qwen2.5-1.5B	+0.00170 [+0.000676, +0.00269]	+	+0.0133 [+0.0104, +0.0161]	+
Gemma-2-2B	+0.246 [+0.218, +0.273]	+	+0.806 [+0.767, +0.844]	+
Qwen2.5-3B	+0.00409 [+0.00308, +0.00512]	+	+0.0195 [+0.0164, +0.0227]	+
Llama-3.2-3B	+0.000976 [+0.000138, +0.00185]	+	+0.0107 [+0.00813, +0.0135]	+

Appendix

Table 28

Figure S5: Local attention-output magnitude as a negative control. Per-layer (A) RMS change and (B) cosine of the local output

pV

for matched flattening (

A^{-}

) and sharpening (

A^{+}

) at

J_{0}

(solid) and

10J_{0}

(dashed), Qwen2.5-1.5B; descriptive aggregates without bootstrap intervals.

Figure S6: Similar probability redistribution, model-dependent losses. (A) Signed leakage

L_{0.9}

of matched flattening (filled) and sharpening (hatched) at

J_{0}

. (B) Local

pV

RMS change for the same conditions, with the sign of the resolved

D_{T}

above each pair. (C) Original

D_{T}(J_{0})

(circles) and the contrast after two-region mass restoration (hollow squares), 95% paired intervals, symlog axis;

G_{\mathrm{gap}}

is their signed difference. (D)

L_{0.9}

K=21

for

R=1

(dotted) and

R=4

(filled). Panels A, B and D are descriptive, without intervals.

Model	$D_{\mathrm{rescue}}(10J_{0})$ [95% CI]	$G_{\mathrm{gap}}(10J_{0})$ [95% CI]
Qwen2.5-0.5B	0.00249 [ $-$ 0.000190, 0.00518]	0.0275 [0.0252, 0.0297]
Llama-3.2-1B	$-$ 0.00862 [ $-$ 0.0112, $-$ 0.00604]	0.00842 [0.00657, 0.0103]
Qwen2.5-1.5B	$-$ 0.00278 [ $-$ 0.00510, $-$ 0.000565]	0.0161 [0.0144, 0.0178]
Gemma-2-2B	$-$ 0.0512 [ $-$ 0.0854, $-$ 0.0175]	0.857 [0.825, 0.889]
Qwen2.5-3B	0.00204 [ $-$ 0.0000977, 0.00417]	0.0175 [0.0156, 0.0194]
Llama-3.2-3B	0.00371 [0.00142, 0.00607]	0.00700 [0.00491, 0.00916]

Appendix

Table D3: Primary endpoint

G_{\mathrm{gap}}(10J_{0})

. Rescued contrast

D_{\mathrm{rescue}}(10J_{0})

and the signed change

G_{\mathrm{gap}}(10J_{0})=D_{T}(10J_{0})-D_{\mathrm{rescue}}(10J_{0})

(original

D_{T}(10J_{0})

in Table D.3 ); 97 blocks, paired block bootstrap of the per-block difference, 5,000 replicates, seed 0. A positive

G_{\mathrm{gap}}

is not necessarily a smaller

|D|

Figure S7: Rowmax-PoT across models and tail cutoffs. (A) Ten-model perplexity increase of the Rowmax-PoT

K_{\max}=20

clamp condition with 95% intervals and the 1% line; Llama-3.1-70B is the 1.513% exception. (B) Qwen2.5-1.5B

K_{\max}

sweep, clamp and drop tail modes. Supporting characterization, not the Rowmax-H15 kernel result.

Method	Object	Anchor	Calibration / training	Evaluation
FQ-ViT	post-softmax probabilities, 4-bit $\log_{2}$	absolute ( $p=1$ )	none	ViT accuracy; no fused kernel
RepQ-ViT	post-softmax probabilities, $\log_{\sqrt{2}}$ ( $\log_{2}$ at inference)	absolute, layer-wise scale	percentile calibration	ViT accuracy; no fused kernel
ITA	integer streaming softmax, power of two of the floor-rounded distance	row maximum (streaming)	clipping threshold by QAT	22 nm ASIC
EXAQ	coarse exponential	row maximum	clipping range from activation statistics	frozen LLMs; isolated softmax timing
IntAttention	32-entry exponential lookup table	row maximum	clipping range fixed offline	frozen LLMs; integer attention on Armv8 CPUs
FA4	cubic exp2 emulation for part of each row	running maximum (exact exponential)	none; pointwise error	fused kernel on B200

Appendix

Table E1: Closest methods. What each approximates, where its lattice or reference is anchored, what is calibrated or trained, and where it is evaluated.

Explore similar work

Sep 9, 2026cs.LG

EFQ-Softmax: Exp-Free Quantization for Softmax

Low-bit attention accelerates Transformer inference by moving the

QK^\top

and

PV

matrix multiplications to FP8 or FP4 matrix engines. However, the softmax path often evaluates shifted-score exponentials in higher precision, forms a temporary probability block, and quantizes it before low-bit

PV

multiplication. This exp-then-quantize path creates a mismatch between a high-precision probability producer and a low-bit matrix consumer. We propose EFQ-Softmax (Exp-Free Quantization for Softmax), a low-bit probability-generation method that directly maps shifted attention scores to block-scaled E2M1 operands. For each microscaling block, EFQ-Softmax selects an exponent-only scale from the local maximum, maps the shifted scores to a normalized residual domain, and generates nonnegative E2M1 probability codes using a single affine rule. The resulting operand is used consistently in both the

\widetilde{P}V

numerator update and the

\widetilde{P}\mathbf{1}

denominator update. The FlashAttention-style row-maximum update, historical rescaling, high-precision accumulation, and final normalization remain unchanged. We evaluate end-to-end quality on Qwen3-8B, Qwen3-VL-8B-Instruct, and WAN2.2-TI2V-5B, and separately measure kernel-level performance on the A5 vector unit. EFQ-Softmax improves the Qwen3-8B seven-task mean from 0.6749 with MXFP4 to 0.6773 and the Qwen3-VL nine-task mean from 0.7826 to 0.8000. On WAN2.2, it maintains temporal consistency and visual quality comparable to the FP16 and MXFP4 baselines under VBench. On the A5 vector unit, EFQ-Softmax reduces the vector-stage latency of the fused probability-generation kernel by 40.33% on average across sequence lengths from 16K to 128K. These results show that direct low-bit probability generation can replace the conventional exp-then-quantize path while preserving end-to-end model quality.

Haohui Han, Yuming Wan, Hongni Wang +4

Xi'an Jiaotong University · Xi’an Jiaotong University · Huawei Technologies Co., Ltd +1

Gumbel-Softmax Relaxation Block Sparse Flash Attention Matrix Multiplication Transformer Architectures Activation Functions

Ask PDF

Sep 27, 2026cs.LG

Pretraining Transformers with Quantized Softmax in Attention

Low-precision Transformer systems increasingly quantize attention matrix multiplications, while softmax often remains at higher precision. During pretraining, an approximate softmax changes the gradients that train the model as well as its forward computation. We study this interaction with K-interval attention, which approximates the exponential using K+1 grid values. We vary per-row grid calibration, interpolation versus hard rounding, and the placement of a straight-through surrogate relative to normalization. We derive the corresponding backward rules, including calibration derivatives, and compare these choices in pretraining experiments matched on model, data, and optimizer. Detaching the row extrema leaves the forward computation unchanged but produces a delayed increase in validation loss. With hard rounding at K=4, min-max calibration and a pre-normalization surrogate incur a large loss gap; changing either choice substantially reduces it. At 124M parameters and 2.5B training tokens, fixed-window calibration with a post-normalization surrogate yields a validation loss gap of +0.019 nats relative to softmax at K=4, and with a pre-normalization surrogate yields +0.004 nats at K=16.

Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski

University of Illinois Urbana-Champaign

Transformer Attention Gumbel-Softmax Relaxation Generative Pretrained Transformers Transformer Architectures Matrix Multiplication Activation Functions

Ask PDF

Sep 21, 2026cs.LG

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention

Transformer models built on the attention mechanism have become a central building block in modern deep learning, yet softmax attention remains a major bottleneck for long-context workloads. While FlashAttention makes the forward and first backward passes I/O-efficient, it does not support backward-over-backward (BoB), which enables exact differentiation through the backward pass for applications such as second-order optimization, test-time training, gradient-based memory, and meta-learning. Existing BoB implementations either materialize large intermediate tensors or exhaust GPU memory at long sequence lengths. We present FlashBoB, an exact, I/O-efficient algorithm for BoB in softmax attention that keeps computation within on-chip tiles and avoids all

N \times N

intermediate tensors, where

N

is the sequence length. The key insight is a hierarchical affine structure in the softmax double backward: two row-wise scalars determine all outputs through affine transformations. This yields a two-pass schedule with bounded on-chip static random-access memory (SRAM) usage and minimal off-chip high-bandwidth memory (HBM) traffic. FlashBoB achieves

Θ(N^2 d^2/M)

HBM traffic (

d

is the head dimension and

M

is the memory size) and, within the standard FlashAttention-style score-recomputation model, matches the inherited large-cache lower bound for exact forward attention. Empirically, it scales exact attention BoB to

N=262\text{K}

on a single A100 80GB GPU, where prior PyTorch exact baselines fail by

N=16\text{K}

, and is up to

6.3\times

faster than FlashBack. These results make exact second-order attention practical at long-context sequence lengths where prior implementations cannot run efficiently.

Anthony Givans, Michael Crawshaw, Mingrui Liu

Block Sparse Flash Attention Gumbel-Softmax Relaxation Backward Passes Activation Functions Gradient-Based Optimization

Ask PDF