[PyTorch] Fuse MoE chunk sorting and padding

[bzantium]

Description

Sorting dispatched MoE tokens into expert order and padding them for grouped GEMM currently requires separate copies. Add a fused sort/pad operation that avoids the sorted, unpadded intermediate and copies tokens and probabilities together. A matching restore operation removes padding and restores the original order in one copy.

Chunk order and padded sizes are explicit, so callers can group rank chunks by expert and choose an alignment without depending on a GEMM backend. Routing metadata stays on the GPU; an explicit output size supports CUDA graph replay with changing routing at fixed buffer sizes. Token counts are runtime kernel arguments, allowing compiled kernels to be reused when routing changes the number of tokens.

Type of change

• Documentation change
• Bug fix
• New feature
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• Add moe_sort_chunks_and_pad and moe_unpad_and_restore_chunks for FP32, FP16 and BF16 tensors, with optional probabilities, autograd and torch.compile support.
• Save a row map for restoration and backward. Padding is zero-filled and contributes zero gradient.
• Add API documentation, output/gradient and graph replay tests, and a reproducible benchmark. Quantized tensors are outside this change.

Validation

On one B200, with BF16 tokens, FP32 probabilities, hidden size 2048, eight rank chunks per expert, 16 experts and alignment 128:

Dispatched tokens	Separate sort/pad/restore	Fused path	Time reduction	Extra peak allocation, separate → fused
8,192	0.3251 ms	0.1717 ms	47.2%	132.10 → 68.07 MiB
32,768	0.3497 ms	0.2029 ms	42.0%	516.88 → 260.76 MiB

These are layout roundtrip measurements, excluding GEMM and communication. Each entry is the median of two run medians in A-B-B-A order, with 50 CUDA-event samples after warmup. Memory is incremental peak PyTorch allocation, not total GPU memory. The baseline uses existing TE chunk sorting and Fp8Padding/Fp8Unpadding. Both paths run on a full source build of this branch (TE 2.20.0.dev0), with PyTorch 2.11.0a0 / CUDA 13.1 / cuDNN 9.20.

A separate cold-cache probe over 25 routed token counts uses one map kernel and two copy kernels, compared with 25 and 50 variants when token counts are specialized.

All 63 new GPU tests and 12 selected existing permutation tests pass against the freshly built package and extension. Coverage includes exact outputs and gradients, empty chunks, mixed probability dtypes, and changed-routing graph replay with and without compilation. The build disables optional NCCL-EP. Test runs emit dependency deprecations and a pytest cache-path warning. Repository-wide pre-commit, production Python lint and license checks pass.

uv run python -m pytest tests/pytorch/test_chunk_padding.py -q
uv run python benchmarks/benchmark_chunk_padding.py --tokens 8192 32768 --output chunk-padding.json

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the prior missing-CI-coverage concern is fully fixed and no actionable new failures were found.

Findings

1.  Tests Missing From CI ▶

Summary

• Adds a GPU row-map kernel and fused forward/reverse copy kernel.
• Exposes the new PyTorch APIs and documents their metadata and padding contracts.
• Covers output values, gradients, dynamic routing, compilation, and CUDA graph replay.
• Adds the dedicated chunk-padding test suite to the L0 PyTorch CI runner, fully addressing the previous review finding.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Input tokens and optional probabilities] --> B[Build padded destination-to-source row map]
    M[Chunk sizes, order, and padded sizes] --> B
    B --> C[Fused sort, copy, and zero padding]
    C --> D[Expert-ordered padded tokens]
    C --> E[Padded probabilities]
    D --> F[Grouped expert computation]
    F --> G[Unpad and restore using row map]
    B --> G
    G --> H[Original token order]

Loading

_{Reviews (2) · Last reviewed commit: "test: run chunk padding tests in PyTorch..."}

[greptile-apps]

Tests Missing From CI

The new output, gradient, compilation, and CUDA graph tests are not included in the explicitly enumerated L0 PyTorch test runner. As a result, the main CI gate will not execute this dedicated test suite, allowing future regressions in this feature to go unnoticed. Please add test_chunk_padding.py to qa/L0_pytorch_unittest/test.sh.

Knowledge Base Used: Verification and CI matrix

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[bzantium]

Added test_chunk_padding.py to the L0 PyTorch runner in 856167e, alongside the existing permutation tests. Shell syntax validation and all pre-commit checks pass.