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

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 19:00 UTC

[MM][CG] Support ViT CG for Qwen2-VL (#41736)

Signed-off-by: John Calderon <jcalderon@nvidia.com> Co-authored-by: Roger Wang <hey@rogerw.io>

John Calderon committed May 13, 2026 at 13:52 UTC b3c69595a63f14243388d8632e263490457d6126

4 files changed +315 -21

docs/design/cuda_graphs_multimodal.md

+2 -1

@@ -85,13 +85,14 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra

87 | ------------ | ------ | ------------ | ------------ |

88 +| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |

89 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |

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

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

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.

94 - For Qwen2.5-VL only FA2 and FA3 has been tested.

95 + For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested.

97 ## Configuration

examples/generate/multimodal/vision_language_offline.py

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

2557 "qwen2_5_vl",

2558 "qwen3_vl",

2559 "qwen3_vl_moe",

2560 + "qwen2_vl",

2561 "qwen3_5",

2562 "qwen3_5_moe",

2563 ]

tests/models/multimodal/generation/test_vit_cudagraph.py

+12

@@ -78,6 +78,18 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {

78 needs_video_metadata=True,

79 marks=[pytest.mark.core_model],

80 ),

81 + "qwen2_vl": VitCudagraphTestConfig(

82 + model="Qwen/Qwen2-VL-2B-Instruct",

83 + image_prompt=qwen_vl_chat_template(

84 + "<|vision_start|><|image_pad|><|vision_end|>What is in this image?"

85 + ),

86 + video_prompt=qwen_vl_chat_template(

87 + "<|vision_start|><|video_pad|><|vision_end|>"

88 + "Describe this video in one sentence."

89 + ),

90 + needs_video_metadata=False,

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

92 + ),

93 }

vllm/model_executor/models/qwen2_vl.py

+300 -20

@@ -89,9 +89,11 @@ from vllm.sequence import IntermediateTensors

89 from vllm.tokenizers import TokenizerLike

90 from vllm.utils.tensor_schema import TensorSchema, TensorShape

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

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

94 from .interfaces import (

95 MultiModalEmbeddings,

96 + SupportsEncoderCudaGraph,

97 SupportsLoRA,

98 SupportsMRoPE,

99 SupportsMultiModal,

@@ -646,38 +648,84 @@ class Qwen2VisionTransformer(nn.Module):

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

649 return max_seqlen

650

651 + def prepare_encoder_metadata(

652 + self,

653 + grid_thw: list[list[int]],

654 + *,

655 + max_batch_size: int | None = None,

656 + max_frames_per_batch: int | None = None,

657 + max_seqlen_override: int | None = None,

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

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

660 + if device is None:

661 + device = self.device

662 +

663 + # Position embeddings.

664 + rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw)

665 +

666 + grid_thw_np = np.array(grid_thw, dtype=np.int32)

667 + cu_seqlens = np.repeat(

668 + grid_thw_np[:, 1] * grid_thw_np[:, 2],

669 + grid_thw_np[:, 0],

670 + ).cumsum(dtype=np.int32)

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

672 + cu_seqlens = torch.from_numpy(cu_seqlens)

673 +

674 + # Keep cu_seqlens shape stable across capture/replay.

675 + pad_to = (

676 + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size

677 + )

678 + if pad_to is not None:

679 + num_seqs = len(cu_seqlens) - 1

680 + if num_seqs < pad_to:

681 + cu_seqlens = torch.cat(

682 + (

683 + cu_seqlens,

684 + torch.full(

685 + (pad_to - num_seqs,),

686 + cu_seqlens[-1],

687 + dtype=cu_seqlens.dtype,

688 + device=cu_seqlens.device,

689 + ),

690 + )

691 + )

692 +

693 + # Compute (or override) max seqlen used by FA-style backends.

694 + if max_seqlen_override is None:

695 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

696 + else:

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

698 +

699 + return {

700 + "rotary_pos_emb_cos": rotary_pos_emb_cos,

701 + "rotary_pos_emb_sin": rotary_pos_emb_sin,

702 + "cu_seqlens": cu_seqlens.to(device=device, non_blocking=True),

703 + "max_seqlen": max_seqlen,

704 + }

705 +

706 def forward(

707 self,

708 x: torch.Tensor,

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

710 + *,

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

712 ) -> torch.Tensor:

713 # patchify

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

715 x = self.patch_embed(x)

716

658 - if isinstance(grid_thw, list):

659 - grid_thw_list = grid_thw

660 - grid_thw = np.array(grid_thw, dtype=np.int32)

661 - else:

662 - grid_thw_list = grid_thw.tolist()

663 - grid_thw = grid_thw.numpy()

717 + grid_thw_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist()

718

665 - # compute position embedding

666 - rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)

719 + if encoder_metadata is None:

720 + encoder_metadata = self.prepare_encoder_metadata(grid_thw_list)

721

668 - # compute cu_seqlens

669 - cu_seqlens = np.repeat(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(

670 - axis=0, dtype=np.int32

671 - )

672 - cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])

673 - cu_seqlens = torch.from_numpy(cu_seqlens)

722 + rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"]

723 + rotary_pos_emb_sin = encoder_metadata["rotary_pos_emb_sin"]

724 + cu_seqlens = encoder_metadata["cu_seqlens"]

725 + max_seqlen = encoder_metadata["max_seqlen"]

726

727 # transformers

728 x = x.unsqueeze(1)

677 -

678 - # pre-compute seqlens for attn mask to reduce cuMemcpy operations

679 - max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

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

729 for blk in self.blocks:

730 x = blk(

731 x,

@@ -1127,7 +1175,12 @@ class Qwen2VLMultiModalProcessor(BaseMultiModalProcessor[Qwen2VLProcessingInfo])

1175 dummy_inputs=Qwen2VLDummyInputsBuilder,

1176 )

1177 class Qwen2VLForConditionalGeneration(

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

1178 + nn.Module,

1179 + SupportsMultiModal,

1180 + SupportsLoRA,

1181 + SupportsPP,

1182 + SupportsMRoPE,

1183 + SupportsEncoderCudaGraph,

1184 ):

1185 # To ensure correct weight loading and mapping.

