fix(fp8): round onto the e4m3 grid before the native float8e4nv downcast

python/freetoken/kernel/triton/dsv4/fp8_linear.py

@@ -71,10 +71,12 @@ def _act_quant_fp8_kernel(
     e = _log2_ceil(amax * (1.0 / 448.0))                # [BLOCK_M]
     s = tl.exp2(e.to(tl.float32))
     y = tl.clamp(x / s[:, None], -448.0, 448.0)
+    # See fp8_block_linear: triton's float8e4nv downcast double-rounds one-sided
+    # (toward zero), so round onto the grid in fp32 first and let the (now exact)
+    # downcast follow.
+    y = round_e4m3(y)  # emu path keeps these grid values in the wrapper's bf16 buffer
     if e4m3_native_cx():
         y = y.to(tl.float8e4nv)
-    else:
-        y = round_e4m3(y)  # e4m3-grid values into the wrapper's bf16 buffer
     tl.store(
         y_ptr + offs_m[:, None] * stride_ym + offs_k[None, :] * stride_yn,
         y, mask=m_mask[:, None],
@@ -130,9 +132,9 @@ def _act_quant_inplace_kernel(
     q = tl.clamp(x / s[:, None], FP8_MIN, FP8_MAX)
     if FP4:
         q = _round_fp4(q)
-    elif e4m3_native_cx():
-        q = q.to(tl.float8e4nv).to(tl.float32)
     else:
+        # No fp8 round-trip on the native path: it only ever *re*-quantized an
+        # already-grid value, and triton's downcast double-rounds rather than RNE.
         q = round_e4m3(q)
     optrs = o_ptr + offs_m[:, None] * stride_om + offs_k[None, :] * stride_on
     y = (q * s[:, None]).to(optrs.dtype.element_ty)

python/freetoken/kernel/triton/fp8_block_linear.py

@@ -53,10 +53,15 @@ def _act_quant_kernel(
     amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-10)
     s = amax / 448.0  # [BLOCK_M] fp32 per-token-group scale (e4m3 finite max = 448)
     y = tl.clamp(x / s[:, None], -448.0, 448.0)
+    # Round onto the e4m3 grid in fp32 FIRST, on both paths. triton's fp32 ->
+    # 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.
+    y = round_e4m3(y)
     if e4m3_native_cx():
         y = y.to(tl.float8e4nv)
-    else:
-        y = round_e4m3(y)  # e4m3-grid values into the wrapper's bf16 buffer
     tl.store(y_ptr + offs_m[:, None] * stride_ym + offs_k[None, :] * stride_yk, y, mask=m_mask[:, None])
     tl.store(s_ptr + offs_m * stride_sm + pid_k * stride_sk, s, mask=m_mask)
 

python/freetoken/kernel/triton/fp8_pertensor_linear.py

@@ -24,6 +24,7 @@
 import torch
 import triton
 import triton.language as tl
+from freetoken.kernel.triton.e4m3_compat import round_e4m3
 from freetoken.layers import BaseOP
 from freetoken.layers.base import _concat_prefix
 
@@ -204,7 +205,9 @@ def _static_quant_kernel(x_ptr, out_ptr, scale_ptr, n_elements, BLOCK: tl.conste
     inv = 1.0 / tl.load(scale_ptr).to(tl.float32)
     v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv
     v = tl.minimum(tl.maximum(v, -448.0), 448.0)
-    tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask)
+    # RNE onto the e4m3 grid first: triton's float8e4nv downcast double-rounds
+    # one-sided (see fp8_block_linear), biasing 0.38% of values 1 ULP toward zero.
+    tl.store(out_ptr + offs, round_e4m3(v).to(tl.float8e4nv), mask=mask)
 
 
 def _static_quant(a: torch.Tensor, input_scale: torch.Tensor) -> torch.Tensor:

tests/kernels/test_e4m3_compat.py

@@ -30,11 +30,17 @@
 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
+# bf16xbf16 tl.dot to different MMA shapes, so the accumulator's fp32 reduction tree
+# differs and the two land ~1 output-ULP apart. Bit-exact on H100 for dsv4_gemm /
+# moe_prefill_fp8; NOT on sm_89 (RTX 6000 Ada) -- that was lowering luck.
+#
+# The bound is on the tensor SCALE, not elementwise: moe_prefill_fp8 chains two such
+# GEMMs and then top-k-sums outputs of magnitude ~8e3 down to ~2e2, so a difference
+# worth 2e-3 of the GEMM scale reads as a 13% *elementwise* relative error. Comparing
+# against max|ref| measures the GEMM's own error instead of the cancellation's.
 _MMA_TOL_KEYS = {"blk_gemm", "dsv4_gemm", "moe_prefill_fp8"}
-_MMA_TOL = dict(rtol=5e-2, atol=0.5)
+_MMA_REL = 1e-2  # of max|ref|; observed worst case 4.6e-3 (blk_gemm, sm_89)
 
 
 def _native_cc() -> bool:
@@ -112,6 +118,40 @@ def k(x_ptr, y_ptr, N, BLOCK: tl.constexpr):
         ref = x.to(FP8).to(torch.float32)
         assert int(((y != ref) & ~(y.isnan() & ref.isnan())).sum()) == 0
 
+    def test_native_downcast_needs_the_grid_round(self):
+        """``round_e4m3(x).to(float8e4nv)`` is RNE; the bare ``.to(float8e4nv)`` is not.
+
+        triton lowers fp32 -> float8e4nv as a double-round (via fp16, RTZ), so a value
+        a hair above a grid midpoint collapses onto the midpoint and then ties to even
+        -- one-sided, always toward zero. Every native-fp8 quantizer therefore rounds
+        onto the grid in fp32 first; this pins that the pre-round makes it exact.
+        """
+        import triton
+        import triton.language as tl
+        from freetoken.kernel.triton.e4m3_compat import round_e4m3
+
+        @triton.jit
+        def k(x_ptr, raw_ptr, fix_ptr, N, BLOCK: tl.constexpr):
+            offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+            m = offs < N
+            x = tl.load(x_ptr + offs, mask=m, other=0.0)
+            tl.store(raw_ptr + offs, x.to(tl.float8e4nv), mask=m)
+            tl.store(fix_ptr + offs, round_e4m3(x).to(tl.float8e4nv), mask=m)
+
+        torch.manual_seed(0)
+        x = torch.empty(1 << 22, device="cuda").uniform_(-448, 448).contiguous()
+        raw = torch.empty_like(x, dtype=FP8)
+        fix = torch.empty_like(raw)
+        k[(triton.cdiv(x.numel(), 1024),)](x, raw, fix, x.numel(), BLOCK=1024)
+        ref = x.to(FP8)
+        assert torch.equal(fix, ref), "pre-rounding onto the grid must make the downcast exact"
+        # Not asserting raw != ref: a future triton may lower this correctly, at which
+        # point the pre-round is redundant rather than wrong. Report the bias instead.
+        bad = raw.float() != ref.float()
+        if bad.any():
+            assert (raw.float()[bad].abs() < ref.float()[bad].abs()).all(), \
+                "bare downcast should only ever err toward zero"
+
 
 # ======================================================================================
 # Shared emit path: every affected wrapper, deterministic inputs.
@@ -262,7 +302,10 @@ def test_forced_emu_matches_native(tmp_path):
     assert a["native"] and not b["native"]
     for k in [k for k in a if k != "native"]:
         if k in _MMA_TOL_KEYS:
-            torch.testing.assert_close(a[k], b[k], **_MMA_TOL)
+            ref = a[k].float()
+            err = (ref - b[k].float()).abs().max().item()
+            lim = _MMA_REL * ref.abs().max().item()
+            assert err <= lim, f"{k}: EMU differs from native by {err:.4g} > {lim:.4g}"
         else:
             assert torch.equal(a[k], b[k]), f"{k}: EMU output differs from native"