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 7, 00:15 UTC
[BE] Rename `should_torch_compile_mm_vit` to `should_torch_compile_mm_encoder` (#36281)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Mar 9, 2026 at 11:22 UTC
3fd03f1ec29cf9ac20584ad68156fc7279387979
6 files changed
+23
-17
docs/design/torch_compile_multimodal.md
+1
-1
@@ -26,7 +26,7 @@ This feature is off by default, but can be enabled by setting `compile_mm_encode
26
27
To compile a multimodal component such as an encoder, we follow the same mechanism as the LLM text backbone, with a few additional scaffoldings:
28
29
-1. The `@support_torch_compile` decorator should include `enable_if=should_torch_compile_mm_vit`. This will gate the compilation behind our
29
+1. The `@support_torch_compile` decorator should include `enable_if=should_torch_compile_mm_encoder`. This will gate the compilation behind our
30
`compile_mm_encoder` configuration
31
32
2. `with set_model_tag("<component_name>", is_encoder=True)` context manager should be used around the nn.Module's instantiation. Since torch.compile
vllm/compilation/decorators.py
+5
@@ -47,6 +47,11 @@ IGNORE_COMPILE_KEY = "_ignore_compile_vllm"
47
_T = TypeVar("_T", bound=nn.Module)
48
49
50
+def should_torch_compile_mm_encoder(vllm_config: VllmConfig) -> bool:
51
+ """Callable to be passed to `@support_torch_compile`'s `enable_if` argument."""
52
+ return vllm_config.compilation_config.compile_mm_encoder
53
+
54
+
55
def ignore_torch_compile(cls: type[_T]) -> type[_T]:
56
"""
57
A decorator to ignore support_torch_compile decorator
vllm/model_executor/models/lfm2_siglip2.py
+5
-3
@@ -10,7 +10,10 @@ from torch import nn
10
from torch.nn import functional as F
11
from transformers import Siglip2VisionConfig
12
13
-from vllm.compilation.decorators import support_torch_compile
13
+from vllm.compilation.decorators import (
14
+ should_torch_compile_mm_encoder,
15
+ support_torch_compile,
16
+)
17
from vllm.distributed import get_tensor_model_parallel_world_size
18
from vllm.model_executor.layers.activation import get_act_fn
19
from vllm.model_executor.layers.attention import MMEncoderAttention
@@ -25,7 +28,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
from .vision import (
29
is_vit_use_data_parallel,
30
resolve_visual_encoder_outputs,
28
- should_torch_compile_mm_vit,
31
)
32
33
@@ -269,7 +271,7 @@ class Siglip2MLP(nn.Module):
271
272
@support_torch_compile(
273
dynamic_arg_dims={"hidden_states": [0, 1], "cu_seqlens": 0},
272
- enable_if=should_torch_compile_mm_vit,
274
+ enable_if=should_torch_compile_mm_encoder,
275
)
276
class Siglip2EncoderLayer(nn.Module):
277
def __init__(
vllm/model_executor/models/mllama4.py
+5
-3
@@ -31,7 +31,10 @@ from transformers.models.llama4.image_processing_llama4_fast import (
31
get_best_fit,
32
)
33
34
-from vllm.compilation.decorators import support_torch_compile
34
+from vllm.compilation.decorators import (
35
+ should_torch_compile_mm_encoder,
36
+ support_torch_compile,
37
+)
38
from vllm.config import VllmConfig, set_current_vllm_config
39
from vllm.config.multimodal import BaseDummyOptions
40
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -49,7 +52,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
52
from vllm.model_executor.model_loader.utils import initialize_model
53
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
54
from vllm.model_executor.models.module_mapping import MultiModelKeys
52
-from vllm.model_executor.models.vision import should_torch_compile_mm_vit
55
from vllm.multimodal import MULTIMODAL_REGISTRY
56
from vllm.multimodal.inputs import (
57
MultiModalDataDict,
@@ -454,7 +456,7 @@ class Llama4UnfoldConvolution(nn.Module):
456
457
458
@support_torch_compile(
457
- dynamic_arg_dims={"images_flattened": 0}, enable_if=should_torch_compile_mm_vit
459
+ dynamic_arg_dims={"images_flattened": 0}, enable_if=should_torch_compile_mm_encoder
460
)
461
class Llama4VisionModel(nn.Module):
462
def __init__(
vllm/model_executor/models/qwen2_5_vl.py
+7
-5
@@ -42,7 +42,10 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
42
Qwen2_5_VLVisionConfig,
43
)
44
45
-from vllm.compilation.decorators import support_torch_compile
45
+from vllm.compilation.decorators import (
46
+ should_torch_compile_mm_encoder,
47
+ support_torch_compile,
48
+)
49
from vllm.config import VllmConfig
50
from vllm.distributed import parallel_state
51
from vllm.distributed import utils as dist_utils
@@ -65,7 +68,6 @@ from vllm.model_executor.layers.rotary_embedding.common import (
68
)
69
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
70
from vllm.model_executor.models.module_mapping import MultiModelKeys
68
-from vllm.model_executor.models.vision import should_torch_compile_mm_vit
71
from vllm.multimodal import MULTIMODAL_REGISTRY
72
from vllm.multimodal.evs import (
73
compute_mrope_for_media,
@@ -424,7 +426,7 @@ class Qwen2_5_VisionAttention(nn.Module):
426
"rotary_pos_emb_cos": 0,
427
"rotary_pos_emb_sin": 0,
428
},
427
- enable_if=should_torch_compile_mm_vit,
429
+ enable_if=should_torch_compile_mm_encoder,
430
)
431
class Qwen2_5_VisionBlock(nn.Module):
432
def __init__(
@@ -483,7 +485,7 @@ class Qwen2_5_VisionBlock(nn.Module):
485
dynamic_arg_dims={
486
"x": 0,
487
},
486
- enable_if=should_torch_compile_mm_vit,
488
+ enable_if=should_torch_compile_mm_encoder,
489
)
490
class Qwen2_5_VisionPatchEmbed(nn.Module):
491
def __init__(
@@ -518,7 +520,7 @@ class Qwen2_5_VisionPatchEmbed(nn.Module):
520
dynamic_arg_dims={
521
"x": 0,
522
},
521
- enable_if=should_torch_compile_mm_vit,
523
+ enable_if=should_torch_compile_mm_encoder,
524
)
525
class Qwen2_5_VisionPatchMerger(nn.Module):
526
def __init__(
vllm/model_executor/models/vision.py
-5
@@ -143,11 +143,6 @@ def is_vit_use_data_parallel():
143
return mm_encoder_tp_mode == "data"
144
145
146
-def should_torch_compile_mm_vit(vllm_config: VllmConfig) -> bool:
147
- """Callable to be passed to `@support_torch_compile`'s `enable_if` argument."""
148
- return vllm_config.compilation_config.compile_mm_encoder
149
-
150
-
146
VisionFeatureSelectStrategyStr = Literal["class", "default", "full"]
147
148
VisionFeatureSelectStrategy: TypeAlias = (