Feature: Support Relu2 in FusedMoE fp8 cutlass path #27261

amirkl94 · 2025-10-21T11:02:40Z

Purpose

This PR enables FusedMoE FP8 cutlass path for models using non-gated relu2 models. This new path gives a performance gain of around 20% on output token throughput over the triton path.
This PR requires flashinfer 0.5.0 .

Tests

New parameterization in test_flashinfer.py::test_flashinfer_cutlass_moe_fp8_no_graph to verify the new non-gated activation.

Performance tests:
Ran on a single H100. Started the server (once with VLLM_USE_FLASHINFER_MOE_FP8=1 and once with VLLM_USE_FLASHINFER_MOE_FP8=0):

python3 -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --port 8080 --model <NemotronNanov3-modelopt-fp8-ckpt> --tensor-parallel-size 1 --pipeline-parallel-size 1 --swap-space 16 --max-num-seqs 512 --trust-remote-code --max-model-len 3800 --gpu-memory-utilization 0.9 --max-num-batched-tokens 8192 --enable-chunked-prefill --disable-log-requests

Benchmarked using:

vllm bench serve --backend vllm --host 0.0.0.0 --port $port --model $model_path --num-prompts 1280--trust-remote-code --ignore-eos --max-concurrency 256 --random-input-len 1024 --random-output-len 1024 --no-stream --dataset-name random --num-warmups 20

triton path yielded:

Request throughput (req/s):              3.47
Output token throughput (tok/s):         3551.76
Peak output token throughput (tok/s):    6116.00

cutlass path yielded:

Request throughput (req/s):              5.15
Output token throughput (tok/s):         5270.17
Peak output token throughput (tok/s):    7168.00

~17% perf gain for peak (decode), ~40% perf gain for average token throughput (includes prefill).

vllm/model_executor/layers/quantization/modelopt.py

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

mergify · 2025-11-05T11:21:56Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @amirkl94.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Amir Klein <[email protected]>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2025-11-05T12:01:51Z

vllm/model_executor/layers/quantization/modelopt.py

        self.cutlass_fp8_supported = cutlass_fp8_supported()
        self.flashinfer_moe_backend: FlashinferMoeBackend | None = None
