[PR #440] Add a fused SM90 FP8 MegaMoE: one kernel, a lag-ordered schedule and a ring activation pool #449

New issue

Open

opened 2026-10-26 21:18:00 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented 2026-10-26 21:18:00 +00:00

Owner

Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/440
Author: @zkyue
Created: 9/13/2026
Status: 🔄 Open

Base: nv_dev ← Head: nvdev-fused-v1

---

📝 Commits (3)

• e19160b Add PTX/MMA helpers, a generic barrier wait and a PDL launch flag
• cbd08d9 Add a fused SM90 FP8 MegaMoE with a lag-ordered ring schedule
• 8cf8180 Add tests for the fused SM90 MegaMoE ring layout

📊 Changes

17 files changed (+6870 additions, -110 deletions)

View changed files

📝 README.md (+31 -0)
➕ csrc/apis/sm90_fused_mega.hpp (+299 -0)
📝 csrc/jit/kernel_runtime.hpp (+10 -5)
➕ csrc/jit_kernels/heuristics/sm90_fused_mega_moe.hpp (+710 -0)
➕ csrc/jit_kernels/impls/sm90_fp8_fused_mega_moe.hpp (+400 -0)
📝 csrc/python_api.cpp (+2 -0)
📝 deep_gemm/__init__.py (+1 -0)
📝 deep_gemm/include/deep_gemm/comm/barrier.cuh (+50 -14)
➕ deep_gemm/include/deep_gemm/impls/sm90_fp8_fused_mega_moe.cuh (+3836 -0)
➕ deep_gemm/include/deep_gemm/layout/sm90_fused_mega_moe.cuh (+294 -0)
📝 deep_gemm/include/deep_gemm/mma/sm90.cuh (+8 -0)
📝 deep_gemm/include/deep_gemm/ptx/ld_st.cuh (+65 -0)
➕ deep_gemm/include/deep_gemm/scheduler/sm90_fused_mega_moe.cuh (+655 -0)
📝 deep_gemm/mega/__init__.py (+113 -3)
📝 tests/bench_mega_moe_sm90.py (+131 -81)
📝 tests/test_mega_moe_sm90.py (+32 -7)
➕ tests/test_mega_moe_sm90_fused_ring.py (+233 -0)

📄 Description

Co-authored with @iclementine (the first SM90 ring buffer) and @Hyaloid (cluster and multicast).

This adds a single-kernel SM90 FP8 MegaMoE next to the two-kernel path from #383. One task stream runs both GEMMs with L2 trailing L1 by a fixed lag, and the L1 activation pool becomes a ring sized around that lag instead of the token bound. On 8×H100 EP8 that is 1.44–1.56× at 2048 tokens/rank and 1.33–1.43× at 8192 against the summed device time of #383's two kernels, on 4.3–4.6× less symmetric memory. It is opt-in and additive.

Performance

One 8×H100 node, EP8, this repo's own harness as this PR leaves it. Both arms size the symmetric buffer from the harness's default 8192-token bound rather than from the batch, so every row of a table runs on the same buffer: 3023 → 700 MiB on flash and 5083 → 1113 MiB on pro.

--model-config flash — 4096 / 2048 / 256 / top-6

tokens/rank	split (µs)	fused (µs)	fused/split	speedup
64	369	352	0.968 (0.923–0.975)	1.03×
128	371	362	0.977 (0.944–1.005)	1.02×
256	424	381	0.898 (0.868–1.029)	1.11×
512	527	431	0.817 (0.814–0.820)	1.22×
1024	737	482	0.646 (0.632–0.665)	1.55×
2048	1188	764	0.642 (0.633–0.648)	1.56×
4096	2073	1376	0.664 (0.661–0.676)	1.51×
8192	3703	2578	0.697 (0.690–0.700)	1.43×

--model-config pro — 7168 / 3072 / 384 / top-6

tokens/rank	split (µs)	fused (µs)	fused/split	speedup
64	1186	1120	0.947 (0.945–0.947)	1.06×
128	1175	1134	0.965 (0.963–0.969)	1.04×
256	1217	1173	0.972 (0.964–0.973)	1.03×
512	1303	1238	0.944 (0.936–0.961)	1.06×
1024	1550	1400	0.903 (0.883–0.932)	1.11×
2048	2811	1955	0.695 (0.693–0.706)	1.44×
4096	4854	3295	0.681 (0.671–0.681)	1.47×
8192	8623	6503	0.755 (0.754–0.756)	1.33×

Each cell is bench_kineto device time over that arm's kernels, median of three bracketed rounds. The split arm is the run's own drift oracle: 0.5–6.0 % round-to-round across the sixteen cells. Read every cell at or above 0.93 as "no worse than #383 inside that drift"; the flash rows from 512 up and the pro rows from 2048 up are 22–56 % apart. Both configurations in tests/bench_mega_moe_sm90.py are covered: fused/split is 0.642–0.977 over 64–8192 tokens/rank.

pro at 1024 sits between the two regimes. Each rank streams 3.2 GB of expert weights per call on that shape against 0.8 GB on flash, and at 128 tokens per expert both arms are still within 1.5–1.7× of the HBM floor for those bytes, so there is little exposed L1 → L2 serialisation left for the fusion to hide; flash at 1024 runs 192 tokens per expert and its split arm is 3× its floor. At 2048 tokens (256 per expert) pro leaves the floor and the ratio drops to 0.695.

Symmetric buffer per rank, by the bound it is sized for (--num-max-tokens-per-rank; 64 aligns up to 128 on both arms)

bound	flash split → fused	pro split → fused
128	86 → 72 MiB	175 → 141 MiB
256	132 → 117	253 → 217
512	226 → 209	409 → 371
1024	412 → 179	721 → 386
2048	785 → 240	1344 → 490
4096	1531 → 361	2590 → 698
8192	3023 → 700	5083 → 1113

These are the two size queries the buffers call, which is also where the harness reads the 8192 figures. Below the 1024-token lag threshold the fused buffer keeps the full pool and saves 7–20 %; from 1024 the ring takes over, and the saving grows with the bound because only the L1/L2 activation pools come from the ring while the dispatch inputs and the combine buffer still scale with it.

Benchmark method, the noisy cell, and a third shape

8×H100 80 GB HBM3, EP8, driver 535.230.02, CUDA 13.3 (V13.3.73), PyTorch 2.13.0a0. One node, all eight GPUs held for the whole run with no other compute process on them. This repo's own harness, as this PR leaves it:

python tests/bench_mega_moe_sm90.py --model-config {flash,pro} --batches <M> \
    --num-max-tokens-per-rank 8192 --arms split fused --num-processes 8

Each number is bench_kineto device time summed over that arm's kernels (sm90_fp8_mega_moe_l1_impl + sm90_fp8_mega_moe_l2_impl for split, sm90_fp8_fused_mega_moe_impl for fused), with the harness's 8 GB L2 flush, ~10 ms sleep and dist.barrier before each call — a burst duty cycle, not a steady state. 20 calls per observation, the harness's own three observations per round above 128 tokens/rank and 50 at or below it, max over the eight ranks, then median over observations. Three rounds with the arms bracketed — split fused, then fused split, then split fused — so each cell is the median of three round medians, and the band in parentheses is the smallest and largest of the three per-round ratios. Lag units, the ring they size, and the single-wave order are defined under How it works.

The fused arm passes num_tokens_bound=M on every call, which is what a caller whose ranks all run the same count would pass; the split path has no such argument and needs none, since none of its heuristics key on a rank-uniform count. That keyword is what makes the 64 / 128 / 256 rows take the single-wave order here, and on flash it also runs the 512 to 4096 rows at a 14-unit lag instead of the 29 its ring is sized for (DG_PRINT_CONFIGS reports l2_lag_units 4014 against 4029); pro's buffer lag is 22 at this capacity and its bounded calls keep it.

The one cell where the fused arm is the noisier of the two is flash at 256 tokens/rank, 12 % from a single high round against the split arm's 5.6 % — that row's band (0.868–1.029) is spread and not a stable difference, and no other cell's fused spread is above 4 %.

Every cell runs the default intermediate-scale granularity, 128 on both of these shapes. The fused arm asks for PDL at ≤ 256 tokens/rank while #383's kernels carry no griddepcontrol at all, so on the 64 / 128 / 256 rows the fused Kineto duration contains its own griddepcontrol.wait — part of why those rows show the smallest margins.

A third shape — 4096 / 2048 with 288 experts and top-8, i.e. 36 experts per rank instead of 32 or 48 — was taken once rather than in rounds, at the same 8192-token bound: fused/split is 0.950, 0.968, 0.955, 0.739, 0.613, 0.607, 0.634 and 0.677 at 64, 128, 256, 512, 1024, 2048, 4096 and 8192 tokens per rank, and the symmetric buffer goes from 4000 to 857 MiB. Read the first three rows as "no worse" given the drift above; the rest behaves like the flash table.

How it works

The kernel fixes the order first. The B loader enumerates a stream in which a token block's L2 tasks sit lag units after its L1 tasks, in groups of group units; a unit is up to 8 consecutive m-blocks of one expert. The host picks lag = clamp(round(0.45 × units), 8, 64) and group = 4, with units estimated from the call's token bound and the buffer capacity, not the routed count.

It also differs from #383 numerically: one scale per 64 or 128 contiguous K elements of the intermediate activation (l2_act_sf_gran_k, a buffer parameter) where #383 always writes one per 64, and the raw scale where #383's SM90 kernel rounds it to the next power of two. Both are inside the shared test's tolerance against the same reference.

Why the order and not the pool size, and the lag ladder

The split path materialises every rank's L1 output before L2 starts, so its pool is sized for all tokens of all ranks and the GEMMs never overlap. Fusing them buys the overlap and does not fix the pool, because nothing bounds how far L2 trails L1. Worse, L2's NVLink combine scatter starts while the dispatch pull is still running on the same per-rank egress. A ladder that pins the ring at 356 blocks and varies only the lag still leaves the short-lag calls the slow ones, so what costs is the order and not the pool size.

Event traces on an instrumented build of this kernel (8×H100, 8192 tokens/rank) put the end of the dispatch pull at 2671 µs from kernel start when the lag is 8 units and at 1102 µs when it is 40 — 2541 µs against 1044 µs measured from each run's own head barrier, which completes at 130 µs and at 58 µs. At the lag this picks, 55 % of the L1 tiles are done when the pull ends.

Ring sizing, recycling, and why that granularity

Fixing the order makes the ring size arithmetic: at most 8 × (lag + group) blocks of L1 output are live, so the pool becomes a ring of 8 × (lag + group) + 4 slots, asserted against that live set on the host. L1 epilogues wait on l2_empty before overwriting a slot's previous generation, and dispatch warps publish any pending arrival before blocking, so no wait can point at unpublished work.

kSm90FusedLagRingMargin is the + 4. Recycling reuses the l1/l2_full counters and adds the matching empty counters plus launch-parity banks for the flags that would otherwise need all-rank head and exit barriers. A ring clamped to the full pool never wraps and is exempt from the host assertion. Consumers replay the tile table the B loader publishes — payload first, then an acquire-loaded tag — and the L2 tail after the last L1 task is handed out by an atomic ticket. Calls whose worst-case pool already fits, and buffers below the lag threshold, keep a single-wave order instead.

Also in the diff: arrival publishes batched up to 16 rows, on calls above a token threshold whose 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 still pending; a BLOCK_M=64, 256-wide split-N tile with a sliced combine queue for decode-sized calls; num_tokens_bound; and an XOR bank swizzle on the L2 epilogue staging. The early combine stages whole rows in the pre-barrier shared memory where they fit and in chunks otherwise, and falls back from 2 chunks to 4 when 2 do not fit — which is what keeps 7168 hidden working at decode sizes, where the tile leaves that region smallest.

The scale width defaults to 128 when hidden % 512 == 0 and 2 × intermediate_hidden % 512 == 0 (the condition for the 256-wide decode tile) and to 64 otherwise. 128 is the grouping the non-fused DeepEP + grouped-GEMM path uses for this activation; it 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. #360 reports an end-to-end accuracy issue it suspects came from quantising this same activation per-64 where the rest of the stack does per-128.

The scale itself is sf = max(amax, 1e-10) / 448, the raw scale, which is also what the shared test's reference computes (with a 1e-4 floor on amax instead of 1e-10) and what the sgl-project SM90 kernel this data path comes from uses. The fused arm keeps the full FP8 range, the split one keeps an exactly representable scale.

What to check before merging

Two shared files are rewritten rather than added to, and they are the part of the diff to read hardest. In comm/barrier.cuh (+50/−14) the two waits become generic over the workspace type and, under BarrierTimeoutPolicy::TrapOnly, call handle_*_timeout after leaving the wait loop rather than from inside it. The fused kernel needs that shape: a trap; inside the loop makes ptxas allocate that region's registers against the kernel's launch bound and ignore its setmaxnreg.inc, which spills.

csrc/jit/kernel_runtime.hpp (+10/−5) is the other: a launcher can now force the PDL launch attribute instead of leaving it to the global flag, and no existing launcher passes the new argument.

Register inventory, and the SM100 / #383 codegen check

The 60 s timeout and the handler are unchanged, both policies still abort, and the Diagnostic branch is base's loop verbatim. #383's split SM90 kernel does take TrapOnly, at four sites, and keeps identical REG/STACK on all 25 instantiations while its SASS shifts.

Isolated TU, sm_90a, --register-usage-level=10. The SM100 comparison is a kernel that calls both waits with the default Diagnostic policy: identical SASS and identical REG/STACK against this branch and against nv_dev, so the SM100 callers' codegen is unchanged.

For #383's split SM90 kernel, recompiling all 25 instantiations the two SM90 test files and the benchmark produce gives identical REG / STACK / SHARED / LOCAL on every one of them: 23 at REG:168 STACK:0 SHARED:1024 LOCAL:0, 2 at REG:96 STACK:24 SHARED:1024 LOCAL:0. The SASS text does move, at the same instruction count, with a few ISETP variants and the branch targets shifted around the deadline compare. Its own accuracy test passes unchanged (see Tests).

The PDL change only adds the launch attribute, never removes it, and it applies to a kernel compiled with griddepcontrol, which the fused path is and #383's kernels are not.

num_tokens_bound needs a warning label. It must be identical on every rank and cover every rank's own count, e.g. the max over ranks of the local batch. Each rank rejects a bound below its own count or above the buffer's capacity, but nothing compares the ranks, and 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. None means the buffer capacity, and a bound below it is what lets a large buffer take the single-wave order.

What a per-rank token count changes, and what it does not

The rules that key on a per-rank token count fall back to this rank's own count when the bound is omitted, so ranks with different local batches can pick different orders. Nothing cross-rank rides on that: the ring capacity is the buffer's, and the call's lag comes from the bound -- the capacity when none is given -- clamped to the buffer's, so both hold the same value on every rank.

One rule keys on the token count in a way worth spelling out. block_m is chosen from this rank's token count (num_tokens * num_ranks * num_topk / num_experts >= 64 flips it 64 -> 128), block_m sets num_stages through the per-stage smem cost, and early combine turns on only when hidden / block_k and intermediate_hidden / block_k both exceed num_stages. So it is a function of the shapes only where the shape leaves margin, and passing num_tokens_bound does not make it uniform -- the tile choice reads the raw count. On the shapes here the margin is wide: 4096/2048 gives 16 k-blocks and 7168/3072 gives 24, against 3 stages at both M=2048 and M=8192, and the mode never flips. I had to take intermediate_hidden down to 640, which no shipped config uses, before the flip appeared -- 287 tokens gives 7 stages and mode 0, 288 gives 4 stages and mode 1. What does move per rank on production shapes is block_m, block_n and num_stages, and since those are JIT template arguments like the mode itself, ranks with different local batches compile and launch different instantiations of the same kernel. They still meet at the same pre-combine rendezvous: same signal word, same phase and sign rule, same expected arrival count, one signal per rank either way. If the modes ever did disagree the cost would be throughput, not liveness -- a rank on the early-combine path only early-combines tokens whose experts all live on ranks that took the same path, and the rest fall through to the ordinary combine after the barrier. One gap worth noting on its own: the early-combine copy of the barrier keeps the 60 s deadline on the cross-rank signal wait but not on its two grid syncs, which the shared grid_sync does guard.

The shape domain is a strict subset. intermediate_hidden ≤ 4096, rejected when the buffer is sized, where #383's split path accepts 4224. No shared-expert fusion either, same as #383's SM90 path.

No CI coverage. Neither SM90 MegaMoE test runs in CI, since both need 8 GPUs and NVLink. Everything under Tests below was run by hand on one 8×H100 node, sm_90a, CUDA 13.3.

fused=True at buffer allocation (--arms fused in the benchmark harness, --fused in the accuracy test) selects it; #383's kernel, heuristics and launcher are untouched, and its two Python entry points only gain keywords that default to today's behaviour. The two arms are separate buffers, JIT instantiations and host paths. The shape subset, the quantisation difference and the missing CI are why this is opt-in rather than the default. A default needs a fallback for the shapes it cannot serve, and flipping it would change existing callers' outputs bit for bit without them asking. Say so and I'll flip it: the rounding is one line, the fallback a shape check in the buffer query.

API, Tests, Files

The buffer query takes three new keywords and fp8_mega_moe takes one; of the three, only fused matters to a production caller.

Full API surface

Opting in costs one keyword at allocation and nothing per call: fp8_mega_moe dispatches on the buffer type.

get_symm_buffer_for_sm90_mega_moe(..., fused=False, num_experts_per_wave=None, l2_act_sf_gran_k=None) and fp8_mega_moe(..., num_tokens_bound=None). At the defaults both calls do what they do today; with fused=True the buffer query returns an SM90FusedSymmBuffer and fp8_mega_moe dispatches on its type.

num_experts_per_wave: None = ring from the lag rule (full pool below the threshold), -1 = ring sized for an automatically chosen expert wave, N = capacity for a wave of min(N, E_local) experts, never more than the full pool. It is how the ring test builds its auto-wave and full-pool arms.

l2_act_sf_gran_k pins 64 or 128 instead of taking the shape rule; pinning 128 where the rule says 64 is rejected at launch on decode-sized calls.

get_symm_buffer_size_for_sm90_fused_mega_moe returns (bytes, slicer, num_ring_tokens, l2_lag_encoded), and the ring capacity and lag encoding travel back to the kernel on every launch together with the scale granularity. Token alignment is 128 on both arms.

The rules that fall back to this rank's own count when num_tokens_bound is omitted are whether the call takes the single-wave order instead of the lag order, the arrival-publish batch and its eager variant, and whether a decode-sized call takes PDL. The single-wave order additionally needs the worst-case pool at the call's bound -- the buffer's capacity when none is given -- to fit the ring.

Tests: the accuracy runs, the bound replay, and the eight ring shapes

tests/test_mega_moe_sm90.py --fused runs the file's 39 scenarios at both scale granularities, minus the 4224 intermediate_hidden one the fused path skips: 76 cases, all pass on 8 ranks, largest deviation 1e-4 against the 0.01 tolerance. The same file's split arm still passes unchanged on this branch: 39 scenarios, largest deviation 7e-4 against the same 0.01 tolerance.

tests/test_mega_moe_sm90_fused_ring.py is new: it checks the lag ring bitwise against a full-pool control on eight shapes, four launches per buffer with skewed routing in between, exercising both launch-parity banks and the wrap. A fourth buffer replays those launches under each call's own num_tokens_bound, bit for bit again, then under a bound one alignment above capacity, which has to raise on the host before any rank launches. Sixteen cases pass, at most 1e-4 from the reference against the test's 0.07 tolerance.

The ring test takes ~2 min on 8 GPUs per granularity, and the sixteen cases are its eight shapes at both scale granularities. The shapes are the shared test's hidden=1024 shape at 64, 256, 512, 1024 and 2048 tokens per rank, which straddles the 1024-token lag threshold from both sides; a 4096/2048/288/top-8 shape whose lag ring is 20992 of 135680 pool tokens, next to a 98304-token auto-wave ring; a 1024/1024/64/top-6 shape at a 5120-token capacity, 18944 ring tokens of a 246784-token pool next to an 81920-token auto-wave ring; and a 2-rank subgroup whose requested ring is larger than the pool and clamps to it. The verdict is all-reduced, so a single-rank mismatch reports instead of hanging.

The replay is where both things a bound can change get covered. The decode-sized launch is bounded at 256 tokens, so on the lag rings that one call runs the single-wave order against the same allocation running the lag order without a bound (DG_PRINT_CONFIGS reports l2_lag_units 4008, 4014 or 4016 unbounded and 0 under the bound); and on the 5120-token shape the 3840-token launch runs an 11-unit lag under its own bound where the unbounded allocation runs the 14 its ring is sized for (4011 against 4014). Every one of the sixteen cases reports bitwise=OK replay=OK bounds=...:OK over_cap_rejected=OK.

Both tests spawn one process per GPU themselves, so they need 8 SM90 GPUs and a working NVLink group.

Files, and the naming census

New: deep_gemm/include/deep_gemm/impls/sm90_fp8_fused_mega_moe.cuh (the kernel), layout/sm90_fused_mega_moe.cuh (workspace and tile-table sizing), scheduler/sm90_fused_mega_moe.cuh (the lag scheduler), csrc/apis/sm90_fused_mega.hpp, csrc/jit_kernels/heuristics/sm90_fused_mega_moe.hpp, csrc/jit_kernels/impls/sm90_fp8_fused_mega_moe.hpp, tests/test_mega_moe_sm90_fused_ring.py.

Changed outside the fused path: 13 wrappers in ptx/ld_st.cuh, advance_smem_desc in mma/sm90.cuh, one registration line each in csrc/python_api.cpp and deep_gemm/__init__.py, the fused branch in deep_gemm/mega/__init__.py, a subsection in README.md, a fused arm in tests/bench_mega_moe_sm90.py — which is restructured into arms rather than added to — and tests/test_mega_moe_sm90.py, where the reference now quantises the intermediate activation at the buffer's granularity, on the --fused path only; the split path's reference is unchanged. Plus the two shared files above, comm/barrier.cuh and csrc/jit/kernel_runtime.hpp.

Of the 61 namespace-scope names the new fused files introduce, 57 carry a fused stem (MegaMoESM90FusedConfig, SM90FP8FusedMegaMoERuntime, get_mega_moe_config_sm90_fused, kSm90Fused*, and on the device SM90FusedWorkspace / SM90FusedMegaMoEScheduler). The four that do not are new sizing helpers in layout/sm90_fused_mega_moe.cuh (kNumSM90CandidateBlockMs, get_num_max_pool_tokens_sm90, get_sm90_tile_table_entries, get_sm90_tile_table_entries_compact); none of the four exists anywhere in the tree today, so nothing collides with #383's SM90 names and DG_PRINT_CONFIGS tells the two apart.

This PR adds no tuning environment variable. The only new reads are the two trace instruments DG_SM90_TRACE_PTR and DG_SM90_TRACE_EPI, next to the repo's existing diagnostics (DG_JIT_DEBUG / DG_PRINT_CONFIGS, DG_COMM_KERNEL_DEBUG), which the fused path honours exactly where the split path does.

Relation to other work

#383's review named a fused SM90 kernel as follow-up work; this is one. It does not replace that PR: the split path stays the default.

Overlaps with #422, #411, #393 and #360, and the same kernel against sgl-project

Overlaps to know about:

• #422 brings the sgl-project split SM90 kernel into the paths #383 now occupies. It makes the same two barriers generic over the workspace type, with WorkspaceT after sync_scope_t rather than before, and takes trap-only on SM90 from inside the wait loop, the shape that costs this kernel its setmaxnreg. Whichever lands first, the policy parameter has to go in.
• #411 adds SM90 FP8 × MXFP4 on top of #383's kernel, with the same ld_shared(const uint2*) overload this PR adds, so whichever lands second drops its copy.
• #393 rewrites csrc/apis/*.hpp for TORCH_LIBRARY, including the registration this PR adds; whichever lands first, the other rebases three m.defs.
• #360 is a different fused SM90 design against main, one cooperative kernel whose two math warpgroups N-split a 64×256 tile. It predates #383 and is the closest alternative to this one.

The same kernel is open against sgl-project/DeepGEMM as sgl-project/DeepGEMM#88, where it replaces that fork's existing fused SM90 kernel instead of adding one. Same kernel, not the same bytes: the barriers here take a per-instantiation timeout policy where the fork gates the same trap on the architecture, and this workspace's flag bank has no publish-epoch word. Nothing reads a bank across the two repositories. Don't assume a patch to one applies to the other.

The NVLink data path this kernel keeps comes from sgl-project/DeepGEMM#36. That, the first SM90 ring buffer this builds on, and the cluster and multicast paths are all credited in the commit trailers.

---

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/440 **Author:** [@zkyue](https://github.com/zkyue) **Created:** 9/13/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `nvdev-fused-v1` --- ### 📝 Commits (3) - [`e19160b`](https://github.com/deepseek-ai/DeepGEMM/commit/e19160bfe7e5ffdd210cc98a0dab260bad4c5dc8) Add PTX/MMA helpers, a generic barrier wait and a PDL launch flag - [`cbd08d9`](https://github.com/deepseek-ai/DeepGEMM/commit/cbd08d9f2b1d04df75cb50f2ca5cf62ae02f2309) Add a fused SM90 FP8 MegaMoE with a lag-ordered ring schedule - [`8cf8180`](https://github.com/deepseek-ai/DeepGEMM/commit/8cf8180fe6f97eb65ee6638a14a21497cc1cb5e0) Add tests for the fused SM90 MegaMoE ring layout ### 📊 Changes **17 files changed** (+6870 additions, -110 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+31 -0) ➕ `csrc/apis/sm90_fused_mega.hpp` (+299 -0) 📝 `csrc/jit/kernel_runtime.hpp` (+10 -5) ➕ `csrc/jit_kernels/heuristics/sm90_fused_mega_moe.hpp` (+710 -0) ➕ `csrc/jit_kernels/impls/sm90_fp8_fused_mega_moe.hpp` (+400 -0) 📝 `csrc/python_api.cpp` (+2 -0) 📝 `deep_gemm/__init__.py` (+1 -0) 📝 `deep_gemm/include/deep_gemm/comm/barrier.cuh` (+50 -14) ➕ `deep_gemm/include/deep_gemm/impls/sm90_fp8_fused_mega_moe.cuh` (+3836 -0) ➕ `deep_gemm/include/deep_gemm/layout/sm90_fused_mega_moe.cuh` (+294 -0) 📝 `deep_gemm/include/deep_gemm/mma/sm90.cuh` (+8 -0) 📝 `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` (+65 -0) ➕ `deep_gemm/include/deep_gemm/scheduler/sm90_fused_mega_moe.cuh` (+655 -0) 📝 `deep_gemm/mega/__init__.py` (+113 -3) 📝 `tests/bench_mega_moe_sm90.py` (+131 -81) 📝 `tests/test_mega_moe_sm90.py` (+32 -7) ➕ `tests/test_mega_moe_sm90_fused_ring.py` (+233 -0) </details> ### 📄 Description *Co-authored with @iclementine (the first SM90 ring buffer) and @Hyaloid (cluster and multicast).* This adds a single-kernel SM90 FP8 MegaMoE next to the two-kernel path from #383. One task stream runs both GEMMs with L2 trailing L1 by a fixed lag, and the L1 activation pool becomes a ring sized around that lag instead of the token bound. On 8×H100 EP8 that is **1.44–1.56× at 2048 tokens/rank and 1.33–1.43× at 8192** against the summed device time of #383's two kernels, on 4.3–4.6× less symmetric memory. It is opt-in and additive. ## Performance One 8×H100 node, EP8, this repo's own harness as this PR leaves it. Both arms size the symmetric buffer from the harness's default 8192-token bound rather than from the batch, so every row of a table runs on the same buffer: 3023 → 700 MiB on flash and 5083 → 1113 MiB on pro. **`--model-config flash` — 4096 / 2048 / 256 / top-6** | tokens/rank | split (µs) | fused (µs) | fused/split | speedup | |---:|---:|---:|:---:|:---:| | 64 | 369 | 352 | 0.968 (0.923–0.975) | 1.03× | | 128 | 371 | 362 | 0.977 (0.944–1.005) | 1.02× | | 256 | 424 | 381 | 0.898 (0.868–1.029) | 1.11× | | 512 | 527 | 431 | 0.817 (0.814–0.820) | 1.22× | | 1024 | 737 | 482 | 0.646 (0.632–0.665) | 1.55× | | 2048 | 1188 | 764 | 0.642 (0.633–0.648) | 1.56× | | 4096 | 2073 | 1376 | 0.664 (0.661–0.676) | 1.51× | | 8192 | 3703 | 2578 | 0.697 (0.690–0.700) | 1.43× | **`--model-config pro` — 7168 / 3072 / 384 / top-6** | tokens/rank | split (µs) | fused (µs) | fused/split | speedup | |---:|---:|---:|:---:|:---:| | 64 | 1186 | 1120 | 0.947 (0.945–0.947) | 1.06× | | 128 | 1175 | 1134 | 0.965 (0.963–0.969) | 1.04× | | 256 | 1217 | 1173 | 0.972 (0.964–0.973) | 1.03× | | 512 | 1303 | 1238 | 0.944 (0.936–0.961) | 1.06× | | 1024 | 1550 | 1400 | 0.903 (0.883–0.932) | 1.11× | | 2048 | 2811 | 1955 | 0.695 (0.693–0.706) | 1.44× | | 4096 | 4854 | 3295 | 0.681 (0.671–0.681) | 1.47× | | 8192 | 8623 | 6503 | 0.755 (0.754–0.756) | 1.33× | Each cell is `bench_kineto` device time over that arm's kernels, median of three bracketed rounds. The split arm is the run's own drift oracle: 0.5–6.0 % round-to-round across the sixteen cells. Read every cell at or above 0.93 as "no worse than #383 inside that drift"; the flash rows from 512 up and the pro rows from 2048 up are 22–56 % apart. Both configurations in `tests/bench_mega_moe_sm90.py` are covered: fused/split is 0.642–0.977 over 64–8192 tokens/rank. pro at 1024 sits between the two regimes. Each rank streams 3.2 GB of expert weights per call on that shape against 0.8 GB on flash, and at 128 tokens per expert both arms are still within 1.5–1.7× of the HBM floor for those bytes, so there is little exposed L1 → L2 serialisation left for the fusion to hide; flash at 1024 runs 192 tokens per expert and its split arm is 3× its floor. At 2048 tokens (256 per expert) pro leaves the floor and the ratio drops to 0.695. **Symmetric buffer per rank, by the bound it is sized for** (`--num-max-tokens-per-rank`; 64 aligns up to 128 on both arms) | bound | flash split → fused | pro split → fused | |---:|:---|:---| | 128 | 86 → 72 MiB | 175 → 141 MiB | | 256 | 132 → 117 | 253 → 217 | | 512 | 226 → 209 | 409 → 371 | | 1024 | 412 → 179 | 721 → 386 | | 2048 | 785 → 240 | 1344 → 490 | | 4096 | 1531 → 361 | 2590 → 698 | | 8192 | 3023 → 700 | 5083 → 1113 | These are the two size queries the buffers call, which is also where the harness reads the 8192 figures. Below the 1024-token lag threshold the fused buffer keeps the full pool and saves 7–20 %; from 1024 the ring takes over, and the saving grows with the bound because only the L1/L2 activation pools come from the ring while the dispatch inputs and the combine buffer still scale with it. <details> <summary>Benchmark method, the noisy cell, and a third shape</summary> 8×H100 80 GB HBM3, EP8, driver 535.230.02, CUDA 13.3 (V13.3.73), PyTorch 2.13.0a0. One node, all eight GPUs held for the whole run with no other compute process on them. This repo's own harness, as this PR leaves it: ``` python tests/bench_mega_moe_sm90.py --model-config {flash,pro} --batches <M> \ --num-max-tokens-per-rank 8192 --arms split fused --num-processes 8 ``` Each number is `bench_kineto` device time summed over that arm's kernels (`sm90_fp8_mega_moe_l1_impl` + `sm90_fp8_mega_moe_l2_impl` for split, `sm90_fp8_fused_mega_moe_impl` for fused), with the harness's 8 GB L2 flush, ~10 ms sleep and `dist.barrier` before each call — a burst duty cycle, not a steady state. 20 calls per observation, the harness's own three observations per round above 128 tokens/rank and 50 at or below it, max over the eight ranks, then median over observations. Three rounds with the arms bracketed — `split fused`, then `fused split`, then `split fused` — so each cell is the median of three round medians, and the band in parentheses is the smallest and largest of the three per-round ratios. Lag units, the ring they size, and the single-wave order are defined under How it works. The fused arm passes `num_tokens_bound=M` on every call, which is what a caller whose ranks all run the same count would pass; the split path has no such argument and needs none, since none of its heuristics key on a rank-uniform count. That keyword is what makes the 64 / 128 / 256 rows take the single-wave order here, and on flash it also runs the 512 to 4096 rows at a 14-unit lag instead of the 29 its ring is sized for (`DG_PRINT_CONFIGS` reports `l2_lag_units` 4014 against 4029); pro's buffer lag is 22 at this capacity and its bounded calls keep it. The one cell where the fused arm is the noisier of the two is flash at 256 tokens/rank, 12 % from a single high round against the split arm's 5.6 % — that row's band (0.868–1.029) is spread and not a stable difference, and no other cell's fused spread is above 4 %. Every cell runs the default intermediate-scale granularity, 128 on both of these shapes. The fused arm asks for PDL at ≤ 256 tokens/rank while #383's kernels carry no griddepcontrol at all, so on the 64 / 128 / 256 rows the fused Kineto duration contains its own `griddepcontrol.wait` — part of why those rows show the smallest margins. A third shape — 4096 / 2048 with 288 experts and top-8, i.e. 36 experts per rank instead of 32 or 48 — was taken once rather than in rounds, at the same 8192-token bound: fused/split is 0.950, 0.968, 0.955, 0.739, 0.613, 0.607, 0.634 and 0.677 at 64, 128, 256, 512, 1024, 2048, 4096 and 8192 tokens per rank, and the symmetric buffer goes from 4000 to 857 MiB. Read the first three rows as "no worse" given the drift above; the rest behaves like the flash table. </details> ## How it works The kernel fixes the order first. The B loader enumerates a stream in which a token block's L2 tasks sit `lag` units after its L1 tasks, in groups of `group` units; a unit is up to 8 consecutive m-blocks of one expert. The host picks `lag = clamp(round(0.45 × units), 8, 64)` and `group = 4`, with `units` estimated from the call's token bound and the buffer capacity, not the routed count. It also differs from #383 numerically: one scale per 64 or 128 contiguous K elements of the intermediate activation (`l2_act_sf_gran_k`, a buffer parameter) where #383 always writes one per 64, and the raw scale where #383's SM90 kernel rounds it to the next power of two. Both are inside the shared test's tolerance against the same reference. <details> <summary>Why the order and not the pool size, and the lag ladder</summary> The split path materialises every rank's L1 output before L2 starts, so its pool is sized for all tokens of all ranks and the GEMMs never overlap. Fusing them buys the overlap and does not fix the pool, because nothing bounds how far L2 trails L1. Worse, L2's NVLink combine scatter starts while the dispatch pull is still running on the same per-rank egress. A ladder that pins the ring at 356 blocks and varies only the lag still leaves the short-lag calls the slow ones, so what costs is the order and not the pool size. Event traces on an instrumented build of this kernel (8×H100, 8192 tokens/rank) put the end of the dispatch pull at 2671 µs from kernel start when the lag is 8 units and at 1102 µs when it is 40 — 2541 µs against 1044 µs measured from each run's own head barrier, which completes at 130 µs and at 58 µs. At the lag this picks, 55 % of the L1 tiles are done when the pull ends. </details> <details> <summary>Ring sizing, recycling, and why that granularity</summary> Fixing the order makes the ring size arithmetic: at most `8 × (lag + group)` blocks of L1 output are live, so the pool becomes a ring of `8 × (lag + group) + 4` slots, asserted against that live set on the host. L1 epilogues wait on `l2_empty` before overwriting a slot's previous generation, and dispatch warps publish any pending arrival before blocking, so no wait can point at unpublished work. `kSm90FusedLagRingMargin` is the `+ 4`. Recycling reuses the `l1/l2_full` counters and adds the matching `empty` counters plus launch-parity banks for the flags that would otherwise need all-rank head and exit barriers. A ring clamped to the full pool never wraps and is exempt from the host assertion. Consumers replay the tile table the B loader publishes — payload first, then an acquire-loaded tag — and the L2 tail after the last L1 task is handed out by an atomic ticket. Calls whose worst-case pool already fits, and buffers below the lag threshold, keep a single-wave order instead. Also in the diff: arrival publishes batched up to 16 rows, on calls above a token threshold whose 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 still pending; a BLOCK_M=64, 256-wide split-N tile with a sliced combine queue for decode-sized calls; `num_tokens_bound`; and an XOR bank swizzle on the L2 epilogue staging. The early combine stages whole rows in the pre-barrier shared memory where they fit and in chunks otherwise, and falls back from 2 chunks to 4 when 2 do not fit — which is what keeps 7168 hidden working at decode sizes, where the tile leaves that region smallest. The scale width defaults to 128 when `hidden % 512 == 0` and `2 × intermediate_hidden % 512 == 0` (the condition for the 256-wide decode tile) and to 64 otherwise. 128 is the grouping the non-fused DeepEP + grouped-GEMM path uses for this activation; it 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. #360 reports an end-to-end accuracy issue it suspects came from quantising this same activation per-64 where the rest of the stack does per-128. The scale itself is `sf = max(amax, 1e-10) / 448`, the raw scale, which is also what the shared test's reference computes (with a 1e-4 floor on amax instead of 1e-10) and what the sgl-project SM90 kernel this data path comes from uses. The fused arm keeps the full FP8 range, the split one keeps an exactly representable scale. </details> ## What to check before merging **Two shared files are rewritten rather than added to, and they are the part of the diff to read hardest.** In `comm/barrier.cuh` (+50/−14) the two waits become generic over the workspace type and, under `BarrierTimeoutPolicy::TrapOnly`, call `handle_*_timeout` after leaving the wait loop rather than from inside it. The fused kernel needs that shape: a `trap;` inside the loop makes ptxas allocate that region's registers against the kernel's launch bound and ignore its `setmaxnreg.inc`, which spills. `csrc/jit/kernel_runtime.hpp` (+10/−5) is the other: a launcher can now force the PDL launch attribute instead of leaving it to the global flag, and no existing launcher passes the new argument. <details> <summary>Register inventory, and the SM100 / #383 codegen check</summary> The 60 s timeout and the handler are unchanged, both policies still abort, and the `Diagnostic` branch is base's loop verbatim. #383's split SM90 kernel does take `TrapOnly`, at four sites, and keeps identical REG/STACK on all 25 instantiations while its SASS shifts. Isolated TU, `sm_90a`, `--register-usage-level=10`. The SM100 comparison is a kernel that calls both waits with the default `Diagnostic` policy: identical SASS and identical REG/STACK against this branch and against `nv_dev`, so the SM100 callers' codegen is unchanged. For #383's split SM90 kernel, recompiling all 25 instantiations the two SM90 test files and the benchmark produce gives identical REG / STACK / SHARED / LOCAL on every one of them: 23 at `REG:168 STACK:0 SHARED:1024 LOCAL:0`, 2 at `REG:96 STACK:24 SHARED:1024 LOCAL:0`. The SASS text does move, at the same instruction count, with a few `ISETP` variants and the branch targets shifted around the deadline compare. Its own accuracy test passes unchanged (see Tests). The PDL change only adds the launch attribute, never removes it, and it applies to a kernel compiled with griddepcontrol, which the fused path is and #383's kernels are not. </details> **`num_tokens_bound` needs a warning label.** It must be identical on every rank and cover every rank's own count, e.g. the max over ranks of the local batch. Each rank rejects a bound below its own count or above the buffer's capacity, but nothing compares the ranks, and 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. `None` means the buffer capacity, and a bound below it is what lets a large buffer take the single-wave order. <details> <summary>What a per-rank token count changes, and what it does not</summary> The rules that key on a per-rank token count fall back to this rank's own count when the bound is omitted, so ranks with different local batches can pick different orders. Nothing cross-rank rides on that: the ring capacity is the buffer's, and the call's lag comes from the bound -- the capacity when none is given -- clamped to the buffer's, so both hold the same value on every rank. **One rule keys on the token count in a way worth spelling out.** `block_m` is chosen from this rank's token count (`num_tokens * num_ranks * num_topk / num_experts >= 64` flips it 64 -> 128), `block_m` sets `num_stages` through the per-stage smem cost, and early combine turns on only when `hidden / block_k` and `intermediate_hidden / block_k` both exceed `num_stages`. So it is a function of the shapes only where the shape leaves margin, and passing `num_tokens_bound` does not make it uniform -- the tile choice reads the raw count. On the shapes here the margin is wide: 4096/2048 gives 16 k-blocks and 7168/3072 gives 24, against 3 stages at both M=2048 and M=8192, and the mode never flips. I had to take `intermediate_hidden` down to 640, which no shipped config uses, before the flip appeared -- 287 tokens gives 7 stages and mode 0, 288 gives 4 stages and mode 1. What does move per rank on production shapes is `block_m`, `block_n` and `num_stages`, and since those are JIT template arguments like the mode itself, ranks with different local batches compile and launch different instantiations of the same kernel. They still meet at the same pre-combine rendezvous: same signal word, same phase and sign rule, same expected arrival count, one signal per rank either way. If the modes ever did disagree the cost would be throughput, not liveness -- a rank on the early-combine path only early-combines tokens whose experts all live on ranks that took the same path, and the rest fall through to the ordinary combine after the barrier. One gap worth noting on its own: the early-combine copy of the barrier keeps the 60 s deadline on the cross-rank signal wait but not on its two grid syncs, which the shared `grid_sync` does guard. </details> **The shape domain is a strict subset.** `intermediate_hidden ≤ 4096`, rejected when the buffer is sized, where #383's split path accepts 4224. No shared-expert fusion either, same as #383's SM90 path. **No CI coverage.** Neither SM90 MegaMoE test runs in CI, since both need 8 GPUs and NVLink. Everything under Tests below was run by hand on one 8×H100 node, sm_90a, CUDA 13.3. `fused=True` at buffer allocation (`--arms fused` in the benchmark harness, `--fused` in the accuracy test) selects it; #383's kernel, heuristics and launcher are untouched, and its two Python entry points only gain keywords that default to today's behaviour. The two arms are separate buffers, JIT instantiations and host paths. The shape subset, the quantisation difference and the missing CI are why this is opt-in rather than the default. A default needs a fallback for the shapes it cannot serve, and flipping it would change existing callers' outputs bit for bit without them asking. Say so and I'll flip it: the rounding is one line, the fallback a shape check in the buffer query. ## API, Tests, Files The buffer query takes three new keywords and `fp8_mega_moe` takes one; of the three, only `fused` matters to a production caller. <details> <summary>Full API surface</summary> Opting in costs one keyword at allocation and nothing per call: `fp8_mega_moe` dispatches on the buffer type. `get_symm_buffer_for_sm90_mega_moe(..., fused=False, num_experts_per_wave=None, l2_act_sf_gran_k=None)` and `fp8_mega_moe(..., num_tokens_bound=None)`. At the defaults both calls do what they do today; with `fused=True` the buffer query returns an `SM90FusedSymmBuffer` and `fp8_mega_moe` dispatches on its type. `num_experts_per_wave`: `None` = ring from the lag rule (full pool below the threshold), `-1` = ring sized for an automatically chosen expert wave, `N` = capacity for a wave of `min(N, E_local)` experts, never more than the full pool. It is how the ring test builds its auto-wave and full-pool arms. `l2_act_sf_gran_k` pins 64 or 128 instead of taking the shape rule; pinning 128 where the rule says 64 is rejected at launch on decode-sized calls. `get_symm_buffer_size_for_sm90_fused_mega_moe` returns `(bytes, slicer, num_ring_tokens, l2_lag_encoded)`, and the ring capacity and lag encoding travel back to the kernel on every launch together with the scale granularity. Token alignment is 128 on both arms. The rules that fall back to this rank's own count when `num_tokens_bound` is omitted are whether the call takes the single-wave order instead of the lag order, the arrival-publish batch and its eager variant, and whether a decode-sized call takes PDL. The single-wave order additionally needs the worst-case pool at the call's bound -- the buffer's capacity when none is given -- to fit the ring. </details> <details> <summary>Tests: the accuracy runs, the bound replay, and the eight ring shapes</summary> `tests/test_mega_moe_sm90.py --fused` runs the file's 39 scenarios at both scale granularities, minus the 4224 `intermediate_hidden` one the fused path skips: 76 cases, all pass on 8 ranks, largest deviation 1e-4 against the 0.01 tolerance. The same file's split arm still passes unchanged on this branch: 39 scenarios, largest deviation 7e-4 against the same 0.01 tolerance. `tests/test_mega_moe_sm90_fused_ring.py` is new: it checks the lag ring **bitwise** against a full-pool control on eight shapes, four launches per buffer with skewed routing in between, exercising both launch-parity banks and the wrap. A fourth buffer replays those launches under each call's own `num_tokens_bound`, bit for bit again, then under a bound one alignment above capacity, which has to raise on the host before any rank launches. Sixteen cases pass, at most 1e-4 from the reference against the test's 0.07 tolerance. The ring test takes ~2 min on 8 GPUs per granularity, and the sixteen cases are its eight shapes at both scale granularities. The shapes are the shared test's `hidden=1024` shape at 64, 256, 512, 1024 and 2048 tokens per rank, which straddles the 1024-token lag threshold from both sides; a 4096/2048/288/top-8 shape whose lag ring is 20992 of 135680 pool tokens, next to a 98304-token auto-wave ring; a 1024/1024/64/top-6 shape at a 5120-token capacity, 18944 ring tokens of a 246784-token pool next to an 81920-token auto-wave ring; and a 2-rank subgroup whose requested ring is larger than the pool and clamps to it. The verdict is all-reduced, so a single-rank mismatch reports instead of hanging. The replay is where both things a bound can change get covered. The decode-sized launch is bounded at 256 tokens, so on the lag rings that one call runs the single-wave order against the same allocation running the lag order without a bound (`DG_PRINT_CONFIGS` reports `l2_lag_units` 4008, 4014 or 4016 unbounded and 0 under the bound); and on the 5120-token shape the 3840-token launch runs an 11-unit lag under its own bound where the unbounded allocation runs the 14 its ring is sized for (4011 against 4014). Every one of the sixteen cases reports `bitwise=OK replay=OK bounds=...:OK over_cap_rejected=OK`. Both tests spawn one process per GPU themselves, so they need 8 SM90 GPUs and a working NVLink group. </details> <details> <summary>Files, and the naming census</summary> New: `deep_gemm/include/deep_gemm/impls/sm90_fp8_fused_mega_moe.cuh` (the kernel), `layout/sm90_fused_mega_moe.cuh` (workspace and tile-table sizing), `scheduler/sm90_fused_mega_moe.cuh` (the lag scheduler), `csrc/apis/sm90_fused_mega.hpp`, `csrc/jit_kernels/heuristics/sm90_fused_mega_moe.hpp`, `csrc/jit_kernels/impls/sm90_fp8_fused_mega_moe.hpp`, `tests/test_mega_moe_sm90_fused_ring.py`. Changed outside the fused path: 13 wrappers in `ptx/ld_st.cuh`, `advance_smem_desc` in `mma/sm90.cuh`, one registration line each in `csrc/python_api.cpp` and `deep_gemm/__init__.py`, the fused branch in `deep_gemm/mega/__init__.py`, a subsection in `README.md`, a fused arm in `tests/bench_mega_moe_sm90.py` — which is restructured into arms rather than added to — and `tests/test_mega_moe_sm90.py`, where the reference now quantises the intermediate activation at the buffer's granularity, on the `--fused` path only; the split path's reference is unchanged. Plus the two shared files above, `comm/barrier.cuh` and `csrc/jit/kernel_runtime.hpp`. Of the 61 namespace-scope names the new fused files introduce, 57 carry a fused stem (`MegaMoESM90FusedConfig`, `SM90FP8FusedMegaMoERuntime`, `get_mega_moe_config_sm90_fused`, `kSm90Fused*`, and on the device `SM90FusedWorkspace` / `SM90FusedMegaMoEScheduler`). The four that do not are new sizing helpers in `layout/sm90_fused_mega_moe.cuh` (`kNumSM90CandidateBlockMs`, `get_num_max_pool_tokens_sm90`, `get_sm90_tile_table_entries`, `get_sm90_tile_table_entries_compact`); none of the four exists anywhere in the tree today, so nothing collides with #383's SM90 names and `DG_PRINT_CONFIGS` tells the two apart. This PR adds no tuning environment variable. The only new reads are the two trace instruments `DG_SM90_TRACE_PTR` and `DG_SM90_TRACE_EPI`, next to the repo's existing diagnostics (`DG_JIT_DEBUG` / `DG_PRINT_CONFIGS`, `DG_COMM_KERNEL_DEBUG`), which the fused path honours exactly where the split path does. </details> ## Relation to other work #383's review named a fused SM90 kernel as follow-up work; this is one. It does not replace that PR: the split path stays the default. <details> <summary>Overlaps with #422, #411, #393 and #360, and the same kernel against sgl-project</summary> Overlaps to know about: - #422 brings the sgl-project split SM90 kernel into the paths #383 now occupies. It makes the same two barriers generic over the workspace type, with `WorkspaceT` after `sync_scope_t` rather than before, and takes trap-only on SM90 from inside the wait loop, the shape that costs this kernel its `setmaxnreg`. Whichever lands first, the policy parameter has to go in. - #411 adds SM90 FP8 × MXFP4 on top of #383's kernel, with the same `ld_shared(const uint2*)` overload this PR adds, so whichever lands second drops its copy. - #393 rewrites `csrc/apis/*.hpp` for `TORCH_LIBRARY`, including the registration this PR adds; whichever lands first, the other rebases three `m.def`s. - #360 is a different fused SM90 design against `main`, one cooperative kernel whose two math warpgroups N-split a 64×256 tile. It predates #383 and is the closest alternative to this one. The same kernel is open against sgl-project/DeepGEMM as [sgl-project/DeepGEMM#88](https://github.com/sgl-project/DeepGEMM/pull/88), where it replaces that fork's existing fused SM90 kernel instead of adding one. Same kernel, not the same bytes: the barriers here take a per-instantiation timeout policy where the fork gates the same trap on the architecture, and this workspace's flag bank has no publish-epoch word. Nothing reads a bank across the two repositories. Don't assume a patch to one applies to the other. </details> The NVLink data path this kernel keeps comes from [sgl-project/DeepGEMM#36](https://github.com/sgl-project/DeepGEMM/pull/36). That, the first SM90 ring buffer this builds on, and the cluster and multicast paths are all credited in the commit trailers. --- _{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

bsakaguchi added the

pull-request

label 2026-10-26 21:18:00 +00:00

Sign in to join this conversation.

No Branch/Tag specified

Branches Tags

main

public-sync-20261014

0930-release

nv_dev

fix-mega-moe-fence

public-release-2609

public-release-2607

public-release-2606

mega-update

public-release-260416

ko3n1g/ci/build-wheels

nv_dev_f8e8fb5

nv_dev_950d31a

nv_dev_67fc648

nv_dev_c491439

nv_dev_4ff3f54

v2.1.1.post3

v2.1.1.post2

v2.1.1.post1

v2.1.1

v2.1.0

Labels

Clear labels   

pull-request

Mirrored from GitHub Pull Request

No labels

pull-request

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

1 participant

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepGEMM#449

No description provided.