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

tests/pytorch/test_permutation.py

@@ -1739,6 +1739,43 @@ def test_permutation_mask_map(
     )
 
 
+@pytest.mark.parametrize("torch_dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_permutation_mask_map_capacity_drop(torch_dtype):
+    """num_out_tokens smaller than routing_map.sum(): over-capacity entries must be
+    dropped (as the CUDA index-map path does), not written past the output buffer."""
+    num_tokens, num_expert, hidden_size = 64, 8, 128
+    torch.manual_seed(1234)
+
+    routing_map = (torch.rand(num_tokens, num_expert) > 0.5).bool().cuda()
+    total_routed = int(routing_map.sum())
+    capacity = total_routed - 17
+    assert capacity > 0
+
+    inp = torch.randn((num_tokens, hidden_size)).cuda().to(torch_dtype)
+
+    out, row_id_map = te_permute(inp, routing_map, num_out_tokens=capacity, map_type="mask")
+
+    # reference: expert-major order of routed (expert, token) pairs, first
+    # `capacity` destination rows kept
+    token_idx = torch.nonzero(routing_map.T.contiguous(), as_tuple=False)[:, 1]
+    ref = inp[token_idx[:capacity]]
+    assert out.shape[0] == capacity
+    torch.testing.assert_close(out.float(), ref.float(), atol=0, rtol=0)
+
+    # the row map must mark exactly the over-capacity entries as dropped
+    n_routed = row_id_map[:, 2 * num_expert]
+    assert int(n_routed.sum()) == capacity
+
+    # unpermute consumes the same map: each token gets the sum of its kept copies
+    unperm = te_unpermute(out, row_id_map, restore_shape=inp.shape, map_type="mask")
+    kept_flat = torch.zeros(num_expert * num_tokens, dtype=torch.bool, device=inp.device)
+    routed_flat_idx = torch.nonzero(routing_map.T.contiguous().flatten(), as_tuple=False)[:, 0]
+    kept_flat[routed_flat_idx[:capacity]] = True
+    kept_map = kept_flat.view(num_expert, num_tokens).T.contiguous()
+    ref_unperm = (inp.float().unsqueeze(1) * kept_map.unsqueeze(-1).float()).sum(1)
+    torch.testing.assert_close(unperm.float(), ref_unperm, atol=1e-2, rtol=1e-2)
+
+
 @pytest.mark.parametrize("te_dtype", _te_dtypes)
 @pytest.mark.parametrize("num_out_tokens", [None])
 @pytest.mark.parametrize(

transformer_engine/common/triton/permutation.py

@@ -131,6 +131,7 @@ def _row_id_map_pass_2_kernel(
     workspace_ptr,
     # sizes
     num_tokens,
+    num_out_tokens,
     # strides
     stride_row_id_map_token,
     stride_row_id_map_expert,
@@ -155,6 +156,9 @@ def _row_id_map_pass_2_kernel(
         -1,
         row_id_within_token_block + tl.sum(n_tokens_per_chunk) - 1,
     )
+    # capacity limit: destinations at or past num_out_tokens are dropped, matching
+    # the `idx >= num_out_tokens` branch of moe_permute_row_map in permutation.cu
+    row_id = tl.where(row_id >= num_out_tokens, -1, row_id)
     tl.store(
         row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token,
         row_id,

transformer_engine/jax/triton_extensions/permutation.py

@@ -251,6 +251,9 @@ def lowering(ctx, row_id_map, workspace, *, num_tokens, num_experts, block_size)
             input_output_aliases={0: 0, 1: 1},
             constexprs={
                 "num_tokens": num_tokens,
+                # the JAX path sizes its buffers from num_out_tokens up front, so
+                # disable the kernel's capacity drop here (no entry can reach INT32_MAX)
+                "num_out_tokens": 2**31 - 1,
                 "stride_row_id_map_token": row_id_stride_token,
                 "stride_row_id_map_expert": row_id_stride_expert,
                 "WORKSPACE_LOAD_WIDTH": workspace_load_width,

transformer_engine/pytorch/permutation.py

@@ -318,7 +318,9 @@ def moe_permute_mask_map_forward(
     num_tokens, hidden_size = inp.size()
     num_experts = routing_map.size(1)
 
-    row_id_map = triton_permutation.make_row_id_map(routing_map, num_tokens, num_experts)
+    row_id_map = triton_permutation.make_row_id_map(
+        routing_map, num_tokens, num_experts, num_out_tokens
+    )
 
     # FP8 handling
     fp8 = isinstance(inp, QuantizedTensor)

transformer_engine/pytorch/triton/permutation.py

@@ -25,6 +25,7 @@ def make_row_id_map(
     routing_map: torch.Tensor,
     num_tokens: int,
     num_experts: int,
+    num_out_tokens: int,
 ):
     """
     Prepare the row_id_map for the permutation.
@@ -39,6 +40,9 @@ def make_row_id_map(
         Number of tokens in the input tensor.
     num_experts : int
         Number of experts in the input tensor.
+    num_out_tokens : int
+        Number of rows in the permuted output. Routing-map entries whose destination row falls at
+        or past this limit are dropped (marked -1), like the CUDA index-map path does.
 
     Returns
     -------
@@ -94,6 +98,7 @@ def make_row_id_map(
         row_id_map,
         workspace_tensor,
         num_tokens,
+        num_out_tokens,
         row_id_map.stride(0),
         row_id_map.stride(1),
         triton.next_power_of_2(num_experts * triton.cdiv(num_tokens, block_size)),