[MM][CG] Gemma3 Encoder CUDA Graph (#43591)
Signed-off-by: JisoLya <523420504@qq.com> Signed-off-by: Soyaazz <523420504@qq.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Soyaazz committed
Jun 29, 2026 at 12:52 UTC
4559c43a9526597c00cbcc4f59979496500268d1
3 files changed
+151
-1
docs/design/cuda_graphs_multimodal.md
+1
@@ -127,6 +127,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
127
| Architecture | Models | CG for Image | CG for Video | Dual-Path Graph |
128
| ------------ | ------ | ------------ | ------------ | --------------- |
129
| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
130
+| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
131
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
132
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
133
| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
tests/models/multimodal/generation/test_vit_cudagraph.py
+14
@@ -62,7 +62,21 @@ def step3_vl_chat_template(content: str) -> str:
62
)
63
64
65
+def gemma3_chat_template(content: str) -> str:
66
+ return f"<bos><start_of_turn>user\n{content}<end_of_turn>\n<start_of_turn>model\n"
67
+
68
+
69
MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
70
+ "gemma3": VitCudagraphTestConfig(
71
+ model="google/gemma-3-4b-it",
72
+ modalities=["image"],
73
+ image_prompt=gemma3_chat_template("<start_of_image>What is in this image?"),
74
+ compilation_config_overrides={
75
+ "encoder_cudagraph_token_budgets": [512],
76
+ },
77
+ dtype="bfloat16",
78
+ max_model_len=4096,
79
+ ),
80
"llama4": VitCudagraphTestConfig(
81
model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
82
modalities=["image"],
vllm/model_executor/models/gemma3_mm.py
+136
-1
@@ -39,6 +39,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
39
40
from .interfaces import (
41
MultiModalEmbeddings,
42
+ SupportsEncoderCudaGraph,
43
SupportsLoRA,
44
SupportsMultiModal,
45
SupportsPP,
@@ -467,7 +468,7 @@ class Gemma3MultiModalProjector(nn.Module):
468
dummy_inputs=Gemma3DummyInputsBuilder,
469
)
470
class Gemma3ForConditionalGeneration(
470
- nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
471
+ nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsEncoderCudaGraph
472
):
473
packed_modules_mapping = {
474
"qkv_proj": [
@@ -504,8 +505,12 @@ class Gemma3ForConditionalGeneration(
505
quant_config = vllm_config.quant_config
506
multimodal_config = vllm_config.model_config.multimodal_config
507
self.config = config
508
+ self.model_config = vllm_config.model_config
509
self.quant_config = quant_config
510
self.multimodal_config = multimodal_config
511
+ self.vit_positions_per_patch = (
512
+ self.config.vision_config.image_size // self.config.vision_config.patch_size
513
+ ) ** 2
514
515
self.configure_mm_token_handling(
516
vocab_size=config.text_config.vocab_size,
@@ -682,3 +687,133 @@ class Gemma3ForConditionalGeneration(
687
"""
688
# The Gemma3 connector maintains a 1:1 token mapping
689
return num_vision_tokens
690
+
691
+ def get_encoder_cudagraph_config(self):
692
+ from vllm.v1.worker.encoder_cudagraph_defs import (
693
+ EncoderCudaGraphConfig,
694
+ )
695
+
696
+ return EncoderCudaGraphConfig(
697
+ modalities=["image"],
698
+ buffer_keys=["pixel_values"],
699
+ out_hidden_size=self.config.text_config.hidden_size,
700
+ )
701
+
702
+ def get_encoder_cudagraph_budget_range(
703
+ self,
704
+ vllm_config: VllmConfig,
705
+ ) -> tuple[int, int]:
706
+ min_budget = self.config.mm_tokens_per_image
707
+ max_budget = min(
708
+ vllm_config.scheduler_config.max_num_batched_tokens,
709
+ self.model_config.max_model_len,
710
+ )
711
+ return (min_budget, max_budget)
712
+
713
+ def get_encoder_cudagraph_item_specs(
714
+ self,
715
+ mm_kwargs: dict[str, Any],
716
+ ):
717
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
718
+
719
+ num_patches = mm_kwargs["num_patches"]
720
+ mm_tokens_per_image = self.config.mm_tokens_per_image
721
+
722
+ return [
723
+ EncoderItemSpec(
724
+ input_size=int(np) * self.vit_positions_per_patch,
725
+ output_tokens=int(np) * mm_tokens_per_image,
726
+ )
727
+ for np in num_patches
728
+ ]
729
+
730
+ def select_encoder_cudagraph_items(
731
+ self,
732
+ mm_kwargs: dict[str, Any],
733
+ indices: list[int],
734
+ ) -> dict[str, Any]:
735
+ pixel_values = mm_kwargs["pixel_values"]
736
+ num_patches = mm_kwargs["num_patches"]
737
+
738
+ if len(indices) == 0:
739
+ return {
740
+ "pixel_values": pixel_values[:0],
741
+ "num_patches": num_patches[:0],
742
+ }
743
+ cum_patches = [0]
744
+ for p in num_patches:
745
+ cum_patches.append(cum_patches[-1] + int(p))
746
+
747
+ selected_pv = torch.cat(
748
+ [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
749
+ )
750
+ selected_np = num_patches[indices]
751
+
752
+ return {
753
+ "pixel_values": selected_pv,
754
+ "num_patches": selected_np,
755
+ }
756
+
757
+ def prepare_encoder_cudagraph_capture_inputs(
758
+ self,
759
+ token_budget: int,
760
+ max_batch_size: int,
761
+ max_frames_per_batch: int,
762
+ device: torch.device,
763
+ dtype: torch.dtype,
764
+ ):
765
+ from vllm.v1.worker.encoder_cudagraph_defs import (
766
+ EncoderCudaGraphCaptureInputs,
767
+ )
768
+
769
+ mm_tokens_per_image = self.config.mm_tokens_per_image
770
+ num_images = min(
771
+ token_budget // mm_tokens_per_image,
772
+ max_batch_size,
773
+ )
774
+
775
+ image_size = self.config.vision_config.image_size
776
+ dummy_pixel_values = torch.randn(
777
+ num_images,
778
+ 3,
779
+ image_size,
780
+ image_size,
781
+ device=device,
782
+ dtype=dtype,
783
+ )
784
+ values = {"pixel_values": dummy_pixel_values}
785
+
786
+ return EncoderCudaGraphCaptureInputs(
787
+ values,
788
+ )
789
+
790
+ def prepare_encoder_cudagraph_replay_buffers(
791
+ self,
792
+ mm_kwargs: dict[str, Any],
793
+ max_batch_size: int,
794
+ max_frames_per_batch: int,
795
+ ):
796
+ from vllm.v1.worker.encoder_cudagraph_defs import (
797
+ EncoderCudaGraphReplayBuffers,
798
+ )
799
+
800
+ return EncoderCudaGraphReplayBuffers(
801
+ values={"pixel_values": mm_kwargs["pixel_values"]},
802
+ )
803
+
804
+ def encoder_cudagraph_forward(
805
+ self,
806
+ values: dict[str, torch.Tensor],
807
+ ) -> torch.Tensor:
808
+ pixel_values = values["pixel_values"]
809
+ image_features = self.vision_tower(pixel_values)
810
+ image_features = self.multi_modal_projector(image_features)
811
+ return image_features.flatten(end_dim=1)
812
+
813
+ def encoder_eager_forward(
814
+ self,
815
+ mm_kwargs: dict[str, Any],
816
+ ) -> torch.Tensor:
817
+ image_input = self._parse_and_validate_image_input(**mm_kwargs)
818
+ results = self._process_image_input(image_input)
819
+ return torch.cat(results, dim=0)