Fix MXFP8 fc1 weight shape check for non-gated activations by he-yufeng · Pull Request #2753 · flashinfer-ai/flashinfer

he-yufeng · 2026-03-11T04:25:14Z

The MXFP8 quantization path in getQuantParams() hardcodes inter_size * 2 for the fc1 weight block N-dimension check, assuming gated activations (where gate + up weights are concatenated). This causes the check to fail for non-gated activations where the multiplier should be 1.

Replace the hardcoded * 2 with * (isGatedActivation(base_activation_type) ? 2 : 1) in all three MXFP8 getQuantParams overloads. Also updated the error messages to show the actual expected multiplier.

This matches the existing pattern in the NVFP4 path (line ~1131) which already correctly guards the * 2 with an isGatedActivation check.

Summary by CodeRabbit

Bug Fixes
- Corrected internal size validation to account for gated activations by using the proper multiplier when certain quantization formats are used.
- Updated validation error messages to report the adjusted expected size when gated activations affect parameter layout.

The fc1_weight_block shape check hardcodes a factor of 2 for the N dimension (inter_size * 2), assuming gated activation where gate and up weights are concatenated. For non-gated activations the factor should be 1. Use isGatedActivation(base_activation_type) to select the correct multiplier. This matches the existing pattern at line 1131 (NVFP4 path) which already guards the * 2 with an isGatedActivation check. Fixes flashinfer-ai#2731

gemini-code-assist · 2026-03-11T04:25:28Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves an issue in the MXFP8 quantization path where the fc1 weight shape validation incorrectly assumed gated activations, causing failures for non-gated configurations. The changes introduce conditional logic to the weight dimension check, ensuring it correctly applies the appropriate multiplier based on the activation type, thereby improving the robustness and correctness of the quantization process for various activation functions.

Highlights

MXFP8 Quantization Fix: The fc1 weight block N-dimension check in getQuantParams() for MXFP8 quantization was updated to correctly handle both gated and non-gated activations.
Conditional Multiplier: The hardcoded * 2 multiplier in the weight shape check was replaced with a conditional expression * (isGatedActivation(base_activation_type) ? 2 : 1) to adapt to the activation type.
Improved Error Messages: Error messages for fc1 weight block size mismatches were enhanced to dynamically display the actual expected multiplier (1 or 2) based on whether the activation is gated.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu
- Modified the fc1 weight block size validation logic to correctly account for gated vs. non-gated activations in MXFP8 quantization.
- Enhanced the error message for fc1 weight block size mismatches to reflect the conditional multiplier.

Activity

No human activity has been recorded on this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

coderabbitai · 2026-03-11T04:25:31Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 2a5ce80c-d7e3-43c5-8d64-a52ebcb3d3dd

📥 Commits

Reviewing files that changed from the base of the PR and between a2d226a and da5d5e8.

📒 Files selected for processing (1)

csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu

🚧 Files skipped from review as they are similar to previous changes (1)

csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu

📝 Walkthrough

Walkthrough

Replaced hard-coded 2× multiplier with a conditional fc1_n_mult (2 when activation is gated, else 1) in three fc1 weight-block size validation branches and updated the corresponding error messages in the CUTLASS fused MOE binding.

Changes

Cohort / File(s)	Summary
fc1 Weight Block Size Validation `csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu`	Replaced fixed `* 2` with `fc1_n_mult` (2 if `isGatedActivation(base_activation_type)`, otherwise 1) in three quantization branches' `TVM_FFI_ICHECK` shape validations; updated error message text to reference `inter_size * fc1_n_mult`.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Suggested labels

v0.6.3

Suggested reviewers

aleozlx
yzh119

Poem

🐰 I nibble code with careful cheer,
Two when gated, one when clear,
Checks now match the activation's call,
Align the sizes—no more stall! 🥕

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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 accurately summarizes the main fix: addressing the MXFP8 fc1 weight shape check to handle non-gated activations correctly.
Description check	✅ Passed	The PR description covers what was changed, why it was needed, and which issue it fixes. However, it does not follow the provided template structure with the required sections.
Linked Issues check	✅ Passed	The PR successfully addresses issue `#2731` by replacing the hardcoded multiplier 2 with a conditional check that uses 2 for gated activations and 1 for non-gated activations.
Out of Scope Changes check	✅ Passed	All changes are focused on fixing the quantization shape-validation logic in getQuantParams for MXFP8 paths, directly addressing the issue requirements with no extraneous modifications.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

