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

[morgendave]

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.

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 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 distributed MoE training/backward surface on SM100 with tight buffer aliasing, padding, and quantization paths; regressions would affect gradient correctness and performance at scale.

Overview
Adds fused MegaMoE side-LoRA on SM100 for BF16 and MXFP4 base + BF16 side adapters (rank 128), with dedicated forward/backward JIT paths and Python exports (bf16_mega_moe_side_lora, fp8_fp4_mega_moe_side_lora, matching backward entry points, transform_side_lora_for_mega_moe).

Side-LoRA gets its own phase scheduler and ring/wave sizing (get_num_ring_tokens_for_mega_moe, get_symm_buffer_size_for_mega_moe_side_lora, SideLoraMegaMoEConfig) instead of reusing the ordinary upstream scheduler. Backward fuses the base MoE wave with rank-width adapter math (shared/dense GEMMs, grad-x publish, padding clear, six adapter wgrads) via MegaMoESideLoraBackwardParams.

MegaMoE wgrad gains optional mask_grouped_k_tail and k_alignment_override so K=64 tiles stay safe on padded expert pools (e.g. BLOCK_M=240), plus a corrected 3D grouped output TMA store and strided-input support for side wgrad launches.

Tests add side-kernel compile matrix, SymmBuffer side_lora=True contracts, and a native F+B edge/boundary runner.

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

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