[MM][Perf][CG] Support ViT full cudagraphs for mllama4 (#40660)

Signed-off-by: allgather <[email protected]>
Co-authored-by: Isotr0py <[email protected]>

docs/design/cuda_graphs_multimodal.md

@@ -82,6 +82,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 
 | Architecture | Models | CG for Image | CG for Video |
 | ------------ | ------ | ------------ | ------------ |
+| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - |
 | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
@@ -114,6 +115,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

@@ -2532,6 +2532,7 @@ MODELS_NEED_VIDEO_METADATA = [
 
 
 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
+    "llama4",
     "internvl_chat",
     "qwen2_5_vl",
     "qwen3_vl",

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -55,6 +55,26 @@ 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={
+            "load_format": "dummy",
+            "hf_overrides": partial(
+                dummy_hf_overrides,
+                model_arch="Llama4ForConditionalGeneration",
+            ),
+        },
+        marks=[pytest.mark.core_model],
+    ),
     "internvl": VitCudagraphTestConfig(
         model="OpenGVLab/InternVL3-1B",
         num_video_frames=8,

tests/models/utils.py

@@ -507,12 +507,13 @@ def dummy_hf_overrides(
     # Only set MoE related config when the model has MoE layers.
     # Otherwise all models detected as MoE by _get_transformers_backend_cls.
     if model_arch_config.num_experts > 0:
+        num_experts_per_tok = 1 if model_arch == "Llama4ForConditionalGeneration" else 2
         update_dict.update(
             {
                 "num_experts": num_experts,
-                "num_experts_per_tok": 2,
+                "num_experts_per_tok": num_experts_per_tok,
                 # Kimi uses `num_experts_per_token`.
-                "num_experts_per_token": 2,
+                "num_experts_per_token": num_experts_per_tok,
                 "num_local_experts": num_experts,
                 # Otherwise there will not be any expert layers
                 "first_k_dense_replace": 0,

vllm/model_executor/models/mllama4.py

@@ -19,7 +19,7 @@
 import math
 from collections.abc import Iterable, Mapping
 from itertools import tee
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 import torch
 from torch import nn
@@ -78,6 +78,7 @@ from .interfaces import (
     MixtureOfExperts,
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsEncoderCudaGraph,
     SupportsLoRA,
     SupportsMultiModal,
     SupportsPP,
@@ -105,7 +106,7 @@ class Llama4ImagePatchInputs(TensorSchema):
 
     patches_per_image: Annotated[torch.Tensor, TensorShape("batch_size")]
     """
-    The number of total patches for each image in the batch.
+    The number of chunked image tiles for each image in the batch.
     
     This is used to split the embeddings which has the first two dimensions
     flattened just like `pixel_values`.
@@ -731,6 +732,7 @@ class Llama4ForConditionalGeneration(
     SupportsMultiModal,
     SupportsPP,
     MixtureOfExperts,
+    SupportsEncoderCudaGraph,
     SupportsEagle3,
     SupportsLoRA,
 ):
@@ -828,10 +830,161 @@ class Llama4ForConditionalGeneration(
             num_physical_experts, num_local_physical_experts
         )
 
+    def get_image_patches_per_chunk(self) -> int:
+        return Mllama4ProcessingInfo.get_patch_per_chunk(self.config.vision_config)
+
+    def encode_image_chunks(
+        self,
+        pixel_values: torch.Tensor,
+        *,
+        use_data_parallel: bool,
+    ) -> torch.Tensor:
+        if use_data_parallel:
+            vision_embeddings = run_dp_sharded_vision_model(
+                pixel_values, self.vision_model
+            )
+        else:
+            vision_embeddings = self.vision_model(pixel_values)
+
+        return self.multi_modal_projector(vision_embeddings)
+
+    def get_encoder_cudagraph_config(self):
+        from vllm.v1.worker.encoder_cudagraph_defs import (
+            EncoderCudaGraphConfig,
+        )
+
+        return EncoderCudaGraphConfig(
+            modalities=["image"],
+            buffer_keys=["pixel_values"],
+            out_hidden_size=self.config.text_config.hidden_size,
+        )
+
+    def get_input_modality(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> str:
+        return "image"
+
+    def get_encoder_cudagraph_budget_range(
+        self,
+        vllm_config: VllmConfig,
+    ) -> tuple[int, int]:
+        min_budget = self.get_image_patches_per_chunk()
+        max_budget = min(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            self.vllm_config.model_config.max_model_len,
+        )
+        return (min_budget, max_budget)
+
+    def get_encoder_cudagraph_item_specs(
+        self,
+        mm_kwargs: dict[str, Any],
+    ):
+        from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
+
+        patches_per_chunk = self.get_image_patches_per_chunk()
+        return [
+            EncoderItemSpec(
+                input_size=num_chunks,
+                output_tokens=num_chunks * patches_per_chunk,
+            )
+            for num_chunks in mm_kwargs["patches_per_image"].tolist()
+        ]
+
+    def select_encoder_cudagraph_items(
+        self,
+        mm_kwargs: dict[str, Any],
+        indices: list[int],
+    ) -> dict[str, Any]:
+        pixel_values = mm_kwargs["pixel_values"]
+        patches_per_image = mm_kwargs["patches_per_image"]
+
+        if len(indices) == 0:
+            return {
+                "pixel_values": pixel_values[:0],
+                "patches_per_image": patches_per_image[:0],
+            }
+
+        cum_chunks = [0]
+        for num_chunks in patches_per_image.tolist():
+            cum_chunks.append(cum_chunks[-1] + num_chunks)
+
+        selected_pixel_values = torch.cat(
+            [pixel_values[cum_chunks[i] : cum_chunks[i + 1]] for i in indices],
+            dim=0,
+        )
+
+        return {
+            "pixel_values": selected_pixel_values,
+            "patches_per_image": patches_per_image[indices],
+        }
+
+    def prepare_encoder_cudagraph_capture_inputs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ):
+        from vllm.v1.worker.encoder_cudagraph_defs import (
+            EncoderCudaGraphCaptureInputs,
+        )
+
+        vision_config = self.config.vision_config
+        patches_per_chunk = self.get_image_patches_per_chunk()
+        chunks_per_capture = max(
+            1, (token_budget + patches_per_chunk - 1) // patches_per_chunk
+        )
+        dummy_pixel_values = torch.randn(
+            chunks_per_capture,
+            vision_config.num_channels,
+            vision_config.image_size,
+            vision_config.image_size,
+            device=device,
+            dtype=dtype,
+        )
+
+        return EncoderCudaGraphCaptureInputs(
+            values={"pixel_values": dummy_pixel_values},
+        )
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+    ):
+        from vllm.v1.worker.encoder_cudagraph_defs import (
+            EncoderCudaGraphReplayBuffers,
+        )
+
+        return EncoderCudaGraphReplayBuffers(
+            values={"pixel_values": mm_kwargs["pixel_values"]},
+        )
+
+    def encoder_cudagraph_forward(
+        self,
+        inputs: dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        return self.encode_image_chunks(
+            inputs["pixel_values"],
+            use_data_parallel=False,
+        ).flatten(0, 1)
+
+    def encoder_eager_forward(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> torch.Tensor:
+        return self.encode_image_chunks(
+            mm_kwargs["pixel_values"],
+            use_data_parallel=False,
+        ).flatten(0, 1)
+
     def _parse_and_validate_image_input(
         self, **kwargs: object
     ) -> Llama4ImagePatchInputs | None:
-        # num_images, 1, num_chunks, channel, image_size, image_size
+        # total_num_chunks, channel, image_size, image_size
         pixel_values = kwargs.pop("pixel_values", None)
         if pixel_values is None:
             return None
@@ -853,15 +1006,10 @@ class Llama4ForConditionalGeneration(
         pixel_values = image_input["pixel_values"]
         patches_per_image = image_input["patches_per_image"].tolist()
 
-        # shard image input
-        if self.use_data_parallel:
-            vision_embeddings_flat = run_dp_sharded_vision_model(
-                pixel_values, self.vision_model
-            )
-        else:
-            vision_embeddings_flat = self.vision_model(pixel_values)
-
-        vision_embeddings_flat = self.multi_modal_projector(vision_embeddings_flat)
+        vision_embeddings_flat = self.encode_image_chunks(
+            pixel_values,
+            use_data_parallel=self.use_data_parallel,
+        )
 
         return [
             img.flatten(0, 1)