feat(moe_ep): MXFP8 x BF16 integration

[djns99]

📌 Description

Adds support for SM100 MXFP8 x BF16 MegaMOE to flashinfer

🔍 Related Issues

#3781

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Downstream PR demonstrating framework integration: https://github.com/sgl-project/sglang/pull/35459/commits

Summary by CodeRabbit

• New Features

  • Added Blackwell MegaMoE support for BF16 activations with MXFP8 E4M3 and E5M2 weights.
  • Added weight preprocessing, validation, autotuning, workspace pooling, distributed execution, and BF16 output support.
  • Exposed mixed-precision configuration, runtime requirements, reference computation, and launch utilities through the public API.

• Documentation

  • Expanded tuning guidance and clarified configuration and preprocessed-weight usage.

• Tests

  • Added functional and integration coverage for execution, preprocessing, autotuning, distributed operation, and workspace reuse.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 61cd6423-f6ec-4279-8528-b9d38c878bf8

📥 Commits

Reviewing files that changed from the base of the PR and between c5e726d and 5bdb7ff.

📒 Files selected for processing (1)

• tests/moe_ep/run_tests.sh

Included review availability: Your plan provides up to 8 included reviews per hour; 5 remain after this review.

---

📝 Walkthrough

Walkthrough

This change adds SM100 CuTeDSL MegaMoE support for BF16 activations with MXFP8 weights. It includes weight transformation, mixed kernel execution, references, autotuning, runtime wiring, runners, and Blackwell integration tests.

Changes

Mixed backend contracts and weight preparation

