Repository navigation
[SM90] Ring-buffered FP8 MegaMoE with a lag-ordered L1/L2 schedule - #88
Conversation
5dbc04b to
26fd9bb
Compare
4ee8d12 to
c724f5e
Compare
|
Hi @zkyue , thanks for the great work — I've read through this PR and find the lag-ordered ring design really impressive. I do have two questions though: Is the main performance gain at large batch sizes mostly from avoiding the NVLink egress contention between the dispatch pull and the L2 scatter, i.e. is the lag ordering the dominant factor rather than the smaller memory footprint itself? Does the same contention exist on SM100? Since NVLink bandwidth there is roughly 2× SM90's, I'm curious whether the doubled bandwidth would absorb the pull/scatter overlap, or whether the lag-ordered schedule would still pay off on SM100. |
|
@Helicoplacoidea both are measurable, so: Yes, the ordering, not the footprint. The ring's lower bound is the schedule's live set, so pinning the L1->L2 distance is what shrinks the pool; at a fixed lag, extra ring capacity stayed inside our noise band. Instrumented traces at 8192 tokens/rank put the same remote pull at 76 % of the kernel span at lag 8 against 33 % at lag 40, with ~70 % of L2 tiles starting before it completes against none. Contention is the wrong word, though: in a microbenchmark of these shapes, running pull and scatter at once takes the sum of their standalone times to within 1 %. The link is serial to first order, so overlapping the scatter delays the pull rather than losing egress, and pull completion gates L2. On SM100 we already have it: four files on top of Through upstream's own test against unmodified The doubled link doesn't absorb it, and isn't what the gain rides on either. The ordering is already there implicitly: an L2 task can't be claimed before its own block's L1 tasks are issued, so the schedule warms up on L1 first -- closed-form deadlock avoidance, two waves on all six shapes -- leaving We plan to send that one to |
|
@zkyue , thank you for the detailed and measurable answer — it clears up both questions completely. Really nice work — looking forward to seeing this merged, and to the SM100 follow-up landing upstream as well. Thanks again! |
c724f5e to
4267984
Compare
|
@Fridge003 would you have a chance to take a look when you get a moment? |
|
@Fridge003 force-pushed: rebased onto Follow-ups: the workspace now ends on a 1 KiB boundary so the pools after it start on a 128 B line — 1.4–2.8 % faster from 2048 tokens/rank, output bitwise identical. The unreachable 512/128 epilogue-thread branches are gone; SASS is byte-identical for all 87 suite kernels and the 16 production shapes. The five defects, all present at 4267984: a hang when a wave ring is smaller than the pool; a stale Production shapes: host config, kernel SASS and outputs unchanged. |
The Hopper fused MoE kernel needs shared-memory and system-scope load/store wrappers that the PTX helpers did not have yet, plus a way to advance a WGMMA shared-memory descriptor by a byte offset instead of rebuilding it. This commit adds both, and on SM90 moves the NVLink barrier's timeout trap out of the wait loop: with a `trap;` inside the loop body ptxas allocates the registers of the region containing the loop against the kernel's launch bound and ignores that region's `setmaxnreg.inc`. Other architectures, and SM90 with the verbose timeout, compile the barrier as before. Behaviour is unchanged for every kernel that exists today.
The Hopper fused MoE staged every dispatched row in an activation pool sized for the worst case, so the symmetric buffer grew with the token bound rather than with what the two GEMMs have in flight. This adds a lag-ordered schedule -- a per-CTA tile table that keeps the second GEMM a bounded number of tiles behind the first -- and sizes the activation pool as a ring around that lag, so the activation pools hold a bounded window of the rows in flight instead of the worst case, while the dispatch input, combine and workspace regions still grow with the capacity; the layout follows from the shape alone -- the full pool below 1024 tokens per rank, a lag ring from there up -- so the buffer sizing and every launch derive the same ring capacity and lag limit. The intermediate activation scales are per 128 K, which takes the 256-wide tiles: sizing and launch reject a hidden size that is not a multiple of 512 or an intermediate size that is not a multiple of 256, sizing also rejects more than 32 ranks (the kernel keeps one flag per rank in a lane), and the swapAB decode path, whose L1 epilogue is not bitwise stable run to run, is no longer selected and is removed. A split-M tile sizes its pipeline for the quarter-width L2 staging buffer it uses, so a capacity as large as 65536 tokens per rank, whose tile table takes more shared memory, still keeps two stages. A launch can also pass `num_tokens_bound`, the per-rank token count every rank's batch fits in, and have its schedule picked from that instead of from the buffer capacity; a bound the buffer cannot serve -- negative, below this rank's token count, or above the capacity -- is rejected on the host. Calls of at most 256 tokens per rank -- the caller's bound where there is one, this rank's own count otherwise -- take the programmatic dependent launch attribute and the instantiation that carries the matching griddepcontrol wait/trigger; every other call follows the global flag. The SM90 scheduler and workspace layout move into headers of their own, which is why `SM90Workspace` also disappears from the shared layout header here: the SM90 header defines the ring-aware replacement under the same name. Co-authored-by: Haisha Zhao <33570593+Hyaloid@users.noreply.github.com> Co-authored-by: Clement Chan <16222986+iclementine@users.noreply.github.com>
`LegacyMegaMoEScheduler` was the wave scheduler the Hopper kernel used before the lag-ordered one, and it has no instantiations left. Removing it leaves one scheduling contract per architecture in the shared scheduler header.
The new test runs the fused Hopper kernel over several buffer sizes, four launches per buffer, and requires a full activation pool and a ring to agree bitwise on the launches both can take; the verdict is all-reduced so a mismatch on one rank fails the run. Its shapes include the 1024-token buffer, the smallest one the default layout turns into a ring, and it checks the layout each buffer gets against the host's threshold: from 1024 tokens per rank a ring no larger than the full pool, smaller on at least one shape, and the full pool below. A further buffer repeats the same four launches with each call's own `num_tokens_bound` and must match bit for bit, which covers both schedules a bound can pick -- the wave order on the decode-sized launch and a shorter lag on a top-6 shape the test adds for it -- and that a bound above the capacity is rejected, as is a buffer for an intermediate size that is not a multiple of 256. The existing accuracy test's reference now quantizes the intermediate activation per 128 K, like the kernel, and the test script runs the accuracy layers before the benchmark pass, which only checks the output shape.
|
@Fridge003 Rebased onto While porting I also dropped the two layout knobs, On 8x H100 with a cold JIT cache, |
Co-authored with @iclementine (the first SM90 ring buffer) and @Hyaloid (cluster and multicast).
dev's SM90 fused FP8 MegaMoE runs L1 and L2 as expert waves, so the pool for L1 outputs is sized for all tokens of all ranks and L2's NVLink scatter starts while the dispatch pull is still running on the same per-rank egress. This PR runs one task stream with L2 trailing L1 by a fixed lag and sizes that pool as a ring around the lag. On 8×H100 EP8 over three public model configurations that is 1.36–1.49× at 1024–8192 tokens/rank and 1.05–1.12× at 64–256, with the symmetric buffer per rank going 2788 → 287 MiB at 2048 tokens/rank and 6828 → 857 MiB at 8192 on 4096/2048/288/8. One change alters the output on purpose; it is described below.Performance
GLM-5.3-Flash — 4096 / 2048 / 288 / top-8 (hidden / intermediate / routed experts / top-k)
8×H100 SXM 80 GB, EP8, CUDA 13.3, one session. Each number is the device time of the fused kernel alone — one launch per call, nothing else in the MoE forward — from Kineto over a loop of 256 consecutive calls, median call per rank, slowest of the eight ranks, mean of two rounds. The baseline is
devat cbe1727 with.templateadded at the fourget_base_ptr<float>()call sites nvcc 13.3 rejects, so its kernel JITs. Each arm allocates the buffer its own branch would for the same request:devpads the token bound to 1920 on SM90 and this PR pads to 128, so where this branch builds for 2048 tokensdevbuilds for 3840.This branch runs the per-128 scale in every cell of all three tables. The shared expert is not in these numbers: the SM90 path asserts
num_shared_experts == 0, so it is a separate GEMM either way. The lag order is also not what carries the 64- and 256-token rows — a call of ≤ 256 tokens/rank whose worst-case pool fits takes the single-wave order with the lag disabled, so those rows are the decode tile, the ring and PDL. Decode rows use a 2048-token buffer withnum_tokens_bound= the call's M.The other two configurations, and the benchmark method
DeepSeek-V4-Flash — 4096 / 2048 / 256 / top-6
DeepSeek-V4-Pro — 7168 / 3072 / 384 / top-6
Same method throughout. The 256 calls are issued back to back, with no synchronisation and no L2 flush between them. The baseline is the mean of its two arms across both rounds, and a baseline arm runs first and last in every group, so the gap between those two arms is that cell's noise floor: at most 0.96 % across the 18 cells and 0.5 % or less in 11 of them. At 8192 tokens/rank 256 calls is not quite enough for the clocks to settle — a 512-call loop reads up to 1.6 % slower on both arms and puts the three ratios at 1.42×, 1.42× and 1.34× against 1.41×, 1.41× and 1.36× — so those rows are a lower bound on steady-state latency.
dev's 1920 is the LCM of its candidate tile widths (#86 moved it from 384 by adding a 240-wide tile). PDL follows the launcher's rule (requested for calls of ≤ 256 tokens/rank, otherwise the global flag, which was off), anddev's SM90 kernel has no griddepcontrol at all, so the baseline runs without the attribute at every size and the two decode rows of each table include that difference. A 1024-token buffer is where the lag rule starts, so the 1024 rows run the lag order on a ring, by the same rule as 2048 and above.Symmetric buffer per rank,
dev→ this PR, each pair being what the two branches actually allocate for the same request: on 4096/2048/288/8, 2788 → 287 MiB at 2048 tokens/rank and 6828 → 857 MiB at 8192. The ring accounts for 3.5× of that at 2048 tokens/rank; the rest isdev's 1920-token alignment together with its scale-factor pools, which it sizes for the smallest candidate tile where this branch sizes them for the two it can select. At this branch's alignment and pool layout, a full pool for the same 2048-token bound is 1002 MiB against the ring's 287, and 3920 against 857 at 8192. Only the L1/L2 activation pools are sized from the ring; the dispatch input buffers and the combine buffer stay sized bynum_max_tokens_per_rank, which is why the total still grows with the buffer.The rest of the memory numbers, and where the ring starts
4096/2048/256/6 goes 2104 → 240 MiB at 2048 tokens/rank and 5147 → 700 at 8192; 7168/3072/384/6 goes 3546 → 490 and 8587 → 1113. On 4096/2048/288/8 the combine buffer alone is 128 MiB at 2048 tokens and 512 MiB at 8192.
Buffers of 1024 tokens/rank and above get the lag ring, below that the full pool. At 1024 the ring is bitwise identical to the full pool, 2.5 % faster back to back and 209 instead of 515 MiB, and at 2048 and 8192 the same pair reads 1.5 % and 0.4 % faster. Below the threshold the lag is off and the buffer is the full pool; a wave-sized ring there saved only 272 → 258 MiB at 512 tokens/rank while reading 0.4 % slower, which is why the threshold stops at 1024.
How it works
The B loader enumerates a stream in which a token block's L2 tasks sit
lagunits after its L1 tasks, in groups ofgroupunits; a unit is up to 8 consecutive m-blocks of one expert. The host pickslag = clamp(round(0.45 × units), 8, 64),group = 4, from the call's token bound, withunitsestimated from that bound and the capacity, not the routed count. The ring then follows: at most8 × (lag + group)blocks of L1 output are live, so the pool becomes a ring of8 × (lag + group) + 4slots, asserted against that live set on the host. Buffers below the lag threshold keep the wave order, and so does a call on a lag ring that is small enough by the rule below and whose worst-case pool fits the ring.Why the order and not the pool size, and how slots are recycled
In-kernel event traces on 8×H100 at 8192 tokens/rank, taken on an instrumented build of this branch with the lag forced to 8 units and to 40 so that the order is the only variable, put the end of the dispatch pull at 2671 µs when L2 starts early and at 1102 µs when it is held back.
Consumers replay the tile table the B loader publishes — payload first, then an acquire-loaded tag — and the L2 tail after the last L1 is handed out by an atomic ticket. Recycling reuses the existing
l1/l2_fullcounters; the SM90 workspace gains the matchingemptycounters and launch-parity banks for the flags that used to be protected by the head and exit all-rank barriers, which are gone. L1 epilogues wait onl2_emptybefore overwriting a slot's previous generation, and dispatch warps publish any pending arrival before blocking on a slot, so no wait can point at unpublished work. The+ 4iskSm90LagRingMarginandkSm90AutoLagMinTokensthe 1024-token threshold; a ring clamped to the full pool never wraps and is exempt from the host assertion.Also in the diff, downstream of the order and the ring: arrival publishes batched up to 16 rows on calls of 8192 tokens/rank and up, when the ring covers the batch; an early combine in which math warps reduce tokens whose selected experts have all published while the all-rank barrier is pending; a BLOCK_M=64, 256-wide split-N tile with a sliced combine queue for decode-sized calls;
num_tokens_boundfor a single-wave schedule on a large buffer; and an XOR bank swizzle on the L2 epilogue staging, the same formula as #69's.Before you merge
One default changes: the FP8 scale on the intermediate activation between the two GEMMs goes from one per 64 contiguous K elements to one per 128, which halves the scale bytes. 128 is the grouping the non-fused DeepEP + grouped-GEMM path already uses. Per-128 needs the 256-wide tiles, so SM90 requires
hidden % 512 == 0andintermediate_hidden % 256 == 0, the divisibilitydev's sizer already enforces. Creating the symmetric buffer now also rejects anintermediate_hiddenabove 4096, whichdevrejects only at launch, and more than 32 ranks, which the low-latency head does not support.The 4-warpgroup prefill choice and the swapAB decode tile are gone — the swapAB L1 epilogue is not bitwise stable run to run. Its path is removed from the kernel and the host, along with the per-64 branches the fixed width leaves unreachable.
What the width measures against dev and against torch
The requant scale expression is
dev's throughout (#63: amax clamped at 1e-10,sf = amax·(1/448),sf_inv = 1/sf), so the width is the only difference. An earlier revision of this branch could hold it at 64, and there the output was bitwise identical todev— verified at 64 and at 256 tokens/rank on a 2048-token buffer — and differed only by tile choice above that (rel-L2 ≤ 1.0e-3 at 512–8192). Two rewrites sit inside that claim: anex2.approx.ftz.f32fast-math SiLU, which differs from__expf's non-ftz form only whereexp(-x) < 2^-126and the following1 + erounds that difference away, and anfmaxf/fabsfamax reduction over explicit shuffles, which agrees bit for bit withcute::max/absfor non-NaN values. The width-64 check covered both end to end.At per-128 the difference against
devis 2.1e-2 at every size from 64 to 8192 tokens/rank, and comes from the coarser grouping, not from tile choice; against the torch reference it does not drift with size either — rel-L2 2.4e-2 against an fp8-aware reference and 6.4e-2 against bf16, at both 256 and 2048 tokens/rank. Per-128 also halves the act-SF pool and lets each 128-K step of L2's mainloop take one scale group and one rescale instead of two.num_tokens_boundmust be the same value on every rank — an upper bound on this call's per-rank token count everywhere, e.g. the max over ranks of the local batch.Nonemeans the buffer capacity, and a bound the buffer cannot serve is rejected on the host — which checks this rank's own count and the capacity, not the other ranks. A bound below a peer's count hangs: this rank's pool is then sized for fewer rows than will arrive, so the ring reuses a slot inside one wave.What the bound moves, and what it cannot move
A negative bound is rejected too. The bound stands in for this rank's count in the call's lag, its worst-case ring-fit test, the publish batch, eager publishing, PDL and — at
BLOCK_M128 — the single-wave class test; the tile, the wave size and theBLOCK_M-64 class test stay keyed to the rank's own count. Without a bound the ring-fit test uses the capacity, which fits only a ring about the size of its full pool, so it is supplying the bound that lets a large buffer take the single-wave order on a small call. The ring capacity is the buffer's either way, and the call's lag is its bound's, clamped to the buffer's and asserted on the host, so a call cannot ask for a lag its allocation does not hold.The op schemas change by one defaulted argument. The SM90 size and slice ops keep
dev's argument list — the ring capacity, the lag limit and the scale granularity all follow from the shape, so the sizing and every launch derive the same values — andfp8_mega_moegains a trailingint num_tokens_bound=0. The newget_token_alignment_for_sm90_mega_moe()returns 128 instead of the shared 1920, so an unchanged call allocates a smaller buffer.API surface
get_symm_buffer_for_sm90_mega_moe(...)keeps its signature;fp8_mega_moe(..., num_tokens_bound=None)is the only addition.The layout is a function of the shape: the full pool below 1024 tokens per rank, a lag ring from there up. There is no layout knob, and the generic
get_symm_buffer_for_mega_moeassertsnum_shared_experts == 0on SM90.This PR adds no tuning environment variable of its own; the two trace instruments are its only env reads.
dev'sDG_SM90_FP8_SWAP_ABkill-switch goes with the swapAB path it switched.One shared code path changes, for SM90 only:
comm/barrier.cuhmoves the NVLink-barrier timeout trap to after the wait loop, where inside the loop it made ptxas allocate the SM90 kernel's registers against the launch bound. The change sits under the preprocessor condition the file already used for that trap, so other architectures compiledev's code unchanged; five SM100 MegaMoE instantiations (fp8xfp4 and bf16) give byte-identical SASS withdev's header and with this branch's.Everything else the diff touches outside the SM90 path
Kernel
impls/sm90_fp8_mega_moe.cuh; newlayout/sm90_mega_moe.cuh(SM90 workspace, tile-table sizing) andscheduler/sm90_mega_moe.cuh(the lag scheduler).layout/mega_moe.cuhlosesSM90Workspaceandscheduler/mega_moe.cuhlosesLegacyMegaMoEScheduler; nothing else in the SM100 headers changes. Shared headers otherwise: 13 ptx wrappers inptx/ld_st.cuhandadvance_smem_descinmma/sm90.cuh, both additive.Outside those, the diff is the SM90 host path —
apis/sm90_mega.hpp, the SM90 heuristics and impl headers, theTORCH_LIBRARYentry and the Python wrapper — plus the tests; no SM100 or other-architecture code path changes. The kernel keeps.templateon the dependentget_base_ptr<float>()calls that need it:devat cfedbd3 fails to JIT under nvcc 13.3 at four such call sites, the same four #74 fixes.Neither SM90 MegaMoE test runs in CI: there is no test workflow in this repo, and both spawn one process per GPU, 8 by default, so they need SM90 GPUs and a working NVLink group; the ring test also passes at 4 and 2 processes.
What the two tests cover
sgl_deep_gemm/tests/test_mega_moe_hopper.py --accuracy --layers 1 2 3 4: 29/29 scenarios; the reference quantises the intermediate activation per 128 K, like the kernel.sgl_deep_gemm/tests/test_mega_moe_hopper_ring.pyis new (~3 min on 8 GPUs). It defines its own shapes on both sides of the lag threshold, plus 4096/2048/288/8 @ 2048 tokens and a 2-rank subgroup shape whose lag ring clamps to the pool. Four launches per buffer, the full routing, a smaller skewed one, the full routing again and a decode-sized one, exercise both launch-parity banks and the wrap, and the first launch matches the shared test's reference within tolerance. Where the default layout is a ring, the launches of at most 896 tokens run again on a full-pool buffer (a capacity below the threshold) and must agree bit for bit; where the ring is smaller than the pool, that is the wave order on the pool against the lag order on the ring. A further buffer repeats all four with each call's ownnum_tokens_boundand must match the unbounded arm bit for bit, which covers both schedules a bound can pick (the wave order on the decode-sized call, and a shorter lag on a top-6 shape added for it), that a bound above the capacity is rejected, and that a shape the 256-wide tiles do not divide is rejected at buffer creation. Failure flags are all-reduced, so a single-rank mismatch reports instead of hanging.sgl_deep_gemm/run_tests.shgains the ring test and runs the accuracy pass before the benchmark call.Out of scope. It does not redesign #36's NVLink data path — what gets pulled and scattered, and the per-token arrival flags — or #63's FP8 requant expression, unchanged line for line; the synchronisation around that path does change, and so does the activation pool it writes into. No shared-expert fusion here either.
Relation to other PRs
#69 ports the SM100 interleaved scheduler and fuses the shared expert, and keeps the full pool; on the earlier base, in one session and the same method, this branch is 8–9 % faster at 64–256 tokens/rank and 1.49× at 2048, shared experts off in both. #87 parameterises the SM90 SwiGLU and is orthogonal, but all six of its files are ones this PR also edits, so whichever lands first the other rebases; I'll take that side. #82 and #85 correct full-pool pull sizing on the SM100 path, a separate computation from the SM90 sizing here.
The #69 and #87 details, and the base
The #69 numbers were taken on
devat 60651a7, before #86 moved the SM90 alignment to 1920; #69 keptdev's token alignment there, so a 2048-token request sized 2304 tokens where ours sized 2048. Clocks were level where it matters: 1980 MHz on both branches at 64–256 tokens/rank, and at 2048 both drop but not equally — 1680–1695 MHz for #69 against 1485 for this branch — so the 1.49× is what it is with the clock advantage on #69's side.#74 fixes those four
.templatecall sites among many others; that part is already in this kernel, and nothing else here depends on it.This PR is against
devat 7c4a95f. The tables above were measured on cbe1727:dev's SM90 MegaMoE kernel and host files are byte-identical between the two, and its 18 benchmark kernels rebuild to identical SASS. On this branch, all 18 cells keep the host configuration and the bitwise output of the revision the tables were measured on, and the kernels its tests and production shapes build compile to the same SASS. The tables are those measurements, not a new timing run. The merge of #86 (Public Release 26/09) moved the JIT layer to DeepJIT, so the per-launch PDL request goes throughmake_launch_options(..., enable_pdl)and this PR touches no JIT header of its own.The same kernel is open against the upstream repo as deepseek-ai#440, where it lands as a new kernel instead of replacing an existing fused one. Same kernel, not the same bytes — don't assume a patch to one applies to the other.
Base kernel: #36 by @qiushixiaoyu; the ring counters follow the SM100 design (deepseek-ai#364). The first SM90 ring buffer this builds on and the cluster and multicast paths are credited in the commit trailers too.