📝 Coding Plan

Generate coding plan for human review comments

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

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

gemini-code-assist

Code Review

This pull request correctly fixes a bug in the MXFP8 quantization path where the fc1 weight shape check failed for non-gated activations due to a hardcoded multiplier. The change introduces a dynamic multiplier based on the activation type, which is a clear improvement. My review includes suggestions to refactor the duplicated logic for calculating this multiplier across three different locations. This will enhance code readability and maintainability.

gemini-code-assist · 2026-03-11T04:27:03Z

+                  (isGatedActivation(base_activation_type) ? 2 : 1) &&
          fc1_weight_block.size(2) * FP8_PER_INT32 *
                  TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize ==
              TmaWarpSpecializedGroupedGemmInput::alignToSfDim(
                  hidden_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX))
-          << "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 "
-             "// block_scale_vector_size)";
+          << "fc1 weight block size must be (num_experts_on_rank, inter_size * "
+          << (isGatedActivation(base_activation_type) ? 2 : 1) << ", hidden_size // 4"
+             " // block_scale_vector_size)";


The expression (isGatedActivation(base_activation_type) ? 2 : 1) is repeated within this check. To improve readability and avoid code duplication, consider calculating this multiplier once and storing it in a local variable before the TVM_FFI_ICHECK. This would also align with the pattern used for isNvfp4Quant (around line 1134), which uses an if/else block to handle this logic more cleanly.

gemini-code-assist · 2026-03-11T04:27:03Z

+                  (isGatedActivation(base_activation_type) ? 2 : 1) &&
          fc1_weight_block.size(2) * FP8_PER_INT32 *
                  TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize ==
              TmaWarpSpecializedGroupedGemmInput::alignToSfDim(
                  hidden_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX))
-          << "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 "
-             "// block_scale_vector_size)";
+          << "fc1 weight block size must be (num_experts_on_rank, inter_size * "
+          << (isGatedActivation(base_activation_type) ? 2 : 1) << ", hidden_size // 4"
+             " // block_scale_vector_size)";


The expression (isGatedActivation(base_activation_type) ? 2 : 1) is repeated within this check. To improve readability and avoid code duplication, consider calculating this multiplier once and storing it in a local variable before the TVM_FFI_ICHECK. This would also align with the pattern used for isNvfp4Quant (around line 1134), which uses an if/else block to handle this logic more cleanly.

gemini-code-assist · 2026-03-11T04:27:03Z

+                  (isGatedActivation(base_activation_type) ? 2 : 1) &&
          fc1_weight_block.size(2) * FP8_PER_INT32 *
                  TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize ==
              TmaWarpSpecializedGroupedGemmInput::alignToSfDim(
                  hidden_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX))
-          << "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 "
-             "// block_scale_vector_size)";
+          << "fc1 weight block size must be (num_experts_on_rank, inter_size * "
+          << (isGatedActivation(base_activation_type) ? 2 : 1) << ", hidden_size // 4"
+             " // block_scale_vector_size)";


The expression (isGatedActivation(base_activation_type) ? 2 : 1) is repeated within this check. To improve readability and avoid code duplication, consider calculating this multiplier once and storing it in a local variable before the TVM_FFI_ICHECK. This would also align with the pattern used for isNvfp4Quant (around line 1134), which uses an if/else block to handle this logic more cleanly.

aleozlx · 2026-03-19T21:10:53Z

/bot run

flashinfer-bot · 2026-03-19T21:11:14Z

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

flashinfer-bot · 2026-03-20T02:11:46Z

[FAILED] Pipeline #46554074: 14/20 passed

aleozlx · 2026-04-24T02:01:17Z

looks like this PR is doing the same as this one
#3082

aleozlx · 2026-04-24T02:02:28Z

ok to close as dup tho? thanks for the contrib

