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

[morgendave]

Current review status — 2026-09-30

Draft. Includes merged DeepGEMM main6ac0f15.

Published DeepGEMM head b8e547693fd9e8f230bf87f171c0961df04a0b02 closes the strict-repeatability failure with stable source-rank-local slots only for side training (kHasSideLora && kSaveL1Preact), explicit warp reader/writer synchronization, parallel prefix when scratch fits, and exact stable-scan fallback. Round-robin transport and ordinary/inference atomic dispatch remain unchanged.

A separate oracle bug was corrected: shared A1/A3/B2 reference leaves now accumulate in FP32, with separate BF16 cast nodes per expert use. Forward operands and all thresholds remain unchanged. The activation-rounding-only ablation failed and is retained, not published. The original unchanged strict262K oracle already passed before this reference fix, so repeatability closure does not depend on reference modification.

Final committed-source GPU suite: 6/6 exit0—strict EP4/262144 total tokens; wide-expert shared-counter fallback; actual metadata-capacity fallback; BF16; production forward=True/backward=False; and inference save-L1 false/true bitwise parity. All worker assertions passed; only printed rank0 metrics are claimed. Intermediate reference cosines use their existing thresholds, not a blanket >0.9999 claim. CPU shared-gradient/combined-graph regressions:13 passed.

Local provenance: existing immutable image sha256:922f688e045efcfee26ee096c52af49bafb4348dbe796781467e30d9cf8ba78d, installed003b native extension plus current Python wrappers/JIT header, not a rebuilt b8e local wheel. Header SHA 9f8530f32532aa0fccad50d3b5ad322177ed5bbe355fef453b0fc6d157e3d853; oracle SHA a1dcd22898a5488d3a2c0daf7e578ff87221a751f08909f5ac93e93bce03c220. Fresh-process header hashing and a new strict CUBIN compiled after launch were verified. Full logs: stable-slot-final-committed-gpu.log, stable-slot-final-cpu-combined-graph.log.

Matched native-API ABBA, bee1755→b8e, EP4/T65536 per rank, H4096/I2048/E256/top-k6/rank128, forward=True/backward=False, saved-X reuse, same image, 20 samples ×2 per lane: median forward 17.255408→17.394560ms (+0.8064%), F+B 71.122879→71.246769ms (+0.1742%). All4 runs exit0 and pass finite/boundary/full-reference assertions; no OOM/foreign arrival. Evidence: root-stable-slot-perf-abba-20260930.json. This is not E2E/full-model timing or a replacement for historical speedup measurements.

Integration / release gates

Consumer status at 17:54 UTC: ordinary Fireworks#50925 published585494c is ready for review; side#44111 published618c352 remains draft, and both preserve frozen mainbaffda7d. These heads only add the CPU source-overlay CI fix over c164/f4dc. Ordinary coordinator36750640778 passed CPU1691/51 skips +72 isolated and wheel; prepared GPU workspace is synthetic merge39ceb91 (585 into main1c820801), not head-only. DeepSeek-V4 shard36752517117 passed33/40 skips/60 deselected; GLM-5.2 shard36752520153 passed53/29 skips/4 deselected; general36752513651 failed (3726 passed/74 skipped/1 failed plus distributed failures): generic CI DeepGEMM lacks the new layout API, and the zero-rank worker imports TorchTitan's tests package. Skips are not numerical coverage. Refreshed ARM builds36749151989/36749154150 passed installed import smoke (arm64 digests77322bff…/5820e8e9…). Ordinary public runtime36751701917 passed on exact77322bff:13 tests, requests12+11 tokens, public forward/F+B/optimizer step1; required-Helm cleanup41→48 completed17:56:47UTC for trainerp8st8ksmy4f3cvyp, terminal52. Not262K request/performance coverage; no side runtime. A confirmed consumer-only Torch2.14 private-detector bug has local unpushed two-import fixes ordinaryf1863e9/side4a46cae:26 focused CPU cases pass each, ordinary full CPU1699/46 skips +72 isolated. All shards completed; publication now waits for a scoped CI dependency/import remedy decision, not a test waiver. These production Python fixes are absent from the ARM artifacts; they do not alter this DeepGEMM head/native math. Pinned2.11 delegation and mocked-callback2.14 validation remain separate.

Fireworks latest-main refresh: ordinary c164c68a2171e94c57663da63431af6ab29fb20e, side f4dca68b97fd608b259966bcd3a5a60cc49eb505. Main's shared MXFP4 decoder is adopted while side's existing local W13 interleaved-store kernel is preserved unchanged. Focused source-overlay dequant: side CPU11pass/38CUDA-skips and GPU49pass/no-skips/noOOM/no foreign arrival; ordinary CPU7pass/31CUDA-skips. Image sha256:871b71044633b9546ea485e6d5a58dc07b412e30d82610dbb3feec856942b039, current FireTitan/kernels package mounted. This does not rerun native acceptance or rebuild the PR images; earlier image evidence below retains its exact old source. New main also includes upstream RDMA changes, so fresh image/runtime gates remain open. Logs: side-latest-main-mxfp4-{cpu-03,gpu,supervisor}.log, ordinary-latest-main-mxfp4-cpu.log.

Fireworks side #44111 f4dca68b97fd608b259966bcd3a5a60cc49eb505 pins b8e in its isolated deps Dockerfile and Makefile. Both Fireworks heads include main baffda7d10490bc980a9978f69a8392c534d5b62.

Earlier e836/4c887 exact-head CPU 36662914940 passed:95 consumer cases +8CUDA-only skips per head,45 ordinary /48 side workflow contracts,2 wheel tests each. Skips are not hardware acceptance. Corrected lifetime instrumentation, fail-closed FSDP compatibility and construction-only fake-math fixture isolation preserve real FSDP/DTensor/AC/optimizer and dependency-rejection checks.

