[MM][Mistral] Add compile support for Pixtral vision encoders (#57168)

Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/compile/fullgraph/test_multimodal_compile.py

@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from functools import partial
+
 import pytest
 
+from tests.models.utils import dummy_hf_overrides
 from vllm.compilation.counter import compilation_counter
 from vllm.config import VllmConfig
-from vllm.config.compilation import CompilationMode
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.platforms import current_platform
 
 
@@ -14,6 +17,77 @@ def test_compile():
     assert not vllm_config.compilation_config.compile_mm_encoder
 
 
+@pytest.mark.forked
+@pytest.mark.core_model
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
+@pytest.mark.parametrize(
+    ("model", "tokenizer", "config_format", "tokenizer_mode", "model_arch"),
+    [
+        pytest.param(
+            "mistralai/Pixtral-12B-2409",
+            None,
+            "mistral",
+            "mistral",
+            "PixtralForConditionalGeneration",
+            id="pixtral",
+        ),
+        pytest.param(
+            "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
+            "mistral-community/pixtral-12b",
+            "hf",
+            "auto",
+            "Mistral3ForConditionalGeneration",
+            id="mistral3-hf",
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    ("compile_mm_encoder", "expected_models"),
+    [
+        # The dummy model has two compiled text components. Enabling encoder
+        # compilation adds its single vision transformer block.
+        pytest.param(False, 2, id="disabled"),
+        pytest.param(True, 3, id="enabled"),
+    ],
+)
+def test_pixtral_compilation(
+    vllm_runner,
+    monkeypatch,
+    model: str,
+    tokenizer: str | None,
+    config_format: str,
+    tokenizer_mode: str,
+    model_arch: str,
+    compile_mm_encoder: bool,
+    expected_models: int,
+):
+    """Test compilation for both Pixtral vision encoder implementations."""
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
+
+    with (
+        compilation_counter.expect(num_models_seen=expected_models),
+        vllm_runner(
+            model,
+            tokenizer_name=tokenizer,
+            tokenizer_mode=tokenizer_mode,
+            config_format=config_format,
+            load_format="dummy",
+            hf_overrides=partial(dummy_hf_overrides, model_arch=model_arch),
+            max_model_len=4096,
+            limit_mm_per_prompt={"image": 1},
+            gpu_memory_utilization=0.8,
+            attention_backend="FLASH_ATTN",
+            compilation_config={
+                "mode": CompilationMode.VLLM_COMPILE,
+                "compile_mm_encoder": compile_mm_encoder,
+                "cudagraph_mode": CUDAGraphMode.NONE,
+            },
+        ) as _,
+    ):
+        pass
+
+
 # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
 @pytest.mark.forked
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")

vllm/model_executor/models/pixtral.py

@@ -21,6 +21,10 @@ from transformers.models.pixtral.image_processing_pixtral import (
 )
 from transformers.models.pixtral.modeling_pixtral import apply_rotary_pos_emb
 
+from vllm.compilation.decorators import (
+    should_torch_compile_mm_encoder,
+    support_torch_compile,
+)
 from vllm.config import VllmConfig
 from vllm.config.multimodal import MultiModalDummyOptions
 from vllm.distributed import divide, get_tensor_model_parallel_world_size
@@ -829,6 +833,16 @@ class Attention(nn.Module):
         return out
 
 
+@support_torch_compile(
+    dynamic_arg_dims={
+        "x": 1,
+        "freqs_cis": 0,
+        "cu_seqlens": 0,
+        "sequence_lengths": 0,
+    },
+    enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
+)
 class TransformerBlock(nn.Module):
     def __init__(
         self,
@@ -1298,7 +1312,7 @@ class PixtralHFAttention(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        position_embeddings: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.Tensor,
         max_seqlen: torch.Tensor,
         sequence_lengths: torch.Tensor | None,
@@ -1329,6 +1343,17 @@ class PixtralHFAttention(nn.Module):
         return attn_output, None
 
 
+@support_torch_compile(
+    dynamic_arg_dims={
+        "hidden_states": 1,
+        "rotary_pos_emb_cos": 0,
+        "rotary_pos_emb_sin": 0,
+        "cu_seqlens": 0,
+        "sequence_lengths": 0,
+    },
+    enable_if=should_torch_compile_mm_encoder,
+    is_encoder=True,
+)
 class PixtralHFTransformerBlock(nn.Module):
     def __init__(
         self,
@@ -1355,14 +1380,15 @@ class PixtralHFTransformerBlock(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        position_embeddings: torch.Tensor,
+        rotary_pos_emb_cos: torch.Tensor,
+        rotary_pos_emb_sin: torch.Tensor,
         cu_seqlens: torch.Tensor,
         max_seqlen: torch.Tensor,
         sequence_lengths: torch.Tensor | None,
     ) -> torch.Tensor:
         r, _ = self.attention.forward(
             self.attention_norm(hidden_states),
-            position_embeddings=position_embeddings,
+            position_embeddings=(rotary_pos_emb_cos, rotary_pos_emb_sin),
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
             sequence_lengths=sequence_lengths,
@@ -1403,18 +1429,20 @@ class PixtralHFTransformer(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        position_embeddings: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
         cu_seqlens: torch.Tensor,
         max_seqlen: torch.Tensor,
         sequence_lengths: torch.Tensor | None,
         return_all_hidden_states: bool,
     ) -> torch.Tensor:
         hidden_states_pool = [x]
+        rotary_pos_emb_cos, rotary_pos_emb_sin = position_embeddings
 
         for layer in self.layers:
             x = layer(
                 x,
-                position_embeddings,
+                rotary_pos_emb_cos,
+                rotary_pos_emb_sin,
                 cu_seqlens,
                 max_seqlen,
                 sequence_lengths,