Skip to content

fix(ep): correct recv-slice indexing in DispatchInterNodeRecv (v1 kernel) - #621

Draft
isytwu wants to merge 1 commit into
mainfrom
fix/v1-dispatch-recv-residue
Draft

fix(ep): correct recv-slice indexing in DispatchInterNodeRecv (v1 kernel)#621
isytwu wants to merge 1 commit into
mainfrom
fix/v1-dispatch-recv-residue

Conversation

@isytwu

@isytwu isytwu commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Summary

DispatchInterNodeRecv (inter-node dispatch, kernel_type=v1) splits each
64-token chunk into numRecvBlock=8 interleaved slices so multiple
blocks/warps can drain it in parallel. Which 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 only coincide when rdma_block_num is a multiple of
numRecvBlock (8), since that's the only case where
blockId + n*rdmaBlockNum keeps a constant residue mod 8 for every n.
For any other rdma_block_num, 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 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=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 drop enough tokens to fail the output check.

Fix

One-line change: use bid % numRecvBlock instead of blockId % 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 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

Scope: this only affects kernel_type=v1 (non-LL). InterNodeV1LL uses a
separate kernel (EpDispatchInterNodeV1KernelLowLatency_) and is
unaffected.

… 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.
@isytwu isytwu self-assigned this Aug 31, 2026
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