[Core] Keep attention backends eligible for text-only serving of prefix-LM models (#48796)
Signed-off-by: qtris123 <voquangtri2021@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Tri Vo committed
Jul 25, 2026 at 17:06 UTC
fe5145765f211c9b5abc446b42491f84895d70ba
4 files changed
+165
-7
tests/config/test_multimodal_config.py
+100
@@ -1,10 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from unittest.mock import MagicMock, patch
5
+
6
import pytest
7
+from transformers import PretrainedConfig
8
9
from vllm.config.model import ModelConfig
10
from vllm.config.multimodal import MultiModalConfig
11
+from vllm.transformers_utils.model_arch_config_convertor import (
12
+ ModelArchConfigConvertorBase,
13
+)
14
from vllm.v1.attention.backends.registry import AttentionBackendEnum
15
16
@@ -68,3 +74,97 @@ def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
74
).compute_hash()
75
assert base_hash != fp8_hash
76
assert fp8_hash != fp8_static_hash
77
+
78
+
79
+def _make_mm_prefix_model_config(
80
+ *,
81
+ language_model_only: bool = False,
82
+) -> ModelConfig:
83
+ model_config = MagicMock(spec=ModelConfig)
84
+ model_config.multimodal_config = MultiModalConfig(
85
+ language_model_only=language_model_only
86
+ )
87
+ # Bind real helper methods onto the mock.
88
+ model_config._supports_multimodal_for_mm_prefix = (
89
+ ModelConfig._supports_multimodal_for_mm_prefix.__get__(
90
+ model_config, ModelConfig
91
+ )
92
+ )
93
+ return model_config
94
+
95
+
96
+@pytest.mark.parametrize("supports_mm", [True, False])
97
+def test_supports_multimodal_for_mm_prefix_uses_registry(supports_mm: bool):
98
+ model_config = _make_mm_prefix_model_config()
99
+
100
+ with patch(
101
+ "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
102
+ return_value=supports_mm,
103
+ ) as mocked:
104
+ assert model_config._supports_multimodal_for_mm_prefix() is supports_mm
105
+ mocked.assert_called_once_with(model_config)
106
+
107
+ # Sticky cache — registry must not be consulted again.
108
+ with patch(
109
+ "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
110
+ side_effect=AssertionError("should use cache"),
111
+ ):
112
+ assert model_config._supports_multimodal_for_mm_prefix() is supports_mm
113
+
114
+
115
+def test_supports_multimodal_for_mm_prefix_before_multimodal_config():
116
+ model_config = _make_mm_prefix_model_config()
117
+ model_config.multimodal_config = None
118
+
119
+ assert model_config._supports_multimodal_for_mm_prefix() is True
120
+ assert not hasattr(model_config, "_supports_multimodal_inputs_cached")
121
+
122
+
123
+def test_language_model_only_disables_via_supports_multimodal_inputs():
124
+ """language_model_only zeros all limits, so registry reports text-only."""
125
+ model_config = _make_mm_prefix_model_config(language_model_only=True)
126
+
127
+ with patch(
128
+ "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
129
+ return_value=False,
130
+ ):
131
+ assert model_config._supports_multimodal_for_mm_prefix() is False
132
+
133
+
134
+def test_convertor_clears_mm_prefix_when_multimodal_disabled():
135
+ hf_config = PretrainedConfig(
136
+ model_type="gemma3",
137
+ architectures=["Gemma3ForConditionalGeneration"],
138
+ )
139
+ hf_config.is_mm_prefix_lm = True
140
+ convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
141
+
142
+ assert convertor.is_mm_prefix_lm(supports_multimodal=True) is True
143
+ assert convertor.is_mm_prefix_lm(supports_multimodal=False) is False
144
+
145
+ enabled = convertor.convert(supports_multimodal=True)
146
+ disabled = convertor.convert(supports_multimodal=False)
147
+ assert enabled.is_mm_prefix_lm is True
148
+ assert disabled.is_mm_prefix_lm is False
149
+
150
+
151
+def test_sticky_cache_survives_text_subconfig_regeneration():
152
+ """with_hf_config deepcopies the cached decision onto text submodules."""
153
+ model_config = _make_mm_prefix_model_config()
154
+ with patch(
155
+ "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
156
+ return_value=False,
157
+ ):
158
+ assert model_config._supports_multimodal_for_mm_prefix() is False
159
+
160
+ # Simulate deepcopy onto a Gemma4ForCausalLM-like config that would
161
+ # otherwise fail registry lookup / return False incorrectly.
162
+ text_config = _make_mm_prefix_model_config()
163
+ text_config._supports_multimodal_inputs_cached = (
164
+ model_config._supports_multimodal_inputs_cached
165
+ )
166
+ with patch(
167
+ "vllm.multimodal.MULTIMODAL_REGISTRY.supports_multimodal_inputs",
168
+ side_effect=AssertionError("must not re-query registry"),
169
+ ):
170
+ assert text_config._supports_multimodal_for_mm_prefix() is False
tests/models/utils.py
+6
@@ -507,6 +507,12 @@ def dummy_hf_overrides(
507
hf_config = _hf_config
508
hf_text_config = text_config
509
510
+ # Keep architecture conversion on the default multimodal path; this
511
+ # helper only needs HF-derived fields, not deployment MM limits.
512
+ @staticmethod
513
+ def _supports_multimodal_for_mm_prefix() -> bool:
514
+ return True
515
+
516
model_arch_config = ModelConfig.get_model_arch_config(DummyConfig)
517
# Only set MoE related config when the model has MoE layers.
518
# Otherwise all models detected as MoE by _get_transformers_backend_cls.
vllm/config/model.py
+43
-1
@@ -758,6 +758,10 @@ class ModelConfig:
758
"disable the cache with --mm-processor-cache-gb 0."
759
)
760
761
+ # Rebuild after multimodal_config exists so text-only mm_prefix
762
+ # clearing is applied (and cached for later with_hf_config calls).
763
+ self.model_arch_config = self.get_model_arch_config()
764
+
765
if self.disable_sliding_window:
766
# Set after get_and_verify_max_len to ensure that max_model_len
767
# can be correctly capped to sliding window size
@@ -770,6 +774,42 @@ class ModelConfig:
774
self._verify_cuda_graph()
775
self._verify_bnb_config()
776
777
+ def _supports_multimodal_for_mm_prefix(self) -> bool:
778
+ """Whether multimodal inputs can still appear for this deployment.
779
+
780
+ This runs more than once per config: once early in ``__post_init__``
781
+ (before ``multimodal_config`` exists), again after it is created, and
782
+ then for every ``get_model_arch_config`` regeneration -- notably
783
+ ``with_hf_config``, which deep-copies this ``ModelConfig`` and swaps
784
+ ``hf_config`` for a text-only submodule (e.g. ``Gemma4ForCausalLM``).
785
+
786
+ The result is cached for correctness, not just to save work: on the
787
+ ``with_hf_config`` copy the submodule architecture has no registered
788
+ multimodal processor, so re-querying the registry would raise and be
789
+ treated as text-only, wrongly clearing ``is_mm_prefix_lm`` even when a
790
+ vision modality is still enabled (e.g. ``image=0`` but video allowed).
791
+ The deep-copied cache preserves the top-level decision instead.
792
+ """
793
+ cached = getattr(self, "_supports_multimodal_inputs_cached", None)
794
+ if cached is not None:
795
+ return cached
796
+
797
+ if self.multimodal_config is None:
798
+ # Early call before multimodal init — do not clear mm_prefix yet.
799
+ return True
800
+
801
+ from vllm.multimodal import MULTIMODAL_REGISTRY
802
+
803
+ supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(self)
804
+ self._supports_multimodal_inputs_cached = supports_mm
805
+ if not supports_mm:
806
+ logger.info_once(
807
+ "Disabled mm_prefix attention mode because multimodal inputs "
808
+ "are configuration-disabled. Attention backends without "
809
+ "mm_prefix support may now be selected."
810
+ )
811
+ return supports_mm
812
+
813
def get_model_arch_config(
814
self,
815
) -> ModelArchitectureConfig:
@@ -777,7 +817,9 @@ class ModelConfig:
817
self.hf_config.model_type, ModelArchConfigConvertorBase
818
)
819
convertor = convertor_cls(self.hf_config, self.hf_text_config)
780
- return convertor.convert()
820
+ return convertor.convert(
821
+ supports_multimodal=self._supports_multimodal_for_mm_prefix()
822
+ )
823
824
@field_validator("tokenizer", "max_model_len", mode="wrap")
825
@classmethod
vllm/transformers_utils/model_arch_config_convertor.py
+16
-6
@@ -300,8 +300,16 @@ class ModelArchConfigConvertorBase:
300
)
301
return False
302
303
- def is_mm_prefix_lm(self) -> bool:
304
- """Whether to use bidirectional attention for mm positions."""
303
+ def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
304
+ """Whether to use bidirectional attention for mm positions.
305
+
306
+ ``supports_multimodal`` is False when the deployment is configuration-
307
+ disabled for multimodal inputs (text-only serving). In that case
308
+ mm_prefix is unnecessary and must stay off so attention backends
309
+ without ``supports_mm_prefix()`` remain eligible.
310
+ """
311
+ if not supports_multimodal:
312
+ return False
313
if hasattr(self.hf_config, "is_mm_prefix_lm"):
314
return bool(self.hf_config.is_mm_prefix_lm)
315
# fallback to list of known models
@@ -358,7 +366,7 @@ class ModelArchConfigConvertorBase:
366
derived_max_model_len = tmp_max_len
367
return derived_max_model_len, max_len_key
368
361
- def convert(self) -> ModelArchitectureConfig:
369
+ def convert(self, supports_multimodal: bool = True) -> ModelArchitectureConfig:
370
model_arch_config = ModelArchitectureConfig(
371
architectures=self.get_architectures(),
372
model_type=self.hf_config.model_type,
@@ -372,7 +380,7 @@ class ModelArchConfigConvertorBase:
380
num_experts=self.get_num_experts(),
381
quantization_config=self.get_quantization_config(),
382
is_deepseek_mla=self.is_deepseek_mla(),
375
- is_mm_prefix_lm=self.is_mm_prefix_lm(),
383
+ is_mm_prefix_lm=self.is_mm_prefix_lm(supports_multimodal),
384
rswa_window=self.rswa_window(),
385
derived_max_model_len_and_key=self.derive_max_model_len_and_key(),
386
)
@@ -401,7 +409,7 @@ class CohereAsrModelArchConfigConvertor(ModelArchConfigConvertorBase):
409
)
410
return enc_num_kv_heads
411
404
- def is_mm_prefix_lm(self) -> bool:
412
+ def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
413
return False
414
415
@@ -584,7 +592,9 @@ class Gemma4MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
592
593
594
class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
587
- def is_mm_prefix_lm(self) -> bool:
595
+ def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
596
+ if not supports_multimodal:
597
+ return False
598
return (
599
getattr(self.hf_text_config, "use_bidirectional_attention", None)
600
== "vision"