[Core] Skip broadcasting mm tensor data to workers for prefix-cache-covered items (#52041)
Signed-off-by: Xiaoan (Sean) Liu <shawnliu0327@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Xiaoan (Sean) Liu committed
Aug 19, 2026 at 07:16 UTC
db92053e97b5630a6a36118693b1dffe9b03be36
4 files changed
+130
-2
tests/v1/core/test_output.py
+87
@@ -2,6 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
+from vllm.multimodal.inputs import (
6
+ MultiModalBatchedField,
7
+ MultiModalFeatureSpec,
8
+ MultiModalFieldElem,
9
+ MultiModalKwargsItem,
10
+ PlaceholderRange,
11
+)
12
+from vllm.multimodal.utils import strip_covered_mm_data
13
from vllm.v1.core.sched.output import NewRequestData
14
15
@@ -34,3 +42,82 @@ def test_repr_with_multi_element_tensor() -> None:
42
43
assert "prompt_embeds_shape=torch.Size([10, 768])" in repr(new_requests_data)
44
assert "prompt_embeds_shape=torch.Size([10, 768])" in new_requests_data.anon_repr()
45
+
46
+
47
+def _mm_feature(offset: int, length: int) -> MultiModalFeatureSpec:
48
+ return MultiModalFeatureSpec(
49
+ data=MultiModalKwargsItem.dummy(),
50
+ mm_position=PlaceholderRange(offset=offset, length=length),
51
+ identifier=f"hash_{offset}",
52
+ modality="image",
53
+ )
54
+
55
+
56
+def test_strip_covered_mm_data() -> None:
57
+ """Items fully inside the computed prefix lose their data; items touching
58
+ the uncomputed region keep it; already-None data stays None."""
59
+ from dataclasses import replace
60
+
61
+ covered = _mm_feature(offset=0, length=100)
62
+ boundary = _mm_feature(offset=150, length=100) # ends exactly at 250
63
+ uncovered = _mm_feature(offset=300, length=100)
64
+ already_none = replace(_mm_feature(offset=100, length=50), data=None)
65
+
66
+ stripped = strip_covered_mm_data(
67
+ [covered, boundary, uncovered, already_none], num_computed_tokens=250
68
+ )
69
+
70
+ assert stripped[0].data is None # fully covered -> stripped
71
+ assert stripped[1].data is None # span end == computed -> covered -> stripped
72
+ assert stripped[2].data is not None # extends past prefix -> kept
73
+ assert stripped[3].data is None # was already None
74
+ # non-data fields are preserved
75
+ assert stripped[0].identifier == covered.identifier
76
+ assert stripped[0].mm_position == covered.mm_position
77
+ # original list is not mutated
78
+ assert covered.data is not None
79
+
80
+
81
+def test_strip_covered_mm_data_zero_computed() -> None:
82
+ """With no prefix hit nothing is stripped."""
83
+ features = [_mm_feature(offset=0, length=100)]
84
+ stripped = strip_covered_mm_data(features, num_computed_tokens=0)
85
+ assert stripped[0].data is not None
86
+
87
+
88
+def _mm_feature_mixed(offset: int, length: int) -> MultiModalFeatureSpec:
89
+ data = MultiModalKwargsItem(
90
+ {
91
+ "pixel_values": MultiModalFieldElem(
92
+ data=torch.empty(4), field=MultiModalBatchedField()
93
+ ),
94
+ "image_grid_thw": MultiModalFieldElem(
95
+ data=torch.ones(1, 3, dtype=torch.long),
96
+ field=MultiModalBatchedField(keep_on_cpu=True),
97
+ ),
98
+ }
99
+ )
100
+ return MultiModalFeatureSpec(
101
+ data=data,
102
+ mm_position=PlaceholderRange(offset=offset, length=length),
103
+ identifier=f"hash_{offset}",
104
+ modality="image",
105
+ )
106
+
107
+
108
+def test_strip_covered_mm_data_mrope() -> None:
109
+ """For M-RoPE models, covered items keep their keep_on_cpu metadata fields
110
+ (the worker needs them to compute positions); payload fields are dropped."""
111
+ covered = _mm_feature_mixed(offset=0, length=100)
112
+ uncovered = _mm_feature_mixed(offset=300, length=100)
113
+
114
+ stripped = strip_covered_mm_data(
115
+ [covered, uncovered], num_computed_tokens=250, uses_mrope=True
116
+ )
117
+
118
+ assert stripped[0].data is not None
119
+ assert list(stripped[0].data.keys()) == ["image_grid_thw"]
120
+ assert stripped[1].data is not None
121
+ assert set(stripped[1].data.keys()) == {"pixel_values", "image_grid_thw"}
122
+ # original list is not mutated
123
+ assert set(covered.data.keys()) == {"pixel_values", "image_grid_thw"}
vllm/multimodal/utils.py
+31
@@ -5,6 +5,7 @@ import bisect
5
import mimetypes
6
from collections import defaultdict
7
from collections.abc import Generator, Sequence
8
+from dataclasses import replace
9
from itertools import groupby
10
from typing import TYPE_CHECKING, Any
11
@@ -212,6 +213,36 @@ def _batch_mm_items(
213
}
214
215
216
+def strip_covered_mm_data(
217
+ mm_features: list[MultiModalFeatureSpec],
218
+ num_computed_tokens: int,
219
+ uses_mrope: bool = False,
220
+) -> list[MultiModalFeatureSpec]:
221
+ """Drop the tensor data of mm items whose placeholder span is fully inside
222
+ a prefix-cache-covered region: no encoder run can be scheduled for them,
223
+ so the workers never consume the payload fields. M-RoPE models are the
224
+ exception: the worker computes positions for the whole prompt from the
225
+ CPU-side metadata fields (e.g. grid dims), so those are kept. The
226
+ scheduler-side ``Request`` keeps the full features."""
227
+ if not mm_features or num_computed_tokens == 0:
228
+ return mm_features
229
+
230
+ def maybe_strip(f: MultiModalFeatureSpec) -> MultiModalFeatureSpec:
231
+ if f.data is None or (
232
+ f.mm_position.offset + f.mm_position.length > num_computed_tokens
233
+ ):
234
+ return f
235
+
236
+ data = None
237
+ if uses_mrope:
238
+ data = MultiModalKwargsItem(
239
+ {k: elem for k, elem in f.data.items() if elem.field.keep_on_cpu}
240
+ )
241
+ return replace(f, data=data)
242
+
243
+ return [maybe_strip(f) for f in mm_features]
244
+
245
+
246
def group_and_batch_mm_items(
247
items: Sequence[MultiModalKwargsItem],
248
*,
vllm/v1/core/sched/output.py
+7
-1
@@ -6,6 +6,7 @@ from functools import cached_property
6
from typing import TYPE_CHECKING
7
8
from vllm.config.ec_manager_config import EncoderCacheManagerMetadata
9
+from vllm.multimodal.utils import strip_covered_mm_data
10
11
if TYPE_CHECKING:
12
import numpy as np
@@ -53,11 +54,16 @@ class NewRequestData:
54
request: Request,
55
block_ids: tuple[list[int], ...],
56
prefill_token_ids: list[int] | None = None,
57
+ uses_mrope: bool = False,
58
) -> "NewRequestData":
59
return cls(
60
req_id=request.request_id,
61
prompt_token_ids=request.prompt_token_ids,
60
- mm_features=request.mm_features,
62
+ mm_features=strip_covered_mm_data(
63
+ request.mm_features,
64
+ request.num_computed_tokens,
65
+ uses_mrope=uses_mrope,
66
+ ),
67
sampling_params=request.sampling_params,
68
pooling_params=request.pooling_params,
69
block_ids=block_ids,
vllm/v1/core/sched/scheduler.py
+5
-1
@@ -82,6 +82,7 @@ class Scheduler(SchedulerInterface):
82
self.scheduler_config = vllm_config.scheduler_config
83
self.cache_config = vllm_config.cache_config
84
self.lora_config = vllm_config.lora_config
85
+ self.model_uses_mrope = vllm_config.model_config.uses_mrope
86
self.kv_cache_config = kv_cache_config
87
self.kv_events_config = vllm_config.kv_events_config
88
self.parallel_config = vllm_config.parallel_config
@@ -1199,13 +1200,16 @@ class Scheduler(SchedulerInterface):
1200
req,
1201
req_to_new_blocks[req.request_id].get_block_ids(),
1202
req._all_token_ids,
1203
+ uses_mrope=self.model_uses_mrope,
1204
)
1205
for req in scheduled_new_reqs
1206
]
1207
else:
1208
new_reqs_data = [
1209
NewRequestData.from_request(
1208
- req, req_to_new_blocks[req.request_id].get_block_ids()
1210
+ req,
1211
+ req_to_new_blocks[req.request_id].get_block_ids(),
1212
+ uses_mrope=self.model_uses_mrope,
1213
)
1214
for req in scheduled_new_reqs
1215
]