[ischencheng]

On sm_89 (RTX 40, L4, RTX 6000 Ada) the native fp8 quantizers round some values one e4m3 step toward zero. Triton 3.8 lowers fp32 -> tl.float8e4nv there as fp32 -RTZ-> fp16 -RNE-> e4m3, so a value just above an e4m3 midpoint lands on the midpoint and then ties to even (triton-lang/triton#11753, fixed by triton-lang/triton#11924 after 3.8.0). sm_90+ converts straight from fp32 and is fine. For example 200.04 becomes 192 on L4, while torch's cast and H100 give 208.

This is the bug #85 described in August. That PR was closed by its author and the bare cast is still on main.

f32_to_e4m3 in e4m3_compat.py keeps the plain cast on sm_90+ and on non-CUDA targets, and emits cvt.rn.satfinite.e4m3x2.f32 on sm_89, the same instruction Triton 3.8 uses on sm_90. A FIXME marks the sm_89 asm for removal once the minimum Triton is 3.9. The native branches of all seven fp32 -> e4m3 casts call it:

• fp8_block_linear._act_quant_kernel
• fp8_pertensor_linear._static_quant_kernel
• dsv4/fp8_linear.py: _act_quant_fp8_kernel, _act_quant_inplace_kernel, and the in-kernel act quant of _fp8_gemv_kernel
• dsv4/fused_moe._swiglu_kernel
• dsv41/pack.py: _e4m3_codes, _e4m3_round

Only the first two change values today. On 4096 x 7168 bf16 activations on L4, the block-fp8 act quant changes 74,586 of 29,360,128 output bytes (0.25%). The per-tensor quant scales the whole tensor by one factor, so far fewer inputs land in the gap: 1 byte on the same tensor, and the W8A8 test below is one such case. The DSV4 casts and the DSV4.1 KV pack (fp8 and fp4 rows) give the same bytes before and after on those inputs; most of them divide bf16 by a power of two, which fp16 holds exactly. They call the helper too, so every native cast goes through the same path.

test_round_to_grid_bitexact gets an f32_to_e4m3 case on the same ties, neighbors and subnormals it already runs for round_e4m3.

Tested with torch 2.14.1+cu130, Triton 3.8.0, driver 610.57.04, on top of c63fc13. No checkpoint is involved.

pytest tests/kernels/test_e4m3_compat.py tests/kernels/test_fp8_pertensor_linear.py tests/kernels/test_dsv41_pack.py tests/kernels/test_dsv4_fp8_block_linear.py tests/kernels/test_fp8_blockscale_moe.py tests/dsv4 tests/models/deepseek_v41, slow tests included:

GPU	main	this PR
L4 (sm_89)	2 failed, 122 passed, 1 skipped	1 failed, 124 passed, 1 skipped
H100 (sm_90)	124 passed, 1 skipped	125 passed, 1 skipped
RTX PRO 6000 (sm_120)	124 passed, 1 skipped	125 passed, 1 skipped
A10 (sm_86, emulated branch)	104 passed, 21 skipped	104 passed, 22 skipped

On L4:

• test_round_to_grid_bitexact[f32_to_e4m3]: with only the sm_89 line of f32_to_e4m3 turned back into x.to(tl.float8e4nv), assert 6986 == 0; with this PR it passes.
• test_w8a8_matches_w8a8_reference[part_rows0-False-4]: rel 0.0108 > 0.01 on main, passes with this PR.
• test_forced_emu_matches_native: on main it stops at blk_aq_y: EMU output differs from native. With this PR every act-quant key matches the emulated path bit for bit, and it stops at moe_prefill_fp8 (1 of 512 elements: 8.0 apart at 112, atol 0.5). That key compares an fp8-MMA GEMM with a bf16-MMA one; its native-vs-emulated difference is the same on main (44 of 512 elements differ, max 64 at |y| up to 8704, relative norm 1.9e-3), it was just hidden behind blk_aq_y. So on sm_89 this test still fails, now at that key.

The rest of tests/kernels, tests/moe, tests/dsv4, tests/models/deepseek_v41 and tests/models/qwen4_exp on L4 shows no new failures: main and this PR fail the same 4 tests, none of them fp8 (one is #765).

The cast on L4, 6.3M fp32 values (every e4m3 grid value and midpoint, one fp32 ulp either side of each midpoint, uniform samples over the range and over the subnormals), against torch's .to(torch.float8_e4m3fn): x.to(tl.float8e4nv) differs on 17936, round_e4m3 on 0. H100 and RTX PRO 6000 differ on 0. per_token_group_quant_fp8 on 4096 x 512 activations: 5235 of 2,097,152 outputs differ from torch on main, 0 with this PR.

sm_90 and sm_120: the SASS of every FreeToken kernel the tests above compile is identical between main and this PR (45 kernels, cuobjdump -sass instruction text). Only the test's own kernel changes.

Kernel time on L4, do_bench median, main and this PR alternated in one container, median of 3 runs each:

kernel	M	main (µs)	this PR (µs)
per_token_group_quant_fp8, M x 7168	1 / 64 / 4096	6.14 / 12.29 / 442.37	6.14 / 11.26 / 442.37
act_quant_fp8 (dsv4), M x 7168	1 / 64 / 4096	6.14 / 12.29 / 438.27	6.14 / 11.26 / 438.27
_static_quant (per-tensor), M x 7168	1 / 64 / 4096	5.12 / 11.26 / 392.19	6.14 / 11.26 / 393.22
quant_roundtrip fp8_e8m0_b32 (dsv41), M x 512	1 / 64 / 4096	9.22 / 10.24 / 51.20	9.22 / 10.27 / 51.20

The small sizes sit on the timer's ~1 µs step.

ROCm: @zihaomu checked the non-CUDA guard on a Radeon AI PRO R9700 (gfx1201), see the thread below.

Not tested: model-level accuracy.

[jason-fxz]

Thanks, the fix looks right to me. Two small things:

1. Triton fixes this in triton-lang/triton#11924, which is on release/3.9.x and not in any 3.8.x release. Could you put a FIXME on the sm_89 path so it gets dropped when we move to 3.9? For example:

    if _f32_cast_rounds_once_cx():
        return x.to(tl.float8e4nv)
    # FIXME: remove the sm_89 inline asm once the minimum Triton is 3.9 (triton-lang/triton#11924)

2. _f32_cast_rounds_once_cx() returns False on any non-CUDA target, so on HIP f32_to_e4m3 would emit the NVIDIA-only cvt.rn.satfinite.e4m3x2.f32. Today e4m3_native() keeps ROCm off this path, so nothing breaks. But _static_quant_kernel calls the helper with no e4m3_native_cx() branch of its own. Limiting the asm to CUDA would remove that trap:

@constexpr_function
def _f32_cast_rounds_once_cx():
    return not target_info.is_cuda() or target_info.cuda_capability_geq(9, 0)

@zihaomu could you check whether this change is OK on AMD? Specifically, are there ROCm targets where we plan to let fp8e4nv through e4m3_native(), and does the plain x.to(tl.float8e4nv) round correctly there?

[zihaomu]

Hi @jason-fxz, thanks for bringing this up! I'll test this locally and share the results in a follow-up comment.

[zihaomu]

@ischencheng @jason-fxz AMD validation is complete. The suggested non-CUDA guard works on Radeon AI PRO R9700 (gfx1201), ROCm 10.1.0 / HIP 7.16.26385, PyTorch 2.14.0+rocm10.1.0 and Triton 3.8.0. No checkpoint was involved.

@ischencheng, please feel free to cherry-pick this commit directly into this PR: zihaomu/FreeToken@d0f0a9b. It is based on this PR's current head, e7f22fd, and changes only e4m3_compat.py (2 insertions, 1 deletion):

git fetch https://github.com/zihaomu/FreeToken.git fix/pr805-amd-e4m3-cast
git cherry-pick d0f0a9b4a48bfd9657fde78ec8c9732ab572d633

The commit applies exactly the suggested condition:

@constexpr_function
def _f32_cast_rounds_once_cx():
    return not target_info.is_cuda() or target_info.cuda_capability_geq(9, 0)

It also adds the requested Triton 3.9 FIXME before the sm_89 inline asm.

Validation against the original PR:

• Direct calls to the original helper and _static_quant reach NVIDIA PTX on HIP. In this toolchain LLVM prints could not allocate output register for constraint 'h' but still returns a kernel, and the outputs are incorrect. With the guard, both _static_quant cases match PyTorch byte-for-byte.
• The plain x.to(tl.float8e4nv) and the guarded helper each match PyTorch CPU/GPU FP8 bytes on 3,146,751 finite inputs in [-448, 448], with 0 mismatches, including all finite codes, every midpoint and its adjacent fp32 values, subnormals and signed zero. 200.04 correctly rounds to 208.
• Existing targeted tests: original: 1 failed, 53 passed, 16 skipped; guarded: 54 passed, 16 skipped (5 deselected in both). The failure fixed is TestPrimitives::test_round_to_grid_bitexact[f32_to_e4m3].
• Compile-only checks for gfx1100/1101/1200/1201/942/950: the guarded helper contains no NVIDIA PTX and has the same AMD instruction sequence as the plain cast. On sm89/90/120, the helper's PTX is unchanged apart from debug locations. These foreign-target checks are not hardware execution results.

The existing host and constexpr native gates remain disabled on HIP, so normal wrappers currently keep the emulated path. For future OCP E4M3 native enablement, RDNA4 (gfx1200/1201) and CDNA4 (gfx950) are hardware candidates; gfx942 has native FNUZ FP8, which must not be treated as fp8e4nv. My numerical execution evidence here is for gfx1201; enabling other targets still needs their own hardware validation.

The existing-test command in the matched ROCm environment was:

python -m pytest \
  tests/kernels/test_e4m3_compat.py \
  tests/kernels/test_fp8_pertensor_linear.py \
  tests/kernels/test_dsv41_pack.py \
  tests/kernels/test_dsv4_fp8_block_linear.py \
  tests/kernels/test_fp8_blockscale_moe.py \
  -k "not forced_emu and not compile_gate_foreign_arch" -q

[ischencheng]

@jason-fxz thanks for the review, and thanks @zihaomu for testing on AMD! Cherry-picked d0f0a9b as ab85da0, so both changes are in. Reran the tests on L4 and H100 after it: same results as in the description, and the SASS of all 46 compiled kernels is unchanged.

[jason-fxz]

Thanks, both points are in. LGTM.