Add sgl-project's NVFP4 MegaMoE as an additive deep_gemm.mega_nvfp4 module

[vipulSharma18]

What

This adds sgl-project/DeepGEMM's SM100 NVFP4 × NVFP4 MegaMoE decode kernel and its fused mega_moe_pre_dispatch kernel, taken from ec226bc (sgl deepseek-ai#27, 593e376, deepseek-ai#80, deepseek-ai#81). They are exposed as a new module, deep_gemm.mega_nvfp4, which provides the subset of the sgl_deep_gemm API that Fireworks uses:

• transform_ue4m3_sf_into_required_layout
• transform_weights_for_mega_moe
• get_symm_buffer_for_mega_moe
• mega_moe_pre_dispatch
• nvfp4_mega_moe
• get_block_m_for_mega_moe
• set_num_sms

Why

MiniMax-M3.1 decode (FIREWORKS_DPEP_MEGAMOE=1) currently ships a second DeepGEMM wheel, sgl_deep_gemm, next to ours. With this module in our fork, Fireworks can drop that dependency and its separate build. That also unblocks landing the M3.1 decode stack (fireworks#55221) on main.

Design: purely additive

• One existing file changes: csrc/python_api.cpp gains an include and a register_apis call (+2 lines). Every other file is new: 22 files, +7128/−0.
• No existing header changes. DeepJIT's cache key hashes the kernel source and, recursively, every included header. Editing a shared header would therefore change every existing kernel's key and invalidate precompiled caches.
• Device code is a copy. sgl's headers are copied unchanged into deep_gemm/include/deep_gemm/mega_nvfp4/ (kernel, pre-dispatch, layout, scheduler, barrier, math/utils/tma_copy, the ptx helpers, sym_buffer, and the MmaKind block of types.cuh). The only edits are the namespace deep_gemm::mega_nvfp4 and the include paths. Line numbers are kept, so even the assert-path SASS matches sgl.
• Host code and _C entry points have separate, nvfp4-marked names in a new csrc/apis/mega_nvfp4.hpp, heuristics/nvfp4_mega_moe.hpp and impls/sm100_nvfp4_mega_moe{,_pre_dispatch}.hpp.
• These kernels stay independent of the rest of DeepGEMM, as the separate sgl library was. They keep their own SM count, compile for sgl's sm_100f target and launch without PDL. So deep_gemm.set_num_sms and deep_gemm.set_pdl never affect them, and they never affect the existing kernels.
• Not ported (host side): TRT-LLM weight layout, FP8 combine, and the BF16 and SM90 paths.
• No conflict with Fireworks' build patches. None of the files touched by deepgemm-situ, -mhc-graph-replay, -sigmoid-bf16 or -mhc-capability change. All four apply unchanged on top of this branch.

Base branch

This PR is stacked on #19, which folds Fireworks' four build patches into main, and is rebased onto it (ef210f2). Merge #19 first; this PR then retargets to main. After both merge, Fireworks builds main head with no patches, so any merge style is fine.

Verification (B300)

The first round compared c2e49f5 + 4 patches against this PR + 4 patches. After the rebase onto #19, the agent re-ran three checks: zero new JIT cache keys for existing kernels vs #19, bitwise equality with sgl_deep_gemm at EP2 (92 tensors) and EP8 (356), and tests/test_mega_moe_nvfp4.py passing.

Check	Result
JIT cache-key invariance: baseline with a fresh cache, then branch on the same cache	0 new entries (1240 → 1240). Covers fp8×fp4 MegaMoE (DSv4 SwiGLU and Kimi SiTU), BF16 MegaMoE at EP2/EP8, mega gate/mhc, FP8 and BF16 GEMM. The 40 existing-kernel outputs are bitwise equal, and the installed existing headers and .py files are byte-identical.
API surface (inspect.signature across deep_gemm, deep_gemm.mega, _C)	Only additions: 7 _C.*nvfp4* names. Nothing removed or changed.
Existing tests	Same as baseline: test_mega_moe.py fp8×fp4 and BF16 at 8 ranks, test_mega_moe_compat.py, test_mega_gate.py, test_mega_mhc.py, test_bf16.py, and the dense and m-grouped-contiguous tests from test_fp8_fp4.py.
Equivalence vs sgl_deep_gemm@ec226bc	Bitwise equal pre-dispatch slots and outputs at EP2 and EP8, with hidden 6144, inter 3072, E=256, top-k 8 and 1/16/64/128 tokens. Also equal at E=128, top-k 4 up to 512 tokens, and with set_num_sms(120). That is 860 tensors across 6 runs. The SASS is identical.
New tests	tests/test_mega_moe_nvfp4.py passes. It ports sgl's tests: pre-dispatch bytewise for fp8, mxf4 and nvfp4, and the 2-rank alphas test.
Latency (EP8, best of 3)	This PR 189.5 / 198.2 µs vs sgl 189.4 / 198.8 µs at 16 / 128 tokens.

Not run:

• The full test_fp8_fp4.py on both sides. An outside process killed the baseline run after 33 minutes, so only its dense and m-grouped-contiguous tests ran.
• The mhc reference comparison. Its reference module is unavailable, so it was skipped on both sides.
• test_mega_moe_training.py.
• sgl's SM90 path.

Notes for reviewers

• Maintenance: about 5k lines of sgl device code are now duplicated, so future sgl fixes need re-porting. That is the price of not touching any shared header.
• Licence: the ported tests come from Apache-2.0 files in sgl's tree, as their header notes. The device and host code is MIT, like upstream DeepGEMM.
• First use: like sgl_deep_gemm, the kernels compile on first use, so callers need an eager forward before CUDA-graph capture. Fireworks already does this.

Fireworks follow-up (separate PR)

• In megamoe_nvfp4.py, import deep_gemm.mega_nvfp4 as deep_gemm.
• Pin fw-ai/DeepGEMM main head (Fireworks PR prepared as a draft) with no build patches.
• Drop the sgl_deep_gemm build.
• The existing DeepGEMM precompiled cache stays valid, as the cache-key check above shows.

🤖 Generated with Claude Code

---

Note

Medium Risk
New distributed, symmetric-memory MoE inference path with JIT SM100 kernels and strict buffer layout contracts; isolated from existing APIs but operationally critical for decode once adopted.

Overview
Introduces deep_gemm.mega_nvfp4, a Fireworks port of sgl-project DeepGEMM’s SM100 NVFP4×NVFP4 MegaMoE decode stack so MiniMax-style decode can drop the separate sgl_deep_gemm wheel. The change is additive: only csrc/python_api.cpp is modified (register mega_nvfp4 APIs); everything else is new host glue, JIT launchers, copied device headers under deep_gemm/include/deep_gemm/mega_nvfp4/, and a Python package (SymmBuffer, get_symm_buffer_for_mega_moe, mega_moe_pre_dispatch, nvfp4_mega_moe, weight/SF transforms).

Host entry points live in csrc/apis/mega_nvfp4.hpp with nvfp4-prefixed _C names and namespace deep_gemm::mega_nvfp4, avoiding shadowing existing MegaMoE APIs. Symmetric-buffer sizing slices views for quantized activations, scale factors, top-k metadata, and ring buffers; nvfp4_mega_moe validates weights/layout and dispatches sm100_nvfp4_mega_moe on arch 10 only. Heuristics (nvfp4_mega_moe.hpp) keep a separate SM count (set_num_sms_for_nvfp4_mega_moe), compile these kernels with sm_100f for sgl parity, and launch with PDL disabled so global set_num_sms / set_pdl do not affect them.

Not ported on the host side: TRT-LLM weight layout, FP8 combine (use_fp8_combine is always false), BF16 MegaMoE, and SM90 kernels. Python rejects legacy DG_USE_FP4_ACTS / DG_USE_MXF4_KIND in favor of explicit mma_type.

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

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

[cursor]

Scheduler slot reuse missing fence

High Severity

MegaMoEScheduler::release_task_info arrives the empty barrier without a preceding ptx::fence_acq_rel_cta. In-flight task_infos metadata loads can still be pending when the scheduler reuses that pipeline slot, which races and can corrupt the next task. A later grid_sync or nvlink_barrier does not order these slot reads against the arrive.

 

^{Triggered by learned rule: MegaMoEScheduler must fence before slot reuse}

^{Reviewed by Cursor Bugbot for commit ef210f2. Configure here.}

[cursor]

Risk: high. Not approved: this adds a large new SM100 NVFP4 MegaMoE kernel and public deep_gemm.mega_nvfp4 API (thousands of lines of device/host code), which is above the low-risk auto-approve threshold. Human review is needed; no reviewers were assigned.

  

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