[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>
Isotr0py committed
Jul 5, 2026 at 22:45 UTC
cc1d020d01949d11b7ef70dabb0eb196b3f39f53
5 files changed
+165
-3
.buildkite/test_areas/models_multimodal.yaml
+2
-1
@@ -27,6 +27,7 @@ steps:
27
- tests/models/multimodal
28
commands:
29
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
30
+ - pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model
31
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
32
mirror:
33
amd:
@@ -58,7 +59,7 @@ steps:
59
- vllm/
60
- tests/models/multimodal
61
commands:
61
- - 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
62
+ - 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
63
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
64
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
65
- 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
new
+119
@@ -0,0 +1,119 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Any
5
+
6
+import pytest
7
+import torch
8
+from transformers import AutoModelForImageTextToText
9
+
10
+from vllm.platforms import current_platform
11
+
12
+from ....conftest import HfRunner, ImageTestAssets, VllmRunner
13
+from .vlm_utils import model_utils
14
+
15
+MODEL = "google/gemma-3-4b-it"
16
+PROMPT = (
17
+ "<bos><start_of_turn>user\n"
18
+ "<start_of_image>What is the content in the center of the image?"
19
+ "<end_of_turn>\n<start_of_turn>model\n"
20
+)
21
+
22
+
23
+def _install_prefill_hidden_capture(model):
24
+ model = getattr(model, "module", model)
25
+ model._prefill_hidden = None
26
+
27
+ language_model = model.language_model.model
28
+ original_forward = language_model.forward
29
+
30
+ def forward(*args, **kwargs):
31
+ hidden_states = original_forward(*args, **kwargs)
32
+ if model._prefill_hidden is None and torch.is_tensor(hidden_states):
33
+ model._prefill_hidden = hidden_states.detach().float().cpu()
34
+ return hidden_states
35
+
36
+ language_model.forward = forward
37
+
38
+
39
+def _get_prefill_hidden(model):
40
+ model = getattr(model, "module", model)
41
+ hidden = getattr(model, "_prefill_hidden", None)
42
+ assert hidden is not None
43
+ return hidden
44
+
45
+
46
+def _get_hf_prefill_hidden(hf_model: HfRunner, image: Any):
47
+ inputs = hf_model.get_inputs([PROMPT], images=[image])[0]
48
+ with torch.no_grad():
49
+ outputs = hf_model.model.model(
50
+ **hf_model.wrap_device(inputs),
51
+ use_cache=False,
52
+ )
53
+ return outputs.last_hidden_state[0].detach().float().cpu()
54
+
55
+
56
+def _get_vllm_prefill_hidden(
57
+ vllm_runner: type[VllmRunner],
58
+ image: Any,
59
+ vllm_runner_kwargs: dict[str, Any],
60
+):
61
+ with vllm_runner(
62
+ MODEL,
63
+ max_model_len=4096,
64
+ max_num_seqs=2,
65
+ enforce_eager=True,
66
+ limit_mm_per_prompt={"image": 1},
67
+ **vllm_runner_kwargs,
68
+ ) as vllm_model:
69
+ vllm_model.apply_model(_install_prefill_hidden_capture)
70
+ vllm_model.generate_greedy([PROMPT], max_tokens=1, images=[image])
71
+ return vllm_model.apply_model(_get_prefill_hidden)[0]
72
+
73
+
74
+@pytest.mark.core_model
75
+@pytest.mark.skipif(
76
+ current_platform.is_rocm(), reason="ROCm attention has accuracy issue for this test"
77
+)
78
+def test_mm_prefix_lm_e2e(
79
+ hf_runner: type[HfRunner],
80
+ vllm_runner: type[VllmRunner],
81
+ image_assets: ImageTestAssets,
82
+ monkeypatch: pytest.MonkeyPatch,
83
+):
84
+ """Regression: Gemma3 native prefill must apply image prefix-LM mask."""
85
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
86
+ image = image_assets[0].pil_image
87
+
88
+ vllm_runner_kwargs: dict[str, Any] = {
89
+ "mm_processor_cache_gb": 0,
90
+ "mm_processor_kwargs": {"do_pan_and_scan": True},
91
+ }
92
+ vllm_hidden = _get_vllm_prefill_hidden(vllm_runner, image, vllm_runner_kwargs)
93
+
94
+ hf_model = hf_runner(
95
+ MODEL,
96
+ auto_cls=AutoModelForImageTextToText,
97
+ )
98
+ hf_model = model_utils.gemma3_patch_hf_runner(hf_model)
99
+
100
+ with hf_model:
101
+ hf_hidden = _get_hf_prefill_hidden(hf_model, image)
102
+
103
+ assert vllm_hidden.shape == hf_hidden.shape
104
+
105
+ full_cos = torch.nn.functional.cosine_similarity(
106
+ vllm_hidden.flatten(), hf_hidden.flatten(), dim=0
107
+ )
108
+ image_cos = torch.nn.functional.cosine_similarity(
109
+ vllm_hidden[1:769].flatten(), hf_hidden[1:769].flatten(), dim=0
110
+ )
111
+
112
+ assert full_cos > 0.9, (
113
+ "Gemma3 mm-prefix-LM full prefill hidden states should be close to HF; "
114
+ f"got {full_cos=}"
115
+ )
116
+ assert image_cos > 0.9, (
117
+ "Gemma3 mm-prefix-LM image prefill hidden states should be close to HF; "
118
+ f"got {image_cos=}"
119
+ )
vllm/config/model.py
+1
-1
@@ -1252,7 +1252,7 @@ class ModelConfig:
1252
def is_deepseek_mla(self) -> bool:
1253
return self.model_arch_config.is_deepseek_mla
1254
1255
- @property
1255
+ @cached_property
1256
def is_mm_prefix_lm(self) -> bool:
1257
return self.model_arch_config.is_mm_prefix_lm
1258
vllm/v1/worker/gpu/attn_utils.py
+27
@@ -14,6 +14,7 @@ from vllm.config import (
14
)
15
from vllm.model_executor.layers.attention import Attention
16
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17
+from vllm.multimodal.inputs import MultiModalFeatureSpec
18
from vllm.utils.torch_utils import get_dtype_size
19
from vllm.v1.attention.backend import (
20
AttentionCGSupport,
@@ -487,6 +488,7 @@ def build_attn_metadata(
488
seq_lens_cpu_upper_bound: torch.Tensor | None = None,
489
dcp_local_seq_lens: torch.Tensor | None = None,
490
positions: torch.Tensor | None = None,
491
+ mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
492
model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
493
for_cudagraph_capture: bool = False,
494
causal: bool = True,
@@ -523,6 +525,7 @@ def build_attn_metadata(
525
causal=causal,
526
dcp_local_seq_lens=dcp_local_seq_lens,
527
positions=positions,
528
+ mm_req_doc_ranges=mm_req_doc_ranges,
529
rswa_prefix_lens=rswa_prefix_lens,
530
**common_attn_metadata_extra_kwargs,
531
)
@@ -550,3 +553,27 @@ def build_attn_metadata(
553
for layer_name in attn_group.layer_names:
554
attn_metadata[layer_name] = metadata
555
return attn_metadata
556
+
557
+
558
+def compute_mm_prefix_ranges(
559
+ req_ids: list[str],
560
+ mm_features: dict[str, list[MultiModalFeatureSpec]],
561
+ sliding_window: int | None = None,
562
+) -> dict[int, list[tuple[int, int]]]:
563
+ """Compute PrefixLM bidirectional ranges for multimodal tokens.
564
+
565
+ Ranges exceeding sliding_window are skipped to prevent early tokens
566
+ from attending across the entire image span.
567
+ """
568
+ req_doc_ranges: dict[int, list[tuple[int, int]]] = {}
569
+ for req_idx, req_id in enumerate(req_ids):
570
+ image_doc_ranges = []
571
+ for mm_feature in mm_features.get(req_id, ()):
572
+ if mm_feature.modality not in ("image", "video"):
573
+ continue
574
+ for r in mm_feature.mm_position.extract_embeds_range():
575
+ if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window:
576
+ continue
577
+ image_doc_ranges.append(r)
578
+ req_doc_ranges[req_idx] = image_doc_ranges
579
+ return req_doc_ranges
vllm/v1/worker/gpu/model_states/default.py
+16
-1
@@ -9,7 +9,10 @@ from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
from vllm.v1.core.sched.output import NewRequestData
11
from vllm.v1.kv_cache_interface import KVCacheConfig
12
-from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
12
+from vllm.v1.worker.gpu.attn_utils import (
13
+ build_attn_metadata,
14
+ compute_mm_prefix_ranges,
15
+)
16
from vllm.v1.worker.gpu.input_batch import InputBatch
17
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
18
from vllm.v1.worker.gpu.mm.rope import get_rope_state
@@ -152,6 +155,17 @@ class DefaultModelState(ModelState):
155
max_seq_len = self.max_model_len
156
else:
157
max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()
158
+ req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
159
+ if (
160
+ self.supports_mm_inputs
161
+ and self.encoder_cache is not None
162
+ and self.model_config.is_mm_prefix_lm
163
+ ):
164
+ req_doc_ranges = compute_mm_prefix_ranges(
165
+ req_ids=input_batch.req_ids,
166
+ mm_features=self.encoder_cache.mm_features,
167
+ sliding_window=self.model_config.get_sliding_window(),
168
+ )
169
attn_metadata = build_attn_metadata(
170
attn_groups=attn_groups,
171
num_reqs=num_reqs,
@@ -167,6 +181,7 @@ class DefaultModelState(ModelState):
181
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
182
dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
183
positions=input_batch.positions,
184
+ mm_req_doc_ranges=req_doc_ranges,
185
for_cudagraph_capture=for_capture,
186
rswa_prefix_lens=input_batch.prompt_lens,
187
)