From 20ae222b7864174b7285f8ca91402b7ee7f9167a Mon Sep 17 00:00:00 2001 From: Truong Vu Date: Wed, 5 Aug 2026 10:14:09 +0000 Subject: [PATCH] [PyTorch] Make the Triton mask-map permute respect num_out_tokens The Triton _permute_kernel received num_out_tokens explicitly marked unused, and nothing clamped destination rows against it, while the host sizes the output at exactly num_out_tokens rows: a routing map that routes more tokens than num_out_tokens wrote the excess rows past the end of the allocation. The CUDA index-map path behind the same public API drops over-capacity entries by marking them -1 in the row map (moe_permute_row_map); do the same in row-id-map pass 2, so every consumer of the map (permute, unpermute, backward) skips them via the existing -1 handling. The JAX lowering passes an explicit no-drop sentinel, keeping its current behavior unchanged. Signed-off-by: Truong Vu --- tests/pytorch/test_permutation.py | 37 +++++++++++++++++++ .../common/triton/permutation.py | 4 ++ .../jax/triton_extensions/permutation.py | 3 ++ transformer_engine/pytorch/permutation.py | 4 +- .../pytorch/triton/permutation.py | 5 +++ 5 files changed, 52 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_permutation.py b/tests/pytorch/test_permutation.py index 517d6e1543..e4aeaa13bb 100644 --- a/tests/pytorch/test_permutation.py +++ b/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( diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index b3893843af..cf26ea9317 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/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, diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index 22f983f078..e3e1a4d8c0 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/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, diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index bccc486b4f..5dd2ec477d 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/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) diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index c155d73e1e..4f8a2ad7a4 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/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)),