diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_M_groups.py b/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_M_groups.py index 4f59552767f..93c86f248e2 100644 --- a/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_M_groups.py +++ b/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_M_groups.py @@ -3,201 +3,203 @@ # # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. -# this benchmarking script is a modified version of the original script from: https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/utils/benchmark.py -import itertools -from dataclasses import dataclass -from typing import List +import argparse +import sys import torch -from tabulate import tabulate -from tqdm import tqdm from benchmarks.utils import benchmark_cuda_function_in_microseconds -from torchao.prototype.moe_training.kernels.mxfp8 import ( +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( + _mx_block_rearrange_2d_m_groups_cutedsl, +) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( mx_block_rearrange_2d_M_groups_cuda, torch_to_blocked_2d_M_groups, triton_mx_block_rearrange_2d_M_groups, ) from torchao.prototype.moe_training.utils import generate_jagged_offs -device = torch.device("cuda") - -# Needed since changing args to function causes recompiles -torch._dynamo.config.cache_size_limit = 1000 - - -@dataclass(frozen=True) -class ExperimentConfig: - input_shape: tuple[int] - num_groups: int - chunks_per_tb: int - - -@dataclass(frozen=True) -class ExperimentResult: - torch_time_us: float - triton_time_us: float - cuda_time_us: float - torch_mem_bw_gbps: float - triton_mem_bw_gbps: float - cuda_mem_bw_gbps: float - - -@dataclass(frozen=True) -class Experiment: - config: ExperimentConfig - result: ExperimentResult - -def get_configs() -> List[ExperimentConfig]: - block_size = 32 - input_shapes = [ - (16384, 5120 // block_size), - (131072, 5120 // block_size), - (131072, 2048 // block_size), - (131072, 7168 // block_size), - ( - 131072, - 1408 // block_size, - ), # case for dsv3-16b where scale cols not multiple of 16 - ] - num_groups = [8] - chunks_per_tb_list = [1, 4, 8] - - configs = [] - for shape, groups, chunks_per_tb in itertools.product( - input_shapes, - num_groups, - chunks_per_tb_list, - ): - configs.append( - ExperimentConfig( - input_shape=shape, - num_groups=groups, - chunks_per_tb=chunks_per_tb, - ) +def benchmark_cuda_graph_function_in_microseconds(f, *args, iters=1000, **kwargs): + for _ in range(10): + f(*args, **kwargs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + f(*args, **kwargs) + for _ in range(10): + graph.replay() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / iters + + +def benchmark_function(f, *args, use_cuda_graph=False, graph_iters=1000, **kwargs): + if use_cuda_graph: + return benchmark_cuda_graph_function_in_microseconds( + f, + *args, + iters=graph_iters, + **kwargs, ) - return configs + return benchmark_cuda_function_in_microseconds(f, *args, **kwargs) -def run_experiment(config: ExperimentConfig) -> ExperimentResult: - input_shape, num_groups, chunks_per_tb = ( - config.input_shape, - config.num_groups, - config.chunks_per_tb, - ) +def _parse_shape(spec: str): + """Parses `ROWSxCOLSxGROUPS` or `LABEL:ROWSxCOLSxGROUPS`.""" + label, _, dims = spec.rpartition(":") + parts = dims.split("x") + if len(parts) != 3: + raise SystemExit(f"Malformed shape {spec!r}, expected [LABEL:]ROWSxCOLSxGROUPS") + rows, cols, groups = (int(p) for p in parts) + return (label or None, rows, cols, groups) - input_tensor = torch.randint( - low=0, - high=256, - size=input_shape, - dtype=torch.uint8, - device=device, - ) - Mg, K = input_shape - block_size = 32 - input_group_offsets = generate_jagged_offs(num_groups, Mg, multiple_of=block_size) +def _markdown(rows, columns): + widths = [ + max(len(c), *(len(str(r[c])) for r in rows)) if rows else len(c) + for c in columns + ] + header = "| " + " | ".join(c.ljust(w) for c, w in zip(columns, widths)) + " |" + sep = "|" + "|".join("---:".rjust(w + 2) for w in widths) + "|" + body = [ + "| " + " | ".join(str(r[c]).rjust(w) for c, w in zip(columns, widths)) + " |" + for r in rows + ] + return "\n".join([header, sep, *body]) - # bench torch - compiled_run_torch = torch.compile(torch_to_blocked_2d_M_groups) - torch_out_scales, torch_group_offs = compiled_run_torch( - input_tensor, - input_group_offsets, - block_size=block_size, - ) - torch_time_us = benchmark_cuda_function_in_microseconds( - compiled_run_torch, - input_tensor, - input_group_offsets, - block_size=block_size, - ) - # bench triton - triton_out_scales = triton_mx_block_rearrange_2d_M_groups( - input_tensor, - input_group_offsets, +def main(): + parser = argparse.ArgumentParser( + description="Benchmark torch / triton / cuda / CuTeDSL MXFP8 M-groups scale " + "rearrange. Shapes are scale-tensor dimensions." ) - triton_time_us = benchmark_cuda_function_in_microseconds( - triton_mx_block_rearrange_2d_M_groups, - input_tensor, - input_group_offsets, + parser.add_argument( + "shapes", + nargs="+", + metavar="[LABEL:]ROWSxCOLSxGROUPS", + help="scale shapes to benchmark, e.g. 1024x224x8 or dsv3:131072x224x8", ) - - # bench CUDA pipelined kernel with configured chunk_width and chunks_per_tb - _ = mx_block_rearrange_2d_M_groups_cuda( - input_tensor.view(torch.uint8), - input_group_offsets.to(torch.int32), - chunks_per_tb, + parser.add_argument("--multiple-of", type=int, default=128) + parser.add_argument( + "--chunk-width", + type=int, + default=None, + help="CuTeDSL column chunk width; default lets the kernel choose", ) - cuda_time_us = benchmark_cuda_function_in_microseconds( - mx_block_rearrange_2d_M_groups_cuda, - input_tensor.view(torch.uint8), - input_group_offsets.to(torch.int32), - chunks_per_tb, + parser.add_argument( + "--bench-only", + nargs="+", + choices=("torch", "triton", "cuda", "cutedsl"), + default=None, + help="restrict to one or more backends; default is all", ) - - # mem bw calculations - bytes_per_input_el = torch.finfo(torch.float8_e8m0fnu).bits / 8 - bytes_per_output_el = torch.finfo(torch.float8_e4m3fn).bits / 8 - - read_bytes = input_tensor.numel() * bytes_per_input_el - write_bytes = triton_out_scales.numel() * bytes_per_output_el - - torch_mem_bw_gbps = ((read_bytes + write_bytes) / 1e9) / (torch_time_us / 1e6) - triton_mem_bw_gbps = ((read_bytes + write_bytes) / 1e9) / (triton_time_us / 1e6) - cuda_mem_bw_gbps = ((read_bytes + write_bytes) / 1e9) / (cuda_time_us / 1e6) - - return ExperimentResult( - torch_time_us=torch_time_us, - triton_time_us=triton_time_us, - cuda_time_us=cuda_time_us, - torch_mem_bw_gbps=torch_mem_bw_gbps, - triton_mem_bw_gbps=triton_mem_bw_gbps, - cuda_mem_bw_gbps=cuda_mem_bw_gbps, + parser.add_argument("--deterministic", action="store_true") + parser.add_argument("--cuda-graph-bench", action="store_true") + parser.add_argument("--graph-iters", type=int, default=1000) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is not available") + + shapes = [_parse_shape(s) for s in args.shapes] + labelled = any(label is not None for label, _, _, _ in shapes) + selected = ( + set(args.bench_only) + if args.bench_only + else {"torch", "triton", "cuda", "cutedsl"} ) - - -def print_results(experiments: List[Experiment]): - headers = [ - "input_shape", - "chunks_per_tb", - "torch_time_us", - "triton_time_us", - "cuda_time_us", - "triton_speedup", - "cuda_speedup", - ] - rows = [] - for experiment in experiments: - input_shape = ( - f"({experiment.config.input_shape[0]}, {experiment.config.input_shape[1]})" + if args.cuda_graph_bench and "torch" in selected: + # torch_to_blocked_2d_M_groups calls group_offs.tolist(), a synchronous + # D2H copy that CUDA graph capture never tolerates, for any shape. + if selected == {"torch"}: + raise SystemExit("torch cannot be benchmarked under --cuda-graph-bench") + print( + "warning: excluding torch from --cuda-graph-bench (unpinned D2H copy)", + file=sys.stderr, ) - rows.append( - [ - input_shape, - experiment.config.chunks_per_tb, - f"{experiment.result.torch_time_us:.2f}", - f"{experiment.result.triton_time_us:.2f}", - f"{experiment.result.cuda_time_us:.2f}", - f"{experiment.result.torch_time_us / experiment.result.triton_time_us:.2f}x", - f"{experiment.result.torch_time_us / experiment.result.cuda_time_us:.2f}x", - ] + selected.discard("torch") + rows_out = [] + + for label, rows, cols, groups in shapes: + if args.deterministic: + raw = ( + torch.arange(rows * cols, device="cuda", dtype=torch.int32) + .remainder(251) + .to(torch.uint8) + .reshape(rows, cols) + ) + else: + raw = torch.randint(0, 256, (rows, cols), device="cuda", dtype=torch.uint8) + scales = raw.view(torch.float8_e8m0fnu) + offsets = generate_jagged_offs( + groups, rows, multiple_of=args.multiple_of, device="cuda" ) - print(tabulate(rows, headers=headers)) + row = {"Model layer": label or ""} if labelled else {} + row["Scales RxC"] = f"{rows}x{cols}" + row["Groups"] = groups + + if "torch" in selected: + torch_us = benchmark_function( + torch_to_blocked_2d_M_groups, + scales, + offsets, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Torch us"] = f"{torch_us:.2f}" + if "triton" in selected: + triton_us = benchmark_function( + triton_mx_block_rearrange_2d_M_groups, + scales, + offsets, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Triton us"] = f"{triton_us:.2f}" + if "cuda" in selected: + cuda_us = benchmark_function( + mx_block_rearrange_2d_M_groups_cuda, + scales, + offsets, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["CUDA us"] = f"{cuda_us:.2f}" + if "cutedsl" in selected: + out = _mx_block_rearrange_2d_m_groups_cutedsl( + scales, offsets, args.chunk_width + ) + cutedsl_us = benchmark_function( + _mx_block_rearrange_2d_m_groups_cutedsl, + scales, + offsets, + args.chunk_width, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["CuTeDSL us"] = f"{cutedsl_us:.2f}" + # Rearrange is purely memory bound, so GB/s says how close to roofline + # the kernel gets in a way raw microseconds cannot. + moved_bytes = ( + (scales.numel() + out.numel()) + * torch.finfo(torch.float8_e8m0fnu).bits + / 8 + ) + row["CuTeDSL GB/s"] = f"{(moved_bytes / 1e9) / (cutedsl_us / 1e6):.1f}" + if "triton" in selected and "cutedsl" in selected: + row["CuTeDSL speedup"] = f"{triton_us / cutedsl_us:.2f}" + rows_out.append(row) -def main(): - torch.random.manual_seed(123) - configs = get_configs() - results = [] - for config in tqdm(configs): - result = run_experiment(config) - results.append(Experiment(config=config, result=result)) - - # Use Tabulate to print results - print_results(results) + print(_markdown(rows_out, list(rows_out[0].keys()))) if __name__ == "__main__": diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index 8a4a603ddce..077491aa35a 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -54,6 +54,9 @@ def _is_sm_10x() -> bool: triton_mx_block_rearrange_2d_M_groups, triton_mx_block_rearrange_per_group_3d, ) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( + _mx_block_rearrange_2d_m_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, _mxfp8_cutedsl_kernels_available, @@ -273,6 +276,77 @@ def test_triton_mx_block_rearrange_2d_M_groups( ) +@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", +) +@skip_if_rocm("ROCm enablement in progress") +@pytest.mark.parametrize( + "scale_rows,scale_cols,n_groups,chunk_width", + [ + # chunk_width=None lets the kernel pick from cols (64/32/16) + (256, 37, 8, None), + (256, 45, 8, None), + (512, 64, 8, None), + (1024, 224, 8, None), + (1024, 512, 8, None), + (1024, 1025, 8, None), + # explicit widths: the store path differs per width (ALL_COL_CHUNKS_FULL vs + # COLS_MULTIPLE_OF_16 vs byte-wise), and 128 is unreachable via the default + (512, 128, 4, 16), + (512, 128, 4, 32), + (512, 128, 4, 64), + (512, 128, 4, 128), + (1024, 260, 3, 128), + ], +) +def test_cutedsl_mx_block_rearrange_2d_M_groups( + scale_rows: int, + scale_cols: int, + n_groups: int, + chunk_width: int, +): + device = "cuda" + block_size = 32 + e8m0_scales = torch.randint( + 0, + 256, + (scale_rows, scale_cols), + device=device, + dtype=torch.uint8, + ).view(torch.float8_e8m0fnu) + input_group_offsets = generate_jagged_offs( + n_groups, scale_rows, multiple_of=block_size, device=device + ) + + padded_cols = ((scale_cols + 3) // 4) * 4 + ref_out_scales = e8m0_scales.new_zeros((scale_rows + n_groups * 128, padded_cols)) + input_group_start = 0 + output_group_start = 0 + for input_group_end in input_group_offsets.tolist(): + group_scales = e8m0_scales[input_group_start:input_group_end] + group_rows_padded = ((group_scales.shape[0] + 127) // 128) * 128 + ref_out_scales[output_group_start : output_group_start + group_rows_padded] = ( + to_blocked(group_scales).view(-1, padded_cols) + ) + input_group_start = input_group_end + output_group_start += group_rows_padded + + cutedsl_out_scales = _mx_block_rearrange_2d_m_groups_cutedsl( + e8m0_scales, + input_group_offsets, + chunk_width, + ) + # byte 0xFF is e8m0 NaN, and NaN != NaN would break torch.equal on the raw dtype + assert torch.equal( + ref_out_scales.view(torch.uint8), cutedsl_out_scales.view(torch.uint8) + ), "blocked scales not equal" + + @pytest.mark.skipif( not _is_sm_10x(), reason="MXFP8 requires CUDA SM 10.x", diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py b/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py index 830e450a834..02ab8e59a59 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/cute_utils.py @@ -9,6 +9,8 @@ import importlib.util import inspect +import torch + # Runtime package detection _CUTEDSL_RUNTIME_PACKAGES = { "cuda.bindings.driver": "cuda-python", @@ -47,6 +49,20 @@ def _cutedsl_runtime_available() -> bool: return len(_missing_cutedsl_runtime_packages()) == 0 +def current_cuda_stream(): + """Return the current PyTorch CUDA stream, to pass to a kernel launch explicitly. + + Read it on every launch, because callers can switch device or stream (for + example, during CUDA graph capture). The raw-stream query skips the Python + ``Stream`` object that ``torch.cuda.current_stream()`` builds. + """ + import cuda.bindings.driver as cuda + + return cuda.CUstream( + torch._C._cuda_getCurrentRawStream(torch.cuda.current_device()) + ) + + if _cutedsl_runtime_available(): import cutlass import cutlass.cute as cute diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_m_groups.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_m_groups.py new file mode 100644 index 00000000000..20a2a727ab9 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_m_groups.py @@ -0,0 +1,656 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import functools + +import torch + +from torchao.utils import ceil_div, round_up + +from .cute_utils import current_cuda_stream + + +@functools.cache +def _compile_mx_block_rearrange_2d_m_groups_cutedsl( + num_groups: int, + chunk_width: int, + cols_multiple_of_16: bool, + all_col_chunks_full: bool, +): + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + import cutlass.utils as utils + from cutlass.cute.nvgpu import cpasync + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + NUM_GROUPS = num_groups + SF_ROWS = 128 + SF_COLS = 4 + CHUNK_WIDTH = chunk_width + THREADS_PER_ROW = CHUNK_WIDTH // 16 + THREADS_PER_BLOCK = SF_ROWS * THREADS_PER_ROW + SF_TILE_BYTES = SF_ROWS * SF_COLS + SF_TILES_PER_CHUNK = CHUNK_WIDTH // SF_COLS + CHUNK_SIZE_BYTES = SF_ROWS * CHUNK_WIDTH + COLS_MULTIPLE_OF_16 = cols_multiple_of_16 + ALL_COL_CHUNKS_FULL = all_col_chunks_full + + @cute.struct + class SharedStorage: + group_data: cute.struct.MemRange[cutlass.Int32, NUM_GROUPS * 2] + group_bounds: cute.struct.MemRange[cutlass.Int32, 4] + out_smem: cute.struct.Align[ + cute.struct.MemRange[cutlass.Uint8, CHUNK_SIZE_BYTES], + 128, + ] + + class MXBlockRearrange2dMGroups: + @cute.jit + def _fill_group_data( + self, + input_offsets: cute.Tensor, + group_data: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + if tidx == 0: + chunk_cumsum = cutlass.Int32(0) + prev_offset = cutlass.Int32(0) + for g in range(NUM_GROUPS): + end = cutlass.Int32((input_offsets.iterator + g).load()) + group_size = end - prev_offset + chunks_in_group = ceil_div(group_size, SF_ROWS) + chunk_cumsum = chunk_cumsum + chunks_in_group + (group_data.iterator + g).store(cutlass.Int32(chunks_in_group)) + (group_data.iterator + NUM_GROUPS + g).store( + cutlass.Int32(chunk_cumsum) + ) + prev_offset = end + cute.arch.sync_threads() + + @cute.jit + def _fill_group_bounds( + self, + input_offsets: cute.Tensor, + row_chunk: cutlass.Int32, + group_data: cute.Tensor, + group_bounds: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + if tidx == 0: + group = cutlass.Int32(0) + chunk_idx = cutlass.Int32(0) + active = cutlass.Int32(0) + chunk_cumsum_before = cutlass.Int32(0) + for g in range(NUM_GROUPS): + chunk_cumsum = cutlass.Int32( + (group_data.iterator + NUM_GROUPS + g).load() + ) + if row_chunk < chunk_cumsum and active == 0: + group = cutlass.Int32(g) + chunk_idx = row_chunk - chunk_cumsum_before + active = cutlass.Int32(1) + chunk_cumsum_before = chunk_cumsum + group_start = cutlass.Int32(0) + group_end = cutlass.Int32(0) + input_row_base = cutlass.Int32(0) + if active != 0: + if group > 0: + group_start = cutlass.Int32( + (input_offsets.iterator + group - 1).load() + ) + group_end = cutlass.Int32((input_offsets.iterator + group).load()) + input_row_base = group_start + chunk_idx * SF_ROWS + (group_bounds.iterator + 0).store(group_start) + (group_bounds.iterator + 1).store(group_end) + (group_bounds.iterator + 2).store(input_row_base) + (group_bounds.iterator + 3).store(active) + cute.arch.sync_threads() + + @cute.jit + def _load4( + self, + input_scales: cute.Tensor, + input_offset: cutlass.Int32, + input_col: cutlass.Int32, + cols: cutlass.Int32, + ): + value = cutlass.Uint32(0) + if input_col + 3 < cols and input_offset % 4 == 0: + value = ( + cute.recast_ptr( + (input_scales.iterator + input_offset).align(4), + dtype=cutlass.Uint32, + ) + .load() + .to(cutlass.Uint32) + ) + elif input_col < cols: + value = (input_scales.iterator + input_offset).load().to(cutlass.Uint32) + if input_col + 1 < cols: + value = ( + value + | ( + (input_scales.iterator + input_offset + 1) + .load() + .to(cutlass.Uint32) + << 8 + ) + ).to(cutlass.Uint32) + if input_col + 2 < cols: + value = ( + value + | ( + (input_scales.iterator + input_offset + 2) + .load() + .to(cutlass.Uint32) + << 16 + ) + ).to(cutlass.Uint32) + if input_col + 3 < cols: + value = ( + value + | ( + (input_scales.iterator + input_offset + 3) + .load() + .to(cutlass.Uint32) + << 24 + ) + ).to(cutlass.Uint32) + return value + + @cute.jit + def _process_chunk( + self, + input_scales: cute.Tensor, + input_offsets: cute.Tensor, + output_scales: cute.Tensor, + rows: cutlass.Int32, + cols: cutlass.Int32, + padded_cols: cutlass.Int32, + row_chunk: cutlass.Int32, + group_data: cute.Tensor, + group_bounds: cute.Tensor, + sOUT_tile: cute.Tensor, + bulk_store_atom: cute.CopyAtom, + ): + tidx, _, _ = cute.arch.thread_idx() + col_chunk_pid, _, _ = cute.arch.block_idx() + col_chunk = cutlass.Int32(col_chunk_pid) + row_idx = cutlass.Int32(tidx // THREADS_PER_ROW) + col_lane = cutlass.Int32(tidx - (tidx // THREADS_PER_ROW) * THREADS_PER_ROW) + self._fill_group_bounds( + input_offsets, + row_chunk, + group_data, + group_bounds, + ) + group_end = cutlass.Int32((group_bounds.iterator + 1).load()) + input_row_base = cutlass.Int32((group_bounds.iterator + 2).load()) + active = cutlass.Int32((group_bounds.iterator + 3).load()) + input_row = cutlass.Int32(0) + chunk_col_base = cutlass.Int32(0) + cols_in_chunk = cutlass.Int32(0) + if active != 0: + input_row = input_row_base + row_idx + chunk_col_base = col_chunk * CHUNK_WIDTH + col_lane * 16 + if not ALL_COL_CHUNKS_FULL: + cols_in_chunk = cols - col_chunk * CHUNK_WIDTH + if cols_in_chunk > CHUNK_WIDTH: + cols_in_chunk = cutlass.Int32(CHUNK_WIDTH) + + row_valid = cutlass.Int32(0) + if active != 0 and input_row < group_end and input_row < rows: + row_valid = cutlass.Int32(1) + full_col_chunk = cutlass.Int32(0) + if not ALL_COL_CHUNKS_FULL: + if cols_in_chunk == CHUNK_WIDTH: + full_col_chunk = cutlass.Int32(1) + col_valid = cutlass.Int32(0) + if ALL_COL_CHUNKS_FULL: + if active != 0: + col_valid = cutlass.Int32(1) + else: + if active != 0 and col_lane * 16 < cols_in_chunk: + col_valid = cutlass.Int32(1) + + if col_valid != 0: + input_offset = input_row * cols + chunk_col_base + r_div_32 = row_idx // 32 + r_mod_32 = row_idx - r_div_32 * 32 + blocked_layout_row_offset = r_mod_32 * 16 + r_div_32 * 4 + thread_col_start = col_lane * 16 + smem_offset0 = blocked_layout_row_offset + col_lane * 4 * SF_TILE_BYTES + smem_offset1 = ( + blocked_layout_row_offset + (col_lane * 4 + 1) * SF_TILE_BYTES + ) + smem_offset2 = ( + blocked_layout_row_offset + (col_lane * 4 + 2) * SF_TILE_BYTES + ) + smem_offset3 = ( + blocked_layout_row_offset + (col_lane * 4 + 3) * SF_TILE_BYTES + ) + value0 = cutlass.Uint32(0) + value1 = cutlass.Uint32(0) + value2 = cutlass.Uint32(0) + value3 = cutlass.Uint32(0) + if ALL_COL_CHUNKS_FULL: + if row_valid != 0: + input_ptr = (input_scales.iterator + input_offset).align(16) + value = cute.recast_ptr( + input_ptr, + dtype=cutlass.Int128, + ) + value = value.load() + value0 = value.to(cutlass.Uint32) + value1 = (value >> 32).to(cutlass.Uint32) + value2 = (value >> 64).to(cutlass.Uint32) + value3 = (value >> 96).to(cutlass.Uint32) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset0, dtype=cutlass.Uint32 + ).store(value0) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset1, dtype=cutlass.Uint32 + ).store(value1) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset2, dtype=cutlass.Uint32 + ).store(value2) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset3, dtype=cutlass.Uint32 + ).store(value3) + elif COLS_MULTIPLE_OF_16: + if full_col_chunk != 0: + if row_valid != 0: + input_ptr = (input_scales.iterator + input_offset).align(16) + value = cute.recast_ptr( + input_ptr, + dtype=cutlass.Int128, + ) + value = value.load() + value0 = value.to(cutlass.Uint32) + value1 = (value >> 32).to(cutlass.Uint32) + value2 = (value >> 64).to(cutlass.Uint32) + value3 = (value >> 96).to(cutlass.Uint32) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset0, dtype=cutlass.Uint32 + ).store(value0) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset1, dtype=cutlass.Uint32 + ).store(value1) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset2, dtype=cutlass.Uint32 + ).store(value2) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset3, dtype=cutlass.Uint32 + ).store(value3) + else: + if row_valid != 0: + value0 = self._load4( + input_scales, + input_offset, + chunk_col_base, + cols, + ) + value1 = self._load4( + input_scales, + input_offset + 4, + chunk_col_base + 4, + cols, + ) + value2 = self._load4( + input_scales, + input_offset + 8, + chunk_col_base + 8, + cols, + ) + value3 = self._load4( + input_scales, + input_offset + 12, + chunk_col_base + 12, + cols, + ) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset0, dtype=cutlass.Uint32 + ).store(value0) + if thread_col_start + 4 < cols_in_chunk: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset1, + dtype=cutlass.Uint32, + ).store(value1) + if thread_col_start + 8 < cols_in_chunk: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset2, + dtype=cutlass.Uint32, + ).store(value2) + if thread_col_start + 12 < cols_in_chunk: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset3, + dtype=cutlass.Uint32, + ).store(value3) + else: + if ( + row_valid != 0 + and chunk_col_base + 8 <= cols + and input_offset % 8 == 0 + ): + input_ptr = (input_scales.iterator + input_offset).align(8) + value = cute.recast_ptr(input_ptr, dtype=cutlass.Uint64).load() + value0 = (value & cutlass.Uint64(0xFFFFFFFF)).to(cutlass.Uint32) + value1 = ((value >> 32) & cutlass.Uint64(0xFFFFFFFF)).to( + cutlass.Uint32 + ) + value2 = self._load4( + input_scales, + input_offset + 8, + chunk_col_base + 8, + cols, + ) + value3 = self._load4( + input_scales, + input_offset + 12, + chunk_col_base + 12, + cols, + ) + elif row_valid != 0: + value0 = self._load4( + input_scales, + input_offset, + chunk_col_base, + cols, + ) + value1 = self._load4( + input_scales, + input_offset + 4, + chunk_col_base + 4, + cols, + ) + value2 = self._load4( + input_scales, + input_offset + 8, + chunk_col_base + 8, + cols, + ) + value3 = self._load4( + input_scales, + input_offset + 12, + chunk_col_base + 12, + cols, + ) + cute.recast_ptr( + sOUT_tile.iterator + smem_offset0, dtype=cutlass.Uint32 + ).store(value0) + if thread_col_start + 4 < cols_in_chunk: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset1, dtype=cutlass.Uint32 + ).store(value1) + if thread_col_start + 8 < cols_in_chunk: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset2, dtype=cutlass.Uint32 + ).store(value2) + if thread_col_start + 12 < cols_in_chunk: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset3, dtype=cutlass.Uint32 + ).store(value3) + + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + if tidx == 0 and active != 0: + sf_tiles_per_row = padded_cols // SF_COLS + out_tile_base = ( + row_chunk * sf_tiles_per_row + col_chunk * SF_TILES_PER_CHUNK + ) + if ALL_COL_CHUNKS_FULL: + for tile in range(SF_TILES_PER_CHUNK): + src = cute.make_tensor( + sOUT_tile.iterator + tile * SF_TILE_BYTES, + cute.make_layout((SF_TILE_BYTES,), stride=(1,)), + ) + dst = cute.make_tensor( + output_scales.iterator + + (out_tile_base + tile) * SF_TILE_BYTES, + cute.make_layout((SF_TILE_BYTES,), stride=(1,)), + ) + cute.copy(bulk_store_atom, src, dst) + elif full_col_chunk != 0: + for tile in range(SF_TILES_PER_CHUNK): + src = cute.make_tensor( + sOUT_tile.iterator + tile * SF_TILE_BYTES, + cute.make_layout((SF_TILE_BYTES,), stride=(1,)), + ) + dst = cute.make_tensor( + output_scales.iterator + + (out_tile_base + tile) * SF_TILE_BYTES, + cute.make_layout((SF_TILE_BYTES,), stride=(1,)), + ) + cute.copy(bulk_store_atom, src, dst) + else: + valid_sf_tiles = ceil_div(cols_in_chunk, SF_COLS) + for tile in range(SF_TILES_PER_CHUNK): + if tile < valid_sf_tiles: + src = cute.make_tensor( + sOUT_tile.iterator + tile * SF_TILE_BYTES, + cute.make_layout((SF_TILE_BYTES,), stride=(1,)), + ) + dst = cute.make_tensor( + output_scales.iterator + + (out_tile_base + tile) * SF_TILE_BYTES, + cute.make_layout((SF_TILE_BYTES,), stride=(1,)), + ) + cute.copy(bulk_store_atom, src, dst) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @cute.kernel + def kernel( + self, + input_scales: cute.Tensor, + input_offsets: cute.Tensor, + output_scales: cute.Tensor, + rows: cutlass.Int32, + cols: cutlass.Int32, + padded_cols: cutlass.Int32, + bulk_store_atom: cute.CopyAtom, + ): + _, row_superblock_pid, _ = cute.arch.block_idx() + row_superblock = cutlass.Int32(row_superblock_pid) + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(SharedStorage) + group_data = storage.group_data.get_tensor( + cute.make_layout((NUM_GROUPS * 2,), stride=(1,)) + ) + group_bounds = storage.group_bounds.get_tensor( + cute.make_layout((4,), stride=(1,)) + ) + sOUT_tile = storage.out_smem.get_tensor( + cute.make_layout((CHUNK_SIZE_BYTES,), stride=(1,)) + ) + self._fill_group_data(input_offsets, group_data) + self._process_chunk( + input_scales, + input_offsets, + output_scales, + rows, + cols, + padded_cols, + row_superblock, + group_data, + group_bounds, + sOUT_tile, + bulk_store_atom, + ) + + @cute.jit + def __call__( + self, + input_scales: cute.Tensor, + input_offsets: cute.Tensor, + output_scales: cute.Tensor, + rows: cutlass.Int32, + cols: cutlass.Int32, + num_row_chunks: cutlass.Int32, + num_col_chunks: cutlass.Int32, + padded_cols: cutlass.Int32, + stream: cuda.CUstream, + ): + bulk_store_atom = cute.make_copy_atom( + cpasync.CopyBulkS2GOp(), + cutlass.Uint8, + num_bits_per_copy=SF_TILE_BYTES * 8, + ) + self.kernel( + input_scales, + input_offsets, + output_scales, + rows, + cols, + padded_cols, + bulk_store_atom, + ).launch( + grid=(num_col_chunks, num_row_chunks, 1), + block=(THREADS_PER_BLOCK, 1, 1), + cluster=(1, 1, 1), + smem=SharedStorage.size_in_bytes(), + stream=stream, + ) + + kernel = MXBlockRearrange2dMGroups() + + m = cute.sym_int(divisibility=128) + c = cute.sym_int(divisibility=16) if cols_multiple_of_16 else cute.sym_int() + out_n = cute.sym_int() + offs_stride = cute.sym_int() + out_stride = cute.sym_int() + + fake_input = make_fake_tensor( + cutlass.Uint8, + (m, c), + stride=(c, 1), + assumed_align=16, + ) + fake_output = make_fake_tensor( + cutlass.Uint8, + (out_n,), + stride=(out_stride,), + ) + fake_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_stream = make_fake_stream() + + return cute.compile( + kernel, + input_scales=fake_input, + input_offsets=fake_offsets, + output_scales=fake_output, + rows=0, + cols=0, + num_row_chunks=1, + num_col_chunks=1, + padded_cols=16, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + +def _mx_block_rearrange_2d_m_groups_cutedsl_impl( + scales_tensor: torch.Tensor, + input_offsets: torch.Tensor, + chunk_width: int | None = None, +) -> torch.Tensor: + assert scales_tensor.ndim == 2 + assert scales_tensor.element_size() == 1 + assert input_offsets.ndim == 1 + assert input_offsets.dtype == torch.int32 + assert scales_tensor.is_cuda + assert input_offsets.is_cuda + assert scales_tensor.is_contiguous() + rows, cols = scales_tensor.shape + num_groups = input_offsets.shape[0] + padded_rows = rows + num_groups * 128 + padded_cols = round_up(cols, 4) + if chunk_width is None: + if cols >= 64: + chunk_width = 64 + elif cols >= 32: + chunk_width = 32 + else: + chunk_width = 16 + assert chunk_width in (16, 32, 64, 128) + # zeros, not empty: trailing padded rows beyond the last group's chunks are + # never covered by a launched block, so they are only zero because of this. + output = torch.zeros( + (padded_rows, padded_cols), + device=scales_tensor.device, + dtype=torch.uint8, + ) + compiled = _compile_mx_block_rearrange_2d_m_groups_cutedsl( + num_groups, + chunk_width, + cols % 16 == 0, + cols % chunk_width == 0, + ) + stream = current_cuda_stream() + compiled( + scales_tensor.view(torch.uint8), + input_offsets, + output.view(-1), + int(rows), + int(cols), + # + num_groups: each group's ceil-rounding to 128 rows can add at most one + # extra chunk. Over-estimating keeps the grid independent of the offsets + # (no d2h sync); the surplus blocks early-out on active == 0. + int(ceil_div(rows, 128) + num_groups), + int(ceil_div(cols, chunk_width)), + int(padded_cols), + stream, + ) + return output.view(scales_tensor.dtype) + + +@torch.library.custom_op( + "torchao::mx_block_rearrange_2d_m_groups_cutedsl", + mutates_args=(), +) +def _mx_block_rearrange_2d_m_groups_cutedsl_custom_op( + scales_tensor: torch.Tensor, + input_offsets: torch.Tensor, + chunk_width: int, +) -> torch.Tensor: + return _mx_block_rearrange_2d_m_groups_cutedsl_impl( + scales_tensor, + input_offsets, + None if chunk_width == 0 else chunk_width, + ) + + +@_mx_block_rearrange_2d_m_groups_cutedsl_custom_op.register_fake +def _fake_mx_block_rearrange_2d_m_groups_cutedsl_custom_op( + scales_tensor: torch.Tensor, + input_offsets: torch.Tensor, + chunk_width: int, +) -> torch.Tensor: + assert scales_tensor.ndim == 2 + assert scales_tensor.element_size() == 1 + assert input_offsets.ndim == 1 + rows, cols = scales_tensor.shape + num_groups = input_offsets.shape[0] + padded_rows = rows + num_groups * 128 + padded_cols = round_up(cols, 4) + return scales_tensor.new_empty((padded_rows, padded_cols)) + + +def _mx_block_rearrange_2d_m_groups_cutedsl( + scales_tensor: torch.Tensor, + input_offsets: torch.Tensor, + chunk_width: int | None = None, +) -> torch.Tensor: + return _mx_block_rearrange_2d_m_groups_cutedsl_custom_op( + scales_tensor, + input_offsets, + 0 if chunk_width is None else chunk_width, + )