diff --git a/python/freetoken/kernel/triton/dsv4/fp8_linear.py b/python/freetoken/kernel/triton/dsv4/fp8_linear.py index 47e3c9f5..06c46030 100644 --- a/python/freetoken/kernel/triton/dsv4/fp8_linear.py +++ b/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) diff --git a/python/freetoken/kernel/triton/fp8_block_linear.py b/python/freetoken/kernel/triton/fp8_block_linear.py index 0999c56d..35ef013d 100644 --- a/python/freetoken/kernel/triton/fp8_block_linear.py +++ b/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) diff --git a/python/freetoken/kernel/triton/fp8_pertensor_linear.py b/python/freetoken/kernel/triton/fp8_pertensor_linear.py index 28a54c94..05da97c6 100644 --- a/python/freetoken/kernel/triton/fp8_pertensor_linear.py +++ b/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: diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index 8bcb8e58..6fd8254e 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/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"