fix(trainer): run FA4 through a custom op so compile can trace attention by MarioSieg · Pull Request #3676 · PrimeIntellect-ai/prime-rl

MarioSieg · 2026-09-25T21:19:12Z

Summary

FA4's flash_attn_varlen_func is a Python autograd.Function that calls torch._guards.active_fake_mode(), which Dynamo refuses to trace. With activation checkpointing on, each block is one checkpointed region. A graph break inside it drops the whole block to eager under torch.compile, for every FA4 model, dense or MoE.

Wrap FA4's varlen forward/backward in prime_rl_attn::flash_attn_4_varlen (fake impls + autograd). It's registered as FlashAttention._funcs[4] with the same call signature and (out, lse) return as flash_attn_varlen_func, so _compute_attention is unchanged from main.
Add the prime_rl_attn namespace to the selective AC save list so attention isn't recomputed in backward.

Ring and Ulysses CP replace _compute_attention and keep their own FA4 calls, so they behave as before (not broken, but still eager under AC).

For dense models this is the only graph break in the block. MoE with EP has two more, fixed in #3677.

Validation

Output and grads match flash_attn_varlen_func (out/dq exact, dk/dv within a few bf16 ulps), eager and under torch.compile(fullgraph=True), with and without a sliding window.
CP still works. Qwen3-8B FA4 selective AC, 5 steps: loss at step 5 is 3.5778 (cp=1), 3.5828 (cp=2 ring), 3.5799 (cp=2 Ulysses).
Qwen3-8B dense, 1 node, 8k seq, selective AC, same node: 579 → 477 ms/step (1.21×), 45.9 → 37.0 GiB, loss/grad norm match. Measured together with perf(trainer): use TorchTitan's selective activation checkpointing save policy #3678, not in isolation.

Overlaps with #3575 on the FA4 op. If that lands first, this reduces to adding its op to the save list.

🤖 Generated with Claude Code

Note

Medium Risk
Changes training forward/backward for FA4 attention and checkpoint save policy; correctness depends on custom-op autograd matching FA4, but scope is limited to the FA4 FlashAttention path.

Overview
Routes Flash Attention 4 (FA4) varlen through traceable custom ops so torch.compile no longer graph-breaks inside activation-checkpointed transformer blocks when dense FlashAttention uses FA4.

FA4’s native flash_attn_varlen_func is replaced at the FlashAttention._funcs[4] slot by flash_attn_4_varlen_op, which calls prime_rl_attn::flash_attn_4_varlen / flash_attn_4_varlen_backward (fake kernels + registered autograd) while keeping the same (out, lse) contract as before. prime_rl_attn is added to selective activation checkpointing save namespaces so FA4 attention outputs are stored instead of recomputed on backward.

Ring/Ulysses CP paths that still call FA4 directly are unchanged and may remain eager under AC.

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

