@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9f153aa781

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[MM][Perf][CG] Support ViT full CUDA graph for glm4_1v image and video inference (#40576)

Signed-off-by: grYe99 <guorongye99@gmail.com> Co-authored-by: grYe99 <guorongye99@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

GuorongYe committed Jun 9, 2026 at 11:13 UTC 9f153aa781f07d74b8574b57eab71551374c9498

4 files changed +480 -25

docs/design/cuda_graphs_multimodal.md

@@ -88,6 +88,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra

88 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |

89 | `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |

90 | `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |

91 +| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ |

93 !!! note

94 Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.

examples/generate/multimodal/vision_language_offline.py

@@ -2562,6 +2562,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [

2562 "qwen3_5",

2563 "qwen3_5_moe",

2564 "stepvl",

2565 + "glm4_1v",

2566 ]

2567

2568

tests/models/multimodal/generation/test_vit_cudagraph.py

+22

@@ -137,6 +137,28 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

137 ),

138 },

139 ),

140 + "glm4_1v": VitCudagraphTestConfig(

141 + model="zai-org/GLM-4.1V-9B-Thinking",

142 + image_prompt=(

143 + "[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n"

144 + "<|begin_of_image|><|image|><|end_of_image|>"

145 + "What is in this image?<|assistant|>assistant\n"

146 + ),

147 + video_prompt=(

148 + "[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n"

149 + "<|begin_of_video|><|video|><|end_of_video|>"

150 + "Describe this video in one sentence<|assistant|>assistant\n"

151 + ),

152 + needs_video_metadata=True,

153 + marks=[pytest.mark.core_model],

154 + vllm_runner_kwargs={

155 + "load_format": "dummy",

156 + "hf_overrides": partial(

157 + dummy_hf_overrides,

158 + model_arch="Glm4vForConditionalGeneration",

159 + ),

160 + },

161 + ),

162 }

163

164

vllm/model_executor/models/glm4_1v.py

+456 -25

