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..638de55773 --- /dev/null +++ b/tests/pytorch/references/mxfp4_qat_direct_reference.py @@ -0,0 +1,184 @@ +# 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/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 new file mode 100644 index 0000000000..e304b49a52 --- /dev/null +++ b/tests/pytorch/test_mxfp4_qat.py @@ -0,0 +1,905 @@ +# 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.""" +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 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 + +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}" + + assert torch.equal( + dq, wh32 + ), f"MXFP8 rowwise encoding of the MXFP4-grid weight is not lossless {m}x{n}" + + +def test_dequantize_e8m0_extreme_codes(): + """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) + ) + 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) + 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(): + 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(): + """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) + w_hat = mxfp4_fake_quantize(w) + + q8 = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, columnwise=False).quantize(w_hat) + dq8 = q8.dequantize(dtype=torch.bfloat16) + 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( + 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 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 + + 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): + got_kernel = _intree_kernel().mxfp4_fake_quantize_intree(w.contiguous()) + got_torch = mxfp4_fake_quantize_reference(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 + rows, cols = 8192, 8192 + w = make_weight(rows, cols) + k = _intree_kernel() + for _ in range(3): + k.mxfp4_fake_quantize_intree(w) + mxfp4_fake_quantize_reference(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: mxfp4_fake_quantize_reference(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_kernel_fast_math_immune(): + """The kernel stays bit-identical when compiled with --use_fast_math.""" + 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 = mxfp4_fake_quantize_reference(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 = mxfp4_fake_quantize_reference(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(): + """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 + + 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\)" + 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" + "#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 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" + + +def test_ste_identity_gradient(): + """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})" + _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.""" + 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 = mxfp4_fake_quantize_reference(nc) + assert torch.equal(got, mxfp4_fake_quantize_reference(nc.contiguous())) + + +def test_bf16_exhaustive_bit_patterns(): + """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) + _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 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) + 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(): + """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 + 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_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]) + 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_fp16_pipeline_rejected(): + """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()): + 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_recipe_switch_invalidates_cache(): + """base<->QAT recipe switches must requantize instead of reusing cached weight workspaces.""" + 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_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 + + 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_dequantize_e8m0_extreme_codes, + 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_blockwise_feasibility_enumeration, + test_recipe_switch_invalidates_cache, + test_recipe_field_controls_qat, + test_ops_api_rejected, + test_fp16_pipeline_rejected, + test_kernel_perf, + test_kernel_fast_math_immune, + 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/tests/pytorch/test_mxfp4_qat_direct.py b/tests/pytorch/test_mxfp4_qat_direct.py new file mode 100644 index 0000000000..63ac4d278c --- /dev/null +++ b/tests/pytorch/test_mxfp4_qat_direct.py @@ -0,0 +1,433 @@ +# 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") + + +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_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 1e3c04573b..93db9a9759 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -16,6 +16,8 @@ #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" void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { @@ -47,6 +49,103 @@ 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_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..c0bcfdc0af --- /dev/null +++ b/transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh @@ -0,0 +1,252 @@ +/************************************************************************* + * 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::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. +__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/cast/mxfp4/fake_quantize_mxfp4.cuh b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh new file mode 100644 index 0000000000..6b4c437221 --- /dev/null +++ b/transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh @@ -0,0 +1,127 @@ +/************************************************************************* + * 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); + } +} + +} // namespace fake_quantize_kernel + +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); +} + +} // namespace mxfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 4d6d24ba65..5089fb35e6 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -418,6 +418,29 @@ 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, writing back dequantized values. + * + * \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. + */ +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 8c209ace15..cad298498a 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -168,6 +168,17 @@ def custom(cls): """Whether the given recipe is custom.""" return issubclass(cls, CustomRecipe) + 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) + ) + + 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): @@ -368,6 +379,8 @@ 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" + 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." @@ -437,6 +450,8 @@ 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" + 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" @@ -461,6 +476,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 ( @@ -481,6 +507,34 @@ def _make_repr(self) -> str: ) +@dataclass(repr=False) +class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling): + """ + MXFP8 recipe with MXFP4 weight quantization-aware training. + + 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. + """ + + mxfp4_qat_weights: bool = True + + +@dataclass(repr=False) +class MXFP4QATFloat8BlockScaling(Float8BlockScaling): + """ + Float8 block-scaling recipe with MXFP4 weight quantization-aware training. + + 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. + """ + + mxfp4_qat_weights: bool = True + + @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..5df2d51a9a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -347,6 +347,10 @@ 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, std::optional tensor_offsets, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..64a225a39b 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -556,6 +556,71 @@ 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; +} + +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 7e9d114be8..2c751b4228 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -199,6 +199,15 @@ 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("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 7c11e635c6..8c087919c1 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -46,6 +46,8 @@ _fsdp_gather_tensors, ) 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 @@ -800,8 +802,29 @@ def quantize_weight( ``_fp8_workspaces``. """ + _mxfp4_qat_active = ( + FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat() + 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: " + "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): + 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 +856,11 @@ 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_"): workspace.quantize_(tensor, noop_flag=skip_update_flag) else: @@ -842,6 +870,13 @@ 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: # Ensure the tensor in the cache is an instance of torch.Tensor, # as it persists beyond a single forward pass. @@ -1506,9 +1541,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..57ee4f854c --- /dev/null +++ b/transformer_engine/pytorch/mxfp4_qat.py @@ -0,0 +1,71 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP4 weight fake-quantization for quantization-aware training. + +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, 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 +(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 torch + +__all__ = ["mxfp4_fake_quantize"] + +_MXFP4_BLOCK = 32 + + +class _MXFP4FakeQuantizeSTE(torch.autograd.Function): + """Straight-through estimator: identity gradient through the projection.""" + + @staticmethod + def forward(ctx, weight): # pylint: disable=arguments-differ + import transformer_engine_torch as tex + + 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) + + @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. 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): + raise ValueError( + f"MXFP4 QAT supports bf16/fp32 weights, got {weight.dtype} " + "(fp16 cannot represent the full MXFP4 scale range)" + ) + if weight.shape[1] % _MXFP4_BLOCK != 0: + raise ValueError( + "MXFP4 QAT needs the weight inner dim divisible by " + f"{_MXFP4_BLOCK}, got {weight.shape[1]}" + ) + return _MXFP4FakeQuantizeSTE.apply(weight) diff --git a/transformer_engine/pytorch/mxfp4_qat_direct.py b/transformer_engine/pytorch/mxfp4_qat_direct.py new file mode 100644 index 0000000000..0c11fd3c17 --- /dev/null +++ b/transformer_engine/pytorch/mxfp4_qat_direct.py @@ -0,0 +1,149 @@ +# 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): 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. + +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, 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_"] + + +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 _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. + + 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. + """ + 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: + """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 = _mxfp8_columnwise_from_projection(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__}" + ) 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..01f245fdec 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -488,6 +488,16 @@ 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..0ed4f0fb81 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,14 @@ 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():