[bradhilton]

Agent: Shazeer

Problem. #1135 lets dense canonical adapters share one decoder forward and backward, but it rejects routed experts. MoE, where many small adapters gain the most from fusion, stayed on one pass per adapter. A naive fused MoE pass failed the same-state gradient gate. The combined batch changed:

• router aux/z losses and capacity;
• TE row geometry;
• FLASH mask tile origins;
• the reduction tree of FLA's batch-autotuned gate cumsum.

Behavior. New explicit installers and scopes let several adapters share one MoE decoder pass, while each adapter keeps its native per-adapter numerics:

• The router runs once per adapter, preserving aux/z losses, capacity and padding.
• Dispatched expert rows carry adapter IDs, so each adapter's own LoRA factors use the native grouped or shared-outer kernels.
• TE norms and projections, and attention, run per root with native row geometry and tile origins.
• FLA's scalar chunk-local cumsum runs per root while the recurrence stays batched. The scope survives activation recompute.
• execution(None, numerical_policy="native", deterministic_native=True) scopes PyTorch determinism. FLASH forward is kept, and backward uses deterministic Triton. If execution falls out of compilation, it now raises rather than silently using the eager path.

Nothing changes unless a caller installs these and enters the scopes, with one exception: callers that already enable PyTorch determinism with FLASH now get the Triton backward. On main, strict mode raises NotImplementedError for masked FLASH, and warn-only runs the nondeterministic FLASH backward. Determinism-off behavior is unchanged. OpenPipe/caladan will enable MoE waves only behind STARK_FUSED_WAVES=1, under a versioned native-moe-wave-v1 policy (separate PR).

This PR carries the native-repeat lane's whole deterministic attention change: d17408e0, with 355fccef, 3b4517ec and e23d0482 cherry-picked. That lane doesn't need its own ART PR.

Validation.

• Exact-head GPT-6.1 Sol reviews:
  • The first, at a0bd182a, was NOT CLEAR. Its one P2: deterministic FLASH could silently fall back to eager attention after Dynamo's recompile limit. It also raised a P3 on test portability.
  • The three cherry-picks fix both. The second review, at f47f793c, is CLEAR-WITH-NOTES.
  • CI then failed because three new unit-test files import Megatron, which the backend test step lacks. fab384ed runs them in the Megatron step instead. A third review, at this head fab384ed, is also CLEAR-WITH-NOTES.
• 95 CPU unit tests pass, including the real-MCore router, the expert plan, the GDN scope and contexts, and the deterministic lowering and cache guard. Ruff and format pass.
• The lanes passed these GPU tests on H200: 16 TE/attention/expert, 4 GDN and the deterministic attention suite.
• Full-model Qwen3.6-35B-A3B qualification ran on one H200 (Torch 2.11, TE 2.11, FLA 0.5.0, MCore 0.17). All 24 cells pass: LoRA rank 1/8/32 × 2/4/8/16 adapters × 1K tokens each or 64K bulk. Each ran three continued AdamW steps, and every output, gradient and optimizer delta is bitwise equal across three serial controls and the fused run. At rank 1 with 1K tokens per adapter, a full step is 2.17× faster with 16 adapters and 2.09× with 64.

Limitations.

• Single rank only (TP/CP/EP/PP 1), with no expert bias or global aux loss, qualified on H200.
• The full-model matrix used uniform 1,024- and 2,048-token roots. Unequal roots are covered only by component tests: TE and attention at 257/129 tokens. The GDN tests use equal roots.
• The policy matches its own serial execution, not default native gradients, which still differ.
• The GPU matrix ran before the three cherry-picks and wasn't rerun. They add a guard and tests only.
• Raising on compilation escape could stop a long-lived deterministic trainer with a very diverse compile history. The review found this plausible but not demonstrated. Recompile-limit tuning or monitoring is a follow-up.
• The earlier 7–10× spike estimates used a coalesced proxy, and this PR doesn't reproduce them.

🤖 Generated with Claude Code