diff --git a/benchmarks/benchmark_chunk_padding.py b/benchmarks/benchmark_chunk_padding.py new file mode 100644 index 00000000000..8c7bf5ec23f --- /dev/null +++ b/benchmarks/benchmark_chunk_padding.py @@ -0,0 +1,140 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Compare fused chunk sort/pad/restore with TE chunk sort and fused row padding.""" + +import argparse +import json +import logging +import statistics + +import torch + +from transformer_engine.pytorch.chunk_padding import ( + moe_sort_chunks_and_pad, + moe_unpad_and_restore_chunks, +) +from transformer_engine.pytorch.module.fp8_padding import Fp8Padding +from transformer_engine.pytorch.module.fp8_unpadding import Fp8Unpadding +from transformer_engine.pytorch.permutation import ( + moe_sort_chunks_by_index, + moe_sort_chunks_by_index_with_probs, +) + + +def benchmark(num_tokens, hidden_size, ranks, experts, alignment): + """Time the same layout roundtrip and probability padding in A-B-B-A order.""" + generator = torch.Generator().manual_seed(2026) + counts = torch.bincount( + torch.randint(ranks * experts, (num_tokens,), generator=generator), + minlength=ranks * experts, + ).tolist() + order = [ + rank * experts + expert for expert in reversed(range(experts)) for rank in range(ranks) + ] + sorted_counts = [counts[index] for index in order] + group_sizes = [ + sum(sorted_counts[start : start + ranks]) for start in range(0, len(order), ranks) + ] + padded_counts = sorted_counts.copy() + for index, size in enumerate(group_sizes): + padded_counts[(index + 1) * ranks - 1] += (-size) % alignment + device = torch.device("cuda") + sizes, indices, padded_sizes = [ + torch.tensor(v, device=device, dtype=torch.int32) for v in (counts, order, padded_counts) + ] + sorted_sizes = torch.tensor(sorted_counts, device=device, dtype=torch.int32) + inverse = torch.argsort(indices).to(torch.int32) + inp = torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) + probs = torch.rand(num_tokens, device=device) + padding, unpadding = Fp8Padding(experts, alignment), Fp8Unpadding(experts, alignment) + + def reference(): + sorted_inp, sorted_probs = moe_sort_chunks_by_index_with_probs(inp, probs, sizes, indices) + padded, _ = padding(sorted_inp, group_sizes) + padded_probs, _ = padding(sorted_probs[:, None], group_sizes) + restored = moe_sort_chunks_by_index(unpadding(padded, group_sizes), sorted_sizes, inverse) + return restored, padded_probs.flatten() + + def candidate(): + padded, padded_probs, row_map = moe_sort_chunks_and_pad( + inp, sizes, indices, padded_sizes, sum(padded_counts), probs + ) + return moe_unpad_and_restore_chunks(padded, row_map, num_tokens), padded_probs + + with torch.no_grad(): + for expected, actual in zip(reference(), candidate()): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + records = [] + for name, operation in ( + ("reference", reference), + ("candidate", candidate), + ("candidate", candidate), + ("reference", reference), + ): + for _ in range(10): + operation() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_allocated() + output = operation() + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - before + del output + samples = [] + for _ in range(50): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + start.record() + operation() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + records.append( + dict( + variant=name, + median_ms=statistics.median(samples), + peak_allocated_bytes=peak, + samples_ms=samples, + ) + ) + return dict( + num_tokens=num_tokens, + hidden_size=hidden_size, + ranks=ranks, + experts=experts, + alignment=alignment, + gpu=torch.cuda.get_device_name(), + torch=torch.__version__, + cuda=torch.version.cuda, + reference="TE chunk sort + fused row padding/unpadding", + scope="layout roundtrip only; no GEMM or communication", + runs=records, + ) + + +def main(): + """Write raw measurements as JSON for independent analysis.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=[8192, 32768]) + parser.add_argument("--hidden-size", type=int, default=2048) + parser.add_argument("--ranks", type=int, default=8) + parser.add_argument("--experts", type=int, default=16) + parser.add_argument("--alignment", type=int, default=128) + parser.add_argument("--output", required=True) + args = parser.parse_args() + results = [ + benchmark(count, args.hidden_size, args.ranks, args.experts, args.alignment) + for count in args.tokens + ] + with open(args.output, "w", encoding="utf-8") as output: + json.dump(results, output, indent=2) + output.write("\n") + logging.getLogger(__name__).info("Saved benchmark results to %s", args.output) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + main() diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index a6afa2d0cc8..f1526f4669d 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -100,6 +100,10 @@ Mixture of Experts (MoE) functions .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs +.. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_and_pad + +.. autoapifunction:: transformer_engine.pytorch.moe_unpad_and_restore_chunks + Communication-computation overlap --------------------------------- diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a78a99d7f95..4ae0bcd7ed4 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -56,6 +56,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_selective_activa python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_chunk_padding.xml $TE_PATH/tests/pytorch/test_chunk_padding.py || test_fail "test_chunk_padding.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cross_entropy.xml $TE_PATH/tests/pytorch/test_cross_entropy.py || test_fail "test_cross_entropy.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" diff --git a/tests/pytorch/test_chunk_padding.py b/tests/pytorch/test_chunk_padding.py new file mode 100644 index 00000000000..439de77bd8c --- /dev/null +++ b/tests/pytorch/test_chunk_padding.py @@ -0,0 +1,143 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Output, gradient and graph coverage for fused chunk sorting and padding.""" + +import pytest +import torch +import torch.nn.functional as F + +from transformer_engine.pytorch.chunk_padding import ( + moe_sort_chunks_and_pad, + moe_unpad_and_restore_chunks, +) + + +def reference(inp, probs, sizes, order, output_sizes): + """Independently split, reorder and pad tokens and probabilities.""" + chunks = inp.split(sizes) + prob_chunks = probs.split(sizes) if probs is not None else None + tokens, probabilities = [], [] + for index, size in zip(order, output_sizes): + padding = size - sizes[index] + tokens.append(F.pad(chunks[index], (0, 0, 0, padding))) + if probs is not None: + probabilities.append(F.pad(prob_chunks[index], (0, padding))) + return torch.cat(tokens), torch.cat(probabilities) if probs is not None else None + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("prob_dtype", [None, torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "sizes,order,output_sizes", + [ + ([0, 3, 5], [2, 0, 1], [5, 3, 8]), + ([2, 0, 129, 1], [3, 1, 2, 0], [1, 7, 129, 127]), + ([0, 0, 0], [2, 0, 1], [0, 0, 0]), + ([0, 0, 0], [2, 0, 1], [4, 0, 4]), + ([7], [0], [7]), + ], +) +def test_outputs_and_gradients(dtype, prob_dtype, sizes, order, output_sizes): + """Check every element, including arbitrary gradients at padded rows.""" + count = sum(sizes) + with_probs = prob_dtype is not None + inp = torch.randn(count, 130, device="cuda", dtype=dtype)[:, ::2].requires_grad_() + probs = ( + torch.rand(count * 2, device="cuda", dtype=prob_dtype)[::2].requires_grad_() + if with_probs + else None + ) + expected_inp = inp.detach().clone().requires_grad_() + expected_probs = probs.detach().clone().requires_grad_() if with_probs else None + metadata = [ + torch.tensor(values, device="cuda", dtype=torch.int64) + for values in (sizes, order, output_sizes) + ] + actual, actual_probs, row_map = moe_sort_chunks_and_pad( + inp, *metadata, sum(output_sizes), probs + ) + expected, expected_p = reference(expected_inp, expected_probs, sizes, order, output_sizes) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + restored = moe_unpad_and_restore_chunks(actual, row_map, count) + torch.testing.assert_close(restored, inp, rtol=0, atol=0) + outputs, references = [actual], [expected] + if with_probs: + torch.testing.assert_close(actual_probs, expected_p, rtol=0, atol=0) + outputs.append(actual_probs) + references.append(expected_p) + grads = [torch.randn_like(output) for output in outputs] + torch.autograd.backward(outputs, grads) + torch.autograd.backward(references, grads) + torch.testing.assert_close(inp.grad, expected_inp.grad, rtol=0, atol=0) + if with_probs: + torch.testing.assert_close(probs.grad, expected_probs.grad, rtol=0, atol=0) + expert_output = torch.randn( + sum(output_sizes), 37, device="cuda", dtype=dtype, requires_grad=True + ) + restored = moe_unpad_and_restore_chunks(expert_output, row_map, count) + dy = torch.randn_like(restored) + restored.backward(dy) + valid = row_map >= 0 + torch.testing.assert_close(expert_output.grad[valid], dy[row_map[valid].long()], rtol=0, atol=0) + assert torch.count_nonzero(expert_output.grad[~valid]).item() == 0 + + +def test_probability_only_gradient(): + """An unused token output must not prevent probability differentiation.""" + inp = torch.randn(5, 7, device="cuda", requires_grad=True) + probs = torch.rand(5, device="cuda", requires_grad=True) + meta = [torch.tensor(v, device="cuda") for v in ([2, 3], [1, 0], [4, 4])] + _, output_probs, _ = moe_sort_chunks_and_pad(inp, *meta, 8, probs) + output_probs.sum().backward() + torch.testing.assert_close(probs.grad, torch.ones_like(probs), rtol=0, atol=0) + torch.testing.assert_close(inp.grad, torch.zeros_like(inp), rtol=0, atol=0) + + +@pytest.mark.parametrize("compiled", [False, True]) +def test_dynamic_metadata_graph_replay(compiled): + """Replay with changed routing and values at fixed input/output sizes.""" + inp = torch.randn(8, 65, device="cuda", requires_grad=True) + probs = torch.rand(8, device="cuda", requires_grad=True) + sizes = torch.tensor([2, 3, 3], device="cuda") + order = torch.tensor([2, 0, 1], device="cuda") + output_sizes = torch.tensor([4, 4, 4], device="cuda") + + def run(): + output, output_probs, row_map = moe_sort_chunks_and_pad( + inp, sizes, order, output_sizes, 12, probs + ) + restored = moe_unpad_and_restore_chunks(output, row_map, 8) + return output, output_probs, restored + + fn = torch.compile(run, fullgraph=True) if compiled else run + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + inp.grad = probs.grad = None + output, output_probs, restored = fn() + (output.sum() + output_probs.sum() + restored.sum()).backward() + del output, output_probs, restored + torch.cuda.current_stream().wait_stream(stream) + inp.grad = probs.grad = None + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output, output_probs, restored = fn() + (output.sum() + output_probs.sum() + restored.sum()).backward() + for values, permutation in (([1, 4, 3], [1, 0, 2]), ([4, 0, 4], [2, 1, 0])): + with torch.no_grad(): + inp.normal_() + probs.uniform_() + sizes.copy_(torch.tensor(values, device="cuda")) + order.copy_(torch.tensor(permutation, device="cuda")) + graph.replay() + torch.cuda.synchronize() + expected, expected_p = reference(inp, probs, values, permutation, [4, 4, 4]) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + torch.testing.assert_close(output_probs, expected_p, rtol=0, atol=0) + torch.testing.assert_close(restored, inp, rtol=0, atol=0) + torch.testing.assert_close(inp.grad, torch.full_like(inp, 2), rtol=0, atol=0) + torch.testing.assert_close(probs.grad, torch.ones_like(probs), rtol=0, atol=0) + graph.reset() diff --git a/transformer_engine/common/triton/chunk_padding.py b/transformer_engine/common/triton/chunk_padding.py new file mode 100644 index 00000000000..e2e4113963d --- /dev/null +++ b/transformer_engine/common/triton/chunk_padding.py @@ -0,0 +1,77 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Chunk permutation with zero padding between output chunks.""" + +import triton +import triton.language as tl + + +@triton.jit +def make_padded_chunk_map( + split_sizes, + sorted_indices, + output_split_sizes, + row_map, + num_chunks: tl.constexpr, + num_output_rows, + CHUNKS: tl.constexpr, + ROWS: tl.constexpr, +): + """Map each output row to an input row, or -1 for padding.""" + chunks = tl.arange(0, CHUNKS) + sizes = tl.load(split_sizes + chunks, chunks < num_chunks, other=0).to(tl.int32) + starts = tl.cumsum(sizes) - sizes + order = tl.load(sorted_indices + chunks, chunks < num_chunks, other=0).to(tl.int32) + output_sizes = tl.load(output_split_sizes + chunks, chunks < num_chunks, other=0) + ends = tl.cumsum(output_sizes.to(tl.int32)) + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + output_chunk = tl.sum( + ((rows[:, None] >= ends[None, :]) & (chunks[None, :] < num_chunks)).to(tl.int32), + axis=1, + ) + safe_chunk = tl.minimum(output_chunk, num_chunks - 1) + source_chunk = tl.gather(order, safe_chunk, axis=0) + offset = rows - tl.gather(ends - output_sizes, safe_chunk, axis=0) + source = tl.gather(starts, source_chunk, axis=0) + offset + valid = (output_chunk < num_chunks) & (offset < tl.gather(sizes, source_chunk, axis=0)) + tl.store(row_map + rows, tl.where(valid, source, -1), rows < num_output_rows) + + +@triton.jit +def copy_padded_chunks( + inp, + probs, + row_map, + output, + output_probs, + stride_row, + stride_col, + stride_prob, + num_padded_rows, + hidden_size: tl.constexpr, + REVERSE: tl.constexpr, + WITH_PROBS: tl.constexpr, + BLOCK: tl.constexpr, +): + """Copy tokens and probabilities together; reverse copies discard padding.""" + index = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + row = index // hidden_size + col = index % hidden_size + active = row < num_padded_rows + original = tl.load(row_map + row, active, other=-1).to(tl.int64) + if REVERSE: + source_row, target_row = row, original + else: + source_row, target_row = original, row + value = tl.load( + inp + source_row * stride_row + col * stride_col, active & (original >= 0), other=0 + ) + store = active & (original >= 0) if REVERSE else active + tl.store(output + target_row * hidden_size + col, value, store) + if WITH_PROBS: + probability = tl.load( + probs + source_row * stride_prob, active & (original >= 0) & (col == 0), other=0 + ) + tl.store(output_probs + target_row, probability, store & (col == 0)) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 576fb57c5e5..db4fda9c559 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -42,6 +42,8 @@ moe_unpermute, moe_sort_chunks_by_index, moe_sort_chunks_by_index_with_probs, + moe_sort_chunks_and_pad, + moe_unpad_and_restore_chunks, ) from transformer_engine.pytorch.quantization import fp8_autocast from transformer_engine.pytorch.quantization import fp8_model_init diff --git a/transformer_engine/pytorch/chunk_padding.py b/transformer_engine/pytorch/chunk_padding.py new file mode 100644 index 00000000000..402750e72e0 --- /dev/null +++ b/transformer_engine/pytorch/chunk_padding.py @@ -0,0 +1,233 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused chunk sorting, padding and restoration for dispatched MoE tokens.""" + +from typing import Optional, Tuple + +import torch +import triton + +from transformer_engine.common.triton.chunk_padding import ( + copy_padded_chunks, + make_padded_chunk_map, +) +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor + + +@torch.library.custom_op("te_moe::padded_chunk_map", mutates_args=[]) +def _padded_chunk_map( + splits: torch.Tensor, order: torch.Tensor, output_splits: torch.Tensor, num_rows: int +) -> torch.Tensor: + """Build a destination-to-source map without reading GPU metadata on the host.""" + result = torch.empty(num_rows, device=splits.device, dtype=torch.int32) + if num_rows: + make_padded_chunk_map[(triton.cdiv(num_rows, 128),)]( + splits, + order, + output_splits, + result, + splits.numel(), + num_rows, + triton.next_power_of_2(splits.numel()), + 128, + ) + return result + + +@_padded_chunk_map.register_fake +def _padded_chunk_map_fake( + splits: torch.Tensor, order: torch.Tensor, output_splits: torch.Tensor, num_rows: int +) -> torch.Tensor: + """Infer the map shape from the explicit output size.""" + del order, output_splits + return torch.empty(num_rows, device=splits.device, dtype=torch.int32) + + +@torch.library.custom_op("te_moe::padded_chunk_copy", mutates_args=[]) +def _padded_chunk_copy( + inp: torch.Tensor, + probs: Optional[torch.Tensor], + row_map: torch.Tensor, + num_rows: int, + reverse: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply a padded bijection or its inverse.""" + output = inp.new_empty((num_rows, inp.shape[1])) + output_probs = ( + probs.new_empty(num_rows) if probs is not None else torch.empty(0, device=inp.device) + ) + if row_map.numel(): + copy_padded_chunks[(triton.cdiv(row_map.numel() * inp.shape[1], 1024),)]( + inp, + probs, + row_map, + output, + output_probs, + inp.stride(0), + inp.stride(1), + probs.stride(0) if probs is not None else 0, + row_map.numel(), + inp.shape[1], + reverse, + probs is not None, + 1024, + ) + return output, output_probs + + +@_padded_chunk_copy.register_fake +def _padded_chunk_copy_fake( + inp: torch.Tensor, + probs: Optional[torch.Tensor], + row_map: torch.Tensor, + num_rows: int, + reverse: bool, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Infer token and probability output shapes.""" + del row_map, reverse + return inp.new_empty((num_rows, inp.shape[1])), ( + probs.new_empty(num_rows) if probs is not None else torch.empty(0, device=inp.device) + ) + + +def _copy_setup_context(ctx, inputs, output): + """Retain only the row map and scalar shape metadata for backward.""" + inp, probs, row_map, _, reverse = inputs + ctx.save_for_backward(row_map) + ctx.num_rows = inp.shape[0] + ctx.reverse = reverse + ctx.with_probs = probs is not None + if probs is None: + ctx.mark_non_differentiable(output[1]) + + +def _copy_backward(ctx, grad, grad_probs): + """The adjoint of a zero-padded bijection discards padded rows.""" + (row_map,) = ctx.saved_tensors + dx, dp = _padded_chunk_copy( + grad, grad_probs if ctx.with_probs else None, row_map, ctx.num_rows, not ctx.reverse + ) + return dx, dp if ctx.with_probs else None, None, None, None + + +_padded_chunk_copy.register_autograd(_copy_backward, setup_context=_copy_setup_context) + + +def moe_sort_chunks_and_pad( + inp: torch.Tensor, + split_sizes: torch.Tensor, + sorted_indices: torch.Tensor, + output_split_sizes: torch.Tensor, + num_out_tokens: int, + probs: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + """Sort chunks and append zero padding in a single token copy. + + Parameters + ---------- + inp : torch.Tensor + CUDA FP32, FP16 or BF16 tokens of shape ``[num_tokens, hidden_size]``. + Quantized tensors are not supported. + split_sizes : torch.Tensor + Contiguous CUDA integer chunk sizes in input order, summing to num_tokens. + sorted_indices : torch.Tensor + Contiguous CUDA integer permutation of the chunk indices. + output_split_sizes : torch.Tensor + Contiguous CUDA integer sizes in output order. Each entry must be at least + ``split_sizes[sorted_indices[i]]``; extra rows are filled with zeros. + To align an expert comprising several rank chunks, append its padding to + the last chunk of that expert. For example, chunks of sizes [2, 3] can use + output sizes [2, 6] to align their combined size to eight rows. + num_out_tokens : int + Sum of output_split_sizes. Supplied explicitly to avoid GPU-to-CPU + synchronization and permit CUDA graph capture with a fixed output size. + probs : torch.Tensor, optional + CUDA FP32, FP16 or BF16 probabilities of shape ``[num_tokens]``. + + Returns + ------- + output : torch.Tensor + Sorted, zero-padded tokens of shape ``[num_out_tokens, hidden_size]``. + output_probs : torch.Tensor or None + Sorted, zero-padded probabilities, if provided. + row_map : torch.Tensor + Int32 destination-to-source indices, with -1 at padding positions. + Pass this map and the original token count to moe_unpad_and_restore_chunks. + + Notes + ----- + Metadata values must satisfy the contracts above. Values are not copied to + the host for validation. Metadata may change during graph replay if the + input and output buffer sizes remain fixed. Padding has zero gradient. + """ + dtypes = (torch.float32, torch.float16, torch.bfloat16) + if ( + isinstance(inp, QuantizedTensor) + or not inp.is_cuda + or inp.ndim != 2 + or inp.dtype not in dtypes + or inp.shape[1] == 0 + ): + raise ValueError("Expected a CUDA FP32, FP16 or BF16 token matrix") + if split_sizes.numel() == 0 or any( + tensor.device != inp.device + or tensor.ndim != 1 + or not tensor.is_contiguous() + or tensor.dtype not in (torch.int32, torch.int64) + or tensor.numel() != split_sizes.numel() + for tensor in (split_sizes, sorted_indices, output_split_sizes) + ): + raise ValueError("Expected matching nonempty CUDA integer metadata vectors") + if num_out_tokens < inp.shape[0] or num_out_tokens >= 2**31: + raise ValueError("Output size must cover all input rows and fit in int32") + if probs is not None and ( + isinstance(probs, QuantizedTensor) + or probs.device != inp.device + or probs.shape != (inp.shape[0],) + or probs.dtype not in dtypes + ): + raise ValueError("Expected one floating-point probability per input row") + row_map = _padded_chunk_map(split_sizes, sorted_indices, output_split_sizes, num_out_tokens) + output, output_probs = _padded_chunk_copy(inp, probs, row_map, num_out_tokens, False) + return output, output_probs if probs is not None else None, row_map + + +def moe_unpad_and_restore_chunks( + inp: torch.Tensor, row_map: torch.Tensor, num_tokens: int +) -> torch.Tensor: + """Discard padding and restore the input order of moe_sort_chunks_and_pad. + + Parameters + ---------- + inp : torch.Tensor + Expert outputs of shape ``[num_padded_tokens, hidden_size]``. The hidden + size may differ from the input passed to moe_sort_chunks_and_pad. + row_map : torch.Tensor + The unmodified map returned by moe_sort_chunks_and_pad. + num_tokens : int + Original number of unpadded tokens. + + Returns + ------- + torch.Tensor + Restored tokens. Backward fills padded rows with zero gradients. + """ + if ( + isinstance(inp, QuantizedTensor) + or not inp.is_cuda + or inp.ndim != 2 + or inp.dtype not in (torch.float32, torch.float16, torch.bfloat16) + or inp.shape[1] == 0 + ): + raise ValueError("Expected a CUDA floating-point token matrix") + if ( + row_map.device != inp.device + or row_map.dtype != torch.int32 + or row_map.shape != (inp.shape[0],) + or not row_map.is_contiguous() + or not 0 <= num_tokens <= inp.shape[0] + ): + raise ValueError("Expected the padded row map and original token count") + return _padded_chunk_copy(inp, None, row_map, num_tokens, True)[0] diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index 5e6bdb4dfcd..d809e80b2eb 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -16,11 +16,17 @@ from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor +from transformer_engine.pytorch.chunk_padding import ( + moe_sort_chunks_and_pad, + moe_unpad_and_restore_chunks, +) __all__ = [ "moe_permute", "moe_unpermute", "moe_sort_chunks_by_index", + "moe_sort_chunks_and_pad", + "moe_unpad_and_restore_chunks", ]