Bump Flashinfer version to 0.6.16.post3 (#50892)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
Wei Zhao committed
Aug 9, 2026 at 03:54 UTC
f18e10a7e1c1eb37e898c31989ff625d79791657
5 files changed
+51
-58
docker/Dockerfile
+1
-1
@@ -752,7 +752,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
752
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
753
# https://docs.flashinfer.ai/installation.html
754
# From versions.json: .flashinfer.version
755
-ARG FLASHINFER_VERSION=0.6.15.post1
755
+ARG FLASHINFER_VERSION=0.6.16.post3
756
RUN --mount=type=cache,target=/opt/uv/cache \
757
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
758
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
docker/versions.json
+1
-1
@@ -68,7 +68,7 @@
68
"default": "true"
69
},
70
"FLASHINFER_VERSION": {
71
- "default": "0.6.15.post1"
71
+ "default": "0.6.16.post3"
72
},
73
"GDRCOPY_CUDA_VERSION": {
74
"default": "12.8"
requirements/cuda.txt
+2
-2
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
14
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
15
# install_requires so the published wheel does not carry an unresolvable pin
16
--extra-index-url https://flashinfer.ai/whl/
17
-flashinfer-python==0.6.15.post1
18
-flashinfer-cubin==0.6.15.post1
17
+flashinfer-python==0.6.16.post3
18
+flashinfer-cubin==0.6.16.post3
19
apache-tvm-ffi==0.1.11
20
tilelang==0.1.12
21
nvidia-cudnn-frontend>=1.19.1
vllm/model_executor/warmup/kernel_warmup.py
+36
-51
@@ -229,12 +229,19 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
229
Without autotuning, FlashInfer will rely on heuristics, which may
230
be significantly slower.
231
232
- Tuning is performed only on rank 0. The resulting cache is broadcast
233
- to every rank so all ranks dispatch the same kernel tactic.
232
+ Every rank profiles the same tactics. When distributed, per-tactic
233
+ timings are averaged over the world CPU group so all ranks select the
234
+ same tactic.
235
"""
236
+ from flashinfer.autotuner import AutoTuner, set_autotune_process_group
237
+
238
import vllm.utils.flashinfer as fi_utils
239
from vllm.distributed.parallel_state import get_world_group
240
241
+ world = get_world_group()
242
+ is_leader = world.rank_in_group == 0
243
+ tuner = AutoTuner.get()
244
+
245
autotune_kwargs: dict = {}
246
skip_ops = _flashinfer_autotune_skip_ops(runner)
247
if skip_ops:
@@ -244,25 +251,6 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
251
)
252
autotune_kwargs["skip_ops"] = skip_ops
253
247
- use_persistent_cache = True
248
-
249
- # When distributed, tune on every rank so the collectives stay synchronized.
250
- if get_world_group().world_size > 1:
251
- use_persistent_cache = False
252
-
253
- if not use_persistent_cache:
254
- with torch.inference_mode(), fi_utils.autotune(**autotune_kwargs):
255
- runner._dummy_run(
256
- num_tokens=runner.scheduler_config.max_num_batched_tokens,
257
- skip_eplb=True,
258
- is_profile=True,
259
- )
260
- get_world_group().barrier()
261
- return
262
-
263
- world = get_world_group()
264
- is_leader = world.rank_in_group == 0
265
-
254
cache_path = resolve_flashinfer_autotune_file(runner)
255
if is_leader:
256
logger.info_once("Using FlashInfer autotune cache file: %s", cache_path)
@@ -271,43 +259,40 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
259
# When autotuning with number of tokens m, flashinfer will autotune
260
# operations for all number of tokens up to m, so we only need to
261
# run with the max number of tokens.
262
+ # Randomize inputs to avoid every token pick the same experts,
263
+ # which lead to some EP ranks receiving no tokens and skipping their
264
+ # MoE kernel entirely, and cause hang due to all-reduce collective
265
+ # during synchronized autotuning.
266
dummy_run_kwargs = dict(
267
num_tokens=runner.scheduler_config.max_num_batched_tokens,
268
skip_eplb=True,
269
is_profile=True,
270
+ randomize_inputs=True,
271
)
272
280
- with torch.inference_mode():
281
- if is_leader:
282
- with fi_utils.autotune(
283
- tune_mode=True, cache=str(cache_path), **autotune_kwargs
284
- ):
285
- runner._dummy_run(**dummy_run_kwargs)
286
- else:
287
- runner._dummy_run(**dummy_run_kwargs)
288
-
289
- # Broadcast autotune cache from rank 0 to all other ranks so every
290
- # rank loads the same set of chosen tactics.
291
- tune_results: bytes | None = None
273
+ # Read cached autotune results and broadcast to all ranks.
274
+ cached_results: bytes | None = None
275
if is_leader and cache_path.exists():
276
with open(cache_path, "rb") as f:
294
- tune_results = f.read()
295
-
296
- tune_results = world.broadcast_object(tune_results, src=0)
297
-
298
- if tune_results is None:
299
- logger.warning_once(
300
- "No FlashInfer autotune cache entries found; "
301
- "falling back to default tactics."
302
- )
303
- else:
304
- write_flashinfer_autotune_cache(cache_path, tune_results)
277
+ cached_results = f.read()
278
+ cached_results = world.broadcast_object(cached_results, src=0)
279
+ if cached_results is not None:
280
+ write_flashinfer_autotune_cache(cache_path, cached_results)
281
world.barrier()
306
- from flashinfer.autotuner import AutoTuner
282
+ tuner.load_configs(str(cache_path))
283
+
284
+ group = world.cpu_group if world.world_size > 1 else None
285
+ set_autotune_process_group(group)
286
+ try:
287
+ with (
288
+ torch.inference_mode(),
289
+ fi_utils.autotune(tune_mode=True, **autotune_kwargs),
290
+ ):
291
+ runner._dummy_run(**dummy_run_kwargs)
292
+ finally:
293
+ set_autotune_process_group(None)
294
308
- AutoTuner.get().load_configs(str(cache_path))
309
- logger.info_once(
310
- "FlashInfer autotune cache loaded on rank %d from %s.",
311
- world.rank_in_group,
312
- cache_path,
313
- )
295
+ if world.world_size > 1:
296
+ world.barrier()
297
+ if is_leader:
298
+ tuner.save_configs(str(cache_path))
vllm/v1/worker/gpu_model_runner.py
+11
-3
@@ -5824,7 +5824,10 @@ class GPUModelRunner(
5824
5825
@contextmanager
5826
def maybe_randomize_inputs(
5827
- self, input_ids: torch.Tensor | None, inputs_embeds: torch.Tensor | None
5827
+ self,
5828
+ input_ids: torch.Tensor | None,
5829
+ inputs_embeds: torch.Tensor | None,
5830
+ randomize_inputs: bool = False,
5831
):
5832
"""
5833
Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set.
@@ -5834,7 +5837,9 @@ class GPUModelRunner(
5837
"""
5838
5839
dp_size = self.vllm_config.parallel_config.data_parallel_size
5837
- randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1
5840
+ randomize_inputs = randomize_inputs or (
5841
+ envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1
5842
+ )
5843
if not randomize_inputs:
5844
yield
5845
elif input_ids is not None:
@@ -5911,6 +5916,7 @@ class GPUModelRunner(
5916
is_graph_capturing: bool = False,
5917
num_active_loras: int = 0,
5918
profile_seq_lens: int | None = None,
5919
+ randomize_inputs: bool = False,
5920
) -> tuple[torch.Tensor, torch.Tensor]:
5921
"""
5922
Run a dummy forward pass to warm up/profile run or capture the
@@ -6208,7 +6214,9 @@ class GPUModelRunner(
6214
num_tokens_across_dp[:] = num_tokens_padded
6215
6216
with (
6211
- self.maybe_randomize_inputs(input_ids, inputs_embeds),
6217
+ self.maybe_randomize_inputs(
6218
+ input_ids, inputs_embeds, randomize_inputs=randomize_inputs
6219
+ ),
6220
set_forward_context(
6221
attn_metadata,
6222
self.vllm_config,