[Bugfix][MM] Fix JinaVL processing cache order (#53553)
Signed-off-by: JiataiWang <wangjiatai@proton.me> Co-authored-by: OpenAI Codex <codex@openai.com>
Jiatai Wang committed
Aug 25, 2026 at 00:03 UTC
06ecec7a8424106dc80c5c40bb0cc22bcc6da667
3 files changed
+88
-29
tests/models/multimodal/processing/test_common.py
-8
@@ -390,14 +390,6 @@ def test_processing_correctness(
390
"OpenVLA uses a custom vLLM processor because its HF remote "
391
"processor is incompatible with current Transformers."
392
)
393
- if model_id == "jinaai/jina-reranker-m0":
394
- pytest.skip(
395
- "JinaVL reverses the multi-modal item order inside "
396
- "_apply_hf_processor_main to match its document-before-query "
397
- "score template, which breaks the positional correspondence "
398
- "between multi-modal kwargs and hashes assumed by the "
399
- "processing cache. Needs a fix that keeps both consistent."
400
- )
393
if model_id == "mistralai/Voxtral-Mini-4B-Realtime-2602":
394
pytest.skip(
395
"Voxtral Realtime doesn't make use of any place-holder "
tests/models/multimodal/processing/test_qwen2_vl.py
+51
@@ -4,15 +4,66 @@
4
import pytest
5
import torch
6
from packaging.version import Version
7
+from PIL import Image
8
from transformers import __version__ as TRANSFORMERS_VERSION
9
10
from vllm.model_executor.models.vision import FusedInputNorm
11
from vllm.multimodal import MULTIMODAL_REGISTRY
12
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
13
+from vllm.multimodal.inputs import batched_tensors_equal
14
15
from ....conftest import ImageTestAssets
16
from ...utils import build_model_context
17
18
19
+def test_jina_vl_processing_order() -> None:
20
+ """Jina's document-first prompt keeps cached features and hashes aligned."""
21
+ ctx = build_model_context(
22
+ "jinaai/jina-reranker-m0",
23
+ runner="pooling",
24
+ limit_mm_per_prompt={"image": 2},
25
+ mm_processor_cache_gb=1,
26
+ )
27
+ cache = MultiModalProcessorOnlyCache(ctx.model_config)
28
+ processor = MULTIMODAL_REGISTRY.create_processor(
29
+ ctx.model_config,
30
+ tokenizer=ctx.tokenizer,
31
+ cache=cache,
32
+ )
33
+
34
+ placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
35
+ query_image = Image.new("RGB", (128, 160), color=(255, 0, 0))
36
+ document_image = Image.new("RGB", (192, 128), color=(0, 255, 0))
37
+
38
+ def process(images: list[Image.Image]):
39
+ return processor(
40
+ placeholder * len(images),
41
+ mm_items=processor.info.parse_mm_data({"image": images}),
42
+ )
43
+
44
+ query = process([query_image])
45
+ document = process([document_image])
46
+ pair = process([query_image, document_image])
47
+
48
+ pair_items = pair["mm_kwargs"]["image"]
49
+ assert pair["mm_hashes"]["image"] == [
50
+ document["mm_hashes"]["image"][0],
51
+ query["mm_hashes"]["image"][0],
52
+ ]
53
+ assert batched_tensors_equal(
54
+ pair_items[0].get_data(),
55
+ document["mm_kwargs"]["image"][0].get_data(),
56
+ )
57
+ assert batched_tensors_equal(
58
+ pair_items[1].get_data(),
59
+ query["mm_kwargs"]["image"][0].get_data(),
60
+ )
61
+ assert [item.length for item in pair["mm_placeholders"]["image"]] == [
62
+ document["mm_placeholders"]["image"][0].length,
63
+ query["mm_placeholders"]["image"][0].length,
64
+ ]
65
+
66
+
67
@pytest.mark.parametrize(
68
("image_mean", "image_std", "rescale_factor", "is_identity"),
69
[
vllm/model_executor/models/jina_vl.py
+37
-21
@@ -1,10 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterable, Mapping
3
+from collections.abc import Iterable
4
5
import torch
6
import torch.nn as nn
7
-from transformers import BatchFeature
7
8
from vllm.config import ModelConfig, VllmConfig
9
from vllm.inputs import TokensPrompt
@@ -12,7 +11,12 @@ from vllm.logger import init_logger
11
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
12
from vllm.model_executor.layers.pooler import DispatchPooler
13
from vllm.multimodal import MULTIMODAL_REGISTRY
15
-from vllm.multimodal.parse import MultiModalDataItems
14
+from vllm.multimodal.inputs import MultiModalKwargsItems
15
+from vllm.multimodal.processing.processor import (
16
+ MultiModalProcessingInfo,
17
+ ProcessorInputs,
18
+ TimingContext,
19
+)
20
from vllm.sequence import IntermediateTensors
21
22
from .interfaces import SupportsCrossEncoding, SupportsMultiModal, SupportsScoreTemplate
@@ -55,28 +59,40 @@ class JinaVLScorer(nn.Module):
59
60
61
class JinaVLMultiModalProcessor(Qwen2VLMultiModalProcessor):
58
- def _apply_hf_processor_main(
62
+ def _cached_apply_hf_processor(
63
self,
60
- mm_items: MultiModalDataItems,
61
- hf_processor_mm_kwargs: Mapping[str, object],
62
- ) -> BatchFeature:
63
- valid_mm_items = mm_items.select(
64
- {k for k, c in mm_items.get_all_counts().items() if c > 0}
64
+ inputs: ProcessorInputs,
65
+ timing_ctx: TimingContext,
66
+ ) -> MultiModalProcessingInfo:
67
+ mm_info = super()._cached_apply_hf_processor(inputs, timing_ctx)
68
+
69
+ # Score inputs are query-first, while the prompt template is document-first.
70
+ mm_kwargs = MultiModalKwargsItems(
71
+ {
72
+ modality: list(reversed(items))
73
+ for modality, items in mm_info.kwargs.items()
74
+ }
75
)
66
- mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
67
-
68
- if not mm_data:
69
- return BatchFeature(dict(passthrough_data))
76
+ mm_hashes = {
77
+ modality: list(reversed(hashes))
78
+ for modality, hashes in mm_info.hashes.items()
79
+ }
80
+ mm_prompt_updates = {
81
+ modality: [
82
+ [
83
+ self._recompute_cached_prompt_update(update, item_idx)
84
+ for update in updates
85
+ ]
86
+ for item_idx, updates in enumerate(reversed(item_updates))
87
+ ]
88
+ for modality, item_updates in mm_info.prompt_updates.items()
89
+ }
90
71
- for _, value in mm_data.items():
72
- value.reverse()
73
- processed_data = self.info.ctx.call_hf_processor(
74
- self.info.get_hf_processor(**hf_processor_mm_kwargs),
75
- mm_data,
76
- hf_processor_mm_kwargs,
91
+ return MultiModalProcessingInfo(
92
+ kwargs=mm_kwargs,
93
+ hashes=mm_hashes,
94
+ prompt_updates=mm_prompt_updates,
95
)
78
- processed_data.update(passthrough_data)
79
- return processed_data
96
97
98
@MULTIMODAL_REGISTRY.register_processor(