Align MegaMoE training with upgraded K-grouped TMA and layouts

[morgendave]

2026-09-21: wide-K wgrad recovery

Current head: 0be759300fd5d198ef211e452ccb857217c41a1a.

• Restore K=64 for MegaMoE wgrad, masking final 16-wide MMA atoms outside each expert's padded extent. Preserve physical offsets and accumulation order; no full-width repack scratch or ring/forward-layout changes.
• Generic GEMM defaults remain unchanged. Wgrad compilation is isolated from unrelated process-global grouped alignment; the 224 case reproduces a compile failure before isolation and passes afterward without mutating that setting.
• Independently built ordinary and side wheels: 30/30 GPU regression tests each, including all seven pool tiles, N=128/256, empty/uneven groups, adversarial adjacent data, K=16/32/64 byte equality, invalid K, and global=224.
• Side stack 48e4df7 includes this exact ordinary update. Its EP4/262K ordinary BF16 and MXFP4 complete fixed-input backward tensors are byte-identical across K=16/64/16/64 on every rank, cosine=1/L2=0. All-expert independent FP32 wgrad oracle minimum cosine: 0.9999998386.
• 34 CUDA variants compiled; final focused memcheck 3 passed / 0 errors. Only sanitizer runs disable expandable segments due to allocator permission incompatibility.

Same physical 4×B300, EP4, 65,536 tokens/rank, H4096/I2048/E256/top-k6; native-API timing at the companion side-stack head with side disabled:

Ordinary path	Pre-upgrade F+B	Upgraded K=16 F+B	Optimized F+B
BF16	66.76 ms	87.19 ms	65.48 ms
MXFP4	70.53 ms	88.06 ms	66.55 ms

Direct F+B medians over 3×10 max-rank CUDA-event samples after warmup. W13/W2 wgrad+combine recovered from 39.79 to 13.48 ms BF16 and 39.28 to 12.50 ms MXFP4. Ordinary wgrad output dtype remains BF16. No transforms, AC, optimizer timing, full-model E2E, or ARM certification. Companion ordinary integration #50925 pins this head; #9 includes it and #44111 pins that stack.

Prior alignment validation and historical measurements (superseded heads)

Summary

Align the existing MegaMoE training bridge with the upgraded DeepJIT / K-grouped GEMM implementation on 02452d4 (after #15/#16). This is an upstream-upgrade compatibility update, independent of the side-LoRA extension in #9.

• Build the wgrad output descriptor as [N, M, group], matching the new K-grouped 3D TMA store. The old 2D descriptor traps at SM90_TMA_STORE_3D::copy on the upgraded kernel.
• Query the native BLOCK_M selection in the training oracle instead of duplicating the pre-upgrade heuristic in Python; use the same pool layout for forward and backward.
• Add GPU coverage for seven supported pool block sizes (16/32/64/96/128/192/240), 128/256-column outputs, empty experts, and unequal group sizes.

Validation

• Standalone wgrad reproduces the illegal instruction on unmodified 02452d4; compute-sanitizer identifies the 3D TMA store.
• Updated standalone wgrad: 14/14 GPU tests passed on B300.
• EP4 BF16 exact F+B oracle: passed, all logged forward/dgrad/wgrad/source-combine stages bitwise match the native reference (256 tokens/rank, H=512, I=256, E=16, top-2, fast_math=false).
• EP4 MXFP4 active-saved-pool forward/backward oracle: passed, same shape. Legacy external comparison was unavailable in this image due to duplicate NCCL discovery; this is the self-contained route/backward oracle, not a legacy-baseline claim.

No scheduler, ring-sizing, model-lifetime caching, or side-LoRA behavior changes. No performance or ARM validation claim in this PR. The integration update must consume the native layout query as well.

---

[!NOTE]
Medium Risk
Changes CUDA wgrad correctness (3D TMA ABI and grouped-K tail masking) on the MegaMoE training backward path; mistakes could corrupt expert gradients but scope is limited to that kernel bridge.

Overview
MegaMoE BF16 wgrad is aligned with the upgraded K-grouped SM100 path: the output TMA descriptor is built as [N, M, group] via make_tma_3d_desc instead of a flattened 2D store that traps on the new epilogue.

Wgrad tiling defaults to K=64 (still overridable with DG_BF16_MEGA_MOE_WGRAD_BLOCK_K). When pool_block_m is not divisible by that tile (e.g. 96/240), the host sets mask_grouped_k_tail and k_alignment_override so the kernel skips 16-wide MMA atoms past each expert’s real row count—avoiding cross-expert contamination from wide TMA loads. The BF16 GEMM JIT gains the corresponding kMaskGroupedKTail template knob.

Training tests call get_block_m_for_mega_moe for pool_block_m (including 240) instead of duplicating the old Python heuristic; compile coverage adds a masked-tail wgrad variant; new GPU tests exercise all supported pool sizes, empty experts, and global MK alignment independence.

^{Reviewed by Cursor Bugbot for commit 0be7593. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Risk: medium. Left a non-blocking comment; this automation did not approve and did not assign reviewers. Human review is needed: MegaMoE wgrad TMA/layout changes are above the low-risk auto-approve threshold and need a correctness check of the 3D store descriptor and training oracle alignment.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}

[cursor]

Risk: medium. Left a non-blocking comment; this automation did not approve and did not assign reviewers. Human review is needed: MegaMoE wgrad TMA/layout isolation from global grouped alignment is above the low-risk auto-approve threshold and needs a correctness check of the store descriptors and training oracle alignment.

  

_{Sent by Cursor Approval Agent: Pull Request Approver}