perf(ep): optimize IntraNode dispatch kernel for MI350X

[kudomcho]

Motivation

IntraNode EP8 dispatch on MI350X (gfx950) shows a 12% bandwidth gap against combine at large token counts on the DeepSeek V3 configuration (BF16 same-type, 7168 hidden, top-8 experts, zero-copy combine). This gap was reported as a ~10% performance deficiency in dispatch vs combine for high-token intra-node workloads.

The portable dispatch body (EpDispatchIntraNodeKernel_body in intranode.hpp) has three sources of overhead that do not exist in the combine kernel path:

1. Redundant routing computation: Phase 3 re-reads tokenIndices, re-computes destPe, and re-checks per-PE deduplication — all of which Phase 1 already computed.
2. Serialized completion waits: The grid barrier and slot-drain wait run sequentially, costing cbar + cslot instead of max(cbar, cslot).
3. Suboptimal WarpCopy unrolling: WarpCopy<T, 8> at BF16 hidden=7168 produces a degenerate unrolled loop (elemsPerWarp = 4096 > hiddenDim), falling back to Unroll=1 for the entire copy.

Technical Details

Kernel changes (src/ops/dispatch_combine/intranode.hpp)

1. Cached routing (eliminate Phase 3 redundancy)

Phase 1 now caches each kept token's destPe in dispDestTokIdMap with a sentinel bit (0x40000000). Phase 3 reads the cached value with a single global load + __shfl instead of re-reading tokenIndices (up to topk global loads per pair), re-dividing by numExpertPerRank, and re-checking dedup via __any(). This eliminates ~32768 redundant global memory reads per kernel launch at 4096 tokens.

The sentinel bit (0x40000000 = 2^30) is safe because FlatTokenIndex values are bounded by worldSize * MaxNumTokensToSend(), which at EP8 with 4096 tokens is well below 2^30.

2. Overlapped completion waits

Ported from the gfx1250 dispatch body (intranode_1250x.hpp lines 811-826): the slot-drain wait (ShmemInt32WaitUntilEquals on peer memory) is issued before the grid barrier (ShmemUint32WaitUntilEquals on local memory). The slot read targets uncached peer memory and pays a full xGMI fabric round-trip even when the slot has long been drained — issuing it while the grid barrier is still spinning hides this latency. The two waits are independent (the slot address depends only on destPe, not on the barrier state).

3. WarpCopy Unroll 8 → 2

WarpCopy<T, 8> with BF16 (2 bytes) and wave64 computes elemsPerWarp = 8 × 64 × 8 = 4096. At hiddenDim = 7168, the main unrolled loop body runs only once (covering 4096 of 7168 elements), with the remaining 3072 elements handled by the Unroll=1 fallback — producing suboptimal instruction scheduling for the remote store pipeline.

WarpCopy<T, 2> gives elemsPerWarp = 1024, producing 7 well-pipelined unrolled iterations that better utilize the memory controller's store queue for cross-GPU xGMI writes.

4. Vectorized metadata copies

Weights and indices copies (8 floats + 8 ints per token) are converted from lane-parallel scalar stores (only 8 of 64 lanes active) to WarpCopy calls that use all 64 lanes. This is a code quality improvement; the metadata is too small (64 bytes) relative to the 14 KB payload to produce a measurable BW change.

Benchmark script (tools/bench_dispatch_gap.sh)

New benchmark script for rigorous dispatch-vs-combine comparison:

• Verifies clean GPU state (< 15% VRAM, 0% GPU util) before running
• Records full environment (kernel version, diff, GPU arch, ROCm, geometry)
• Reports mean, stddev, and best BW across 5 benchmark iterations
• Saves raw data and summary to bench_results/
• Configurable dispatch/combine geometry via --dispatch-geo / --combine-geo

Additional finding: missing bf16 dispatch tuning configs

PR #464's gfx950_mi350x_IntraNode_ep8_dispatch.json has only one bf16 entry (64 tokens at 128×4). With MORI_EP_LAUNCH_CONFIG_MODE=AUTO, large-token bf16 dispatch falls back to this geometry, producing 284 GB/s instead of 363 GB/s with the correct 2048×16 geometry. Adding bf16 dispatch tuning entries for 128–524288 tokens is recommended as a follow-up.

Performance Results

All numbers measured on MI350X (gfx950, 256 CUs), BF16 same-type EP8 IntraNode, 7168 hidden, top-8 experts, zero-copy combine. Dispatch geometry: 2048×16, combine geometry: 56×15. Each row is mean ± stddev GB/s over 5 iterations (3 warmup, 10 graph replays per iter) on verified-idle GPUs.

Kernel optimization (same geometry, 2048×16)

Tokens	Baseline Dispatch	Optimized Dispatch	Combine	Baseline Gap	Optimized Gap
4096	362.9 ± 0.4 GB/s	374.0 ± 0.6 GB/s	392.4 ± 0.2 GB/s	8.1%	4.9%
8192	367.0 ± 0.4 GB/s	380.7 ± 0.4 GB/s	401.6 ± 0.1 GB/s	9.4%	5.5%
16384	369.2 ± 0.3 GB/s	383.2 ± 0.4 GB/s	407.0 ± 0.1 GB/s	10.2%	6.2%

Kernel-level improvement: +3.1% dispatch BW, consistent across token counts.

Geometry optimization (256×16 → 2048×16)

Geometry	Warps	4096t Dispatch BW	Gain
256×16	4,096	349.6 GB/s	baseline
512×16	8,192	365 GB/s	+4.4%
1024×16	16,384	368 GB/s	+5.3%
2048×16	32,768	374 GB/s	+7.0%
4096×16	65,536	374 GB/s	saturated

More blocks generate more concurrent xGMI write traffic, better saturating the 7-link fabric. The dispatch kernel uses minimal shared memory per block (3 × 8 × 4B = 96B for s_N/s_base/s_run), so high block counts do not pressure LDS.

Combined improvement

Tokens	Before (256×16, original kernel)	After (2048×16, optimized kernel)	Combine	Gap reduction
4096	349.6 GB/s	374.0 GB/s (+7.0%)	392.4 GB/s	12.1% → 4.9%
8192	357.3 GB/s	380.7 GB/s (+6.5%)	401.6 GB/s	12.3% → 5.5%
16384	361.4 GB/s	383.2 GB/s (+6.0%)	407.0 GB/s	12.7% → 6.2%

Small/mid token results (dispatch faster than combine)

At small token counts, the optimized dispatch is faster than combine — the kernel overhead (Phase 1/2/completion) is amortized and dispatch's push model has lower fixed latency than combine's barrier + accumulation.

Tokens	Dispatch BW (GB/s)	Dispatch Lat (us)	Combine BW (GB/s)	Combine Lat (us)	Gap
256	242.7 ± 0.8	79.4	196.0 ± 0.2	98.5	dispatch 19% faster
512	295.8 ± 0.9	130.5	308.2 ± 0.3	125.8	+4.2%
1024	312.3	242.1	338.1	228.3	+8.3%
2048	356.5 ± 0.8	434.0	375.0 ± 0.2	413.7	+5.2%

Dispatch geometry: 128×16 (256/512t), 256×16 (1024t), 2048×16 (2048t). Combine geometry: per MI350X tuning configs.

Comparison with MI355X reference (docs/EP8.en.md)

Tokens	MI350X Dispatch (optimized)	MI355X Dispatch (reference)	MI355X Gap
4096	374.0 GB/s	359.9 GB/s	16.5%
8192	380.7 GB/s	369.9 GB/s	15.0%
16384	383.2 GB/s	377.4 GB/s	13.8%

MI350X optimized dispatch now exceeds MI355X reference dispatch BW, and the gap vs combine (4.9–6.2%) is better than MI355X's own gap (13.8–16.5%).

Remaining gap (~5%) — profiled root cause

Local-write profiling (redirecting payload to local HBM instead of remote xGMI) measured 213 us of routing overhead out of the total 860 us dispatch time at 4096 tokens:

Component	Time (us)	% of total
Phase 1 (routing + dedup) + Phase 2 (remote atomics) + barriers	213	25%
xGMI payload writes	647	75%
Total dispatch	860	100%
Combine (reference)	791	—

The xGMI writes alone (647 us) are 18% faster than combine (791 us), confirming the fabric transfer itself is efficient. The 213 us routing overhead is the gap — combine has no equivalent phase (it receives pre-routed data and goes straight to P2P reads + FMA accumulation).

Approaches tested to reduce the overhead:

• Single-pass (merge Phase 1+2+3, per-token remote atomics): regressed -2.5% due to atomic contention
• Token-level Phase 1 (process all topk experts per warp via ): GPU crash — intrinsic not available on gfx950 wave64
• Skip drop writes: correctness failure — combine reads stale dispDestTokIdMap
• Local staging + bulk copy: correctness failure — buffer layout conflicts need host-side changes

Further reduction of the 213 us overhead requires either restructuring Phase 1 to process at the token level (needs wave64-compatible dedup), or host-side changes to pre-compute routing on a separate kernel launch.

Test Plan

Hardware required: 8× AMD Instinct MI350X or MI355X (gfx950)

# Build
BUILD_UMBP=OFF pip install .

# Correctness: run IntraNode dispatch/combine stress test
PYTHONPATH=. MORI_SHMEM_HEAP_SIZE=6G python3 tests/python/ops/bench_dispatch_combine.py \
  --max-tokens 4096 --dtype bf16 --hidden-dim 7168 --world-size 8 \
  --num-experts-per-token 8 --kernel-type IntraNode --zero-copy 1 \
  --dispatch-block-num 2048 --dispatch-warp-per-block 16 \
  --cmd stress

# Performance: rigorous dispatch vs combine comparison
bash tools/bench_dispatch_gap.sh --label test --dispatch-geo 2048x16 --tokens "4096 8192 16384"

# Regression: verify FP8 dispatch is unaffected
PYTHONPATH=. MORI_SHMEM_HEAP_SIZE=6G MORI_EP_LAUNCH_CONFIG_MODE=AUTO python3 \
  tests/python/ops/bench_dispatch_combine.py \
  --max-tokens 4096 --dtype fp8_e4m3 --combine-dtype bf16 --hidden-dim 7168 \
  --world-size 8 --num-experts-per-token 8 --kernel-type IntraNode \
  --zero-copy 0 --quant-type fp8_direct_cast --cmd bench

# Full test suite
pytest tests/python/ops/test_dispatch_combine_intranode.py -q

Test Result

Stress test: 200 rounds, 0 failures (MI350X, BF16 EP8, 4096 tokens, 2048×16).

Performance (MI350X, clean GPUs):

Tokens   |  Disp Mean   Disp Std  Disp Best |  Comb Mean   Comb Std  Comb Best |    Gap
---------+----------------------------------+----------------------------------+-------
    4096 |      374.0       0.57      374.7 |      392.4       0.18      392.6 |  +4.9%
    8192 |      380.7       0.40      381.3 |      401.6       0.11      401.8 |  +5.5%
   16384 |      383.2       0.36      383.6 |      407.0       0.11      407.3 |  +6.2%

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
• Kernel changes are architecture-neutral (portable body, not gfx125x-specific)
• No new dependencies introduced
• Correctness verified via stress test (200 rounds)
• Performance measured rigorously (mean ± stddev, clean GPUs, multiple iterations)
• No regression on FP8 dispatch path

[TianDi101]

Great work, thanks @kudomcho !
Here're some minor suggestions:

1. The CI stage pre_commit failed, please format the code to clean this failure
2. Name the magic constant. 0x40000000 appears three times; make it a constexpr index_t with a meaningful name and comments why this magic constant is selected.
3. Comment the load-bearing invariant. The cached-routing scheme is correct only because Phase 1 and Phase 3 share identical loop bounds, stride, and guard, so the same warp writes and reads dispDestTokIdMap[i]. Anyone who later changes one loop desynchronizes them and silently corrupts routing. One comment prevents that.
4. Restore the upstream rationale comment for the wait reorder. intranode_1250x.hpp:825-826 carries 14 lines explaining why the two spin-waits are independent, plus its isolated A/B numbers. The PR replaced it with a one-liner; that comment is exactly what the next reviewer needs.

[kawhil-amd]

@kudomcho Hi, have you tested the performance of intranode_ll dispatch on MI350? On MI355, its peak bandwidth is around ~380 GB/s. Also, do these optimizations bring any performance benefit to intranode_ll dispatch? Could we try porting them to intranode_ll dispatch?

[kudomcho]

@TianDi101 All four items addressed:

1. ✅ clang-format — applied, CI should pass now
2. ✅ Named constant — constexpr index_t kCachedRoutingSentinel = 0x40000000 with comment explaining why the bit is safe (bounded by worldSize * MaxNumTokensToSend(), well below 2^30)
3. ✅ Phase 1/3 loop invariant — added comment at Phase 3 explaining that both loops must share identical bounds, stride, and guard for the cached-routing scheme to work (same warp writes and reads dispDestTokIdMap[i])
4. ✅ Completion wait rationale — restored the full upstream comment from intranode_1250x.hpp:811-826 explaining why the two spin-waits are independent, the overlap strategy, and the isolated A/B measurements (+8.7% at 512, +1.6% at 4096)

Commits: e63ab92, c90694e

[kudomcho]

@kawhil-amd Tested IntraNodeLL dispatch on MI350X (BF16 EP8, AUTO tuning configs from PR #464):

Tokens	Dispatch BW	Combine BW	Gap
512	309.8 GB/s	306.7 GB/s	dispatch 1% faster
1024	342.3 GB/s	338.2 GB/s	dispatch 1.2% faster
2048	358.9 GB/s	373.9 GB/s	+4.2%
4096	368.2 GB/s	391.8 GB/s	+6.4%

MI350X IntraNodeLL dispatch peaks at ~368 GB/s (vs MI355X ~380 GB/s — expected from 256 vs 304 CUs).

Regarding porting the optimizations:

• WarpCopy Unroll=2: already applied in IntraNodeLL (kCopyUnroll = 2 at line 379 of intranode_ll.hpp)
• Cached routing: does not apply — IntraNodeLL uses a single-pass ComputeTokenRoute (one block per token), not the three-phase count→reserve→copy architecture
• Completion wait overlap: applicable — IntraNodeLL has the same non-overlapped pattern (grid barrier at line 427, then slot drain at line 432). Reordering these to overlap should give the same +1-2% at large tokens. I can submit this as a follow-up patch.

The 4-6% gap at large tokens (2048-4096) is similar to IntraNode V1's gap and comes from the same root cause: dispatch's routing + remote atomic slot reservation overhead that combine doesn't have.