Repository navigation
gfx1250 EPv2: EP4 topk-9 geometry; selfFirst follows MORI_EP_TOKOFF_EXT - #726
Merged
Merged
Conversation
…ELF_FIRST) On MTYPE_RW memory a rank's local RMW on its slot allocator word, which the peers' remote RMWs also hit, costs microseconds per call. With selfFirst (default on, gfx125x dispatch only) a sender's slot in a peer is what its remote RMW returns, with no wait on that peer. Each rank publishes, once all its blocks have reserved, how many slots it took in every peer, and places its own tokens above their sum in a pass of its own. A late rank therefore holds up no sender, so the prefill tail measured from the last rank to arrive stays at the original protocol's level. The per-call state is a few hundred bytes of symmetric memory per arena that EpDispatchPlan allocates on the arena's own communicator, found through the new mori.cco.communicator_of_window, so arenas built outside mori need no change. MORI_EP_SELF_FIRST=0 restores the original protocol, the better one where local memory is not mapped MTYPE_RW. Tests: binding tests for the state (one per arena, zeroed, freed with the last plan, refused without a communicator or with a descriptor that misses the local copy); test_op.py RAGGED and NOSELF; bench_ep.py ROUTE=noself and LATE_US.
…y block (MORI_EP_TOKCHUNK) The plain token map fills block b with its warpPerBlock consecutive work units, so a call with fewer units than warps leaves the high blocks idle: 512 tokens on the default 64 x 16 run on 32 blocks. When the units fit in gridDim.x * 8, block b now takes units [8b, 8b + 8) on its first 8 warps. Where the payload pass walks the token map, the 8 warps left without a token then take every metadata run, one peer each, and the token warps none: the tensorcnt wait that ends the metadata pass counts every TDM op of the wave, so a token warp that also sent metadata issued no payload until those stores completed. selfFirst only. Under the original protocol every block that holds tokens also RMWs the rank's own slot word, and there the chunked map measured slower at 512 tokens on EP8 (+2 to +8 us, 3/3 rounds). Geometries with at most 8 warps per block, such as EP4's tuned 64 x 8, and every selfFirst=false kernel compile to the same instructions as before; -DMORI_EP_TOKCHUNK=0 restores the previous code everywhere (all checked by disassembly). EP8 fp4 dispatch, two 4-GPU gfx1250 hosts over UALink, top-k 6, 96 experts per rank, graph-replay median against the selfFirst commit alone, 3 alternating rounds: 16 tokens 33.3 -> 27.2 us, 128: 35.8 -> 30.4, 256: 36.7 -> 30.0, 512: 39.9 -> 34.2, 4096 unchanged. EP4 at its tuned geometry: 16 and 512 tokens run the same instructions, 768 and 1536 are not slower, 4096 is within 0.7 us.
… TokOffExt only on one host selfFirst saves the RMW a rank makes on its own slot allocator word, which only costs anything when the word is in the cco window. With TokOffExt holding it -- the hip backend's default for a single-host gfx1250 EP -- that RMW is cheap and selfFirst is pure overhead: against origin/main, EP4 fp4 dispatch measured 2-3 us slower at 16 tokens and 3-15 us at 768-1536, and up to 29 us with the chunked map compiled in. MORI_EP_SELF_FIRST unset now resolves on exactly when the word is in the window; 1/0 still force either protocol. TokOffExt opens each peer's word by hipIpc handle, which no other host can open, so an EP spanning hosts could not build the op at its defaults (hipIpcOpenMemHandle(pe 4) returned hipError 17). It is now used only when every rank shares this host's kernel boot id; otherwise the word stays in the cco window, which is mapped across hosts, and selfFirst turns on with it. Measured on two 4-GPU gfx1250 hosts over UALink, fp4 dispatch, top-k 6, 96 experts per rank, both variables unset, graph-replay median, 3 alternating rounds: EP4 on one host matches origin/main within 0.7 us at 16-4096 tokens and 1.5 us at 16384; EP8 across the two hosts is 3.1-7.9 us faster than origin/main (run with MORI_EP_TOKOFF_EXT=0, the only way it builds there) at 16-512 tokens, 4096 unchanged.
…ce, bf16 pull loads every source at once fp4 push combine: a run sends four rows where the warp tile holds them (EP8, and EP4 at 64x8). At eight ranks and up with 8-warp blocks, the first four warps of a block load and quantize and the other four only store: a block owns one contiguous range of received rows, cut into 6-row chunks that pass through an LDS ring, and a store warp merges the rows of a chunk that return to one rank into one TDM store, with up to four chunks in flight. Loads and stores then retire on different waves, so no load waits behind a store. Below 64 rows a block the per-warp runs stay. Once every warp has a token, the reduce takes a whole token per item and prepares the next item while the rows of this one load. bf16 pull combine: the rows loaded together before folding cover every source a token can have (min(npes, topk), at least 4), and a chunk is rounded to the fold stride of one 16 B vector a lane where that fits, so no element falls to the scalar tail. EP8, two hosts x 4 gfx1250, topk 6, graph median, 3 rounds, checks pass, against the combine of 3145152: fp4 ct512 / 1536 / 4096 29.9 / 77.7 / 207.9 -> 25.1 / 56.1 / 119.6 us; bf16 53.4 / 119.4 / 291.3 -> 44.9 / 100.4 / 243.6 us. EP4: machine code unchanged except the 64x8 fp4 units (ct <= 256), A/B within 0.11 us.
At eight ranks and up, a 64x16 fp4 combine with e8m0 scales splits each block: eight warps send (four load and quantize, four store through the block's LDS ring) and the other eight reduce. A store warp flags each row at its destination with the call's epoch once the row's stores are done, and a reducing warp reduces a token as soon as every row of it is flagged or comes from a peer that has already arrived at the barrier, which also covers a peer that picked 64x8 and writes no flags. The barrier is arrived at after the send and waited on at the end of the kernel. Once its send is done, each sending warp takes tokens from the reducing warp it is paired with, through a per-warp LDS word cleared with atomicAnd, so the tail of the reduce runs on all sixteen warps. A block sends its rows grouped by origin token, eight groups of numTokens / 8 from a stable counting sort in LDS (a block with more than 1024 rows sends in receive order), so every rank finishes a group at about the same time and the receivers reduce it while the next one is in flight. A 16-warp block that does not overlap splits its send six load warps to ten store warps and reduces through a block-shared LDS ring wherever a warp tile cannot hold a token's rows. The 64x8 reduce at eight ranks and up cuts a small call's tokens into up to seven items spread over the blocks first, reads the row count once a warp and decodes with the packed fp4 converter. EP8 fp4 combine geometry: 64x8 up to 1536 tokens a rank, 64x16 up to 14336, 64x8 above. out_tok gets one uint64 flag a landing row at eight ranks and up. EP8, two hosts x 4 gfx1250, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, equal tokens a rank, graph median, 3 alternating rounds, checks pass, combine us against 084b78e (max 4096 tokens a rank; 16384 for 8192 and 16384): ct 16 / 32 / 64 / 128 / 256 / 512 / 1024: 15.1 / 15.4 / 15.6 / 16.4 / 18.4 / 24.0 / 40.2 -> 14.2 / 14.4 / 14.7 / 15.9 / 18.2 / 23.6 / 39.6 ct 2048 / 4096 / 8192 / 16384: 65.7 / 115.9 / 212.0 / 401.9 -> 60.5 / 100.7 / 181.2 / 395.5 Dispatch code is unchanged, and the EP4 units compile to the same instructions (compared by disassembly).
…ts GPU by LOCAL_RANK COMB_MODE now defaults to push (fp4, PUSH_QUANT), the combine the gfx1250 numbers are taken on. flydsl and parts without that combine need COMB_MODE=pull, which the CI bench step and the README examples now pass. The device is the local rank, so the bench runs across hosts (EP8 on two 4-GPU nodes), where the global rank is not a device index.
…rp, geometry by token count, fp4 own pass on the idle warps Past gridDim * warpPerBlock tokens every warp held two or more tokens and a metadata run, and one wave's TDM ops complete in issue order, so each warp's payload waited for its metadata stores to be acknowledged. There the block's top npes warps now send its metadata, a peer each, and no payload, and the other warps carry the block's tokens; past twice that many tokens those metadata warps also take the whole own pass. Past 22 warps, where an fp4 slab holds one token, metadata warp m works in the W / npes metadata slabs from m * W / npes, the payload warps' included. The split is compiled only where maxTokPerRank exceeds blockNum * warpPerBlock. Below that range an fp4 dispatch (by the dtype the JIT renders into the Cfg) gives the whole own pass to the warps whose metadata run is the own one, idle through the metadata pass, after their own payload, instead of a share to every warp after its payload. fp8 and bf16 move own rows one at a time, and on so few warps fp8 at 512 tokens came out 1.2 us slower, so they keep the even split. EP8 fp4 dispatch geometry by tokens a rank: 64x16 up to 1024, 64x24 up to 1536, 64x32 up to 2048, 64x16 up to 2944, 64x32 up to 4096, 64x16 above; the other dtypes stay on 64x16. MORI_EP_TOKCHUNK becomes the constant kEpTokChunk = 8. EP8, two hosts x 4 gfx1250, bench_ep.py (fp4 push combine), topk 6, 96 experts a rank, SWEEP 512,4096, graph, median of 3 alternating rounds, checks pass, dispatch us against 084b78e: fp4 512 / 4096: 32.86 / 99.88 -> 31.53 / 81.35 fp8 512 / 4096: 39.11 / 132.00 -> 39.26 / 123.08 bf16 512 / 4096: 45.47 / 196.05 -> 45.54 / 187.46 fp4 with a bf16 pull combine (another pair of hosts, 4096 tokens a rank allocated), ct 16 / 512 / 1024 / 1088 / 1536 / 2048 / 3072 / 3328 / 4096: 28.0 / 38.5 / 52.2 / 76.9 / 75.9 / 82.3 / 107.9 / 118.1 / 135.1 -> 25.9 / 37.0 / 50.5 / 54.7 / 63.0 / 75.8 / 97.3 / 102.0 / 122.1 Combine unchanged. The EP4 dispatch and combine units compile to the same instructions (compared by disassembly), and EP4 on one host stays within 0.5 us.
… by default The gfx125x dispatch runs selfFirst wherever the slot allocator word is in the cco window, with no environment override. MORI_EP_TOKOFF_EXT, which moves a single-host EP's word into a separate IPC allocation, is now opt-in (=1): by default a single-host EP keeps the word in the window and runs selfFirst like an EP spanning hosts, and with the opt-in the original protocol runs, as before. The README and the binding test follow.
kPlainMeta, which past gridDim * warpPerBlock tokens gives a block's metadata to its top npes warps, was measured at eight ranks. With selfFirst now the default for a single-host EP it also compiled into the EP4 64x16 dispatch, and there, against the same protocol without it (084b78e with MORI_EP_TOKOFF_EXT=0), fp8 came out 3.9 / 12.8 us slower at 1536 / 4096 tokens a rank and bf16 1.8 / 3.9 us faster. It now needs eight peers or more, so the EP8 units are unchanged and EP4 runs the selfFirst dispatch it had before the split. EP4 on one host, bench_ep.py (fp4 push combine), topk 6, 96 experts a rank, graph, median of 3 alternating rounds on each of two hosts, every point verified: against 084b78e with MORI_EP_TOKOFF_EXT=0 every dispatch cell (bf16 / fp8 / fp4 at 16 / 512 / 1536 / 4096 tokens) is within 0.4 us; fp8 at 4096 was +12.7 / +12.9 us over it and is now -0.03 / +0.12. EP8 across the two hosts against 084b78e, dispatch us: fp4 512 / 4096 32.59 / 100.01 -> 31.75 / 82.07, fp8 39.04 / 131.92 -> 39.34 / 123.33, bf16 45.61 / 196.45 -> 45.67 / 188.16, as before this change. Co-authored-by: Cursor <cursoragent@cursor.com>
…wave scale staging, own pass reads ahead of R The last block to count now tells every peer how many slots this rank takes in it right after the count pass, while warp 0's remote reservations are in flight; it used to wait for the reservations and then the assignment pass. Under selfFirst those per-peer numbers come from the count pass, so the reservations no longer add to destPeTokenCounter. When every warp holds at most one token, the token's scale row is read into registers at the start of the call, and each kept destination stages it with one contiguous whole-wave store instead of eight lanes storing scattered dwords. The own pass now issues every read that does not need R (the staged index, weight and scale rows, the reverse-map entry and the token's top-k row) before it waits for R, so after R it only stores; the scale moves 16 B a lane instead of 4 B, which had cost one round trip per 128 B of scale a lane. A scaled fp8 own pass moves two rows per TDM round trip, which the bf16-wide slab holds. EP4 on one host, bench_ep.py (fp4 push combine), topk 6, 96 experts a rank, graph, median of 3 alternating rounds on each of two hosts, every point verified, dispatch us against 3b02c53 at 16 / 512 / 1536 / 4096 tokens a rank: fp8: 27.95 / 32.50 / 50.24 / 97.55 -> 25.78 / 30.45 / 49.03 / 95.09 (second host 31.08 / 33.14 / 50.94 / 98.69 -> 28.01 / 31.79 / 50.36 / 96.50) fp4: 22.96 / 25.41 / 35.76 / 58.99 -> 20.41 / 23.45 / 34.00 / 59.18 (second host 25.83 / 26.17 / 36.98 / 60.03 -> 22.59 / 25.09 / 35.52 / 60.39) bf16: 27.54 / 37.85 / 66.39 / 132.91 -> 27.44 / 36.56 / 65.98 / 132.26 (second host 29.94 / 38.62 / 67.97 / 133.91 -> 29.65 / 37.65 / 67.35 / 133.07) EP8 across the two hosts, same setup, 16 / 512 / 1024 / 4096: fp4: 27.18 / 31.78 / 41.24 / 81.97 -> 23.42 / 28.63 / 38.43 / 81.30 fp8: 30.91 / 39.19 / 51.06 / 123.07 -> 29.14 / 37.17 / 48.75 / 123.46 bf16: 31.64 / 45.42 / 68.21 / 187.61 -> 31.54 / 44.99 / 65.42 / 187.63 EP4 fp4 at 4096 (+0.2 / +0.4 us) and EP8 fp8 at 4096 (+0.4 us) come out slower; combine is unchanged. test_op.py (bf16, scales, ragged, noself) passes on EP4.
…the count and assignment passes With more than one token a warp, the assignment pass staged each kept (token, destination) pair by reading the token's index, weight and scale rows from global memory and then writing them, the scale four bytes a lane in seven dependent rounds, so at 4096 tokens a rank on 64x16 a warp waited through about forty round trips. Up to four tokens a warp, those rows are now read once at the start of the call, at clamped addresses so that every read issues unconditionally, and the count and assignment passes and the staging take them from registers; each destination gets the scale row as one whole-wave store, its padding zeroed at the store. This replaces the prefetch that covered one token a warp. EP8 across two hosts x 4 gfx1250, bench_ep.py (fp4 push combine), topk 6, 96 experts a rank, graph, median of 3 alternating rounds, every point verified, dispatch us against 0a48447: fp8 2048 / 4096: 75.21 / 123.44 -> 71.35 / 113.11 fp4 4096: 81.35 -> 75.83 bf16 2048 / 4096: 106.69 / 187.82 -> 105.82 / 184.82 16 to 1024 tokens (one token a warp) and the combine are unchanged. EP4 on one host, same setup, on each of two hosts: fp8 1536 / 4096: 48.86 / 95.12 -> 47.45 / 87.83 (second host 50.21 / 96.40 -> 48.45 / 89.32) fp4 4096: 59.00 -> 55.62 (60.28 -> 57.01) bf16 4096: 132.26 -> 129.22 (133.46 -> 130.19) 16 and 512 tokens are unchanged. The combine after a 4096-token EP4 dispatch comes out about 0.4 us slower with its code unchanged. test_op.py (bf16, scales, ragged, noself) passes on EP4.
…token A token's combine rows come from the distinct ranks of its top-k experts, so there are at most min(worldSize, topk) of them, but the fp4 push reduce sized its row arrays and loops by worldSize and took the block-shared ring reduce whenever a warp's tile held fewer than worldSize rows. At 12 and 16 ranks that put the default 64x8 (10 rows a warp) on the ring reduce, and the combine came out about 2.4x slower than at 8 ranks. A token's rows are now min(worldSize, numExpertPerToken) wherever they are staged or counted, so 64x8 keeps the per-warp reduce at any rank count for topk up to 10, and the ws32 fp4 push combine compiles. C rack (4 gfx1250 a host, one UALink vPOD, cross-host load/store), bench_ep.py, bf16 dispatch + fp4 push combine, topk 6, 96 experts a rank, graph, median of 3 alternating rounds, every point verified, combine us against 3b02c53 at 16 / 128 / 512 tokens a rank: EP16 (4 hosts): 36.82 / 41.92 / 88.32 -> 15.35 / 20.21 / 42.31 EP12 (3 hosts): 35.59 / 46.86 / 95.01 -> 15.04 / 19.30 / 39.65 EP8 (2 hosts): 14.51 / 18.05 / 33.73 -> 14.53 / 17.99 / 33.67; on 64x16, the table's choice past 2048 tokens, 3072 / 4096: 135.19 / 174.55 -> 135.20 / 174.66 EP4 at 64x8 / 64x16, and EP8 at topk 8, compile to the same instructions as before. Every combine unit above compiles to the same instructions on this parent as on the measured 3b02c53 with this change. Co-authored-by: Cursor <cursoragent@cursor.com>
… before any rank can call Building the op zeroes this rank's arena (an asynchronous fill) and its selfFirst state, and nothing ordered that against a peer's first dispatch, which already writes into both: its selfFirst count lands in this rank's inbox at the count pass. A zero that ran after that write wiped it, and this rank's first dispatch then waited forever in EpSelfFirstPeerSent for a call number that never came back -- a silent hang of the first call, no error and no timeout, whenever one rank finished building the op later than a peer reached its first count pass. reset() zeroes the same state and had the same race. The constructor and reset() now synchronize the device and barrier on the communicator before returning, as test_op.py already did on the caller side. C rack (gfx1250, one UALink vPOD), bench_ep.py with one rank sleeping 3 s before zeroing its selfFirst state: without this change the first dispatch hung every time (EP12 once; EP8 twice, with the sleeping rank on either host), every rank stuck in the synchronize after it; with it, both EP8 cases pass. Unchanged runs pass on EP4 (one host) and EP8 (two hosts, graph). Kernels and the dispatch / combine paths are unchanged; the internode path and EP12 / EP16 were not rerun with this change. Co-authored-by: Cursor <cursoragent@cursor.com>
…4x16 through 16384 tokens The 64x16 fp4 combine at eight ranks reduces under the send because every block sends its rows grouped by origin token: every rank finishes a group at about the same time, and the receivers reduce it while the next one is in flight. The counting sort that groups them took at most 1024 rows a block, and a block with more sent in receive order, which leaves the receivers no complete token to reduce until the send is over; that is why the table handed the combine back to 64x8 past 14336 tokens a rank, where a block holds about 1000 rows. The sort now takes up to 2048. Up to 1024 rows it runs as before, with every row's group in registers; past that it reads the rows' sources 1024 at a time, every load of a batch in flight at once, and keeps each row's group as one byte at the head of the LDS ring, which no row reaches before the sort is done, for the scatter to read back. EP8 fp4 combine geometry: 64x8 up to 1536 tokens a rank, 64x16 up to 16384, 64x8 above. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us against 9b66e35: 16384 tokens a rank (16384 allocated): combine 394.80 -> 341.61 (-53.19 / -53.21 / -53.33); the dispatch that follows it, whose code is unchanged, 280.58 -> 281.45 (+1.18 / +0.84 / +0.63). 4096 (4096 allocated): combine 100.40 -> 100.49 (+0.10 / -0.13 / +0.01), measured with both arms' EP8 fp4 dispatch at 64x16 instead of the table's 64x32 at 4096; that dispatch kernel stopped with HSA_STATUS_ERROR_ILLEGAL_INSTRUCTION three times on c3-1 the same day, once on 9b66e35 itself. Every 64x8 and EP4 combine unit and every dispatch unit compiles to the same instructions as on 9b66e35; the 64x16 units keep 147 VGPRs and use no scratch. Co-authored-by: Cursor <cursoragent@cursor.com>
… rows as its tile holds Below 64 rows a block the fp4 push combine sends per warp, in runs: a run's rows are loaded and quantized into the warp's LDS tile, their stores go out together and are waited once, and every further run of the warp costs one more row load and one more wait for its stores. A run held 4 rows wherever the tile had room for them, while an 8-warp tile holds 6 ahead of the staged bf16 row (6 x 3808 + 14336 of 38080 bytes), so at EP8 512 tokens a rank, about 4.4 rows a warp, every warp with a fifth row ran a second time for it. A run now holds as many rows as the tile leaves room for, at most 10, so that a run's stores and the next row's load stay within the 11 TDM ops a wave keeps in flight; 16-warp tiles keep their 2. The run size no longer has to divide the receive capacity: the rows of a run are strided by the grid's warps, so nothing tiles the capacity by it. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, combine us against 9b66e35, whose 64x8 units compile to the same instructions as this parent's: EP8 512 tokens a rank (4096 allocated): 23.67 -> 23.31 (-0.36 / -0.28 / -0.40). EP4 on one host, 128 / 256 tokens a rank (256 allocated): 14.60 / 16.04 -> 14.55 / 16.03 (-0.11, +0.02, -0.03 / -0.06, +0.03, 0.00); the dispatch at 256, whose code is unchanged, read +0.42 / +0.24 / +0.18. Only the 64x8 combine units (and the 80x4 single-shot fallback) change; the 64x16 units and every dispatch unit compile to the same instructions as on the parent. Co-authored-by: Cursor <cursoragent@cursor.com>
… warps The 64x16 fp4 combine at eight ranks groups a block's rows by origin token before it sends them, and every wave of the block waited at the barrier behind that counting sort while warp 0 did it alone: at 16384 tokens a rank, about 1133 rows a block, warp 0's sort and the barrier took 16 us of a 341 us combine (MORI_COMB_SEG_PROBE phase probe, 3 rounds, both hosts). Every warp now takes one contiguous run of the block's 32-row chunks (at most 4 at the 2048-row limit), counts its rows a group and leaves the counts at the head of the LDS ring, which no row reaches before the sort is done; after one more barrier each warp scatters its rows from the offsets the warps before it leave. A row lands where one warp walking every chunk in turn would put it, so the send order, and everything after it, is unchanged. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us against 33732a0: 16384 tokens a rank (16384 allocated): combine 341.29 -> 329.39 (-12.01 / -11.90 / -11.64); the dispatch, whose code is unchanged, +0.64 / -0.32 / +0.10. 4096 (4096 allocated): combine 100.61 -> 97.90 (-2.61 / -2.71 / -2.70); the dispatch +0.14 / +0.15 / +0.10. Both arms ran the EP8 fp4 dispatch at 64x16 instead of the table's 64x32 at 4096, as for e7d2d65. 512 (4096 allocated): combine 23.32 -> 23.27 (-0.05 / -0.02 / -0.07), a 64x8 unit. Only the two EP8 64x16 combine units change (10140 -> 4489 instructions, 147 VGPRs, no scratch); every 64x8 and EP4 combine unit and every dispatch unit compiles to the same instructions as on 33732a0. Co-authored-by: Cursor <cursoragent@cursor.com>
…kens, two producers a SIMD The overlapped fp4 combine at eight ranks sends with half its waves -- producers that load bf16 rows and quantize them into an LDS ring, and four store waves that push the ring to the source ranks -- and reduces with the other half. At 16384 tokens a rank its four producers decide the kernel: each quantizes about 283 rows at about 1 us a row (MORI_COMB_SEG_PROBE phase probe). With two producers on every SIMD a row took only 14% longer, so a producer is bound by its own dependent chains and LDS reads rather than by the SIMD's issue rate, but in a 16-warp block that left too few waves to reduce. The overlapped path now also takes 24-warp blocks: eight producers, the same four store waves and twelve reducing waves, the ring two slots shorter (nine) to make room for the four more staging rows. The EP8 fp4 combine table runs 64x24 from 4097 to 16384 tokens a rank. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us against b77f654: 16384 tokens a rank (16384 allocated): combine 329.55 -> 314.35 (-15.20 / -15.31 / -14.99); the dispatch that follows it, whose code is unchanged, 282.25 -> 280.19 (-2.28 / -2.06 / -1.22). An earlier 3 rounds with the 64x24 unit built from the probe header without the probe, the same instructions as this header's, read -15.33 / -14.86 / -14.65. 4096 (4096 allocated), both arms' EP8 fp4 dispatch at 64x16 instead of the table's 64x32: a 64x24 combine tied with 64x16, 97.91 -> 97.85 (-0.06 / -0.10 / -0.03), so the table keeps 64x16 through 4096; 512 unchanged (+0.04 / 0.00 / +0.08). The new 64x24 unit has 147 VGPRs and no scratch. Every existing combine and dispatch unit, EP4 included, compiles to the same instructions as on b77f654. Co-authored-by: Cursor <cursoragent@cursor.com>
…past four tokens a warp The count and assignment passes of the intranode dispatch keep a warp's tokens' expert indices, weights and scale rows in registers only while the warp holds at most kPreK = 4 tokens. Past that every token was staged from global memory: each kept destination by a group of eight lanes that read its indices, its weights and, four bytes a lane in seven trips, its scale row, every read waiting on the store before it. At EP8 16384 tokens a rank on 64x16 a warp holds 16 tokens, and the assignment took ~67 us of the 277 us dispatch (MORI_DISP_SEG_PROBE phase probe) on ~19 dependent loads a token. Past kPreK tokens the same registers now hold a window of one token that slides: each pass reads its next token's row into slot 1 while it handles slot 0, and the indices, weights and the whole-wave scale row are staged from registers, as the kPreK path already did. Only units that can see more than kPreK tokens a warp (maxTokPerRank > 4 x blockNum x warpPerBlock) with 16 or more warps a block build it; the tables take 8-warp blocks only for small token counts, where a first version's extra code alone cost EP4 fp4 0.3-0.5 us at 128 / 512 tokens. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, topk 6, 96 experts a rank, fp4 push combine, graph, 3 alternating rounds, every point verified, us against 93e471b, the op built for 16384 tokens a rank: fp4 dispatch: 16384 279.66 -> 249.53 (-30.79 / -29.99 / -30.13); 1024 +0.13 / +0.22 / -0.04; 512 -0.36 / +0.35 / -0.12; 128 -0.11 / +0.58 / -0.08. fp8 dispatch: 16384 419.61 -> 390.05 (-29.56 / -29.65 / -28.88); 1024 +0.27 / -0.09 / +0.43; 512 +0.05 / +0.02 / +0.14; 128 +0.37 / +0.64 / +0.48. The combine after the 16384-token dispatch, its code unchanged: -1.03 / -1.42 / -1.37 (fp4). The indices, weights and scale rows every rank received were compared row by row with the senders' (bench_ep.py plus a metadata check, fp4, 1024 / 8192 / 12001 / 16384 tokens: no mismatch), and test_op.py at EP4 passes at 16384 tokens (128x16, eight tokens a warp), scales and ragged included. The four 16384-token EP8 dispatch units drop the old per-group fallback (64x16 fp4: 5583 instead of 6222 instructions), keep their VGPR counts and have no scratch. Every unit built for 4096 tokens or fewer, EP4 included, and every combine unit compile to the same instructions as on 93e471b. Co-authored-by: Cursor <cursoragent@cursor.com>
…otating four-token window Past kPreK = 4 tokens a warp the count and assignment passes now run their own loops instead of hooking into the shared ones: the warp's tokens go through the four prefetch slots in turn, each pass handling its k-th token from slot k % 4 and refilling the slot with token k + 4 after its last read, unrolled by four so every slot index is a constant and every load lands in its slot's own registers (a refill issued before the slot's last read was loaded into a temporary and copied back, and the copy waited for the load). The assignment keeps the first lane of each destination with eight ballots, as the count pass does, rather than __match_any_sync, which loops over a shuffle per distinct key, and stages each kept destination in one pass of the whole wave that reads the destination and slot with readlane instead of __shfl: the lanes below topk write the indices and weights, every lane its two dwords of the scale row. The loops shared with smaller token counts are back to their 93e471b text. The window's loads are unconditional, so a warp takes this path only with a scale buffer whenever the scale row is laid out for the wave. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, topk 6, 96 experts a rank, fp4 push combine, graph, 3 alternating rounds, every point verified, us against cd4c9f0, the op built for 16384 tokens a rank: fp4 dispatch: 16384 249.63 -> 232.61 (-17.02 / -16.99 / -17.16); 1024 -0.05 / +0.09 / +0.07; 512 -0.09 / +0.29 / -0.07; 128 +0.18 / +0.51 / +0.10. fp8 dispatch: 16384 390.06 -> 371.41 (-18.81 / -18.68 / -18.37); 1024 -0.54 / -0.59 / -0.50; 512 -0.05 / 0.00 / -0.14; 128 -0.57 / -0.36 / -0.47. The combine after the 16384-token dispatch, its code unchanged: fp4 -0.70 / -0.44 / +0.15, fp8 +0.45 / +0.76 / +1.11. The indices, weights and scale rows every rank received match the senders' row by row (bench_ep.py plus a metadata check, fp4, 1024 / 8192 / 12001 / 16384 tokens), and test_op.py at EP4 passes at 16384 tokens (128x16, eight tokens a warp), scales and ragged included. The four 16384-token EP8 dispatch units keep their VGPR counts and have no scratch. Every unit built for 4096 tokens or fewer, EP4 included, and every combine unit compile to the same instructions as on 93e471b. Co-authored-by: Cursor <cursoragent@cursor.com>
… 1024 tokens: 64x24 to 1536 (fp8 / fp4), 64x22 to 1408 (bf16) EP12 and EP16 had no entry in the HIP geometry table and ran 64x16 at every token count. Past 1024 tokens a 64x16 grid hands the extra tokens to the low blocks, two a warp, and those blocks set the time: on the C rack EP16 64x16 dispatch went from 157 to 238 us between 1024 and 1280 tokens for bf16. 64x24 keeps one token a warp up to 1536 tokens and 64x22 (the most warps a bf16 block holds) up to 1408; past that 64x24 reaches the metadata split at EP12 and two tokens a warp at EP16, and 64x22 two tokens a warp, all slower than 64x16, so the entries go back to 64x16 there. bf16 has its own key, so a dtype without one keeps 64x16. The combine half is the 64x8 default, listed so the dispatch schedule has one; 64x16 measured 0.7-6.9 us slower at EP12. C rack (gfx1250, one UALink vPOD), bench_ep.py, topk 6, 96 experts a rank, fp4 push combine, graph, the op built for 4096 tokens a rank, rank 0's median over the timed replays, us, 64x16 -> this table, 3 alternating rounds, every round the same sign, every point verified (36 runs); EP16 on C16 + C5 + C1 + C9 / EP12 on C5 + C1 + C9: bf16 dispatch: 1280 238.0 -> 189.0 / 211.8 -> 167.4; 1408 240.2 -> 205.4 / 214.9 -> 185.8. fp8 dispatch: 1280 138.9 -> 111.0 / 124.3 -> 102.5; 1408 140.8 -> 120.2 / 125.2 -> 108.0; 1536 142.5 -> 129.2 / 126.6 -> 114.2. fp4 dispatch: 1280 93.8 -> 78.9 / 85.6 -> 72.6; 1408 94.8 -> 81.7 / 85.9 -> 74.0; 1536 95.5 -> 85.2 / 86.1 -> 76.2. 16 / 128 / 512 / 1024 / 2048 / 4096 tokens: within 1.7 us, both signs; combine within 0.6 us. Measured on 93e471b. The twelve EP12 / EP16 dispatch units this table selects, built for 4096 tokens, compile to the same instructions on 9f391fe; built for 8192 or more they take its rotating-window path, which these numbers do not cover. EP4 / EP8 lookups are unchanged.
…ry block, ceil(units / blocks) a block instead of eight Under the chunked token map (selfFirst, more than kEpTokChunk = 8 warps a block, a call of at most gridDim.x * 8 units) block b took units [8b, 8b + 8) on its first eight warps, so 16 tokens ran eight to a block on two blocks and 128 on sixteen. A block's TDM stores leave through its own CU, and at these sizes the call waited on those few blocks: at 16 tokens the dispatch time followed the dtype (EP16 bf16 63.5 / fp8 48.3 / fp4 39.9 us) though the call moves a few hundred KB. Block b now takes t = ceil(units / gridDim.x) units, [b * t, b * t + t), on its first t warps; the warps at or past t take the ids left over as before, so the map stays a bijection and every pass walks the same token set, and the metadata and own-pass warps (at or past kEpTokChunk) are unchanged. At 512 units on 64 blocks t is 8, the old map. C rack (gfx1250, one UALink vPOD), bench_ep.py plus a row-by-row check of the indices, weights and scale rows every rank received, topk 6, 96 experts a rank, fp4 push combine, the HIP table's geometry, graph, the op built for 4096 tokens a rank, rank 0's median over the timed replays, dispatch us, before -> after, 3 alternating rounds, every round the same sign; EP8 on C16 + C9 / EP12 on C5 + C1 + C9 / EP16 on C16 + C5 + C1 + C9: fp4 16: 31.7 -> 25.1 / 36.9 -> 28.7 / 39.9 -> 29.3; 64: 33.1 -> 25.6 / 39.4 -> 30.9 / 42.2 -> 31.6; 128: 34.1 -> 27.1 / 40.4 -> 34.2 / 42.8 -> 36.6; 256: -2.8 / -0.6 / -1.7. fp8 16: 39.1 -> 25.5 / 44.5 -> 29.0 / 48.3 -> 30.6; 64: 40.9 -> 28.2 / 47.9 -> 32.1 / 51.2 -> 33.5; 128: 41.7 -> 32.1 / 49.2 -> 39.0 / 51.9 -> 41.8; 256: -5.6 / -4.6 / -4.0. bf16 16: 52.0 -> 26.0 / 59.0 -> 30.5 / 63.5 -> 31.8; 64: 53.9 -> 29.5 / 63.6 -> 33.9 / 68.2 -> 35.6; 128: 55.2 -> 35.4 / 65.0 -> 42.5 / 68.8 -> 46.7; 256: -11.9 / -10.0 / -9.8. 512 and 4096 tokens within 1.5 us, both signs; combine within 0.8 us. Every point verified: dispatch bytes and the fp4 combine output, and the indices, weights and scale rows row by row (EP8 fp8, EP12 and EP16 all three dtypes, EP16 also at the ragged counts 1 / 7 / 100 / 300 / 511: 7.0 M rows, no mismatch). The measured build was this change behind a compile switch; with the switch on it compiles to the same instructions as this commit for 13 dispatch units (EP4 / EP8 / EP12 / EP16, 64x8 to 64x24, built for 4096 and 16384 tokens), VGPR / SGPR counts unchanged and no scratch. Without selfFirst (EP4's default) or with at most eight warps a block the map stays the plain one at compile time.
…flight on 24-warp blocks On the overlapped fp4 combine's 24-warp blocks (EP8, 4097 to 16384 tokens a rank) the stores bound the push: at 16384 tokens the eight producers waited 94 us for a free ring slot (MORI_COMB_SEG_PROBE phase probe). With the producers' row loads, the overlapped reduce and the row flags all switched off the push still ran at about 1.05 TB/s, against 1.49 TB/s for the dispatch's payload over the same links, while each store wave kept up to four 6-row chunks in flight (about 370 KB a CU; the dispatch keeps about 125 KB). One chunk in flight a store wave pushed faster than four, both in that stripped build (stores done 263.0 -> 256.4 us, two chunks 254.0) and in the full kernel (301.5 -> 294.8 us, one round each), so the chunks behind the first were only queueing. A store wave on a 24-warp block now retires its chunk before it sends the next. 16-warp blocks, whose four producers bound the push, keep four: one in flight tied there. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us against 9f391fe: 16384 tokens a rank (op built for 16384): combine 313.17 -> 309.03 (-3.72 / -4.20 / -4.60); the dispatch, its code unchanged, +0.10 / -0.12 / +0.22. 8192 (op built for 8192): combine 170.98 -> 167.30 (-3.11 / -3.82 / -3.68); dispatch +0.06 / +0.07 / -0.05. Both were measured on 9f391fe with one chunk in flight on 16-warp blocks too; in the same run its 2048 and 4096 points (64x16 combine, both arms' EP8 fp4 dispatch at 64x16 instead of the table's 64x32) read -0.01 / -0.12 / -0.22 and +0.01 / -0.09 / +0.03. On this commit's parent c036041 the change builds the 64x24 combine unit to the same instructions as the measured one (147 VGPRs, no scratch) and every other unit, the 64x16 combine units, every dispatch unit and EP4 included, to the same instructions as c036041; the combine units of c036041 and 9f391fe are the same. Co-authored-by: Cursor <cursoragent@cursor.com>
…nk, not only past 4096 93e471b kept 64x16 through 4096 tokens because 64x24 only tied it there (97.91 against 97.85 us). Since 48916b6 the 24-warp blocks' store waves keep one chunk in flight, and at 4096 tokens 64x24 is now ahead. The EP8 fp4 combine table runs 64x24 from 1537 to 16384 tokens a rank; 1536 and below keep 64x8. Table only, no kernel source changed. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us, both arms on 48916b6's kernels, the op built for 4096 tokens a rank, both arms' EP8 fp4 dispatch at 64x16 instead of the table's 64x32 at 4096: 4096: combine 97.76 -> 96.00 (-1.76 / -1.76 / -1.71); the dispatch after it, its code unchanged, -0.55 / -0.43 / -0.43. 2048: combine 61.29 -> 60.98 (-0.72 / -0.31 / -0.27); dispatch -0.68 / -0.60 / -0.96. Co-authored-by: Cursor <cursoragent@cursor.com>
…roducers on 24-warp blocks On the overlapped fp4 combine's 24-warp blocks each SIMD runs one store wave next to two producers that quantize, and the producers wait on the ring slots the store wave frees. The store wave now sends at the highest wave priority (s_setprio 3) and drops back to 0 before it helps reduce. In the phase probe (16384 tokens, one round) the store waves' own execution fell 59.1 -> 55.3 us and their stores finished 1.1 us earlier. Only the 64x24 combine unit changes (the two s_setprio); every other unit compiles to the same instructions as c239dad. C rack c3 (two hosts x 4 gfx1250, one UALink vPOD), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us against c239dad, the op built for 16384 tokens a rank, both arms' EP8 fp4 dispatch at 64x16 instead of the table's 64x32 at 4096: 16384: combine 309.66 -> 308.33 (-1.27 / -1.00 / -1.45); the dispatch, its code unchanged, +0.20 / +0.09 / +0.30. 4096: combine 96.30 -> 95.83 (-0.23 / -0.63 / -0.49); dispatch -0.02 / +0.17 / -0.04. Co-authored-by: Cursor <cursoragent@cursor.com>
EP sizes that are not a power of two are not supported for now, and their entries only add to what has to be kept in step. EP12 goes back to the single-shot default (64x16 dispatch, 64x8 combine), as before e2b87c1; the EP16 entries stay. The EP16 fp4 push combine entry was the 64x8 default as well and goes too: lookup falls back to _COMBINE_TABLE, whose EP16 64x8 line the dispatch schedule needs anyway. Checked by loading this file and its parent with the device stubbed to gfx1250 and looking up every world size (4 / 8 / 12 / 16 / 32), dispatch dtype (bf16 / fp8 / fp4) and combine quant (none / fp4_blockwise): only EP12 changes, to no schedule; every other lookup is identical.
The EP16 entries chose 64x24 (fp8 / fp4) up to 1536 and 64x22 (bf16) up to 1408 tokens a rank and 64x16 everywhere else, so they only acted on 1025-1536 tokens a rank: too narrow a range to be worth the rows. EP16 goes back to the single-shot default (64x16 dispatch, 64x8 combine), and with the EP12 entries gone in 111a424 the table is as before e2b87c1 apart from the EP8 fp4 combine changes made since. Checked by loading this file and its parent with the device stubbed to gfx1250 and looking up every world size (4 / 8 / 12 / 16 / 32), dispatch dtype (bf16 / fp8 / fp4) and combine quant (none / fp4_blockwise): only EP16 changes, to no schedule; every other lookup is identical.
…24 from 2048 tokens a rank The overlapped fp4 combine -- producers quantize received rows into an LDS ring, store waves push them to the source ranks, and the other half of the block reduces each token as its rows land -- was compiled in only from eight ranks up, so EP4 ran the per-warp send and a separate reduce at every size: at 16384 tokens a rank it took 362 us for 54040 rows, where EP8 takes 308 us for 72461. EpFp4Ovl now also takes 24-warp blocks at four ranks (EP4's 8- and 16-warp units keep the per-warp send), the split send's chunk size follows the overlapped path, and the host sizes the row flags in out_tok from four ranks up (they were allocated only from eight). The EP4 fp4 combine table runs 64x24 from 2048 to 16384 tokens a rank; 257 to 2047 keep 64x16, and so does everything above 16384, which was not measured. C rack c3-2 (4 gfx1250, EP4 on one host), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us, 64x16 against 64x24: 16384: combine 361.86 -> 252.43 (-109.72 / -109.38 / -108.94); 8192: 183.35 -> 137.42 (-46.23 / -30.24 / -53.57); 4096: 97.10 -> 80.33 (-17.07 / -16.77 / -23.57); 2048: 53.96 -> 52.41 (-1.89 / -5.45 / -1.26). Below 2048 64x24 is behind: 1792 +0.48 / +0.82 / +0.91, 1536 +3.94 / +4.11 / +3.92, 1280 +4.77 / +4.58 / +4.22, 1024 +8.09 / +8.14 / +8.15, 512 +11.10 / +11.11 / +10.98. The dispatch, its code unchanged, moved by less than 0.5 at every point but single slow replays. Below 2048, this commit against d8e5a2f: combine 128 / 512 / 1024 / 1536 within 0.2 every round; 2047 +0.38 / +0.39 / +0.35 in one run and -0.53 / +0.27 / -0.11 in the next (an arm with only the branch's hip_backend.py, no EP4 flags, read -0.41 / +0.23 / -0.49 there). The combine output was also checked at 256, 2047, 2048, 2049, 5000, 12001 and 16384 tokens a rank. Every existing unit (37: all of EP8, EP4's 64x8 and 64x16 combine, every dispatch) compiles to the same instructions as on d8e5a2f. Co-authored-by: Cursor <cursoragent@cursor.com>
…ght send metadata on small calls Under the chunked token map a 64 x 16 block sends its metadata from warps 8-15 only, so at EP16 eight warps carry the fifteen remote runs and seven of them send two. A wave issues its remote TDM stores one after another and one run takes a few us to issue, so those warps set the length of the metadata pass, the largest part of a small dispatch (EP16 bf16 at 16 tokens, probe build: 12.5 of the 27.4 us in the kernel). The same map hands tokens only to the first t = ceil(units / gridDim.x) warps of a block (t = 1 at 16 tokens, 2 at 128), and warps t..7 sat idle. Where a block has fewer warps past kEpTokChunk than remote peers but at least one warp a peer (kMetaLo: EP16 at 16 to 22 warps a block), the metadata runs now start at warp t, and the peers are numbered from the rank after this one so that the skipped own run comes last: at 16 tokens every remote peer has a warp of its own. EP32 at 64 x 16 would keep two runs on a wave and is left as it was until it is measured. C rack (gfx1250, one UALink vPOD), EP16 on C13 + C5 + C1 + C9, bench_ep_pr.py plus a row-by-row check of the indices, weights and scale rows every rank received, topk 6, 96 experts a rank, the default 64 x 16 geometry, graph, the op built for 4096 tokens a rank, rank 0's median over the timed replays, dispatch us, before -> after, 3 alternating rounds: bf16 16: 31.7 -> 29.6; 32: 32.4 -> 30.6; 64: 35.5 -> 33.7; 128: 45.8 -> 43.9 (every round -1.4 to -2.5). fp8 16: 30.7 -> 27.0; 32: 31.1 -> 27.9; 64: 33.9 -> 30.6; 128: 42.1 -> 40.6; 256: 52.0 -> 51.5 (every round -0.4 to -3.9). fp4 16: 29.3 -> 26.3; 32: 29.9 -> 26.7; 64: 31.2 -> 28.4; 128: 36.2 -> 35.1; 256: 44.1 -> 43.7 (every round -0.3 to -3.4). bf16 at 256 and every dtype from 512 to 4096 tokens within 1.8 us, both signs (fp4 at 512 and 2048 +0.2 to +1.1 in all three rounds); combine within 0.6 us. Every point passed the dispatch byte check and the row check with no mismatch. The measured build compiles to the same instructions as this commit for EP16 64 x 16 in all three dtypes, built for 4096 and 512 tokens a rank. EP4 at 64 x 8 and 64 x 16, EP8 at 64 x 16, 64 x 24 and 64 x 32, EP16 at 64 x 24, EP32 at 64 x 16 in all three dtypes and the EP16 fp4 combine compile to the same instructions as the parent.
…he next chunk's rows ahead The fp4 combine's send (EpFp4CombineSend) took warpId as threadIdx.x / WS, which lives in a VGPR, so every loop the send runs over it -- chunk counters, in-flight counts, the EpWaitTensorAtMost switch -- was compiled as exec-masked branches. warpId now comes through readfirstlane. On the overlapped 64x24 path, where a store wave keeps one chunk in flight, the store wave also reads the next chunk's rows (sort table) and their ranks (receive map) before it waits for the chunk, and keeps the in-flight chunk's rows and ranks in registers for its retire, instead of reading back after the wait what the producer left in LDS. c3-1 + c3-2 (EP8: 8 gfx1250 on two hosts in one UALink team; EP4: c3-2 alone), bench_ep.py, fp4 dispatch and fp4 push combine, topk 6, 96 experts a rank, graph, 3 alternating rounds, every point verified, us, 37723c9 -> 37723c9 with this change; per round: dispatch + combine, which the bench times back to back in one graph. On d05ab8a, its parent here, every EP8 / EP4 unit measured compiles to the same instructions as in the measured build. EP8 16384: combine 308.30 -> 302.94 (-3.8 / -4.6 / -4.8); 8192: 166.65 -> 163.29 (-3.0 / -2.7 / -3.4); 4096: 95.71 -> 93.22 (-2.2 / -2.6 / -2.3); 2048: 59.04 -> 57.42 (-1.6 / -1.8 / -1.6); 512: 23.27 -> 23.13 (-0.1 / 0.0 / -0.3). EP8 2048 and 4096 were measured with the table's 64x32 dispatch rows replaced by 64x16 on both arms. EP4 16384: combine 252.02 -> 251.00 (-1.4 / -0.8 / -1.2); 4096: 80.25 -> 79.54 (-0.7 / -1.0 / -0.7); 2048: 52.34 -> 51.54 (-1.2 / -0.9 / -1.3); 512: 19.29 -> 19.23 (0.0); 128: 15.26 -> 15.24 (0.0). Every dispatch unit compiles to the same instructions as on 37723c9. VGPRs: EP8 64x24 combine 147 -> 144, 64x8 148 -> 150; EP4 64x24 147 -> 138, 64x16 146 -> 143, 64x8 145 -> 142; no scratch anywhere. Co-authored-by: Cursor <cursoragent@cursor.com>
…ocks past 2048 tokens a rank Past 2048 tokens a rank the table no longer holds the grid at 64 blocks for a co-resident GEMM: the EP8 dispatch (every dtype) takes 128x16, the EP8 and EP4 fp4 push combine 128x24, and the EP4 bf16 / fp8 dispatch moves its 128x16 edge from 4096 down to 2048, where the EP4 fp4 dispatch already was. Up to 2048 tokens nothing changes; the EP8 fp4 dispatch's 64x32 row at 2945-4096 is gone. C rack C12 + C11 (EP8: 8 gfx1250 on two hosts in one UALink vPOD; EP4: C12 alone), bench_ep.py with the metadata check added (every received index / weight / scale row compared with the sender's), topk 6, 96 experts a rank, fp4 push combine, graph, 3 alternating rounds, us, the combine-send commit before this one -> this commit (both arms at 4096 run the EP8 fp4 dispatch as 64x16 rather than the old 64x32 row). Every point verified: payload bytes, fp4 combine output, 0 metadata mismatches in 18.3 M EP8 and 6.8 M EP4 rows; no fabric error on either node. Each per-round difference has the same sign in all three rounds. EP8 dispatch at 4096 / 8192 / 16384: fp4 123.0 -> 82.5, 219.5 -> 135.1, 415.4 -> 240.3; fp8 181.2 -> 120.0, 333.4 -> 209.5, 637.7 -> 387.4; bf16 309.6 -> 198.8, 591.0 -> 363.1, 1151.0 -> 692.0. EP8 combine (every dispatch dtype within 2 us of these): 164.5 -> 97.9, 315.4 -> 178.4, 609.5 -> 340.2. EP4 combine (fp4; fp8 / bf16 within 2 us): 127.2 -> 76.0, 243.1 -> 139.2, 474.5 -> 262.9. EP4 dispatch at 4096: fp8 129.1 -> 86.1, bf16 212.2 -> 134.3; every other EP4 dispatch point keeps its geometry and moved by -1.7 to +1.1. Co-authored-by: Cursor <cursoragent@cursor.com>
…or op, and reads the cco window as constant memory Since d05ab8a the EP16 fp4 dispatch (64 x 16) stopped with HSA_STATUS_ERROR_ILLEGAL_INSTRUCTION in about one bench run in three at 2048 and 4096 tokens a rank (9 of 29 runs). The ROCm debug agent put every faulting wave (16 waves in 9 runs, on three hosts) at the head of the metadata loop over a run's chunks, on a legal s_mov_b32. The 7 dumps read back were identical register for register and show the wave had finished the first chunk, issued the chunk's last tensor_store_from_lds, and reached the head without running the loop latch: the back edge right after that store behaved as if it had landed on the s_branch in front of the latch. Shifting the whole kernel by 4 and by 8 bytes moved the fault with the code, so it is neither memory corrupted at a fixed address, nor the head's 64-byte alignment, nor the store straddling a 64-byte line. Why the parent's code, with the same store and back edge in this loop, never faulted in 10 such runs is not known. TdmIssueLoad / TdmIssueStore now follow every tensor op with an s_nop, so the instruction after it is never a branch; the same reproducer then ran 60 times without a fault (36 runs with the s_nop alone on d05ab8a, 24 with the dispatch of this commit). The asm s_nop alone cost the payload loop: the compiler takes an asm statement as touching memory, so the dispatch reloaded the cco window's winBase / stride4G before every TDM store (bf16 at 4096 tokens +34 us). The dispatch now reads the window through an address-space-4 pointer (EpWinC): the window does not change while a kernel runs, the loads become scalar and are reused across the asm, and dispatch VGPRs drop (EP16 fp4 64 x 16 122 -> 115, bf16 96 -> 81, EP8 fp4 64 x 32 128 -> 116). The combine keeps the generic reads, which measured 0.1-0.3 us faster there at small ct. C rack (gfx1250, one UALink vPOD), EP16 on C14 + C5 + C1 + C9, EP8 on C14 + C5, EP4 on C9, bench_ep_pr.py plus the row-by-row check of the indices, weights and scale rows every rank received, topk 6, 96 experts a rank, graph, rank 0's median over the timed replays, us, before -> after, 3 alternating rounds. On d05ab8a: EP16 dispatch bf16 16: 30.22 -> 29.88; 256: 54.23 -> 53.33; 4096: 373.71 -> 373.42. EP16 dispatch fp8 16: 27.51 -> 27.25; 256: 46.73 -> 45.74; 4096: 220.79 -> 220.64. EP16 dispatch fp4 16: 26.19 -> 25.69; 256: 43.09 -> 41.64; 1024: 61.56 -> 61.51 (the parent faults from 2048; at 2048 against 37723c9: 93.69 -> 93.82). EP8 bf16 dispatch 256: 42.56 -> 41.93; 1024: 94.96 -> 93.91. EP4 fp4 dispatch, 16 to 4096 tokens: -0.10 to -0.45 in every round. On 49ca9d7, the parent here, fp4: EP8 dispatch 4096: 82.27 -> 81.55; 8192: 135.44 -> 134.72. EP8 combine 4096: 97.69 -> 97.86; 8192: 178.62 -> 178.76. EP4 dispatch 8192: 99.75 -> 99.27. EP4 combine 4096: 76.02 -> 76.17; 8192: 138.38 -> 138.95. EP16 dispatch 256: 42.81 -> 42.07; EP16 combine within 0.25 us. Every other combine point within 0.4 us. Every point passed the dispatch byte check and the row check with no mismatch. The dispatch of this commit compiles to the same instructions as the measured builds in every unit checked: EP16 bf16 / fp8 / fp4 64 x 16, EP16 fp4 64 x 24, EP8 fp4 64 x 16 and 128 x 16, EP8 bf16 128 x 16, EP4 fp4 128 x 16 and EP32 bf16 64 x 16 (the measured builds also match each other in EP16 fp4 64 x 32, EP8 fp4 64 x 32, EP8 bf16 64 x 24, EP4 fp4 64 x 16 and EP4 bf16 128 x 16). No unit checked spills.
No code change: the units checked (EP16 bf16 / fp8 / fp4 64 x 16, EP8 fp4 128 x 16, EP4 fp4 64 x 16 and EP32 bf16 64 x 16 dispatch; EP16 fp4 / bf16 64 x 8, EP8 fp4 128 x 24 and EP4 fp4 64 x 24 combine) compile to the same instructions as 87ccc60.
No code change: the units checked (EP16 bf16 / fp8 / fp4 64 x 16, EP8 fp4 128 x 16, EP8 bf16 64 x 16, EP4 fp4 64 x 16 and EP32 bf16 64 x 16 dispatch; EP16 fp4 64 x 8 and EP8 fp4 128 x 24 combine) compile to the same instructions as 78b252b.
…XT is the only one The gfx125x dispatch runs selfFirst whenever the slot allocator word lives in the cco window, which is the default. The request no longer carries a selfFirst flag: EpRequest.selfFirst becomes EpRequest.tokOffExt (the word lives outside the window), and MakeEpCfg derives Cfg.selfFirst from it. The Python self_first kwarg and its default (ep_plans.self_first_enabled) are gone; the hip backend passes tok_off_ext=1 only when MORI_EP_TOKOFF_EXT=1 (default off) and every rank shares one host, and there the original protocol runs as before. EpCfg and the kernel source are unchanged, and on the Python path Cfg.selfFirst resolves to the same value as before, so the rendered source is the same. Tested on one gfx1250 host (4 GPUs), hip backend: - test_jit_binding.py: 23 passed, 0 skipped. - bench_ep.py fp4 dispatch, topk 6, 96 experts a rank, 512 tokens a rank, graph, CHECK=1: passes at the defaults and with MORI_EP_TOKOFF_EXT=1. The defaults compile dispatch kernels with selfFirst=true, MORI_EP_TOKOFF_EXT=1 compiles separate ones without it. - Same point, CHECK=0, ITERS=50, WARMUP=10, 3 alternating rounds, 7114be0 -> this commit, in us: dispatch 23.51 -> 23.40 (per round -0.29 / +0.05 / -0.11), combine 18.78 -> 18.77 (0.00 / 0.00 / -0.03).
…p8-selffirst-default # Conflicts: # include/mori/ops/dispatch_combine_v2/ep_spec.hpp # python/mori/ops/dispatch_combine_v2/README.md # python/mori/ops/dispatch_combine_v2/ep_plans.py # python/mori/ops/dispatch_combine_v2/hip_backend.py # src/ops/dispatch_combine_v2/ep_spec.cpp # tests/python/ops/dispatch_combine_v2/test_jit_binding.py
…s geometry topk 9 had no entry, so it ran the single-shot 64x16 dispatch and 64x8 fp4 combine, and 64x8 never takes the overlapped fp4 combine path at EP4. Run-internal median (us), graph, fp4 dispatch, 256 experts, 16384 tokens a rank allocated, 3 alternating rounds on 4x gfx1250, main 28d4973 build: ct dispatch fp4 combine 256 36.3 -> 30.4 17.2 -> 17.3 512 35.4 -> 31.4 22.3 -> 20.4 1024 39.2 -> 39.0 36.5 -> 32.3 1536 55.2 -> 42.5 52.2 -> 44.3 2048 57.5 -> 46.0 69.1 -> 53.8 4096 73.9 -> 70.7 125.1 -> 61.5 8192 106.2 -> 100.6 236.8 -> 97.1 16384 192.5 -> 164.0 457.4 -> 170.6 bf16 / fp8 dispatch are unmeasured and keep the default.
This was referenced Oct 8, 2026
Closed
zhangfei829
added a commit
that referenced
this pull request
Oct 11, 2026
Conflicts were all #726's selfFirst-follows-TOKOFF_EXT, which this branch already replaced with aggReserve on the same tokOffExt field; kept the branch side. #725's EP4 topk 8/9 dispatch rows are taken, and its topk 9 row is folded into the one (4, 7168, 9) key so #726's fp4 row no longer shadows it. #735 is taken: the gfx1250 dispatch still stores totalRecvTokenNum.
Merged
zhangfei829
added a commit
that referenced
this pull request
Oct 11, 2026
gfx1250 EPv2: one-RMW slot reservation, EP16 fp4 dispatch and push combine, fp8 push combine
## Summary
- Dispatch slot reservation: a rank's blocks add up their token counts and the last block reserves once per destination (one RMW on the slot word instead of one per block). Replaces selfFirst; `MORI_EP_TOKOFF_EXT=1` keeps the per-block protocol for a single-host EP.
- EP16 fp4 dispatch: 128x16 sends each destination's rows from a pinned set of warps.
- fp4 push combine: the landing offset rides the 64-bit dispatch signal, and the landing area is M x min(world, topk) rows instead of world x world x M, so EP16 push fits at 16384 tokens a rank.
- fp8 push combine (`fp8_blockwise`).
- EP16 64x16 bf16/fp8: a metadata warp holding exactly two runs loads both, then stores both. Built into those instances and chosen at run time, with no switch.
- Tuning tables: EP16 dispatch and fp4 push combine, EP4 fp8 push combine. Removed the `MORI_V2_EP_{DISP,COMB}_{BLOCKS,WARPS}` overrides.
- Merged main: #725 and #735 taken; #726's selfFirst change superseded. The duplicate `(4, 7168, 9)` dispatch key on main is folded into one, so #725's topk 9 row takes effect.
## Test
CHECK=1 passes on EP4 (topk 6 and 9), EP8 (2 hosts) and EP16 (4 hosts) at 1 / 64 / 512 / 4096 / 8192 / 8193 / 16384 tokens a rank, for fp4 push, fp4 pull, bf16 and fp8 push.
EP16 (4 hosts x 4 gfx1250), fp4 dispatch, topk 6, 96 experts/rank, HIP graph, ITERS=50, 3 alternating rounds, medians in us. The baseline has main's kernel and its untuned 64x16 geometry; main cannot allocate the fp4 push landing area at EP16 with 16384 tokens, so its combine is pull. The fp4 machine code is identical to this branch's head.
| tokens/rank | dispatch main → branch | combine main (pull) → branch (fp4 push) |
|---:|---:|---:|
| 64 | 27.39 → 26.56 | 17.95 → 16.02 |
| 512 | 38.83 → 31.66 | 50.57 → 28.33 |
| 4096 | 113.03 → 92.57 | 275.81 → 87.99 |
| 16384 | 378.46 → 292.98 | 1032.22 → 302.42 |
EP4 is about even with main: topk 9 fp4 push, same node, 3 alternating rounds, dispatch +0.1..+1.2 us, combine +0.2..+0.7 us (+2.5 at 1536).
jhchouuu
added a commit
to jhchouuu/mori
that referenced
this pull request
Oct 12, 2026
Conflicts in hip_tuning_configs.py, resolved to this branch's tables: - topk-9 fp4 dispatch: this branch's sweep (main kept ROCm#726's row under "fp4_disp_bf16_comb"; main also un-shadowed the ROCm#725 row, as this branch did). - topk-6 fp4 dispatch: same buckets, keyed "fp4" as on this branch. - the combine-table comment: main's (EP8 and EP16). Main's new EP16 dispatch row is re-keyed from "fp4_disp_bf16_comb" to "fp4", since lookup() now reads that spelling as "fp4"; under the old key the row would never match. Checked against main over every table key, dispatch dtype and quant type (112 lookups): every key but topk 9 returns what main returns for "bf16", "fp8" and "fp4_disp_bf16_comb". Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
jhchouuu
added a commit
to jhchouuu/aiter
that referenced
this pull request
Oct 27, 2026
Four fixes to the mori dispatch/combine path (stage1_fused/stage2_fused off), each from driving mori's plans without its op layer: - dispatch: since ROCm/mori#726 the plan takes tok_off_ext and derives selfFirst itself; self_first_enabled() is gone, so importing it fails. - tuning lookup: key it on the dispatch wire the way mori's op layer does ("fp4_disp_bf16_comb" for fp4). Under the hard-coded "bf16" key topk 9 matched no entry and both dispatch and the fp4 combine fell back to the single-shot default. - comb_inp: add the overlapped fp4 combine's u64 flag per landing row (four ranks up). Its 24-warp geometries write them past the region, i.e. past the arena: a page fault at 16384 tokens a rank. - close(): synchronize first, so a launch still in flight does not read a window that has been freed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
JiaoliangYu
pushed a commit
to ROCm/aiter
that referenced
this pull request
Oct 30, 2026
…tches, mori combine path, current mori EPv2 (#6430) * MegaMoE: stage1_fused/stage2_fused switches, mori combine path * bench_mega_moe: fix kernel table on multi-stream graphs * MegaMoE: follow mori tokoff-ext and selfFirst defaults * bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them * MegaMoE: follow mori#726 and size mori's fp4 combine flags Four fixes to the mori dispatch/combine path (stage1_fused/stage2_fused off), each from driving mori's plans without its op layer: - dispatch: since ROCm/mori#726 the plan takes tok_off_ext and derives selfFirst itself; self_first_enabled() is gone, so importing it fails. - tuning lookup: key it on the dispatch wire the way mori's op layer does ("fp4_disp_bf16_comb" for fp4). Under the hard-coded "bf16" key topk 9 matched no entry and both dispatch and the fp4 combine fell back to the single-shot default. - comb_inp: add the overlapped fp4 combine's u64 flag per landing row (four ranks up). Its 24-warp geometries write them past the region, i.e. past the arena: a page fault at 16384 tokens a rank. - close(): synchronize first, so a launch still in flight does not read a window that has been freed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: zero-copy hand-off from fused_moe to mori's combine With stage2_fused off, fused_moe's rows went into a fresh tensor and then into mori's staging rows (comb_inp): the bf16 combine copied them in its own kernel, the fp4 one through a host-issued copy whenever the recv bound was below max_recv -- every decode step under ATOM. Have fused_moe write them there directly. fused_moe(output=) already existed but the gfx1250 grouped path ignored it and copied at the end; thread it down to the gather-reduce. mori's bf16 combine then sees its input at the staging base and skips the copy, and the fp4 send reads the staging rows in place. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: follow mori#745/#746 on the mori dispatch/combine path mori's gfx1250 dispatch now reserves its slots once per destination (the aggregated reservation, ROCm/mori#745) and hands the push combine its landing offsets through the dispatch. MegaMoE drives the plans directly, so it has to size and pass what mori's op layer does: - dest_pe_token_counter holds 2 * world_size ints: send counts, then per destination the reservation base and the push combine's landing offset. It held world_size; the dispatch wrote the second half into allocator padding. - recv_num signals are 8 bytes a source. The region held 4 bytes a source; the arena's 256-byte alignment hid the overrun. - The combine reads its landing offsets from dest_pe_token_counter, which MegaMoE did not pass, so the mxfp4 combine (stage2_fused=False) got a null pointer there. - MORI_EP_TOKOFF_EXT=1 turns the aggregated reservation off, so a mori push combine is refused with it, as mori's op layer refuses it. - The push landing rows come from mori's push_land_rows(): one per (token, destination rank), max_tokens * min(world_size, topk) rather than world_size * max_recv. - The mori dispatch launcher drops total_recv.zero_(): gfx1250's dispatch overwrites the count (since ROCm/mori#599), and mori's op layer skips the same zero there. Against a mori from before #745 the landing rows keep the old count and the two larger buffers have unused tails. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging buffer (ROCm/mori#703) mori#703 moves the gfx1250 dispatch's staging arrays out of each kernel module's .bss into one buffer the caller allocates and binds as staging_base; an unbound plan is refused on the host. MegaMoE drives EpDispatchPlan directly, so it allocates ep_staging_bytes() once and binds it to every dispatch plan (they never run at the same time), and drops it after the plans on close(). A mori without ep_staging_bytes keeps the .bss arrays and is left as is. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: keep accepting dispatch_backend stage1_fused now picks the dispatch, and the dispatch_backend argument and config field went away with that. Callers still pass it (ATOM main passes dispatch_backend="mori" on a quantizing wire; ROCm/ATOM#2384 passes "flydsl" or "mori" alongside stage1_fused) and log mega._config.dispatch_backend, so both would fail on this aiter. Accept the argument again and check it against stage1_fused: "mori" is stage1_fused=False, "flydsl" is the compact plan. A "flydsl" without stage1_fused asked for the FlyDSL token-major dispatch, which no longer exists, so it is refused rather than run on mori. MegaMoEConfig gets a read-only dispatch_backend derived from stage1_fused. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * ci: run the MegaMoE bench on mori's dispatch stage1_fused=false with stage2_fused=true is mori's fp4 dispatch with the fused combine (the bench's auto wire is fp4 for a4w4_mxfp4), the configuration DeepSeek-R1 serves with. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: gfx1250 behaviour tests ahead of the refactor test_mega_moe_gfx1250.py pins what a change to MegaMoEGfx1250's dispatch/combine wiring must keep, through the constructor and members its callers use today: - every dispatch/combine pairing matches the fp32 reference, and a captured graph replays bit for bit what eager produced; - the pairings MegaMoE rejects fail before the symmetric arena is allocated; - mori's combine reads fused_moe's rows in place (zero-copy staging); - close() releases the plans before the window, the window before its memory; - the members ATOM's Triton experts reach (_dispatch, _combine, _arena, _config). It spawns its own ranks and skips (exit 0) off gfx1250, below two GPUs or without mori, so CI's MULTIGPU job can run it on gfx950. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: build from a MegaMoERecipe and a MegaMoEGeometry A caller now says what to run once, as a recipe, and MegaMoE builds it: recipe = MegaMoERecipe(dispatch=DispatchKind.MORI, combine=CombineKind.FUSED_SCATTER, dispatch_wire="fp4", combine_wires=("bf16", "fp4"), activation=..., quant_type=...) geometry = MegaMoEGeometry(rank, world_size, model_dim, inter_dim, experts, topk, max_tokens_per_rank) mega = MegaMoEGfx1250.from_recipe(recipe, geometry, communicator=comm) - MegaMoERecipe holds the transport (dispatch and combine kind, dispatch wire, the combine wires to build) and the expert-GEMM recipe; frozen and hashable, so a caller can key its cache of instances on it. MegaMoEGeometry holds the model and EP shape. - MegaMoERecipe.from_legacy is the one reader of stage1_fused, stage2_fused, dispatch_backend and the $MEGA_DISPATCH_WIRE / $MEGA_COMBINE_QUANT defaults. The keyword constructor goes through it, so both entries build the same thing. - MegaMoEConfig is now that recipe resolved against a geometry: frozen, with the same properties as before (stage1_fused, stage2_fused and dispatch_backend included, for the callers that read them). The dispatch schedule it used to have written into it lives on the instance; the dispatch_block_num / dispatch_warp_num_per_block fallback it carried was never reached. - forward() takes combine_wire= (bf16 | fp8 | fp4), the recipe's vocabulary, alongside combine_quant=. Every check and its message is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: probe mori once, and decide everything before the arena exists _MoriCapabilities gathers what MegaMoE needs to know about the installed mori -- the plan classes, the tuning lookup, and each helper a mori PR added (scale_stride_bytes #593, push_wire_nbytes #719, push_land_rows #745, ep_staging_bytes #703, TokOffExt.wanted #726, self_first_enabled) -- into one cached probe, and the decisions that hang off them into its methods: the tuned schedules, the landing-row sizes, the staging size, where the slot word lives. The ImportError/hasattr branches are gone from the pipeline; the RuntimeError messages for an mori that is too old are unchanged. _plan_transport runs all of it, plus the fused reduce's tiles and mori combine's block/epoch-slot check, before the symmetric arena is allocated. A configuration that cannot run -- mori missing, MORI_EP_TOKOFF_EXT=1 with mori's fp4 combine, no reduce tile, too many combine blocks -- now fails with nothing allocated instead of leaking the arena. TokOffExt.wanted, collective when MORI_EP_TOKOFF_EXT is on, moves ahead of the arena's own collective allocation on every rank alike. The compact dispatch with the fused combine still never loads mori's ops. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: one arena layout, and required_vmm_bytes from it The arena's regions now come from one formula, _arena_layout: the dispatch says what it writes (_dispatch_regions: the routing both kinds share, then its own rows), the combine what it stages through (_combine_regions), and _ArenaLayout.merge lays them down in the order the pipeline always used, so every pairing keeps the offsets it had. SymmetricArena takes the layout (or a region list, as before) and frees its memory if the window registration fails. MegaMoEGfx1250.required_vmm_bytes(recipe, geometry) reports the same layout, rounded the way ccoMemAlloc rounds an allocation -- the per_rank_vmm a caller needs for this MegaMoE, before any communicator exists, and exactly what the constructor then allocates. It is MegaMoE's share only; a communicator that also carries mori's own EP ops needs theirs on top. required_vmm_bytes(world_size=, hidden_dim=, max_tokens_per_rank=, experts=, topk=, stage1_fused=), the keyword form #6525 proposed and ATOM's yanbo/mega already calls, still works: it names neither wires nor combine nor inter_dim, so it returns the largest any recipe with that dispatch takes. _compact_gemm_align_m reads activation, quant type and inter_dim off the config instead of taking them separately. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: dispatch and combine stages under one dispatcher; linear forward The transport is now two stages the recipe picks, over one shared arena: _MoriDispatch / _CompactDispatch tokens out to their experts' ranks _MoriCombine / _FusedScatterCombine results back, summed per token Each declares its arena regions (regions(), pure -- what _arena_layout and required_vmm_bytes merge), plans in __init__ with nothing allocated, and allocates in build(). They meet only through _EpResources (the arena, the routing state both ends read, TokOffExt), the _DispatchBatch the dispatch hands the expert GEMM, and the _ComputeTarget the combine has it write into. _EpDispatcher picks the two from tables, runs every check before the arena is allocated, and frees what was built when a later step fails -- the arena and plans used to leak. close() goes: synchronize, combine plans, dispatch plans, staging, arena, TokOffExt. forward() is now: pick the step's combine wire, check the inputs, dispatch, prepare, compute, finish. The compact plan's machinery -- slots, side stream, per-step geometry (_CompactStep, now immutable) -- lives in _CompactDispatch, and the next layer's plan launches at the end of its dispatch, the same place on the GPU streams as before. The recv views are made once at build instead of per step. The TDM dispatch keeps only its compact arm; the token-major arms were unreachable once stage1_fused=False went through mori. Kept for the callers that reach past forward(), ATOM's Triton experts among them: _config, _arena, _dispatch() and its 5-tuple, _combine(), _recv_dispatch_scales(), prefetch_compact_plan(), warmup_compact_plan(). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: size mori's regions from describe(); the bench sizes cco from the recipe On a mori with Plan.describe() (ROCm/mori#844), the mori stages ask it for their arena regions and the dispatch's staging bytes -- the numbers mori's own kernels and its bind-time size check use -- instead of keeping formulas that mirror them. Each stage builds its plan request in one place (request()), which both describe() and the plan take. describe() compiles nothing and resolves against gfx1250 by name, so required_vmm_bytes stays off the GPU; the slot word's kwargs stay out of it, since where that word lives sizes nothing and asking would be collective. An older mori keeps the formulas. SymmetricArena.size() lets a #844 plan check its regions at bind. bench_mega_moe builds from a MegaMoERecipe and sizes its communicator with required_vmm_bytes(recipe, geometry) instead of a fixed 16 GiB. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: CPU tests for the recipe and the arena it sizes tests/test_mega_moe_gfx1250_recipe.py needs no GPU and skips where MegaMoE does not import. It checks how the legacy flags map to a recipe and that the rejections keep their wording, and lays the arena out over a grid of geometries and every pairing against the region list #6430 built inline -- through a fake mori both with and without Plan.describe(), the describe one following mori's rules, which add no overlapped-reduce flags on a wide EP. required_vmm_bytes is held to the layout as cco rounds it, and its keyword form to at least that. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging as a tensor A mori with ROCm/mori#844 checks a staging tensor against the plan's stagingBytes, device and alignment at bind, and holds it; an address is taken on trust. An older mori reads the tensor's address at each launch, as it did the int. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * TEMP: grouped GEMM takes a strided scale view again (fixup for #6946) #6946 flattens the grouped GEMM's scale arguments with .view(-1), which a non-contiguous tensor cannot do. MegaMoE's compact rows hand over exactly that -- each row's e8m0 scales sit inside its wire row -- so every stage1_fused MegaMoE step raises. The kernel only takes the base address and walks the rows by a_scale_row_stride_bytes, so a strided scale view becomes a one-dimensional view over the same base, as it was before #6946. Temporary: drop this commit once main fixes it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * bench_mega_moe: take out the kernel-table and --fused_shared_experts changes They are bench fixes rather than part of the MegaMoE recipe work, and move to #7074. This reverts "bench_mega_moe: fix kernel table on multi-stream graphs" and "bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them". Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * FlyDSL grouped GEMM: pass a strided scale view at its base address Replaces the temporary workaround for the cached launcher (#6946), which flattens the grouped GEMM's scale arguments with .view(-1). A non-contiguous tensor cannot do that, and MegaMoE's expert-grouped dispatch hands over exactly one: each row's e8m0 scales sit inside its wire row, read by a_scale_row_stride_bytes. The kernel reads these arguments only through their base pointer. tensor_shim gains flat_arg(), beside ptr_arg(): a contiguous tensor is .view(-1), as before, so no other caller changes; a strided one becomes the run of elements from its first to its last, at the same base. tests/test_grouped_gemm_strided_scales.py checks flat_arg and what the launcher hands the compiled kernel, without a GPU. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: stage contracts, mori sizing from describe() only, from_recipe first The recipe - The expert GEMM's fields move into ExpertComputeConfig, nested as MegaMoERecipe.compute; the recipe's own fields are the transport's. - Which pairings run is declared by the stages: a dispatch names the wires it carries and the BatchLayout it delivers (token-major or expert-grouped), a combine the layouts it accepts and the wires it offers. The recipe checks against those, and its errors name recipe fields (dispatch=, combine=, dispatch_wire=, combine_wires=). Each stage's geometry check (check()) runs from MegaMoEConfig. - from_recipe is the entry point. The keyword constructor warns DeprecationWarning; the underscored members ATOM's Triton experts use are documented as kept for them. The facade reads geometry and compute fields through properties rather than copying them. The stages - DispatchStage and CombineStage (Protocols) say what each provides; every stage has regions(config, caps) and check(config), and is built as Stage(config, caps). The compact- and mori-only attributes are gone from the shared base. - Each dispatch stage owns the stride its e8m0 rows arrive at (scale_stride), replacing MegaMoEConfig.dispatch_scale_dst_nbytes. - The batch carries its layout explicitly, and an expert-grouped one an ExpertGrouping; the combine hands the grouped GEMM its dispatch context in the _ComputeTarget, so the expert GEMM does not branch on the layout. - Whether the combine needs mori's aggregated reservation and whether the dispatch makes it are each stage's to say; the dispatcher refuses the pair before the arena exists. - One schedule pick (_pick) serves both mori plans and the compact kernels; _quantize returns the scale rows' address instead of leaving it in state. - _EpDispatcher has dispatch/prepare/combine/prefetch/warmup and the recv views, so forward() does not reach into the stages, and close() releases the rest even when one part raises. mori - _MoriCapabilities needs Plan.describe() (ROCm/mori#844). Regions, staging bytes and the scale stride come from there, so the copies of mori's formulas and the older-mori branches are gone; such a mori also overwrites total_recv in its dispatch, which is why nothing zeroes it. - _XDB_FLAG_SLOTS, TokOffExt and the plan constructors are reached only through _MoriCapabilities. Checks - An out-of-range recv_token_bound (and next_recv_token_bound) raises on the compact dispatch as on mori's, in dispatch, prefetch and warmup. - a1_scale on a quantizing dispatch wire raises ValueError (was an assert). - The keyword form of required_vmm_bytes raises when no candidate recipe runs at the geometry, instead of returning 0. - _MoriCombine.prepare documents why writing fused_moe's rows straight into comb_inp is safe: the next dispatch does not finish until every peer has dispatched, which each does only after its own combine. Tests: the CPU test fakes mori from describe() alone and covers the stage contract messages, the bound checks and the keyword-bound refusal; the multi-GPU test builds through from_recipe and checks the deprecated keyword constructor against it. bench_moe_gemm_a4w4_cudagraph builds through from_recipe. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: split mega_moe.py into recipe, mori_caps, arena, stages, transport A move, no change in behaviour: each block of mega_moe.py lands in one module as it was, and mega_moe.py keeps only MegaMoEGfx1250. - recipe.py: the wires, DispatchKind / CombineKind, ExpertComputeConfig, MegaMoERecipe, MegaMoEGeometry, MegaMoEConfig. - mori_caps.py: _MoriCapabilities and mori's region names. - arena.py: SymmetricArena, the region order, _ArenaLayout, cco's rounding. - stages/base.py: DispatchStage / CombineStage and what the stages share (routing, resources, batch, compute target, schedule pick). - stages/dispatch_mori.py, dispatch_compact.py, combine_fused.py, combine_mori.py: one stage each; stages/__init__.py maps each kind to its stage. - transport.py: _EpDispatcher and the one arena layout. The only edit: recipe.py looks the stages up when asked for (the stages import recipe.py). mega_moe.py still exports the package's public names, and Routing, which a bench imports from there. The CPU test imports the private names it checks from the modules that define them. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: compact dispatch bound and plan fixes, wire-picked GEMM A dtype - Compact dispatch: with no recv_token_bound the step is sized for the arena's capacity; a plan started for another alignment is refused; the plan's segment and wave limits are checked before the arena exists; plans are compiled when created, ahead of any graph capture. - A quantizing dispatch wire sets fused_moe's quant_dtype_a. - bench_mega_moe runs on gfx1250 only. - Tests: fp8 dispatch and fp8 combine cases, rank-local failures raised on every rank, CPU tests skip only without flydsl. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE CPU test: skip without mori as well - The package's kernel modules import mori.cco, so the test skips when mori is not installed instead of failing collection. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
solinzby1
pushed a commit
to ROCm/aiter
that referenced
this pull request
Nov 2, 2026
…tches, mori combine path, current mori EPv2 (#6430) * MegaMoE: stage1_fused/stage2_fused switches, mori combine path * bench_mega_moe: fix kernel table on multi-stream graphs * MegaMoE: follow mori tokoff-ext and selfFirst defaults * bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them * MegaMoE: follow mori#726 and size mori's fp4 combine flags Four fixes to the mori dispatch/combine path (stage1_fused/stage2_fused off), each from driving mori's plans without its op layer: - dispatch: since ROCm/mori#726 the plan takes tok_off_ext and derives selfFirst itself; self_first_enabled() is gone, so importing it fails. - tuning lookup: key it on the dispatch wire the way mori's op layer does ("fp4_disp_bf16_comb" for fp4). Under the hard-coded "bf16" key topk 9 matched no entry and both dispatch and the fp4 combine fell back to the single-shot default. - comb_inp: add the overlapped fp4 combine's u64 flag per landing row (four ranks up). Its 24-warp geometries write them past the region, i.e. past the arena: a page fault at 16384 tokens a rank. - close(): synchronize first, so a launch still in flight does not read a window that has been freed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: zero-copy hand-off from fused_moe to mori's combine With stage2_fused off, fused_moe's rows went into a fresh tensor and then into mori's staging rows (comb_inp): the bf16 combine copied them in its own kernel, the fp4 one through a host-issued copy whenever the recv bound was below max_recv -- every decode step under ATOM. Have fused_moe write them there directly. fused_moe(output=) already existed but the gfx1250 grouped path ignored it and copied at the end; thread it down to the gather-reduce. mori's bf16 combine then sees its input at the staging base and skips the copy, and the fp4 send reads the staging rows in place. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: follow mori#745/#746 on the mori dispatch/combine path mori's gfx1250 dispatch now reserves its slots once per destination (the aggregated reservation, ROCm/mori#745) and hands the push combine its landing offsets through the dispatch. MegaMoE drives the plans directly, so it has to size and pass what mori's op layer does: - dest_pe_token_counter holds 2 * world_size ints: send counts, then per destination the reservation base and the push combine's landing offset. It held world_size; the dispatch wrote the second half into allocator padding. - recv_num signals are 8 bytes a source. The region held 4 bytes a source; the arena's 256-byte alignment hid the overrun. - The combine reads its landing offsets from dest_pe_token_counter, which MegaMoE did not pass, so the mxfp4 combine (stage2_fused=False) got a null pointer there. - MORI_EP_TOKOFF_EXT=1 turns the aggregated reservation off, so a mori push combine is refused with it, as mori's op layer refuses it. - The push landing rows come from mori's push_land_rows(): one per (token, destination rank), max_tokens * min(world_size, topk) rather than world_size * max_recv. - The mori dispatch launcher drops total_recv.zero_(): gfx1250's dispatch overwrites the count (since ROCm/mori#599), and mori's op layer skips the same zero there. Against a mori from before #745 the landing rows keep the old count and the two larger buffers have unused tails. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging buffer (ROCm/mori#703) mori#703 moves the gfx1250 dispatch's staging arrays out of each kernel module's .bss into one buffer the caller allocates and binds as staging_base; an unbound plan is refused on the host. MegaMoE drives EpDispatchPlan directly, so it allocates ep_staging_bytes() once and binds it to every dispatch plan (they never run at the same time), and drops it after the plans on close(). A mori without ep_staging_bytes keeps the .bss arrays and is left as is. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: keep accepting dispatch_backend stage1_fused now picks the dispatch, and the dispatch_backend argument and config field went away with that. Callers still pass it (ATOM main passes dispatch_backend="mori" on a quantizing wire; ROCm/ATOM#2384 passes "flydsl" or "mori" alongside stage1_fused) and log mega._config.dispatch_backend, so both would fail on this aiter. Accept the argument again and check it against stage1_fused: "mori" is stage1_fused=False, "flydsl" is the compact plan. A "flydsl" without stage1_fused asked for the FlyDSL token-major dispatch, which no longer exists, so it is refused rather than run on mori. MegaMoEConfig gets a read-only dispatch_backend derived from stage1_fused. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * ci: run the MegaMoE bench on mori's dispatch stage1_fused=false with stage2_fused=true is mori's fp4 dispatch with the fused combine (the bench's auto wire is fp4 for a4w4_mxfp4), the configuration DeepSeek-R1 serves with. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: gfx1250 behaviour tests ahead of the refactor test_mega_moe_gfx1250.py pins what a change to MegaMoEGfx1250's dispatch/combine wiring must keep, through the constructor and members its callers use today: - every dispatch/combine pairing matches the fp32 reference, and a captured graph replays bit for bit what eager produced; - the pairings MegaMoE rejects fail before the symmetric arena is allocated; - mori's combine reads fused_moe's rows in place (zero-copy staging); - close() releases the plans before the window, the window before its memory; - the members ATOM's Triton experts reach (_dispatch, _combine, _arena, _config). It spawns its own ranks and skips (exit 0) off gfx1250, below two GPUs or without mori, so CI's MULTIGPU job can run it on gfx950. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: build from a MegaMoERecipe and a MegaMoEGeometry A caller now says what to run once, as a recipe, and MegaMoE builds it: recipe = MegaMoERecipe(dispatch=DispatchKind.MORI, combine=CombineKind.FUSED_SCATTER, dispatch_wire="fp4", combine_wires=("bf16", "fp4"), activation=..., quant_type=...) geometry = MegaMoEGeometry(rank, world_size, model_dim, inter_dim, experts, topk, max_tokens_per_rank) mega = MegaMoEGfx1250.from_recipe(recipe, geometry, communicator=comm) - MegaMoERecipe holds the transport (dispatch and combine kind, dispatch wire, the combine wires to build) and the expert-GEMM recipe; frozen and hashable, so a caller can key its cache of instances on it. MegaMoEGeometry holds the model and EP shape. - MegaMoERecipe.from_legacy is the one reader of stage1_fused, stage2_fused, dispatch_backend and the $MEGA_DISPATCH_WIRE / $MEGA_COMBINE_QUANT defaults. The keyword constructor goes through it, so both entries build the same thing. - MegaMoEConfig is now that recipe resolved against a geometry: frozen, with the same properties as before (stage1_fused, stage2_fused and dispatch_backend included, for the callers that read them). The dispatch schedule it used to have written into it lives on the instance; the dispatch_block_num / dispatch_warp_num_per_block fallback it carried was never reached. - forward() takes combine_wire= (bf16 | fp8 | fp4), the recipe's vocabulary, alongside combine_quant=. Every check and its message is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: probe mori once, and decide everything before the arena exists _MoriCapabilities gathers what MegaMoE needs to know about the installed mori -- the plan classes, the tuning lookup, and each helper a mori PR added (scale_stride_bytes #593, push_wire_nbytes #719, push_land_rows #745, ep_staging_bytes #703, TokOffExt.wanted #726, self_first_enabled) -- into one cached probe, and the decisions that hang off them into its methods: the tuned schedules, the landing-row sizes, the staging size, where the slot word lives. The ImportError/hasattr branches are gone from the pipeline; the RuntimeError messages for an mori that is too old are unchanged. _plan_transport runs all of it, plus the fused reduce's tiles and mori combine's block/epoch-slot check, before the symmetric arena is allocated. A configuration that cannot run -- mori missing, MORI_EP_TOKOFF_EXT=1 with mori's fp4 combine, no reduce tile, too many combine blocks -- now fails with nothing allocated instead of leaking the arena. TokOffExt.wanted, collective when MORI_EP_TOKOFF_EXT is on, moves ahead of the arena's own collective allocation on every rank alike. The compact dispatch with the fused combine still never loads mori's ops. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: one arena layout, and required_vmm_bytes from it The arena's regions now come from one formula, _arena_layout: the dispatch says what it writes (_dispatch_regions: the routing both kinds share, then its own rows), the combine what it stages through (_combine_regions), and _ArenaLayout.merge lays them down in the order the pipeline always used, so every pairing keeps the offsets it had. SymmetricArena takes the layout (or a region list, as before) and frees its memory if the window registration fails. MegaMoEGfx1250.required_vmm_bytes(recipe, geometry) reports the same layout, rounded the way ccoMemAlloc rounds an allocation -- the per_rank_vmm a caller needs for this MegaMoE, before any communicator exists, and exactly what the constructor then allocates. It is MegaMoE's share only; a communicator that also carries mori's own EP ops needs theirs on top. required_vmm_bytes(world_size=, hidden_dim=, max_tokens_per_rank=, experts=, topk=, stage1_fused=), the keyword form #6525 proposed and ATOM's yanbo/mega already calls, still works: it names neither wires nor combine nor inter_dim, so it returns the largest any recipe with that dispatch takes. _compact_gemm_align_m reads activation, quant type and inter_dim off the config instead of taking them separately. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: dispatch and combine stages under one dispatcher; linear forward The transport is now two stages the recipe picks, over one shared arena: _MoriDispatch / _CompactDispatch tokens out to their experts' ranks _MoriCombine / _FusedScatterCombine results back, summed per token Each declares its arena regions (regions(), pure -- what _arena_layout and required_vmm_bytes merge), plans in __init__ with nothing allocated, and allocates in build(). They meet only through _EpResources (the arena, the routing state both ends read, TokOffExt), the _DispatchBatch the dispatch hands the expert GEMM, and the _ComputeTarget the combine has it write into. _EpDispatcher picks the two from tables, runs every check before the arena is allocated, and frees what was built when a later step fails -- the arena and plans used to leak. close() goes: synchronize, combine plans, dispatch plans, staging, arena, TokOffExt. forward() is now: pick the step's combine wire, check the inputs, dispatch, prepare, compute, finish. The compact plan's machinery -- slots, side stream, per-step geometry (_CompactStep, now immutable) -- lives in _CompactDispatch, and the next layer's plan launches at the end of its dispatch, the same place on the GPU streams as before. The recv views are made once at build instead of per step. The TDM dispatch keeps only its compact arm; the token-major arms were unreachable once stage1_fused=False went through mori. Kept for the callers that reach past forward(), ATOM's Triton experts among them: _config, _arena, _dispatch() and its 5-tuple, _combine(), _recv_dispatch_scales(), prefetch_compact_plan(), warmup_compact_plan(). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: size mori's regions from describe(); the bench sizes cco from the recipe On a mori with Plan.describe() (ROCm/mori#844), the mori stages ask it for their arena regions and the dispatch's staging bytes -- the numbers mori's own kernels and its bind-time size check use -- instead of keeping formulas that mirror them. Each stage builds its plan request in one place (request()), which both describe() and the plan take. describe() compiles nothing and resolves against gfx1250 by name, so required_vmm_bytes stays off the GPU; the slot word's kwargs stay out of it, since where that word lives sizes nothing and asking would be collective. An older mori keeps the formulas. SymmetricArena.size() lets a #844 plan check its regions at bind. bench_mega_moe builds from a MegaMoERecipe and sizes its communicator with required_vmm_bytes(recipe, geometry) instead of a fixed 16 GiB. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: CPU tests for the recipe and the arena it sizes tests/test_mega_moe_gfx1250_recipe.py needs no GPU and skips where MegaMoE does not import. It checks how the legacy flags map to a recipe and that the rejections keep their wording, and lays the arena out over a grid of geometries and every pairing against the region list #6430 built inline -- through a fake mori both with and without Plan.describe(), the describe one following mori's rules, which add no overlapped-reduce flags on a wide EP. required_vmm_bytes is held to the layout as cco rounds it, and its keyword form to at least that. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: bind mori's dispatch staging as a tensor A mori with ROCm/mori#844 checks a staging tensor against the plan's stagingBytes, device and alignment at bind, and holds it; an address is taken on trust. An older mori reads the tensor's address at each launch, as it did the int. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * TEMP: grouped GEMM takes a strided scale view again (fixup for #6946) #6946 flattens the grouped GEMM's scale arguments with .view(-1), which a non-contiguous tensor cannot do. MegaMoE's compact rows hand over exactly that -- each row's e8m0 scales sit inside its wire row -- so every stage1_fused MegaMoE step raises. The kernel only takes the base address and walks the rows by a_scale_row_stride_bytes, so a strided scale view becomes a one-dimensional view over the same base, as it was before #6946. Temporary: drop this commit once main fixes it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * bench_mega_moe: take out the kernel-table and --fused_shared_experts changes They are bench fixes rather than part of the MegaMoE recipe work, and move to #7074. This reverts "bench_mega_moe: fix kernel table on multi-stream graphs" and "bench_mega_moe: --fused_shared_experts as ATOM DP+EP runs them". Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * FlyDSL grouped GEMM: pass a strided scale view at its base address Replaces the temporary workaround for the cached launcher (#6946), which flattens the grouped GEMM's scale arguments with .view(-1). A non-contiguous tensor cannot do that, and MegaMoE's expert-grouped dispatch hands over exactly one: each row's e8m0 scales sit inside its wire row, read by a_scale_row_stride_bytes. The kernel reads these arguments only through their base pointer. tensor_shim gains flat_arg(), beside ptr_arg(): a contiguous tensor is .view(-1), as before, so no other caller changes; a strided one becomes the run of elements from its first to its last, at the same base. tests/test_grouped_gemm_strided_scales.py checks flat_arg and what the launcher hands the compiled kernel, without a GPU. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: stage contracts, mori sizing from describe() only, from_recipe first The recipe - The expert GEMM's fields move into ExpertComputeConfig, nested as MegaMoERecipe.compute; the recipe's own fields are the transport's. - Which pairings run is declared by the stages: a dispatch names the wires it carries and the BatchLayout it delivers (token-major or expert-grouped), a combine the layouts it accepts and the wires it offers. The recipe checks against those, and its errors name recipe fields (dispatch=, combine=, dispatch_wire=, combine_wires=). Each stage's geometry check (check()) runs from MegaMoEConfig. - from_recipe is the entry point. The keyword constructor warns DeprecationWarning; the underscored members ATOM's Triton experts use are documented as kept for them. The facade reads geometry and compute fields through properties rather than copying them. The stages - DispatchStage and CombineStage (Protocols) say what each provides; every stage has regions(config, caps) and check(config), and is built as Stage(config, caps). The compact- and mori-only attributes are gone from the shared base. - Each dispatch stage owns the stride its e8m0 rows arrive at (scale_stride), replacing MegaMoEConfig.dispatch_scale_dst_nbytes. - The batch carries its layout explicitly, and an expert-grouped one an ExpertGrouping; the combine hands the grouped GEMM its dispatch context in the _ComputeTarget, so the expert GEMM does not branch on the layout. - Whether the combine needs mori's aggregated reservation and whether the dispatch makes it are each stage's to say; the dispatcher refuses the pair before the arena exists. - One schedule pick (_pick) serves both mori plans and the compact kernels; _quantize returns the scale rows' address instead of leaving it in state. - _EpDispatcher has dispatch/prepare/combine/prefetch/warmup and the recv views, so forward() does not reach into the stages, and close() releases the rest even when one part raises. mori - _MoriCapabilities needs Plan.describe() (ROCm/mori#844). Regions, staging bytes and the scale stride come from there, so the copies of mori's formulas and the older-mori branches are gone; such a mori also overwrites total_recv in its dispatch, which is why nothing zeroes it. - _XDB_FLAG_SLOTS, TokOffExt and the plan constructors are reached only through _MoriCapabilities. Checks - An out-of-range recv_token_bound (and next_recv_token_bound) raises on the compact dispatch as on mori's, in dispatch, prefetch and warmup. - a1_scale on a quantizing dispatch wire raises ValueError (was an assert). - The keyword form of required_vmm_bytes raises when no candidate recipe runs at the geometry, instead of returning 0. - _MoriCombine.prepare documents why writing fused_moe's rows straight into comb_inp is safe: the next dispatch does not finish until every peer has dispatched, which each does only after its own combine. Tests: the CPU test fakes mori from describe() alone and covers the stage contract messages, the bound checks and the keyword-bound refusal; the multi-GPU test builds through from_recipe and checks the deprecated keyword constructor against it. bench_moe_gemm_a4w4_cudagraph builds through from_recipe. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: split mega_moe.py into recipe, mori_caps, arena, stages, transport A move, no change in behaviour: each block of mega_moe.py lands in one module as it was, and mega_moe.py keeps only MegaMoEGfx1250. - recipe.py: the wires, DispatchKind / CombineKind, ExpertComputeConfig, MegaMoERecipe, MegaMoEGeometry, MegaMoEConfig. - mori_caps.py: _MoriCapabilities and mori's region names. - arena.py: SymmetricArena, the region order, _ArenaLayout, cco's rounding. - stages/base.py: DispatchStage / CombineStage and what the stages share (routing, resources, batch, compute target, schedule pick). - stages/dispatch_mori.py, dispatch_compact.py, combine_fused.py, combine_mori.py: one stage each; stages/__init__.py maps each kind to its stage. - transport.py: _EpDispatcher and the one arena layout. The only edit: recipe.py looks the stages up when asked for (the stages import recipe.py). mega_moe.py still exports the package's public names, and Routing, which a bench imports from there. The CPU test imports the private names it checks from the modules that define them. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE: compact dispatch bound and plan fixes, wire-picked GEMM A dtype - Compact dispatch: with no recv_token_bound the step is sized for the arena's capacity; a plan started for another alignment is refused; the plan's segment and wave limits are checked before the arena exists; plans are compiled when created, ahead of any graph capture. - A quantizing dispatch wire sets fused_moe's quant_dtype_a. - bench_mega_moe runs on gfx1250 only. - Tests: fp8 dispatch and fp8 combine cases, rank-local failures raised on every rank, CPU tests skip only without flydsl. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * MegaMoE CPU test: skip without mori as well - The package's kernel modules import mori.cco, so the test skips when mori is not installed instead of failing collection. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <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.
gfx1250 EPv2: EP4 topk-9 geometry; selfFirst follows MORI_EP_TOKOFF_EXT
Summary
hip_tuning_configs.pyentries. Without them it ran the single-shot 64x16 dispatch and 64x8 fp4 combine, and 64x8 never takes the overlapped fp4 combine path at EP4. The entries reuse topk 6's fp4 buckets (dispatch 64x8 / 64x16 / 128x16, fp4 combine 64x8 / 64x16 / 64x24 / 128x24). bf16 / fp8 dispatch are unmeasured and keep the default.MORI_EP_TOKOFF_EXTis the only one. The rendered kernel source is unchanged (64b1384).Test
4x gfx1250 single host (EP4), fp4 dispatch + fp4 push combine, topk 9, 64 experts/rank, HIP graph, ITERS=50, 3 alternating rounds, run-internal medians in us. Every arm passed CHECK=1 first, including fp32-scale fp4 combine and bf16 pull combine. Baseline is main 28d4973; the other arm is the same build with this table. Dispatch at 1024 and combine at 256 keep their old geometry.