DeepSeek-V4.1's fp4_e4m3_b16 pack rounds a tied scale up on pre-sm_89 GPUs #673

Before you start

• I have read the FAQ and my problem is not answered there.
• I have read the Roadmap and this is not already planned there.
• I have searched existing issues and found no duplicate.
• I am on the latest release, or on a freshly rebuilt main when building from source.

What happened

On an RTX 3090 Ti (sm_86), test_roundtrip_matches_the_reference_quantizer[128-fp4_e4m3_b16] and [512-fp4_e4m3_b16] in tests/kernels/test_dsv41_pack.py fail with fp4_e4m3_b16: 0.015625. The kernel stores some group scales one e4m3 step above the reference's: 2 of 2,400 groups at dim 128 and 16 of 9,600 at dim 512. In every one, amax / 6 lies exactly on an e4m3 midpoint. For example, amax 0.052734375 gives 0.0087890625, halfway between the subnormals 0.0078125 and 0.009765625. .to(torch.float8_e4m3fn) rounds to the even 0.0078125, but the kernel stores 0.009765625.

The emulated path does this. _pack_rows_kernel rounds the scale with _e4m3_round(amax * (1.0 / 6.0)), and round_e4m3's subnormal branch is (x + 24576.0) - 24576.0. Triton fuses the multiply into that add: a kernel computing round_e4m3(amax * (1.0 / 6.0)) compiles the add to fma.rn.f32 %r8, %r1, 0f3E2AAAAB, 0f46C00000. The add therefore rounds the unrounded product, which lies just above the midpoint because 1/6 rounds up in fp32, and gets the odd code. The same value loaded from memory rounds to 0.0078125. The native branch converts in hardware and does not take this path.

I tested three fixes here; each rounds both tied values I hit (amax 0.052734375 and 0.076171875) as torch's cast does:

1. Divide as the reference does: s = _e4m3_round(tl.math.div_rn(amax, 6.0)), as the kernel already does for the codes. With this change the test file passes, 10 of 10.
2. Make round_e4m3's subnormal branch unfusable with libdevice.sub_rn(libdevice.add_rn(x, 24576.0), 24576.0). That keeps its round-to-nearest-even result whatever the caller computes its argument with. The other callers clamp before rounding, so no multiply reaches the add there today.
3. Compile with enable_fp_fusion=False. Triton 3.9 turns fusion off by default (Upgrade to torch 2.14.1, Triton 3.8 and sglang-kernel 0.4.9 #629), which would hide this.

Expected: every group scale equals the reference quantizer's, so all 10 cases of tests/kernels/test_dsv41_pack.py pass.

Actual: the two fp4_e4m3_b16 cases fail with fp4_e4m3_b16: 0.015625, a tied scale one e4m3 step high (the log below).

How did you install FreeToken

Built from source

FreeToken version

9c635a1

OS

WSL2 on Windows 11

OS details

Ubuntu 24.04, Python 3.12, torch 2.14.0+cu130, triton 3.8.0. Native Windows 11 (Python 3.13, triton-windows 3.8.0, with the Windows PRs from #384) gives the same scales.

GPU and driver

RTX 3090 Ti 24 GB, driver 617.14

CPU and system RAM

AMD Ryzen 9 9950X3D, 256 GB

Checkpoint

None; the kernel test runs without one.

Command

pytest -q --tb=line -rf tests/kernels/test_dsv41_pack.py

Full log

..F...F...                                                               [100%]
=================================== FAILURES ===================================
/home/yevhenii/FreeToken/tests/kernels/test_dsv41_pack.py:73: AssertionError: fp4_e4m3_b16: 0.015625
/home/yevhenii/FreeToken/tests/kernels/test_dsv41_pack.py:73: AssertionError: fp4_e4m3_b16: 0.015625
=========================== short test summary info ============================
FAILED tests/kernels/test_dsv41_pack.py::test_roundtrip_matches_the_reference_quantizer[128-fp4_e4m3_b16]
FAILED tests/kernels/test_dsv41_pack.py::test_roundtrip_matches_the_reference_quantizer[512-fp4_e4m3_b16]
2 failed, 8 passed in 3.93s

Anything else

FREETOKEN_FORCE_E4M3_EMU=1 selects the same emulated path, so it should reproduce on a newer GPU too; I could only run it on sm_86.