@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 777b01d1a8

[MM][CG] Support ViT full CUDA graph for Ernie-4.5-VL image inference (#45254)

Signed-off-by: Qiuyang Yue <yueqiuyang1389@gmail.com>

Qiuyang Yue committed Aug 5, 2026 at 22:04 UTC 777b01d1a86e4036fdd7368a63ca0c4d4803b488

4 files changed +289 -13

docs/design/cuda_graphs_multimodal.md

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

119 | ------------ | ------ | ------------ | ------------ | --------------- |

120 | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |

121 +| `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ |

122 | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ |

123 | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |

124 | `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ |

examples/generate/multimodal/vision_language_offline.py

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

2426 "stepvl",

2427 "glm4_1v",

2428 "deepseek_ocr",

2429 + "ernie45_vl",

2430 ]

2431

2432

tests/models/multimodal/generation/test_vit_cudagraph.py

+27

@@ -66,6 +66,14 @@ def gemma3_chat_template(content: str) -> str:

66 return f"<bos><start_of_turn>user\n{content}<end_of_turn>\n<start_of_turn>model\n"

69 +def ernie45_vl_chat_template(content: str) -> str:

70 + return (

71 + f"<|begin_of_sentence|>User: {content}"

72 + "Picture 1:<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>\n"

73 + "Assistant: <think></think>"

74 + )

75 +

76 +

77 MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

78 "gemma3": VitCudagraphTestConfig(

79 model="google/gemma-3-4b-it",

@@ -208,6 +216,25 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

216 ),

217 },

218 ),

219 + "ernie45_vl": VitCudagraphTestConfig(

220 + model="baidu/ERNIE-4.5-VL-28B-A3B-PT",

221 + # Image only: Ernie's resampler applies a temporal conv for video

222 + # that changes the output token count, so video uses the eager path.

223 + modalities=["image"],

224 + image_prompt=ernie45_vl_chat_template("What is in this image?"),

225 + # Shrink to 1 text + 1 vision layer with random weights so the test

226 + # runs on any CI GPU and skips the ~56 GiB weight download. The test

227 + # only validates encoder CG capture/replay, not output quality.

228 + vllm_runner_kwargs={

229 + "load_format": "dummy",

230 + "trust_remote_code": True,

231 + "revision": "refs/pr/17",

232 + "hf_overrides": partial(

233 + dummy_hf_overrides,

234 + model_arch="Ernie4_5_VLMoeForConditionalGeneration",

235 + ),

236 + },

237 + ),

