[Bug] Triton mask-map moe_permute ignores num_out_tokens

Summary

The Triton path behind moe_permute(..., map_type="mask") sizes its output at exactly num_out_tokens rows but never enforces that limit: _permute_kernel receives num_out_tokens explicitly marked # pylint: disable=unused-argument, the row-id map assigns destination rows as an unbounded cumulative sum over the routing map, and nothing clamps them. If the routing map routes more tokens than num_out_tokens, the excess rows are written past the end of the allocation.

The CUDA path behind the same public API (map_type="index") implements the capacity limit — moe_permute_row_map has an explicit idx >= num_out_tokens branch that marks over-capacity entries -1, and every consumer skips -1 (common/permutation/permutation.cu lines 28–34) — so the two backends of one entry point disagree, and a caller expecting the documented capacity-drop behaviour from the mask path gets out-of-bounds writes. #2876 added a positivity guard for num_out_tokens, but a positive value smaller than routing_map.sum() still reaches the kernel unchecked, and the tests always construct the map so that routing_map.sum() == num_out_tokens, which is why this has not been seen.

The kernel was added in #1373.

Steps to reproduce

import torch
import transformer_engine.pytorch  # noqa
from transformer_engine.pytorch.permutation import moe_permute

routing_map = torch.zeros(8, 4, dtype=torch.int32, device="cuda")
routing_map[:, 0] = 1
routing_map[:4, 1] = 1          # 12 routed entries
inp = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16)
out, row_id_map = moe_permute(inp, routing_map, num_out_tokens=9, map_type="mask")
# out has 9 rows; the kernel writes 12

Under compute-sanitizer --tool memcheck (with PYTORCH_NO_CUDA_MEMORY_CACHING=1 so the caching allocator's slack does not absorb it), this reports 48 invalid global writes at the permute store (transformer_engine/common/triton/permutation.py:303) — the three over-capacity rows times the 16-element hidden size — and the run then ends in an unspecified launch failure. After the fix: 0 errors, the in-capacity rows are unchanged, and the row map's per-token n_routed sums to num_out_tokens.

Environment

• TransformerEngine commit 07e281f2 (main); transformer-engine-cu12 2.17.0 wheel + source checkout for the Python/Triton code
• NVIDIA L40S (sm_89), driver CUDA 12.4
• torch 2.13.0+cu126, triton 3.7.1, Python 3.12

Start in transformer_engine/common/triton/permutation.py at _permute_kernel and the permute store around line 303, then compare the capacity handling in common/permutation/permutation.cu. Reproduce the issue with the provided 12-entry routing map and num_out_tokens=9 under compute-sanitizer. Done means no invalid writes, unchanged in-capacity rows, and row_id_map n_routed sums to num_out_tokens.