[Multimodal] Optimize encoder cudagraph padding logic (#58925)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Shanshan Shen <467638484@qq.com>

tests/v1/cudagraph/test_encoder_cudagraph.py

@@ -245,6 +245,138 @@ class TestFindBudgetGraph:
         assert mgr.get_num_graphs_to_capture() == 3
 
 
+# ---------------------------------------------------------------------------
+# _execute_local greedy packing with budget-cliff deferral (no GPU required)
+# ---------------------------------------------------------------------------
+
+
+def _run_packing(
+    mgr: EncoderCudaGraphManager, specs: list[EncoderItemSpec]
+) -> list[tuple[list[int], str, int]]:
+    """Run _execute_local with graph replay and postprocess mocked out.
+
+    Returns:
+        (batch_indices, path, token_budget) for every graph replay.
+
+    """
+    mgr.model = MagicMock()
+
+    runs: list[tuple[list[int], str, int]] = []
+
+    def fake_run(mm_kwargs, token_budget, path="default", axis_keys=()):
+        runs.append((list(mm_kwargs["indices"]), path, token_budget))
+        n_tokens = sum(
+            specs[i].get_path_output_tokens(path) for i in mm_kwargs["indices"]
+        )
+        return torch.zeros(n_tokens, 32)
+
+    def fake_postprocess(
+        graph_outputs,
+        batch_indices,
+        per_item_out_tokens,
+        outputs_by_orig_idx,
+        clone,
+        batch_mm_kwargs,
+    ):
+        for i in batch_indices:
+            outputs_by_orig_idx[i] = torch.zeros(per_item_out_tokens[i], 32)
+
+    mgr.model.postprocess_encoder_output = fake_postprocess
+    with (
+        patch.object(mgr, "_get_item_specs", lambda mm_kwargs: specs),
+        patch.object(
+            mgr,
+            "_select_items",
+            lambda mm_kwargs, indices: ({"indices": indices}, ()),
+        ),
+        patch.object(mgr, "_run_budget_graph", fake_run),
+    ):
+        result = mgr._execute_local({})
+    assert len(result) == len(specs)
+    return runs
+
+
+class TestExecuteLocalPacking:
+    """Greedy packing defers items that cross a budget cliff."""
+
+    @staticmethod
+    def _spec(tokens: int) -> EncoderItemSpec:
+        return EncoderItemSpec(input_size=tokens, output_tokens=tokens)
+
+    def test_cliff_deferral_splits_batch(self):
+        # Merging would force 4200 tokens into the 8192 budget; deferring
+        # the 4000-token item costs 256 + 4096 instead.
+        mgr = _make_manager_with_budgets([256, 4096, 8192])
+        runs = _run_packing(mgr, [self._spec(4000), self._spec(200)])
+        assert runs == [([1], "default", 256), ([0], "default", 4096)]
+
+    def test_tie_does_not_defer(self):
+        # 4096 + 4096 ties the merged 8192 budget; merging wins because it
+        # saves one replay.
+        mgr = _make_manager_with_budgets([256, 4096, 8192])
+        runs = _run_packing(mgr, [self._spec(4000), self._spec(4000)])
+        assert runs == [([0, 1], "default", 8192)]
+
+    def test_deferral_requires_item_outweighing_batch(self):
+        # The third 1366-token item would save budget if deferred (4096+2048
+        # < 8192) but is smaller than the current batch (4096), so it merges:
+        # deferring mid-size items fragments flood packing into extra replays.
+        mgr = _make_manager_with_budgets([256, 2048, 4096, 8192])
+        runs = _run_packing(mgr, [self._spec(1366)] * 3)
+        assert runs == [([0, 1, 2], "default", 8192)]
+
+    def test_packing_without_cliff_unchanged(self):
+        mgr = _make_manager_with_budgets([256, 4096])
+        runs = _run_packing(mgr, [self._spec(100), self._spec(100)])
+        assert runs == [([0, 1], "default", 256)]
+
+    def test_small_items_split_to_avoid_cliff(self):
+        # Merging 3+ items of 100 tokens crosses 256 -> 4096; splitting
+        # into pairs costs 256 + 256 < 4096.
+        mgr = _make_manager_with_budgets([256, 4096])
+        runs = _run_packing(mgr, [self._spec(100)] * 4)
+        assert runs == [([0, 1], "default", 256), ([2, 3], "default", 256)]
+
+    def test_max_batch_size_split_unchanged(self):
+        mgr = _make_manager_with_budgets([1024, 4096])
+        mgr.max_batch_size = 3
+        runs = _run_packing(mgr, [self._spec(100)] * 4)
+        assert runs == [([0, 1, 2], "default", 1024), ([3], "default", 1024)]
+
+    def test_oversized_item_still_falls_back_to_eager(self):
+        mgr = _make_manager_with_budgets([256, 4096])
+        runs = _run_packing(mgr, [self._spec(100), self._spec(9000)])
+        assert runs == [([0], "default", 256)]
+        assert mgr.graph_misses == 1
+
+    def test_multi_path_deferral_skips_zero_token_paths(self):
+        mgr = _make_manager_with_budgets([256, 4096, 8192])
+        mgr.path_token_budgets = {
+            "a": [0, 256, 4096, 8192],
+            "b": [256, 4096, 8192],
+        }
+        specs = [
+            EncoderItemSpec(
+                input_size=300,
+                output_tokens=300,
+                path_output_tokens={"a": 200, "b": 100},
+            ),
+            EncoderItemSpec(
+                input_size=4000,
+                output_tokens=4000,
+                path_output_tokens={"a": 0, "b": 4000},
+            ),
+        ]
+        runs = _run_packing(mgr, specs)
+        # Path b alone triggers deferral (256 + 4096 < 8192); the deferred
+        # item has zero tokens on path a, so that path is skipped.
+        assert runs == [
+            ([0], "a", 256),
+            ([0], "b", 256),
+            ([1], "b", 4096),
+        ]
+
+
 # ---------------------------------------------------------------------------
 # get_cumulative_stats
 # ---------------------------------------------------------------------------

vllm/model_executor/models/kimi_k25.py

@@ -479,30 +479,16 @@ class KimiK25ForConditionalGeneration(
     # Image-only (t == 1). Video chunks (t > 1) fall back to eager.      #
     # ------------------------------------------------------------------ #
 
-    @property
-    def _encoder_cudagraph_pad_totals(self) -> dict[int, int]:
-        """Row count of each captured buffer set, keyed by cu_seqlens ptr."""
-        totals = self.__dict__.get("_encoder_cg_pad_totals")
-        if totals is None:
-            totals = {}
-            self.__dict__["_encoder_cg_pad_totals"] = totals
-        return totals
-
     def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig":
         from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
 
-        pad_totals = self._encoder_cudagraph_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())
+            # Zero-length padding (same as Qwen): tail slots repeat the last
+            # real offset, so FlashAttn schedules no work for padding rows.
             n = min(src.shape[0], dst.shape[0])
             dst[:n].copy_(src[:n])
-            dst[n:] = total if total is not None else src[-1]
+            if n < dst.shape[0]:
+                dst[n:] = src[-1] if n else 0
 
         return EncoderCudaGraphConfig(
             modalities=["vision_chunk"],
@@ -633,8 +619,8 @@ class KimiK25ForConditionalGeneration(
 
         # max_seqlen must cover the worst case: one item consuming the full
         # budget, i.e. token_budget * kh * kw patches.
-        # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can append
-        # a padding sequence covering rows the real batch does not fill.
+        # max_batch_size + 1 leaves a spare cu_seqlens slot so replay can
+        # pad the tail with zero-length sequences.
         metadata = self.vision_tower.prepare_encoder_cudagraph_metadata(
             grid_thw_list,
             max_batch_size=max_batch_size + 1,
@@ -645,10 +631,6 @@ class KimiK25ForConditionalGeneration(
         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})
 
-        cu_seqlens = values.get("cu_seqlens")
-        if cu_seqlens is not None:
-            self._encoder_cudagraph_pad_totals[cu_seqlens.data_ptr()] = total_patches
-
         return EncoderCudaGraphCaptureInputs(values=values)
 
     def prepare_encoder_cudagraph_replay_buffers(
@@ -663,8 +645,8 @@ class KimiK25ForConditionalGeneration(
         grid_thw_list = self._get_grid_thw_list(mm_kwargs)
         pixel_values = mm_kwargs["pixel_values"]
 
-        # Unpadded: pad_cu_seqlens completes the tail with the padding sequence
-        # so cu_seqlens[-1] matches the captured buffer's row count.
+        # Unpadded: pad_cu_seqlens pads the tail with zero-length sequences
+        # whose entries repeat the real total.
         metadata = self.vision_tower.prepare_encoder_cudagraph_metadata(
             grid_thw_list,
             max_batch_size=None,

vllm/models/deepseek_v41/common/vl_cudagraph.py

@@ -94,8 +94,6 @@ class DeepseekV4VLEncoderCudaGraphMixin:
             grid = grid.tolist()
         return [[int(v) for v in row] for row in grid]
 
-    _encoder_cg_pad_totals: dict[int, int] = {}
-
     # -- SupportsEncoderCudaGraph protocol --
 
     def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
@@ -129,24 +127,13 @@ class DeepseekV4VLEncoderCudaGraphMixin:
                 "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 is 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 at capture "
-                    "time; the manager must replay into the capture-time "
-                    "buffers."
-                )
+            # Zero-length padding (same as Qwen): tail slots repeat the last
+            # real offset, so FlashAttn schedules no work for padding rows.
             n = min(src.shape[0], dst.shape[0])
             dst[:n].copy_(src[:n])
-            dst[n:] = total
+            if n < dst.shape[0]:
+                dst[n:] = src[-1] if n else 0
 
         return EncoderCudaGraphConfig(
             modalities=["image"],
@@ -268,8 +255,8 @@ class DeepseekV4VLEncoderCudaGraphMixin:
             max_seqlen_override=total_patches,
             cached=False,
         )
-        # Spare cu_seqlens slots let replay append a padding sequence that
-        # covers the rows a smaller real batch does not fill.
+        # Spare cu_seqlens slots: replay pads smaller real batches with
+        # zero-length sequences (tail slots repeat the last real offset).
         real_cu = metadata.pop("cu_seqlens")
         cu_seqlens = torch.full(
             (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device
@@ -278,8 +265,6 @@ class DeepseekV4VLEncoderCudaGraphMixin:
 
         merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype)
 
-        self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches
-
         return EncoderCudaGraphCaptureInputs(
             values={
                 "patches": patches,
@@ -303,7 +288,7 @@ class DeepseekV4VLEncoderCudaGraphMixin:
             patches = patches.to(dtype)
 
         # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and
-        # pad_cu_seqlens appends the padding sequence covering the tail rows.
+        # pad_cu_seqlens pads the tail with zero-length sequences.
         metadata = build_packed_vit_metadata(
             vit_grid,
             rope_dim=self.vision.rope_dim,

vllm/utils/math_utils.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Math utility functions for vLLM."""
 
+from bisect import bisect_left
+
 # Approximate value of 1/ln(2), used for log/exp base conversion
 # Best FP32 approximation: 1.4426950216 (hex 0x3FB8AA3B)
 RCP_LN2 = 1.4426950216
@@ -30,3 +32,9 @@ def round_down(x: int, y: int) -> int:
 def largest_power_of_2_divisor(n: int) -> int:
     """Return the largest power-of-2 that divides *n* (isolate lowest set bit)."""
     return n & (-n)
+
+
+def round_up_to(values: list[int], n: int) -> int | None:
+    """Round `n` up to the smallest member of sorted `values` >= n, or None."""
+    i = bisect_left(values, n)
+    return values[i] if i < len(values) else None

vllm/v1/worker/encoder_cudagraph.py

@@ -23,6 +23,7 @@ from vllm.model_executor.models.interfaces import (
 from vllm.model_executor.models.utils import scatter_output_slices
 from vllm.model_executor.models.vision import get_load_balance_assignment
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
+from vllm.utils.math_utils import round_up_to
 from vllm.utils.torch_utils import current_stream
 from vllm.v1.worker.encoder_cudagraph_defs import (
     ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG,
@@ -347,17 +348,17 @@ class EncoderCudaGraphManager:
     def _find_smallest_fitting_budget_given_tokens(
         self, total_tokens: int, budgets: list[int] | None = None
     ) -> int | None:
-        """Find smallest budget >= total_tokens.
-
-        Returns:
-            Token budget if found, None if no fitting budget.
-
-        """
+        """Find smallest budget >= total_tokens, or None if none fits."""
         budgets = budgets if budgets is not None else self.token_budgets
-        for budget in budgets:
-            if budget >= total_tokens:
-                return budget
-        return None
+        return round_up_to(budgets, total_tokens)
+
+    def _path_fitting_budget(self, path: str, tokens: int) -> int | None:
+        """Smallest captured budget on a path; 0 when idle, None if oversized."""
+        if tokens == 0:
+            return 0
+        return self._find_smallest_fitting_budget_given_tokens(
+            tokens, self.path_token_budgets[path]
+        )
 
     def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:
         """Get item specs from the model."""
@@ -479,24 +480,38 @@ class EncoderCudaGraphManager:
             if not current_batch:
                 return
             path_budgets = {
-                path: (
-                    0
-                    if current_tokens[path] == 0
-                    else self._find_smallest_fitting_budget_given_tokens(
-                        current_tokens[path], self.path_token_budgets[path]
-                    )
-                )
+                path: self._path_fitting_budget(path, current_tokens[path])
                 for path in paths
             }
             batches.append((list(current_batch), path_budgets))
 
+        def should_defer_item(item_tokens: dict[str, int]) -> bool:
+            """Whether closing the batch strictly reduces summed budgets."""
+            for path in paths:
+                item = item_tokens[path]
+                if item == 0:
+                    continue
+                cur = current_tokens[path]
+                b_cur = self._path_fitting_budget(path, cur)
+                b_alone = self._path_fitting_budget(path, item)
+                b_merged = self._path_fitting_budget(path, cur + item)
+                # Only called when the item fits; lookups land on a budget.
+                assert (
+                    b_cur is not None and b_alone is not None and b_merged is not None
+                )
+                # Only defer items that alone outweigh the whole current
+                # batch; splitting smaller ones fragments flood packing.
+                if b_cur + b_alone < b_merged and b_alone >= b_cur:
+                    return True
+            return False
+
         for orig_idx in sorted_indices:
             item_tokens = {path: per_item_path_tokens[path][orig_idx] for path in paths}
             fits = len(current_batch) < self.max_batch_size and all(
                 current_tokens[path] + item_tokens[path] <= max_path_budgets[path]
                 for path in paths
             )
-            if current_batch and not fits:
+            if current_batch and (not fits or should_defer_item(item_tokens)):
                 append_current_batch()
                 current_batch = []
                 current_tokens = dict.fromkeys(paths, 0)