@@ -97,10 +97,12 @@ from vllm.multimodal.processing import (

97 from vllm.sequence import IntermediateTensors

98 from vllm.utils.tensor_schema import TensorSchema, TensorShape

99 from vllm.v1.attention.backends.registry import AttentionBackendEnum

100 +from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

101

102 from ..layers.activation import SiluAndMul

103 from .interfaces import (

104 MultiModalEmbeddings,

105 + SupportsEncoderCudaGraph,

106 SupportsLoRA,

107 SupportsMRoPE,

108 SupportsMultiModal,

@@ -626,6 +628,11 @@ class Glm4vVisionTransformer(nn.Module):

628 ) -> None:

629 super().__init__()

630

631 + use_data_parallel = is_vit_use_data_parallel()

632 + self.tp_size = (

633 + 1 if use_data_parallel else get_tensor_model_parallel_world_size()

634 + )

635 +

636 patch_size = vision_config.patch_size

637 temporal_patch_size = vision_config.temporal_patch_size

638 in_channels = vision_config.in_channels

@@ -701,7 +708,7 @@ class Glm4vVisionTransformer(nn.Module):

708 return self.patch_embed.proj.weight.device

709

710 def rot_pos_emb(

704 - self, grid_thw: torch.Tensor

711 + self, grid_thw: list[list[int]]

712 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

713 pos_ids = []

714 for t, h, w in grid_thw:

@@ -729,7 +736,7 @@ class Glm4vVisionTransformer(nn.Module):

736 )

737 pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

738 pos_ids = torch.cat(pos_ids, dim=0)

732 - max_grid_size = grid_thw[:, 1:].max()

739 + max_grid_size = max(max(h, w) for _, h, w in grid_thw)

740

741 # Use pre-computed cos_sin_cache from RotaryEmbedding

742 cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

@@ -752,45 +759,192 @@ class Glm4vVisionTransformer(nn.Module):

759 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

760 return max_seqlen

761

762 + def pos_embeds_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:

763 + """Pre-compute absolute position embeddings for all input samples.

764 + The original `self.embeddings` fused token embeddings and position embeddings

765 + in one call, which prevented preparing position embeddings as static metadata

766 + required by CUDA graph capture / replay. This method decouples the two by

767 + feeding an all-zero token tensor to `self.embeddings`. The module therefore only

768 + performs bicubic interpolation based on the coordinates and returns pure

769 + position embeddings. These are cached in `prepare_encoder_metadata` and later

770 + added to the patch tokens in `forward` via `x = x + pos_embeds`, keeping the

771 + forward graph compatible with CUDA graph replay. Coordinate generation matches

772 + `rot_pos_emb` exactly to guarantee spatial alignment.

773 + """

774 +

775 + device = self.embeddings.position_embedding.weight.device

776 + dtype = self.dtype

777 + all_embeds = []

778 +

779 + for t, h, w in grid_thw:

780 + # Use the same coordinate generation logic as rot_pos_emb

781 + # to ensure consistent positional embedding interpolation

782 + h_coords = torch.arange(h).unsqueeze(1).expand(-1, w)

783 + w_coords = torch.arange(w).unsqueeze(0).expand(h, -1)

784 + h_coords = (

785 + h_coords.reshape(

786 + h // self.spatial_merge_size,

787 + self.spatial_merge_size,

788 + w // self.spatial_merge_size,

789 + self.spatial_merge_size,

790 + )

791 + .permute(0, 2, 1, 3)

792 + .flatten()

793 + )

794 + w_coords = (

795 + w_coords.reshape(

796 + h // self.spatial_merge_size,

797 + self.spatial_merge_size,

798 + w // self.spatial_merge_size,

799 + self.spatial_merge_size,

800 + )

801 + .permute(0, 2, 1, 3)

802 + .flatten()

803 + )

804 +

805 + lengths = [h * w] * t

806 + image_shapes = torch.tensor([[t, h, w]], device=device)

807 +

808 + h_coords_repeated = h_coords.repeat(t)

809 + w_coords_repeated = w_coords.repeat(t)

810 +

811 + embeds = self.embeddings(

812 + embeddings=torch.zeros(

813 + h * w * t, self.hidden_size, device=device, dtype=dtype

814 + ),

815 + lengths=lengths,

816 + image_shapes=image_shapes,

817 + h_coords=h_coords_repeated,

818 + w_coords=w_coords_repeated,

819 + )

820 + all_embeds.append(embeds)

821 +

822 + return torch.cat(all_embeds, dim=0).to(dtype)

823 +

824 + def prepare_encoder_metadata(

825 + self,

826 + grid_thw_list: list[list[int]],

827 + *,

828 + max_batch_size: int | None = None,

829 + max_frames_per_batch: int | None = None,

830 + max_seqlen_override: int | None = None,

831 + device: torch.device | None = None,

832 + ) -> dict[str, torch.Tensor | None]:

833 + """Compute encoder metadata from grid_thw_list.

834 +

835 + Shared by the eager forward path, CUDA graph capture, and

836 + CUDA graph replay to avoid duplicated implementation.

837 +

838 + Args:

839 + grid_thw_list: Grid configurations as list of [t, h, w].

840 + max_batch_size: If set, pad cu_seqlens to this size

841 + (needed for CUDA graph capture/replay).

842 + max_frames_per_batch: If set, overrides max_batch_size for

843 + cu_seqlens padding. For video inputs each item contributes

844 + T attention sequences (frames); this sizes the buffer to

845 + the total frame budget so video replays never overflow.

846 + max_seqlen_override: If set, use this value for max_seqlen

847 + instead of computing from cu_seqlens (needed for CUDA

848 + graph capture to cover worst-case replay scenarios).

849 + device: Device to place tensors on. Defaults to self.device.

850 + """

851 + if device is None:

852 + device = self.device

853 +

854 + metadata: dict[str, torch.Tensor | None] = {}

855 +

856 + # Positional embeddings

857 + metadata["pos_embeds"] = self.pos_embeds_interpolate(grid_thw_list)

858 + rotary_cos, rotary_sin, _ = self.rot_pos_emb(grid_thw_list)

859 + metadata["rotary_pos_emb_cos"] = rotary_cos

860 + metadata["rotary_pos_emb_sin"] = rotary_sin

861 +

862 + # cu_seqlens from grid_thw

863 + grid_thw_np = np.array(grid_thw_list, dtype=np.int32)

864 + patches_per_frame = grid_thw_np[:, 1] * grid_thw_np[:, 2]

865 + cu_seqlens = np.repeat(patches_per_frame, grid_thw_np[:, 0]).cumsum(

866 + dtype=np.int32

867 + )

868 + cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])

