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
[Perf] Warmup FlashInfer attention during startup (#23439)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni001@gmail.com>
Michael Goin committed
Sep 10, 2025 at 18:03 UTC
fba785658170e8723ccba762654f30c54b0217a4
3 files changed
+55
-20
vllm/model_executor/warmup/kernel_warmup.py
+27
-1
@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING
10
import torch
11
12
import vllm.envs as envs
13
+from vllm.logger import init_logger
14
from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
15
from vllm.platforms import current_platform
16
from vllm.utils.deep_gemm import is_deep_gemm_supported
@@ -19,6 +20,8 @@ if TYPE_CHECKING:
20
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
21
from vllm.v1.worker.gpu_worker import Worker
22
23
+logger = init_logger(__name__)
24
+
25
26
def kernel_warmup(worker: "Worker"):
27
# Deep GEMM warmup
@@ -30,10 +33,33 @@ def kernel_warmup(worker: "Worker"):
33
max_tokens = worker.scheduler_config.max_num_batched_tokens
34
deep_gemm_warmup(model, max_tokens)
35
33
- # FlashInfer autotune for Blackwell (SM 10.0) GPUs
36
+ # FlashInfer kernel autotune for Blackwell (SM 10.0) GPUs
37
if has_flashinfer() and current_platform.is_device_capability(100):
38
flashinfer_autotune(worker.model_runner)
39
40
+ # FlashInfer attention warmup
41
+ # Only warmup if the model has FlashInfer attention groups
42
+ # and is not a pooling model
43
+ def _is_flashinfer_backend(backend):
44
+ try:
45
+ return backend.get_name() == "FLASHINFER_VLLM_V1"
46
+ except NotImplementedError:
47
+ return False
48
+
49
+ if not worker.model_runner.is_pooling_model and all(
50
+ _is_flashinfer_backend(group.backend)
51
+ for groups in worker.model_runner.attn_groups for group in groups):
52
+ logger.info("Warming up FlashInfer attention.")
53
+ # Warmup with mixed batch containing both prefill and decode tokens
54
+ # This is to warm up both prefill and decode attention kernels
55
+ worker.model_runner._dummy_run(
56
+ num_tokens=16,
57
+ skip_eplb=True,
58
+ is_profile=True,
59
+ force_attention=True,
60
+ create_mixed_batch=True,
61
+ )
62
+
63
64
def flashinfer_autotune(runner: "GPUModelRunner") -> None:
65
"""
vllm/v1/attention/backends/flashinfer.py
-16
@@ -549,22 +549,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
549
)
550
return attn_metadata
551
552
- def build_for_cudagraph_capture(
553
- self, common_attn_metadata: CommonAttentionMetadata):
554
- """
555
- This method builds the metadata for full cudagraph capture.
556
- Currently, only decode is supported for full cudagraphs with FlashInfer.
557
- """
558
- m = common_attn_metadata
559
-
560
- assert m.num_reqs == m.num_actual_tokens, \
561
- "FlashInfer only supports decode-only full CUDAGraph capture. " \
562
- "Make sure all cudagraph capture sizes <= max_num_seq."
563
-
564
- m.max_query_len = 1 # decode-only
565
-
566
- return self.build(0, m)
567
-
552
def use_cascade_attention(self, *args, **kwargs) -> bool:
553
if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
554
# TODO: The cascade wrapper currently does not support setting
vllm/v1/worker/gpu_model_runner.py
+28
-3
@@ -2578,6 +2578,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2578
uniform_decode: bool = False,
2579
skip_eplb: bool = False,
2580
is_profile: bool = False,
2581
+ create_mixed_batch: bool = False,
2582
remove_lora: bool = True,
2583
) -> tuple[torch.Tensor, torch.Tensor]:
2584
"""
@@ -2596,6 +2597,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2597
uniform_decode: If True, the batch is a uniform decode batch.
2598
skip_eplb: If True, skip EPLB state update.
2599
is_profile: If True, this is a profile run.
2600
+ create_mixed_batch: If True, create a mixed batch with both decode
2601
+ (1 token) and prefill (multiple tokens) requests.
2602
remove_lora: If False, dummy LoRAs are not destroyed after the run
2603
"""
2604
assert cudagraph_runtime_mode in {
@@ -2627,7 +2630,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2630
# has num_tokens in total.
2631
assert num_tokens <= self.scheduler_config.max_num_batched_tokens
2632
max_num_reqs = self.scheduler_config.max_num_seqs
2630
- if uniform_decode:
2633
+ if create_mixed_batch:
2634
+ assert not uniform_decode
2635
+ # Create mixed batch:
2636
+ # first half decode tokens, second half one prefill
2637
+ num_decode_tokens = num_tokens // 2
2638
+ num_prefill_tokens = num_tokens - num_decode_tokens
2639
+ num_reqs = num_decode_tokens + 1
2640
+
2641
+ # Create decode requests (1 token each) followed by prefill request
2642
+ num_scheduled_tokens_list = [1] * num_decode_tokens + [
2643
+ num_prefill_tokens
2644
+ ]
2645
+ # Note: Overriding max_query_len to be the prefill tokens
2646
+ max_query_len = num_prefill_tokens
2647
+ elif uniform_decode:
2648
+ assert not create_mixed_batch
2649
num_reqs = cdiv(num_tokens, max_query_len)
2650
assert num_reqs <= max_num_reqs, \
2651
"Do not capture num_reqs > max_num_reqs for uniform batch"
@@ -2652,8 +2670,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2670
if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
2671
attn_metadata = {}
2672
2655
- # Make sure max_model_len is used at the graph capture time.
2656
- self.seq_lens.np[:num_reqs] = self.max_model_len
2673
+ if create_mixed_batch:
2674
+ # In the mixed batch mode (used for FI warmup), we use
2675
+ # shorter sequence lengths to run faster.
2676
+ # TODO(luka) better system for describing dummy batches
2677
+ seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
2678
+ else:
2679
+ # Make sure max_model_len is used at the graph capture time.
2680
+ seq_lens = self.max_model_len
2681
+ self.seq_lens.np[:num_reqs] = seq_lens
2682
self.seq_lens.np[num_reqs:] = 0
2683
self.seq_lens.copy_to_gpu()
2684