[Bugfix][Spec Decode] Profile adaptive verification tail on a schedulable batch shape

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

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]

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)

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(