[Bugfix][MM] Fix JinaVL processing cache order (#53553)

Signed-off-by: JiataiWang <wangjiatai@proton.me>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/models/multimodal/processing/test_common.py

@@ -390,14 +390,6 @@ def test_processing_correctness(
             "OpenVLA uses a custom vLLM processor because its HF remote "
             "processor is incompatible with current Transformers."
         )
-    if model_id == "jinaai/jina-reranker-m0":
-        pytest.skip(
-            "JinaVL reverses the multi-modal item order inside "
-            "_apply_hf_processor_main to match its document-before-query "
-            "score template, which breaks the positional correspondence "
-            "between multi-modal kwargs and hashes assumed by the "
-            "processing cache. Needs a fix that keeps both consistent."
-        )
     if model_id == "mistralai/Voxtral-Mini-4B-Realtime-2602":
         pytest.skip(
             "Voxtral Realtime doesn't make use of any place-holder "

tests/models/multimodal/processing/test_qwen2_vl.py

@@ -4,15 +4,66 @@
 import pytest
 import torch
 from packaging.version import Version
+from PIL import Image
 from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.model_executor.models.vision import FusedInputNorm
 from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
+from vllm.multimodal.inputs import batched_tensors_equal
 
 from ....conftest import ImageTestAssets
 from ...utils import build_model_context
 
 
+def test_jina_vl_processing_order() -> None:
+    """Jina's document-first prompt keeps cached features and hashes aligned."""
+    ctx = build_model_context(
+        "jinaai/jina-reranker-m0",
+        runner="pooling",
+        limit_mm_per_prompt={"image": 2},
+        mm_processor_cache_gb=1,
+    )
+    cache = MultiModalProcessorOnlyCache(ctx.model_config)
+    processor = MULTIMODAL_REGISTRY.create_processor(
+        ctx.model_config,
+        tokenizer=ctx.tokenizer,
+        cache=cache,
+    )
+
+    placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    query_image = Image.new("RGB", (128, 160), color=(255, 0, 0))
+    document_image = Image.new("RGB", (192, 128), color=(0, 255, 0))
+
+    def process(images: list[Image.Image]):
+        return processor(
+            placeholder * len(images),
+            mm_items=processor.info.parse_mm_data({"image": images}),
+        )
+
+    query = process([query_image])
+    document = process([document_image])
+    pair = process([query_image, document_image])
+
+    pair_items = pair["mm_kwargs"]["image"]
+    assert pair["mm_hashes"]["image"] == [
+        document["mm_hashes"]["image"][0],
+        query["mm_hashes"]["image"][0],
+    ]
+    assert batched_tensors_equal(
+        pair_items[0].get_data(),
+        document["mm_kwargs"]["image"][0].get_data(),
+    )
+    assert batched_tensors_equal(
+        pair_items[1].get_data(),
+        query["mm_kwargs"]["image"][0].get_data(),
+    )
+    assert [item.length for item in pair["mm_placeholders"]["image"]] == [
+        document["mm_placeholders"]["image"][0].length,
+        query["mm_placeholders"]["image"][0].length,
+    ]
+
+
 @pytest.mark.parametrize(
     ("image_mean", "image_std", "rescale_factor", "is_identity"),
     [

vllm/model_executor/models/jina_vl.py

@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature
 
 from vllm.config import ModelConfig, VllmConfig
 from vllm.inputs import TokensPrompt
@@ -12,7 +11,12 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.pooler import DispatchPooler
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.inputs import MultiModalKwargsItems
+from vllm.multimodal.processing.processor import (
+    MultiModalProcessingInfo,
+    ProcessorInputs,
+    TimingContext,
+)
 from vllm.sequence import IntermediateTensors
 
 from .interfaces import SupportsCrossEncoding, SupportsMultiModal, SupportsScoreTemplate
@@ -55,28 +59,40 @@ class JinaVLScorer(nn.Module):
 
 
 class JinaVLMultiModalProcessor(Qwen2VLMultiModalProcessor):
-    def _apply_hf_processor_main(
+    def _cached_apply_hf_processor(
         self,
-        mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-    ) -> BatchFeature:
-        valid_mm_items = mm_items.select(
-            {k for k, c in mm_items.get_all_counts().items() if c > 0}
-        )
-        mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
+        inputs: ProcessorInputs,
+        timing_ctx: TimingContext,
+    ) -> MultiModalProcessingInfo:
+        mm_info = super()._cached_apply_hf_processor(inputs, timing_ctx)
 
-        if not mm_data:
-            return BatchFeature(dict(passthrough_data))
+        # Score inputs are query-first, while the prompt template is document-first.
+        mm_kwargs = MultiModalKwargsItems(
+            {
+                modality: list(reversed(items))
+                for modality, items in mm_info.kwargs.items()
+            }
+        )
+        mm_hashes = {
+            modality: list(reversed(hashes))
+            for modality, hashes in mm_info.hashes.items()
+        }
+        mm_prompt_updates = {
+            modality: [
+                [
+                    self._recompute_cached_prompt_update(update, item_idx)
+                    for update in updates
+                ]
+                for item_idx, updates in enumerate(reversed(item_updates))
+            ]
+            for modality, item_updates in mm_info.prompt_updates.items()
+        }
 
-        for _, value in mm_data.items():
-            value.reverse()
-        processed_data = self.info.ctx.call_hf_processor(
-            self.info.get_hf_processor(**hf_processor_mm_kwargs),
-            mm_data,
-            hf_processor_mm_kwargs,
+        return MultiModalProcessingInfo(
+            kwargs=mm_kwargs,
+            hashes=mm_hashes,
+            prompt_updates=mm_prompt_updates,
         )
-        processed_data.update(passthrough_data)
-        return processed_data
 
 
 @MULTIMODAL_REGISTRY.register_processor(