Mario/overlap moe kernel integration by MarioSieg · Pull Request #3651 · PrimeIntellect-ai/prime-rl

MarioSieg · 2026-09-23T21:27:25Z

Integration of our closed mega-moe based on Sonic Moe's abstractions

Note

High Risk
Introduces a new distributed training path with persistent cross-rank kernels, strict NCCL/SM and buffer sizing constraints, and a proprietary deep_gemm dependency; misconfiguration can deadlock or fail at runtime.

Overview
Adds Mega MoE as a new expert-parallel MoE dispatch path (model.moe.dispatch.type = "mega_moe") that fuses token dispatch, SwiGLU expert GEMMs, and combine into prime-mega-moe deep_gemm kernels, overlapping EP communication with compute. Config covers precision (bf16 / mxfp8 / fp8 / fp8_bwd), max_tokens_per_rank (auto-derived from the micro batch in RL/SFT), and num_reserved_sms with a startup check that NCCL_MAX_CTAS ≤ reserved SMs to avoid deadlocks with FSDP NCCL.

Trainer wiring uses FusedTokenDispatcher plus MegaMoEExpertCompute (custom autograd ops, symmetric buffers, optional MXFP8 casting); model.moe.compute is ignored when Mega MoE is selected. GroupedExperts now accepts fused compute that takes routing tensors directly. Docs, a multi-GPU benchmark vs the bf16 Torch dispatcher path, activation-checkpointing registration, SFT timing sync, and a unit test that Mega MoE requires ep > 1 are included.

^{Reviewed by Cursor Bugbot for commit e433a59. Bugbot is set up for automated code reviews on this repo. Configure here.}

Add a mega_moe MoE dispatch config and MegaMoeTokenDispatcher that runs DeepGEMM's fused dispatch + SwiGLU expert compute + combine kernels (forward and backward) over the expert-parallel group. Includes a benchmark against the bf16 TorchTokenDispatcher path.

…nel-integration

garrett361 · 2026-09-24T18:15:27Z

nice, benchmark numbers?

MarioSieg · 2026-09-24T18:17:11Z

nice, benchmark numbers?

mega-moe is 1.55× faster (35.4% lower fwd+bwd time), while using slightly less peak memory.

Note: we are currently investgating an MFU issue unrelated to this kernels so the speedup might change.
That's why this PR is a draft until we get 100% right bench numbers .- speedup might be a bit higher or lower.

…o/overlap-moe-kernel-integration Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Mega MoE is now a `MegaMoEExpertCompute` installed on `GroupedExperts.compute` by `configure_moe_runtime`, like the other expert compute backends, instead of a token dispatcher that injected a `fused` callback into `GroupedExperts.forward`. Because the kernel fuses dispatch and combine, it consumes undispatched tokens and their routing: `FusedDispatchExpertCompute` describes such a compute, and `FusedTokenDispatcher` hands the tokens to it without communicating. `GroupedExperts.forward(x, *routing)` forwards either `num_tokens_per_expert` or `top_scores, selected_experts_indices` to its compute. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…e extra Mega MoE kernels come from prime-mega-moe, a closed-source DeepGEMM fork. Its wheel is built by hand and attached to the prime-kernels release; prime-rl pins the URL and never builds it from source. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

SonicMoE landed on main squashed (#3535) with an ExpertCompute.validate hook and GroupedExperts.set_compute. Mega MoE adopts both: MegaMoEExpertCompute.validate carries the expert-structure checks and configure_moe_runtime installs it through set_compute. uv.lock was regenerated with the previous deep_gemm wheel URL because the prime-mega-moe wheel is not uploaded yet; re-run uv lock once it is. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…ck compiles with fullgraph The fused forward was a custom op wrapped in a Python autograd.Function that saved x.to(bfloat16).contiguous() for backward. When the block compiles with fullgraph under activation checkpointing, dynamo rejects that saved tensor because it aliases the autograd.Function's input inside the checkpoint subgraph. Make the backward a custom op too and register autograd on prime_rl::mega_moe_forward, saving the op's own inputs. Both ops take the caller's dtypes, run the kernels in bf16, and return gradients in the inputs' dtypes, so the expert dW dtype is unchanged.

cursor

Stale Bugbot comment from a previous run.

Main moved the DeepSeek V4 clamped SwiGLU to `activations.ClampedSilu`; the helper still imported a `ClampedSwiglu` from `deepseek_v4.moe` that no longer exists, so constructing `MegaMoEExpertCompute` failed for every model. `ClampedSwiglu` in `activations.py` is a different function (sigmoid * 1.702, `up + 1`) that the kernel does not compute, so it stays rejected.

garrett361

Looks really good, just had organization comments and questions, mainly.

cursor

Stale Bugbot comment from a previous run.

cursor

Stale Bugbot comment from a previous run.

…nel-integration # Conflicts: # packages/prime-rl-configs/src/prime_rl/configs/trainer.py

cursor

Stale Bugbot comment from a previous run.

garrett361 · 2026-10-15T19:50:11Z

+            )
+        self.validate(experts)
+        hidden = experts.down_proj.shape[1]
+        reserve_sms_for_comm(num_reserved_sms)


This one @MarioSieg ?

garrett361 · 2026-10-15T20:09:54Z

@MarioSieg response to this?

cursor

Stale Bugbot comment from a previous run.

…nel-integration

garrett361

Looks good, thanks for all the iteration on this! I'd still like to clean up some of the internals of mega_moe.py and the kernel names at some point (still requires a number of fairly technical and non-obvious mappings to follow), but let's get this in and clean up later.

MarioSieg · 2026-10-19T16:26:08Z

