[PyTorch] Fuse MoE chunk sorting and padding

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()

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
 ---------------------------------

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"

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()

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))

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