1186 hf_to_vllm_mapper = WeightsMapper(

@@ -1234,7 +1287,7 @@ class Qwen2VLForConditionalGeneration(

1287 config: Qwen2VLConfig = vllm_config.model_config.hf_config

1288 quant_config = vllm_config.quant_config

1289 multimodal_config = vllm_config.model_config.multimodal_config

1237 -

1290 + self.model_config = vllm_config.model_config

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

1292 self.config = config

1293 self.multimodal_config = multimodal_config

@@ -1396,6 +1449,233 @@ class Qwen2VLForConditionalGeneration(

1449

1450 return multimodal_embeddings

1451

1452 + # -- SupportsEncoderCudaGraph protocol methods --

1453 +

1454 + def get_encoder_cudagraph_config(self):

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

1456 + EncoderCudaGraphConfig,

1457 + )

1458 +

1459 + return EncoderCudaGraphConfig(

1460 + modalities=["image", "video"],

1461 + input_key_by_modality={

1462 + "image": "pixel_values",

1463 + "video": "pixel_values_videos",

1464 + },

1465 + buffer_keys=[

1466 + "rotary_pos_emb_cos",

1467 + "rotary_pos_emb_sin",

1468 + "cu_seqlens",

1469 + "max_seqlen",

1470 + ],

1471 + out_hidden_size=self.visual.out_hidden_size,

1472 + )

1473 +

1474 + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:

1475 + if "image_grid_thw" in mm_kwargs:

1476 + return "image"

1477 + return "video"

1478 +

1479 + def get_max_frames_per_video(self) -> int:

1480 + mm_registry = MULTIMODAL_REGISTRY

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

1482 + max_frames_per_video = info.get_num_frames_with_most_features(

1483 + seq_len=self.model_config.max_model_len,

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

1485 + )

1486 + return max_frames_per_video

1487 +

1488 + def get_encoder_cudagraph_budget_range(

1489 + self,

1490 + vllm_config: VllmConfig,

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

1492 + # Min: estimated smallest possible encoder input.

1493 + # 224x224 image -> 16x16 patches (patch_size=14)

1494 + # spatial_merge_size=2 -> 8x8 = 64 tokens

1495 + min_budget = 64

1496 + # Max: capped by max_num_batched_tokens

1497 + max_budget = min(

1498 + vllm_config.scheduler_config.max_num_batched_tokens,

1499 + self.model_config.max_model_len,

1500 + )

1501 + return (min_budget, max_budget)

1502 +

1503 + def _get_pixel_values_by_modality(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:

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

1505 + pixel_values = mm_kwargs["pixel_values"]

1506 + else:

1507 + pixel_values = mm_kwargs["pixel_values_videos"]

1508 + return pixel_values

1509 +

1510 + def _get_grid_thw_by_modality(self, mm_kwargs: dict[str, Any]) -> list[list[int]]:

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

1512 + grid_thw = mm_kwargs[grid_thw_key]

1513 + if not isinstance(grid_thw, list):

1514 + grid_thw = grid_thw.tolist()

1515 + return grid_thw

1516 +

1517 + def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:

1518 + return len(self._get_grid_thw_by_modality(mm_kwargs))

1519 +

1520 + def get_encoder_cudagraph_per_item_output_tokens(

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

1522 + ) -> list[int]:

1523 + m = self.visual.spatial_merge_size

1524 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1525 + return [t * (h // m) * (w // m) for t, h, w in grid_thw]

1526 +

1527 + def get_encoder_cudagraph_per_item_input_sizes(

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

1529 + ) -> list[int]:

1530 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1531 + return [t * h * w for t, h, w in grid_thw]

1532 +

1533 + def select_encoder_cudagraph_items(

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

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

1536 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1537 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1538 +

1539 + if len(indices) == 0:

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

1541 + return {

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

1543 + "image_grid_thw": [],

1544 + }

1545 + else:

1546 + return {

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

1548 + "video_grid_thw": [],

1549 + }

1550 +

1551 + # Compute cumulative patch offsets for slicing pixel_values.

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

1553 + cum_patches = [0]

1554 + for p in patches_per_item:

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

1556 +

1557 + selected_pv = torch.cat(

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

1559 + )

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

1561 +

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

1563 + return {

1564 + "pixel_values": selected_pv,

1565 + "image_grid_thw": selected_grid,

1566 + }

1567 + else:

1568 + return {

1569 + "pixel_values_videos": selected_pv,

1570 + "video_grid_thw": selected_grid,

1571 + }

1572 +

1573 + def prepare_encoder_cudagraph_capture_inputs(

1574 + self,

1575 + token_budget: int,

1576 + max_batch_size: int,

1577 + max_frames_per_batch: int,

1578 + device: torch.device,

1579 + dtype: torch.dtype,

1580 + ):

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

1582 + EncoderCudaGraphCaptureInputs,

1583 + )

1584 +

1585 + spatial_merge_size = self.visual.spatial_merge_size

1586 + # Use ceil so captured capacity is never smaller than token_budget.

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

1588 +

1589 + frames_per_item = max_frames_per_batch // max_batch_size

1590 + if frames_per_item > 1:

1591 + tokens_per_frame = (

1592 + per_mm_item_output + frames_per_item - 1

1593 + ) // frames_per_item

1594 + grid_config = [

1595 + [

1596 + frames_per_item,

1597 + spatial_merge_size,

1598 + tokens_per_frame * spatial_merge_size,

1599 + ]

1600 + for _ in range(max_batch_size)

1601 + ]

1602 + else:

1603 + grid_config = [

1604 + [1, spatial_merge_size, per_mm_item_output * spatial_merge_size]

1605 + for _ in range(max_batch_size)

1606 + ]

1607 +

1608 + # Create dummy pixel_values.

1609 + patch_embed = self.visual.patch_embed

1610 + in_channels = patch_embed.proj.in_channels

1611 + patch_size = patch_embed.patch_size

1612 + temporal_patch_size = patch_embed.temporal_patch_size

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

1614 + flattened_patch_size = (

1615 + in_channels * temporal_patch_size * patch_size * patch_size

1616 + )

1617 + dummy_pixel_values = torch.randn(

1618 + total_patches, flattened_patch_size, device=device, dtype=dtype

1619 + )

1620 +

1621 + # max_seqlen.item() gets baked into the CUDA graph at capture time.

1622 + buffers = self.visual.prepare_encoder_metadata(

1623 + grid_config,

1624 + max_batch_size=max_batch_size,

1625 + max_frames_per_batch=max_frames_per_batch,

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

1627 + device=device,

1628 + )

1629 +

1630 + # Capture with image-format kwargs; pixel_values shape is compatible with

1631 + # both image and video replay paths.

1632 + mm_kwargs = {

1633 + "pixel_values": dummy_pixel_values,

1634 + "image_grid_thw": grid_config,

1635 + }

1636 +

1637 + return EncoderCudaGraphCaptureInputs(

1638 + mm_kwargs=mm_kwargs,

1639 + buffers=buffers,

1640 + )

1641 +

1642 + def prepare_encoder_cudagraph_replay_buffers(

1643 + self,

1644 + mm_kwargs: dict[str, Any],

1645 + max_batch_size: int,

1646 + max_frames_per_batch: int,

1647 + ) -> EncoderCudaGraphReplayBuffers:

1648 + modality = self.get_input_modality(mm_kwargs)

1649 + grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)

1650 +

1651 + if modality == "image":

1652 + buffers = self.visual.prepare_encoder_metadata(

1653 + grid_thw_list,

1654 + max_batch_size=max_batch_size,

1655 + )

1656 + else:

1657 + buffers = self.visual.prepare_encoder_metadata(

1658 + grid_thw_list,

1659 + max_frames_per_batch=max_frames_per_batch,

1660 + )

1661 +

1662 + return EncoderCudaGraphReplayBuffers(buffers=buffers)

1663 +

1664 + def encoder_cudagraph_forward(

1665 + self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]

1666 + ) -> torch.Tensor:

1667 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1668 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1669 + return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)

1670 +

1671 + def encoder_eager_forward(

1672 + self,

1673 + mm_kwargs: dict[str, Any],

1674 + ) -> torch.Tensor:

1675 + pixel_values = self._get_pixel_values_by_modality(mm_kwargs)

1676 + grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1677 + return self.visual(pixel_values, grid_thw)

1678 +

1679 def forward(

1680 self,

1681 input_ids: torch.Tensor | None,