A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 12:00 UTC
[Perf] Eliminate padding and slicing op for GPT-OSS with Flashinfer MXFP4 MXFP8 MoE (#30647)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Mar 18, 2026 at 23:01 UTC
296839a1b07e63daecca67bfce80375614b5b863
6 files changed
+40
-3
tests/compile/fusions_e2e/conftest.py
+4
@@ -82,6 +82,10 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
82
f"attention backend '{attn_backend.backend.name}'"
83
)
84
85
+ # TODO: remove this after finishing migration from envs to model kwargs
86
+ if model_name == "openai/gpt-oss-20b":
87
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
88
+
89
# Disable, compile cache to make sure custom passes run.
90
# Otherwise, we can't verify fusion happened through the logs.
91
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
tests/compile/fusions_e2e/models.py
+9
@@ -162,3 +162,12 @@ deepseek_v3_fp8 = ModelFusionInfo(
162
# async_tp=n_layers * 2,
163
),
164
)
165
+
166
+gpt_oss_20b = ModelFusionInfo(
167
+ model_name="openai/gpt-oss-20b",
168
+ matches=lambda n_layers: Matches(
169
+ ar_rms_fusion=n_layers * 2 + 1,
170
+ sequence_parallel=n_layers * 2 + 1,
171
+ async_tp=n_layers * 2,
172
+ ),
173
+)
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+2
-1
@@ -20,6 +20,7 @@ from .models import (
20
FLASHINFER_MLA_ATTN,
21
TRITON_ATTN,
22
deepseek_v3_fp8,
23
+ gpt_oss_20b,
24
llama3_8b,
25
llama3_8b_fp4,
26
llama3_8b_fp8,
@@ -158,7 +159,7 @@ def test_tp2_ar_rms_fp4_fusions(
159
@multi_gpu_test(num_gpus=2)
160
@pytest.mark.parametrize(
161
"model_name, matches_fn, model_kwargs, hf_overrides",
161
- [llama3_8b, qwen3_a3b],
162
+ [llama3_8b, qwen3_a3b, gpt_oss_20b],
163
)
164
@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
165
@pytest.mark.parametrize("n_layers", [4])
vllm/model_executor/layers/fused_moe/fused_moe_method_base.py
+5
@@ -101,6 +101,11 @@ class FusedMoEMethodBase(QuantizeMethodBase):
101
return self.moe_kernel.prepare_finalize.topk_indices_dtype()
102
return None
103
104
+ @property
105
+ def skip_forward_padding(self) -> bool:
106
+ """Whether to skip the padding in the forward before applying the moe method."""
107
+ return False
108
+
109
@property
110
def supports_eplb(self) -> bool:
111
return False
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+4
-1
@@ -415,7 +415,10 @@ class DefaultMoERunner(MoERunner):
415
416
# This is the dimension after transform (for routed expert output slicing)
417
transformed_hidden_dim = hidden_states.shape[-1]
418
- if self.moe_config.hidden_dim != transformed_hidden_dim:
418
+ if (
419
+ not self.quant_method.skip_forward_padding
420
+ and self.moe_config.hidden_dim != transformed_hidden_dim
421
+ ):
422
hidden_states = F.pad(
423
hidden_states,
424
(0, self.moe_config.hidden_dim - transformed_hidden_dim),
vllm/model_executor/layers/quantization/mxfp4.py
+16
-1
@@ -294,6 +294,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
294
# Initialized in process_weights_after_loading for CUTLASS/SM90 backends
295
self.moe_kernel: mk.FusedMoEKernel | None = None
296
297
+ @property
298
+ def skip_forward_padding(self) -> bool:
299
+ # SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant
300
+ # so can skip the padding in the forward before applying the moe method
301
+ return self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
302
+
303
def create_weights(
304
self,
305
layer: torch.nn.Module,
@@ -1130,9 +1136,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1136
elif self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM:
1137
from flashinfer import mxfp8_quantize
1138
1133
- x_quant, x_scale = mxfp8_quantize(x, False) # to mxfp8
1139
+ # x_quant is padded in hidden dimension with alignment=256
1140
+ x_quant, x_scale = mxfp8_quantize(
1141
+ x,
1142
+ is_sf_swizzled_layout=False,
1143
+ alignment=256,
1144
+ )
1145
x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)
1146
1147
+ # output with original unpadded hidden size
1148
+ output = torch.empty_like(x)
1149
+
1150
trtllm_gen_output = trtllm_fp4_block_scale_moe(
1151
routing_logits=router_logits.to(torch.bfloat16),
1152
routing_bias=None,
@@ -1161,6 +1175,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1175
routing_method_type=1 if layer.renormalize else 0,
1176
do_finalize=True,
1177
tune_max_num_tokens=max(self.max_capture_size, 1),
1178
+ output=output,
1179
)[0]
1180
return trtllm_gen_output
1181
elif self.mxfp4_backend == Mxfp4Backend.CK: