perf(EPv2): optimize epv2 disp/comb kernel performance - #520
Draft
kawhil-amd wants to merge 6 commits into
Draft
Conversation
Rework the intranode dispatch phase-1 token-embedding scatter from per-(token,expert) redundant loads to a warp-per-token load-once / store-many scheme: the local token tile is loaded ONCE and fanned out to every distinct dest PE, removing the per-k redundant HBM reloads of the same token. The main copy loop issues 4 concurrent vec4 load streams (4-way memory-level parallelism) before storing to each published dest. Also bump combine gather inner-unroll to 4 (_unroll=4). Peak intranode bandwidth (EP4, hidden=7168, topk=6, bf16, DISP 256x8w): dispatch 4096: 756 GB/s 8192: 931 GB/s 16384: 1108 GB/s Co-authored-by: Cursor <cursoragent@cursor.com>
…e vs original), both 4-way
make_dispatch gains a compile-time `load_once` switch that emits one of two
Phase-1 implementations (each compiled variant contains only ONE path, so
there is no VGPR union / occupancy penalty):
- load_once=True : warp-per-token load-once / store-many (local tile read
ONCE, fanned out to every distinct dest) — wins at large token counts.
- load_once=False: the original warp-per-(token,expert) path — finer work
granularity fills the grid at small/mid token counts.
Both paths now use the SAME 4-way (4 vec4 streams) load MLP, so the two
variants differ ONLY in load-once/store-many vs per-work-item reload (the
original path was previously 2-way, which unfairly halved its copy width).
Host-side selection: EpDispatchCombineOp precompiles both algorithm variants
per (block, warp) spec (keyed by (block, warp, load_once)) and picks one at
launch from input token count via cfg.load_once_threshold (default 4096:
tokens > threshold use load-once, else the original path). Variants are
JIT-compiled lazily, so an unused algorithm never pays compile cost.
bench: add LOAD_ONCE env to pin the dispatch algorithm in non-AUTO mode.
Measured crossover (EP4, hidden=7168, topk=6, bf16, each at its best geometry):
1024 tok: original 256x32w 480 GB/s > load-once 256x8w 424 (+13%)
4096 tok: ~tie (~760 GB/s)
16384 tok: load-once 1078 > original 973 (+11%)
Co-authored-by: Cursor <cursoragent@cursor.com>
…r contention Replace the combine gather cross-block sync (comb_bar) with per-block private flag counters. Each block owns xdb_flag[bid] (all in lockstep == call count), bumps only its own (single writer, no atomic), and independently polls the shared local xdb slots; block 0 still does the npes cross-device pushes. - Delete the exit barrier: the former (block_num-1)-way contended atomic_add on comb_bar was the dominant combine latency, especially at small token counts. - Poll with >= (spin_until_ge_i64) not ==: every block reads the slot, so a faster peer can lap and overwrite its monotonic push before a slow rank's late blocks read it; >= stays deadlock-free (peer-ahead is the safe direction). - No acquire fence after the poll: gather reads peer out_tok via non-temporal (cache-bypassing) loads and the spin is a control dependency, so observing the flag already implies the data is ready. - cross_device_flag grows to max(combine block_num) i64s; the last block fills the unused tail so counters stay synced across calls with different block_num. EP4 hidden=7168 topk=6 bf16 gather combine (GRAPH), same geometry: 256: 342->129us 1024: 351->140us 4096: 465->239us 8192: 663->377us 8192 @ 256x8w peaks at ~1022 GB/s (vs 855 GB/s at the old 256x4w default). Co-authored-by: Cursor <cursoragent@cursor.com>
… combine Rebuild _GFX1250_SCHED_BF16_T6 from a full block x warp sweep with the load-once/store-many dispatch + per-block-flag combine kernels. Combine now wants block=256 (= CU count, 1 block/CU) with warp ramping 4->8; the old "<CU, 192 ceiling" guardrail was too conservative. Dispatch flips 256x32 (original) -> 256x8 (load-once) across the 4096 threshold. Measured GB/s disp/comb: 256=140/155 1024=428/378 4096=756/741 8192=968/968 16384=1108/1149. Co-authored-by: Cursor <cursoragent@cursor.com>
- Stage 1 scatter copy: vec4 × multi-stream load (matching dispatch pattern) - Stage 3 local reduce: vec4 × _unroll pre-load (matching gather combine pattern) - _COMB_NSTREAMS: 4 for gfx1250 (wave32), 2 for others (wave64) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
kawhil-amd
changed the title
kawhil-amd
marked this pull request as draft
August 7, 2026 10:13
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
1 participant
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
Kernel-level optimizations for EPv2 intranode dispatch and combine, targeting memory-level parallelism, synchronization overhead, and architecture-adaptive tuning.
Changes
1. Dispatch load-once/store-many with multi-stream MLP
Rework dispatch Phase-1 from per-(token, expert) redundant HBM reads to a load-once / store-many scheme: each warp loads the token tile once and fans out to all destination PEs. The copy loop issues 4 concurrent vec4 (16B) load streams for MLP. Host-side algorithm selection:
load_once_threshold(default 4096) picks between load-once (large batches) and the original per-work-item path (small batches) at launch time. Both variants share the same 4-way vec4 MLP. Variants are JIT-compiled lazily — unused algorithms never pay compile cost.2. Combine gather: per-block flag entry barrier, exit barrier removed
Replace
comb_baratomic sync with per-block monotonic flag counters — each block owns a private counter (single writer, no atomic). Exit barrier deleted entirely (was the dominant latency source).3. Re-tune EP4 bf16 topk6 schedule
Full block × warp sweep with new kernels. Combine now uses
block=256(= CU count) with warp ramping 4→8.4. Architecture-conditional stream count
Dispatch and combine unroll counts adapt to wave size: 4 streams for gfx1250 (wave32), 2 for gfx942/gfx950 (wave64).
5. Scatter combine vec4 + unroll
Apply the same vec4 + multi-stream optimizations to the scatter combine path (previously scalar-only): Stage 1 scatter copy uses vec4 × multi-stream loads; Stage 3 local reduce uses vec4 × unroll pre-load with batched accum.