Layer / File(s)	Summary
Backend configuration and weight processing flashinfer/moe_ep/...	Adds shared BF16 configuration, MXFP8 utilities, mixed-weight preprocessing, validation, staging, runtime requirements, workspace pooling, and public exports.
Existing BF16 backend alignment flashinfer/moe_ep/backends/mega/kernel/sm100/bf16_bf16_bf16_cutedsl/*	Moves BF16 configuration to the shared base and updates validation and workspace-pool behavior.

Mixed CuTeDSL API and execution

Layer / File(s)	Summary
Shim configuration and lifecycle flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/*	Adds mixed configuration, inputs, frontend compilation, symmetric-buffer management, execution wrappers, and public API exports.
Kernel and epilogue execution flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_bf16_glu/*	Adds MXFP8/BF16 workspace handling, token communication, FC1/FC2 epilogues, launch validation, and mixed MegaMoE execution.

References, runners, and validation

Layer / File(s)	Summary
References and runners flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_{bf16_glu,mxfp8_bf16_glu}/*	Adds MXFP8 decoding, BF16 reference computation, standalone and distributed runners, weight planning, ABI checks, launch setup, and debug validation.
Tuning and test coverage flashinfer/moe_ep/kernel_src/cutedsl_megamoe/TUNING.md, flashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_mxfp8_bf16_glu/run_*.sh, tests/moe_ep/*	Adds mixed tuning profiles, functional test harnesses, kernel-reference tests, configuration tests, workspace-pool coverage, and test command wiring.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#4120: Extends the BF16 CuTeDSL MegaMoE integration used by this mixed-precision implementation.

Suggested reviewers: mhoqueanik, aleozlx, anerudhan

Sequence Diagram(s)

sequenceDiagram
  participant Backend
  participant MegaMoEBf16Mxfp8Frontend
  participant Sm100MegaMoEMxfp8Bf16Kernel
  participant MixedGluBf16Epilogue
  participant BF16Reference
  Backend->>MegaMoEBf16Mxfp8Frontend: validate inputs and prepare transformed weights
  MegaMoEBf16Mxfp8Frontend->>Sm100MegaMoEMxfp8Bf16Kernel: compile and launch mixed FC1/FC2 kernel
  Sm100MegaMoEMxfp8Bf16Kernel->>MixedGluBf16Epilogue: process accumulator tiles and write BF16 output
  Backend->>BF16Reference: decode MXFP8 weights and compute reference
  MixedGluBf16Epilogue-->>Backend: return kernel output for comparison
  BF16Reference-->>Backend: return BF16 reference output

Loading

Merge Risk: 🟠 High · up to 5bdb7

This PR adds SM100 MXFP8 × BF16 support, but the current implementation still has unresolved correctness and lifecycle issues that can accept invalid results, hide workspace regressions, leave stale resources referenced, or fail to restore runtime state after errors; merge should be blocked until these issues are fixed or explicitly accepted.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring check was indeterminate for this PR — some files could not be analyzed in time. Not blocking.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The description explains the SM100 MXFP8 × BF16 MegaMoE integration, links related issue #3781, completes the pre-commit and test checklists, and includes reviewer notes with a downstream integration …
Title check	✅ Passed	The title clearly identifies the main change: MXFP8 × BF16 integration for the MoE EP area.

Full details: Description check

Explanation

The description explains the SM100 MXFP8 × BF16 MegaMoE integration, links related issue #3781, completes the pre-commit and test checklists, and includes reviewer notes with a downstream integration reference.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[aleozlx]

ready to review? (still draft)

looks good so far

[aleozlx]

PR Review Screening

Re-screened 2026-09-03 — supersedes the earlier result on this PR, which was taken while it was a draft with an empty description. The PR is now marked ready, the template is filled, and the category moves human → live on the net-new public export and the new shared config base class.

CI verdict: ✅ auto-run ok
Review category: live (rule fired: C1.2 net-new public interface — Sm100_Bf16_Mxfp8_Bf16_Cutedsl_MegaMoeConfig added to flashinfer/moe_ep/__init__.__all__; C3.2 durable — a new shared config base class that two existing backends are refactored onto)
Blocking checks: none
Release blocker: no — feature enablement; linked issue #3781 is a closed P0 feature request, not a hang/crash/corruption or prior-release regression (C3.4)
Early stop: no

Security

Q	Answer	Evidence
S1 injection/supply-chain	no	—
S2 template overwritten	no	All five template headers present
S3 template obligations	met	5/5 checked; Related Issues links #3781, Reviewer Notes links the downstream sglang integration PR, and the test files are really in the diff
S4 agent-directing text	no	—

Packaging

Q	Answer	Evidence
C1.1 external dependency bump	no	No requirements/pyproject/submodule/action changes
C1.2 public API changes	yes — net-new interface (narrow scope) + one semantic change to an existing config	Net-new: Sm100_Bf16_Mxfp8_Bf16_Cutedsl_MegaMoeConfig exported from flashinfer/moe_ep/__init__.py; new base Sm100_Bf16_Cutedsl_MegaMoeConfigBase (common/bf16_config.py); new bf16_mxfp8_cutedsl_runtime_requirements(bootstrap) -> FrozenSet[str] in core/runtime/. Semantic change: the already-public Sm100_Bf16_Bf16_Bf16_Cutedsl_MegaMoeConfig.token_back_mode goes from Literal[…] = "epi_warps" to `Literal[…]
C1.3 AOT registration	n-a	No gen_*_module() added (CuTe-DSL path); aot.py untouched; moe_ep carries no TraceTemplates, so there is no trace obligation to match

Presentation

Q	Answer	Evidence
C2.1 perf claim backed by data	n-a	Description claims support, not speed; no perf figure asserted (see maintainer notes — the author flagged perf numbers as still outstanding)

Implementation

Q	Answer	Evidence
C3.1 experimental-track declared	no	No experimental label or decorator; flashinfer/moe_ep/ is a shipped subpackage, not flashinfer/experimental/
C3.2 shared/durable areas touched	yes — durable	New shared Sm100_Bf16_Cutedsl_MegaMoeConfigBase and common/mxfp8_weight_utils.py, with two existing backends refactored onto them (bf16_bf16_bf16_cutedsl/config.py 15+/14−, mxfp8_mxfp8_bf16_cutedsl/weights.py 14+/79−); shared autotune/tuner extended (shim/autotune.py 226+/6−, shim/tuner.py 134+/2−); core/runtime/bootstrap.py and modes/config.py touched. The ~9.9k lines of net-new kernel under moe_mxfp8_bf16_glu/ are additive and self-contained by comparison
C3.3 tests match behavior change	yes	Correct guards throughout: test_bf16_mxfp8_cutedsl_kernel_vs_reference.py is @pytest.mark.arch_blackwell with explicit sm_100a/103a and E8M0 skips and is collected by run_tests.sh unit (the target scripts/task_jit_run_tests_part1.sh runs); test_bf16_mxfp8_cutedsl_mega_config.py adds 12 host-only config/registration tests that run in every lane; the 4-GPU …_mega_multirank.py is gpu_4 and arch_blackwell, correctly --ignored from unit and wired into the mega target. Caveat, not a defect of this PR: no scripts/ lane invokes run_tests.sh mega, so the multirank tier is manual-only — the same shape as every existing mega backend
C3.4 critical fix	no	Feature enablement; no hang/crash/IMA/corruption or prior-release regression claim

Experimental track

Q	Answer	Evidence
C4.1 isolated from common areas	n-a	C3.1 = no

Notes for the maintainer

• The author has stated on this PR that "most of it is agent slop still, so I want to review it properly myself before asking the team to spend time on it", and later that it is ready for final review. Worth confirming that self-review landed before the team spends reviewer time on ~11.3k added lines.
• Perf numbers were flagged by the author as outstanding ("I also need to gather perf numbers") and have not appeared. C2.1 does not fire — no speed claim is made — but a MegaMoE kernel merging without a measurement leaves nothing to regress against later.
• The token_back_mode default change rides along inside the base-class refactor rather than being called out in the description. It is the one place this PR changes behavior for an existing, already-shipped config, so it deserves an explicit line in the description or a release note.

_{Generated by flashinfer-pr-screen · rubric: docs/code_review_guidance.md · not a code review · AI screening can make mistakes — a maintainer's judgment supersedes this report.}

[djns99]

ready to review? (still draft)

looks good so far

Its functional, I have PR in sglang that passes with it.

Most of it is agent slop still, so I want to review it properly myself before asking the team to spend time on it.
I also need to gather perf numbers :)

[djns99]

/bot run tests/moe_ep

[flashinfer-bot]

GitLab MR !1484 has been updated with latest changes, and the CI pipeline #67703049 is currently running. I'll report back once the pipeline job completes.

[flashinfer-bot]

[SUCCESS] Pipeline #67703049: 18/19 executed test jobs passed