From abfd43c861687e6a28657185f5a952c81e17226e Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 01:39:58 -0700 Subject: [PATCH 1/2] Add NVFP4 four-over-six (row-wise) quantization to the training prototype Four-over-six is an adaptive NVFP4 block-scaling recipe: each 1x16 (or 16x16) block is encoded twice -- the standard map-to-6 encoding and a map-to-4 encoding whose E4M3 block scale is expanded by 1.5x -- and the candidate with the lower MAE/MSE dequantization error is stored. The global scale bound is reduced to 256 by default to leave headroom for the 1.5x expansion, and activations optionally take one FP32 global scale per row (row-wise) instead of per tensor. This adds: - four_over_six_quantize: pure-PyTorch quantizer, transcribed operation-for-operation from the reference CUDA kernels; codes and scales verified bitwise against them on GB200 (384/384 comparisons across 1x16/16x16, MAE/MSE, bound 256/448, per-tensor/row-scaled, rowwise/columnwise, bf16/fp32, 4 shapes). - four_over_six_mm / four_over_six_linear: training autograd function. Forward GEMM operands use four-over-six (activations 1x16, weights 16x16). Backward uses standard-NVFP4 RTNE gradients with the saved columnwise four-over-six operands, except in row-scaled mode where the backward runs in bf16 (a row-scaled four-over-six tensor has no columnwise form, so the quantized wgrad operand cannot be produced). - NVFP4FourOverSixLinear: drop-in nn.Linear. Co-Authored-By: Claude Fable 5 --- .../nvfp4_training/test_four_over_six.py | 247 ++++++++ .../nvfp4_training/four_over_six.py | 541 ++++++++++++++++++ 2 files changed, 788 insertions(+) create mode 100644 test/prototype/moe_training/nvfp4_training/test_four_over_six.py create mode 100644 torchao/prototype/moe_training/nvfp4_training/four_over_six.py diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py new file mode 100644 index 0000000000..5b50a3e64f --- /dev/null +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py @@ -0,0 +1,247 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + + +import pytest +import torch + +from torchao.float8.float8_utils import compute_error +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + NVFP4FourOverSixLinear, + four_over_six_global_encode_scale, + four_over_six_linear, + four_over_six_quantize, +) +from torchao.prototype.mx_formats.kernels import f4_unpacked_to_f32, unpack_uint4 +from torchao.utils import is_sm_at_least_100, torch_version_at_least + +_skip_no_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) +_skip_no_sm100 = pytest.mark.skipif( + not ( + torch.cuda.is_available() + and is_sm_at_least_100() + and torch_version_at_least("2.10.0") + ), + reason="requires SM100+ and PyTorch 2.10+ (FP4 scaled_mm)", +) + + +def _dequantize(codes, scales, global_amax, e4m3_scale_bound): + """Reconstruct FP32 values from packed codes, block scales, and global amax.""" + rows = codes.shape[0] + values = f4_unpacked_to_f32(unpack_uint4(codes)).view(rows, -1, 16) + s_dec = 1.0 / four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) + if s_dec.dim() == 1: + s_dec = s_dec.view(rows, 1, 1) + return (values * scales.to(torch.float32).unsqueeze(-1) * s_dec).view(rows, -1) + + +def _map6_reference(x, global_amax, e4m3_scale_bound): + """Standard (map-to-6 only) encoding with the four-over-six scale chain.""" + from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + _FP32_MAX, + FP4_E2M1_MAX, + FP8_E4M3_MAX, + _fp4_rtne, + ) + + rows, cols = x.shape + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) + fp4_max = torch.full((), FP4_E2M1_MAX, dtype=torch.float32, device=x.device) + base = (xf.abs().amax(dim=-1) / fp4_max) * s_enc + scale6 = base.clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + inv6 = (1.0 / (scale6.to(torch.float32) * (1.0 / s_enc))).clamp(max=_FP32_MAX) + _, values6 = _fp4_rtne(xf * inv6.unsqueeze(-1)) + s_dec = (1.0 / s_enc).view(-1, 1, 1) if s_enc.dim() == 1 else 1.0 / s_enc + dequant6 = values6 * scale6.to(torch.float32).unsqueeze(-1) * s_dec + return dequant6.view(rows, cols) + + +@_skip_no_cuda +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +def test_scales_are_candidate_scales(err_mode, e4m3_scale_bound, block): + """Every stored block scale is one of the two candidate scales.""" + torch.manual_seed(0) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + _, scales = four_over_six_quantize( + x, amax, block=block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + + xf = x.float().view(128, 16, 16) + if block == "16x16": + tiles = x.float().abs().view(8, 16, 16, 16) + block_amax = tiles.amax(dim=(1, 3)).repeat_interleave(16, dim=0) + else: + block_amax = xf.abs().amax(dim=-1) + s_enc = four_over_six_global_encode_scale(amax, e4m3_scale_bound) + fp4_max = torch.full((), 6.0, dtype=torch.float32, device="cuda") + base = (block_amax / fp4_max) * s_enc + scale6 = base.clamp(max=448.0).to(torch.float8_e4m3fn).view(torch.uint8) + scale4 = (base * 1.5).clamp(max=448.0).to(torch.float8_e4m3fn).view(torch.uint8) + got = scales.view(torch.uint8) + assert ((got == scale6) | (got == scale4)).all() + + +@_skip_no_cuda +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +def test_selection_not_worse_than_map6(e4m3_scale_bound): + """Per-block MAE of the stored encoding <= the map-to-6-only encoding.""" + torch.manual_seed(0) + x = torch.randn(256, 512, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize( + x, amax, block="1x16", err_mode="mae", e4m3_scale_bound=e4m3_scale_bound + ) + dq = _dequantize(codes, scales, amax, e4m3_scale_bound) + dq6 = _map6_reference(x, amax, e4m3_scale_bound) + xf = x.float() + err = (dq - xf).abs().view(256, -1, 16).sum(dim=-1).double() + err6 = (dq6 - xf).abs().view(256, -1, 16).sum(dim=-1).double() + # Selection minimizes the FP32 sequential-sum error; allow FP32-vs-FP64 + # summation slack on ties. + assert (err <= err6 + 1e-4).all() + # And the recipe must actually engage: some blocks pick map-to-4. + assert (err < err6 - 1e-4).any() + + +@_skip_no_cuda +def test_row_scaled_matches_per_row_quantization(): + """Row-scaled output == each row quantized alone with its own scalar amax.""" + torch.manual_seed(0) + x = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") + row_amax = x.abs().amax(dim=1).to(torch.float32) + codes, scales = four_over_six_quantize(x, row_amax, block="1x16") + for r in range(0, 64, 17): + codes_r, scales_r = four_over_six_quantize(x[r : r + 1], row_amax[r].view(())) + torch.testing.assert_close(codes[r : r + 1], codes_r, atol=0, rtol=0) + torch.testing.assert_close( + scales[r : r + 1].view(torch.uint8), + scales_r.view(torch.uint8), + atol=0, + rtol=0, + ) + + +@_skip_no_cuda +def test_row_scaled_rejects_16x16(): + x = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") + row_amax = x.abs().amax(dim=1).to(torch.float32) + with pytest.raises(ValueError, match="1x16 blocks only"): + four_over_six_quantize(x, row_amax, block="16x16") + + +@_skip_no_cuda +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +def test_dequant_sqnr(block): + torch.manual_seed(0) + x = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize(x, amax, block=block) + dq = _dequantize(codes, scales, amax, 256) + assert compute_error(x.float(), dq).item() > 14.0 + + +@_skip_no_sm100 +@pytest.mark.parametrize("row_scaled_activation", [False, True]) +@pytest.mark.parametrize("bias", [False, True]) +def test_linear_forward_backward(row_scaled_activation, bias): + torch.manual_seed(0) + M, K, N = 256, 512, 384 + x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = (torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1).requires_grad_( + True + ) + b = ( + torch.randn(N, dtype=torch.bfloat16, device="cuda", requires_grad=True) + if bias + else None + ) + y = four_over_six_linear(x, w, b, "mae", 256, row_scaled_activation) + assert y.shape == (M, N) + dy = torch.randn_like(y) + y.backward(dy) + + y_ref = x.detach().float() @ w.detach().float().t() + if bias: + y_ref = y_ref + b.detach().float() + dx_ref = dy.float() @ w.detach().float() + dw_ref = dy.float().t() @ x.detach().float() + assert compute_error(y_ref, y.float()).item() > 14.0 + assert compute_error(dx_ref, x.grad.float()).item() > 14.0 + assert compute_error(dw_ref, w.grad.float()).item() > 14.0 + if bias: + # grad_bias is reduced in bf16, matching nvfp4_linear. + torch.testing.assert_close(b.grad, dy.sum(dim=0)) + + +@_skip_no_sm100 +def test_linear_module(): + torch.manual_seed(0) + lin = NVFP4FourOverSixLinear(512, 384, device="cuda", dtype=torch.bfloat16) + x = torch.randn(128, 512, dtype=torch.bfloat16, device="cuda", requires_grad=True) + y = lin(x) + y.sum().backward() + assert y.shape == (128, 384) + assert lin.weight.grad is not None + + +@_skip_no_sm100 +def test_linear_rejects_unaligned_dims(): + x = torch.randn(100, 512, dtype=torch.bfloat16, device="cuda") + w = torch.randn(384, 512, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="divisible by 128"): + four_over_six_linear(x, w, None, "mae", 256, False) + + +@_skip_no_cuda +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +@pytest.mark.parametrize("row_scaled", [False, True]) +def test_bitwise_parity_with_transformer_engine( + err_mode, e4m3_scale_bound, block, row_scaled +): + """Bitwise codes/scales vs TransformerEngine's 4over6 kernels, if available.""" + te = pytest.importorskip("transformer_engine.pytorch") + if row_scaled and block == "16x16": + pytest.skip("row-scaled is 1x16 only") + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + if not te.is_nvfp4_available(): + pytest.skip("NVFP4 not available in this TransformerEngine build") + + torch.manual_seed(0) + M, N = 256, 512 + x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda") + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=(block == "16x16"), + row_scaled_nvfp4=row_scaled, + nvfp4_use_4over6=True, + nvfp4_e4m3_max=e4m3_scale_bound, + nvfp4_4over6_err_mode=err_mode.upper(), + ) + t = quantizer(x) + if row_scaled: + amax = x.abs().amax(dim=1).to(torch.float32) + else: + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize( + x, amax, block=block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + te_codes = t._rowwise_data.view(torch.uint8)[:, : N // 2] + te_scales = t._rowwise_scale_inv[:M, : N // 16].view(torch.uint8) + torch.testing.assert_close(te_codes, codes, atol=0, rtol=0) + torch.testing.assert_close(te_scales, scales.view(torch.uint8), atol=0, rtol=0) diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py new file mode 100644 index 0000000000..056d3aa81b --- /dev/null +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -0,0 +1,541 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""NVFP4 four-over-six quantization and the linear layer that consumes it. + +Four-over-six is an adaptive NVFP4 block-scaling recipe: every quantization +block is encoded twice and the candidate with the lower dequantization error +is stored. + +* The **map-to-6** candidate is the standard NVFP4 encoding: the E4M3 block + scale maps the block amax to FP4 value 6. +* The **map-to-4** candidate expands the E4M3 block scale by 1.5x, so FP4 + value 4 reaches the range that value 6 reaches in the standard encoding. + The FP4 grid is denser around 4 than around 6, which lowers error for + blocks whose mass sits below the amax. + +Errors are compared per block with a configurable metric (mean-absolute or +mean-squared, computed in the input domain); ties select map-to-6. To leave +E4M3 headroom for the 1.5x scale expansion, the global (per-tensor) scale is +derived from a reduced E4M3 bound of 256 by default instead of 448. + +Two global-scale granularities are supported for activations: + +* per-tensor: one FP32 scale for the whole tensor (the default), and +* row-wise: one FP32 scale per tensor row, derived from that row's amax. + +The arithmetic is transcribed operation for operation from TransformerEngine's +``quantize_4over6_nvfp4.cuh`` so codes and scales are bitwise identical to +that recipe (see ``compute_scale_pair`` and ``accumulate_dequant_error`` +there). Two details are load-bearing: + +* The block-scale association is ``(block_amax / 6) * S_enc`` — one division + then one multiply. The standard NVFP4 path uses + ``block_amax * (S_enc * (1/6))``, which rounds differently on a fraction of + blocks. +* The per-block error is accumulated sequentially in element order with FP32 + round-to-nearest adds, and 16x16 tiles reduce their 16 row-group errors in + a pairwise halving tree. Both orders affect candidate selection on ties + near the FP32 rounding boundary. + +``four_over_six_linear`` mirrors the recipe's training semantics: + +* forward GEMM: activations quantized 1x16 four-over-six (optionally + row-scaled), weights quantized 16x16 four-over-six; +* backward with per-tensor activations: gradients use standard NVFP4 + round-to-nearest-even (four-over-six never applies to gradients), and the + saved columnwise activation/weight codes are four-over-six; +* backward with row-scaled activations: high-precision (bf16) GEMMs. A + row-scaled four-over-six tensor has no columnwise form — the per-row scales + do not transpose — so the quantized wgrad operand cannot be produced. +""" + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from torchao.prototype.mx_formats.kernels import ( + f4_unpacked_to_f32, + f32_to_f4_unpacked, + pack_uint4, +) +from torchao.prototype.mx_formats.utils import to_blocked + +FP4_E2M1_MAX = 6.0 +FP8_E4M3_MAX = 448.0 +_FP32_MAX = torch.finfo(torch.float32).max + +__all__ = [ + "four_over_six_global_encode_scale", + "four_over_six_quantize", + "four_over_six_mm", + "four_over_six_linear", + "NVFP4FourOverSixLinear", +] + + +def four_over_six_global_encode_scale( + global_amax: torch.Tensor, e4m3_scale_bound: int = 256 +) -> torch.Tensor: + """``compute_global_encode_scaling_factor_FP4``: bound * 6 / amax. + + ``global_amax`` may be a scalar (per-tensor) or a 1-D per-row vector. + ``amax == 0`` gives inf and an enormous amax underflows the scale to + zero; both fall back to the identity scale. + """ + amax = global_amax.to(torch.float32) + candidate = torch.full_like(amax, float(e4m3_scale_bound) * FP4_E2M1_MAX) / amax + candidate = candidate.clamp(max=_FP32_MAX) + return torch.where( + (amax == 0.0) | (candidate == 0.0), torch.ones_like(candidate), candidate + ) + + +def _fp4_rtne(scaled: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Round-to-nearest-even FP4 codes and their exact FP32 values. + + Reproduces ``cvt.rn.satfinite.e2m1x2.f32`` followed by + ``cvt.rn.f16x2.e2m1x2`` (E2M1 values are exact in FP16 and FP32). + """ + clamped = scaled.clamp(-FP4_E2M1_MAX, FP4_E2M1_MAX) + codes = f32_to_f4_unpacked(clamped) + return codes, f4_unpacked_to_f32(codes) + + +def _candidate_error( + values: torch.Tensor, + scale_fp8: torch.Tensor, + xf: torch.Tensor, + global_amax: torch.Tensor, + err_mode: str, + e4m3_scale_bound: int, +) -> torch.Tensor: + """``accumulate_dequant_error``: FP32 adds in element order, per 1x16 group. + + values/xf: (rows, num_groups, 16); scale_fp8: (rows, num_groups, 1); + global_amax broadcastable against (rows, num_groups). + """ + sf = scale_fp8.to(torch.float32)[..., 0] + # The denominator must be a tensor: dividing by a python scalar lowers to a + # multiply by its (inexact) reciprocal, which double-rounds and flips + # candidate picks near error ties. Tensor-tensor division is a true + # correctly-rounded FP32 division, matching the kernel's __fdiv_rn. + err_denom = torch.full( + (), + FP4_E2M1_MAX * float(e4m3_scale_bound), + dtype=torch.float32, + device=xf.device, + ) + err = torch.zeros_like(xf[..., 0]) + for idx in range(16): + val = ((values[..., idx] * sf) * global_amax) / err_denom + diff = val - xf[..., idx] + if err_mode == "mse": + err = err + diff * diff + else: + err = err + diff.abs() + return err + + +def _tile_error_tree_sum(err: torch.Tensor) -> torch.Tensor: + """Reduce 16 row-group errors per 16x16 tile in the warp-shuffle tree order. + + err: (rows, num_groups) with rows % 16 == 0 -> (rows // 16, num_groups). + """ + rows = err.view(err.shape[0] // 16, 16, err.shape[1]) + rows = rows[:, 0:8] + rows[:, 8:16] + rows = rows[:, 0:4] + rows[:, 4:8] + rows = rows[:, 0:2] + rows[:, 2:4] + return rows[:, 0] + rows[:, 1] + + +def four_over_six_quantize( + x: torch.Tensor, + global_amax: torch.Tensor, + *, + block: str = "1x16", + err_mode: str = "mae", + e4m3_scale_bound: int = 256, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2-D tensor to NVFP4 with four-over-six block selection. + + Args: + x: (R, C) bfloat16 or float32, C % 16 == 0 (R % 16 == 0 for 16x16). + global_amax: scalar FP32 amax, or a (R,) per-row amax vector for the + row-scaled variant (1x16 blocks only). + block: "1x16" (activations/gradient operands) or "16x16" (weights). + err_mode: "mae" or "mse" candidate-selection error metric. + e4m3_scale_bound: 256 (default, leaves map-to-4 headroom) or 448. + + Returns: + (codes, scales): (R, C//2) uint8 packed FP4 codes (low nibble = even + element) and (R, C//16) float8_e4m3fn block scales. + """ + if block not in ("1x16", "16x16"): + raise ValueError(f"block must be '1x16' or '16x16', got {block!r}") + if err_mode not in ("mae", "mse"): + raise ValueError(f"err_mode must be 'mae' or 'mse', got {err_mode!r}") + if e4m3_scale_bound not in (256, 448): + raise ValueError(f"e4m3_scale_bound must be 256 or 448, got {e4m3_scale_bound}") + rows, cols = x.shape + if cols % 16: + raise ValueError(f"columns must be divisible by 16, got {cols}") + if block == "16x16" and rows % 16: + raise ValueError(f"16x16 blocks need rows divisible by 16, got {rows}") + row_scaled = global_amax.dim() == 1 and global_amax.numel() == rows + if row_scaled and block != "1x16": + raise ValueError("row-scaled four-over-six supports 1x16 blocks only") + if not row_scaled and global_amax.numel() != 1: + raise ValueError( + f"global_amax must be a scalar or a ({rows},) row vector, " + f"got shape {tuple(global_amax.shape)}" + ) + + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) + if row_scaled: + s_enc = s_enc.view(rows, 1) + err_amax = global_amax.to(torch.float32).view(rows, 1) + else: + err_amax = global_amax.to(torch.float32) + + if block == "16x16": + tiles = xf.abs().view(rows // 16, 16, cols // 16, 16) + block_amax = tiles.amax(dim=(1, 3)).repeat_interleave(16, dim=0) + else: + block_amax = xf.abs().amax(dim=-1) + + # compute_scale_pair: base = (block_amax / 6) * S_enc, then the 1.5x + # map-to-4 expansion; both capped at the full E4M3 range. The divisor is a + # tensor for a true correctly-rounded FP32 division (a python-scalar + # divisor lowers to a reciprocal multiply, which double-rounds). + fp4_max = torch.full((), FP4_E2M1_MAX, dtype=torch.float32, device=xf.device) + base = (block_amax / fp4_max) * s_enc + scale6 = base.clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + scale4 = (base * 1.5).clamp(max=FP8_E4M3_MAX).to(torch.float8_e4m3fn) + s_dec = 1.0 / s_enc + inv6 = (1.0 / (scale6.to(torch.float32) * s_dec)).clamp(max=_FP32_MAX) + inv4 = (1.0 / (scale4.to(torch.float32) * s_dec)).clamp(max=_FP32_MAX) + + codes6, values6 = _fp4_rtne(xf * inv6.unsqueeze(-1)) + codes4, values4 = _fp4_rtne(xf * inv4.unsqueeze(-1)) + err6 = _candidate_error( + values6, scale6.unsqueeze(-1), xf, err_amax, err_mode, e4m3_scale_bound + ) + err4 = _candidate_error( + values4, scale4.unsqueeze(-1), xf, err_amax, err_mode, e4m3_scale_bound + ) + if block == "16x16": + pick4 = ( + _tile_error_tree_sum(err4) < _tile_error_tree_sum(err6) + ).repeat_interleave(16, dim=0) + else: + pick4 = err4 < err6 + + codes = torch.where(pick4.unsqueeze(-1), codes4, codes6) + scales = torch.where(pick4, scale4, scale6) + return pack_uint4(codes.view(rows, cols)), scales + + +def _standard_rtne_quantize( + x: torch.Tensor, global_amax: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Standard NVFP4 1x16 round-to-nearest-even quantize for gradient operands. + + The gradient scale chain keeps the standard association + ``block_amax * (S_enc * (1/6))`` and the full 448 E4M3 bound; only + non-gradient four-over-six tensors use the ``(block_amax / 6) * S_enc`` + association above. + """ + rows, cols = x.shape + xf = x.float().view(rows, cols // 16, 16) + s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound=448) + block_amax = xf.abs().amax(dim=-1) + scales = ( + (block_amax * (s_enc * (1.0 / FP4_E2M1_MAX))) + .clamp(max=FP8_E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + enc = (1.0 / (scales.to(torch.float32) * (1.0 / s_enc))).clamp(max=_FP32_MAX) + codes, _ = _fp4_rtne(xf * enc.unsqueeze(-1)) + return pack_uint4(codes.view(rows, cols)), scales + + +def _global_decode_scale(amax: torch.Tensor, e4m3_scale_bound: int) -> torch.Tensor: + """Per-tensor decode scale consumed by the GEMM: amax / (bound * 6).""" + return amax.to(torch.float32) / (float(e4m3_scale_bound) * FP4_E2M1_MAX) + + +def _scaled_mm_nvfp4( + a_codes: torch.Tensor, + a_scales: torch.Tensor, + a_global: torch.Tensor, + b_codes_t: torch.Tensor, + b_scales: torch.Tensor, + b_global: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Block-scaled FP4 GEMM with per-tensor second-level scales. + + a_codes: (M, K//2) uint8; b_codes_t: (K//2, N) transposed uint8 view; + a_scales/b_scales: plain (rows, K//16) float8 block scales (swizzled here). + """ + return F.scaled_mm( + a_codes.view(torch.float4_e2m1fn_x2), + b_codes_t.view(torch.float4_e2m1fn_x2), + scale_a=[to_blocked(a_scales).flatten(), a_global], + scale_recipe_a=[F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise], + scale_b=[to_blocked(b_scales).flatten(), b_global], + scale_recipe_b=[F.ScalingType.BlockWise1x16, F.ScalingType.TensorWise], + swizzle_a=[F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE], + swizzle_b=[F.SwizzleType.SWIZZLE_32_4_4, F.SwizzleType.NO_SWIZZLE], + output_dtype=out_dtype, + ) + + +@torch._dynamo.allow_in_graph +class four_over_six_mm(torch.autograd.Function): + """NVFP4 four-over-six quantized matmul. + + 3 GEMMs: + forward: x_row @ W.T = output (1x16 four-over-six x, 16x16 four-over-six W) + backward: dy_row @ W.T = grad_input (standard-NVFP4 dy; saved columnwise W) + backward: dy_col.T @ x_col = grad_weight (standard-NVFP4 dy; saved columnwise x) + + With row-scaled activations the backward runs in bf16 instead (see the + module docstring), saving the high-precision operands. + + Requires: M % 128 == 0, K % 128 == 0, N % 128 == 0. + """ + + @staticmethod + def forward( + ctx, + input_hp: torch.Tensor, + weight_hp: torch.Tensor, + bias: Optional[torch.Tensor], + err_mode: str = "mae", + e4m3_scale_bound: int = 256, + row_scaled_activation: bool = False, + ): + M = input_hp.shape[:-1].numel() + K = input_hp.shape[-1] + N = weight_hp.shape[0] + if input_hp.dtype != torch.bfloat16: + input_hp = input_hp.to(torch.bfloat16) + if weight_hp.dtype != torch.bfloat16: + weight_hp = weight_hp.to(torch.bfloat16) + if M % 128 != 0 or K % 128 != 0 or N % 128 != 0: + raise ValueError( + f"four_over_six_mm requires M, K, N all divisible by 128; " + f"got M={M}, K={K}, N={N}" + ) + input_2d = input_hp.reshape(-1, K).contiguous() + + if row_scaled_activation: + x_amax = input_2d.abs().amax(dim=1).to(torch.float32) + else: + x_amax = input_2d.abs().amax().to(torch.float32) + w_amax = weight_hp.abs().amax().to(torch.float32) + + x_codes, x_scales = four_over_six_quantize( + input_2d, + x_amax, + block="1x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_codes, w_scales = four_over_six_quantize( + weight_hp, + w_amax, + block="16x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_global = _global_decode_scale(w_amax, e4m3_scale_bound) + + if row_scaled_activation: + # The GEMM's per-tensor slot cannot hold a per-row scale: run it + # with the constant 1/(6*bound) factor, then apply the raw per-row + # amaxes on the FP32 output before the bf16 cast. + x_global = torch.full( + (), + 1.0 / (FP4_E2M1_MAX * float(e4m3_scale_bound)), + dtype=torch.float32, + device=input_2d.device, + ) + output = _scaled_mm_nvfp4( + x_codes, + x_scales, + x_global, + w_codes.t(), + w_scales, + w_global, + torch.float32, + ) + output = (output * x_amax.view(-1, 1)).to(torch.bfloat16) + else: + x_global = _global_decode_scale(x_amax, e4m3_scale_bound) + output = _scaled_mm_nvfp4( + x_codes, + x_scales, + x_global, + w_codes.t(), + w_scales, + w_global, + torch.bfloat16, + ) + output = output.reshape(*input_hp.shape[:-1], N) + if bias is not None: + output = output + bias + + if row_scaled_activation: + ctx.save_for_backward(input_2d, weight_hp) + else: + x_col_codes, x_col_scales = four_over_six_quantize( + input_2d.t().contiguous(), + x_amax, + block="1x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + w_col_codes, w_col_scales = four_over_six_quantize( + weight_hp.t().contiguous(), + w_amax, + block="16x16", + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + ) + ctx.save_for_backward( + x_col_codes, + x_col_scales, + x_amax, + w_col_codes, + w_col_scales, + w_amax, + ) + ctx.row_scaled_activation = row_scaled_activation + ctx.e4m3_scale_bound = e4m3_scale_bound + ctx.input_orig_shape = input_hp.shape + ctx.has_bias = bias is not None + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + grad_output = grad_output.contiguous() + grad_output_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + if ctx.row_scaled_activation: + input_2d, weight_hp = ctx.saved_tensors + grad_input = (grad_output_2d @ weight_hp).reshape(ctx.input_orig_shape) + grad_weight = grad_output_2d.t() @ input_2d + else: + ( + x_col_codes, + x_col_scales, + x_amax, + w_col_codes, + w_col_scales, + w_amax, + ) = ctx.saved_tensors + dy_amax = grad_output_2d.abs().amax().to(torch.float32) + dy_row_codes, dy_row_scales = _standard_rtne_quantize( + grad_output_2d, dy_amax + ) + dy_col_codes, dy_col_scales = _standard_rtne_quantize( + grad_output_2d.t().contiguous(), dy_amax + ) + dy_global = _global_decode_scale(dy_amax, 448) + grad_input = _scaled_mm_nvfp4( + dy_row_codes, + dy_row_scales, + dy_global, + w_col_codes.t(), + w_col_scales, + _global_decode_scale(w_amax, ctx.e4m3_scale_bound), + torch.bfloat16, + ).reshape(ctx.input_orig_shape) + grad_weight = _scaled_mm_nvfp4( + dy_col_codes, + dy_col_scales, + dy_global, + x_col_codes.t(), + x_col_scales, + _global_decode_scale(x_amax, ctx.e4m3_scale_bound), + torch.bfloat16, + ) + + grad_bias = ( + grad_output.sum(dim=tuple(range(grad_output.dim() - 1))) + if ctx.has_bias + else None + ) + return grad_input, grad_weight, grad_bias, None, None, None + + +four_over_six_linear = four_over_six_mm.apply + + +class NVFP4FourOverSixLinear(nn.Linear): + """Linear layer with NVFP4 four-over-six quantized GEMMs. + + Drop-in replacement for nn.Linear implementing the four-over-six recipe: + forward GEMM operands use four-over-six NVFP4, gradients use standard + NVFP4 (or bf16 when ``row_scaled_activation`` is set — see the module + docstring for why row-scaled has no quantized backward). + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + err_mode: str = "mae", + e4m3_scale_bound: int = 256, + row_scaled_activation: bool = False, + device=None, + dtype=None, + ): + super().__init__(in_features, out_features, bias, device=device, dtype=dtype) + self.err_mode = err_mode + self.e4m3_scale_bound = e4m3_scale_bound + self.row_scaled_activation = row_scaled_activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return four_over_six_linear( + x, + self.weight, + self.bias, + self.err_mode, + self.e4m3_scale_bound, + self.row_scaled_activation, + ) + + @classmethod + def from_linear( + cls, + mod: nn.Linear, + err_mode: str = "mae", + e4m3_scale_bound: int = 256, + row_scaled_activation: bool = False, + ) -> "NVFP4FourOverSixLinear": + new = cls( + mod.in_features, + mod.out_features, + mod.bias is not None, + err_mode=err_mode, + e4m3_scale_bound=e4m3_scale_bound, + row_scaled_activation=row_scaled_activation, + device=mod.weight.device, + dtype=mod.weight.dtype, + ) + if mod.weight.device != torch.device("meta"): + new.weight = mod.weight + if mod.bias is not None: + new.bias = mod.bias + return new From 6f8c6dbe957f6053de38d39ab6642537a3a6645a Mon Sep 17 00:00:00 2001 From: Hanlin Bi Date: Sun, 23 Aug 2026 03:16:47 -0700 Subject: [PATCH 2/2] Add a CuTe DSL fast path for NVFP4 four-over-six quantization One SM100 kernel behind torchao::four_over_six_quantize_cutedsl, dispatched from four_over_six_quantize when the input is eligible (CUDA bf16/fp32, contiguous, C % 64 == 0); everything else silently falls through to the pure-PyTorch body, which stays intact as the fallback and bitwise oracle. The kernel is an op-for-op reimplementation of the reference arithmetic with every rounding pinned by inline PTX: real div.rn.f32 divisions everywhere (the (block_amax / 6) * S_enc association, S_enc/S_dec, the encode reciprocals, and the error denominator - never reciprocal-multiply), cvt.rn.satfinite.e4m3x2.f32 scale casts with exact e4m3->f16->f32 decode, cvt.rn.satfinite.e2m1x2.f32 FP4 casts with exact cvt.rn.f16x2.e2m1x2 error dequant, NaN-dropping max.f32 amaxes, strictly sequential per-group error accumulation in element order, and the exact width-16 shuffle-down error tree for 16x16 tiles. One (128, 64) tile per CTA, one tile row per thread (16-lane segments = 16 consecutive rows for the 2D mode), TMA G2S in and TMA S2G out with row clipping, u32-vectorized scale stores. Verified bitwise against the pure-torch reference: 1920/1920 over {mae,mse} x {256,448} x {1x16,16x16} x {per-tensor,row-scaled} x {bf16,fp32} x 5 shapes x 16 seeds, plus zeros / Inf / subnormal / near-448 / wide-dynamic / amax==0 special values, and 12/12 bitwise against TransformerEngine's quantize_4over6_nvfp4 kernels. NaN inputs follow TE semantics (NaN-dropping amax, NaN -> +6 codes), which the torch body cannot reproduce since torch.amax propagates NaN; documented in the op docstring and pinned by a dedicated test. 26-30x over the pure-torch body on DSV3-671B shapes (1200 MHz SM): 6.2 ms vs 174 ms at 131072x7168. Co-Authored-By: Claude Fable 5 --- .../bench_cutedsl_four_over_six_quantize.py | 160 +++++ .../nvfp4_training/test_four_over_six.py | 154 ++++ .../nvfp4_training/four_over_six.py | 13 + .../nvfp4_training/four_over_six_cutedsl.py | 676 ++++++++++++++++++ 4 files changed, 1003 insertions(+) create mode 100644 benchmarks/prototype/moe_training/nvfp4_training/bench_cutedsl_four_over_six_quantize.py create mode 100644 torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py diff --git a/benchmarks/prototype/moe_training/nvfp4_training/bench_cutedsl_four_over_six_quantize.py b/benchmarks/prototype/moe_training/nvfp4_training/bench_cutedsl_four_over_six_quantize.py new file mode 100644 index 0000000000..2e05bdaf86 --- /dev/null +++ b/benchmarks/prototype/moe_training/nvfp4_training/bench_cutedsl_four_over_six_quantize.py @@ -0,0 +1,160 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. +# this benchmarking script is a modified version of the original script from: https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/utils/benchmark.py + +import itertools +from dataclasses import dataclass +from typing import List + +import torch +from tabulate import tabulate +from tqdm import tqdm + +import torchao.prototype.moe_training.nvfp4_training.four_over_six as four_over_six_module +from benchmarks.utils import benchmark_cuda_function_in_microseconds +from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( + four_over_six_quantize, +) + +device = torch.device("cuda") + + +@dataclass(frozen=True) +class ExperimentConfig: + input_shape: tuple[int, int] + block: str + row_scaled: bool + + +@dataclass(frozen=True) +class ExperimentResult: + # time + cutedsl_us: float + torch_ref_us: float + # mem bw + cutedsl_gbps: float + torch_ref_gbps: float + + +@dataclass(frozen=True) +class Experiment: + config: ExperimentConfig + result: ExperimentResult + + +def get_configs() -> List[ExperimentConfig]: + input_shapes = [ + # DeepSeekV3 671b shapes + (8192, 2048), + (8192, 7168), + (32768, 2048), + (32768, 7168), + (131072, 2048), + (131072, 7168), + ] + cases = [ + ("1x16", False), # activations, per-tensor amax + ("1x16", True), # activations, row-scaled amax + ("16x16", False), # weights + ] + configs = [] + for shape, (block, row_scaled) in itertools.product(input_shapes, cases): + configs.append( + ExperimentConfig(input_shape=shape, block=block, row_scaled=row_scaled) + ) + return configs + + +def _reference_quantize(x, global_amax, **kwargs): + """Pure-PyTorch four_over_six_quantize body (dispatch gate disabled).""" + orig = four_over_six_module._cutedsl_quantize_eligible + four_over_six_module._cutedsl_quantize_eligible = lambda t: False + try: + return four_over_six_quantize(x, global_amax, **kwargs) + finally: + four_over_six_module._cutedsl_quantize_eligible = orig + + +def run_experiment(config: ExperimentConfig) -> ExperimentResult: + x = torch.randn(*config.input_shape, dtype=torch.bfloat16, device=device) + if config.row_scaled: + amax = x.abs().amax(dim=1).to(torch.float32) + else: + amax = x.abs().amax().to(torch.float32) + + quantize_kwargs = dict(block=config.block, err_mode="mae", e4m3_scale_bound=256) + + # Correctness first: the CuTe DSL fast path must be bitwise identical. + assert four_over_six_module._cutedsl_quantize_eligible(x) + codes, scales = four_over_six_quantize(x, amax, **quantize_kwargs) + ref_codes, ref_scales = _reference_quantize(x, amax, **quantize_kwargs) + assert torch.equal(codes, ref_codes) + assert torch.equal(scales.view(torch.uint8), ref_scales.view(torch.uint8)) + + cutedsl_us = benchmark_cuda_function_in_microseconds( + four_over_six_quantize, x, amax, **quantize_kwargs + ) + torch_ref_us = benchmark_cuda_function_in_microseconds( + _reference_quantize, x, amax, **quantize_kwargs + ) + + bytes_per_input_el = torch.finfo(torch.bfloat16).bits / 8 + read_bytes = x.numel() * bytes_per_input_el + write_bytes = codes.numel() * 1 + scales.numel() * 1 + cutedsl_gbps = ((read_bytes + write_bytes) / 1e9) / (cutedsl_us / 1e6) + torch_ref_gbps = ((read_bytes + write_bytes) / 1e9) / (torch_ref_us / 1e6) + + return ExperimentResult( + cutedsl_us=cutedsl_us, + torch_ref_us=torch_ref_us, + cutedsl_gbps=cutedsl_gbps, + torch_ref_gbps=torch_ref_gbps, + ) + + +def print_results(experiments: List[Experiment]): + headers = [ + "input_shape", + "block", + "row_scaled", + "cutedsl_us", + "torch_ref_us", + "speedup", + "cutedsl_gbps", + "torch_ref_gbps", + ] + rows = [] + for experiment in experiments: + speedup = experiment.result.torch_ref_us / experiment.result.cutedsl_us + rows.append( + [ + str(experiment.config.input_shape), + experiment.config.block, + experiment.config.row_scaled, + f"{experiment.result.cutedsl_us:.2f}", + f"{experiment.result.torch_ref_us:.2f}", + f"{speedup:.2f}x", + f"{experiment.result.cutedsl_gbps:.1f}", + f"{experiment.result.torch_ref_gbps:.1f}", + ] + ) + print(tabulate(rows, headers=headers)) + + +def main(): + torch.random.manual_seed(123) + configs = get_configs() + results = [] + for config in tqdm(configs): + result = run_experiment(config) + results.append(Experiment(config=config, result=result)) + + # Use Tabulate to print results + print_results(results) + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py index 5b50a3e64f..96f6272f78 100644 --- a/test/prototype/moe_training/nvfp4_training/test_four_over_six.py +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py @@ -8,6 +8,7 @@ import pytest import torch +import torchao.prototype.moe_training.nvfp4_training.four_over_six as four_over_six_module from torchao.float8.float8_utils import compute_error from torchao.prototype.moe_training.nvfp4_training.four_over_six import ( NVFP4FourOverSixLinear, @@ -15,6 +16,9 @@ four_over_six_linear, four_over_six_quantize, ) +from torchao.prototype.moe_training.nvfp4_training.four_over_six_cutedsl import ( + _cutedsl_quantize_available, +) from torchao.prototype.mx_formats.kernels import f4_unpacked_to_f32, unpack_uint4 from torchao.utils import is_sm_at_least_100, torch_version_at_least @@ -29,6 +33,22 @@ ), reason="requires SM100+ and PyTorch 2.10+ (FP4 scaled_mm)", ) +_cutedsl_available = torch.cuda.is_available() and _cutedsl_quantize_available() +_skip_no_cutedsl = pytest.mark.skipif( + not _cutedsl_available, + reason="requires SM100+ and the CuTe DSL runtime packages", +) + + +def _reference_quantize(x, global_amax, **kwargs): + """Run the pure-PyTorch four_over_six_quantize body (the bitwise oracle) + by disabling the CuTe DSL dispatch gate for the duration of the call.""" + orig = four_over_six_module._cutedsl_quantize_eligible + four_over_six_module._cutedsl_quantize_eligible = lambda t: False + try: + return four_over_six_quantize(x, global_amax, **kwargs) + finally: + four_over_six_module._cutedsl_quantize_eligible = orig def _dequantize(codes, scales, global_amax, e4m3_scale_bound): @@ -245,3 +265,137 @@ def test_bitwise_parity_with_transformer_engine( te_scales = t._rowwise_scale_inv[:M, : N // 16].view(torch.uint8) torch.testing.assert_close(te_codes, codes, atol=0, rtol=0) torch.testing.assert_close(te_scales, scales.view(torch.uint8), atol=0, rtol=0) + # Also run the CuTe DSL kernel explicitly against TE (the call above + # already dispatches to it when eligible; this pins the op itself). + if _cutedsl_available: + dsl_codes, dsl_scales = torch.ops.torchao.four_over_six_quantize_cutedsl( + x, amax, block, err_mode, e4m3_scale_bound + ) + torch.testing.assert_close(te_codes, dsl_codes, atol=0, rtol=0) + torch.testing.assert_close( + te_scales, dsl_scales.view(torch.uint8), atol=0, rtol=0 + ) + + +@_skip_no_cutedsl +@pytest.mark.parametrize("err_mode", ["mae", "mse"]) +@pytest.mark.parametrize("e4m3_scale_bound", [256, 448]) +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +@pytest.mark.parametrize("row_scaled", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_cutedsl_bitwise_matches_reference( + err_mode, e4m3_scale_bound, block, row_scaled, dtype +): + """CuTe DSL fast path is bitwise identical to the pure-PyTorch body. + + Shapes cover R < the kernel's 128-row tile (TMA-clipped stores), R not a + multiple of 16 (1x16 only), and multi-tile rows/columns. + """ + if row_scaled and block == "16x16": + pytest.skip("row-scaled is 1x16 only") + shapes = [(128, 256), (64, 1024)] + if block == "1x16": + shapes.append((100, 320)) + for shape in shapes: + torch.manual_seed(0) + x = torch.randn(*shape, dtype=dtype, device="cuda") + amax = (x.abs().amax(dim=1) if row_scaled else x.abs().amax()).to(torch.float32) + assert four_over_six_module._cutedsl_quantize_eligible(x) + codes, scales = four_over_six_quantize( + x, amax, block=block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + ref_codes, ref_scales = _reference_quantize( + x, amax, block=block, err_mode=err_mode, e4m3_scale_bound=e4m3_scale_bound + ) + torch.testing.assert_close(codes, ref_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales.view(torch.uint8), ref_scales.view(torch.uint8), atol=0, rtol=0 + ) + + +@_skip_no_cutedsl +@pytest.mark.parametrize("block", ["1x16", "16x16"]) +def test_cutedsl_special_values(block): + """Zeros, Inf injections, and amax==0 rows stay bitwise vs the reference.""" + torch.manual_seed(0) + # all zeros: S_enc falls back to 1.0, zero scales, zero codes + x = torch.zeros(64, 256, dtype=torch.bfloat16, device="cuda") + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize(x, amax, block=block) + ref_codes, ref_scales = _reference_quantize(x, amax, block=block) + torch.testing.assert_close(codes, ref_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales.view(torch.uint8), ref_scales.view(torch.uint8), atol=0, rtol=0 + ) + # Inf injections: block scale caps at 448, Inf encodes as +/-6, both + # candidate errors go Inf and the tie picks map-to-6 + x = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") + x[7, 32] = float("inf") + x[23, 100] = float("-inf") + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize(x, amax, block=block) + ref_codes, ref_scales = _reference_quantize(x, amax, block=block) + torch.testing.assert_close(codes, ref_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales.view(torch.uint8), ref_scales.view(torch.uint8), atol=0, rtol=0 + ) + if block == "1x16": + # row-scaled with amax == 0 rows over nonzero data: identity S_enc + x = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda") + row_amax = x.abs().amax(dim=1).to(torch.float32) + row_amax[::3] = 0.0 + codes, scales = four_over_six_quantize(x, row_amax, block=block) + ref_codes, ref_scales = _reference_quantize(x, row_amax, block=block) + torch.testing.assert_close(codes, ref_codes, atol=0, rtol=0) + torch.testing.assert_close( + scales.view(torch.uint8), ref_scales.view(torch.uint8), atol=0, rtol=0 + ) + + +@_skip_no_cutedsl +def test_cutedsl_nan_semantics(): + """NaN inputs follow the TE kernel semantics, which the pure-PyTorch body + cannot reproduce (torch.amax propagates NaN into the block scales while + the kernel's fmaxf drops it): an all-NaN group gets amax 0 -> scale byte + 0x00, and NaN elements encode to +6 (satfinite), i.e. code bytes 0x77.""" + torch.manual_seed(0) + x = torch.randn(32, 256, dtype=torch.bfloat16, device="cuda") + x[3, 32:48] = float("nan") # group (3, 2) + # a NaN-free global amax, as TE's own NaN-dropping amax kernel produces + amax = torch.nan_to_num(x.float(), nan=0.0).abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize(x, amax, block="1x16") + assert scales.view(torch.uint8)[3, 2].item() == 0x00 + assert (codes[3, 16:24] == 0x77).all() + # NaN-free groups are still bitwise vs the reference + ref_codes, ref_scales = _reference_quantize(x, amax, block="1x16") + keep = torch.ones_like(codes, dtype=torch.bool) + keep[3, 16:24] = False + torch.testing.assert_close(codes[keep], ref_codes[keep], atol=0, rtol=0) + + +@_skip_no_cutedsl +def test_cutedsl_ineligible_falls_back(): + """Ineligible shapes/layouts silently use the pure-PyTorch body.""" + x = torch.randn(64, 272, dtype=torch.bfloat16, device="cuda") # C % 64 != 0 + assert not four_over_six_module._cutedsl_quantize_eligible(x) + amax = x.abs().amax().to(torch.float32) + codes, scales = four_over_six_quantize(x, amax) + assert codes.shape == (64, 136) and scales.shape == (64, 17) + x_t = torch.randn(64, 256, dtype=torch.bfloat16, device="cuda").t() + assert not four_over_six_module._cutedsl_quantize_eligible(x_t) + + +@_skip_no_sm100 +@pytest.mark.skipif(not _cutedsl_available, reason="requires the CuTe DSL runtime") +def test_cutedsl_linear_compile(): + """torch.compile traces through the dispatch (custom op + fake impl).""" + torch.manual_seed(0) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(384, 256, dtype=torch.bfloat16, device="cuda") * 0.1 + + def fn(x, w): + return four_over_six_linear(x, w, None, "mae", 256, False) + + y_eager = fn(x, w) + y_compiled = torch.compile(fn, fullgraph=True)(x, w) + torch.testing.assert_close(y_compiled, y_eager, atol=0, rtol=0) diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py index 056d3aa81b..4330f91099 100644 --- a/torchao/prototype/moe_training/nvfp4_training/four_over_six.py +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -66,6 +66,11 @@ ) from torchao.prototype.mx_formats.utils import to_blocked +from .four_over_six_cutedsl import ( + _cutedsl_quantize_eligible, + four_over_six_quantize_cutedsl, +) + FP4_E2M1_MAX = 6.0 FP8_E4M3_MAX = 448.0 _FP32_MAX = torch.finfo(torch.float32).max @@ -196,6 +201,14 @@ def four_over_six_quantize( f"got shape {tuple(global_amax.shape)}" ) + # Fast path: the CuTe DSL kernel is an op-for-op reimplementation of the + # arithmetic below with bitwise-identical codes and scales; ineligible + # shapes/dtypes silently fall through to the pure-PyTorch body. + if _cutedsl_quantize_eligible(x): + return four_over_six_quantize_cutedsl( + x, global_amax, block, err_mode, e4m3_scale_bound + ) + xf = x.float().view(rows, cols // 16, 16) s_enc = four_over_six_global_encode_scale(global_amax, e4m3_scale_bound) if row_scaled: diff --git a/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py b/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py new file mode 100644 index 0000000000..64bc383652 --- /dev/null +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six_cutedsl.py @@ -0,0 +1,676 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""CuTe DSL NVFP4 four-over-six quantize kernel (SM100+). + +Fast path for ``four_over_six_quantize``: an op-for-op reimplementation of the +pure-PyTorch body (itself transcribed from TransformerEngine's +``quantize_4over6_nvfp4.cuh``), producing bitwise-identical codes and scales. +The load-bearing arithmetic is pinned with inline PTX so no compiler lowering +choice can change a rounding: + +* every division is a real ``div.rn.f32`` (the ``(block_amax / 6) * S_enc`` + association, ``S_enc``/``S_dec``, the encode reciprocals, and the + dequant-error denominator) — never a reciprocal multiply; +* the E4M3 scale cast is ``cvt.rn.satfinite.e4m3x2.f32`` and its exact decode + ``cvt.rn.f16x2.e4m3x2``; +* the FP4 cast is ``cvt.rn.satfinite.e2m1x2.f32`` (NaN -> +6) and the error + path dequantizes with the exact ``cvt.rn.f16x2.e2m1x2``; +* block amaxes use NaN-dropping ``max.f32`` (``fmaxf``: an all-NaN group + yields amax 0) and the 448 caps use NaN-dropping ``min.f32`` (``fminf``); +* the per-group dequant error is accumulated strictly sequentially in element + order 0..15 with scalar FP32 round-to-nearest adds (no reduction trees), and + 16x16 tiles reduce their 16 row errors with the exact width-16 + shuffle-down tree ``(((e0+e8)+(e4+e12))+((e2+e10)+(e6+e14))) + + (((e1+e9)+(e5+e13))+((e3+e11)+(e7+e15)))`` broadcast from the segment base. + +Tiling: one (128, 64) input tile per CTA, 128 threads, one tile row per +thread, so lanes 0-15 / 16-31 of each warp hold 16 consecutive rows — the +16-lane segments the 16x16 (2D) mode reduces over. The tile is loaded with +one TMA G2S copy, packed FP4 codes leave via TMA S2G (rows past R are +clipped by the TMA bounds check; zero-filled OOB rows compute garbage that +is never stored), and the four scale bytes per tile row leave as one u32. + +Eligibility (the dispatch gate in ``four_over_six_quantize``): CUDA bf16 or +fp32 input, contiguous, C % 64 == 0 (the tile width; also makes the u32 scale +store aligned). Ineligible calls fall through to the pure-PyTorch body. +""" + +import functools +from typing import Tuple + +import torch + +from torchao.utils import ceil_div + +from .hadamard_cutedsl_utils import ( + cutedsl_nvfp4_kernels_available, + raise_if_cutedsl_nvfp4_unavailable, +) + +TILE_ROWS = 128 +TILE_COLS = 64 # 4 1x16 groups per tile row; the C % 64 dispatch gate +_FP32_MAX = torch.finfo(torch.float32).max +# shfl.sync c-operand for width-16 segments: segmask 0x10 | clamp 0x1f, the +# same ((32 - width) << 8) | 0x1f encoding __shfl_down_sync(..., 16) uses. +_SHFL_WIDTH16 = 0x101F + + +@functools.cache +def _cutedsl_quantize_available() -> bool: + """Cached ``cutedsl_nvfp4_kernels_available()`` for the per-call dispatch gate.""" + return cutedsl_nvfp4_kernels_available() + + +def _cutedsl_quantize_eligible(x: torch.Tensor) -> bool: + """True iff ``four_over_six_quantize`` may dispatch ``x`` to the CuTe DSL kernel. + + Shape/dtype gates come first so FakeTensor tracing takes the same branch + as eager. Ineligible inputs silently use the pure-PyTorch body. + """ + return ( + x.is_cuda + and x.dtype in (torch.bfloat16, torch.float32) + and x.shape[0] > 0 + and x.shape[0] <= 65535 * TILE_ROWS # grid.y limit + and x.shape[1] > 0 + and x.shape[1] % TILE_COLS == 0 + and x.is_contiguous() + and _cutedsl_quantize_available() + ) + + +@functools.cache +def _compile_four_over_six_quantize_cutedsl( + input_dtype_name: str, + block_16x16: bool, + err_mode: str, + e4m3_scale_bound: int, + row_scaled: bool, + device_index: int, +): + """Compile one (dtype, block, err_mode, bound, row_scaled) kernel variant.""" + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + import cutlass.utils as utils + from cutlass._mlir.dialects import llvm + from cutlass.cute.nvgpu import cpasync, tcgen05 + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + from cutlass.cutlass_dsl import T, dsl_user_op + + from ._cutedsl_kernels_impl import ( + _abs_f32, + _cvt_rn_satfinite_e2m1x2_f32_pack4, + _min_f32, + ) + + if input_dtype_name == "torch.float32": + INPUT_CUTLASS_DTYPE = cutlass.Float32 + elif input_dtype_name == "torch.bfloat16": + INPUT_CUTLASS_DTYPE = cutlass.BFloat16 + else: + raise ValueError( + f"Unsupported input dtype for CuTe DSL four_over_six_quantize: {input_dtype_name}" + ) + + BLOCK_16X16 = block_16x16 + USE_MSE = err_mode == "mse" + ROW_SCALED = row_scaled + # bound * 6: the S_enc numerator AND the slow-path error denominator + # (1536.0 for bound 256, 2688.0 for bound 448) — both exact in FP32. + BOUND_TIMES_FP4MAX = float(e4m3_scale_bound) * 6.0 + + GROUPS_PER_ROW = TILE_COLS // 16 + CODE_TILE_BYTES = TILE_COLS // 2 + THREADS_PER_BLOCK = TILE_ROWS # one tile row per thread + input_elem_bytes = INPUT_CUTLASS_DTYPE.width // 8 + TILE_COPY_BYTES = TILE_ROWS * TILE_COLS * input_elem_bytes + + @dsl_user_op + def _div_rn_f32( + a: cutlass.Float32, b: cutlass.Float32, *, loc=None, ip=None + ) -> cutlass.Float32: + """Correctly-rounded FP32 division (``__fdiv_rn``), pinned with inline PTX. + + The bitwise contract needs real divisions — a reciprocal multiply + double-rounds and flips candidate picks near error ties — so we never + rely on how the DSL lowers ``/``. + """ + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + @dsl_user_op + def _fmax_f32( + a: cutlass.Float32, b: cutlass.Float32, *, loc=None, ip=None + ) -> cutlass.Float32: + # Plain max.f32 = fmaxf: NaN inputs are silently dropped, so an + # all-NaN group yields amax 0 (the TE amax semantics). This is NOT the + # NaN-propagating _max_f32 the RHT kernels use. + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "max.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + @dsl_user_op + def _mul_rn_f32( + a: cutlass.Float32, b: cutlass.Float32, *, loc=None, ip=None + ) -> cutlass.Float32: + """``__fmul_rn`` pinned with inline PTX so it can never fuse into an FMA. + + Used for the MSE ``diff * diff`` feeding the sequential error adds — + the one mul+add pair a contraction would double-round. + """ + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + @dsl_user_op + def _e4m3x2_rn_satfinite_with_decode( + c6: cutlass.Float32, c4: cutlass.Float32, *, loc=None, ip=None + ): + """E4M3-cast both candidate scales and decode them back, all exactly. + + ``cvt.rn.satfinite.e4m3x2.f32`` packs e4m3(c4) in the high byte and + e4m3(c6) in the low byte (matching ``__nv_cvt_float_to_fp8`` RN + satfinite); ``cvt.rn.f16x2.e4m3x2`` + ``cvt.f32.f16`` recover the + exact FP32 value of each scale byte. Returns (packed_u32, f6, f4). + """ + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.f32(), T.f32()]), + [c6.ir_value(loc=loc, ip=ip), c4.ir_value(loc=loc, ip=ip)], + ( + "{\n" + ".reg .b16 sp, s6, s4;\n" + ".reg .b32 dp;\n" + "cvt.rn.satfinite.e4m3x2.f32 sp, $4, $3;\n" + "cvt.u32.u16 $0, sp;\n" + "cvt.rn.f16x2.e4m3x2 dp, sp;\n" + "mov.b32 {s6, s4}, dp;\n" + "cvt.f32.f16 $1, s6;\n" + "cvt.f32.f16 $2, s4;\n" + "}" + ), + "=r,=f,=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + pair = cutlass.Uint32(llvm.extractvalue(T.i32(), res, [0])) + f6 = cutlass.Float32(llvm.extractvalue(T.f32(), res, [1])) + f4 = cutlass.Float32(llvm.extractvalue(T.f32(), res, [2])) + return pair, f6, f4 + + @dsl_user_op + def _dequant_e2m1x8_f32(w: cutlass.Uint32, *, loc=None, ip=None): + """Exact dequant of 8 packed FP4 codes to 8 FP32s in element order. + + ``cvt.rn.f16x2.e2m1x2`` per byte (low nibble -> low half) then exact + ``cvt.f32.f16`` — the slow-path error dequant. + """ + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32()] * 8), + [w.ir_value(loc=loc, ip=ip)], + ( + "{\n" + ".reg .b8 q0, q1, q2, q3;\n" + ".reg .b32 r0, r1, r2, r3;\n" + ".reg .b16 ha, hb;\n" + "mov.b32 {q0, q1, q2, q3}, $8;\n" + "cvt.rn.f16x2.e2m1x2 r0, q0;\n" + "mov.b32 {ha, hb}, r0;\n" + "cvt.f32.f16 $0, ha;\n" + "cvt.f32.f16 $1, hb;\n" + "cvt.rn.f16x2.e2m1x2 r1, q1;\n" + "mov.b32 {ha, hb}, r1;\n" + "cvt.f32.f16 $2, ha;\n" + "cvt.f32.f16 $3, hb;\n" + "cvt.rn.f16x2.e2m1x2 r2, q2;\n" + "mov.b32 {ha, hb}, r2;\n" + "cvt.f32.f16 $4, ha;\n" + "cvt.f32.f16 $5, hb;\n" + "cvt.rn.f16x2.e2m1x2 r3, q3;\n" + "mov.b32 {ha, hb}, r3;\n" + "cvt.f32.f16 $6, ha;\n" + "cvt.f32.f16 $7, hb;\n" + "}" + ), + "=f,=f,=f,=f,=f,=f,=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return tuple( + cutlass.Float32(llvm.extractvalue(T.f32(), res, [i])) for i in range(8) + ) + + @cute.struct + class SharedStorage: + tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1] + in_smem: cute.struct.Align[ + cute.struct.MemRange[INPUT_CUTLASS_DTYPE, TILE_ROWS * TILE_COLS], 128 + ] + out_smem: cute.struct.Align[ + cute.struct.MemRange[cutlass.Uint8, TILE_ROWS * CODE_TILE_BYTES], 128 + ] + + # The helpers below are plain (undecorated) python functions traced inline + # from the kernel body — the _cutedsl_kernels_impl convention — so their + # loops unroll at trace time (the sequential error adds stay 16 dependent + # scalar FP32 RN adds) and python tuples stay indexable. + + def _load_group(sIN, m_rel, k_base): + """Load one 1x16 group from SMEM, upconverting to FP32 exactly once.""" + raw = cute.make_rmem_tensor((16,), INPUT_CUTLASS_DTYPE) + cute.autovec_copy( + cute.make_tensor( + (sIN.iterator + (m_rel * TILE_COLS + k_base)).align(16), + cute.make_layout(16), + ), + raw, + ) + vals = cute.make_rmem_tensor((16,), cutlass.Float32) + for i in range(16): + vals[i] = cutlass.Float32(raw[i]) + return vals + + def _group_amax(vals): + """fmaxf chain over |vals| from 0.0 (NaN-dropping; order-insensitive).""" + amax = cutlass.Float32(0.0) + for i in range(16): + amax = _fmax_f32(amax, _abs_f32(vals[i])) + return amax + + def _tile_reduce_max16(v): + """16-lane-segment max: shuffle-down offsets 8/4/2/1, broadcast from base.""" + for delta in (8, 4, 2, 1): + v = _fmax_f32( + v, cute.arch.shuffle_sync_down(v, delta, mask_and_clamp=_SHFL_WIDTH16) + ) + return cute.arch.shuffle_sync(v, 0, mask_and_clamp=_SHFL_WIDTH16) + + def _tile_reduce_sum16(v): + """The exact TE error tree: value(l) += value(l+8), +=(l+4), +=(l+2), + +=(l+1) (clamped shuffle-down, NOT butterfly), broadcast from the + segment base. Summation shape + (((e0+e8)+(e4+e12))+((e2+e10)+(e6+e14))) + (((e1+e9)+(e5+e13))+((e3+e11)+(e7+e15))).""" + for delta in (8, 4, 2, 1): + v = v + cute.arch.shuffle_sync_down(v, delta, mask_and_clamp=_SHFL_WIDTH16) + return cute.arch.shuffle_sync(v, 0, mask_and_clamp=_SHFL_WIDTH16) + + def _scale_pair(block_amax, s_enc, s_dec): + """compute_scale_pair: base = (block_amax / 6) * S_enc — a real div.rn + THEN mul.rn — the 1.5x map-to-4 expansion, both capped at the full + 448 E4M3 range, and the exact-division encode multipliers.""" + base = _div_rn_f32(block_amax, cutlass.Float32(6.0)) * s_enc + c6 = _min_f32(base, cutlass.Float32(448.0)) + c4 = _min_f32(base * cutlass.Float32(1.5), cutlass.Float32(448.0)) + pair, f6, f4 = _e4m3x2_rn_satfinite_with_decode(c6, c4) + inv6 = _min_f32( + _div_rn_f32(cutlass.Float32(1.0), f6 * s_dec), + cutlass.Float32(_FP32_MAX), + ) + inv4 = _min_f32( + _div_rn_f32(cutlass.Float32(1.0), f4 * s_dec), + cutlass.Float32(_FP32_MAX), + ) + b6 = pair & 0xFF + b4 = (pair >> 8) & 0xFF + return b6, b4, f6, f4, inv6, inv4 + + def _encode_and_error(vals, inv, f_dec, gamax): + """One candidate: FP4-encode 16 values and accumulate the dequant + error strictly sequentially in element order 0..15 (16 dependent + FP32 RN adds — never a reduction tree). Returns (w0, w1, err).""" + q = cute.make_rmem_tensor((16,), cutlass.Float32) + for i in range(16): + q[i] = vals[i] * inv + # Even element -> low nibble: byte j = cvt(hi=q[2j+1], lo=q[2j]). + w0 = _cvt_rn_satfinite_e2m1x2_f32_pack4( + q[0], q[2], q[4], q[6], q[1], q[3], q[5], q[7] + ) + w1 = _cvt_rn_satfinite_e2m1x2_f32_pack4( + q[8], q[10], q[12], q[14], q[9], q[11], q[13], q[15] + ) + err = cutlass.Float32(0.0) + for half in range(2): + dq = _dequant_e2m1x8_f32(w0 if half == 0 else w1) + for i in range(8): + # val = div.rn(mul.rn(mul.rn(dequant, (f32)scale_byte), + # global_amax), 6*bound) in the input domain. + val = _div_rn_f32( + (dq[i] * f_dec) * gamax, + cutlass.Float32(BOUND_TIMES_FP4MAX), + ) + diff = val - vals[half * 8 + i] + if USE_MSE: + err = err + _mul_rn_f32(diff, diff) + else: + err = err + _abs_f32(diff) + return w0, w1, err + + class FourOverSixQuantizeKernel: + @cute.kernel + def kernel( + self, + tma_atom_in: cute.CopyAtom, + tma_tensor_in: cute.Tensor, + tma_atom_out: cute.CopyAtom, + tma_tensor_out: cute.Tensor, + scales_u32: cute.Tensor, + amax_f32: cute.Tensor, + R: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, bidy, _ = cute.arch.block_idx() + + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(SharedStorage) + tma_mbar_ptr = storage.tma_mbar_ptr.data_ptr() + + smem_layout_in = cute.make_layout( + (TILE_ROWS, TILE_COLS), stride=(TILE_COLS, 1) + ) + smem_layout_out = cute.make_layout( + (TILE_ROWS, CODE_TILE_BYTES), stride=(CODE_TILE_BYTES, 1) + ) + sIN = storage.in_smem.get_tensor(smem_layout_in) + sOUT = storage.out_smem.get_tensor(smem_layout_out) + # u64 view of the code tile: one 8-byte store per 1x16 group. + sOUT_u64 = cute.recast_tensor(sOUT, cutlass.Uint64) + + if tidx == 0: + cpasync.prefetch_descriptor(tma_atom_in) + cpasync.prefetch_descriptor(tma_atom_out) + cute.arch.mbarrier_init(tma_mbar_ptr, 1) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + m_tile = cutlass.Int64(bidy) + c_tile = cutlass.Int64(bidx) + + gIN_tile = cute.local_tile( + tma_tensor_in, (TILE_ROWS, TILE_COLS), (m_tile, c_tile) + ) + if warp_idx == 0: + cta_layout = cute.make_layout((1,)) + tINs, tINg = cpasync.tma_partition( + tma_atom_in, + 0, + cta_layout, + cute.group_modes(sIN, 0, 2), + cute.group_modes(gIN_tile, 0, 2), + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_mbar_ptr, TILE_COPY_BYTES + ) + cute.copy(tma_atom_in, tINg[None], tINs[None], tma_bar_ptr=tma_mbar_ptr) + + m_rel = tidx + m = m_tile * TILE_ROWS + m_rel + + # Global scale chain, once per thread (per row when row-scaled): + # S_enc = fminf(bound*6 / global_amax, FLT_MAX), falling back to + # 1.0 when global_amax == 0 or S_enc == 0; S_dec = 1 / S_enc. + gamax = cutlass.Float32(0.0) + if cutlass.const_expr(ROW_SCALED): + if m < R: + gamax = amax_f32[m] + else: + gamax = amax_f32[0] + s_enc_raw = _min_f32( + _div_rn_f32(cutlass.Float32(BOUND_TIMES_FP4MAX), gamax), + cutlass.Float32(_FP32_MAX), + ) + s_enc = s_enc_raw + if gamax == 0.0: + s_enc = cutlass.Float32(1.0) + if s_enc_raw == 0.0: + s_enc = cutlass.Float32(1.0) + s_dec = _div_rn_f32(cutlass.Float32(1.0), s_enc) + + cute.arch.mbarrier_wait(tma_mbar_ptr, 0) + + scale_word = cutlass.Uint32(0) + for gc in cutlass.range_constexpr(GROUPS_PER_ROW): + vals = _load_group(sIN, m_rel, gc * 16) + block_amax = _group_amax(vals) + if cutlass.const_expr(BLOCK_16X16): + block_amax = _tile_reduce_max16(block_amax) + b6, b4, f6, f4, inv6, inv4 = _scale_pair(block_amax, s_enc, s_dec) + w6_0, w6_1, err6 = _encode_and_error(vals, inv6, f6, gamax) + w4_0, w4_1, err4 = _encode_and_error(vals, inv4, f4, gamax) + if cutlass.const_expr(BLOCK_16X16): + err6 = _tile_reduce_sum16(err6) + err4 = _tile_reduce_sum16(err4) + # Strict < with map4 on the left: ties and NaN errors pick map6. + w0 = w6_0 + w1 = w6_1 + sbyte = b6 + if err4 < err6: + w0 = w4_0 + w1 = w4_1 + sbyte = b4 + sOUT_u64[m_rel, gc] = cutlass.Uint64(w0) | (cutlass.Uint64(w1) << 32) + scale_word = scale_word | (sbyte << (8 * gc)) + + if m < R: + scales_u32[m, c_tile] = scale_word + + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + if warp_idx == 0: + gOUT_tile = cute.local_tile( + tma_tensor_out, (TILE_ROWS, CODE_TILE_BYTES), (m_tile, c_tile) + ) + cta_layout = cute.make_layout((1,)) + tOUTs, tOUTg = cpasync.tma_partition( + tma_atom_out, + 0, + cta_layout, + cute.group_modes(sOUT, 0, 2), + cute.group_modes(gOUT_tile, 0, 2), + ) + cute.copy(tma_atom_out, tOUTs[None], tOUTg[None]) + + @cute.jit + def __call__( + self, + inp_rc: cute.Tensor, + out_codes: cute.Tensor, + scales_u32: cute.Tensor, + amax_f32: cute.Tensor, + R: cutlass.Int32, + r_tiles: cutlass.Int32, + c_tiles: cutlass.Int32, + stream: cuda.CUstream, + ): + smem_layout_in = cute.make_layout( + (TILE_ROWS, TILE_COLS), stride=(TILE_COLS, 1) + ) + smem_layout_out = cute.make_layout( + (TILE_ROWS, CODE_TILE_BYTES), stride=(CODE_TILE_BYTES, 1) + ) + g2s_op = cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE) + tma_atom_in, tma_tensor_in = cpasync.make_tiled_tma_atom( + g2s_op, + inp_rc, + smem_layout_in, + (TILE_ROWS, TILE_COLS), + ) + tma_atom_out, tma_tensor_out = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + out_codes, + smem_layout_out, + (TILE_ROWS, CODE_TILE_BYTES), + ) + self.kernel( + tma_atom_in, + tma_tensor_in, + tma_atom_out, + tma_tensor_out, + scales_u32, + amax_f32, + R, + ).launch( + grid=(c_tiles, r_tiles, 1), + block=(THREADS_PER_BLOCK, 1, 1), + cluster=(1, 1, 1), + smem=SharedStorage.size_in_bytes(), # pyrefly: ignore [missing-attribute] + stream=stream, + ) + + kernel = FourOverSixQuantizeKernel() + + r = cute.sym_int() + c = cute.sym_int(divisibility=64) + ch = cute.sym_int(divisibility=32) + cs = cute.sym_int() + fake_inp = make_fake_tensor( + INPUT_CUTLASS_DTYPE, + (r, c), + stride=(cute.sym_int(), cute.sym_int()), + ) + fake_out = make_fake_tensor( + cutlass.Uint8, + (r, ch), + stride=(cute.sym_int(), cute.sym_int()), + ) + fake_scales = make_fake_tensor( + cutlass.Uint32, + (r, cs), + stride=(cute.sym_int(), cute.sym_int()), + ) + fake_amax = make_fake_tensor( + cutlass.Float32, + (cute.sym_int(),), + stride=(cute.sym_int(),), + ) + fake_stream = make_fake_stream() + + return cute.compile( + kernel, + inp_rc=fake_inp, + out_codes=fake_out, + scales_u32=fake_scales, + amax_f32=fake_amax, + R=0, + r_tiles=1, + c_tiles=1, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + +@torch.library.custom_op("torchao::four_over_six_quantize_cutedsl", mutates_args=()) +def four_over_six_quantize_cutedsl( + x: torch.Tensor, + global_amax: torch.Tensor, + block: str, + err_mode: str, + e4m3_scale_bound: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """NVFP4 four-over-six quantize (CuTe DSL, SM100+). + + Bitwise-identical to the pure-PyTorch ``four_over_six_quantize`` body for + every (block, err_mode, e4m3_scale_bound, per-tensor/row-scaled) mode, + with one documented exception: NaN inputs follow the TE kernel semantics + (NaN-dropping block amax, NaN -> +6 FP4 codes) while torch's ``amax`` + propagates NaN into the reference's block scales. + + Args: + x: (R, C) bfloat16 or float32, contiguous, C % 64 == 0. + global_amax: scalar FP32 amax, or (R,) per-row amax (1x16 only). + block: "1x16" or "16x16". + err_mode: "mae" or "mse". + e4m3_scale_bound: 256 or 448. + + Returns: + (codes, scales): (R, C//2) uint8 packed FP4 codes (low nibble = even + element) and (R, C//16) float8_e4m3fn block scales. + """ + raise_if_cutedsl_nvfp4_unavailable("four_over_six_quantize_cutedsl") + if x.ndim != 2: + raise ValueError("x must be 2-D") + rows, cols = x.shape + if cols % TILE_COLS != 0: + raise ValueError( + f"four_over_six_quantize_cutedsl requires C % {TILE_COLS} == 0, got {cols}" + ) + if block == "16x16" and rows % 16: + raise ValueError(f"16x16 blocks need rows divisible by 16, got {rows}") + row_scaled = global_amax.dim() == 1 and global_amax.numel() == rows + + codes = torch.empty_strided( + (rows, cols // 2), (cols // 2, 1), device=x.device, dtype=torch.uint8 + ) + scales_u8 = torch.empty_strided( + (rows, cols // 16), (cols // 16, 1), device=x.device, dtype=torch.uint8 + ) + scales_u32 = scales_u8.view(torch.uint32) + amax_f32 = global_amax.to(torch.float32).reshape(-1).contiguous() + + device_index = x.device.index + if device_index is None: + device_index = torch.cuda.current_device() + compiled = _compile_four_over_six_quantize_cutedsl( + str(x.dtype), + block == "16x16", + err_mode, + int(e4m3_scale_bound), + row_scaled, + device_index, + ) + + import cuda.bindings.driver as cuda + + with torch.cuda.device(x.device): + stream = cuda.CUstream(int(torch.cuda.current_stream().cuda_stream)) + compiled( + x, + codes, + scales_u32, + amax_f32, + int(rows), + int(ceil_div(rows, TILE_ROWS)), + int(cols // TILE_COLS), + stream, + ) + return codes, scales_u8.view(torch.float8_e4m3fn) + + +@four_over_six_quantize_cutedsl.register_fake +def _(x, global_amax, block, err_mode, e4m3_scale_bound): + rows, cols = x.shape + codes = x.new_empty((rows, cols // 2), dtype=torch.uint8) + scales = x.new_empty((rows, cols // 16), dtype=torch.float8_e4m3fn) + return codes, scales