fix(ep): correct recv-slice indexing in DispatchInterNodeRecv (v1 kernel) - #621
Draft
isytwu wants to merge 1 commit into
Draft
fix(ep): correct recv-slice indexing in DispatchInterNodeRecv (v1 kernel)#621isytwu wants to merge 1 commit into
isytwu wants to merge 1 commit into
Conversation
… recv tokens in DispatchInterNodeRecv DispatchInterNodeRecv splits each 64-token chunk into numRecvBlock=8 interleaved slices so multiple blocks/warps can drain it in parallel. The slice a given loop iteration should handle is `bid % numRecvBlock`, where `bid` is the (chunk, slice) index that iteration is actually processing -- CombineInterNodeTyped, which decomposes the identical grid-stride loop, does exactly this a few hundred lines down. DispatchInterNodeRecv instead used `blockId % numRecvBlock`: the block's own fixed identity, constant across every bid value that block visits. The two coincide only when rdmaBlockNum is a multiple of numRecvBlock (8), because that is the only case where `blockId + n*rdmaBlockNum` keeps a constant residue mod 8 for every n. For any other rdmaBlockNum, a block's later iterations are computing a different real chunk (bid/8 changes) but sourcing tokens from the wrong slice of it -- some slices get processed by multiple blocks (duplicate atomicAdd, inflating the recv-slot counter), others by none (their real tokens are simply never copied out). Whether this is visible depends on how much real data lands in the mis-sliced range, which is why it looked geometry-dependent rather than a clean on/off switch: rdma_block_num=2 crashed at warp_per_block=8/12 (device assert `destTokId < MaxNumTokensToRecv()`, the recv-slot counter overrun from the duplicate atomicAdds) but not at warp_per_block=4/6, and rdma_block_num=3/4/5/6/7 (equally non-multiples of 8) didn't crash at all in the same sweep -- all consistent with the same defect, just not always enough real traffic in the affected slices to overrun capacity or to be missing enough tokens to fail the output check. Reproduced and fixed on a 2-node MI300X EP16 setup (kernel_type=v1, hidden 6144, bf16): `--cmd bench --block-num 8 --warp-per-block 8 --rdma-block-num 2 --max-tokens 16` hit the device assert on every run before this change. After rebuilding with the fix, the same command and five other previously-tested geometries (warp_per_block 4/6/8/12 at rdma_block_num=2, rdma_block_num=3/5 at warp_per_block=8) all pass the warmup round's dispatch/combine correctness check. Re-ran the full 105 config v1 tuning sweep at max_tokens=16 (the sweep that originally hit this assert at combo 7/105, block_num=8/warp_per_block=8/rdma_block_num=2) to completion with zero assertions.
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.
Summary
DispatchInterNodeRecv(inter-node dispatch,kernel_type=v1) splits each64-token chunk into
numRecvBlock=8interleaved slices so multipleblocks/warps can drain it in parallel. Which slice a given loop iteration
should handle is
bid % numRecvBlock, wherebidis the (chunk, slice)index that iteration is actually processing —
CombineInterNodeTyped,which decomposes the identical grid-stride loop, does exactly this a few
hundred lines down.
DispatchInterNodeRecvinstead usedblockId % numRecvBlock: the block's own fixed identity, constant acrossevery
bidvalue that block visits.The two only coincide when
rdma_block_numis a multiple ofnumRecvBlock(8), since that's the only case whereblockId + n*rdmaBlockNumkeeps a constant residue mod 8 for everyn.For any other
rdma_block_num, a block's later iterations are computing adifferent real chunk (
bid/8changes) but sourcing tokens from the wrongslice of it: some slices get processed by multiple blocks (duplicate
atomicAdd, inflating the recv-slot counter), others by none (their realtokens are never copied out). Whether this is observable depends on how
much real data lands in the mis-sliced range, which is why it looked
geometry-dependent rather than a clean on/off switch:
rdma_block_num=2crashed atwarp_per_block=8/12(device assertdestTokId < MaxNumTokensToRecv(), the recv-slot counter overrun from theduplicate atomicAdds) but not at
warp_per_block=4/6, andrdma_block_num=3/4/5/6/7(equally non-multiples of 8) didn't crash atall in the same sweep — all consistent with the same defect, just not
always enough real traffic in the affected slices to overrun capacity or
to drop enough tokens to fail the output check.
Fix
One-line change: use
bid % numRecvBlockinstead ofblockId % numRecvBlock,matching the sibling combine implementation.
Test plan
Reproduced and fixed on a 2-node MI300X EP16 setup (
kernel_type=v1,hidden_dim=6144, bf16):
--cmd bench --block-num 8 --warp-per-block 8 --rdma-block-num 2 --max-tokens 16hit the device assert on every run before this changepreviously-tested geometries (
warp_per_block4/6/8/12 atrdma_block_num=2,rdma_block_num3/5 atwarp_per_block=8) allpass the warmup round's dispatch/combine correctness check
max_tokens=16(thesweep that originally hit this assert at combo 7/105,
block_num=8/warp_per_block=8/rdma_block_num=2) to completion withzero assertions
Scope: this only affects
kernel_type=v1(non-LL).InterNodeV1LLuses aseparate kernel (
EpDispatchInterNodeV1KernelLowLatency_) and isunaffected.