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 13, 20:45 UTC
[BugFix][Performance] Restore flashinfer autotuning for all scenarios (#27904)
Varun Sundar Rabindranath committed
Nov 4, 2025 at 02:56 UTC
4022a9d279d09efe1b8a36ff3531bf1d4c8f08ca
4 files changed
+14
-44
tests/quantization/test_blackwell_moe.py
+2
-14
@@ -172,21 +172,9 @@ def test_gptoss_mxfp4mxfp8_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch
172
can_initialize("openai/gpt-oss-20b", hf_overrides=HF_OVERRIDE_TEXT)
173
174
175
-def test_gptoss_dp2_mxfp4mxfp8_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
176
- monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
177
- monkeypatch.setenv("VLLM_ALL2ALL_BACKEND", "deepep_high_throughput")
178
- can_initialize(
179
- "openai/gpt-oss-20b",
180
- extra_args=["--data-parallel-size", "2", "--enable-expert-parallel"],
181
- hf_overrides=HF_OVERRIDE_TEXT,
182
- )
183
-
184
-
185
-def test_gptoss_dp2_mxfp4bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
186
- monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16", "1")
187
- monkeypatch.setenv("VLLM_ALL2ALL_BACKEND", "deepep_high_throughput")
175
+def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch):
176
can_initialize(
177
"openai/gpt-oss-20b",
190
- extra_args=["--data-parallel-size", "2", "--enable-expert-parallel"],
178
hf_overrides=HF_OVERRIDE_TEXT,
179
+ extra_args=["--enforce-eager"],
180
)
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+9
-2
@@ -127,10 +127,17 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
127
"routing_method_type": 1,
128
"do_finalize": True,
129
"output": output,
130
- "tune_max_num_tokens": self.max_capture_size,
130
+ "tune_max_num_tokens": max(self.max_capture_size, 1),
131
}
132
133
from flashinfer import trtllm_fp4_block_scale_routed_moe
134
135
- trtllm_fp4_block_scale_routed_moe(**kwargs)
135
+ from vllm.utils.flashinfer import autotune
136
+
137
+ with autotune(False):
138
+ # Enable autotune when,
139
+ # https://github.com/flashinfer-ai/flashinfer/issues/2023 is
140
+ # resolved.
141
+ trtllm_fp4_block_scale_routed_moe(**kwargs)
142
+
143
return output
vllm/model_executor/layers/quantization/mxfp4.py
+2
-2
@@ -1047,7 +1047,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1047
None,
1048
1 if renormalize else 0, # routing_method_type, renormalize
1049
True, # do finalize
1050
- tune_max_num_tokens=self.max_capture_size,
1050
+ tune_max_num_tokens=max(self.max_capture_size, 1),
1051
)[0]
1052
return trtllm_gen_output
1053
elif (
@@ -1122,7 +1122,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1122
tp_rank=self.moe.tp_rank,
1123
ep_size=self.moe.ep_size,
1124
ep_rank=self.moe.ep_rank,
1125
- tune_max_num_tokens=self.max_capture_size,
1125
+ tune_max_num_tokens=max(self.max_capture_size, 1),
1126
**extra_kwargs,
1127
)
1128
vllm/model_executor/warmup/kernel_warmup.py
+1
-26
@@ -11,7 +11,6 @@ from typing import TYPE_CHECKING
11
import torch
12
13
import vllm.envs as envs
14
-from vllm.config import CUDAGraphMode, VllmConfig
14
from vllm.logger import init_logger
15
from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
16
from vllm.platforms import current_platform
@@ -25,26 +24,6 @@ if TYPE_CHECKING:
24
logger = init_logger(__name__)
25
26
28
-def flashinfer_autotune_supported(vllm_config: VllmConfig) -> bool:
29
- """
30
- Record known issues with vllm + flashinfer autotune here. Return True if
31
- and only if flashinfer autotune will run through without issues.
32
- """
33
- is_tp_or_dp = (vllm_config.parallel_config.data_parallel_size > 1) or (
34
- vllm_config.parallel_config.tensor_parallel_size > 1
35
- )
36
- is_fi_mxfp4_backend = (
37
- envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
38
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16
39
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS
40
- ) or (
41
- current_platform.is_cuda() and current_platform.is_device_capability(100)
42
- ) # on >=sm100, default mxfp4 backend is flashinfer
43
- is_eager = vllm_config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
44
-
45
- return not (is_tp_or_dp and is_fi_mxfp4_backend and is_eager)
46
-
47
-
27
def kernel_warmup(worker: "Worker"):
28
# Deep GEMM warmup
29
do_deep_gemm_warmup = (
@@ -58,11 +37,7 @@ def kernel_warmup(worker: "Worker"):
37
deep_gemm_warmup(model, max_tokens)
38
39
# FlashInfer autotune for Hopper (SM 9.0) and Blackwell (SM 10.0) GPUs
61
- if (
62
- has_flashinfer()
63
- and current_platform.has_device_capability(90)
64
- and flashinfer_autotune_supported(worker.vllm_config)
65
- ):
40
+ if has_flashinfer() and current_platform.has_device_capability(90):
41
flashinfer_autotune(worker.model_runner)
42
43
# FlashInfer attention warmup