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 7, 14:00 UTC
[Model Runner V2] Warmup kernels (#35172)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 27, 2026 at 10:43 UTC
b1d9f5372d802513e9e009a5d572dd594a09e1dc
3 files changed
+133
-6
vllm/sampling_params.py
+18
@@ -840,6 +840,24 @@ class SamplingParams(
840
f"extra_args={self.extra_args})"
841
)
842
843
+ @staticmethod
844
+ def for_sampler_warmup() -> "SamplingParams":
845
+ """Set parameters to exercise all sampler logic."""
846
+ return SamplingParams(
847
+ temperature=0.9,
848
+ top_p=0.9,
849
+ top_k=50,
850
+ min_p=0.1,
851
+ frequency_penalty=0.5,
852
+ presence_penalty=0.5,
853
+ repetition_penalty=1.2,
854
+ min_tokens=2,
855
+ logit_bias={0: -1.0, 1: 0.5},
856
+ _bad_words_token_ids=[[0], [1, 2]],
857
+ logprobs=5,
858
+ prompt_logprobs=1,
859
+ )
860
+
861
862
class BeamSearchParams(
863
msgspec.Struct,
vllm/v1/worker/gpu/warmup.py
new
+105
@@ -0,0 +1,105 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import numpy as np
5
+import torch
6
+
7
+from vllm import PoolingParams, SamplingParams
8
+from vllm.v1.core.sched.output import (
9
+ CachedRequestData,
10
+ GrammarOutput,
11
+ NewRequestData,
12
+ SchedulerOutput,
13
+)
14
+from vllm.v1.request import Request
15
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
16
+
17
+
18
+@torch.inference_mode()
19
+def warmup_kernels(model_runner: GPUModelRunner) -> None:
20
+ """Run two execute_model + sample_tokens iterations to JIT compile
21
+ triton kernels.
22
+
23
+ The first iteration simulates a prefill with requests of 2 prompt
24
+ tokens each. The second iteration simulates a decode step with all
25
+ requests generating 1 token each.
26
+ """
27
+ prompt_token_ids = [0, 1]
28
+ prompt_len = len(prompt_token_ids)
29
+ num_reqs = min(
30
+ model_runner.scheduler_config.max_num_seqs,
31
+ model_runner.scheduler_config.max_num_batched_tokens // prompt_len,
32
+ )
33
+
34
+ num_kv_cache_groups = len(model_runner.kv_cache_config.kv_cache_groups)
35
+ req_ids = [f"_warmup_{i}_" for i in range(num_reqs)]
36
+
37
+ # SamplingParams exercising all sampling features.
38
+ if model_runner.is_pooling_model:
39
+ sampling_params = None
40
+ pooling_params = PoolingParams()
41
+ else:
42
+ sampling_params = SamplingParams.for_sampler_warmup()
43
+ pooling_params = None
44
+
45
+ # Step 1: Prefill all requests with 2 prompt tokens each.
46
+ new_reqs = [
47
+ NewRequestData.from_request(
48
+ Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params),
49
+ # Each request uses a distinct block per KV cache group.
50
+ block_ids=tuple([i] for _ in range(num_kv_cache_groups)),
51
+ prefill_token_ids=prompt_token_ids,
52
+ )
53
+ for i in range(num_reqs)
54
+ ]
55
+
56
+ prefill_output = SchedulerOutput.make_empty()
57
+ prefill_output.scheduled_new_reqs = new_reqs
58
+ prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids}
59
+ prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs
60
+ prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
61
+
62
+ # Disable KV connector for warmup run.
63
+ model_runner.kv_connector.set_disabled(True)
64
+ model_runner.execute_model(prefill_output)
65
+
66
+ if not model_runner.is_pooling_model:
67
+ # Warm up sampler and perform a decode step for non-pooling models.
68
+
69
+ grammar_output = None
70
+ if model_runner.is_last_pp_rank:
71
+ # Build a GrammarOutput to exercise the structured output bitmask
72
+ # kernel during the prefill step.
73
+ vocab_size = model_runner.model_config.get_vocab_size()
74
+ bitmask_width = (vocab_size + 31) // 32
75
+ grammar_bitmask = np.full(
76
+ (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32
77
+ )
78
+ grammar_output = GrammarOutput(
79
+ structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask
80
+ )
81
+
82
+ model_runner.sample_tokens(grammar_output)
83
+
84
+ # Step 2: Decode all requests with 1 token each.
85
+ cached_req_data = CachedRequestData.make_empty()
86
+ cached_req_data.req_ids = list(req_ids)
87
+ cached_req_data.new_block_ids = [None] * num_reqs
88
+ cached_req_data.num_computed_tokens = [prompt_len] * num_reqs
89
+ cached_req_data.num_output_tokens = [1] * num_reqs
90
+
91
+ decode_output = SchedulerOutput.make_empty()
92
+ decode_output.scheduled_cached_reqs = cached_req_data
93
+ decode_output.num_scheduled_tokens = {rid: 1 for rid in req_ids}
94
+ decode_output.total_num_scheduled_tokens = num_reqs
95
+ decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
96
+
97
+ model_runner.execute_model(decode_output)
98
+ model_runner.sample_tokens(None)
99
+
100
+ # Clean up - process finish_req_ids.
101
+ cleanup_output = SchedulerOutput.make_empty()
102
+ cleanup_output.finished_req_ids = set(req_ids)
103
+ model_runner.execute_model(cleanup_output)
104
+ model_runner.kv_connector.set_disabled(False)
105
+ torch.cuda.synchronize()
vllm/v1/worker/gpu_worker.py
+10
-6
@@ -61,6 +61,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp
61
from vllm.v1.worker.worker_base import WorkerBase
62
from vllm.v1.worker.workspace import init_workspace_manager
63
64
+from .gpu.warmup import warmup_kernels
65
from .utils import request_memory
66
67
logger = init_logger(__name__)
@@ -558,12 +559,15 @@ class Worker(WorkerBase):
559
560
logger.debug(msg)
561
561
- # Warm up sampler and preallocate memory buffer for logits and other
562
- # sampling related tensors of max possible shape to avoid memory
563
- # fragmentation issue.
564
- # NOTE: This is called after `capture_model` on purpose to prevent
565
- # memory buffers from being cleared by `torch.cuda.empty_cache`.
566
- if get_pp_group().is_last_rank:
562
+ if self.use_v2_model_runner:
563
+ # V2: Run full execute_model + sample_tokens to JIT compile triton kernels.
564
+ warmup_kernels(self.model_runner)
565
+ elif get_pp_group().is_last_rank:
566
+ # V1: Warm up sampler and preallocate memory buffer for logits and other
567
+ # sampling related tensors of max possible shape to avoid memory
568
+ # fragmentation issue.
569
+ # NOTE: This is called after `capture_model` on purpose to prevent
570
+ # memory buffers from being cleared by `torch.cuda.empty_cache`.
571
max_num_reqs = min(
572
self.scheduler_config.max_num_seqs,
573
self.scheduler_config.max_num_batched_tokens,