From 50001a42bde69205de9365f8623388ee9b3e3878 Mon Sep 17 00:00:00 2001 From: protonu Date: Mon, 31 Aug 2026 19:24:49 +0000 Subject: [PATCH 1/2] [moe_training] Add a fused 1x32 + 32x1 MXFP8 quantization kernel Quantizing a 2D tensor along both axes currently takes three kernels and two full reads of the same tensor: the 1x32 CuTeDSL cast along N, the 32x1 CUDA cast along M, and a standalone Triton K-groups swizzle of the colwise scales. Add a CuTeDSL kernel that reads the input once and emits all four outputs, with both scale tensors already in the blocked layouts their consumers expect. Byte accounting drops from 6.125 to 4.0625 bytes per input element. The kernel puts one CTA on each (128, 128) input tile, stages it into SMEM with cp.async, and runs the rowwise and colwise passes sequentially at full CTA width. The arithmetic is packed 16-bit throughout -- max.u16x2 on sign-masked bf16 for amax, cvt.rn.satfinite.e4m3x2.bf16x2 for the cast -- and prmt.b32 byte permutes keep the logically transposed colwise store fully coalesced. bf16 and rceil only; both dimensions and every group size must be multiples of 128. Exposed as mxfp8_quantize_2d_1x32_32x1_cutedsl, registered as a custom op with a fake impl, alongside the two kernels it composes. Uncovered tail. offs[-1] < total_M is supported and is the norm rather than a corner case: a dropless MoE dispatcher sizes its buffer for worst-case padding and fills only a prefix. Rows past offs[-1] are still quantized, but their colwise scale bytes are zeroed to match the K-groups swizzle, which iterates groups and returns new_zeros. Falling through to the placement math instead aims the store into group 0's slab and races every n-tile CTA on one 512-byte line. Bug-compatible where the two kernels it replaces disagree with each other, since the point is to be a drop-in: * A block whose amax is subnormal enough to reach E8M0 scale byte 0 gets a 2**127 multiplier on both outputs, matching the 254 - scale_biased reciprocal both use. * Inf: the rowwise kernel converts NOSAT, so Inf poisons the block to scale 255 and all-NaN data; the colwise CUDA kernel converts SATFINITE, so Inf clamps to scale 254 with finite data. Both are reproduced. The one deliberate divergence is NaN on the colwise output. The CUDA kernel drops NaN from its amax entirely -- an all-NaN block still gets the scale its finite elements imply -- and reproducing that would need a per-element NaN test in the hottest loop in order to imitate a behavior that hides NaN corruption rather than surfacing it. The fused kernel poisons instead, and the test asserts the difference so it fails loudly if either side changes. Tests compare raw bytes at atol=0/rtol=0 against the three-kernel composition, not against to_mx, because being a drop-in for that composition is the contract. They cover shapes and group counts, uncovered tails including the zero-fill, degenerate and mixed-magnitude blocks, Inf/NaN, dirty destination buffers so a never-written region cannot hide behind a zeroed page, determinism across repeated launches, and input validation. The rowwise half is additionally checked against the shared MXFP8 semantic contract. The benchmark follows the existing 1x32 and 32x1 scripts, but times both sides under CUDA-graph replay rather than eagerly. torchao's triton_op dispatch costs on the order of 100us per call regardless of problem size and the baseline pays it three times, which timed eagerly swamps every shape below roughly 0.5 GB and reports 15x on small shapes -- launch overhead, not a kernel win. It also checks that the kernel does not capture as an empty graph, which is what happens when a CuTeDSL kernel ignores the caller's stream and which would read as a 100x win. Geomean is 2.6x across the swept shapes, 2.0x on the largest, at up to 95% of the measured SM bandwidth roof on a GB200 at 2062 MHz. Co-Authored-By: Claude Opus 5 (1M context) --- ...nch_cutedsl_quantize_2d_1x32_32x1_fused.py | 303 ++++++++ test/prototype/moe_training/test_kernels.py | 419 ++++++++++- .../moe_training/kernels/mxfp8/__init__.py | 1 + .../cutedsl_quantize_2d_1x32_32x1_fused.py | 666 ++++++++++++++++++ .../moe_training/kernels/mxfp8/quant.py | 122 ++++ 5 files changed, 1505 insertions(+), 6 deletions(-) create mode 100644 benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/cutedsl_quantize_2d_1x32_32x1_fused.py 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..8f089f0a7b4 --- /dev/null +++ b/benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.py @@ -0,0 +1,303 @@ +# 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 +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, + 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[:, :used], exp[:, :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, offs, block_size=BLOCK_SIZE, scaling_mode="rceil") + 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, offs, block_size=BLOCK_SIZE, scaling_mode="rceil" + ) + 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) + 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..7ad156323d1 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, @@ -513,10 +522,8 @@ def test_cuda_mx_3d_cutedsl_numerics(E, N, K, input_dtype, scaling_mode, variant y_ref = y_ref.transpose(-2, -1) s_ref = s_ref.transpose(-2, -1) s_rows, s_cols = K, N // block_size - undo_scale = ( - lambda scale: from_blocked(scale, s_rows, s_cols) - .transpose(-2, -1) - .contiguous() + undo_scale = lambda scale: ( + from_blocked(scale, s_rows, s_cols).transpose(-2, -1).contiguous() ) else: x_tiles = ( @@ -1682,3 +1689,403 @@ 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. Its contract is therefore to be a DROP-IN for +# that composition, so every test below compares raw bytes against the +# composition itself rather than against `to_mx`. That distinction matters: the +# fused kernel is deliberately bug-compatible with the colwise CUDA kernel's +# degenerate-block convention, so a `to_mx` reference would disagree on inputs +# that no random tensor can produce. +# ============================================================================= + +_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_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): + """Bytewise, atol=0/rtol=0. A single wrong E8M0 byte is a 2x error.""" + expected = _fused_1x32_32x1_reference(x, offs) + names = ("q_row", "s_row", "q_col", "s_col") + num_groups = offs.numel() + for name, got, exp in zip(names, actual, expected): + got = got.view(torch.uint8) + if name == "s_col": + # Both sides over-allocate the colwise scale buffer by num_groups * 4 + # columns of per-group padding slack that the grouped GEMM never + # reads. Compare the region that is actually consumed. + used = exp.shape[1] - num_groups * 4 + got, exp = got[:, :used], exp[:, :used] + torch.testing.assert_close(got, exp, rtol=0, atol=0, msg=f"{name} mismatch") + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +@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 = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") + offs = generate_jagged_offs(num_groups, total_M, multiple_of=128, device="cuda").to( + torch.int32 + ) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, 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) + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +@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 = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + # The colwise scale bytes past the covered prefix must be ZEROED, matching + # triton_mx_block_rearrange_2d_K_groups, which iterates groups and returns + # new_zeros. Asserted directly because a fresh torch.empty can hide a + # never-written region whenever the allocator hands back zeroed pages. + # + # Indexed FLAT, not by column: the K-groups blocked layout stores each + # 4-scale-column block as N contiguous bytes, so scale columns + # [covered // 32, total_M // 32) are the flat byte range below and are NOT a + # 2D column slice of the (padded_N, num_scale_cols) view. + s_col_flat = out[3].view(torch.uint8).reshape(-1) + uncovered = s_col_flat[(covered // 32) * N : (total_M // 32) * N] + assert uncovered.numel() == (tail // 32) * N + assert torch.all(uncovered == 0), "uncovered s_col bytes must be zero-filled" + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +@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 = (torch.randn(total_M, N, device="cuda") * scale).to(torch.bfloat16) + offs = generate_jagged_offs(num_groups, total_M, multiple_of=128, device="cuda").to( + torch.int32 + ) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +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 = (torch.randn(total_M, N, device="cuda") * 100).to(torch.bfloat16) + x[::4] = (x[::4].float() * 1e-36).to(torch.bfloat16) + x[1::4] = (x[1::4].float() * 1e28).to(torch.bfloat16) + offs = generate_jagged_offs(num_groups, total_M, multiple_of=128, device="cuda").to( + torch.int32 + ) + out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +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 = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") + + 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 buf in bufs: + buf.view(torch.uint8).fill_(0xFF) + + quantize_fused( + x, + 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) + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_determinism(): + """Three launches on one input must agree byte for byte. + + Mis-placed colwise scales showed up as several CTAs racing on a single + address, so the output differed run to run. A comparison against a reference + catches that only probabilistically; this catches it directly. + """ + total_M, N, num_groups = 2048, 512, 4 + offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 512) + x = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") + + first = [ + t.view(torch.uint8).clone() + for t in mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + ] + for _ in range(2): + again = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + for name, a, b in zip(("q_row", "s_row", "q_col", "s_col"), first, again): + assert torch.equal(a, b.view(torch.uint8)), f"{name} is nondeterministic" + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_special_value_semantics(): + """Subnormals, saturation and NaN/Inf, against the shared MXFP8 contract. + + The tile is laid out one semantic case per ROW, so it pins the ROWWISE half. + The colwise half reads each 32-row block as a mixture of cases, which has no + contract; test_..._inf_and_nan below covers what can be said about it. + """ + 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") + x[:num_cases] = cases.inputs.repeat(1, 4) + offs = torch.tensor([128], dtype=torch.int32, device="cuda") + + q_row, s_row = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs)[:2] + + scales = from_blocked(s_row, 128, 4) + for block_idx in range(4): + assert_mxfp8_semantics( + q_row[:num_cases, block_idx * 32 : (block_idx + 1) * 32], + scales[:num_cases, block_idx : block_idx + 1], + cases, + ) + + # And bytewise against the kernel this half replaces. + ref_q, ref_s = mxfp8_quantize_2d_1x32_cutedsl( + x, block_size=32, scaling_mode="rceil", offs=offs + ) + torch.testing.assert_close( + q_row.view(torch.uint8), ref_q.view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close( + s_row.view(torch.uint8), ref_s.view(torch.uint8), rtol=0, atol=0 + ) + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +@pytest.mark.parametrize("poison", ("inf", "nan", "inf_and_nan")) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_inf_and_nan(poison): + """Inf/NaN, where the two kernels being replaced do not agree with each other. + + The amax here is an unsigned max over sign-masked bf16 bit patterns, which + orders Inf and NaN as merely large integers, so this is the one input class + that needs an explicit branch -- and the two references want different + branches: + + rowwise (CuTeDSL, NOSAT) poisons the block: scale 255, every element NaN. + colwise (CUDA, SATFINITE) clamps Inf to scale 254 and keeps finite data, + and DROPS NaN from the amax entirely -- an all-NaN block still gets the + scale its finite elements imply. + + The fused kernel reproduces the rowwise half exactly and the colwise half + exactly for Inf. It deliberately does NOT reproduce the colwise kernel's + NaN-dropping amax: matching it would need a per-element NaN test in the + hottest loop, to imitate a behavior that hides NaN corruption rather than + surfacing it. That divergence is asserted rather than skipped, so it fails + loudly if either side changes. + """ + x = torch.ones((128, 128), 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([128], dtype=torch.int32, device="cuda") + + q_row, s_row, q_col, s_col = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + + # Rowwise: bit-exact with the kernel it replaces, always. + ref_q, ref_s = mxfp8_quantize_2d_1x32_cutedsl( + x, block_size=32, scaling_mode="rceil", offs=offs + ) + torch.testing.assert_close( + q_row.view(torch.uint8), ref_q.view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close( + s_row.view(torch.uint8), ref_s.view(torch.uint8), rtol=0, atol=0 + ) + + 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, + ) + fused_col_scale = int(s_col.view(torch.uint8).reshape(-1)[0]) + ref_col_scale = int(mx.scale.view(torch.uint8).reshape(-1)[0]) + + if poison == "inf": + # Inf clamps to 2**127 on both sides, data included. + assert fused_col_scale == ref_col_scale == 254 + torch.testing.assert_close( + q_col.view(torch.uint8), mx.qdata.view(torch.uint8), rtol=0, atol=0 + ) + else: + # Known, intentional divergence: we poison, the CUDA kernel ignores NaN. + assert fused_col_scale == 255, "fused colwise must poison a NaN block" + assert ref_col_scale != 255, ( + "the CUDA colwise kernel used to drop NaN from its amax; if it no " + "longer does, this divergence can be removed" + ) + + +@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +) +def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_input_validation(): + offs = torch.tensor([256], dtype=torch.int32, device="cuda") + + with pytest.raises(AssertionError, match="bf16"): + mxfp8_quantize_2d_1x32_32x1_cutedsl( + torch.randn(256, 256, dtype=torch.float32, device="cuda"), offs + ) + with pytest.raises(AssertionError, match="multiple of 128"): + mxfp8_quantize_2d_1x32_32x1_cutedsl( + torch.randn(256, 160, dtype=torch.bfloat16, device="cuda"), offs + ) + with pytest.raises(AssertionError, match="multiple of 128"): + mxfp8_quantize_2d_1x32_32x1_cutedsl( + torch.randn(160, 256, dtype=torch.bfloat16, device="cuda"), + torch.tensor([160], dtype=torch.int32, device="cuda"), + ) + with pytest.raises(AssertionError, match="contiguous"): + mxfp8_quantize_2d_1x32_32x1_cutedsl( + torch.randn(256, 256, dtype=torch.bfloat16, device="cuda").t(), offs + ) + with pytest.raises(AssertionError, match="block_size"): + mxfp8_quantize_2d_1x32_32x1_cutedsl( + torch.randn(256, 256, dtype=torch.bfloat16, device="cuda"), + offs, + block_size=16, + ) 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/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..2cf651613df --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_quantize_2d_1x32_32x1_fused.py @@ -0,0 +1,666 @@ +# 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 + 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 + +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. + * total_M % 128 == 0 and N % 128 == 0. These are the tile dimensions, not a + convenience -- one CTA owns a whole (128, 128) tile and neither edge is + predicated. + * offs must be 1D int32, contiguous, on device, holding ascending group END + offsets along 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. `_scale` runs here with the colwise Inf convention, which differs +from the rowwise one. + +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: it 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. + +COMPILATION +----------- +Compiled per concrete (total_M, N, G) -- all three 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, degenerate_ib: Uint32, nosat_inf: cutlass.Constexpr[bool]): + # 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^-118 (0x04e0). + e = es if bits >= Uint32(0x04E00000) else Int32(0) + ib = Uint32((Int32(254) - e) << Int32(23)) + if e == Int32(0): + ib = Uint32(0x3F800000) if bits == Uint32(0) else degenerate_ib + # 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. + # + # The two kernels being replaced DISAGREE here, so this does too, the + # same way it carries their degenerate-block convention: + # + # rowwise (cvt.rp.NOSAT.ue8m0x2.f32, cute_utils.py) maps Inf to 0xff, + # so Inf and NaN alike give scale 255 and a NaN reciprocal, and every + # element in the block comes out NaN. + # colwise (SATFINITE, mxfp8_quantize.cuh) CLAMPS Inf to 0xfe = 2^127 + # and only NaN reaches 0xff. Its reciprocal is the E8M0 byte + # 254 - scale, so scale 254 gives byte 0 = 2^-127 -- NOT the 0.0 that + # this function's (254 - e) << 23 would produce for e == 254. + # + # Neither branch is reachable from randn, so both need constructed input. + if bits > Uint32(0x7F800000): + e = Int32(255) + ib = Uint32(0x7FC00000) + elif bits == Uint32(0x7F800000): + if cutlass.const_expr(nosat_inf): + e = Int32(255) + ib = Uint32(0x7FC00000) + else: + e = Int32(254) + ib = Uint32(0x00800000) + return e, ib + + @cute.kernel + def _kernel( + mX: cute.Tensor, + mRQ: cute.Tensor, + mCQ: cute.Tensor, + mRS: cute.Tensor, + mCS: cute.Tensor, + mOffs: cute.Tensor, + kblocks: cutlass.Constexpr, + mblocks: cutlass.Constexpr, + groups: cutlass.Constexpr, + threads: cutlass.Constexpr, + ): + 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 tid == 0: + validate_group_sizes(mOffs) + + 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) + ) + + if mt >= mblocks: + # One extra CTA per group clears its 4-column over-allocation slab. + if kt == 0: + slab = mt - mblocks + base = ( + mblocks * TM // 32 * kblocks * TN + slab * 4 * kblocks * TN + ) // 4 + i = tid + while i < kblocks * TN: + mCS[base + i] = Uint32(0) + i += threads + else: + 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, Uint32(0x7F000000), True) + 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, Uint32(0x7F000000), True) + 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), Uint32(0x7F000000), False) + e1, ib1 = _scale(pk & Uint32(0xFFFF0000), Uint32(0x7F000000), False) + 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), Uint32(0x7F000000), False) + e1, ib1 = _scale(acc & Uint32(0xFFFF0000), Uint32(0x7F000000), False) + 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(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 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): + M = cutlass.const_expr(tX.shape[0]) + N = cutlass.const_expr(tX.shape[1]) + 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): + def w(t): + return from_dlpack(t, assumed_align=16, enable_tvm_ffi=True) + + return cute.compile( + _launch, + w(x), + w(offs), + w(q_row), + w(s_row), + w(q_col), + w(s_col), + make_fake_stream(), + options="--enable-tvm-ffi", + ) + + def run(x, offs, q_row, s_row, q_col, s_col): + key = (x.shape, offs.shape) + 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().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. + + Over-allocated by G*4 columns: the true padded width depends on the group + sizes, and computing it on the host would cost a D2H sync. The slack bytes + are zeroed by this kernel. + """ + 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. Copying `offs` (a handful of + ints) is free; doing the same to `x` would defeat the whole kernel, so `x` + is asserted instead. + """ + if t.data_ptr() % 16 == 0: + return t + return t.clone() + + +def quantize_fused( + x: torch.Tensor, + offs: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", + 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. N % 128 == 0, total_M % 128 == 0. + offs: (G,) int32 CUDA tensor of group END offsets along total_M, + ascending, every group size a multiple of 128 -- checked with a + device-side assert (see module docstring), 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. + block_size: MX block size, 32 only. + scaling_mode: "rceil" only (what TorchTitan runs). + *_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. + + 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 and offs.is_cuda, "x and offs must be CUDA tensors" + 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" + 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 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() + + s_col_rows, s_col_cols = s_col_shape(total_M, N, G) + + 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 + ) + + run(x, _aligned16(offs), 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..2fa4c4a6201 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/quant.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/quant.py @@ -1098,6 +1098,35 @@ 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, + offs: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + 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, + offs, + block_size=block_size, + scaling_mode=scaling_mode, + ) + # 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, @@ -1160,6 +1189,44 @@ def _fake_mxfp8_quantize_2d_32x1_cutedsl_custom_op( 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, + offs: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + assert x.ndim == 2, "input tensor must be 2D" + assert block_size == 32, "Only block_size=32 is supported" + total_m, n = x.shape + num_groups = offs.shape[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 +1643,61 @@ def mxfp8_quantize_2d_32x1_cutedsl( return qdata, scales +def mxfp8_quantize_2d_1x32_32x1_cutedsl( + x: torch.Tensor, + offs: torch.Tensor, + block_size: int = 32, + scaling_mode: str = "rceil", +) -> 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. + + Args: + x: Input tensor of shape (total_M, N), bfloat16, row-major contiguous. + Both dimensions must be multiples of 128. + offs: Group end offsets along total_M, int32, ascending, 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. + block_size: Block size for quantization (only 32 supported) + scaling_mode: Scaling mode (only "rceil" supported) + + 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 K-groups blocked. + """ + 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, + offs, + block_size=block_size, + scaling_mode=scaling_mode, + ) + + # ============================================================================= # FlyDSL MXFP8 quantize kernels (AMD CDNA3+ via the FlyDSL stack). # From 86f7e2aaa6bf5611ce3eed9fb50dd41fea4280d9 Mon Sep 17 00:00:00 2001 From: Sanket Jayant Purandare Date: Mon, 12 Oct 2026 11:53:46 -0700 Subject: [PATCH 2/2] [moe_training] Match MXFP8 dense and grouped contracts Performnace of this fused cute-dsl quantizer is better than cuda, a single API for dense and grouped mxfp8 linears, no need for scale rearrangements since swizzled layouts are directly output. AI-assisted implementation details: > Keep Protonu's original authored fused kernel from pytorch/ao#4895 as the > parent commit and consolidate current-main correctness fixes in this PR. > Credit: Protonu (@protonu), original commit > 1d4e8d3d4eace3a7a8b705c997ca23523228cbae; Aleksandar Samardzic's #4702 > independently addressed the single-axis FakeTensor metadata issue. > > Match the public/custom-op/fake argument order to surrounding APIs. None > offsets select dense placement in the same kernel, with compact blocked > scales and no synthetic endpoint or offset search. Tensor offsets retain > grouped placement, empty experts and uncovered-tail behavior. > > Correct the scale0/1 RCEIL cutoff to 448*2^-127, preserve NOSAT nonfinite > semantics, and compare both axes directly to to_mx with group unblocking. > Cover exact cutoff neighbors, mixed magnitudes, special values, empty groups, > dirty destinations, actual compile/fake contracts and device selection. > > Validate aligned destination metadata and nonoverlapping byte intervals. > Enable device assertions explicitly for grouped compilation; CuTeDSL strips > them by default. Reuse the endpoint validation loop for ordering and bounds. > Remove unreachable CTA work and redundant reciprocal conventions. > > Test plan: affected fused/single-axis CUDA tests, repository Ruff checks, > strict API documentation rendering, matched replay timings and traces, and > real TorchTitan FSDP/MXFP8 eager/capture consumer tests. Pull Request resolved: https://github.com/pytorch/ao/pull/4995 --- ...nch_cutedsl_quantize_2d_1x32_32x1_fused.py | 15 +- test/prototype/moe_training/test_kernels.py | 596 +++++++++++------- .../moe_training/kernels/mxfp8/cute_utils.py | 13 +- .../cutedsl_quantize_2d_1x32_32x1_fused.py | 538 ++++++++-------- .../moe_training/kernels/mxfp8/quant.py | 55 +- 5 files changed, 728 insertions(+), 489 deletions(-) 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 index 8f089f0a7b4..6d4d6c48518 100644 --- 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 @@ -37,6 +37,7 @@ import itertools import os +import random from dataclasses import dataclass from typing import List @@ -147,7 +148,7 @@ def assert_graph_capture_is_real(x: torch.Tensor, offs: torch.Tensor): def run(): quantize_fused( x, - offs, + offs=offs, block_size=BLOCK_SIZE, scaling_mode="rceil", q_row_out=bufs[0], @@ -205,7 +206,10 @@ def validate_outputs(fused, unfused, num_groups: int): 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[:, :used], exp[:, :used] + 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") @@ -218,14 +222,16 @@ def run_experiment(config: ExperimentConfig) -> ExperimentResult: torch.int32 ) - fused_out = quantize_fused(x, offs, block_size=BLOCK_SIZE, scaling_mode="rceil") + 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, offs, block_size=BLOCK_SIZE, scaling_mode="rceil" + quantize_fused, x, block_size=BLOCK_SIZE, scaling_mode="rceil", offs=offs ) unfused_time_us = benchmark_cuda_graph_in_microseconds(unfused_baseline, x, offs) @@ -290,6 +296,7 @@ def print_results(experiments: List[Experiment]): def main(): torch.random.manual_seed(123) + random.seed(123) configs = get_configs() results = [] for config in tqdm(configs): diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index 7ad156323d1..e8a315dbd6b 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -522,8 +522,10 @@ def test_cuda_mx_3d_cutedsl_numerics(E, N, K, input_dtype, scaling_mode, variant y_ref = y_ref.transpose(-2, -1) s_ref = s_ref.transpose(-2, -1) s_rows, s_cols = K, N // block_size - undo_scale = lambda scale: ( - from_blocked(scale, s_rows, s_cols).transpose(-2, -1).contiguous() + undo_scale = ( + lambda scale: from_blocked(scale, s_rows, s_cols) + .transpose(-2, -1) + .contiguous() ) else: x_tiles = ( @@ -1696,14 +1698,15 @@ def test_amd_mx_3d_flydsl_numerics( # # 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. Its contract is therefore to be a DROP-IN for -# that composition, so every test below compares raw bytes against the -# composition itself rather than against `to_mx`. That distinction matters: the -# fused kernel is deliberately bug-compatible with the colwise CUDA kernel's -# degenerate-block convention, so a `to_mx` reference would disagree on inputs -# that no random tensor can produce. +# -- 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. @@ -1737,7 +1740,7 @@ def _fused_1x32_32x1_offs_with_tail(total_M, num_groups, tail, device="cuda"): return ends.to(device=device, dtype=torch.int32) -def _fused_1x32_32x1_reference(x, offs): +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. @@ -1764,32 +1767,60 @@ def _fused_1x32_32x1_reference(x, offs): def _assert_fused_1x32_32x1_matches_reference(x, offs, actual): - """Bytewise, atol=0/rtol=0. A single wrong E8M0 byte is a 2x error.""" - expected = _fused_1x32_32x1_reference(x, offs) + """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") - num_groups = offs.numel() - for name, got, exp in zip(names, actual, expected): - got = got.view(torch.uint8) - if name == "s_col": - # Both sides over-allocate the colwise scale buffer by num_groups * 4 - # columns of per-group padding slack that the grouped GEMM never - # reads. Compare the region that is actually consumed. - used = exp.shape[1] - num_groups * 4 - got, exp = got[:, :used], exp[:, :used] + 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") -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) +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 = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") - offs = generate_jagged_offs(num_groups, total_M, multiple_of=128, device="cuda").to( - torch.int32 - ) - out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + 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) @@ -1800,12 +1831,12 @@ def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_numerics(total_M, N, num_groups): 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) -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) +@_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. @@ -1822,29 +1853,12 @@ def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_uncovered_tail(tail): offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, tail) assert int(offs[-1]) == covered < total_M - x = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") - out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + 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) - # The colwise scale bytes past the covered prefix must be ZEROED, matching - # triton_mx_block_rearrange_2d_K_groups, which iterates groups and returns - # new_zeros. Asserted directly because a fresh torch.empty can hide a - # never-written region whenever the allocator hands back zeroed pages. - # - # Indexed FLAT, not by column: the K-groups blocked layout stores each - # 4-scale-column block as N contiguous bytes, so scale columns - # [covered // 32, total_M // 32) are the flat byte range below and are NOT a - # 2D column slice of the (padded_N, num_scale_cols) view. - s_col_flat = out[3].view(torch.uint8).reshape(-1) - uncovered = s_col_flat[(covered // 32) * N : (total_M // 32) * N] - assert uncovered.numel() == (tail // 32) * N - assert torch.all(uncovered == 0), "uncovered s_col bytes must be zero-filled" - - -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) + +@_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. @@ -1856,36 +1870,28 @@ def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_degenerate_blocks(scale): exactly how a 2x error survived here once already. """ total_M, N, num_groups = 512, 256, 4 - x = (torch.randn(total_M, N, device="cuda") * scale).to(torch.bfloat16) - offs = generate_jagged_offs(num_groups, total_M, multiple_of=128, device="cuda").to( - torch.int32 + x = (_fused_1x32_32x1_randn(total_M, N, dtype=torch.float32) * scale).to( + torch.bfloat16 ) - out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + 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) -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) +@_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 = (torch.randn(total_M, N, device="cuda") * 100).to(torch.bfloat16) + 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 = generate_jagged_offs(num_groups, total_M, multiple_of=128, device="cuda").to( - torch.int32 - ) - out = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + 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) -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) +@_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. @@ -1900,7 +1906,7 @@ def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_dirty_output_buffers(): total_M, N, num_groups = 2048, 512, 4 offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 512) - x = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") + x = _fused_1x32_32x1_randn(total_M, N) s_col_rows, s_col_cols = s_col_shape(total_M, N, num_groups) bufs = ( @@ -1909,183 +1915,337 @@ def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_dirty_output_buffers(): 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 buf in bufs: - buf.view(torch.uint8).fill_(0xFF) - - quantize_fused( - x, - 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) - - -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) -def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_determinism(): - """Three launches on one input must agree byte for byte. - - Mis-placed colwise scales showed up as several CTAs racing on a single - address, so the output differed run to run. A comparison against a reference - catches that only probabilistically; this catches it directly. - """ - total_M, N, num_groups = 2048, 512, 4 - offs = _fused_1x32_32x1_offs_with_tail(total_M, num_groups, 512) - x = torch.randn(total_M, N, dtype=torch.bfloat16, device="cuda") - - first = [ - t.view(torch.uint8).clone() - for t in mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) - ] - for _ in range(2): - again = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) - for name, a, b in zip(("q_row", "s_row", "q_col", "s_col"), first, again): - assert torch.equal(a, b.view(torch.uint8)), f"{name} is nondeterministic" + 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) -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) -def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_special_value_semantics(): - """Subnormals, saturation and NaN/Inf, against the shared MXFP8 contract. - - The tile is laid out one semantic case per ROW, so it pins the ROWWISE half. - The colwise half reads each 32-row block as a mixture of cases, which has no - contract; test_..._inf_and_nan below covers what can be said about it. - """ +@_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") - x[:num_cases] = cases.inputs.repeat(1, 4) + 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") - q_row, s_row = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs)[:2] + 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(s_row, 128, 4) + scales = from_blocked(blocked_scales.flatten()[:512], 128, 4) for block_idx in range(4): assert_mxfp8_semantics( - q_row[:num_cases, block_idx * 32 : (block_idx + 1) * 32], + q_data[:num_cases, block_idx * 32 : (block_idx + 1) * 32], scales[:num_cases, block_idx : block_idx + 1], cases, ) - # And bytewise against the kernel this half replaces. - ref_q, ref_s = mxfp8_quantize_2d_1x32_cutedsl( - x, block_size=32, scaling_mode="rceil", offs=offs - ) - torch.testing.assert_close( - q_row.view(torch.uint8), ref_q.view(torch.uint8), rtol=0, atol=0 - ) - torch.testing.assert_close( - s_row.view(torch.uint8), ref_s.view(torch.uint8), rtol=0, atol=0 - ) + _assert_fused_1x32_32x1_matches_reference(x, offs, out) -@pytest.mark.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" -) +@_FUSED_CUTEDSL_REQUIRED @pytest.mark.parametrize("poison", ("inf", "nan", "inf_and_nan")) -def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_inf_and_nan(poison): - """Inf/NaN, where the two kernels being replaced do not agree with each other. - - The amax here is an unsigned max over sign-masked bf16 bit patterns, which - orders Inf and NaN as merely large integers, so this is the one input class - that needs an explicit branch -- and the two references want different - branches: - - rowwise (CuTeDSL, NOSAT) poisons the block: scale 255, every element NaN. - colwise (CUDA, SATFINITE) clamps Inf to scale 254 and keeps finite data, - and DROPS NaN from the amax entirely -- an all-NaN block still gets the - scale its finite elements imply. - - The fused kernel reproduces the rowwise half exactly and the colwise half - exactly for Inf. It deliberately does NOT reproduce the colwise kernel's - NaN-dropping amax: matching it would need a per-element NaN test in the - hottest loop, to imitate a behavior that hides NaN corruption rather than - surfacing it. That divergence is asserted rather than skipped, so it fails - loudly if either side changes. - """ - x = torch.ones((128, 128), dtype=torch.bfloat16, device="cuda") +@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([128], dtype=torch.int32, device="cuda") + offs = torch.tensor([shape[0]], dtype=torch.int32, device="cuda") - q_row, s_row, q_col, s_col = mxfp8_quantize_2d_1x32_32x1_cutedsl(x, offs) + 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 - # Rowwise: bit-exact with the kernel it replaces, always. - ref_q, ref_s = mxfp8_quantize_2d_1x32_cutedsl( - x, block_size=32, scaling_mode="rceil", offs=offs - ) - torch.testing.assert_close( - q_row.view(torch.uint8), ref_q.view(torch.uint8), rtol=0, atol=0 - ) - torch.testing.assert_close( - s_row.view(torch.uint8), ref_s.view(torch.uint8), rtol=0, atol=0 - ) - 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, - ) - fused_col_scale = int(s_col.view(torch.uint8).reshape(-1)[0]) - ref_col_scale = int(mx.scale.view(torch.uint8).reshape(-1)[0]) +@_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) + - if poison == "inf": - # Inf clamps to 2**127 on both sides, data included. - assert fused_col_scale == ref_col_scale == 254 +@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( - q_col.view(torch.uint8), mx.qdata.view(torch.uint8), rtol=0, atol=0 + left.view(torch.uint8).flatten(), right.view(torch.uint8), rtol=0, atol=0 ) - else: - # Known, intentional divergence: we poison, the CUDA kernel ignores NaN. - assert fused_col_scale == 255, "fused colwise must poison a NaN block" - assert ref_col_scale != 255, ( - "the CUDA colwise kernel used to drop NaN from its amax; if it no " - "longer does, this divergence can be removed" + + +@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.skipif(not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x") -@pytest.mark.skipif( - not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" +@pytest.mark.parametrize("fake", [False, True]) +@pytest.mark.parametrize( + "kwargs,match", + [({"block_size": 16}, "block_size"), ({"scaling_mode": "floor"}, "rceil")], ) -def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_input_validation(): - offs = torch.tensor([256], dtype=torch.int32, device="cuda") +@_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 - with pytest.raises(AssertionError, match="bf16"): - mxfp8_quantize_2d_1x32_32x1_cutedsl( - torch.randn(256, 256, dtype=torch.float32, device="cuda"), offs - ) - with pytest.raises(AssertionError, match="multiple of 128"): - mxfp8_quantize_2d_1x32_32x1_cutedsl( - torch.randn(256, 160, dtype=torch.bfloat16, device="cuda"), offs - ) - with pytest.raises(AssertionError, match="multiple of 128"): - mxfp8_quantize_2d_1x32_32x1_cutedsl( - torch.randn(160, 256, dtype=torch.bfloat16, device="cuda"), - torch.tensor([160], dtype=torch.int32, device="cuda"), - ) - with pytest.raises(AssertionError, match="contiguous"): - mxfp8_quantize_2d_1x32_32x1_cutedsl( - torch.randn(256, 256, dtype=torch.bfloat16, device="cuda").t(), offs + 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 ) - with pytest.raises(AssertionError, match="block_size"): - mxfp8_quantize_2d_1x32_32x1_cutedsl( - torch.randn(256, 256, dtype=torch.bfloat16, device="cuda"), - offs, - block_size=16, + 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/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 index 2cf651613df..fe29fa8968b 100644 --- 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 @@ -6,12 +6,17 @@ """Fused 1x32 + 32x1 MXFP8 quantization -- one read of x, four outputs. - x (total_M, N) bf16 + offs (G,) int32, group sizes % 128 == 0 + 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 @@ -22,11 +27,11 @@ * bf16 only. * x must be 2D, row-major contiguous, and 16-byte aligned. - * total_M % 128 == 0 and N % 128 == 0. These are the tile dimensions, not a + * 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. - * offs must be 1D int32, contiguous, on device, holding ascending group END - offsets along total_M. + * 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 @@ -72,8 +77,8 @@ 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. `_scale` runs here with the colwise Inf convention, which differs -from the rowwise one. +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 @@ -82,14 +87,15 @@ 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: it scans `offs` for the group owning this tile and places its 512 bytes at +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. +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) -- all three are constant-folded into the +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. """ @@ -199,42 +205,23 @@ def _umax(a: Uint32, b: Uint32) -> Uint32: return a if a > b else b @cute.jit - def _scale(bits: Uint32, degenerate_ib: Uint32, nosat_inf: cutlass.Constexpr[bool]): + 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^-118 (0x04e0). - e = es if bits >= Uint32(0x04E00000) else Int32(0) + # 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)) - if e == Int32(0): - ib = Uint32(0x3F800000) if bits == Uint32(0) else degenerate_ib # 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. # - # The two kernels being replaced DISAGREE here, so this does too, the - # same way it carries their degenerate-block convention: - # - # rowwise (cvt.rp.NOSAT.ue8m0x2.f32, cute_utils.py) maps Inf to 0xff, - # so Inf and NaN alike give scale 255 and a NaN reciprocal, and every - # element in the block comes out NaN. - # colwise (SATFINITE, mxfp8_quantize.cuh) CLAMPS Inf to 0xfe = 2^127 - # and only NaN reaches 0xff. Its reciprocal is the E8M0 byte - # 254 - scale, so scale 254 gives byte 0 = 2^-127 -- NOT the 0.0 that - # this function's (254 - e) << 23 would produce for e == 254. - # - # Neither branch is reachable from randn, so both need constructed input. - if bits > Uint32(0x7F800000): + # Match NOSAT in both the rowwise CuTeDSL and native CUDA references. + if bits >= Uint32(0x7F800000): e = Int32(255) ib = Uint32(0x7FC00000) - elif bits == Uint32(0x7F800000): - if cutlass.const_expr(nosat_inf): - e = Int32(255) - ib = Uint32(0x7FC00000) - else: - e = Int32(254) - ib = Uint32(0x00800000) return e, ib @cute.kernel @@ -244,20 +231,22 @@ def _kernel( mCQ: cute.Tensor, mRS: cute.Tensor, mCS: cute.Tensor, - mOffs: 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 tid == 0: - validate_group_sizes(mOffs) + 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. @@ -282,175 +271,162 @@ def _kernel( cute.recast_ptr(scol, dtype=cutlass.Uint8), cute.make_layout(512) ) - if mt >= mblocks: - # One extra CTA per group clears its 4-column over-allocation slab. - if kt == 0: - slab = mt - mblocks - base = ( - mblocks * TM // 32 * kblocks * TN + slab * 4 * kblocks * TN - ) // 4 - i = tid - while i < kblocks * TN: - mCS[base + i] = Uint32(0) - i += threads + 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: - 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, Uint32(0x7F000000), True) - 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, Uint32(0x7F000000), True) - 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 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)], ) - 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), Uint32(0x7F000000), False) - e1, ib1 = _scale(pk & Uint32(0xFFFF0000), Uint32(0x7F000000), False) - 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), Uint32(0x7F000000), False) - e1, ib1 = _scale(acc & Uint32(0xFFFF0000), Uint32(0x7F000000), False) + 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], cf[4 * g + 1], invw) - ub = cvt2(cf[4 * g + 2], cf[4 * g + 3], invw) + 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 * 2, mchunk, None)]) - cute.autovec_copy(oo, mCQ[(kt * TN + jq * 2 + 1, mchunk, None)]) + 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] = 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) + 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) @@ -462,36 +438,37 @@ def _kernel( 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 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 @@ -503,9 +480,13 @@ def _kernel( @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]) - G = cutlass.const_expr(tOffs.shape[0]) + 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), @@ -536,8 +517,15 @@ def _launch(tX, tOffs, tRQ, tRS, tCQ, tCS, 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): - return from_dlpack(t, assumed_align=16, enable_tvm_ffi=True) + """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, @@ -548,11 +536,16 @@ def w(t): w(q_col), w(s_col), make_fake_stream(), - options="--enable-tvm-ffi", + # 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): - key = (x.shape, offs.shape) + """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) @@ -564,16 +557,16 @@ def run(x, offs, q_row, s_row, q_col, s_col): s_row, q_col, s_col, - _cuda_drv.CUstream(int(torch.cuda.current_stream().cuda_stream)), + _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. - Over-allocated by G*4 columns: the true padded width depends on the group - sizes, and computing it on the host would cost a D2H sync. The slack bytes - are zeroed by this kernel. + 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 @@ -582,9 +575,8 @@ 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. Copying `offs` (a handful of - ints) is free; doing the same to `x` would defeat the whole kernel, so `x` - is asserted instead. + 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 @@ -593,9 +585,10 @@ def _aligned16(t: torch.Tensor) -> torch.Tensor: def quantize_fused( x: torch.Tensor, - offs: 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, @@ -604,10 +597,13 @@ def quantize_fused( """Both MXFP8 casts of grad_out from one pass over `x`. Args: - x: (total_M, N) bfloat16, row-major. N % 128 == 0, total_M % 128 == 0. - offs: (G,) int32 CUDA tensor of group END offsets along total_M, - ascending, every group size a multiple of 128 -- checked with a - device-side assert (see module docstring), which raises an async + 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 @@ -615,11 +611,10 @@ def quantize_fused( 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. - block_size: MX block size, 32 only. - scaling_mode: "rceil" only (what TorchTitan runs). *_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, @@ -631,21 +626,27 @@ def quantize_fused( + ", ".join(_missing_cutedsl_runtime_packages()) ) - assert x.is_cuda and offs.is_cuda, "x and offs must be CUDA tensors" + 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" - assert offs.dtype == torch.int32 and offs.dim() == 1, "offs must be 1D int32" - assert offs.is_contiguous(), "offs must be contiguous" + 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() + 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( @@ -662,5 +663,38 @@ def quantize_fused( (s_col_rows, s_col_cols), device=x.device, dtype=torch.uint8 ) - run(x, _aligned16(offs), q_row_out, s_row_out, q_col_out, s_col_out) + 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 2fa4c4a6201..a621a3c29f6 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/quant.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/quant.py @@ -1103,19 +1103,20 @@ def _mxfp8_quantize_2d_32x1_cutedsl_custom_op( ) def _mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( x: torch.Tensor, - offs: 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, - offs, 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. @@ -1173,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, ) @@ -1182,8 +1183,13 @@ 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 @@ -1192,14 +1198,27 @@ def _fake_mxfp8_quantize_2d_32x1_cutedsl_custom_op( @_mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op.register_fake def _fake_mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( x: torch.Tensor, - offs: 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 - num_groups = offs.shape[0] + 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), @@ -1645,9 +1664,9 @@ def mxfp8_quantize_2d_32x1_cutedsl( def mxfp8_quantize_2d_1x32_32x1_cutedsl( x: torch.Tensor, - offs: 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, @@ -1658,12 +1677,20 @@ def mxfp8_quantize_2d_1x32_32x1_cutedsl( 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. + 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 multiples of 128. - offs: Group end offsets along total_M, int32, ascending, every group + 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 @@ -1671,12 +1698,12 @@ def mxfp8_quantize_2d_1x32_32x1_cutedsl( 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. - block_size: Block size for quantization (only 32 supported) - scaling_mode: Scaling mode (only "rceil" supported) + 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 K-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() @@ -1692,9 +1719,9 @@ def mxfp8_quantize_2d_1x32_32x1_cutedsl( ) return _mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( x, - offs, block_size=block_size, scaling_mode=scaling_mode, + offs=offs, )