feat(mega): add fused shared side-LoRA kernels - #9
morgendave wants to merge 20 commits into
Conversation
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
f323595 to
3e45539
Compare
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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_mathdefault), 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
There was a problem hiding this comment.
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
1e1b4f7 to
ccaa442
Compare
There was a problem hiding this comment.
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
|
Acceptance audit at exact head
No additional full-model runtime sweep is required for this kernel acceptance seal. |
|
Repeatability attribution is now covered at head
This isolates the small replay drift to the shared atomic dispatch ordering rather than the side-LoRA backward contractions. |
There was a problem hiding this comment.
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.
|
Accepted backward stack is now on this PR branch at 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. |
…pr9-deepjit-align-20260920



2026-09-21: restore wide-K wgrad safely
Kernel heads: ordinary
0be759300fd5d198ef211e452ccb857217c41a1a(#17), side stack48e4df7b5c964257987d2a431b5b3540ecbfc5a7(#9). Integration pins: ordinary8d27eeb21aa48e995fb1a3f0e354190c7a70c7ff, sidea98db521e7003f018c0c5c114ce73dd43a08e157.Implementation
Matched production-shape performance
Same physical 4×B300 GPUs, EP4, 65,536 tokens/rank / 262,144 total, H=4096, I=2048, E=256, top-k=6, side rank=128. Old
d6fbae4, upgraded pre-optimization118a7f9, optimized48e4df7. Each path retains its corresponding integration math/ring policy.F+B is directly timed, not the sum of separate medians. Side rows include shared-factor EP gradient reductions. Each phase: 3 warmups, 3 rounds × 10 measured samples, CUDA events, max rank per sample, then median. These final-head values are separate from the two earlier candidate runs that reproduced the recovery.
Ordinary W13/W2 wgrad+combine dropped from approximately 39.79 → 13.48 ms BF16 and 39.28 → 12.50 ms MXFP4 in separate component diagnostics. The previous 25–31% ordinary F+B regression is resolved at this native-API fixture; this does not certify whole-model E2E.
Accuracy and regression tests
cuMemCreatepermission error; only sanitizer reruns disabled expandable segments. Production accuracy/performance retained their allocator configuration.Scope: native single-layer API/core; excludes base transforms, routing-plan construction, AC, optimizer timing, model E2E, and ARM runtime. Existing full-redispatch MXFP4 gradient-noise limits are not relaxed or declared fixed. The 224 change isolates MegaMoE wgrad only, not every generic DeepGEMM caller.
Prior alignment validation and historical measurements (superseded heads)
2026-09-20: upstream-upgrade alignment (Draft)
Current head:
118a7f9a379d56409e9b7608145d5b02729bb650, based on main02452d4(#15/#16), including the ordinary-training alignment in #17.What changed
side_lora=Truebuffer API without adding the BF16 source plane to ordinary MXFP4 inference/training.Fresh validation
Matched old/new performance (2026-09-21): ordinary backward regression — NOT sealed
Same physical 4×B300 GPUs (4–7), EP4, 65,536 tokens/rank / 262,144 total, H=4096, I=2048, E=256, top-k=6; side rank=128. Old
d6fbae4versus new118a7f9. Each version retains its native layout/ring policy and integration math flags. Ordinary rows compute full BF16-output W13/W2 gradients; side rows compute adapter gradients.Each phase uses 3 warmups and 3×10 CUDA-event samples, max rank per sample, then median. F+B is timed directly, not added from F and B medians. Side B/F+B below include shared-factor EP gradient reductions.
Separate backward-component diagnostics localize the ordinary regression to W13/W2 wgrad+combine, not dgrad:
These diagnostic events run in separate passes, outside primary timing. Source inspection shows old pool BLOCK_M=192 selects wgrad K=64, whereas new BLOCK_M=240 selects K=16. This is a concrete candidate mechanism, not yet an isolated causal proof; do not force K=64 over incompatible expert boundaries.
Scope is single-layer native API/core, not model E2E: no base transforms, routing-plan construction, AC, optimizer, or integration-memory timing. Selected cards were free before each case; the unrelated service on GPU 0 was left untouched. BF16 side being approximately flat does not imply its backward improved. No kernel code was changed to obtain these numbers, and this update does not certify full-model/ARM runtime or close existing repeatability/memory gates.
Reverse-order confirmation (new then old) reproduced the result: ordinary BF16 F+B +31.72%, ordinary MXFP4 +24.71%, side BF16 -0.73%, side MXFP4 -3.98%. Both complete matrices finished successfully; all tested outputs passed the benchmark's finite checks.
Four-mode numerical coverage (2026-09-20)
The previous direct comparison covered side-LoRA only. Coverage now also includes ordinary MegaMoE with full W13/W2 weight gradients.
Matched old
d6fbae4versus new118a7f9: EP4, 4 B300 GPUs per run, 65,536 tokens/rank (262,144 total), H=4096, I=2048, E=256, top-k=6. Side rank=128. Two full F+B repetitions per version/form/precision; all-rank full tensors compared with input/weight SHA-256 checks and FP64 metric accumulation.Direct old/new numerical comparison (2026-09-20)
Compared old
d6fbae484c4cc7068aec00fe77ace9b1e273eb43against new118a7f9a379d56409e9b7608145d5b02729bb650, not two independent cosine-to-oracle numbers.fast_math=True; MXFP4 forwardTrue, backwardFalse. Native tile selection is retained: old BLOCK_M=192, new BLOCK_M=240.||new-old||₂ / ||old||₂. Table values are the worst rank/tensor across both comparisons, not rank averages.BF16 adapter gradients are also bitwise identical, both cross-version and within each version. MXFP4 adapter gradients are not bitwise: maximum absolute difference is 0.0078125 for shared A1/A3/B2 and 0.001953125 for expert-local B1/B3/A2. The same-version control's worst gradient relative L2 is 1.6450903e-4 old / 1.8329473e-4 new. Cross-version differences are slightly larger, but remain at the same small numerical scale; this does not relax the existing 1e-4 redispatch gate or claim universal determinism.
Scope: synthetic single-layer production-shape kernel parity, not a 43-layer checkpoint, AC/optimizer trajectory, full-model performance or memory seal. Accuracy comparisons ran on separate free 4-GPU groups where overlapped; no timing claim is derived from these runs.
Not sealed / do not reuse historical claims for this head
d6fbae4control also fails at ~1.62e-4. Both versions retain bitwise outputs, grad-x/router gradients, and all gradients when saved inputs are held fixed. This control establishes a pre-existing failure of this gate, not a claim of end-to-end bitwise backward.Review order: ordinary alignment #17, then this side-kernel update; FireTitan integration is tracked in fw-ai/fireworks#44111.
Historical implementation, measurements and validation
Summary
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:
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, 2048These 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:
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.
With the identical non-reentrant whole-MoE checkpoint wrapper:
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
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)
main/ fix(layout): support strided PSUM scale packing #114 passedHistorical source head:
b56ba1f57c89e1b6ac228bb9db4d409960e04484(acceptedc4d23b1+ff3ca21stack). 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.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.DG_MEGA_MOE_SM_HEADROOMdefaults to 0 and rounds up to an even reservation; the explicitDG_BF16_MEGA_MOE_NUM_SMSforward override retains precedence. Ordinary dense adapter GEMMs are not indiscriminately capped.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.test_mega_moe_native_side_lora.py+test_mega_moe_headroom.py, CUDA hidden).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.