@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 27757dde02

[DSV4.1] Add encoder cuda graph support for deepseek-v4.1-flash (#56625)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

Isotr0py committed Sep 20, 2026 at 21:51 UTC 27757dde020ecda4f9b0e2c5ca2df1c29badc82f

6 files changed +599 -63

tests/models/multimodal/generation/test_vit_cudagraph.py

+23

@@ -55,6 +55,10 @@ def kimi_vl_chat_template(content: str) -> str:

55 )

58 +def deepseek_v41_chat_template(content: str) -> str:

59 + return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜></think>"

60 +

61 +

62 def step3_vl_chat_template(content: str) -> str:

63 return (

64 "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "

@@ -221,6 +225,25 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

225 vllm_runner_kwargs={"gpu_memory_utilization": 0.80},

226 marks=[pytest.mark.core_model],

227 ),

228 + "deepseek_v41": VitCudagraphTestConfig(

229 + model="deepseek-ai/DeepSeek-V4.1-Flash",

230 + modalities=["image"],

231 + image_prompt=deepseek_v41_chat_template(

232 + "<｜deepseek_image｜>\n\nWhat is in this image?"

233 + ),

234 + compilation_config_overrides={

235 + "encoder_cudagraph_token_budgets": [1024],

236 + },

237 + vllm_runner_kwargs={

238 + "load_format": "dummy",

239 + "attention_backend": "FLASHMLA_SPARSE_DSV41",

240 + "hf_overrides": partial(

241 + dummy_hf_overrides,

242 + model_arch="DeepseekV41ForCausalLM",

243 + exist_overrides={"vision_n_layers": 1},

244 + ),

245 + },

246 + ),

247 "step3_vl": VitCudagraphTestConfig(

248 model="stepfun-ai/Step3-VL-10B",

249 modalities=["image"],

tests/models/multimodal/processing/test_tensor_schema.py

@@ -218,8 +218,14 @@ def test_model_tensor_schema(model_id: str):

218

219 factories = model_cls._processor_factory

220

221 + # Capture helpers return capture-buffer containers (e.g.

222 + # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs.

223 + capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"}

224 +

225 inputs_parse_methods = []

226 for attr_name in dir(model_cls):

227 + if attr_name in capture_helpers:

228 + continue

229 attr = getattr(model_cls, attr_name)

230 if hasattr(attr, "__annotations__"):

231 return_type = attr.__annotations__.get("return", None)

vllm/models/deepseek_v4/common/vision.py

+156 -7

@@ -40,8 +40,7 @@ from vllm.model_executor.models.vision import (

40 )

43 -@lru_cache(8)

44 -def get_vision_cos_sin(

43 +def _compute_vision_cos_sin(

44 n_h: int, n_w: int, dim: int, theta: float

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

46 inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))

@@ -52,6 +51,13 @@ def get_vision_cos_sin(

51 return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)

54 +@lru_cache(8)

55 +def get_vision_cos_sin(

56 + n_h: int, n_w: int, dim: int, theta: float

57 +) -> tuple[torch.Tensor, torch.Tensor]:

58 + return _compute_vision_cos_sin(n_h, n_w, dim, theta)

59 +

60 +

61 def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:

62 dtype = x.dtype

63 x1, x2 = x.float().chunk(2, dim=-1)

@@ -123,15 +129,23 @@ class DeepseekV4VisionAttention(nn.Module):

129 )

130

131 def forward(

126 - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor

132 + self,

133 + x: torch.Tensor,

134 + cos: torch.Tensor,

135 + sin: torch.Tensor,

136 + cu_seqlens: torch.Tensor | None = None,

137 + max_seqlen: torch.Tensor | None = None,

138 ) -> torch.Tensor:

139 n = x.size(0)

140 qkv, _ = self.wqkv(x)

141 q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1))

142 q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d)

143 k = apply_rotary(k, cos, sin).unsqueeze(0)

