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 2, 13:30 UTC
[MM][CG] Avoid over-padding Qwen2.5-VL encoder cudagraph window metadata (#42796)
Signed-off-by: Hua Huang <huah@nvidia.com>
Hua Huang committed
May 28, 2026 at 11:22 UTC
9006204e90d3f086222a581cf96b2c517ba8202e
3 files changed
+108
-16
vllm/model_executor/models/qwen2_5_vl.py
+82
-11
@@ -122,6 +122,39 @@ from .vision import (
122
123
logger = init_logger(__name__)
124
125
+
126
+def _pad_cumulative_seqlens_buffer(
127
+ dst: torch.Tensor,
128
+ src: torch.Tensor,
129
+) -> None:
130
+ n = src.shape[0]
131
+ dst.zero_()
132
+ dst[:n].copy_(src)
133
+ if n < dst.shape[0]:
134
+ dst[n:] = src[-1]
135
+
136
+
137
+def _pad_flashinfer_cu_seqlens_buffer(
138
+ dst: torch.Tensor,
139
+ src: torch.Tensor,
140
+) -> None:
141
+ src_mid = src.shape[0] // 2
142
+ dst_mid = dst.shape[0] // 2
143
+ assert src_mid <= dst_mid, (
144
+ f"FlashInfer cu_seqlens replay buffer is larger than capture buffer: "
145
+ f"src_section={src_mid}, dst_section={dst_mid}"
146
+ )
147
+
148
+ dst.zero_()
149
+ dst[:src_mid].copy_(src[:src_mid])
150
+ if src_mid < dst_mid:
151
+ dst[src_mid:dst_mid] = src[src_mid - 1]
152
+
153
+ dst[dst_mid : dst_mid + src_mid].copy_(src[src_mid:])
154
+ if dst_mid + src_mid < dst.shape[0]:
155
+ dst[dst_mid + src_mid :] = src[-1]
156
+
157
+
158
# === Vision Inputs === #
159
160
@@ -796,6 +829,38 @@ class Qwen2_5_VisionTransformer(nn.Module):
829
inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)
830
return inv
831
832
+ def get_encoder_cudagraph_max_window_seqs(
833
+ self,
834
+ token_budget: int,
835
+ max_batch_size: int,
836
+ max_frames_per_batch: int,
837
+ ) -> int:
838
+ # token_budget is an upper bound on the total number of merged vision
839
+ # tokens replayed by this encoder CUDA graph. cu_window_seqlens, however,
840
+ # is sized by the number of window-attention sequences (non-empty local
841
+ # windows), not by the number of tokens. Using max_num_batched_tokens as
842
+ # this sequence count can over-pad cu_window_seqlens and make FlashAttention
843
+ # launch thousands of empty CTAs during replay.
844
+ vit_merger_window_size = (
845
+ self.window_size // self.spatial_merge_size // self.patch_size
846
+ )
847
+ max_sequence_units = max(max_batch_size, max_frames_per_batch)
848
+
849
+ # Each local window covers vit_merger_window_size tokens along one merged
850
+ # spatial axis. The largest number of non-empty windows for a fixed token
851
+ # budget comes from a thin strip that advances along only one axis, so
852
+ # ceil(token_budget / window_side) is a safe geometry-driven bound. Multiple
853
+ # images or video frames can fragment that strip at item/frame boundaries,
854
+ # so add max_sequence_units to cover one extra partial window per sequence.
855
+ max_strip_windows = (
856
+ token_budget + vit_merger_window_size - 1
857
+ ) // vit_merger_window_size
858
+
859
+ # A non-empty window must contain at least one merged vision token, so the
860
+ # number of window sequences can never exceed token_budget. This final
861
+ # clamp keeps the bound tight for tiny budgets while remaining safe.
862
+ return min(token_budget, max_sequence_units + max_strip_windows)
863
+
864
def prepare_encoder_metadata(
865
self,
866
grid_thw: list[list[int]],
@@ -1640,6 +1705,11 @@ class Qwen2_5_VLForConditionalGeneration(
1705
modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]
1706
1707
max_frames = self.get_max_frames_per_video() if "video" in modalities else 1
1708
+ cu_seqlens_padding = (
1709
+ _pad_flashinfer_cu_seqlens_buffer
1710
+ if self.visual.attn_backend == AttentionBackendEnum.FLASHINFER
1711
+ else _pad_cumulative_seqlens_buffer
1712
+ )
1713
return EncoderCudaGraphConfig(
1714
modalities=modalities,
1715
input_key_by_modality={
@@ -1658,6 +1728,10 @@ class Qwen2_5_VLForConditionalGeneration(
1728
"sequence_lengths_full",
1729
"sequence_lengths_window",
1730
],
1731
+ padding_logics={
1732
+ "cu_seqlens": cu_seqlens_padding,
1733
+ "cu_window_seqlens": cu_seqlens_padding,
1734
+ },
1735
out_hidden_size=self.visual.out_hidden_size,
1736
max_frames_per_video=max_frames,
1737
)
@@ -1791,9 +1865,8 @@ class Qwen2_5_VLForConditionalGeneration(
1865
)
1866
1867
spatial_merge_size = self.visual.spatial_merge_size
1794
- max_window_seqs_per_batch = min(
1795
- self.vllm_config.scheduler_config.max_num_batched_tokens,
1796
- self.model_config.max_model_len,
1868
+ max_window_seqs_per_batch = self.visual.get_encoder_cudagraph_max_window_seqs(
1869
+ token_budget, max_batch_size, max_frames_per_batch
1870
)
1871
# Use ceil here (not floor) so total captured capacity is never smaller
1872
# than token_budget when token_budget is not divisible by max_batch_size
@@ -1888,23 +1961,21 @@ class Qwen2_5_VLForConditionalGeneration(
1961
modality = self.get_input_modality(mm_kwargs)
1962
grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
1963
1964
+ # Keep replay metadata sized to the actual batch. The captured buffers
1965
+ # may be larger, but EncoderCudaGraphManager fills the remaining
1966
+ # cu*_seqlens entries with the last cumulative offset to represent empty
1967
+ # sequences. Padding cu_window_seqlens here would require a static upper
1968
+ # bound and can over-pad window attention into many empty FlashAttention
1969
+ # CTAs.
1970
if modality == "image":
1971
buffers = self.visual.prepare_encoder_metadata(
1972
grid_thw_list,
1973
max_batch_size=max_batch_size,
1895
- max_window_seqs_per_batch=min(
1896
- self.vllm_config.scheduler_config.max_num_batched_tokens,
1897
- self.model_config.max_model_len,
1898
- ),
1974
)
1975
elif modality == "video":
1976
buffers = self.visual.prepare_encoder_metadata(
1977
grid_thw_list,
1978
max_frames_per_batch=max_frames_per_batch,
1904
- max_window_seqs_per_batch=min(
1905
- self.vllm_config.scheduler_config.max_num_batched_tokens,
1906
- self.model_config.max_model_len,
1907
- ),
1979
)
1980
else:
1981
raise AssertionError("This line should be unreachable.")
vllm/v1/worker/encoder_cudagraph.py
+15
-4
@@ -45,7 +45,9 @@ class BudgetGraphMetadata:
45
# The input tensor updated before replay (e.g. pixel_values)
46
input_buffer: torch.Tensor
47
# Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata).
48
- # Before replay the manager zeros then slice-copies new data into these.
48
+ # Before replay the manager updates these in-place. By default buffers are
49
+ # zeroed before slice-copying the actual values; model-specific padding
50
+ # behavior is provided by EncoderCudaGraphConfig.padding_logics.
51
metadata_buffers: dict[str, torch.Tensor]
52
# Output written by graph, read after replay
53
output_buffer: torch.Tensor
@@ -259,6 +261,14 @@ class EncoderCudaGraphManager:
261
"""Get per-item output token counts as plain ints."""
262
return [spec.output_tokens for spec in self._get_item_specs(mm_kwargs)]
263
264
+ @staticmethod
265
+ def _copy_padded_buffer(
266
+ dst: torch.Tensor,
267
+ src: torch.Tensor,
268
+ ) -> None:
269
+ dst.zero_()
270
+ dst[: src.shape[0]].copy_(src)
271
+
272
def _run_budget_graph(
273
self,
274
mm_kwargs: dict[str, Any],
@@ -302,9 +312,10 @@ class EncoderCudaGraphManager:
312
if src.ndim == 0:
313
buf.copy_(src)
314
else:
305
- n = src.shape[0]
306
- buf.zero_()
307
- buf[:n].copy_(src)
315
+ padding_logic = self.config.padding_logics.get(
316
+ key, self._copy_padded_buffer
317
+ )
318
+ padding_logic(buf, src)
319
320
graph_meta.graph.replay()
321
vllm/v1/worker/encoder_cudagraph_defs.py
+11
-1
@@ -2,11 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Data transfer objects for encoder CUDA graph management."""
4
5
-from dataclasses import dataclass
5
+from collections.abc import Callable
6
+from dataclasses import dataclass, field
7
from typing import Any
8
9
import torch
10
11
+EncoderCudaGraphPaddingLogic = Callable[[torch.Tensor, torch.Tensor], None]
12
+
13
14
@dataclass
15
class EncoderItemSpec:
@@ -51,6 +54,13 @@ class EncoderCudaGraphConfig:
54
"""Output hidden dim of the vision encoder.
55
Used for DP gather buffer allocation."""
56
57
+ padding_logics: dict[str, EncoderCudaGraphPaddingLogic] = field(
58
+ default_factory=dict
59
+ )
60
+ """Optional per-buffer replay padding/copy logic.
61
+ If absent for a key, the manager zeros the capture buffer and slice-copies
62
+ the replay buffer into it."""
63
+
64
max_frames_per_video: int = 1
65
"""Maximum number of frames per video.
66
Only relevant when "video" is in ``modalities``.