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 14, 05:30 UTC
Update Flashinfer to 0.2.14.post1 (#23537)
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com> Signed-off-by: siyuanf <siyuanf@nvidia.com> Signed-off-by: Weiliang Liu <weiliangl@nvidia.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Siyuan Fu <siyuanf@nvidia.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
weiliang committed
Aug 26, 2025 at 09:30 UTC
ae067888d6803b0fe0a2201ae9b464a848a0de01
5 files changed
+14
-7
docker/Dockerfile
+1
-1
@@ -373,7 +373,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
373
# Install FlashInfer from source
374
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
375
# Keep this in sync with "flashinfer" extra in setup.py
376
-ARG FLASHINFER_GIT_REF="v0.2.12"
376
+ARG FLASHINFER_GIT_REF="v0.2.14.post1"
377
# Flag to control whether to compile FlashInfer AOT kernels
378
# Set to "true" to enable AOT compilation:
379
# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
setup.py
+1
-1
@@ -694,7 +694,7 @@ setup(
694
"mistral_common[audio]"], # Required for audio processing
695
"video": [], # Kept for backwards compatibility
696
# FlashInfer should be updated together with the Dockerfile
697
- "flashinfer": ["flashinfer-python==0.2.12"],
697
+ "flashinfer": ["flashinfer-python==0.2.14.post1"],
698
# Optional deps for AMD FP4 quantization support
699
"petit-kernel": ["petit-kernel"],
700
},
vllm/compilation/collective_fusion.py
+2
-1
@@ -465,7 +465,8 @@ if flashinfer_comm is not None:
465
quant_out=quant_out,
466
scale_out=scale_out,
467
# in vllm we only support swizzled layout
468
- layout_code=flashinfer_comm.FP4QuantizationSFLayout.SWIZZLED,
468
+ layout_code=flashinfer_comm.QuantizationSFLayout.
469
+ SWIZZLED_128x4,
470
scale_factor=scale_factor,
471
)
472
else:
vllm/model_executor/layers/quantization/mxfp4.py
+6
-1
@@ -6,6 +6,7 @@ import torch
6
from torch.nn.parameter import Parameter
7
8
from vllm import envs
9
+from vllm.config import get_current_vllm_config
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
12
FusedMoEMethodBase)
@@ -113,6 +114,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
114
self.topk_indices_dtype = None
115
self.moe = moe
116
self.use_marlin = self._should_use_marlin()
117
+ self.max_capture_size = get_current_vllm_config(
118
+ ).compilation_config.max_capture_size
119
120
if current_platform.is_device_capability(100) and not has_flashinfer():
121
logger.warning_once(
@@ -520,7 +523,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
523
x_scale = None
524
else:
525
x_quant, x_scale = mxfp8_quantize(x, False) # to mxfp8
523
- x_scale = x_scale.view(torch.float8_e4m3fn).reshape(-1)
526
+ x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
527
+ *x.shape[:-1], -1)
528
trtllm_gen_output = trtllm_fp4_block_scale_moe(
529
router_logits.to(torch.bfloat16),
530
None, # routing_bias
@@ -549,6 +553,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
553
self._get_tile_tokens_dim(x, top_k),
554
1 if renormalize else 0, # routing_method_type, renormalize
555
True, # do finalize
556
+ tune_max_num_tokens=self.max_capture_size,
557
)[0]
558
return trtllm_gen_output
559
else:
vllm/v1/worker/gpu_worker.py
+4
-3
@@ -310,6 +310,10 @@ class Worker(WorkerBase):
310
logger.info("Compile and warming up model for size %d", size)
311
self.model_runner._dummy_run(size, skip_eplb=True)
312
313
+ # Warmup and tune the kernels used during model execution before
314
+ # cuda graph capture.
315
+ kernel_warmup(self)
316
+
317
if not self.model_config.enforce_eager:
318
self.model_runner.capture_model()
319
@@ -334,9 +338,6 @@ class Worker(WorkerBase):
338
self.model_runner._dummy_sampler_run(
339
hidden_states=last_hidden_states)
340
337
- # Warmup kernels used during model execution
338
- kernel_warmup(self)
339
-
341
# Reset the seed to ensure that the random state is not affected by
342
# the model initialization and profiling.
343
set_random_seed(self.model_config.seed)