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 new file mode 100644 index 0000000000..96f6272f78 --- /dev/null +++ b/test/prototype/moe_training/nvfp4_training/test_four_over_six.py @@ -0,0 +1,401 @@ +# 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 + +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, + four_over_six_global_encode_scale, + 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 + +_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)", +) +_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): + """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) + # 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 new file mode 100644 index 0000000000..4330f91099 --- /dev/null +++ b/torchao/prototype/moe_training/nvfp4_training/four_over_six.py @@ -0,0 +1,554 @@ +# 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 + +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 + +__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)}" + ) + + # 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: + 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 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