[MM][Perf][CG] Support ViT full cudagraphs for mllama4

docs/design/cuda_graphs_multimodal.md

@@ -86,6 +86,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 
 | Architecture | Models | CG for Image | CG for Video |
 | ------------ | ------ | ------------ | ------------ |
+| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - |
 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
@@ -116,6 +117,14 @@ vllm serve Qwen/Qwen3-VL-32B \
   --compilation-config '{"cudagraph_mm_encoder": true}'
 ```
 
+For `Llama 4` (image only):
+
+```bash
+vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct \
+  --limit-mm-per-prompt '{"image": 1}' \
+  --compilation-config '{"cudagraph_mm_encoder": true}'
+```
+
 With explicit budgets:
 
 ```bash

examples/generate/multimodal/vision_language_offline.py

@@ -2554,6 +2554,7 @@ def run_tarsier2(questions: list[str], modality: str) -> ModelRequestData:
 
 
 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
+    "llama4",
     "qwen2_5_vl",
     "qwen3_vl",
     "qwen3_vl_moe",

tests/models/multimodal/generation/test_common.py

@@ -697,7 +697,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model):
         dtype="bfloat16",
         auto_cls=AutoModelForImageTextToText,
         tensor_parallel_size=4,
-        marks=multi_gpu_marks(num_gpus=4),
+        marks=[pytest.mark.core_model, *multi_gpu_marks(num_gpus=4)],
     ),
     "llava_next": VLMTestInfo(
         models=["llava-hf/llava-v1.6-mistral-7b-hf"],

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -10,7 +10,8 @@
 from vllm.platforms import current_platform
 
 from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS
-from ...utils import dummy_hf_overrides
+
+from ....utils import create_new_process_for_each_test, multi_gpu_marks, dummy_hf_overrides
 from .vlm_utils.builders import sample_frames_with_video_metadata
 
 
@@ -51,6 +52,23 @@ def step3_vl_chat_template(content: str) -> str:
 
 
 MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
+    "llama4": VitCudagraphTestConfig(
+        model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
+        modalities=["image"],
+        image_prompt=(
+            "<|begin_of_text|><|header_start|>user<|header_end|>\n\n"
+            "<|image|>What is in this image?<|eot|>"
+            "<|header_start|>assistant<|header_end|>\n\n"
+        ),
+        max_model_len=4096,
+        max_tokens=32,
+        max_num_seqs=2,
+        vllm_runner_kwargs={
+            "distributed_executor_backend": "mp",
+            "tensor_parallel_size": 4,
+        },
+        marks=[pytest.mark.core_model, *multi_gpu_marks(num_gpus=4)],
+    ),
     "qwen2_5_vl": VitCudagraphTestConfig(
         model="Qwen/Qwen2.5-VL-3B-Instruct",
         image_prompt=qwen_vl_chat_template(

tests/models/multimodal/processing/test_mllama4.py

@@ -2,10 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for mllama's multimodal preprocessing and profiling."""
 
+from types import SimpleNamespace
+
 import pytest
+import torch
 from torch import prod
 from transformers import Llama4Config
 
+from vllm.model_executor.models.mllama4 import Llama4ForConditionalGeneration
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
 from ...utils import build_model_context
@@ -53,3 +57,111 @@ def test_profiling(model_id: str, max_model_len: int):
     assert total_tokens == sum(
         placeholder.length for placeholder in mm_inputs["mm_placeholders"]["image"]
     )
+
+
+class VisionModelStub(torch.nn.Module):
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        return pixel_values.unsqueeze(-1)
+
+
+class ProjectorStub(torch.nn.Module):
+    def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor:
+        return vision_embeddings + 1
+
+
+def make_test_model() -> Llama4ForConditionalGeneration:
+    model = object.__new__(Llama4ForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(
+        vision_config=SimpleNamespace(
+            image_size=8,
+            patch_size=4,
+            pixel_shuffle_ratio=1.0,
+            num_channels=3,
+        ),
+        text_config=SimpleNamespace(hidden_size=1),
+    )
+    model.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=20),
+    )
+    model.vision_model = VisionModelStub()
+    model.multi_modal_projector = ProjectorStub()
+    model.use_data_parallel = False
+    return model
+
+
+def test_encoder_cudagraph_metadata():
+    model = make_test_model()
+    patches_per_chunk = model.get_image_patches_per_chunk()
+    mm_kwargs = {
+        "patches_per_image": torch.tensor([2, 3], dtype=torch.int32),
+    }
+    vllm_config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32),
+    )
+
+    config = model.get_encoder_cudagraph_config()
+
+    assert config.modalities == ["image"]
+    assert config.input_key_by_modality == {"image": "pixel_values"}
+    assert config.buffer_keys == []
+    assert config.out_hidden_size == 1
+    assert patches_per_chunk == 4
+    assert model.get_encoder_cudagraph_budget_range(vllm_config) == (
+        patches_per_chunk,
+        20,
+    )
+    assert model.get_encoder_cudagraph_num_items(mm_kwargs) == 2
+    assert model.get_encoder_cudagraph_per_item_input_sizes(mm_kwargs) == [2, 3]
+    assert model.get_encoder_cudagraph_per_item_output_tokens(mm_kwargs) == [
+        2 * patches_per_chunk,
+        3 * patches_per_chunk,
+    ]
+
+
+def test_select_encoder_cudagraph_items():
+    model = make_test_model()
+    mm_kwargs = {
+        "pixel_values": torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0]]),
+        "patches_per_image": torch.tensor([2, 3], dtype=torch.int32),
+    }
+
+    selected = model.select_encoder_cudagraph_items(mm_kwargs, [1])
+    empty = model.select_encoder_cudagraph_items(mm_kwargs, [])
+
+    assert torch.equal(selected["pixel_values"], torch.tensor([[2.0], [3.0], [4.0]]))
+    assert torch.equal(
+        selected["patches_per_image"], torch.tensor([3], dtype=torch.int32)
+    )
+    assert empty["pixel_values"].shape == (0, 1)
+    assert empty["patches_per_image"].shape == (0,)
+
+
+def test_prepare_encoder_cudagraph_capture_inputs_rounds_up():
+    model = make_test_model()
+    patches_per_chunk = model.get_image_patches_per_chunk()
+
+    capture_inputs = model.prepare_encoder_cudagraph_capture_inputs(
+        token_budget=patches_per_chunk + 1,
+        max_batch_size=2,
+        max_frames_per_batch=0,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+    )
+
+    assert capture_inputs.mm_kwargs["pixel_values"].shape == (2, 3, 8, 8)
+    assert capture_inputs.buffers == {}
+
+
+def test_encoder_cudagraph_forward_matches_eager():
+    model = make_test_model()
+    mm_kwargs = {
+        "pixel_values": torch.tensor([[1.0], [2.0]]),
+    }
+
+    eager = model.encoder_eager_forward(mm_kwargs)
+    cg = model.encoder_cudagraph_forward(mm_kwargs, buffers={})
+
+    expected = torch.tensor([[2.0], [3.0]])
+    assert torch.equal(eager, expected)
+    assert torch.equal(cg, expected)