[feat] Add SITU trtllmgen MOE

[jiahanc]

📌 Description

Add SiTU activation support for TRTLLM-Gen MoE for MXFP4 × MXFP8 and NVFP4 × NVFP4

🔍 Related Issues

🚀 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

Summary by CodeRabbit

• New Features
  • Added SiTU gated activation support for FP4 MoE in routed and non-routed execution, with configurable per-expert gemm1_alpha, gemm1_beta, and gemm1_clamp_limit.
  • Extended FP4 MoE trace templates with an activation_type input to select SiTU.
• Bug Fixes
  • Improved activation-type mapping/validation so SiTU uses the correct gated behavior.
  • Made MoE kernel/tile selection dtype-aware (activation + weight) and corrected per-token scaling eligibility during config generation.
• Tests
  • Expanded MoE coverage for SiTU and optional per-expert parameters, including additional routed-logit and autotune regression checks.
• Chores
  • Updated the remote BMM artifact reference and checksum.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[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

📝 Walkthrough

Walkthrough

FP4 block-scale MoE now supports SiTU activation with validated GEMM1 parameters, updated activation mappings, dtype-aware tile selection, trace propagation, artifact metadata updates, and regression tests.

Changes

FP4 SiTU MoE support

Layer / File(s)	Summary
Activation contracts and kernel mappings include/flashinfer/trtllm/..., flashinfer/tllm_enums.py, csrc/trtllm_fused_moe_runner.cu	Adds SiTU activation enums, serialization, gated-activation handling, and fused GEMM mapping.
FP4 activation parameters and tile dispatch flashinfer/fused_moe/core.py, csrc/trtllm_fused_moe_kernel_launcher.cu, csrc/trtllm_batched_gemm_runner.cu, flashinfer/artifacts.py	Validates and materializes SiTU parameters, updates dtype-dependent FP4 tiles and scaling filters, and changes BMM artifact metadata.
Trace schemas and reference execution flashinfer/trace/templates/moe.py	Propagates activation type and per-expert parameters through FP4 schemas, routing references, initialization, and SiTU computation.
Runtime references and regression tests tests/moe/*	Adds SiTU reference math, parameter wiring, routed and non-routed coverage, and NVFP4 tactic regression checks.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant FP4MoEEntryPoint
  participant TraceOrReference
  participant FP4Kernel
  participant Autotuner
  Caller->>FP4MoEEntryPoint: select activation_type=Situ
  FP4MoEEntryPoint->>FP4MoEEntryPoint: validate GEMM1 parameters
  FP4MoEEntryPoint->>TraceOrReference: provide activation parameters
  TraceOrReference->>FP4Kernel: compute or launch SiTU activation
  Autotuner->>FP4Kernel: query dtype-aware tile configurations

Loading

Possibly related PRs

• flashinfer-ai/flashinfer#4009: Adds related SiTU activation plumbing across fused MoE paths.
• flashinfer-ai/flashinfer#4168: Updates overlapping MoE valid-configuration and per-token scaling logic.
• flashinfer-ai/flashinfer#2014: Updates the same BMM artifact path and checksum constants.

Suggested labels: op: moe

Suggested reviewers: yzh119, nv-yunzheq, samuellees, iwakurarein, feih-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 56.25% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: adding SiTU support for TRTLLM-Gen MoE.
Description check	✅ Passed	The description matches the template and includes the required sections, checklist items, and test status.
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.

✨ Finishing Touches

🧪 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]

code review: looks good

[flashinfer-bot]

GitLab MR !1062 has been created, and the CI pipeline #59978566 is currently running. I'll report back once the pipeline job completes.

[jiahanc]

/bot run tests/moe

[flashinfer-bot]

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

[IwakuraRein]

/bot run tests/moe

[flashinfer-bot]

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

[jiahanc]

/bot run tests/moe

[flashinfer-bot]

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