[Bugfix] Keep grid dims for XD-RoPE models on a prefix-cache hit (#53456)
Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com>
drakosha committed
Aug 25, 2026 at 04:14 UTC
9bf0d4717ecd5ad9e6c5fe3eb933cf654a256321
5 files changed
+32
-4
tests/v1/core/test_output.py
+18
@@ -105,6 +105,24 @@ def _mm_feature_mixed(offset: int, length: int) -> MultiModalFeatureSpec:
105
)
106
107
108
+def test_strip_covered_mm_data_xdrope() -> None:
109
+ """XD-RoPE models (e.g. HunyuanOCR) compute positions the same way, so a
110
+ covered item must keep its grid dims: stripping them made the worker index
111
+ an empty ``image_grid_thw`` and crash the engine on the second request with
112
+ an identical prompt."""
113
+ covered = _mm_feature_mixed(offset=0, length=100)
114
+ uncovered = _mm_feature_mixed(offset=300, length=100)
115
+
116
+ stripped = strip_covered_mm_data(
117
+ [covered, uncovered], num_computed_tokens=250, uses_xdrope=True
118
+ )
119
+
120
+ assert stripped[0].data is not None
121
+ assert list(stripped[0].data.keys()) == ["image_grid_thw"]
122
+ assert stripped[1].data is not None
123
+ assert set(stripped[1].data.keys()) == {"pixel_values", "image_grid_thw"}
124
+
125
+
126
def test_strip_covered_mm_data_mrope() -> None:
127
"""For M-RoPE models, covered items keep their keep_on_cpu metadata fields
128
(the worker needs them to compute positions); payload fields are dropped."""
vllm/config/model.py
+4
@@ -1843,6 +1843,10 @@ class ModelConfig:
1843
def uses_xdrope_dim(self) -> int:
1844
return uses_xdrope_dim(self.hf_config)
1845
1846
+ @property
1847
+ def uses_xdrope(self) -> bool:
1848
+ return self.uses_xdrope_dim > 0
1849
+
1850
@property
1851
def is_multimodal_model(self) -> bool:
1852
return self.multimodal_config is not None
vllm/multimodal/utils.py
+5
-4
@@ -216,12 +216,13 @@ def strip_covered_mm_data(
216
mm_features: list[MultiModalFeatureSpec],
217
num_computed_tokens: int,
218
uses_mrope: bool = False,
219
+ uses_xdrope: 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,
222
- so the workers never consume the payload fields. M-RoPE models are the
223
- exception: the worker computes positions for the whole prompt from the
224
- CPU-side metadata fields (e.g. grid dims), so those are kept. The
223
+ so the workers never consume the payload fields. M-RoPE and XD-RoPE models
224
+ are the exception: the worker computes positions for the whole prompt from
225
+ the 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
@@ -233,7 +234,7 @@ def strip_covered_mm_data(
234
return f
235
236
data = None
236
- if uses_mrope:
237
+ if uses_mrope or uses_xdrope:
238
data = MultiModalKwargsItem(
239
{k: elem for k, elem in f.data.items() if elem.field.keep_on_cpu}
240
)
vllm/v1/core/sched/output.py
+2
@@ -55,6 +55,7 @@ class NewRequestData:
55
block_ids: tuple[list[int], ...],
56
prefill_token_ids: list[int] | None = None,
57
uses_mrope: bool = False,
58
+ uses_xdrope: bool = False,
59
) -> "NewRequestData":
60
return cls(
61
req_id=request.request_id,
@@ -63,6 +64,7 @@ class NewRequestData:
64
request.mm_features,
65
request.num_computed_tokens,
66
uses_mrope=uses_mrope,
67
+ uses_xdrope=uses_xdrope,
68
),
69
sampling_params=request.sampling_params,
70
pooling_params=request.pooling_params,
vllm/v1/core/sched/scheduler.py
+3
@@ -87,6 +87,7 @@ class Scheduler(SchedulerInterface):
87
self.cache_config = vllm_config.cache_config
88
self.lora_config = vllm_config.lora_config
89
self.model_uses_mrope = vllm_config.model_config.uses_mrope
90
+ self.model_uses_xdrope = vllm_config.model_config.uses_xdrope
91
self.kv_cache_config = kv_cache_config
92
self.kv_events_config = vllm_config.kv_events_config
93
self.parallel_config = vllm_config.parallel_config
@@ -1223,6 +1224,7 @@ class Scheduler(SchedulerInterface):
1224
req_to_new_blocks[req.request_id].get_block_ids(),
1225
req._all_token_ids,
1226
uses_mrope=self.model_uses_mrope,
1227
+ uses_xdrope=self.model_uses_xdrope,
1228
)
1229
for req in scheduled_new_reqs
1230
]
@@ -1232,6 +1234,7 @@ class Scheduler(SchedulerInterface):
1234
req,
1235
req_to_new_blocks[req.request_id].get_block_ids(),
1236
uses_mrope=self.model_uses_mrope,
1237
+ uses_xdrope=self.model_uses_xdrope,
1238
)
1239
for req in scheduled_new_reqs
1240
]