Earlier e836/4c887 composed GPU follow-up passed: all5 fixture commands exit0—8 byte-exact BF16 checkpoint VJPs; DSV4+GLM BF16 full-AC/FSDP two-step parity with bitwise outputs;21 fused-MXFP4-merge/expanded-FP8 consumer cases; side BF16 three layers ×3 AC/FP32AdamW steps with13 trace captures per rank; side MXFP4 H1792/I768 three layers ×3 AC/AdamW steps with disabled/category-masked/all-masked routes. No OOM or foreign GPU arrivals. Sources tested: ordinary e836, side4c887 with b8e native wrapper/header; later side d6b0036 changes only the external auth harness, so model/kernel paths remain identical. Existing immutable local images with current source/helper mounts are not rebuilt current-head images. Logs: current-ordinary-integration-gpu.log, current-side-integration-gpu.log and matching supervisors.

Earlier ARM native deps b8e 36662911986 passed (wheel2.8.0+b8e5476; ARM sha256:1830de36b3e337e98da105547145ed4b0cf48be996af25d41f05f9a30260bb99). Side trainer 36663689528 passed, including installed canonical BF16/native imports; ARM sha256:47ee1b9ab851b48086c366e7515ca06f0555689c37169f96579c99961afe8aa8, index sha256:1d3155ff7d1fc4aba304a88c807f7ce06ee701b7fb13fe099c64242ec60ace12. Built-for-e836 ordinary trainer 36663008182 passed, including installed canonical BF16/native dependency closure smoke; ARM sha256:a1d0d29b78b08a57c72ba0731f9c8216b7e476e19133fe9c0121a4532dcce169, index sha256:500813d8f64245c3355414037476a662c074f665787384aa20dd4f017442736a.

Older d3630/main47 ordinary runtime 36661470294 reached readiness (Temporal event19 at03:19:23UTC), then failed with sanitized AuthenticationError before public operations. Exact CI Tinker0.16.1 offline wire test reproduced duplicate API-key headers from a case-sensitive merge; canonical X-API-Key fixes it. Regression failed before and full harness passed12 tests/1runtime skip afterward; published side d6b0036 changes only this external harness. Actual required-Helm cleanup completed event43 at03:20:19UTC, not a DELETED-only waiver. One auth-fixed serialized retry 36664741744, validation635d264fa10, created trainer bugxxuflwguaz5dn at03:33:14UTC against the then-current e836 ordinary ARM image above, but failed before readiness at03:51:16UTC with explicit JOB_STATE_FAILED (12 local harness tests passed,1runtime failed). No public operations ran. Temporal confirms readiness activity17 failed at03:51:08 with Pod trainer-training-rlor-bugxxuflwguaz5dn-0 failed: unknown (PREREQUISITE_NOT_MET, non-retryable). Required-Helm cleanup23→25 completed at03:51:12UTC with PreserveDebugInfo=false; cleanup is proved. Direct pod get/log/events access is denied, but existing read-only GCP Logging recovered202 exact-pod events: 25 distinct pod UIDs and25 explicit Preempted events, from03:34:16 through03:51:11UTC. Preemptions at03:51:06 and03:51:11 bracket the Temporal failure. This establishes repeated preemptible scheduling/capacity interruption; it does not pass runtime acceptance or establish a kernel/auth failure. Evidence: ordinary-bugxx-central-events.json and ordinary-bugxx-preemption-proof.json. Full log: ordinary-e836-auth-fixed-runtime-full.log. Side runtime remains unallocated pending capacity coordination; no blind retry or changes to region/reservation/priority. Region remains omitted/backend-selected. No public-operation acceptance yet; configured262K capacity is not a262K public request. Other matrix jobs with skipped work steps receive no hardware credit.

Remaining: fresh current-head Fireworks image/build verification, completed exact-image ARM public F+B/optimizer acceptance, existing matched full-model performance/memory acceptance, and fresh human review. DeepGEMM #9 remains unmerged; side must repin the merged SHA after separately authorized human merge. No approval, merge, rollout or tolerance waiver.

Superseded acceptance summary (frozen; current status above governs)

Summary

Native BF16/MXFP4 MegaMoE side-LoRA forward/backward, including shared A1/A3/B2 factors, expert-local B1/B3/A2, bounded scratch and native adapter gradients. Latest published fix makes public side backward BLOCK_M rank-uniform before native launch; callers may explicitly retain a proven uniform forward tile.

Published head: bee17552ffa4cd0d10a2ead4885eed824a720587. Includes DeepGEMM main 6ac0f152af8fd22dce9bb3445d2eda4afeb57263. Draft; clean mergeability is not completed acceptance. No merge or fresh human approval claimed.

Dependencies and review order

Review this kernel PR before Fireworks #44111, which explicitly pins this unmerged SHA. Ordinary Fireworks #50925 uses merged DeepGEMM 6ac0f15 and is independent of this side-specific numerical gate.

Acceptance at the published source

Check	Result and scope
Rank-uniform review fix	10 CPU kernel contracts + 6 integration contracts passed; review thread resolved
Real NCCL / capture	EP2 empty/short/full local-token-derived tile reduction and explicit uniform CUDA graph capture/replay passed
BF16 and MXFP4 native F+B	EP2 injected backward tiles 16/128 passed; uniform forward fixture, not full uneven-input kernel certification
Exact local provenance	Existing immutable003b extension with read-only bee1755 backward.py; native C++/CUDA unchanged, not a rebuilt local wheel
ARM deps / trainer	Actual builds 36648021683 and 36648850649 succeeded; aarch64 wheel2.8.0+bee1755, source assertion passed
ARM artifact identity	Trainer tag dev-megamoe-side-bee1755-c6b7f9ef-20260929-arm; OCI/GAR index sha256:a0d391dd505bdcb249cb508d4e2bbe78d3f4be6e86f02bef3a7cf8e1f75faa9c; ARM manifest sha256:5a4133f50a97145b4fc789415700a4ee100d4c4b7bb03664000babe2d50be086
Strict 262K repeatability	Open / historical FAIL at unchanged <1e-4. Current head not recertified
Exploratory timing / separate failure	Stable-slot remains an unpushed strict-repeatability candidate, not an accepted optimization. Two short EP2 production-math pairs: forward+2.76%/+3.07%, F+B−1.48%/−0.50%, no matched allocation delta. Different GPU pairs, not pooled. Both baseline and candidate exit1 at unchanged full-reference cosine>0.9999 (~0.99989585); finite/native-boundary checks passed, not full correctness. Omitted-reuse run excluded and retained.
Small causal diagnostic	EP2 two redispatches: canonical operands identical; all native contraction replays exact; row-order-only replay reproduced every drift; canonical replay exact. Mechanism evidence only, not262K closure
Review	Zero unresolved threads refreshed; historical approval does not certify current head. Draft skips are not hardware passes

