diff --git a/scripts/validate_pad_unpad_cutedsl.py b/scripts/validate_pad_unpad_cutedsl.py new file mode 100644 index 00000000000..7d342ee56e7 --- /dev/null +++ b/scripts/validate_pad_unpad_cutedsl.py @@ -0,0 +1,298 @@ +# 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.cute_utils import ( + _missing_cutedsl_runtime_packages, +) +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_pad_token_groups import ( + pad_token_groups_cutedsl, + unpad_token_groups_cutedsl, +) +from torchao.prototype.moe_training.kernels.mxfp8.quant import ( + _mxfp8_cutedsl_kernels_available, + fused_pad_token_groups_cuda, + fused_unpad_token_groups_cuda, + torch_pad_token_groups, + torch_unpad_token_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 maybe_benchmark_function( + name, + f, + *args, + use_cuda_graph=False, + graph_iters=1000, + graph_safe=True, + **kwargs, +): + if use_cuda_graph and not graph_safe: + print(f"{name}=skipped_cuda_graph_unsafe") + return None + value = benchmark_function( + f, + *args, + use_cuda_graph=use_cuda_graph, + graph_iters=graph_iters, + **kwargs, + ) + print(f"{name}={value:.2f}") + return value + + +def cuda_reference_supports_shape(dtype: str, dim: int) -> bool: + return dtype == "bf16" and (dim < 8 or dim % 8 == 0) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, default=16384) + parser.add_argument("--dim", type=int, default=7168) + parser.add_argument("--groups", type=int, default=8) + parser.add_argument("--multiple-of", type=int, default=1) + parser.add_argument("--alignment-size", type=int, default=32) + parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16") + parser.add_argument("--mode", choices=("pad", "unpad", "both"), default="both") + parser.add_argument("--bench-only", choices=("torch", "cuda", "cutedsl")) + 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: + missing = _missing_cutedsl_runtime_packages() + detail = ( + f"missing package(s): {', '.join(missing)}" + if missing + else "requires CUDA 12.8+ on SM 10.x" + ) + raise SystemExit(f"CuteDSL mxfp8 kernels unavailable ({detail})") + if args.bench_only in (None, "cuda") and not cuda_reference_supports_shape( + args.dtype, + args.dim, + ): + raise SystemExit( + "CUDA pad/unpad reference requires bf16 with dim < 8 or dim % 8 == 0" + ) + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + if args.deterministic: + inputs = ( + torch.arange(args.tokens * args.dim, device="cuda", dtype=torch.float32) + .remainder(251) + .reshape(args.tokens, args.dim) + .to(dtype) + ) + else: + torch.manual_seed(0) + inputs = torch.randn(args.tokens, args.dim, dtype=dtype, device="cuda") + offsets = generate_jagged_offs( + args.groups, + args.tokens, + multiple_of=args.multiple_of, + device="cuda", + dtype=torch.int32, + ) + + print( + f"shape={tuple(inputs.shape)} dtype={args.dtype} groups={args.groups} alignment_size={args.alignment_size} mode={args.mode}" + ) + + if args.mode in ("pad", "both"): + if args.bench_only is None: + out_torch, start_torch, end_torch = torch_pad_token_groups( + inputs, + offsets, + args.alignment_size, + ) + out_cuda, start_cuda, end_cuda = fused_pad_token_groups_cuda( + inputs, + offsets, + args.alignment_size, + ) + out_cutedsl, start_cutedsl, end_cutedsl = pad_token_groups_cutedsl( + inputs, + offsets, + args.alignment_size, + ) + torch.cuda.synchronize() + print(f"pad_cuda_equal={torch.equal(out_torch, out_cuda)}") + print(f"pad_cutedsl_equal={torch.equal(out_torch, out_cutedsl)}") + print(f"pad_cuda_start_equal={torch.equal(start_torch, start_cuda)}") + print(f"pad_cutedsl_start_equal={torch.equal(start_torch, start_cutedsl)}") + print(f"pad_cuda_end_equal={torch.equal(end_torch, end_cuda)}") + print(f"pad_cutedsl_end_equal={torch.equal(end_torch, end_cutedsl)}") + if not torch.equal(out_torch, out_cutedsl): + diff = out_torch != out_cutedsl + print(f"pad_cutedsl_diff_count={diff.sum().item()}") + coords = diff.nonzero()[:16] + idx = coords[:, 0] * out_torch.shape[1] + coords[:, 1] + torch_flat = out_torch.flatten() + cutedsl_flat = out_cutedsl.flatten() + print(f"pad_first_diff_coords={coords.cpu().tolist()}") + print(f"pad_first_diff_flat_indices={idx.cpu().tolist()}") + print(f"pad_torch_values={torch_flat[idx].float().cpu().tolist()}") + print(f"pad_cutedsl_values={cutedsl_flat[idx].float().cpu().tolist()}") + print(f"offsets={offsets.cpu().tolist()}") + if args.bench_only in (None, "torch"): + maybe_benchmark_function( + "pad_torch_us", + torch_pad_token_groups, + inputs, + offsets, + args.alignment_size, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + graph_safe=False, + ) + if args.bench_only in (None, "cuda"): + maybe_benchmark_function( + "pad_cuda_us", + fused_pad_token_groups_cuda, + inputs, + offsets, + args.alignment_size, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + if args.bench_only in (None, "cutedsl"): + maybe_benchmark_function( + "pad_cutedsl_us", + pad_token_groups_cutedsl, + inputs, + offsets, + args.alignment_size, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + + if args.mode in ("unpad", "both"): + padded_inputs, padded_group_start_offsets, _ = torch_pad_token_groups( + inputs, + offsets, + args.alignment_size, + ) + if args.bench_only is None: + out_torch = torch_unpad_token_groups( + padded_inputs, + offsets, + padded_group_start_offsets, + args.tokens, + args.alignment_size, + ) + out_cuda = fused_unpad_token_groups_cuda( + padded_inputs, + offsets, + padded_group_start_offsets, + args.tokens, + args.alignment_size, + ) + out_cutedsl = unpad_token_groups_cutedsl( + padded_inputs, + offsets, + padded_group_start_offsets, + args.tokens, + args.alignment_size, + ) + torch.cuda.synchronize() + print(f"unpad_cuda_equal={torch.equal(out_torch, out_cuda)}") + print(f"unpad_cutedsl_equal={torch.equal(out_torch, out_cutedsl)}") + if not torch.equal(out_torch, out_cutedsl): + diff = out_torch != out_cutedsl + print(f"unpad_cutedsl_diff_count={diff.sum().item()}") + coords = diff.nonzero()[:16] + idx = coords[:, 0] * out_torch.shape[1] + coords[:, 1] + torch_flat = out_torch.flatten() + cutedsl_flat = out_cutedsl.flatten() + print(f"unpad_first_diff_coords={coords.cpu().tolist()}") + print(f"unpad_first_diff_flat_indices={idx.cpu().tolist()}") + print(f"unpad_torch_values={torch_flat[idx].float().cpu().tolist()}") + print( + f"unpad_cutedsl_values={cutedsl_flat[idx].float().cpu().tolist()}" + ) + print(f"offsets={offsets.cpu().tolist()}") + if args.bench_only in (None, "torch"): + maybe_benchmark_function( + "unpad_torch_us", + torch_unpad_token_groups, + padded_inputs, + offsets, + padded_group_start_offsets, + args.tokens, + args.alignment_size, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + graph_safe=False, + ) + if args.bench_only in (None, "cuda"): + maybe_benchmark_function( + "unpad_cuda_us", + fused_unpad_token_groups_cuda, + padded_inputs, + offsets, + padded_group_start_offsets, + args.tokens, + args.alignment_size, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + if args.bench_only in (None, "cutedsl"): + maybe_benchmark_function( + "unpad_cutedsl_us", + unpad_token_groups_cutedsl, + padded_inputs, + offsets, + padded_group_start_offsets, + args.tokens, + args.alignment_size, + use_cuda_graph=args.cuda_graph_bench, + graph_iters=args.graph_iters, + ) + + +if __name__ == "__main__": + main() diff --git a/test/prototype/moe_training/test_kernels.py b/test/prototype/moe_training/test_kernels.py index 7c2edbc7340..d2ee4d11ceb 100644 --- a/test/prototype/moe_training/test_kernels.py +++ b/test/prototype/moe_training/test_kernels.py @@ -54,6 +54,10 @@ 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_pad_token_groups import ( + pad_token_groups_cutedsl, + unpad_token_groups_cutedsl, +) from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_rearrange_2d_m_groups import ( mx_block_rearrange_2d_m_groups_cutedsl, ) @@ -785,7 +789,7 @@ def test_cuda_mx_dim1_2d_numerics_32x1( @pytest.mark.parametrize("num_tokens", [128, 157, 4096, 16392]) @pytest.mark.parametrize("dim", [7168]) @pytest.mark.parametrize("num_groups", [1, 2, 4, 8]) -@pytest.mark.parametrize("alignment_size", [32]) +@pytest.mark.parametrize("alignment_size", [32, 128]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) def test_cuda_fused_pad_token_groups( num_tokens: int, dim: int, num_groups: int, alignment_size: int, dtype: torch.dtype @@ -832,7 +836,7 @@ def test_cuda_fused_pad_token_groups( @pytest.mark.parametrize("num_tokens", [128, 157, 4096]) @pytest.mark.parametrize("dim", [7168]) @pytest.mark.parametrize("num_groups", [1, 2, 4, 8]) -@pytest.mark.parametrize("alignment_size", [32]) +@pytest.mark.parametrize("alignment_size", [32, 128]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) def test_cuda_fused_unpad_token_groups( num_tokens: int, dim: int, num_groups: int, alignment_size: int, dtype: torch.dtype @@ -882,6 +886,78 @@ def test_cuda_fused_unpad_token_groups( ) +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, + reason="MXFP8 cutedsl kernels not available", +) +@skip_if_rocm("ROCm enablement in progress") +@pytest.mark.parametrize("num_tokens", [128, 157, 4096]) +@pytest.mark.parametrize("dim", [32, 7168]) +@pytest.mark.parametrize("num_groups", [1, 8]) +@pytest.mark.parametrize("alignment_size", [32]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_cutedsl_pad_token_groups( + num_tokens: int, dim: int, num_groups: int, alignment_size: int, dtype: torch.dtype +): + device = "cuda" + inputs = torch.randn(num_tokens, dim, dtype=dtype, device=device) + group_offsets = generate_jagged_offs( + num_groups, num_tokens, multiple_of=1, device=device + ) + + ref_padded_tokens, ref_start_offsets, ref_end_offsets = torch_pad_token_groups( + inputs, group_offsets, alignment_size + ) + cutedsl_padded_tokens, cutedsl_start_offsets, cutedsl_end_offsets = ( + pad_token_groups_cutedsl(inputs, group_offsets, alignment_size) + ) + + assert torch.equal(ref_padded_tokens, cutedsl_padded_tokens) + assert torch.equal(ref_start_offsets, cutedsl_start_offsets) + assert torch.equal(ref_end_offsets, cutedsl_end_offsets) + + +@pytest.mark.skipif( + not _mxfp8_cutedsl_kernels_available, + reason="MXFP8 cutedsl kernels not available", +) +@skip_if_rocm("ROCm enablement in progress") +@pytest.mark.parametrize("num_tokens", [128, 157, 4096]) +@pytest.mark.parametrize("dim", [32, 7168]) +@pytest.mark.parametrize("num_groups", [1, 8]) +@pytest.mark.parametrize("alignment_size", [32]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_cutedsl_unpad_token_groups( + num_tokens: int, dim: int, num_groups: int, alignment_size: int, dtype: torch.dtype +): + device = "cuda" + inputs = torch.randn(num_tokens, dim, dtype=dtype, device=device) + group_offsets = generate_jagged_offs( + num_groups, num_tokens, multiple_of=1, device=device + ) + padded_tokens, padded_group_start_offsets, _ = torch_pad_token_groups( + inputs, group_offsets, alignment_size + ) + + ref_unpadded_tokens = torch_unpad_token_groups( + padded_tokens, + group_offsets, + padded_group_start_offsets, + num_tokens, + alignment_size, + ) + cutedsl_unpadded_tokens = unpad_token_groups_cutedsl( + padded_tokens, + group_offsets, + padded_group_start_offsets, + num_tokens, + alignment_size, + ) + + assert torch.equal(ref_unpadded_tokens, cutedsl_unpadded_tokens) + assert torch.equal(inputs, cutedsl_unpadded_tokens) + + @pytest.mark.parametrize("round_scales_to_power_of_2", [True, False]) @pytest.mark.parametrize( "m,k", diff --git a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py index 0ed217b43a4..442ef7598d2 100644 --- a/torchao/prototype/moe_training/kernels/mxfp8/__init__.py +++ b/torchao/prototype/moe_training/kernels/mxfp8/__init__.py @@ -1,3 +1,10 @@ +from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_pad_token_groups import ( # noqa: F401 + 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_pad_token_groups.py b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_pad_token_groups.py new file mode 100644 index 00000000000..88d74c3e7e2 --- /dev/null +++ b/torchao/prototype/moe_training/kernels/mxfp8/cutedsl_pad_token_groups.py @@ -0,0 +1,426 @@ +# 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 + + +@functools.cache +def _compile_copy_token_groups_cutedsl( + num_groups: int, + element_size: int, + is_pad: bool, + row_bytes_aligned: bool, + alignment_size: int, +): + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass.cute.runtime import make_fake_stream, make_fake_tensor + + NUM_GROUPS = num_groups + ELEMENT_SIZE = element_size + IS_PAD = is_pad + ROW_BYTES_ALIGNED = row_bytes_aligned + ALIGNMENT_SIZE = alignment_size + ALIGNMENT_MASK = alignment_size - 1 + VEC_ELEMS = 16 // ELEMENT_SIZE + WARP_SIZE = 32 + WARPS_PER_BLOCK = 8 + THREADS_PER_BLOCK = WARP_SIZE * WARPS_PER_BLOCK + UNROLL = 4 + LANE_STRIDE = WARP_SIZE * VEC_ELEMS + UNROLL_STRIDE = UNROLL * LANE_STRIDE + ELEMENT_DTYPE = cutlass.BFloat16 if ELEMENT_SIZE == 2 else cutlass.Float32 + + class CopyTokenGroups: + @cute.jit + def _find_group( + self, + row: cutlass.Int32, + group_end_offsets: cute.Tensor, + ): + group = cutlass.Int32(0) + keep_looking = cutlass.Int32(1) + while group < NUM_GROUPS - 1 and keep_looking != 0: + end = cutlass.Int32((group_end_offsets.iterator + group).load()) + if row >= end: + group = group + cutlass.Int32(1) + else: + keep_looking = cutlass.Int32(0) + return group + + @cute.jit + def _group_start( + self, + group: cutlass.Int32, + group_end_offsets: cute.Tensor, + ): + group_start = cutlass.Int32(0) + if group > 0: + group_start = cutlass.Int32( + (group_end_offsets.iterator + group - 1).load() + ) + return group_start + + @cute.jit + def _padded_start( + self, + group: cutlass.Int32, + group_end_offsets: cute.Tensor, + ): + cumulative = cutlass.Int32(0) + previous = cutlass.Int32(0) + for g in range(NUM_GROUPS): + group_end = cutlass.Int32((group_end_offsets.iterator + g).load()) + group_size = group_end - previous + padded_size = ( + (group_size + cutlass.Int32(ALIGNMENT_MASK)) + // cutlass.Int32(ALIGNMENT_SIZE) + ) * cutlass.Int32(ALIGNMENT_SIZE) + if g < group: + cumulative = cumulative + padded_size + previous = group_end + return cumulative + + @cute.jit + def _fill_padded_offsets( + self, + tidx: cutlass.Int32, + group_end_offsets: cute.Tensor, + padded_group_start_offsets: cute.Tensor, + padded_group_end_offsets: cute.Tensor, + ): + if tidx < NUM_GROUPS: + group = cutlass.Int32(tidx) + group_start = self._group_start(group, group_end_offsets) + group_end = cutlass.Int32((group_end_offsets.iterator + group).load()) + group_size = group_end - group_start + padded_size = ( + (group_size + cutlass.Int32(ALIGNMENT_MASK)) + // cutlass.Int32(ALIGNMENT_SIZE) + ) * cutlass.Int32(ALIGNMENT_SIZE) + padded_start = self._padded_start(group, group_end_offsets) + (padded_group_start_offsets.iterator + group).store(padded_start) + (padded_group_end_offsets.iterator + group).store( + padded_start + padded_size + ) + + @cute.jit + def _copy_vec_or_tail( + self, + inputs: cute.Tensor, + input_offset: cutlass.Int32, + elem_col: cutlass.Int32, + output: cute.Tensor, + output_offset: cutlass.Int32, + dim: cutlass.Int32, + ): + if cutlass.const_expr(ROW_BYTES_ALIGNED): + value = cute.recast_ptr( + (inputs.iterator + input_offset).align(16), + dtype=cutlass.Int128, + ).load() + cute.recast_ptr( + (output.iterator + output_offset).align(16), + dtype=cutlass.Int128, + ).store(value) + else: + for i in range(VEC_ELEMS): + if elem_col + i < dim: + value = (inputs.iterator + input_offset + i).load() + (output.iterator + output_offset + i).store(value) + + @cute.kernel + def kernel( + self, + inputs: cute.Tensor, + group_end_offsets: cute.Tensor, + padded_group_start_offsets: cute.Tensor, + padded_group_end_offsets: cute.Tensor, + output: cute.Tensor, + num_tokens: cutlass.Int32, + dim: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + lane_id = cutlass.Int32(tidx % WARP_SIZE) + warp_id = cutlass.Int32(tidx // WARP_SIZE) + if cutlass.const_expr(IS_PAD): + if block_idx == 0: + self._fill_padded_offsets( + tidx, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + ) + row = cutlass.Int32(block_idx * WARPS_PER_BLOCK + warp_id) + if row < num_tokens: + group_start = cutlass.Int32(0) + padded_start = cutlass.Int32(0) + if lane_id == 0: + if cutlass.const_expr(IS_PAD): + group = cutlass.Int32(0) + previous = cutlass.Int32(0) + cumulative = cutlass.Int32(0) + keep_looking = cutlass.Int32(1) + while group < NUM_GROUPS and keep_looking != 0: + group_end = cutlass.Int32( + (group_end_offsets.iterator + group).load() + ) + group_size = group_end - previous + padded_size = ( + (group_size + cutlass.Int32(ALIGNMENT_MASK)) + // cutlass.Int32(ALIGNMENT_SIZE) + ) * cutlass.Int32(ALIGNMENT_SIZE) + if row < group_end: + group_start = previous + padded_start = cumulative + keep_looking = cutlass.Int32(0) + else: + cumulative = cumulative + padded_size + previous = group_end + group = group + cutlass.Int32(1) + else: + group = self._find_group(row, group_end_offsets) + group_start = self._group_start(group, group_end_offsets) + padded_start = cutlass.Int32( + (padded_group_start_offsets.iterator + group).load() + ) + group_start = cute.arch.shuffle_sync(group_start, 0) + padded_start = cute.arch.shuffle_sync(padded_start, 0) + offset_in_group = row - group_start + padded_row = padded_start + offset_in_group + input_row = row + output_row = padded_row + if cutlass.const_expr(not IS_PAD): + input_row = padded_row + output_row = row + elem_col = lane_id * VEC_ELEMS + input_offset = input_row * dim + elem_col + output_offset = output_row * dim + elem_col + if cutlass.const_expr(ROW_BYTES_ALIGNED): + while elem_col + (UNROLL - 1) * LANE_STRIDE + VEC_ELEMS <= dim: + vals = [] + for k in cutlass.range_constexpr(UNROLL): + vals.append( + cute.recast_ptr( + ( + inputs.iterator + input_offset + k * LANE_STRIDE + ).align(16), + dtype=cutlass.Int128, + ).load() + ) + for k in cutlass.range_constexpr(UNROLL): + cute.recast_ptr( + ( + output.iterator + output_offset + k * LANE_STRIDE + ).align(16), + dtype=cutlass.Int128, + ).store(vals[k]) + elem_col = elem_col + UNROLL_STRIDE + input_offset = input_offset + UNROLL_STRIDE + output_offset = output_offset + UNROLL_STRIDE + while elem_col < dim: + self._copy_vec_or_tail( + inputs, + input_offset, + elem_col, + output, + output_offset, + dim, + ) + elem_col = elem_col + WARP_SIZE * VEC_ELEMS + input_offset = input_offset + WARP_SIZE * VEC_ELEMS + output_offset = output_offset + WARP_SIZE * VEC_ELEMS + + @cute.jit + def __call__( + self, + inputs: cute.Tensor, + group_end_offsets: cute.Tensor, + padded_group_start_offsets: cute.Tensor, + padded_group_end_offsets: cute.Tensor, + output: cute.Tensor, + num_tokens: cutlass.Int32, + dim: cutlass.Int32, + num_blocks: cutlass.Int32, + stream: cuda.CUstream, + ): + self.kernel( + inputs, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + output, + num_tokens, + dim, + ).launch( + grid=(num_blocks, 1, 1), + block=(THREADS_PER_BLOCK, 1, 1), + cluster=(1, 1, 1), + stream=stream, + ) + + kernel = CopyTokenGroups() + + n = cute.sym_int() + d = cute.sym_int() + out_n = cute.sym_int() + offs_stride = cute.sym_int() + assumed_align = 16 if ROW_BYTES_ALIGNED else 1 + fake_input = make_fake_tensor( + ELEMENT_DTYPE, + (n, d), + stride=(d, 1), + assumed_align=assumed_align, + ) + fake_output = make_fake_tensor( + ELEMENT_DTYPE, + (out_n, d), + stride=(d, 1), + assumed_align=assumed_align, + ) + fake_group_end_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_padded_group_start_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_padded_group_end_offsets = make_fake_tensor( + cutlass.Int32, + (NUM_GROUPS,), + stride=(offs_stride,), + ) + fake_stream = make_fake_stream() + + return cute.compile( + kernel, + inputs=fake_input, + group_end_offsets=fake_group_end_offsets, + padded_group_start_offsets=fake_padded_group_start_offsets, + padded_group_end_offsets=fake_padded_group_end_offsets, + output=fake_output, + num_tokens=0, + dim=0, + num_blocks=1, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + +def pad_token_groups_cutedsl( + inputs: torch.Tensor, + group_end_offsets: torch.Tensor, + alignment_size: int = 32, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert inputs.ndim == 2 + assert inputs.is_cuda + assert inputs.is_contiguous() + assert inputs.dtype in (torch.float32, torch.bfloat16) + assert group_end_offsets.ndim == 1 + assert group_end_offsets.is_cuda + assert group_end_offsets.dtype == torch.int32 + assert alignment_size in (32, 128) + num_tokens, dim = inputs.shape + num_groups = group_end_offsets.shape[0] + assert num_groups <= 32 + padded_group_start_offsets = torch.empty( + (num_groups,), + dtype=group_end_offsets.dtype, + device=group_end_offsets.device, + ) + padded_group_end_offsets = torch.empty( + (num_groups,), + dtype=group_end_offsets.dtype, + device=group_end_offsets.device, + ) + output_rows = num_tokens + num_groups * alignment_size + output_rows = ( + (output_rows + alignment_size - 1) // alignment_size + ) * alignment_size + output = torch.zeros( + (output_rows, dim), + dtype=inputs.dtype, + device=inputs.device, + ) + compiled = _compile_copy_token_groups_cutedsl( + num_groups, + inputs.element_size(), + True, + dim * inputs.element_size() % 16 == 0, + alignment_size, + ) + import cuda.bindings.driver as cuda + + stream = cuda.CUstream(int(torch.cuda.current_stream().cuda_stream)) + compiled( + inputs, + group_end_offsets, + padded_group_start_offsets, + padded_group_end_offsets, + output, + int(num_tokens), + int(dim), + int((num_tokens + 7) // 8), + stream, + ) + return output, padded_group_start_offsets, padded_group_end_offsets + + +def unpad_token_groups_cutedsl( + padded_inputs: torch.Tensor, + group_end_offsets: torch.Tensor, + padded_group_start_offsets: torch.Tensor, + num_tokens: int, + alignment_size: int = 32, +) -> torch.Tensor: + assert padded_inputs.ndim == 2 + assert padded_inputs.is_cuda + assert padded_inputs.is_contiguous() + assert padded_inputs.dtype in (torch.float32, torch.bfloat16) + assert group_end_offsets.ndim == 1 + assert group_end_offsets.is_cuda + assert group_end_offsets.dtype == torch.int32 + assert padded_group_start_offsets.ndim == 1 + assert padded_group_start_offsets.is_cuda + assert padded_group_start_offsets.dtype == torch.int32 + assert alignment_size in (32, 128) + dim = padded_inputs.shape[1] + num_groups = group_end_offsets.shape[0] + assert num_groups <= 32 + output = torch.empty( + (num_tokens, dim), + dtype=padded_inputs.dtype, + device=padded_inputs.device, + ) + compiled = _compile_copy_token_groups_cutedsl( + num_groups, + padded_inputs.element_size(), + False, + dim * padded_inputs.element_size() % 16 == 0, + alignment_size, + ) + import cuda.bindings.driver as cuda + + stream = cuda.CUstream(int(torch.cuda.current_stream().cuda_stream)) + compiled( + padded_inputs, + group_end_offsets, + padded_group_start_offsets, + padded_group_start_offsets, + output, + int(num_tokens), + int(dim), + int((num_tokens + 7) // 8), + stream, + ) + return output