fix(kernel): convert fp32 to e4m3 in one rounding on sm_89

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

@@ -33,6 +33,7 @@
     e4m3_kernel_view,
     e4m3_native_cx,
     e4m3_u8_to_f32,
+    f32_to_e4m3,
     round_e4m3,
 )
 
@@ -74,7 +75,7 @@ def _act_quant_fp8_kernel(
     s = tl.exp2(e.to(tl.float32))
     y = tl.clamp(x / s[:, None], -448.0, 448.0)
     if e4m3_native_cx():
-        y = y.to(tl.float8e4nv)
+        y = f32_to_e4m3(y)
     else:
         y = round_e4m3(y)  # e4m3-grid values into the wrapper's bf16 buffer
     tl.store(
@@ -133,7 +134,7 @@ def _act_quant_inplace_kernel(
     if FP4:
         q = _round_fp4(q)
     elif e4m3_native_cx():
-        q = q.to(tl.float8e4nv).to(tl.float32)
+        q = f32_to_e4m3(q).to(tl.float32)
     else:
         q = round_e4m3(q)
     optrs = o_ptr + offs_m[:, None] * stride_om + offs_k[None, :] * stride_on
@@ -344,7 +345,7 @@ def _fp8_gemv_kernel(
                     sa = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32))  # [NSLAB]
                     q = tl.clamp(xs / sa[:, None], -448.0, 448.0)
                     if e4m3_native_cx():
-                        q = q.to(tl.float8e4nv).to(tl.float32)
+                        q = f32_to_e4m3(q).to(tl.float32)
                     else:
                         q = round_e4m3(q).to(tl.float32)
                 else:

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

@@ -25,7 +25,7 @@
 import triton.language as tl
 
 from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil
-from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, round_e4m3
+from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, f32_to_e4m3, round_e4m3
 
 _E2M1_VALUES = [
     0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
@@ -158,7 +158,7 @@ def _swiglu_kernel(
         sc = tl.exp2(_log2_ceil(amax * (1.0 / 448.0)).to(tl.float32))
         q = tl.clamp(xs / sc[:, None], -448.0, 448.0)
         if e4m3_native_cx():
-            q = q.to(tl.float8e4nv).to(tl.float32)
+            q = f32_to_e4m3(q).to(tl.float32)
         else:
             q = round_e4m3(q)
         act = tl.reshape(q * sc[:, None], (BLOCK,))

python/freetoken/kernel/triton/dsv41/pack.py

@@ -22,7 +22,7 @@
 import triton.language as tl
 
 from freetoken.kernel.triton.dsv4.fp8_linear import _log2_ceil, _round_fp4
-from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8, e4m3_native_cx, round_e4m3
+from freetoken.kernel.triton.e4m3_compat import e4m3_f32_to_u8, e4m3_native_cx, f32_to_e4m3, round_e4m3
 from freetoken.kvcache.dsv4.v41_row_format import BF16, RowFormat
 
 from .row_format import fp4_f32_to_nibble, load_rows
@@ -32,7 +32,7 @@
 def _e4m3_codes(q):
     """fp32 (|q| <= 448) -> e4m3 bit pattern, RNE."""
     if e4m3_native_cx():
-        return q.to(tl.float8e4nv).to(tl.uint8, bitcast=True)
+        return f32_to_e4m3(q).to(tl.uint8, bitcast=True)
     else:
         return e4m3_f32_to_u8(round_e4m3(q))
 
@@ -41,7 +41,7 @@ def _e4m3_codes(q):
 def _e4m3_round(q):
     """fp32 -> nearest e4m3 value, as fp32."""
     if e4m3_native_cx():
-        return q.to(tl.float8e4nv).to(tl.float32)
+        return f32_to_e4m3(q).to(tl.float32)
     else:
         return round_e4m3(q)
 

python/freetoken/kernel/triton/e4m3_compat.py

@@ -3,7 +3,7 @@
 Triton rejects the fp8e4nv type anywhere in a kernel compiled for sm < 89 -- the
 check sits in ``dtype.to_ir``, so even an fp8 *pointer argument* is illegal.
 Affected kernels branch on :func:`e4m3_native_cx` (a compile-time constexpr): the
-native branch stays byte-identical on sm_89+, the emulated branch is dead-code
+native branch casts fp32 through :func:`f32_to_e4m3`, the emulated branch is dead-code
 eliminated there. When the emulated branch is active, wrappers must pass e4m3
 tensors as ``.view(torch.uint8)`` and allocate act-quant outputs as bf16 -- use
 the host-side twin :func:`e4m3_native` for those decisions.
@@ -83,6 +83,11 @@ def e4m3_act_dtype() -> torch.dtype:
     return torch.float8_e4m3fn if e4m3_native() else torch.bfloat16
 
 
+@constexpr_function
+def _f32_cast_rounds_once_cx():
+    return not target_info.is_cuda() or target_info.cuda_capability_geq(9, 0)
+
+
 @constexpr_function
 def e4m3_native_cx():
     """Compile-time: does the compilation target have native fp8e4nv (sm_89+)?
@@ -140,3 +145,18 @@ def round_e4m3(x):
     y_norm = ((b + 524287 + lsb) & 0xFFF00000).to(tl.float32, bitcast=True)
     y_sub = (x + 24576.0) - 24576.0
     return tl.where(tl.abs(x) >= 0.015625, y_norm, y_sub)
+
+
+@jit
+def f32_to_e4m3(x):
+    """Native-branch fp32 -> fp8e4nv cast with one RNE rounding (satfinite).
+
+    Triton 3.8 lowers this cast on sm_89 as fp32 -RTZ-> fp16 -RNE-> e4m3, which double-rounds a
+    value just above an e4m3 midpoint (triton-lang/triton#11753); emit the direct cvt there."""
+    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).
+    # cvt packs its first source into the high byte, so the low byte is fp8(x)
+    h = tl.inline_asm_elementwise(
+        "cvt.rn.satfinite.e4m3x2.f32 $0, 0f00000000, $1;", "=h,r", [x], dtype=tl.uint16, is_pure=True, pack=1)
+    return h.to(tl.uint8).to(tl.float8e4nv, bitcast=True)

python/freetoken/kernel/triton/fp8_block_linear.py

@@ -22,6 +22,7 @@
     e4m3_kernel_view,
     e4m3_native_cx,
     e4m3_u8_to_f32,
+    f32_to_e4m3,
     round_e4m3,
 )
 
@@ -53,7 +54,7 @@ def _act_quant_kernel(
     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)
     if e4m3_native_cx():
-        y = y.to(tl.float8e4nv)
+        y = f32_to_e4m3(y)
     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])

python/freetoken/kernel/triton/fp8_pertensor_linear.py

@@ -31,6 +31,7 @@
     e4m3_native,
     e4m3_native_cx,
     e4m3_u8_to_f32,
+    f32_to_e4m3,
 )
 
 FP8 = torch.float8_e4m3fn
