[torch.compile][BE][Multimodal] Remove requirement to set_model_tag to avoid...

[torch.compile][BE][Multimodal] Remove requirement to set_model_tag to avoid cache conflict (#37345)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>

docs/design/torch_compile_multimodal.md

@@ -29,10 +29,9 @@ To compile a multimodal component such as an encoder, we follow the same mechani
 1. The `@support_torch_compile` decorator should include `enable_if=should_torch_compile_mm_encoder`. This will gate the compilation behind our
 `compile_mm_encoder` configuration
 
-2. `with set_model_tag("<component_name>", is_encoder=True)` context manager should be used around the nn.Module's instantiation. Since torch.compile
-relies on caching artifacts to reduce start time, we must properly propagate the `<component_name>` information to the cache in order to avoid collisions
-with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
-components (see Compile Range Integration).
+2. The `@support_torch_compile` decorator should include `is_encoder=True` for encoder components. This is needed for compile range integration
+(see Compile Range Integration). The decorator automatically uses the class name as the cache directory prefix, avoiding collisions between
+independently compiled sub-modules (e.g. vision encoder components vs the text backbone).
 
 ### CompilationConfig
 
@@ -57,8 +56,8 @@ tradeoff
 ### Compile ranges
 
 The torch.compile integration will try to rely on max_batch_size to infer compilation ranges for dynamic shapes; however, for modules used in the encoder, this
-shape can be difficult to infer due to the unspecified range of shapes the encoder may see as input. Therefore, we rely on `is_encoder=True` in the `set_model_tag`
-to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).
+shape can be difficult to infer due to the unspecified range of shapes the encoder may see as input. Therefore, we rely on `is_encoder=True` in the
+`@support_torch_compile` decorator to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).
 
 !!! note
     We may seek to tighten this range for better performance in the future

vllm/compilation/decorators.py

@@ -118,6 +118,7 @@ def support_torch_compile(
     dynamic_arg_dims: dict[str, int | list[int]] | None = None,
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
+    is_encoder: bool = False,
     shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,
 ) -> Callable[[type[_T]], type[_T]] | type[_T]:
     """
@@ -177,6 +178,11 @@ def support_torch_compile(
     enforce that dynamo does not specialize on 0/1 values in the case of dummy input
     such as for vision model compilation
 
+    `is_encoder` marks this module as a portion of an multimodal encoder.
+    When True, the compile range upper bound is set to MAX_INT32 instead of
+    max_num_batched_tokens, since encoder input shapes are unpredictable.
+    This is typically used for vision encoder sub-modules in multimodal models.
+
     `shape_invariants` is a function that gets compiled right before forward.
     The function should have the torch._check calls that are needed to set
     the relationships between different input sizes. For example:
@@ -226,6 +232,7 @@ def support_torch_compile(
             inferred_dynamic_arg_dims,
             mark_unbacked_dims,
             enable_if,
+            is_encoder,
             shape_invariants,
         )
 
@@ -316,6 +323,7 @@ def _support_torch_compile(
     dynamic_arg_dims: dict[str, int | list[int]],
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
+    is_encoder: bool = False,
     shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,
 ) -> type[_T]:
     """
