From 88890e8fd08d0bc66bacffa26d40d6db5d2d1fb0 Mon Sep 17 00:00:00 2001 From: kakiuwang-ui Date: Thu, 27 Aug 2026 23:16:55 +0800 Subject: [PATCH 1/2] [Bugfix][Spec Decode] Profile adaptive verification tail on a schedulable batch shape Adaptive verification seeds its cost tables from dummy runs at every captured size plus a doubling tail up to max_num_batched_tokens. The tail sizes were handed to _dummy_run as a bare token count, which splits them evenly over min(num_tokens, max_num_seqs) requests. With max_num_seqs=128 and num_speculative_tokens=7 the 16384-token tail was therefore timed as 128 requests of 128 queries each, while a fully occupied decode step is 128 * (1 + 7) = 1024 tokens. Every tail sample above 1024 measured a shape the scheduler cannot produce, and those samples are what the verify curve is fitted and extrapolated from. Shape the tail as the mixed step it has to be: the surplus over the decode ceiling can only be prefill, so it is cut into chunks sized by VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN, the freed request slots stay decode requests of exactly 1 + num_speculative_tokens queries, and the split is passed down to InputBatch.make_dummy. Captured sizes keep the even split, since those runs must match the descriptor their graph was captured with, and make_dummy only honours an explicit split when cudagraph dispatch did not pad or reshape the batch. Signed-off-by: kakiuwang-ui --- .../spec_decode/test_adaptive_verification.py | 77 +++++++++++++++++++ vllm/v1/worker/gpu/input_batch.py | 40 +++++++--- vllm/v1/worker/gpu/model_runner.py | 37 +++++++-- .../gpu/spec_decode/adaptive_verification.py | 55 +++++++++++-- 4 files changed, 185 insertions(+), 24 deletions(-) diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index d92f78a71535..4dd3105a814c 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -275,3 +275,80 @@ def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): # (request, position) keys, so the kernel can mask rows the compacted # device layout no longer has room for. assert mapping == [0, 1, 2, 3, 4, 5, 6] + + +def make_profiling_manager( + max_num_reqs: int = 128, + max_num_batched_tokens: int = 16384, + num_speculative_steps: int = 7, +) -> AdaptiveVerificationManager: + manager = AdaptiveVerificationManager.__new__(AdaptiveVerificationManager) + manager.req_states = SimpleNamespace( + max_num_batched_tokens=max_num_batched_tokens, max_num_reqs=max_num_reqs + ) + manager.num_speculative_steps = num_speculative_steps + manager.num_bonus_tokens = 1 + return manager + + +def test_profile_split_left_alone_while_the_step_still_fits_in_decode(): + # 128 requests x (1 + 7) queries is the largest decode step, so at or below + # it the caller's even split is already the real shape. + manager = make_profiling_manager() + assert manager.profile_batch_split(1024) is None + assert manager.profile_batch_split(512) is None + + +def test_profile_split_shapes_the_surplus_above_decode_as_prefill(monkeypatch): + # The case from the report: max_num_seqs=128, k=7, ctx=8192, 16384 tokens. + # An even split would time 128 requests of 128 queries each, which no + # decode step can be; the surplus over 128 * 8 has to be prefill. + monkeypatch.setattr( + adaptive_module.envs, "VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN", 8192 + ) + manager = make_profiling_manager() + + split = manager.profile_batch_split(16384) + + assert split is not None + assert sum(split) == 16384 + assert len(split) == 128 + # 126 decode requests keep the real decode query length; the two prefill + # slots absorb everything else. + assert split[:126] == [8] * 126 + assert split[126:] == [7688, 7688] + + +def test_profiled_tail_batches_are_shapes_the_scheduler_can_produce(monkeypatch): + monkeypatch.setattr( + adaptive_module.envs, "VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN", 8192 + ) + max_num_reqs = 128 + decode_query_len = 8 + manager = make_profiling_manager(max_num_reqs=max_num_reqs) + capture_sizes = [8, 1024] + + for batch in manager.batches_to_profile(capture_sizes): + num_tokens = batch["num_tokens"] + split = batch.get("num_scheduled_tokens_per_req") + if num_tokens in capture_sizes or num_tokens <= max_num_reqs * decode_query_len: + # Captured shapes must keep the even split the graph was captured + # with, and decode-sized steps are already realistic. + assert split is None + continue + assert split is not None + assert sum(split) == num_tokens + assert len(split) <= max_num_reqs + assert min(split) > 0 + # Every request is either a real decode query or a prefill chunk. The + # even split this replaces gave all 128 requests num_tokens / 128 + # queries, so there would be no decode requests left at all. + decode_lens = [n for n in split if n <= decode_query_len] + prefill_lens = [n for n in split if n > decode_query_len] + assert decode_lens == [decode_query_len] * len(decode_lens) + assert len(decode_lens) >= 1 + assert len(prefill_lens) >= 1 + assert len(decode_lens) + len(prefill_lens) == max_num_reqs + # Prefill chunks stay near-equal and within one scheduler step. + assert max(prefill_lens) - min(prefill_lens) <= 1 + assert max(prefill_lens) <= manager.req_states.max_num_batched_tokens diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 9a574a4faac5..c25bafd58427 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -119,7 +119,14 @@ def make_dummy( num_tokens: int, input_buffers: InputBuffers, max_query_len: int | None = None, + num_scheduled_tokens: np.ndarray | None = None, ) -> "InputBatch": + """Build a dummy batch of *num_tokens* over *num_reqs* requests. + + The tokens are spread evenly unless *num_scheduled_tokens* pins an + explicit per-request split, which callers use when the uniform shape + would not be one the scheduler can produce. + """ assert 0 < num_reqs <= num_tokens device = input_buffers.device @@ -129,21 +136,30 @@ def make_dummy( expanded_idx_mapping = idx_mapping expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device) - # Distribute the remainder evenly so that no dummy request exceeds + # seq_len equals to query_len. Without an explicit split, distribute the + # remainder evenly so that no dummy request exceeds # ceil(num_tokens / num_reqs) <= max_model_len tokens. Varlen graphs - # accept any split with non-empty slots, so this shape works for them - # too; attention metadata is built from the promised max_query_len. - base_tokens = num_tokens // num_reqs - num_extra = num_tokens % num_reqs - assert max_query_len is None or base_tokens + (num_extra > 0) <= max_query_len - num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32) - if num_extra > 0: - num_scheduled_tokens[-num_extra:] += 1 + # accept any split with non-empty slots, so either shape works for them; + # attention metadata is built from the promised max_query_len. + if num_scheduled_tokens is None: + base_tokens = num_tokens // num_reqs + num_extra = num_tokens % num_reqs + max_split = base_tokens + (num_extra > 0) + assert max_query_len is None or max_split <= max_query_len + num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32) + if num_extra > 0: + num_scheduled_tokens[-num_extra:] += 1 + input_buffers.seq_lens[: num_reqs - num_extra] = base_tokens + input_buffers.seq_lens[num_reqs - num_extra : num_reqs] = base_tokens + 1 + else: + num_scheduled_tokens = num_scheduled_tokens.astype(np.int32, copy=True) + assert num_scheduled_tokens.shape == (num_reqs,) + assert num_scheduled_tokens.min() > 0 + max_split = int(num_scheduled_tokens.max()) + assert max_query_len is None or max_split <= max_query_len + input_buffers.seq_lens[:num_reqs] = torch.from_numpy(num_scheduled_tokens) assert int(num_scheduled_tokens.sum()) == num_tokens - # seq_len equals to query_len - input_buffers.seq_lens[: num_reqs - num_extra] = base_tokens - input_buffers.seq_lens[num_reqs - num_extra : num_reqs] = base_tokens + 1 # Pad for full CUDA graph mode. input_buffers.seq_lens[num_reqs:] = 0 seq_lens = input_buffers.seq_lens[:num_reqs] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f92de74ed9e4..d29f57eca576 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -680,6 +680,7 @@ def _dummy_run( skip_attn: bool = False, uniform_decode: bool = False, context_len: int = 0, + num_scheduled_tokens_per_req: list[int] | None = None, skip_eplb: bool = False, is_profile: bool = False, **kwargs, @@ -699,12 +700,22 @@ def _dummy_run( num_tokens = max(num_tokens, self.decode_query_len) num_reqs = num_tokens // self.decode_query_len assert num_tokens % self.decode_query_len == 0 - # Distribute the remainder evenly so no dummy request exceeds - # ceil(num_tokens / num_reqs) <= max_model_len tokens. - num_tokens_per_request = [ - num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs) - for i in range(num_reqs) - ] + if num_scheduled_tokens_per_req is not None: + # An explicit split, for callers that need the dummy batch to have a + # specific shape rather than a uniform one (e.g. adaptive + # verification profiling a mixed decode/prefill step). + assert not uniform_decode + num_tokens_per_request = list(num_scheduled_tokens_per_req) + num_reqs = len(num_tokens_per_request) + assert num_reqs <= self.max_num_reqs + assert all(n > 0 for n in num_tokens_per_request) + else: + # Distribute the remainder evenly so no dummy request exceeds + # ceil(num_tokens / num_reqs) <= max_model_len tokens. + num_tokens_per_request = [ + num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs) + for i in range(num_reqs) + ] assert sum(num_tokens_per_request) == num_tokens num_scheduled_tokens = { @@ -1099,7 +1110,8 @@ def gather_batch_req_state( max_query_len = max(scheduler_output.num_scheduled_tokens.values()) if dummy_run: - # Dummy batches are uniform by construction. + # Dummy batches are uniform unless the caller pinned a split, and + # get_uniform_decode_token_count returns None for those. return None, get_uniform_decode_token_count( num_reqs, num_toks, max_query_len, has_prefill=False ) @@ -1594,11 +1606,22 @@ def execute_model( else: # No actual tokens to run. A dummy run for DP or memory profiling. dummy_num_reqs = batch_desc.num_reqs or num_reqs + # Honour the caller's split when cudagraph dispatch did not pad or + # reshape the batch; otherwise the shape belongs to the descriptor + # and make_dummy falls back to its even split. + dummy_split = None + if dummy_num_reqs == num_reqs and batch_desc.num_tokens == num_toks: + dummy_split = np.fromiter( + scheduler_output.num_scheduled_tokens.values(), + dtype=np.int32, + count=num_reqs, + ) input_batch = InputBatch.make_dummy( dummy_num_reqs, batch_desc.num_tokens, self.input_buffers, max_query_len=batch_desc.max_query_len, + num_scheduled_tokens=dummy_split, ) if not skip_attn_for_dummy_run: block_tables, slot_mappings = self.prepare_dummy_attn(input_batch) diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index d3b3743d67fd..f8386c58622d 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -13,6 +13,7 @@ from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.math_utils import cdiv from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream @@ -170,7 +171,43 @@ def add_request(self, req_idx: int) -> None: self._pending_resets.append(req_idx) self._confidence_probs[req_idx].fill_(1.0) - def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int]]: + def profile_batch_split(self, num_tokens: int) -> list[int] | None: + """Per-request query lengths for a profiled step of *num_tokens*. + + A decode request contributes ``1 + num_speculative_steps`` queries, so a + fully occupied decode step is ``max_num_reqs * decode_query_len`` tokens + and nothing above that can be decode. Splitting such a step evenly over + ``max_num_reqs`` requests would time a batch the scheduler cannot + produce, so the surplus is shaped as the prefill it has to be. + + Returns ``None`` when the step still fits in decode, leaving the caller's + even split (which already gives every request at most + ``decode_query_len`` queries) alone. + """ + max_num_reqs = self.req_states.max_num_reqs + decode_query_len = self.num_speculative_steps + self.num_bonus_tokens + max_decode_tokens = max_num_reqs * decode_query_len + if num_tokens <= max_decode_tokens: + return None + + # ``context_len`` is the profiler's stand-in for a sequence length, so it + # also sizes the prefill chunks the surplus is cut into. + chunk_len = max(envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN, 1) + surplus = num_tokens - max_decode_tokens + num_prefill_reqs = min(max(cdiv(surplus, chunk_len), 1), max_num_reqs) + # Prefill slots come out of the request budget, freeing their decode + # tokens back to the chunks. + num_decode_reqs = max_num_reqs - num_prefill_reqs + prefill_tokens = num_tokens - num_decode_reqs * decode_query_len + + split = [decode_query_len] * num_decode_reqs + base, extra = divmod(prefill_tokens, num_prefill_reqs) + split += [base + (i < extra) for i in range(num_prefill_reqs)] + return split + + def batches_to_profile( + self, capture_sizes: list[int] + ) -> Iterator[dict[str, int | list[int]]]: """Dummy-run kwargs whose step timings seed the cost tables. Run these inside StepTimingCollector.collect(), then hand the block's @@ -188,11 +225,19 @@ def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int tail_sizes.add(size) tail_sizes -= set(capture_sizes) for num_tokens in capture_sizes + sorted(tail_sizes): + batch: dict[str, int | list[int]] = { + "num_tokens": num_tokens, + "context_len": envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN, + } + # Only the tail is reshaped: captured sizes have to keep the even + # split so the run matches the descriptor the graph was captured + # with. + if num_tokens in tail_sizes: + split = self.profile_batch_split(num_tokens) + if split is not None: + batch["num_scheduled_tokens_per_req"] = split for _ in range(_PROFILE_REPLAYS): - yield { - "num_tokens": num_tokens, - "context_len": envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN, - } + yield dict(batch) def set_initial_cost_curves(self, samples: list[StepTimingSample]) -> None: def median_curve( From 4e3df1ab758b9084cdaf845516f9a629ed3c5de2 Mon Sep 17 00:00:00 2001 From: kakiuwang-ui Date: Thu, 15 Oct 2026 01:49:00 +0800 Subject: [PATCH 2/2] Type the profiled dummy-run kwargs as a TypedDict mypy rejects the test's arithmetic on batch["num_tokens"] because the yielded mapping was typed dict[str, int | list[int]], which makes every value a union. A TypedDict gives each key its own type and keeps the optional split optional. Signed-off-by: kakiuwang-ui --- .../gpu/spec_decode/adaptive_verification.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 5d9cecbe93e3..e7948da902fd 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -4,7 +4,7 @@ from collections import defaultdict from collections.abc import Iterable, Iterator -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NotRequired, TypedDict import numpy as np import torch @@ -129,6 +129,18 @@ def build_table(limit: int, curve: list[tuple[int, float]]) -> np.ndarray: return draft_table, verify_table +class ProfileBatch(TypedDict): + """`_dummy_run` kwargs for one profiled step. + + `num_scheduled_tokens_per_req` is present only for a step whose even split + would not be a shape the scheduler can produce. + """ + + num_tokens: int + context_len: int + num_scheduled_tokens_per_req: NotRequired[list[int]] + + class AdaptiveVerificationManager: def __init__( self, @@ -220,9 +232,7 @@ def profile_batch_split(self, num_tokens: int) -> list[int] | None: split += [base + (i < extra) for i in range(num_prefill_reqs)] return split - def batches_to_profile( - self, capture_sizes: list[int] - ) -> Iterator[dict[str, int | list[int]]]: + def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[ProfileBatch]: """Dummy-run kwargs whose step timings seed the cost tables. Run these inside StepTimingCollector.collect(), then hand the block's @@ -240,7 +250,7 @@ def batches_to_profile( tail_sizes.add(size) tail_sizes -= set(capture_sizes) for num_tokens in capture_sizes + sorted(tail_sizes): - batch: dict[str, int | list[int]] = { + batch: ProfileBatch = { "num_tokens": num_tokens, "context_len": envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN, } @@ -252,7 +262,7 @@ def batches_to_profile( if split is not None: batch["num_scheduled_tokens_per_req"] = split for _ in range(_PROFILE_REPLAYS): - yield dict(batch) + yield batch.copy() def set_initial_cost_curves(self, samples: list[StepTimingSample]) -> None: def median_curve(