diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.py b/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.py new file mode 100644 index 00000000000..6d4d6c48518 --- /dev/null +++ b/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.py @@ -0,0 +1,310 @@ +# 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 + +"""One fused MXFP8 cast of grad_out vs the three kernels it replaces. + + python benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.py + + MXFP8_BENCH_VALIDATE=1 also check every output bytewise against the baseline + MXFP8_BENCH_EXTRA_NON_NICE=1 add small and awkward shapes + +Baseline is the composition a grouped-GEMM backward runs today: the 1x32 CuTeDSL +cast for dgrad, the 32x1 CUDA cast for wgrad, and a standalone Triton K-groups +swizzle of the colwise scales. The fused kernel reads the input once and emits +all four outputs, moving 4.0625 bytes per input element against the +composition's 6.125. + +A NOTE ON WHAT THIS MEASURES. Both sides are timed under CUDA-GRAPH REPLAY, not +eagerly, which is the difference between measuring the kernels and measuring +torch's dispatcher. + +torchao's triton_op dispatch costs on the order of 100us per call independent of +problem size, and the baseline pays it THREE times against the fused kernel's +one. Timed eagerly that overhead swamps every shape below roughly 0.5 GB and +reports speedups of 15x on small shapes -- an artifact of launch overhead, not a +kernel win. Graph replay removes dispatch from both sides and leaves the +kernels, which is also what a real training loop sees, since those capture +graphs. + +The sibling benchmarks in this directory use eager do_bench. This one +deliberately does not: a 3-kernel baseline against a 1-kernel candidate is +exactly the case where dispatch overhead dominates the number being reported. +""" + +import itertools +import os +import random +from dataclasses import dataclass +from typing import List + +import torch +from tabulate import tabulate +from tqdm import tqdm +from triton.testing import do_bench_cudagraph + +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_quantize_2d_1x32_32x1_fused import ( + quantize_fused, + s_col_shape, +) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( + mxfp8_quantize_2d_1x32_cutedsl, + triton_mx_block_rearrange_2d_K_groups, +) +from torchao.prototype.moe_training.utils import generate_jagged_offs +from torchao.prototype.mx_formats.config import ( + KernelPreference, + MXFP8Dim1CastKernelChoice, +) +from torchao.prototype.mx_formats.mx_tensor import ScaleCalculationMode +from torchao.prototype.mx_formats.utils import _to_mxfp8_dim1_kernel_wrapper + +device = torch.device("cuda") +VALIDATE = os.environ.get("MXFP8_BENCH_VALIDATE", "0") == "1" +EXTRA_NON_NICE = os.environ.get("MXFP8_BENCH_EXTRA_NON_NICE", "0") == "1" + +BLOCK_SIZE = 32 + +# Needed since changing args to function causes recompiles +torch._dynamo.config.cache_size_limit = 1000 + + +@dataclass(frozen=True) +class ExperimentConfig: + input_shape: tuple[int, int] + num_groups: int + + +@dataclass(frozen=True) +class ExperimentResult: + # time + fused_us: float + unfused_us: float + # mem bw + fused_gbps: float + unfused_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), + ] + if EXTRA_NON_NICE: + input_shapes += [ + (256, 128), + (1152, 1408), + ] + num_groups_list = [4, 8] + configs = [] + for shape, num_groups in itertools.product(input_shapes, num_groups_list): + total_M, N = shape + # The kernel requires both dims to be multiples of 128, and every group + # size to be a multiple of 128 as well. + if total_M % 128 or N % 128 or num_groups > total_M // 128: + continue + configs.append(ExperimentConfig(input_shape=shape, num_groups=num_groups)) + return configs + + +def benchmark_cuda_graph_in_microseconds(f, *args, **kwargs) -> float: + """Median device time of one call, under CUDA-graph replay.""" + return do_bench_cudagraph(lambda: f(*args, **kwargs), return_mode="median") * 1e3 + + +def assert_graph_capture_is_real(x: torch.Tensor, offs: torch.Tensor): + """Fail loudly if the fused kernel captures as an EMPTY graph. + + A CuTeDSL kernel that launches on the default stream instead of the caller's + captures nothing: no error, and the replay then times at about 1us, which + reads as a 100x win. Every campaign winner for this kernel shipped with that + bug before it was patched, so check rather than assume. Replay into zeroed + buffers and compare against the eager result. + """ + total_M, N = x.shape + num_groups = offs.numel() + s_col_rows, s_col_cols = s_col_shape(total_M, N, num_groups) + bufs = ( + torch.empty((total_M, N), device=x.device, dtype=torch.float8_e4m3fn), + torch.empty((total_M, N // BLOCK_SIZE), device=x.device, dtype=torch.uint8), + torch.empty((N, total_M), device=x.device, dtype=torch.float8_e4m3fn), + torch.empty((s_col_rows, s_col_cols), device=x.device, dtype=torch.uint8), + ) + + def run(): + quantize_fused( + x, + offs=offs, + block_size=BLOCK_SIZE, + scaling_mode="rceil", + q_row_out=bufs[0], + s_row_out=bufs[1], + q_col_out=bufs[2], + s_col_out=bufs[3], + ) + + run() + eager = [b.view(torch.uint8).clone() for b in bufs] + + for b in bufs: + b.view(torch.uint8).zero_() + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + run() + for b in bufs: + b.view(torch.uint8).zero_() + g.replay() + torch.cuda.synchronize() + + for name, got, exp in zip(("q_row", "s_row", "q_col", "s_col"), bufs, eager): + if not torch.equal(got.view(torch.uint8), exp): + raise RuntimeError( + f"CUDA-graph replay != eager on {name}. The kernel almost " + "certainly ignores the caller's stream, so it captured as an " + "empty graph and any timing below would be meaningless." + ) + + +def unfused_baseline(x: torch.Tensor, group_offs: torch.Tensor): + """The three kernels a grouped-GEMM backward dispatches for grad_out today.""" + q_row, s_row = mxfp8_quantize_2d_1x32_cutedsl( + x, block_size=BLOCK_SIZE, scaling_mode="rceil", offs=group_offs + ) + mx = _to_mxfp8_dim1_kernel_wrapper( + x, + BLOCK_SIZE, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=x.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + ) + s_col = triton_mx_block_rearrange_2d_K_groups(mx.scale, group_offs // BLOCK_SIZE) + return q_row, s_row, mx.qdata, s_col + + +def validate_outputs(fused, unfused, num_groups: int): + names = ("q_row", "s_row", "q_col", "s_col") + for name, got, exp in zip(names, fused, unfused): + got = got.view(torch.uint8) + exp = exp.view(torch.uint8) + if name == "s_col": + # Trailing per-group padding slack, never read by the grouped GEMM. + used = exp.shape[1] - num_groups * 4 + got, exp = ( + got.flatten()[: exp.shape[0] * used], + exp.flatten()[: exp.shape[0] * used], + ) + torch.testing.assert_close(got, exp, rtol=0, atol=0, msg=f"{name} mismatch") + + +def run_experiment(config: ExperimentConfig) -> ExperimentResult: + total_M, N = config.input_shape + num_groups = config.num_groups + + x = torch.randn(total_M, N, dtype=torch.bfloat16, device=device) + offs = generate_jagged_offs(num_groups, total_M, multiple_of=128, device=device).to( + torch.int32 + ) + + fused_out = quantize_fused( + x, block_size=BLOCK_SIZE, scaling_mode="rceil", offs=offs + ) + if VALIDATE: + validate_outputs(fused_out, unfused_baseline(x, offs), num_groups) + + assert_graph_capture_is_real(x, offs) + + fused_time_us = benchmark_cuda_graph_in_microseconds( + quantize_fused, x, block_size=BLOCK_SIZE, scaling_mode="rceil", offs=offs + ) + unfused_time_us = benchmark_cuda_graph_in_microseconds(unfused_baseline, x, offs) + + bytes_per_input_el = torch.finfo(torch.bfloat16).bits / 8 + bytes_per_output_el = torch.finfo(torch.float8_e4m3fn).bits / 8 + bytes_per_scale_el = torch.finfo(torch.float8_e8m0fnu).bits / 8 + + q_row, s_row, q_col, s_col = fused_out + write_bytes = ( + q_row.numel() * bytes_per_output_el + + q_col.numel() * bytes_per_output_el + + s_row.numel() * bytes_per_scale_el + + s_col.numel() * bytes_per_scale_el + ) + # The two sides move DIFFERENT byte volumes: the baseline reads x twice and + # rewrites the colwise scales, so each is charged its own traffic. That makes + # the GB/s columns comparable against a roofline but NOT against each other; + # the speedup column carries the win. + fused_read_bytes = x.numel() * bytes_per_input_el + unfused_read_bytes = 2 * fused_read_bytes + s_col.numel() * bytes_per_scale_el + unfused_write_bytes = write_bytes + s_col.numel() * bytes_per_scale_el + + fused_gbps = ((fused_read_bytes + write_bytes) / 1e9) / (fused_time_us / 1e6) + unfused_gbps = ((unfused_read_bytes + unfused_write_bytes) / 1e9) / ( + unfused_time_us / 1e6 + ) + + return ExperimentResult( + fused_us=fused_time_us, + unfused_us=unfused_time_us, + fused_gbps=fused_gbps, + unfused_gbps=unfused_gbps, + ) + + +def print_results(experiments: List[Experiment]): + headers = [ + "input_shape", + "num_groups", + "fused_us", + "unfused_us", + "speedup", + "fused_gbps", + "unfused_gbps", + ] + rows = [] + for experiment in experiments: + speedup = experiment.result.unfused_us / experiment.result.fused_us + rows.append( + [ + str(experiment.config.input_shape), + experiment.config.num_groups, + f"{experiment.result.fused_us:.2f}", + f"{experiment.result.unfused_us:.2f}", + f"{speedup:.2f}x", + f"{experiment.result.fused_gbps:.1f}", + f"{experiment.result.unfused_gbps:.1f}", + ] + ) + print(tabulate(rows, headers=headers)) + + +def main(): + torch.random.manual_seed(123) + random.seed(123) + configs = get_configs() + results = [] + for config in tqdm(configs): + result = run_experiment(config) + results.append(Experiment(config=config, result=result)) + + print_results(results) + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index 8a4a603ddce..e8a315dbd6b 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -8,7 +8,7 @@ import torch # FP8 MoE kernels require FP8-capable hardware (SM 10.x on CUDA, MI300+ on ROCm) -from torchao.utils import is_MI300, is_MI350 +from torchao.utils import ceil_div, is_MI300, is_MI350 def _is_sm_10x() -> bool: @@ -39,6 +39,7 @@ def _is_sm_10x() -> bool: fused_pad_token_groups_cuda, fused_unpad_token_groups_cuda, mx_block_rearrange_2d_M_groups_cuda, + mxfp8_quantize_2d_1x32_32x1_cutedsl, mxfp8_quantize_2d_1x32_cutedsl, mxfp8_quantize_2d_1x32_flydsl, mxfp8_quantize_2d_32x1_cutedsl, @@ -66,9 +67,17 @@ def _is_sm_10x() -> bool: torch_to_float8_per_group_colwise, torch_to_float8_per_group_rowwise, ) +from torchao.prototype.mx_formats.config import ( + KernelPreference, + MXFP8Dim1CastKernelChoice, +) from torchao.prototype.mx_formats.kernels import triton_mx_block_rearrange from torchao.prototype.mx_formats.mx_tensor import ScaleCalculationMode, to_mx -from torchao.prototype.mx_formats.utils import from_blocked, to_blocked +from torchao.prototype.mx_formats.utils import ( + _to_mxfp8_dim1_kernel_wrapper, + from_blocked, + to_blocked, +) from torchao.testing._mxfp8_test_utils import ( assert_mxfp8_semantics, make_mxfp8_semantic_cases, @@ -1682,3 +1691,561 @@ def test_amd_mx_3d_flydsl_numerics( assert y.stride() == y_ref.stride(), "quantized tensor strides do not match" assert y.dtype == torch.float8_e4m3fn assert s.dtype == torch.float8_e8m0fnu + + +# ============================================================================= +# Fused 1x32 + 32x1 MXFP8 quantization (cutedsl_quantize_2d_1x32_32x1_fused.py) +# +# The kernel replaces three kernels in a grouped-GEMM backward -- the 1x32 +# rowwise cast, the 32x1 colwise cast, and the standalone K-groups scale swizzle +# -- with one pass over the input. to_mx(RCEIL) is the numerical oracle for both +# axes; per-group unblocking separately checks the dynamic column-scale layout. +# ============================================================================= + +_FUSED_CUTEDSL_REQUIRED = pytest.mark.skipif( + not (_is_sm_10x() and _mxfp8_cutedsl_kernels_available), + reason="Requires SM10.x and the CuTeDSL runtime", +) + +_FUSED_1X32_32X1_SHAPES = ( + # (total_M, N, num_groups). Both dims must be multiples of 128. Each new + # triple costs a JIT compile, so this list is deliberately short. + (512, 256, 4), + (1024, 512, 8), + (2048, 512, 4), + (4096, 1024, 8), + # (M // 128) * (N // 128) = 64 * 9 = 576 CTAs, past SMALL_GRID_CTAS (512) in + # _launch -- the kernel switches from 512 to 256 threads/CTA here, a + # separately-coded path (threads == 256 branches in _kernel) that none of + # the shapes above reach. Without this shape every other test in this file + # only exercises the 512-thread branch. + (8192, 1152, 8), +) + + +def _fused_1x32_32x1_offs_with_tail(total_M, num_groups, tail, device="cuda"): + """Group end offsets covering total_M - tail rows, every size a multiple of 128. + + Unlike generate_jagged_offs, this deliberately stops SHORT of total_M so the + uncovered-tile path is exercised. + """ + covered = total_M - tail + assert covered % 128 == 0 + n_blocks = covered // 128 + assert n_blocks >= num_groups + sizes = [n_blocks // num_groups] * num_groups + for i in range(n_blocks % num_groups): + sizes[i] += 1 + ends = torch.cumsum(torch.tensor(sizes, dtype=torch.int32), 0) * 128 + return ends.to(device=device, dtype=torch.int32) + + +def _fused_1x32_32x1_composition_reference(x, offs): + """The three kernels the fused kernel replaces, in the order backward runs them. + + Returns the fused kernel's four outputs as raw uint8, in the same order. + """ + q_row, s_row = mxfp8_quantize_2d_1x32_cutedsl( + x, block_size=32, scaling_mode="rceil", offs=offs + ) + mx = _to_mxfp8_dim1_kernel_wrapper( + x, + 32, + elem_dtype=torch.float8_e4m3fn, + hp_dtype=x.dtype, + kernel_preference=KernelPreference.AUTO, + cast_kernel_choice=MXFP8Dim1CastKernelChoice.CUDA, + scale_calculation_mode=ScaleCalculationMode.RCEIL, + ) + s_col = triton_mx_block_rearrange_2d_K_groups(mx.scale, offs // 32) + return ( + q_row.view(torch.uint8), + s_row.view(torch.uint8), + mx.qdata.view(torch.uint8), + s_col.view(torch.uint8), + ) + + +def _assert_fused_1x32_32x1_matches_reference(x, offs, actual): + """Check direct RCEIL bytes after unblocking; FP8 NaN signs are equivalent.""" + m, n = x.shape + row_scales, row_data = to_mx(x, torch.float8_e4m3fn, 32, ScaleCalculationMode.RCEIL) + col_scales, col_data = to_mx( + x.t().contiguous(), torch.float8_e4m3fn, 32, ScaleCalculationMode.RCEIL + ) + expected = (row_data, row_scales, col_data, col_scales.view(torch.uint8).clone()) + logical_row = from_blocked(actual[1].view(torch.uint8), m, n // 32) + flat_col = actual[3].view(torch.uint8).flatten() + if offs is None: + logical_col = from_blocked(flat_col, n, m // 32) + else: + logical_col = torch.zeros_like(expected[3]) + start = 0 + for end in offs.tolist(): + if end > start: + slab = flat_col[(start // 32) * n : (end // 32) * n] + logical_col[:, start // 32 : end // 32] = from_blocked( + slab, n, (end - start) // 32 + ) + start = end + expected[3][:, start // 32 :].zero_() + # 128-aligned groups have no internal padding: the consumed prefix is + # M*N/32 bytes, followed by legacy G*4 scale-column slack. + assert torch.all(flat_col[m * n // 32 :] == 0), ( + "unused scale slack must be zero" + ) + unpacked = (actual[0], logical_row, actual[2], logical_col) + names = ("q_row", "s_row", "q_col", "s_col") + for name, got, exp in zip(names, unpacked, expected, strict=True): + got, exp = ( + got.contiguous().view(torch.uint8), + exp.contiguous().view(torch.uint8), + ) + if name.startswith("q_"): + got, exp = ( + torch.where(got == 0xFF, 0x7F, got), + torch.where(exp == 0xFF, 0x7F, exp), + ) + torch.testing.assert_close(got, exp, rtol=0, atol=0, msg=f"{name} mismatch") + + +def _fused_1x32_32x1_randn(*shape, dtype=torch.bfloat16, device="cuda"): + """Create reproducible inputs without changing process-global RNG state.""" + generator = torch.Generator(device=device).manual_seed(42) + return torch.randn(*shape, dtype=dtype, device=device, generator=generator) + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("total_M,N,num_groups", _FUSED_1X32_32X1_SHAPES) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_numerics(total_M, N, num_groups): + x = _fused_1x32_32x1_randn(total_M, N) + offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 0) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + + q_row, s_row, q_col, s_col = out + assert q_row.shape == (total_M, N) + assert s_row.shape == (total_M, N // 32) + assert q_col.shape == (N, total_M) + assert q_row.stride() == (N, 1) + assert q_col.stride() == (total_M, 1) + assert s_col.shape[0] == ceil_div(N, 128) * 128 + + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + expected = _fused_1x32_32x1_composition_reference(x, offs) + for got, reference in zip(out, expected, strict=True): + torch.testing.assert_close(got.view(torch.uint8), reference, rtol=0, atol=0) + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("tail", (128, 512, 1024)) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_uncovered_tail(tail): + """offs[-1] < total_M, which for a dropless MoE dispatcher is the NORM. + + Such a dispatcher sizes this buffer for worst-case per-expert padding and + fills only a variable prefix, so rows that no group owns are expected rather + than exceptional. Every offset generator in this file ends at exactly + total_M, so nothing else here covers this path -- and getting it wrong is not + a rounding error: the store lands inside group 0's slab and every n-tile CTA + races on one 512-byte line, which surfaces downstream as grad_norm=inf. + """ + total_M, N, num_groups = 2048, 512, 4 + covered = total_M - tail + offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, tail) + assert int(offs[-1]) == covered < total_M + + x = _fused_1x32_32x1_randn(total_M, N) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("scale", (1e-38, 1e-30, 1e30)) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_degenerate_blocks(scale): + """Blocks whose amax is subnormal enough to land on E8M0 scale byte 0. + + Both halves apply a 2**127 multiplier there, matching the reciprocal + `254 - scale_biased` of the kernels being replaced. `randn` structurally + cannot reach this branch, so it has to be constructed: a suite built only on + random tensors passes while leaving the multiplier untested, which is + exactly how a 2x error survived here once already. + """ + total_M, N, num_groups = 512, 256, 4 + x = (_fused_1x32_32x1_randn(total_M, N, dtype=torch.float32) * scale).to( + torch.bfloat16 + ) + offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 0) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_mixed_magnitude(): + """Interleaved tiny and huge rows, so degenerate and saturating blocks + coexist in one launch and neither branch can be tuned for in isolation.""" + total_M, N, num_groups = 512, 256, 4 + x = _fused_1x32_32x1_randn(total_M, N) + x[::4] = (x[::4].float() * 1e-36).to(torch.bfloat16) + x[1::4] = (x[1::4].float() * 1e28).to(torch.bfloat16) + offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 0) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_dirty_output_buffers(): + """Every output byte must be written, including regions no group covers. + + Uses the destination-passing entry point with buffers pre-filled to 0xFF. A + fresh torch.empty hides a never-written region whenever the allocator hands + back zeroed pages, so this is the only test here that can catch it. + """ + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_quantize_2d_1x32_32x1_fused import ( + quantize_fused, + s_col_shape, + ) + + total_M, N, num_groups = 2048, 512, 4 + offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 512) + x = _fused_1x32_32x1_randn(total_M, N) + + s_col_rows, s_col_cols = s_col_shape(total_M, N, num_groups) + bufs = ( + torch.empty((total_M, N), device="cuda", dtype=torch.float8_e4m3fn), + torch.empty((total_M, N // 32), device="cuda", dtype=torch.uint8), + torch.empty((N, total_M), device="cuda", dtype=torch.float8_e4m3fn), + torch.empty((s_col_rows, s_col_cols), device="cuda", dtype=torch.uint8), + ) + for _ in range(3): + for buf in bufs: + buf.view(torch.uint8).fill_(0xFF) + quantize_fused( + x, + offs=offs, + q_row_out=bufs[0], + s_row_out=bufs[1], + q_col_out=bufs[2], + s_col_out=bufs[3], + ) + _assert_fused_1x32_32x1_matches_reference(x, offs, bufs) + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("axis", ["rowwise", "colwise"]) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_special_value_semantics(axis): + """Pin both orientations to the independent shared MXFP8 special-value contract.""" + cases = make_mxfp8_semantic_cases(torch.bfloat16, "rceil", device="cuda") + num_cases = len(cases.names) + x = torch.zeros((128, 128), dtype=torch.bfloat16, device="cuda") + if axis == "rowwise": + x[:num_cases] = cases.inputs.repeat(1, 4) + else: + x[:, :num_cases] = cases.inputs.t().repeat(4, 1) + offs = torch.tensor([128], dtype=torch.int32, device="cuda") + + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + q_data, blocked_scales = out[:2] if axis == "rowwise" else out[2:] + + scales = from_blocked(blocked_scales.flatten()[:512], 128, 4) + for block_idx in range(4): + assert_mxfp8_semantics( + q_data[:num_cases, block_idx * 32 : (block_idx + 1) * 32], + scales[:num_cases, block_idx : block_idx + 1], + cases, + ) + + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("poison", ("inf", "nan", "inf_and_nan")) +@pytest.mark.parametrize("shape", [(128, 128), (8192, 2048)]) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_inf_and_nan(poison, shape): + """Both NOSAT references poison nonfinite blocks, including quantized data.""" + # Exercise both the 512-thread small-grid and 256-thread large-grid paths. + x = torch.ones(shape, dtype=torch.bfloat16, device="cuda") + if poison in ("inf", "inf_and_nan"): + x[0].fill_(float("inf")) + if poison in ("nan", "inf_and_nan"): + x[1].fill_(float("nan")) + offs = torch.tensor([shape[0]], dtype=torch.int32, device="cuda") + + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + assert int(out[3].view(torch.uint8).reshape(-1)[0]) == 255 + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("shape", [(128, 256), (256, 128)]) +@pytest.mark.parametrize("blocked_scale_output", [False, True]) +def test_mxfp8_32x1_cutedsl_fake_metadata(shape, blocked_scale_output): + """Fake metadata must match real strides and both scale layout modes.""" + x = _fused_1x32_32x1_randn(*shape) + torch.library.opcheck( + torch.ops.torchao.mxfp8_quantize_2d_32x1_cutedsl.default, + (x,), + {"blocked_scale_output": blocked_scale_output}, + test_utils=["test_schema", "test_faketensor"], + ) + + +@_FUSED_CUTEDSL_REQUIRED +@pytest.mark.parametrize("shape", [(128, 256), (256, 128)]) +@pytest.mark.parametrize("use_offsets", [False, True]) +def test_mxfp8_1x32_32x1_cutedsl_fake_metadata(shape, use_offsets): + """The fused custom op must expose all four real output layouts.""" + x = _fused_1x32_32x1_randn(*shape) + offs = ( + torch.tensor([shape[0]], device="cuda", dtype=torch.int32) + if use_offsets + else None + ) + torch.library.opcheck( + torch.ops.torchao.mxfp8_quantize_2d_1x32_32x1_cutedsl.default, + (x,), + {"offs": offs}, + test_utils=["test_schema", "test_faketensor"], + ) + + +@pytest.mark.parametrize("shape", [(128, 128), (8192, 1152)]) +@pytest.mark.parametrize("axis", ["rowwise", "colwise"]) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_scale_zero_boundary(shape, axis): + """Check exact amax neighbors of the RCEIL scale0/1 boundary on both CTA paths.""" + values = torch.tensor( + [0x045F, 0x0460, 0x0461, 0x04DF, 0x04E0, 0x04E1], + device="cuda", + dtype=torch.uint16, + ).view(torch.bfloat16) + m, n = shape + if axis == "rowwise": + x = values.repeat(ceil_div(m, len(values)))[:m, None].expand(m, n).contiguous() + else: + x = values.repeat(ceil_div(n, len(values)))[None, :n].expand(m, n).contiguous() + for offs in [None, torch.tensor([m], dtype=torch.int32, device="cuda")]: + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@pytest.mark.parametrize( + "offsets,m", [([0, 256, 256, 512], 512), ([512, 512], 512), ([0], 256), ([], 256)] +) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_empty_groups(offsets, m): + """Empty experts and an empty offset tensor cover no rows, unlike dense None.""" + x = _fused_1x32_32x1_randn(m, 256) + offs = torch.tensor(offsets, device="cuda", dtype=torch.int32) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs=offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@pytest.mark.parametrize("shape", [(512, 256), (8192, 1152)]) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_dense_layout(shape): + """Dense mode preserves defined grouped bytes without synthetic offsets or slack.""" + x = _fused_1x32_32x1_randn(*shape) + dense = mxfp8_quantize_2d_1x32_32x1_cutedsl(x) + grouped = mxfp8_quantize_2d_1x32_32x1_cutedsl( + x, offs=torch.tensor([shape[0]], dtype=torch.int32, device="cuda") + ) + assert dense[3].numel() == x.numel() // 32 + _assert_fused_1x32_32x1_matches_reference(x, None, dense) + for index, (left, right) in enumerate(zip(dense, grouped, strict=True)): + right = right.flatten()[: left.numel()] if index == 3 else right.flatten() + torch.testing.assert_close( + left.view(torch.uint8).flatten(), right.view(torch.uint8), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("use_offsets", [False, True]) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_compile(use_offsets): + """Exercise actual Inductor custom-op/fake contracts with two matrix shapes.""" + compiled = torch.compile( + mxfp8_quantize_2d_1x32_32x1_cutedsl, fullgraph=True, dynamic=True + ) + for m in [256, 512]: + x = _fused_1x32_32x1_randn(m, 256) + offs = ( + torch.tensor([m], dtype=torch.int32, device="cuda") if use_offsets else None + ) + _assert_fused_1x32_32x1_matches_reference(x, offs, compiled(x, offs=offs)) + + +@pytest.mark.parametrize("fake", [False, True]) +@pytest.mark.parametrize( + "kwargs,match", + [({"block_size": 16}, "block_size"), ({"scaling_mode": "floor"}, "rceil")], +) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_real_fake_api_contract( + fake, kwargs, match +): + """New public/fake signatures reject unsupported parameters consistently.""" + from contextlib import nullcontext + + from torch._subclasses.fake_tensor import FakeTensorMode + + with FakeTensorMode() if fake else nullcontext(): + x = torch.empty((128, 128), device="cuda", dtype=torch.bfloat16) + with pytest.raises(AssertionError, match=match): + mxfp8_quantize_2d_1x32_32x1_cutedsl(x, **kwargs) + + +@pytest.mark.parametrize( + "name,shape,dtype", + [ + ("q_row_out", (128, 128), torch.float8_e4m3fn), + ("s_col_out", (128, 4), torch.uint8), + ], +) +@pytest.mark.parametrize( + "invalid", ["shape", "dtype", "strides", "alignment", "device"] +) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_output_contract( + monkeypatch, name, shape, dtype, invalid +): + """Invalid destinations fail before any raw aligned GPU write.""" + from torchao.prototype.moe_training.kernels.mxfp8 import ( + cutedsl_quantize_2d_1x32_32x1_fused as kernel, + ) + + x = _fused_1x32_32x1_randn(128, 128) + output = torch.empty(shape, dtype=dtype, device="cuda") + if invalid == "shape": + output = torch.empty((shape[0], shape[1] + 1), dtype=dtype, device="cuda") + elif invalid == "dtype": + output = output.float() + elif invalid == "strides": + output = torch.empty(shape[::-1], dtype=dtype, device="cuda").t() + elif invalid == "alignment": + output = torch.empty(output.numel() + 1, dtype=dtype, device="cuda")[1:].view( + shape + ) + else: + output = output.cpu() + monkeypatch.setattr( + kernel, "run", lambda *args: pytest.fail("invalid output reached a GPU launch") + ) + with pytest.raises(AssertionError, match=name): + kernel.quantize_fused(x, **{name: output}) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two CUDA devices") +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_input_device_stream(): + """Launch on the input's nondefault stream even when another device is current.""" + x = _fused_1x32_32x1_randn(256, 256, device="cuda:1") + with torch.cuda.device(1): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + x.add_(1) + with torch.cuda.device(0): + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x) + stream.synchronize() + _assert_fused_1x32_32x1_matches_reference(x, None, out) + + +@pytest.mark.parametrize("alias", ["input", "another_output", "disjoint_views"]) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_output_overlap(monkeypatch, alias): + """Overlapping destinations fail before writes; disjoint arena views are valid.""" + from torchao.prototype.moe_training.kernels.mxfp8 import ( + cutedsl_quantize_2d_1x32_32x1_fused as kernel, + ) + + x = _fused_1x32_32x1_randn(128, 128) + if alias == "input": + output = x.view(torch.float8_e4m3fn).flatten()[: x.numel()].view(x.shape) + buffers = {"q_row_out": output} + elif alias == "another_output": + output = torch.empty_like(x, dtype=torch.float8_e4m3fn) + buffers = {"q_row_out": output, "q_col_out": output} + else: + arena = torch.empty( + 2 * x.numel() + 2 * x.numel() // 32, dtype=torch.uint8, device=x.device + ) + buffers, start = {}, 0 + for name, shape, dtype in [ + ("q_row_out", (128, 128), torch.float8_e4m3fn), + ("s_row_out", (128, 4), torch.uint8), + ("q_col_out", (128, 128), torch.float8_e4m3fn), + ("s_col_out", (128, 4), torch.uint8), + ]: + size = shape[0] * shape[1] + buffers[name] = arena[start : start + size].view(dtype).view(shape) + start += size + out = kernel.quantize_fused(x, **buffers) + _assert_fused_1x32_32x1_matches_reference(x, None, out) + return + monkeypatch.setattr( + kernel, "run", lambda *args: pytest.fail("aliased output reached a GPU launch") + ) + with pytest.raises(AssertionError, match="must not overlap"): + kernel.quantize_fused(x, **buffers) + + +def _run_mxfp8_fused_offset_validation(offsets): + """Run only device validation: no raw stores, async errors isolated in a child.""" + import cuda.bindings.driver as driver + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack, make_fake_stream + + from torchao.prototype.moe_training.kernels.mxfp8.cute_utils import ( + validate_group_sizes, + ) + + @cute.kernel + def check(tensor): + tid, _, _ = cute.arch.thread_idx() + if tid == 0: + validate_group_sizes(tensor, max_offset=256) + + @cute.jit + def launch(tensor, stream): + check(tensor).launch(grid=[1, 1, 1], block=[32, 1, 1], stream=stream) + + tensor = torch.tensor(offsets, dtype=torch.int32, device="cuda") + compiled = cute.compile( + launch, + from_dlpack(tensor, assumed_align=16, enable_tvm_ffi=True), + make_fake_stream(), + options="--enable-tvm-ffi --enable-assertions", + ) + compiled(tensor, driver.CUstream(int(torch.cuda.current_stream().cuda_stream))) + torch.cuda.synchronize() + + +@pytest.mark.parametrize( + "offsets,match", + [ + ([-128, 256], "Offsets must be nonnegative"), + ([256, 128], "Offsets must be nondecreasing"), + ([384], "Offset exceeds input rows"), + ([0, 128, 128, 256], None), + ], +) +@_FUSED_CUTEDSL_REQUIRED +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_offset_bounds(offsets, match): + """Reject negative, descending and too-large aligned endpoints on the device.""" + import subprocess + import sys + + script = ( + "from test.prototype.moe_training.test_kernels import _run_mxfp8_fused_offset_validation; _run_mxfp8_fused_offset_validation(" + + repr(offsets) + + ")" + ) + process = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=60 + ) + error = process.stdout + process.stderr + if match is None: + assert process.returncode == 0, error + else: + assert ( + process.returncode != 0 + and f"Assertion Error: {match}" in error + and "CUDA error:" in error + ), error diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index 0ed217b43a4..5a924bb30e8 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -4,6 +4,7 @@ fused_pad_token_groups_cuda, # noqa: F401 fused_unpad_token_groups_cuda, # noqa: F401 mx_block_rearrange_2d_M_groups_cuda, # noqa: F401 + mxfp8_quantize_2d_1x32_32x1_cutedsl, # noqa: F401 mxfp8_quantize_2d_1x32_cutedsl, # noqa: F401 mxfp8_quantize_2d_1x32_flydsl, # noqa: F401 mxfp8_quantize_2d_32x1_cutedsl, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py b/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py index 830e450a834..165be896ba2 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py @@ -340,7 +340,8 @@ def compute_scale_from_amax( return compute_scale_floor(amax) @cute.jit - def validate_group_sizes(offs: cute.Tensor): + def validate_group_sizes(offs: cute.Tensor, max_offset: cutlass.Constexpr = None): + """Check aligned, nondecreasing endpoints and an optional row bound.""" # Only first thread validates to avoid redundant work num_groups = offs.shape[0] @@ -351,6 +352,11 @@ def validate_group_sizes(offs: cute.Tensor): first_group_size % 128 == 0, "Group sizes must be multiples of 128", ) + cute.testing.assert_(first_group_size >= 0, "Offsets must be nonnegative") + if cutlass.const_expr(max_offset is not None): + cute.testing.assert_( + first_group_size <= max_offset, "Offset exceeds input rows" + ) # Validate subsequent groups for i in range(1, num_groups): @@ -361,3 +367,8 @@ def validate_group_sizes(offs: cute.Tensor): group_size % 128 == 0, "Group sizes must be multiples of 128", ) + cute.testing.assert_(curr_end >= prev_end, "Offsets must be nondecreasing") + if cutlass.const_expr(max_offset is not None): + cute.testing.assert_( + curr_end <= max_offset, "Offset exceeds input rows" + ) diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_quantize_2d_1x32_32x1_fused.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_quantize_2d_1x32_32x1_fused.py new file mode 100644 index 00000000000..fe29fa8968b --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_quantize_2d_1x32_32x1_fused.py @@ -0,0 +1,700 @@ +# 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. + +"""Fused 1x32 + 32x1 MXFP8 quantization -- one read of x, four outputs. + + x (total_M, N) bf16 + optional offs (G,) int32, group sizes % 128 == 0 + -> q_row (total_M, N) e4m3 [1x32 blocks along N] + s_row (total_M, N//32) uint8, M-groups blocked + q_col (N, total_M) e4m3 [32x1 blocks along M] + s_col (N, total_M//32 + G*4) uint8, K-groups blocked + +With offs=None, the same kernel uses dense placement, skips offset validation +and search, and emits compact s_col (N, total_M//32). Tensor offsets retain +grouped behavior; an empty offset tensor covers no rows rather than selecting +dense mode. + +LIMITATIONS +----------- +Hardware and runtime, enforced by the public wrapper in quant.py rather than +here: SM 10.x, CUDA 12.8+, and the CuTeDSL runtime (`nvidia-cutlass-dsl`, +`apache-tvm-ffi`). Without them this module raises NotImplementedError. + +Input, all asserted in `quantize_fused`: + + * bf16 only. + * x must be 2D, row-major contiguous, and 16-byte aligned. + * Both axes are positive multiples of 128. These are the tile dimensions, not a + convenience -- one CTA owns a whole (128, 128) tile and neither edge is + predicated. + * When supplied, offs must be 1D int32, contiguous, on x.device, holding + nondecreasing group END offsets in [0, total_M]. + * block_size 32 and scaling_mode "rceil" only. + +Every group size (consecutive difference in offs) must also be a multiple of +128; the s_col blocked layout addresses each group's slab with a stride +derived from that group's size, so a group that violates this would misplace +scales silently instead of failing. Checked with a device-side assert (thread +0 of each CTA, mirroring `cutedsl_quantize_2d_1x32.py` / +`cutedsl_quantize_2d_32x1.py`), not on the host, so it still surfaces as an +async CUDA error rather than a Python exception at the call site. +`offs[-1] < total_M` is allowed and is in fact the normal case; see +`quantize_fused` for what happens to that tail. + +DESIGN +------ +Parallelization. One CTA per (128, 128) tile of x: the grid is (N/128, M/128), +blockIdx.x selects the column slice and blockIdx.y the row slice. A CTA produces +everything its tile owns -- both quantized outputs and both sets of scales -- so +the two casts share one global read and CTAs never communicate. Block size is +picked in `_launch`: 512 threads when the whole grid is at most SMALL_GRID_CTAS +tiles, 256 otherwise, so a small grid still keeps enough threads in flight. +`threads` is a constexpr, and the two sizes are separately coded paths through +the kernel. + +Staging. The tile is copied global->shared once with `cp.async`, 16 bytes per +thread per issue (4 issues at 512 threads, 8 at 256), then a single +`cp_async_wait_group(0)` and one barrier. SMEM holds the tile as u32 words, two +bf16 per word, in four 32-row groups. The group stride is 2056 words rather than +2048; that 8-word gap is what keeps the column pass free of bank conflicts. +Three views (`sW`, `sC256`, `sC512`) alias the same buffer -- the rowwise pass +reads it row-major, the colwise pass reads the same bytes strided down M, and +neither goes back to HBM. + +Rowwise pass, 1x32 along N. Four consecutive lanes cooperate on one 32-element +block, 8 elements each. A lane folds its 8 with `max.u16x2` over sign-masked +bf16 -- non-negative floats order identically as unsigned integers, so that is +two elements per instruction -- and two butterfly shuffles finish the reduction +across the quad. `_scale` turns the amax into an E8M0 exponent plus a +reciprocal, `cvt.rn.satfinite.e4m3x2.bf16x2` scales and converts two elements at +a time, and the lead lane of each quad keeps the scale byte. + +Colwise pass, 32x1 along M. The staged tile is re-read under a fresh +thread->data map, again at full CTA width, so neither pass strands the other. A +lane pair covers one 32-row block: each lane reduces 16 rows, and one butterfly +shuffle joins the halves. Since a u32 word holds two adjacent columns, the pair +produces two neighbouring 32x1 blocks at once -- hence two exponents per +iteration. Both passes use the native NOSAT scale convention: Inf and NaN +poison the block with scale byte 255 and NaN quantized data. + +Write-back. q_row and q_col go straight from registers to HBM. q_col is a +logical transpose and would scatter, so `prmt.b32` byte permutes de-interleave +the two columns first, leaving each store 16 contiguous bytes down M and fully +coalesced. Scales take a different route: each pass parks its bytes in a +512-byte SMEM buffer, and after a barrier 128 threads flush s_row while another +128 flush s_col, one coalesced store each. s_row lands at a flat offset -- +128-aligned groups make its blocked layout the identity. s_col is the awkward +one: grouped mode scans `offs` for the group owning this tile and places its 512 bytes at +that group's dynamic stride, or zeroes those columns when no group owns the tile. +CTAs in the first row additionally clear the G*4-column over-allocation slab at +the tail of s_col while the rest of the grid is still computing. Dense mode +uses static placement and has neither the scan nor the slack clear. + +COMPILATION +----------- +Compiled per concrete (total_M, N, G) and dense/grouped mode -- these are constant-folded into the +kernel and the result is cached in `_cache`, so a caller whose total_M varies +step to step pays a JIT compile on each new shape. +""" + +from typing import Optional, Tuple + +import torch + +from torchao.utils import ceil_div + +from .cute_utils import ( + _cutedsl_runtime_available, + _missing_cutedsl_runtime_packages, +) + +_RUNTIME_OK = _cutedsl_runtime_available() + + +if _RUNTIME_OK: + import cuda.bindings.driver as _cuda_drv + import cutlass + import cutlass.cute as cute + import cutlass.utils as utils + import torch + from cutlass import Int32, Uint32 + from cutlass._mlir.dialects import llvm + from cutlass.cute.nvgpu import cpasync + from cutlass.cute.runtime import from_dlpack, make_fake_stream + from cutlass.cutlass_dsl import T, dsl_user_op + + from .cute_utils import validate_group_sizes + + TM = 128 + TN = 128 + SMALL_GRID_CTAS = 512 + + @dsl_user_op + def amax_acc(acc: Uint32, a: Uint32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + Uint32(acc).ir_value(loc=loc, ip=ip), + Uint32(a).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b32 t;\n" + " and.b32 t, $2, 0x7fff7fff;\n" + " max.u16x2 $0, $1, t; }", + "=r,r,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op + def cvt2(a: Uint32, b: Uint32, invx2: Uint32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + Uint32(a).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + Uint32(invx2).ir_value(loc=loc, ip=ip), + ], + "{ .reg .b16 l, h; .reg .b32 t;\n" + " mul.rn.bf16x2 t, $1, $3;\n" + " cvt.rn.satfinite.e4m3x2.bf16x2 l, t;\n" + " mul.rn.bf16x2 t, $2, $3;\n" + " cvt.rn.satfinite.e4m3x2.bf16x2 h, t;\n" + " mov.b32 $0, {l, h}; }", + "=r,r,r,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + ) + + def _prmt_op(sel): + @dsl_user_op + def f(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + Uint32(a).ir_value(loc=loc, ip=ip), + Uint32(b).ir_value(loc=loc, ip=ip), + ], + "prmt.b32 $0, $1, $2, %d;" % sel, + "=r,r,r", + has_side_effects=False, + asm_dialect=0, + loc=loc, + ip=ip, + ) + ) + + return f + + prmt_even = _prmt_op(0x6420) + prmt_odd = _prmt_op(0x7531) + + @cute.jit + def _umax(a: Uint32, b: Uint32) -> Uint32: + return a if a > b else b + + @cute.jit + def _scale(bits: Uint32): + """Return the RCEIL exponent and reciprocal bits, poisoning nonfinite amax.""" + # Exact RCEIL for bf16 amax represented in the high half of an f32 word. + w = bits + Uint32(0x1FFFFF) + es = Int32(w >> Uint32(23)) - Int32(8) + # cvt.rp.satfinite maps a subnormal descale to E8M0 byte 0. Since + # descale=amax/448, the exact BF16 cutoff is 1.75*2^-119 (0x0460). + e = es if bits >= Uint32(0x04600000) else Int32(0) + ib = Uint32((Int32(254) - e) << Int32(23)) + # Inf/NaN. The amax above is an UNSIGNED max over sign-masked bf16 bit + # patterns, which orders Inf and NaN as merely very large integers, so + # without this a poisoned block would quantize to a healthy-looking 448. + # + # Match NOSAT in both the rowwise CuTeDSL and native CUDA references. + if bits >= Uint32(0x7F800000): + e = Int32(255) + ib = Uint32(0x7FC00000) + return e, ib + + @cute.kernel + def _kernel( + mX: cute.Tensor, + mRQ: cute.Tensor, + mCQ: cute.Tensor, + mRS: cute.Tensor, + mCS: cute.Tensor, + mOffs: Optional[cute.Tensor], + kblocks: cutlass.Constexpr, + mblocks: cutlass.Constexpr, + groups: cutlass.Constexpr, + threads: cutlass.Constexpr, + ): + """Quantize one aligned 128x128 tile into both data and blocked-scale layouts.""" + tid, _, _ = cute.arch.thread_idx() + kt, mt, _ = cute.arch.block_idx() + + # Validate group sizes are multiples of 128; a group that violates + # this would otherwise silently misplace s_col scales (see module + # docstring). + if cutlass.const_expr(mOffs is not None): + if tid == 0: + validate_group_sizes(mOffs, max_offset=mblocks * TM) + + smem = utils.SmemAllocator() + # The 8-u32 gap between 32-row groups removes column-pass bank conflicts. + sbuf = smem.allocate_array(cutlass.Uint32, 4 * 2056, byte_alignment=16) + srow = smem.allocate_array(cutlass.Uint32, 128, byte_alignment=16) + scol = smem.allocate_array(cutlass.Uint32, 128, byte_alignment=16) + sW = cute.make_tensor( + sbuf, cute.make_layout((4, 2, 16, 16, 4), stride=(2056, 1028, 64, 4, 1)) + ) + sC256 = cute.make_tensor( + sbuf, cute.make_layout((4, 2, 16, 32, 2), stride=(2056, 1028, 64, 2, 1)) + ) + sC512 = cute.make_tensor( + sbuf, cute.make_layout((4, 2, 16, 64, 1), stride=(2056, 1028, 64, 1, 1)) + ) + sR32 = cute.make_tensor(srow, cute.make_layout(128)) + sC32 = cute.make_tensor(scol, cute.make_layout(128)) + sR8 = cute.make_tensor( + cute.recast_ptr(srow, dtype=cutlass.Uint8), cute.make_layout(512) + ) + sS8 = cute.make_tensor( + cute.recast_ptr(scol, dtype=cutlass.Uint8), cute.make_layout(512) + ) + + lane16 = tid % 16 + rbase = tid // 16 + kchunk = kt * 16 + lane16 + quad = lane16 // 4 + is_lead = (lane16 % 4) == 0 + sr_base = rbase * 16 + quad + + g2s = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.GLOBAL), + cutlass.Uint32, + num_bits_per_copy=128, + ) + if cutlass.const_expr(threads == 256): + for r in cutlass.range_constexpr(8): + lr = r * 16 + rbase + cute.copy( + g2s, + mX[(mt * TM + lr, kchunk, None)], + sW[(lr // 32, (lr % 32) // 16, lr % 16, lane16, None)], + ) + else: + for r in cutlass.range_constexpr(4): + lr = r * 32 + rbase + cute.copy( + g2s, + mX[(mt * TM + lr, kchunk, None)], + sW[(lr // 32, (lr % 32) // 16, lr % 16, lane16, None)], + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + frg = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), Uint32) + oq = cute.make_rmem_tensor(cute.make_layout((2,), stride=(1,)), Uint32) + if cutlass.const_expr(threads == 256): + for r in cutlass.range_constexpr(8): + lr = r * 16 + rbase + cute.autovec_copy( + sW[(lr // 32, (lr % 32) // 16, lr % 16, lane16, None)], frg + ) + acc = amax_acc(Uint32(0), frg[0]) + acc = amax_acc(acc, frg[1]) + acc = amax_acc(acc, frg[2]) + acc = amax_acc(acc, frg[3]) + v = _umax(acc << Uint32(16), acc & Uint32(0xFFFF0000)) + v = _umax(v, Uint32(cute.arch.shuffle_sync_bfly(v, 1))) + v = _umax(v, Uint32(cute.arch.shuffle_sync_bfly(v, 2))) + e, ib = _scale(v) + invx2 = ib | (ib >> Uint32(16)) + oq[0] = cvt2(frg[0], frg[1], invx2) + oq[1] = cvt2(frg[2], frg[3], invx2) + cute.autovec_copy(oq, mRQ[(mt * TM + lr, kchunk, None)]) + if is_lead: + sR8[sr_base + ((r % 2) * 256 + (r // 2) * 4)] = cutlass.Uint8(e) + else: + for r in cutlass.range_constexpr(4): + lr = r * 32 + rbase + cute.autovec_copy( + sW[(lr // 32, (lr % 32) // 16, lr % 16, lane16, None)], frg + ) + acc = amax_acc(Uint32(0), frg[0]) + acc = amax_acc(acc, frg[1]) + acc = amax_acc(acc, frg[2]) + acc = amax_acc(acc, frg[3]) + v = _umax(acc << Uint32(16), acc & Uint32(0xFFFF0000)) + v = _umax(v, Uint32(cute.arch.shuffle_sync_bfly(v, 1))) + v = _umax(v, Uint32(cute.arch.shuffle_sync_bfly(v, 2))) + e, ib = _scale(v) + invx2 = ib | (ib >> Uint32(16)) + oq[0] = cvt2(frg[0], frg[1], invx2) + oq[1] = cvt2(frg[2], frg[3], invx2) + cute.autovec_copy(oq, mRQ[(mt * TM + lr, kchunk, None)]) + if is_lead: + sR8[sr_base + r * 4] = cutlass.Uint8(e) + + # Each lane pair owns rows of one 32x1 group; byte permutes produce + # fully coalesced q_col stores despite the logical transpose. + mg = (tid % 8) // 2 + half = tid % 2 + jq = tid // 8 + if cutlass.const_expr(threads == 256): + cf = cute.make_rmem_tensor(cute.make_layout((16, 2), stride=(2, 1)), Uint32) + for i in cutlass.range_constexpr(16): + cute.autovec_copy(sC256[(mg, half, i, jq, None)], cf[(i, None)]) + + a0 = amax_acc(Uint32(0), cf[0, 0]) + a1 = amax_acc(Uint32(0), cf[0, 1]) + b0 = amax_acc(Uint32(0), cf[1, 0]) + b1 = amax_acc(Uint32(0), cf[1, 1]) + for i in cutlass.range_constexpr(1, 8): + a0 = amax_acc(a0, cf[2 * i, 0]) + a1 = amax_acc(a1, cf[2 * i, 1]) + b0 = amax_acc(b0, cf[2 * i + 1, 0]) + b1 = amax_acc(b1, cf[2 * i + 1, 1]) + a0 = amax_acc(a0, b0) + a1 = amax_acc(a1, b1) + a0 = amax_acc(a0, Uint32(cute.arch.shuffle_sync_bfly(a0, 1))) + a1 = amax_acc(a1, Uint32(cute.arch.shuffle_sync_bfly(a1, 1))) + + oe = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), Uint32) + oo = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), Uint32) + mchunk = mt * 8 + mg * 2 + half + ss_base = (jq % 8) * 64 + (jq // 8) * 4 + mg + for w in cutlass.range_constexpr(2): + pk = a0 if w == 0 else a1 + e0, ib0 = _scale(pk << Uint32(16)) + e1, ib1 = _scale(pk & Uint32(0xFFFF0000)) + invw = (ib0 >> Uint32(16)) | (ib1 & Uint32(0xFFFF0000)) + for g in cutlass.range_constexpr(4): + ua = cvt2(cf[4 * g, w], cf[4 * g + 1, w], invw) + ub = cvt2(cf[4 * g + 2, w], cf[4 * g + 3, w], invw) + oe[g] = prmt_even(ua, ub) + oo[g] = prmt_odd(ua, ub) + cute.autovec_copy(oe, mCQ[(kt * TN + jq * 4 + 2 * w, mchunk, None)]) + cute.autovec_copy(oo, mCQ[(kt * TN + jq * 4 + 2 * w + 1, mchunk, None)]) + if half == 0: + sS8[ss_base + (2 * w) * 16] = cutlass.Uint8(e0) + sS8[ss_base + (2 * w + 1) * 16] = cutlass.Uint8(e1) + else: + cf = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), Uint32) + for i in cutlass.range_constexpr(16): + cf[i] = sC512[(mg, half, i, jq, 0)] + + acc = amax_acc(Uint32(0), cf[0]) + for i in cutlass.range_constexpr(1, 16): + acc = amax_acc(acc, cf[i]) + acc = amax_acc(acc, Uint32(cute.arch.shuffle_sync_bfly(acc, 1))) + + oe = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), Uint32) + oo = cute.make_rmem_tensor(cute.make_layout((4,), stride=(1,)), Uint32) + mchunk = mt * 8 + mg * 2 + half + ss_base = (jq % 16) * 32 + (jq // 16) * 4 + mg + e0, ib0 = _scale(acc << Uint32(16)) + e1, ib1 = _scale(acc & Uint32(0xFFFF0000)) + invw = (ib0 >> Uint32(16)) | (ib1 & Uint32(0xFFFF0000)) + for g in cutlass.range_constexpr(4): + ua = cvt2(cf[4 * g], cf[4 * g + 1], invw) + ub = cvt2(cf[4 * g + 2], cf[4 * g + 3], invw) + oe[g] = prmt_even(ua, ub) + oo[g] = prmt_odd(ua, ub) + cute.autovec_copy(oe, mCQ[(kt * TN + jq * 2, mchunk, None)]) + cute.autovec_copy(oo, mCQ[(kt * TN + jq * 2 + 1, mchunk, None)]) + if half == 0: + sS8[ss_base] = cutlass.Uint8(e0) + sS8[ss_base + 16] = cutlass.Uint8(e1) + + cute.arch.barrier() + if tid < 128: + mRS[(mt * kblocks + kt) * 128 + tid] = sR32[tid] + elif tid < 256: + # Map this 512-byte col-scale tile into its dynamic group-major slab. + m0 = mt * TM + group_start = Int32(0) + group_cols = Int32(mblocks * 4) + found = Int32(1) + if cutlass.const_expr(mOffs is not None): + group_cols = Int32(0) + found = Int32(0) + prev = Int32(0) + for g in cutlass.range_constexpr(groups): + end = Int32(mOffs[g]) + if found == 0: + if m0 < end: + group_start = prev + group_cols = (end - prev) // 32 + found = Int32(1) + prev = end + if found == 0: + # UNCOVERED tile: m0 >= offs[-1], so no group owns it. A dropless + # MoE dispatcher sizes grad_out for worst-case padding and fills + # only a prefix, so this is the common case. torchao's K-groups + # swizzle never writes these scale columns and its buffer is + # new_zeros, so the contract here is ZERO. + # + # Falling through to the placement math below is doubly wrong: + # group_start=0 aims the store into GROUP 0's slab, and + # group_cols=0 kills the kt term so all kblocks CTAs race on one + # 512B line. That is what produced grad_norm=inf in TorchTitan. + # + # This tile owns exactly scale columns [m0//32, m0//32+4) == 4*N + # bytes; the kblocks n-tiles split them 512B each, so the fill is + # perfectly partitioned. Uncovered tiles together cover + # [offs[-1]//32, total_M//32) exactly, since every group size -- + # hence offs[-1] -- is a multiple of 128. + zero_base = (m0 // 32) * kblocks * TN + kt * 512 + mCS[zero_base // 4 + tid - 128] = Uint32(0) + else: + local_cb = (m0 - group_start) // TM + byte_base = ( + (group_start // 32) * kblocks * TN + + kt * (group_cols // 4) * 512 + + local_cb * 512 + ) + mCS[byte_base // 4 + tid - 128] = sC32[tid - 128] + + # The first real CTA row clears the disjoint over-allocation tail while + # the rest of the grid is still computing, avoiding G extra CTA rows. + if cutlass.const_expr(mOffs is not None): + if mt == 0: + tail_i = kt * threads + tid + tail_words = groups * kblocks * TN + tail_base = mblocks * kblocks * 128 + tail_stride = kblocks * threads + while tail_i < tail_words: + mCS[tail_base + tail_i] = Uint32(0) + tail_i += tail_stride + + @cute.jit + def _launch(tX, tOffs, tRQ, tRS, tCQ, tCS, stream): + """Launch the tile grid, specializing None offsets as compact dense placement.""" + M = cutlass.const_expr(tX.shape[0]) + N = cutlass.const_expr(tX.shape[1]) + if cutlass.const_expr(tOffs is None): + G = 0 + else: + G = cutlass.const_expr(tOffs.shape[0]) + u32 = Uint32 + mX = cute.make_tensor( + cute.recast_ptr(tX.iterator, dtype=u32), + cute.make_layout((M, N // 8, 4), stride=(N // 2, 4, 1)), + ) + mRQ = cute.make_tensor( + cute.recast_ptr(tRQ.iterator, dtype=u32), + cute.make_layout((M, N // 8, 2), stride=(N // 4, 2, 1)), + ) + mCQ = cute.make_tensor( + cute.recast_ptr(tCQ.iterator, dtype=u32), + cute.make_layout((N, M // 16, 4), stride=(M // 4, 4, 1)), + ) + mRS = cute.make_tensor( + cute.recast_ptr(tRS.iterator, dtype=u32), cute.make_layout(M * N // 128) + ) + mCS = cute.make_tensor( + cute.recast_ptr(tCS.iterator, dtype=u32), + cute.make_layout(N * (M // 32 + G * 4) // 4), + ) + threads = cutlass.const_expr( + 512 if (M // TM) * (N // TN) <= SMALL_GRID_CTAS else 256 + ) + _kernel(mX, mRQ, mCQ, mRS, mCS, tOffs, N // TN, M // TM, G, threads).launch( + grid=[N // TN, M // TM, 1], block=[threads, 1, 1], stream=stream + ) + + _cache = {} + + def _compile(x, offs, q_row, s_row, q_col, s_col): + """Compile shape/mode-specific code with active grouped endpoint checks.""" + + def w(t): + """Wrap aligned tensor views while preserving literal None offsets.""" + return ( + from_dlpack(t, assumed_align=16, enable_tvm_ffi=True) + if t is not None + else None + ) + + return cute.compile( + _launch, + w(x), + w(offs), + w(q_row), + w(s_row), + w(q_col), + w(s_col), + make_fake_stream(), + # CuTeDSL strips cf.assert by default; grouped placement requires + # these endpoint checks to remain active in the device program. + options="--enable-tvm-ffi --enable-assertions" + if offs is not None + else "--enable-tvm-ffi", + ) + + def run(x, offs, q_row, s_row, q_col, s_col): + """Launch cached code into validated destinations on x.device's current stream.""" + key = (x.shape, offs.shape if offs is not None else None) + fn = _cache.get(key) + if fn is None: + fn = _compile(x, offs, q_row, s_row, q_col, s_col) + _cache[key] = fn + fn( + x, + offs, + q_row, + s_row, + q_col, + s_col, + _cuda_drv.CUstream(int(torch.cuda.current_stream(x.device).cuda_stream)), + ) + + +def s_col_shape(total_M: int, N: int, num_groups: int) -> Tuple[int, int]: + """Shape of the K-groups blocked colwise scale buffer. + + Grouped mode preserves the legacy G*4-column zero-filled slack. Aligned + groups consume only N*total_M/32 bytes. Dense mode uses num_groups=0 and + has no slack. + """ + return ceil_div(N, 128) * 128, total_M // 32 + num_groups * 4 + + +def _aligned16(t: torch.Tensor) -> torch.Tensor: + """`run` hands tensors to `from_dlpack(..., assumed_align=16)`. + + Torch's caching allocator returns 512-byte-aligned blocks, so this only ever + bites on a view with an odd storage offset. Clone only misaligned offsets; + the large input stays zero-copy and is checked instead. + """ + if t.data_ptr() % 16 == 0: + return t + return t.clone() + + +def quantize_fused( + x: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", + offs: Optional[torch.Tensor] = None, + *, + q_row_out: Optional[torch.Tensor] = None, + s_row_out: Optional[torch.Tensor] = None, + q_col_out: Optional[torch.Tensor] = None, + s_col_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Both MXFP8 casts of grad_out from one pass over `x`. + + Args: + x: (total_M, N) bfloat16, row-major, with positive 128-aligned axes. + block_size: MX block size, 32 only. + scaling_mode: "rceil" only (what TorchTitan runs). + offs: None selects dense placement with compact scales. Otherwise, + (G,) int32 CUDA tensor of group END offsets along total_M, + nondecreasing and bounded by total_M. Every group size is a + multiple of 128, checked with a device assertion, which raises an async + CUDA error rather than failing here. `offs[-1]` MAY be less than + total_M -- a dropless MoE dispatcher sizes this buffer for + worst-case padding and fills only a prefix, so an uncovered tail + is the norm. Rows >= offs[-1] are still quantized into + q_row/q_col/s_row (torchao's 1x32 kernel does the same), but + their s_col scale columns are ZEROED to match torchao's + K-groups swizzle. + *_out: optional destination buffers (destination-passing, so a benchmark + harness can hoist allocation out of the timed region). None of them + need to be zeroed -- every byte of every output is written here. + Buffers must not overlap x or each other; disjoint views are allowed. + + Returns: + q_row (total_M, N) e4m3, s_row (total_M, N//32) uint8, + q_col (N, total_M) e4m3, s_col (ceil(N,128)*128, total_M//32+G*4) uint8 + """ + if not _RUNTIME_OK: + raise NotImplementedError( + "missing CuTeDSL runtime package(s): " + + ", ".join(_missing_cutedsl_runtime_packages()) + ) + + assert x.is_cuda, "x must be a CUDA tensor" + assert x.dim() == 2 and x.is_contiguous(), "x must be 2D row-major contiguous" + assert x.dtype == torch.bfloat16, "x must be bf16" + assert block_size == 32, "only block_size=32 is supported" + assert scaling_mode == "rceil", "only rceil is supported" + if offs is not None: + assert offs.device == x.device, "offs must be on x.device" + assert offs.dtype == torch.int32 and offs.dim() == 1, "offs must be 1D int32" + assert offs.is_contiguous(), "offs must be contiguous" + + total_M, N = x.shape + assert total_M > 0 and N > 0, "input axes must be positive" + assert N % 128 == 0, "N must be a multiple of 128" + assert total_M % 128 == 0, "total_M must be a multiple of 128" + assert x.data_ptr() % 16 == 0, "x must be 16-byte aligned" + G = offs.numel() if offs is not None else 0 + + s_col_rows, s_col_cols = s_col_shape(total_M, N, G) + has_supplied_output = any( + output is not None for output in (q_row_out, s_row_out, q_col_out, s_col_out) + ) + + if q_row_out is None: + q_row_out = torch.empty( + (total_M, N), device=x.device, dtype=torch.float8_e4m3fn + ) + if s_row_out is None: + s_row_out = torch.empty((total_M, N // 32), device=x.device, dtype=torch.uint8) + if q_col_out is None: + q_col_out = torch.empty( + (N, total_M), device=x.device, dtype=torch.float8_e4m3fn + ) + if s_col_out is None: + s_col_out = torch.empty( + (s_col_rows, s_col_cols), device=x.device, dtype=torch.uint8 + ) + + for name, output, shape, dtype in ( + ("q_row_out", q_row_out, (total_M, N), torch.float8_e4m3fn), + ("s_row_out", s_row_out, (total_M, N // 32), torch.uint8), + ("q_col_out", q_col_out, (N, total_M), torch.float8_e4m3fn), + ("s_col_out", s_col_out, (s_col_rows, s_col_cols), torch.uint8), + ): + assert output.shape == shape and output.dtype == dtype, ( + f"{name} must have shape {shape} and dtype {dtype}" + ) + assert output.device == x.device and output.is_contiguous(), ( + f"{name} must be contiguous on x.device" + ) + assert output.data_ptr() % 16 == 0, f"{name} must be 16-byte aligned" + + if has_supplied_output: + # Contiguous tensors occupy exact byte intervals, including storage-offset + # views. Disjoint slices of one allocation are valid destinations. + ranges = [] + for output in (x, q_row_out, s_row_out, q_col_out, s_col_out): + start = output.data_ptr() + end = start + output.numel() * output.element_size() + assert all(end <= left or start >= right for left, right in ranges), ( + "output buffers must not overlap input or each other" + ) + ranges.append((start, end)) + + run( + x, + _aligned16(offs) if offs is not None else None, + q_row_out, + s_row_out, + q_col_out, + s_col_out, + ) + return q_row_out, s_row_out, q_col_out, s_col_out diff --git a/torchao/prototype/moe_training/kernels/mxfp8/quant.py b/torchao/prototype/moe_training/kernels/mxfp8/quant.py index d5e42e6d74a..a621a3c29f6 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/quant.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/quant.py @@ -1098,6 +1098,36 @@ def _mxfp8_quantize_2d_32x1_cutedsl_custom_op( ) +@torch.library.custom_op( + "torchao::mxfp8_quantize_2d_1x32_32x1_cutedsl", mutates_args=() +) +def _mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( + x: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", + offs: torch.Tensor = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Return both quantized orientations with typed blocked E8M0 scales.""" + from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_quantize_2d_1x32_32x1_fused import ( + quantize_fused, + ) + + q_row, s_row, q_col, s_col = quantize_fused( + x, + block_size=block_size, + scaling_mode=scaling_mode, + offs=offs, + ) + # The kernel writes raw E8M0 bytes; hand them back typed, matching what the + # 1x32 and 32x1 ops above return. + return ( + q_row, + s_row.view(torch.float8_e8m0fnu), + q_col, + s_col.view(torch.float8_e8m0fnu), + ) + + @_mxfp8_quantize_2d_1x32_cutedsl_custom_op.register_fake def _fake_mxfp8_quantize_2d_1x32_cutedsl_custom_op( x: torch.Tensor, @@ -1144,7 +1174,7 @@ def _fake_mxfp8_quantize_2d_32x1_cutedsl_custom_op( # For 32x1 scaling, output data has same dimensions as input (no padding) q_data = torch.empty_strided( (m, k), - (k, 1), + (1, m), device=x.device, dtype=torch.float8_e4m3fn, ) @@ -1153,13 +1183,69 @@ def _fake_mxfp8_quantize_2d_32x1_cutedsl_custom_op( m_blocks = m // block_size padded_scale_rows = ceil_div(m_blocks, 4) * 4 padded_scale_cols = ceil_div(k, 128) * 128 + scale_shape = ( + (padded_scale_cols * padded_scale_rows,) + if blocked_scale_output + else (k, m_blocks) + ) scales = x.new_empty( - (padded_scale_cols * padded_scale_rows,), + scale_shape, dtype=torch.float8_e8m0fnu, ) return q_data, scales +@_mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op.register_fake +def _fake_mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( + x: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", + offs: torch.Tensor = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Match dense/grouped runtime layouts without reading offset values.""" + assert x.ndim == 2, "input tensor must be 2D" + assert block_size == 32, "Only block_size=32 is supported" + assert scaling_mode == "rceil", "Only rceil is supported" + assert x.dtype == torch.bfloat16 and x.is_contiguous(), ( + "input must be contiguous bf16" + ) + total_m, n = x.shape + assert total_m > 0 and n > 0 and total_m % 128 == n % 128 == 0, ( + "input axes must be positive multiples of 128" + ) + if offs is not None: + assert offs.ndim == 1 and offs.dtype == torch.int32, "offs must be 1D int32" + assert offs.device == x.device and offs.is_contiguous(), ( + "offs must be contiguous on x.device" + ) + num_groups = offs.shape[0] if offs is not None else 0 + + q_row = torch.empty_strided( + (total_m, n), + (n, 1), + device=x.device, + dtype=torch.float8_e4m3fn, + ) + # Rowwise groups run along ROWS and every group base is a multiple of 128, so + # the blocked layout needs no per-group padding and this is exactly + # to_blocked() of the whole (total_M, N // 32) array. + s_row = x.new_empty((total_m, n // block_size), dtype=torch.float8_e8m0fnu) + + q_col = torch.empty_strided( + (n, total_m), + (total_m, 1), + device=x.device, + dtype=torch.float8_e4m3fn, + ) + # Colwise groups run along COLUMNS, so each is blocked independently and the + # group structure survives as per-group padding. See s_col_shape(). + s_col_rows = ceil_div(n, 128) * 128 + s_col_cols = total_m // block_size + num_groups * 4 + s_col = x.new_empty((s_col_rows, s_col_cols), dtype=torch.float8_e8m0fnu) + + return q_row, s_row, q_col, s_col + + if _mxfp8_cutedsl_kernels_available: @register_sharding(torch.ops.torchao.mxfp8_quantize_2d_1x32_cutedsl.default) @@ -1576,6 +1662,69 @@ def mxfp8_quantize_2d_32x1_cutedsl( return qdata, scales +def mxfp8_quantize_2d_1x32_32x1_cutedsl( + x: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", + offs: torch.Tensor = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Quantize a 2D tensor of shape (total_M, N) along BOTH axes in one pass, + reading the input from HBM once. + + Emits what a grouped-GEMM backward needs from grad_output: the 1x32 rowwise + form along N (for dgrad) and the 32x1 colwise form along total_M (for + wgrad), each with its scales already in the blocked layout its consumer + expects. Bit-exact with mxfp8_quantize_2d_1x32_cutedsl composed with the + colwise cast and a standalone K-groups swizzle, degenerate blocks included, + while moving 4.0625 bytes per input element instead of 6.125 for that older + three-kernel composition. The existing dense CUDA API already fuses both + casts, but requires separate blocked-scale rearrangements. + AO grouped-GEMM wiring is separate from this kernel API; dense consumers + such as TorchTitan can use it directly. + + Args: + x: Input tensor of shape (total_M, N), bfloat16, row-major contiguous. + Both dimensions must be positive multiples of 128. + block_size: Block size for quantization (only 32 supported) + scaling_mode: Scaling mode (only "rceil" supported) + offs: None selects a dense matrix without offset processing or scale + slack. Otherwise, group end offsets along total_M, int32, nondecreasing, + within [0, total_M], every group + size a multiple of 128 -- checked with a device-side assert, which + surfaces as an async CUDA error rather than a Python exception + here. offs[-1] MAY be less than total_M: a dropless MoE + dispatcher sizes this buffer for worst-case padding and fills only + a prefix, so an uncovered tail is the norm. Rows past offs[-1] are + still quantized into q_row/q_col/s_row, but their s_col scale + columns are zeroed, matching the K-groups swizzle. + An empty offset tensor covers no rows; it is not dense mode. + + Returns: + q_row (total_M, N) row-major, s_row (total_M, N // 32) M-groups blocked, + q_col (N, total_M) row-major, s_col (N, total_M // 32 + 4*G) + K-groups blocked; G=0 in dense mode. Scales use E8M0 dtype. + """ + if not _mxfp8_cutedsl_kernels_available: + missing_packages = _missing_cutedsl_runtime_packages() + if missing_packages: + missing = ", ".join(missing_packages) + raise NotImplementedError( + "mxfp8_quantize_2d_1x32_32x1 requires additional Python " + f"runtime package(s): {missing}. Please install " + "`nvidia-cutlass-dsl` and `apache-tvm-ffi`." + ) + raise NotImplementedError( + "mxfp8_quantize_2d_1x32_32x1 requires CUDA, SM 10.x, and CUDA 12.8+." + ) + return _mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( + x, + block_size=block_size, + scaling_mode=scaling_mode, + offs=offs, + ) + + # ============================================================================= # FlyDSL MXFP8 quantize kernels (AMD CDNA3+ via the FlyDSL stack). #