diff --git a/tests/compile/fullgraph/test_multimodal_compile.py b/tests/compile/fullgraph/test_multimodal_compile.py index c5dc6f96b2a5..8f76a18c69c3 100644 --- a/tests/compile/fullgraph/test_multimodal_compile.py +++ b/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") diff --git a/vllm/model_executor/models/pixtral.py b/vllm/model_executor/models/pixtral.py index 521df62a9217..fb10026ad19a 100644 --- a/vllm/model_executor/models/pixtral.py +++ b/vllm/model_executor/models/pixtral.py @@ -21,6 +21,10 @@ ) 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 @@ def forward( 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 @@ def __init__( 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 @@ def forward( 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 @@ def __init__( 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 @@ def __init__( 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,