feat(mega): add fused shared side-LoRA kernels

[morgendave]

Summary

• add dedicated SM100 BF16 MegaMoE and MXFP4-base/BF16-side forward kernels for rank-128 shared side-LoRA
• add a native backward that returns grad-x, router gradients, and all six adapter gradients without frozen W1/W2/W3 gradients or full-width side-delta scratch
• implement the current shared-factor contract: W1/W3 share A, W2 shares B; B1/B3/A2 remain expert-local
• reuse transformed weights and bounded scratch while zeroing expert padding and distributed rank-uniform pool suffixes before shared contractions
• keep the established accurate MXFP4 backward default and explicitly thread SwiGLU/GeGLU, activation clamp, and fast-math selection
• reject unsupported post-down side-LoRA and undersized scratch before launching kernels

Correctness coverage

The numerical harness checks BF16 and MXFP4 forward/backward against independent references and checks every adapter contraction again at the exact saved native kernel boundaries. Coverage includes:

• EP1/EP2/EP4, local and remote routes, masked routes, empty experts, and skewed rank-uniform pools
• top-k 1/2/6
• single token plus BM-1/BM/BM+1 scheduler boundaries
• BF16 and MXFP4 SwiGLU/GeGLU, activation clamp, and zero adapter scale
• reusable output/source buffers, source-combined grad-x/router gradients, default scratch, and short-scratch rejection

The exact DSV4 Flash production-width sweep uses H=4096, I=2048, 256 experts, top-k 6, remote routing, and 20% masked routes. BF16 and MXFP4 cover all 15 token counts per rank:

1, 15, 16, 17, 90, 91, 176, 177, 346, 347, 688, 689, 1029, 1030, 2048

These straddle every relevant BLOCK_M scheduling transition. The exact native-boundary adapter contractions remain above 0.9999 cosine. The independently dequantized PyTorch graph is retained as a separate approximate end-to-end quantization reference and is not substituted for the native-boundary gate.

At the exact 262,144-token EP4 FireTitan control shape:

metric	BF16 base	native MXFP4 base
forward cosine	0.999996719	0.999862155
grad-x cosine	0.999982967	0.999979778
grad-route cosine	0.999997287	0.999996063

All fused tensors are finite. Zero-scale BF16 and MXFP4 outputs remain bitwise identical to original MegaMoE.

Performance

B300 exact DSV4 Flash shape: EP4, 65,536 tokens/rank, 262,144 total, H=4096, I=2048, 256 experts, top-k 6. Control and candidate both use side-LoRA and include shared-factor EP reductions.

base compute	trunk side-LoRA forward	fused forward	speedup
MXFP4-dequantized BF16	52.82 ms	24.76 ms	2.134x
native MXFP4	45.58 ms	22.46 ms	2.029x

With the identical non-reentrant whole-MoE checkpoint wrapper:

base compute	trunk side-LoRA F+B	fused F+B	speedup
MXFP4-dequantized BF16	192.86 ms	114.03 ms	1.691x
native MXFP4	188.21 ms	109.47 ms	1.719x

Timing provenance remains the previously recorded B300 production-shape run; no old timing is relabeled as a post-rebase measurement.

Bounded MXFP4 scale storage

The symmetric arena now sizes scale-factor rows only for BLOCK_M regimes reachable at each live-token count and asserts the launch-time requirement. DSV4 EP4 one-wave provisioning uses 350,208 scale rows rather than the unreachable BLOCK_M=8 worst case. Together with the FireTitan one-wave token ring, the exact production arena falls from 10.794 GiB to 5.606 GiB without changing kernel arithmetic. Direct EP2 repeatability is bitwise exact for fused output and for all gradients when the saved forward boundary is held fixed.

Integration stack

• base BF16/native-MXFP4 FireTitan re-land: fw-ai/fireworks#44095
• generic torchtitan expert-plan hook: fw-ai/torchtitan#142
• stacked private FireTitan side-LoRA integration: fw-ai/fireworks#44111
• superseded public design copy removal: fw-ai/DeepGEMM#10

This PR owns the standalone kernel ABI, edge-case correctness, and kernel-level performance contract. The private integration remains disabled by default and requires a trainer image containing this revision.

Historical validation (before the headroom update)

• clean extension rebuild from current main / fix(layout): support strided PSUM scale packing #11
• compact BF16/MXFP4 forward+backward edge matrix: passed
• exact EP4 BF16/MXFP4 30-case production-width boundary sweep: passed
• host API contracts: 4 passed
• exact 262K FireTitan forward/backward accuracy comparison for both base-compute modes
• Python compile and diff checks

Historical source head: b56ba1f57c89e1b6ac228bb9db4d409960e04484 (accepted c4d23b1 + ff3ca21 stack). The numerical/performance tables above are retained historical measurements, not a new headroom-on benchmark.