133 - # One image per call: a dense batch, no varlen packing metadata.

134 - o = self.attn(q, k, v.unsqueeze(0))

144 + # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens

145 + # set): packed multi-image batch for encoder CUDA graph replay.

146 + o = self.attn(

147 + q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen

148 + )

149 out, _ = self.wo(o.reshape(n, -1))

150 return out

151

@@ -173,9 +187,14 @@ class DeepseekV4VisionBlock(nn.Module):

187 self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp")

188

189 def forward(

176 - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor

190 + self,

191 + x: torch.Tensor,

192 + cos: torch.Tensor,

193 + sin: torch.Tensor,

194 + cu_seqlens: torch.Tensor | None = None,

195 + max_seqlen: torch.Tensor | None = None,

196 ) -> torch.Tensor:

178 - x = x + self.attn(self.norm1(x), cos, sin)

197 + x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen)

198 return x + self.mlp(self.norm2(x))

199

200

@@ -206,6 +225,115 @@ class DeepseekV4ViT(nn.Module):

225 x = block(x, cos, sin)

226 return self.norm(x)

227

228 + def forward_packed(

229 + self,

230 + patches: torch.Tensor,

231 + cos: torch.Tensor,

232 + sin: torch.Tensor,

233 + cu_seqlens: torch.Tensor,

234 + max_seqlen: torch.Tensor,

235 + ) -> torch.Tensor:

236 + """Varlen path for encoder CUDA graphs: multiple images packed along

237 + rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed

238 + on device. Pure-tensor (no host reads, no H2D), safe to capture."""

239 + x = self.patch_embed(patches)

240 + for block in self.blocks:

241 + x = block(x, cos, sin, cu_seqlens, max_seqlen)

242 + return self.norm(x)

243 +

244 +

245 +def build_packed_vit_metadata(

246 + grids: list[list[int]] | list[tuple[int, int]],

247 + *,

248 + rope_dim: int,

249 + rope_theta: float,

250 + device: torch.device,

251 + max_seqlen_override: int | None = None,

252 + cached: bool = True,

253 +) -> dict[str, torch.Tensor]:

254 + """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``.

255 +

256 + Args:

257 + grids: ``[n_vit_h, n_vit_w]`` per image, in packing order.

258 + rope_dim: RoPE embedding dimension of the ViT.

259 + rope_theta: RoPE base frequency of the ViT.

260 + device: Device for the returned tensors (except ``max_seqlen``,

261 + which stays on the host).

262 + max_seqlen_override: Worst-case value baked in at CUDA graph capture

263 + (the attention wrapper reads ``max_seqlen`` on the host, so the

264 + capture-time value becomes a graph constant).

265 + cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time

266 + dummy grids pass ``False`` to avoid evicting real entries.

267 +

268 + """

269 + cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin

270 + cos_list: list[torch.Tensor] = []

271 + sin_list: list[torch.Tensor] = []

272 + cu_seqlens = [0]

273 + max_seqlen = 0

274 + for n_h, n_w in grids:

275 + cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta)

276 + cos_list.append(cos)

277 + sin_list.append(sin)

278 + n = n_h * n_w

279 + cu_seqlens.append(cu_seqlens[-1] + n)

280 + max_seqlen = max(max_seqlen, n)

281 + if cos_list:

282 + cos = torch.cat(cos_list).to(device)

283 + sin = torch.cat(sin_list).to(device)

284 + else:

285 + cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)

286 + sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)

287 + if max_seqlen_override is not None:

288 + max_seqlen = max_seqlen_override

289 + return {

290 + "vit_cos": cos,

291 + "vit_sin": sin,

292 + "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device),

293 + # Read on the host by the attention wrapper; keep on CPU.

294 + "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32),

295 + }

296 +

297 +

298 +def build_packed_merge_metadata(

299 + grids: list[list[int]] | list[tuple[int, int]],

300 + downsample_ratio: int,

301 + *,

302 + device: torch.device,

303 + dtype: torch.dtype,

304 +) -> dict[str, torch.Tensor]:

305 + """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``.

306 +

307 + For each aligner output row, the ``r x r`` source patch rows in the

308 + packed ViT output; positions past the image edge (the eager path pads

309 + them with zeros) get mask 0 and a clamped in-range index.

310 + """

311 + r = downsample_ratio

312 + idx_list: list[torch.Tensor] = []

313 + mask_list: list[torch.Tensor] = []

314 + offset = 0

315 + for n_h, n_w in grids:

316 + rows = -(-n_h // r)

317 + cols = -(-n_w // r)

318 + bi = torch.arange(rows).repeat_interleave(cols)

319 + bj = torch.arange(cols).repeat(rows)

320 + di = torch.arange(r).repeat_interleave(r)

321 + dj = torch.arange(r).repeat(r)

322 + hi = bi[:, None] * r + di[None, :]

323 + wj = bj[:, None] * r + dj[None, :]

324 + valid = (hi < n_h) & (wj < n_w)

325 + idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1)

326 + idx_list.append(idx)

327 + mask_list.append(valid)

328 + offset += n_h * n_w

329 + if idx_list:

330 + merge_idx = torch.cat(idx_list).to(device)

331 + merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype)

332 + else:

333 + merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device)

334 + merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device)

335 + return {"merge_idx": merge_idx, "merge_mask": merge_mask}

336 +

337

338 class DeepseekV4Aligner(nn.Module):

339 """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector."""

@@ -240,6 +368,27 @@ class DeepseekV4Aligner(nn.Module):

368 out, _ = self.w2(F.gelu(hidden))

369 return out

370

371 + def forward_packed(

372 + self,

373 + x: torch.Tensor,

374 + merge_idx: torch.Tensor,

375 + merge_mask: torch.Tensor,

376 + ) -> torch.Tensor:

377 + """Packed multi-image path for encoder CUDA graphs.

378 +

379 + ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``;

380 + gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge

381 + exactly (unfold rows are channel-major, so the gathered patch rows

382 + are transposed before flattening).

383 + """

384 + r2 = self.downsample_ratio**2

385 + m = merge_idx.shape[0]

386 + gathered = (x[merge_idx] * merge_mask).view(m, r2, -1)

387 + gathered = gathered.permute(0, 2, 1).reshape(m, -1)

388 + hidden, _ = self.w1(gathered)

389 + out, _ = self.w2(F.gelu(hidden))

390 + return out

391 +

392

393 def run_dp_sharded_vision_tower(

394 vision_model: DeepseekV4ViT,

vllm/models/deepseek_v41/amd/vl_model.py

+12 -28

@@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size

26 from vllm.model_executor.models.interfaces import (

27 MultiModalEmbeddings,

28 SupportsEagle3,

29 + SupportsEncoderCudaGraph,

30 SupportsMultiModal,

31 SupportsPP,

32 )

@@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY

45 from vllm.utils.tensor_schema import TensorSchema, TensorShape

47 from ..common.mm_preprocess import (

47 - IMAGE,

48 - IMAGE_END,

49 - IMAGE_NEW_LINE,

48 IMAGE_PLACEHOLDER,

51 - IMAGE_START,

49 DeepseekV4VLDummyInputsBuilder,

50 DeepseekV4VLMultiModalProcessor,

51 DeepseekV4VLProcessingInfo,

52 )

53 +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin

54 from .model import (

55 DeepseekV41LLMForCausalLM,

56 _linear_scale_param_name,

@@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper(

113 info=DeepseekV4VLProcessingInfo,

114 dummy_inputs=DeepseekV4VLDummyInputsBuilder,

115 )

118 -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3):

116 +class DeepseekV41ForCausalLM(

117 + nn.Module,

118 + DeepseekV4VLEncoderCudaGraphMixin,

119 + SupportsMultiModal,

120 + SupportsEncoderCudaGraph,

121 + SupportsPP,

122 + SupportsEagle3,

123 +):

124 """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower.

125

126 ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)

127 delegates through ``language_model`` via the protocol defaults.

128 + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is

129 + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``.

130 """

131

132 supports_encoder_tp_data = True

@@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports

203 resolve_bindings={"p": self.config.vision_patch_size},

204 )