869 +

870 + # Pad cu_seqlens to the required number of sequences.

871 + # For videos each item contributes T frames = T attention sequences,

872 + # so the total can exceed max_batch_size. max_frames_per_batch

873 + # overrides the pad target when set.

874 + pad_to = (

875 + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size

876 + )

877 + if pad_to is not None:

878 + num_seqs = len(cu_seqlens) - 1

879 + if num_seqs < pad_to:

880 + cu_seqlens = np.concatenate(

881 + [

882 + cu_seqlens,

883 + np.full(

884 + pad_to - num_seqs,

885 + cu_seqlens[-1],

886 + dtype=np.int32,

887 + ),

888 + ]

889 + )

890 +

891 + # sequence_lengths (backend-specific)

892 + metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(

893 + self.attn_backend, cu_seqlens, device

894 + )

895 +

896 + # max_seqlen

897 + if max_seqlen_override is not None:

898 + max_seqlen_val = max_seqlen_override

899 + else:

900 + max_seqlen_val = MMEncoderAttention.compute_max_seqlen(

901 + self.attn_backend, cu_seqlens

902 + )

903 + # Keep max_seqlen on CPU: attention wrappers call .item() on it,

904 + # and having it on GPU would capture a wasteful D2H copy in CUDA

905 + # graphs without changing behavior (the scalar is baked at capture).

906 + metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)

907 +

908 + # Recompute cu_seqlens (backend-specific transformation)

909 + metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(

910 + self.attn_backend,

911 + cu_seqlens,

912 + self.hidden_size,

913 + self.tp_size,

914 + device,

915 + )

916 +

917 + return metadata

918 +

919 def forward(

920 self,

921 x: torch.Tensor,

922 grid_thw: torch.Tensor | list[list[int]],

923 + *,

924 + encoder_metadata: dict[str, torch.Tensor] | None = None,

925 ) -> torch.Tensor:

760 - if isinstance(grid_thw, list):

761 - grid_thw = torch.tensor(grid_thw, dtype=torch.int32)

926 + if encoder_metadata is None:

927 + if not isinstance(grid_thw, list):

928 + grid_thw = grid_thw.tolist()

929 + encoder_metadata = self.prepare_encoder_metadata(grid_thw)

930

931 # patchify

932 x = x.to(device=self.device, dtype=self.dtype)

933 x = self.patch_embed(x)

934 x = self.post_conv_layernorm(x)

935

768 - # compute position embedding

769 - rotary_pos_emb_cos, rotary_pos_emb_sin, image_type_ids = self.rot_pos_emb(

770 - grid_thw

771 - )

772 - # compute cu_seqlens

773 - cu_seqlens = torch.repeat_interleave(

774 - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

775 - ).cumsum(dim=0, dtype=torch.int32)

776 - cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])

777 - # pre-compute max_seqlen for attn mask to reduce cuMemcpy operations

778 - max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

779 - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

780 - cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)

781 - x = self.embeddings(

782 - x, seqlens, grid_thw, image_type_ids[:, 0], image_type_ids[:, 1]

783 - )

936 + pos_embeds = encoder_metadata["pos_embeds"]

937 + x = x + pos_embeds

938

939 # transformers

940 x = x.unsqueeze(1)

941 for blk in self.blocks:

942 x = blk(

943 x,

790 - cu_seqlens=cu_seqlens,

791 - rotary_pos_emb_cos=rotary_pos_emb_cos,

792 - rotary_pos_emb_sin=rotary_pos_emb_sin,

793 - max_seqlen=max_seqlen,

944 + cu_seqlens=encoder_metadata["cu_seqlens"],

945 + rotary_pos_emb_cos=encoder_metadata["rotary_pos_emb_cos"],

946 + rotary_pos_emb_sin=encoder_metadata["rotary_pos_emb_sin"],

947 + max_seqlen=encoder_metadata["max_seqlen"],

948 )