Failure ledger (exact heads/images, failures, control attribution and remaining work): /shared/morgendave/megamoe-pr-ready-20260929/failure-ledger.json. Full logs, immutable image IDs, commands, boundary captures and source hashes: /shared/morgendave/megamoe-pr-ready-20260929/validation-metadata.json and adjacent logs. Earlier multi-layer AC/optimizer and production-shape evidence is retained below with original heads; not relabeled as fresh acceptance.

Remaining gates

• Code/numerics: strict262K remains failed at unchanged <1e-4. Local stable-slot prototype has small correctness evidence only; no262K recertification. Separate unchanged short full-reference cosine failure occurs on published baseline and prototype and remains unwaived. Historical strict oracle uses forward=False/backward=False; production-math timing uses True/False and must not be mixed into that repeatability fixture.
• Platform/runtime: ARM runtime is unaccepted. Previous Taiwan trainer accounts/training/rlorTrainerJobs/ncrt470fjzihociy is a retained DELETED row and its shape is404, but actual GPU workload teardown still lacks independent proof. No new remote allocation or DELETED-only bypass.
• Packaging resolved, not runtime acceptance: original exact-artifact probe36652678355 proved helper omission. Trainer-only builds36654263588/36654261840 and independent exact-image probe36656018004 now pass both artifacts' installed module identity, canonical helper hash and actual DeepGEMM closure. Current side ARM sha256:39ff45929fece8b2704f5718198d9a35163899dcd4a3474817271322a48eb468, existing native deps unchanged; no GPU/source copies in probe.
• Composed acceptance: matched current full-model performance/memory and affected shared-consumer GPU coverage. Historical CPU17failures are5merged_w1 KeyErrors,4storage-pointer KeyErrors and8checkpoint _user_hooks AttributeErrors, not one proven cause. Corrected tracking now passes45policy CPU cases on both exact production ARM2.11images (36657101050), with8CUDA skips and no hardware credit. Separate declared Torch2.14CPU compatibility failures remain under validation-only repair; main reproduction is not acceptance.
• Review: human current-head review, then dependency/integration review. No merge, tolerance waiver, skipped-job acceptance, or rollout claimed.

Historical evidence and progress (frozen snapshots; current status above supersedes old head/progress claims)

2026-09-30: rank-uniform backward fix; acceptance still incomplete

New bounded numerical mechanism evidence (not 262K closure)

One small EP2 MXFP4 boundary capture/replay completed at T128/rank, H1024/I512, E8, top-k2, rank128, remote routes, 20% masks, clamp10, pre-down, forward fast_math=True / backward=False. On both ranks, the two identical-input redispatches produced different physical row orders but all eight canonically aligned BF16 contraction operands were bitwise identical. Actual prelude dYe was cloned before scratch reuse and verified against source grad_y. All six native contraction replays matched original native gradients bitwise. Changing only first-capture row order to the second reproduced every observed gradient difference (max relative L2: rank0 3.8310434e-5, rank1 2.9796895e-5); fixed-order value differences and canonical-order repeats were exactly zero. This establishes a row-order cause for this fixture only, not for the original 262K failure.

No production/native kernel edit or port was made. The diagnostic is an uncommitted change to the existing oracle; exact-source patch, full log row-order-mechanism-gpu-002.log, and 27 MiB of boundary captures/reports under row-order-mechanism-20260930-001 are retained in the evidence directory. Initial entrypoint-only failure001 occurred before Python/GPU work; corrected002 is the only actual diagnostic run. Existing immutable003b extension plus bee1755 Python wrapper provenance applies. No timing claim. A possible minimal follow-up is stable MXFP4 side-training source-slot assignment, retaining deterministic round-robin transport; no automatic rank-major port. Two extra prefix/grid-sync phases may cost latency, with prefix scratch reused from metadata, so performance must be measured. The strict262K gate, ARM runtime and current matched full-model acceptance remain open.

Current ARM deps refresh succeeded: 36648021683, native aarch64 deep_gemm-2.8.0+bee1755, source assertion passed; OCI/GAR index sha256:d63965a4ea96dfd9f65fd823d53ff2f5cd7fd547c31b4a896227d17351ee7869, ARM manifest sha256:79778882c996cb9f75fbed13e0aefbb6abb0ac4308e12a0e42a1c1b50f368666. Trainer 36648850649 is building the same pinned validation source and unique tag; no runtime allocation.

Current published patch and previous focused acceptance

Current kernel head bee17552ffa4cd0d10a2ead4885eed824a720587 includes DeepGEMM main 6ac0f152af8fd22dce9bb3445d2eda4afeb57263. Composed #44111 head ff6ea894f8b8b5c1517c10d0c153ec06ea77ca82 includes Fireworks main snapshot a5f4e8ec250221b7893e0be9ebcd201d37ec79fb. Latest main-only refresh does not change scoped training/kernel sources.

Fixed review finding 4138744373: both side backward wrappers now MAX-reduce rank-local BLOCK_M before native prelude/wave by default. Keyword-only rank_uniform_block_m=True is an explicit caller contract for a retained, proven rank-uniform forward tile; FireTitan supplies that tile, including activation recomputation. No C++/CUDA math change. Review reply 4139508949; thread resolved, zero unresolved threads at this snapshot. Historical approval is not fresh approval of this head.

Focused verification: 10 CPU kernel contracts + 6 FireTitan CPU integration contracts passed. Real two-rank NCCL tile-uniformity cases (0,4096), (7,4096), (4096,4096) passed, as did explicit fast-path CUDA graph capture/replay. BF16 and MXFP4 EP2 native F+B oracles passed with injected backward local tiles16/128, remote routing,20% masked slots and saved-X reuse (T128/rank,H1024/I512,E8,topk2). Forward fixture remains uniform: this is not full uneven-input/empty-source native-kernel certification. Initial helper launch failed Docker hostname rendezvous; explicit localhost retry passed.