Looks good, thanks for all the iteration on this! I'd still like to clean up some of the internals of mega_moe.py and the kernel names at some point (still requires a number of fairly technical and non-obvious mappings to follow), but let's get this in and clean up later.

LMK what you would cleanup / refractor I can finish it today and we can merge tomorrow then.

cursor

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit fecda7f. Configure here.}

…nel-integration

S1ro1 and others added 11 commits September 12, 2026 16:09

feat(moe): add SonicMoE BF16 expert compute backend

4b7857c

refactor(moe): consolidate GEMM backends under expert compute

49f6407

fix: pin SAC-safe QuACK and EP-aware SonicMoE autotuning

35faa90

Integration of mega-moe backward

5e3e0e2

Merge remote-tracking branch 'origin/main' into mario/overlap-moe-ker…

4ed24fe

…nel-integration

Dusal layout support

27fb710

Fix weight storage for FSDP all gather

5ee3819

Merge remote-tracking branch 'origin/main' into mario/overlap-moe-ker…

876ca96

…nel-integration

Update from main

512d88d

AC buffer registration

9466bfa

MarioSieg and others added 3 commits September 24, 2026 20:17

Merge branch 'main' into mario/overlap-moe-kernel-integration

29a961f

Merge remote-tracking branch 'origin/feat/sonicmoe-bf16-ep' into mari…

ce28728

…o/overlap-moe-kernel-integration Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

MarioSieg changed the base branch from main to feat/sonicmoe-bf16-ep September 24, 2026 21:40

AC buffer registration

15bd003

S1ro1 force-pushed the feat/sonicmoe-bf16-ep branch 2 times, most recently from 1f1ff90 to 0609006 Compare September 29, 2026 14:53

Base automatically changed from feat/sonicmoe-bf16-ep to main September 30, 2026 20:42

MarioSieg and others added 2 commits October 2, 2026 06:39

MarioSieg mentioned this pull request Oct 2, 2026

perf(trainer): compile MoE blocks with fullgraph by keeping EP experts in the block's FSDP unit #3752

Closed

This was referenced Oct 3, 2026

chore(configs): add the Qwen3-235B SFT benchmark config matched to torchtitan #3765

Draft

docs: performance.md with the Qwen3-235B numbers vs torchtitan and the max-MFU benchmark configs #3766

Draft

MarioSieg added 2 commits October 5, 2026 08:17

Remove natural layout flags

36552d4

Add mxfp8 integration for mega moe

d6434c0

MarioSieg marked this pull request as ready for review October 6, 2026 16:27

cursor Bot reviewed Oct 7, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/layers/mega_moe.py Resolved

garrett361 previously approved these changes Oct 7, 2026

View reviewed changes

Fixes

1d49fee

MarioSieg dismissed garrett361’s stale review via 1d49fee October 8, 2026 14:36

cursor Bot reviewed Oct 8, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/layers/mega_moe.py Resolved

Fix comments of integration

821d6a8

cursor Bot reviewed Oct 8, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/layers/mega_moe.py Resolved

MarioSieg added 2 commits October 13, 2026 14:43

Fp8 integration

916ad8e

Merge remote-tracking branch 'origin/main' into mario/overlap-moe-ker…

19138e9

…nel-integration # Conflicts: # packages/prime-rl-configs/src/prime_rl/configs/trainer.py

cursor Bot reviewed Oct 15, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/layers/mega_moe.py Resolved

garrett361 reviewed Oct 15, 2026

View reviewed changes

Improve integration, bug fixes and comment fixes

a4cc078

cursor Bot reviewed Oct 19, 2026

View reviewed changes

Comment thread benchmarks/mega_moe/bench_mega_moe.py Resolved

MarioSieg added 3 commits October 19, 2026 13:11

Merge remote-tracking branch 'origin/main' into mario/overlap-moe-ker…

a7070dd

…nel-integration

Auto destroy buffers

e63dded

Make SM allocation count local

ecbab89

garrett361 previously approved these changes Oct 19, 2026

View reviewed changes

garrett361 reviewed Oct 19, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/layers/mega_moe.py Resolved

Comment thread src/prime_rl/trainer/models/layers/mega_moe.py Resolved

MarioSieg dismissed garrett361’s stale review via fecda7f October 20, 2026 10:20

cursor Bot reviewed Oct 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/rl/train.py Resolved

Comment thread src/prime_rl/trainer/parallel_dims.py Resolved

MarioSieg added 2 commits October 20, 2026 10:58

CLearer naming

d86fd76

Merge remote-tracking branch 'origin/main' into mario/overlap-moe-ker…

af43116

…nel-integration

MarioSieg force-pushed the mario/overlap-moe-kernel-integration branch from a5c0048 to af43116 Compare October 20, 2026 11:01

MarioSieg added 2 commits October 20, 2026 11:02

Merge remote-tracking branch 'origin/main' into mario/overlap-moe-ker…

405877b

…nel-integration

Cleanup symmmetric buffer allocation

e433a59

garrett361 approved these changes Oct 20, 2026

View reviewed changes

MarioSieg merged commit a74bcb8 into main Oct 20, 2026
19 of 20 checks passed

MarioSieg deleted the mario/overlap-moe-kernel-integration branch October 20, 2026 11:21

Conversation

MarioSieg commented Sep 23, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

garrett361 commented Sep 24, 2026

Uh oh!

MarioSieg commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

garrett361 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

garrett361 Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

garrett361 Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

garrett361 left a comment

Choose a reason for hiding this comment

Uh oh!

MarioSieg commented Oct 19, 2026

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

MarioSieg commented Sep 23, 2026 •

edited by cursor Bot

Loading

MarioSieg commented Sep 24, 2026 •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading

cursor Bot left a comment •

edited

Loading