@@ -345,8 +353,7 @@ def _support_torch_compile(
             vllm_config = get_current_vllm_config()
 
         # NOTE: to support multimodal models (such as encoder),
-        # we may not have vllm_config so we may need to patch
-        # it
+        # we may not have vllm_config so we may need to patch it
         sig = inspect.signature(old_init)
         if "vllm_config" in sig.parameters:
             kwargs["vllm_config"] = vllm_config
@@ -374,7 +381,11 @@ def _support_torch_compile(
         self.compiled = False
 
         # Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class
-        TorchCompileWithNoGuardsWrapper.__init__(self)
+        TorchCompileWithNoGuardsWrapper.__init__(
+            self,
+            compile_prefix=cls.__name__ if is_encoder else "",
+            is_encoder=is_encoder,
+        )
 
     cls.__init__ = __init__
 

vllm/compilation/wrapper.py

@@ -75,8 +75,14 @@ class TorchCompileWithNoGuardsWrapper:
             return ctx.result
         return callable_fn(*args, **kwargs)
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        compile_prefix: str = "",
+        is_encoder: bool = False,
+    ) -> None:
         self.compiled = False
+        self._compile_prefix = compile_prefix
+        self._is_encoder = is_encoder
 
         vllm_config = get_current_vllm_config()
         self.vllm_config = vllm_config
@@ -87,7 +93,9 @@ class TorchCompileWithNoGuardsWrapper:
         if mode is None:
             raise RuntimeError("Compilation mode cannot be NO_COMPILATION")
 
-        backend = vllm_config.compilation_config.init_backend(vllm_config)
+        backend = vllm_config.compilation_config.init_backend(
+            vllm_config, prefix=compile_prefix, is_encoder=is_encoder
+        )
         options = {}
 
         if isinstance(backend, str) and backend == "inductor":
@@ -332,4 +340,8 @@ def reset_compile_wrapper(model: torch.nn.Module) -> None:
     compilation_config.local_cache_dir = ""
 
     model.__class__.forward.__code__ = model.original_code_object()
-    TorchCompileWithNoGuardsWrapper.__init__(model)
+    TorchCompileWithNoGuardsWrapper.__init__(
+        model,
+        compile_prefix=model._compile_prefix,
+        is_encoder=model._is_encoder,
+    )

vllm/config/compilation.py

@@ -909,11 +909,19 @@ class CompilationConfig:
         if self.backend == "":
             self.backend = current_platform.get_compile_backend()
 
-    def init_backend(self, vllm_config: "VllmConfig") -> str | Callable:
+    def init_backend(
+        self,
+        vllm_config: "VllmConfig",
+        prefix: str = "",
+        is_encoder: bool = False,
+    ) -> str | Callable:
         """
         Initialize the backend for the compilation config from a vllm config.
         Arguments:
             vllm_config: The vllm config to initialize the backend from.
+            prefix: Cache directory prefix for this compiled module.
+            is_encoder: Whether this module is used in an encoder (as
+                opposed to a text backbone).
         Returns:
             The backend for the compilation config.
         """
@@ -943,9 +951,7 @@ class CompilationConfig:
 
         from vllm.compilation.backends import VllmBackend
 
-        # TODO[@lucaskabela]: See if we can forward prefix
-        # https://github.com/vllm-project/vllm/issues/27045
-        return VllmBackend(vllm_config)
+        return VllmBackend(vllm_config, prefix=prefix, is_encoder=is_encoder)
 
     def post_init_cudagraph_sizes(self) -> None:
         """To complete the initialization after cudagraph related

vllm/model_executor/models/lfm2_siglip2.py

@@ -272,6 +272,7 @@ class Siglip2MLP(nn.Module):
 @support_torch_compile(
     dynamic_arg_dims={"hidden_states": [0, 1], "cu_seqlens": 0},
     enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
 )
 class Siglip2EncoderLayer(nn.Module):
     def __init__(
@@ -395,10 +396,6 @@ class Siglip2VisionTransformer(nn.Module):
         embed_dim = config.hidden_size
         self.config = config
         self.embeddings = Siglip2VisionEmbeddings(config)
-        # Keep the import local to avoid circular dependencies during model init.
-        from vllm.compilation.backends import set_model_tag
-
-        with set_model_tag("Siglip2Encoder", is_encoder=True):
         self.encoder = Siglip2Encoder(
             config,
             quant_config=quant_config,

vllm/model_executor/models/mllama4.py

@@ -453,7 +453,9 @@ class Llama4UnfoldConvolution(nn.Module):
 
 
 @support_torch_compile(
-    dynamic_arg_dims={"images_flattened": 0}, enable_if=should_torch_compile_mm_encoder
+    dynamic_arg_dims={"images_flattened": 0},
+    enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
 )
 class Llama4VisionModel(nn.Module):
     def __init__(
@@ -754,12 +756,7 @@ class Llama4ForConditionalGeneration(
         self.multimodal_config = multimodal_config
 
         with self._mark_tower_model(vllm_config, "image"):
-            from vllm.compilation.backends import set_model_tag
-
-            with (
-                set_current_vllm_config(vllm_config),
-                set_model_tag("Llama4VisionModel", is_encoder=True),
-            ):
+            with set_current_vllm_config(vllm_config):
                 self.vision_model = Llama4VisionModel(
                     config=config.vision_config,
                     quant_config=None,

vllm/model_executor/models/qwen2_5_vl.py

@@ -427,6 +427,7 @@ class Qwen2_5_VisionAttention(nn.Module):
         "rotary_pos_emb_sin": 0,
     },
     enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
 )
 class Qwen2_5_VisionBlock(nn.Module):
     def __init__(
@@ -486,6 +487,7 @@ class Qwen2_5_VisionBlock(nn.Module):
         "x": 0,
     },
     enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
 )
 class Qwen2_5_VisionPatchEmbed(nn.Module):
     def __init__(
@@ -521,6 +523,7 @@ class Qwen2_5_VisionPatchEmbed(nn.Module):
         "x": 0,
     },
     enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
 )
 class Qwen2_5_VisionPatchMerger(nn.Module):
     def __init__(
@@ -592,12 +595,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
         self.spatial_merge_size = vision_config.spatial_merge_size
         self.fullatt_block_indexes = vision_config.fullatt_block_indexes
         self.spatial_merge_unit = self.spatial_merge_size**2
-        # TODO[@lucaskabela]: Investigate fixing this usage
-        # see https://github.com/vllm-project/vllm/issues/27044
-        # DO NOT MOVE THIS IMPORT
-        from vllm.compilation.backends import set_model_tag
-
-        with set_model_tag("Qwen2_5_VisionPatchEmbed", is_encoder=True):
         self.patch_embed = Qwen2_5_VisionPatchEmbed(
             patch_size=patch_size,
             temporal_patch_size=temporal_patch_size,
@@ -619,7 +616,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
             dtype=torch.get_default_dtype(),
         )
 
-        with set_model_tag("Qwen2_5_VisionBlock", is_encoder=True):
         self.blocks = nn.ModuleList(
             [
                 Qwen2_5_VisionBlock(
@@ -635,7 +631,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
             ]
         )
 
-        with set_model_tag("Qwen2_5_VisionPatchMerger", is_encoder=True):
         self.merger = Qwen2_5_VisionPatchMerger(
             d_model=vision_config.out_hidden_size,
             context_dim=self.hidden_size,