FA4's flash_attn_varlen_func is a Python autograd.Function that calls torch._guards.active_fake_mode(), which Dynamo refuses to trace. Under activation checkpointing each block is one checkpointed region, so that graph break drops the whole block to eager for every FA4 model, dense or MoE. Wrap FA4's varlen forward/backward in prime_rl_attn::flash_attn_4_varlen (with fake impls and autograd) and add the namespace to the selective checkpointing save list so attention is not recomputed in backward. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…_funcs Keep _compute_attention identical to main and register the custom op as the FA4 entry of _funcs, with the same call signature and (out, lse) return value as flash_attn_varlen_func. Code that reads or swaps self.func / _funcs keeps working, and CP (ring, ulysses) still replaces _compute_attention as before. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Adds `model.moe.compute.backend = "sonicmoe"` for BF16 MoE expert computation. Sonic consumes already-dispatched expert rows and fused gate/up weights; dispatchers retain routing-score application and combination. `GroupedExperts` owns its parameters and delegates execution through the structural `ExpertCompute` protocol. Concrete backends do not inherit from the protocol. `set_compute()` validates compatibility before assignment: Sonic requires fused gate/up weights, standard SwiGLU, and bias-free experts. Unsupported structures, including DeepSeek V4's clamped SwiGLU, raise during setup without replacing the active backend. Native BF16, DeepGemmFP8, and MXFP8 share the grouped-GEMM implementation. Sonic uses its fused expert pipeline. Selective checkpointing saves `quack::gemm_out`, `quack::gemm_gated_out`, and `quack::gemm_act_out`; compiled expert traces verify that saving eliminates the repeated gate/up GEMM. ```toml [trainer.model] name = "Qwen/Qwen3-30B-A3B" ep = 8 [trainer.model.moe.compute] type = "bf16" backend = "sonicmoe" [trainer.model.moe.dispatch] type = "torch" ``` Install with the `sonic-moe` extra, also included in `all`. Uses default gate/up fusion and retains optimization/reduction dtype defaults. Rebased onto main `9eacd477f`. Pins SonicMoE's EP-aware token bucketing at `0bf06ac` and [QuACK's checkpoint-isolation fix](Dao-AILab/quack#204) at `f62009f`, based directly on QuACK 0.6.5 and compatible with main's vLLM 0.30 / CUTLASS 4.7.1 stack. Validation: - After rebase and dependency update: 224 MoE, checkpoint, and configuration tests pass; Ruff, formatting, whitespace, lockfile consistency, and all-extras installation pass. - After the QuACK 0.6.5 update: six realistic-shape GPU comparisons pass at 32,896 and 65,664 routed rows (16 experts, dimensions 2048/768), including uneven routing, eager execution, and fullgraph compiled selective checkpointing with/without GEMM saving. Zero graph breaks; maximum output and gradient relative errors were 0.454% and 0.402%. - Before this rebase: 18 realistic expert-shard comparisons passed across 8,192–65,664 routed rows, 16 local experts, dimensions 2048/768, eager and compiled selective-checkpoint paths. Maximum relative errors were 0.454% for outputs and 0.402% for gradients. - Paired full-model B300 EP8 runs completed 20 Qwen3-30B-A3B math SFT steps with sequence length 8192, batch 8, seed 42, selective checkpointing, and SignSGD 1e-5. FA4, FA2, and a separate FA4 experiment with #3676 all produced finite losses/gradients and readable final checkpoints. These E2E runs used the pre-rebase QuACK 0.6.4 stack; #3676 was only an experimental overlay and is not included here. Performance follow-up: Sonic's later-step overhead was 14.3% with original FA4, 16.9% with FA2, and 21.1% in the #3676 FA4 experiment. These are separate pairs of single runs, not controlled cross-backend attention comparisons. Shared compilation breaks prevented compiled checkpointed-block execution: #3676 removed the attention break, then both expert backends hit the FP32 router's nested FSDP hook. This limits conclusions about effective Sonic GEMM saving in E2E. It does not establish that fixing the shared breaks eliminates all Sonic overhead. Sonic also retains long first-step autotuning and occasional approximately 37-second bucket-transition steps without a Dynamo recompile. These remain performance follow-ups; no correctness blocker was found in the realistic-shape and E2E checks.  --- > [!NOTE] > **Medium Risk** > Changes core MoE forward/backward paths and adds a new optional kernel backend with strict structural requirements; misconfiguration fails at setup, but numeric parity depends on the pinned Sonic/QuACK stack. > > **Overview** > Adds an optional **SonicMoE** path for BF16 routed-expert compute via `model.moe.compute.backend = "sonicmoe"` (default remains `"torch"`). Install with the new `sonic-moe` extra; `quack-kernels` is pinned from git for Sonic’s dependency stack. > > MoE execution is refactored from injectable **grouped GEMM** objects to an **`ExpertCompute`** protocol: `GroupedExperts` delegates forward to `set_compute()`, which runs **`validate()`** before switching backends. Torch BF16, DeepGEMM FP8, and MXFP8 share **`GroupedGemmExpertCompute`**; Sonic uses a fused **`moe_general_routing_inputs`** path and rejects unfused gate/up, non-SwIGLU activations, and expert biases at setup. > > Runtime wiring in **`configure_moe_runtime`** picks the backend from config; selective activation checkpointing now retains **`quack::gemm_*`** ops for recomputation with Sonic/QuACK. Docs and unit tests cover config defaults, Sonic compatibility checks, and unchanged BF16 fallback when layers are excluded by **`apply_to`**. > > Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 0609006. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).  --------- Co-authored-by: matejsirovatka <matejsirovatka@users.noreply.github.com>

MarioSieg · 2026-10-02T12:22:34Z

Superseded by #3754 (same change on mario/optimize-fa4-custom-op).

…ion (#3754) Same change as #3676 on a `mario/optimize-*` branch, so the fullgraph work lands as independent PRs; supersedes #3676. ## Why FA4's `flash_attn_varlen_func` is a Python `autograd.Function` that calls `torch._guards.active_fake_mode()`, which dynamo refuses to trace. Under activation checkpointing each block is one checkpointed region, so that graph break drops the whole block to eager for every FA4 model, dense or MoE, and `compile.fullgraph=true` fails (job 1175). ## What - Wrap FA4's varlen forward and backward in a `prime_rl_attn::flash_attn_4_varlen` custom op with fake impls and registered autograd, dispatched through `FlashAttention._funcs` like the other versions. - Add the `prime_rl_attn` namespace to the selective activation checkpointing save list so attention is not recomputed in backward. Needed, together with the sibling `mario/optimize-fsdp-expert-fold` PR, for the Qwen3-235B result of 4.40 s/step vs torchtitan's 5.01 s at seq 16k on 4 nodes.  --- > [!NOTE] > **Medium Risk** > Changes the FA4 attention forward/backward path and autograd behavior; incorrect custom-op wiring would affect training correctness, but scope is limited to FA4 + compile/checkpointing. > > **Overview** > Flash Attention 4 varlen is no longer called via FA4’s untraceable Python `autograd.Function`. Forward and backward are wrapped in **`prime_rl_attn::flash_attn_4_varlen`** (and a matching backward op) with fake kernels and registered autograd, and **`FlashAttention`** version 4 now dispatches through **`flash_attn_4_varlen_op`**. > > That avoids Dynamo graph breaks inside activation-checkpointed blocks (which previously forced whole blocks to eager and broke **`compile.fullgraph=true`**) and lets selective activation checkpointing treat attention as an opaque op. > > **`prime_rl_attn`** is added to the selective checkpoint **save** namespace list so FA4 attention outputs can be saved instead of recomputed on backward. > > Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 37cdea5. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).  --------- Co-authored-by: sami jaghouar <sami.jaghouar@gmail.com>

… compile with fullgraph (#3753) ## Why Under expert parallelism the EP experts were their own FSDP unit nested inside each transformer block. FSDP2 wraps its hooks in `torch._dynamo.disable`, so dynamo graph-broke at the experts and the MoE dispatch/combine glue and the activation-checkpointing policy ran eagerly. On the Qwen3-235B SFT profile against torchtitan that was ~2.7 s/step of unfused elementwise and index kernels plus 3.4x the kernel launches, and `compile.fullgraph=true` failed outright (job 1180). torchtitan avoids this by sharding the experts over the EP-complement mesh from the block's own `fully_shard` call. ## What - The block's `shard_placement_fn` returns a `ShardPlacementResult` with the `dp_shard_mod_ep` mesh for expert parameters. FSDP2 in torch 2.13 groups one unit's parameters by process group, so the experts still all-gather and reduce-scatter over the EP-complement mesh, with no nested unit and no hook inside the compiled block. - The gradient divide factor is set on the block. With an fp32 reduce it is the dense parameters' default divisor (`ReduceOp.AVG` over the shard group), so dense numerics are unchanged; the experts keep dividing by the full data-parallel size as before. - Explicit EP prefetch lists drop the experts module. - `compile.fullgraph=true` on a custom MoE model now requires `moe_router_dtype="bfloat16"`: the fp32 router is still its own FSDP unit inside the block. `setup_fsdp` raises early with that message. The fp32-router default is unchanged. ## Benchmark Qwen3-235B-A22B-Thinking-2507 SFT, 4 nodes x 8 B300, FSDP over 32 GPUs, ep=8, seq 16384 (one packed sample per GPU per step), balanced routing, AdamW fp32 states, full AC, 20 steps, steady state = steps 5..20, same MFU formula in every row (280.8 GFLOP/token). Measured together with the FA4 custom op (#3676 / the sibling `mario/optimize-fa4-custom-op` PR), which fullgraph also needs on B300. | run | s/step | MFU | tok/s/GPU | peak GiB | |---|---|---|---|---| | torchtitan 9e159aed7, fullgraph | 5.01 | 40.9% | 3273 | 144 | | prime-rl before, graph breaks at the expert unit | 6.98 | 29.3% | 2347 | 140 | | **prime-rl after, fullgraph** | **4.40** | 46.5% | 3726 | 141 | | prime-rl before, real routing (job 1184) | 9.51 | 21.5% | 1724 | 142 | | **prime-rl after, fullgraph, real routing (job 1461)** | **6.11** | 31.2% | 2681 | 141 | Per-step losses match the previous runs on the same data order (balanced and real routing). A rank-0 torch-profiler trace of the new step (4.46 s) vs the old one (7.26 s): eager MoE glue went from 2.95 s elementwise + 1.09 s index kernels to 0.62 s of inductor kernels, kernel launches from 33.9k to 9.9k, and the per-op Python AC dispatch mode is gone. The remaining step is GEMM 47%, the EP all-to-all 22% (still on the compute stream, fully exposed), fused glue 14%, attention 7%. prime-rl is ahead of torchtitan because torchtitan's compute stream idles 1.55 s waiting on its 1.34 s all-to-all, while prime-rl's inline all-to-all takes 0.99 s. Caveat: the MFU formula charges dense causal attention, but packed documents make attention 0.2 to 0.3 s of the step in both stacks, so absolute MFU is about 2x overstated for every row; the comparison is unaffected. A 1-node Qwen3-30B-A3B fullgraph smoke run and `tests/unit/train/models/test_moe.py` pass. ## Open - fp32 router under fullgraph needs the router out of its own FSDP unit (ignored params with a replicated fp32 gate and a manual grad all-reduce, or a bf16 gate). Not decided here.  --- > [!NOTE] > **Medium Risk** > Changes distributed FSDP sharding and gradient divisors for EP MoE blocks, which affects training numerics and performance paths; misconfiguration now errors at setup instead of silent graph breaks. > > **Overview** > **MoE + expert parallelism (EP) no longer uses a nested FSDP unit on `experts`**, so `torch.compile(fullgraph=True)` can capture whole transformer blocks without dynamo breaking on FSDP hooks inside the block. > > Under EP, expert weights still shard on the EP-complement mesh via a per-block `shard_placement_fn` that returns `ShardPlacementResult` for expert parameters, while dense block params use the usual placement. Gradient scaling moves to the **transformer block** (`set_gradient_divide_factor`) instead of the removed experts submodule. EP forward/backward **prefetch** lists drop explicit `experts` prefetch and only add the fp32 **router** when it remains its own FSDP module. > > **Config/docs:** `CompileConfig.fullgraph` documents that custom MoE + fullgraph needs `moe_router_dtype="bfloat16"`. **`setup_fsdp` fails fast** if `compile.fullgraph=true` is combined with the default fp32 router, since that router must stay a separate FSDP unit inside the compiled block. > > Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit d2f7635. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).

This was referenced Sep 25, 2026

fix(trainer): remove MoE graph breaks inside checkpointed blocks #3677

Closed

perf(trainer): use TorchTitan's selective activation checkpointing save policy #3678

Draft

samsja mentioned this pull request Sep 25, 2026

fix(trainer): rework FA4 custom op to the FA2/FA3 varlen signature, dispatched via _funcs #3679

Closed

MarioSieg marked this pull request as ready for review September 29, 2026 00:08

samsja approved these changes Sep 29, 2026

View reviewed changes

S1ro1 mentioned this pull request Sep 29, 2026

feat(moe): add SonicMoE BF16 expert compute backend #3535

Merged

MarioSieg closed this Oct 2, 2026

Navigation Menu

Repository navigation

fix(trainer): run FA4 through a custom op so compile can trace attention - #3676

fix(trainer): run FA4 through a custom op so compile can trace attention#3676
MarioSieg wants to merge 2 commits into
mainfrom
mario/fa4-custom-op

MarioSieg commented Sep 25, 2026 •

edited by cursor Bot

Loading

Uh oh!

MarioSieg commented Oct 2, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Summary

Validation

Uh oh!

MarioSieg commented Oct 2, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

MarioSieg commented Sep 25, 2026 •

edited by cursor Bot

Loading