@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7769b58307

Fork 0

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 6, 09:30 UTC

[torch.compile][BE][Multimodal] Remove requirement to set_model_tag to avoid cache conflict (#37345)

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

Lucas Kabela committed Mar 19, 2026 at 10:26 UTC 7769b58307c9604ac833ba790f511cea3989c0e6

7 files changed +86 -69

docs/design/torch_compile_multimodal.md

+5 -6

@@ -29,10 +29,9 @@ To compile a multimodal component such as an encoder, we follow the same mechani

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

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

33 -relies on caching artifacts to reduce start time, we must properly propagate the `<component_name>` information to the cache in order to avoid collisions

34 -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

35 -components (see Compile Range Integration).

32 +2. The `@support_torch_compile` decorator should include `is_encoder=True` for encoder components. This is needed for compile range integration

33 +(see Compile Range Integration). The decorator automatically uses the class name as the cache directory prefix, avoiding collisions between

34 +independently compiled sub-modules (e.g. vision encoder components vs the text backbone).

36 ### CompilationConfig

@@ -57,8 +56,8 @@ tradeoff

56 ### Compile ranges

58 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

60 -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`

61 -to alert torch.compile to the fact that this range cannot be inferred, and we default to the range (1, MAX_INT).

59 +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

60 +`@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).

62 !!! note

63 We may seek to tighten this range for better performance in the future

vllm/compilation/decorators.py

+14 -3

@@ -118,6 +118,7 @@ def support_torch_compile(

118 dynamic_arg_dims: dict[str, int | list[int]] | None = None,

119 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

120 enable_if: Callable[[VllmConfig], bool] | None = None,

121 + is_encoder: bool = False,

122 shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,

123 ) -> Callable[[type[_T]], type[_T]] | type[_T]:

124 """

@@ -177,6 +178,11 @@ def support_torch_compile(

178 enforce that dynamo does not specialize on 0/1 values in the case of dummy input

179 such as for vision model compilation

180

181 + `is_encoder` marks this module as a portion of an multimodal encoder.

182 + When True, the compile range upper bound is set to MAX_INT32 instead of

183 + max_num_batched_tokens, since encoder input shapes are unpredictable.

184 + This is typically used for vision encoder sub-modules in multimodal models.

185 +

186 `shape_invariants` is a function that gets compiled right before forward.

187 The function should have the torch._check calls that are needed to set

188 the relationships between different input sizes. For example:

@@ -226,6 +232,7 @@ def support_torch_compile(

232 inferred_dynamic_arg_dims,

233 mark_unbacked_dims,

234 enable_if,

235 + is_encoder,

236 shape_invariants,

237 )

238

@@ -316,6 +323,7 @@ def _support_torch_compile(

323 dynamic_arg_dims: dict[str, int | list[int]],

324 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

325 enable_if: Callable[[VllmConfig], bool] | None = None,

326 + is_encoder: bool = False,

327 shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,

328 ) -> type[_T]:

329 """

