From d98e30437e6a8abe3c6eb08b2295e73863acf081 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Samard=C5=BEi=C4=87?= Date: Wed, 5 Aug 2026 17:55:55 +0000 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- scripts/validate_m_groups_cutedsl.py | 201 ++++++ test/prototype/moe_training/test_kernels.py | 61 ++ .../mxfp8/cutedsl_rearrange_2d_m_groups.py | 654 ++++++++++++++++++ 3 files changed, 916 insertions(+) create mode 100644 scripts/validate_m_groups_cutedsl.py create mode 100644 torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_m_groups.py diff --git a/scripts/validate_m_groups_cutedsl.py b/scripts/validate_m_groups_cutedsl.py new file mode 100644 index 00000000000..a76cca5d298 --- /dev/null +++ b/scripts/validate_m_groups_cutedsl.py @@ -0,0 +1,201 @@ +# 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 argparse + +import torch + +from benchmarks.utils import benchmark_cuda_function_in_microseconds +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_cutedsl_kernels_available, + mx_block_rearrange_2d_M_groups_cuda, + triton_mx_block_rearrange_2d_M_groups, +) +from torchao.prototype.moe_training.utils import generate_jagged_offs + + +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) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--rows", type=int, default=131072) + parser.add_argument("--cols", type=int, default=224) + parser.add_argument("--groups", type=int, default=8) + parser.add_argument("--multiple-of", type=int, default=128) + parser.add_argument( + "--chunk-width", type=int, default=0, choices=(0, 16, 32, 64, 128) + ) + parser.add_argument( + "--bench-only", + choices=("cuda", "triton", "cutedsl"), + default=None, + ) + 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") + if torch.cuda.get_device_capability()[0] != 10: + raise SystemExit(f"Requires SM 10.x; got {torch.cuda.get_device_capability()}") + if not _mxfp8_cutedsl_kernels_available: + raise SystemExit("CuteDSL mxfp8 kernels are unavailable") + + if args.deterministic: + scales = ( + torch.arange(args.rows * args.cols, device="cuda", dtype=torch.int32) + .remainder(251) + .to(torch.uint8) + .reshape(args.rows, args.cols) + .view(torch.float8_e8m0fnu) + ) + else: + scales = torch.randint( + 0, + 255, + (args.rows, args.cols), + device="cuda", + dtype=torch.uint8, + ).view(torch.float8_e8m0fnu) + scales_ref = scales.contiguous() + offsets = generate_jagged_offs( + args.groups, + args.rows, + multiple_of=args.multiple_of, + device="cuda", + dtype=torch.int32, + ) + + chunk_width = None if args.chunk_width == 0 else args.chunk_width + + print( + f"shape={tuple(scales.shape)} strides={scales.stride()} groups={args.groups} chunk_width={chunk_width or 'auto'}" + ) + if args.bench_only is None: + out_cuda = mx_block_rearrange_2d_M_groups_cuda(scales_ref, offsets) + out_triton = triton_mx_block_rearrange_2d_M_groups(scales_ref, offsets) + out_cutedsl = mx_block_rearrange_2d_m_groups_cutedsl( + scales, + offsets, + chunk_width=chunk_width, + ) + torch.cuda.synchronize() + + print(f"triton_equal={torch.equal(out_cuda, out_triton)}") + print(f"cutedsl_equal={torch.equal(out_cuda, out_cutedsl)}") + if not torch.equal(out_cuda, out_cutedsl): + diff = out_cuda.view(torch.uint8) != out_cutedsl.view(torch.uint8) + print(f"cutedsl_diff_count={diff.sum().item()}") + coords = diff.nonzero()[:16] + idx = coords[:, 0] * out_cuda.shape[1] + coords[:, 1] + cuda_flat = out_cuda.view(torch.uint8).flatten() + cutedsl_flat = out_cutedsl.view(torch.uint8).flatten() + print(f"first_diff_coords_v2={coords.cpu().tolist()}") + print(f"first_diff_flat_indices_v2={idx.cpu().tolist()}") + print(f"cuda_values={cuda_flat[idx].cpu().tolist()}") + print(f"cutedsl_values={cutedsl_flat[idx].cpu().tolist()}") + print(f"offsets={offsets.cpu().tolist()}") + if args.deterministic: + first_flat = int(idx[0].item()) + sf_tile = first_flat // 512 + tile_byte = first_flat % 512 + sf_tiles_per_row = out_cuda.shape[1] // 4 + row_chunk = sf_tile // sf_tiles_per_row + sf_col_tile = sf_tile % sf_tiles_per_row + local_col = sf_col_tile * 4 + local_row = None + local_byte = None + for r in range(128): + off = (r % 32) * 16 + (r // 32) * 4 + if off <= tile_byte < off + 4: + local_row = r + local_byte = tile_byte - off + break + offsets_cpu = offsets.cpu().tolist() + prev = 0 + chunk_base = 0 + input_row = None + for g, end in enumerate(offsets_cpu): + chunks = (end - prev + 127) // 128 + if row_chunk < chunk_base + chunks: + input_row = prev + (row_chunk - chunk_base) * 128 + local_row + break + chunk_base += chunks + prev = end + print( + f"decoded_first_diff=row_chunk:{row_chunk} sf_col_tile:{sf_col_tile} local_row:{local_row} local_col:{local_col} local_byte:{local_byte} input_row:{input_row}" + ) + if input_row is not None: + src = scales.view(torch.uint8)[ + input_row, local_col : local_col + 16 + ] + print(f"decoded_source16={src.cpu().tolist()}") + + if args.bench_only in (None, "cuda"): + cuda_us = benchmark_function( + mx_block_rearrange_2d_M_groups_cuda, + scales_ref, + offsets, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + print(f"cuda_us={cuda_us:.2f}") + if args.bench_only in (None, "triton"): + triton_us = benchmark_function( + triton_mx_block_rearrange_2d_M_groups, + scales_ref, + offsets, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + print(f"triton_us={triton_us:.2f}") + if args.bench_only in (None, "cutedsl"): + cutedsl_us = benchmark_function( + mx_block_rearrange_2d_m_groups_cutedsl, + scales, + offsets, + chunk_width, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + print(f"cutedsl_us={cutedsl_us:.2f}") + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index e40ac9b4f1e..7a93584b4b7 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, @@ -269,6 +272,64 @@ 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", + [ + (256, 37, 8), + (256, 45, 8), + (512, 64, 8), + (1024, 224, 8), + (1024, 512, 8), + (1024, 1025, 8), + ], +) +def test_cutedsl_mx_block_rearrange_2d_M_groups( + scale_rows: int, + scale_cols: int, + n_groups: int, +): + device = "cuda" + block_size = 32 + e8m0_scales = torch.randint( + 0, + 255, + (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, + ) + assert torch.equal(ref_out_scales, cutedsl_out_scales), "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/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..fbe6b7d6bdb --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_rearrange_2d_m_groups.py @@ -0,0 +1,654 @@ +# 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 + + +@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 _ceil_div(self, a: cutlass.Int32, b: cutlass.Int32): + return (a + b - 1) // b + + @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 = self._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, 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 = (cols_in_chunk + SF_COLS - 1) // 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 = ceil_div(cols, 4) * 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) + 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, + ) + import cuda.bindings.driver as cuda + + stream = cuda.CUstream(int(torch.cuda.current_stream().cuda_stream)) + compiled( + scales_tensor.view(torch.uint8), + input_offsets, + output.view(-1), + int(rows), + int(cols), + 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 = ceil_div(cols, 4) * 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, + ) From 067cc44d619fdf93bd3d4d705220d6ed200e610e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleksandar=20Samard=C5=BEi=C4=87?= Date: Sat, 15 Aug 2026 11:38:27 +0200 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- torchao/prototype/mx_formats/mx_funcs.py | 43 ++++++++++++ .../codebook_groupwise/__init__.py | 2 + .../gguf/gguf_quantized_tensor.py | 12 ++++ torchao/quantization/quant_primitives.py | 70 +++++++++++++++++++ .../quantize_/workflows/nf4/nf4_tensor.py | 10 +++ torchao/quantization/utils.py | 22 ++++++ 6 files changed, 159 insertions(+) create mode 100644 torchao/prototype/mx_formats/mx_funcs.py diff --git a/torchao/prototype/mx_formats/mx_funcs.py b/torchao/prototype/mx_formats/mx_funcs.py new file mode 100644 index 00000000000..13e8ef6da37 --- /dev/null +++ b/torchao/prototype/mx_formats/mx_funcs.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +This file defines the top level torch ops that are extended by MXTensor +See: https://docs.pytorch.org/docs/stable/notes/extending.html#extending-torch-with-a-tensor-wrapper-type +for more details. +""" + +from typing import Any, Dict + +import torch + +from torchao.prototype.mx_formats.mx_ops import _addmm_mx_dispatch +from torchao.prototype.mx_formats.mx_tensor import ( # noqa: E501 + MXTensor, +) + +aten = torch.ops.aten + +MX_FUNC_TABLE: Dict[Any, Any] = {} + + +def implements_func(torch_ops): + """Register torch ops to the mx op table for torch function""" + + def decorator(func): + for op in torch_ops: + MX_FUNC_TABLE[op] = func + return func + + return decorator + + +@implements_func([aten.linear.default]) +def mx_linear(func, types, args, kwargs): + a, b = args[0], args[1] + assert isinstance(a, MXTensor) and isinstance(b, MXTensor) + bias = args[2] if len(args) == 3 else None + return _addmm_mx_dispatch(a, b.t(), func, bias=bias) diff --git a/torchao/prototype/quantization/codebook_groupwise/__init__.py b/torchao/prototype/quantization/codebook_groupwise/__init__.py index 5683c627131..8cf56240cd7 100644 --- a/torchao/prototype/quantization/codebook_groupwise/__init__.py +++ b/torchao/prototype/quantization/codebook_groupwise/__init__.py @@ -4,4 +4,6 @@ __all__ = [ "CodebookQuantizedPackedTensor", "GroupwiseLutWeightConfig", + "QuantizedLutEmbedding", + "EmbeddingLutQuantizer", ] diff --git a/torchao/prototype/quantization/gguf/gguf_quantized_tensor.py b/torchao/prototype/quantization/gguf/gguf_quantized_tensor.py index 4c8486d2a4c..369afaa049e 100644 --- a/torchao/prototype/quantization/gguf/gguf_quantized_tensor.py +++ b/torchao/prototype/quantization/gguf/gguf_quantized_tensor.py @@ -66,6 +66,18 @@ def __init__( self.quantized_block_min = quantized_block_min self.int_data = int_data + def _apply_fn_to_data(self, fn): + return self.__class__( + self.n_blocks_per_superblock, + fn(self.super_block_scale_scale), + fn(self.super_block_min_sclae), + fn(self.quantized_block_scale), + fn(self.quantized_block_min), + fn(self.int_data), + self.shape, + dtype=self.dtype, + ) + def __tensor_flatten__(self): return [ "super_block_scale_scale", diff --git a/torchao/quantization/quant_primitives.py b/torchao/quantization/quant_primitives.py index 9ca28f56fd0..15a1b2ad556 100644 --- a/torchao/quantization/quant_primitives.py +++ b/torchao/quantization/quant_primitives.py @@ -11,6 +11,11 @@ import torch +from torchao.prototype.custom_fp_utils import ( + _f32_to_floatx_unpacked, + _floatx_unpacked_to_f32, + _n_ones, +) from torchao.utils import ( _register_custom_op, _register_meta_op, @@ -26,6 +31,7 @@ "TorchAODType", "_choose_qparams_affine_tinygemm", "_choose_qparams_affine_dont_preserve_zero", + "_choose_qparams_affine_floatx", "_choose_qparams_and_quantize_affine_hqq", "_choose_qparams_and_quantize_scale_only_hqq", "_choose_qparams_and_quantize_scale_only_sinq", @@ -33,10 +39,12 @@ "_choose_qparams_gguf", "_quantize_affine_no_zero_point", "_quantize_affine_tinygemm", + "_quantize_affine_floatx", "_quantize_affine_float8", "_quantize_gguf", "_dequantize_affine_no_zero_point", "_dequantize_affine_tinygemm", + "_dequantize_affine_floatx", "_dequantize_affine_float8", "_dequantize_gguf", "_fake_quantize_affine", @@ -207,6 +215,8 @@ class TorchAODType(Enum, metaclass=_TorchAODTypeMeta): _GGUF_QK_K = 256 +_ONES_TABLE = [_n_ones(i) for i in range(8)] + quant_lib = torch.library.Library("torchao", "FRAGMENT") register_custom_op = _register_custom_op(quant_lib) @@ -2169,6 +2179,66 @@ def _choose_qparams_and_quantize_scale_only_sinq( return qdata, scale_row, scale_col +def _choose_qparams_affine_floatx( + tensor: torch.Tensor, ebits: int, mbits: int +) -> torch.Tensor: + """Choose quantization parameters for floatx quantization. + + Calculates scale parameter for quantizing to custom floating point format. + + Args: + tensor: Input tensor to quantize (float32, float16, or bfloat16) + ebits: Number of exponent bits in target floatx format + mbits: Number of mantissa bits in target floatx format + + Returns: + Scale tensor for floatx quantization + + Note: + Uses global lookup table as workaround for torch.compile() compatibility + since _n_ones() is not compatible due to << operator. + """ + # _n_ones() is not compatible with torch.compile() due to << operator + # https://github.com/pytorch/pytorch/issues/119152 + # exp_bias = _n_ones(ebits - 1) + # max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2 ** mbits)) + + # workaround: global lookup table + exp_bias = _ONES_TABLE[ebits - 1] + max_normal = 2 ** (_ONES_TABLE[ebits] - exp_bias) * ( + _ONES_TABLE[mbits + 1] / (2**mbits) + ) + + dtype = tensor.dtype + tensor = tensor.float() + scale = tensor.abs().amax(1).clamp(min=1e-12) / max_normal + return scale.to(dtype) + + +def _quantize_affine_floatx( + tensor: torch.Tensor, scale: torch.Tensor, ebits: int, mbits: int +) -> torch.Tensor: + """Quantizes the float32 high precision floating point tensor to low precision floating point number and + converts the result to unpacked floating point format with the format of 00SEEEMM (for fp6_e3m2) where S means sign bit, e means exponent bit and m means mantissa bit + """ + tensor = tensor.float() + tensor_floatx = _f32_to_floatx_unpacked(tensor / scale.view(-1, 1), ebits, mbits) + return tensor_floatx + + +def _dequantize_affine_floatx( + tensor: torch.Tensor, + scale: torch.Tensor, + ebits: int, + mbits: int, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + tensor = _floatx_unpacked_to_f32(tensor, ebits, mbits) + tensor = tensor * scale.float().view(-1, 1) + tensor = tensor.to(dtype=output_dtype) + return tensor + + @register_custom_op def _choose_scale_float8( tensor: torch.Tensor, diff --git a/torchao/quantization/quantize_/workflows/nf4/nf4_tensor.py b/torchao/quantization/quantize_/workflows/nf4/nf4_tensor.py index 9cd1d7f0a24..f1f95308ab5 100644 --- a/torchao/quantization/quantize_/workflows/nf4/nf4_tensor.py +++ b/torchao/quantization/quantize_/workflows/nf4/nf4_tensor.py @@ -137,6 +137,16 @@ def apply_to_inner_tensors(nf4tensor: "NF4Tensor", aten_op, args, kwargs): return attr_to_tensor +# __torch_function__ utils: call tensor ops from inner tensors +def call_from_inner_tensors(nf4tensor: "NF4Tensor", method_name: str, args, kwargs): + attr_to_tensor = {} + for attr in _INNER_TENSOR_NAMES_FOR_SHARDING: + inner_tensor = getattr(nf4tensor, attr) + func = getattr(inner_tensor, method_name) + attr_to_tensor[attr] = func(*args, **kwargs) + return attr_to_tensor + + class CompareOp(Enum): EQ = auto() LT = auto() diff --git a/torchao/quantization/utils.py b/torchao/quantization/utils.py index 4a772d53b6c..b7b69155850 100644 --- a/torchao/quantization/utils.py +++ b/torchao/quantization/utils.py @@ -111,6 +111,28 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): return rs +class _MultiInput: + def __init__(self, inputs): + self.values = list(inputs) + + def add_input(self, input): + self.values.append(input) + return self + + def __getitem__(self, slice): + return _MultiInput(self.values[slice]) + + def cuda(self): + self.values = [ + val.cuda() if isinstance(val, torch.Tensor) else val for val in self.values + ] + + def xpu(self): + self.values = [ + val.xpu() if isinstance(val, torch.Tensor) else val for val in self.values + ] + + def _guard_dtype_size(tensor_arg, arg_name, dtype=None, size=None): if dtype is not None and tensor_arg.dtype != dtype: raise ValueError(