205

199 - def _encode_image(

200 - self,

201 - patches: torch.Tensor,

202 - n_vit_h: int,

203 - n_vit_w: int,

204 - ) -> torch.Tensor:

205 - # Aligner rows in reading order, one per IMAGE slot.

206 - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)

207 -

208 - def _build_image_span(

209 - self, image_embeds: torch.Tensor, types: torch.Tensor

210 - ) -> torch.Tensor:

211 - """Full image span: aligner rows at IMAGE slots, the learned

212 - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""

213 - types = types.to(image_embeds.device)

214 - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])

215 - dtype = image_embeds.dtype

216 - span[types == IMAGE_START] = self.image_start.to(dtype)

217 - span[types == IMAGE_END] = self.image_end.to(dtype)

218 - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)

219 - span[types == IMAGE] = image_embeds

220 - return span

221 -

206 def _process_image_input(

207 self,

208 image_input: DeepseekV4VLImagePixelInputs,

vllm/models/deepseek_v41/common/vl_cudagraph.py new

+390

@@ -0,0 +1,390 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Encoder CUDA graph support for the DeepSeek-V4.1 vision tower.

4 +

5 +Mixin implementing the ``SupportsEncoderCudaGraph`` protocol

6 +(``vllm/model_executor/models/interfaces.py``) on top of the shared

7 +``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole

8 +batch of images into one varlen run: ViT blocks attend per image via

9 +``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask

10 +against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/

11 +IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is

12 +batched over all items of the packed batch.

13 +

14 +Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the

15 +manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP

16 +sharding and eager fallback.

17 +"""

18 +

19 +from collections.abc import Hashable

20 +from typing import TYPE_CHECKING, Any

21 +

22 +import torch

23 +

24 +from vllm.models.deepseek_v4.common.vision import (

25 + build_packed_merge_metadata,

26 + build_packed_vit_metadata,

27 +)

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

29 + EncoderCudaGraphCaptureInputs,

30 + EncoderCudaGraphConfig,

31 + EncoderCudaGraphReplayBuffers,

32 + EncoderItemSpec,

33 +)

34 +

35 +from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START

36 +

37 +if TYPE_CHECKING:

38 + from vllm.config import VllmConfig

39 + from vllm.config.multimodal import MultiModalConfig

40 + from vllm.models.deepseek_v4.common.vision import (

41 + DeepseekV4Aligner,

42 + DeepseekV4ViT,

43 + )

44 + from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config

45 +

46 +

47 +class DeepseekV4VLEncoderCudaGraphMixin:

48 + """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``.

49 +

50 + Expects ``self.vision``/``self.aligner``/``self.config``/

51 + ``self.multimodal_config`` and the ``image_start``/``image_newline``/

52 + ``image_end`` parameters from the host model.

53 + """

54 +

55 + # Provided by the host model; declared here so mixin methods type-check.

56 + vision: "DeepseekV4ViT"

57 + aligner: "DeepseekV4Aligner"

58 + config: "DeepseekV41Config"

59 + multimodal_config: "MultiModalConfig | None"

60 + image_start: torch.nn.Parameter

61 + image_end: torch.nn.Parameter

62 + image_newline: torch.nn.Parameter

63 +

64 + # -- helpers shared with the eager multimodal path --

65 +

66 + def _encode_image(

67 + self,

68 + patches: torch.Tensor,

69 + n_vit_h: int,

70 + n_vit_w: int,

71 + ) -> torch.Tensor:

72 + # Aligner rows in reading order, one per IMAGE slot.

73 + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w)

74 +