@@ -345,8 +353,7 @@ def _support_torch_compile(

353 vllm_config = get_current_vllm_config()

354

355 # NOTE: to support multimodal models (such as encoder),

348 - # we may not have vllm_config so we may need to patch

349 - # it

356 + # we may not have vllm_config so we may need to patch it

357 sig = inspect.signature(old_init)

358 if "vllm_config" in sig.parameters:

359 kwargs["vllm_config"] = vllm_config

@@ -374,7 +381,11 @@ def _support_torch_compile(

381 self.compiled = False

382

383 # Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class

377 - TorchCompileWithNoGuardsWrapper.__init__(self)

384 + TorchCompileWithNoGuardsWrapper.__init__(

385 + self,

386 + compile_prefix=cls.__name__ if is_encoder else "",

387 + is_encoder=is_encoder,

388 + )

389

390 cls.__init__ = __init__

391

vllm/compilation/wrapper.py

+15 -3

@@ -75,8 +75,14 @@ class TorchCompileWithNoGuardsWrapper:

75 return ctx.result

76 return callable_fn(*args, **kwargs)

78 - def __init__(self) -> None:

78 + def __init__(

79 + self,

80 + compile_prefix: str = "",

81 + is_encoder: bool = False,

82 + ) -> None:

83 self.compiled = False

84 + self._compile_prefix = compile_prefix

85 + self._is_encoder = is_encoder

87 vllm_config = get_current_vllm_config()

88 self.vllm_config = vllm_config

@@ -87,7 +93,9 @@ class TorchCompileWithNoGuardsWrapper:

93 if mode is None:

94 raise RuntimeError("Compilation mode cannot be NO_COMPILATION")

90 - backend = vllm_config.compilation_config.init_backend(vllm_config)

96 + backend = vllm_config.compilation_config.init_backend(

97 + vllm_config, prefix=compile_prefix, is_encoder=is_encoder

98 + )

99 options = {}

100

101 if isinstance(backend, str) and backend == "inductor":

@@ -332,4 +340,8 @@ def reset_compile_wrapper(model: torch.nn.Module) -> None:

340 compilation_config.local_cache_dir = ""

341

342 model.__class__.forward.__code__ = model.original_code_object()

335 - TorchCompileWithNoGuardsWrapper.__init__(model)

343 + TorchCompileWithNoGuardsWrapper.__init__(

344 + model,

345 + compile_prefix=model._compile_prefix,

346 + is_encoder=model._is_encoder,

347 + )

vllm/config/compilation.py

+10 -4

@@ -909,11 +909,19 @@ class CompilationConfig:

909 if self.backend == "":

910 self.backend = current_platform.get_compile_backend()

911

912 - def init_backend(self, vllm_config: "VllmConfig") -> str | Callable:

912 + def init_backend(

913 + self,

914 + vllm_config: "VllmConfig",

915 + prefix: str = "",

916 + is_encoder: bool = False,

917 + ) -> str | Callable:

918 """

919 Initialize the backend for the compilation config from a vllm config.

920 Arguments:

921 vllm_config: The vllm config to initialize the backend from.

922 + prefix: Cache directory prefix for this compiled module.

923 + is_encoder: Whether this module is used in an encoder (as

924 + opposed to a text backbone).

925 Returns:

926 The backend for the compilation config.

927 """

@@ -943,9 +951,7 @@ class CompilationConfig:

951

952 from vllm.compilation.backends import VllmBackend

953

946 - # TODO[@lucaskabela]: See if we can forward prefix

947 - # https://github.com/vllm-project/vllm/issues/27045

948 - return VllmBackend(vllm_config)

954 + return VllmBackend(vllm_config, prefix=prefix, is_encoder=is_encoder)

955

956 def post_init_cudagraph_sizes(self) -> None:

957 """To complete the initialization after cudagraph related

vllm/model_executor/models/lfm2_siglip2.py

+7 -10

@@ -272,6 +272,7 @@ class Siglip2MLP(nn.Module):

272 @support_torch_compile(

273 dynamic_arg_dims={"hidden_states": [0, 1], "cu_seqlens": 0},

274 enable_if=should_torch_compile_mm_encoder,

275 + is_encoder=True,

276 )

277 class Siglip2EncoderLayer(nn.Module):

278 def __init__(

@@ -395,16 +396,12 @@ class Siglip2VisionTransformer(nn.Module):

396 embed_dim = config.hidden_size

397 self.config = config

398 self.embeddings = Siglip2VisionEmbeddings(config)

398 - # Keep the import local to avoid circular dependencies during model init.

399 - from vllm.compilation.backends import set_model_tag

400 -

401 - with set_model_tag("Siglip2Encoder", is_encoder=True):

402 - self.encoder = Siglip2Encoder(

403 - config,

404 - quant_config=quant_config,

405 - num_hidden_layers_override=num_hidden_layers_override,

406 - prefix=f"{prefix}.encoder",

407 - )

399 + self.encoder = Siglip2Encoder(

400 + config,

401 + quant_config=quant_config,

402 + num_hidden_layers_override=num_hidden_layers_override,

403 + prefix=f"{prefix}.encoder",

404 + )

405 num_hidden_layers = config.num_hidden_layers

406 if len(self.encoder.layers) > config.num_hidden_layers:

407 raise ValueError(

vllm/model_executor/models/mllama4.py

+4 -7

@@ -453,7 +453,9 @@ class Llama4UnfoldConvolution(nn.Module):

453

454

455 @support_torch_compile(

456 - dynamic_arg_dims={"images_flattened": 0}, enable_if=should_torch_compile_mm_encoder

456 + dynamic_arg_dims={"images_flattened": 0},

457 + enable_if=should_torch_compile_mm_encoder,

458 + is_encoder=True,

459 )

460 class Llama4VisionModel(nn.Module):

461 def __init__(

@@ -754,12 +756,7 @@ class Llama4ForConditionalGeneration(

756 self.multimodal_config = multimodal_config

757

758 with self._mark_tower_model(vllm_config, "image"):

757 - from vllm.compilation.backends import set_model_tag

758 -

759 - with (

760 - set_current_vllm_config(vllm_config),

761 - set_model_tag("Llama4VisionModel", is_encoder=True),

762 - ):

759 + with set_current_vllm_config(vllm_config):

760 self.vision_model = Llama4VisionModel(

761 config=config.vision_config,

762 quant_config=None,

vllm/model_executor/models/qwen2_5_vl.py

+31 -36

@@ -427,6 +427,7 @@ class Qwen2_5_VisionAttention(nn.Module):

427 "rotary_pos_emb_sin": 0,

428 },

429 enable_if=should_torch_compile_mm_encoder,

430 + is_encoder=True,

431 )

432 class Qwen2_5_VisionBlock(nn.Module):

433 def __init__(

@@ -486,6 +487,7 @@ class Qwen2_5_VisionBlock(nn.Module):

487 "x": 0,

488 },

489 enable_if=should_torch_compile_mm_encoder,

490 + is_encoder=True,

491 )

492 class Qwen2_5_VisionPatchEmbed(nn.Module):

493 def __init__(

@@ -521,6 +523,7 @@ class Qwen2_5_VisionPatchEmbed(nn.Module):

523 "x": 0,

524 },

525 enable_if=should_torch_compile_mm_encoder,

526 + is_encoder=True,

527 )

528 class Qwen2_5_VisionPatchMerger(nn.Module):

529 def __init__(

@@ -592,18 +595,12 @@ class Qwen2_5_VisionTransformer(nn.Module):

595 self.spatial_merge_size = vision_config.spatial_merge_size

596 self.fullatt_block_indexes = vision_config.fullatt_block_indexes

597 self.spatial_merge_unit = self.spatial_merge_size**2

595 - # TODO[@lucaskabela]: Investigate fixing this usage

596 - # see https://github.com/vllm-project/vllm/issues/27044

597 - # DO NOT MOVE THIS IMPORT

598 - from vllm.compilation.backends import set_model_tag

599 -

600 - with set_model_tag("Qwen2_5_VisionPatchEmbed", is_encoder=True):

601 - self.patch_embed = Qwen2_5_VisionPatchEmbed(

602 - patch_size=patch_size,

603 - temporal_patch_size=temporal_patch_size,

604 - in_channels=in_channels,

605 - hidden_size=self.hidden_size,

606 - )

598 + self.patch_embed = Qwen2_5_VisionPatchEmbed(

599 + patch_size=patch_size,

600 + temporal_patch_size=temporal_patch_size,

601 + in_channels=in_channels,

602 + hidden_size=self.hidden_size,

603 + )

604

605 norm_layer = partial(RMSNorm, eps=norm_eps)

606 head_dim = self.hidden_size // self.num_heads

@@ -619,31 +616,29 @@ class Qwen2_5_VisionTransformer(nn.Module):

616 dtype=torch.get_default_dtype(),

617 )

618

622 - with set_model_tag("Qwen2_5_VisionBlock", is_encoder=True):

623 - self.blocks = nn.ModuleList(

624 - [

625 - Qwen2_5_VisionBlock(

626 - dim=self.hidden_size,

627 - num_heads=self.num_heads,

628 - mlp_hidden_dim=vision_config.intermediate_size,

629 - act_fn=get_act_and_mul_fn(vision_config.hidden_act),

630 - norm_layer=norm_layer,

631 - quant_config=quant_config,

632 - prefix=f"{prefix}.blocks.{layer_idx}",

633 - )

634 - for layer_idx in range(depth)

635 - ]

636 - )

619 + self.blocks = nn.ModuleList(

620 + [

621 + Qwen2_5_VisionBlock(

622 + dim=self.hidden_size,

623 + num_heads=self.num_heads,

624 + mlp_hidden_dim=vision_config.intermediate_size,

625 + act_fn=get_act_and_mul_fn(vision_config.hidden_act),

626 + norm_layer=norm_layer,

627 + quant_config=quant_config,

628 + prefix=f"{prefix}.blocks.{layer_idx}",

629 + )

630 + for layer_idx in range(depth)

631 + ]

632 + )

633

638 - with set_model_tag("Qwen2_5_VisionPatchMerger", is_encoder=True):

639 - self.merger = Qwen2_5_VisionPatchMerger(

640 - d_model=vision_config.out_hidden_size,

641 - context_dim=self.hidden_size,

642 - norm_layer=norm_layer,

643 - spatial_merge_size=self.spatial_merge_size,

644 - quant_config=quant_config,

645 - prefix=f"{prefix}.merger",

646 - )

634 + self.merger = Qwen2_5_VisionPatchMerger(

635 + d_model=vision_config.out_hidden_size,

636 + context_dim=self.hidden_size,

637 + norm_layer=norm_layer,

638 + spatial_merge_size=self.spatial_merge_size,

639 + quant_config=quant_config,

640 + prefix=f"{prefix}.merger",

641 + )

642

643 @property

644 def dtype(self) -> torch.dtype: