Fix to VisionEncoder instead of any one layer

Signed-off-by: Lucas Kabela <[email protected]>

Author: Lucaskabela

vllm/model_executor/models/mllama4.py

@@ -197,7 +197,7 @@ def pixel_shuffle(input_tensor, shuffle_ratio):
     return output_tensor
 
 
-@support_torch_compile
+@support_torch_compile(dynamic_arg_dims={"encoded_patches": 0})
 class Llama4VisionPixelShuffleMLP(nn.Module):
     def __init__(
         self,
@@ -228,7 +228,6 @@ def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:
         return self.mlp(encoded_patches)
 
 
-@support_torch_compile
 class Llama4VisionAttention(nn.Module):
     def __init__(
         self,
@@ -323,7 +322,6 @@ def forward(
         return attn_output
 
 
-@support_torch_compile
 class Llama4VisionEncoderLayer(nn.Module):
     def __init__(
         self,
@@ -377,6 +375,7 @@ def forward(
         return outputs
 
 
+@support_torch_compile(dynamic_arg_dims={"hidden_states": 0})
 class Llama4VisionEncoder(nn.Module):
     def __init__(
         self,
@@ -387,20 +386,17 @@ def __init__(
     ):
         super().__init__()
         self.config = config
-        from vllm.compilation.backends import set_model_tag
-
-        with set_model_tag("Llama4VisionEncoderLayer"):
-            self.layers = nn.ModuleList(
-                [
-                    Llama4VisionEncoderLayer(
-                        config=config,
-                        quant_config=quant_config,
-                        prefix=f"{prefix}.layers.{layer_idx}",
-                        use_data_parallel=use_data_parallel,
-                    )
-                    for layer_idx in range(config.num_hidden_layers)
-                ]
-            )
+        self.layers = nn.ModuleList(
+            [
+                Llama4VisionEncoderLayer(
+                    config=config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{layer_idx}",
+                    use_data_parallel=use_data_parallel,
+                )
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
 
     def forward(
         self,
@@ -488,14 +484,16 @@ def __init__(
         self.layernorm_post = nn.LayerNorm(self.hidden_size, eps=1e-5)
 
         # encoders
-        self.model = Llama4VisionEncoder(
-            config=config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.model",
-            use_data_parallel=use_data_parallel,
-        )
         from vllm.compilation.backends import set_model_tag
 
+        with set_model_tag("Llama4VisionEncoderLayer"):
+            self.model = Llama4VisionEncoder(
+                config=config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.model",
+                use_data_parallel=use_data_parallel,
+            )
+
         with set_model_tag("Llama4VisionPixelShuffleMLP"):
             self.vision_adapter = Llama4VisionPixelShuffleMLP(
                 config=config,