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 computes grad-x, router gradients, and the six adapter gradients without frozen W1/W2/W3 gradients or full-width side-delta scratch
• use shared A1/A3/B2 and expert-local B1/B3/A2 layouts, with reusable transformed weights and scratch buffers
• zero both per-expert padding and distributed rank-uniform pool suffixes before shared wgrads, preventing sparse routes from contaminating reductions with allocator NaNs
• size default MXFP4 readiness scratch for the minimum supported BLOCK_M=8, and require the MXFP4 saved-down scratch to cover the full rank-uniform pool before any prelude/kernel launch
• fail fast for unsupported post-down side-LoRA backward instead of returning invalid gradients

Correctness coverage

The numerical harness checks BF16 and MXFP4 forward/backward against PyTorch references and independently checks adapter contractions at the exact saved native boundaries. Coverage includes EP1/EP2/EP4, masked and remote routes, empty experts, top-k 1/2/6, SwiGLU/GeGLU, activation clamping, zero adapter scale, reusable outputs, source-combined grad-x/router gradients, and sparse rank-uniform route pools.

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

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

These straddle every BLOCK_M scheduler transition. The original sparse failures at 1 and 15-17 tokens are fixed.

• BF16, 17 tokens/rank: shared native-boundary gradient minimum cosine 0.99999995, maximum relative L2 3.03e-5, maximum absolute error 1.22e-4. The independent PyTorch graph is also reported; at very small route counts its different BF16 reduction order bottoms out near 0.9933 cosine.
• MXFP4 across the 15-point sweep: minimum full-reference adapter cosine 0.9999400; minimum native-boundary adapter cosine 0.99999990; maximum native-boundary relative L2 2.47e-4.
• Zero-scale BF16 and MXFP4 outputs remain bitwise identical to original MegaMoE.

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. The current trunk control emitted nonfinite shared-factor gradients in this production comparison (A1/B2 for BF16 and A1/A3 for MXFP4), so this PR does not manufacture a shared-gradient cosine against that invalid control; the independent native-boundary checks above cover the fused contractions directly.

Performance

B300 exact DSV4 Flash shape: EP4, 65,536 tokens/rank, 262,144 total, H=4096, I=2048, 256 experts, top-k 6. The control and candidate are both side-LoRA, and shared-factor EP reductions are included. Warmup 2, five measured iterations, maximum-rank CUDA latency.

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

The production-relevant F+B comparison applies the identical non-reentrant whole-MoE activation-checkpoint wrapper to the latest-trunk control and the fused FireTitan custom-autograd integration:

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

Under the matched checkpointed contract, native MXFP4 is 2.47% faster than BF16 compute on the trunk control and 4.16% faster on the fused path. The non-checkpointed custom-autograd diagnostic is 84.09 ms BF16 and 87.49 ms MXFP4: with only one forward charged, MXFP4's backward dequantization cost outweighs its forward saving. That diagnostic is not used as the production headline. The separate original-MegaMoE/full-base-wgrad comparison is also diagnostic only and is not used for these speedups.

Tests

• local extension rebuild
• python tests/run_mega_moe_side_lora_edge_matrix.py --ep-processes 2
• exact EP4 BF16/MXFP4 30-case production boundary sweep
• exact 262K latest-trunk FireTitan forward/backward accuracy comparison for BF16 and MXFP4 base compute
• exact 262K BF16 production performance guard, warmup 1, iterations 3
• default MXFP4 scratch-allocation regression and short saved-down rejection regression
• Python compile and diff checks

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

[cursor]

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

^{❌ 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 1564125. 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]

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}