fix(moe): pass an explicit output size to the EP all-to-all ops by garrett361 · Pull Request #4086 · PrimeIntellect-ai/prime-rl

garrett361 · 2026-10-14T13:45:25Z

This PR passes the EP all-to-all's output length into the op, instead of creating a dynamic size inside the op's fake, so MoE layers compile under the default fullgraph = false. On main, that fake breaks the graph inside the activation checkpoint and every MoE layer runs eagerly; with this PR they match fullgraph = true speed, 1.3x to 2.0x faster on B300.

The dispatcher passes int(output_splits.sum()) from CPU splits; _mxfp8_all_to_all gets the same change.
fullgraph = true already compiled these layers because it implicitly turns on capture_dynamic_output_shape_ops.
The combine reuses the dispatch length, so both all-to-alls share one symbolic size instead of two.
At cp 2 each rank recompiles once more, on the ULYSSES_PARAMS guard (one 8.1 s step); fix(cp): pass cp attention params as forward args #4068 removes it.

1 x 8 B300, ep 8, full AC, default compile	seq len per sample	main (s)	this PR (s)	speedup
Qwen3-30B-A3B 12 layers, fake data	16384	1.067	0.703	1.52x
mini-glm-moe, fake data	8192	1.066	0.542	1.97x
GLM-4.5-Air 8 layers, real data, cp 1	16384, 1-3 docs, median 12.2k tokens	1.892	1.427	1.33x
GLM-4.5-Air 8 layers, real data, cp 2	16384 (8192 per rank), 1-3 docs, median 12.2k tokens	2.246	1.718	1.31x
Qwen3-30B-A3B 12 layers, MXFP8 experts and transport	16384	1.093	0.716	1.53x

This PR has 0 graph breaks in every arm, against 1 on main, and loss agrees within same-code repeat noise. Main with fullgraph = true runs the MXFP8 arm in 0.716 s. Baseline fea1267ef; harness and full results in REPORT.md at 82984c8a8, removed in the last commit.

Qwen3-30B-A3B layer 2 forward, rank 0, on H200 (an earlier version of this PR with the same fix). Eager elementwise, copy and index kernels (orange) become a few fused Inductor kernels (green): 16.3 ms to 11.0 ms.

🤖 Generated with Claude Code

The all_to_all_single and mxfp8_all_to_all fakes declared a data-dependent output length, which Dynamo only captures with fullgraph=True. With the default fullgraph=False every MoE block under EP graph-broke inside the activation checkpoint and ran eagerly. The dispatcher now reads the received-token count from the CPU splits (a captured scalar) and the combine reuses the pre-dispatch routed length, so the fakes have static output shapes. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…sis commands Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…e the proxy claim Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…4068 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…eport Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

MarioSieg · 2026-10-22T15:41:31Z

lgtm!

MarioSieg

lgtm!

