[MRV2] Enable mm prefix bidi attention support on MRV2 (#46942)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <2037008807@qq.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

.buildkite/test_areas/models_multimodal.yaml

@@ -27,6 +27,7 @@ steps:
   - tests/models/multimodal
   commands:
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
+    - pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
   mirror:
     amd:
@@ -58,7 +59,7 @@ steps:
   - vllm/
   - tests/models/multimodal
   commands:
-    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
+    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
     - pytest models/multimodal/generation/test_memory_leak.py -m core_model
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work

tests/models/multimodal/generation/test_mm_prefix_lm.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import pytest
+import torch
+from transformers import AutoModelForImageTextToText
+
+from vllm.platforms import current_platform
+
+from ....conftest import HfRunner, ImageTestAssets, VllmRunner
+from .vlm_utils import model_utils
+
+MODEL = "google/gemma-3-4b-it"
+PROMPT = (
+    "<bos><start_of_turn>user\n"
+    "<start_of_image>What is the content in the center of the image?"
+    "<end_of_turn>\n<start_of_turn>model\n"
+)
+
+
+def _install_prefill_hidden_capture(model):
+    model = getattr(model, "module", model)
+    model._prefill_hidden = None
+
+    language_model = model.language_model.model
+    original_forward = language_model.forward
+
+    def forward(*args, **kwargs):
+        hidden_states = original_forward(*args, **kwargs)
+        if model._prefill_hidden is None and torch.is_tensor(hidden_states):
+            model._prefill_hidden = hidden_states.detach().float().cpu()
+        return hidden_states
+
+    language_model.forward = forward
+
+
+def _get_prefill_hidden(model):
+    model = getattr(model, "module", model)
+    hidden = getattr(model, "_prefill_hidden", None)
+    assert hidden is not None
+    return hidden
+
+
+def _get_hf_prefill_hidden(hf_model: HfRunner, image: Any):
+    inputs = hf_model.get_inputs([PROMPT], images=[image])[0]
+    with torch.no_grad():
+        outputs = hf_model.model.model(
+            **hf_model.wrap_device(inputs),
+            use_cache=False,
+        )
+    return outputs.last_hidden_state[0].detach().float().cpu()
+
+
+def _get_vllm_prefill_hidden(
+    vllm_runner: type[VllmRunner],
+    image: Any,
+    vllm_runner_kwargs: dict[str, Any],
+):
+    with vllm_runner(
+        MODEL,
+        max_model_len=4096,
+        max_num_seqs=2,
+        enforce_eager=True,
+        limit_mm_per_prompt={"image": 1},
+        **vllm_runner_kwargs,
+    ) as vllm_model:
+        vllm_model.apply_model(_install_prefill_hidden_capture)
+        vllm_model.generate_greedy([PROMPT], max_tokens=1, images=[image])
+        return vllm_model.apply_model(_get_prefill_hidden)[0]
+
+
+@pytest.mark.core_model
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="ROCm attention has accuracy issue for this test"
+)
+def test_mm_prefix_lm_e2e(
+    hf_runner: type[HfRunner],
+    vllm_runner: type[VllmRunner],
+    image_assets: ImageTestAssets,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Regression: Gemma3 native prefill must apply image prefix-LM mask."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    image = image_assets[0].pil_image
+
+    vllm_runner_kwargs: dict[str, Any] = {
+        "mm_processor_cache_gb": 0,
+        "mm_processor_kwargs": {"do_pan_and_scan": True},
+    }
+    vllm_hidden = _get_vllm_prefill_hidden(vllm_runner, image, vllm_runner_kwargs)
+
+    hf_model = hf_runner(
+        MODEL,
+        auto_cls=AutoModelForImageTextToText,
+    )
+    hf_model = model_utils.gemma3_patch_hf_runner(hf_model)
+
+    with hf_model:
+        hf_hidden = _get_hf_prefill_hidden(hf_model, image)
+
+    assert vllm_hidden.shape == hf_hidden.shape
+
+    full_cos = torch.nn.functional.cosine_similarity(
+        vllm_hidden.flatten(), hf_hidden.flatten(), dim=0
+    )
+    image_cos = torch.nn.functional.cosine_similarity(
+        vllm_hidden[1:769].flatten(), hf_hidden[1:769].flatten(), dim=0
+    )
+
+    assert full_cos > 0.9, (
+        "Gemma3 mm-prefix-LM full prefill hidden states should be close to HF; "
+        f"got {full_cos=}"
+    )
+    assert image_cos > 0.9, (
+        "Gemma3 mm-prefix-LM image prefill hidden states should be close to HF; "
+        f"got {image_cos=}"
+    )

vllm/config/model.py

@@ -1252,7 +1252,7 @@ class ModelConfig:
     def is_deepseek_mla(self) -> bool:
         return self.model_arch_config.is_deepseek_mla
 
-    @property
+    @cached_property
     def is_mm_prefix_lm(self) -> bool:
         return self.model_arch_config.is_mm_prefix_lm
 

vllm/v1/worker/gpu/attn_utils.py

@@ -14,6 +14,7 @@ from vllm.config import (
 )
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
@@ -487,6 +488,7 @@ def build_attn_metadata(
     seq_lens_cpu_upper_bound: torch.Tensor | None = None,
     dcp_local_seq_lens: torch.Tensor | None = None,
     positions: torch.Tensor | None = None,
+    mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
     model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     for_cudagraph_capture: bool = False,
     causal: bool = True,
@@ -523,6 +525,7 @@ def build_attn_metadata(
             causal=causal,
             dcp_local_seq_lens=dcp_local_seq_lens,
             positions=positions,
+            mm_req_doc_ranges=mm_req_doc_ranges,
             rswa_prefix_lens=rswa_prefix_lens,
             **common_attn_metadata_extra_kwargs,
         )
@@ -550,3 +553,27 @@ def build_attn_metadata(
             for layer_name in attn_group.layer_names:
                 attn_metadata[layer_name] = metadata
     return attn_metadata
+
+
+def compute_mm_prefix_ranges(
+    req_ids: list[str],
+    mm_features: dict[str, list[MultiModalFeatureSpec]],
+    sliding_window: int | None = None,
+) -> dict[int, list[tuple[int, int]]]:
+    """Compute PrefixLM bidirectional ranges for multimodal tokens.
+
+    Ranges exceeding sliding_window are skipped to prevent early tokens
+    from attending across the entire image span.
+    """
+    req_doc_ranges: dict[int, list[tuple[int, int]]] = {}
+    for req_idx, req_id in enumerate(req_ids):
+        image_doc_ranges = []
+        for mm_feature in mm_features.get(req_id, ()):
+            if mm_feature.modality not in ("image", "video"):
+                continue
+            for r in mm_feature.mm_position.extract_embeds_range():
+                if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window:
+                    continue
+                image_doc_ranges.append(r)
+        req_doc_ranges[req_idx] = image_doc_ranges
+    return req_doc_ranges

vllm/v1/worker/gpu/model_states/default.py

@@ -9,7 +9,10 @@ from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import KVCacheConfig
-from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
+from vllm.v1.worker.gpu.attn_utils import (
+    build_attn_metadata,
+    compute_mm_prefix_ranges,
+)
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.rope import get_rope_state
@@ -152,6 +155,17 @@ class DefaultModelState(ModelState):
             max_seq_len = self.max_model_len
         else:
             max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()
+        req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
+        if (
+            self.supports_mm_inputs
+            and self.encoder_cache is not None
+            and self.model_config.is_mm_prefix_lm
+        ):
+            req_doc_ranges = compute_mm_prefix_ranges(
+                req_ids=input_batch.req_ids,
+                mm_features=self.encoder_cache.mm_features,
+                sliding_window=self.model_config.get_sliding_window(),
+            )
         attn_metadata = build_attn_metadata(
             attn_groups=attn_groups,
             num_reqs=num_reqs,
@@ -167,6 +181,7 @@ class DefaultModelState(ModelState):
             seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
             dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
             positions=input_batch.positions,
+            mm_req_doc_ranges=req_doc_ranges,
             for_cudagraph_capture=for_capture,
             rswa_prefix_lens=input_batch.prompt_lens,
         )