949

950 # adapter

@@ -1587,7 +1741,12 @@ class Glm4vMultiModalProcessor(BaseMultiModalProcessor[Glm4vProcessingInfo]):

1741 dummy_inputs=Glm4vDummyInputsBuilder,

1742 )

1743 class Glm4vForConditionalGeneration(

1590 - nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, SupportsMRoPE

1744 + nn.Module,

1745 + SupportsMultiModal,

1746 + SupportsEncoderCudaGraph,

1747 + SupportsLoRA,

1748 + SupportsPP,

1749 + SupportsMRoPE,

1750 ):

1751 packed_modules_mapping = {

1752 "qkv_proj": [

@@ -1625,8 +1784,12 @@ class Glm4vForConditionalGeneration(

1784 multimodal_config = vllm_config.model_config.multimodal_config

1785

1786 self.config = config

1787 + self.model_config = vllm_config.model_config

1788 self.multimodal_config = multimodal_config

1789 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1790 + self.is_multimodal_pruning_enabled = (

1791 + multimodal_config.is_multimodal_pruning_enabled()

1792 + )

1793

1794 with self._mark_tower_model(vllm_config, {"image", "video"}):

1795 self.visual = Glm4vVisionTransformer(

@@ -1752,6 +1915,274 @@ class Glm4vForConditionalGeneration(

1915 sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1916 return video_embeds.split(sizes)

1917

1918 + # -- SupportsEncoderCudaGraph protocol methods --

1919 +

1920 + def get_encoder_cudagraph_config(self):

1921 + from vllm.v1.worker.encoder_cudagraph_defs import (

1922 + EncoderCudaGraphConfig,

1923 + )

1924 +

1925 + # When EVS pruning is enabled, embed_multimodal post-processes both

1926 + # image and video embeddings (mrope positions are appended for image,

1927 + # prune+append for video). The encoder CUDA graph path bypasses that

1928 + # post-process, producing inconsistent embedding formats vs eager. So

1929 + # disable CUDA graph for all modalities when pruning is on.

1930 + modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]

1931 +

1932 + # Compute max_frames_per_video for budget sizing.

1933 + max_frames = self.get_max_frames_per_video() if "video" in modalities else 1

1934 +

1935 + return EncoderCudaGraphConfig(

1936 + modalities=modalities,

1937 + buffer_keys=[

1938 + "pixel_values",

1939 + "pos_embeds",

1940 + "rotary_pos_emb_cos",

1941 + "rotary_pos_emb_sin",

1942 + "cu_seqlens",

1943 + "max_seqlen",

1944 + "sequence_lengths",

1945 + ],

1946 + out_hidden_size=self.visual.out_hidden_size,

1947 + max_frames_per_video=max_frames,

1948 + )

1949 +

1950 + def get_input_modality(

1951 + self,

1952 + mm_kwargs: dict[str, Any],

1953 + ) -> str:

1954 + if "image_grid_thw" in mm_kwargs:

1955 + return "image"

1956 + elif "video_grid_thw" in mm_kwargs:

1957 + return "video"

1958 + raise AssertionError("This line should be unreachable.")

1959 +

1960 + def get_max_frames_per_video(self) -> int:

1961 + mm_registry = MULTIMODAL_REGISTRY

1962 + info = mm_registry.get_processing_info(self.model_config)

1963 + max_frames_per_video = info.get_num_frames_with_most_features(

1964 + seq_len=self.model_config.max_model_len,

1965 + mm_counts={"video": self.multimodal_config.get_limit_per_prompt("video")},

1966 + )

1967 + # Small 'max_frames_per_video' will cause 'tensor mismatch' in PR#43403

1968 + # 16 is the default 'num_frames' of '_get_vision_info'

1969 + return max(max_frames_per_video, 16)

1970 +

1971 + def get_encoder_cudagraph_budget_range(

1972 + self,

1973 + vllm_config,

1974 + ) -> tuple[int, int]:

1975 + # Min: estimated smallest possible encoder input.

1976 + # 224x224 image → 16x16 patches (patch_size=14)

1977 + # spatial_merge_size=2 → 8x8 = 64 tokens

1978 + min_budget = 64

1979 + # Max: capped by max_num_batched_tokens

1980 + max_budget = min(

1981 + vllm_config.scheduler_config.max_num_batched_tokens,

1982 + vllm_config.model_config.max_model_len,

1983 + )

1984 + return (min_budget, max_budget)

1985 +

1986 + def _get_pixel_values_by_modality(

1987 + self,

1988 + mm_kwargs: dict[str, Any],

1989 + ) -> torch.Tensor:

1990 + if self.get_input_modality(mm_kwargs) == "image":

1991 + pixel_values = mm_kwargs["pixel_values"]

1992 + else:

1993 + pixel_values = mm_kwargs["pixel_values_videos"]

1994 + return pixel_values

1995 +

1996 + def _get_grid_thw_by_modality(

1997 + self,

1998 + mm_kwargs: dict[str, Any],

1999 + ) -> list[tuple[int, int, int]]:

2000 + grid_thw_key = f"{self.get_input_modality(mm_kwargs)}_grid_thw"

2001 + grid_thw = mm_kwargs[grid_thw_key]

2002 + if not isinstance(grid_thw, list):

2003 + grid_thw = grid_thw.tolist()

2004 + return grid_thw

2005 +

2006 + def get_encoder_cudagraph_item_specs(

2007 + self,

2008 + mm_kwargs: dict[str, Any],

2009 + ):

2010 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

2011 +

2012 + m = self.visual.spatial_merge_size

2013 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

2014 + return [

2015 + EncoderItemSpec(

2016 + input_size=t * h * w,

2017 + output_tokens=t * (h // m) * (w // m),

2018 + )

2019 + for t, h, w in grid_thw

2020 + ]

2021 +

2022 + def select_encoder_cudagraph_items(

2023 + self,

2024 + mm_kwargs: dict[str, Any],

2025 + indices: list[int],

2026 + ) -> dict[str, Any]:

2027 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

2028 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

2029 +

2030 + if len(indices) == 0:

2031 + if self.get_input_modality(mm_kwargs) == "image":

2032 + return {

2033 + "pixel_values": pixel_values[:0],

2034 + "image_grid_thw": [],

2035 + }

2036 + else:

2037 + return {

2038 + "pixel_values_videos": pixel_values[:0],

2039 + "video_grid_thw": [],

2040 + }

2041 +

2042 + # Compute cumulative patch offsets for slicing pixel_values

2043 + patches_per_item = [t * h * w for t, h, w in grid_thw]

2044 + cum_patches = [0]

2045 + for p in patches_per_item:

2046 + cum_patches.append(cum_patches[-1] + p)

2047 +

2048 + selected_pv = torch.cat(

2049 + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]

2050 + )

2051 + selected_grid = [grid_thw[i] for i in indices]

2052 +

2053 + if self.get_input_modality(mm_kwargs) == "image":

2054 + return {

2055 + "pixel_values": selected_pv,

2056 + "image_grid_thw": selected_grid,

2057 + }

2058 + else:

2059 + return {

2060 + "pixel_values_videos": selected_pv,

2061 + "video_grid_thw": selected_grid,

2062 + }

2063 +

2064 + def prepare_encoder_cudagraph_capture_inputs(

2065 + self,

2066 + token_budget: int,

2067 + max_batch_size: int,

2068 + max_frames_per_batch: int,

2069 + device: torch.device,

2070 + dtype: torch.dtype,

2071 + ):

2072 + from vllm.v1.worker.encoder_cudagraph_defs import (

2073 + EncoderCudaGraphCaptureInputs,

2074 + )

2075 +

2076 + spatial_merge_size = self.visual.spatial_merge_size

2077 + per_mm_item_output = token_budget // max_batch_size

2078 +

2079 + frames_per_item = max_frames_per_batch // max_batch_size

2080 + if frames_per_item > 1:

2081 + # Build the capture grid using a video-format layout so that

2082 + # cu_seqlens is sized for video replays from the start.

2083 + # cu_seqlens has one entry per attention sequence (one per frame),

2084 + # so using T > 1 per item makes the buffer large enough without

2085 + # relying solely on padding.

2086 + # Ceiling ensures frames_per_item * tokens_per_frame >= per_mm_item_output

2087 + # so the pixel_values buffer covers any valid single-item replay.

2088 + tokens_per_frame = (

2089 + per_mm_item_output + frames_per_item - 1

2090 + ) // frames_per_item

2091 + # Video-format grid_config (T=frames_per_item).

2092 + grid_config = [

2093 + [

2094 + frames_per_item,

2095 + spatial_merge_size,

2096 + tokens_per_frame * spatial_merge_size,

2097 + ]

2098 + for _ in range(max_batch_size)

2099 + ]

2100 + else:

2101 + # Image-format grid_config (T=1).

2102 + grid_config = [

2103 + [1, spatial_merge_size, per_mm_item_output * spatial_merge_size]

2104 + for _ in range(max_batch_size)

2105 + ]

2106 +

2107 + # Create dummy pixel_values

2108 + patch_embed = self.visual.patch_embed

2109 + in_channels = patch_embed.proj.in_channels

2110 + patch_size = patch_embed.patch_size

2111 + temporal_patch_size = patch_embed.temporal_patch_size

2112 + total_patches = sum(t * h * w for t, h, w in grid_config)

2113 + flattened_patch_size = (

2114 + in_channels * temporal_patch_size * patch_size * patch_size

2115 + )

2116 + dummy_pixel_values = torch.randn(

2117 + total_patches, flattened_patch_size, device=device, dtype=dtype

2118 + )

2119 +

2120 + # Override max_seqlen with a safe upper bound for capture.

2121 + # max_seqlen.item() gets baked into the CUDA graph (not replayed),

2122 + # so the capture value must cover any replay scenario.

2123 + # Worst case: 1 item consuming the full budget ->

2124 + # seq_len = token_budget * spatial_merge_size^2.

2125 + metadata = self.visual.prepare_encoder_metadata(

2126 + grid_config,

2127 + max_batch_size=max_batch_size,

2128 + max_frames_per_batch=max_frames_per_batch,

2129 + max_seqlen_override=token_budget * (spatial_merge_size**2),

2130 + device=device,

2131 + )

2132 +

2133 + # Just use image-modality dummy input_buffer for capturing, since it's also

2134 + # compatible for video inputs (has the same shape: [num_patches, C*T*P*P]).

2135 + values = metadata | {

2136 + "pixel_values": dummy_pixel_values,

2137 + }

2138 +

2139 + return EncoderCudaGraphCaptureInputs(

2140 + values=values,

2141 + )

2142 +

2143 + def prepare_encoder_cudagraph_replay_buffers(

2144 + self,

2145 + mm_kwargs: dict[str, Any],

2146 + max_batch_size: int,

2147 + max_frames_per_batch: int,

2148 + ):

2149 + modality = self.get_input_modality(mm_kwargs)

2150 + grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)

2151 +

2152 + if modality == "image":

2153 + metadata = self.visual.prepare_encoder_metadata(

2154 + grid_thw_list,

2155 + max_batch_size=max_batch_size,

2156 + )

2157 + elif modality == "video":

2158 + metadata = self.visual.prepare_encoder_metadata(

2159 + grid_thw_list,

2160 + max_frames_per_batch=max_frames_per_batch,

2161 + )

2162 + else:

2163 + raise AssertionError("This line should be unreachable.")

2164 +

2165 + values = metadata | {

2166 + "pixel_values": self._get_pixel_values_by_modality(mm_kwargs),

2167 + }

2168 + return EncoderCudaGraphReplayBuffers(values=values)

2169 +

2170 + def encoder_cudagraph_forward(

2171 + self,

2172 + values: dict[str, torch.Tensor],

2173 + ) -> torch.Tensor:

2174 + pixel_values = values.pop("pixel_values")

2175 + metadata = values

2176 + return self.visual(pixel_values, None, encoder_metadata=metadata)

2177 +

2178 + def encoder_eager_forward(

2179 + self,

2180 + mm_kwargs: dict[str, Any],

2181 + ) -> torch.Tensor:

2182 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

2183 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

2184 + return self.visual(pixel_values, grid_thw)

2185 +

2186 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

2187 mm_input_by_modality = {}

2188