238 "glm4_1v": VitCudagraphTestConfig(

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

240 image_prompt=(

vllm/model_executor/models/ernie45_vl.py

+260 -13

@@ -77,12 +77,18 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum

77 from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM

78 from .interfaces import (

79 MultiModalEmbeddings,

80 + SupportsEncoderCudaGraph,

81 SupportsLoRA,

82 SupportsMRoPE,

83 SupportsMultiModal,

84 SupportsPP,

85 )

85 -from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix

86 +from .utils import (

87 + AutoWeightsLoader,

88 + WeightsMapper,

89 + maybe_prefix,

90 + scatter_output_slices,

91 +)

92 from .vision import get_vit_attn_backend

94 logger = init_logger(__name__)

@@ -453,13 +459,43 @@ class Ernie4_5_VisionTransformer(nn.Module):

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

460 return max_seqlen

461

456 - def forward(

457 - self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, num_pad=0

458 - ) -> torch.Tensor:

459 - hidden_states = self.patch_embed(hidden_states)

460 -

461 - rotary_pos_emb = self.rot_pos_emb(grid_thw)

462 - rotary_pos_emb = rotary_pos_emb.to(hidden_states.device)

462 + def prepare_encoder_metadata(

463 + self,

464 + grid_thw: torch.Tensor | list[list[int]],

465 + *,

466 + num_pad: int = 0,

467 + max_batch_size: int | None = None,

468 + max_seqlen_override: int | None = None,

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

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

471 + """Compute encoder metadata outside the CUDA graph.

472 +

473 + Splits the rotary embeddings, ``cu_seqlens`` and ``max_seqlen`` out of

474 + `forward` so they can be precomputed on the host side and fed into a

475 + captured graph through fixed buffers. Shared by the eager path, CUDA

476 + graph capture and CUDA graph replay

477 +

478 + Args:

479 + grid_thw: Per-frame grid sizes as a ``[num_frames, 3]`` tensor or

480 + an equivalent list of ``[t, h, w]``.

481 + num_pad: Legacy ``cu_seqlens`` padding

482 + max_batch_size: If set, pad ``cu_seqlens`` to ``max_batch_size + 1``

483 + entries for a fixed graph buffer shape.

484 + max_seqlen_override: If set, use this value for ``max_seqlen``

485 + instead of deriving it from ``cu_seqlens``. CUDA graph capture

486 + bakes ``max_seqlen`` at capture time, so it must cover the

487 + worst-case replay.

488 + device: Device for the returned tensors. Defaults to ``self.device``.

489 +

490 + Returns:

491 + Dict with ``rotary_pos_emb``, ``cu_seqlens`` and ``max_seqlen``.

492 + """

493 + if device is None:

494 + device = self.device

495 + if not isinstance(grid_thw, torch.Tensor):

496 + grid_thw = torch.tensor(grid_thw, dtype=torch.int32)

497 +

498 + rotary_pos_emb = self.rot_pos_emb(grid_thw).to(device)

499

500 cu_seqlens = torch.repeat_interleave(

501 grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

@@ -472,14 +508,60 @@ class Ernie4_5_VisionTransformer(nn.Module):

508 else:

509 cu_seqlens = torch.cat([zeros, cu_seqlens])

510

511 + # Pad cu_seqlens to a fixed number of sequences so the CUDA graph

512 + # replay buffer keeps a constant shape across batches.

513 + if max_batch_size is not None:

514 + num_seqs = cu_seqlens.numel() - 1

515 + if num_seqs < max_batch_size:

516 + cu_seqlens = torch.cat(

517 + [

518 + cu_seqlens,

519 + cu_seqlens.new_full(

520 + (max_batch_size - num_seqs,), cu_seqlens[-1]

521 + ),

522 + ]

523 + )

524 +

525 + # max_seqlen is baked into the graph at capture time, so capture passes

526 + # a worst-case override. compute_attn_mask_seqlen keeps it on CPU

527 + # (consumed via .item()), avoiding a captured D2H copy.

528 + if max_seqlen_override is not None:

529 + max_seqlen = torch.tensor(max_seqlen_override, dtype=torch.int32)

530 + else:

531 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

532 +

533 + cu_seqlens = cu_seqlens.to(device)

534 +

535 + return {

536 + "rotary_pos_emb": rotary_pos_emb,

537 + "cu_seqlens": cu_seqlens,

538 + "max_seqlen": max_seqlen,

539 + }

540 +

541 + def forward(

542 + self,

543 + hidden_states: torch.Tensor,

544 + grid_thw: torch.Tensor | None = None,

545 + num_pad: int = 0,

546 + *,

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

548 + ) -> torch.Tensor:

549 + hidden_states = self.patch_embed(hidden_states)

550 +

551 + if encoder_metadata is None:

552 + # Eager path: compute metadata inline

553 + encoder_metadata = self.prepare_encoder_metadata(

554 + grid_thw, num_pad=num_pad, device=hidden_states.device

555 + )

556 + rotary_pos_emb = encoder_metadata["rotary_pos_emb"]

557 + cu_seqlens = encoder_metadata["cu_seqlens"]

558 + max_seqlen = encoder_metadata["max_seqlen"]

559 +

560 # add batch size

561 if hidden_states.ndim == 2:

562 hidden_states = hidden_states.unsqueeze(dim=1)

563

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

480 - max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

481 -

482 - for i, blk in enumerate(self.blocks):

564 + for blk in self.blocks:

565 hidden_states = blk(

566 hidden_states,

567 cu_seqlens=cu_seqlens,

@@ -1256,7 +1338,12 @@ class Ernie4_5_VLDummyInputsBuilder(BaseDummyInputsBuilder[Ernie4_5_VLProcessing

1338 dummy_inputs=Ernie4_5_VLDummyInputsBuilder,

1339 )

1340 class Ernie4_5_VLMoeForConditionalGeneration(

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

1341 + nn.Module,

1342 + SupportsMultiModal,

1343 + SupportsLoRA,

1344 + SupportsPP,

1345 + SupportsMRoPE,

1346 + SupportsEncoderCudaGraph,

1347 ):

1348 packed_modules_mapping = {

1349 "qkv_proj": [

@@ -1489,6 +1576,166 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1576 video_grid_thw=video_grid_thw,

1577 )

1578

1579 + # -- SupportsEncoderCudaGraph protocol methods --

1580 + #

1581 + # Only the vision transformer (self.vision_model) is captured into the

1582 + # CUDA graph. The VariableResolutionResamplerModel (spatial merge +

1583 + # projection) runs eagerly in encoder_eager_forward / postprocess_

1584 + # encoder_output because its temporal path performs host-side (numpy)

1585 + # indexing that cannot be captured. Image only: video uses a temporal

1586 + # conv that changes the output token count, so it falls back to the

1587 + # eager multimodal path for now.

1588 +

1589 + def get_encoder_cudagraph_config(self):

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

1591 +

1592 + return EncoderCudaGraphConfig(

1593 + modalities=["image"],

1594 + # Ernie consumes a single rotary freqs tensor (not cos/sin)

1595 + buffer_keys=[

1596 + "pixel_values",

1597 + "rotary_pos_emb",

1598 + "cu_seqlens",

1599 + "max_seqlen",

1600 + ],

1601 + # Post-merge embeddings, produced in postprocess_encoder_output,

1602 + # are at the LM hidden dim, used only for DP gather sizing.

1603 + out_hidden_size=self.config.hidden_size,

1604 + )

1605 +

1606 + def get_encoder_cudagraph_budget_range(self, vllm_config) -> tuple[int, int]:

1607 + # Min: a 224x224 image -> 16x16 patches (patch_size=14),

1608 + # spatial_merge_size=2 -> 8x8 = 64 output tokens.

1609 + min_budget = 64

1610 + max_budget = min(

1611 + vllm_config.scheduler_config.max_num_batched_tokens,

1612 + vllm_config.model_config.max_model_len,

1613 + )

1614 + return (min_budget, max_budget)

1615 +

1616 + def get_encoder_cudagraph_item_specs(self, mm_kwargs: dict[str, Any]):

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

1618 +

1619 + m = self.vision_model.spatial_merge_size

1620 + grid_thw = mm_kwargs["image_grid_thw"].tolist()

1621 + return [

1622 + EncoderItemSpec(

1623 + input_size=t * h * w,

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

1625 + )

1626 + for t, h, w in grid_thw

1627 + ]

1628 +

1629 + def select_encoder_cudagraph_items(

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

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

1632 + grid_thw = mm_kwargs["image_grid_thw"]

1633 + pixel_values = mm_kwargs["pixel_values"]

1634 +

1635 + if len(indices) == 0:

1636 + return {"pixel_values": pixel_values[:0], "image_grid_thw": grid_thw[:0]}

1637 +

1638 + # Cumulative patch offsets for slicing the concatenated pixel_values.

1639 + patches_per_item = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).tolist()

1640 + cum_patches = [0]

1641 + for p in patches_per_item:

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

1643 + selected_pv = torch.cat(

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

1645 + )

1646 + return {"pixel_values": selected_pv, "image_grid_thw": grid_thw[indices]}

1647 +

1648 + def prepare_encoder_cudagraph_capture_inputs(

1649 + self,

1650 + token_budget,

1651 + max_batch_size,

1652 + max_frames_per_batch,

1653 + device,

1654 + dtype,

1655 + path: str = "default",

1656 + ):

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

1658 + EncoderCudaGraphCaptureInputs,

1659 + )

1660 +

1661 + m = self.vision_model.spatial_merge_size

1662 + # Ceil so the buffer fits the worst case of one item using the full

1663 + # budget.

1664 + per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size

1665 + # Image-format grid (T=1): h=m, w=per_mm_item_output*m, so the merged

1666 + # token count per item equals per_mm_item_output.

1667 + grid_config = [[1, m, per_mm_item_output * m] for _ in range(max_batch_size)]

1668 +

1669 + patch_embed = self.vision_model.patch_embed

1670 + in_channels = patch_embed.in_channels

1671 + patch_size = patch_embed.patch_size

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

1673 + flattened_patch_size = in_channels * patch_size * patch_size

1674 + dummy_pixel_values = torch.randn(

1675 + total_patches, flattened_patch_size, device=device, dtype=dtype

1676 + )

1677 +

1678 + # max_seqlen is baked at capture: worst case is one item consuming the

1679 + # full budget - seq_len = token_budget * spatial_merge_size**2.

1680 + metadata = self.vision_model.prepare_encoder_metadata(

1681 + grid_config,

1682 + max_batch_size=max_batch_size,

1683 + max_seqlen_override=token_budget * (m**2),

1684 + device=device,

1685 + )

1686 + values = metadata | {"pixel_values": dummy_pixel_values}

1687 + return EncoderCudaGraphCaptureInputs(values=values)

1688 +

1689 + def prepare_encoder_cudagraph_replay_buffers(

1690 + self, mm_kwargs, max_batch_size, max_frames_per_batch, path: str = "default"

1691 + ):

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

1693 + EncoderCudaGraphReplayBuffers,

1694 + )

1695 +

1696 + metadata = self.vision_model.prepare_encoder_metadata(

1697 + mm_kwargs["image_grid_thw"].tolist(), max_batch_size=max_batch_size

1698 + )

1699 + values = metadata | {

1700 + "pixel_values": mm_kwargs["pixel_values"],

1701 + }

1702 + return EncoderCudaGraphReplayBuffers(values=values)

1703 +

1704 + def encoder_cudagraph_forward(

1705 + self, values: dict[str, torch.Tensor], path: str = "default"

1706 + ) -> torch.Tensor:

1707 + # Graph captures the ViT only, and the resampler runs in

1708 + # postprocess_encoder_output.

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

1710 + return self.vision_model(pixel_values, encoder_metadata=values)

1711 +

1712 + def encoder_eager_forward(

1713 + self, mm_kwargs: dict[str, Any], path: str = "default"

1714 + ) -> torch.Tensor:

1715 + # Eager fallback: run the full pipeline (ViT + resampler). The result

1716 + # is scattered directly, so it must be the post-merge embeddings.

1717 + pixel_values = mm_kwargs["pixel_values"].type(self.vision_model.dtype)

1718 + grid_thw = mm_kwargs["image_grid_thw"].to(self.vision_model.device)

1719 + image_features = self.vision_model(pixel_values, grid_thw)

1720 + return self.resampler_model(image_features, grid_thw)

1721 +

1722 + def postprocess_encoder_output(

1723 + self,

1724 + output: torch.Tensor,

1725 + indices: list[int],

1726 + per_item_out_tokens: list[int],

1727 + dest,

1728 + clone: bool = False,

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

1730 + ) -> None:

1731 + # The graph output is the raw ViT features (pre-merge), padded to the

1732 + # token budget. Run the resampler eagerly on the valid portion using

1733 + # the actual batch grid_thw, then scatter the post-merge embeddings.

1734 + grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device)

1735 + num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum())

1736 + image_embeds = self.resampler_model(output[:num_valid], grid_thw)

1737 + scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone)

1738 +

1739 def _process_image_input(

1740 self, image_input: Ernie4_5_VLImageInputs

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