Repository navigation
[PyTorch] Make the Triton mask-map permute respect num_out_tokens - #3347
Conversation
github-actions
Bot
added
the
community-contribution
 [Medium impact] The PR appears safe to merge; no actionable issues were found. SummaryThe PyTorch Triton mask-map path now drops destinations at or past
No actionable issues were found. Tests were not run during this review. Reviews (2) · Last reviewed commit: "[PyTorch] Harden the mask-map capacity-d..." · Reviewed by Greptile |
denera
left a comment
denera
left a comment
There was a problem hiding this comment.
Thanks for the fix! Everything looks good aside from two particular nits about hardening the unit test for coverage and tolerance. We can launch CI once you make those changes and merge If it comes back clean.
| 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 |
There was a problem hiding this comment.
num_tokens=64 here stays under one 1024-token block so the cross-block prefix is never exercised. Can you parameterize this with an additional num_tokens > 1024 case as well?
| 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) |
There was a problem hiding this comment.
atol=rtol=1e-2 seems very relaxed for the FP32 case. We have an assert_close() implementation in the test utils that scales tolerances with dtype. I would recommend using that, or justifying why these tolerances should be kept.
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 <truongvu0911nd@gmail.com>
Add a 4096-token case so the row-map kernels span several 1024-token blocks, and compare the unpermuted output with the test-utils assert_close, which scales tolerances with dtype. Signed-off-by: Alp Dener <adener@nvidia.com>
20ae222 to
d5efe56
Compare
|
/te-ci pytorch |
Description
Fixes #3346.
The Triton mask-map path of
moe_permutesizes its output at exactlynum_out_tokensrows, but_permute_kernelreceivesnum_out_tokensmarked unused and the row-id map assigns destination rows as an unbounded cumsum — so a routing map with more routed entries thannum_out_tokenswrites 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-1in its row map (moe_permute_row_map'sidx >= num_out_tokensbranch).This PR gives the Triton path the same semantics at the same place:
_row_id_map_pass_2_kerneltakesnum_out_tokensand marks destinations at or past it as-1— pass 3 then excludes them fromn_routed, and every consumer of the map (permute, unpermute, and their backwards) already skips dropped entries via the existing-1handling. No compute kernel changes.make_row_id_mapand its caller threadnum_out_tokensthrough.2**31 - 1), keeping JAX behavior unchanged; wiringtoken_dispatch'snum_out_tokensinto it can follow separately if desired.test_permutation_mask_map_capacity_drop(fp32/fp16/bf16): capacity 17 belowrouting_map.sum(), checking the kept rows against an expert-major reference permutation, the row map'sn_routedaccounting, 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 FP8DelayedScalingcases 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
make_row_id_map).Environment