Skip to content

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

Draft
morgendave wants to merge 20 commits into
mainfrom
codex/megamoe-native-side-lora
Draft

morgendave wants to merge 20 commits into
mainfrom
codex/megamoe-native-side-lora

Conversation

@morgendave

@morgendave morgendave commented Aug 13, 2026

Copy link
Copy Markdown

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

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 changes to distributed MoE training kernels, backward numerics, and shared wgrad GEMM behavior; the PR notes a substantial ordinary backward slowdown tied to tile/K selection.

Overview
Adds fused rank-128 side-LoRA paths on top of MegaMoE for BF16 and MXFP4 (BF16 side adapters), exposed through new C++/pybind entry points (*_side_lora, *_side_lora_backward, side symmetric-buffer sizing) and Python exports.

Forward work includes dedicated JIT launchers, a separate wave scheduler (experts_per_wave, ring capacity), shrink/expand phases around expert GEMMs, and buffer layout helpers (refactored get_num_ring_tokens_for_mega_moe, get_symm_buffer_size_for_mega_moe_side_lora).

Backward fuses the base MegaMoE dgrad wave with six adapter gradients via shared MegaMoESideLoraBackwardParams, rank-128 grouped/dense GEMMs, grad-x publishing, padding clears, and optional kernel trace—without full-width side dgrad scratch.

BF16 grouped wgrad gains mask_grouped_k_tail and k_alignment_override in sm100_bf16_gemm (plus 3D output TMA in wgrad) so padded pool tiles (e.g. BLOCK_M 240) do not bleed across expert boundaries; side wgrad can use row-strided inputs and custom kernel names.

Reviewed by Cursor Bugbot for commit 48e4df7. Bugbot is set up for automated code reviews on this repo. Configure here.

Comment thread deep_gemm/mega/__init__.py
Comment thread csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

@cursor
cursor Bot requested review from ipiszy and vipulSharma18 August 13, 2026 17:54

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

@morgendave
morgendave force-pushed the codex/megamoe-native-side-lora branch from f323595 to 3e45539 Compare August 17, 2026 23:34
Comment thread csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp
Comment thread csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

Comment thread deep_gemm/mega/backward.py Outdated

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

@morgendave
morgendave force-pushed the codex/megamoe-native-side-lora branch from 1e1b4f7 to ccaa442 Compare August 24, 2026 19:33
Comment thread csrc/jit_kernels/impls/sm100_bf16_mega_moe_side_lora_backward.hpp

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Open in Web View Automation 

Sent by Cursor Approval Agent: Pull Request Approver

@morgendave

Copy link
Copy Markdown
Author

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

Copy link
Copy Markdown
Author

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 Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Fix All in Cursor

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

Comment thread tests/test_mega_moe_native_side_lora.py Outdated
@morgendave

Copy link
Copy Markdown
Author

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.

@morgendave
morgendave marked this pull request as draft September 20, 2026 20:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants