A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 16:00 UTC
Don't compile vision encoder for Transformers backend (#30518)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Apr 2, 2026 at 13:42 UTC
dfe5e31689a0b2511e4deaa16e16a823b734130d
5 files changed
+189
-52
vllm/compilation/decorators.py
+17
-2
@@ -205,6 +205,8 @@ def support_torch_compile(
205
if v.annotation in [
206
torch.Tensor,
207
torch.Tensor | None,
208
+ torch.FloatTensor,
209
+ torch.FloatTensor | None,
210
IntermediateTensors,
211
IntermediateTensors | None,
212
]:
@@ -346,7 +348,7 @@ def _support_torch_compile(
348
349
def __init__(
350
self: _T,
349
- *,
351
+ *args,
352
vllm_config: VllmConfig | None = None,
353
prefix: str = "",
354
**kwargs: Any,
@@ -357,11 +359,24 @@ def _support_torch_compile(
359
# NOTE: to support multimodal models (such as encoder),
360
# we may not have vllm_config so we may need to patch it
361
sig = inspect.signature(old_init)
362
+ # Check that any positional arguments match the old_init method signature
363
+ annotations = [p.annotation for p in sig.parameters.values()]
364
+ for arg, annotation in zip(args, annotations):
365
+ if annotation is inspect._empty:
366
+ continue
367
+ if not isinstance(arg, annotation):
368
+ init = f"'{type(self).__name__}.__init__'"
369
+ arg_type = f"'{type(arg).__name__}'"
370
+ raise TypeError(
371
+ f"{init} received a positional argument of type {arg_type}, "
372
+ "but no parameter of that type was found in the method signature. "
373
+ f"Please either annotate {init} or pass it as a keyword argument."
374
+ )
375
if "vllm_config" in sig.parameters:
376
kwargs["vllm_config"] = vllm_config
377
if "prefix" in sig.parameters:
378
kwargs["prefix"] = prefix
364
- old_init(self, **kwargs)
379
+ old_init(self, *args, **kwargs)
380
381
self.vllm_config = vllm_config
382
self.compilation_config = self.vllm_config.compilation_config
vllm/config/compilation.py
+4
-3
@@ -495,9 +495,10 @@ class CompilationConfig:
495
If empty list [], no ops are excluded (suitable for full cudagraphs)."""
496
compile_mm_encoder: bool = False
497
"""Whether or not to compile the multimodal encoder.
498
- Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models
499
- on selected platforms. Disabled by default until more models
500
- are supported/tested to work."""
498
+ Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models on selected
499
+ platforms. It may also work for models loaded with the Transformers modeling backend
500
+ if the encoder is compilable. Disabled by default until more models are
501
+ supported/tested to work."""
502
503
# Vision encoder CUDA graph
504
cudagraph_mm_encoder: bool = False
vllm/model_executor/models/transformers/__init__.py
-21
@@ -16,13 +16,11 @@
16
# limitations under the License.
17
"""Wrapper around `transformers` models"""
18
19
-from vllm.compilation.decorators import support_torch_compile
19
from vllm.model_executor.models.transformers.base import Base
20
from vllm.model_executor.models.transformers.causal import CausalMixin
21
from vllm.model_executor.models.transformers.legacy import LegacyMixin
22
from vllm.model_executor.models.transformers.moe import MoEMixin
23
from vllm.model_executor.models.transformers.multimodal import (
25
- DYNAMIC_ARG_DIMS,
24
MultiModalDummyInputsBuilder,
25
MultiModalMixin,
26
MultiModalProcessingInfo,
@@ -32,16 +30,13 @@ from vllm.model_executor.models.transformers.pooling import (
30
EmbeddingMixin,
31
SequenceClassificationMixin,
32
)
35
-from vllm.model_executor.models.transformers.utils import can_enable_torch_compile
33
from vllm.multimodal import MULTIMODAL_REGISTRY
34
35
36
# Text only models
40
-@support_torch_compile(enable_if=can_enable_torch_compile)
37
class TransformersForCausalLM(CausalMixin, Base): ...
38
39
44
-@support_torch_compile(enable_if=can_enable_torch_compile)
40
class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
41
42
@@ -51,9 +46,6 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ...
46
info=MultiModalProcessingInfo,
47
dummy_inputs=MultiModalDummyInputsBuilder,
48
)
54
-@support_torch_compile(
55
- dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
56
-)
49
class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
50
51
@@ -62,20 +54,15 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ...
54
info=MultiModalProcessingInfo,
55
dummy_inputs=MultiModalDummyInputsBuilder,
56
)
65
-@support_torch_compile(
66
- dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
67
-)
57
class TransformersMultiModalMoEForCausalLM(
58
MoEMixin, MultiModalMixin, CausalMixin, Base
59
): ...
60
61
62
# Embedding models
74
-@support_torch_compile(enable_if=can_enable_torch_compile)
63
class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ...
64
65
78
-@support_torch_compile(enable_if=can_enable_torch_compile)
66
class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
67
68
@@ -84,20 +71,15 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ...
71
info=MultiModalProcessingInfo,
72
dummy_inputs=MultiModalDummyInputsBuilder,
73
)
87
-@support_torch_compile(
88
- dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
89
-)
74
class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ...
75
76
77
# Sequence classification models
94
-@support_torch_compile(enable_if=can_enable_torch_compile)
78
class TransformersForSequenceClassification(
79
SequenceClassificationMixin, LegacyMixin, Base
80
): ...
81
82
100
-@support_torch_compile(enable_if=can_enable_torch_compile)
83
class TransformersMoEForSequenceClassification(
84
SequenceClassificationMixin, MoEMixin, Base
85
): ...
@@ -108,9 +90,6 @@ class TransformersMoEForSequenceClassification(
90
info=MultiModalProcessingInfo,
91
dummy_inputs=MultiModalDummyInputsBuilder,
92
)
111
-@support_torch_compile(
112
- dynamic_arg_dims=DYNAMIC_ARG_DIMS, enable_if=can_enable_torch_compile
113
-)
93
class TransformersMultiModalForSequenceClassification(
94
SequenceClassificationMixin, MultiModalMixin, Base
95
): ...
vllm/model_executor/models/transformers/base.py
+101
-17
@@ -16,6 +16,7 @@
16
# limitations under the License.
17
"""Transformers modeling backend base class."""
18
19
+import sys
20
from collections.abc import Callable, Iterable
21
from itertools import chain
22
from operator import attrgetter
@@ -29,6 +30,7 @@ from torch import nn
30
from transformers import AutoModel
31
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
32
33
+from vllm.compilation.decorators import support_torch_compile
34
from vllm.config.utils import getattr_iter
35
from vllm.distributed import get_pp_group, get_tp_group
36
from vllm.distributed.utils import get_pp_indices
@@ -47,6 +49,7 @@ from vllm.model_executor.models.interfaces import (
49
)
50
from vllm.model_executor.models.interfaces_base import VllmModel
51
from vllm.model_executor.models.transformers.utils import (
52
+ can_enable_torch_compile,
53
get_feature_request_tip,
54
init_on_device_without_buffers,
55
log_replacement,
@@ -117,6 +120,7 @@ class Base(
120
self.config = vllm_config.model_config.hf_config
121
self.text_config = self.config.get_text_config()
122
self.cache_config = vllm_config.cache_config
123
+ self.compilation_config = vllm_config.compilation_config
124
self.device_config = vllm_config.device_config
125
self.model_config = vllm_config.model_config
126
self.parallel_config = vllm_config.parallel_config
@@ -155,14 +159,16 @@ class Base(
159
if "gptq" in quant_method_name:
160
self.ignore_unexpected_suffixes.append(".bias")
161
158
- # Patch config and init on "meta" to delay allocating GPU tensors
162
self._patch_config()
163
+ from_config_kwargs = dict(
164
+ config=self.config,
165
+ dtype=self.model_config.dtype,
166
+ trust_remote_code=self.model_config.trust_remote_code,
167
+ )
168
+ self._decorate_for_torch_compile(**from_config_kwargs)
169
+ # Init on "meta" to delay allocating GPU tensors
170
with init_on_device_without_buffers("meta"):
161
- self.model: PreTrainedModel = AutoModel.from_config(
162
- self.config,
163
- dtype=self.model_config.dtype,
164
- trust_remote_code=self.model_config.trust_remote_code,
165
- )
171
+ self.model: PreTrainedModel = AutoModel.from_config(**from_config_kwargs)
172
173
# Create weight name to module qualname mapper
174
self._create_hf_to_vllm_mapper()
@@ -218,6 +224,82 @@ class Base(
224
if sub_config.dtype != (dtype := self.config.dtype):
225
sub_config.dtype = dtype
226
227
+ def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]:
228
+ """
229
+ Get the decoder class from the model.
230
+
231
+ Args:
232
+ kwargs: The kwargs to create the model.
233
+
234
+ Returns:
235
+ The decoder class.
236
+ """
237
+ with torch.device("meta"):
238
+ model: PreTrainedModel = AutoModel.from_config(**kwargs)
239
+ decoder_cls = type(model.get_decoder())
240
+ logger.debug("Identified decoder class as: %s", decoder_cls)
241
+ del model
242
+ return decoder_cls
243
+
244
+ def _decorate_cls_for_torch_compile(
245
+ self,
246
+ cls: type[PreTrainedModel],
247
+ dynamic_arg_dims: dict[str, int] | None,
248
+ enable_if: Callable[["VllmConfig"], bool],
249
+ is_encoder: bool,
250
+ ):
251
+ """
252
+ Decorate `cls` to indicate to vLLM that it supports torch compile.
253
+
254
+ Args:
255
+ cls: The PreTrainedModel class to decorate.
256
+ dynamic_arg_dims: A mapping from argument name to the dynamic dimensions
257
+ of the argument. If None, default dynamic arg dims will be used. See
258
+ [`support_torch_compile`][vllm.compilation.decorators.support_torch_compile]
259
+ for more details.
260
+ enable_if: A function which takes in the vLLM config and returns whether
261
+ torch compile should be enabled for this class.
262
+ is_encoder: Whether the class being decorated is an encoder.
263
+ """
264
+ logger.debug(
265
+ "Decorating `%s` as %s for torch compile with dynamic_arg_dims of %s",
266
+ cls.__name__,
267
+ "encoder" if is_encoder else "decoder",
268
+ dynamic_arg_dims,
269
+ )
270
+
271
+ @support_torch_compile(
272
+ dynamic_arg_dims=dynamic_arg_dims,
273
+ enable_if=enable_if,
274
+ is_encoder=is_encoder,
275
+ )
276
+ class SupportTorchCompileWrapper(cls): ...
277
+
278
+ # Patch the class in its module
279
+ module = sys.modules[cls.__module__]
280
+ setattr(module, cls.__name__, SupportTorchCompileWrapper)
281
+
282
+ def _decorate_for_torch_compile(self, **kwargs: dict):
283
+ """
284
+ Decorate the model's decoder class to indicate to vLLM that it supports torch
285
+ compile if `can_enable_torch_compile` is True.
286
+
287
+ Args:
288
+ kwargs: The kwargs to create the model, which are needed to get the decoder
289
+ class.
290
+ """
291
+ self._decorate_cls_for_torch_compile(
292
+ cls=self._get_decoder_cls(**kwargs),
293
+ # Applied to a PreTrainedModel so the batch dimension will exist
294
+ dynamic_arg_dims=dict[str, int](
295
+ input_ids=1, # shape: [1, seq_len]
296
+ inputs_embeds=1, # shape: [1, seq_len, hidden_size]
297
+ position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope
298
+ ),
299
+ enable_if=can_enable_torch_compile,
300
+ is_encoder=False,
301
+ )
302
+
303
def _create_hf_to_vllm_mapper(self):
304
"""
305
Create a WeightsMapper to map checkpoint weight names to module qualnames.
@@ -553,11 +635,6 @@ class Base(
635
input_ids = None
636
inputs_embeds = intermediate_tensors["hidden_states"]
637
556
- if input_ids is not None:
557
- input_ids = input_ids[None, ...]
558
- if inputs_embeds is not None:
559
- inputs_embeds = inputs_embeds[None, ...]
560
-
638
# If the model scales embeddings inside the input embedding layer we must
639
# ensure they are scaled here since VocabParallelEmbedding will not do it
640
if (
@@ -568,22 +645,29 @@ class Base(
645
inputs_embeds = self.embed_input_ids(input_ids)
646
input_ids = None
647
571
- if self.model_config.uses_mrope:
572
- position_ids = positions[:, None]
573
- else:
574
- position_ids = positions[None, ...]
648
+ # Add batch dimension before entering Transformers model
649
+ if input_ids is not None and input_ids.ndim == 1:
650
+ # [seq_len] -> [1, seq_len]
651
+ input_ids = input_ids[None, ...]
652
+ if inputs_embeds is not None and inputs_embeds.ndim == 2:
653
+ # [seq_len, hidden_size] -> [1, seq_len, hidden_size]
654
+ inputs_embeds = inputs_embeds[None, ...]
655
+ if positions.ndim == 1:
656
+ # [seq_len] -> [1, seq_len]
657
+ positions = positions[None, ...]
658
659
outputs = self.model(
660
input_ids=input_ids,
661
inputs_embeds=inputs_embeds,
662
use_cache=False,
580
- position_ids=position_ids,
663
+ position_ids=positions,
664
attention_instances=self.attention_instances,
665
return_dict=False,
666
**self._output_aux_hidden_states_kwargs,
667
**kwargs,
668
)
586
- # We must remove the batch dimension from these outputs
669
+
670
+ # Remove batch dimension after exiting Transformers model
671
hidden_states = outputs[0][0, ...]
672
if self._output_aux_hidden_states_kwargs:
673
aux_hidden_states = [x[0][0, ...] for x in outputs[1:]]
vllm/model_executor/models/transformers/multimodal.py
+67
-9
@@ -20,7 +20,9 @@ from collections.abc import Mapping
20
from typing import TYPE_CHECKING
21
22
import torch
23
+from transformers import AutoModel
24
25
+from vllm.compilation.decorators import should_torch_compile_mm_encoder
26
from vllm.config.utils import getattr_iter
27
from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input
28
from vllm.logger import init_logger
@@ -46,19 +48,11 @@ from vllm.platforms import current_platform
48
from vllm.sequence import IntermediateTensors
49
50
if TYPE_CHECKING:
49
- from transformers import BatchFeature
51
+ from transformers import BatchFeature, PreTrainedModel
52
53
from vllm.config import VllmConfig
54
from vllm.config.multimodal import BaseDummyOptions
55
54
-DYNAMIC_ARG_DIMS = {
55
- "input_ids": 0,
56
- # set `positions` to last dim to support Qwen-mrope
57
- "positions": -1,
58
- "intermediate_tensors": 0,
59
- "inputs_embeds": 0,
60
-}
61
-
56
logger = init_logger(__name__)
57
58
@@ -274,6 +268,66 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
268
# Skip SupportsMRoPE.__init__ and call the next class in MRO
269
super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
270
271
+ def _get_encoder_cls(
272
+ self, modality: str = "image", **kwargs: dict
273
+ ) -> type["PreTrainedModel"]:
274
+ """
275
+ Get the encoder class from the model.
276
+
277
+ Args:
278
+ kwargs: The kwargs to create the model.
279
+
280
+ Returns:
281
+ The encoder class.
282
+ """
283
+ with torch.device("meta"):
284
+ model: PreTrainedModel = AutoModel.from_config(**kwargs)
285
+ encoder_cls = type(model.get_encoder(modality=modality))
286
+ logger.debug("Identified encoder class as: %s", encoder_cls)
287
+ if type(model) is encoder_cls:
288
+ raise ValueError(
289
+ "Unable to infer vision encoder class from the model. "
290
+ "You must either: update the model so that "
291
+ "https://huggingface.co/docs/transformers/en/main_classes/model#transformers.PreTrainedModel.get_encoder"
292
+ " can detect the vision encoder correctly, or remove "
293
+ "'compile_mm_encoder'."
294
+ )
295
+ del model
296
+ return encoder_cls
297
+
298
+ def _decorate_for_torch_compile(self, **kwargs: dict):
299
+ """
300
+ Decorate the model's decoder and encoder classes to indicate to vLLM that they
301
+ support torch compile if `can_enable_torch_compile` and
302
+ `should_torch_compile_mm_encoder` are True respectively.
303
+
304
+ Args:
305
+ kwargs: The kwargs to create the model, which are needed to get the decoder
306
+ and encoder classes.
307
+ """
308
+ super()._decorate_for_torch_compile(**kwargs)
309
+ # Decorate the vision encoder model class to support torch compile if needed
310
+ if self.compilation_config.compile_mm_encoder:
311
+ self.check_version("5.0.0", "multimodal encoder compilation support")
312
+ logger.warning_once(
313
+ "Multimodal encoder compilation with the Transformers modeling backend "
314
+ "is an experimental feature. It relies on:\n"
315
+ "- The vision encoder being torch compilable.\n"
316
+ "- All vision encoder tensor inputs must be type hinted as either "
317
+ "`torch.Tensor` or `torch.FloatTensor`.\n"
318
+ "- The 0-th dimension of all tensor inputs to the vision encoder being "
319
+ "the dynamic dimension (i.e., sequence length or number of patches).\n"
320
+ "Please report any issues you encounter to help us improve it."
321
+ )
322
+ self._decorate_cls_for_torch_compile(
323
+ cls=self._get_encoder_cls(**kwargs),
324
+ # TODO: properly infer dynamic_arg_dims based on the encoder's forward
325
+ # method signature. Currently we assume dim 0 for all tensor inputs.
326
+ dynamic_arg_dims=None,
327
+ enable_if=should_torch_compile_mm_encoder,
328
+ is_encoder=True,
329
+ )
330
+
331
def forward(
332
self,
333
input_ids: torch.Tensor | None,
@@ -285,6 +339,10 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
339
# Gemma3 and PaliGemma needs `token_type_ids` to work correctly
340
# Other models will not have `token_type_ids` in kwargs
341
kwargs = {k: v for k, v in kwargs.items() if k == "token_type_ids"}
342
+ # Positions shape handling for MRoPE models
343
+ if self.model_config.uses_mrope:
344
+ # [3, seq_len] -> [3, 1, seq_len]
345
+ positions = positions[:, None]
346
model_output = super().forward(
347
input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
348
)