From 8352b2427f704652da1910f0d53f3fdcc2fa2b0d Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Thu, 8 Oct 2026 11:38:09 +0800 Subject: [PATCH] [Multimodal] Optimize encoder cudagraph padding logic (#58925) Signed-off-by: Isotr0py Co-authored-by: Shanshan Shen <467638484@qq.com> --- tests/v1/cudagraph/test_encoder_cudagraph.py | 132 ++++++++++++++++++ vllm/model_executor/models/kimi_k25.py | 34 ++--- .../deepseek_v41/common/vl_cudagraph.py | 29 +--- vllm/utils/math_utils.py | 8 ++ vllm/v1/worker/encoder_cudagraph.py | 51 ++++--- 5 files changed, 188 insertions(+), 66 deletions(-) diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 749f9aa20d..3489443878 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index d66bed3642..8cd354b903 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/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, diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py index bd8ed5a3fd..19baf339eb 100644 --- a/vllm/models/deepseek_v41/common/vl_cudagraph.py +++ b/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, diff --git a/vllm/utils/math_utils.py b/vllm/utils/math_utils.py index 8fb0009d60..9854249af2 100644 --- a/vllm/utils/math_utils.py +++ b/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 diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 3b587ad71f..51c832001c 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/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)