Ulysses context parallelism published the full-sequence `cu_seqlens` and `max_seqlen` through a module global, `ULYSSES_PARAMS`, which attention read at call time. Compiled blocks guard on that global, so a micro batch with a different document count recompiles every rank. This PR drops the global: attention uses the `cu_seqlens` and `max_seqlen` it is already passed, which the models build from the pre-shard `seq_lens` and so cover the full sequence each rank holds after the all-to-all. - `ULYSSES_PARAMS`, `update_ulysses_params` and `setup_cp_attention_params` are gone, so `setup_cp_params` only shards inputs. - Both trainers check once per micro batch, on the CPU `seq_lens`, that they sum to the full sequence length. This also covers the Qwen3.5 VLM path that defers sharding to the model, and never syncs with the GPU (main's check made three host syncs per micro batch). - GPT-OSS's Ulysses attention gets the same change. - Rebased onto main after #4134 removed ring CP, so this is now Ulysses only; the ring half of the original PR is gone with it. The recompile shows only where blocks compile. On main, EP MoE blocks run eagerly under the default `fullgraph = false` (fixed by #4086), so on the MoE workload below this PR alone is speed-neutral. On top of #4086 it removes the recompile and its slow step: | 1 x 8 B300, GLM-4.5-Air 8 layers, ep 8, cp 2, real data | median (min-max) s | recompiles per rank | `ULYSSES_PARAMS` guard failures | | --- | --- | --- | --- | | main | 2.246 (2.183-2.344) | 1 | 0 | | this PR | 2.231 (2.191-2.339) | 1 | 0 | | #4086 | 1.718 (1.672-8.126) | 3 | 2 | | #4086 + this PR | 1.702 (1.664-1.819) | 1 | 0 | Measured at `b57635d84`; the later commit only moves the `seq_lens` check to the CPU in the trainers. Loss agrees across arms within 0.003 over steps 1-9 and within the repeat spread after that. Baseline: [`fea1267ef`](fea1267). Harness, configs and results: [`REPORT.md` at f3a57c5](https://github.com/PrimeIntellect-ai/prime-rl/blob/f3a57c55f5b7922df3a6d0f2d30c8497df6603ea/benchmarks/scripts/cp-ulysses-params-b300/REPORT.md), removed in the last commit. The ring-era results are in the earlier report at [`47070d3d2`](https://github.com/PrimeIntellect-ai/prime-rl/blob/47070d3d281452d748cd64cfa562f05bb9c5eddd/benchmarks/scripts/cp-ring-params/REPORT.md). GPU unit tests on the same node pass and match main: `test_moe.py`, `test_checkpointing.py`, the Qwen3.5 CP setup test and the 2-rank GPT-OSS CP attention test. Unrelated bugs found along the way, not fixed here: - GPT-OSS calls FA4 directly because the FA4 custom op has no sink argument, so it graph-breaks and cannot use fullgraph. - The color-codeword RL env fails in the orchestrator with `Inline image count does not match expanded placeholder runs`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

#4177 removed model.moe_router_dtype, which broke base.toml's fp32 pin and the bf16_router feature row. The base now keeps the ModelConfig defaults, and a fullgraph=True row on qwen3_moe's cp2_ep4 layout turns any graph break in the EP all-to-all (#4086) or Ulysses CP (#4068) paths into a hard failure. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

garrett361 force-pushed the fix/ep-a2a-static-size branch from c6b3073 to 63cf029 Compare October 14, 2026 14:39

garrett361 changed the base branch from main to fix/cp-attn-params-as-args October 14, 2026 14:40

garrett361 added this pull request to stack #4089 October 14, 2026 14:40

garrett361 marked this pull request as ready for review October 14, 2026 16:33

garrett361 marked this pull request as draft October 14, 2026 18:53

garrett361 and others added 18 commits October 22, 2026 13:37

chore(bench): add EP all-to-all compile probe

5053006

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add EP all-to-all A/B harness

e924d03

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add trace, fullgraph repeat and reverse-text RL arms

cafdeca

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add CP arms combining the CP attention-params fix

a306f94

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): start the CP combo arm while the RL arms run

fe19d11

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

docs(bench): record EP all-to-all A/B results, trace figure and analy…

88bfe66

…sis commands Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add Qwen3-30B-A3B and GLM-4.5-Air end-to-end arms, scop…

10181c3

…e the proxy claim Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

docs(bench): record Qwen3-30B-A3B and GLM-4.5-Air end-to-end results

c764aed

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add GLM-4.5-Air cp=2 arms crossing this fix with #4068

6138df5

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): restore the executable bit on run_e2e_arm.sh

1ed550c

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

docs(bench): record GLM-4.5-Air cp=2 results crossing this fix with #…

1dc6642

…4068 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add an untraced-only end-to-end driver

cd03c63

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

docs(bench): record end-to-end results against #4068 as the parent

a8fea76

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): remove the EP all-to-all benchmark harness

4b48880

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): record the MXFP8 EP all-to-all smoke check on B300

8a2962b

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): add the B300 A/B harness and results against main

82984c8

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

chore(bench): remove the B300 A/B harness and the stale MXFP8 smoke r…

b0b231f

…eport Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

garrett361 force-pushed the fix/ep-a2a-static-size branch from 938ad25 to b0b231f Compare October 22, 2026 14:19

garrett361 removed this pull request from stack #4089 October 22, 2026 14:20

garrett361 changed the base branch from fix/cp-attn-params-as-args to main October 22, 2026 14:20

garrett361 mentioned this pull request Oct 22, 2026

fix(cp): pass cp attention params as forward args #4068

Merged

garrett361 marked this pull request as ready for review October 22, 2026 14:29

MarioSieg approved these changes Oct 22, 2026

View reviewed changes

garrett361 merged commit 666acf0 into main Oct 22, 2026
24 of 25 checks passed

garrett361 deleted the fix/ep-a2a-static-size branch October 22, 2026 15:45

yafshar mentioned this pull request Nov 3, 2026

fix(trainer): import the trainer without the kernels extra #3741

Open

Navigation Menu

Repository navigation

fix(moe): pass an explicit output size to the EP all-to-all ops - #4086

fix(moe): pass an explicit output size to the EP all-to-all ops#4086
garrett361 merged 18 commits into
mainfrom
fix/ep-a2a-static-size

garrett361 commented Oct 14, 2026 •

edited

Loading

Uh oh!

MarioSieg commented Oct 22, 2026

Uh oh!

MarioSieg left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

garrett361 commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MarioSieg commented Oct 22, 2026

Uh oh!

MarioSieg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

garrett361 commented Oct 14, 2026 •

edited

Loading