From e7f22fd7454fcd6a3a58f6cf05372a05261187d1 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:54:45 -0400 Subject: [PATCH 1/2] fix(kernel): convert fp32 to e4m3 in one rounding on sm_89 --- .../kernel/triton/dsv4/fp8_linear.py | 7 ++++--- .../freetoken/kernel/triton/dsv4/fused_moe.py | 4 ++-- python/freetoken/kernel/triton/dsv41/pack.py | 6 +++--- python/freetoken/kernel/triton/e4m3_compat.py | 21 ++++++++++++++++++- .../kernel/triton/fp8_block_linear.py | 3 ++- .../kernel/triton/fp8_pertensor_linear.py | 3 ++- tests/kernels/test_e4m3_compat.py | 13 +++++++----- 7 files changed, 41 insertions(+), 16 deletions(-) diff --git a/python/freetoken/kernel/triton/dsv4/fp8_linear.py b/python/freetoken/kernel/triton/dsv4/fp8_linear.py index fb17a35175..49eee88f01 100644 --- a/python/freetoken/kernel/triton/dsv4/fp8_linear.py +++ b/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: diff --git a/python/freetoken/kernel/triton/dsv4/fused_moe.py b/python/freetoken/kernel/triton/dsv4/fused_moe.py index 0adfbdd34b..4df808a15c 100644 --- a/python/freetoken/kernel/triton/dsv4/fused_moe.py +++ b/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,)) diff --git a/python/freetoken/kernel/triton/dsv41/pack.py b/python/freetoken/kernel/triton/dsv41/pack.py index c0fca83162..4dc119f938 100644 --- a/python/freetoken/kernel/triton/dsv41/pack.py +++ b/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) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 14a7acd835..670113d083 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/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 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,17 @@ 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) + # 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) diff --git a/python/freetoken/kernel/triton/fp8_block_linear.py b/python/freetoken/kernel/triton/fp8_block_linear.py index b79aa5b2f6..ce1384d99a 100644 --- a/python/freetoken/kernel/triton/fp8_block_linear.py +++ b/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]) diff --git a/python/freetoken/kernel/triton/fp8_pertensor_linear.py b/python/freetoken/kernel/triton/fp8_pertensor_linear.py index 5561a86a6b..0701e3c31e 100644 --- a/python/freetoken/kernel/triton/fp8_pertensor_linear.py +++ b/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: diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index 1532feb406..37629614b6 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/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 From ab85da0e93930d71537e820e82a5af74eac9805b Mon Sep 17 00:00:00 2001 From: zihaomu Date: Fri, 23 Oct 2026 14:08:45 +0800 Subject: [PATCH 2/2] fix(kernel): avoid nvidia fp8 cast asm on hip --- python/freetoken/kernel/triton/e4m3_compat.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 670113d083..60cb35a51e 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -85,7 +85,7 @@ def e4m3_act_dtype() -> torch.dtype: @constexpr_function def _f32_cast_rounds_once_cx(): - return target_info.cuda_capability_geq(9, 0) + return not target_info.is_cuda() or target_info.cuda_capability_geq(9, 0) @constexpr_function @@ -155,6 +155,7 @@ def f32_to_e4m3(x): 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)