fix(fp8): round onto the e4m3 grid before the native float8e4nv downcast - #85
fix(fp8): round onto the e4m3 grid before the native float8e4nv downcast#85gdevenyi wants to merge 1 commit into
Conversation
The store kernel's int branch rounds half-away-from-zero before its cast because the float->int cast truncates. The fp8 branch went straight to .to(float8e4nv), which does not round to nearest either: triton lowers fp32 -> float8e4nv as a double-round (fp32 -> fp16 RTZ -> e4m3), so a value just above a grid midpoint collapses onto the midpoint and then ties to even, always downward. On sm_89 that made two of the PR's own tests fail: test_store_kernel_matches_the_reference_quantizer[256-fp8_e4m3] and [512-fp8_e4m3]. q8_0 was unaffected, which is the tell -- only the fp8 path takes that cast. round_e4m3 puts the value on the grid in fp32 first, after which the cast is exact. 53/53 of the PR's tests pass on sm_89 with this. Same root cause as FlashML-org#85. Deployment branch only.
triton lowers fp32 -> float8e4nv as a double-round (fp32 -> fp16 RTZ -> e4m3), so a value a hair above a grid midpoint collapses onto the midpoint and then ties to even. The error is one-sided: over a uniform [-448, 448] sweep of 2^22 values, 16127 (0.38%) disagree with torch's RNE and every one of them lands 1 ULP *toward zero*. No fp_downcast_rounding setting changes it. Every native-fp8 activation quantizer now rounds with round_e4m3 in fp32 first; the downcast is then exact and the lowering's rounding mode stops mattering. The emulated (pre-sm_89) path already did this, which is why the forced-EMU A/B in test_e4m3_compat caught it. - fp8_block_linear._act_quant_kernel - fp8_pertensor_linear._static_quant_kernel - dsv4/fp8_linear._act_quant_fp8_kernel - dsv4/fp8_linear._act_quant_inplace_kernel: the native branch's fp8 round-trip only ever re-quantized an already-grid value, so it collapses into the shared round_e4m3 call. test_forced_emu_matches_native's per-tensor bit-equality now holds for the activation quantizers. The three fp8-MMA-vs-bf16-MMA GEMM keys still differ by ~1 output ULP -- triton picks different MMA shapes for fp8xfp8 and bf16xbf16 tl.dot, so the fp32 reduction trees differ. That was bit-exact on H100 by luck; on sm_89 it is not. moe_prefill_fp8 chains two such GEMMs and then top-k-sums ~8e3-magnitude rows down to ~2e2, so an error worth 2e-3 of the GEMM scale reads as 13% elementwise. The bound is now on max|ref| rather than elementwise, which measures the GEMM's error instead of the cancellation's (worst observed 4.6e-3, limit 1e-2). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun (cherry picked from commit dcc2f60a6a9f3c9fb102c3466c24102745aad219)
298bf3c to
d3eb4d3
Compare
|
Rebased onto Worth noting because it looked superseded and is not: main does call Full 🤖 Generated with Claude Code |
There was a problem hiding this comment.
🟢 Approval recommended
The functional changes align with the stated bug and are backed by targeted regression coverage, with only minor documentation wording nits noted.
Pull request overview
This PR fixes a native sm_89+ Triton FP8 (e4m3) downcast accuracy bug by explicitly rounding FP32 values onto the e4m3 grid before converting to tl.float8e4nv, eliminating a one-sided double-rounding bias (fp32 -> fp16 RTZ -> e4m3) in native-fp8 activation quantization paths.
Changes:
- Pre-round to the e4m3 grid (in FP32) before
tl.float8e4nvdowncast in FP8 quantization kernels (block, per-tensor, and dsv4 paths). - Simplify dsv4 inplace quant path to avoid the native fp8 round-trip and rely on
round_e4m3for deterministic grid rounding. - Add/adjust kernel compatibility tests: pin the “pre-round makes native downcast exact” invariant and switch fp8-MMA vs bf16-MMA comparisons to a scale-based bound.
File summaries
| File | Description |
|---|---|
tests/kernels/test_e4m3_compat.py |
Adds a regression test that pins the pre-round invariant; updates fp8-MMA vs bf16-MMA comparison to a scale-based tolerance. |
python/freetoken/kernel/triton/fp8_pertensor_linear.py |
Rounds to the e4m3 grid before native float8e4nv storage in the static activation quant kernel. |
python/freetoken/kernel/triton/fp8_block_linear.py |
Rounds activations onto the e4m3 grid before the native downcast to avoid Triton’s biased double-rounding. |
python/freetoken/kernel/triton/dsv4/fp8_linear.py |
Applies the same pre-round in dsv4 activation quant and removes the native fp8 round-trip in the inplace kernel. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| # float8e4nv downcast double-rounds (fp32 -> fp16 RTZ -> e4m3), so a value | ||
| # just above a grid midpoint collapses onto the midpoint and then ties to | ||
| # even -- always downward. 0.38% of a uniform [-448,448] sweep lands 1 ULP | ||
| # low, never high, and no fp_downcast_rounding setting changes it. Once the | ||
| # value is already a grid point the downcast is exact and this cannot bite. |
|
|
||
| pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") | ||
|
|
||
| FP8 = torch.float8_e4m3fn | ||
| # fp8-MMA vs bf16-MMA GEMMs: identical grid values, but the MMA-internal fp32 | ||
| # reduction order may differ (dsv4_gemm / moe_prefill_fp8 happen to be bit-exact | ||
| # on H100 -- that is lowering luck, not a guarantee). | ||
| # fp8-MMA vs bf16-MMA GEMMs: identical grid values, but triton lowers fp8xfp8 and |
|
Re-tested against current main (
Both passes merge this PR's head onto main first. The CPU pass hides the GPUs on purpose — while the box is serving, 65-95 GPU tests fail on main itself with It also fixes a test that is currently failing on main: For reference, main's own GPU baseline in this window was 6 failures, so a nonzero count here is not automatically this PR's doing — the comparison above is against that set, not against zero. 🤖 Generated with Claude Code |
The bug
Triton lowers
fp32 -> tl.float8e4nvas a double rounding: fp32 to fp16 with truncation, then fp16 to e4m3. A value just above a grid midpoint collapses onto the midpoint in the first step, and round-half-to-even then sends it down. The error is one-sided.Measured on sm_89, 2^22 samples drawn uniformly from [-448, 448], against torch's RNE cast:
Setting
fp_downcast_rounding="rtne"does not change it.So every native-fp8 activation quantizer shrinks about 0.4% of its values by one unit in the last place, always toward zero. On the per-token-group quantizer that is a 0.02% bias in magnitude:
mean|dequant| / mean|x|moves from 0.998937 to 0.999169. The bias is small, but it is one-sided, and it costs nothing to remove.The emulated path for pre-sm_89 GPUs never had this. It goes through
round_e4m3, whose docstring already warns about an fp32 to fp16 to 3-bit chain that double-rounds on a tie. The hazard was known when the emulation was written. Only the native branch skipped the guard.The fix
Round onto the e4m3 grid in fp32 first, on both paths. The downcast is then exact and the lowering's rounding mode stops mattering.
fp8_block_linear._act_quant_kernelfp8_pertensor_linear._static_quant_kerneldsv4/fp8_linear._act_quant_fp8_kerneldsv4/fp8_linear._act_quant_inplace_kernel, where the native branch's fp8 round trip only ever re-quantized a value already on the grid, so it collapses into the sharedround_e4m3call.Test changes
tests/kernels/test_e4m3_compat.py::test_forced_emu_matches_nativefails onmainon sm_89 hardware, withblk_aq_y: EMU output differs from native, 52 of 18944 elements. The forced-emulation A/B is what catches this, because the emulation was the correct side.Two changes there:
test_native_downcast_needs_the_grid_roundpins the invariant directly, so a later change cannot make the emulation match the native bug instead. It deliberately does not assert that the bare downcast is wrong: a future Triton may lower it correctly, and the pre-round would then be redundant rather than incorrect.The three GEMM keys that compare an fp8 MMA against a bf16 MMA get a bound on the tensor scale instead of an elementwise one. Triton selects different MMA shapes for
fp8 x fp8andbf16 x bf16, so the fp32 reduction trees differ and the results land about one output ULP apart. The file's own comment predicted this: bit-exact on H100 is "lowering luck, not a guarantee", and it is not bit-exact on sm_89.moe_prefill_fp8chains two such GEMMs and then sums top-k outputs of magnitude 8e3 down to about 2e2, so an error worth 2e-3 of the GEMM scale reads as 13% elementwise. Bounding againstmax|ref|measures the GEMM's own error rather than the cancellation:blk_gemmdsv4_gemmmoe_prefill_fp8The limit is 1e-2.
Testing
RTX 6000 Ada, sm_89.
tests/kernels/test_e4m3_compat.pyonmainwith this PR: 9 passed. Onmainwithout it: 1 failed, 7 passed.Found while reviewing #48. That PR does not cause the failure; I reverted it and reproduced the failure on
main.