class MiniMaxM3EncoderCudaGraphMixin(SupportsEncoderCudaGraph):
"""SupportsEncoderCudaGraph for hosts with ``self.vision_tower``."""
vision_tower: "MiniMaxVLVisionModel"
multimodal_config: "MultiModalConfig | None"
_encoder_cg_pad_totals: dict[int, int] = {}
def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig":
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
if (
self.multimodal_config is not None
and self.multimodal_config.get_limit_per_prompt("image") == 0
):
raise ValueError(
"cudagraph_mm_encoder requires the MiniMax M3 vision tower, "
"but the image limit is 0. Disable cudagraph_mm_encoder."
)
mm_cfg = self.multimodal_config
if (
mm_cfg is not None
and mm_cfg.mm_encoder_attn_dtype == "fp8"
and mm_cfg.mm_encoder_fp8_scale_path is None
):
raise ValueError(
"cudagraph_mm_encoder requires static FP8 scales for the "
"MiniMax M3 vision encoder. Set --mm-encoder-fp8-scale-path "
"or disable cudagraph_mm_encoder."
)
# FlashInfer reads max_seqlen on the host and TORCH_SDPA calls
# .tolist() on CUDA cu_seqlens; neither can be captured.
backend = self.vision_tower.vision_model.attn_backend
if backend not in (
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.ROCM_AITER_FA,
):
raise ValueError(
f"cudagraph_mm_encoder is not supported with the "
f"{backend.name} ViT attention backend for MiniMax M3. Set "
"--mm-encoder-attn-backend FLASH_ATTN or ROCM_AITER_FA "
"(ROCm), or disable cudagraph_mm_encoder."
)
pad_totals = self._encoder_cg_pad_totals
def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None:
# Varlen attention requires cu_seqlens[-1] to equal the number of
# rows actually passed in. The captured buffers are sized for the
# full token budget, so a smaller real batch has to be completed
# with one trailing padding sequence; declaring fewer rows than the
# buffer holds is undefined behaviour and returns NaN on FlashAttn.
total = pad_totals.get(dst.data_ptr())
if total is None:
raise RuntimeError("cu_seqlens replay buffer was not registered")
n = src.shape[0]
dst[:n].copy_(src[:n])
dst[n:] = total
def pad_rope_planes(dst: torch.Tensor, src: torch.Tensor) -> None:
# cos/sin planes are (3, N, half_rot_dim): the row axis is dim 1,
# not dim 0 as the manager's default slice-copy assumes.
dst.zero_()
dst[:, : src.shape[1]].copy_(src)
return EncoderCudaGraphConfig(
modalities=["image"],
buffer_keys=[
"pixel_values",
"rotary_cos",
"rotary_sin",
"cu_seqlens",
"max_seqlen",
],
out_hidden_size=self.vision_tower.out_hidden_size,
padding_logics={
"cu_seqlens": pad_cu_seqlens,
"rotary_cos": pad_rope_planes,
"rotary_sin": pad_rope_planes,
},
)
def get_encoder_cudagraph_budget_range(
self,
vllm_config: "VllmConfig",
) -> tuple[int, int]:
min_budget = 64
max_budget = min(
vllm_config.scheduler_config.max_num_batched_tokens,
vllm_config.model_config.max_model_len,
)
return (min_budget, max_budget)
@staticmethod
def _get_grid_thw_list(mm_kwargs: dict[str, Any]) -> list[list[int]]:
grid_thw = mm_kwargs["image_grid_thw"]
if isinstance(grid_thw, torch.Tensor):
return [[int(x) for x in row] for row in grid_thw.tolist()]
return [[int(x) for x in row] for row in grid_thw]
@staticmethod
def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor:
pixel_values = mm_kwargs["pixel_values"]
if isinstance(pixel_values, list):
pixel_values = torch.cat(pixel_values)
return pixel_values
def get_encoder_cudagraph_item_specs(
self, mm_kwargs: dict[str, Any]
) -> list["EncoderItemSpec"]:
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
merge = self.vision_tower.spatial_merge_size
return [
EncoderItemSpec(
input_size=t * h * w,
output_tokens=t * h * w // (merge * merge),
)
for t, h, w in self._get_grid_thw_list(mm_kwargs)
]
def select_encoder_cudagraph_items(
self, mm_kwargs: dict[str, Any], indices: list[int]
) -> dict[str, Any]:
grid_thw_list = self._get_grid_thw_list(mm_kwargs)
pixel_values = self._get_pixel_values(mm_kwargs)
source_grid = mm_kwargs["image_grid_thw"]
if len(indices) == 0:
empty_grid = (
source_grid[:0]
if isinstance(source_grid, torch.Tensor)
else pixel_values.new_zeros((0, 3), dtype=torch.long)
)
return {"pixel_values": pixel_values[:0], "image_grid_thw": empty_grid}
patch_counts = [t * h * w for t, h, w in grid_thw_list]
cum = [0]
for pc in patch_counts:
cum.append(cum[-1] + pc)
selected_pv = torch.cat(
[pixel_values[cum[i] : cum[i + 1]] for i in indices], dim=0
)
grid_device = (
source_grid.device if isinstance(source_grid, torch.Tensor) else None
)
selected_grid = torch.tensor(
[grid_thw_list[i] for i in indices],
dtype=torch.long,
device=grid_device,
)
return {"pixel_values": selected_pv, "image_grid_thw": selected_grid}
def prepare_encoder_cudagraph_capture_inputs(
self,
token_budget: int,
max_batch_size: int,
max_frames_per_batch: int,
device: torch.device,
dtype: torch.dtype,
path: str = "default",
axis_keys: tuple[Hashable, ...] | None = None,
) -> "EncoderCudaGraphCaptureInputs":
from vllm.v1.worker.encoder_cudagraph_defs import (
EncoderCudaGraphCaptureInputs,
)
merge = self.vision_tower.spatial_merge_size
# Output tokens per item in the dummy grid (ceiling so total >= budget).
per_item_out = (token_budget + max_batch_size - 1) // max_batch_size
grid_thw_list = [
[1, merge, per_item_out * merge] for _ in range(max_batch_size)
]
embeddings = self.vision_tower.vision_model.embeddings
patch_dim = (
embeddings.num_channels
* embeddings.temporal_patch_size
* embeddings.patch_size**2
)
total_patches = sum(t * h * w for t, h, w in grid_thw_list)
# Match the tower dtype exactly: MiniMaxVLPatchEmbed.forward mutates
# the module on dtype mismatch, which must not happen during capture.
dummy_pixel_values = torch.zeros(
total_patches,
patch_dim,
device=device,
dtype=self.vision_tower.dtype,
)
# The padding sequence can contain nearly all captured rows when a
# small image replays this graph. Leave one spare cu_seqlens slot.
metadata = self.vision_tower.vision_model.prepare_encoder_metadata(
grid_thw_list,
device=device,
max_batch_size=max_batch_size + 1,
max_seqlen_override=total_patches,
)
values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values}
values.update({k: v for k, v in metadata.items() if v is not None})
self._encoder_cg_pad_totals[values["cu_seqlens"].data_ptr()] = total_patches
return EncoderCudaGraphCaptureInputs(values=values)
def prepare_encoder_cudagraph_replay_buffers(
self,
mm_kwargs: dict[str, Any],
max_batch_size: int,
max_frames_per_batch: int,
path: str = "default",
) -> "EncoderCudaGraphReplayBuffers":
from vllm.v1.worker.encoder_cudagraph_defs import (
EncoderCudaGraphReplayBuffers,
)
pixel_values = self._get_pixel_values(mm_kwargs)
# Unpadded: pad_cu_seqlens completes the tail with the padding sequence
# so cu_seqlens[-1] matches the captured buffer's row count.
metadata = self.vision_tower.vision_model.prepare_encoder_metadata(
self._get_grid_thw_list(mm_kwargs),
device=pixel_values.device,
)
# FlashAttention reads this CPU scalar during capture, so the captured
# launch must retain the bound for the padded buffer.
metadata.pop("max_seqlen")
values: dict[str, torch.Tensor | None] = {"pixel_values": pixel_values}
values.update(metadata)
return EncoderCudaGraphReplayBuffers(values=values)
def encoder_cudagraph_forward(
self,
inputs: dict[str, torch.Tensor],
path: str = "default",
) -> torch.Tensor:
pixel_values = inputs.pop("pixel_values")
# Remaining keys are consumed as encoder_metadata.
return self.vision_tower(pixel_values, grid_thw=None, encoder_metadata=inputs)
def encoder_eager_forward(
self,
mm_kwargs: dict[str, Any],
path: str = "default",
) -> torch.Tensor:
return self.vision_tower(
pixel_values=self._get_pixel_values(mm_kwargs).type(
self.vision_tower.dtype
),
grid_thw=self._get_grid_thw_list(mm_kwargs),
)