2026-09-10: synchronize with main and cover side-LoRA headroom

Current source head: d6fbae484c4cc7068aec00fe77ace9b1e273eb43.

• Merge main 0913be4 (including feat(mega): capture MXFP4 L1 pool across ring wraps #12 capture changes and mega_moe: configurable SM headroom on cooperative mega grids (DG_MEGA_MOE_SM_HEADROOM) #13 configurable SM headroom), preserving the accepted side-LoRA stack.
• Wire all four dedicated persistent launch sites through the shared headroom policy: BF16/MXFP4 forward and BF16/MXFP4 backward. A plain main merge alone did not cover these side-specific launchers.
• Preserve main's semantics: DG_MEGA_MOE_SM_HEADROOM defaults to 0 and rounds up to an even reservation; the explicit DG_BF16_MEGA_MOE_NUM_SMS forward override retains precedence. Ordinary dense adapter GEMMs are not indiscriminately capped.
• Preserve the existing fused-tail synchronization and saved-expert-input reuse. No Merge upstream DeepGEMM through PR #432 #15 DeepJIT/scheduler/layout migration, new ring sizing, or trainer pin change is included.
• CPU regression: python tests/test_mega_moe_headroom.py -v — 4 tests passed, including 11 host-policy cases compiled from the actual C++ helper and launch-wiring checks.
• Clean extension build against the existing training ABI: passed.
• CPU/API suite: 10 tests and 11 parameter subtests passed (test_mega_moe_native_side_lora.py + test_mega_moe_headroom.py, CUDA hidden).
• B300 headroom=8, NCCL min/max channels=8, NVLS disabled: all 11 compact BF16/MXFP4 F+B edge cases passed (EP1/EP2, SwiGLU/GeGLU, masks, empty experts, remote routes, scratch contracts, and zero-scale base preservation).
• EP4 production-width BF16 and MXFP4 F+B: both passed at H=4096, I=2048, 256 experts, top-k=6, rank=128, 2,048 tokens/rank (8,192 total), remote routing and 20% masked routes, headroom=8. This is a production-width oracle, NOT a 262K/full-model run. Minimum native-boundary cosine: 0.999999969 for BF16 shared-adapter contractions; 0.999999937 for all six MXFP4 adapter contractions. The approximate independently dequantized reference remains separate from this native-boundary gate.
• EP2 MXFP4 saved-x reuse/repeatability: passed with headroom=7 (rounded to 8), H=1024, I=512, 8 experts, top-k=2, 128 tokens/rank, remote routing and 20% masked routes. With fixed saved inputs, all six adapter gradients, grad-x and router gradients are bitwise identical. Repeating the entire forward/backward keeps outputs, grad-x and router gradients bitwise identical; adapter gradients are NOT bitwise identical across redispatch (maximum relative L2 6.80e-6, maximum absolute difference 3.82e-6), within the existing 1e-4 repeatability gate.
• Total fresh GPU validation: 14 passing cases (11 compact + 2 production-width + 1 repeatability). Generated side-kernel JIT sources confirm the reduced 140-SM configuration on B300.

Headroom remains a deployment-configured coexistence measure, not a guarantee against arbitrary concurrent kernels. This update does not claim a new NCCL overlap stress-test or full-model performance result.

Private trainer integration and trainer-dependent E2E sources remain outside this public repository.

---

Note

High Risk
Large new GPU training paths (distributed symmetric buffers, MXFP4 quantization, backward aliasing/memory modes) touch core MoE forward/backward; regressions would show up as silent wrong grads or OOM from buffer sizing.

Overview
Adds fused MegaMoE + rank-128 shared side-LoRA paths on SM100 for BF16 and MXFP4 base / BF16 adapters, exposed through new C++/pybind entry points and Python mega APIs (forward, backward, weight transform, and an extra symmetric-buffer side_lora_source plane for the MXFP4 specialization while keeping the legacy buffer ABI).

The backward stack is a multi-phase native pipeline (persistent dgrad wave plus grouped/dense GEMMs, grad-x publish, adapter wgrads, padding clears) that returns grad-x, router grads, and six adapter gradients without full-width side scratch or frozen expert W13/W2 wgrads in the side-LoRA contract.

MXFP4 scale-factor ring storage no longer assumes every candidate BLOCK_M fills the full token ring; sizing follows reachable live-batch regimes via get_num_max_required_sf_ring_tokens_for_mega_moe, with a launch-time assert in heuristics—reducing symmetric arena footprint for production shapes.

JIT builds also pick up an extra CUTLASS include path, and sm100_bf16_mega_moe_wgrad_1sm gains optional row-strided inputs and a custom kernel name for side-LoRA wgrad reuse.

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

[cursor]

Stale comment

Risk: high. Not approving: Cursor Bugbot finished skipped/neutral with 2 unresolved high-severity findings (undersized MXFP4 ready buffer; unsafe down-unweighted scratch reuse), and this ~16k-line fused side-LoRA CUDA change exceeds the low-risk auto-approval threshold. Human review is required for the MXFP4 buffer contract and backward scratch aliasing.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: high. Not approving: Cursor Bugbot completed skipped/neutral with 2 unresolved high-severity findings (undersized MXFP4 ready buffer; unsafe down-unweighted scratch reuse), and this ~16k-line fused side-LoRA CUDA change exceeds the low-risk auto-approval threshold. Human review is required; existing reviewer requests were left in place.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: high. Not approving: Cursor Bugbot passed with prior findings resolved and Security Agent was not running, but this large fused side-LoRA CUDA training change exceeds the low-risk auto-approval threshold. Human review is required; existing reviewer requests were left in place.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: high. Left a non-blocking comment and did not approve: Cursor Bugbot completed as skipped with 2 unresolved findings (MXFP4 GeGLU backward; missing post-down fail-fast), and this fused side-LoRA CUDA training change exceeds the low-risk auto-approval threshold. Human review is required; existing reviewer requests were left in place.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: high. Left a non-blocking comment and did not approve: Cursor Bugbot completed as skipped with an unresolved medium finding (inconsistent MXFP4 fast_math default), and this fused side-LoRA CUDA training change exceeds the low-risk auto-approval threshold. Human review is required; existing reviewer requests were left in place.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Stale comment

Risk: high. Left a non-blocking comment and did not approve: Cursor Bugbot passed with prior findings resolved and Security Agent was not running, but this fused side-LoRA CUDA training change exceeds the low-risk auto-approval threshold. Human review is required; existing reviewer requests were left in place.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Risk: high. Left a non-blocking comment and did not approve: Cursor Bugbot completed as skipped with 1 unresolved high-severity finding (BF16 adapter grads use wrong planes), and this fused side-LoRA CUDA training change exceeds the low-risk auto-approval threshold. Human review is required; existing reviewer requests were left in place.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[morgendave]

Acceptance audit at exact head ccaa442e9f71ba7a83e9acafd21189034536ebad:

• Performance is sealed by the production-shape B300 tables in the PR body. The controls are the correct same-precision trunk side-LoRA paths; these are retained measurements, not relabeled reruns.
• Accuracy is sealed by the exact 262,144-token EP4 table plus the native saved-boundary gates. All fused tensors are finite; zero-scale BF16/MXFP4 remains bitwise equal to original MegaMoE.
• Exact-head B300 DeepGEMM matrix: 41/41 cases completed, including 30 BF16/MXFP4 DSV4 production scheduler-boundary cases; zero traceback/assertion/CUDA/runtime error signatures.
• Review audit: 0 unresolved threads (6/6 resolved); no outstanding DeepGEMM kernel bug report.

No additional full-model runtime sweep is required for this kernel acceptance seal.

[morgendave]

Repeatability attribution is now covered at head 1b92209.

• MXFP4 side-LoRA backward repeated on identical saved forward boundaries is bitwise exact for all six adapter grads, grad-x, and route grad.
• Full side forward+backward replay changes atomic dispatch row metadata (representative exact fraction 45.90%) but keeps adapter-grad relative L2 below 1e-4.
• Ordinary MegaMoE forward shows the same row reorder (45.16% metadata exact) while its final output is bitwise exact.
• The new --check-repeatability regression asserts those contracts; 4xB300, EP4, H=4096, I=2048, 256 experts, top-k 6, 2,048 tokens/rank passed.

This isolates the small replay drift to the shared atomic dispatch ordering rather than the side-LoRA backward contractions.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

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

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

[morgendave]

Accepted backward stack is now on this PR branch at b56ba1f57c89e1b6ac228bb9db4d409960e04484: c4d23b1 (safe fused side-LoRA backward tail) followed by ff3ca21 (reuse saved expert input). The only tree difference from the original ff3ca21 stack head is the pre-existing expanded legacy-buffer ABI regression test. 28992c2, b0d0dcc, and the unsealed 11-loop candidate are not ancestors of this head.

Validation: host contracts 4/4; clean SM100 wheel; exact 4xB300 BF16 EP4 oracle at 262,144 total tokens/top-k 6/rank 128 is finite with min adapter-grad cosine 0.999942746, native-boundary shared-grad cosine 0.999999075, and grad-x cosine 0.999986239. FireTitan integration additionally passed two repeated 3-step BF16 F/B+AdamW trials bitwise exactly after production shared-factor gradient synchronization.