75 + def _build_image_span(

76 + self, image_embeds: torch.Tensor, types: torch.Tensor

77 + ) -> torch.Tensor:

78 + """Full image span: aligner rows at IMAGE slots, the learned

79 + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END."""

80 + types = types.to(image_embeds.device)

81 + span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1])

82 + dtype = image_embeds.dtype

83 + span[types == IMAGE_START] = self.image_start.to(dtype)

84 + span[types == IMAGE_END] = self.image_end.to(dtype)

85 + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)

86 + span[types == IMAGE] = image_embeds

87 + return span

88 +

89 + @staticmethod

90 + def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]:

91 + grid = mm_kwargs[key]

92 + if isinstance(grid, torch.Tensor):

93 + grid = grid.tolist()

94 + return [[int(v) for v in row] for row in grid]

95 +

96 + _encoder_cg_pad_totals: dict[int, int] = {}

97 +

98 + # -- SupportsEncoderCudaGraph protocol --

99 +

100 + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:

101 + mm_cfg = self.multimodal_config

102 + if (

103 + mm_cfg is not None

104 + and mm_cfg.mm_encoder_attn_dtype == "fp8"

105 + and mm_cfg.mm_encoder_fp8_scale_path is None

106 + ):

107 + raise ValueError(

108 + "cudagraph_mm_encoder is incompatible with dynamic FP8 "

109 + "scaling for the vision encoder (the amax history is "

110 + "host-side state that cannot be captured). Provide a static "

111 + "scale file via --mm-encoder-fp8-scale-path or disable "

112 + "cudagraph_mm_encoder."

113 + )

114 +

115 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

116 +

117 + backend = self.vision.blocks[0].attn.attn.attn_backend

118 + # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls

119 + # .tolist() on CUDA cu_seqlens; neither can be captured.

120 + if backend in (

121 + AttentionBackendEnum.FLASHINFER,

122 + AttentionBackendEnum.TORCH_SDPA,

123 + ):

124 + raise ValueError(

125 + f"cudagraph_mm_encoder is not supported with the {backend.name} "

126 + "ViT attention backend for DeepSeek-V4.1. Set "

127 + "--mm-encoder-attn-backend FLASH_ATTN or disable "

128 + "cudagraph_mm_encoder."

129 + )

130 +

131 + pad_totals = self._encoder_cg_pad_totals

132 +

133 + def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None:

134 + # Varlen attention requires cu_seqlens[-1] to equal the number of

135 + # rows actually passed in. The captured buffers are sized for the

136 + # full token budget, so a smaller real batch is completed with one

137 + # trailing padding sequence; declaring fewer rows than the buffer

138 + # holds is undefined behaviour and returns NaN on FlashAttn.

139 + total = pad_totals.get(dst.data_ptr())

140 + if total is None:

141 + raise RuntimeError(

142 + "cu_seqlens replay buffer was not registered at capture "

143 + "time; the manager must replay into the capture-time "

144 + "buffers."

145 + )

146 + n = min(src.shape[0], dst.shape[0])

147 + dst[:n].copy_(src[:n])

148 + dst[n:] = total

149 +

150 + return EncoderCudaGraphConfig(

151 + modalities=["image"],

152 + buffer_keys=[

153 + "patches",

154 + "vit_cos",

155 + "vit_sin",

156 + "cu_seqlens",

157 + "max_seqlen",

158 + "merge_idx",

159 + "merge_mask",

160 + ],

161 + out_hidden_size=self.config.hidden_size,

162 + padding_logics={"cu_seqlens": pad_cu_seqlens},

163 + )

164 +

