[Multimodal][torch.compile] Add compilation config field for turning off ViT/MM compile (#28242)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com>

vllm/model_executor/models/qwen2_5_vl.py

@@ -67,6 +67,9 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.module_mapping import MultiModelKeys
+from vllm.model_executor.models.transformers.utils import (
+    should_torch_compile_mm_vit,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.evs import (
     compute_mrope_for_media,
@@ -464,6 +467,7 @@ class Qwen2_5_VisionAttention(nn.Module):
         "seqlens": 0,
     },
     mark_unbacked_dims={"seqlens": 0},
+    enable_if=should_torch_compile_mm_vit,
 )
 class Qwen2_5_VisionBlock(nn.Module):
     def __init__(
@@ -529,7 +533,8 @@ class Qwen2_5_VisionBlock(nn.Module):
 @support_torch_compile(
     dynamic_arg_dims={
         "x": 0,
-    }
+    },
+    enable_if=should_torch_compile_mm_vit,
 )
 class Qwen2_5_VisionPatchEmbed(nn.Module):
     def __init__(
@@ -560,7 +565,8 @@ class Qwen2_5_VisionPatchEmbed(nn.Module):
 @support_torch_compile(
     dynamic_arg_dims={
         "x": 0,
-    }
+    },
+    enable_if=should_torch_compile_mm_vit,
 )
 class Qwen2_5_VisionPatchMerger(nn.Module):
     def __init__(

vllm/model_executor/models/transformers/utils.py

@@ -205,3 +205,14 @@ def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool:
     # Dynamic rope scaling is not compatible with torch.compile
     rope_scaling: dict = getattr(text_config, "rope_scaling", None) or {}
     return rope_scaling.get("rope_type") != "dynamic"
+
+
+def should_torch_compile_mm_vit(vllm_config: "VllmConfig") -> bool:
+    """
+    Callable to be passed to `@support_torch_compile`'s `enable_if` argument.
+
+    Defaults to `True` but is disabled in the following situations:
+
+    - The model uses dynamic rope scaling.
+    """
+    return vllm_config.compilation_config.compile_mm_encoder