@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 55de40a2fc

[Minimax-M3] Add Encoder CUDA graph support (#58673)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Codex <codex@openai.com>

Thien Tran committed Sep 28, 2026 at 12:17 UTC 55de40a2fc1ad2e54f13711daa83303d18f80247

5 files changed +397 -17

tests/models/multimodal/generation/test_vit_cudagraph.py

+18

@@ -175,6 +175,24 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

175 },

176 marks=[pytest.mark.core_model],

177 ),

178 + "minimax_m3": VitCudagraphTestConfig(

179 + model="MiniMaxAI/MiniMax-M3",

180 + modalities=["image"],

181 + image_prompt=("]~b]user\n]<]image[>[What is in this image?[e~[\n]~b]ai\n"),

182 + compilation_config_overrides={

183 + "encoder_cudagraph_token_budgets": [512],

184 + },

185 + vllm_runner_kwargs={

186 + "trust_remote_code": True,

187 + "load_format": "dummy",

188 + "hf_overrides": partial(

189 + dummy_hf_overrides,

190 + model_arch="MiniMaxM3SparseForConditionalGeneration",

191 + ),

192 + "mm_processor_kwargs": {"max_long_side_pixel": 336},

193 + },

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

195 + ),

196 "qwen3_vl": VitCudagraphTestConfig(

197 model="Qwen/Qwen3-VL-2B-Instruct",

198 image_prompt=qwen_vl_chat_template(

vllm/models/minimax_m3/amd/model.py

+8 -1

@@ -106,6 +106,9 @@ from vllm.models.minimax_m3.amd.sparse_attention_msa import (

106 MiniMaxM3SparseAiterPAImpl,

107 MiniMaxM3SparseAiterPAPrefillMetadata,

108 )

109 +from vllm.models.minimax_m3.common.encoder_cudagraph import (

110 + MiniMaxM3EncoderCudaGraphMixin,

111 +)

112 from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer

113 from vllm.models.minimax_m3.common.mm_preprocess import (

114 MiniMaxM3VLDummyInputsBuilder,

@@ -1685,7 +1688,11 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):

1688 dummy_inputs=MiniMaxM3VLDummyInputsBuilder,

1689 )

1690 class MiniMaxM3SparseForConditionalGeneration(

1688 - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3

1691 + nn.Module,

1692 + SupportsMultiModal,

1693 + MiniMaxM3EncoderCudaGraphMixin,

1694 + SupportsPP,

1695 + SupportsEagle3,

1696 ):

1697 """Top-level (VL) entry point for MiniMax M3.

1698

vllm/models/minimax_m3/common/encoder_cudagraph.py new

+281

@@ -0,0 +1,281 @@

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

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

3 +"""Encoder CUDA graph support for the MiniMax M3 vision tower.

4 +

5 +Mixin implementing the SupportsEncoderCudaGraph protocol for any host class

6 +that holds the vision tower as ``self.vision_tower`` (a

7 +``MiniMaxVLVisionModel``). The manager routes image items through this mixin;

8 +video items use the eager encoder path.

9 +"""

10 +

11 +from collections.abc import Hashable

12 +from typing import TYPE_CHECKING, Any

13 +

14 +import torch

15 +

16 +from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph

17 +

18 +if TYPE_CHECKING:

19 + from vllm.config import VllmConfig

20 + from vllm.config.multimodal import MultiModalConfig

21 + from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel

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

23 + EncoderCudaGraphCaptureInputs,

24 + EncoderCudaGraphConfig,

25 + EncoderCudaGraphReplayBuffers,

26 + EncoderItemSpec,

27 + )

28 +

29 +

30 +class MiniMaxM3EncoderCudaGraphMixin(SupportsEncoderCudaGraph):

31 + """SupportsEncoderCudaGraph for hosts with ``self.vision_tower``."""

32 +

33 + vision_tower: "MiniMaxVLVisionModel"

34 + multimodal_config: "MultiModalConfig | None"

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

36 +

37 + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig":

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

39 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig

40 +

41 + if (

42 + self.multimodal_config is not None

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

44 + ):

45 + raise ValueError(

46 + "cudagraph_mm_encoder requires the MiniMax M3 vision tower, "

47 + "but the image limit is 0. Disable cudagraph_mm_encoder."

48 + )

49 +

50 + mm_cfg = self.multimodal_config

51 + if (

52 + mm_cfg is not None

53 + and mm_cfg.mm_encoder_attn_dtype == "fp8"

54 + and mm_cfg.mm_encoder_fp8_scale_path is None

55 + ):

56 + raise ValueError(

57 + "cudagraph_mm_encoder requires static FP8 scales for the "

58 + "MiniMax M3 vision encoder. Set --mm-encoder-fp8-scale-path "

59 + "or disable cudagraph_mm_encoder."

60 + )

61 +

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

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

64 + backend = self.vision_tower.vision_model.attn_backend

65 + if backend not in (

66 + AttentionBackendEnum.FLASH_ATTN,

67 + AttentionBackendEnum.ROCM_AITER_FA,

68 + ):

69 + raise ValueError(

70 + f"cudagraph_mm_encoder is not supported with the "

71 + f"{backend.name} ViT attention backend for MiniMax M3. Set "

72 + "--mm-encoder-attn-backend FLASH_ATTN or ROCM_AITER_FA "

73 + "(ROCm), or disable cudagraph_mm_encoder."

74 + )

75 +

76 + pad_totals = self._encoder_cg_pad_totals

77 +

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

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

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

81 + # full token budget, so a smaller real batch has to be completed

82 + # with one trailing padding sequence; declaring fewer rows than the

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

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

85 + if total is None:

86 + raise RuntimeError("cu_seqlens replay buffer was not registered")

87 + n = src.shape[0]

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

89 + dst[n:] = total

90 +

91 + def pad_rope_planes(dst: torch.Tensor, src: torch.Tensor) -> None:

92 + # cos/sin planes are (3, N, half_rot_dim): the row axis is dim 1,

93 + # not dim 0 as the manager's default slice-copy assumes.

94 + dst.zero_()

95 + dst[:, : src.shape[1]].copy_(src)

96 +

97 + return EncoderCudaGraphConfig(

98 + modalities=["image"],

99 + buffer_keys=[

100 + "pixel_values",

101 + "rotary_cos",

102 + "rotary_sin",

103 + "cu_seqlens",

104 + "max_seqlen",

105 + ],

106 + out_hidden_size=self.vision_tower.out_hidden_size,

107 + padding_logics={

108 + "cu_seqlens": pad_cu_seqlens,

109 + "rotary_cos": pad_rope_planes,

110 + "rotary_sin": pad_rope_planes,

111 + },

112 + )

113 +

114 + def get_encoder_cudagraph_budget_range(

115 + self,

116 + vllm_config: "VllmConfig",

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

118 + min_budget = 64

119 + max_budget = min(

120 + vllm_config.scheduler_config.max_num_batched_tokens,

121 + vllm_config.model_config.max_model_len,

122 + )

123 + return (min_budget, max_budget)

124 +

125 + @staticmethod

126 + def _get_grid_thw_list(mm_kwargs: dict[str, Any]) -> list[list[int]]:

127 + grid_thw = mm_kwargs["image_grid_thw"]

128 + if isinstance(grid_thw, torch.Tensor):

129 + return [[int(x) for x in row] for row in grid_thw.tolist()]

130 + return [[int(x) for x in row] for row in grid_thw]

131 +

132 + @staticmethod

133 + def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor:

134 + pixel_values = mm_kwargs["pixel_values"]

135 + if isinstance(pixel_values, list):

136 + pixel_values = torch.cat(pixel_values)

137 + return pixel_values

138 +

139 + def get_encoder_cudagraph_item_specs(

140 + self, mm_kwargs: dict[str, Any]

141 + ) -> list["EncoderItemSpec"]:

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

143 +

144 + merge = self.vision_tower.spatial_merge_size

145 + return [

146 + EncoderItemSpec(

147 + input_size=t * h * w,

148 + output_tokens=t * h * w // (merge * merge),

149 + )

150 + for t, h, w in self._get_grid_thw_list(mm_kwargs)

151 + ]

152 +

153 + def select_encoder_cudagraph_items(

154 + self, mm_kwargs: dict[str, Any], indices: list[int]

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

156 + grid_thw_list = self._get_grid_thw_list(mm_kwargs)

157 + pixel_values = self._get_pixel_values(mm_kwargs)

158 + source_grid = mm_kwargs["image_grid_thw"]

159 +

160 + if len(indices) == 0:

161 + empty_grid = (

162 + source_grid[:0]

163 + if isinstance(source_grid, torch.Tensor)

164 + else pixel_values.new_zeros((0, 3), dtype=torch.long)

165 + )

166 + return {"pixel_values": pixel_values[:0], "image_grid_thw": empty_grid}

167 +

168 + patch_counts = [t * h * w for t, h, w in grid_thw_list]

169 + cum = [0]

170 + for pc in patch_counts:

171 + cum.append(cum[-1] + pc)

172 + selected_pv = torch.cat(

173 + [pixel_values[cum[i] : cum[i + 1]] for i in indices], dim=0

174 + )

175 + grid_device = (

176 + source_grid.device if isinstance(source_grid, torch.Tensor) else None

177 + )

178 + selected_grid = torch.tensor(

179 + [grid_thw_list[i] for i in indices],

180 + dtype=torch.long,

181 + device=grid_device,

182 + )

183 + return {"pixel_values": selected_pv, "image_grid_thw": selected_grid}

184 +

185 + def prepare_encoder_cudagraph_capture_inputs(

186 + self,

187 + token_budget: int,

188 + max_batch_size: int,

189 + max_frames_per_batch: int,

190 + device: torch.device,

191 + dtype: torch.dtype,

192 + path: str = "default",

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

194 + ) -> "EncoderCudaGraphCaptureInputs":

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

196 + EncoderCudaGraphCaptureInputs,

197 + )

198 +

199 + merge = self.vision_tower.spatial_merge_size

200 + # Output tokens per item in the dummy grid (ceiling so total >= budget).

201 + per_item_out = (token_budget + max_batch_size - 1) // max_batch_size

202 + grid_thw_list = [

203 + [1, merge, per_item_out * merge] for _ in range(max_batch_size)

204 + ]

205 +

206 + embeddings = self.vision_tower.vision_model.embeddings

207 + patch_dim = (

208 + embeddings.num_channels

209 + * embeddings.temporal_patch_size

210 + * embeddings.patch_size**2

211 + )

212 + total_patches = sum(t * h * w for t, h, w in grid_thw_list)

213 + # Match the tower dtype exactly: MiniMaxVLPatchEmbed.forward mutates

214 + # the module on dtype mismatch, which must not happen during capture.

215 + dummy_pixel_values = torch.zeros(

216 + total_patches,

217 + patch_dim,

218 + device=device,

219 + dtype=self.vision_tower.dtype,

220 + )

221 +

222 + # The padding sequence can contain nearly all captured rows when a

223 + # small image replays this graph. Leave one spare cu_seqlens slot.

224 + metadata = self.vision_tower.vision_model.prepare_encoder_metadata(

225 + grid_thw_list,

226 + device=device,

227 + max_batch_size=max_batch_size + 1,

228 + max_seqlen_override=total_patches,

229 + )

230 +

231 + values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values}

232 + values.update({k: v for k, v in metadata.items() if v is not None})

233 + self._encoder_cg_pad_totals[values["cu_seqlens"].data_ptr()] = total_patches

234 + return EncoderCudaGraphCaptureInputs(values=values)

235 +

236 + def prepare_encoder_cudagraph_replay_buffers(

237 + self,

238 + mm_kwargs: dict[str, Any],

239 + max_batch_size: int,

240 + max_frames_per_batch: int,

241 + path: str = "default",

242 + ) -> "EncoderCudaGraphReplayBuffers":

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

244 + EncoderCudaGraphReplayBuffers,

245 + )

246 +

247 + pixel_values = self._get_pixel_values(mm_kwargs)

248 + # Unpadded: pad_cu_seqlens completes the tail with the padding sequence

249 + # so cu_seqlens[-1] matches the captured buffer's row count.

250 + metadata = self.vision_tower.vision_model.prepare_encoder_metadata(

251 + self._get_grid_thw_list(mm_kwargs),

252 + device=pixel_values.device,

253 + )

254 + # FlashAttention reads this CPU scalar during capture, so the captured

255 + # launch must retain the bound for the padded buffer.

256 + metadata.pop("max_seqlen")

257 +

258 + values: dict[str, torch.Tensor | None] = {"pixel_values": pixel_values}

259 + values.update(metadata)

260 + return EncoderCudaGraphReplayBuffers(values=values)

261 +

262 + def encoder_cudagraph_forward(

263 + self,

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

265 + path: str = "default",

266 + ) -> torch.Tensor:

267 + pixel_values = inputs.pop("pixel_values")

268 + # Remaining keys are consumed as encoder_metadata.

269 + return self.vision_tower(pixel_values, grid_thw=None, encoder_metadata=inputs)

270 +

271 + def encoder_eager_forward(

272 + self,

273 + mm_kwargs: dict[str, Any],

274 + path: str = "default",

275 + ) -> torch.Tensor:

276 + return self.vision_tower(

277 + pixel_values=self._get_pixel_values(mm_kwargs).type(

278 + self.vision_tower.dtype

279 + ),

280 + grid_thw=self._get_grid_thw_list(mm_kwargs),

281 + )

vllm/models/minimax_m3/common/vision_tower.py

+82 -15

@@ -392,19 +392,33 @@ class MiniMaxVLVisionTransformer(nn.Module):

392 out.append([min(max_f, t - i), h, w])

393 return out

394

395 - # ── Forward ──────────────────────────────────────────────────────────

395 + # ── Encoder metadata (shared by eager forward and CUDA graph replay) ──

396

397 - def forward(

397 + def prepare_encoder_metadata(

398 self,

399 - pixel_values: torch.Tensor,

399 grid_thw: list[list[int]],

401 - ) -> torch.Tensor:

402 - # pixel_values: (total_N, C * temporal_patch_size * patch_size²)

403 - # Output: (total_N, hidden_size)

404 -

405 - hidden = self.embeddings(pixel_values) # (total_N, hidden_size)

406 - hidden = self.pre_layrnorm(hidden)

407 -

400 + *,

401 + device: torch.device,

402 + max_batch_size: int | None = None,

403 + max_seqlen_override: int | None = None,

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

405 + """Build all grid-dependent forward inputs as standalone tensors.

406 +

407 + The eager path calls this per forward; encoder CUDA graphs call it at

408 + capture (dummy grids, padded to fixed buffer sizes) and before each

409 + replay (real grids, unpadded — the manager's padding logic completes

410 + the captured buffers).

411 +

412 + Args:

413 + grid_thw: per-item [t, h, w] patch grids (host metadata).

414 + device: device for the attention metadata tensors.

415 + max_batch_size: pad cu_seqlens to this many segments with empty

416 + (zero-length) trailing segments; None leaves it unpadded.

417 + max_seqlen_override: bake this max_seqlen instead of deriving it

418 + from the grids; CUDA graph capture passes a worst-case value

419 + because max_seqlen is a CPU scalar baked into the graph.

420 +

421 + """

422 limited = self._apply_max_frames_limit(grid_thw)

423

424 # Token-level cumulative sequence lengths (one segment per limited grid).

@@ -412,15 +426,23 @@ class MiniMaxVLVisionTransformer(nn.Module):

426 cu_seqlens_np = np.zeros(len(lens) + 1, dtype=np.int32)

427 np.cumsum(np.array(lens, dtype=np.int32), out=cu_seqlens_np[1:])

428

429 + if max_batch_size is not None and len(lens) < max_batch_size:

430 + pad = np.full(max_batch_size - len(lens), cu_seqlens_np[-1], dtype=np.int32)

431 + cu_seqlens_np = np.concatenate([cu_seqlens_np, pad])

432 +

433 # Backend-specific encoder metadata. For FLASH_ATTN this returns the raw

434 # token cu_seqlens, the max segment length, and sequence_lengths=None;

435 # for FLASHINFER (cuDNN) it repacks cu_seqlens into element-offset

436 # indptrs, buckets max_seqlen, and builds padded per-sequence lengths.

437 sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(

420 - self.attn_backend, cu_seqlens_np, hidden.device

438 + self.attn_backend, cu_seqlens_np, device

439 )

440 max_seqlen = torch.tensor(

423 - MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens_np),

441 + max_seqlen_override

442 + if max_seqlen_override is not None

443 + else MMEncoderAttention.compute_max_seqlen(

444 + self.attn_backend, cu_seqlens_np

445 + ),

446 dtype=torch.int32,

447 )

448 cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(

@@ -428,7 +450,7 @@ class MiniMaxVLVisionTransformer(nn.Module):

450 cu_seqlens_np,

451 self.hidden_size,

452 self.tp_size,

431 - hidden.device,

453 + device,

454 )

455

456 # 3D RoPE cos/sin: (3, total_N, half_rot_dim) fp32 t/h/w planes for

@@ -442,6 +464,45 @@ class MiniMaxVLVisionTransformer(nn.Module):

464 self.spatial_merge_size,

465 )

466

467 + return {

468 + "rotary_cos": rotary_cos,

469 + "rotary_sin": rotary_sin,

470 + "cu_seqlens": cu_seqlens,

471 + "max_seqlen": max_seqlen,

472 + "sequence_lengths": sequence_lengths,

473 + }

474 +

475 + # ── Forward ──────────────────────────────────────────────────────────

476 +

477 + def forward(

478 + self,

479 + pixel_values: torch.Tensor,

480 + grid_thw: list[list[int]] | None = None,

481 + *,

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

483 + ) -> torch.Tensor:

484 + # pixel_values: (total_N, C * temporal_patch_size * patch_size²)

485 + # Output: (total_N, hidden_size)

486 +

487 + hidden = self.embeddings(pixel_values) # (total_N, hidden_size)

488 + hidden = self.pre_layrnorm(hidden)

489 +

490 + if encoder_metadata is None:

491 + assert grid_thw is not None

492 + encoder_metadata = self.prepare_encoder_metadata(

493 + grid_thw, device=hidden.device

494 + )

495 +

496 + cu_seqlens = encoder_metadata["cu_seqlens"]

497 + rotary_cos = encoder_metadata["rotary_cos"]

498 + rotary_sin = encoder_metadata["rotary_sin"]

499 + max_seqlen = encoder_metadata["max_seqlen"]

500 + sequence_lengths = encoder_metadata.get("sequence_lengths")

501 + assert cu_seqlens is not None

502 + assert rotary_cos is not None

503 + assert rotary_sin is not None

504 + assert max_seqlen is not None

505 +

506 # Encoder expects (N, 1, hidden_size) — add batch dim

507 hidden = hidden.unsqueeze(1)

508

@@ -600,9 +661,15 @@ class MiniMaxVLVisionModel(nn.Module):

661 def forward(

662 self,

663 pixel_values: torch.Tensor,

603 - grid_thw: list[list[int]],

664 + grid_thw: list[list[int]] | None = None,

665 + *,

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

667 ) -> torch.Tensor:

605 - hidden = self.vision_model(pixel_values=pixel_values, grid_thw=grid_thw)

668 + hidden = self.vision_model(

669 + pixel_values=pixel_values,

670 + grid_thw=grid_thw,

671 + encoder_metadata=encoder_metadata,

672 + )

673 if hidden.dim() == 3:

674 hidden = hidden.squeeze(0)

675 hidden = self.multi_modal_projector(hidden)

vllm/models/minimax_m3/nvidia/model.py

+8 -1

@@ -70,6 +70,9 @@ from vllm.model_executor.models.utils import (

70 maybe_prefix,

71 )

72 from vllm.model_executor.models.vision import run_dp_sharded_mrope_vision_model

73 +from vllm.models.minimax_m3.common.encoder_cudagraph import (

74 + MiniMaxM3EncoderCudaGraphMixin,

75 +)

76 from vllm.models.minimax_m3.common.indexer import MiniMaxM3Indexer

77 from vllm.models.minimax_m3.common.mm_preprocess import (

78 MiniMaxM3VLDummyInputsBuilder,

@@ -1076,7 +1079,11 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):

1079 dummy_inputs=MiniMaxM3VLDummyInputsBuilder,

1080 )

1081 class MiniMaxM3SparseForConditionalGeneration(

1079 - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3

1082 + nn.Module,

1083 + SupportsMultiModal,

1084 + MiniMaxM3EncoderCudaGraphMixin,

1085 + SupportsPP,

1086 + SupportsEagle3,

1087 ):

1088 """Top-level (VL) entry point for MiniMax M3.

1089