165 + def get_encoder_cudagraph_budget_range(

166 + self,

167 + vllm_config: "VllmConfig",

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

169 + # Min: comfortably above the smallest possible image span (4 tokens

170 + # for a 1x1 grid); also sets max_batch_size = 64 via the manager's

171 + # auto-inference (max_batch_size <= min_budget invariant).

172 + min_budget = 64

173 + max_budget = min(

174 + vllm_config.scheduler_config.max_num_batched_tokens,

175 + vllm_config.model_config.max_model_len,

176 + )

177 + return (min_budget, max_budget)

178 +

179 + def get_encoder_cudagraph_item_specs(

180 + self,

181 + mm_kwargs: dict[str, Any],

182 + ) -> list[EncoderItemSpec]:

183 + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")

184 + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid")

185 + return [

186 + EncoderItemSpec(

187 + input_size=h * w,

188 + output_tokens=lh * (lw + 1) + 2,

189 + )

190 + for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True)

191 + ]

192 +

193 + def select_encoder_cudagraph_items(

194 + self,

195 + mm_kwargs: dict[str, Any],

196 + indices: list[int],

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

198 + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")

199 + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid")

200 + patches = mm_kwargs["patches"]

201 + types = mm_kwargs["types"]

202 +

203 + if len(indices) == 0:

204 + return {

205 + "patches": patches[:0],

206 + "vit_grid": torch.zeros((0, 2), dtype=torch.int64),

207 + "llm_grid": torch.zeros((0, 2), dtype=torch.int64),

208 + "types": types[:0],

209 + }

210 +

211 + cum_patches = [0]

212 + for h, w in vit_grid:

213 + cum_patches.append(cum_patches[-1] + h * w)

214 + cum_spans = [0]

215 + for lh, lw in llm_grid:

216 + cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2)

217 +

218 + return {

219 + "patches": torch.cat(

220 + [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices]

221 + ),

222 + "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64),

223 + "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64),

224 + "types": torch.cat(

225 + [types[cum_spans[i] : cum_spans[i + 1]] for i in indices]

226 + ),

227 + }

228 +

229 + def prepare_encoder_cudagraph_capture_inputs(

230 + self,

231 + token_budget: int,

232 + max_batch_size: int,

233 + max_frames_per_batch: int,

234 + device: torch.device,

235 + dtype: torch.dtype,

236 + path: str = "default",

237 + axis_keys: tuple[Hashable, ...] | None = None,

238 + ) -> EncoderCudaGraphCaptureInputs:

239 + if (

240 + self.multimodal_config is not None

241 + and self.multimodal_config.get_limit_per_prompt("image") == 0

242 + ):

243 + raise RuntimeError(

244 + "cudagraph_mm_encoder requires the vision tower, but it was "

245 + "stubbed out because the image limit is 0. Disable "

246 + "cudagraph_mm_encoder."

247 + )

248 +

249 + config = self.config

250 + r = config.vision_downsample_ratio

251 + # Span rows per dummy item (ceiling so the buffers tile the budget).

252 + # Each dummy image is a single aligner row of per_item_out blocks:

253 + # r x (r * per_item_out) patches. One real image using the whole

254 + # budget needs at most r^2 * token_budget patches, which this fits.

255 + per_item_out = -(-token_budget // max_batch_size)

256 + grids = [[r, r * per_item_out]] * max_batch_size

257 + total_patches = r * r * per_item_out * max_batch_size

258 +

259 + p = config.vision_patch_size

260 + patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype)

261 +

262 + metadata = build_packed_vit_metadata(

263 + grids,

264 + rope_dim=self.vision.rope_dim,

265 + rope_theta=self.vision.rope_theta,

266 + device=device,

267 + max_seqlen_override=total_patches,

268 + cached=False,

269 + )

270 + # Spare cu_seqlens slots let replay append a padding sequence that

271 + # covers the rows a smaller real batch does not fill.

272 + real_cu = metadata.pop("cu_seqlens")

273 + cu_seqlens = torch.full(

274 + (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device

275 + )

276 + cu_seqlens[: real_cu.numel()] = real_cu

277 +

278 + merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype)

279 +

280 + self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches

281 +

