Numeric tweaks to fp8

[drisspg]

Numeric tweaks to fp8 rescale and offset

Addressing: #2577

This however has some tradeoffs ill show. And this is choice is not yet finalized;

Summary

The summary in that issue is pretty much spot on and I agree with its fix although there is some interesting perf/ numeric trade offs, and fundamentally the optimal choice is data dependent so we might wanna think about exposing this somehow.

TLDR: With offset=8 and rescale_threshold=4, the max intermediate value in the first P block is 2^8 = 256. A later block can have an unshifted max up to 3.999 log2 units above the tracked max. Because this increase is below the threshold, we keep the old M and skip rescaling the accumulator. The second block's max intermediate P value can then approach 2^(8+4) = 4096.

This exceeds E4M3's max value of 448. The denominator l_i is accumulated in FP32, but P is cast to E4M3 before P @ V, causing these large values to saturate. The denominator includes the original values while the numerator includes only their clipped contributions. Final normalization therefore can destroy these values that should contribute to the output

lets say we set rescale threshold to zero, then any increase in the block max also updates M. The previous accumulator and l_i are rescaled by alpha = 2^(M_old - M_new), and the new block is evaluated relative to M_new. Its max intermediate P value is therefore 2^8 = 256, rather than 4096.

The current Pr uses a threshold of 0.75; it allows the maximum to reach 2^8.75 ≈ 431, which remains below 448, while still avoiding some accumulator rescales, but there are many ways to accomplish the same affect

Perf Numeric Summary

-- agent notes --

Pareto frontier

The full confirmation supports max_offset=8, rescale_threshold=0 as the best global policy measured. Relative to upstream (4,4):

Path	Median latency	Time-weighted latency	Median relative-L2 / upstream
Production decode, auto-SplitKV	+0.11%	+0.01%	0.9945
Decode diagnostic, forced one split	+0.94%	+0.94%	0.8917
Paged causal prefill	-2.70%	-3.80%	0.9382

Lower is better. Threshold zero produced the stable prefill improvement. Offsets 6–8 had no reproducible performance ordering in the focused recheck, so offset 8 is preferred for its wider retained probability-tail range and slightly better aggregate prefill numerics. The tradeoff is roughly 1% slower forced-one-split decode; production auto-SplitKV decode remained neutral.

Implementation note: the current diff uses (8, 0) for FP8 inputs, matching the measured global policy.

Why (8, 0) vs (8, 0.75)

0.75 does genuinely help decode: it skips accumulator corrections when the running maximum increases by less than the threshold. Relative to (8,0), (8,0.75) was 0.12% faster time-weighted across the broad 72-case auto-SplitKV decode set and 0.61% faster across the 72 forced-one-split cases. In a focused two-case long-context recheck, the gains were larger—1.67% and 4.26%, respectively—showing that correction cost can be exposed in long unsplit decode.

Threshold zero, however, is a distinct compile-time specialization that removes the positive-threshold compare/select sequence. A codegen isolation test compared 0 with 1e-6: their BF16 outputs matched bit-for-bit in 8/8 cases, but 1e-6 was 7.66% slower time-weighted for prefill. 0.75 recovered some work by skipping real corrections, but remained 6.24% slower than zero for prefill. It also had slightly worse aggregate numerics than zero.

I am going to stick with (8, 0) for now but int eh future this should maybe be user configurable ..

Methodology

• Discovery: 38 distinct offset/threshold pairs, 54 matched cases each — 2,052 cells.
• Confirmation: 9 selected pairs, 216 matched cases each — 1,944 cells.
• Focused recheck: 5 finalists across 8 long-context cases with rotated policy order — 40 cells.
• Total: 4,036 matched policy/workload cells.
• Captured Qwen2.5-1.5B and Qwen2.5-7B attention tensors using C4 and instruction/GSM8K text, layers 0/16/27, and realistic MHA/GQA geometry.
• Paged causal decode used Q=1, SK={129,512,2048,8192}, both auto-SplitKV and forced one split. Paged causal prefill used Q=K, S={512,2048,4096,8192}.
• NVIDIA B200 with SM clock locked to 1500 MHz. Warm fixed-pointer CUDA-graph timing; quantization/setup excluded.
• Numerics compare FA4 against FP32 attention over the exact dequantized vLLM-style E4M3 inputs. This measures isolated FA4 attention, not end-to-end vLLM serving.

Full 2D discovery heatmaps

Green is better than upstream (4,4). Unsafe pairs are retained as diagnostic controls and show the expected rapid saturation.

E5M2 follow-up

I repeated the performance/numerics sweep with the captured inputs quantized to E5M2. This supports retaining max_offset=8 and changing E5M2 from rescale_threshold=4 to rescale_threshold=0. Relative to inherited (8,4):

Path	Median latency	Time-weighted latency	Median relative-L2 / (8,4)
Production decode, auto-SplitKV	-0.52%	-0.55%	1.0000
Decode diagnostic, forced one split	+0.08%	+0.11%	0.9191
Paged causal prefill	-3.49%	-4.17%	0.9270

Threshold zero again removes the positive-threshold compare/select path and improves prefill while remaining slightly faster for production auto-SplitKV decode. A focused 8K forced-one-split decode recheck exposed the path-specific downside more clearly: (8,0) was about 3.1% slower than (8,4) there.

E5M2's extra exponent range is still being used at offset 8: its effective probability floor is 2^-24, versus 2^-17 for E4M3 at the same offset. Threshold-zero offsets 8, 12, and 15 produced essentially identical numerics, so spending more saturation headroom on a larger offset did not help. The grid did not test offsets below 8; this supports retaining the existing offset, not claiming that 8 is the global E5M2 optimum for every context length.

E5M2 methodology

• Discovery: 40 pairs (34 safe candidates plus 6 saturation controls) × 54 cases — 2,160 cells.
• Confirmation: 8 selected pairs × 216 cases — 1,728 cells.
• Focused recheck: 6 policies × 8 long-context cases — 48 cells.
• Total: 3,936 matched policy/workload cells.
• Same Qwen2.5-1.5B/Qwen2.5-7B captures, prefill/decode shapes, B200 timing contract, and FP32-reference methodology as the E4M3 study; only the input quantization format changed to E5M2.

Full E5M2 2D discovery heatmaps

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 16fe0cd0db

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[MatthewBonanni]

0.75 is what I used here and it works well: vllm-project#166

[drisspg]

@MatthewBonanni hey thanks for pinging, I did a bunch of tests and came to a slightly different conclusion.. although I think we can re-assess. I am working on a fwd config like system and honestly this feels like it should be user overridable in the limit

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 16fe0cd0db

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".