[MM][CG] Avoid over-padding Qwen2.5-VL encoder cudagraph window metadata

vllm/model_executor/models/qwen2_5_vl.py

@@ -796,6 +796,38 @@ def invert_permutation(perm: torch.Tensor) -> torch.Tensor:
         inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)
         return inv
 
+    def get_encoder_cudagraph_max_window_seqs(
+        self,
+        token_budget: int,
+        max_batch_size: int,
+        max_frames_per_batch: int,
+    ) -> int:
+        # token_budget is an upper bound on the total number of merged vision
+        # tokens replayed by this encoder CUDA graph. cu_window_seqlens, however,
+        # is sized by the number of window-attention sequences (non-empty local
+        # windows), not by the number of tokens. Using max_num_batched_tokens as
+        # this sequence count can over-pad cu_window_seqlens and make FlashAttention
+        # launch thousands of empty CTAs during replay.
+        vit_merger_window_size = (
+            self.window_size // self.spatial_merge_size // self.patch_size
+        )
+        max_sequence_units = max(max_batch_size, max_frames_per_batch)
+
+        # Each local window covers vit_merger_window_size tokens along one merged
+        # spatial axis. The largest number of non-empty windows for a fixed token
+        # budget comes from a thin strip that advances along only one axis, so
+        # ceil(token_budget / window_side) is a safe geometry-driven bound. Multiple
+        # images or video frames can fragment that strip at item/frame boundaries,
+        # so add max_sequence_units to cover one extra partial window per sequence.
+        max_strip_windows = (
+            token_budget + vit_merger_window_size - 1
+        ) // vit_merger_window_size
+
+        # A non-empty window must contain at least one merged vision token, so the
+        # number of window sequences can never exceed token_budget. This final
+        # clamp keeps the bound tight for tiny budgets while remaining safe.
+        return min(token_budget, max_sequence_units + max_strip_windows)
+
     def prepare_encoder_metadata(
         self,
         grid_thw: list[list[int]],
@@ -1791,9 +1823,8 @@ def prepare_encoder_cudagraph_capture_inputs(
         )
 
         spatial_merge_size = self.visual.spatial_merge_size
-        max_window_seqs_per_batch = min(
-            self.vllm_config.scheduler_config.max_num_batched_tokens,
-            self.model_config.max_model_len,
+        max_window_seqs_per_batch = self.visual.get_encoder_cudagraph_max_window_seqs(
+            token_budget, max_batch_size, max_frames_per_batch
         )
         # Use ceil here (not floor) so total captured capacity is never smaller
         # than token_budget when token_budget is not divisible by max_batch_size
@@ -1888,23 +1919,21 @@ def prepare_encoder_cudagraph_replay_buffers(
         modality = self.get_input_modality(mm_kwargs)
         grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
 
+        # Keep replay metadata sized to the actual batch. The captured buffers
+        # may be larger, but EncoderCudaGraphManager fills the remaining
+        # cu*_seqlens entries with the last cumulative offset to represent empty
+        # sequences. Padding cu_window_seqlens here would require a static upper
+        # bound and can over-pad window attention into many empty FlashAttention
+        # CTAs.
         if modality == "image":
             buffers = self.visual.prepare_encoder_metadata(
                 grid_thw_list,
                 max_batch_size=max_batch_size,
-                max_window_seqs_per_batch=min(
-                    self.vllm_config.scheduler_config.max_num_batched_tokens,
-                    self.model_config.max_model_len,
-                ),
             )
         elif modality == "video":
             buffers = self.visual.prepare_encoder_metadata(
                 grid_thw_list,
                 max_frames_per_batch=max_frames_per_batch,
-                max_window_seqs_per_batch=min(
-                    self.vllm_config.scheduler_config.max_num_batched_tokens,
-                    self.model_config.max_model_len,
-                ),
             )
         else:
             raise AssertionError("This line should be unreachable.")

vllm/v1/worker/encoder_cudagraph.py

@@ -45,7 +45,12 @@ class BudgetGraphMetadata:
     # The input tensor updated before replay (e.g. pixel_values)
     input_buffer: torch.Tensor
     # Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata).
-    # Before replay the manager zeros then slice-copies new data into these.
+    # Before replay the manager updates these in-place. Most buffers are zeroed
+    # before slice-copying the actual values. cu*_seqlens buffers are special:
+    # their padded tail repeats the final cumulative offset so the extra entries
+    # represent empty sequences rather than invalid zero-length ranges. FlashInfer
+    # stores two cumulative-offset sections back-to-back; each section must be
+    # padded independently.
     metadata_buffers: dict[str, torch.Tensor]
     # Output written by graph, read after replay
     output_buffer: torch.Tensor
@@ -259,6 +264,40 @@ def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:
         """Get per-item output token counts as plain ints."""
         return [spec.output_tokens for spec in self._get_item_specs(mm_kwargs)]
 
+    @staticmethod
+    def _is_cu_seqlens_buffer(key: str) -> bool:
+        return key.startswith("cu_") and "seqlens" in key
+
+    @staticmethod
+    def _is_flashinfer_cu_seqlens(
+        src: torch.Tensor,
+        dst: torch.Tensor,
+    ) -> bool:
+        if src.ndim != 1 or dst.ndim != 1:
+            return False
+        return src.shape[0] >= 4 and src.shape[0] % 2 == 0 and dst.shape[0] % 2 == 0
+
+    @staticmethod
+    def _copy_padded_flashinfer_cu_seqlens(
+        dst: torch.Tensor,
+        src: torch.Tensor,
+    ) -> None:
+        src_mid = src.shape[0] // 2
+        dst_mid = dst.shape[0] // 2
+        assert src_mid <= dst_mid, (
+            f"FlashInfer cu_seqlens replay buffer is larger than capture buffer: "
+            f"src_section={src_mid}, dst_section={dst_mid}"
+        )
+
+        dst.zero_()
+        dst[:src_mid].copy_(src[:src_mid])
+        if src_mid < dst_mid:
+            dst[src_mid:dst_mid] = src[src_mid - 1]
+
+        dst[dst_mid : dst_mid + src_mid].copy_(src[src_mid:])
+        if dst_mid + src_mid < dst.shape[0]:
+            dst[dst_mid + src_mid :] = src[-1]
+
     def _run_budget_graph(
         self,
         mm_kwargs: dict[str, Any],
@@ -282,6 +321,10 @@ def _run_budget_graph(
             return None
 
         graph_meta = self.budget_graphs[token_budget]
+        has_flashinfer_metadata = any(
+            replay_buffers.get(key) is not None
+            for key in ("sequence_lengths_full", "sequence_lengths_window")
+        )
 
         # Copy the input tensor. Buffers are sized for the full budget;
         # actual inputs may be smaller. Zero then slice-copy so padded
@@ -303,8 +346,21 @@ def _run_budget_graph(
                 buf.copy_(src)
             else:
                 n = src.shape[0]
-                buf.zero_()
-                buf[:n].copy_(src)
+                is_cu_seqlens = self._is_cu_seqlens_buffer(key)
+                is_flashinfer_cu_seqlens = (
+                    has_flashinfer_metadata
+                    and is_cu_seqlens
+                    and self._is_flashinfer_cu_seqlens(src, buf)
+                )
+                if is_flashinfer_cu_seqlens:
+                    self._copy_padded_flashinfer_cu_seqlens(buf, src)
+                else:
+                    buf.zero_()
+                    buf[:n].copy_(src)
+                if n < buf.shape[0] and is_cu_seqlens and not is_flashinfer_cu_seqlens:
+                    # Cumulative sequence-length buffers encode empty padded
+                    # sequences by repeating the final offset, not by zeros.
+                    buf[n:] = src[-1]
 
         graph_meta.graph.replay()