diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 84ba8da55ce9..ff6c21241c2e 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -306,3 +306,80 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits(): assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype # The prefill keeps its scheduled tokens; only drafts are dropped. assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32)) + + +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 c05bb1a431a6..89110b8628a8 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -131,7 +131,14 @@ def make_dummy( input_buffers: InputBuffers, max_query_len: int | None = None, is_padding: bool = True, + 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 @@ -141,21 +148,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 2d30fbb612cf..408a602224b1 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -833,6 +833,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, valid_dummy_state_slots: bool = False, @@ -873,13 +874,28 @@ def _dummy_run( f"({self.max_num_tokens})" ) - # Distribute the remainder evenly so no dummy request exceeds - # ceil(num_tokens / num_reqs) <= max_model_len tokens. + if num_scheduled_tokens_per_req is not None and ( + sum(num_scheduled_tokens_per_req) == num_tokens + and len(num_scheduled_tokens_per_req) <= num_reqs + ): + # 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). The clamps + # above can lower num_tokens or num_reqs for the draft budget; a + # pinned split that no longer fits them is stale, so it gives way to + # the even split for the shape that will actually run. + assert not uniform_decode + num_tokens_per_request = list(num_scheduled_tokens_per_req) + num_reqs = len(num_tokens_per_request) + 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) - ] + 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 = { @@ -1294,7 +1310,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, decode_graph_eligible=True ) @@ -1848,11 +1865,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, # Profiling and warmup must route the dummy tokens to experts # so MoE memory is measured and MoE kernels are exercised. is_padding=not is_profile, diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 5f91af59e592..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 @@ -14,6 +14,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.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.utils import CpuGpuBuffer @@ -128,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, @@ -185,7 +198,41 @@ 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[ProfileBatch]: """Dummy-run kwargs whose step timings seed the cost tables. Run these inside StepTimingCollector.collect(), then hand the block's @@ -203,11 +250,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: ProfileBatch = { + "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 batch.copy() def set_initial_cost_curves(self, samples: list[StepTimingSample]) -> None: def median_curve(