From cc1d020d01949d11b7ef70dabb0eb196b3f39f53 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Sun, 5 Jul 2026 22:45:29 +0800 Subject: [PATCH] [MRV2] Enable mm prefix bidi attention support on MRV2 (#46942) Signed-off-by: Isotr0py Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Nick Hill --- .buildkite/test_areas/models_multimodal.yaml | 3 +- .../generation/test_mm_prefix_lm.py | 119 ++++++++++++++++++ vllm/config/model.py | 2 +- vllm/v1/worker/gpu/attn_utils.py | 27 ++++ vllm/v1/worker/gpu/model_states/default.py | 17 ++- 5 files changed, 165 insertions(+), 3 deletions(-) create mode 100644 tests/models/multimodal/generation/test_mm_prefix_lm.py diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 9ecfe01d400..a721efa6067 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.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 diff --git a/tests/models/multimodal/generation/test_mm_prefix_lm.py b/tests/models/multimodal/generation/test_mm_prefix_lm.py new file mode 100644 index 00000000000..8d3f5b77b71 --- /dev/null +++ b/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 = ( + "user\n" + "What is the content in the center of the image?" + "\nmodel\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=}" + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index e19c9408914..2ec754f9346 100644 --- a/vllm/config/model.py +++ b/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 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 5fcc9053bf4..aecce3c575c 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index e5e89da2b2e..854b71b69fc 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/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, )