Repository navigation
poc(cco): MegaMoE stage1 dispatch over SDMA - #1
Merged
kawhil-amd merged 7 commits intoNov 6, 2026
Merged
Conversation
aiter MegaMoEV2 stage1 dispatch moved onto the cco SDMA copy engines: * one resident prepare kernel plans, pushes the route table, quantises each token once per destination rank into a send buffer and puts the first chunks; the GEMM1 kernel's comm workgroup puts the rest * GEMM1 reads its A rows and e8m0 scales through a rowmap and starts a tile as soon as the chunks its expert needs have landed, with no per-tile cache invalidation * tests compare byte for byte with aiter (rows, route table, stage1 output, final layer output) and include a poisoned stress run Stage1 output is bit-exact. Measured 1.03-1.15x on prepare + stage1 against aiter main without ROCm/aiter#7001; probably behind that PR at large batch, not measured. One call at a time: no protection of S and R between back-to-back calls yet. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…ks sooner * quant is driven by tokens, not by rows: a token's bf16 row is read and quantised once and stored to its row in every destination's run, instead of once per destination (~4.6x fewer reads of the input) * a grid rendezvous separates the quant pass that fills the early chunks from the rest, so the second pass no longer races the first for bandwidth and delays the puts; the comm workgroup then puts those chunks with nothing to check, which drops the per-chunk completion counters entirely * `--early-chunks` sets how many chunks of each run prepare puts itself * clock stamps in both kernels behind `--trace`, plus `--stress` and `--stall-test` in the GEMM1 test, and a check that the fanout-group segments (which the route push does not cover) are empty At bs 8192 the first put moves from 177 to 106 us and the first tile's wait for data from 53 to 7 us; prepare pays ~25 us for the rendezvous. Flash, prepare + stage1 against aiter main without ROCm/aiter#7001: 1.36x at 512 tokens/rank, 1.17x at 1024, 1.14x at 2048, 1.06x at 4096, 1.00x at 8192. Two things measured and rejected rather than kept: writing through only the rows a put reads while the kernel runs (no change), and walking the early rows off `row_tok` instead of scanning the tokens (slower -- up to 75% of the emissions are duplicates at small batch). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
* gemm1_sdma: aiter's _load_row_map_subgroup takes a subgroup_width since #6259; pass padded_groups_per_row, as aiter's own indexed A-scale stage does. * stage1_sdma: comm_ops.traced only exists when example 11 has been imported first (it installs the fallback); test_prepare.py imports stage1_sdma alone and failed at import. Install the same fallback here. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The route push wrote each route to its final expert-major row on the destination, numbered by aiter's compact plan (pair_order, pair_base, local_hist, my_base). That plan has to run every call, 0.10-0.15 ms whatever the batch, and with it timed the SDMA stage1 was 0.92-1.04x of aiter's (test_gemm1.py leaves it out and reports 0.98-1.33x). With Stage1SdmaPrepare.run(..., recv_layout=True) the receiver builds its layout itself: * the sender puts every route of its, in no order, into a staging area of the destination partitioned by source (Arena.off_stage: local expert, rowmap value, weight, srcmap); slots are reserved per block (LDS counts, one atomic per block and destination); the last block out writes the counts (off_scount) and bumps the flag by epr, the total arm() expects; * build_layout (64 resident blocks, between the prepare and GEMM1) waits for every source, counts the routes per expert, pads each expert to the sort block and writes rowmap and aiter's srcmap_em / wts_em / sorted_expert_ids / tile_row_base / tile_input_base / num_valid / expert_tile_end / max_expert_tiles, so aiter's stage2 runs on them as they are. No runtime fanout pairs: layout() zeroes fanout_pair_config (aiter's output is the same with and without them; the test checks it). Arena.max_routes now leaves room for the padding (up to epr * 128 rows past the routes; it wrote into the next array before). build_layout closes over its window offsets as plain ints: read through `arena` they are not in FlyDSL's cache key, and a kernel cached for another hidden size with the same experts / top-k / cap was reused (V4-Pro's ran for Flash and faulted). test_recv_layout.py: final output bit-exact against aiter over 20 poisoned gated calls, and new prepare + layout + gated GEMM1 timed against aiter's compact prepare + fused stage1 (eager, 8x MI355X, aiter 46cdc4c52): tokens/rank 256 512 1024 4096 8192 16384 GLM-5.2 1.20x 1.18x 1.19x 1.10x 1.05x 1.00x V4-Pro 1.11x 1.07x 1.08x 1.04x 1.01x 0.97x V4.1-Flash 1.31x 1.24x 1.07x 1.01x 0.98x 0.98x Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ile slice * FlyDSL's cache key holds a closure's scalars only, not the Arena, so a kernel compiled for one chunk_rows / early_chunks / hidden could be reused for another from the same cache directory (gemm1_sdma, prepare, post_rest, wait). They now close over repr(arena). * arm() takes a device tensor of chunk counts, so a whole call (arm + prepare + layout + gated GEMM1) can be captured in a CUDA graph. test_recv_layout.py --graph times both sides as graphs after the eager numbers and checks poisoned replays bit for bit against aiter: 1.02-1.40x (GLM-5.2 / V4-Pro / Flash, 256-16384 tokens/rank), eager 0.96-1.29x. * test_recv_layout.py --trace: gate wait and compute per tenth of GEMM1's tile order, and when each tenth first passed its gate; plus GEMM1 alone, ungated and gated with every row landed. * --stream-group (experimental, kept for reference): prepare puts every chunk itself as soon as the quant pass that fills it is done, per-pass counters and no rendezvous. Removes the GEMM1 head wait but prepare and layout grow by as much; ~1% at bs 16384. * recv_layout's per-destination route cursors move to cursor[E + d], off the plan's counters. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Counting buffer_wbl2 (whole-L2 writeback) and buffer_inv (whole-L2 invalidate) in the ISA, not timing guesses, found most of this. On gfx950 every fence_*_release is a writeback and every acquire load is an invalidate, and the kernels had 9 of them. * GEMM1: the per-tile gate loaded `need` with acquire=True, so every TILE invalidated the whole L2 -- with the expert weights in it -- and the start gate did one per poll iteration. Both are redundant: `need` is published before the source bumps its flag, the kernel already does one system acquire after the start gate, and the loads are system-scope. buffer_inv 4 -> 1. * tree_barrier takes `release`: the writeback is only needed when what the barrier publishes was written with PLAIN stores. Barrier 1 publishes `hist` via atomics and barrier 3 publishes SEND_CM write-through stores, so only the plan's barrier 2 keeps it. prepare buffer_wbl2 3 -> 1. * layout: the flags spin invalidated the L2 on every poll (now one acquire after it), and neither of its barriers needs a writeback -- #1 publishes atomics and ROCm#2 is a rendezvous before block 0 zeroes the counters, with the tables it wrote read by the NEXT kernel. buffer_wbl2 2 -> 0, buffer_inv 4 -> 3. * the route push published scount and the flags from ONE thread walking the 8 peers: 16 serialised cross-device round trips, ~10 us. One lane per peer instead, ~2.5 us. Cross-device accesses are uncached, so the cost is the number of serialised trips, not the bytes -- packing its four 4 B stores into one 16 B store changed nothing. * the slot reservation skips its atomic when a block has no route for a destination (255 blocks x 8 destinations of atomics to place ~1500 routes), and the blocks-done counter is two-level like tree_barrier. New knobs, all off by default except where noted: `--split N` (N waves share a token, each taking a slice of its hidden; 3 is the measured best and is what the numbers use), `--prep-block`, `--stage-first` (push the routes before the quant so the peer's layout overlaps the rest of prepare), `--big-rest`, `--quant-sweep` (one CUDA graph per config in a single process, each checked bit-exact), and a prepare phase breakdown under `--trace` (plan / staging / publish / quant). V4.1-Flash stage1 against the FASTER of aiter's two dispatch modes, CUDA graph, median of 30 replays, bit-exact: 1.19x at 256, 1.40x at 512, 1.20x at 1024, 1.14x at 2048, 1.11x at 4096 tokens/rank (was 1.11 / 1.29 / 1.13 / 1.10 / 1.08x). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
kawhil-amd
changed the title
kawhil-amd
changed the title
…he quant stores Two mechanisms, both off by default, both bit-exact against aiter. `--ready-first`: GEMM1 used to gate every tile on the SDMA signals itself, once per tile. Wave 0 of the comm workgroup now reads each source's signal level once per round and publishes m-tiles into a ready queue in the order they become ready; a worker claims a ticket with one atomic and reads the queue, and never touches the signals. Needs `--fuse-layout`, since the scheduler starts when the layout in the same kernel is done. `--quant-lds-store`: the quantised row left each lane as two 16-byte pieces 32 bytes apart, so every store instruction covered half of each 128 B line of an uncached, write-through window. A wave now stages its round of 64 groups in LDS (2304 B per wave) and reads it back so each lane stores one contiguous 16 B -- 1024 contiguous bytes per wave instruction -- with the e8m0 scales packed four per dword. Same bytes, same cache modifier. ROCm#904 measured the same change at 3.1x on its own dispatch; here `prepare` at bs 4096 goes from 114.3 us to 73.4 (quant pass 1: 54.3 -> 29.7). Flash, 8x MI355X, CUDA graph, interleaved A/B, 3 reps, medians (ms): | tokens/rank | base | C1 | C1 + swizzle | aiter compact | |---|---|---|---|---| | 1024 | 0.389 | 0.354 | 0.345 | 0.463 | | 2048 | 0.586 | 0.505 | 0.485 | 0.670 | | 4096 | 0.845 | 0.811 | 0.778 | 0.910 | | 8192 | 1.541 | 1.479 | - | 1.536 | C1 is `--own-tokens --fuse-layout --ready-first`. At bs 256 the scheduler is a small loss (0.222 against 0.218): the first publish waits for the layout and for the slowest peer's first chunk, which no reordering shortens. Under C1 the tile loop is compute-bound -- 354 us of `do_tile` in a 365 us loop at bs 2048, 5 us of queue wait -- and the front end is gated by the slowest peer's first chunk (put to landing is 70-85 us, ~28 of it the wire time for 1.35 MB), so the remaining front-end budget is 20-30 us. Measured and not kept, with numbers, in the README: fusing the prepare into GEMM1 (+-1%), a dynamic arrival-ordered layout (-90 us at bs 2048), local-first and arrival-order row ordering, layout on dedicated blocks or merged into the scheduler block, more quant blocks (511 is net zero, 767 deadlocks), the plan on fewer blocks, unroll, quant-skip, a pre-built doorbell, and `--tile-n 256` (which gives the ~14 us per-tile fixed cost that rules out variable-M tiles). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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.
What
Example 13 (
examples/cco/python/13_flydsl_mega_moe_stage1_sdma/): MegaMoEV2 stage1 dispatch over SDMA on gfx950 / MI355X, EP8. The sender quantises each token once into a send buffer grouped by destination, the copy enginesputeach run to its destination, and GEMM1 (aiter's owndo_tile) reads its A rows in place through arowmap. No CU is spent on the copy and no aiter plan runs. aiter's stage2 + combine run unchanged on the tables it produces; the layer output is bit for bit aiter's.Commits:
8db36383prototype: SDMA stage1 dispatch, GEMM1 gated on SDMA signals.a7deffdeach token read and quantised once and stored to every rank it goes to; early chunks put after one rendezvous, the rest from GEMM1's comm workgroup.bcfb887run against current aiter.45309dfrecv_layout: no aiter plan at all — senders stage routes on the destination, the receiver builds the expert-major tables.049e600CUDA-graph timing, FlyDSL cache-key fix,--traceby tile slice.f4465decut the whole-L2 cache operations out of stage1 (counted in the ISA).d2d72f2publish tiles in readiness order (--ready-first), coalesce the quant stores through LDS.dd181d2keep what measured, add--no-plan, drop the rest: proven mechanisms baked in, measured-dead switches deleted with their code (options 41 + 3 env → 17;stage1_sdma.py2904 → 1971 lines), the scripts that drove the deleted paths removed.How it works now
puts the early chunks.puts the rest; the workers build the receive layout, then run aiter's tile, either in a static order gated per tile on the SDMA signals, or with--ready-first, where one scheduler wave publishes m-tiles in readiness order and workers take tickets.--no-plan(new): no counting sort — one reservation atomic per worker and run, no histogram / prefix sum / grid rendezvous — and tiles wait for all of a source's chunks. Free while a run fits in the early chunks (small batch).--fuse-prepareruns the prepare inside the GEMM1 kernel.Switches left:
--ready-first,--no-plan,--fuse-prepare,--chunk-rows,--early-chunks,--head-rows,--sbm,--tile-n.Where we stand — 2026-11-09
8x MI355X, CUDA graph, median of 30 replays, interleaved A/B with 3 reps, every run bit-exact. Stage1 = quant + dispatch + GEMM1 on both sides. "aiter" = aiter's compact mode measured in the same process, except at bs 256 where aiter's fixed-slot mode is faster and is measured alone (standalone runs of aiter read 1-5% slower than in-process ones on V4-Pro; the in-process number is used).
V4.1-Flash
--no-plan--ready-first --no-plan--ready-first--ready-first--ready-first--ready-firstV4-Pro (fixed-slot is slower than compact at every size measured)
--ready-first--ready-firstWhat each mechanism is worth (same-run A/B): the LDS-swizzled stores take Flash's prepare at bs 4096 from 114.3 to 73.4 us and scale with the bytes staged (−3 us at 256 … −78 us at 8192);
--ready-firstis −35 … −85 us on Flash from 1024 up and a loss at 256 (+11 us) and on V4-Pro 2048 (+54 us, not explained);--no-planis −12 us at Flash 256, −5 at 512, +23 at 1024.The simplified tree was checked against the previous one on the same node: Flash 256
--no-plan0.200 / 0.201 ms, 512--ready-first --no-plan0.255 / 0.253, 2048--ready-first0.481 / 0.474, 4096 0.774 / 0.770, V4-Pro 2048 0.741 / 0.739; bit-exact in 8 configurations including--fuse-prepare.Why the speedup falls with the batch: what we save over aiter is roughly fixed (≈130-150 us from bs 2048 up: GEMM1 on 255 CUs instead of 224, and a shorter front end), while both sides run the same
do_tile, which is 80%+ of the time at bs 4096. Each tile carries ~14 us that does not depend on N (--tile-n 256vs 512).Measured and not kept
Fusing the prepare into GEMM1 as the default (±1%), a dynamic arrival-ordered layout (−90 us at bs 2048), local-first / arrival-ordered rows, the layout on dedicated workgroups or in the scheduler workgroup, more quant workgroups (767 deadlocks), the plan on fewer workgroups, unroll, skipping tokens in pass 1, a pre-built doorbell,
stream, one copy for the rest, chunk-size sweeps,do_tile'spair_k/mfma_amajor=False/async_a_copy=False, the LLVM unclustered-high-RP flag. Numbers in the README.Caveats
sort_block_m.--ready-first, first run in a fresh FlyDSL cache); not reproduced.Test
🤖 Generated with Claude Code