@aleozlx

) Fixes #2731. ## What's broken? When using the CUTLASS fused MoE backend with **non-gated activations** (e.g., Relu2, Gelu, Silu) and MXFP8 quantization, the fc1 weight shape validation unconditionally rejects the input — even when the shape is correct. ## Who is affected? Anyone using the **CUTLASS fused MoE** path with: - **Quantization**: `WMxfp8AMxfp8`, `WMxfp4AFp8`, or `WMxfp4AMxfp8` - **Activation**: any non-gated type (Relu2, Gelu, Silu, etc.) Not affected: gated activations (Swiglu, Geglu, SwigluBias), or other quant modes (NVFP4 already handles this correctly). ## Where is the bug? `csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu`, inside `getQuantParams()` — the fc1 weight block N-dimension check hardcodes `* 2` at three MXFP8 branches (~L898, ~L1004, ~L1063). ## Why does it happen? PR #2581 introduced MXFP8 support when only gated activations (Swiglu) existed, so `inter_size * 2` was correct. Later, non-gated activation support was added to the trtllm-gen backend (PR #2707), but the CUTLASS backend's validation was never updated. The NVFP4 path in the same file (line ~1131) already handles this correctly with an `if (isGatedActivation(...))` guard. ## How did we fix it? For each of the 3 MXFP8 quant branches: 1. Extract `int const fc1_n_mult = isGatedActivation(base_activation_type) ? 2 : 1;` 2. Replace the hardcoded `* 2` with `* fc1_n_mult` 3. Update error messages: gated shows `"inter_size * 2"`, non-gated shows `"inter_size"` **Before:** ```cpp fc1_weight_block.size(1) == alignToSfDim(inter_size, ...) * 2 ``` **After:** ```cpp int const fc1_n_mult = isGatedActivation(base_activation_type) ? 2 : 1; fc1_weight_block.size(1) == alignToSfDim(inter_size, ...) * fc1_n_mult ``` ## How do we know it works? - `pre-commit run` passes (clang-format, lint, etc.) - Gated activations (default Swiglu): `fc1_n_mult = 2` — identical to old behavior, no regression - Non-gated activations: `fc1_n_mult = 1` — shape check now accepts correct `inter_size` dimension - Full GPU test suite requires CI (`@flashinfer-bot run`) ## Related - Builds on the approach identified in #2753 (stale ~27 days, CI unresolved). - Addresses the Gemini review feedback from #2753 by extracting the multiplier to a local variable before the validation checks. cc @aleozlx @nv-yunzheq  ## Summary by CodeRabbit * **Bug Fixes** * Fixed weight block size validation for Mixture of Experts (MOE) to correctly handle both gated and non-gated activation types, ensuring proper support across different activation configurations.  Signed-off-by: Yiyang Liu <37043548+ianliuy@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

he-yufeng requested review from jimmyzho, nv-yunzheq and yzh119 as code owners March 11, 2026 04:25

gemini-code-assist Bot reviewed Mar 11, 2026

View reviewed changes

Extract fc1_n_mult local variable for readability

a2d226a

aleozlx added the op: moe label Mar 19, 2026

aleozlx self-assigned this Mar 19, 2026

style: fix clang-format for error message lines

da5d5e8

he-yufeng requested review from IwakuraRein, aleozlx and jiahanc as code owners March 20, 2026 02:47

This was referenced Apr 15, 2026

Potentially superfluous check that disables non gated activations in the cutlass fused moe API #2731

Closed

fix: guard MXFP8 fc1 weight shape check for non-gated activations #3082

Merged

aleozlx approved these changes Apr 24, 2026

View reviewed changes

aleozlx added the duplicate This issue or pull request already exists label Apr 24, 2026

aleozlx closed this Apr 24, 2026

Conversation

he-yufeng commented Mar 11, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

gemini-code-assist Bot commented Mar 11, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

coderabbitai Bot commented Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

Poem

❌ Failed checks (1 warning)

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

aleozlx commented Mar 19, 2026

Uh oh!

flashinfer-bot commented Mar 19, 2026

Uh oh!

flashinfer-bot commented Mar 20, 2026

Uh oh!

aleozlx commented Apr 24, 2026

Uh oh!

aleozlx commented Apr 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

he-yufeng commented Mar 11, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Mar 11, 2026 •

edited

Loading