From 32c39e0930829828a4f263e330df6d191b1d60ab Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 25 Jul 2026 03:20:44 -0700 Subject: [PATCH 01/14] Add MXFP4 weight QAT recipes on MXFP8 and FP8 block-scaling hosts MXFP4QATMXFP8BlockScaling / MXFP4QATFloat8BlockScaling project weights onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid before the host recipe quantizes them; activations and gradients are untouched. The projection is a fused fake-quantization with three bit-identical implementations (CUDA kernel, CuTe DSL, PyTorch reference) dispatched via NVTE_MXFP4_QAT_IMPL, with an identity-STE gradient. - scale contract matches the TileKernels deployment path (floor 2^-126, cap 2^125); integer-bit amax/scale-exponent derivation and non-FTZ PTX keep results bit-exact under --use_fast_math builds - the rowwise MXFP8 and 128x128 blockwise weight encodings of the projected weight are bitwise lossless (raw payload/scale verified); the columnwise 32x1 encoding is bounded - fix ptx.cuh exp2f expansion of UE8M0 code 0 (2^-127) and code 255 (NaN) - reject MXFP4 QAT loudly on surfaces that bypass the weight-quantization hook (te.ops, Userbuffers, fused grouped MLP, quantized primary weights); project the weight in FSDP2/GTP backward rematerialization; invalidate cached weight workspaces on base<->QAT recipe switches - tests: per-step bitwise losslessness, bf16-exhaustive and fp32 bit-fuzz oracles, RTNE midpoint/threshold vectors, fast-math immunity build, STE, misaligned/non-contiguous inputs, and a TileKernels/CuTe-DSL/CUDA/torch four-way bitwise matrix plus a 24-config e2e backward-override matrix --- tests/pytorch/test_mxfp4_qat.py | 1139 +++++++++++++++++ transformer_engine/common/cast/cast.cu | 32 + .../common/cast/mxfp4/fake_quantize_mxfp4.cuh | 129 ++ .../common/include/transformer_engine/cast.h | 11 + transformer_engine/common/recipe/__init__.py | 52 + transformer_engine/common/util/ptx.cuh | 2 + transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 19 + .../pytorch/csrc/extensions/pybind.cpp | 3 + transformer_engine/pytorch/module/base.py | 26 +- .../pytorch/module/layernorm_linear.py | 3 + .../pytorch/module/layernorm_mlp.py | 16 +- transformer_engine/pytorch/module/linear.py | 5 + transformer_engine/pytorch/mxfp4_qat.py | 217 ++++ .../pytorch/mxfp4_qat_cute_dsl.py | 217 ++++ .../pytorch/ops/basic/basic_linear.py | 10 + .../pytorch/ops/basic/grouped_linear.py | 7 + .../pytorch/ops/fused/grouped_mlp.py | 6 + .../ops/fused/userbuffers_forward_linear.py | 4 + 19 files changed, 1892 insertions(+), 8 deletions(-) create mode 100644 tests/pytorch/test_mxfp4_qat.py create mode 100644 transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh create mode 100644 transformer_engine/pytorch/mxfp4_qat.py create mode 100644 transformer_engine/pytorch/mxfp4_qat_cute_dsl.py diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py new file mode 100644 index 0000000000..1e33e907a5 --- /dev/null +++ b/tests/pytorch/test_mxfp4_qat.py @@ -0,0 +1,1139 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Per-step and end-to-end tests for the MXFP4 weight-QAT recipes. + +Step coverage (letters refer to the design notes): + D/E/F losslessness: fake-quant exact in bf16; MXFP8 rowwise and 128x128 + blockwise weight encodings of the MXFP4-grid weight are bit-exact. + A/B/C backward_override weight semantics through real fwd/bwd GEMMs, + with discriminators (closer to w_hat vs closer to the master). + G STE: weight gradients land on the module parameter / main_grad. +""" +import ctypes +import os + +for _n in ("libnvrtc.so.13", "libnvrtc.so.12", "libnvrtc.so"): + try: + ctypes.CDLL(_n, mode=ctypes.RTLD_GLOBAL) + break + except OSError: + continue + +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import ( + MXFP4QATFloat8BlockScaling, + MXFP4QATMXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8_autocast +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + +torch.manual_seed(1234) +DEV = "cuda" +RUN_PERF = os.environ.get("NVTE_MXFP4_QAT_PERF", "0") == "1" +E2M1_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=DEV) +SHAPES = [(128, 256), (192, 448), (512, 1024)] + +FP8_TOL_1OP = 0.045 +FP8_TOL_2OP = 0.065 +BF16_TOL = 0.005 + + +def make_weight(m, n, dtype=torch.bfloat16, zero_blocks=0, outliers=0): + w = (torch.randn(m, n, device=DEV, dtype=torch.float32) * 0.02).to(dtype) + if outliers: + idx = torch.randint(0, m * n, (outliers,), device=DEV) + w.view(-1)[idx] *= 500.0 + if zero_blocks: + for i in range(zero_blocks): + r = torch.randint(0, m, (1,)).item() + k = torch.randint(0, n // 32, (1,)).item() + w[r, k * 32 : (k + 1) * 32] = 0.0 + return w + + +def rel_err(got, ref): + got, ref = got.to(torch.float64), ref.to(torch.float64) + return ((got - ref).norm() / ref.norm().clamp(min=1e-30)).item() + + +def fake_quant_ref_fp64(w): + """Independent fp64 reference: nearest-with-ties-to-even onto the E2M1 grid.""" + m, n = w.shape + w64 = w.to(torch.float64).view(m, n // 32, 32) + amax = w64.abs().amax(dim=-1, keepdim=True) + bits = amax.to(torch.float32).view(torch.int32) + exp_field = bits >> 23 + mant = bits & 0x7FFFFF + exp = exp_field - 129 + (mant > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)).clamp(-126, 125) + exp = torch.where(amax > 0, exp, torch.zeros_like(exp)) + scale = torch.ldexp(torch.ones_like(amax), exp) + y = (w64 / scale).clamp(min=-6.0, max=6.0) + grid = E2M1_GRID.to(torch.float64) + cand = torch.cat([-grid.flip(0)[:-1], grid]) + d = (y.unsqueeze(-1) - cand).abs() + near = d.argmin(dim=-1) + even = torch.tensor( + [abs(v) in (0.0, 1.0, 2.0, 4.0) for v in cand.tolist()], device=w.device + ) + dmin = d.gather(-1, near.unsqueeze(-1)).squeeze(-1) + tie_even = d.eq(dmin.unsqueeze(-1)) & even + has_even_tie = tie_even.any(dim=-1) + near_even = torch.where( + has_even_tie, tie_even.to(torch.uint8).argmax(dim=-1), near + ) + q = cand[near_even] + return (q * scale).view(m, n).to(torch.float64) + + + + +def test_fake_quant_grid_scale_rtne(): + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w_hat = mxfp4_fake_quantize(w) + assert w_hat.dtype == w.dtype + + wh32 = w_hat.to(torch.float32).view(m, n // 32, 32) + amax = w.to(torch.float32).abs().view(m, n // 32, 32).amax(dim=-1, keepdim=True) + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mant = bits & 0x7FFFFF + exp = exp_field - 129 + (mant > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)).clamp(-126, 125) + scale = torch.ldexp(torch.ones_like(amax), exp) + scale = torch.where(amax > 0, scale, torch.ones_like(scale)) + + vals = (wh32 / scale).abs() + assert torch.isin(vals, E2M1_GRID).all(), "values off the E2M1 grid" + nz = (amax.double() > 6.0 * 2.0**-126) & (amax.double() <= 6.0 * 2.0**125) + assert ((scale.double() * 6.0)[nz] >= amax.double()[nz]).all() + assert (((scale.double() / 2.0) * 6.0)[nz] < amax.double()[nz]).all() + ref = fake_quant_ref_fp64(w) + assert torch.equal(w_hat.to(torch.float64), ref), f"RTNE mismatch {m}x{n}" + assert torch.equal(mxfp4_fake_quantize(w_hat), w_hat), "not idempotent" + + +def test_D_lossless_in_bf16(): + for m, n in SHAPES: + w = make_weight(m, n, outliers=4) + w_hat_bf16 = mxfp4_fake_quantize(w) + w_hat_fp32 = mxfp4_fake_quantize(w.to(torch.float32)) + assert torch.equal(w_hat_bf16.to(torch.float32), w_hat_fp32), ( + "MXFP4-grid weight is not exact in bf16" + ) + + + + +def test_E_mxfp8_rowwise_lossless(): + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w[0, 0] = 2.0**-128 + w_hat = mxfp4_fake_quantize(w) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(w_hat) + dq = q.dequantize(dtype=torch.float32) + wh32 = w_hat.to(torch.float32) + + data = q._rowwise_data.view(torch.uint8)[:m, :n] + codes = q._rowwise_scale_inv.view(torch.uint8)[:m, : n // 32].to(torch.int32) + raw_vals = data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) + raw_scale = torch.ldexp(torch.ones_like(codes, dtype=torch.float32), codes - 127) + raw_dq = (raw_vals * raw_scale.unsqueeze(-1)).view(m, n) + assert torch.equal(raw_dq, wh32), ( + f"RAW MXFP8 encoding of the MXFP4-grid weight is not lossless {m}x{n}" + ) + + flushed = wh32.abs() == 2.0**-127 + assert flushed.any() and not flushed.all(), "artifact case not exercised" + assert torch.equal(dq[~flushed], wh32[~flushed]), ( + f"MXFP8 rowwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" + ) + if not (dq[flushed] == 0).all(): + assert torch.equal(dq, wh32), "dequant neither flushed nor exact?" + print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") + + +def test_G_blockwise_lossless_and_transpose(): + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4) + w_hat = mxfp4_fake_quantize(w) + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, + ).quantize(w_hat) + dq = q.dequantize(dtype=torch.float32) + assert torch.equal(dq, w_hat.to(torch.float32)), ( + f"128x128 blockwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" + ) + q.update_usage(rowwise_usage=False, columnwise_usage=True) + dqc = q.dequantize(dtype=torch.float32) + assert torch.equal(dqc, w_hat.to(torch.float32)), ( + f"columnwise 128x128 dequant differs from the MXFP4-grid weight {m}x{n}" + ) + + +def test_lossless_dequant_to_bf16(): + """The bwd override paths dequantize to activation dtype (bf16): both weight + encodings must reproduce the MXFP4-grid weight bit-exactly at bf16 too.""" + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w_hat = mxfp4_fake_quantize(w) + + q8 = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(w_hat) + dq8 = q8.dequantize(dtype=torch.bfloat16) + flushed = w_hat.to(torch.float32).abs() == 2.0**-127 + assert flushed.any(), "artifact case not exercised" + assert torch.equal(dq8[~flushed], w_hat[~flushed]), ( + f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" + ) + assert (dq8[flushed] == 0).all() + + w2_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4)) + qb = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, + force_pow_2_scales=True, block_scaling_dim=2, + ).quantize(w2_hat) + assert torch.equal(qb.dequantize(dtype=torch.bfloat16), w2_hat), ( + f"blockwise 128x128 dequant to bf16 not lossless {m}x{n}" + ) + + +def test_C_mxfp8_colwise_bound(): + for m, n in SHAPES: + w_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4, outliers=4)) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3).quantize(w_hat) + q.update_usage(rowwise_usage=False, columnwise_usage=True) + dq = q.dequantize(dtype=torch.float32) + ref = w_hat.to(torch.float32) + blk = ref.view(m // 32, 32, n) + bound = blk.abs().amax(dim=1, keepdim=True) * 2.0**-3 + err = (dq.view(m // 32, 32, n) - blk).abs() + assert (err <= bound + 1e-30).all(), f"colwise 32x1 requant error above bound {m}x{n}" + + + + + + +_INTREE_KERNELS = {} + + +def _intree_kernel(fast_math=False): + """JIT-compile the in-tree TE kernel for testing (the product path is the + tex.mxfp4_fake_quantize binding compiled with the wheel). fast_math=True + builds the SAME source with --use_fast_math to prove flag immunity.""" + if fast_math not in _INTREE_KERNELS: + import os + from torch.utils.cpp_extension import load_inline + + src_dir = os.environ.get( + "NVTE_MXFP4_QAT_TEST_SRC", + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", "..", "transformer_engine", "common", + ), + ) + cuda_src = r""" +#include +#include +#include +#include "cast/mxfp4/fake_quantize_mxfp4.cuh" + +torch::Tensor mxfp4_fake_quantize_intree(torch::Tensor w) { + TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.numel() % 32 == 0); + auto out = torch::empty_like(w); + auto stream = at::cuda::getCurrentCUDAStream(); + if (w.scalar_type() == at::kBFloat16) { + transformer_engine::dispatch::mxfp4::fake_quantize_mxfp4_launch( + reinterpret_cast(w.data_ptr()), + reinterpret_cast(out.data_ptr()), w.numel(), stream); + } else if (w.scalar_type() == at::kFloat) { + transformer_engine::dispatch::mxfp4::fake_quantize_mxfp4_launch( + w.data_ptr(), out.data_ptr(), w.numel(), stream); + } else { + TORCH_CHECK(false, "bf16/fp32 only"); + } + return out; +} +""" + _INTREE_KERNELS[fast_math] = load_inline( + name="nvte_mxfp4_qat_intree_test" + ("_fastmath" if fast_math else ""), + cpp_sources="torch::Tensor mxfp4_fake_quantize_intree(torch::Tensor w);", + cuda_sources=cuda_src, + functions=["mxfp4_fake_quantize_intree"], + extra_include_paths=[src_dir], + extra_cuda_cflags=[ + "-O3", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + ] + (["--use_fast_math"] if fast_math else []), + verbose=False, + ) + return _INTREE_KERNELS[fast_math] + + +def _both_paths(w): + import transformer_engine.pytorch.mxfp4_qat as M + got_kernel = _intree_kernel().mxfp4_fake_quantize_intree(w.contiguous()) + got_torch = M._mxfp4_fake_quantize_torch(w) + return got_kernel, got_torch + + +def _assert_bits_equal(a, b, msg): + assert a.dtype == b.dtype, f"{msg}: dtype mismatch" + ib = torch.int16 if a.dtype == torch.bfloat16 else torch.int32 + assert torch.equal(a.view(ib), b.view(ib)), f"{msg}: raw bits differ" + + +def _assert_same_with_nan(a, b, msg): + an, bn = torch.isnan(a), torch.isnan(b) + assert torch.equal(an, bn), f"{msg}: NaN masks differ" + _assert_bits_equal( + torch.where(an, torch.zeros_like(a), a), + torch.where(bn, torch.zeros_like(b), b), + msg, + ) + + +PARITY_SHAPES = SHAPES + [(3, 416), (1, 32)] + + +def test_kernel_matches_torch_reference(): + for m, n in PARITY_SHAPES: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) + gk, gt = _both_paths(w) + _assert_bits_equal(gk, gt, f"kernel != torch reference {m}x{n} {dtype}") + + +def test_edge_value_domains(): + m, n = 64, 128 + w = make_weight(m, n) + + z = torch.zeros(m, n, device=DEV, dtype=torch.bfloat16) + gk, gt = _both_paths(z) + assert torch.equal(gk, z) and torch.equal(gt, z) + + wz = w.clone(); wz[0, :32] = -0.0 + gk, gt = _both_paths(wz) + assert (gk[0, :32] == 0).all() and torch.signbit(gk[0, :32]).all(), "-0.0 sign lost" + _assert_bits_equal(gk, gt, "-0 block") + + ws = w.clone(); ws[1, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + gk, gt = _both_paths(ws) + assert torch.isfinite(gk[1, :32]).all() and torch.equal(gk, gt) + ref = fake_quant_ref_fp64(ws) + assert torch.equal(gt.to(torch.float64), ref), "subnormal block deviates from fp64 ref" + + wp = w.clone(); wp[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + gk, gt = _both_paths(wp) + assert (gk[5, :32].to(torch.float32) == 2.0**-127).all(), "2^-127 grid point lost" + _assert_bits_equal(gk, gt, "2^-127 grid point") + ref = fake_quant_ref_fp64(wp) + assert torch.equal(gt.to(torch.float64), ref), "2^-127 block deviates from fp64 ref" + + wh = w.clone(); wh[2, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) + wh[2, 0] = -3.0e38 + gk, gt = _both_paths(wh) + assert torch.isfinite(gk[2, :32]).all() and torch.equal(gk, gt) + ref = fake_quant_ref_fp64(wh) + assert torch.equal(gt.to(torch.float64), ref), "huge block deviates from fp64 ref" + + for bad in (float("inf"), float("-inf"), float("nan")): + wb = w.clone(); wb[3, 40] = bad + gk, gt = _both_paths(wb) + _assert_same_with_nan(gk, gt, f"bad={bad}") + assert torch.isnan(gk[3, 32:64]).all(), f"{bad}: block not NaN-poisoned" + assert torch.isfinite(gk[3, :32]).all(), f"{bad}: neighbor block affected" + clean_rows = torch.ones(m, dtype=torch.bool); clean_rows[3] = False + assert torch.isfinite(gk[clean_rows]).all() + + try: + mxfp4_fake_quantize(torch.zeros(32, 32, device=DEV, dtype=torch.float16)) + raise SystemExit("fp16 should be rejected") + except ValueError: + pass + + +def test_blockwise_recipe_is_128x128(): + r = MXFP4QATFloat8BlockScaling() + assert r.x_block_scaling_dim == 1 and r.w_block_scaling_dim == 2 + assert r.grad_block_scaling_dim == 1 + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, + force_pow_2_scales=True, block_scaling_dim=2, + ) + assert q.block_len == 128 + w_hat = mxfp4_fake_quantize(make_weight(256, 512)) + qt = q.quantize(w_hat) + assert tuple(qt._rowwise_scale_inv.shape)[0] == 256 // 128 + + +def test_blockwise_over_ratio_is_bounded_not_lossless(): + w = make_weight(128, 128) + w[0, :32] = torch.tensor(2.0**-14, dtype=torch.bfloat16) + w[64, 0] = 100.0 + w_hat = mxfp4_fake_quantize(w) + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, + force_pow_2_scales=True, block_scaling_dim=2, + ).quantize(w_hat) + dq = q.dequantize(dtype=torch.float32) + assert torch.isfinite(dq).all() + err = (dq - w_hat.to(torch.float32)).abs().max().item() + tile_amax = w_hat.to(torch.float32).abs().max().item() + assert err <= tile_amax * 2.0**-9, "over-ratio loss larger than FP8 headroom" + + +def test_kernel_perf(): + if not RUN_PERF: + print(" SKIP (set NVTE_MXFP4_QAT_PERF=1 to run benchmarks)") + return + import transformer_engine.pytorch.mxfp4_qat as M + + rows, cols = 8192, 8192 + w = make_weight(rows, cols) + k = _intree_kernel() + for _ in range(3): + k.mxfp4_fake_quantize_intree(w) + M._mxfp4_fake_quantize_torch(w) + torch.cuda.synchronize() + + def timeit(fn, iters=20): + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters * 1e3 + + t_kernel = timeit(lambda: k.mxfp4_fake_quantize_intree(w)) + t_torch = timeit(lambda: M._mxfp4_fake_quantize_torch(w)) + moved = rows * cols * 2 * 2 + print( + f" kernel: {t_kernel:.1f} us ({moved / t_kernel / 1e3:.0f} GB/s), " + f"torch: {t_torch:.1f} us ({moved / t_torch / 1e3:.0f} GB/s)" + ) + assert t_kernel < t_torch, "kernel slower than the torch reference" + + +def test_cute_dsl_version(): + try: + from transformer_engine.pytorch.mxfp4_qat_cute_dsl import mxfp4_fake_quantize_cute_dsl + except Exception as e: + print(f" SKIP (cutlass DSL unavailable: {type(e).__name__}: {e})") + return + import transformer_engine.pytorch.mxfp4_qat as M + + for m, n in SHAPES: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) + gd = mxfp4_fake_quantize_cute_dsl(w) + gt = M._mxfp4_fake_quantize_torch(w) + _assert_bits_equal(gd, gt, f"cute-dsl != torch ref {m}x{n} {dtype}") + + w = make_weight(64, 128) + w[0, :32] = 0.0 + w[1, 5] = float("inf") + w[2, 9] = float("nan") + w[3, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + w[4, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) + w[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + gd = mxfp4_fake_quantize_cute_dsl(w) + gt = M._mxfp4_fake_quantize_torch(w) + _assert_same_with_nan(gd, gt, "cute-dsl edge") + assert (gd[5, :32].to(torch.float32) == 2.0**-127).all(), "cute-dsl: 2^-127 lost" + + if not RUN_PERF: + print(" parity OK; SKIP perf (set NVTE_MXFP4_QAT_PERF=1)") + return + + rows, cols = 8192, 8192 + wb = make_weight(rows, cols) + k = _intree_kernel() + for _ in range(3): + mxfp4_fake_quantize_cute_dsl(wb) + k.mxfp4_fake_quantize_intree(wb) + torch.cuda.synchronize() + + def timeit(fn, iters=20): + start, end = torch.cuda.Event(True), torch.cuda.Event(True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters * 1e3 + + t_dsl = timeit(lambda: mxfp4_fake_quantize_cute_dsl(wb)) + t_cuda = timeit(lambda: k.mxfp4_fake_quantize_intree(wb)) + moved = rows * cols * 2 * 2 + print( + f" cute-dsl: {t_dsl:.1f} us ({moved / t_dsl / 1e3:.0f} GB/s), " + f"cuda kernel: {t_cuda:.1f} us ({moved / t_cuda / 1e3:.0f} GB/s)" + ) + + +def test_kernel_fast_math_immune(): + """The kernel must stay bit-identical when compiled with --use_fast_math. + + Guards the hazard class from the bitwise.md PR list (flashinfer#3387 etc.): + approximate division / ftz under fast-math flags silently changing + quantization results. The kernel uses only integer exponent math, exact + power-of-two multiplies, and rintf — no fp division, no MUFU. + """ + import transformer_engine.pytorch.mxfp4_qat as M + + k = _intree_kernel(fast_math=True) + for m, n in PARITY_SHAPES: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) + got = k.mxfp4_fake_quantize_intree(w.contiguous()) + ref = M._mxfp4_fake_quantize_torch(w) + _assert_bits_equal(got, ref, f"fast-math build diverged {m}x{n} {dtype}") + + w = make_weight(64, 128) + w[0, :32] = 0.0 + w[1, 5] = float("inf") + w[2, 9] = float("nan") + w[3, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + w[4, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) + w[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + got = k.mxfp4_fake_quantize_intree(w.contiguous()) + ref = M._mxfp4_fake_quantize_torch(w) + _assert_same_with_nan(got, ref, "fast-math edge") + assert (got[5, :32].to(torch.float32) == 2.0**-127).all(), ( + "fast-math build flushed the 2^-127 grid point" + ) + + + +def test_exp2f_e8m0_all_codes(): + """UE8M0 scale expansion: code 0 is 2^-127 (fp32 subnormal), code 255 NaN. + + Compiles the exp2f/exp2f_rcp FUNCTION TEXT extracted verbatim from + common/util/ptx.cuh (the mxfp8 dequant kernels' scale expansion) and checks + all 256 codes. Guards the code-0 fix: the old bit-shift construction gave + +0.0 and zeroed whole blocks on software dequant.""" + import re as _re + from torch.utils.cpp_extension import load_inline + + src_dir = os.environ.get( + "NVTE_MXFP4_QAT_TEST_SRC", + os.path.join(os.path.dirname(os.path.abspath(__file__)), + "..", "..", "transformer_engine", "common"), + ) + ptx_src = open(os.path.join(src_dir, "util", "ptx.cuh")).read() + m_exp = _re.search(r"__device__ __forceinline__ float exp2f\(e8m0_t biased_exp\) \{.*?\n\}", ptx_src, _re.S) + m_rcp = _re.search(r"template <>\n__device__ __forceinline__ float exp2f_rcp\(e8m0_t biased_exp\) \{.*?\n\}", ptx_src, _re.S) + assert m_exp and m_rcp, "could not extract exp2f/exp2f_rcp from ptx.cuh" + cuda_src = ( + "#include \n#include \n" + "#include \n" + "namespace teptx {\n" + "typedef unsigned char e8m0_t;\n" + "constexpr unsigned FP32_MANTISSA_BITS = 23;\n" + "template __device__ __forceinline__ T exp2f_rcp(e8m0_t);\n" + + m_rcp.group(0) + "\n" + m_exp.group(0) + "\n}\n" + "__global__ void all_codes_kernel(float *e, float *r) {\n" + " const int b = threadIdx.x;\n" + " e[b] = teptx::exp2f(static_cast(b));\n" + " r[b] = teptx::exp2f_rcp(static_cast(b));\n" + "}\n" + "void run_all_codes(torch::Tensor e, torch::Tensor r) {\n" + " all_codes_kernel<<<1, 256, 0, at::cuda::getCurrentCUDAStream()>>>(\n" + " e.data_ptr(), r.data_ptr());\n" + "}\n" + ) + mod = load_inline( + name="nvte_mxfp4_qat_exp2f_test", + cpp_sources="void run_all_codes(torch::Tensor e, torch::Tensor r);", + cuda_sources=cuda_src, functions=["run_all_codes"], verbose=False, + ) + e = torch.empty(256, dtype=torch.float32, device=DEV) + r = torch.empty_like(e) + mod.run_all_codes(e, r) + torch.cuda.synchronize() + codes = torch.arange(256, dtype=torch.int32) + exp_ref = torch.ldexp(torch.ones(256), codes - 127) + rcp_ref = torch.ldexp(torch.ones(256), 127 - codes) + assert torch.isnan(e[255]) and torch.isnan(r[255]), "code 255 must be NaN" + assert torch.equal(e[:255].cpu(), exp_ref[:255]), "exp2f wrong (code 0 == 2^-127?)" + assert torch.equal(r[:255].cpu(), rcp_ref[:255]), "exp2f_rcp wrong" + + +def test_ste_identity_gradient(): + """Public mxfp4_fake_quantize must carry an identity (STE) gradient on both + implementation paths — the QAT semantics for d(w_hat)/dw.""" + import transformer_engine.pytorch.mxfp4_qat as M + + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(64, 128, dtype=dtype).requires_grad_(True) + up = torch.randn_like(w) + out = mxfp4_fake_quantize(w) + out.backward(up) + assert torch.equal(w.grad, up), f"STE gradient is not identity ({dtype})" + + w2 = w.detach().clone().requires_grad_(True) + old = M._torch_path_forced + M._torch_path_forced = True + try: + out2 = mxfp4_fake_quantize(w2) + out2.backward(up) + finally: + M._torch_path_forced = old + assert torch.equal(w2.grad, up), f"torch-path STE gradient is not identity ({dtype})" + _assert_bits_equal(out.detach(), out2.detach(), "STE fwd parity") + + +def test_misaligned_and_noncontiguous(): + """Storage-offset views (contiguous but not 16B-aligned) and non-contiguous + inputs must produce the same bits as an aligned contiguous copy.""" + import transformer_engine.pytorch.mxfp4_qat as M + + for dtype, off in ((torch.bfloat16, 4), (torch.float32, 2)): + m, n = 64, 128 + flat = (torch.randn(m * n + off, device=DEV, dtype=torch.float32) * 0.02).to(dtype) + v = flat[off:].view(m, n) + assert v.is_contiguous() and v.data_ptr() % 16 != 0 + got = mxfp4_fake_quantize(v) + ref = mxfp4_fake_quantize(v.clone()) + _assert_bits_equal(got, ref, f"misaligned view {dtype}") + + nc = (torch.randn(64, 256, device=DEV, dtype=torch.float32) * 0.02)[:, ::2] + assert not nc.is_contiguous() + got = M._mxfp4_fake_quantize_torch(nc) + assert torch.equal(got, M._mxfp4_fake_quantize_torch(nc.contiguous())) + + +def test_bf16_exhaustive_bit_patterns(): + """All 65536 bf16 bit patterns (as 2048 blocks of 32): kernel/torch bit + parity everywhere, fp64 oracle on the all-finite blocks.""" + bits = torch.arange(65536, dtype=torch.int32, device=DEV).to(torch.int16) + w = bits.view(torch.bfloat16).view(2048, 32) + gk, gt = _both_paths(w) + _assert_same_with_nan(gk, gt, "bf16 exhaustive") + finite = torch.isfinite(w.to(torch.float32)).all(dim=-1) + ref = fake_quant_ref_fp64(w[finite]) + assert torch.equal(gt[finite].to(torch.float64), ref), "bf16 exhaustive vs fp64 ref" + + +def test_fp32_bit_fuzz(): + """Random raw fp32 bit patterns (subnormals/NaN/Inf included).""" + g = torch.Generator().manual_seed(99) + for _ in range(4): + bits = torch.randint(-2**31, 2**31 - 1, (4096, 32), generator=g, dtype=torch.int64) + w = bits.to(torch.int32).cuda().view(torch.float32) + gk, gt = _both_paths(w) + _assert_same_with_nan(gk, gt, "fp32 fuzz") + finite = torch.isfinite(w).all(dim=-1) + if finite.any(): + ref = fake_quant_ref_fp64(w[finite]) + assert torch.equal(gt[finite].to(torch.float64), ref), "fp32 fuzz vs fp64 ref" + + +def test_scale_threshold_and_rtne_midpoints(): + """amax = 1.5*2^k boundary (one ulp below / exact / one ulp above) and every + RTNE midpoint of the E2M1 grid, against explicit expected values.""" + rows = [] + for k in (-100, -10, 0, 42, 100): + a = torch.tensor(1.5 * 2.0**k, dtype=torch.float32) + ab = a.view(torch.int32) + for delta in (-1, 0, 1): + row = torch.zeros(32, dtype=torch.float32) + row[0] = (ab + delta).view(torch.float32) + rows.append(row) + w = torch.stack(rows).cuda() + gk, gt = _both_paths(w) + _assert_bits_equal(gk, gt, "threshold parity") + assert torch.equal(gt.to(torch.float64), fake_quant_ref_fp64(w)), "threshold vs fp64 ref" + + mids = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]) + expected = torch.tensor([0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0]) + for t in (-60, 0, 60): + sc = 2.0**t + block = torch.zeros(32); block[:7] = mids * sc; block[7] = 6.0 * sc + w = block.view(1, 32).to(torch.float32).cuda() + gk, gt = _both_paths(w) + _assert_bits_equal(gk, gt, f"midpoints t={t}") + exp_row = torch.zeros(32); exp_row[:7] = expected * sc; exp_row[7] = 6.0 * sc + assert torch.equal(gt[0].cpu().to(torch.float32), exp_row.to(torch.float32)), f"midpoint values t={t}" + + +def test_deployment_top_cap_domain(): + """Top-of-range contract: for amax <= 6*2^125 the fake-quant is bitwise the + TileKernels deployment result. Above that, TileKernels moves to scale 2^126 + (packed payload + UE8M0 byte 253 -- representable in packed form only, since + {4,6}*2^126 overflow bf16/fp32), while this path pins satfinite at the cap + 2^125. Restricted-domain divergence, documented and unreachable for real + weights (requires |w| > 2.55e38).""" + s125 = 6.0 * 2.0**125 + w = torch.zeros(2, 32, dtype=torch.float32) + w[0, 0] = s125 + just_above = (torch.tensor(s125, dtype=torch.float32).view(torch.int32) + 1).view(torch.float32) + w[1, 0] = just_above + wc = w.cuda() + gk, gt = _both_paths(wc) + _assert_bits_equal(gk, gt, "top-cap parity") + assert gk[0, 0].item() == s125, "boundary block must be exact" + assert gk[1, 0].item() == s125, "above-boundary block must satfinite at 6*2^125" + assert torch.equal(gt.to(torch.float64), fake_quant_ref_fp64(wc)), "top-cap vs fp64 ref" + + +def test_tilekernels_cross_parity(): + """Cross-backend golden: TileKernels' torch reference (deployment semantics, + (1,32) e2m1 cast -> dequant) must reproduce mxfp4_fake_quantize bit-for-bit + on the finite domain with amax <= 6*2^125.""" + import sys + tk_root = os.environ.get("NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels")) + if os.path.isdir(tk_root) and tk_root not in sys.path: + sys.path.insert(0, tk_root) + try: + from tile_kernels.torch.cast import cast as tk_cast, cast_back as tk_cast_back + except Exception as e: + print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") + return + + torch.manual_seed(21) + for m, n in SHAPES: + w = make_weight(m, n, zero_blocks=4, outliers=4) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w[1, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + w[2, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) + data, sf = tk_cast(w, "e2m1", block_size=(1, 32), round_sf=True) + tk_dq = tk_cast_back((data, sf), "bf16", block_size=(1, 32)) + ours = mxfp4_fake_quantize(w) + _assert_bits_equal(tk_dq, ours, f"TileKernels cross parity {m}x{n}") + + +def test_blockwise_feasibility_enumeration(): + """128x128 exactness contract, enumerated: for exponent spread d = 0..16 and + every signed E2M1 payload, predict from the E4M3 lattice (subnormals + included) whether p*2^(t-d) survives the tile encode at the tile's actual + scale, and check TE's quantizer agrees element-for-element.""" + payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + boundary = None + for d in range(17): + w = torch.zeros(128, 128, dtype=torch.float32) + w[0, 0] = 6.0 + vals = payloads * 2.0**(-d) + w[1, :7] = vals + w[1, 7:14] = -vals + wc = w.cuda() + assert torch.equal(mxfp4_fake_quantize(wc), wc), f"tile not on MXFP4 grid (d={d})" + q = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False, + force_pow_2_scales=True, block_scaling_dim=2, + ).quantize(wc) + S = q._rowwise_scale_inv.reshape(-1)[0].item() + dq = q.dequantize(dtype=torch.float32) + pred = ((wc / S).to(torch.float8_e4m3fn).to(torch.float32) * S) == wc + got = dq == wc + assert torch.equal(pred, got), f"E4M3-lattice prediction mismatch at spread d={d}" + if boundary is None and not bool(pred.all()): + boundary = d + assert boundary is not None and boundary >= 12, ( + f"full-payload exactness should hold at least through spread 2^11, got {boundary}" + ) + print(f" full-grid exact through spread d={boundary - 1}; first loss at d={boundary} " + "(E4M3 subnormal lattice bound)") + + +def test_fourway_matrix(): + """TileKernels deployment reference vs CuTe DSL vs CUDA kernel vs torch + reference, bitwise, over every previously-hit edge domain and every + quant/dequant step: fake-quant (bf16+fp32), mxfp8 row encode/dequant, + mxfp8 col encode/dequant, blockwise 128x128 encode/dequant, bf16+fp32 + dequant targets, bf16-exhaustive and fp32 bit-fuzz sweeps.""" + import sys + tk_root = os.environ.get( + "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") + ) + if os.path.isdir(tk_root) and tk_root not in sys.path: + sys.path.insert(0, tk_root) + try: + from tile_kernels.torch.cast import ( + cast as tk_cast, + cast_back as tk_cast_back, + get_min_clamp_val as tk_min_clamp, + ) + except Exception as e: + print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") + return + try: + from transformer_engine.pytorch.mxfp4_qat_cute_dsl import ( + mxfp4_fake_quantize_cute_dsl as dsl, + ) + except Exception as e: + print(f" SKIP (cute dsl unavailable: {type(e).__name__}: {e})") + return + import transformer_engine.pytorch.mxfp4_qat as M + + kern = _intree_kernel().mxfp4_fake_quantize_intree + tref = M._mxfp4_fake_quantize_torch + CAP = 6.0 * 2.0**125 + + def tk_roundtrip(x, block, fmt): + pair = tk_cast(x, "e2m1" if fmt == "e2m1" else "e4m3", block_size=block, round_sf=True) + return tk_cast_back(pair, "fp32", block), tk_cast_back(pair, "bf16", block) + + torch.manual_seed(77) + w = make_weight(64, 256, zero_blocks=2, outliers=4) + w[0, :32] = 0.0 + w[1, :32] = -0.0 + w[2, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w[2, 0] = 2.0**-128 + w[3, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + w[4, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) + w[5, :7] = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], dtype=torch.bfloat16) + w[5, 7] = 6.0 + a15 = torch.tensor(1.5, dtype=torch.bfloat16).view(torch.int16) + w[6, 0] = (a15 - 1).view(torch.bfloat16) + w[6, 32] = torch.tensor(1.5, dtype=torch.bfloat16) + w[6, 64] = (a15 + 1).view(torch.bfloat16) + + for W in (w, w.to(torch.float32)): + t = tref(W) + _assert_bits_equal(kern(W.contiguous()), t, "fake-quant CUDA vs torch") + _assert_bits_equal(dsl(W), t, "fake-quant DSL vs torch") + tk32, tkb16 = tk_roundtrip(W, (1, 32), "e2m1") + _assert_bits_equal(tk32, t.to(torch.float32), "fake-quant TK vs torch (fp32)") + _assert_bits_equal(tkb16, t.to(torch.bfloat16), "fake-quant TK vs torch (bf16)") + print(" fake-quant 4-way bitwise: PASS (bf16 + fp32 corpus)") + + wnf = w.clone() + wnf[8, 40] = float("inf") + wnf[9, 33] = float("-inf") + wnf[10, 5] = float("nan") + t = tref(wnf) + _assert_same_with_nan(kern(wnf.contiguous()), t, "nonfinite CUDA") + _assert_same_with_nan(dsl(wnf), t, "nonfinite DSL") + assert torch.isnan(t[8, 32:64]).all() and torch.isnan(t[9, 32:64]).all() + assert torch.isnan(t[10, :32]).all() and torch.isfinite(t[11].to(torch.float32)).all() + print(" nonfinite 3-way (whole-block NaN poisoning): PASS (TK policy differs, excluded)") + + wcap = torch.zeros(1, 32, dtype=torch.bfloat16, device=DEV) + wcap[0, 0] = torch.tensor(3.2e38, dtype=torch.bfloat16) + t = tref(wcap) + assert t[0, 0].item() == CAP, "satfinite cap at 6*2^125" + _assert_bits_equal(kern(wcap.contiguous()), t, "cap CUDA") + _assert_bits_equal(dsl(wcap), t, "cap DSL") + tk32, _ = tk_roundtrip(wcap, (1, 32), "e2m1") + assert not torch.equal(tk32, t.to(torch.float32)), "TK must diverge above 6*2^125" + print(f" above-cap domain: ours satfinite {CAP:.3e}, TK diverges as documented " + f"(payload at 2^126 -> {tk32[0, 0].item()}): PASS") + + w_hat = tref(w) + wh32 = w_hat.to(torch.float32) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=True).quantize(w_hat) + m, n = w_hat.shape + data = q._rowwise_data.view(torch.uint8)[:m, :n] + codes = q._rowwise_scale_inv.view(torch.uint8)[:m, : n // 32].to(torch.int32) + raw = (data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) + * torch.ldexp(torch.ones_like(codes, dtype=torch.float32), codes - 127).unsqueeze(-1)) + assert torch.equal(raw.view(m, n), wh32), "mxfp8 row RAW encode not lossless" + e4m3_floor = float(tk_min_clamp(torch.float8_e4m3fn)) + blk_amax = wh32.abs().view(m, n // 32, 32).amax(dim=-1) + ok = ((blk_amax >= e4m3_floor) | (blk_amax == 0)) + okm = ok.unsqueeze(-1).expand(m, n // 32, 32).reshape(m, n) + tk32, tkb16 = tk_roundtrip(w_hat, (1, 32), "e4m3") + assert torch.equal(tk32.view(torch.int32)[okm], wh32.view(torch.int32)[okm]), ( + "mxfp8 row TK roundtrip (fp32)" + ) + assert torch.equal(tkb16.view(torch.int16)[okm], w_hat.view(torch.int16)[okm]), ( + "mxfp8 row TK roundtrip (bf16)" + ) + assert (tk32[~okm] == 0).all(), "TK e4m3 sub-floor blocks must flush to zero" + assert (~ok).sum() > 0, "TK e4m3 floor case not exercised" + dq = q.dequantize(dtype=torch.float32) + flushed = wh32.abs() == 2.0**-127 + assert torch.equal(dq[~flushed], wh32[~flushed]), "mxfp8 row TE software dequant" + if flushed.any() and not (dq[flushed] == 0).all(): + assert torch.equal(dq, wh32) + print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") + print(f" mxfp8 row: TE raw encode lossless incl 2^-127; TK roundtrip bitwise on " + f"amax >= {e4m3_floor:g} blocks ({int((~ok).sum())} sub-floor blocks flush in TK " + "e4m3, a TK-only floor); TE software dequant exact outside the code-0 wheel bug: PASS") + + q.update_usage(rowwise_usage=False, columnwise_usage=True) + te_col = q.dequantize(dtype=torch.float32) + tk_col32, _ = tk_roundtrip(w_hat, (32, 1), "e4m3") + col_amax = wh32.abs().view(m // 32, 32, n).amax(dim=1) + cok = ((col_amax >= e4m3_floor) | (col_amax == 0)) + cokm = cok.unsqueeze(1).expand(m // 32, 32, n).reshape(m, n) + nzm = (wh32 != 0) & cokm + zm = (wh32 == 0) & cokm + assert torch.equal(te_col.view(torch.int32)[nzm], tk_col32.view(torch.int32)[nzm]), ( + "mxfp8 col TE vs TK (nonzero)" + ) + assert (te_col[zm] == 0).all() and (tk_col32[zm] == 0).all(), "mxfp8 col zeros" + negz = zm & torch.signbit(wh32) + assert negz.any() + assert torch.signbit(tk_col32)[negz].all(), "TK col must preserve -0" + assert not torch.signbit(te_col)[negz].any(), "TE col canonicalizes -0 -> +0" + blk = wh32.view(m // 32, 32, n) + bound = blk.abs().amax(dim=1, keepdim=True) * 2.0**-3 + assert ((te_col.view(m // 32, 32, n) - blk).abs() <= bound + 1e-30).all() + print(" mxfp8 col (32x1): TE == TK bitwise on nonzeros, zeros value-equal " + "(TE col canonicalizes -0 -> +0, TK preserves the sign), bounded vs w_hat: PASS") + + wb_hat = tref(make_weight(128, 256, zero_blocks=4)) + wb32 = wb_hat.to(torch.float32) + qb = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, + force_pow_2_scales=True, block_scaling_dim=2, + ).quantize(wb_hat) + assert torch.equal(qb.dequantize(dtype=torch.float32), wb32) + assert torch.equal(qb.dequantize(dtype=torch.bfloat16), wb_hat) + tk32, tkb16 = tk_roundtrip(wb_hat, (128, 128), "e4m3") + _assert_bits_equal(tk32, wb32, "blockwise TK roundtrip (fp32)") + _assert_bits_equal(tkb16, wb_hat, "blockwise TK roundtrip (bf16)") + qb.update_usage(rowwise_usage=False, columnwise_usage=True) + assert torch.equal(qb.dequantize(dtype=torch.float32), wb32), "blockwise col transpose" + print(" blockwise 128x128: TE row+col dequant == TK roundtrip == w_hat bitwise " + "(fp32 + bf16 targets): PASS") + + bits = torch.arange(65536, dtype=torch.int32, device=DEV).to(torch.int16) + we = bits.view(torch.bfloat16).view(2048, 32) + t = tref(we) + _assert_same_with_nan(kern(we.contiguous()), t, "exhaustive CUDA") + _assert_same_with_nan(dsl(we), t, "exhaustive DSL") + we32 = we.to(torch.float32) + ok_rows = torch.isfinite(we32).all(dim=-1) & (we32.abs().amax(dim=-1) <= CAP) + tk32, _ = tk_roundtrip(we[ok_rows], (1, 32), "e2m1") + _assert_bits_equal(tk32, t[ok_rows].to(torch.float32), "exhaustive TK") + print(f" bf16 exhaustive 65536 patterns: CUDA/DSL/torch 3-way full, " + f"TK on {int(ok_rows.sum())}/2048 finite below-cap rows, bitwise: PASS") + + g = torch.Generator().manual_seed(123) + fb = torch.randint(-2**31, 2**31 - 1, (512, 32), generator=g, dtype=torch.int64) + wf = fb.to(torch.int32).cuda().view(torch.float32) + t = tref(wf) + _assert_same_with_nan(kern(wf.contiguous()), t, "fuzz CUDA") + _assert_same_with_nan(dsl(wf), t, "fuzz DSL") + ok_rows = torch.isfinite(wf).all(dim=-1) & (wf.abs().amax(dim=-1) <= CAP) + if ok_rows.any(): + tk32, _ = tk_roundtrip(wf[ok_rows], (1, 32), "e2m1") + _assert_bits_equal(tk32, t[ok_rows], "fuzz TK") + print(f" fp32 bit-fuzz 512 blocks: 3-way full, TK on {int(ok_rows.sum())} " + f"finite below-cap rows, bitwise: PASS") + + +def test_recipe_switch_invalidates_cache(): + """base <-> QAT recipe switch must invalidate cached weight workspaces: a + non-first microbatch right after the switch has to requantize under the new + semantics instead of silently reusing the other recipe's weight.""" + import warnings as _warnings + from transformer_engine.common.recipe import MXFP8BlockScaling + + torch.manual_seed(3) + x = torch.randn(64, 256, device=DEV, dtype=torch.bfloat16) + mod = te.Linear(256, 128, bias=False, params_dtype=torch.bfloat16, device=DEV) + base, qat = MXFP8BlockScaling(), MXFP4QATMXFP8BlockScaling() + + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + with fp8_autocast(enabled=True, fp8_recipe=base): + mod(x, is_first_microbatch=True) + with fp8_autocast(enabled=True, fp8_recipe=qat): + y_switch = mod(x, is_first_microbatch=False) + y_fresh = mod(x, is_first_microbatch=True) + assert torch.equal(y_switch, y_fresh), ( + "base->QAT switch reused an UNPROJECTED cached weight" + ) + with fp8_autocast(enabled=True, fp8_recipe=base): + y2_switch = mod(x, is_first_microbatch=False) + y2_fresh = mod(x, is_first_microbatch=True) + assert torch.equal(y2_switch, y2_fresh), ( + "QAT->base switch reused a PROJECTED cached weight" + ) + + +def test_ops_api_rejected(): + """The operations API bypasses the quantize_weight QAT hook, so it must + reject MXFP4 QAT recipes loudly instead of silently training without the + projection.""" + import transformer_engine.pytorch.ops as te_ops + + op = te_ops.Linear(256, 128, bias=False, device=DEV, dtype=torch.bfloat16) + x = torch.randn(64, 256, device=DEV, dtype=torch.bfloat16) + for recipe in (MXFP4QATMXFP8BlockScaling(), MXFP4QATFloat8BlockScaling()): + try: + with fp8_autocast(enabled=True, fp8_recipe=recipe): + op(x) + raise SystemExit(f"ops API must reject {recipe.__class__.__name__}") + except NotImplementedError: + pass + + + + +def _run_module(module_kind, recipe, override, fuse): + torch.manual_seed(7) + m_splits = [192, 320] + tokens, in_f, out_f = sum(m_splits), 256, 512 + recipe.backward_override = override + + if module_kind == "linear": + mod = te.Linear( + in_f, out_f, bias=False, params_dtype=torch.bfloat16, device=DEV, + fuse_wgrad_accumulation=fuse, + ) + params = [mod.weight] + else: + mod = te.GroupedLinear( + 2, in_f, out_f, bias=False, params_dtype=torch.bfloat16, device=DEV, + fuse_wgrad_accumulation=fuse, + ) + params = [mod.weight0, mod.weight1] + with torch.no_grad(): + for p in params: + p.copy_((torch.randn_like(p, dtype=torch.float32) * 0.02).to(p.dtype)) + w_hats = [mxfp4_fake_quantize(p.detach()) for p in params] + if fuse: + for p in params: + p.main_grad = torch.zeros_like(p, dtype=torch.float32) + p.grad_added_to_main_grad = False + + x = (torch.randn(tokens, in_f, device=DEV, dtype=torch.float32) * 0.5).to( + torch.bfloat16 + ) + x.requires_grad_(True) + with fp8_autocast(enabled=True, fp8_recipe=recipe): + if module_kind == "linear": + out = mod(x) + else: + out = mod(x, m_splits) + + offs = [0, m_splits[0], tokens] if module_kind == "grouped" else [0, tokens] + nsplit = len(params) + + def cat_ref(weights, mat, transpose): + return torch.cat( + [ + mat[offs[i] : offs[i + 1]].to(torch.float32) + @ (weights[i].detach().to(torch.float32).t() if transpose + else weights[i].detach().to(torch.float32)) + for i in range(nsplit) + ] + ) + + ref_out_hat = cat_ref(w_hats, x, True) + ref_out_master = cat_ref(params, x, True) + e_hat, e_master = rel_err(out, ref_out_hat), rel_err(out, ref_out_master) + assert e_hat < FP8_TOL_1OP, f"fwd err {e_hat}" + assert e_hat < e_master, "fwd closer to master than to the MXFP4-grid weight" + + g = (torch.randn_like(out, dtype=torch.float32) * 0.1).to(out.dtype) + out.backward(g) + + ref_dx_hat = cat_ref(w_hats, g, False) + ref_dx_master = cat_ref(params, g, False) + e_h, e_m = rel_err(x.grad, ref_dx_hat), rel_err(x.grad, ref_dx_master) + if override == "high_precision": + assert e_m < BF16_TOL, f"hp dgrad err vs master {e_m}" + assert e_m < e_h, "hp dgrad must use the unquantized master (B)" + elif override == "dequantized": + assert e_h < BF16_TOL, f"dequantized dgrad err vs w_hat {e_h}" + assert e_h < e_m, "dequantized dgrad must use the MXFP4-grid weight (A)" + else: + assert e_h < FP8_TOL_1OP, f"mode-None dgrad err {e_h}" + assert e_h < e_m, "mode-None dgrad must use the MXFP4-grid weight (C)" + + if module_kind == "linear": + with fp8_autocast(enabled=True, fp8_recipe=recipe): + out_c1 = mod(x.detach(), is_first_microbatch=True) + out_c2 = mod(x.detach(), is_first_microbatch=False) + assert rel_err(out_c2, ref_out_hat) < FP8_TOL_1OP, "cached-weight fwd drifted" + assert rel_err(out_c1, out_c2) < 1e-6, "workspace update changed the weights" + + wgrad_tol = BF16_TOL if override == "high_precision" else FP8_TOL_2OP + for i, p in enumerate(params): + ref_wg = ( + g[offs[i] : offs[i + 1]].to(torch.float32).t() + @ x[offs[i] : offs[i + 1]].to(torch.float32) + ) + if fuse: + assert p.grad_added_to_main_grad is True, "fused-wgrad flag not set (G)" + e = rel_err(p.main_grad, ref_wg) + assert e < wgrad_tol, f"main_grad err {e} (G)" + else: + assert p.grad is not None, "STE did not deliver a weight grad (G)" + e = rel_err(p.grad, ref_wg) + assert e < wgrad_tol, f"wgrad err {e} (G)" + + +def test_e2e_matrix(): + for module_kind in ("linear", "grouped"): + for recipe_cls in (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling): + for override in (None, "dequantized", "high_precision"): + for fuse in (False, True): + _run_module(module_kind, recipe_cls(), override, fuse) + print( + f" e2e OK {module_kind} {recipe_cls.__name__} " + f"override={override} fuse={fuse}" + ) + + +TESTS = [ + test_fake_quant_grid_scale_rtne, + test_D_lossless_in_bf16, + test_E_mxfp8_rowwise_lossless, + test_G_blockwise_lossless_and_transpose, + test_C_mxfp8_colwise_bound, + test_lossless_dequant_to_bf16, + test_kernel_matches_torch_reference, + test_edge_value_domains, + test_blockwise_recipe_is_128x128, + test_blockwise_over_ratio_is_bounded_not_lossless, + test_exp2f_e8m0_all_codes, + test_ste_identity_gradient, + test_misaligned_and_noncontiguous, + test_bf16_exhaustive_bit_patterns, + test_fp32_bit_fuzz, + test_scale_threshold_and_rtne_midpoints, + test_deployment_top_cap_domain, + test_tilekernels_cross_parity, + test_blockwise_feasibility_enumeration, + test_recipe_switch_invalidates_cache, + test_ops_api_rejected, + test_fourway_matrix, + test_kernel_perf, + test_kernel_fast_math_immune, + test_cute_dsl_version, + test_e2e_matrix, +] + +if __name__ == "__main__": + import sys + import traceback + + failed = 0 + for t in TESTS: + try: + t() + torch.cuda.synchronize() + print(f"PASS {t.__name__}") + except Exception: + failed += 1 + print(f"FAIL {t.__name__}") + traceback.print_exc() + print(f"{len(TESTS) - failed}/{len(TESTS)} passed") + sys.exit(1 if failed else 0) diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 1e3c04573b..35c6d48fcb 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -16,6 +16,7 @@ #include "../utils.cuh" #include "dispatch/dequantize.cuh" #include "dispatch/quantize.cuh" +#include "mxfp4/fake_quantize_mxfp4.cuh" #include "transformer_engine/transpose.h" void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { @@ -47,6 +48,37 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, dispatch::quantize_fwd_helper(input, output, quant_config, stream); } +void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp4_fake_quantize); + using namespace transformer_engine; + const Tensor &input_tensor = *convertNVTETensorCheck(input); + Tensor &output_tensor = *convertNVTETensorCheck(output); + NVTE_CHECK(input_tensor.dtype() == DType::kBFloat16 || input_tensor.dtype() == DType::kFloat32, + "MXFP4 fake-quantize supports BF16/FP32 inputs, got ", to_string(input_tensor.dtype())); + NVTE_CHECK(output_tensor.dtype() == input_tensor.dtype(), + "MXFP4 fake-quantize output dtype must match the input dtype."); + NVTE_CHECK(input_tensor.data.shape == output_tensor.data.shape, + "MXFP4 fake-quantize output shape must match the input shape."); + const size_t numel = input_tensor.numel(); + NVTE_CHECK(numel % 32 == 0, "MXFP4 fake-quantize needs the element count divisible by 32."); + NVTE_CHECK(!input_tensor.data.shape.empty() && input_tensor.data.shape.back() % 32 == 0, + "MXFP4 fake-quantize needs the innermost dimension divisible by 32 " + "(1x32 blocks must not cross rows)."); + if (numel == 0) { + return; + } + if (input_tensor.dtype() == DType::kBFloat16) { + dispatch::mxfp4::fake_quantize_mxfp4_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(output_tensor.data.dptr), numel, stream); + } else { + dispatch::mxfp4::fake_quantize_mxfp4_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(output_tensor.data.dptr), numel, stream); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dequantize); using namespace transformer_engine; diff --git a/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh new file mode 100644 index 0000000000..5d96ec21e5 --- /dev/null +++ b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh @@ -0,0 +1,129 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + + +#ifndef TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_ +#define TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_ + +#include +#include + +#include "../../utils.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp4 { +namespace fake_quantize_kernel { + +constexpr size_t MXFP4_BLOCK_SIZE = 32; +constexpr size_t THREADS_PER_CHUNK = 256; +constexpr float E2M1_MAX = 6.0f; +constexpr int SCALE_EXP_MIN = -126; +constexpr int SCALE_EXP_MAX = 125; + +__device__ __forceinline__ float round_e2m1_magnitude(const float y) { + if (y <= 2.0f) { + return rintf(y * 2.0f) * 0.5f; + } + if (y <= 4.0f) { + return rintf(y); + } + return rintf(y * 0.5f) * 2.0f; +} + +__device__ __forceinline__ int block_scale_exponent_from_bits(const int bits) { + const int exp_field = bits >> 23; + const int mantissa = bits & 0x7FFFFF; + int e = exp_field - 129 + (mantissa > 0x400000 ? 1 : 0); + if (exp_field == 0) { + e = SCALE_EXP_MIN; + } + return max(SCALE_EXP_MIN, min(SCALE_EXP_MAX, e)); +} + +__device__ __forceinline__ float exp2_from_int(const int e) { + return __int_as_float((e + 127) << 23); +} + +__device__ __forceinline__ float mul_rn_no_ftz(const float a, const float b) { + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + fake_quantize_mxfp4_kernel(const IType *__restrict__ input, IType *__restrict__ output, + const size_t num_vectors) { + constexpr size_t VEC_ELTS = 16 / sizeof(IType); + constexpr int LANES_PER_BLOCK = MXFP4_BLOCK_SIZE / VEC_ELTS; + + const size_t stride = static_cast(gridDim.x) * blockDim.x; + for (size_t vec_idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + vec_idx < num_vectors; vec_idx += stride) { + Vec in_vec; + in_vec.load_from(input, vec_idx); + + float elts[VEC_ELTS]; + int amax_bits = 0; +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + elts[i] = static_cast(in_vec.data.elt[i]); + amax_bits = max(amax_bits, __float_as_int(elts[i]) & 0x7fffffff); + } + const unsigned lane = threadIdx.x & 31u; + const unsigned group_mask = + (LANES_PER_BLOCK == 32) + ? 0xffffffffu + : (((1u << LANES_PER_BLOCK) - 1u) << (lane & ~(LANES_PER_BLOCK - 1u))); +#pragma unroll + for (int offset = LANES_PER_BLOCK / 2; offset > 0; offset >>= 1) { + amax_bits = max(amax_bits, __shfl_xor_sync(group_mask, amax_bits, offset)); + } + const int nonfinite = amax_bits >= 0x7f800000; + + Vec out_vec; + if (nonfinite) { + const float qnan = __int_as_float(0x7fc00000); +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + out_vec.data.elt[i] = static_cast(qnan); + } + } else { + const int e = (amax_bits > 0) ? block_scale_exponent_from_bits(amax_bits) : 0; + const float scale = exp2_from_int(e); + const float inv_scale = exp2_from_int(-e); +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + const float av = __int_as_float(__float_as_int(elts[i]) & 0x7fffffff); + const float y = fminf(mul_rn_no_ftz(av, inv_scale), E2M1_MAX); + const float q = copysignf(round_e2m1_magnitude(y), elts[i]); + out_vec.data.elt[i] = static_cast(mul_rn_no_ftz(q, scale)); + } + } + out_vec.store_to(output, vec_idx); + } +} + +} + +template +void fake_quantize_mxfp4_launch(const IType *input, IType *output, const size_t numel, + cudaStream_t stream) { + using namespace fake_quantize_kernel; + constexpr size_t VEC_ELTS = 16 / sizeof(IType); + const size_t num_vectors = numel / VEC_ELTS; + const size_t blocks_needed = (num_vectors + THREADS_PER_CHUNK - 1) / THREADS_PER_CHUNK; + const int grid = static_cast(blocks_needed < 65535 ? blocks_needed : 65535); + fake_quantize_mxfp4_kernel + <<>>(input, output, num_vectors); +} + +} +} +} + +#endif diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 4d6d24ba65..c6ca1cf2f0 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -418,6 +418,17 @@ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, */ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Project a BF16/FP32 tensor onto the MXFP4 (E2M1) grid with 1x32 + * power-of-two block scales and write back the dequantized values + * (weight fake-quantization for MXFP4 QAT). Non-finite values poison + * their whole 1x32 block with NaN. + * + * \param[in] input Input tensor (BF16/FP32, element count divisible by 32). + * \param[in,out] output Output tensor with the same dtype and shape. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); + /*! \brief Casts input grouped tensor from reduced to higher precision. * In case of the MXFP8 dequantization, the dequantized values are stored to the rowwise * data of the output tensor, regardless of whether the row- or columnwise scaling is used. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8c209ace15..1a223a414c 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -168,6 +168,11 @@ def custom(cls): """Whether the given recipe is custom.""" return issubclass(cls, CustomRecipe) + @classmethod + def mxfp4_qat(cls): + """Whether the given recipe applies MXFP4 weight QAT on top of its base recipe.""" + return issubclass(cls, (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling)) + @dataclass(repr=False) class DelayedScaling(Recipe): @@ -481,6 +486,53 @@ def _make_repr(self) -> str: ) +@dataclass(repr=False) +class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling): + """ + MXFP8 recipe with MXFP4 weight quantization-aware training. + + Identical to :class:`MXFP8BlockScaling` except that weights are first + projected onto the MXFP4 (E2M1) grid with a 1x32 power-of-two scale + before the regular MXFP8 weight quantization. Because MXFP4 blocks are + 1x32-aligned with MXFP8 blocks and the E2M1 grid is a subset of E4M3, + the rowwise MXFP8 weight representation is an exact (lossless) encoding + of the MXFP4-grid weight; the columnwise (32x1) representation used by + dgrad is quantized from the same MXFP4-grid weight. Activations and + gradients are handled exactly as in the base recipe, and + ``backward_override`` keeps its base-recipe meaning ('high_precision' + uses the original unquantized weight for dgrad). + """ + + +@dataclass(repr=False) +class MXFP4QATFloat8BlockScaling(Float8BlockScaling): + """ + Float8 block-scaling recipe with MXFP4 weight quantization-aware training. + + Identical to :class:`Float8BlockScaling` except that weights are first + projected onto the MXFP4 (E2M1) grid with a 1x32 power-of-two scale + before the regular 128x128 blockwise FP8 weight quantization. The FP8 + weight encoding is exact (lossless) whenever every 1x32 MXFP4 scale + within a 128x128 tile lies within the FP8 dynamic-range headroom of the + tile's maximum scale; weights are quantized in square 128x128 blocks so + no rowwise/columnwise distinction applies (the transpose is exact). + Activations and gradients are handled exactly as in the base recipe. + """ + + def __post_init__(self) -> None: + super().__post_init__() + if self.w_block_scaling_dim != 2: + raise ValueError( + "MXFP4 QAT requires 128x128 (2D) weight scaling blocks, got " + f"w_block_scaling_dim={self.w_block_scaling_dim}." + ) + if not self.fp8_quant_fwd_weight.power_2_scale: + raise ValueError( + "MXFP4 QAT requires power-of-two weight scales; " + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 is incompatible." + ) + + @dataclass(repr=False) class NVFP4BlockScaling(Recipe): """ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 2814aa3490..94a82990e3 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -380,6 +380,8 @@ __device__ __forceinline__ bf16 exp2f_rcp(e8m0_t biased_exp) { } __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { + if (biased_exp == 0) return __int_as_float(0x00400000); + if (biased_exp == 255) return __int_as_float(0x7fffffff); return __int_as_float(biased_exp << FP32_MANTISSA_BITS); } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..2a77895ead 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -347,6 +347,8 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz py::object dequantize(const py::handle &input, DType otype); +at::Tensor mxfp4_fake_quantize(const at::Tensor &input); + py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, std::optional tensor_offsets, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..6ff1662d83 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -556,6 +556,25 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, py::cast(std::move(dbias_torch))); } +at::Tensor mxfp4_fake_quantize(const at::Tensor &input) { + TORCH_CHECK(input.is_cuda(), "mxfp4_fake_quantize: input must be a CUDA tensor."); + TORCH_CHECK(input.dim() == 2, "mxfp4_fake_quantize: input must be 2D, got dim=", input.dim()); + TORCH_CHECK(input.size(1) % 32 == 0, + "mxfp4_fake_quantize: innermost dimension must be divisible by 32, got ", + input.size(1)); + const c10::cuda::CUDAGuard device_guard(input.device()); + auto input_contiguous = input.contiguous(); + if (reinterpret_cast(input_contiguous.data_ptr()) % 16 != 0) { + input_contiguous = input_contiguous.clone(); + } + auto output = at::empty_like(input_contiguous); + auto input_cpp = makeTransformerEngineTensor(input_contiguous); + auto output_cpp = makeTransformerEngineTensor(output); + nvte_mxfp4_fake_quantize(input_cpp.data(), output_cpp.data(), + at::cuda::getCurrentCUDAStream()); + return output; +} + py::object dequantize(const py::handle &input, transformer_engine::DType otype) { init_extension(); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7e9d114be8..e28844c11d 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -199,6 +199,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none(), py::arg("noop") = py::none()); + m.def("mxfp4_fake_quantize", &transformer_engine::pytorch::mxfp4_fake_quantize, + "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)", + py::arg("input"), py::call_guard()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); m.def("create_empty_quantized_tensor", diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 7c11e635c6..f9ecaaa143 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -46,6 +46,7 @@ _fsdp_gather_tensors, ) from ..constants import dist_group_type +from ..mxfp4_qat import mxfp4_fake_quantize from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer @@ -800,8 +801,20 @@ def quantize_weight( ``_fp8_workspaces``. """ + _mxfp4_qat_active = ( + FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + if FP8GlobalStateManager.is_fp8_enabled() + else False + ) + # Already-quantized weight (primary FP8 parameters) if isinstance(tensor, QuantizedTensor): + if _mxfp4_qat_active: + raise NotImplementedError( + "MXFP4 QAT recipes do not support primary quantized weights: the " + "high-precision master weight is required to project onto the " + "MXFP4 grid." + ) update_rowwise = True if quantizer.rowwise_usage else None update_columnwise = True if quantizer.columnwise_usage else None tensor.update_usage( @@ -833,6 +846,8 @@ def quantize_weight( if update_workspace: if tensor is None: raise ValueError("tensor kwarg must be provided to update FP8 workspace") + if _mxfp4_qat_active: + tensor = mxfp4_fake_quantize(tensor) if hasattr(workspace, "quantize_"): workspace.quantize_(tensor, noop_flag=skip_update_flag) else: @@ -842,6 +857,8 @@ def quantize_weight( # Cache miss — create new workspace if tensor is None or quantizer is None: raise ValueError("tensor and quantizer kwargs must be provided to construct FP8 workspace") + if _mxfp4_qat_active: + tensor = mxfp4_fake_quantize(tensor) if cache: # Ensure the tensor in the cache is an instance of torch.Tensor, # as it persists beyond a single forward pass. @@ -1506,9 +1523,12 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() _current_recipe = meta["recipe"] - if _original_recipe is not None and not ( - issubclass(_current_recipe.__class__, _original_recipe.__class__) - or issubclass(_original_recipe.__class__, _current_recipe.__class__) + if _original_recipe is not None and ( + not ( + issubclass(_current_recipe.__class__, _original_recipe.__class__) + or issubclass(_original_recipe.__class__, _current_recipe.__class__) + ) + or _current_recipe.mxfp4_qat() != _original_recipe.mxfp4_qat() ): warnings.warn( f"Recipe type changed from {_original_recipe.__class__.__name__} " diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a588e21a0c..d17e221325 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -32,6 +32,7 @@ _2X_ACC_WGRAD, ) from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..mxfp4_qat import mxfp4_fake_quantize from ..utils import ( assert_dim_for_fp8_exec, cast_if_needed, @@ -810,6 +811,8 @@ def backward( # fsdp2 quantized-tensor hooks when workspace was not saved. weight = saved_weight elif ctx.weight_quantizer is not None: + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + saved_weight = mxfp4_fake_quantize(saved_weight) ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight = ctx.weight_quantizer(saved_weight) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 19e20d775c..dfa9f7f0f8 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -33,6 +33,7 @@ _2X_ACC_WGRAD, ) from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..mxfp4_qat import mxfp4_fake_quantize from ..jit import ( bias_gelu_fused, bgrad_dgelu_fused, @@ -247,11 +248,12 @@ def _forward( backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override else: backward_override = None - assert backward_override is None, ( - "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." - " Replace LayerNormMLP with LayerNormLinear + Linear to enable" - " high_precision/dequantized backward." - ) + if backward_override is not None: + raise NotImplementedError( + "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in" + " LayerNormMLP. Replace LayerNormMLP with LayerNormLinear + Linear to enable" + " high_precision/dequantized backward." + ) # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take if is_grad_enabled and not recompute_for_bwd: @@ -1228,6 +1230,8 @@ def backward( if isinstance(origin_fc2_weight, QuantizedTensorStorage): fc2_weight = origin_fc2_weight elif ctx.fc2_weight_quantizer is not None: + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + origin_fc2_weight = mxfp4_fake_quantize(origin_fc2_weight) ctx.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) fc2_weight = ctx.fc2_weight_quantizer(origin_fc2_weight) @@ -1522,6 +1526,8 @@ def fc2_wgrad_gemm( if isinstance(origin_fc1_weight, QuantizedTensorStorage): fc1_weight = origin_fc1_weight elif ctx.fc1_weight_quantizer is not None: + if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat(): + origin_fc1_weight = mxfp4_fake_quantize(origin_fc1_weight) ctx.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) fc1_weight = ctx.fc1_weight_quantizer(origin_fc1_weight) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 6b0f941bc4..0fe378b11a 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -15,6 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from ..mxfp4_qat import mxfp4_fake_quantize from transformer_engine.pytorch.torch_version import torch_version from .base import ( @@ -1003,6 +1004,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. # fsdp2 quantized-tensor hooks when workspace was not saved. weight_fp8 = saved_weight elif bwd_args.weight_quantizer is not None: + if bwd_args.fp8 and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + saved_weight = mxfp4_fake_quantize(saved_weight) bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(saved_weight) elif ( @@ -1013,6 +1016,8 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. ): # Distributed weight re-gathered a BF16 weight: quantize with the layer quantizer # so the dgrad operand isn't cast by the delayed recipe. + if FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + weight_fp8 = mxfp4_fake_quantize(weight_fp8) bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) weight_fp8 = bwd_args.weight_quantizer(weight_fp8) diff --git a/transformer_engine/pytorch/mxfp4_qat.py b/transformer_engine/pytorch/mxfp4_qat.py new file mode 100644 index 0000000000..c072121503 --- /dev/null +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -0,0 +1,217 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP4 weight fake-quantization for quantization-aware training. + +Projects a weight onto the MXFP4 (E2M1) grid with 1x32 power-of-two (UE8M0) +block scales and returns the dequantized result in the input dtype — i.e. the +fused form of ``bf16 -> mxfp4(row) -> dequantize -> bf16``. The returned values +are exactly representable in bf16/fp32 (E2M1 needs at most two significant +mantissa bits and the scales are powers of two), so composing with the base +recipe's weight quantization stays lossless: + +* MXFP8 (1x32, E8M0): MXFP4 blocks are 1x32-aligned and the E2M1 grid is a + subset of E4M3 — the rowwise MXFP8 encoding of the projected weight is exact. +* Float8 block scaling (128x128, power-of-two scales): exact whenever every + 1x32 MXFP4 scale within a tile lies within the FP8 dynamic-range headroom of + the tile's maximum scale. + +Three bit-identical implementations are dispatched via ``NVTE_MXFP4_QAT_IMPL`` +(``auto``/``cuda``/``cute_dsl``/``torch``, default ``auto`` = CUDA binding, +then CuTe DSL, then the PyTorch reference): the CUDA kernel in +``common/cast/mxfp4/fake_quantize_mxfp4.cuh`` (``tex.mxfp4_fake_quantize``), +the CuTe DSL kernel in ``mxfp4_qat_cute_dsl.py``, and +``_mxfp4_fake_quantize_torch`` below. + +Value-domain semantics (identical in all implementations): +* zero block (amax == 0): scale 1, output zeros — exact. +* tiny amax (<= 6*2^-126): the scale floors at 2^-126 (TileKernels deployment + contract). The 0.5 payload then dequantizes to 2^-127, a bf16/fp32 + subnormal. NOTE: TE's MXFP8 software dequantize flushes that grid point to + zero until rebuilt with the ptx.cuh UE8M0 code-0 fix (raw encoding and + hardware GEMM are exact). +* non-finite amax (inf/NaN anywhere in a 1x32 block): the WHOLE block becomes + NaN. A single inf must not silently rescale its 31 neighbors to zero, and + NaN must not silently fall back to scale 1 — corruption is made visible. +* huge amax (> 6*2^125): the scale is capped at 2^125 so the dequantized grid + stays representable in bf16/fp32; values saturate at 6*2^125 (satfinite). +* fp16 weights are rejected: even capped scales overflow fp16's range. +""" +import os +import warnings + +import torch + +__all__ = ["mxfp4_fake_quantize"] + +_E2M1_MAX = 6.0 +_MXFP4_BLOCK = 32 + +_IMPL_CHOICE = os.getenv("NVTE_MXFP4_QAT_IMPL", "auto") +_torch_path_forced = os.getenv("NVTE_MXFP4_QAT_DISABLE_CUDA_KERNEL", "0") == "1" +_kernel_required = os.getenv("NVTE_MXFP4_QAT_REQUIRE_KERNEL", "0") == "1" +_missing_kernel_warned = False +_cute_dsl_import_error = None + + +class _MXFP4FakeQuantizeSTE(torch.autograd.Function): + """Straight-through estimator: identity gradient through the projection.""" + + @staticmethod + def forward(ctx, weight, impl): # pylint: disable=arguments-differ + return impl(weight) + + @staticmethod + def backward(ctx, grad_output): # pylint: disable=arguments-differ + return grad_output, None + + +def _round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: + """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}. + + Input must be non-negative and <= 6. Within each binade the grid is + uniform, so torch.round (RTNE) on the rescaled value reproduces IEEE RTNE + exactly, including ties across binade boundaries (2.5 -> 2, 3.5 -> 4, 5 -> 4). + """ + fine = torch.round(y * 2.0) * 0.5 + mid = torch.round(y) + coarse = torch.round(y * 0.5) * 2.0 + return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) + + +def _mxfp4_fake_quantize_torch(weight: torch.Tensor) -> torch.Tensor: + """Bit-identical PyTorch reference for the CUDA and CuTe DSL kernels.""" + rows, cols = weight.shape + w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) + amax = w32.abs().amax(dim=-1, keepdim=True) + nonfinite = ~torch.isfinite(amax) + + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mantissa = bits & 0x7FFFFF + exp = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)) + exp = exp.clamp(min=-126, max=125) + scale = torch.ldexp(torch.ones_like(amax), exp) + scale = torch.where(amax > 0, scale, torch.ones_like(scale)) + + y = (w32 / scale).abs().clamp(max=_E2M1_MAX) + q = _round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) + q = torch.copysign(q, w32) + out = q * scale + out = torch.where(nonfinite.expand_as(out), torch.full_like(out, float("nan")), out) + return out.view(rows, cols).to(weight.dtype) + + +def _cuda_impl(): + import transformer_engine_torch as tex + + if not hasattr(tex, "mxfp4_fake_quantize"): + return None + + def impl(w): + w = w.contiguous() + if w.data_ptr() % 16 != 0: + w = w.clone() + return tex.mxfp4_fake_quantize(w) + + return impl + + +def _cute_dsl_impl(): + global _cute_dsl_import_error + if _cute_dsl_import_error is not None: + return None + try: + from .mxfp4_qat_cute_dsl import mxfp4_fake_quantize_cute_dsl + except Exception as exc: # pylint: disable=broad-except + _cute_dsl_import_error = exc + return None + return mxfp4_fake_quantize_cute_dsl + + +def _resolve_impl(weight: torch.Tensor): + if not weight.is_cuda: + return _mxfp4_fake_quantize_torch + if _IMPL_CHOICE == "torch": + return _mxfp4_fake_quantize_torch + if _IMPL_CHOICE == "cuda": + impl = _cuda_impl() + if impl is None: + raise RuntimeError( + "NVTE_MXFP4_QAT_IMPL=cuda but transformer_engine_torch was built " + "without mxfp4_fake_quantize. Rebuild Transformer Engine." + ) + return impl + if _IMPL_CHOICE == "cute_dsl": + impl = _cute_dsl_impl() + if impl is None: + raise RuntimeError( + "NVTE_MXFP4_QAT_IMPL=cute_dsl but the cutlass CuTe DSL is " + f"unavailable: {_cute_dsl_import_error!r}" + ) + return impl + if _IMPL_CHOICE != "auto": + raise ValueError( + f"NVTE_MXFP4_QAT_IMPL={_IMPL_CHOICE!r} is not one of " + "'auto', 'cuda', 'cute_dsl', 'torch'" + ) + if _torch_path_forced: + return _mxfp4_fake_quantize_torch + impl = _cuda_impl() + if impl is not None: + return impl + impl = _cute_dsl_impl() + if impl is not None: + return impl + if _kernel_required: + raise RuntimeError( + "NVTE_MXFP4_QAT_REQUIRE_KERNEL=1 but neither the CUDA binding nor " + "the CuTe DSL kernel is available. Rebuild Transformer Engine or " + "install the cutlass CuTe DSL." + ) + global _missing_kernel_warned + if not _missing_kernel_warned: + _missing_kernel_warned = True + warnings.warn( + "transformer_engine_torch was built without mxfp4_fake_quantize and " + "the CuTe DSL is unavailable; falling back to the (slower) PyTorch " + "reference. Rebuild Transformer Engine to use the CUDA kernel.", + stacklevel=2, + ) + return _mxfp4_fake_quantize_torch + + +def mxfp4_fake_quantize(weight: torch.Tensor) -> torch.Tensor: + """Project ``weight`` onto the MXFP4 grid and return it in the input dtype. + + Per 1x32 block: scale = 2^ceil(log2(amax / 6)) (all-zero blocks use 1.0), + values are divided by the scale, rounded RTNE onto the E2M1 grid with + saturation at +-6, and multiplied back. Blocks containing non-finite + values become NaN. The gradient is the straight-through estimator + (identity), the QAT-correct semantics. + + Env knobs: ``NVTE_MXFP4_QAT_IMPL`` picks the implementation + (``auto``/``cuda``/``cute_dsl``/``torch``); + ``NVTE_MXFP4_QAT_REQUIRE_KERNEL=1`` makes a kernel-less fallback a hard + error instead of a warn-once (recommended in production so a stale wheel + cannot silently train on the slow path). + """ + if weight.dim() != 2: + raise ValueError(f"MXFP4 QAT expects a 2D weight, got {tuple(weight.shape)}") + if weight.dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + f"MXFP4 QAT supports bf16/fp32 weights, got {weight.dtype} " + "(fp16 cannot represent the full MXFP4 scale range)" + ) + rows, cols = weight.shape + if cols % _MXFP4_BLOCK != 0: + raise ValueError( + f"MXFP4 QAT needs the weight inner dim divisible by {_MXFP4_BLOCK}, got {cols}" + ) + + impl = _resolve_impl(weight) + if torch.is_grad_enabled() and weight.requires_grad: + return _MXFP4FakeQuantizeSTE.apply(weight, impl) + return impl(weight) diff --git a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py new file mode 100644 index 0000000000..e4f07774ec --- /dev/null +++ b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py @@ -0,0 +1,217 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""CuTe DSL implementation of MXFP4 weight fake-quantization for QAT. + +Bit-identical to the CUDA kernel in +``common/cast/mxfp4/fake_quantize_mxfp4.cuh`` and to the PyTorch reference in +``transformer_engine/pytorch/mxfp4_qat.py``: the block-scale exponent is +derived in integer arithmetic and the value path uses explicit non-FTZ PTX, +so results are insensitive to compilation modes. Requires the cutlass CuTe +DSL package; selected via ``NVTE_MXFP4_QAT_IMPL=cute_dsl`` or as the ``auto`` +fallback when the CUDA binding is unavailable. +""" +import functools + +import torch + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import llvm + +__all__ = ["mxfp4_fake_quantize_cute_dsl"] + +_MXFP4_BLOCK = 32 +_THREADS = 256 + +_ASM_KW = dict( + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, +) + + +@dsl_user_op +def _is_nonfinite_f32(x: Float32, *, loc=None, ip=None) -> Int32: + """1 if x is inf/NaN else 0 (PTX testp.finite).""" + return Int32( + llvm.inline_asm( + T.i32(), + [Float32(x).ir_value(loc=loc, ip=ip)], + "{\n" + " .reg .pred p;\n" + " testp.finite.f32 p, $1;\n" + " selp.s32 $0, 0, 1, p;\n" + "}", + "=r,f", + loc=loc, + ip=ip, + **_ASM_KW, + ) + ) + + +@dsl_user_op +def _block_scale_from_amax(amax: Float32, *, loc=None, ip=None) -> Float32: + """scale = 2^clamp(ceil(log2(amax/6)), -126, 125); amax == 0 -> 1.0.""" + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(amax).ir_value(loc=loc, ip=ip)], + "{\n" + " .reg .s32 b, k, mnt, e, t;\n" + " .reg .pred pz, pm, psub;\n" + " setp.eq.f32 pz, $1, 0f00000000;\n" + " mov.b32 b, $1;\n" + " shr.s32 k, b, 23;\n" + " and.b32 mnt, b, 8388607;\n" + " setp.gt.s32 pm, mnt, 4194304;\n" + " selp.s32 t, 1, 0, pm;\n" + " add.s32 e, k, -129;\n" + " add.s32 e, e, t;\n" + " setp.eq.s32 psub, k, 0;\n" + " selp.s32 e, -126, e, psub;\n" + " max.s32 e, e, -126;\n" + " min.s32 e, e, 125;\n" + " add.s32 e, e, 127;\n" + " shl.b32 e, e, 23;\n" + " mov.b32 $0, e;\n" + " @pz mov.f32 $0, 0f3F800000;\n" + "}", + "=f,f", + loc=loc, + ip=ip, + **_ASM_KW, + ) + ) + + +@dsl_user_op +def _fake_quant_elem(x: Float32, scale: Float32, nonfinite: Int32, *, loc=None, ip=None) -> Float32: + """RTNE onto the E2M1 grid at the given power-of-two scale; nonfinite -> NaN.""" + return Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(x).ir_value(loc=loc, ip=ip), + Float32(scale).ir_value(loc=loc, ip=ip), + Int32(nonfinite).ir_value(loc=loc, ip=ip), + ], + "{\n" + " .reg .f32 y, ay, f, m, c, q, r, inv;\n" + " .reg .s32 ib;\n" + " .reg .pred p2, p4, pnf;\n" + " setp.ne.s32 pnf, $3, 0;\n" + " mov.b32 ib, $2;\n" + " sub.s32 ib, 2130706432, ib;\n" + " mov.b32 inv, ib;\n" + " mul.rn.f32 y, $1, inv;\n" + " abs.f32 ay, y;\n" + " min.f32 ay, ay, 0f40C00000;\n" + " mul.rn.f32 f, ay, 0f40000000;\n" + " cvt.rni.f32.f32 f, f;\n" + " mul.rn.f32 f, f, 0f3F000000;\n" + " cvt.rni.f32.f32 m, ay;\n" + " mul.rn.f32 c, ay, 0f3F000000;\n" + " cvt.rni.f32.f32 c, c;\n" + " mul.rn.f32 c, c, 0f40000000;\n" + " setp.le.f32 p2, ay, 0f40000000;\n" + " setp.le.f32 p4, ay, 0f40800000;\n" + " selp.f32 q, m, c, p4;\n" + " selp.f32 q, f, q, p2;\n" + " copysign.f32 r, $1, q;\n" + " mul.rn.f32 r, r, $2;\n" + " selp.f32 $0, 0f7FC00000, r, pnf;\n" + "}", + "=f,f,f,r", + loc=loc, + ip=ip, + **_ASM_KW, + ) + ) + + +class _MXFP4FakeQuantCuteDsl: + """Thread-per-1x32-block kernel; loads/stores via autovec_copy fragments.""" + + def __init__(self, dtype): + self.dtype = dtype + + @cute.jit + def __call__(self, mX: cute.Tensor, mO: cute.Tensor, num_blocks: Int32, stream): + self.kernel(mX, mO, num_blocks).launch( + grid=[cute.ceil_div(num_blocks, _THREADS), 1, 1], + block=[_THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel(self, mX: cute.Tensor, mO: cute.Tensor, num_blocks: Int32): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + idx = bidx * _THREADS + tidx + if idx < num_blocks: + x_row = mX[idx, None] + o_row = mO[idx, None] + x_frag = cute.make_fragment_like(x_row) + cute.autovec_copy(x_row, x_frag) + o_frag = cute.make_fragment_like(o_row) + + amax = Float32(0.0) + nonfinite = Int32(0) + for i in cutlass.range_constexpr(_MXFP4_BLOCK): + v = Float32(x_frag[i]) + nonfinite = nonfinite + _is_nonfinite_f32(v) + amax = cute.arch.fmax(amax, cute.arch.fmax(v, Float32(0.0) - v)) + + scale = _block_scale_from_amax(amax) + for i in cutlass.range_constexpr(_MXFP4_BLOCK): + v = Float32(x_frag[i]) + o_frag[i] = self.dtype(_fake_quant_elem(v, scale, nonfinite)) + cute.autovec_copy(o_frag, o_row) + + +@functools.lru_cache(maxsize=4) +def _compiled_kernel(dtype_key: str): + cutlass_dtype = cutlass.BFloat16 if dtype_key == "bf16" else cutlass.Float32 + kernel_obj = _MXFP4FakeQuantCuteDsl(cutlass_dtype) + sym_blocks = cute.sym_int() + x_fake = cute.runtime.make_fake_compact_tensor( + cutlass_dtype, (sym_blocks, _MXFP4_BLOCK), stride_order=(1, 0), assumed_align=16 + ) + o_fake = cute.runtime.make_fake_compact_tensor( + cutlass_dtype, (sym_blocks, _MXFP4_BLOCK), stride_order=(1, 0), assumed_align=16 + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel_obj, x_fake, o_fake, Int32(1), stream_fake, options="--enable-tvm-ffi" + ) + + +def mxfp4_fake_quantize_cute_dsl(weight: torch.Tensor) -> torch.Tensor: + """Project ``weight`` onto the MXFP4 grid and return it in the input dtype.""" + if not weight.is_cuda: + raise ValueError("MXFP4 QAT CuTe DSL kernel expects a CUDA tensor") + if weight.dim() != 2: + raise ValueError(f"MXFP4 QAT expects a 2D weight, got {tuple(weight.shape)}") + if weight.dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + f"MXFP4 QAT supports bf16/fp32 weights, got {weight.dtype} " + "(fp16 cannot represent the full MXFP4 scale range)" + ) + rows, cols = weight.shape + if cols % _MXFP4_BLOCK != 0: + raise ValueError( + f"MXFP4 QAT needs the weight inner dim divisible by {_MXFP4_BLOCK}, got {cols}" + ) + w = weight.contiguous() + if w.data_ptr() % 16 != 0: + w = w.clone() + out = torch.empty_like(w) + num_blocks = w.numel() // _MXFP4_BLOCK + fn = _compiled_kernel("bf16" if weight.dtype == torch.bfloat16 else "fp32") + fn(w.view(-1, _MXFP4_BLOCK), out.view(-1, _MXFP4_BLOCK), num_blocks) + return out diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index cb429055a4..2ceec72b5e 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -558,6 +558,16 @@ def _functional_forward( # Check weight tensor w = weight + if ( + with_quantized_compute + and FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + ): + raise NotImplementedError( + "MXFP4 QAT recipes are not implemented for the operations API " + "(te.ops): weight quantization here bypasses the QAT projection. " + "Use the module API (te.Linear/te.GroupedLinear/...)." + ) if not with_quantized_compute: w = maybe_dequantize(w, dtype) elif with_quantized_compute and not is_quantized_tensor(w): diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index d3e2dc0a13..45cec84acc 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -488,6 +488,13 @@ def _quantize_weights( ) -> Sequence[torch.Tensor]: """Construct quantized weight tensors.""" + if FP8GlobalStateManager.is_fp8_enabled() and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + raise NotImplementedError( + "MXFP4 QAT recipes do not support quantized primary weights in the " + "operations API: the high-precision master is required for the QAT " + "projection." + ) + # Manually construct MXFP8 weights if isinstance(quantizers[0], MXFP8Quantizer): return self._quantize_weights_mxfp8(weights, quantizers) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 83954a9b3d..554317b6f2 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -26,6 +26,7 @@ finalize_weight_grads, ) from ...module.base import _2X_ACC_WGRAD +from ...fp8 import FP8GlobalStateManager from ...quantization import Recipe from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ...tensor.grouped_tensor import GroupedTensor @@ -960,6 +961,11 @@ def fuser_forward( fc2_input_quantizer = fc2_op.get_quantizer("forward", 0) fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1) fc2_grad_output_quantizer = fc2_op.get_quantizer("backward", 0) + if FP8GlobalStateManager.is_fp8_enabled() and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + raise NotImplementedError( + "MXFP4 QAT recipes are not implemented for fused grouped-MLP ops: " + "weight quantization here bypasses the QAT projection." + ) # Extract split sizes from extra input fc1_split_sizes = basic_op_extra_inputs[0][0] diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 3a8ff5438d..ee03a63f06 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -319,6 +319,10 @@ def fuser_forward( raise RuntimeError( f"Unsupported recipe for Userbuffers ({recipe.__class__.__name__})" ) + if recipe.mxfp4_qat(): + raise NotImplementedError( + "MXFP4 QAT recipes are not implemented for Userbuffers fused ops." + ) # Get autocast dtype if needed if torch.is_autocast_enabled(): From 000b637ede9c221d01067e5cdf4f0953da3ab7cd Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 25 Jul 2026 10:20:55 -0700 Subject: [PATCH 02/14] Reject fp16 activation/dequantize dtype under MXFP4 QAT The MXFP4 grid reaches 6*2^125, far beyond fp16 range, so an fp16 pipeline cannot represent the projected weight or its dequantized form. quantize_weight now raises when a QAT recipe is active and the workspace/activation dtype is fp16 (fp16 weights were already rejected by the projection itself); documented in both recipe docstrings. --- tests/pytorch/test_mxfp4_qat.py | 16 ++++++++++++++++ transformer_engine/common/recipe/__init__.py | 5 ++++- transformer_engine/pytorch/module/base.py | 6 ++++++ transformer_engine/pytorch/mxfp4_qat.py | 3 ++- 4 files changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index 1e33e907a5..e346b3a107 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -756,6 +756,21 @@ def test_blockwise_feasibility_enumeration(): "(E4M3 subnormal lattice bound)") +def test_fp16_pipeline_rejected(): + """fp16 cannot represent the MXFP4 grid (top 6*2^125 >> fp16 max): both the + fp16-weight path and the fp16 activation/dequantize dtype must fail loudly + under MXFP4 QAT instead of silently overflowing.""" + mod = te.Linear(256, 128, bias=False, params_dtype=torch.float16, device=DEV) + x = torch.randn(64, 256, device=DEV, dtype=torch.float16) + for recipe in (MXFP4QATMXFP8BlockScaling(), MXFP4QATFloat8BlockScaling()): + try: + with fp8_autocast(enabled=True, fp8_recipe=recipe): + mod(x) + raise SystemExit(f"fp16 pipeline must be rejected ({recipe.__class__.__name__})") + except (NotImplementedError, ValueError): + pass + + def test_fourway_matrix(): """TileKernels deployment reference vs CuTe DSL vs CUDA kernel vs torch reference, bitwise, over every previously-hit edge domain and every @@ -1114,6 +1129,7 @@ def test_e2e_matrix(): test_blockwise_feasibility_enumeration, test_recipe_switch_invalidates_cache, test_ops_api_rejected, + test_fp16_pipeline_rejected, test_fourway_matrix, test_kernel_perf, test_kernel_fast_math_immune, diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 1a223a414c..d649a2fd9e 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -500,7 +500,8 @@ class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling): dgrad is quantized from the same MXFP4-grid weight. Activations and gradients are handled exactly as in the base recipe, and ``backward_override`` keeps its base-recipe meaning ('high_precision' - uses the original unquantized weight for dgrad). + uses the original unquantized weight for dgrad). Only bf16/fp32 weights + and activation dtypes are supported: the MXFP4 grid exceeds fp16 range. """ @@ -517,6 +518,8 @@ class MXFP4QATFloat8BlockScaling(Float8BlockScaling): tile's maximum scale; weights are quantized in square 128x128 blocks so no rowwise/columnwise distinction applies (the transpose is exact). Activations and gradients are handled exactly as in the base recipe. + Only bf16/fp32 weights and activation dtypes are supported: the MXFP4 + grid exceeds fp16 range. """ def __post_init__(self) -> None: diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f9ecaaa143..5447aa676f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -806,6 +806,12 @@ def quantize_weight( if FP8GlobalStateManager.is_fp8_enabled() else False ) + if _mxfp4_qat_active and workspace_dtype == torch.float16: + raise NotImplementedError( + "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " + "the MXFP4 grid (values up to 6*2^125) exceeds fp16 range. Use bf16 " + "or fp32." + ) # Already-quantized weight (primary FP8 parameters) if isinstance(tensor, QuantizedTensor): diff --git a/transformer_engine/pytorch/mxfp4_qat.py b/transformer_engine/pytorch/mxfp4_qat.py index c072121503..14322ef457 100644 --- a/transformer_engine/pytorch/mxfp4_qat.py +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -36,7 +36,8 @@ NaN must not silently fall back to scale 1 — corruption is made visible. * huge amax (> 6*2^125): the scale is capped at 2^125 so the dequantized grid stays representable in bf16/fp32; values saturate at 6*2^125 (satfinite). -* fp16 weights are rejected: even capped scales overflow fp16's range. +* fp16 is rejected for both weights and the activation/dequantize dtype: + even capped scales overflow fp16's range. """ import os import warnings From 20b367641eed245a3d47e405f786359dfc0aa7b7 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 25 Jul 2026 10:33:01 -0700 Subject: [PATCH 03/14] Tighten UE8M0 dequantize verification - exp2f JIT test now checks all 256 codes bitwise including the code-255 NaN payload (0x7FFFFFFF) - add an end-to-end nvte MXFP8 software-dequantize test with planted extreme scale codes through the real kernel: demands code 0 -> 2^-127 and code 255 -> NaN on a fixed build, and pins the pre-fix wheel behavior (flush to zero / Inf) until then - the bf16-dequant losslessness test auto-detects a fixed build instead of hard-asserting the pre-fix flush --- tests/pytorch/test_mxfp4_qat.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index e346b3a107..55772590f3 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -162,6 +162,31 @@ def test_E_mxfp8_rowwise_lossless(): print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") +def test_dequantize_e8m0_extreme_codes(): + """End-to-end nvte MXFP8 software dequantize with planted extreme UE8M0 + scale codes through the REAL kernel: code 0 must expand to 2^-127 and code + 255 to NaN once the ptx.cuh exp2f fix is in the built library; the pre-fix + wheel flushes code 0 to zero and expands code 255 to Inf. Auto-detects and + reports which behavior the installed build has.""" + w = make_weight(32, 64) + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize( + mxfp4_fake_quantize(w) + ) + data = q._rowwise_data.view(torch.uint8) + codes = q._rowwise_scale_inv.view(torch.uint8) + data[0, :32] = 56 + codes[0, 0] = 0 + data[1, :32] = 56 + codes[1, 0] = 255 + dq = q.dequantize(dtype=torch.float32) + if (dq[0, :32] == 2.0**-127).all() and torch.isnan(dq[1, :32]).all(): + print(" fixed build: code 0 -> 2^-127, code 255 -> NaN") + else: + assert (dq[0, :32] == 0).all(), "code 0 neither 2^-127 (fixed) nor 0 (pre-fix)" + assert torch.isinf(dq[1, :32]).all(), "code 255 neither NaN (fixed) nor Inf (pre-fix)" + print(" pre-fix wheel: code 0 flushed to 0, code 255 -> Inf (rebuild pending)") + + def test_G_blockwise_lossless_and_transpose(): for m, n in SHAPES: w = make_weight(m, n, zero_blocks=4) @@ -199,7 +224,9 @@ def test_lossless_dequant_to_bf16(): assert torch.equal(dq8[~flushed], w_hat[~flushed]), ( f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" ) - assert (dq8[flushed] == 0).all() + if not (dq8[flushed] == 0).all(): + assert torch.equal(dq8, w_hat), "dequant neither flushed nor exact?" + print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") w2_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4)) qb = Float8BlockQuantizer( @@ -575,6 +602,8 @@ def test_exp2f_e8m0_all_codes(): exp_ref = torch.ldexp(torch.ones(256), codes - 127) rcp_ref = torch.ldexp(torch.ones(256), 127 - codes) assert torch.isnan(e[255]) and torch.isnan(r[255]), "code 255 must be NaN" + assert e.view(torch.int32)[255].item() == 0x7FFFFFFF, "code 255 NaN payload" + assert r.view(torch.int32)[255].item() == 0x7FFFFFFF, "code 255 rcp NaN payload" assert torch.equal(e[:255].cpu(), exp_ref[:255]), "exp2f wrong (code 0 == 2^-127?)" assert torch.equal(r[:255].cpu(), rcp_ref[:255]), "exp2f_rcp wrong" @@ -1111,6 +1140,7 @@ def test_e2e_matrix(): test_fake_quant_grid_scale_rtne, test_D_lossless_in_bf16, test_E_mxfp8_rowwise_lossless, + test_dequantize_e8m0_extreme_codes, test_G_blockwise_lossless_and_transpose, test_C_mxfp8_colwise_bound, test_lossless_dequant_to_bf16, From d4d6fd06d72ddf2b30b9a58aeada63bfc5a5c30b Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 25 Jul 2026 11:08:54 -0700 Subject: [PATCH 04/14] Condense docstrings to terse functional descriptions --- tests/pytorch/test_mxfp4_qat.py | 84 ++++--------------- .../common/include/transformer_engine/cast.h | 6 +- transformer_engine/common/recipe/__init__.py | 29 ++----- transformer_engine/pytorch/mxfp4_qat.py | 62 +++----------- .../pytorch/mxfp4_qat_cute_dsl.py | 14 ++-- 5 files changed, 43 insertions(+), 152 deletions(-) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index 55772590f3..7e210284f0 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -2,15 +2,7 @@ # # See LICENSE for license information. -"""Per-step and end-to-end tests for the MXFP4 weight-QAT recipes. - -Step coverage (letters refer to the design notes): - D/E/F losslessness: fake-quant exact in bf16; MXFP8 rowwise and 128x128 - blockwise weight encodings of the MXFP4-grid weight are bit-exact. - A/B/C backward_override weight semantics through real fwd/bwd GEMMs, - with discriminators (closer to w_hat vs closer to the master). - G STE: weight gradients land on the module parameter / main_grad. -""" +"""Per-step and end-to-end tests for the MXFP4 weight-QAT recipes.""" import ctypes import os @@ -163,11 +155,7 @@ def test_E_mxfp8_rowwise_lossless(): def test_dequantize_e8m0_extreme_codes(): - """End-to-end nvte MXFP8 software dequantize with planted extreme UE8M0 - scale codes through the REAL kernel: code 0 must expand to 2^-127 and code - 255 to NaN once the ptx.cuh exp2f fix is in the built library; the pre-fix - wheel flushes code 0 to zero and expands code 255 to Inf. Auto-detects and - reports which behavior the installed build has.""" + """Real MXFP8 software dequantize with planted UE8M0 codes 0/255; auto-detects pre-fix wheel vs fixed build.""" w = make_weight(32, 64) q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize( mxfp4_fake_quantize(w) @@ -210,8 +198,7 @@ def test_G_blockwise_lossless_and_transpose(): def test_lossless_dequant_to_bf16(): - """The bwd override paths dequantize to activation dtype (bf16): both weight - encodings must reproduce the MXFP4-grid weight bit-exactly at bf16 too.""" + """Both weight encodings dequantize to bf16 bit-exactly (bwd override target dtype).""" for m, n in SHAPES: w = make_weight(m, n, zero_blocks=4, outliers=4) w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) @@ -259,9 +246,7 @@ def test_C_mxfp8_colwise_bound(): def _intree_kernel(fast_math=False): - """JIT-compile the in-tree TE kernel for testing (the product path is the - tex.mxfp4_fake_quantize binding compiled with the wheel). fast_math=True - builds the SAME source with --use_fast_math to prove flag immunity.""" + """JIT-compile the in-tree kernel; fast_math=True builds the same source with --use_fast_math.""" if fast_math not in _INTREE_KERNELS: import os from torch.utils.cpp_extension import load_inline @@ -519,13 +504,7 @@ def timeit(fn, iters=20): def test_kernel_fast_math_immune(): - """The kernel must stay bit-identical when compiled with --use_fast_math. - - Guards the hazard class from the bitwise.md PR list (flashinfer#3387 etc.): - approximate division / ftz under fast-math flags silently changing - quantization results. The kernel uses only integer exponent math, exact - power-of-two multiplies, and rintf — no fp division, no MUFU. - """ + """The kernel stays bit-identical when compiled with --use_fast_math.""" import transformer_engine.pytorch.mxfp4_qat as M k = _intree_kernel(fast_math=True) @@ -553,12 +532,7 @@ def test_kernel_fast_math_immune(): def test_exp2f_e8m0_all_codes(): - """UE8M0 scale expansion: code 0 is 2^-127 (fp32 subnormal), code 255 NaN. - - Compiles the exp2f/exp2f_rcp FUNCTION TEXT extracted verbatim from - common/util/ptx.cuh (the mxfp8 dequant kernels' scale expansion) and checks - all 256 codes. Guards the code-0 fix: the old bit-shift construction gave - +0.0 and zeroed whole blocks on software dequant.""" + """All 256 UE8M0 codes through exp2f/exp2f_rcp extracted verbatim from ptx.cuh.""" import re as _re from torch.utils.cpp_extension import load_inline @@ -609,8 +583,7 @@ def test_exp2f_e8m0_all_codes(): def test_ste_identity_gradient(): - """Public mxfp4_fake_quantize must carry an identity (STE) gradient on both - implementation paths — the QAT semantics for d(w_hat)/dw.""" + """mxfp4_fake_quantize carries an identity (STE) gradient on both implementation paths.""" import transformer_engine.pytorch.mxfp4_qat as M for dtype in (torch.bfloat16, torch.float32): @@ -633,8 +606,7 @@ def test_ste_identity_gradient(): def test_misaligned_and_noncontiguous(): - """Storage-offset views (contiguous but not 16B-aligned) and non-contiguous - inputs must produce the same bits as an aligned contiguous copy.""" + """Misaligned storage-offset views and non-contiguous inputs match an aligned copy bitwise.""" import transformer_engine.pytorch.mxfp4_qat as M for dtype, off in ((torch.bfloat16, 4), (torch.float32, 2)): @@ -653,8 +625,7 @@ def test_misaligned_and_noncontiguous(): def test_bf16_exhaustive_bit_patterns(): - """All 65536 bf16 bit patterns (as 2048 blocks of 32): kernel/torch bit - parity everywhere, fp64 oracle on the all-finite blocks.""" + """All 65536 bf16 bit patterns: kernel/torch parity, fp64 oracle on finite blocks.""" bits = torch.arange(65536, dtype=torch.int32, device=DEV).to(torch.int16) w = bits.view(torch.bfloat16).view(2048, 32) gk, gt = _both_paths(w) @@ -679,8 +650,7 @@ def test_fp32_bit_fuzz(): def test_scale_threshold_and_rtne_midpoints(): - """amax = 1.5*2^k boundary (one ulp below / exact / one ulp above) and every - RTNE midpoint of the E2M1 grid, against explicit expected values.""" + """amax = 1.5*2^k ulp triples and all E2M1 RTNE midpoints vs explicit expected values.""" rows = [] for k in (-100, -10, 0, 42, 100): a = torch.tensor(1.5 * 2.0**k, dtype=torch.float32) @@ -707,12 +677,7 @@ def test_scale_threshold_and_rtne_midpoints(): def test_deployment_top_cap_domain(): - """Top-of-range contract: for amax <= 6*2^125 the fake-quant is bitwise the - TileKernels deployment result. Above that, TileKernels moves to scale 2^126 - (packed payload + UE8M0 byte 253 -- representable in packed form only, since - {4,6}*2^126 overflow bf16/fp32), while this path pins satfinite at the cap - 2^125. Restricted-domain divergence, documented and unreachable for real - weights (requires |w| > 2.55e38).""" + """amax <= 6*2^125 is deployment-bitwise; above, satfinite at the cap while TileKernels moves to scale 2^126.""" s125 = 6.0 * 2.0**125 w = torch.zeros(2, 32, dtype=torch.float32) w[0, 0] = s125 @@ -727,9 +692,7 @@ def test_deployment_top_cap_domain(): def test_tilekernels_cross_parity(): - """Cross-backend golden: TileKernels' torch reference (deployment semantics, - (1,32) e2m1 cast -> dequant) must reproduce mxfp4_fake_quantize bit-for-bit - on the finite domain with amax <= 6*2^125.""" + """Bitwise parity with the TileKernels torch reference (round_sf=True) on the finite below-cap domain.""" import sys tk_root = os.environ.get("NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels")) if os.path.isdir(tk_root) and tk_root not in sys.path: @@ -753,10 +716,7 @@ def test_tilekernels_cross_parity(): def test_blockwise_feasibility_enumeration(): - """128x128 exactness contract, enumerated: for exponent spread d = 0..16 and - every signed E2M1 payload, predict from the E4M3 lattice (subnormals - included) whether p*2^(t-d) survives the tile encode at the tile's actual - scale, and check TE's quantizer agrees element-for-element.""" + """Per-element E4M3-lattice prediction == TE 128x128 quantizer for exponent spreads d=0..16.""" payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) boundary = None for d in range(17): @@ -786,9 +746,7 @@ def test_blockwise_feasibility_enumeration(): def test_fp16_pipeline_rejected(): - """fp16 cannot represent the MXFP4 grid (top 6*2^125 >> fp16 max): both the - fp16-weight path and the fp16 activation/dequantize dtype must fail loudly - under MXFP4 QAT instead of silently overflowing.""" + """fp16 weights and fp16 activation dtypes must fail loudly under MXFP4 QAT.""" mod = te.Linear(256, 128, bias=False, params_dtype=torch.float16, device=DEV) x = torch.randn(64, 256, device=DEV, dtype=torch.float16) for recipe in (MXFP4QATMXFP8BlockScaling(), MXFP4QATFloat8BlockScaling()): @@ -801,11 +759,7 @@ def test_fp16_pipeline_rejected(): def test_fourway_matrix(): - """TileKernels deployment reference vs CuTe DSL vs CUDA kernel vs torch - reference, bitwise, over every previously-hit edge domain and every - quant/dequant step: fake-quant (bf16+fp32), mxfp8 row encode/dequant, - mxfp8 col encode/dequant, blockwise 128x128 encode/dequant, bf16+fp32 - dequant targets, bf16-exhaustive and fp32 bit-fuzz sweeps.""" + """TileKernels vs CuTe DSL vs CUDA vs torch, bitwise, over all edge domains and every quant/dequant hop.""" import sys tk_root = os.environ.get( "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") @@ -981,9 +935,7 @@ def tk_roundtrip(x, block, fmt): def test_recipe_switch_invalidates_cache(): - """base <-> QAT recipe switch must invalidate cached weight workspaces: a - non-first microbatch right after the switch has to requantize under the new - semantics instead of silently reusing the other recipe's weight.""" + """base<->QAT recipe switches must requantize instead of reusing cached weight workspaces.""" import warnings as _warnings from transformer_engine.common.recipe import MXFP8BlockScaling @@ -1011,9 +963,7 @@ def test_recipe_switch_invalidates_cache(): def test_ops_api_rejected(): - """The operations API bypasses the quantize_weight QAT hook, so it must - reject MXFP4 QAT recipes loudly instead of silently training without the - projection.""" + """te.ops bypasses the QAT weight hook and must reject MXFP4 QAT recipes.""" import transformer_engine.pytorch.ops as te_ops op = te_ops.Linear(256, 128, bias=False, device=DEV, dtype=torch.bfloat16) diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index c6ca1cf2f0..e19f40d3fd 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -419,11 +419,9 @@ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Project a BF16/FP32 tensor onto the MXFP4 (E2M1) grid with 1x32 - * power-of-two block scales and write back the dequantized values - * (weight fake-quantization for MXFP4 QAT). Non-finite values poison - * their whole 1x32 block with NaN. + * power-of-two block scales, writing back dequantized values. * - * \param[in] input Input tensor (BF16/FP32, element count divisible by 32). + * \param[in] input Input tensor (BF16/FP32, innermost dim divisible by 32). * \param[in,out] output Output tensor with the same dtype and shape. * \param[in] stream CUDA stream used for the operation. */ diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index d649a2fd9e..bee777b134 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -491,17 +491,10 @@ class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling): """ MXFP8 recipe with MXFP4 weight quantization-aware training. - Identical to :class:`MXFP8BlockScaling` except that weights are first - projected onto the MXFP4 (E2M1) grid with a 1x32 power-of-two scale - before the regular MXFP8 weight quantization. Because MXFP4 blocks are - 1x32-aligned with MXFP8 blocks and the E2M1 grid is a subset of E4M3, - the rowwise MXFP8 weight representation is an exact (lossless) encoding - of the MXFP4-grid weight; the columnwise (32x1) representation used by - dgrad is quantized from the same MXFP4-grid weight. Activations and - gradients are handled exactly as in the base recipe, and - ``backward_override`` keeps its base-recipe meaning ('high_precision' - uses the original unquantized weight for dgrad). Only bf16/fp32 weights - and activation dtypes are supported: the MXFP4 grid exceeds fp16 range. + Weights are projected onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid + before the regular MXFP8 weight quantization; the rowwise encoding of the + projected weight is lossless. Activations, gradients and + ``backward_override`` behave as in the base recipe. bf16/fp32 only. """ @@ -510,16 +503,10 @@ class MXFP4QATFloat8BlockScaling(Float8BlockScaling): """ Float8 block-scaling recipe with MXFP4 weight quantization-aware training. - Identical to :class:`Float8BlockScaling` except that weights are first - projected onto the MXFP4 (E2M1) grid with a 1x32 power-of-two scale - before the regular 128x128 blockwise FP8 weight quantization. The FP8 - weight encoding is exact (lossless) whenever every 1x32 MXFP4 scale - within a 128x128 tile lies within the FP8 dynamic-range headroom of the - tile's maximum scale; weights are quantized in square 128x128 blocks so - no rowwise/columnwise distinction applies (the transpose is exact). - Activations and gradients are handled exactly as in the base recipe. - Only bf16/fp32 weights and activation dtypes are supported: the MXFP4 - grid exceeds fp16 range. + Weights are projected onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid + before the regular 128x128 blockwise FP8 weight quantization (lossless + within each tile's FP8 dynamic-range headroom). Activations, gradients and + ``backward_override`` behave as in the base recipe. bf16/fp32 only. """ def __post_init__(self) -> None: diff --git a/transformer_engine/pytorch/mxfp4_qat.py b/transformer_engine/pytorch/mxfp4_qat.py index 14322ef457..66695aafb6 100644 --- a/transformer_engine/pytorch/mxfp4_qat.py +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -4,40 +4,13 @@ """MXFP4 weight fake-quantization for quantization-aware training. -Projects a weight onto the MXFP4 (E2M1) grid with 1x32 power-of-two (UE8M0) -block scales and returns the dequantized result in the input dtype — i.e. the -fused form of ``bf16 -> mxfp4(row) -> dequantize -> bf16``. The returned values -are exactly representable in bf16/fp32 (E2M1 needs at most two significant -mantissa bits and the scales are powers of two), so composing with the base -recipe's weight quantization stays lossless: - -* MXFP8 (1x32, E8M0): MXFP4 blocks are 1x32-aligned and the E2M1 grid is a - subset of E4M3 — the rowwise MXFP8 encoding of the projected weight is exact. -* Float8 block scaling (128x128, power-of-two scales): exact whenever every - 1x32 MXFP4 scale within a tile lies within the FP8 dynamic-range headroom of - the tile's maximum scale. - -Three bit-identical implementations are dispatched via ``NVTE_MXFP4_QAT_IMPL`` -(``auto``/``cuda``/``cute_dsl``/``torch``, default ``auto`` = CUDA binding, -then CuTe DSL, then the PyTorch reference): the CUDA kernel in -``common/cast/mxfp4/fake_quantize_mxfp4.cuh`` (``tex.mxfp4_fake_quantize``), -the CuTe DSL kernel in ``mxfp4_qat_cute_dsl.py``, and -``_mxfp4_fake_quantize_torch`` below. - -Value-domain semantics (identical in all implementations): -* zero block (amax == 0): scale 1, output zeros — exact. -* tiny amax (<= 6*2^-126): the scale floors at 2^-126 (TileKernels deployment - contract). The 0.5 payload then dequantizes to 2^-127, a bf16/fp32 - subnormal. NOTE: TE's MXFP8 software dequantize flushes that grid point to - zero until rebuilt with the ptx.cuh UE8M0 code-0 fix (raw encoding and - hardware GEMM are exact). -* non-finite amax (inf/NaN anywhere in a 1x32 block): the WHOLE block becomes - NaN. A single inf must not silently rescale its 31 neighbors to zero, and - NaN must not silently fall back to scale 1 — corruption is made visible. -* huge amax (> 6*2^125): the scale is capped at 2^125 so the dequantized grid - stays representable in bf16/fp32; values saturate at 6*2^125 (satfinite). -* fp16 is rejected for both weights and the activation/dequantize dtype: - even capped scales overflow fp16's range. +Projects weights onto the MXFP4 (E2M1) grid with 1x32 power-of-two block +scales and returns the dequantized values in the input dtype. The output is +exact in bf16/fp32, so the host recipe's weight quantization (MXFP8 rowwise, +128x128 blockwise) encodes it losslessly. Scale floor 2^-126 / cap 2^125 +(TileKernels deployment contract); non-finite values NaN-poison their 1x32 +block; fp16 is rejected. Implementation picked by NVTE_MXFP4_QAT_IMPL +(auto/cuda/cute_dsl/torch). """ import os import warnings @@ -69,12 +42,7 @@ def backward(ctx, grad_output): # pylint: disable=arguments-differ def _round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: - """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}. - - Input must be non-negative and <= 6. Within each binade the grid is - uniform, so torch.round (RTNE) on the rescaled value reproduces IEEE RTNE - exactly, including ties across binade boundaries (2.5 -> 2, 3.5 -> 4, 5 -> 4). - """ + """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" fine = torch.round(y * 2.0) * 0.5 mid = torch.round(y) coarse = torch.round(y * 0.5) * 2.0 @@ -187,17 +155,9 @@ def _resolve_impl(weight: torch.Tensor): def mxfp4_fake_quantize(weight: torch.Tensor) -> torch.Tensor: """Project ``weight`` onto the MXFP4 grid and return it in the input dtype. - Per 1x32 block: scale = 2^ceil(log2(amax / 6)) (all-zero blocks use 1.0), - values are divided by the scale, rounded RTNE onto the E2M1 grid with - saturation at +-6, and multiplied back. Blocks containing non-finite - values become NaN. The gradient is the straight-through estimator - (identity), the QAT-correct semantics. - - Env knobs: ``NVTE_MXFP4_QAT_IMPL`` picks the implementation - (``auto``/``cuda``/``cute_dsl``/``torch``); - ``NVTE_MXFP4_QAT_REQUIRE_KERNEL=1`` makes a kernel-less fallback a hard - error instead of a warn-once (recommended in production so a stale wheel - cannot silently train on the slow path). + Per 1x32 block: scale = 2^clamp(ceil(log2(amax / 6)), -126, 125), RTNE onto + E2M1, saturate at 6, multiply back. Gradient is the straight-through + estimator. ``NVTE_MXFP4_QAT_REQUIRE_KERNEL=1`` forbids the torch fallback. """ if weight.dim() != 2: raise ValueError(f"MXFP4 QAT expects a 2D weight, got {tuple(weight.shape)}") diff --git a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py index e4f07774ec..08624116bd 100644 --- a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py +++ b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py @@ -2,15 +2,11 @@ # # See LICENSE for license information. -"""CuTe DSL implementation of MXFP4 weight fake-quantization for QAT. - -Bit-identical to the CUDA kernel in -``common/cast/mxfp4/fake_quantize_mxfp4.cuh`` and to the PyTorch reference in -``transformer_engine/pytorch/mxfp4_qat.py``: the block-scale exponent is -derived in integer arithmetic and the value path uses explicit non-FTZ PTX, -so results are insensitive to compilation modes. Requires the cutlass CuTe -DSL package; selected via ``NVTE_MXFP4_QAT_IMPL=cute_dsl`` or as the ``auto`` -fallback when the CUDA binding is unavailable. +"""CuTe DSL implementation of MXFP4 weight fake-quantization. + +Bit-identical to the CUDA kernel and the PyTorch reference: integer-bit scale +derivation and non-FTZ inline PTX on the value path. Requires the cutlass +CuTe DSL package. """ import functools From 7a820af1b5ad592fbaa04cbca85f44d497526acb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 18:12:51 +0000 Subject: [PATCH 05/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_mxfp4_qat.py | 292 +++++++++++------- transformer_engine/common/cast/cast.cu | 3 +- .../common/cast/mxfp4/fake_quantize_mxfp4.cuh | 16 +- .../pytorch/csrc/extensions/cast.cpp | 3 +- .../pytorch/csrc/extensions/pybind.cpp | 4 +- .../pytorch/mxfp4_qat_cute_dsl.py | 6 +- .../pytorch/ops/basic/grouped_linear.py | 5 +- .../pytorch/ops/fused/grouped_mlp.py | 5 +- 8 files changed, 195 insertions(+), 139 deletions(-) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index 7e210284f0..68a6dc43b5 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -72,21 +72,15 @@ def fake_quant_ref_fp64(w): cand = torch.cat([-grid.flip(0)[:-1], grid]) d = (y.unsqueeze(-1) - cand).abs() near = d.argmin(dim=-1) - even = torch.tensor( - [abs(v) in (0.0, 1.0, 2.0, 4.0) for v in cand.tolist()], device=w.device - ) + even = torch.tensor([abs(v) in (0.0, 1.0, 2.0, 4.0) for v in cand.tolist()], device=w.device) dmin = d.gather(-1, near.unsqueeze(-1)).squeeze(-1) tie_even = d.eq(dmin.unsqueeze(-1)) & even has_even_tie = tie_even.any(dim=-1) - near_even = torch.where( - has_even_tie, tie_even.to(torch.uint8).argmax(dim=-1), near - ) + near_even = torch.where(has_even_tie, tie_even.to(torch.uint8).argmax(dim=-1), near) q = cand[near_even] return (q * scale).view(m, n).to(torch.float64) - - def test_fake_quant_grid_scale_rtne(): for m, n in SHAPES: w = make_weight(m, n, zero_blocks=4, outliers=4) @@ -118,11 +112,9 @@ def test_D_lossless_in_bf16(): w = make_weight(m, n, outliers=4) w_hat_bf16 = mxfp4_fake_quantize(w) w_hat_fp32 = mxfp4_fake_quantize(w.to(torch.float32)) - assert torch.equal(w_hat_bf16.to(torch.float32), w_hat_fp32), ( - "MXFP4-grid weight is not exact in bf16" - ) - - + assert torch.equal( + w_hat_bf16.to(torch.float32), w_hat_fp32 + ), "MXFP4-grid weight is not exact in bf16" def test_E_mxfp8_rowwise_lossless(): @@ -140,15 +132,15 @@ def test_E_mxfp8_rowwise_lossless(): raw_vals = data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) raw_scale = torch.ldexp(torch.ones_like(codes, dtype=torch.float32), codes - 127) raw_dq = (raw_vals * raw_scale.unsqueeze(-1)).view(m, n) - assert torch.equal(raw_dq, wh32), ( - f"RAW MXFP8 encoding of the MXFP4-grid weight is not lossless {m}x{n}" - ) + assert torch.equal( + raw_dq, wh32 + ), f"RAW MXFP8 encoding of the MXFP4-grid weight is not lossless {m}x{n}" flushed = wh32.abs() == 2.0**-127 assert flushed.any() and not flushed.all(), "artifact case not exercised" - assert torch.equal(dq[~flushed], wh32[~flushed]), ( - f"MXFP8 rowwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" - ) + assert torch.equal( + dq[~flushed], wh32[~flushed] + ), f"MXFP8 rowwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" if not (dq[flushed] == 0).all(): assert torch.equal(dq, wh32), "dequant neither flushed nor exact?" print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") @@ -187,14 +179,14 @@ def test_G_blockwise_lossless_and_transpose(): block_scaling_dim=2, ).quantize(w_hat) dq = q.dequantize(dtype=torch.float32) - assert torch.equal(dq, w_hat.to(torch.float32)), ( - f"128x128 blockwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" - ) + assert torch.equal( + dq, w_hat.to(torch.float32) + ), f"128x128 blockwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" q.update_usage(rowwise_usage=False, columnwise_usage=True) dqc = q.dequantize(dtype=torch.float32) - assert torch.equal(dqc, w_hat.to(torch.float32)), ( - f"columnwise 128x128 dequant differs from the MXFP4-grid weight {m}x{n}" - ) + assert torch.equal( + dqc, w_hat.to(torch.float32) + ), f"columnwise 128x128 dequant differs from the MXFP4-grid weight {m}x{n}" def test_lossless_dequant_to_bf16(): @@ -208,21 +200,24 @@ def test_lossless_dequant_to_bf16(): dq8 = q8.dequantize(dtype=torch.bfloat16) flushed = w_hat.to(torch.float32).abs() == 2.0**-127 assert flushed.any(), "artifact case not exercised" - assert torch.equal(dq8[~flushed], w_hat[~flushed]), ( - f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" - ) + assert torch.equal( + dq8[~flushed], w_hat[~flushed] + ), f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" if not (dq8[flushed] == 0).all(): assert torch.equal(dq8, w_hat), "dequant neither flushed nor exact?" print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") w2_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4)) qb = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=True, block_scaling_dim=2, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, ).quantize(w2_hat) - assert torch.equal(qb.dequantize(dtype=torch.bfloat16), w2_hat), ( - f"blockwise 128x128 dequant to bf16 not lossless {m}x{n}" - ) + assert torch.equal( + qb.dequantize(dtype=torch.bfloat16), w2_hat + ), f"blockwise 128x128 dequant to bf16 not lossless {m}x{n}" def test_C_mxfp8_colwise_bound(): @@ -238,10 +233,6 @@ def test_C_mxfp8_colwise_bound(): assert (err <= bound + 1e-30).all(), f"colwise 32x1 requant error above bound {m}x{n}" - - - - _INTREE_KERNELS = {} @@ -255,7 +246,10 @@ def _intree_kernel(fast_math=False): "NVTE_MXFP4_QAT_TEST_SRC", os.path.join( os.path.dirname(os.path.abspath(__file__)), - "..", "..", "transformer_engine", "common", + "..", + "..", + "transformer_engine", + "common", ), ) cuda_src = r""" @@ -294,7 +288,8 @@ def _intree_kernel(fast_math=False): "-U__CUDA_NO_HALF2_OPERATORS__", "-U__CUDA_NO_BFLOAT16_OPERATORS__", "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", - ] + (["--use_fast_math"] if fast_math else []), + ] + + (["--use_fast_math"] if fast_math else []), verbose=False, ) return _INTREE_KERNELS[fast_math] @@ -302,6 +297,7 @@ def _intree_kernel(fast_math=False): def _both_paths(w): import transformer_engine.pytorch.mxfp4_qat as M + got_kernel = _intree_kernel().mxfp4_fake_quantize_intree(w.contiguous()) got_torch = M._mxfp4_fake_quantize_torch(w) return got_kernel, got_torch @@ -342,25 +338,29 @@ def test_edge_value_domains(): gk, gt = _both_paths(z) assert torch.equal(gk, z) and torch.equal(gt, z) - wz = w.clone(); wz[0, :32] = -0.0 + wz = w.clone() + wz[0, :32] = -0.0 gk, gt = _both_paths(wz) assert (gk[0, :32] == 0).all() and torch.signbit(gk[0, :32]).all(), "-0.0 sign lost" _assert_bits_equal(gk, gt, "-0 block") - ws = w.clone(); ws[1, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) + ws = w.clone() + ws[1, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) gk, gt = _both_paths(ws) assert torch.isfinite(gk[1, :32]).all() and torch.equal(gk, gt) ref = fake_quant_ref_fp64(ws) assert torch.equal(gt.to(torch.float64), ref), "subnormal block deviates from fp64 ref" - wp = w.clone(); wp[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + wp = w.clone() + wp[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) gk, gt = _both_paths(wp) assert (gk[5, :32].to(torch.float32) == 2.0**-127).all(), "2^-127 grid point lost" _assert_bits_equal(gk, gt, "2^-127 grid point") ref = fake_quant_ref_fp64(wp) assert torch.equal(gt.to(torch.float64), ref), "2^-127 block deviates from fp64 ref" - wh = w.clone(); wh[2, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) + wh = w.clone() + wh[2, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) wh[2, 0] = -3.0e38 gk, gt = _both_paths(wh) assert torch.isfinite(gk[2, :32]).all() and torch.equal(gk, gt) @@ -368,12 +368,14 @@ def test_edge_value_domains(): assert torch.equal(gt.to(torch.float64), ref), "huge block deviates from fp64 ref" for bad in (float("inf"), float("-inf"), float("nan")): - wb = w.clone(); wb[3, 40] = bad + wb = w.clone() + wb[3, 40] = bad gk, gt = _both_paths(wb) _assert_same_with_nan(gk, gt, f"bad={bad}") assert torch.isnan(gk[3, 32:64]).all(), f"{bad}: block not NaN-poisoned" assert torch.isfinite(gk[3, :32]).all(), f"{bad}: neighbor block affected" - clean_rows = torch.ones(m, dtype=torch.bool); clean_rows[3] = False + clean_rows = torch.ones(m, dtype=torch.bool) + clean_rows[3] = False assert torch.isfinite(gk[clean_rows]).all() try: @@ -388,8 +390,11 @@ def test_blockwise_recipe_is_128x128(): assert r.x_block_scaling_dim == 1 and r.w_block_scaling_dim == 2 assert r.grad_block_scaling_dim == 1 q = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=True, block_scaling_dim=2, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, ) assert q.block_len == 128 w_hat = mxfp4_fake_quantize(make_weight(256, 512)) @@ -403,8 +408,11 @@ def test_blockwise_over_ratio_is_bounded_not_lossless(): w[64, 0] = 100.0 w_hat = mxfp4_fake_quantize(w) q = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=True, block_scaling_dim=2, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, ).quantize(w_hat) dq = q.dequantize(dtype=torch.float32) assert torch.isfinite(dq).all() @@ -525,10 +533,9 @@ def test_kernel_fast_math_immune(): got = k.mxfp4_fake_quantize_intree(w.contiguous()) ref = M._mxfp4_fake_quantize_torch(w) _assert_same_with_nan(got, ref, "fast-math edge") - assert (got[5, :32].to(torch.float32) == 2.0**-127).all(), ( - "fast-math build flushed the 2^-127 grid point" - ) - + assert ( + got[5, :32].to(torch.float32) == 2.0**-127 + ).all(), "fast-math build flushed the 2^-127 grid point" def test_exp2f_e8m0_all_codes(): @@ -538,12 +545,20 @@ def test_exp2f_e8m0_all_codes(): src_dir = os.environ.get( "NVTE_MXFP4_QAT_TEST_SRC", - os.path.join(os.path.dirname(os.path.abspath(__file__)), - "..", "..", "transformer_engine", "common"), + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "..", "transformer_engine", "common" + ), ) ptx_src = open(os.path.join(src_dir, "util", "ptx.cuh")).read() - m_exp = _re.search(r"__device__ __forceinline__ float exp2f\(e8m0_t biased_exp\) \{.*?\n\}", ptx_src, _re.S) - m_rcp = _re.search(r"template <>\n__device__ __forceinline__ float exp2f_rcp\(e8m0_t biased_exp\) \{.*?\n\}", ptx_src, _re.S) + m_exp = _re.search( + r"__device__ __forceinline__ float exp2f\(e8m0_t biased_exp\) \{.*?\n\}", ptx_src, _re.S + ) + m_rcp = _re.search( + r"template <>\n__device__ __forceinline__ float exp2f_rcp\(e8m0_t biased_exp\)" + r" \{.*?\n\}", + ptx_src, + _re.S, + ) assert m_exp and m_rcp, "could not extract exp2f/exp2f_rcp from ptx.cuh" cuda_src = ( "#include \n#include \n" @@ -552,7 +567,10 @@ def test_exp2f_e8m0_all_codes(): "typedef unsigned char e8m0_t;\n" "constexpr unsigned FP32_MANTISSA_BITS = 23;\n" "template __device__ __forceinline__ T exp2f_rcp(e8m0_t);\n" - + m_rcp.group(0) + "\n" + m_exp.group(0) + "\n}\n" + + m_rcp.group(0) + + "\n" + + m_exp.group(0) + + "\n}\n" "__global__ void all_codes_kernel(float *e, float *r) {\n" " const int b = threadIdx.x;\n" " e[b] = teptx::exp2f(static_cast(b));\n" @@ -566,7 +584,9 @@ def test_exp2f_e8m0_all_codes(): mod = load_inline( name="nvte_mxfp4_qat_exp2f_test", cpp_sources="void run_all_codes(torch::Tensor e, torch::Tensor r);", - cuda_sources=cuda_src, functions=["run_all_codes"], verbose=False, + cuda_sources=cuda_src, + functions=["run_all_codes"], + verbose=False, ) e = torch.empty(256, dtype=torch.float32, device=DEV) r = torch.empty_like(e) @@ -639,7 +659,7 @@ def test_fp32_bit_fuzz(): """Random raw fp32 bit patterns (subnormals/NaN/Inf included).""" g = torch.Generator().manual_seed(99) for _ in range(4): - bits = torch.randint(-2**31, 2**31 - 1, (4096, 32), generator=g, dtype=torch.int64) + bits = torch.randint(-(2**31), 2**31 - 1, (4096, 32), generator=g, dtype=torch.int64) w = bits.to(torch.int32).cuda().view(torch.float32) gk, gt = _both_paths(w) _assert_same_with_nan(gk, gt, "fp32 fuzz") @@ -668,12 +688,18 @@ def test_scale_threshold_and_rtne_midpoints(): expected = torch.tensor([0.0, 1.0, 1.0, 2.0, 2.0, 4.0, 4.0]) for t in (-60, 0, 60): sc = 2.0**t - block = torch.zeros(32); block[:7] = mids * sc; block[7] = 6.0 * sc + block = torch.zeros(32) + block[:7] = mids * sc + block[7] = 6.0 * sc w = block.view(1, 32).to(torch.float32).cuda() gk, gt = _both_paths(w) _assert_bits_equal(gk, gt, f"midpoints t={t}") - exp_row = torch.zeros(32); exp_row[:7] = expected * sc; exp_row[7] = 6.0 * sc - assert torch.equal(gt[0].cpu().to(torch.float32), exp_row.to(torch.float32)), f"midpoint values t={t}" + exp_row = torch.zeros(32) + exp_row[:7] = expected * sc + exp_row[7] = 6.0 * sc + assert torch.equal( + gt[0].cpu().to(torch.float32), exp_row.to(torch.float32) + ), f"midpoint values t={t}" def test_deployment_top_cap_domain(): @@ -694,7 +720,10 @@ def test_deployment_top_cap_domain(): def test_tilekernels_cross_parity(): """Bitwise parity with the TileKernels torch reference (round_sf=True) on the finite below-cap domain.""" import sys - tk_root = os.environ.get("NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels")) + + tk_root = os.environ.get( + "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") + ) if os.path.isdir(tk_root) and tk_root not in sys.path: sys.path.insert(0, tk_root) try: @@ -722,14 +751,17 @@ def test_blockwise_feasibility_enumeration(): for d in range(17): w = torch.zeros(128, 128, dtype=torch.float32) w[0, 0] = 6.0 - vals = payloads * 2.0**(-d) + vals = payloads * 2.0 ** (-d) w[1, :7] = vals w[1, 7:14] = -vals wc = w.cuda() assert torch.equal(mxfp4_fake_quantize(wc), wc), f"tile not on MXFP4 grid (d={d})" q = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False, - force_pow_2_scales=True, block_scaling_dim=2, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + force_pow_2_scales=True, + block_scaling_dim=2, ).quantize(wc) S = q._rowwise_scale_inv.reshape(-1)[0].item() dq = q.dequantize(dtype=torch.float32) @@ -738,11 +770,13 @@ def test_blockwise_feasibility_enumeration(): assert torch.equal(pred, got), f"E4M3-lattice prediction mismatch at spread d={d}" if boundary is None and not bool(pred.all()): boundary = d - assert boundary is not None and boundary >= 12, ( - f"full-payload exactness should hold at least through spread 2^11, got {boundary}" + assert ( + boundary is not None and boundary >= 12 + ), f"full-payload exactness should hold at least through spread 2^11, got {boundary}" + print( + f" full-grid exact through spread d={boundary - 1}; first loss at d={boundary} " + "(E4M3 subnormal lattice bound)" ) - print(f" full-grid exact through spread d={boundary - 1}; first loss at d={boundary} " - "(E4M3 subnormal lattice bound)") def test_fp16_pipeline_rejected(): @@ -761,6 +795,7 @@ def test_fp16_pipeline_rejected(): def test_fourway_matrix(): """TileKernels vs CuTe DSL vs CUDA vs torch, bitwise, over all edge domains and every quant/dequant hop.""" import sys + tk_root = os.environ.get( "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") ) @@ -835,8 +870,10 @@ def tk_roundtrip(x, block, fmt): _assert_bits_equal(dsl(wcap), t, "cap DSL") tk32, _ = tk_roundtrip(wcap, (1, 32), "e2m1") assert not torch.equal(tk32, t.to(torch.float32)), "TK must diverge above 6*2^125" - print(f" above-cap domain: ours satfinite {CAP:.3e}, TK diverges as documented " - f"(payload at 2^126 -> {tk32[0, 0].item()}): PASS") + print( + f" above-cap domain: ours satfinite {CAP:.3e}, TK diverges as documented " + f"(payload at 2^126 -> {tk32[0, 0].item()}): PASS" + ) w_hat = tref(w) wh32 = w_hat.to(torch.float32) @@ -844,20 +881,21 @@ def tk_roundtrip(x, block, fmt): m, n = w_hat.shape data = q._rowwise_data.view(torch.uint8)[:m, :n] codes = q._rowwise_scale_inv.view(torch.uint8)[:m, : n // 32].to(torch.int32) - raw = (data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) - * torch.ldexp(torch.ones_like(codes, dtype=torch.float32), codes - 127).unsqueeze(-1)) + raw = data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) * torch.ldexp( + torch.ones_like(codes, dtype=torch.float32), codes - 127 + ).unsqueeze(-1) assert torch.equal(raw.view(m, n), wh32), "mxfp8 row RAW encode not lossless" e4m3_floor = float(tk_min_clamp(torch.float8_e4m3fn)) blk_amax = wh32.abs().view(m, n // 32, 32).amax(dim=-1) - ok = ((blk_amax >= e4m3_floor) | (blk_amax == 0)) + ok = (blk_amax >= e4m3_floor) | (blk_amax == 0) okm = ok.unsqueeze(-1).expand(m, n // 32, 32).reshape(m, n) tk32, tkb16 = tk_roundtrip(w_hat, (1, 32), "e4m3") - assert torch.equal(tk32.view(torch.int32)[okm], wh32.view(torch.int32)[okm]), ( - "mxfp8 row TK roundtrip (fp32)" - ) - assert torch.equal(tkb16.view(torch.int16)[okm], w_hat.view(torch.int16)[okm]), ( - "mxfp8 row TK roundtrip (bf16)" - ) + assert torch.equal( + tk32.view(torch.int32)[okm], wh32.view(torch.int32)[okm] + ), "mxfp8 row TK roundtrip (fp32)" + assert torch.equal( + tkb16.view(torch.int16)[okm], w_hat.view(torch.int16)[okm] + ), "mxfp8 row TK roundtrip (bf16)" assert (tk32[~okm] == 0).all(), "TK e4m3 sub-floor blocks must flush to zero" assert (~ok).sum() > 0, "TK e4m3 floor case not exercised" dq = q.dequantize(dtype=torch.float32) @@ -866,21 +904,23 @@ def tk_roundtrip(x, block, fmt): if flushed.any() and not (dq[flushed] == 0).all(): assert torch.equal(dq, wh32) print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") - print(f" mxfp8 row: TE raw encode lossless incl 2^-127; TK roundtrip bitwise on " - f"amax >= {e4m3_floor:g} blocks ({int((~ok).sum())} sub-floor blocks flush in TK " - "e4m3, a TK-only floor); TE software dequant exact outside the code-0 wheel bug: PASS") + print( + " mxfp8 row: TE raw encode lossless incl 2^-127; TK roundtrip bitwise on " + f"amax >= {e4m3_floor:g} blocks ({int((~ok).sum())} sub-floor blocks flush in TK " + "e4m3, a TK-only floor); TE software dequant exact outside the code-0 wheel bug: PASS" + ) q.update_usage(rowwise_usage=False, columnwise_usage=True) te_col = q.dequantize(dtype=torch.float32) tk_col32, _ = tk_roundtrip(w_hat, (32, 1), "e4m3") col_amax = wh32.abs().view(m // 32, 32, n).amax(dim=1) - cok = ((col_amax >= e4m3_floor) | (col_amax == 0)) + cok = (col_amax >= e4m3_floor) | (col_amax == 0) cokm = cok.unsqueeze(1).expand(m // 32, 32, n).reshape(m, n) nzm = (wh32 != 0) & cokm zm = (wh32 == 0) & cokm - assert torch.equal(te_col.view(torch.int32)[nzm], tk_col32.view(torch.int32)[nzm]), ( - "mxfp8 col TE vs TK (nonzero)" - ) + assert torch.equal( + te_col.view(torch.int32)[nzm], tk_col32.view(torch.int32)[nzm] + ), "mxfp8 col TE vs TK (nonzero)" assert (te_col[zm] == 0).all() and (tk_col32[zm] == 0).all(), "mxfp8 col zeros" negz = zm & torch.signbit(wh32) assert negz.any() @@ -889,14 +929,19 @@ def tk_roundtrip(x, block, fmt): blk = wh32.view(m // 32, 32, n) bound = blk.abs().amax(dim=1, keepdim=True) * 2.0**-3 assert ((te_col.view(m // 32, 32, n) - blk).abs() <= bound + 1e-30).all() - print(" mxfp8 col (32x1): TE == TK bitwise on nonzeros, zeros value-equal " - "(TE col canonicalizes -0 -> +0, TK preserves the sign), bounded vs w_hat: PASS") + print( + " mxfp8 col (32x1): TE == TK bitwise on nonzeros, zeros value-equal " + "(TE col canonicalizes -0 -> +0, TK preserves the sign), bounded vs w_hat: PASS" + ) wb_hat = tref(make_weight(128, 256, zero_blocks=4)) wb32 = wb_hat.to(torch.float32) qb = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=True, block_scaling_dim=2, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, ).quantize(wb_hat) assert torch.equal(qb.dequantize(dtype=torch.float32), wb32) assert torch.equal(qb.dequantize(dtype=torch.bfloat16), wb_hat) @@ -905,8 +950,10 @@ def tk_roundtrip(x, block, fmt): _assert_bits_equal(tkb16, wb_hat, "blockwise TK roundtrip (bf16)") qb.update_usage(rowwise_usage=False, columnwise_usage=True) assert torch.equal(qb.dequantize(dtype=torch.float32), wb32), "blockwise col transpose" - print(" blockwise 128x128: TE row+col dequant == TK roundtrip == w_hat bitwise " - "(fp32 + bf16 targets): PASS") + print( + " blockwise 128x128: TE row+col dequant == TK roundtrip == w_hat bitwise " + "(fp32 + bf16 targets): PASS" + ) bits = torch.arange(65536, dtype=torch.int32, device=DEV).to(torch.int16) we = bits.view(torch.bfloat16).view(2048, 32) @@ -917,11 +964,13 @@ def tk_roundtrip(x, block, fmt): ok_rows = torch.isfinite(we32).all(dim=-1) & (we32.abs().amax(dim=-1) <= CAP) tk32, _ = tk_roundtrip(we[ok_rows], (1, 32), "e2m1") _assert_bits_equal(tk32, t[ok_rows].to(torch.float32), "exhaustive TK") - print(f" bf16 exhaustive 65536 patterns: CUDA/DSL/torch 3-way full, " - f"TK on {int(ok_rows.sum())}/2048 finite below-cap rows, bitwise: PASS") + print( + " bf16 exhaustive 65536 patterns: CUDA/DSL/torch 3-way full, " + f"TK on {int(ok_rows.sum())}/2048 finite below-cap rows, bitwise: PASS" + ) g = torch.Generator().manual_seed(123) - fb = torch.randint(-2**31, 2**31 - 1, (512, 32), generator=g, dtype=torch.int64) + fb = torch.randint(-(2**31), 2**31 - 1, (512, 32), generator=g, dtype=torch.int64) wf = fb.to(torch.int32).cuda().view(torch.float32) t = tref(wf) _assert_same_with_nan(kern(wf.contiguous()), t, "fuzz CUDA") @@ -930,8 +979,10 @@ def tk_roundtrip(x, block, fmt): if ok_rows.any(): tk32, _ = tk_roundtrip(wf[ok_rows], (1, 32), "e2m1") _assert_bits_equal(tk32, t[ok_rows], "fuzz TK") - print(f" fp32 bit-fuzz 512 blocks: 3-way full, TK on {int(ok_rows.sum())} " - f"finite below-cap rows, bitwise: PASS") + print( + f" fp32 bit-fuzz 512 blocks: 3-way full, TK on {int(ok_rows.sum())} " + "finite below-cap rows, bitwise: PASS" + ) def test_recipe_switch_invalidates_cache(): @@ -951,15 +1002,13 @@ def test_recipe_switch_invalidates_cache(): with fp8_autocast(enabled=True, fp8_recipe=qat): y_switch = mod(x, is_first_microbatch=False) y_fresh = mod(x, is_first_microbatch=True) - assert torch.equal(y_switch, y_fresh), ( - "base->QAT switch reused an UNPROJECTED cached weight" - ) + assert torch.equal( + y_switch, y_fresh + ), "base->QAT switch reused an UNPROJECTED cached weight" with fp8_autocast(enabled=True, fp8_recipe=base): y2_switch = mod(x, is_first_microbatch=False) y2_fresh = mod(x, is_first_microbatch=True) - assert torch.equal(y2_switch, y2_fresh), ( - "QAT->base switch reused a PROJECTED cached weight" - ) + assert torch.equal(y2_switch, y2_fresh), "QAT->base switch reused a PROJECTED cached weight" def test_ops_api_rejected(): @@ -977,8 +1026,6 @@ def test_ops_api_rejected(): pass - - def _run_module(module_kind, recipe, override, fuse): torch.manual_seed(7) m_splits = [192, 320] @@ -987,13 +1034,22 @@ def _run_module(module_kind, recipe, override, fuse): if module_kind == "linear": mod = te.Linear( - in_f, out_f, bias=False, params_dtype=torch.bfloat16, device=DEV, + in_f, + out_f, + bias=False, + params_dtype=torch.bfloat16, + device=DEV, fuse_wgrad_accumulation=fuse, ) params = [mod.weight] else: mod = te.GroupedLinear( - 2, in_f, out_f, bias=False, params_dtype=torch.bfloat16, device=DEV, + 2, + in_f, + out_f, + bias=False, + params_dtype=torch.bfloat16, + device=DEV, fuse_wgrad_accumulation=fuse, ) params = [mod.weight0, mod.weight1] @@ -1006,9 +1062,7 @@ def _run_module(module_kind, recipe, override, fuse): p.main_grad = torch.zeros_like(p, dtype=torch.float32) p.grad_added_to_main_grad = False - x = (torch.randn(tokens, in_f, device=DEV, dtype=torch.float32) * 0.5).to( - torch.bfloat16 - ) + x = (torch.randn(tokens, in_f, device=DEV, dtype=torch.float32) * 0.5).to(torch.bfloat16) x.requires_grad_(True) with fp8_autocast(enabled=True, fp8_recipe=recipe): if module_kind == "linear": @@ -1023,8 +1077,11 @@ def cat_ref(weights, mat, transpose): return torch.cat( [ mat[offs[i] : offs[i + 1]].to(torch.float32) - @ (weights[i].detach().to(torch.float32).t() if transpose - else weights[i].detach().to(torch.float32)) + @ ( + weights[i].detach().to(torch.float32).t() + if transpose + else weights[i].detach().to(torch.float32) + ) for i in range(nsplit) ] ) @@ -1060,9 +1117,8 @@ def cat_ref(weights, mat, transpose): wgrad_tol = BF16_TOL if override == "high_precision" else FP8_TOL_2OP for i, p in enumerate(params): - ref_wg = ( - g[offs[i] : offs[i + 1]].to(torch.float32).t() - @ x[offs[i] : offs[i + 1]].to(torch.float32) + ref_wg = g[offs[i] : offs[i + 1]].to(torch.float32).t() @ x[offs[i] : offs[i + 1]].to( + torch.float32 ) if fuse: assert p.grad_added_to_main_grad is True, "fused-wgrad flag not set (G)" diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 35c6d48fcb..bf44e3bffb 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -54,7 +54,8 @@ void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStr const Tensor &input_tensor = *convertNVTETensorCheck(input); Tensor &output_tensor = *convertNVTETensorCheck(output); NVTE_CHECK(input_tensor.dtype() == DType::kBFloat16 || input_tensor.dtype() == DType::kFloat32, - "MXFP4 fake-quantize supports BF16/FP32 inputs, got ", to_string(input_tensor.dtype())); + "MXFP4 fake-quantize supports BF16/FP32 inputs, got ", + to_string(input_tensor.dtype())); NVTE_CHECK(output_tensor.dtype() == input_tensor.dtype(), "MXFP4 fake-quantize output dtype must match the input dtype."); NVTE_CHECK(input_tensor.data.shape == output_tensor.data.shape, diff --git a/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh index 5d96ec21e5..6b4c437221 100644 --- a/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh +++ b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh @@ -4,7 +4,6 @@ * See LICENSE for license information. ************************************************************************/ - #ifndef TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_ #define TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_ @@ -75,10 +74,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) amax_bits = max(amax_bits, __float_as_int(elts[i]) & 0x7fffffff); } const unsigned lane = threadIdx.x & 31u; - const unsigned group_mask = - (LANES_PER_BLOCK == 32) - ? 0xffffffffu - : (((1u << LANES_PER_BLOCK) - 1u) << (lane & ~(LANES_PER_BLOCK - 1u))); + const unsigned group_mask = (LANES_PER_BLOCK == 32) ? 0xffffffffu + : (((1u << LANES_PER_BLOCK) - 1u) + << (lane & ~(LANES_PER_BLOCK - 1u))); #pragma unroll for (int offset = LANES_PER_BLOCK / 2; offset > 0; offset >>= 1) { amax_bits = max(amax_bits, __shfl_xor_sync(group_mask, amax_bits, offset)); @@ -108,7 +106,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } } -} +} // namespace fake_quantize_kernel template void fake_quantize_mxfp4_launch(const IType *input, IType *output, const size_t numel, @@ -122,8 +120,8 @@ void fake_quantize_mxfp4_launch(const IType *input, IType *output, const size_t <<>>(input, output, num_vectors); } -} -} -} +} // namespace mxfp4 +} // namespace dispatch +} // namespace transformer_engine #endif diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 6ff1662d83..41dcd03e7f 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -570,8 +570,7 @@ at::Tensor mxfp4_fake_quantize(const at::Tensor &input) { auto output = at::empty_like(input_contiguous); auto input_cpp = makeTransformerEngineTensor(input_contiguous); auto output_cpp = makeTransformerEngineTensor(output); - nvte_mxfp4_fake_quantize(input_cpp.data(), output_cpp.data(), - at::cuda::getCurrentCUDAStream()); + nvte_mxfp4_fake_quantize(input_cpp.data(), output_cpp.data(), at::cuda::getCurrentCUDAStream()); return output; } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index e28844c11d..facbb0c489 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -200,8 +200,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none(), py::arg("noop") = py::none()); m.def("mxfp4_fake_quantize", &transformer_engine::pytorch::mxfp4_fake_quantize, - "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)", - py::arg("input"), py::call_guard()); + "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)", py::arg("input"), + py::call_guard()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); m.def("create_empty_quantized_tensor", diff --git a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py index 08624116bd..98d3f5941f 100644 --- a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py +++ b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py @@ -37,11 +37,7 @@ def _is_nonfinite_f32(x: Float32, *, loc=None, ip=None) -> Int32: llvm.inline_asm( T.i32(), [Float32(x).ir_value(loc=loc, ip=ip)], - "{\n" - " .reg .pred p;\n" - " testp.finite.f32 p, $1;\n" - " selp.s32 $0, 0, 1, p;\n" - "}", + "{\n .reg .pred p;\n testp.finite.f32 p, $1;\n selp.s32 $0, 0, 1, p;\n}", "=r,f", loc=loc, ip=ip, diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 45cec84acc..01f245fdec 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -488,7 +488,10 @@ def _quantize_weights( ) -> Sequence[torch.Tensor]: """Construct quantized weight tensors.""" - if FP8GlobalStateManager.is_fp8_enabled() and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + if ( + FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + ): raise NotImplementedError( "MXFP4 QAT recipes do not support quantized primary weights in the " "operations API: the high-precision master is required for the QAT " diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 554317b6f2..0ed4f0fb81 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -961,7 +961,10 @@ def fuser_forward( fc2_input_quantizer = fc2_op.get_quantizer("forward", 0) fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1) fc2_grad_output_quantizer = fc2_op.get_quantizer("backward", 0) - if FP8GlobalStateManager.is_fp8_enabled() and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat(): + if ( + FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + ): raise NotImplementedError( "MXFP4 QAT recipes are not implemented for fused grouped-MLP ops: " "weight quantization here bypasses the QAT projection." From e83e43ae7f2f5113de23406e5b8313ebdfed5479 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 25 Jul 2026 11:41:00 -0700 Subject: [PATCH 06/14] Add mxfp4_qat_weights recipe field (NVTE_MXFP4_QAT env default) NVFP4-style control: MXFP4 weight QAT is now a field on the MXFP8 and Float8 block-scaling host recipes, defaulting from NVTE_MXFP4_QAT, so a stock --fp8-recipe mxfp8/blockwise launch enables QAT without framework changes. The QAT subclasses remain as the explicit API (field pinned True); blockwise validation moves to the host, guarded on the field. --- tests/pytorch/test_mxfp4_qat.py | 29 +++++++++++++++ transformer_engine/common/recipe/__init__.py | 37 +++++++++++--------- 2 files changed, 50 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index 68a6dc43b5..371b1ff1af 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -1011,6 +1011,34 @@ def test_recipe_switch_invalidates_cache(): assert torch.equal(y2_switch, y2_fresh), "QAT->base switch reused a PROJECTED cached weight" +def test_recipe_field_controls_qat(): + """mxfp4_qat_weights field (env NVTE_MXFP4_QAT) turns base recipes into QAT recipes.""" + import warnings as _warnings + from transformer_engine.common.recipe import Float8BlockScaling, MXFP8BlockScaling + + base = MXFP8BlockScaling(mxfp4_qat_weights=True) + assert base.mxfp4_qat() and MXFP4QATMXFP8BlockScaling().mxfp4_qat() + assert not MXFP8BlockScaling().mxfp4_qat() + + torch.manual_seed(9) + x = torch.randn(64, 256, device=DEV, dtype=torch.bfloat16) + mod = te.Linear(256, 128, bias=False, params_dtype=torch.bfloat16, device=DEV) + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + with fp8_autocast(enabled=True, fp8_recipe=base): + y_field = mod(x, is_first_microbatch=True) + with fp8_autocast(enabled=True, fp8_recipe=MXFP4QATMXFP8BlockScaling()): + y_class = mod(x, is_first_microbatch=True) + assert torch.equal(y_field, y_class), "field-driven QAT differs from class-driven QAT" + + assert Float8BlockScaling(mxfp4_qat_weights=True).mxfp4_qat() + try: + Float8BlockScaling(mxfp4_qat_weights=True, w_block_scaling_dim=1) + raise SystemExit("field-driven blockwise QAT must validate 2D weight blocks") + except ValueError: + pass + + def test_ops_api_rejected(): """te.ops bypasses the QAT weight hook and must reject MXFP4 QAT recipes.""" import transformer_engine.pytorch.ops as te_ops @@ -1164,6 +1192,7 @@ def test_e2e_matrix(): test_tilekernels_cross_parity, test_blockwise_feasibility_enumeration, test_recipe_switch_invalidates_cache, + test_recipe_field_controls_qat, test_ops_api_rejected, test_fp16_pipeline_rejected, test_fourway_matrix, diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index bee777b134..a7224f1e63 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -168,10 +168,11 @@ def custom(cls): """Whether the given recipe is custom.""" return issubclass(cls, CustomRecipe) - @classmethod - def mxfp4_qat(cls): - """Whether the given recipe applies MXFP4 weight QAT on top of its base recipe.""" - return issubclass(cls, (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling)) + def mxfp4_qat(self): + """Whether the recipe applies MXFP4 weight QAT on top of its base recipe.""" + return bool(getattr(self, "mxfp4_qat_weights", False)) or isinstance( + self, (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling) + ) @dataclass(repr=False) @@ -373,6 +374,7 @@ class MXFP8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + mxfp4_qat_weights: bool = os.getenv("NVTE_MXFP4_QAT", "0") == "1" def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -442,6 +444,7 @@ class Float8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + mxfp4_qat_weights: bool = os.getenv("NVTE_MXFP4_QAT", "0") == "1" def __post_init__(self) -> None: assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" @@ -466,6 +469,17 @@ def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + if self.mxfp4_qat_weights: + if self.w_block_scaling_dim != 2: + raise ValueError( + "MXFP4 QAT requires 128x128 (2D) weight scaling blocks, got " + f"w_block_scaling_dim={self.w_block_scaling_dim}." + ) + if not self.fp8_quant_fwd_weight.power_2_scale: + raise ValueError( + "MXFP4 QAT requires power-of-two weight scales; " + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 is incompatible." + ) def _make_repr(self) -> str: return ( @@ -497,6 +511,8 @@ class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling): ``backward_override`` behave as in the base recipe. bf16/fp32 only. """ + mxfp4_qat_weights: bool = True + @dataclass(repr=False) class MXFP4QATFloat8BlockScaling(Float8BlockScaling): @@ -509,18 +525,7 @@ class MXFP4QATFloat8BlockScaling(Float8BlockScaling): ``backward_override`` behave as in the base recipe. bf16/fp32 only. """ - def __post_init__(self) -> None: - super().__post_init__() - if self.w_block_scaling_dim != 2: - raise ValueError( - "MXFP4 QAT requires 128x128 (2D) weight scaling blocks, got " - f"w_block_scaling_dim={self.w_block_scaling_dim}." - ) - if not self.fp8_quant_fwd_weight.power_2_scale: - raise ValueError( - "MXFP4 QAT requires power-of-two weight scales; " - "NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 is incompatible." - ) + mxfp4_qat_weights: bool = True @dataclass(repr=False) From 8b40ab417703650a030149d4a6a63fe53257fe8d Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 25 Jul 2026 18:51:09 -0700 Subject: [PATCH 07/14] State the weight-encoding contract precisely: decoded-value exact, not raw-tuple identity The host encoder re-canonicalizes (payload, scale): TE-native normalizes the block amax into (224, 448] while a fixed-shift direct converter keeps the original factorization, so the raw E4M3/UE8M0 bytes generally differ between the two even though every decoded value is identical. 'Losslessly' in the module docstring was readable as raw-bit identity; say exactly what holds. --- transformer_engine/pytorch/mxfp4_qat.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/mxfp4_qat.py b/transformer_engine/pytorch/mxfp4_qat.py index 66695aafb6..ad161b2e75 100644 --- a/transformer_engine/pytorch/mxfp4_qat.py +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -6,8 +6,12 @@ Projects weights onto the MXFP4 (E2M1) grid with 1x32 power-of-two block scales and returns the dequantized values in the input dtype. The output is -exact in bf16/fp32, so the host recipe's weight quantization (MXFP8 rowwise, -128x128 blockwise) encodes it losslessly. Scale floor 2^-126 / cap 2^125 +exact in bf16/fp32, and the host recipe's weight quantization (MXFP8 rowwise; +128x128 blockwise within the tile scale-spread bound) is DECODED-VALUE exact +on this finite QAT domain. It is not raw-tuple identical to a direct +MXFP4->MXFP8 converter: the host encoder re-canonicalizes (payload, scale) +per block, so equal values may carry different E4M3/UE8M0 bytes, and the +original MXFP4 scale metadata is not preserved. Scale floor 2^-126 / cap 2^125 (TileKernels deployment contract); non-finite values NaN-poison their 1x32 block; fp16 is rejected. Implementation picked by NVTE_MXFP4_QAT_IMPL (auto/cuda/cute_dsl/torch). From cd0d68854972ddeb63a11a9d041dcde2fa58a6a0 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sun, 26 Jul 2026 11:18:51 -0700 Subject: [PATCH 08/14] Restructure the PyTorch reference and backends to upstream conventions - Move the composite-torch-ops reference implementation to tests/pytorch/references/mxfp4_qat_reference.py (the numerical oracle, matching the existing blockwise reference layout); the main module no longer carries a runtime fallback path. - transformer_engine/pytorch/mxfp4_qat.py now only validates inputs, applies the straight-through-estimator wrapper, and calls the tex.mxfp4_fake_quantize binding; a missing binding is a hard error and the NVTE_MXFP4_QAT_IMPL/REQUIRE_KERNEL/DISABLE_CUDA_KERNEL knobs are removed. - Drop the CuTe DSL implementation: the cast layer has no DSL precedent and TE carries no cutlass-DSL runtime dependency. - Tests compare the kernel against the reference directly and demand the fixed UE8M0 code-0/255 dequantize semantics unconditionally. --- .../pytorch/references/mxfp4_qat_reference.py | 49 ++++ tests/pytorch/test_mxfp4_qat.py | 161 +++----------- transformer_engine/pytorch/mxfp4_qat.py | 159 ++----------- .../pytorch/mxfp4_qat_cute_dsl.py | 209 ------------------ 4 files changed, 98 insertions(+), 480 deletions(-) create mode 100644 tests/pytorch/references/mxfp4_qat_reference.py delete mode 100644 transformer_engine/pytorch/mxfp4_qat_cute_dsl.py diff --git a/tests/pytorch/references/mxfp4_qat_reference.py b/tests/pytorch/references/mxfp4_qat_reference.py new file mode 100644 index 0000000000..d6e41e2ad5 --- /dev/null +++ b/tests/pytorch/references/mxfp4_qat_reference.py @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pure-PyTorch reference for MXFP4 weight fake-quantization. + +Bit-identical specification of transformer_engine.pytorch.mxfp4_qat +(tex.mxfp4_fake_quantize): per 1x32 block, scale = 2^clamp(ceil(log2(amax/6)), +-126, 125) derived from the amax bit pattern, RTNE onto the E2M1 grid, +saturate at 6, multiply back. Non-finite values NaN-poison their whole block; +the sign of -0 is preserved; the scale floor/cap keep every grid value exactly +representable in bf16/fp32. +""" +import torch + +_E2M1_MAX = 6.0 +_MXFP4_BLOCK = 32 + + +def round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: + """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" + fine = torch.round(y * 2.0) * 0.5 + mid = torch.round(y) + coarse = torch.round(y * 0.5) * 2.0 + return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) + + +def mxfp4_fake_quantize_reference(weight: torch.Tensor) -> torch.Tensor: + """Numerical oracle for the mxfp4_fake_quantize kernel (composite torch ops).""" + rows, cols = weight.shape + w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) + amax = w32.abs().amax(dim=-1, keepdim=True) + nonfinite = ~torch.isfinite(amax) + + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mantissa = bits & 0x7FFFFF + exp = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) + exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)) + exp = exp.clamp(min=-126, max=125) + scale = torch.ldexp(torch.ones_like(amax), exp) + scale = torch.where(amax > 0, scale, torch.ones_like(scale)) + + y = (w32 / scale).abs().clamp(max=_E2M1_MAX) + q = round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) + q = torch.copysign(q, w32) + out = q * scale + out = torch.where(nonfinite.expand_as(out), torch.full_like(out, float("nan")), out) + return out.view(rows, cols).to(weight.dtype) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index 371b1ff1af..f964332a2f 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -23,6 +23,7 @@ ) from transformer_engine.pytorch import fp8_autocast from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize +from references.mxfp4_qat_reference import mxfp4_fake_quantize_reference from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer @@ -136,18 +137,13 @@ def test_E_mxfp8_rowwise_lossless(): raw_dq, wh32 ), f"RAW MXFP8 encoding of the MXFP4-grid weight is not lossless {m}x{n}" - flushed = wh32.abs() == 2.0**-127 - assert flushed.any() and not flushed.all(), "artifact case not exercised" assert torch.equal( - dq[~flushed], wh32[~flushed] + dq, wh32 ), f"MXFP8 rowwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" - if not (dq[flushed] == 0).all(): - assert torch.equal(dq, wh32), "dequant neither flushed nor exact?" - print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") def test_dequantize_e8m0_extreme_codes(): - """Real MXFP8 software dequantize with planted UE8M0 codes 0/255; auto-detects pre-fix wheel vs fixed build.""" + """Real MXFP8 software dequantize with planted UE8M0 codes 0/255.""" w = make_weight(32, 64) q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize( mxfp4_fake_quantize(w) @@ -159,12 +155,8 @@ def test_dequantize_e8m0_extreme_codes(): data[1, :32] = 56 codes[1, 0] = 255 dq = q.dequantize(dtype=torch.float32) - if (dq[0, :32] == 2.0**-127).all() and torch.isnan(dq[1, :32]).all(): - print(" fixed build: code 0 -> 2^-127, code 255 -> NaN") - else: - assert (dq[0, :32] == 0).all(), "code 0 neither 2^-127 (fixed) nor 0 (pre-fix)" - assert torch.isinf(dq[1, :32]).all(), "code 255 neither NaN (fixed) nor Inf (pre-fix)" - print(" pre-fix wheel: code 0 flushed to 0, code 255 -> Inf (rebuild pending)") + assert (dq[0, :32] == 2.0**-127).all(), "UE8M0 code 0 must dequantize to 2^-127" + assert torch.isnan(dq[1, :32]).all(), "UE8M0 code 255 must dequantize to NaN" def test_G_blockwise_lossless_and_transpose(): @@ -198,14 +190,7 @@ def test_lossless_dequant_to_bf16(): q8 = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(w_hat) dq8 = q8.dequantize(dtype=torch.bfloat16) - flushed = w_hat.to(torch.float32).abs() == 2.0**-127 - assert flushed.any(), "artifact case not exercised" - assert torch.equal( - dq8[~flushed], w_hat[~flushed] - ), f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" - if not (dq8[flushed] == 0).all(): - assert torch.equal(dq8, w_hat), "dequant neither flushed nor exact?" - print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") + assert torch.equal(dq8, w_hat), f"mxfp8 rowwise dequant to bf16 not lossless {m}x{n}" w2_hat = mxfp4_fake_quantize(make_weight(m, n, zero_blocks=4)) qb = Float8BlockQuantizer( @@ -296,10 +281,8 @@ def _intree_kernel(fast_math=False): def _both_paths(w): - import transformer_engine.pytorch.mxfp4_qat as M - got_kernel = _intree_kernel().mxfp4_fake_quantize_intree(w.contiguous()) - got_torch = M._mxfp4_fake_quantize_torch(w) + got_torch = mxfp4_fake_quantize_reference(w) return got_kernel, got_torch @@ -425,14 +408,12 @@ def test_kernel_perf(): if not RUN_PERF: print(" SKIP (set NVTE_MXFP4_QAT_PERF=1 to run benchmarks)") return - import transformer_engine.pytorch.mxfp4_qat as M - rows, cols = 8192, 8192 w = make_weight(rows, cols) k = _intree_kernel() for _ in range(3): k.mxfp4_fake_quantize_intree(w) - M._mxfp4_fake_quantize_torch(w) + mxfp4_fake_quantize_reference(w) torch.cuda.synchronize() def timeit(fn, iters=20): @@ -445,7 +426,7 @@ def timeit(fn, iters=20): return start.elapsed_time(end) / iters * 1e3 t_kernel = timeit(lambda: k.mxfp4_fake_quantize_intree(w)) - t_torch = timeit(lambda: M._mxfp4_fake_quantize_torch(w)) + t_torch = timeit(lambda: mxfp4_fake_quantize_reference(w)) moved = rows * cols * 2 * 2 print( f" kernel: {t_kernel:.1f} us ({moved / t_kernel / 1e3:.0f} GB/s), " @@ -454,73 +435,14 @@ def timeit(fn, iters=20): assert t_kernel < t_torch, "kernel slower than the torch reference" -def test_cute_dsl_version(): - try: - from transformer_engine.pytorch.mxfp4_qat_cute_dsl import mxfp4_fake_quantize_cute_dsl - except Exception as e: - print(f" SKIP (cutlass DSL unavailable: {type(e).__name__}: {e})") - return - import transformer_engine.pytorch.mxfp4_qat as M - - for m, n in SHAPES: - for dtype in (torch.bfloat16, torch.float32): - w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) - gd = mxfp4_fake_quantize_cute_dsl(w) - gt = M._mxfp4_fake_quantize_torch(w) - _assert_bits_equal(gd, gt, f"cute-dsl != torch ref {m}x{n} {dtype}") - - w = make_weight(64, 128) - w[0, :32] = 0.0 - w[1, 5] = float("inf") - w[2, 9] = float("nan") - w[3, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) - w[4, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) - w[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) - gd = mxfp4_fake_quantize_cute_dsl(w) - gt = M._mxfp4_fake_quantize_torch(w) - _assert_same_with_nan(gd, gt, "cute-dsl edge") - assert (gd[5, :32].to(torch.float32) == 2.0**-127).all(), "cute-dsl: 2^-127 lost" - - if not RUN_PERF: - print(" parity OK; SKIP perf (set NVTE_MXFP4_QAT_PERF=1)") - return - - rows, cols = 8192, 8192 - wb = make_weight(rows, cols) - k = _intree_kernel() - for _ in range(3): - mxfp4_fake_quantize_cute_dsl(wb) - k.mxfp4_fake_quantize_intree(wb) - torch.cuda.synchronize() - - def timeit(fn, iters=20): - start, end = torch.cuda.Event(True), torch.cuda.Event(True) - start.record() - for _ in range(iters): - fn() - end.record() - torch.cuda.synchronize() - return start.elapsed_time(end) / iters * 1e3 - - t_dsl = timeit(lambda: mxfp4_fake_quantize_cute_dsl(wb)) - t_cuda = timeit(lambda: k.mxfp4_fake_quantize_intree(wb)) - moved = rows * cols * 2 * 2 - print( - f" cute-dsl: {t_dsl:.1f} us ({moved / t_dsl / 1e3:.0f} GB/s), " - f"cuda kernel: {t_cuda:.1f} us ({moved / t_cuda / 1e3:.0f} GB/s)" - ) - - def test_kernel_fast_math_immune(): """The kernel stays bit-identical when compiled with --use_fast_math.""" - import transformer_engine.pytorch.mxfp4_qat as M - k = _intree_kernel(fast_math=True) for m, n in PARITY_SHAPES: for dtype in (torch.bfloat16, torch.float32): w = make_weight(m, n, dtype=dtype, zero_blocks=4, outliers=4) got = k.mxfp4_fake_quantize_intree(w.contiguous()) - ref = M._mxfp4_fake_quantize_torch(w) + ref = mxfp4_fake_quantize_reference(w) _assert_bits_equal(got, ref, f"fast-math build diverged {m}x{n} {dtype}") w = make_weight(64, 128) @@ -531,7 +453,7 @@ def test_kernel_fast_math_immune(): w[4, :32] = torch.tensor(3.0e38, dtype=torch.bfloat16) w[5, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) got = k.mxfp4_fake_quantize_intree(w.contiguous()) - ref = M._mxfp4_fake_quantize_torch(w) + ref = mxfp4_fake_quantize_reference(w) _assert_same_with_nan(got, ref, "fast-math edge") assert ( got[5, :32].to(torch.float32) == 2.0**-127 @@ -603,32 +525,18 @@ def test_exp2f_e8m0_all_codes(): def test_ste_identity_gradient(): - """mxfp4_fake_quantize carries an identity (STE) gradient on both implementation paths.""" - import transformer_engine.pytorch.mxfp4_qat as M - + """mxfp4_fake_quantize carries an identity (STE) gradient and matches the reference forward.""" for dtype in (torch.bfloat16, torch.float32): w = make_weight(64, 128, dtype=dtype).requires_grad_(True) up = torch.randn_like(w) out = mxfp4_fake_quantize(w) out.backward(up) assert torch.equal(w.grad, up), f"STE gradient is not identity ({dtype})" - - w2 = w.detach().clone().requires_grad_(True) - old = M._torch_path_forced - M._torch_path_forced = True - try: - out2 = mxfp4_fake_quantize(w2) - out2.backward(up) - finally: - M._torch_path_forced = old - assert torch.equal(w2.grad, up), f"torch-path STE gradient is not identity ({dtype})" - _assert_bits_equal(out.detach(), out2.detach(), "STE fwd parity") + _assert_bits_equal(out.detach(), mxfp4_fake_quantize_reference(w.detach()), "STE fwd") def test_misaligned_and_noncontiguous(): """Misaligned storage-offset views and non-contiguous inputs match an aligned copy bitwise.""" - import transformer_engine.pytorch.mxfp4_qat as M - for dtype, off in ((torch.bfloat16, 4), (torch.float32, 2)): m, n = 64, 128 flat = (torch.randn(m * n + off, device=DEV, dtype=torch.float32) * 0.02).to(dtype) @@ -640,8 +548,8 @@ def test_misaligned_and_noncontiguous(): nc = (torch.randn(64, 256, device=DEV, dtype=torch.float32) * 0.02)[:, ::2] assert not nc.is_contiguous() - got = M._mxfp4_fake_quantize_torch(nc) - assert torch.equal(got, M._mxfp4_fake_quantize_torch(nc.contiguous())) + got = mxfp4_fake_quantize_reference(nc) + assert torch.equal(got, mxfp4_fake_quantize_reference(nc.contiguous())) def test_bf16_exhaustive_bit_patterns(): @@ -792,8 +700,8 @@ def test_fp16_pipeline_rejected(): pass -def test_fourway_matrix(): - """TileKernels vs CuTe DSL vs CUDA vs torch, bitwise, over all edge domains and every quant/dequant hop.""" +def test_cross_impl_matrix(): + """TileKernels vs CUDA kernel vs torch reference, bitwise, over all edge domains and every quant/dequant hop.""" import sys tk_root = os.environ.get( @@ -810,17 +718,8 @@ def test_fourway_matrix(): except Exception as e: print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") return - try: - from transformer_engine.pytorch.mxfp4_qat_cute_dsl import ( - mxfp4_fake_quantize_cute_dsl as dsl, - ) - except Exception as e: - print(f" SKIP (cute dsl unavailable: {type(e).__name__}: {e})") - return - import transformer_engine.pytorch.mxfp4_qat as M - kern = _intree_kernel().mxfp4_fake_quantize_intree - tref = M._mxfp4_fake_quantize_torch + tref = mxfp4_fake_quantize_reference CAP = 6.0 * 2.0**125 def tk_roundtrip(x, block, fmt): @@ -845,11 +744,10 @@ def tk_roundtrip(x, block, fmt): for W in (w, w.to(torch.float32)): t = tref(W) _assert_bits_equal(kern(W.contiguous()), t, "fake-quant CUDA vs torch") - _assert_bits_equal(dsl(W), t, "fake-quant DSL vs torch") tk32, tkb16 = tk_roundtrip(W, (1, 32), "e2m1") _assert_bits_equal(tk32, t.to(torch.float32), "fake-quant TK vs torch (fp32)") _assert_bits_equal(tkb16, t.to(torch.bfloat16), "fake-quant TK vs torch (bf16)") - print(" fake-quant 4-way bitwise: PASS (bf16 + fp32 corpus)") + print(" fake-quant 3-way bitwise: PASS (bf16 + fp32 corpus)") wnf = w.clone() wnf[8, 40] = float("inf") @@ -857,17 +755,15 @@ def tk_roundtrip(x, block, fmt): wnf[10, 5] = float("nan") t = tref(wnf) _assert_same_with_nan(kern(wnf.contiguous()), t, "nonfinite CUDA") - _assert_same_with_nan(dsl(wnf), t, "nonfinite DSL") assert torch.isnan(t[8, 32:64]).all() and torch.isnan(t[9, 32:64]).all() assert torch.isnan(t[10, :32]).all() and torch.isfinite(t[11].to(torch.float32)).all() - print(" nonfinite 3-way (whole-block NaN poisoning): PASS (TK policy differs, excluded)") + print(" nonfinite 2-way (whole-block NaN poisoning): PASS (TK policy differs, excluded)") wcap = torch.zeros(1, 32, dtype=torch.bfloat16, device=DEV) wcap[0, 0] = torch.tensor(3.2e38, dtype=torch.bfloat16) t = tref(wcap) assert t[0, 0].item() == CAP, "satfinite cap at 6*2^125" _assert_bits_equal(kern(wcap.contiguous()), t, "cap CUDA") - _assert_bits_equal(dsl(wcap), t, "cap DSL") tk32, _ = tk_roundtrip(wcap, (1, 32), "e2m1") assert not torch.equal(tk32, t.to(torch.float32)), "TK must diverge above 6*2^125" print( @@ -899,15 +795,11 @@ def tk_roundtrip(x, block, fmt): assert (tk32[~okm] == 0).all(), "TK e4m3 sub-floor blocks must flush to zero" assert (~ok).sum() > 0, "TK e4m3 floor case not exercised" dq = q.dequantize(dtype=torch.float32) - flushed = wh32.abs() == 2.0**-127 - assert torch.equal(dq[~flushed], wh32[~flushed]), "mxfp8 row TE software dequant" - if flushed.any() and not (dq[flushed] == 0).all(): - assert torch.equal(dq, wh32) - print(" NOTE: TE software dequant now handles UE8M0 code 0 -> fully exact") + assert torch.equal(dq, wh32), "mxfp8 row TE software dequant" print( " mxfp8 row: TE raw encode lossless incl 2^-127; TK roundtrip bitwise on " f"amax >= {e4m3_floor:g} blocks ({int((~ok).sum())} sub-floor blocks flush in TK " - "e4m3, a TK-only floor); TE software dequant exact outside the code-0 wheel bug: PASS" + "e4m3, a TK-only floor); TE software dequant exact incl 2^-127: PASS" ) q.update_usage(rowwise_usage=False, columnwise_usage=True) @@ -959,13 +851,12 @@ def tk_roundtrip(x, block, fmt): we = bits.view(torch.bfloat16).view(2048, 32) t = tref(we) _assert_same_with_nan(kern(we.contiguous()), t, "exhaustive CUDA") - _assert_same_with_nan(dsl(we), t, "exhaustive DSL") we32 = we.to(torch.float32) ok_rows = torch.isfinite(we32).all(dim=-1) & (we32.abs().amax(dim=-1) <= CAP) tk32, _ = tk_roundtrip(we[ok_rows], (1, 32), "e2m1") _assert_bits_equal(tk32, t[ok_rows].to(torch.float32), "exhaustive TK") print( - " bf16 exhaustive 65536 patterns: CUDA/DSL/torch 3-way full, " + " bf16 exhaustive 65536 patterns: CUDA/torch full, " f"TK on {int(ok_rows.sum())}/2048 finite below-cap rows, bitwise: PASS" ) @@ -974,13 +865,12 @@ def tk_roundtrip(x, block, fmt): wf = fb.to(torch.int32).cuda().view(torch.float32) t = tref(wf) _assert_same_with_nan(kern(wf.contiguous()), t, "fuzz CUDA") - _assert_same_with_nan(dsl(wf), t, "fuzz DSL") ok_rows = torch.isfinite(wf).all(dim=-1) & (wf.abs().amax(dim=-1) <= CAP) if ok_rows.any(): tk32, _ = tk_roundtrip(wf[ok_rows], (1, 32), "e2m1") _assert_bits_equal(tk32, t[ok_rows], "fuzz TK") print( - f" fp32 bit-fuzz 512 blocks: 3-way full, TK on {int(ok_rows.sum())} " + f" fp32 bit-fuzz 512 blocks: CUDA/torch full, TK on {int(ok_rows.sum())} " "finite below-cap rows, bitwise: PASS" ) @@ -1195,10 +1085,9 @@ def test_e2e_matrix(): test_recipe_field_controls_qat, test_ops_api_rejected, test_fp16_pipeline_rejected, - test_fourway_matrix, + test_cross_impl_matrix, test_kernel_perf, test_kernel_fast_math_immune, - test_cute_dsl_version, test_e2e_matrix, ] diff --git a/transformer_engine/pytorch/mxfp4_qat.py b/transformer_engine/pytorch/mxfp4_qat.py index ad161b2e75..57ee4f854c 100644 --- a/transformer_engine/pytorch/mxfp4_qat.py +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -12,157 +12,50 @@ MXFP4->MXFP8 converter: the host encoder re-canonicalizes (payload, scale) per block, so equal values may carry different E4M3/UE8M0 bytes, and the original MXFP4 scale metadata is not preserved. Scale floor 2^-126 / cap 2^125 -(TileKernels deployment contract); non-finite values NaN-poison their 1x32 -block; fp16 is rejected. Implementation picked by NVTE_MXFP4_QAT_IMPL -(auto/cuda/cute_dsl/torch). +(deployment contract); non-finite values NaN-poison their 1x32 block; fp16 is +rejected. The numerical specification lives in +tests/pytorch/references/mxfp4_qat_reference.py. """ -import os -import warnings - import torch __all__ = ["mxfp4_fake_quantize"] -_E2M1_MAX = 6.0 _MXFP4_BLOCK = 32 -_IMPL_CHOICE = os.getenv("NVTE_MXFP4_QAT_IMPL", "auto") -_torch_path_forced = os.getenv("NVTE_MXFP4_QAT_DISABLE_CUDA_KERNEL", "0") == "1" -_kernel_required = os.getenv("NVTE_MXFP4_QAT_REQUIRE_KERNEL", "0") == "1" -_missing_kernel_warned = False -_cute_dsl_import_error = None - class _MXFP4FakeQuantizeSTE(torch.autograd.Function): """Straight-through estimator: identity gradient through the projection.""" @staticmethod - def forward(ctx, weight, impl): # pylint: disable=arguments-differ - return impl(weight) - - @staticmethod - def backward(ctx, grad_output): # pylint: disable=arguments-differ - return grad_output, None - - -def _round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: - """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" - fine = torch.round(y * 2.0) * 0.5 - mid = torch.round(y) - coarse = torch.round(y * 0.5) * 2.0 - return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) - - -def _mxfp4_fake_quantize_torch(weight: torch.Tensor) -> torch.Tensor: - """Bit-identical PyTorch reference for the CUDA and CuTe DSL kernels.""" - rows, cols = weight.shape - w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) - amax = w32.abs().amax(dim=-1, keepdim=True) - nonfinite = ~torch.isfinite(amax) + def forward(ctx, weight): # pylint: disable=arguments-differ + import transformer_engine_torch as tex - bits = amax.view(torch.int32) - exp_field = bits >> 23 - mantissa = bits & 0x7FFFFF - exp = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) - exp = torch.where(exp_field > 0, exp, torch.full_like(exp, -126)) - exp = exp.clamp(min=-126, max=125) - scale = torch.ldexp(torch.ones_like(amax), exp) - scale = torch.where(amax > 0, scale, torch.ones_like(scale)) - - y = (w32 / scale).abs().clamp(max=_E2M1_MAX) - q = _round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) - q = torch.copysign(q, w32) - out = q * scale - out = torch.where(nonfinite.expand_as(out), torch.full_like(out, float("nan")), out) - return out.view(rows, cols).to(weight.dtype) - - -def _cuda_impl(): - import transformer_engine_torch as tex - - if not hasattr(tex, "mxfp4_fake_quantize"): - return None - - def impl(w): - w = w.contiguous() + if not hasattr(tex, "mxfp4_fake_quantize"): + raise RuntimeError( + "transformer_engine_torch was built without mxfp4_fake_quantize; " + "rebuild Transformer Engine." + ) + w = weight.contiguous() if w.data_ptr() % 16 != 0: + # logically-contiguous storage-offset views can be misaligned for + # the kernel's 16-byte vector accesses w = w.clone() return tex.mxfp4_fake_quantize(w) - return impl - - -def _cute_dsl_impl(): - global _cute_dsl_import_error - if _cute_dsl_import_error is not None: - return None - try: - from .mxfp4_qat_cute_dsl import mxfp4_fake_quantize_cute_dsl - except Exception as exc: # pylint: disable=broad-except - _cute_dsl_import_error = exc - return None - return mxfp4_fake_quantize_cute_dsl - - -def _resolve_impl(weight: torch.Tensor): - if not weight.is_cuda: - return _mxfp4_fake_quantize_torch - if _IMPL_CHOICE == "torch": - return _mxfp4_fake_quantize_torch - if _IMPL_CHOICE == "cuda": - impl = _cuda_impl() - if impl is None: - raise RuntimeError( - "NVTE_MXFP4_QAT_IMPL=cuda but transformer_engine_torch was built " - "without mxfp4_fake_quantize. Rebuild Transformer Engine." - ) - return impl - if _IMPL_CHOICE == "cute_dsl": - impl = _cute_dsl_impl() - if impl is None: - raise RuntimeError( - "NVTE_MXFP4_QAT_IMPL=cute_dsl but the cutlass CuTe DSL is " - f"unavailable: {_cute_dsl_import_error!r}" - ) - return impl - if _IMPL_CHOICE != "auto": - raise ValueError( - f"NVTE_MXFP4_QAT_IMPL={_IMPL_CHOICE!r} is not one of " - "'auto', 'cuda', 'cute_dsl', 'torch'" - ) - if _torch_path_forced: - return _mxfp4_fake_quantize_torch - impl = _cuda_impl() - if impl is not None: - return impl - impl = _cute_dsl_impl() - if impl is not None: - return impl - if _kernel_required: - raise RuntimeError( - "NVTE_MXFP4_QAT_REQUIRE_KERNEL=1 but neither the CUDA binding nor " - "the CuTe DSL kernel is available. Rebuild Transformer Engine or " - "install the cutlass CuTe DSL." - ) - global _missing_kernel_warned - if not _missing_kernel_warned: - _missing_kernel_warned = True - warnings.warn( - "transformer_engine_torch was built without mxfp4_fake_quantize and " - "the CuTe DSL is unavailable; falling back to the (slower) PyTorch " - "reference. Rebuild Transformer Engine to use the CUDA kernel.", - stacklevel=2, - ) - return _mxfp4_fake_quantize_torch + @staticmethod + def backward(ctx, grad_output): # pylint: disable=arguments-differ + return grad_output def mxfp4_fake_quantize(weight: torch.Tensor) -> torch.Tensor: """Project ``weight`` onto the MXFP4 grid and return it in the input dtype. Per 1x32 block: scale = 2^clamp(ceil(log2(amax / 6)), -126, 125), RTNE onto - E2M1, saturate at 6, multiply back. Gradient is the straight-through - estimator. ``NVTE_MXFP4_QAT_REQUIRE_KERNEL=1`` forbids the torch fallback. + E2M1, saturate at 6, multiply back. The gradient is the straight-through + estimator (identity). """ + if not weight.is_cuda: + raise ValueError("MXFP4 QAT expects a CUDA weight tensor") if weight.dim() != 2: raise ValueError(f"MXFP4 QAT expects a 2D weight, got {tuple(weight.shape)}") if weight.dtype not in (torch.bfloat16, torch.float32): @@ -170,13 +63,9 @@ def mxfp4_fake_quantize(weight: torch.Tensor) -> torch.Tensor: f"MXFP4 QAT supports bf16/fp32 weights, got {weight.dtype} " "(fp16 cannot represent the full MXFP4 scale range)" ) - rows, cols = weight.shape - if cols % _MXFP4_BLOCK != 0: + if weight.shape[1] % _MXFP4_BLOCK != 0: raise ValueError( - f"MXFP4 QAT needs the weight inner dim divisible by {_MXFP4_BLOCK}, got {cols}" + "MXFP4 QAT needs the weight inner dim divisible by " + f"{_MXFP4_BLOCK}, got {weight.shape[1]}" ) - - impl = _resolve_impl(weight) - if torch.is_grad_enabled() and weight.requires_grad: - return _MXFP4FakeQuantizeSTE.apply(weight, impl) - return impl(weight) + return _MXFP4FakeQuantizeSTE.apply(weight) diff --git a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py b/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py deleted file mode 100644 index 98d3f5941f..0000000000 --- a/transformer_engine/pytorch/mxfp4_qat_cute_dsl.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""CuTe DSL implementation of MXFP4 weight fake-quantization. - -Bit-identical to the CUDA kernel and the PyTorch reference: integer-bit scale -derivation and non-FTZ inline PTX on the value path. Requires the cutlass -CuTe DSL package. -""" -import functools - -import torch - -import cutlass -import cutlass.cute as cute -from cutlass import Float32, Int32 -from cutlass.cutlass_dsl import T, dsl_user_op -from cutlass._mlir.dialects import llvm - -__all__ = ["mxfp4_fake_quantize_cute_dsl"] - -_MXFP4_BLOCK = 32 -_THREADS = 256 - -_ASM_KW = dict( - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, -) - - -@dsl_user_op -def _is_nonfinite_f32(x: Float32, *, loc=None, ip=None) -> Int32: - """1 if x is inf/NaN else 0 (PTX testp.finite).""" - return Int32( - llvm.inline_asm( - T.i32(), - [Float32(x).ir_value(loc=loc, ip=ip)], - "{\n .reg .pred p;\n testp.finite.f32 p, $1;\n selp.s32 $0, 0, 1, p;\n}", - "=r,f", - loc=loc, - ip=ip, - **_ASM_KW, - ) - ) - - -@dsl_user_op -def _block_scale_from_amax(amax: Float32, *, loc=None, ip=None) -> Float32: - """scale = 2^clamp(ceil(log2(amax/6)), -126, 125); amax == 0 -> 1.0.""" - return Float32( - llvm.inline_asm( - T.f32(), - [Float32(amax).ir_value(loc=loc, ip=ip)], - "{\n" - " .reg .s32 b, k, mnt, e, t;\n" - " .reg .pred pz, pm, psub;\n" - " setp.eq.f32 pz, $1, 0f00000000;\n" - " mov.b32 b, $1;\n" - " shr.s32 k, b, 23;\n" - " and.b32 mnt, b, 8388607;\n" - " setp.gt.s32 pm, mnt, 4194304;\n" - " selp.s32 t, 1, 0, pm;\n" - " add.s32 e, k, -129;\n" - " add.s32 e, e, t;\n" - " setp.eq.s32 psub, k, 0;\n" - " selp.s32 e, -126, e, psub;\n" - " max.s32 e, e, -126;\n" - " min.s32 e, e, 125;\n" - " add.s32 e, e, 127;\n" - " shl.b32 e, e, 23;\n" - " mov.b32 $0, e;\n" - " @pz mov.f32 $0, 0f3F800000;\n" - "}", - "=f,f", - loc=loc, - ip=ip, - **_ASM_KW, - ) - ) - - -@dsl_user_op -def _fake_quant_elem(x: Float32, scale: Float32, nonfinite: Int32, *, loc=None, ip=None) -> Float32: - """RTNE onto the E2M1 grid at the given power-of-two scale; nonfinite -> NaN.""" - return Float32( - llvm.inline_asm( - T.f32(), - [ - Float32(x).ir_value(loc=loc, ip=ip), - Float32(scale).ir_value(loc=loc, ip=ip), - Int32(nonfinite).ir_value(loc=loc, ip=ip), - ], - "{\n" - " .reg .f32 y, ay, f, m, c, q, r, inv;\n" - " .reg .s32 ib;\n" - " .reg .pred p2, p4, pnf;\n" - " setp.ne.s32 pnf, $3, 0;\n" - " mov.b32 ib, $2;\n" - " sub.s32 ib, 2130706432, ib;\n" - " mov.b32 inv, ib;\n" - " mul.rn.f32 y, $1, inv;\n" - " abs.f32 ay, y;\n" - " min.f32 ay, ay, 0f40C00000;\n" - " mul.rn.f32 f, ay, 0f40000000;\n" - " cvt.rni.f32.f32 f, f;\n" - " mul.rn.f32 f, f, 0f3F000000;\n" - " cvt.rni.f32.f32 m, ay;\n" - " mul.rn.f32 c, ay, 0f3F000000;\n" - " cvt.rni.f32.f32 c, c;\n" - " mul.rn.f32 c, c, 0f40000000;\n" - " setp.le.f32 p2, ay, 0f40000000;\n" - " setp.le.f32 p4, ay, 0f40800000;\n" - " selp.f32 q, m, c, p4;\n" - " selp.f32 q, f, q, p2;\n" - " copysign.f32 r, $1, q;\n" - " mul.rn.f32 r, r, $2;\n" - " selp.f32 $0, 0f7FC00000, r, pnf;\n" - "}", - "=f,f,f,r", - loc=loc, - ip=ip, - **_ASM_KW, - ) - ) - - -class _MXFP4FakeQuantCuteDsl: - """Thread-per-1x32-block kernel; loads/stores via autovec_copy fragments.""" - - def __init__(self, dtype): - self.dtype = dtype - - @cute.jit - def __call__(self, mX: cute.Tensor, mO: cute.Tensor, num_blocks: Int32, stream): - self.kernel(mX, mO, num_blocks).launch( - grid=[cute.ceil_div(num_blocks, _THREADS), 1, 1], - block=[_THREADS, 1, 1], - stream=stream, - ) - - @cute.kernel - def kernel(self, mX: cute.Tensor, mO: cute.Tensor, num_blocks: Int32): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - idx = bidx * _THREADS + tidx - if idx < num_blocks: - x_row = mX[idx, None] - o_row = mO[idx, None] - x_frag = cute.make_fragment_like(x_row) - cute.autovec_copy(x_row, x_frag) - o_frag = cute.make_fragment_like(o_row) - - amax = Float32(0.0) - nonfinite = Int32(0) - for i in cutlass.range_constexpr(_MXFP4_BLOCK): - v = Float32(x_frag[i]) - nonfinite = nonfinite + _is_nonfinite_f32(v) - amax = cute.arch.fmax(amax, cute.arch.fmax(v, Float32(0.0) - v)) - - scale = _block_scale_from_amax(amax) - for i in cutlass.range_constexpr(_MXFP4_BLOCK): - v = Float32(x_frag[i]) - o_frag[i] = self.dtype(_fake_quant_elem(v, scale, nonfinite)) - cute.autovec_copy(o_frag, o_row) - - -@functools.lru_cache(maxsize=4) -def _compiled_kernel(dtype_key: str): - cutlass_dtype = cutlass.BFloat16 if dtype_key == "bf16" else cutlass.Float32 - kernel_obj = _MXFP4FakeQuantCuteDsl(cutlass_dtype) - sym_blocks = cute.sym_int() - x_fake = cute.runtime.make_fake_compact_tensor( - cutlass_dtype, (sym_blocks, _MXFP4_BLOCK), stride_order=(1, 0), assumed_align=16 - ) - o_fake = cute.runtime.make_fake_compact_tensor( - cutlass_dtype, (sym_blocks, _MXFP4_BLOCK), stride_order=(1, 0), assumed_align=16 - ) - stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - return cute.compile( - kernel_obj, x_fake, o_fake, Int32(1), stream_fake, options="--enable-tvm-ffi" - ) - - -def mxfp4_fake_quantize_cute_dsl(weight: torch.Tensor) -> torch.Tensor: - """Project ``weight`` onto the MXFP4 grid and return it in the input dtype.""" - if not weight.is_cuda: - raise ValueError("MXFP4 QAT CuTe DSL kernel expects a CUDA tensor") - if weight.dim() != 2: - raise ValueError(f"MXFP4 QAT expects a 2D weight, got {tuple(weight.shape)}") - if weight.dtype not in (torch.bfloat16, torch.float32): - raise ValueError( - f"MXFP4 QAT supports bf16/fp32 weights, got {weight.dtype} " - "(fp16 cannot represent the full MXFP4 scale range)" - ) - rows, cols = weight.shape - if cols % _MXFP4_BLOCK != 0: - raise ValueError( - f"MXFP4 QAT needs the weight inner dim divisible by {_MXFP4_BLOCK}, got {cols}" - ) - w = weight.contiguous() - if w.data_ptr() % 16 != 0: - w = w.clone() - out = torch.empty_like(w) - num_blocks = w.numel() // _MXFP4_BLOCK - fn = _compiled_kernel("bf16" if weight.dtype == torch.bfloat16 else "fp32") - fn(w.view(-1, _MXFP4_BLOCK), out.view(-1, _MXFP4_BLOCK), num_blocks) - return out From bb452614a3e06f514790fe87b79e4494b4e84687 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sun, 26 Jul 2026 11:12:40 -0700 Subject: [PATCH 09/14] Add direct MXFP4->FP8 converter mode (TileKernels canonicalization) Direct conversion of the MXFP4 decomposition into the host FP8 representation without a bf16 bridge tensor: fixed-shift-6 rowwise MXFP8 (deployment-canonical, invertible bytes), host-amax-rule columnwise (dgrad byte parity with the bridge), and per-tile exponent-folding 128x128 blockwise with RTNE onto E4M3 (exact through spread 2^14, no device assert). In-tree CUDA kernels (direct_convert_mxfp4.cuh) with C APIs and bindings, a composite-torch reference/fallback, recipe field mxfp4_qat_direct (NVTE_MXFP4_QAT_DIRECT), quantize_weight routing, and a parity suite incl. bitwise e2e equality against the bridge mode. --- .../references/mxfp4_qat_direct_reference.py | 186 +++++++ tests/pytorch/test_mxfp4_qat_direct.py | 483 ++++++++++++++++++ transformer_engine/common/cast/cast.cu | 66 +++ .../cast/mxfp4/direct_convert_mxfp4.cuh | 255 +++++++++ .../common/include/transformer_engine/cast.h | 14 + transformer_engine/common/recipe/__init__.py | 7 + transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 47 ++ .../pytorch/csrc/extensions/pybind.cpp | 6 + transformer_engine/pytorch/module/base.py | 10 + .../pytorch/mxfp4_qat_direct.py | 199 ++++++++ 11 files changed, 1275 insertions(+) create mode 100644 tests/pytorch/references/mxfp4_qat_direct_reference.py create mode 100644 tests/pytorch/test_mxfp4_qat_direct.py create mode 100644 transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh create mode 100644 transformer_engine/pytorch/mxfp4_qat_direct.py diff --git a/tests/pytorch/references/mxfp4_qat_direct_reference.py b/tests/pytorch/references/mxfp4_qat_direct_reference.py new file mode 100644 index 0000000000..8a72dced1d --- /dev/null +++ b/tests/pytorch/references/mxfp4_qat_direct_reference.py @@ -0,0 +1,186 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pure-PyTorch reference for the direct MXFP4 -> FP8 weight converters. + +Bit-identical specification of tex.mxfp4_direct_mxfp8_rowwise (TileKernels +fixed-shift-6 canonicalization: scale code = fp4_exp + 121 clamped at UE8M0 +code 0 with exact payload absorption) and tex.mxfp4_direct_blockwise +(per-tile exponent folding: tile scale = 2^(max block exponent - 6) clamped +at 2^-127, payloads RTNE onto E4M3 with subnormals -- exact through scale +spread 2^14, bounded beyond, no device assert). All-zero blocks pin the fp4 +exponent at -126 (TileKernels clamp); non-finite 1x32 blocks emit E4M3 NaN +payloads (0x7F). +""" +import torch + +_E4M3_NAN = 0x7F +_MXFP4_BLOCK = 32 + + +def _roundup(x: int, m: int) -> int: + return (x + m - 1) // m * m + + +def _round_to_e2m1_grid(y): + """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" + fine = torch.round(y * 2.0) * 0.5 + mid = torch.round(y) + coarse = torch.round(y * 0.5) * 2.0 + return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) + + +def _mxfp4_decompose(weight: torch.Tensor): + """Per-1x32-block MXFP4 decomposition of a bf16/fp32 weight. + + Returns (e, q, nonfinite): block scale exponents e (int32, [R, B, 1], + zero blocks pinned at -126 per the TileKernels clamp), signed E2M1 payload + values q (fp32, [R, B, 32]), and the non-finite block mask ([R, B, 1]). + """ + rows, cols = weight.shape + if cols % _MXFP4_BLOCK != 0: + raise ValueError(f"inner dim must be divisible by {_MXFP4_BLOCK}, got {cols}") + w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) + amax = w32.abs().amax(dim=-1, keepdim=True) + nonfinite = ~torch.isfinite(amax) + + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mantissa = bits & 0x7FFFFF + e = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) + e = torch.where(exp_field > 0, e, torch.full_like(e, -126)) + e = e.clamp(min=-126, max=125) + # all-zero block: TileKernels clamps amax at 6*2^-126 -> fp4 exponent -126 + # (deployment-canonical; payload 0 regardless, so the value is exact zero) + e = torch.where(amax > 0, e, torch.full_like(e, -126)) + e = torch.where(nonfinite, torch.full_like(e, -126), e) + + scale = torch.ldexp(torch.ones_like(amax), e) + y = (w32 / scale).abs().clamp(max=6.0) + q = _round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) + q = torch.copysign(q, w32) + return e, q, nonfinite + + +def _encode_e4m3(values: torch.Tensor) -> torch.Tensor: + """RTNE encode fp32 values into E4M3 payload bytes (uint8).""" + return values.to(torch.float8_e4m3fn).view(torch.uint8) + + +def _direct_mxfp8_rowwise(e, q, nonfinite, rows, cols): + """TileKernels fixed-shift-6 lift: payload p*2^6 at scale code e+121.""" + ones = torch.ones_like(e, dtype=torch.float32) + biased = e + 121 # e - 6 + 127 + code = biased.clamp(min=0) + # below UE8M0 code 0 the scale pins at 2^-127 and the payload absorbs the + # rest: p * 2^(e+127) (e in [-126, -122] -> multiplier 2..32, exact) + shift = torch.where(biased >= 0, torch.full_like(e, 6), e + 127) + payload_vals = q * torch.ldexp(ones, shift) + data = _encode_e4m3(payload_vals) + data = torch.where( + nonfinite.expand_as(data), torch.full_like(data, _E4M3_NAN), data + ).view(rows, cols) + code = torch.where(nonfinite, torch.full_like(code, 127), code) + + scale_inv = torch.zeros( + (_roundup(rows, 128), _roundup(cols // _MXFP4_BLOCK, 4)), + dtype=torch.uint8, + device=data.device, + ) + scale_inv[:rows, : cols // _MXFP4_BLOCK] = code.view(rows, -1).to(torch.uint8) + return data, scale_inv + + +def _host_amax_scale_exp(amax: torch.Tensor) -> torch.Tensor: + """Host MXFP8 encoder scale rule: 2^ceil(log2(amax/448)), floored at code 0. + + ceil(log2(m*2^k / 448)) = k - 8 + (m > 1.75), computed from the bits. + """ + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mantissa = bits & 0x7FFFFF + k = exp_field - 127 + ec = k - 8 + (mantissa > 0x600000).to(torch.int32) # 1.75 mantissa = 0x600000 + ec = torch.where(exp_field > 0, ec, torch.full_like(ec, -127)) # subnormal amax + ec = ec.clamp(min=-127, max=127) + ec = torch.where(amax > 0, ec, torch.zeros_like(ec)) # zero block: scale 1 + return ec + + +def _direct_mxfp8_columnwise(e, q, nonfinite, rows, cols): + """Fresh 32x1 quantization of the grid values with the host amax rule. + + Matches the bridge path's columnwise bytes (same scale rule, same RTNE + lattice) so mode-None dgrad is bit-identical across the two converters. + """ + ones = torch.ones_like(e, dtype=torch.float32) + vals = (q * torch.ldexp(ones, e)).view(rows, cols) + nf_elem = nonfinite.expand(-1, -1, _MXFP4_BLOCK).reshape(rows, cols) + + vc = vals.view(rows // 32, 32, cols) + amax_c = vc.abs().amax(dim=1, keepdim=True) + finite_c = torch.isfinite(amax_c) + amax_c = torch.where(finite_c, amax_c, torch.zeros_like(amax_c)) + ec = _host_amax_scale_exp(amax_c) + inv = torch.ldexp(torch.ones_like(ec, dtype=torch.float32), -ec) + # + 0.0 canonicalizes -0 to +0: the host colwise kernel drops the sign of + # zero (its rowwise kernel keeps it) and byte-parity with the bridge is the + # contract here; the value is unchanged either way. + payload = _encode_e4m3(vc * inv + 0.0).view(rows, cols) + payload = torch.where(nf_elem, torch.full_like(payload, _E4M3_NAN), payload) + + code = (ec + 127).clamp(0, 254).to(torch.uint8) + scale_inv = torch.zeros( + (_roundup(rows // 32, 4), _roundup(cols, 128)), + dtype=torch.uint8, + device=payload.device, + ) + scale_inv[: rows // 32, :cols] = code.view(rows // 32, cols) + return payload, scale_inv + + +def _direct_blockwise_128(e, q, nonfinite, rows, cols): + """per_block_cast_lossless exponent folding, RTNE onto E4M3 (no assert). + + Tile scale = 2^(max block exponent - 6), clamped to >= 2^-127 so the + fp32-stored scale stays on the E8M0 lattice; folded payloads round RTNE + (E4M3 subnormals included): exact through spread 2^14, bounded beyond. + """ + if rows % 128 != 0 or cols % 128 != 0: + raise ValueError(f"blockwise direct conversion needs 128-divisible dims, got {rows}x{cols}") + tiles_m, tiles_n = rows // 128, cols // 128 + + e_eff = torch.where(nonfinite, torch.full_like(e, -126), e) + # [R, B] -> [Tm, 128, Tn, 4] -> max over the 512 blocks of each tile + e_t = e_eff.view(rows, cols // _MXFP4_BLOCK).view(tiles_m, 128, tiles_n, 128 // _MXFP4_BLOCK) + max_e = e_t.amax(dim=(1, 3), keepdim=True) + s_exp = (max_e - 6).clamp(min=-127) + + fold = (e_t - s_exp).view(rows, cols // _MXFP4_BLOCK, 1).to(torch.float32) + payload_vals = q * torch.ldexp(torch.ones_like(fold), fold.to(torch.int32)) + data = _encode_e4m3(payload_vals) + data = torch.where( + nonfinite.expand_as(data), torch.full_like(data, _E4M3_NAN), data + ).view(rows, cols) + + s_exp2 = s_exp.view(tiles_m, tiles_n) + scale_inv = torch.zeros( + (tiles_m, _roundup(tiles_n, 4)), dtype=torch.float32, device=data.device + ) + scale_inv[:, :tiles_n] = torch.ldexp( + torch.ones_like(s_exp2, dtype=torch.float32), s_exp2 + ) + return data, scale_inv + + +def mxfp4_direct_mxfp8_rowwise_reference(weight: torch.Tensor): + """(payload bytes (M, N) uint8, padded UE8M0 scale codes) for a weight.""" + rows, cols = weight.shape + return _direct_mxfp8_rowwise(*_mxfp4_decompose(weight), rows, cols) + + +def mxfp4_direct_blockwise_reference(weight: torch.Tensor): + """(payload bytes (M, N) uint8, padded fp32 tile scales) for a weight.""" + rows, cols = weight.shape + return _direct_blockwise_128(*_mxfp4_decompose(weight), rows, cols) diff --git a/tests/pytorch/test_mxfp4_qat_direct.py b/tests/pytorch/test_mxfp4_qat_direct.py new file mode 100644 index 0000000000..7f81520211 --- /dev/null +++ b/tests/pytorch/test_mxfp4_qat_direct.py @@ -0,0 +1,483 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the direct MXFP4 -> FP8 converter (TileKernels canonicalization). + +Contracts: + - rowwise MXFP8 bytes follow the deployment fixed-shift-6 lift (invertible, + scale code = fp4_exp + 121 clamped at 0 with exact payload absorption); + - decoded values equal the bridge path (mxfp4_fake_quantize) bit for bit; + - columnwise MXFP8 bytes equal the bridge quantizer's (host amax rule); + - blockwise 128x128 uses tile scale 2^(max fp4 exp - 6) with RTNE folding: + exact through spread 2^14, bounded beyond, no device assert; + - end to end, direct and bridge recipes produce bitwise-identical outputs + and gradients (equivalent FP8 factorizations through the real GEMMs). +""" +import ctypes +import os + +for _n in ("libnvrtc.so.13", "libnvrtc.so.12", "libnvrtc.so"): + try: + ctypes.CDLL(_n, mode=ctypes.RTLD_GLOBAL) + break + except OSError: + continue + +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import ( + MXFP4QATFloat8BlockScaling, + MXFP4QATMXFP8BlockScaling, +) +from transformer_engine.pytorch import fp8_autocast +from transformer_engine.pytorch.mxfp4_qat import mxfp4_fake_quantize +from transformer_engine.pytorch.mxfp4_qat_direct import mxfp4_qat_direct_quantize +from references.mxfp4_qat_direct_reference import ( + _direct_blockwise_128, + _direct_mxfp8_rowwise, + _mxfp4_decompose, +) +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + +torch.manual_seed(4321) +DEV = "cuda" +RUN_PERF = os.environ.get("NVTE_MXFP4_QAT_PERF", "0") == "1" + + +def make_weight(m, n, dtype=torch.bfloat16, zero_blocks=0, outliers=0): + w = (torch.randn(m, n, device=DEV, dtype=torch.float32) * 0.02).to(dtype) + if outliers: + idx = torch.randint(0, m * n, (outliers,), device=DEV) + w.view(-1)[idx] *= 500.0 + if zero_blocks: + for _ in range(zero_blocks): + r = torch.randint(0, m, (1,)).item() + k = torch.randint(0, n // 32, (1,)).item() + w[r, k * 32 : (k + 1) * 32] = 0.0 + return w + + +def _decode_mxfp8_rowwise(qt, m, n): + data = qt._rowwise_data.view(torch.uint8)[:m, :n] + codes = qt._rowwise_scale_inv.view(torch.uint8)[:m, : n // 32].to(torch.int32) + vals = data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) + scale = torch.ldexp(torch.ones_like(codes, dtype=torch.float32), codes - 127) + return (vals * scale.unsqueeze(-1)).view(m, n) + + +def _decode_blockwise_rowwise(qt, m, n): + data = qt._rowwise_data.view(torch.uint8)[:m, :n] + scales = qt._rowwise_scale_inv[: m // 128, : n // 128] + vals = data.view(torch.float8_e4m3fn).to(torch.float32) + up = scales.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1) + return vals * up + + +def test_rowwise_canonical_bytes_and_decode(): + m, n = 64, 256 + w = make_weight(m, n, zero_blocks=3, outliers=3) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) # floor-scale block + w[1, :32] = 0.0 + w[1, 0] = -0.0 + w[2, :32] = 0.0 + w[2, 0] = 3.25 # quantizes to payload 3 at fp4 exp 0 -> canonical divergence case + w[3, :32] = 0.0 + w[3, 0] = 6.0 # payload 6 at fp4 exp 0 -> canonical agreement case + + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False) + qt = mxfp4_qat_direct_quantize(w, q) + what = mxfp4_fake_quantize(w) + + # decoded values == bridge values, bitwise, everywhere (incl. 2^-127) + dec = _decode_mxfp8_rowwise(qt, m, n) + assert torch.equal(dec, what.to(torch.float32)), "direct decode != bridge values" + + data = qt._rowwise_data.view(torch.uint8) + codes = qt._rowwise_scale_inv.view(torch.uint8) + + # fixed-shift canonicalization: scale code = e + 121 (recomputed independently) + e, _, _ = _mxfp4_decompose(w) + expect_codes = (e.view(m, n // 32) + 121).clamp(min=0).to(torch.uint8) + assert torch.equal(codes[:m, : n // 32], expect_codes), "scale codes not fixed-shift" + + # pmax=3 block: TileKernels bytes (192 @ 2^-6), differing from the bridge's + assert int(codes[2, 0]) == 121 and int(data[2, 0]) == 0x74, "pmax=3 canonical bytes" + bridge_qt = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(what) + bdata = bridge_qt._rowwise_data.view(torch.uint8) + bcodes = bridge_qt._rowwise_scale_inv.view(torch.uint8) + assert (int(bcodes[2, 0]), int(bdata[2, 0])) == (120, 0x7C), "bridge canonical changed?" + # pmax=6 block: the two canonicalizations coincide + assert int(codes[3, 0]) == int(bcodes[3, 0]) and int(data[3, 0]) == int(bdata[3, 0]) + + # floor block: code 0, payload absorbs (0.5 * 2^(e+127) = 1.0 -> byte 56) + assert int(codes[0, 0]) == 0 and int(data[0, 0]) == 56, "code-0 absorption" + # zero block: TileKernels clamp convention -> code 0, payload 0; -0 keeps sign + assert int(codes[1, 0]) == 0 and int(data[1, 1]) == 0 + assert int(data[1, 0]) == 0x80, "-0 payload sign lost" + print(" rowwise canonical bytes + decode: PASS") + + +def test_colwise_matches_bridge(): + m, n = 128, 256 + w = make_weight(m, n, zero_blocks=4, outliers=4) + what = mxfp4_fake_quantize(w) + + qd = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + dt = mxfp4_qat_direct_quantize(w, qd) + qb = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + bt = qb.quantize(what) + + dcol = dt._columnwise_data.view(torch.uint8) + bcol = bt._columnwise_data.view(torch.uint8) + assert torch.equal(dcol, bcol), "columnwise payload bytes differ from bridge" + dsc = dt._columnwise_scale_inv.view(torch.uint8)[: m // 32, :n] + bsc = bt._columnwise_scale_inv.view(torch.uint8)[: m // 32, :n] + # zero columns may legitimately carry different (unused) scale codes + nz = dcol.view(m // 32, 32, n).ne(0).any(dim=1) + assert torch.equal(dsc[nz], bsc[nz]), "columnwise scale codes differ from bridge" + print(" colwise bytes == bridge: PASS") + + +def test_blockwise_folding_and_decode(): + def mk_quantizer(): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, + force_pow_2_scales=True, block_scaling_dim=2, + ) + + # spread sweep: exact through d<=14, first loss at d=15 (E4M3 subnormal bound) + payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + boundary = None + for d in range(17): + w = torch.zeros(128, 128, dtype=torch.float32, device=DEV) + w[0, 0] = 6.0 + w[1, :7] = payloads.to(DEV) * 2.0**-d + w[1, 7:14] = -w[1, :7] + qt = mxfp4_qat_direct_quantize(w, mk_quantizer()) + assert qt._rowwise_scale_inv[0, 0].item() == 2.0 ** (0 - 6), f"tile scale d={d}" + dec = _decode_blockwise_rowwise(qt, 128, 128) + exact = torch.equal(dec, mxfp4_fake_quantize(w).to(torch.float32)) + if boundary is None and not exact: + boundary = d + assert boundary == 15, f"expected first loss at spread 15, got {boundary}" + + # normal weights: decode + TE dequantize + columnwise transpose all exact + w = make_weight(256, 256, zero_blocks=4).to(torch.bfloat16) + what = mxfp4_fake_quantize(w) + qt = mxfp4_qat_direct_quantize(w, mk_quantizer()) + assert torch.equal(_decode_blockwise_rowwise(qt, 256, 256), what.to(torch.float32)) + assert torch.equal(qt.dequantize(dtype=torch.float32), what.to(torch.float32)) + qt.update_usage(rowwise_usage=False, columnwise_usage=True) + assert torch.equal(qt.dequantize(dtype=torch.float32), what.to(torch.float32)) + + # all-zero tile: clamped scale, zero payloads + wz = torch.zeros(128, 128, dtype=torch.bfloat16, device=DEV) + qz = mxfp4_qat_direct_quantize(wz, mk_quantizer()) + assert (qz._rowwise_data.view(torch.uint8) == 0).all() + assert qz._rowwise_scale_inv[0, 0].item() == 2.0**-127 + print(f" blockwise folding: exact through d=14, first loss at d={boundary}: PASS") + + +def test_e2e_direct_matches_bridge(): + torch.manual_seed(7) + tokens, in_f, out_f = 192, 256, 256 + x0 = torch.randn(tokens, in_f, device=DEV, dtype=torch.bfloat16) + + for recipe_cls in (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling): + for override in (None, "dequantized", "high_precision"): + results = {} + for direct in (False, True): + torch.manual_seed(11) + mod = te.Linear(in_f, out_f, bias=False, params_dtype=torch.bfloat16, device=DEV) + recipe = recipe_cls() + recipe.backward_override = override + recipe.mxfp4_qat_direct = direct + x = x0.clone().requires_grad_(True) + outs = [] + for first in (True, False, True): # miss, cache-hit, update paths + with fp8_autocast(enabled=True, fp8_recipe=recipe): + out = mod(x, is_first_microbatch=first) + outs.append(out.detach().clone()) + out.backward(torch.ones_like(out)) + results[direct] = ( + outs, + x.grad.detach().clone(), + mod.weight.grad.detach().clone(), + ) + for i, (ob, od) in enumerate(zip(results[False][0], results[True][0])): + assert torch.equal(ob, od), ( + f"fwd differs (bridge vs direct) {recipe_cls.__name__} " + f"override={override} step={i}" + ) + assert torch.equal(results[False][1], results[True][1]), ( + f"dgrad differs {recipe_cls.__name__} override={override}" + ) + assert torch.equal(results[False][2], results[True][2]), ( + f"wgrad differs {recipe_cls.__name__} override={override}" + ) + print(" e2e direct == bridge (fwd/dgrad/wgrad bitwise, 2 recipes x 3 overrides): PASS") + + +def test_direct_flag_plumbing(): + r = MXFP4QATMXFP8BlockScaling() + assert r.mxfp4_qat() and not r.mxfp4_qat_direct_conversion() + r.mxfp4_qat_direct = True + assert r.mxfp4_qat_direct_conversion() + from transformer_engine.common.recipe import MXFP8BlockScaling + + base = MXFP8BlockScaling() + base.mxfp4_qat_direct = True # direct without QAT is inert + assert not base.mxfp4_qat_direct_conversion() + print(" recipe flag plumbing: PASS") + + + +def _edge_weight(): + w = make_weight(128, 256) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) + w[1, :32] = 0.0 + w[1, 0] = -0.0 + w[2, :32] = 0.0 + w[2, 0] = 3.25 + w[3, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) + w[4, 5] = float("inf") + w[5, 9] = float("nan") + return w + + +_DIRECT_KERNELS = {} + + +def _intree_direct(fast_math=False): + """JIT-compile the in-tree CUDA direct kernels for testing.""" + if fast_math not in _DIRECT_KERNELS: + from torch.utils.cpp_extension import load_inline + + src_dir = os.environ.get( + "NVTE_MXFP4_QAT_TEST_SRC", + os.path.join(os.path.dirname(os.path.abspath(__file__)), + "..", "..", "transformer_engine", "common"), + ) + cuda_src = r""" +#include +#include +#include +#include "cast/mxfp4/direct_convert_mxfp4.cuh" + +std::vector direct_row(torch::Tensor w) { + TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.dim() == 2 && w.size(1) % 32 == 0); + const int64_t rows = w.size(0), cols = w.size(1), bpr = cols / 32; + auto data = torch::empty({rows, cols}, w.options().dtype(torch::kUInt8)); + auto scales = torch::zeros({(rows + 127) / 128 * 128, (bpr + 3) / 4 * 4}, + w.options().dtype(torch::kUInt8)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int ss = scales.size(1); + if (w.scalar_type() == at::kBFloat16) { + transformer_engine::dispatch::mxfp4::direct_mxfp8_rowwise_launch( + reinterpret_cast(w.data_ptr()), data.data_ptr(), + scales.data_ptr(), rows, cols, ss, stream); + } else { + transformer_engine::dispatch::mxfp4::direct_mxfp8_rowwise_launch( + w.data_ptr(), data.data_ptr(), scales.data_ptr(), rows, cols, + ss, stream); + } + return {data, scales}; +} + +std::vector direct_block(torch::Tensor w) { + TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.dim() == 2); + TORCH_CHECK(w.size(0) % 128 == 0 && w.size(1) % 128 == 0); + const int64_t rows = w.size(0), cols = w.size(1), tm = rows / 128, tn = cols / 128; + auto data = torch::empty({rows, cols}, w.options().dtype(torch::kUInt8)); + auto scales = torch::zeros({tm, (tn + 3) / 4 * 4}, w.options().dtype(torch::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int ss = scales.size(1); + if (w.scalar_type() == at::kBFloat16) { + transformer_engine::dispatch::mxfp4::direct_blockwise_launch( + reinterpret_cast(w.data_ptr()), data.data_ptr(), + scales.data_ptr(), rows, cols, ss, stream); + } else { + transformer_engine::dispatch::mxfp4::direct_blockwise_launch( + w.data_ptr(), data.data_ptr(), scales.data_ptr(), rows, cols, ss, + stream); + } + return {data, scales}; +} +""" + _DIRECT_KERNELS[fast_math] = load_inline( + name="nvte_mxfp4_direct_intree" + ("_fastmath" if fast_math else ""), + cpp_sources=( + "#include \n#include \n" + "std::vector direct_row(torch::Tensor w);\n" + "std::vector direct_block(torch::Tensor w);" + ), + cuda_sources=cuda_src, + functions=["direct_row", "direct_block"], + extra_include_paths=[src_dir], + extra_cuda_cflags=[ + "-O3", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + ] + (["--use_fast_math"] if fast_math else []), + verbose=False, + ) + return _DIRECT_KERNELS[fast_math] + + +def test_cuda_kernels_parity(): + """In-tree CUDA direct kernels: bitwise parity, normal and fast-math builds.""" + for fast_math in (False, True): + k = _intree_direct(fast_math=fast_math) + tag = "fast-math" if fast_math else "normal" + for m, n in [(1, 32), (3, 416), (128, 256), (192, 448)]: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=2, outliers=2) + dd, ds = k.direct_row(w.contiguous()) + td, ts = _direct_mxfp8_rowwise(*_mxfp4_decompose(w), m, n) + assert torch.equal(dd, td) and torch.equal(ds, ts), ( + f"cuda row {tag} differs {m}x{n} {dtype}" + ) + wx = _edge_weight() + dd, ds = k.direct_row(wx.contiguous()) + td, ts = _direct_mxfp8_rowwise(*_mxfp4_decompose(wx), 128, 256) + assert torch.equal(dd, td) and torch.equal(ds, ts), f"cuda row {tag} edge differs" + + for m, n in [(128, 128), (256, 384)]: + for dtype in (torch.bfloat16, torch.float32): + w = make_weight(m, n, dtype=dtype, zero_blocks=3, outliers=3) + dd, ds = k.direct_block(w.contiguous()) + td, ts = _direct_blockwise_128(*_mxfp4_decompose(w), m, n) + assert torch.equal(dd, td) and torch.equal(ds, ts), ( + f"cuda blockwise {tag} differs {m}x{n} {dtype}" + ) + payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + for d in (0, 11, 14, 15): + w = torch.zeros(128, 128, dtype=torch.float32, device=DEV) + w[0, 0] = 6.0 + w[1, :7] = payloads.to(DEV) * 2.0**-d + dd, ds = k.direct_block(w) + td, ts = _direct_blockwise_128(*_mxfp4_decompose(w), 128, 128) + assert torch.equal(dd, td) and torch.equal(ds, ts), f"cuda blockwise {tag} d={d}" + wx2 = _edge_weight() + dd, ds = k.direct_block(wx2) + td, ts = _direct_blockwise_128(*_mxfp4_decompose(wx2), 128, 256) + assert torch.equal(dd, td) and torch.equal(ds, ts), f"cuda blockwise {tag} edge differs" + + if RUN_PERF: + k = _intree_direct() + rows = cols = 8192 + wb = make_weight(rows, cols) + for _ in range(3): + k.direct_row(wb) + k.direct_block(wb) + torch.cuda.synchronize() + + def timeit(fn, iters=20): + s0, e0 = torch.cuda.Event(True), torch.cuda.Event(True) + s0.record() + for _ in range(iters): + fn() + e0.record() + torch.cuda.synchronize() + return s0.elapsed_time(e0) / iters * 1e3 + + t_row = timeit(lambda: k.direct_row(wb)) + t_blk = timeit(lambda: k.direct_block(wb)) + moved = rows * cols * 2 + rows * cols + rows * cols // 32 + print( + f" cuda direct row: {t_row:.1f} us ({moved / t_row / 1e3:.0f} GB/s), " + f"blockwise: {t_blk:.1f} us" + ) + print(" cuda kernels (normal + fast-math): PASS") + + + +def test_tilekernels_byte_parity(): + """Direct kernels vs the real TileKernels quantize+lift chain, byte for byte. + + round_sf=True is the deployment semantics (power-of-two scales; the False + default emits raw amax/6 fp32 scales that the lift then truncates). + Exclusions: blocks whose lifted scale is UE8M0 code 0 -- TileKernels' + fp32-scale store materializes code 0 as 0.0 (their analog of the + ptx::exp2f defect fixed in TE); the payload bytes still match there. + """ + import sys + + tk_root = os.environ.get( + "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") + ) + if os.path.isdir(tk_root) and tk_root not in sys.path: + sys.path.insert(0, tk_root) + try: + from tile_kernels.quant.per_token_cast_kernel import per_token_cast + from tile_kernels.quant.per_block_cast_lossless_kernel import per_block_cast_lossless + except Exception as e: + print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") + return + if not hasattr(tex, "mxfp4_direct_mxfp8_rowwise"): + print(" SKIP (direct bindings not built)") + return + + torch.manual_seed(9) + m, n = 256, 512 + w = make_weight(m, n, outliers=4) + w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) # code-0 block + w[1, :32] = 0.0 + w[1, 0] = 3.25 # pmax=3 canonical block + w[2, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) + + fp4 = per_token_cast(w, "e2m1", 32, round_sf=True) + tk_d, tk_s = per_block_cast_lossless(fp4, "e4m3", (1, 32), (1, 32), round_sf=True) + od, os_ = tex.mxfp4_direct_mxfp8_rowwise(w) + assert torch.equal(od, tk_d.view(torch.uint8)), "rowwise payload bytes differ" + oc = os_[:m, : n // 32].to(torch.int32) + tk_sf = tk_s.float()[:m, : n // 32] + defect = tk_sf == 0.0 # their code-0 fp32-scale store defect + assert torch.equal(defect, oc == 0), "TK zero-sf defect not confined to code-0 blocks" + assert defect.any(), "code-0 case not exercised" + codes = (torch.log2(tk_sf.clamp(min=1e-45)) + 127).round().to(torch.int32) + assert torch.equal(oc[~defect], codes[~defect]), "rowwise scale codes differ" + + w2 = (torch.randn(256, 256, device=DEV, dtype=torch.float32) * 0.02 + 0.05).to(torch.bfloat16) + fp42 = per_token_cast(w2, "e2m1", 32, round_sf=True) + tk_bd, tk_bs = per_block_cast_lossless(fp42, "e4m3", (1, 32), (128, 128), round_sf=True) + obd, obs = tex.mxfp4_direct_blockwise(w2) + assert torch.equal(obd, tk_bd.view(torch.uint8)), "blockwise payload bytes differ" + assert torch.equal(obs[:, :2], tk_bs.float()), "blockwise tile scales differ" + print(" byte parity vs TileKernels chain (payloads everywhere, scales outside " + "their code-0 store defect): PASS") + + +TESTS = [ + test_rowwise_canonical_bytes_and_decode, + test_cuda_kernels_parity, + test_colwise_matches_bridge, + test_blockwise_folding_and_decode, + test_direct_flag_plumbing, + test_tilekernels_byte_parity, + test_e2e_direct_matches_bridge, +] + +if __name__ == "__main__": + import sys + import traceback + + failed = 0 + for t in TESTS: + try: + t() + torch.cuda.synchronize() + print(f"PASS {t.__name__}") + except Exception: + failed += 1 + print(f"FAIL {t.__name__}") + traceback.print_exc() + print(f"{len(TESTS) - failed}/{len(TESTS)} passed") + sys.exit(1 if failed else 0) diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index bf44e3bffb..93db9a9759 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -16,6 +16,7 @@ #include "../utils.cuh" #include "dispatch/dequantize.cuh" #include "dispatch/quantize.cuh" +#include "mxfp4/direct_convert_mxfp4.cuh" #include "mxfp4/fake_quantize_mxfp4.cuh" #include "transformer_engine/transpose.h" @@ -80,6 +81,71 @@ void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStr NVTE_CHECK_CUDA(cudaGetLastError()); } +void nvte_mxfp4_direct_mxfp8_rowwise(const NVTETensor input, NVTETensor data, NVTETensor scales, + cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp4_direct_mxfp8_rowwise); + using namespace transformer_engine; + const auto &input_tensor = *convertNVTETensorCheck(input); + auto &data_tensor = *convertNVTETensorCheck(data); + auto &scales_tensor = *convertNVTETensorCheck(scales); + if (input_tensor.data.numel() == 0) return; + NVTE_CHECK(input_tensor.data.shape.size() == 2, "MXFP4 direct rowwise expects a 2D input."); + const size_t rows = input_tensor.data.shape[0], cols = input_tensor.data.shape[1]; + NVTE_CHECK(cols % 32 == 0, "MXFP4 direct rowwise needs the innermost dim divisible by 32."); + NVTE_CHECK(data_tensor.data.dtype == DType::kByte && scales_tensor.data.dtype == DType::kByte, + "MXFP4 direct rowwise outputs must be uint8."); + const int scale_stride = static_cast(scales_tensor.data.shape[1]); + switch (input_tensor.data.dtype) { + case DType::kBFloat16: + dispatch::mxfp4::direct_mxfp8_rowwise_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(data_tensor.data.dptr), + reinterpret_cast(scales_tensor.data.dptr), rows, cols, scale_stride, stream); + break; + case DType::kFloat32: + dispatch::mxfp4::direct_mxfp8_rowwise_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(data_tensor.data.dptr), + reinterpret_cast(scales_tensor.data.dptr), rows, cols, scale_stride, stream); + break; + default: + NVTE_ERROR("MXFP4 direct conversion supports BF16/FP32 inputs only."); + } +} + +void nvte_mxfp4_direct_blockwise(const NVTETensor input, NVTETensor data, NVTETensor scales, + cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp4_direct_blockwise); + using namespace transformer_engine; + const auto &input_tensor = *convertNVTETensorCheck(input); + auto &data_tensor = *convertNVTETensorCheck(data); + auto &scales_tensor = *convertNVTETensorCheck(scales); + if (input_tensor.data.numel() == 0) return; + NVTE_CHECK(input_tensor.data.shape.size() == 2, "MXFP4 direct blockwise expects a 2D input."); + const size_t rows = input_tensor.data.shape[0], cols = input_tensor.data.shape[1]; + NVTE_CHECK(rows % 128 == 0 && cols % 128 == 0, + "MXFP4 direct blockwise needs 128-divisible dimensions."); + NVTE_CHECK(scales_tensor.data.dtype == DType::kFloat32, + "MXFP4 direct blockwise scales must be fp32."); + const int scale_stride = static_cast(scales_tensor.data.shape[1]); + switch (input_tensor.data.dtype) { + case DType::kBFloat16: + dispatch::mxfp4::direct_blockwise_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(data_tensor.data.dptr), + reinterpret_cast(scales_tensor.data.dptr), rows, cols, scale_stride, stream); + break; + case DType::kFloat32: + dispatch::mxfp4::direct_blockwise_launch( + reinterpret_cast(input_tensor.data.dptr), + reinterpret_cast(data_tensor.data.dptr), + reinterpret_cast(scales_tensor.data.dptr), rows, cols, scale_stride, stream); + break; + default: + NVTE_ERROR("MXFP4 direct conversion supports BF16/FP32 inputs only."); + } +} + void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dequantize); using namespace transformer_engine; diff --git a/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh new file mode 100644 index 0000000000..16f05ff1ce --- /dev/null +++ b/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh @@ -0,0 +1,255 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file direct_convert_mxfp4.cuh + * \brief Direct MXFP4 -> FP8 weight conversion kernels (MXFP4 QAT). + * + * Converts a bf16/fp32 master weight straight into the host recipe's FP8 + * representation without materializing the dequantized MXFP4-grid weight: + * + * - rowwise MXFP8: the TileKernels deployment canonicalization. Per 1x32 + * block, scale code = fp4_exp + 121 (fixed shift 6) clamped at UE8M0 + * code 0 with exact payload absorption; E4M3 payload = E2M1 payload + * * 2^shift, in [32, 384] (or [1, 192] in the clamped domain). + * - blockwise 128x128: the per_block_cast_lossless exponent folding. Tile + * scale = 2^(max block exponent - 6) clamped at 2^-127 (the fp32-stored + * scale stays on the E8M0 lattice Blackwell extracts); folded payloads + * round RTNE onto E4M3 (subnormals included) -- no 2^11 device assert: + * exact through spread 2^14, bounded rounding beyond. + * + * Value-domain semantics (bit-identical to the torch reference in + * transformer_engine/pytorch/mxfp4_qat_direct.py): all-zero blocks pin the + * fp4 exponent at -126 (TileKernels clamp); non-finite 1x32 blocks emit + * E4M3 NaN payloads (0x7F) with scale code 127 (rowwise); the rowwise + * payload keeps the sign of -0. + */ + +#ifndef TRANSFORMER_ENGINE_CAST_MXFP4_DIRECT_CONVERT_MXFP4_CUH_ +#define TRANSFORMER_ENGINE_CAST_MXFP4_DIRECT_CONVERT_MXFP4_CUH_ + +#include +#include +#include + +#include "fake_quantize_mxfp4.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp4 { +namespace direct_convert_kernel { + +using fake_quantize_kernel::E2M1_MAX; +using fake_quantize_kernel::MXFP4_BLOCK_SIZE; +using fake_quantize_kernel::THREADS_PER_CHUNK; +using fake_quantize_kernel::block_scale_exponent_from_bits; +using fake_quantize_kernel::exp2_from_int; +using fake_quantize_kernel::mul_rn_no_ftz; +using fake_quantize_kernel::round_e2m1_magnitude; + +// 2^k as fp32 for k in [-149, 127]: normal via exponent field, subnormal via +// mantissa bit, exact zero below the subnormal range. Integer construction. +__device__ __forceinline__ float pow2_f32_ext(const int k) { + if (k >= -126) { + return exp2_from_int(k); + } + if (k >= -149) { + return __int_as_float(1 << (k + 149)); + } + return 0.0f; +} + +// E4M3 byte of one exact payload value (sign of zero preserved, RTNE): the +// hardware cvt.rn.satfinite.e4m3x2 instruction via the CUDA fp8 intrinsics. +__device__ __forceinline__ void store_e4m3_pair(uint8_t *dst, const float lo, const float hi) { + const __nv_fp8x2_storage_t two = + __nv_cvt_float2_to_fp8x2(make_float2(lo, hi), __NV_SATFINITE, __NV_E4M3); + dst[0] = static_cast(two & 0xFF); + dst[1] = static_cast(two >> 8); +} + +// One element: MXFP4 RTNE at scale 2^e, then payload value * 2^shift. +template +__device__ __forceinline__ float direct_payload(const IType v, const float inv_scale, + const float pscale) { + const float vf = static_cast(v); + const float av = __int_as_float(__float_as_int(vf) & 0x7fffffff); + const float y = fminf(mul_rn_no_ftz(av, inv_scale), E2M1_MAX); + const float q = copysignf(round_e2m1_magnitude(y), vf); + return mul_rn_no_ftz(q, pscale); +} + +// ---------------------------------------------------------------- rowwise +// Same shape of work as the fake-quantize kernel: each thread owns one +// 16-byte input vector; aligned lane groups reduce the block amax over the +// fp32 bit patterns; every lane emits its payload bytes and the group leader +// writes the scale code into the padded (roundup(M,128), roundup(N/32,4)) +// layout. +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + direct_mxfp8_rowwise_kernel(const IType *__restrict__ input, uint8_t *__restrict__ data, + uint8_t *__restrict__ scales, const size_t num_vectors, + const int blocks_per_row, const int scale_stride) { + constexpr size_t VEC_ELTS = 16 / sizeof(IType); + constexpr int LANES_PER_BLOCK = MXFP4_BLOCK_SIZE / VEC_ELTS; + + const size_t stride = static_cast(gridDim.x) * blockDim.x; + for (size_t vec_idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + vec_idx < num_vectors; vec_idx += stride) { + Vec in_vec; + in_vec.load_from(input, vec_idx); + + int amax_bits = 0; +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + amax_bits = max(amax_bits, __float_as_int(static_cast(in_vec.data.elt[i])) & 0x7fffffff); + } + const unsigned lane = threadIdx.x & 31u; + const unsigned group_mask = + (LANES_PER_BLOCK == 32) + ? 0xffffffffu + : (((1u << LANES_PER_BLOCK) - 1u) << (lane & ~(LANES_PER_BLOCK - 1u))); +#pragma unroll + for (int offset = LANES_PER_BLOCK / 2; offset > 0; offset >>= 1) { + amax_bits = max(amax_bits, __shfl_xor_sync(group_mask, amax_bits, offset)); + } + const bool nonfinite = amax_bits >= 0x7f800000; + + // fp4 exponent (zero/subnormal amax pins at -126 = the TileKernels clamp); + // fixed-shift-6 lift with exact payload absorption below UE8M0 code 0 + const int e = block_scale_exponent_from_bits(amax_bits); + const int biased = e + 121; + const int shift = biased >= 0 ? 6 : e + 127; + const float inv_scale = exp2_from_int(-e); + const float pscale = exp2_from_int(shift); + + Vec out_vec; + if (nonfinite) { +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); ++i) { + out_vec.data.elt[i] = 0x7F; // E4M3 NaN payload + } + } else { +#pragma unroll + for (int i = 0; i < static_cast(VEC_ELTS); i += 2) { + store_e4m3_pair(&out_vec.data.elt[i], + direct_payload(in_vec.data.elt[i], inv_scale, pscale), + direct_payload(in_vec.data.elt[i + 1], inv_scale, pscale)); + } + } + out_vec.store_to(data, vec_idx); + + if ((vec_idx % LANES_PER_BLOCK) == 0) { + const size_t block_idx = vec_idx / LANES_PER_BLOCK; + const size_t r = block_idx / blocks_per_row; + const size_t c = block_idx % blocks_per_row; + scales[r * scale_stride + c] = + nonfinite ? static_cast(127) : static_cast(max(biased, 0)); + } + } +} + +// ---------------------------------------------------------------- blockwise +// One CTA per 128x128 tile, one thread per 1x32 block (512 threads). Stage 1 +// derives every block's fp4 exponent; a shared-memory tree reduce yields the +// tile maximum; stage 2 folds each block's payloads by 2^(e - tile_exp) and +// encodes RTNE onto E4M3 (subnormal payload codes included). +template +__global__ void __launch_bounds__(512) + direct_blockwise_kernel(const IType *__restrict__ input, uint8_t *__restrict__ data, + float *__restrict__ scales, const size_t cols, + const int scale_stride) { + constexpr int BLOCKS_PER_TILE_ROW = 128 / MXFP4_BLOCK_SIZE; // 4 + + const int t = threadIdx.x; + const size_t row = static_cast(blockIdx.y) * 128 + (t / BLOCKS_PER_TILE_ROW); + const size_t col0 = static_cast(blockIdx.x) * 128 + + static_cast(t % BLOCKS_PER_TILE_ROW) * MXFP4_BLOCK_SIZE; + const IType *src = input + row * cols + col0; + + float elts[MXFP4_BLOCK_SIZE]; + int amax_bits = 0; +#pragma unroll + for (int i = 0; i < static_cast(MXFP4_BLOCK_SIZE); ++i) { + elts[i] = static_cast(src[i]); + amax_bits = max(amax_bits, __float_as_int(elts[i]) & 0x7fffffff); + } + const bool nonfinite = amax_bits >= 0x7f800000; + const int e = block_scale_exponent_from_bits(amax_bits); + + __shared__ int e_sh[512]; + e_sh[t] = nonfinite ? -126 : e; // non-finite blocks do not raise the tile max + __syncthreads(); +#pragma unroll + for (int s = 256; s > 0; s >>= 1) { + if (t < s) { + e_sh[t] = max(e_sh[t], e_sh[t + s]); + } + __syncthreads(); + } + const int s_exp = max(e_sh[0] - 6, -127); + + uint8_t *dst = data + row * cols + col0; + if (nonfinite) { +#pragma unroll + for (int i = 0; i < static_cast(MXFP4_BLOCK_SIZE); ++i) { + dst[i] = 0x7F; + } + } else { + const float inv_scale = exp2_from_int(-e); + const float pscale = pow2_f32_ext(e - s_exp); +#pragma unroll + for (int i = 0; i < static_cast(MXFP4_BLOCK_SIZE); i += 2) { + uint8_t pair[2]; + const float av0 = __int_as_float(__float_as_int(elts[i]) & 0x7fffffff); + const float av1 = __int_as_float(__float_as_int(elts[i + 1]) & 0x7fffffff); + const float q0 = + copysignf(round_e2m1_magnitude(fminf(mul_rn_no_ftz(av0, inv_scale), E2M1_MAX)), elts[i]); + const float q1 = copysignf( + round_e2m1_magnitude(fminf(mul_rn_no_ftz(av1, inv_scale), E2M1_MAX)), elts[i + 1]); + store_e4m3_pair(pair, mul_rn_no_ftz(q0, pscale), mul_rn_no_ftz(q1, pscale)); + dst[i] = pair[0]; + dst[i + 1] = pair[1]; + } + } + + if (t == 0) { + scales[static_cast(blockIdx.y) * scale_stride + blockIdx.x] = + (s_exp == -127) ? __int_as_float(0x00400000) : exp2_from_int(s_exp); + } +} + +} // namespace direct_convert_kernel + +// Raw-pointer launchers. Callers guarantee contiguous tensors, 16-byte +// aligned input/data pointers, cols % 32 == 0 (rowwise) or rows/cols % 128 +// (blockwise), and zero-initialized padded scale buffers. +template +void direct_mxfp8_rowwise_launch(const IType *input, uint8_t *data, uint8_t *scales, + const size_t rows, const size_t cols, const int scale_stride, + cudaStream_t stream) { + using namespace direct_convert_kernel; + constexpr size_t VEC_ELTS = 16 / sizeof(IType); + const size_t num_vectors = rows * cols / VEC_ELTS; + const size_t blocks_needed = (num_vectors + THREADS_PER_CHUNK - 1) / THREADS_PER_CHUNK; + const int grid = static_cast(blocks_needed < 65535 ? blocks_needed : 65535); + direct_mxfp8_rowwise_kernel<<>>( + input, data, scales, num_vectors, static_cast(cols / MXFP4_BLOCK_SIZE), scale_stride); +} + +template +void direct_blockwise_launch(const IType *input, uint8_t *data, float *scales, const size_t rows, + const size_t cols, const int scale_stride, cudaStream_t stream) { + using namespace direct_convert_kernel; + const dim3 grid(static_cast(cols / 128), static_cast(rows / 128)); + direct_blockwise_kernel<<>>(input, data, scales, cols, + scale_stride); +} + +} // namespace mxfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_CAST_MXFP4_DIRECT_CONVERT_MXFP4_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index e19f40d3fd..5089fb35e6 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -427,6 +427,20 @@ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t str */ void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Direct MXFP4 -> MXFP8 rowwise weight conversion (fixed-shift-6 + * TileKernels canonicalization; E4M3 payload bytes + UE8M0 scale + * codes in the padded compact layout). Innermost dim must be /32. + */ +void nvte_mxfp4_direct_mxfp8_rowwise(const NVTETensor input, NVTETensor data, NVTETensor scales, + cudaStream_t stream); + +/*! \brief Direct MXFP4 -> FP8 128x128 blockwise weight conversion + * (per-tile exponent folding; fp32 power-of-two tile scales in the + * padded layout). Dimensions must be /128. + */ +void nvte_mxfp4_direct_blockwise(const NVTETensor input, NVTETensor data, NVTETensor scales, + cudaStream_t stream); + /*! \brief Casts input grouped tensor from reduced to higher precision. * In case of the MXFP8 dequantization, the dequantized values are stored to the rowwise * data of the output tensor, regardless of whether the row- or columnwise scaling is used. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index a7224f1e63..cad298498a 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -174,6 +174,11 @@ def mxfp4_qat(self): self, (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling) ) + def mxfp4_qat_direct_conversion(self): + """Whether QAT weights use the direct MXFP4->FP8 converter (TileKernels + canonicalization, no bf16 bridge tensor) instead of the fake-quant bridge.""" + return self.mxfp4_qat() and bool(getattr(self, "mxfp4_qat_direct", False)) + @dataclass(repr=False) class DelayedScaling(Recipe): @@ -375,6 +380,7 @@ class MXFP8BlockScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) mxfp4_qat_weights: bool = os.getenv("NVTE_MXFP4_QAT", "0") == "1" + mxfp4_qat_direct: bool = os.getenv("NVTE_MXFP4_QAT_DIRECT", "0") == "1" def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -445,6 +451,7 @@ class Float8BlockScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) mxfp4_qat_weights: bool = os.getenv("NVTE_MXFP4_QAT", "0") == "1" + mxfp4_qat_direct: bool = os.getenv("NVTE_MXFP4_QAT_DIRECT", "0") == "1" def __post_init__(self) -> None: assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 2a77895ead..5df2d51a9a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -348,6 +348,8 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz py::object dequantize(const py::handle &input, DType otype); at::Tensor mxfp4_fake_quantize(const at::Tensor &input); +std::tuple mxfp4_direct_mxfp8_rowwise(const at::Tensor &input); +std::tuple mxfp4_direct_blockwise(const at::Tensor &input); py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims, std::optional last_dims, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 41dcd03e7f..2108035cdb 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -574,6 +574,53 @@ at::Tensor mxfp4_fake_quantize(const at::Tensor &input) { return output; } +static void check_direct_input(const at::Tensor &input, int mod) { + TORCH_CHECK(input.is_cuda(), "mxfp4 direct conversion: input must be a CUDA tensor."); + TORCH_CHECK(input.dim() == 2, "mxfp4 direct conversion: input must be 2D."); + TORCH_CHECK(input.size(1) % mod == 0, "mxfp4 direct conversion: inner dim must be divisible by ", + mod); +} + +std::tuple mxfp4_direct_mxfp8_rowwise(const at::Tensor &input) { + check_direct_input(input, 32); + const c10::cuda::CUDAGuard device_guard(input.device()); + auto input_contiguous = input.contiguous(); + if (reinterpret_cast(input_contiguous.data_ptr()) % 16 != 0) { + input_contiguous = input_contiguous.clone(); + } + const int64_t rows = input.size(0), cols = input.size(1); + const int64_t bpr = cols / 32; + auto data = at::empty({rows, cols}, input.options().dtype(at::kByte)); + auto scales = at::zeros({(rows + 127) / 128 * 128, (bpr + 3) / 4 * 4}, + input.options().dtype(at::kByte)); + auto in_cpp = makeTransformerEngineTensor(input_contiguous); + auto data_cpp = makeTransformerEngineTensor(data); + auto scales_cpp = makeTransformerEngineTensor(scales); + nvte_mxfp4_direct_mxfp8_rowwise(in_cpp.data(), data_cpp.data(), scales_cpp.data(), + at::cuda::getCurrentCUDAStream()); + return {data, scales}; +} + +std::tuple mxfp4_direct_blockwise(const at::Tensor &input) { + check_direct_input(input, 128); + TORCH_CHECK(input.size(0) % 128 == 0, "mxfp4 direct blockwise: rows must be divisible by 128"); + const c10::cuda::CUDAGuard device_guard(input.device()); + auto input_contiguous = input.contiguous(); + if (reinterpret_cast(input_contiguous.data_ptr()) % 16 != 0) { + input_contiguous = input_contiguous.clone(); + } + const int64_t rows = input.size(0), cols = input.size(1); + const int64_t tm = rows / 128, tn = cols / 128; + auto data = at::empty({rows, cols}, input.options().dtype(at::kByte)); + auto scales = at::zeros({tm, (tn + 3) / 4 * 4}, input.options().dtype(at::kFloat)); + auto in_cpp = makeTransformerEngineTensor(input_contiguous); + auto data_cpp = makeTransformerEngineTensor(data); + auto scales_cpp = makeTransformerEngineTensor(scales); + nvte_mxfp4_direct_blockwise(in_cpp.data(), data_cpp.data(), scales_cpp.data(), + at::cuda::getCurrentCUDAStream()); + return {data, scales}; +} + py::object dequantize(const py::handle &input, transformer_engine::DType otype) { init_extension(); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index facbb0c489..30376ff151 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -202,6 +202,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("mxfp4_fake_quantize", &transformer_engine::pytorch::mxfp4_fake_quantize, "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)", py::arg("input"), py::call_guard()); + m.def("mxfp4_direct_mxfp8_rowwise", &transformer_engine::pytorch::mxfp4_direct_mxfp8_rowwise, + "Direct MXFP4->MXFP8 rowwise conversion (fixed-shift-6 canonicalization)", + py::arg("input"), py::call_guard()); + m.def("mxfp4_direct_blockwise", &transformer_engine::pytorch::mxfp4_direct_blockwise, + "Direct MXFP4->FP8 128x128 blockwise conversion (exponent folding)", + py::arg("input"), py::call_guard()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); m.def("create_empty_quantized_tensor", diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5447aa676f..ee47ab55b5 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -47,6 +47,7 @@ ) from ..constants import dist_group_type from ..mxfp4_qat import mxfp4_fake_quantize +from ..mxfp4_qat_direct import mxfp4_qat_direct_quantize, mxfp4_qat_direct_update_ from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer @@ -806,6 +807,7 @@ def quantize_weight( if FP8GlobalStateManager.is_fp8_enabled() else False ) + _mxfp4_qat_direct = _mxfp4_qat_active and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat_direct_conversion() if _mxfp4_qat_active and workspace_dtype == torch.float16: raise NotImplementedError( "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " @@ -852,6 +854,9 @@ def quantize_weight( if update_workspace: if tensor is None: raise ValueError("tensor kwarg must be provided to update FP8 workspace") + if _mxfp4_qat_direct: + mxfp4_qat_direct_update_(tensor, workspace) + return workspace, None if _mxfp4_qat_active: tensor = mxfp4_fake_quantize(tensor) if hasattr(workspace, "quantize_"): @@ -863,6 +868,11 @@ def quantize_weight( # Cache miss — create new workspace if tensor is None or quantizer is None: raise ValueError("tensor and quantizer kwargs must be provided to construct FP8 workspace") + if _mxfp4_qat_direct: + out = mxfp4_qat_direct_quantize(tensor, quantizer) + if cache: + return out, out + return out, None if _mxfp4_qat_active: tensor = mxfp4_fake_quantize(tensor) if cache: diff --git a/transformer_engine/pytorch/mxfp4_qat_direct.py b/transformer_engine/pytorch/mxfp4_qat_direct.py new file mode 100644 index 0000000000..cbd50724e9 --- /dev/null +++ b/transformer_engine/pytorch/mxfp4_qat_direct.py @@ -0,0 +1,199 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Direct MXFP4 -> FP8 weight conversion for MXFP4 QAT (TileKernels canonicalization). + +Converts the MXFP4 decomposition of a bf16/fp32 master weight straight into +the host recipe's FP8 representation without a bf16 bridge tensor: + +* MXFP8 rowwise (tex.mxfp4_direct_mxfp8_rowwise): the deployment fixed-shift-6 + canonicalization -- scale code = fp4_exp + 121 clamped at UE8M0 code 0 with + exact payload absorption; the bytes are invertible back to packed FP4. +* MXFP8 columnwise (backward_override=None only): a fresh lossy 32x1 + quantization of the grid values with the HOST encoder's amax rule + (scale = 2^ceil(log2(amax/448))), so dgrad bytes match the bridge path bit + for bit. This is the sole implementation of the columnwise direction (a + composite-torch computation, not a fallback for a kernel). +* Float8 blockwise 128x128 (tex.mxfp4_direct_blockwise): per-tile exponent + folding, RTNE onto E4M3 with subnormals -- exact through scale spread 2^14, + bounded beyond, no device assert. + +Decoded values equal the bridge path (mxfp4_fake_quantize) everywhere both +are defined; only the rowwise/blockwise raw bytes differ (fixed-shift vs +amax canonicalization). The numerical specification lives in +tests/pytorch/references/mxfp4_qat_direct_reference.py. +""" +import torch + +import transformer_engine_torch as tex + +from .mxfp4_qat import _MXFP4_BLOCK +from .tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from .tensor.float8_blockwise_tensor import Float8BlockQuantizer, Float8BlockwiseQTensor + +__all__ = ["mxfp4_qat_direct_quantize", "mxfp4_qat_direct_update_"] + +_E4M3_NAN = 0x7F + + +def _require_kernel(name: str): + fn = getattr(tex, name, None) + if fn is None: + raise RuntimeError( + f"transformer_engine_torch was built without {name}; rebuild Transformer Engine." + ) + return fn + + +def _prepare(weight: torch.Tensor) -> torch.Tensor: + w = weight.contiguous() + if w.data_ptr() % 16 != 0: + w = w.clone() + return w + + +def _roundup(x: int, m: int) -> int: + return (x + m - 1) // m * m + + +def _round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: + """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" + fine = torch.round(y * 2.0) * 0.5 + mid = torch.round(y) + coarse = torch.round(y * 0.5) * 2.0 + return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) + + +def _direct_mxfp8_columnwise(weight: torch.Tensor): + """Sole (composite-torch) implementation of the columnwise direction. + + MXFP4-decomposes the weight, then quantizes the grid values in 32x1 + blocks with the host encoder's amax rule so the bytes match the bridge + path's columnwise representation bit for bit. + """ + rows, cols = weight.shape + w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) + amax = w32.abs().amax(dim=-1, keepdim=True) + nonfinite = ~torch.isfinite(amax) + + bits = amax.view(torch.int32) + exp_field = bits >> 23 + mantissa = bits & 0x7FFFFF + e = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) + e = torch.where(exp_field > 0, e, torch.full_like(e, -126)) + e = e.clamp(min=-126, max=125) + e = torch.where(amax > 0, e, torch.full_like(e, -126)) + e = torch.where(nonfinite, torch.full_like(e, -126), e) + + scale = torch.ldexp(torch.ones_like(amax), e) + y = (w32 / scale).abs().clamp(max=6.0) + q = _round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) + q = torch.copysign(q, w32) + + vals = (q * scale).view(rows, cols) + nf_elem = nonfinite.expand(-1, -1, _MXFP4_BLOCK).reshape(rows, cols) + + vc = vals.view(rows // 32, 32, cols) + amax_c = vc.abs().amax(dim=1, keepdim=True) + amax_c = torch.where(torch.isfinite(amax_c), amax_c, torch.zeros_like(amax_c)) + cb = amax_c.view(torch.int32) + kf = (cb >> 23) - 127 + ec = kf - 8 + ((cb & 0x7FFFFF) > 0x600000).to(torch.int32) # host rule: ceil(log2(amax/448)) + ec = torch.where((cb >> 23) > 0, ec, torch.full_like(ec, -127)) + ec = ec.clamp(min=-127, max=127) + ec = torch.where(amax_c > 0, ec, torch.zeros_like(ec)) + inv = torch.ldexp(torch.ones_like(ec, dtype=torch.float32), -ec) + # + 0.0 canonicalizes -0 to +0 (the host colwise kernel drops the sign of zero) + payload = (vc * inv + 0.0).to(torch.float8_e4m3fn).view(torch.uint8).view(rows, cols) + payload = torch.where(nf_elem, torch.full_like(payload, _E4M3_NAN), payload) + + code = (ec + 127).clamp(0, 254).to(torch.uint8) + scale_inv = torch.zeros( + (_roundup(rows // 32, 4), _roundup(cols, 128)), dtype=torch.uint8, device=payload.device + ) + scale_inv[: rows // 32, :cols] = code.view(rows // 32, cols) + return payload, scale_inv + + +def mxfp4_qat_direct_quantize(weight: torch.Tensor, quantizer) -> torch.Tensor: + """Convert a bf16/fp32 weight straight to the host FP8 representation.""" + if not weight.is_cuda: + raise ValueError("MXFP4 QAT direct conversion expects a CUDA weight tensor") + if weight.dim() != 2: + raise ValueError(f"expected a 2D weight, got {tuple(weight.shape)}") + if weight.dtype not in (torch.bfloat16, torch.float32): + raise ValueError(f"bf16/fp32 weights only, got {weight.dtype}") + rows, cols = weight.shape + if cols % _MXFP4_BLOCK != 0: + raise ValueError(f"inner dim must be divisible by {_MXFP4_BLOCK}, got {cols}") + + if isinstance(quantizer, MXFP8Quantizer): + data, scale_inv = _require_kernel("mxfp4_direct_mxfp8_rowwise")(_prepare(weight)) + col_data, col_scale_inv = None, None + if quantizer.columnwise_usage: + if rows % 32 != 0: + raise ValueError(f"columnwise 32x1 blocks need rows % 32 == 0, got {rows}") + col_data, col_scale_inv = _direct_mxfp8_columnwise(weight) + return MXFP8Tensor( + shape=weight.shape, + dtype=weight.dtype, + rowwise_data=data if quantizer.rowwise_usage else None, + rowwise_scale_inv=scale_inv if quantizer.rowwise_usage else None, + columnwise_data=col_data, + columnwise_scale_inv=col_scale_inv, + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + with_gemm_swizzled_scales=False, + device=weight.device, + ) + + if isinstance(quantizer, Float8BlockQuantizer): + if getattr(quantizer, "block_scaling_dim", 2) != 2: + raise ValueError("direct blockwise conversion supports 128x128 (2D) scaling only") + data, scale_inv = _require_kernel("mxfp4_direct_blockwise")(_prepare(weight)) + out = Float8BlockwiseQTensor( + shape=weight.shape, + dtype=weight.dtype, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise_data=data, + rowwise_scale_inv=scale_inv, + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=quantizer, + is_2D_scaled=True, + device=weight.device, + ) + if quantizer.columnwise_usage: + out._create_columnwise() # exact transpose of the rowwise encoding + return out + + raise NotImplementedError( + f"MXFP4 QAT direct conversion does not support {type(quantizer).__name__}" + ) + + +def mxfp4_qat_direct_update_(weight: torch.Tensor, workspace) -> None: + """Requantize ``weight`` into an existing workspace's buffers in place.""" + quantizer = workspace._quantizer + fresh = mxfp4_qat_direct_quantize(weight, quantizer) + if isinstance(workspace, MXFP8Tensor): + if workspace._rowwise_data is not None: + workspace._rowwise_data.copy_(fresh._rowwise_data) + workspace._rowwise_scale_inv.copy_(fresh._rowwise_scale_inv) + if workspace._columnwise_data is not None: + workspace._columnwise_data.copy_(fresh._columnwise_data) + workspace._columnwise_scale_inv.copy_(fresh._columnwise_scale_inv) + return + if isinstance(workspace, Float8BlockwiseQTensor): + if workspace._rowwise_data is not None: + workspace._rowwise_data.copy_(fresh._rowwise_data) + workspace._rowwise_scale_inv.copy_(fresh._rowwise_scale_inv) + if workspace._columnwise_data is not None: + fresh._create_columnwise() + workspace._columnwise_data.copy_(fresh._columnwise_data) + workspace._columnwise_scale_inv.copy_(fresh._columnwise_scale_inv) + return + raise NotImplementedError( + f"MXFP4 QAT direct update does not support {type(workspace).__name__}" + ) From 2f40ca8695257c1c5ff1b80f46cc6db921a6e82b Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sun, 26 Jul 2026 23:02:53 -0700 Subject: [PATCH 10/14] Drop the TileKernels cross-check tests They imported an external TileKernels checkout with a developer-machine default path and can therefore never run in CI; the cross-backend comparison lives on as a standalone harness outside the tree. The in-tree suite keeps the kernel-vs-reference parity, exhaustive bf16 and fp32 fuzz coverage, and the raw-byte encode oracles. --- tests/pytorch/test_mxfp4_qat.py | 204 -------------------------------- 1 file changed, 204 deletions(-) diff --git a/tests/pytorch/test_mxfp4_qat.py b/tests/pytorch/test_mxfp4_qat.py index f964332a2f..e304b49a52 100644 --- a/tests/pytorch/test_mxfp4_qat.py +++ b/tests/pytorch/test_mxfp4_qat.py @@ -625,33 +625,6 @@ def test_deployment_top_cap_domain(): assert torch.equal(gt.to(torch.float64), fake_quant_ref_fp64(wc)), "top-cap vs fp64 ref" -def test_tilekernels_cross_parity(): - """Bitwise parity with the TileKernels torch reference (round_sf=True) on the finite below-cap domain.""" - import sys - - tk_root = os.environ.get( - "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") - ) - if os.path.isdir(tk_root) and tk_root not in sys.path: - sys.path.insert(0, tk_root) - try: - from tile_kernels.torch.cast import cast as tk_cast, cast_back as tk_cast_back - except Exception as e: - print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") - return - - torch.manual_seed(21) - for m, n in SHAPES: - w = make_weight(m, n, zero_blocks=4, outliers=4) - w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) - w[1, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) - w[2, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) - data, sf = tk_cast(w, "e2m1", block_size=(1, 32), round_sf=True) - tk_dq = tk_cast_back((data, sf), "bf16", block_size=(1, 32)) - ours = mxfp4_fake_quantize(w) - _assert_bits_equal(tk_dq, ours, f"TileKernels cross parity {m}x{n}") - - def test_blockwise_feasibility_enumeration(): """Per-element E4M3-lattice prediction == TE 128x128 quantizer for exponent spreads d=0..16.""" payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) @@ -700,181 +673,6 @@ def test_fp16_pipeline_rejected(): pass -def test_cross_impl_matrix(): - """TileKernels vs CUDA kernel vs torch reference, bitwise, over all edge domains and every quant/dequant hop.""" - import sys - - tk_root = os.environ.get( - "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") - ) - if os.path.isdir(tk_root) and tk_root not in sys.path: - sys.path.insert(0, tk_root) - try: - from tile_kernels.torch.cast import ( - cast as tk_cast, - cast_back as tk_cast_back, - get_min_clamp_val as tk_min_clamp, - ) - except Exception as e: - print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") - return - kern = _intree_kernel().mxfp4_fake_quantize_intree - tref = mxfp4_fake_quantize_reference - CAP = 6.0 * 2.0**125 - - def tk_roundtrip(x, block, fmt): - pair = tk_cast(x, "e2m1" if fmt == "e2m1" else "e4m3", block_size=block, round_sf=True) - return tk_cast_back(pair, "fp32", block), tk_cast_back(pair, "bf16", block) - - torch.manual_seed(77) - w = make_weight(64, 256, zero_blocks=2, outliers=4) - w[0, :32] = 0.0 - w[1, :32] = -0.0 - w[2, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) - w[2, 0] = 2.0**-128 - w[3, :32] = torch.tensor(2.0**-133, dtype=torch.bfloat16) - w[4, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) - w[5, :7] = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], dtype=torch.bfloat16) - w[5, 7] = 6.0 - a15 = torch.tensor(1.5, dtype=torch.bfloat16).view(torch.int16) - w[6, 0] = (a15 - 1).view(torch.bfloat16) - w[6, 32] = torch.tensor(1.5, dtype=torch.bfloat16) - w[6, 64] = (a15 + 1).view(torch.bfloat16) - - for W in (w, w.to(torch.float32)): - t = tref(W) - _assert_bits_equal(kern(W.contiguous()), t, "fake-quant CUDA vs torch") - tk32, tkb16 = tk_roundtrip(W, (1, 32), "e2m1") - _assert_bits_equal(tk32, t.to(torch.float32), "fake-quant TK vs torch (fp32)") - _assert_bits_equal(tkb16, t.to(torch.bfloat16), "fake-quant TK vs torch (bf16)") - print(" fake-quant 3-way bitwise: PASS (bf16 + fp32 corpus)") - - wnf = w.clone() - wnf[8, 40] = float("inf") - wnf[9, 33] = float("-inf") - wnf[10, 5] = float("nan") - t = tref(wnf) - _assert_same_with_nan(kern(wnf.contiguous()), t, "nonfinite CUDA") - assert torch.isnan(t[8, 32:64]).all() and torch.isnan(t[9, 32:64]).all() - assert torch.isnan(t[10, :32]).all() and torch.isfinite(t[11].to(torch.float32)).all() - print(" nonfinite 2-way (whole-block NaN poisoning): PASS (TK policy differs, excluded)") - - wcap = torch.zeros(1, 32, dtype=torch.bfloat16, device=DEV) - wcap[0, 0] = torch.tensor(3.2e38, dtype=torch.bfloat16) - t = tref(wcap) - assert t[0, 0].item() == CAP, "satfinite cap at 6*2^125" - _assert_bits_equal(kern(wcap.contiguous()), t, "cap CUDA") - tk32, _ = tk_roundtrip(wcap, (1, 32), "e2m1") - assert not torch.equal(tk32, t.to(torch.float32)), "TK must diverge above 6*2^125" - print( - f" above-cap domain: ours satfinite {CAP:.3e}, TK diverges as documented " - f"(payload at 2^126 -> {tk32[0, 0].item()}): PASS" - ) - - w_hat = tref(w) - wh32 = w_hat.to(torch.float32) - q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=True).quantize(w_hat) - m, n = w_hat.shape - data = q._rowwise_data.view(torch.uint8)[:m, :n] - codes = q._rowwise_scale_inv.view(torch.uint8)[:m, : n // 32].to(torch.int32) - raw = data.view(torch.float8_e4m3fn).to(torch.float32).view(m, n // 32, 32) * torch.ldexp( - torch.ones_like(codes, dtype=torch.float32), codes - 127 - ).unsqueeze(-1) - assert torch.equal(raw.view(m, n), wh32), "mxfp8 row RAW encode not lossless" - e4m3_floor = float(tk_min_clamp(torch.float8_e4m3fn)) - blk_amax = wh32.abs().view(m, n // 32, 32).amax(dim=-1) - ok = (blk_amax >= e4m3_floor) | (blk_amax == 0) - okm = ok.unsqueeze(-1).expand(m, n // 32, 32).reshape(m, n) - tk32, tkb16 = tk_roundtrip(w_hat, (1, 32), "e4m3") - assert torch.equal( - tk32.view(torch.int32)[okm], wh32.view(torch.int32)[okm] - ), "mxfp8 row TK roundtrip (fp32)" - assert torch.equal( - tkb16.view(torch.int16)[okm], w_hat.view(torch.int16)[okm] - ), "mxfp8 row TK roundtrip (bf16)" - assert (tk32[~okm] == 0).all(), "TK e4m3 sub-floor blocks must flush to zero" - assert (~ok).sum() > 0, "TK e4m3 floor case not exercised" - dq = q.dequantize(dtype=torch.float32) - assert torch.equal(dq, wh32), "mxfp8 row TE software dequant" - print( - " mxfp8 row: TE raw encode lossless incl 2^-127; TK roundtrip bitwise on " - f"amax >= {e4m3_floor:g} blocks ({int((~ok).sum())} sub-floor blocks flush in TK " - "e4m3, a TK-only floor); TE software dequant exact incl 2^-127: PASS" - ) - - q.update_usage(rowwise_usage=False, columnwise_usage=True) - te_col = q.dequantize(dtype=torch.float32) - tk_col32, _ = tk_roundtrip(w_hat, (32, 1), "e4m3") - col_amax = wh32.abs().view(m // 32, 32, n).amax(dim=1) - cok = (col_amax >= e4m3_floor) | (col_amax == 0) - cokm = cok.unsqueeze(1).expand(m // 32, 32, n).reshape(m, n) - nzm = (wh32 != 0) & cokm - zm = (wh32 == 0) & cokm - assert torch.equal( - te_col.view(torch.int32)[nzm], tk_col32.view(torch.int32)[nzm] - ), "mxfp8 col TE vs TK (nonzero)" - assert (te_col[zm] == 0).all() and (tk_col32[zm] == 0).all(), "mxfp8 col zeros" - negz = zm & torch.signbit(wh32) - assert negz.any() - assert torch.signbit(tk_col32)[negz].all(), "TK col must preserve -0" - assert not torch.signbit(te_col)[negz].any(), "TE col canonicalizes -0 -> +0" - blk = wh32.view(m // 32, 32, n) - bound = blk.abs().amax(dim=1, keepdim=True) * 2.0**-3 - assert ((te_col.view(m // 32, 32, n) - blk).abs() <= bound + 1e-30).all() - print( - " mxfp8 col (32x1): TE == TK bitwise on nonzeros, zeros value-equal " - "(TE col canonicalizes -0 -> +0, TK preserves the sign), bounded vs w_hat: PASS" - ) - - wb_hat = tref(make_weight(128, 256, zero_blocks=4)) - wb32 = wb_hat.to(torch.float32) - qb = Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - force_pow_2_scales=True, - block_scaling_dim=2, - ).quantize(wb_hat) - assert torch.equal(qb.dequantize(dtype=torch.float32), wb32) - assert torch.equal(qb.dequantize(dtype=torch.bfloat16), wb_hat) - tk32, tkb16 = tk_roundtrip(wb_hat, (128, 128), "e4m3") - _assert_bits_equal(tk32, wb32, "blockwise TK roundtrip (fp32)") - _assert_bits_equal(tkb16, wb_hat, "blockwise TK roundtrip (bf16)") - qb.update_usage(rowwise_usage=False, columnwise_usage=True) - assert torch.equal(qb.dequantize(dtype=torch.float32), wb32), "blockwise col transpose" - print( - " blockwise 128x128: TE row+col dequant == TK roundtrip == w_hat bitwise " - "(fp32 + bf16 targets): PASS" - ) - - bits = torch.arange(65536, dtype=torch.int32, device=DEV).to(torch.int16) - we = bits.view(torch.bfloat16).view(2048, 32) - t = tref(we) - _assert_same_with_nan(kern(we.contiguous()), t, "exhaustive CUDA") - we32 = we.to(torch.float32) - ok_rows = torch.isfinite(we32).all(dim=-1) & (we32.abs().amax(dim=-1) <= CAP) - tk32, _ = tk_roundtrip(we[ok_rows], (1, 32), "e2m1") - _assert_bits_equal(tk32, t[ok_rows].to(torch.float32), "exhaustive TK") - print( - " bf16 exhaustive 65536 patterns: CUDA/torch full, " - f"TK on {int(ok_rows.sum())}/2048 finite below-cap rows, bitwise: PASS" - ) - - g = torch.Generator().manual_seed(123) - fb = torch.randint(-(2**31), 2**31 - 1, (512, 32), generator=g, dtype=torch.int64) - wf = fb.to(torch.int32).cuda().view(torch.float32) - t = tref(wf) - _assert_same_with_nan(kern(wf.contiguous()), t, "fuzz CUDA") - ok_rows = torch.isfinite(wf).all(dim=-1) & (wf.abs().amax(dim=-1) <= CAP) - if ok_rows.any(): - tk32, _ = tk_roundtrip(wf[ok_rows], (1, 32), "e2m1") - _assert_bits_equal(tk32, t[ok_rows], "fuzz TK") - print( - f" fp32 bit-fuzz 512 blocks: CUDA/torch full, TK on {int(ok_rows.sum())} " - "finite below-cap rows, bitwise: PASS" - ) - - def test_recipe_switch_invalidates_cache(): """base<->QAT recipe switches must requantize instead of reusing cached weight workspaces.""" import warnings as _warnings @@ -1079,13 +877,11 @@ def test_e2e_matrix(): test_fp32_bit_fuzz, test_scale_threshold_and_rtne_midpoints, test_deployment_top_cap_domain, - test_tilekernels_cross_parity, test_blockwise_feasibility_enumeration, test_recipe_switch_invalidates_cache, test_recipe_field_controls_qat, test_ops_api_rejected, test_fp16_pipeline_rejected, - test_cross_impl_matrix, test_kernel_perf, test_kernel_fast_math_immune, test_e2e_matrix, From 7a4e47e153f2822f98c3b8559739def6d4e30acd Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sun, 26 Jul 2026 23:03:24 -0700 Subject: [PATCH 11/14] Drop the TileKernels byte-parity test from the direct suite Same rationale as the bridge suite: it imported an external TileKernels checkout via a developer-machine default path and cannot run in CI. The direct converters keep kernel-vs-reference bitwise parity, canonical-byte assertions, and the e2e bitwise equality against the bridge mode. --- tests/pytorch/test_mxfp4_qat_direct.py | 57 -------------------------- 1 file changed, 57 deletions(-) diff --git a/tests/pytorch/test_mxfp4_qat_direct.py b/tests/pytorch/test_mxfp4_qat_direct.py index 7f81520211..220100d223 100644 --- a/tests/pytorch/test_mxfp4_qat_direct.py +++ b/tests/pytorch/test_mxfp4_qat_direct.py @@ -399,69 +399,12 @@ def timeit(fn, iters=20): -def test_tilekernels_byte_parity(): - """Direct kernels vs the real TileKernels quantize+lift chain, byte for byte. - - round_sf=True is the deployment semantics (power-of-two scales; the False - default emits raw amax/6 fp32 scales that the lift then truncates). - Exclusions: blocks whose lifted scale is UE8M0 code 0 -- TileKernels' - fp32-scale store materializes code 0 as 0.0 (their analog of the - ptx::exp2f defect fixed in TE); the payload bytes still match there. - """ - import sys - - tk_root = os.environ.get( - "NVTE_MXFP4_QAT_TILEKERNELS", os.path.expanduser("~/Desktop/v4/TileKernels") - ) - if os.path.isdir(tk_root) and tk_root not in sys.path: - sys.path.insert(0, tk_root) - try: - from tile_kernels.quant.per_token_cast_kernel import per_token_cast - from tile_kernels.quant.per_block_cast_lossless_kernel import per_block_cast_lossless - except Exception as e: - print(f" SKIP (tile_kernels unavailable: {type(e).__name__}: {e})") - return - if not hasattr(tex, "mxfp4_direct_mxfp8_rowwise"): - print(" SKIP (direct bindings not built)") - return - - torch.manual_seed(9) - m, n = 256, 512 - w = make_weight(m, n, outliers=4) - w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) # code-0 block - w[1, :32] = 0.0 - w[1, 0] = 3.25 # pmax=3 canonical block - w[2, :32] = torch.tensor(2.0**120, dtype=torch.bfloat16) - - fp4 = per_token_cast(w, "e2m1", 32, round_sf=True) - tk_d, tk_s = per_block_cast_lossless(fp4, "e4m3", (1, 32), (1, 32), round_sf=True) - od, os_ = tex.mxfp4_direct_mxfp8_rowwise(w) - assert torch.equal(od, tk_d.view(torch.uint8)), "rowwise payload bytes differ" - oc = os_[:m, : n // 32].to(torch.int32) - tk_sf = tk_s.float()[:m, : n // 32] - defect = tk_sf == 0.0 # their code-0 fp32-scale store defect - assert torch.equal(defect, oc == 0), "TK zero-sf defect not confined to code-0 blocks" - assert defect.any(), "code-0 case not exercised" - codes = (torch.log2(tk_sf.clamp(min=1e-45)) + 127).round().to(torch.int32) - assert torch.equal(oc[~defect], codes[~defect]), "rowwise scale codes differ" - - w2 = (torch.randn(256, 256, device=DEV, dtype=torch.float32) * 0.02 + 0.05).to(torch.bfloat16) - fp42 = per_token_cast(w2, "e2m1", 32, round_sf=True) - tk_bd, tk_bs = per_block_cast_lossless(fp42, "e4m3", (1, 32), (128, 128), round_sf=True) - obd, obs = tex.mxfp4_direct_blockwise(w2) - assert torch.equal(obd, tk_bd.view(torch.uint8)), "blockwise payload bytes differ" - assert torch.equal(obs[:, :2], tk_bs.float()), "blockwise tile scales differ" - print(" byte parity vs TileKernels chain (payloads everywhere, scales outside " - "their code-0 store defect): PASS") - - TESTS = [ test_rowwise_canonical_bytes_and_decode, test_cuda_kernels_parity, test_colwise_matches_bridge, test_blockwise_folding_and_decode, test_direct_flag_plumbing, - test_tilekernels_byte_parity, test_e2e_direct_matches_bridge, ] From 7808523401ea5315af91a966032604bcdfbc19f5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 07:42:34 +0000 Subject: [PATCH 12/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../references/mxfp4_qat_direct_reference.py | 16 +++---- tests/pytorch/test_mxfp4_qat_direct.py | 45 +++++++++++-------- .../cast/mxfp4/direct_convert_mxfp4.cuh | 25 +++++------ .../pytorch/csrc/extensions/cast.cpp | 4 +- .../pytorch/csrc/extensions/pybind.cpp | 8 ++-- transformer_engine/pytorch/module/base.py | 4 +- 6 files changed, 53 insertions(+), 49 deletions(-) diff --git a/tests/pytorch/references/mxfp4_qat_direct_reference.py b/tests/pytorch/references/mxfp4_qat_direct_reference.py index 8a72dced1d..638de55773 100644 --- a/tests/pytorch/references/mxfp4_qat_direct_reference.py +++ b/tests/pytorch/references/mxfp4_qat_direct_reference.py @@ -78,9 +78,9 @@ def _direct_mxfp8_rowwise(e, q, nonfinite, rows, cols): shift = torch.where(biased >= 0, torch.full_like(e, 6), e + 127) payload_vals = q * torch.ldexp(ones, shift) data = _encode_e4m3(payload_vals) - data = torch.where( - nonfinite.expand_as(data), torch.full_like(data, _E4M3_NAN), data - ).view(rows, cols) + data = torch.where(nonfinite.expand_as(data), torch.full_like(data, _E4M3_NAN), data).view( + rows, cols + ) code = torch.where(nonfinite, torch.full_like(code, 127), code) scale_inv = torch.zeros( @@ -160,17 +160,15 @@ def _direct_blockwise_128(e, q, nonfinite, rows, cols): fold = (e_t - s_exp).view(rows, cols // _MXFP4_BLOCK, 1).to(torch.float32) payload_vals = q * torch.ldexp(torch.ones_like(fold), fold.to(torch.int32)) data = _encode_e4m3(payload_vals) - data = torch.where( - nonfinite.expand_as(data), torch.full_like(data, _E4M3_NAN), data - ).view(rows, cols) + data = torch.where(nonfinite.expand_as(data), torch.full_like(data, _E4M3_NAN), data).view( + rows, cols + ) s_exp2 = s_exp.view(tiles_m, tiles_n) scale_inv = torch.zeros( (tiles_m, _roundup(tiles_n, 4)), dtype=torch.float32, device=data.device ) - scale_inv[:, :tiles_n] = torch.ldexp( - torch.ones_like(s_exp2, dtype=torch.float32), s_exp2 - ) + scale_inv[:, :tiles_n] = torch.ldexp(torch.ones_like(s_exp2, dtype=torch.float32), s_exp2) return data, scale_inv diff --git a/tests/pytorch/test_mxfp4_qat_direct.py b/tests/pytorch/test_mxfp4_qat_direct.py index 220100d223..63ac4d278c 100644 --- a/tests/pytorch/test_mxfp4_qat_direct.py +++ b/tests/pytorch/test_mxfp4_qat_direct.py @@ -145,8 +145,11 @@ def test_colwise_matches_bridge(): def test_blockwise_folding_and_decode(): def mk_quantizer(): return Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True, - force_pow_2_scales=True, block_scaling_dim=2, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + block_scaling_dim=2, ) # spread sweep: exact through d<=14, first loss at d=15 (E4M3 subnormal bound) @@ -213,12 +216,12 @@ def test_e2e_direct_matches_bridge(): f"fwd differs (bridge vs direct) {recipe_cls.__name__} " f"override={override} step={i}" ) - assert torch.equal(results[False][1], results[True][1]), ( - f"dgrad differs {recipe_cls.__name__} override={override}" - ) - assert torch.equal(results[False][2], results[True][2]), ( - f"wgrad differs {recipe_cls.__name__} override={override}" - ) + assert torch.equal( + results[False][1], results[True][1] + ), f"dgrad differs {recipe_cls.__name__} override={override}" + assert torch.equal( + results[False][2], results[True][2] + ), f"wgrad differs {recipe_cls.__name__} override={override}" print(" e2e direct == bridge (fwd/dgrad/wgrad bitwise, 2 recipes x 3 overrides): PASS") @@ -235,7 +238,6 @@ def test_direct_flag_plumbing(): print(" recipe flag plumbing: PASS") - def _edge_weight(): w = make_weight(128, 256) w[0, :32] = torch.tensor(2.0**-127, dtype=torch.bfloat16) @@ -259,8 +261,13 @@ def _intree_direct(fast_math=False): src_dir = os.environ.get( "NVTE_MXFP4_QAT_TEST_SRC", - os.path.join(os.path.dirname(os.path.abspath(__file__)), - "..", "..", "transformer_engine", "common"), + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "transformer_engine", + "common", + ), ) cuda_src = r""" #include @@ -325,7 +332,8 @@ def _intree_direct(fast_math=False): "-U__CUDA_NO_HALF2_OPERATORS__", "-U__CUDA_NO_BFLOAT16_OPERATORS__", "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", - ] + (["--use_fast_math"] if fast_math else []), + ] + + (["--use_fast_math"] if fast_math else []), verbose=False, ) return _DIRECT_KERNELS[fast_math] @@ -341,9 +349,9 @@ def test_cuda_kernels_parity(): w = make_weight(m, n, dtype=dtype, zero_blocks=2, outliers=2) dd, ds = k.direct_row(w.contiguous()) td, ts = _direct_mxfp8_rowwise(*_mxfp4_decompose(w), m, n) - assert torch.equal(dd, td) and torch.equal(ds, ts), ( - f"cuda row {tag} differs {m}x{n} {dtype}" - ) + assert torch.equal(dd, td) and torch.equal( + ds, ts + ), f"cuda row {tag} differs {m}x{n} {dtype}" wx = _edge_weight() dd, ds = k.direct_row(wx.contiguous()) td, ts = _direct_mxfp8_rowwise(*_mxfp4_decompose(wx), 128, 256) @@ -354,9 +362,9 @@ def test_cuda_kernels_parity(): w = make_weight(m, n, dtype=dtype, zero_blocks=3, outliers=3) dd, ds = k.direct_block(w.contiguous()) td, ts = _direct_blockwise_128(*_mxfp4_decompose(w), m, n) - assert torch.equal(dd, td) and torch.equal(ds, ts), ( - f"cuda blockwise {tag} differs {m}x{n} {dtype}" - ) + assert torch.equal(dd, td) and torch.equal( + ds, ts + ), f"cuda blockwise {tag} differs {m}x{n} {dtype}" payloads = torch.tensor([0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) for d in (0, 11, 14, 15): w = torch.zeros(128, 128, dtype=torch.float32, device=DEV) @@ -398,7 +406,6 @@ def timeit(fn, iters=20): print(" cuda kernels (normal + fast-math): PASS") - TESTS = [ test_rowwise_canonical_bytes_and_decode, test_cuda_kernels_parity, diff --git a/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh index 16f05ff1ce..c0bcfdc0af 100644 --- a/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh +++ b/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh @@ -41,13 +41,13 @@ namespace dispatch { namespace mxfp4 { namespace direct_convert_kernel { -using fake_quantize_kernel::E2M1_MAX; -using fake_quantize_kernel::MXFP4_BLOCK_SIZE; -using fake_quantize_kernel::THREADS_PER_CHUNK; using fake_quantize_kernel::block_scale_exponent_from_bits; +using fake_quantize_kernel::E2M1_MAX; using fake_quantize_kernel::exp2_from_int; using fake_quantize_kernel::mul_rn_no_ftz; +using fake_quantize_kernel::MXFP4_BLOCK_SIZE; using fake_quantize_kernel::round_e2m1_magnitude; +using fake_quantize_kernel::THREADS_PER_CHUNK; // 2^k as fp32 for k in [-149, 127]: normal via exponent field, subnormal via // mantissa bit, exact zero below the subnormal range. Integer construction. @@ -104,13 +104,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) int amax_bits = 0; #pragma unroll for (int i = 0; i < static_cast(VEC_ELTS); ++i) { - amax_bits = max(amax_bits, __float_as_int(static_cast(in_vec.data.elt[i])) & 0x7fffffff); + amax_bits = + max(amax_bits, __float_as_int(static_cast(in_vec.data.elt[i])) & 0x7fffffff); } const unsigned lane = threadIdx.x & 31u; - const unsigned group_mask = - (LANES_PER_BLOCK == 32) - ? 0xffffffffu - : (((1u << LANES_PER_BLOCK) - 1u) << (lane & ~(LANES_PER_BLOCK - 1u))); + const unsigned group_mask = (LANES_PER_BLOCK == 32) ? 0xffffffffu + : (((1u << LANES_PER_BLOCK) - 1u) + << (lane & ~(LANES_PER_BLOCK - 1u))); #pragma unroll for (int offset = LANES_PER_BLOCK / 2; offset > 0; offset >>= 1) { amax_bits = max(amax_bits, __shfl_xor_sync(group_mask, amax_bits, offset)); @@ -134,8 +134,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } else { #pragma unroll for (int i = 0; i < static_cast(VEC_ELTS); i += 2) { - store_e4m3_pair(&out_vec.data.elt[i], - direct_payload(in_vec.data.elt[i], inv_scale, pscale), + store_e4m3_pair(&out_vec.data.elt[i], direct_payload(in_vec.data.elt[i], inv_scale, pscale), direct_payload(in_vec.data.elt[i + 1], inv_scale, pscale)); } } @@ -159,8 +158,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) template __global__ void __launch_bounds__(512) direct_blockwise_kernel(const IType *__restrict__ input, uint8_t *__restrict__ data, - float *__restrict__ scales, const size_t cols, - const int scale_stride) { + float *__restrict__ scales, const size_t cols, const int scale_stride) { constexpr int BLOCKS_PER_TILE_ROW = 128 / MXFP4_BLOCK_SIZE; // 4 const int t = threadIdx.x; @@ -244,8 +242,7 @@ void direct_blockwise_launch(const IType *input, uint8_t *data, float *scales, c const size_t cols, const int scale_stride, cudaStream_t stream) { using namespace direct_convert_kernel; const dim3 grid(static_cast(cols / 128), static_cast(rows / 128)); - direct_blockwise_kernel<<>>(input, data, scales, cols, - scale_stride); + direct_blockwise_kernel<<>>(input, data, scales, cols, scale_stride); } } // namespace mxfp4 diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 2108035cdb..64a225a39b 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -591,8 +591,8 @@ std::tuple mxfp4_direct_mxfp8_rowwise(const at::Tensor & const int64_t rows = input.size(0), cols = input.size(1); const int64_t bpr = cols / 32; auto data = at::empty({rows, cols}, input.options().dtype(at::kByte)); - auto scales = at::zeros({(rows + 127) / 128 * 128, (bpr + 3) / 4 * 4}, - input.options().dtype(at::kByte)); + auto scales = + at::zeros({(rows + 127) / 128 * 128, (bpr + 3) / 4 * 4}, input.options().dtype(at::kByte)); auto in_cpp = makeTransformerEngineTensor(input_contiguous); auto data_cpp = makeTransformerEngineTensor(data); auto scales_cpp = makeTransformerEngineTensor(scales); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 30376ff151..2c751b4228 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -203,11 +203,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)", py::arg("input"), py::call_guard()); m.def("mxfp4_direct_mxfp8_rowwise", &transformer_engine::pytorch::mxfp4_direct_mxfp8_rowwise, - "Direct MXFP4->MXFP8 rowwise conversion (fixed-shift-6 canonicalization)", - py::arg("input"), py::call_guard()); + "Direct MXFP4->MXFP8 rowwise conversion (fixed-shift-6 canonicalization)", py::arg("input"), + py::call_guard()); m.def("mxfp4_direct_blockwise", &transformer_engine::pytorch::mxfp4_direct_blockwise, - "Direct MXFP4->FP8 128x128 blockwise conversion (exponent folding)", - py::arg("input"), py::call_guard()); + "Direct MXFP4->FP8 128x128 blockwise conversion (exponent folding)", py::arg("input"), + py::call_guard()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); m.def("create_empty_quantized_tensor", diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index ee47ab55b5..8c087919c1 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -807,7 +807,9 @@ def quantize_weight( if FP8GlobalStateManager.is_fp8_enabled() else False ) - _mxfp4_qat_direct = _mxfp4_qat_active and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat_direct_conversion() + _mxfp4_qat_direct = ( + _mxfp4_qat_active and FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat_direct_conversion() + ) if _mxfp4_qat_active and workspace_dtype == torch.float16: raise NotImplementedError( "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: " From d40e13cc8e2ec02fc918ba25e40fc8d8f3c7be42 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Mon, 27 Jul 2026 02:04:32 -0700 Subject: [PATCH 13/14] Columnwise via the projection and the existing MXFP8 columnwise quantizer The backward_override=None columnwise representation is the literal spec chain bf16 -> mxfp4(row) -> dequantized bf16 -> host MXFP8 columnwise quantize, and both steps already have production kernels (tex.mxfp4_fake_quantize and the MXFP8 columnwise encoder). Use them instead of the composite-torch replication of the host amax rule: the columnwise buffers are now the bridge path's byte for byte by construction, and the module carries no torch numerics. --- .../pytorch/mxfp4_qat_direct.py | 98 +++++-------------- 1 file changed, 27 insertions(+), 71 deletions(-) diff --git a/transformer_engine/pytorch/mxfp4_qat_direct.py b/transformer_engine/pytorch/mxfp4_qat_direct.py index cbd50724e9..5db64a52fe 100644 --- a/transformer_engine/pytorch/mxfp4_qat_direct.py +++ b/transformer_engine/pytorch/mxfp4_qat_direct.py @@ -10,11 +10,11 @@ * MXFP8 rowwise (tex.mxfp4_direct_mxfp8_rowwise): the deployment fixed-shift-6 canonicalization -- scale code = fp4_exp + 121 clamped at UE8M0 code 0 with exact payload absorption; the bytes are invertible back to packed FP4. -* MXFP8 columnwise (backward_override=None only): a fresh lossy 32x1 - quantization of the grid values with the HOST encoder's amax rule - (scale = 2^ceil(log2(amax/448))), so dgrad bytes match the bridge path bit - for bit. This is the sole implementation of the columnwise direction (a - composite-torch computation, not a fallback for a kernel). +* MXFP8 columnwise (backward_override=None only): the literal + projection-then-requantize chain on existing kernels -- + tex.mxfp4_fake_quantize produces the dequantized MXFP4-grid weight and the + host MXFP8 columnwise quantizer encodes it, so the columnwise buffers are + the bridge path's, byte for byte, by construction. * Float8 blockwise 128x128 (tex.mxfp4_direct_blockwise): per-tile exponent folding, RTNE onto E4M3 with subnormals -- exact through scale spread 2^14, bounded beyond, no device assert. @@ -28,14 +28,12 @@ import transformer_engine_torch as tex -from .mxfp4_qat import _MXFP4_BLOCK +from .mxfp4_qat import _MXFP4_BLOCK, mxfp4_fake_quantize from .tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor from .tensor.float8_blockwise_tensor import Float8BlockQuantizer, Float8BlockwiseQTensor __all__ = ["mxfp4_qat_direct_quantize", "mxfp4_qat_direct_update_"] -_E4M3_NAN = 0x7F - def _require_kernel(name: str): fn = getattr(tex, name, None) @@ -53,67 +51,19 @@ def _prepare(weight: torch.Tensor) -> torch.Tensor: return w -def _roundup(x: int, m: int) -> int: - return (x + m - 1) // m * m - - -def _round_to_e2m1_grid(y: torch.Tensor) -> torch.Tensor: - """RTNE onto the E2M1 magnitude grid {0, .5, 1, 1.5, 2, 3, 4, 6}; input in [0, 6].""" - fine = torch.round(y * 2.0) * 0.5 - mid = torch.round(y) - coarse = torch.round(y * 0.5) * 2.0 - return torch.where(y <= 2.0, fine, torch.where(y <= 4.0, mid, coarse)) - +def _mxfp8_columnwise_from_projection(weight: torch.Tensor): + """Columnwise (backward_override=None) buffers via the spec-literal chain + bf16 -> mxfp4(row) -> dequantized bf16 -> host MXFP8 columnwise quantize. -def _direct_mxfp8_columnwise(weight: torch.Tensor): - """Sole (composite-torch) implementation of the columnwise direction. - - MXFP4-decomposes the weight, then quantizes the grid values in 32x1 - blocks with the host encoder's amax rule so the bytes match the bridge - path's columnwise representation bit for bit. + Both steps run existing kernels (the QAT projection and the production + columnwise encoder), so the result is bit-identical to the bridge path's + columnwise representation by construction. """ - rows, cols = weight.shape - w32 = weight.contiguous().to(torch.float32).view(rows, cols // _MXFP4_BLOCK, _MXFP4_BLOCK) - amax = w32.abs().amax(dim=-1, keepdim=True) - nonfinite = ~torch.isfinite(amax) - - bits = amax.view(torch.int32) - exp_field = bits >> 23 - mantissa = bits & 0x7FFFFF - e = exp_field - 129 + (mantissa > 0x400000).to(torch.int32) - e = torch.where(exp_field > 0, e, torch.full_like(e, -126)) - e = e.clamp(min=-126, max=125) - e = torch.where(amax > 0, e, torch.full_like(e, -126)) - e = torch.where(nonfinite, torch.full_like(e, -126), e) - - scale = torch.ldexp(torch.ones_like(amax), e) - y = (w32 / scale).abs().clamp(max=6.0) - q = _round_to_e2m1_grid(torch.where(nonfinite, torch.zeros_like(y), y)) - q = torch.copysign(q, w32) - - vals = (q * scale).view(rows, cols) - nf_elem = nonfinite.expand(-1, -1, _MXFP4_BLOCK).reshape(rows, cols) - - vc = vals.view(rows // 32, 32, cols) - amax_c = vc.abs().amax(dim=1, keepdim=True) - amax_c = torch.where(torch.isfinite(amax_c), amax_c, torch.zeros_like(amax_c)) - cb = amax_c.view(torch.int32) - kf = (cb >> 23) - 127 - ec = kf - 8 + ((cb & 0x7FFFFF) > 0x600000).to(torch.int32) # host rule: ceil(log2(amax/448)) - ec = torch.where((cb >> 23) > 0, ec, torch.full_like(ec, -127)) - ec = ec.clamp(min=-127, max=127) - ec = torch.where(amax_c > 0, ec, torch.zeros_like(ec)) - inv = torch.ldexp(torch.ones_like(ec, dtype=torch.float32), -ec) - # + 0.0 canonicalizes -0 to +0 (the host colwise kernel drops the sign of zero) - payload = (vc * inv + 0.0).to(torch.float8_e4m3fn).view(torch.uint8).view(rows, cols) - payload = torch.where(nf_elem, torch.full_like(payload, _E4M3_NAN), payload) - - code = (ec + 127).clamp(0, 254).to(torch.uint8) - scale_inv = torch.zeros( - (_roundup(rows // 32, 4), _roundup(cols, 128)), dtype=torch.uint8, device=payload.device - ) - scale_inv[: rows // 32, :cols] = code.view(rows // 32, cols) - return payload, scale_inv + what = mxfp4_fake_quantize(weight) + col = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=False, columnwise=True + ).quantize(what) + return col._columnwise_data, col._columnwise_scale_inv def mxfp4_qat_direct_quantize(weight: torch.Tensor, quantizer) -> torch.Tensor: @@ -129,12 +79,16 @@ def mxfp4_qat_direct_quantize(weight: torch.Tensor, quantizer) -> torch.Tensor: raise ValueError(f"inner dim must be divisible by {_MXFP4_BLOCK}, got {cols}") if isinstance(quantizer, MXFP8Quantizer): - data, scale_inv = _require_kernel("mxfp4_direct_mxfp8_rowwise")(_prepare(weight)) + data, scale_inv = _require_kernel("mxfp4_direct_mxfp8_rowwise")( + _prepare(weight) + ) col_data, col_scale_inv = None, None if quantizer.columnwise_usage: if rows % 32 != 0: - raise ValueError(f"columnwise 32x1 blocks need rows % 32 == 0, got {rows}") - col_data, col_scale_inv = _direct_mxfp8_columnwise(weight) + raise ValueError( + f"columnwise 32x1 blocks need rows % 32 == 0, got {rows}" + ) + col_data, col_scale_inv = _mxfp8_columnwise_from_projection(weight) return MXFP8Tensor( shape=weight.shape, dtype=weight.dtype, @@ -150,7 +104,9 @@ def mxfp4_qat_direct_quantize(weight: torch.Tensor, quantizer) -> torch.Tensor: if isinstance(quantizer, Float8BlockQuantizer): if getattr(quantizer, "block_scaling_dim", 2) != 2: - raise ValueError("direct blockwise conversion supports 128x128 (2D) scaling only") + raise ValueError( + "direct blockwise conversion supports 128x128 (2D) scaling only" + ) data, scale_inv = _require_kernel("mxfp4_direct_blockwise")(_prepare(weight)) out = Float8BlockwiseQTensor( shape=weight.shape, From beb731c573c118853989704b9b96b8b6fa385937 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 09:09:44 +0000 Subject: [PATCH 14/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/mxfp4_qat_direct.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/mxfp4_qat_direct.py b/transformer_engine/pytorch/mxfp4_qat_direct.py index 5db64a52fe..0c11fd3c17 100644 --- a/transformer_engine/pytorch/mxfp4_qat_direct.py +++ b/transformer_engine/pytorch/mxfp4_qat_direct.py @@ -60,9 +60,9 @@ def _mxfp8_columnwise_from_projection(weight: torch.Tensor): columnwise representation by construction. """ what = mxfp4_fake_quantize(weight) - col = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=False, columnwise=True - ).quantize(what) + col = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=False, columnwise=True).quantize( + what + ) return col._columnwise_data, col._columnwise_scale_inv @@ -79,15 +79,11 @@ def mxfp4_qat_direct_quantize(weight: torch.Tensor, quantizer) -> torch.Tensor: raise ValueError(f"inner dim must be divisible by {_MXFP4_BLOCK}, got {cols}") if isinstance(quantizer, MXFP8Quantizer): - data, scale_inv = _require_kernel("mxfp4_direct_mxfp8_rowwise")( - _prepare(weight) - ) + data, scale_inv = _require_kernel("mxfp4_direct_mxfp8_rowwise")(_prepare(weight)) col_data, col_scale_inv = None, None if quantizer.columnwise_usage: if rows % 32 != 0: - raise ValueError( - f"columnwise 32x1 blocks need rows % 32 == 0, got {rows}" - ) + raise ValueError(f"columnwise 32x1 blocks need rows % 32 == 0, got {rows}") col_data, col_scale_inv = _mxfp8_columnwise_from_projection(weight) return MXFP8Tensor( shape=weight.shape, @@ -104,9 +100,7 @@ def mxfp4_qat_direct_quantize(weight: torch.Tensor, quantizer) -> torch.Tensor: if isinstance(quantizer, Float8BlockQuantizer): if getattr(quantizer, "block_scaling_dim", 2) != 2: - raise ValueError( - "direct blockwise conversion supports 128x128 (2D) scaling only" - ) + raise ValueError("direct blockwise conversion supports 128x128 (2D) scaling only") data, scale_inv = _require_kernel("mxfp4_direct_blockwise")(_prepare(weight)) out = Float8BlockwiseQTensor( shape=weight.shape,