[PyTorch] Make the Triton mask-map permute respect num_out_tokens

[truong-v]

Description

Fixes #3346.

The Triton mask-map path of moe_permute sizes its output at exactly num_out_tokens rows, but _permute_kernel receives num_out_tokens marked unused and the row-id map assigns destination rows as an unbounded cumsum — so a routing map with more routed entries than num_out_tokens writes the excess rows past the end of the allocation (48 invalid global writes under compute-sanitizer in the linked issue's repro, ending in an unspecified launch failure; 0 errors after this change). The CUDA index-map path behind the same public API implements the capacity limit by marking over-capacity entries -1 in its row map (moe_permute_row_map's idx >= num_out_tokens branch).

This PR gives the Triton path the same semantics at the same place:

• _row_id_map_pass_2_kernel takes num_out_tokens and marks destinations at or past it as -1 — pass 3 then excludes them from n_routed, and every consumer of the map (permute, unpermute, and their backwards) already skips dropped entries via the existing -1 handling. No compute kernel changes.
• The PyTorch make_row_id_map and its caller thread num_out_tokens through.
• The JAX lowering passes an explicit no-drop sentinel (2**31 - 1), keeping JAX behavior unchanged; wiring token_dispatch's num_out_tokens into it can follow separately if desired.
• New test_permutation_mask_map_capacity_drop (fp32/fp16/bf16): capacity 17 below routing_map.sum(), checking the kept rows against an expert-major reference permutation, the row map's n_routed accounting, and the unpermute of the dropped map. All three fail without the fix and pass with it; the rest of the mask-map suite is unchanged — 151 passed / 113 skipped / 16 failed, with the same 16 failing ids, both with and without the change (those 16 are FP8 DelayedScaling cases that already fail at this commit in my prebuilt-wheel setup, on an NVRTC/CUDA-header version mismatch that has nothing to do with permutation).

Checklist

• The documentation is up to date with these changes (docstring of make_row_id_map).
• The tests are up to date with these changes.

Environment

• Base commit: 07e281f (main); prebuilt transformer-engine-cu12 2.17.0 + source checkout for Python/Triton
• NVIDIA L40S (sm_89), driver CUDA 12.4
• torch 2.13.0+cu126, triton 3.7.1, Python 3.12

[greptile-apps]

Greptile Summary

The PR bounds Triton mask-map row IDs by the PyTorch output capacity while preserving JAX’s existing no-drop behavior.

• Threads num_out_tokens through the PyTorch row-map generation path.
• Marks over-capacity destinations as dropped before map compaction.
• Adds forward, row-count, and unpermute coverage for capacity-limited routing.

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, security, or compatibility issue was identified.

The new bound is applied before map compaction, all downstream kernels consume only compacted valid routes, PyTorch callers provide the real allocation size, and JAX explicitly retains its existing full-buffer semantics.

Important Files Changed

Filename	Overview
transformer_engine/common/triton/permutation.py	Adds the capacity check at row-map construction so all compacted-map consumers consistently skip dropped routes.
transformer_engine/pytorch/triton/permutation.py	Extends the row-map wrapper with the output capacity and forwards it in the correct kernel argument position.
transformer_engine/pytorch/permutation.py	Threads the custom-op capacity into row-map generation without changing the public custom-op signature.
transformer_engine/jax/triton_extensions/permutation.py	Supplies an explicit INT32_MAX sentinel in the lowering to preserve the current full-capacity JAX behavior.
tests/pytorch/test_permutation.py	Adds numerical and map-accounting coverage for routes exceeding a smaller output capacity.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Routing mask] --> B[Pass 1: block counts]
  B --> C[Pass 2: global row IDs]
  C --> D{row ID below num_out_tokens?}
  D -- Yes --> E[Keep destination]
  D -- No --> F[Mark as -1]
  E --> G[Pass 3: compact valid routes]
  F --> G
  G --> H[Permute / unpermute / backward]

Loading

_{Reviews (1): Last reviewed commit: "[PyTorch] Make the Triton mask-map permu..." | Re-trigger Greptile}