Skip to content

fix(ep): AsyncLL slot assignment double-allocates when top-k does not divide warpSize - #505

Merged
TianDi101 merged 2 commits into
mainfrom
fix/asyncll-slotassign-lane-tiling
Jul 31, 2026
Merged

fix(ep): AsyncLL slot assignment double-allocates when top-k does not divide warpSize#505
TianDi101 merged 2 commits into
mainfrom
fix/asyncll-slotassign-lane-tiling

Conversation

@TianDi101

@TianDi101 TianDi101 commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Problem

EpDispatchLowLatencyAsyncSendCopySlotAssign tiles a wavefront into
tokensPerWarp = warpSize / numExpertPerToken token groups. If top-k does not divide
warpSize, the trailing lanes form an incomplete group that resolves to the next warp's
first token:

warpSize=64, numEpt=6 -> tokensPerWarp=10
warp g   lanes 60..63 -> token g*10 + 10
warp g+1 lanes 0..3   -> token (g+1)*10     // same entry

Both warps take the atomicAdd branch, so one entry consumes two slots while
dispDestTokIdMap[i] keeps only one. The surplus slot is orphaned — SendCopyMultiBlock
never writes it, but RecvTransfer signals the inflated count, so the receiver reads
untouched memory as a token. Its index bytes give pe = id / numExpertPerRank out of
range, 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), which
exposed it. Top-k 9 was already covered but _AsyncLLCombineOnlyTestCase skips
check_dispatch_result at non-multiple-of-8 top-k, so it could not catch this.

Fix

if ((inWarpTokIdx >= tokensPerWarp) || (tokenId >= args.curRankNumToken)) continue;

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. This tiling
pattern appears at one site in the tree.

Test

test_dispatch_slot_assign_lane_tiling checks three top-k agnostic invariants, so it works
where 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-k 8, 6, 9, 10, 12 (8 is the clean-tiling
control).

Verified on 8x MI355X:

New test, before fix FAILED — Rank[0] received 587 dispatched tokens, expected 547
New test, after fix 5 passed
Full test_dispatch_combine_async_ll.py, after fix 194 passed, 233 skipped, 0 failed

ditian12 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
TianDi101 force-pushed the fix/asyncll-slotassign-lane-tiling branch from 5eff956 to cd9d2bd Compare July 29, 2026 08:19
@TianDi101
TianDi101 merged commit 12d1bc3 into main Jul 31, 2026
31 of 35 checks passed
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.
@QizhouZhang97
QizhouZhang97 deleted the fix/asyncll-slotassign-lane-tiling branch September 3, 2026 08:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant