diff --git a/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_K_groups.py b/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_K_groups.py index eeda0faab31..7810f6cf6df 100644 --- a/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_K_groups.py +++ b/benchmarks/prototype/moe_training/mxfp8/bench_mx_block_rearrange_2d_K_groups.py @@ -3,180 +3,195 @@ # # 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 ( - torch_to_blocked_2d_K_groups, + _mx_block_rearrange_2d_k_groups_cutedsl, triton_mx_block_rearrange_2d_K_groups, ) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( + torch_to_blocked_2d_K_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 - - -@dataclass(frozen=True) -class ExperimentResult: - torch_time_us: float - triton_time_us: float - torch_mem_bw_gbps: float - triton_mem_bw_gbps: float +from torchao.utils import ceil_div + + +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 benchmark_cuda_function_in_microseconds(f, *args, **kwargs) -@dataclass(frozen=True) -class Experiment: - config: ExperimentConfig - result: ExperimentResult +def _parse_shape(spec: str): + """Parses `MxTOTALKxGROUPS` or `LABEL:MxTOTALKxGROUPS`.""" + label, _, dims = spec.rpartition(":") + parts = dims.split("x") + if len(parts) != 3: + raise SystemExit(f"Malformed shape {spec!r}, expected [LABEL:]MxTOTALKxGROUPS") + m, total_k, groups = (int(p) for p in parts) + return (label or None, m, total_k, groups) -def get_configs() -> List[ExperimentConfig]: - block_size = 32 - # DeepSeekV3 671b shapes. - # total tokens and groups are along the contracting dimension - input_shapes = [ - (7168, 131072 // block_size), - (2048, 131072 // block_size), - (7168, 65536 // block_size), - (2048, 65536 // block_size), - (7168, 32768 // block_size), - (2048, 32768 // 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 ] - num_groups_list = [4, 8] - - configs = [] - for shape, groups in itertools.product( - input_shapes, - num_groups_list, - ): - configs.append( - ExperimentConfig( - input_shape=shape, - num_groups=groups, - ) - ) - return configs - - -def run_experiment(config: ExperimentConfig) -> ExperimentResult: - input_shape, num_groups = config.input_shape, config.num_groups + 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]) - input_tensor = torch.randint( - low=0, - high=256, - size=input_shape, - dtype=torch.uint8, - device=device, - ) - M, scale_cols = input_shape - block_size = 32 - # total_K is the total number of elements along K dimension - total_K = scale_cols * block_size - - # Generate group end offsets along total_K, then divide by block_size to get scale group end offsets - input_group_offsets = generate_jagged_offs( - num_groups, total_K, multiple_of=block_size - ) - scale_group_offsets = input_group_offsets // block_size - - # bench torch (torch.compile hangson this, can't use) - torch_out_scales, torch_group_offs = torch_to_blocked_2d_K_groups( - input_tensor, - scale_group_offsets, - block_size=block_size, +def main(): + parser = argparse.ArgumentParser( + description="Benchmark torch / triton / CuTeDSL MXFP8 K-groups scale rearrange." ) - torch_time_us = benchmark_cuda_function_in_microseconds( - torch_to_blocked_2d_K_groups, - input_tensor, - scale_group_offsets, - block_size=block_size, + parser.add_argument( + "shapes", + nargs="+", + metavar="[LABEL:]MxTOTALKxGROUPS", + help="shapes to benchmark, e.g. 2048x4096x8 or dsv3:2048x4096x8", ) - - # bench triton - triton_out_scales = triton_mx_block_rearrange_2d_K_groups( - input_tensor, - scale_group_offsets, + parser.add_argument("--multiple-of", type=int, default=1) + parser.add_argument( + "--chunk-width", + type=int, + default=None, + help="CuTeDSL column chunk width; default lets the kernel choose", ) - triton_time_us = benchmark_cuda_function_in_microseconds( - triton_mx_block_rearrange_2d_K_groups, - input_tensor, - scale_group_offsets, + parser.add_argument( + "--bench-only", + nargs="+", + choices=("torch", "triton", "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) - - return ExperimentResult( - torch_time_us=torch_time_us, - triton_time_us=triton_time_us, - torch_mem_bw_gbps=torch_mem_bw_gbps, - triton_mem_bw_gbps=triton_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", "cutedsl"} ) - - -def print_results(experiments: List[Experiment]): - headers = [ - "input_shape", - "num_groups", - "torch_time_us", - "triton_time_us", - "triton_speedup", - "torch_mem_bw_gbps", - "triton_mem_bw_gbps", - ] - 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_K_groups calls .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.num_groups, - f"{experiment.result.torch_time_us:.2f}", - f"{experiment.result.triton_time_us:.2f}", - f"{experiment.result.torch_time_us / experiment.result.triton_time_us:.2f}x", - f"{experiment.result.torch_mem_bw_gbps:.2f}", - f"{experiment.result.triton_mem_bw_gbps:.2f}", - ] + selected.discard("torch") + block_size = 32 + rows_out = [] + + for label, m, total_k, groups in shapes: + scale_cols = ceil_div(total_k, block_size) + if args.deterministic: + raw = ( + torch.arange(m * scale_cols, device="cuda", dtype=torch.int32) + .remainder(251) + .to(torch.uint8) + .reshape(m, scale_cols) + ) + else: + raw = torch.randint( + 0, 256, (m, scale_cols), device="cuda", dtype=torch.uint8 + ) + scales = raw.view(torch.float8_e8m0fnu) + offsets = generate_jagged_offs( + groups, scale_cols, multiple_of=args.multiple_of, device="cuda" ) - print(tabulate(rows, headers=headers)) + row = {"Model layer": label or ""} if labelled else {} + row["M x TotalK"] = f"{m}x{total_k}" + row["Groups"] = groups + + if "torch" in selected: + torch_us = benchmark_function( + torch_to_blocked_2d_K_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_K_groups, + scales, + offsets, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + row["Triton us"] = f"{triton_us:.2f}" + if "cutedsl" in selected: + out = _mx_block_rearrange_2d_k_groups_cutedsl( + scales, offsets, args.chunk_width + ) + cutedsl_us = benchmark_function( + _mx_block_rearrange_2d_k_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 582b6e6ba73..068d1b5a4c4 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -9,7 +9,7 @@ from torch._subclasses.fake_tensor import FakeTensorMode # 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: @@ -59,6 +59,9 @@ def _is_sm_10x() -> bool: _pad_token_groups_cutedsl, _unpad_token_groups_cutedsl, ) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_k_groups import ( + _mx_block_rearrange_2d_k_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( _mx_block_rearrange_2d_m_groups_cutedsl, ) @@ -431,41 +434,69 @@ def test_mxfp8_per_group_blocked_scales_3d( ) +# Cases listed per impl rather than crossed, so merging the triton and cutedsl +# tests preserves exactly the coverage each had. Scales are raw bytes: a rearrange +# only moves bytes, so their values are irrelevant, and both impls share the same +# torch reference. +_K_GROUPS_CASES = [ + *( + ("triton", m, total_k, n_groups) + for m in (256, 512, 1024, 5120) + for total_k in (512, 1024, 2048, 4096, 8192, 16384) + for n_groups in (1, 4, 8, 16) + ), + # 4097 exercises scale_cols not a multiple of 4, and can yield empty groups + *( + ("cutedsl", m, total_k, n_groups) + for m, total_k, n_groups in ( + (256, 512, 1), + (256, 1024, 4), + (512, 2048, 8), + (1024, 4096, 8), + (1024, 4097, 8), + ) + ), +] + + @skip_if_rocm("ROCm enablement in progress") -@pytest.mark.parametrize("m", [256, 512, 1024, 5120]) -@pytest.mark.parametrize("total_k", [512, 1024, 2048, 4096, 8192, 16384]) -@pytest.mark.parametrize("n_groups", [1, 4, 8, 16]) -def test_triton_mx_block_rearrange_2d_K_groups( +@pytest.mark.parametrize(("impl", "m", "total_k", "n_groups"), _K_GROUPS_CASES) +def test_mx_block_rearrange_2d_K_groups( + impl: str, m: int, total_k: int, n_groups: int, ): + if impl == "cutedsl" and not _mxfp8_cutedsl_kernels_available: + pytest.skip("CuteDSL MXFP8 kernels are unavailable") device = "cuda" block_size = 32 - input_data = torch.randn(m, total_k, device=device) - - e8m0_scales, _ = to_mx( - input_data, elem_dtype=torch.float8_e4m3fn, block_size=block_size + scale_cols = ceil_div(total_k, block_size) + e8m0_scales = ( + torch.arange(m * scale_cols, device=device, dtype=torch.int32) + .remainder(251) + .to(torch.uint8) + .reshape(m, scale_cols) + .view(torch.float8_e8m0fnu) ) - - # Generate group end offsets along total_K, then divide by block_size to get scale group end offsets - input_group_offsets = generate_jagged_offs( - n_groups, total_k, multiple_of=block_size, device=device + scale_group_offsets = generate_jagged_offs( + n_groups, + scale_cols, + multiple_of=1, + device=device, ) - scale_group_offsets = input_group_offsets // block_size - # torch reference - ref_out_scales, ref_start_cols_after_padding = torch_to_blocked_2d_K_groups( + ref_out_scales, _ = torch_to_blocked_2d_K_groups( e8m0_scales, scale_group_offsets, ) - - # triton kernel - triton_out_scales = triton_mx_block_rearrange_2d_K_groups( - e8m0_scales, - scale_group_offsets, + kernel = ( + _mx_block_rearrange_2d_k_groups_cutedsl + if impl == "cutedsl" + else triton_mx_block_rearrange_2d_K_groups ) - assert torch.equal(ref_out_scales, triton_out_scales), "blocked scales not equal" + out_scales = kernel(e8m0_scales, scale_group_offsets) + assert torch.equal(ref_out_scales, out_scales), "blocked scales not equal" @pytest.mark.skipif( diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index ceb0c1459e2..f51fb575584 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -2,6 +2,9 @@ _pad_token_groups_cutedsl, _unpad_token_groups_cutedsl, ) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_k_groups import ( # noqa: F401 + _mx_block_rearrange_2d_k_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.quant import ( _mxfp8_cuda_kernels_available, # noqa: F401 _mxfp8_flydsl_kernels_available, # noqa: F401 diff --git a/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_k_groups.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_k_groups.py new file mode 100644 index 00000000000..9f1e8a28006 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_k_groups.py @@ -0,0 +1,450 @@ +# 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_k_groups_cutedsl( + num_groups: int, + chunk_width: int, + cols_multiple_of_16: 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 + + @cute.struct + class SharedStorage: + out_smem: cute.struct.Align[ + cute.struct.MemRange[cutlass.Uint8, CHUNK_SIZE_BYTES], + 128, + ] + + class MXBlockRearrange2dKGroups: + @cute.jit + def _group_start_after_padding( + self, + input_offsets: cute.Tensor, + group: cutlass.Int32, + ): + start_after_padding = cutlass.Int32(0) + prev = cutlass.Int32(0) + for g in range(NUM_GROUPS): + end = cutlass.Int32((input_offsets.iterator + g).load()) + if cutlass.Int32(g) < group: + group_size = end - prev + start_after_padding = start_after_padding + round_up( + group_size, SF_COLS + ) + prev = end + return start_after_padding + + @cute.jit + def _load4( + self, + input_scales: cute.Tensor, + input_offset: cutlass.Int32, + input_col: cutlass.Int32, + group_end: cutlass.Int32, + ): + value = cutlass.Uint32(0) + if input_col + 3 < group_end and input_offset % 4 == 0: + value = ( + cute.recast_ptr( + input_scales.iterator + input_offset, + dtype=cutlass.Uint32, + ) + .load() + .to(cutlass.Uint32) + ) + elif input_col < group_end: + value = (input_scales.iterator + input_offset).load().to(cutlass.Uint32) + if input_col + 1 < group_end: + value = ( + value + | ( + (input_scales.iterator + input_offset + 1) + .load() + .to(cutlass.Uint32) + << 8 + ) + ).to(cutlass.Uint32) + if input_col + 2 < group_end: + value = ( + value + | ( + (input_scales.iterator + input_offset + 2) + .load() + .to(cutlass.Uint32) + << 16 + ) + ).to(cutlass.Uint32) + if input_col + 3 < group_end: + value = ( + value + | ( + (input_scales.iterator + input_offset + 3) + .load() + .to(cutlass.Uint32) + << 24 + ) + ).to(cutlass.Uint32) + return value + + @cute.kernel + def kernel( + self, + input_scales: cute.Tensor, + input_offsets: cute.Tensor, + output_scales: cute.Tensor, + rows: cutlass.Int32, + cols: cutlass.Int32, + padded_rows: cutlass.Int32, + bulk_store_atom: cute.CopyAtom, + ): + flat_col_chunk_pid, row_block_pid, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + flat_col_chunk = cutlass.Int32(flat_col_chunk_pid) + row_idx = cutlass.Int32(tidx // THREADS_PER_ROW) + col_lane = cutlass.Int32(tidx - (tidx // THREADS_PER_ROW) * THREADS_PER_ROW) + row = cutlass.Int32(row_block_pid * SF_ROWS + row_idx) + + group = cutlass.Int32(0) + group_start = cutlass.Int32(0) + group_end = cutlass.Int32(0) + col_chunk = cutlass.Int32(0) + active = cutlass.Int32(0) + chunk_cumsum = cutlass.Int32(0) + prev_end = cutlass.Int32(0) + for g in range(NUM_GROUPS): + end = cutlass.Int32((input_offsets.iterator + g).load()) + group_size_for_scan = end - prev_end + chunks_in_group = ceil_div(group_size_for_scan, CHUNK_WIDTH) + next_chunk_cumsum = chunk_cumsum + chunks_in_group + if flat_col_chunk < next_chunk_cumsum and active == 0: + group = cutlass.Int32(g) + group_start = prev_end + group_end = end + col_chunk = flat_col_chunk - chunk_cumsum + active = cutlass.Int32(1) + chunk_cumsum = next_chunk_cumsum + prev_end = end + + group_size = group_end - group_start + group_col_blocks = ceil_div(group_size, SF_COLS) + + smem_allocator = utils.SmemAllocator() + storage = smem_allocator.allocate(SharedStorage) + sOUT_tile = storage.out_smem.get_tensor( + cute.make_layout((SF_TILE_BYTES,), stride=(1,)) + ) + + value = cutlass.Uint32(0) + value1 = cutlass.Uint32(0) + value2 = cutlass.Uint32(0) + value3 = cutlass.Uint32(0) + input_col = group_start + col_chunk * CHUNK_WIDTH + col_lane * 16 + if active != 0 and row < rows: + input_offset = row * cols + input_col + if ( + COLS_MULTIPLE_OF_16 + and input_col + 15 < group_end + and input_offset % 16 == 0 + ): + input_ptr = (input_scales.iterator + input_offset).align(16) + packed = cute.recast_ptr(input_ptr, dtype=cutlass.Int128).load() + value = packed.to(cutlass.Uint32) + value1 = (packed >> 32).to(cutlass.Uint32) + value2 = (packed >> 64).to(cutlass.Uint32) + value3 = (packed >> 96).to(cutlass.Uint32) + else: + value = self._load4( + input_scales, + input_offset, + input_col, + group_end, + ) + value1 = self._load4( + input_scales, + input_offset + 4, + input_col + 4, + group_end, + ) + value2 = self._load4( + input_scales, + input_offset + 8, + input_col + 8, + group_end, + ) + value3 = self._load4( + input_scales, + input_offset + 12, + input_col + 12, + group_end, + ) + + r_div_32 = row_idx // 32 + r_mod_32 = row_idx - r_div_32 * 32 + smem_offset = r_mod_32 * 16 + r_div_32 * 4 + smem_offset0 = smem_offset + col_lane * 4 * SF_TILE_BYTES + smem_offset1 = smem_offset + (col_lane * 4 + 1) * SF_TILE_BYTES + smem_offset2 = smem_offset + (col_lane * 4 + 2) * SF_TILE_BYTES + smem_offset3 = smem_offset + (col_lane * 4 + 3) * SF_TILE_BYTES + cute.recast_ptr( + sOUT_tile.iterator + smem_offset0, dtype=cutlass.Uint32 + ).store(value) + if input_col + 4 < group_end: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset1, dtype=cutlass.Uint32 + ).store(value1) + if input_col + 8 < group_end: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset2, dtype=cutlass.Uint32 + ).store(value2) + if input_col + 12 < group_end: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset3, dtype=cutlass.Uint32 + ).store(value3) + if input_col + 4 >= group_end and col_lane * 16 + 4 < CHUNK_WIDTH: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset1, dtype=cutlass.Uint32 + ).store(cutlass.Uint32(0)) + if input_col + 8 >= group_end and col_lane * 16 + 8 < CHUNK_WIDTH: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset2, dtype=cutlass.Uint32 + ).store(cutlass.Uint32(0)) + if input_col + 12 >= group_end and col_lane * 16 + 12 < CHUNK_WIDTH: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset3, dtype=cutlass.Uint32 + ).store(cutlass.Uint32(0)) + if input_col >= group_end: + cute.recast_ptr( + sOUT_tile.iterator + smem_offset0, dtype=cutlass.Uint32 + ).store(cutlass.Uint32(0)) + + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.sync_threads() + if tidx == 0 and active != 0: + group_start_after_padding = self._group_start_after_padding( + input_offsets, + group, + ) + out_group_base = group_start_after_padding * padded_rows + out_base = ( + out_group_base + + row_block_pid * group_col_blocks * SF_TILE_BYTES + + col_chunk * SF_TILES_PER_CHUNK * SF_TILE_BYTES + ) + valid_tiles = group_col_blocks - col_chunk * SF_TILES_PER_CHUNK + if valid_tiles > SF_TILES_PER_CHUNK: + valid_tiles = cutlass.Int32(SF_TILES_PER_CHUNK) + for tile in range(SF_TILES_PER_CHUNK): + if tile < valid_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_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.jit + def __call__( + self, + input_scales: cute.Tensor, + input_offsets: cute.Tensor, + output_scales: cute.Tensor, + rows: cutlass.Int32, + cols: cutlass.Int32, + padded_rows: cutlass.Int32, + num_row_blocks: cutlass.Int32, + max_active_col_chunks: 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_rows, + bulk_store_atom, + ).launch( + grid=(max_active_col_chunks, num_row_blocks, 1), + block=(THREADS_PER_BLOCK, 1, 1), + cluster=(1, 1, 1), + smem=SharedStorage.size_in_bytes(), + stream=stream, + ) + + kernel = MXBlockRearrange2dKGroups() + + 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, + padded_rows=128, + num_row_blocks=1, + max_active_col_chunks=1, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + +def _mx_block_rearrange_2d_k_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 = round_up(rows, 128) + padded_cols = cols + num_groups * 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: inter-group column padding and rows beyond the last + # group's chunks are never written by a launched block, so they are only + # zero because of this allocation. + output = torch.zeros( + (padded_rows, padded_cols), + device=scales_tensor.device, + dtype=torch.uint8, + ) + compiled = _compile_mx_block_rearrange_2d_k_groups_cutedsl( + num_groups, + chunk_width, + cols % 16 == 0, + ) + stream = current_cuda_stream() + compiled( + scales_tensor.view(torch.uint8), + input_offsets, + output.view(-1), + int(rows), + int(cols), + int(padded_rows), + int(ceil_div(rows, 128)), + # + num_groups: each group's chunk count rounds up independently, so the + # column-chunk grid is over-estimated to stay offset-independent (no d2h + # sync). Surplus blocks find active == 0 and skip the store. + int(ceil_div(cols, chunk_width) + num_groups), + stream, + ) + return output.view(scales_tensor.dtype) + + +@torch.library.custom_op( + "torchao::mx_block_rearrange_2d_k_groups_cutedsl", + mutates_args=(), +) +def _mx_block_rearrange_2d_k_groups_cutedsl_custom_op( + scales_tensor: torch.Tensor, + input_offsets: torch.Tensor, + chunk_width: int, +) -> torch.Tensor: + return _mx_block_rearrange_2d_k_groups_cutedsl_impl( + scales_tensor, + input_offsets, + None if chunk_width == 0 else chunk_width, + ) + + +@_mx_block_rearrange_2d_k_groups_cutedsl_custom_op.register_fake +def _fake_mx_block_rearrange_2d_k_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 = round_up(rows, 128) + padded_cols = cols + num_groups * 4 + return scales_tensor.new_empty((padded_rows, padded_cols)) + + +def _mx_block_rearrange_2d_k_groups_cutedsl( + scales_tensor: torch.Tensor, + input_offsets: torch.Tensor, + chunk_width: int | None = None, +) -> torch.Tensor: + return _mx_block_rearrange_2d_k_groups_cutedsl_custom_op( + scales_tensor, + input_offsets, + 0 if chunk_width is None else chunk_width, + )