A high-throughput and memory-efficient inference and serving engine for LLMs
[MM][Perf][CG] Support ViT full cudagraphs for mllama4 (#40660)
Signed-off-by: allgather <all2allops@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
allgather committed
Jun 11, 2026 at 22:17 UTC
39dee1114a2cd183a9fb72b561808b385b6c9daa
5 files changed
+193
-14
docs/design/cuda_graphs_multimodal.md
+9
@@ -82,6 +82,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
82
83
| Architecture | Models | CG for Image | CG for Video |
84
| ------------ | ------ | ------------ | ------------ |
85
+| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - |
86
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
87
| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
88
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
@@ -114,6 +115,14 @@ vllm serve Qwen/Qwen3-VL-32B \
115
--compilation-config '{"cudagraph_mm_encoder": true}'
116
```
117
118
+For `Llama 4` (image only):
119
+
120
+```bash
121
+vllm serve meta-llama/Llama-4-Scout-17B-16E-Instruct \
122
+ --limit-mm-per-prompt '{"image": 1}' \
123
+ --compilation-config '{"cudagraph_mm_encoder": true}'
124
+```
125
+
126
With explicit budgets:
127
128
```bash
examples/generate/multimodal/vision_language_offline.py
+1
@@ -2532,6 +2532,7 @@ MODELS_NEED_VIDEO_METADATA = [
2532
2533
2534
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2535
+ "llama4",
2536
"internvl_chat",
2537
"qwen2_5_vl",
2538
"qwen3_vl",
tests/models/multimodal/generation/test_vit_cudagraph.py
+20
@@ -55,6 +55,26 @@ def step3_vl_chat_template(content: str) -> str:
55
56
57
MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
58
+ "llama4": VitCudagraphTestConfig(
59
+ model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
60
+ modalities=["image"],
61
+ image_prompt=(
62
+ "<|begin_of_text|><|header_start|>user<|header_end|>\n\n"
63
+ "<|image|>What is in this image?<|eot|>"
64
+ "<|header_start|>assistant<|header_end|>\n\n"
65
+ ),
66
+ max_model_len=4096,
67
+ max_tokens=32,
68
+ max_num_seqs=2,
69
+ vllm_runner_kwargs={
70
+ "load_format": "dummy",
71
+ "hf_overrides": partial(
72
+ dummy_hf_overrides,
73
+ model_arch="Llama4ForConditionalGeneration",
74
+ ),
75
+ },
76
+ marks=[pytest.mark.core_model],
77
+ ),
78
"internvl": VitCudagraphTestConfig(
79
model="OpenGVLab/InternVL3-1B",
80
num_video_frames=8,
tests/models/utils.py
+3
-2
@@ -507,12 +507,13 @@ def dummy_hf_overrides(
507
# Only set MoE related config when the model has MoE layers.
508
# Otherwise all models detected as MoE by _get_transformers_backend_cls.
509
if model_arch_config.num_experts > 0:
510
+ num_experts_per_tok = 1 if model_arch == "Llama4ForConditionalGeneration" else 2
511
update_dict.update(
512
{
513
"num_experts": num_experts,
513
- "num_experts_per_tok": 2,
514
+ "num_experts_per_tok": num_experts_per_tok,
515
# Kimi uses `num_experts_per_token`.
515
- "num_experts_per_token": 2,
516
+ "num_experts_per_token": num_experts_per_tok,
517
"num_local_experts": num_experts,
518
# Otherwise there will not be any expert layers
519
"first_k_dense_replace": 0,
vllm/model_executor/models/mllama4.py
+160
-12
@@ -19,7 +19,7 @@
19
import math
20
from collections.abc import Iterable, Mapping
21
from itertools import tee
22
-from typing import Annotated, Literal
22
+from typing import Annotated, Any, Literal
23
24
import torch
25
from torch import nn
@@ -78,6 +78,7 @@ from .interfaces import (
78
MixtureOfExperts,
79
MultiModalEmbeddings,
80
SupportsEagle3,
81
+ SupportsEncoderCudaGraph,
82
SupportsLoRA,
83
SupportsMultiModal,
84
SupportsPP,
@@ -105,7 +106,7 @@ class Llama4ImagePatchInputs(TensorSchema):
106
107
patches_per_image: Annotated[torch.Tensor, TensorShape("batch_size")]
108
"""
108
- The number of total patches for each image in the batch.
109
+ The number of chunked image tiles for each image in the batch.
110
111
This is used to split the embeddings which has the first two dimensions
112
flattened just like `pixel_values`.
@@ -731,6 +732,7 @@ class Llama4ForConditionalGeneration(
732
SupportsMultiModal,
733
SupportsPP,
734
MixtureOfExperts,
735
+ SupportsEncoderCudaGraph,
736
SupportsEagle3,
737
SupportsLoRA,
738
):
@@ -828,10 +830,161 @@ class Llama4ForConditionalGeneration(
830
num_physical_experts, num_local_physical_experts
831
)
832
833
+ def get_image_patches_per_chunk(self) -> int:
834
+ return Mllama4ProcessingInfo.get_patch_per_chunk(self.config.vision_config)
835
+
836
+ def encode_image_chunks(
837
+ self,
838
+ pixel_values: torch.Tensor,
839
+ *,
840
+ use_data_parallel: bool,
841
+ ) -> torch.Tensor:
842
+ if use_data_parallel:
843
+ vision_embeddings = run_dp_sharded_vision_model(
844
+ pixel_values, self.vision_model
845
+ )
846
+ else:
847
+ vision_embeddings = self.vision_model(pixel_values)
848
+
849
+ return self.multi_modal_projector(vision_embeddings)
850
+
851
+ def get_encoder_cudagraph_config(self):
852
+ from vllm.v1.worker.encoder_cudagraph_defs import (
853
+ EncoderCudaGraphConfig,
854
+ )
855
+
856
+ return EncoderCudaGraphConfig(
857
+ modalities=["image"],
858
+ buffer_keys=["pixel_values"],
859
+ out_hidden_size=self.config.text_config.hidden_size,
860
+ )
861
+
862
+ def get_input_modality(
863
+ self,
864
+ mm_kwargs: dict[str, Any],
865
+ ) -> str:
866
+ return "image"
867
+
868
+ def get_encoder_cudagraph_budget_range(
869
+ self,
870
+ vllm_config: VllmConfig,
871
+ ) -> tuple[int, int]:
872
+ min_budget = self.get_image_patches_per_chunk()
873
+ max_budget = min(
874
+ vllm_config.scheduler_config.max_num_batched_tokens,
875
+ self.vllm_config.model_config.max_model_len,
876
+ )
877
+ return (min_budget, max_budget)
878
+
879
+ def get_encoder_cudagraph_item_specs(
880
+ self,
881
+ mm_kwargs: dict[str, Any],
882
+ ):
883
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
884
+
885
+ patches_per_chunk = self.get_image_patches_per_chunk()
886
+ return [
887
+ EncoderItemSpec(
888
+ input_size=num_chunks,
889
+ output_tokens=num_chunks * patches_per_chunk,
890
+ )
891
+ for num_chunks in mm_kwargs["patches_per_image"].tolist()
892
+ ]
893
+
894
+ def select_encoder_cudagraph_items(
895
+ self,
896
+ mm_kwargs: dict[str, Any],
897
+ indices: list[int],
898
+ ) -> dict[str, Any]:
899
+ pixel_values = mm_kwargs["pixel_values"]
900
+ patches_per_image = mm_kwargs["patches_per_image"]
901
+
902
+ if len(indices) == 0:
903
+ return {
904
+ "pixel_values": pixel_values[:0],
905
+ "patches_per_image": patches_per_image[:0],
906
+ }
907
+
908
+ cum_chunks = [0]
909
+ for num_chunks in patches_per_image.tolist():
910
+ cum_chunks.append(cum_chunks[-1] + num_chunks)
911
+
912
+ selected_pixel_values = torch.cat(
913
+ [pixel_values[cum_chunks[i] : cum_chunks[i + 1]] for i in indices],
914
+ dim=0,
915
+ )
916
+
917
+ return {
918
+ "pixel_values": selected_pixel_values,
919
+ "patches_per_image": patches_per_image[indices],
920
+ }
921
+
922
+ def prepare_encoder_cudagraph_capture_inputs(
923
+ self,
924
+ token_budget: int,
925
+ max_batch_size: int,
926
+ max_frames_per_batch: int,
927
+ device: torch.device,
928
+ dtype: torch.dtype,
929
+ ):
930
+ from vllm.v1.worker.encoder_cudagraph_defs import (
931
+ EncoderCudaGraphCaptureInputs,
932
+ )
933
+
934
+ vision_config = self.config.vision_config
935
+ patches_per_chunk = self.get_image_patches_per_chunk()
936
+ chunks_per_capture = max(
937
+ 1, (token_budget + patches_per_chunk - 1) // patches_per_chunk
938
+ )
939
+ dummy_pixel_values = torch.randn(
940
+ chunks_per_capture,
941
+ vision_config.num_channels,
942
+ vision_config.image_size,
943
+ vision_config.image_size,
944
+ device=device,
945
+ dtype=dtype,
946
+ )
947
+
948
+ return EncoderCudaGraphCaptureInputs(
949
+ values={"pixel_values": dummy_pixel_values},
950
+ )
951
+
952
+ def prepare_encoder_cudagraph_replay_buffers(
953
+ self,
954
+ mm_kwargs: dict[str, Any],
955
+ max_batch_size: int,
956
+ max_frames_per_batch: int,
957
+ ):
958
+ from vllm.v1.worker.encoder_cudagraph_defs import (
959
+ EncoderCudaGraphReplayBuffers,
960
+ )
961
+
962
+ return EncoderCudaGraphReplayBuffers(
963
+ values={"pixel_values": mm_kwargs["pixel_values"]},
964
+ )
965
+
966
+ def encoder_cudagraph_forward(
967
+ self,
968
+ inputs: dict[str, torch.Tensor],
969
+ ) -> torch.Tensor:
970
+ return self.encode_image_chunks(
971
+ inputs["pixel_values"],
972
+ use_data_parallel=False,
973
+ ).flatten(0, 1)
974
+
975
+ def encoder_eager_forward(
976
+ self,
977
+ mm_kwargs: dict[str, Any],
978
+ ) -> torch.Tensor:
979
+ return self.encode_image_chunks(
980
+ mm_kwargs["pixel_values"],
981
+ use_data_parallel=False,
982
+ ).flatten(0, 1)
983
+
984
def _parse_and_validate_image_input(
985
self, **kwargs: object
986
) -> Llama4ImagePatchInputs | None:
834
- # num_images, 1, num_chunks, channel, image_size, image_size
987
+ # total_num_chunks, channel, image_size, image_size
988
pixel_values = kwargs.pop("pixel_values", None)
989
if pixel_values is None:
990
return None
@@ -853,15 +1006,10 @@ class Llama4ForConditionalGeneration(
1006
pixel_values = image_input["pixel_values"]
1007
patches_per_image = image_input["patches_per_image"].tolist()
1008
856
- # shard image input
857
- if self.use_data_parallel:
858
- vision_embeddings_flat = run_dp_sharded_vision_model(
859
- pixel_values, self.vision_model
860
- )
861
- else:
862
- vision_embeddings_flat = self.vision_model(pixel_values)
863
-
864
- vision_embeddings_flat = self.multi_modal_projector(vision_embeddings_flat)
1009
+ vision_embeddings_flat = self.encode_image_chunks(
1010
+ pixel_values,
1011
+ use_data_parallel=self.use_data_parallel,
1012
+ )
1013
1014
return [
1015
img.flatten(0, 1)