[ModelRunner V2] Warmup cross-attn properly in encoder-decoder case (#47308)
Nick Hill committed
Jul 1, 2026 at 20:36 UTC
8cfeb84dba41a0c56570334757d921abd71e5288
1 file changed
+27
-2
vllm/v1/worker/gpu/warmup.py
+27
-2
@@ -10,6 +10,7 @@ import torch
10
11
from vllm import PoolingParams, SamplingParams
12
from vllm.logger import init_logger
13
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
14
from vllm.utils.math_utils import cdiv
15
from vllm.v1.core.sched.output import (
16
CachedRequestData,
@@ -17,7 +18,7 @@ from vllm.v1.core.sched.output import (
18
NewRequestData,
19
SchedulerOutput,
20
)
20
-from vllm.v1.kv_cache_interface import MambaSpec
21
+from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
22
from vllm.v1.request import Request
23
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
24
@@ -177,8 +178,26 @@ def warmup_kernels(
178
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
179
num_kv_cache_groups = len(kv_cache_groups)
180
181
+ # Encoder-decoder models: give each warmup request a dummy encoder input so
182
+ # cross-attention warms up over a realistic, non-empty key sequence.
183
+ # The dummy mm_feature is registered in the encoder cache and only its encoder
184
+ # length is read (not the inputs themselves); the encoder itself is not scheduled.
185
+ max_encoder_len = getattr(model_runner.model_state, "max_encoder_len", 0)
186
+ warmup_mm_features: list[MultiModalFeatureSpec] = []
187
+ if model_runner.is_encoder_decoder and max_encoder_len:
188
+ warmup_mm_features = [
189
+ MultiModalFeatureSpec(
190
+ data=None,
191
+ modality="",
192
+ identifier="_warmup_encoder",
193
+ mm_position=PlaceholderRange(offset=0, length=max_encoder_len),
194
+ )
195
+ ]
196
+
197
# Compute per-request block counts for each KV cache group.
198
def _warmup_block_count(num_tokens: int, spec: Any) -> int:
199
+ if isinstance(spec, CrossAttentionSpec):
200
+ num_tokens = max_encoder_len
201
num_blocks = cdiv(num_tokens, spec.block_size)
202
if isinstance(spec, MambaSpec) and spec.mamba_cache_mode == "align":
203
# Align mode reserves extra blocks beyond the token range for the
@@ -222,7 +241,13 @@ def warmup_kernels(
241
# Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each.
242
new_reqs = [
243
NewRequestData.from_request(
225
- Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params),
244
+ Request(
245
+ req_ids[i],
246
+ prompt_token_ids,
247
+ sampling_params,
248
+ pooling_params,
249
+ mm_features=warmup_mm_features,
250
+ ),
251
block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
252
prefill_token_ids=prompt_token_ids,
253
)