-        if (
-            envs.VLLM_USE_FLASHINFER_MOE_FP8
-            and has_flashinfer_moe()
-            and self.moe.is_act_and_mul
-        ):
+        if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
            self.flashinfer_moe_backend = get_flashinfer_moe_backend()
            logger.info_once(


Avoid enabling TensorRT flashinfer for relu2 activations

Removing the self.moe.is_act_and_mul guard means a flashinfer backend is now enabled whenever VLLM_USE_FLASHINFER_MOE_FP8 is set, regardless of the activation. If the user selects the latency backend (TensorRT‑LLM) and runs a relu2_no_mul model, apply() will hit the hard assertion activation == "silu" and abort instead of falling back to the existing non‑flashinfer path, which previously worked (albeit slower). Consider only enabling flashinfer when either the model is gated or the chosen backend is CUTLASS; otherwise leave flashinfer_moe_backend as None so non‑gated models continue to run.

Useful? React with 👍 / 👎.

Signed-off-by: Amir Klein <[email protected]>

tlrmchlsmth · 2025-11-05T12:25:52Z

vllm/model_executor/layers/quantization/modelopt.py

-            g1_alphas=(layer.w13_weight_scale * layer.w13_input_scale).squeeze(),
+            g1_alphas=layer.output1_scales_gate_scalar.squeeze()


It looks like output1_scales_gate_scalar and output2_scales_scalar are only used in flashinfer_trtllm_moe. It's not clear from the vLLM source code how these are getting set. Does flashinfer add them to the layer's member variables?

I think we should try to keep the quantization format decoupled from the kernels used for the implementation

These factors are registered into the layer during process_weights_after_loading (line 566) .
They are registered to the layer when we're using cutlass or trtllm backends, using register_moe_scaling_factors .
The problem I see with trying to make this decouple is that the cutlass path uses the layers' get_fused_moe_quant_config to get the relevant scaling factors here .
There are 2 options I see here let me know which one you prefer:

I can remove the if else when setting the quantization and set it the same for all paths.

I can build the needed quantization for the flashinfer cutlass kernel in here , and not change the quantization in the ModelOptFusedMoEFP8 object.

I can remove the if else when setting the quantization and set it the same for all paths.
I think I'd prefer to just provide the same information in all cases. The kernel can choose to ignore it

mgoin · 2025-11-11T23:17:37Z

vllm/model_executor/layers/quantization/modelopt.py

+            if (
+                self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+                and not self.moe.is_act_and_mul
+            ):
+                logger.info_once(
+                    "Non-gated MoE is not supported for min-latency mode,"
+                    "falling back to high-throughput mode"
+                )


It seems you are missing the override of self.flashinfer_moe_backend here

mgoin · 2025-11-11T23:28:00Z

vllm/model_executor/layers/quantization/modelopt.py

-            g1_alphas=(layer.w13_weight_scale * layer.w13_input_scale).squeeze(),
+            g1_alphas=layer.output1_scales_gate_scalar.squeeze()


I can remove the if else when setting the quantization and set it the same for all paths.
I think I'd prefer to just provide the same information in all cases. The kernel can choose to ignore it

mgoin · 2025-11-11T23:28:20Z

vllm/model_executor/layers/quantization/modelopt.py

+            if (
+                self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
+                and not self.moe.is_act_and_mul
+            ):
+                logger.info_once(
+                    "Non-gated MoE is not supported for min-latency mode,"
+                    "falling back to high-throughput mode"
+                )


It looks like the self.flashinfer_moe_backend override was left out

Signed-off-by: Amir Klein <[email protected]>

vllm/model_executor/layers/quantization/modelopt.py

Signed-off-by: Michael Goin <[email protected]>

mgoin

PTAL at the failing blackwell test

mgoin · 2025-11-14T15:37:50Z

tests/kernels/moe/test_flashinfer.py

+        if activation != "relu2_no_mul":
+            is_gated = False


This doesn't seem right as it breaks test_flashinfer_per_tensor_moe_fp8_no_graph on blackwell
https://buildkite.com/vllm/ci/builds/38920/steps/canvas?jid=019a7fad-270b-4d40-8820-e3a1e75dc35e#019a7fad-270b-4d40-8820-e3a1e75dc35e/102-2387

Yeah it should be if activation == "relu2_no_mul: , I originally wrote it as a one liner but the pre-commit hook complained, and I fixed it incorrectly. Changed it back and this should be correct.

Signed-off-by: Amir Klein <[email protected]>

mergify · 2025-11-15T14:38:05Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @amirkl94.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Amir Klein <[email protected]>

mgoin

Thank you!

Signed-off-by: Bram Wasti <[email protected]>

tomeras91 reviewed Oct 21, 2025

View reviewed changes

vllm/model_executor/layers/quantization/modelopt.py Show resolved Hide resolved

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py Show resolved Hide resolved

amirkl94 force-pushed the feat/relu2-cutlass branch from fb950dd to c595cbb Compare October 23, 2025 18:35

amirkl94 force-pushed the feat/relu2-cutlass branch 2 times, most recently from 98b3df9 to ed78454 Compare November 2, 2025 14:23

mergify bot added the needs-rebase label Nov 5, 2025

Enable cutlass fp8 kernels

ef94e83

Signed-off-by: Amir Klein <[email protected]>

amirkl94 force-pushed the feat/relu2-cutlass branch from ed78454 to ef94e83 Compare November 5, 2025 11:30

amirkl94 requested a review from tomeras91 November 5, 2025 11:30

mergify bot removed the needs-rebase label Nov 5, 2025

Add accuracy test

a928a4d

Signed-off-by: Amir Klein <[email protected]>

amirkl94 marked this pull request as ready for review November 5, 2025 11:57

amirkl94 requested review from WoosukKwon, mgoin, pavanimajety, robertgshaw2-redhat, tlrmchlsmth and yewentao256 as code owners November 5, 2025 11:57

Pre-commit fix

57634c0

Signed-off-by: Amir Klein <[email protected]>

chatgpt-codex-connector bot reviewed Nov 5, 2025

View reviewed changes

mgoin added the nvidia label Nov 5, 2025

github-project-automation bot added this to NVIDIA Nov 5, 2025

Fallback to cutlass when trying non-gated and min-latency mode

9ce625d

Signed-off-by: Amir Klein <[email protected]>

tlrmchlsmth reviewed Nov 5, 2025

View reviewed changes

Merge branch 'main' into feat/relu2-cutlass

bb25647

amirkl94 requested a review from tlrmchlsmth November 10, 2025 09:20

mgoin reviewed Nov 11, 2025

View reviewed changes

CR Fixes

43d242d

Signed-off-by: Amir Klein <[email protected]>

amirkl94 requested a review from mgoin November 12, 2025 08:50

Small fix

76037a1

Signed-off-by: Amir Klein <[email protected]>

mgoin approved these changes Nov 13, 2025

View reviewed changes

github-project-automation bot moved this to In review in NVIDIA Nov 13, 2025

mgoin added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 13, 2025

mgoin reviewed Nov 13, 2025

View reviewed changes

vllm/model_executor/layers/quantization/modelopt.py Outdated Show resolved Hide resolved

mgoin added 2 commits November 13, 2025 13:33

Update vllm/model_executor/layers/quantization/modelopt.py

ad29cb4

Signed-off-by: Michael Goin <[email protected]>

Merge branch 'main' into feat/relu2-cutlass

f57cb15

mgoin requested changes Nov 14, 2025

View reviewed changes

Correct condition

8b3edda

Signed-off-by: Amir Klein <[email protected]>

mergify bot added the needs-rebase label Nov 15, 2025

amirkl94 added 2 commits November 15, 2025 16:38

Simplify

05daf09

Signed-off-by: Amir Klein <[email protected]>

Merge branch 'main' into feat/relu2-cutlass

e6fd46d

Signed-off-by: Amir Klein <[email protected]>

mergify bot removed the needs-rebase label Nov 15, 2025

Merge branch 'main' into feat/relu2-cutlass

67242e7

amirkl94 requested a review from mgoin November 16, 2025 11:11

mgoin approved these changes Nov 16, 2025

View reviewed changes

mgoin merged commit 03ee481 into vllm-project:main Nov 16, 2025
53 checks passed

github-project-automation bot moved this from In review to Done in NVIDIA Nov 16, 2025

bwasti pushed a commit to bwasti/vllm that referenced this pull request Nov 17, 2025

Feature: Support Relu2 in FusedMoE fp8 cutlass path (vllm-project#27261)

e29fbff

Signed-off-by: Bram Wasti <[email protected]>

		g1_alphas=(layer.w13_weight_scale * layer.w13_input_scale).squeeze(),
		g1_alphas=layer.output1_scales_gate_scalar.squeeze()

Uh oh!

Feature: Support Relu2 in FusedMoE fp8 cutlass path #27261

Feature: Support Relu2 in FusedMoE fp8 cutlass path #27261

Uh oh!

Conversation

amirkl94 commented Oct 21, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Tests

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Nov 5, 2025

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mgoin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Nov 15, 2025

Uh oh!

mgoin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

amirkl94 commented Oct 21, 2025 •

edited by github-actions bot

Loading