@@ -253,7 +254,7 @@ 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)
+    tl.store(out_ptr + offs, f32_to_e4m3(v), mask=mask)
 
 
 def _static_quant(a: torch.Tensor, input_scale: torch.Tensor) -> torch.Tensor:

tests/kernels/test_e4m3_compat.py

@@ -106,16 +106,19 @@ def k(v_ptr, o_ptr, BLOCK: tl.constexpr):
         ok = ~ref.isnan()
         assert torch.equal(out[ok], ref[ok])
 
-    def test_round_to_grid_bitexact(self):
+    @pytest.mark.parametrize("native_cast", [False, True], ids=["round_e4m3", "f32_to_e4m3"])
+    def test_round_to_grid_bitexact(self, native_cast):
         import triton
         import triton.language as tl
-        from freetoken.kernel.triton.e4m3_compat import round_e4m3
+        from freetoken.kernel.triton.e4m3_compat import f32_to_e4m3, round_e4m3
 
         @triton.jit
-        def k(x_ptr, y_ptr, N, BLOCK: tl.constexpr):
+        def k(x_ptr, y_ptr, N, BLOCK: tl.constexpr, NATIVE_CAST: tl.constexpr):
             offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
             m = offs < N
-            tl.store(y_ptr + offs, round_e4m3(tl.load(x_ptr + offs, mask=m, other=0.0)), mask=m)
+            x = tl.load(x_ptr + offs, mask=m, other=0.0)
+            y = f32_to_e4m3(x).to(tl.float32) if NATIVE_CAST else round_e4m3(x)
+            tl.store(y_ptr + offs, y, mask=m)
 
         torch.manual_seed(0)
         grid = torch.arange(256, dtype=torch.uint8, device="cuda").view(FP8).to(torch.float32)
@@ -131,7 +134,7 @@ def k(x_ptr, y_ptr, N, BLOCK: tl.constexpr):
         ]
         x = torch.cat(parts).clamp(-448, 448).contiguous()
         y = torch.empty_like(x)
-        k[(triton.cdiv(x.numel(), 1024),)](x, y, x.numel(), BLOCK=1024)
+        k[(triton.cdiv(x.numel(), 1024),)](x, y, x.numel(), BLOCK=1024, NATIVE_CAST=native_cast)
         ref = x.to(FP8).to(torch.float32)
         assert int(((y != ref) & ~(y.isnan() & ref.isnan())).sum()) == 0