fix(ep): AsyncLL slot assignment double-allocates when top-k does not divide warpSize - #505
Merged
Merged
Conversation
added 2 commits
July 29, 2026 08:19
…opk != 0 SlotAssign tiles a wavefront into tokensPerWarp = warpSize / numExpertPerToken token groups. When top-k does not divide warpSize the trailing lanes form an incomplete group that aliases the next warp's first token, so warp g lanes [tokensPerWarp*numEpt, warpSize) and warp g+1 lanes [0, warpSize % numEpt) compute the same entry. Both take the atomicAdd branch, so one entry consumes two slots on destPeTokenCounter while dispDestTokIdMap keeps only the last writer's. The surplus slot is orphaned: SendCopyMultiBlock never writes it, but RecvTransfer still signals the inflated count, so the receiver reads untouched memory as a token -- giving pe = id / numExpertPerRank out of range and tripping the assert in RecvCopyMultiBlock. This corrupts dispatch before it ever crashes: some tokens are routed twice with one copy reading stale memory, so runs that complete without asserting are also affected. Latent because 64 % 8 == 0. DeepSeek-V4 uses top-k 6 (64 % 6 == 4), which exposed it. Small batches also masked it, since orphan slots still held valid ids from earlier iterations until the batch grew past that high-water mark. Active lanes only __shfl from baseLane + j < tokensPerWarp * numEpt, so masking these lanes cannot disturb the dedup shuffle. No-op where top-k divides warpSize.
…% topk != 0
check_dispatch_result cannot run at non-multiple-of-8 top-k, and
_AsyncLLCombineOnlyTestCase skips it with a comment asserting the dispatch data
is correct there -- exactly the assumption the lane-tiling bug violates. That is
why the existing top-k=9 coverage, which also has a straggler lane, passed.
The new checks are top-k agnostic: received count must equal the deduplicated
host-side expectation, source positions must be unique, and each must decode to
an in-range (src_rank, src_id). An orphan slot fails all three.
top-k 8 is the clean-tiling control; 6 is DeepSeek-V4's and exposed this; 9/10/12
cover other remainders. max_num_inp_token_per_rank=128 exceeds tokensPerWarp so
adjacent warps actually collide rather than relying on a rare boundary hit.
Verified on 8x MI355X: unfixed -> FAILED ("Rank[0] received 587 dispatched
tokens, expected 547"); fixed -> 5 passed.
TianDi101
force-pushed
the
fix/asyncll-slotassign-lane-tiling
branch
from
July 29, 2026 08:19
5eff956 to
cd9d2bd
Compare
Oseltamivir
added a commit
to SemiAnalysisAI/InferenceX-app
that referenced
this pull request
Aug 30, 2026
…opk-6 assert) Bisected on-metal with a pure mori.ops probe on mia1-p01-g18: AsyncLL passes at top-k 8 with both 256 and 384 experts and dies at top-k 6 with both, on the device assert (pe >= 0) && (pe < worldSize) in EpDispatchLowLatencyAsyncRecvCopyMultiBlock. Known upstream: ROCm/mori#505 ('AsyncLL slot assignment double-allocates when top-k does not divide warpSize', merged 2026-07-31); every shipped mi35x image predates it (newest mori-0706). Cells flip green with the image bump; no new upstream issue needed.
Oseltamivir
added a commit
to SemiAnalysisAI/InferenceX
that referenced
this pull request
Aug 31, 2026
…(mori >= #505) AsyncLL under the DSv4-Pro workload (topk 6) dies on every shipped mi35x-mori image: ROCm/mori#505 (AsyncLL slot assignment double-allocates when top-k does not divide warpSize, fixed upstream 2026-07-31) postdates them all. On-metal probe of this date-stamped nightly (job 41563, pure mori.ops): AsyncLL topk6@384 and topk8@256 both pass, split-phase recv API present. mi355x only — the -tw SKUs keep their image until separately validated on their docker path.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
EpDispatchLowLatencyAsyncSendCopySlotAssigntiles a wavefront intotokensPerWarp = warpSize / numExpertPerTokentoken groups. If top-k does not dividewarpSize, the trailing lanes form an incomplete group that resolves to the next warp'sfirst token:
Both warps take the
atomicAddbranch, so one entry consumes two slots whiledispDestTokIdMap[i]keeps only one. The surplus slot is orphaned —SendCopyMultiBlocknever writes it, but
RecvTransfersignals the inflated count, so the receiver readsuntouched memory as a token. Its index bytes give
pe = id / numExpertPerRankout ofrange, tripping the assert in
RecvCopyMultiBlock.Note this corrupts dispatch before it ever crashes: some tokens are routed twice with one
copy reading stale memory, so runs that complete without asserting are also affected.
Hidden until now because
64 % 8 == 0. DeepSeek-V4 uses top-k 6 (64 % 6 == 4), whichexposed it. Top-k 9 was already covered but
_AsyncLLCombineOnlyTestCaseskipscheck_dispatch_resultat non-multiple-of-8 top-k, so it could not catch this.Fix
Active lanes only
__shflfrombaseLane + j < tokensPerWarp * numEpt, so masking theselanes cannot disturb the dedup shuffle. No-op where top-k divides
warpSize. This tilingpattern appears at one site in the tree.
Test
test_dispatch_slot_assign_lane_tilingchecks three top-k agnostic invariants, so it workswhere the positional check cannot: received count matches the deduplicated host-side
expectation, source positions are unique, and each decodes to an in-range
(src_rank, src_id). Parametrized over top-k8, 6, 9, 10, 12(8 is the clean-tilingcontrol).
Verified on 8x MI355X:
Rank[0] received 587 dispatched tokens, expected 547test_dispatch_combine_async_ll.py, after fix