282 + return EncoderCudaGraphCaptureInputs(

283 + values={

284 + "patches": patches,

285 + **metadata,

286 + "cu_seqlens": cu_seqlens,

287 + **merge,

288 + }

289 + )

290 +

291 + def prepare_encoder_cudagraph_replay_buffers(

292 + self,

293 + mm_kwargs: dict[str, Any],

294 + max_batch_size: int,

295 + max_frames_per_batch: int,

296 + path: str = "default",

297 + ) -> EncoderCudaGraphReplayBuffers:

298 + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")

299 + patches = mm_kwargs["patches"]

300 + dtype = self.aligner.w1.weight.dtype

301 + if patches.dtype != dtype:

302 + patches = patches.to(dtype)

303 +

304 + # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and

305 + # pad_cu_seqlens appends the padding sequence covering the tail rows.

306 + metadata = build_packed_vit_metadata(

307 + vit_grid,

308 + rope_dim=self.vision.rope_dim,

309 + rope_theta=self.vision.rope_theta,

310 + device=patches.device,

311 + )

312 + merge = build_packed_merge_metadata(

313 + vit_grid,

314 + self.config.vision_downsample_ratio,

315 + device=patches.device,

316 + dtype=dtype,

317 + )

318 + return EncoderCudaGraphReplayBuffers(

319 + values={"patches": patches, **metadata, **merge}

320 + )

321 +

322 + def encoder_cudagraph_forward(

323 + self,

324 + inputs: dict[str, torch.Tensor],

325 + path: str = "default",

326 + ) -> torch.Tensor:

327 + vit_out = self.vision.forward_packed(

328 + inputs["patches"],

329 + inputs["vit_cos"],

330 + inputs["vit_sin"],

331 + inputs["cu_seqlens"],

332 + inputs["max_seqlen"],

333 + )

334 + return self.aligner.forward_packed(

335 + vit_out, inputs["merge_idx"], inputs["merge_mask"]

336 + )

337 +

338 + def encoder_eager_forward(

339 + self,

340 + mm_kwargs: dict[str, Any],

341 + path: str = "default",

342 + ) -> torch.Tensor:

343 + patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype)

344 + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid")

345 + outs: list[torch.Tensor] = []

346 + offset = 0

347 + for n_vit_h, n_vit_w in vit_grid:

348 + n_vit = n_vit_h * n_vit_w

349 + outs.append(

350 + self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w)

351 + )

352 + offset += n_vit

353 + if not outs:

354 + return patches.new_zeros((0, self.config.hidden_size))

355 + return torch.cat(outs, dim=0)

356 +

357 + def postprocess_encoder_output(

358 + self,

359 + outputs: dict[str, torch.Tensor],

360 + indices: list[int],

361 + per_item_out_tokens: list[int],

362 + dest: dict[int, torch.Tensor] | list[torch.Tensor | None],

363 + clone: bool = False,

364 + batch_mm_kwargs: dict[str, Any] | None = None,

365 + ) -> None:

366 + if not indices:

367 + return

368 + assert batch_mm_kwargs is not None

369 + aligner_out = outputs["default"]

370 + r = self.config.vision_downsample_ratio

371 + vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid")

372 + llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid")

373 + types = batch_mm_kwargs["types"].to(aligner_out.device)

374 +

375 + n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid)

376 + span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid]

377 +

378 + # Batched span assembly: one masked fill per role across all items.

379 + dtype = aligner_out.dtype

380 + span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size)

381 + span[types == IMAGE] = aligner_out[:n_rows]

382 + span[types == IMAGE_START] = self.image_start.to(dtype)

383 + span[types == IMAGE_END] = self.image_end.to(dtype)

384 + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype)

385 +

386 + # Freshly allocated, so later replays cannot clobber the results.

387 + offset = 0

388 + for idx, span_len in zip(indices, span_lens):

389 + dest[idx] = span[offset : offset + span_len]

390 + offset += span_len

vllm/models/deepseek_v41/nvidia/vl_model.py