@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70009fb934

[MM][CG] Support ViT CUDA Graph for Gemma-4 (#46837)

Signed-off-by: Anthony Su <xsuanthony@gmail.com> Co-authored-by: Linkun Chen <github@lkchen.net>

anthonsu committed Jul 25, 2026 at 13:02 UTC 70009fb9344d2a7ba642e68369e0a64a6252e8bc
4 files changed +476 -1
docs/design/cuda_graphs_multimodal.md
+1
@@ -129,6 +129,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
129 | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
130 | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |
131 | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
132 +| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ |
133 | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
134 | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
135 | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
examples/generate/multimodal/vision_language_offline.py
+41
@@ -503,6 +503,45 @@ def run_gemma3n(questions: list[str], modality: str) -> ModelRequestData:
503 )
504
505
506 +# Gemma 4
507 +def run_gemma4(questions: list[str], modality: str) -> ModelRequestData:
508 + assert modality in ("image", "video")
509 + model_name = "google/gemma-4-31B-it"
510 +
511 + # NOTE: Gemma-4-31B is a large model. Users running into Out-Of-Memory (OOM)
512 + # errors might need to set `tensor_parallel_size` to > 1.
513 + engine_args = EngineArgs(
514 + model=model_name,
515 + max_model_len=4096,
516 + max_num_seqs=2,
517 + limit_mm_per_prompt={modality: 1},
518 + )
519 +
520 + if modality == "image":
521 + prompts = [
522 + (
523 + "<bos><start_of_turn>user\n"
524 + f"<|image|>\n{question}<end_of_turn>\n"
525 + "<start_of_turn>model\n"
526 + )
527 + for question in questions
528 + ]
529 + else: # video
530 + prompts = [
531 + (
532 + "<bos><start_of_turn>user\n"
533 + f"<|video|>\n{question}<end_of_turn>\n"
534 + "<start_of_turn>model\n"
535 + )
536 + for question in questions
537 + ]
538 +
539 + return ModelRequestData(
540 + engine_args=engine_args,
541 + prompts=prompts,
542 + )
543 +
544 +
545 # GLM-4v
546 def run_glm4v(questions: list[str], modality: str) -> ModelRequestData:
547 assert modality == "image"
@@ -2303,6 +2342,7 @@ model_example_map = {
2342 "exaone4_5": run_exaone4_5,
2343 "gemma3": run_gemma3,
2344 "gemma3n": run_gemma3n,
2345 + "gemma4": run_gemma4,
2346 "glm4v": run_glm4v,
2347 "glm4_1v": run_glm4_1v,
2348 "glm4_5v": run_glm4_5v,
@@ -2374,6 +2414,7 @@ MODELS_NEED_VIDEO_METADATA = [
2414
2415 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2416 "llama4",
2417 + "gemma4",
2418 "qwen2_vl",
2419 "qwen2_5_vl",
2420 "qwen3_vl",
tests/models/multimodal/generation/test_vit_cudagraph.py
+13
@@ -249,6 +249,19 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
249 },
250 skip=True, # TODO: Re-enable this once OOM issues are resolved on CI.
251 ),
252 + "gemma4": VitCudagraphTestConfig(
253 + model="google/gemma-4-E2B-it",
254 + image_prompt=(
255 + "<bos><start_of_turn>user\n<|image|>\nWhat is in this image?<end_of_turn>\n"
256 + "<start_of_turn>model\n"
257 + ),
258 + video_prompt=(
259 + "<bos><start_of_turn>user\n<|video|>\nDescribe this video in one sentence."
260 + "<end_of_turn>\n<start_of_turn>model\n"
261 + ),
262 + needs_video_metadata=True,
263 + marks=[pytest.mark.core_model],
264 + ),
265 }
266
267
vllm/model_executor/models/gemma4_mm.py
+421 -1
@@ -16,7 +16,7 @@ reason about temporal order.
16
17 import math
18 from collections.abc import Iterable, Mapping, Sequence
19 -from typing import TYPE_CHECKING, Annotated, Any, Literal
19 +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal
20
21 import numpy as np
22 import torch
@@ -72,6 +72,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
72 from .interfaces import (
73 MultiModalEmbeddings,
74 SupportsEagle3,
75 + SupportsEncoderCudaGraph,
76 SupportsLoRA,
77 SupportsMultiModal,
78 SupportsPP,
@@ -86,6 +87,12 @@ from .utils import (
87
88 if TYPE_CHECKING:
89 from vllm.model_executor.layers.quantization import QuantizationConfig
90 + from vllm.v1.worker.encoder_cudagraph_defs import (
91 + EncoderCudaGraphCaptureInputs,
92 + EncoderCudaGraphConfig,
93 + EncoderCudaGraphReplayBuffers,
94 + EncoderItemSpec,
95 + )
96
97 logger = init_logger(__name__)
98
@@ -980,7 +987,9 @@ class Gemma4ForConditionalGeneration(
987 SupportsPP,
988 SupportsLoRA,
989 SupportsEagle3,
990 + SupportsEncoderCudaGraph,
991 ):
992 + supports_encoder_cudagraph: ClassVar[Literal[True]] = True
993 # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window
994 # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model
995 # runner reads this to keep image bidirectional ranges that exceed the
@@ -1025,6 +1034,7 @@ class Gemma4ForConditionalGeneration(
1034 self.quant_config = quant_config
1035 self.multimodal_config = multimodal_config
1036 self.model_dtype = vllm_config.model_config.dtype
1037 + self.vllm_config = vllm_config
1038
1039 # Only quantize towers when the quant method supports their
1040 # dimensions. BNB/torchao handle arbitrary sizes; other methods
@@ -1525,6 +1535,416 @@ class Gemma4ForConditionalGeneration(
1535
1536 return multimodal_embeddings
1537
1538 + # ------------------------------------------------------------------ #
1539 + # EncoderCudaGraph protocol methods
1540 + # ------------------------------------------------------------------ #
1541 +
1542 + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig":
1543 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
1544 +
1545 + def pad_pixel_values(dst: torch.Tensor, src: torch.Tensor) -> None:
1546 + dst.zero_()
1547 + batch_size, num_patches = src.shape[0], src.shape[1]
1548 + dst[:batch_size, :num_patches].copy_(src)
1549 +
1550 + def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None:
1551 + dst.fill_(-1)
1552 + batch_size, num_patches = src.shape[0], src.shape[1]
1553 + dst[:batch_size, :num_patches].copy_(src)
1554 +
1555 + return EncoderCudaGraphConfig(
1556 + modalities=["image", "video"],
1557 + buffer_keys=[
1558 + "pixel_values",
1559 + "pixel_position_ids",
1560 + "gather_indices",
1561 + ],
1562 + out_hidden_size=self.config.text_config.hidden_size,
1563 + max_frames_per_video=_VIDEO_MAX_FRAMES,
1564 + padding_logics={
1565 + "pixel_values": pad_pixel_values,
1566 + "pixel_position_ids": pad_pixel_position_ids,
1567 + },
1568 + )
1569 +
1570 + def get_encoder_cudagraph_budget_range(
1571 + self,
1572 + vllm_config: VllmConfig,
1573 + ) -> tuple[int, int]:
1574 + min_budget = _SUPPORTED_SOFT_TOKENS[0]
1575 + max_budget = min(
1576 + vllm_config.scheduler_config.max_num_batched_tokens,
1577 + vllm_config.model_config.max_model_len,
1578 + )
1579 + return (min_budget, max_budget)
1580 +
1581 + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
1582 + if "pixel_values" in mm_kwargs:
1583 + return "image"
1584 + elif "pixel_values_videos" in mm_kwargs:
1585 + return "video"
1586 + raise ValueError("Unsupported modality in mm_kwargs")
1587 +
1588 + def get_max_frames_per_video(self) -> int:
1589 + return _VIDEO_MAX_FRAMES
1590 +
1591 + def get_encoder_cudagraph_item_specs(
1592 + self,
1593 + mm_kwargs: dict[str, Any],
1594 + ) -> list["EncoderItemSpec"]:
1595 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
1596 +
1597 + vision_cfg = self.vision_tower.config
1598 + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2
1599 +
1600 + modality = self.get_input_modality(mm_kwargs)
1601 + if modality == "image":
1602 + pixel_values = mm_kwargs["pixel_values"]
1603 + if isinstance(pixel_values, list):
1604 + return [
1605 + EncoderItemSpec(
1606 + input_size=pv.shape[0],
1607 + output_tokens=pv.shape[0] // pool_ratio,
1608 + )
1609 + for pv in pixel_values
1610 + ]
1611 + else:
1612 + return [
1613 + EncoderItemSpec(
1614 + input_size=pixel_values.shape[1],
1615 + output_tokens=pixel_values.shape[1] // pool_ratio,
1616 + )
1617 + for _ in range(pixel_values.shape[0])
1618 + ]
1619 + elif modality == "video":
1620 + pixel_values_videos = mm_kwargs["pixel_values_videos"]
1621 + video_frame_counts = mm_kwargs["video_frame_counts"]
1622 + fc_list = (
1623 + video_frame_counts.tolist()
1624 + if isinstance(video_frame_counts, torch.Tensor)
1625 + else list(video_frame_counts)
1626 + )
1627 + np_patches = pixel_values_videos.shape[1]
1628 + return [
1629 + EncoderItemSpec(
1630 + input_size=fc * np_patches,
1631 + output_tokens=fc * (np_patches // pool_ratio),
1632 + )
1633 + for fc in fc_list
1634 + ]
1635 + raise ValueError(f"Unknown modality: {modality}")
1636 +
1637 + def select_encoder_cudagraph_items(
1638 + self,
1639 + mm_kwargs: dict[str, Any],
1640 + indices: list[int],
1641 + ) -> dict[str, Any]:
1642 + modality = self.get_input_modality(mm_kwargs)
1643 + if modality == "image":
1644 + pixel_values = mm_kwargs["pixel_values"]
1645 + pixel_position_ids = mm_kwargs["pixel_position_ids"]
1646 + if len(indices) == 0:
1647 + is_pv_list = isinstance(pixel_values, list)
1648 + is_pp_list = isinstance(pixel_position_ids, list)
1649 + return {
1650 + "pixel_values": ([] if is_pv_list else pixel_values[:0]),
1651 + "pixel_position_ids": (
1652 + [] if is_pp_list else pixel_position_ids[:0]
1653 + ),
1654 + }
1655 + if isinstance(pixel_values, list):
1656 + return {
1657 + "pixel_values": [pixel_values[i] for i in indices],
1658 + "pixel_position_ids": [pixel_position_ids[i] for i in indices],
1659 + }
1660 + return {
1661 + "pixel_values": pixel_values[indices],
1662 + "pixel_position_ids": pixel_position_ids[indices],
1663 + }
1664 + elif modality == "video":
1665 + pixel_values_videos = mm_kwargs["pixel_values_videos"]
1666 + pixel_position_ids_videos = mm_kwargs["pixel_position_ids_videos"]
1667 + video_frame_counts = mm_kwargs["video_frame_counts"]
1668 +
1669 + if len(indices) == 0:
1670 + is_fc_tensor = isinstance(video_frame_counts, torch.Tensor)
1671 + return {
1672 + "pixel_values_videos": pixel_values_videos[:0],
1673 + "pixel_position_ids_videos": pixel_position_ids_videos[:0],
1674 + "video_frame_counts": (
1675 + video_frame_counts[:0] if is_fc_tensor else []
1676 + ),
1677 + }
1678 +
1679 + fc_list = (
1680 + video_frame_counts.tolist()
1681 + if isinstance(video_frame_counts, torch.Tensor)
1682 + else list(video_frame_counts)
1683 + )
1684 + cum_frames = [0]
1685 + for fc in fc_list:
1686 + cum_frames.append(cum_frames[-1] + fc)
1687 +
1688 + selected_pv = torch.cat(
1689 + [
1690 + pixel_values_videos[cum_frames[i] : cum_frames[i + 1]]
1691 + for i in indices
1692 + ],
1693 + dim=0,
1694 + )
1695 + selected_pp = torch.cat(
1696 + [
1697 + pixel_position_ids_videos[cum_frames[i] : cum_frames[i + 1]]
1698 + for i in indices
1699 + ],
1700 + dim=0,
1701 + )
1702 + selected_fc = (
1703 + video_frame_counts[indices]
1704 + if isinstance(video_frame_counts, torch.Tensor)
1705 + else [video_frame_counts[i] for i in indices]
1706 + )
1707 + return {
1708 + "pixel_values_videos": selected_pv,
1709 + "pixel_position_ids_videos": selected_pp,
1710 + "video_frame_counts": selected_fc,
1711 + }
1712 + raise ValueError(f"Unknown modality: {modality}")
1713 +
1714 + def prepare_encoder_cudagraph_capture_inputs(
1715 + self,
1716 + token_budget: int = 256,
1717 + max_batch_size: int = 4,
1718 + max_frames_per_batch: int = 1,
1719 + device: torch.device | str = "cpu",
1720 + dtype: torch.dtype | None = None,
1721 + path: str = "default",
1722 + **kwargs: Any,
1723 + ) -> "EncoderCudaGraphCaptureInputs":
1724 + from vllm.v1.worker.encoder_cudagraph_defs import (
1725 + EncoderCudaGraphCaptureInputs,
1726 + )
1727 +
1728 + dtype = dtype or torch.float32
1729 +
1730 + max_size = max(max_batch_size, max_frames_per_batch)
1731 +
1732 + vision_cfg = self.vision_tower.config
1733 + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2
1734 +
1735 + # Retrieve the model's actual configured maximum tokens:
1736 + configured_max_tokens = getattr(
1737 + self.config.vision_config,
1738 + "num_soft_tokens",
1739 + _SUPPORTED_SOFT_TOKENS[2],
1740 + )
1741 + # Dynamically compute the slot capacity per item bounded by both the
1742 + # current graph budget and the user's maximum config:
1743 + per_item_output = min(token_budget, configured_max_tokens)
1744 + # Satisfy k^2 * per_item_output = per_item_patches
1745 + per_item_patches = per_item_output * pool_ratio
1746 +
1747 + patch_size = self.vision_tower.config.patch_size
1748 + num_channels = getattr(self.vision_tower.config, "num_channels", 3)
1749 + patch_pixels = (patch_size**2) * num_channels
1750 +
1751 + dummy_pixel_values = torch.zeros(
1752 + (max_size, per_item_patches, patch_pixels),
1753 + device=device,
1754 + dtype=dtype,
1755 + )
1756 + dummy_pixel_position_ids = torch.full(
1757 + (max_size, per_item_patches, 2),
1758 + -1,
1759 + device=device,
1760 + dtype=torch.long,
1761 + )
1762 + dummy_gather_indices = torch.zeros(
1763 + (token_budget,),
1764 + device=device,
1765 + dtype=torch.long,
1766 + )
1767 +
1768 + return EncoderCudaGraphCaptureInputs(
1769 + values={
1770 + "pixel_values": dummy_pixel_values,
1771 + "pixel_position_ids": dummy_pixel_position_ids,
1772 + "gather_indices": dummy_gather_indices,
1773 + }
1774 + )
1775 +
1776 + def prepare_encoder_cudagraph_replay_buffers(
1777 + self,
1778 + mm_kwargs: dict[str, Any],
1779 + max_batch_size: int = 4,
1780 + max_frames_per_batch: int = 1,
1781 + path: str = "default",
1782 + **kwargs: Any,
1783 + ) -> "EncoderCudaGraphReplayBuffers":
1784 + from vllm.v1.worker.encoder_cudagraph_defs import (
1785 + EncoderCudaGraphReplayBuffers,
1786 + )
1787 +
1788 + modality = self.get_input_modality(mm_kwargs)
1789 + if modality == "image":
1790 + pixel_values = mm_kwargs["pixel_values"]
1791 + pixel_position_ids = mm_kwargs["pixel_position_ids"]
1792 + elif modality == "video":
1793 + pixel_values = mm_kwargs["pixel_values_videos"]
1794 + pixel_position_ids = mm_kwargs["pixel_position_ids_videos"]
1795 + else:
1796 + raise ValueError(f"Unsupported modality: {modality}")
1797 +
1798 + if isinstance(pixel_values, list):
1799 + max_patches = max(pv.shape[0] for pv in pixel_values)
1800 + batch_size = len(pixel_values)
1801 + pv_tensor = torch.zeros(
1802 + (batch_size, max_patches, pixel_values[0].shape[1]),
1803 + dtype=pixel_values[0].dtype,
1804 + device=pixel_values[0].device,
1805 + )
1806 + pp_tensor = torch.full(
1807 + (batch_size, max_patches, 2),
1808 + -1,
1809 + dtype=pixel_position_ids[0].dtype,
1810 + device=pixel_position_ids[0].device,
1811 + )
1812 + for i, (pv, pp) in enumerate(zip(pixel_values, pixel_position_ids)):
1813 + pv_tensor[i, : pv.shape[0]].copy_(pv)
1814 + pp_tensor[i, : pp.shape[0]].copy_(pp)
1815 + pixel_values = pv_tensor
1816 + pixel_position_ids = pp_tensor
1817 +
1818 + item_specs = self.get_encoder_cudagraph_item_specs(mm_kwargs)
1819 + per_item_out_tokens = [spec.output_tokens for spec in item_specs]
1820 + total_tokens = sum(per_item_out_tokens)
1821 +
1822 + device = pixel_values.device
1823 + vision_cfg = self.vision_tower.config
1824 + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2
1825 + per_item_output = pixel_values.shape[1] // pool_ratio
1826 +
1827 + # ONLY allocate an array of exact size `total_tokens`.
1828 + # DO NOT pad it. The upstream Graph Manager handles the padding securely.
1829 + gather_indices = torch.zeros((total_tokens,), dtype=torch.long, device=device)
1830 +
1831 + if modality == "image":
1832 + dst_offset = 0
1833 + for i, n_tok in enumerate(per_item_out_tokens):
1834 + safe_n_tok = min(n_tok, per_item_output)
1835 + src_start = i * per_item_output
1836 + src_end = src_start + safe_n_tok
1837 + gather_indices[dst_offset : dst_offset + safe_n_tok] = torch.arange(
1838 + src_start, src_end, dtype=torch.long, device=device
1839 + )
1840 + dst_offset += safe_n_tok
1841 + elif modality == "video":
1842 + video_frame_counts = mm_kwargs["video_frame_counts"]
1843 + fc_list = (
1844 + video_frame_counts.tolist()
1845 + if isinstance(video_frame_counts, torch.Tensor)
1846 + else list(video_frame_counts)
1847 + )
1848 + vision_cfg = self.vision_tower.config
1849 + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2
1850 + np_patches = pixel_values.shape[1]
1851 + frame_output_tokens = np_patches // pool_ratio
1852 + safe_frame_output_tokens = min(frame_output_tokens, per_item_output)
1853 +
1854 + dst_offset = 0
1855 + frame_idx = 0
1856 + for fc in fc_list:
1857 + for f in range(fc):
1858 + src_start = (frame_idx + f) * per_item_output
1859 + src_end = src_start + safe_frame_output_tokens
1860 + gather_indices[
1861 + dst_offset : dst_offset + safe_frame_output_tokens
1862 + ] = torch.arange(
1863 + src_start, src_end, dtype=torch.long, device=device
1864 + )
1865 + dst_offset += safe_frame_output_tokens
1866 + frame_idx += fc
1867 +
1868 + return EncoderCudaGraphReplayBuffers(
1869 + values={
1870 + "pixel_values": pixel_values,
1871 + "pixel_position_ids": pixel_position_ids,
1872 + "gather_indices": gather_indices,
1873 + }
1874 + )
1875 +
1876 + def encoder_cudagraph_forward(
1877 + self,
1878 + inputs: dict[str, torch.Tensor],
1879 + path: str = "default",
1880 + **kwargs: Any,
1881 + ) -> torch.Tensor:
1882 + pixel_values = inputs["pixel_values"]
1883 + pixel_position_ids = inputs["pixel_position_ids"]
1884 + gather_indices = inputs["gather_indices"]
1885 +
1886 + pad_tensor = (pixel_position_ids == -1).all(dim=-1)
1887 +
1888 + vt = self.vision_tower
1889 + inputs_embeds = vt.patch_embedder(
1890 + pixel_values,
1891 + pixel_position_ids,
1892 + pad_tensor,
1893 + ).to(self.model_dtype)
1894 +
1895 + encoder_outputs = vt.encoder(
1896 + inputs_embeds=inputs_embeds,
1897 + attention_mask=~pad_tensor,
1898 + pixel_position_ids=pixel_position_ids,
1899 + )
1900 + hidden_states = encoder_outputs.last_hidden_state
1901 +
1902 + pool_ratio = getattr(vt.config, "pooling_kernel_size", 2) ** 2
1903 + per_item_output = pixel_values.shape[1] // pool_ratio
1904 +
1905 + pooled_states, _ = vt.pooler(
1906 + hidden_states=hidden_states,
1907 + pixel_position_ids=pixel_position_ids,
1908 + padding_positions=pad_tensor,
1909 + output_length=per_item_output,
1910 + )
1911 +
1912 + if getattr(vt.config, "standardize", False):
1913 + pooled_states = (pooled_states - vt.std_bias) * vt.std_scale
1914 +
1915 + flat_pooled = pooled_states.reshape(-1, pooled_states.shape[-1])
1916 + gathered_states = flat_pooled[gather_indices]
1917 +
1918 + # Cast to the projection layer's dtype to resolve mixed-precision crash
1919 + target_dtype = self.embed_vision.embedding_projection.weight.dtype
1920 + gathered_states = gathered_states.to(target_dtype)
1921 +
1922 + flat_proj_embs = self.embed_vision(
1923 + inputs_embeds=gathered_states.unsqueeze(0)
1924 + ).squeeze(0)
1925 +
1926 + return flat_proj_embs
1927 +
1928 + def encoder_eager_forward(
1929 + self,
1930 + mm_kwargs: dict[str, Any],
1931 + path: str = "default",
1932 + **kwargs: Any,
1933 + ) -> torch.Tensor:
1934 + modality = self.get_input_modality(mm_kwargs)
1935 + if modality == "image":
1936 + image_input = self._parse_and_validate_image_input(**mm_kwargs)
1937 + assert image_input is not None
1938 + embeddings = self._process_image_input(image_input)
1939 + elif modality == "video":
1940 + video_input = self._parse_and_validate_video_input(**mm_kwargs)
1941 + assert video_input is not None
1942 + embeddings = self._process_video_input(video_input)
1943 + else:
1944 + raise ValueError(f"Unsupported modality: {modality}")
1945 +
1946 + return torch.cat(embeddings, dim=0)
1947 +
1948 def embed_input_ids(
1949 self,
1950 input_ids: torch.Tensor,