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 3, 10:30 UTC
[Perf] Re-enable flashinfer autotune by default and cleanup (#42857)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
May 18, 2026 at 12:12 UTC
8c296de63b47664fc5979831e1ae2d2a14a05b1a
5 files changed
+112
-80
vllm/config/vllm.py
+2
-6
@@ -223,9 +223,7 @@ OPTIMIZATION_LEVEL_01 = {
223
"use_inductor_graph_partition": False,
224
},
225
"kernel_config": {
226
- # Disabled for now due to correctness issues:
227
- # https://github.com/flashinfer-ai/flashinfer/issues/3197
228
- "enable_flashinfer_autotune": False,
226
+ "enable_flashinfer_autotune": True,
227
},
228
}
229
OPTIMIZATION_LEVEL_02 = {
@@ -246,9 +244,7 @@ OPTIMIZATION_LEVEL_02 = {
244
"use_inductor_graph_partition": False,
245
},
246
"kernel_config": {
249
- # Disabled for now due to correctness issues:
250
- # https://github.com/flashinfer-ai/flashinfer/issues/3197
251
- "enable_flashinfer_autotune": False,
247
+ "enable_flashinfer_autotune": True,
248
},
249
}
250
OPTIMIZATION_LEVEL_03 = {
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py
+18
-21
@@ -149,24 +149,21 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
149
# The functional API expects x_sf with trailing dim: (M, K//16, 1).
150
x_sf = a1q_scale.unsqueeze(-1)
151
152
- from vllm.utils.flashinfer import _is_fi_autotuning, autotune
153
-
154
- with autotune(_is_fi_autotuning):
155
- flashinfer_cute_dsl_fused_moe_nvfp4(
156
- x=hidden_states,
157
- x_sf=x_sf,
158
- token_selected_experts=topk_ids.to(torch.int32),
159
- token_final_scales=topk_weights.float(),
160
- w1_weight=w1,
161
- w1_weight_sf=self.w1_scale,
162
- w1_alpha=self.g1_alphas,
163
- fc2_input_scale=self.a2_gscale,
164
- w2_weight=w2,
165
- w2_weight_sf=self.w2_scale,
166
- w2_alpha=self.g2_alphas,
167
- num_experts=self.global_num_experts,
168
- top_k=self.topk,
169
- num_local_experts=self.local_num_experts,
170
- local_expert_offset=self.local_expert_offset,
171
- moe_output=output,
172
- )
152
+ flashinfer_cute_dsl_fused_moe_nvfp4(
153
+ x=hidden_states,
154
+ x_sf=x_sf,
155
+ token_selected_experts=topk_ids.to(torch.int32),
156
+ token_final_scales=topk_weights.float(),
157
+ w1_weight=w1,
158
+ w1_weight_sf=self.w1_scale,
159
+ w1_alpha=self.g1_alphas,
160
+ fc2_input_scale=self.a2_gscale,
161
+ w2_weight=w2,
162
+ w2_weight_sf=self.w2_scale,
163
+ w2_alpha=self.g2_alphas,
164
+ num_experts=self.global_num_experts,
165
+ top_k=self.topk,
166
+ num_local_experts=self.local_num_experts,
167
+ local_expert_offset=self.local_expert_offset,
168
+ moe_output=output,
169
+ )
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+31
-37
@@ -194,39 +194,36 @@ class TrtLlmMxfp4ExpertsMonolithic(
194
device=hidden_states.device,
195
)
196
197
- from vllm.utils.flashinfer import _is_fi_autotuning, autotune
198
-
199
- with autotune(_is_fi_autotuning):
200
- trtllm_fp4_block_scale_moe(
201
- routing_logits=router_logits.to(torch.bfloat16),
202
- routing_bias=None,
203
- hidden_states=x_quant,
204
- hidden_states_scale=x_scale,
205
- gemm1_weights=w1,
206
- gemm1_weights_scale=self.w1_scale,
207
- gemm1_bias=self.w1_bias,
208
- gemm1_alpha=self.gemm1_alpha,
209
- gemm1_beta=self.gemm1_beta,
210
- gemm1_clamp_limit=self.gemm1_clamp_limit,
211
- gemm2_weights=w2,
212
- gemm2_weights_scale=self.w2_scale,
213
- gemm2_bias=self.w2_bias,
214
- output1_scale_scalar=None,
215
- output1_scale_gate_scalar=None,
216
- output2_scale_scalar=None,
217
- num_experts=global_num_experts,
218
- top_k=self.topk,
219
- n_group=None,
220
- topk_group=None,
221
- intermediate_size=self.intermediate_size_per_partition,
222
- local_expert_offset=self.ep_rank * self.local_num_experts,
223
- local_num_experts=self.local_num_experts,
224
- routed_scaling_factor=None,
225
- routing_method_type=self.routing_method_type,
226
- do_finalize=True,
227
- tune_max_num_tokens=max(self.max_capture_size, 1),
228
- output=output,
229
- )
197
+ trtllm_fp4_block_scale_moe(
198
+ routing_logits=router_logits.to(torch.bfloat16),
199
+ routing_bias=None,
200
+ hidden_states=x_quant,
201
+ hidden_states_scale=x_scale,
202
+ gemm1_weights=w1,
203
+ gemm1_weights_scale=self.w1_scale,
204
+ gemm1_bias=self.w1_bias,
205
+ gemm1_alpha=self.gemm1_alpha,
206
+ gemm1_beta=self.gemm1_beta,
207
+ gemm1_clamp_limit=self.gemm1_clamp_limit,
208
+ gemm2_weights=w2,
209
+ gemm2_weights_scale=self.w2_scale,
210
+ gemm2_bias=self.w2_bias,
211
+ output1_scale_scalar=None,
212
+ output1_scale_gate_scalar=None,
213
+ output2_scale_scalar=None,
214
+ num_experts=global_num_experts,
215
+ top_k=self.topk,
216
+ n_group=None,
217
+ topk_group=None,
218
+ intermediate_size=self.intermediate_size_per_partition,
219
+ local_expert_offset=self.ep_rank * self.local_num_experts,
220
+ local_num_experts=self.local_num_experts,
221
+ routed_scaling_factor=None,
222
+ routing_method_type=self.routing_method_type,
223
+ do_finalize=True,
224
+ tune_max_num_tokens=max(self.max_capture_size, 1),
225
+ output=output,
226
+ )
227
228
return output
229
@@ -350,9 +347,6 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
347
348
from flashinfer import trtllm_fp4_block_scale_routed_moe
349
353
- from vllm.utils.flashinfer import _is_fi_autotuning, autotune
354
-
355
- with autotune(_is_fi_autotuning):
356
- trtllm_fp4_block_scale_routed_moe(**kwargs)
350
+ trtllm_fp4_block_scale_routed_moe(**kwargs)
351
352
return output
vllm/model_executor/warmup/kernel_warmup.py
+61
-15
@@ -87,23 +87,69 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
87
future calls to FlashInfer will use the best implementation.
88
Without autotuning, FlashInfer will rely on heuristics, which may
89
be significantly slower.
90
+
91
+ Tuning is performed only on rank 0. The resulting cache is broadcast
92
+ to every rank so all ranks dispatch the same kernel tactic.
93
"""
94
+ import os
95
+ import tempfile
96
+
97
import vllm.utils.flashinfer as fi_utils
98
+ from vllm.distributed.parallel_state import get_world_group
99
93
- with torch.inference_mode(), fi_utils.autotune():
94
- # Certain FlashInfer kernels (e.g. nvfp4 routed moe) are
95
- # incompatible with autotuning. This state is used to skip
96
- # those kernels during the autotuning process.
97
- fi_utils._is_fi_autotuning = True
98
-
99
- # We skip EPLB here since we don't want to record dummy metrics
100
- # When autotuning with number of tokens m, flashinfer will autotune
101
- # operations for all number of tokens up to m.
102
- # So we only need to run with the max number of tokens.
103
- runner._dummy_run(
104
- runner.scheduler_config.max_num_batched_tokens,
105
- skip_eplb=True,
106
- is_profile=True,
100
+ world = get_world_group()
101
+ is_leader = world.rank_in_group == 0
102
+
103
+ cache_dir = tempfile.mkdtemp(prefix="vllm_flashinfer_autotune_")
104
+ cache_path = os.path.join(cache_dir, "autotune_cache.json")
105
+
106
+ # We skip EPLB here since we don't want to record dummy metrics.
107
+ # When autotuning with number of tokens m, flashinfer will autotune
108
+ # operations for all number of tokens up to m, so we only need to
109
+ # run with the max number of tokens.
110
+ dummy_run_kwargs = dict(
111
+ num_tokens=runner.scheduler_config.max_num_batched_tokens,
112
+ skip_eplb=True,
113
+ is_profile=True,
114
+ )
115
+
116
+ with torch.inference_mode():
117
+ if is_leader:
118
+ with fi_utils.autotune(tune_mode=True, cache=cache_path):
119
+ runner._dummy_run(**dummy_run_kwargs)
120
+ else:
121
+ runner._dummy_run(**dummy_run_kwargs)
122
+
123
+ # Broadcast autotune cache from rank 0 to all other ranks so every
124
+ # rank loads the same set of chosen tactics.
125
+ tune_results: bytes | None = None
126
+ if is_leader and os.path.exists(cache_path):
127
+ with open(cache_path, "rb") as f:
128
+ tune_results = f.read()
129
+
130
+ tune_results = world.broadcast_object(tune_results, src=0)
131
+
132
+ if tune_results is None:
133
+ logger.warning(
134
+ "No FlashInfer autotune cache entries found."
135
+ "Falling back to default tactics."
136
+ )
137
+ else:
138
+ if not is_leader:
139
+ with open(cache_path, "wb") as f:
140
+ f.write(tune_results)
141
+ from flashinfer.autotuner import AutoTuner
142
+
143
+ AutoTuner.get().load_configs(cache_path)
144
+ logger.info(
145
+ "FlashInfer autotune cache loaded on rank %d from %s.",
146
+ world.rank_in_group,
147
+ cache_path,
148
)
149
109
- fi_utils._is_fi_autotuning = False
150
+ try:
151
+ if os.path.exists(cache_path):
152
+ os.unlink(cache_path)
153
+ os.rmdir(cache_dir)
154
+ except OSError:
155
+ pass
vllm/utils/flashinfer.py
-1
@@ -144,7 +144,6 @@ autotune = _lazy_import_wrapper(
144
"autotune",
145
fallback_fn=lambda *args, **kwargs: contextlib.nullcontext(),
146
)
147
-_is_fi_autotuning: bool = False
147
148
149
@functools.cache