Exact local provenance: immutable image sha256:922f688e045efcfee26ee096c52af49bafb4348dbe796781467e30d9cf8ba78d, installed003b9dd C++ extension, newbee1755 backward.py mounted read-only; C++/CUDA tree unchanged. Not a rebuilt bee1755 local wheel. Correctness only, no timing claim. Full logs in /shared/morgendave/megamoe-pr-ready-20260929: kernel-block-m-contracts-final.log, side-block-m-integration-cpu.log, side-block-m-nccl-capture-gpu-retry.log, side-block-m-{bf16,mxfp4}-gpu.log.

Current native ARM artifact refresh: deps 36648021683, validation source 2f5f5490042fea2eb90999b28e68192045270869, unique tag dev-megamoe-side-bee1755-c6b7f9ef-20260929-arm; build in progress. Older003b artifacts below remain historical. No ARM runtime acceptance or additional allocation: prior trainer tombstone does not establish actual GPU teardown; independent workload absence/successful cleanup evidence is required.

Strict262K gate remains failed at unchanged <1e-4; current head is not recertified. CPU metric-precision audit of independent retained full adapter captures confirms both old/new self-repeat cases exceed1e-4 using FP32 and FP64 norms (maximum metric change about1.45e-11); metric accumulation precision is not the explanation. These are separate captures, not the original in-process failure tensors. BF16 stable route compaction versus MXFP4 atomic destination ordering supports a possible reduction-order mechanism, not an isolated causal proof. Existing oracle calc_diff is not strict norm-relative L2 and cannot replace that gate. This remains an engineering/numerical blocker, not merely pending approval.

Matched current full-model performance/memory, ARM runtime and fresh human review remain open. PR stays draft/unmerged; #44111 pins this unmerged dependency. No tolerance relaxation or skipped-job acceptance. This section supersedes older current-head/progress statements below; historical evidence is retained with its source provenance.

---

Historical composed freshness: #44111 head a759279e40d608103eb0981939f380dd9ad31c31 includes Fireworks main b7f286837d121423bf78fcf81c63bb15b96575dd; latest two main commits are unrelated DB artifacts/perfagent. Kernel head and all numerical/runtime acceptance limits below are unchanged. #9 stays draft and unmerged.

2026-09-29 23:32 UTC: final composed status — no runtime waiver

Kernel head remains 003b9ddfcef7d56d7ebce04d90a389bb84d7ee75, main 6ac0f152af8fd22dce9bb3445d2eda4afeb57263, no new kernel edits. Current integration #44111 9b0175d28c0eb0a0328903e24d960ad156f0b920 includes Fireworks main c68ac3aad25da7fc8e054f5cb681bc2e03cee5af; refreshed CI/client contracts54passed1runtime-skip. Production model/kernel tree unchanged; new main-owned test-harness behavior is not relabeled as earlier GPU coverage.

Exact native ARM deps/trainer builds below succeeded. Corrected runtime 36644573237 failed before trainer creation on prior-resource absence verification. Read-only observation 36645277143 confirms the old trainer is a retainedDELETED row and its shape is404, but actual GPU workload absence/successful cleanup-activity evidence remains unavailable. ARM runtime not accepted; no additional remote allocation. See #44111 for full provenance and infrastructure blocker; this is not a kernel failure.

Strict262K gate still fails its original threshold as detailed below. Matched full-model performance/memory and fresh human review remain open. No waiver, merge, approval or dependency repin claimed.

---

2026-09-29 23:20 UTC: exact ARM artifacts built; runtime still pending

Kernel head remains 003b9ddfcef7d56d7ebce04d90a389bb84d7ee75; DeepGEMM main remains 6ac0f152af8fd22dce9bb3445d2eda4afeb57263. No kernel edit since the validated tree below. Current composed #44111 head is 9f9dd5fcbd6e12b99f679a0f4d5ca4a8d7cbfee0, merged Fireworks main e3b1c842a0f42c60a5126ba9def7af8e57e91dfa with relevant model/kernel sources unchanged.

Actual native ARM deps 36637332135 and trainer 36638653713 builds succeeded. Logs verify aarch64 DeepGEMM 2.8.0+003b9dd and source assertion. Trainer dev-megamoe-side-643983b8-20260929-arm, immutable OCI/GAR index sha256:59c204e3d0718ccb5e8aeb1dfef72d1b26fbd46613fdd0265d3e94d7e6f41d58, ARM manifest sha256:e440863a0f47b4216e01c6632a3e72ce4d98be6c37b6a2af365e45ded4bbaf02.

First side runtime reached RUNNING but its deprecated-route client waiter never submitted public operations. Cancelled run36640664360 performed its own cleanup successfully. Fixed gateway-proxy client has 6 CPU passes,1 runtime skip. Corrected bounded runtime 36644573237 uses the same immutable artifact and requires prior trainer/shape404 before another GPU allocation. No ARM runtime acceptance yet. Full provenance/status is maintained on #44111; skipped matrix wrappers are not tests.

Strict262K gate below remains failed at the unchanged threshold. Full-model performance/memory acceptance and fresh human review remain open. #9 is unmerged; #44111's pin explicitly targets this PR head. No merge, approval, tolerance relaxation, or waiver claimed.

---

2026-09-29: merge-main refresh and explicit remaining acceptance gates

Head 003b9ddfcef7d56d7ebce04d90a389bb84d7ee75 merges DeepGEMM main 6ac0f152af8fd22dce9bb3445d2eda4afeb57263 (merged #17). Tree identical to previously tested 0cbfc76f5306a41e2b28c031734ddb80e0c90e45; no kernel behavior edit. Source contracts rerun: 6 passed.

Exact-source C++ extension rebuilt using FireTitan's isolated MegaMoE Dockerfile. Local side overlay sha256:922f688e045efcfee26ee096c52af49bafb4348dbe796781467e30d9cf8ba78d labels this exact003b source. Composed with refreshed FireTitan #44111 643983b8cb72086914cd7deba919a88c8312cb22 / merged TorchTitan deepseek-ai#142 035f2754: actual 2×B300, BF16 and MXFP4(H1792/I768),3layers×3steps, non-reentrant AC, FP32-master/state AdamW all passed exit0; BF1613traces/rank,23-site schema; frozen-base, masks and zero-route coverage passed. Full commands, base digest and scope in #44111 September29 section. Logs /shared/morgendave/megamoe-pr-ready-20260929/{kernel-source-contracts,side-overlay-build,side-bf16-ac-gpu,side-mxfp4-ac-gpu}.log.

Strict 262K repeatability gate remains failed — no waiver

Historical EP4, 65,536 tokens/rank (262,144 total), H4096/I2048/E256/top-k6/rank128 full-redispatch same-version repeat results:

Historical source	Maximum adapter-gradient relative L2	Max absolute tensor difference	Existing strict <1e-4 gate
old d6fbae484c4cc7068aec00fe77ace9b1e273eb43	0.00016242989638682117	0.0078125	FAIL
candidate 118a7f9a379d56409e9b7608145d5b02729bb650	0.00016423464036255476	0.0078125	FAIL

Absolute difference between the two self-repeat maxima: 0.00000180474397573359. This is not a direct cross-version tensor-difference metric. Outputs/grad-x/router were bitwise; fixed-saved-boundary backward gradients were exact. Both historical logs end in AssertionError at the unchanged threshold. These heads are not current 003b9dd; this is not a current-head pass. Reviewer disposition remains required, with no invented waiver or threshold relaxation.

Raw logs: /shared/home_workdirs/morgendave/megamoe-upstream-align-0920/pr9-old-side-mxfp4-262k.log and pr9-side-mxfp4-262k.log; provenance in pr9-body-upstream.md, pr9-body-parity.md, and parity-report.md in that directory.

ARM exact-head runtime, matched full-model performance/memory and current-head human review remain open. September29 tests are focused B300 composed-stack correctness, not fresh262K strict repeatability/full-model performance certification. Earlier reviews cover older heads. No merge, approval, tolerance relaxation, default rollout, or waiver claimed. #44111 explicitly pins this immutable PR head, not a merged-main dependency.

---

2026-09-25: address review correctness findings

Current side-kernel head: 0cbfc76f5306a41e2b28c031734ddb80e0c90e45.

• Fix side trace allocation/validation at 23 sites, keeping ordinary BF16 backward at 22; the side wrapper now forwards its correctly sized trace tensor.
• Reject write_grad_x_pool=False in Python and host code before launch, rather than silently forcing writes into a caller's zero-sized pool. Default is explicitly True.
• Explicitly require prepared gate/up inputs for the MXFP4 side specialization; the unsupported recompute specialization is rejected at generation and compilation.
• Initialize scheduler valid_value; stride the stable prefix pass over experts beyond the dispatch-warp count; validate the complete side symmetric-buffer extent.
• Remove the two duplicate Python direct-grad-x plane clears. Retain the grad-y producer completion synchronization and the publisher's masked-slot clear + NVLink barrier; no host-sync deletion or integrated-prelude experiment.

Fresh final-source validation

• C++ extension build passed against trainer ABI 0.638.490; 6 source-contract tests passed.
• BF16/MXFP4 EP2 native F+B numerical oracles passed: T128/rank, H1024/I512, E8, top-k2, remote routing, 20% masked slots, clamp10, saved-X reuse. Existing tolerances unchanged.
• Separate zero-scale tests in both precisions: exactly equal ordinary MegaMoE output, cosine=1, relative L2=0.
• MXFP4 fixed-saved-boundary repeat: all six adapter gradients, grad-x and router gradients exactly equal. Full redispatch on this small case: adapter max relative L2 3.05873e-5, within the existing 1e-4 gate. This does not close the previously failing 262K gate.
• BF16 1024-expert regression passed at EP2, T256/rank, H2048/I1024, remote routing and 20% masks, exercising experts above the 592 dispatch warps. The initial H1024 candidate was correctly rejected by the existing scratch-capacity static assertion; it is not counted as passing.
• Companion integration: two B300 GPUs, 3 layers x 3 non-reentrant AC/FP32-master AdamW steps passed for BF16 and unequal-width MXFP4 (H1792/I768). Includes actual parameter updates, frozen packed-weight preservation, disabled/masked adapters and entirely masked routes. BF16 trace captured 13 launches/rank with populated site 22.

No new full-model/262K performance, ARM runtime or immutable-image claim. Shared-core deduplication remains a follow-up; existing rollout and repeatability gates stay open.

2026-09-21: restore wide-K wgrad safely

Kernel heads: ordinary 0be759300fd5d198ef211e452ccb857217c41a1a (#17), side stack 48e4df7b5c964257987d2a431b5b3540ecbfc5a7 (#9). Integration pins: ordinary 8d27eeb21aa48e995fb1a3f0e354190c7a70c7ff, side a98db521e7003f018c0c5c114ce73dd43a08e157.

Implementation

• The upgraded 240-row pool had forced K=16 wgrad, increasing TMA/barrier work. Restore K=64 loads while skipping any complete 16-wide MMA atoms beyond the current expert's padded end. Keep true physical expert offsets; do not round group starts into neighboring experts.
• This is an opt-in MegaMoE wgrad specialization. Generic GEMM defaults are unchanged. No repacking buffer, new full-width scratch, forward-layout/ring change, or changed accumulation order.
• Isolate the count-based/no-SF wgrad compilation alignment from unrelated global grouped-GEMM alignment. Setting the latter to 224 reproduced a compile failure on the first candidate; the final implementation passes and leaves the process-global value unchanged.

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-optimization 118a7f9, optimized 48e4df7. Each path retains its corresponding integration math/ring policy.

Path	Pre-upgrade F+B	Regressed F+B	Optimized F+B	Time vs pre-upgrade
Ordinary BF16	66.76 ms	87.19 ms	65.48 ms	-1.92%
Side BF16	72.12 ms	71.62 ms	70.63 ms	-2.07%
Ordinary MXFP4	70.53 ms	88.06 ms	66.55 ms	-5.64%
Side MXFP4	75.00 ms	71.99 ms	71.52 ms	-4.63%

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

• Independently rebuilt ordinary and side wheels: 30/30 GPU tests each. Cover pool blocks 16/32/64/96/128/192/240, N=128/256, uneven expert sizes, empty first/middle/last experts, contrasting adjacent-expert data, K=16/32/64 byte equality, invalid K rejection, and global alignment=224 isolation.
• 34 CUDA variants compiled, including masked-tail wgrad with and without fused combine; the final alignment-only follow-up changes host specialization selection, not these kernel sources.
• Final side-stack head, all four production-shape paths on all four ranks: alternate K=16/64/16/64 backward on identical saved forward inputs. Complete output, grad-x, router, and base/adapter gradients are byte-identical: cosine=1, relative L2=0.
• Both ordinary paths additionally passed independent FP32 contractions for every expert at the native wgrad inputs; minimum cosine 0.9999998386. Ordinary wgrad output dtype remains BF16.
• Both side paths: four FP32 AdamW updates from the compared gradients produced byte-identical finite parameters, first moments, and second moments. This is a fixed-boundary single-layer optimizer check, not a full-model trajectory.
• Final compute-sanitizer focused tests: 3 passed, 0 errors. Initial sanitizer run hit allocator cuMemCreate permission 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.

Historical section: Prior alignment validation and historical measurements (superseded heads)

2026-09-20: upstream-upgrade alignment (Draft)

Current head: 118a7f9a379d56409e9b7608145d5b02729bb650, based on main 02452d4 (#15/#16), including the ordinary-training alignment in #17.

What changed

• Port the dedicated BF16/MXFP4 side-LoRA host launchers to DeepJIT, preserving the accepted synchronized backward tail and saved expert-input reuse.
• Use the upgraded native tile selection, UE8M0 helpers, workspace offsets, and shared-expert-compatible v3 buffer ABI.
• Keep side-LoRA's distinct shrink/expand phase scheduler explicit; expose an opt-in side_lora=True buffer API without adding the BF16 source plane to ordinary MXFP4 inference/training.
• Preserve configurable SM headroom on persistent side launchers.
• Include Align MegaMoE training with upgraded K-grouped TMA and layouts #17's 3D K-grouped output descriptor alignment for ordinary and side wgrad.
• No host-sync deletion, integrated prelude, or unsealed I-cache candidate.

Fresh validation

• Clean x86 extension/wheel build: passed.
• 32 CUDA variants compiled: ordinary/shared-expert BF16 and FP8/FP4/FP8 forwards, training/backward variants, side forwards/backwards, wgrad/combine and layout helpers; includes BLOCK_M=240.
• CPU/API/headroom: 18 passed, plus 11 headroom subtests.
• 11/11 compact GPU F+B edge cases passed on B300, including EP1/EP4 remote routes, masks, empty experts, SwiGLU/GeGLU, clamps, default/short scratch and zero-scale base preservation.
• EP4 small-shape MXFP4 repeated F+B: passed the existing gate; fixed saved-boundary gradients and repeated outputs are bitwise identical.
• BF16 and MXFP4 production-shape F+B accuracy oracles passed: EP4, 65,536 tokens/rank (262,144 total), H=4096, I=2048, E=256, top-k=6, adapter rank=128, activation clamp=10. Minimum native-boundary adapter cosine: BF16 shared factors 0.999998853, MXFP4 all six 0.999998942. These are kernel-level oracles, not a full-model run.

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 d6fbae4 versus new 118a7f9. 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.

Path	F old → new (ms)	B old → new (ms)	F+B old → new (ms)	F+B time change
Ordinary BF16	19.72 → 17.47	49.88 → 73.15	66.76 → 87.19	+30.60% slower
Side BF16	24.97 → 23.49	49.18 → 51.23	72.12 → 71.62	-0.70%, approximately unchanged
Ordinary MXFP4	10.08 → 8.37	60.91 → 80.67	70.53 → 88.06	+24.84% slower
Side MXFP4	17.99 → 17.16	58.02 → 56.11	75.00 → 71.99	-4.01%

Separate backward-component diagnostics localize the ordinary regression to W13/W2 wgrad+combine, not dgrad:

Ordinary path	dgrad old → new (ms)	W2 wgrad+combine old → new (ms)	W13 wgrad+combine old → new (ms)
BF16	36.01 → 32.77	4.46 → 13.15	8.71 → 26.63
MXFP4	47.95 → 41.31	4.10 → 12.98	7.99 → 26.30

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 d6fbae4 versus new 118a7f9: 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.

Path	Worst gradient cosine, old vs new	Worst gradient relative L2	Output / grad-x / router
Ordinary BF16	1.0	0	Bitwise identical
Side-LoRA BF16	1.0	0	Bitwise identical
Ordinary MXFP4	0.999999992784	1.2013293e-4 (0.01201%)	Bitwise identical
Side-LoRA MXFP4, including shared-factor EP reductions	0.999999960357	2.8157900e-4 (0.02816%)	Bitwise identical

• BF16's complete adapter/base weight gradients are also byte-identical across versions and repetitions.
• MXFP4 same-version worst gradient relative L2: ordinary 8.52255e-5 old / 9.23880e-5 new; side with EP reductions 2.43134e-4 old / 2.52399e-4 new. Cross-version differences remain small but are not claimed bitwise or strictly equal to the self-noise floor.
• Ordinary W13/W2 gradients additionally passed independent FP32 PyTorch contractions against every expert's native backward inputs, for both versions and both repeats. Native wgrad outputs are BF16, not a certification of the separate FP32-persistent-gradient integration path.
• This unified fixture replicates shared A1/A3/B2 across EP and reduces their gradients. Its side numbers are not the same measurement as the earlier EP-local-gradient fixture below.
• Preserve each integration's math settings and ring policy: ordinary BF16 fast-math off, side BF16 on; MXFP4 forward on / backward default. Old ordinary BF16 and side use the old minimum ring; old ordinary MXFP4 uses its two-wave ring. New paths use the upgraded integration's native policies. A supplemental old library-default-ring BF16 control also passed, but is not the production-ring comparison.
• 2026-09-21 update: the matched four-mode timing matrix and reverse-order confirmation completed. See the new performance section above; ordinary backward has a substantial regression, so performance acceptance is not passed.
• Scope remains native single-layer API/core calls. Weight transforms, routing-plan construction, allocation/lifetime integration, AC, optimizer, and whole-model E2E are not timed by this harness.
• Input manifests, per-rank comparison metrics, logs, scripts and SHA-256 digests are retained. The 195.20 GiB regenerable raw tensor outputs were pruned after comparisons to avoid filling the shared disk.

Direct old/new numerical comparison (2026-09-20)

Compared old d6fbae484c4cc7068aec00fe77ace9b1e273eb43 against new 118a7f9a379d56409e9b7608145d5b02729bb650, not two independent cosine-to-oracle numbers.

• 8 production-shape kernel F+B runs passed their numerical oracles: two independent runs per version and precision. Each run uses 4 B300 GPUs, EP4, 65,536 tokens/rank / 262,144 total, H=4096, I=2048, E=256, top-k=6, LoRA rank=128, balanced routes, SwiGLU clamp=10, side scale=0.25 and saved-X reuse.
• Integration math flags: BF16 forward/backward fast_math=True; MXFP4 forward True, backward False. Native tile selection is retained: old BLOCK_M=192, new BLOCK_M=240.
• SHA-256 checks verify identical inputs, raw and packed base weights/scales, side weights, routes, and upstream gradients on every corresponding rank. The new upstream FP4 fixture utility changed midpoint rounding; both versions therefore use the same new RNE fixture quantizer, without modifying either MegaMoE kernel. Comparing separately generated old/new packed weights was rejected.
• Compare complete source-ordered outputs, grad-x, router gradients and all six returned adapter gradients on all four ranks, without sampling. Metrics use CPU FP64 accumulation; relative L2 is ||new-old||₂ / ||old||₂. Table values are the worst rank/tensor across both comparisons, not rank averages.

Direct old → new	Minimum gradient cosine	Maximum gradient relative L2	Output / grad-x / router
BF16	1.0	0	Bitwise identical
MXFP4	0.999999980737	1.9628123e-4 (0.01963%)	Bitwise identical

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

• At 262K, the full-redispatch adapter-gradient repeatability gate fails: max relative L2 ~1.64e-4 vs the unchanged 1e-4 threshold. An exact old-head d6fbae4 control 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.
• The old bounded-SF binary-search policy assumes a monotonic tile heuristic, which the new upstream scheduler does not have. This port currently uses the conservative upstream SF sizing and a side-safe ring bound. The previous 5.606 GiB arena claim is not applicable to this head; memory/performance re-sealing remains open.
• No new full-model timing, ARM build/runtime, or arbitrary NCCL-overlap certification. Historical timings below are retained only as provenance.

Review order: ordinary alignment #17, then this side-kernel update; FireTitan integration is tracked in fw-ai/fireworks#44111.

Historical section: Historical implementation, measurements and validation

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 fused training kernels on SM100 with custom symmetric-buffer sizing, multi-rank BLOCK_M policy, and adapter-gradient numerics; incorrect buffer or tile contracts could corrupt distributed MoE backward.

Overview
Adds fused rank-128 side-LoRA to SM100 MegaMoE for BF16 and MXFP4 base / BF16 side, with dedicated forward and backward JIT paths that run shrink/expand around the expert GEMMs instead of the ordinary upstream scheduler.

Host / Python: New fp8_fp4_mega_moe_side_lora, bf16_mega_moe_side_lora, and matching backward entry points; get_symm_buffer_size_for_mega_moe_side_lora and extracted get_num_ring_tokens_for_mega_moe; SymmBuffer(..., side_lora=True) and exports (transform_side_lora_for_mega_moe, MegaMoESideLoraBackwardResult). Side backward wrappers MAX-reduce rank-local BLOCK_M before launch unless callers pass a proven uniform tile.

Kernels: New heuristics (num_experts_per_wave), SideLoraMegaMoEScheduler phases, CUDA params/layout helpers, and large BF16/MXFP4 forward plus unified backward (wave, grad-x publish, adapter wgrads, padding clear). sm100_bf16_mega_moe_wgrad_1sm gains optional row-strided inputs and a custom kernel name for side wgrads.

Tests: Compile matrix and edge-case GPU matrix for native F+B; buffer contract tests; NCCL tile-uniformity / CUDA-graph regression for side backward BLOCK_M.

^{Reviewed by Cursor Bugbot for commit b8e5476. 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]

Stale Bugbot comment from a previous run.

[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.

[mayinghan]

Code review (from-scratch, code-only)

Reviewed the full diff (this PR is the superset of merged #17; the #17-specific k_alignment_override / mask_grouped_k_tail / 3-D TMA changes were line-verified correct). Findings on this PR's own changes, most severe first:

1. [High] deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh:5664 — trace site 22 vs a 22-site host schema → OOB writes

The kernel grew to kTraceSiteCount == 23 (line 1844, with DG_STATIC_ASSERT at 1885 saying "update the host trace-site schema") and traces site 22 via full_grid_phase_barrier(22) at line 5664, but the host still pins 22 sites (csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp:952-960) and Python allocates (22, num_sms, 5) (deep_gemm/mega/backward.py:15, 585). A trace-enabled run writes kernel_trace + (22*kNumSMs + blockIdx.x)*5 — ~6 KB past the tensor end on every launch (silent corruption of the neighbouring allocation). Conversely a schema-correct 23-row tensor is rejected by the host assert. The base backward (sm100_fp8_fp4_mega_moe_backward.hpp:1594 vs its .cuh:1450) is consistent at 22; this file kept 22 while the kernel grew to 23.

2. [High] csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp:1159 — write_grad_x_pool=false is accepted but ignored → OOB writes

The host asserts a 0-row pool is legal for write_grad_x_pool == false + direct_remote_grad_x == true (asserts {0, hidden} at 683-687; requires only write_grad_x_pool || direct_remote_grad_x at 716), but the persistent wave is always launched with .write_grad_x_pool = true (1159; fp8 twin at 1964) while the kernel writes grad_x_pool_output under if constexpr (kWriteGradXPool). The Python wrappers additionally overwrite the caller's flag with True (deep_gemm/mega/backward.py:1088, 1208). A pybind/C++ caller following the asserted contract (0-row torch.empty((0, hidden)) pool) gets the whole wave writing into a zero-element tensor — out-of-bounds writes into whatever the caching allocator placed next. The flag looks honored but has no effect at either layer; please either honor it or remove it and tighten the host assert.

3. [Medium] deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_backward.cuh:2664 — epilogue call written against a foreign signature (compile landmine)

In the if constexpr (!kBF16Mode && !kGateUpPrepared) recompute branch, the epilogue::sm100_store_cd_swap_ab call passes 11 explicit template args to this tree's 12-parameter template (so GemmType::Normal lands on kNumOverlappedTmemCols and the type cd_dtype_t on bool kWithAccumulation), only 12 of 15 function arguments, and a PatternVisitor where an smem_t exposing cd_dtype is required. Any instantiation with kBF16Mode=false && kGateUpPrepared=false (the base-recompute config the reference sm100_fp8_fp4_mega_moe_backward.cuh still builds) is a hard compile error. Currently dead only because every generation site sets bf16_mode=true (hpp:1060) or gate_up_prepared=true (hpp:1971).

4. [Medium] deep_gemm/include/deep_gemm/scheduler/mega_moe_side_lora.cuh:92 — uninitialized valid_value (UB, latent)

get_num_tokens declares uint32_t valid_value; uninitialized, while upstream scheduler/mega_moe.cuh:255 initializes it to = 0 (upstream fixed this in 559d79f, merged 8d785f7 2026-09-12 — after this fork's copy was written). When no lane owns the requested expert index (e.g. kNumExpertsPerRank a multiple of 32, the usual EP shape), ptx::exchange shuffles lane 0's indeterminate register into current_num_tokens. Traced to have no live consumer today (the wave-end loop exits before it's read), so latent UB — but it should just be = 0 like upstream.

5. [Medium] deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe_side_lora_forward.cuh:469 — stable-compaction prefix pass guarded by global_warp_idx < kNumExperts (inherited, latent)

The parallel prefix pass is one-warp-per-expert (global_warp_idx < kNumExperts && lane_idx == 0), so when the global expert count exceeds kNumSMs * kNumDispatchWarps (~560-592) experts above that index never get an exclusive prefix — their per-warp slot ranges overlap → routes overwritten/lost in the expert pool, wrong preactivations and wrong adapter grads, no assert. Trigger is arithmetically satisfiable (1024 global experts at ~560 warps needs only ~2.5k tokens/rank), and the branch is live for production DSV4-Flash today. Not reachable in current wired configs (256 experts; qwen3_5 at 512 sits ~4 SMs of headroom from the boundary), and the block is a verbatim copy of the pre-existing upstream code (882a9a4) — but the copy in this PR should get the missing guard/static assert so the boundary can't be crossed silently.

6. [Low] csrc/apis/mega_moe.hpp:949 — bf16 side-LoRA buffer-size check subtracts the route plane

bf16_mega_moe_side_lora asserts expanded_num_required_bytes - num_max_tokens_per_rank*num_topk*sizeof(float) (copied from the legacy ABI in the non-side bf16 entry at 642-646), while the slicer from get_symm_buffer_size_for_mega_moe_side_lora asserts the full base_bytes + extra_bytes (~302) and the fp8 entry point uses the exact size. A legacy-sized buffer passes the relaxed assert at 952 and then trips the stricter assert inside slice() with a misleading message. Since the training backward does require sym_buffer.backward_grad_route, the relaxed bound is wrong for this path.

7. [Perf] deep_gemm/mega/backward.py:1105 (bf16) / :1252 (fp8) — double memset of the direct grad-x plane

Both wrappers _direct_grad_x_planes(sym_buffer).zero_() the full num_topk * num_max_tokens_per_rank * hidden bf16 region before launch, but sm100_bf16_mega_moe_side_lora_grad_x_impl already clears the identical extent in-kernel under if constexpr (kDirectRemoteGradX) (.cuh:1332-1348, same combine_buffer base). FireTitan always sets direct_remote_grad_x=True for EP>1 → ~3.76 GB duplicated DRAM writes per layer, ~230 GB per step across 61 layers, plus an extra elementwise kernel serialized on the main stream. On the local path (flag False) the planes are never read either, so the Python zero is unneeded in both configurations.

8. [Maintainability] ~96%-identical forks of the existing kernels

The new backward .cuh is 5782 lines with the prelude 851/856 lines byte-identical to sm100_fp8_fp4_mega_moe_backward.cuh:463 and the wave impl 95.9% identical (3919/4088) to sm100_fp8_fp4_mega_moe_backward_wave_impl; the two forward kernels fork their base kernels 54-70% identically. Additionally: 4 copies of the rank-uniform num_config_tokens MAX all-reduce (deep_gemm/mega/__init__.py:311, 406, 488, 597) with already-drifting guards (a missed copy means ranks disagree on BLOCK_M and deadlock at the persistent grid barrier), 6 copies of the mode-name helpers across the new hpp files, and _allocate_side_lora_backward_outputs's write_grad_x_pool argument is always True (its empty-tensor arm unreachable). Shared-mainloop fixes — including the K-grouped tail/alignment/3-D TMA fixes in this same PR — must now be applied twice, and the duplicated template names in two .cuh files are an ODR hazard if they ever land in one TU. Worth at least a follow-up to extract the shared prelude/route-reduce/wave core or parameterize the base kernels with a side-LoRA trait (the forwards already carry 30+ policy template params for exactly this).

Checked and correct (no action): Python↔C++ argument orders for all new entry points, the wgrad 3-D TMA store descriptor, the kMaskGroupedKTail predicate for every allowed pool_block_m×kBlockK pair, k_alignment_override = kBlockK (no global-state leak), the scale_grads selector chain, and the differing gate/up interleave between the bf16 and fp8 hosts.

🤖 Generated with Claude Code

[morgendave]

Addressed the correctness findings in 0cbfc76 (now pushed). Items 1–6: separate side23/ordinary22 trace schemas; reject unsupported omitted grad-x pools at both API layers; explicit prepared-gate/up specialization contract; initialized scheduler state; strided prefix over all experts; full side-buffer size validation. Item 7: removed only redundant Python clears, preserving producer completion and publisher clear/barrier synchronization. C++ build, six source regressions, BF16/MXFP4 EP2 numerical oracles, exact zero-scale base preservation, actual 23-site BF16 traces, AC/AdamW integration and the 1024-expert BF16 regression all passed. Full test scope and unchanged 262K/ARM/performance gates are in the new description section. Item 8 (shared-core extraction) remains a follow-up; not claimed resolved by these targeted fixes.

[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 003b9dd. Configure here.}