[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>
Jakub Zakrzewski committed
Sep 30, 2026 at 14:27 UTC
effdea9a5303f02759ba243ee624bf0a9b2576b0
2 files changed
+108
-6
tests/compile/fullgraph/test_multimodal_compile.py
+75
-1
@@ -1,10 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from functools import partial
4
+
5
import pytest
6
7
+from tests.models.utils import dummy_hf_overrides
8
from vllm.compilation.counter import compilation_counter
9
from vllm.config import VllmConfig
7
-from vllm.config.compilation import CompilationMode
10
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
11
from vllm.platforms import current_platform
12
13
@@ -14,6 +17,77 @@ def test_compile():
17
assert not vllm_config.compilation_config.compile_mm_encoder
18
19
20
+@pytest.mark.forked
21
+@pytest.mark.core_model
22
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
23
+@pytest.mark.parametrize(
24
+ ("model", "tokenizer", "config_format", "tokenizer_mode", "model_arch"),
25
+ [
26
+ pytest.param(
27
+ "mistralai/Pixtral-12B-2409",
28
+ None,
29
+ "mistral",
30
+ "mistral",
31
+ "PixtralForConditionalGeneration",
32
+ id="pixtral",
33
+ ),
34
+ pytest.param(
35
+ "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
36
+ "mistral-community/pixtral-12b",
37
+ "hf",
38
+ "auto",
39
+ "Mistral3ForConditionalGeneration",
40
+ id="mistral3-hf",
41
+ ),
42
+ ],
43
+)
44
+@pytest.mark.parametrize(
45
+ ("compile_mm_encoder", "expected_models"),
46
+ [
47
+ # The dummy model has two compiled text components. Enabling encoder
48
+ # compilation adds its single vision transformer block.
49
+ pytest.param(False, 2, id="disabled"),
50
+ pytest.param(True, 3, id="enabled"),
51
+ ],
52
+)
53
+def test_pixtral_compilation(
54
+ vllm_runner,
55
+ monkeypatch,
56
+ model: str,
57
+ tokenizer: str | None,
58
+ config_format: str,
59
+ tokenizer_mode: str,
60
+ model_arch: str,
61
+ compile_mm_encoder: bool,
62
+ expected_models: int,
63
+):
64
+ """Test compilation for both Pixtral vision encoder implementations."""
65
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
66
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
67
+
68
+ with (
69
+ compilation_counter.expect(num_models_seen=expected_models),
70
+ vllm_runner(
71
+ model,
72
+ tokenizer_name=tokenizer,
73
+ tokenizer_mode=tokenizer_mode,
74
+ config_format=config_format,
75
+ load_format="dummy",
76
+ hf_overrides=partial(dummy_hf_overrides, model_arch=model_arch),
77
+ max_model_len=4096,
78
+ limit_mm_per_prompt={"image": 1},
79
+ gpu_memory_utilization=0.8,
80
+ attention_backend="FLASH_ATTN",
81
+ compilation_config={
82
+ "mode": CompilationMode.VLLM_COMPILE,
83
+ "compile_mm_encoder": compile_mm_encoder,
84
+ "cudagraph_mode": CUDAGraphMode.NONE,
85
+ },
86
+ ) as _,
87
+ ):
88
+ pass
89
+
90
+
91
# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
92
@pytest.mark.forked
93
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
vllm/model_executor/models/pixtral.py
+33
-5
@@ -21,6 +21,10 @@ from transformers.models.pixtral.image_processing_pixtral import (
21
)
22
from transformers.models.pixtral.modeling_pixtral import apply_rotary_pos_emb
23
24
+from vllm.compilation.decorators import (
25
+ should_torch_compile_mm_encoder,
26
+ support_torch_compile,
27
+)
28
from vllm.config import VllmConfig
29
from vllm.config.multimodal import MultiModalDummyOptions
30
from vllm.distributed import divide, get_tensor_model_parallel_world_size
@@ -829,6 +833,16 @@ class Attention(nn.Module):
833
return out
834
835
836
+@support_torch_compile(
837
+ dynamic_arg_dims={
838
+ "x": 1,
839
+ "freqs_cis": 0,
840
+ "cu_seqlens": 0,
841
+ "sequence_lengths": 0,
842
+ },
843
+ enable_if=should_torch_compile_mm_encoder,
844
+ is_encoder=True,
845
+)
846
class TransformerBlock(nn.Module):
847
def __init__(
848
self,
@@ -1298,7 +1312,7 @@ class PixtralHFAttention(nn.Module):
1312
def forward(
1313
self,
1314
hidden_states: torch.Tensor,
1301
- position_embeddings: torch.Tensor,
1315
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
1316
cu_seqlens: torch.Tensor,
1317
max_seqlen: torch.Tensor,
1318
sequence_lengths: torch.Tensor | None,
@@ -1329,6 +1343,17 @@ class PixtralHFAttention(nn.Module):
1343
return attn_output, None
1344
1345
1346
+@support_torch_compile(
1347
+ dynamic_arg_dims={
1348
+ "hidden_states": 1,
1349
+ "rotary_pos_emb_cos": 0,
1350
+ "rotary_pos_emb_sin": 0,
1351
+ "cu_seqlens": 0,
1352
+ "sequence_lengths": 0,
1353
+ },
1354
+ enable_if=should_torch_compile_mm_encoder,
1355
+ is_encoder=True,
1356
+)
1357
class PixtralHFTransformerBlock(nn.Module):
1358
def __init__(
1359
self,
@@ -1355,14 +1380,15 @@ class PixtralHFTransformerBlock(nn.Module):
1380
def forward(
1381
self,
1382
hidden_states: torch.Tensor,
1358
- position_embeddings: torch.Tensor,
1383
+ rotary_pos_emb_cos: torch.Tensor,
1384
+ rotary_pos_emb_sin: torch.Tensor,
1385
cu_seqlens: torch.Tensor,
1386
max_seqlen: torch.Tensor,
1387
sequence_lengths: torch.Tensor | None,
1388
) -> torch.Tensor:
1389
r, _ = self.attention.forward(
1390
self.attention_norm(hidden_states),
1365
- position_embeddings=position_embeddings,
1391
+ position_embeddings=(rotary_pos_emb_cos, rotary_pos_emb_sin),
1392
cu_seqlens=cu_seqlens,
1393
max_seqlen=max_seqlen,
1394
sequence_lengths=sequence_lengths,
@@ -1403,18 +1429,20 @@ class PixtralHFTransformer(nn.Module):
1429
def forward(
1430
self,
1431
x: torch.Tensor,
1406
- position_embeddings: torch.Tensor,
1432
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
1433
cu_seqlens: torch.Tensor,
1434
max_seqlen: torch.Tensor,
1435
sequence_lengths: torch.Tensor | None,
1436
return_all_hidden_states: bool,
1437
) -> torch.Tensor:
1438
hidden_states_pool = [x]
1439
+ rotary_pos_emb_cos, rotary_pos_emb_sin = position_embeddings
1440
1441
for layer in self.layers:
1442
x = layer(
1443
x,
1417
- position_embeddings,
1444
+ rotary_pos_emb_cos,
1445
+ rotary_pos_emb_sin,
1446
cu_seqlens,
1447
max_seqlen,
1448
sequence_lengths,