[rebel-jinhwan]

Overview

The full-sequence admission gate (allocate_slots(..., full_sequence_must_fit=True), which the scheduler uses whenever scheduler_reserve_full_isl is on, the default) checks the prompt alone. With speculative decoding, the request's last prefill chunk also allocates the lookahead slots, and the decode step right after the prefill adds the sampled token and the draft tokens. This PR makes the gate cover that step.

Claims

• With speculative decoding, a request admitted under full-ISL reservation now has room through its first decode step: num_tokens + 1 + num_spec_tokens + num_lookahead_tokens, capped at max_model_len. Before, a prompt that fit only without those slots passed the gate, ran its prefill, and was then preempted on its last chunk (the lookahead) or on its first decode step (the sampled and draft tokens). FCFS preemption pops the newest running request, which is that request, so the prefill is thrown away; in the scheduler test below this happens three times in a row.
• No change without speculative decoding. Async KV loads keep their spec decode margin in reserved_blocks (see [Bugfix][Core] Reserve only the blocks a spec decode step adds when admitting async KV loads #63136): the scheduler passes neither the lookahead nor the decode step to the gate there.

Validation

• tests/v1/core/test_prefix_caching.py::test_full_sequence_admission_counts_the_first_decode_step: four usable 16-token blocks; prompt / decode-step tokens / lookahead of 64/0/0, 62/0/3, 57/4/3 and 58/4/3.
• tests/v1/core/test_scheduler.py::test_full_sequence_admission_covers_the_first_decode_step: real Scheduler, ngram k=3. A 30-token prompt arrives while two decodes hold all but two 16-token blocks; its prefill fits, its first decode step (30 + 1 + 3 slots) needs a third block.

# first revision of this PR (8e1f01430, lookahead only), new scheduler test applied
FAILED test_full_sequence_admission_covers_the_first_decode_step - assert 3 == 0   (num_preemptions)

# this PR
test_full_sequence_admission_covers_the_first_decode_step                 passed
test_full_sequence_admission_counts_the_first_decode_step                 4 passed

tests/v1/core/test_prefix_caching.py, test_kv_cache_utils.py, test_scheduler.py, test_priority_preemption_bug.py: 511 passed. Two tests in test_scheduler.py fail the same way without this change, for local reasons: the run used --noconftest (missing caplog_vllm fixture) and the host has no Hugging Face Hub access.

Environment: CPU-only host, this branch's sources on top of a vllm 0.30.0+cpu environment.

End to end on 1x A10 (driver 570, CUDA 13 forward-compatibility libraries), vllm 0.32.1rc1.dev368+g32551991b (the nightly wheel of this PR's base commit) with main's or this PR's kv_cache_manager.py and scheduler.py. Qwen2.5-1.5B-Instruct with Qwen2.5-0.5B-Instruct as a draft_model drafter (num_speculative_tokens=3, lookahead 3), block_size=64, a 7-block pool, max_num_batched_tokens=32, no prefix caching, greedy. Two decodes hold four blocks; v arrives with two blocks (128 slots) free.

v	Slots it then needs	main	this PR
125 tokens	132, first decode step	125 prefill tokens recomputed (1 preemption)	none
128 tokens	131, last chunk's lookahead	105 prefill tokens recomputed (1 preemption)	none

d0/d1 outputs are identical in every run, and the 125-token runs were repeated twice with the same result. With the fix v waits for a decode to free a block instead, so it finishes 1-3 engine steps later in this pool. v's own output follows prefill chunking: its first token is a tie (tokens 12 and 220 at -4.871 for the 125-token prompt with 32-token chunks).

Reproducer (MODE=decode_step or MODE=lookahead)

"""E2E reproducer: a prompt admitted under full-ISL reservation while decodes
hold the rest of the pool, sized so that one spec decode margin crosses into a
block the gate did not count.

MODE=decode_step (default): `v` ends `num_lookahead_tokens` short of its last
block. Its final prefill chunk (prompt + lookahead) fits exactly, but its first
decode step (prompt + 1 + num_spec_tokens + lookahead) needs one more block.
main and the lookahead-only revision admit it and preempt it on that step; the
fixed gate makes it wait.

MODE=lookahead: `v` fills its blocks exactly, so its final prefill chunk needs
one more block for the lookahead (the case the first revision covered).

ONLY_V=1 runs `v` alone; NUM_BLOCKS and V_BLOCKS resize the pool and `v`;
VARIANT only labels the output. Prints one JSON line: `v`'s preemptions split
by phase, the prefill tokens they threw away, and every request's output.
"""

import json
import os
import sys

os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

from vllm import LLM, SamplingParams  # noqa: E402
from vllm.inputs import TokensPrompt  # noqa: E402
from vllm.v1.core.sched.scheduler import Scheduler  # noqa: E402

BLOCK = 64
NUM_SPEC_TOKENS = 3
NUM_BLOCKS = int(os.environ.get("NUM_BLOCKS", "7"))
MODE = os.environ.get("MODE", "decode_step")

preemptions = []
original_preempt = Scheduler._preempt_request


def counting_preempt(self, request, *args, **kwargs):
    preemptions.append(
        {
            "id": request.request_id,
            # Set when the request's last scheduled chunk was still prefill.
            "phase": "prefill" if request.is_prefill_chunk else "decode_step",
            "computed": request.num_computed_tokens,
            "outputs": len(request.output_token_ids),
        }
    )
    return original_preempt(self, request, *args, **kwargs)


Scheduler._preempt_request = counting_preempt


def main() -> None:
    if MODE not in ("decode_step", "lookahead"):
        sys.exit(f"MODE must be decode_step or lookahead, got {MODE!r}")
    llm = LLM(
        model="Qwen/Qwen2.5-1.5B-Instruct",
        speculative_config={
            "method": "draft_model",
            "model": "Qwen/Qwen2.5-0.5B-Instruct",
            "num_speculative_tokens": NUM_SPEC_TOKENS,
        },
        block_size=BLOCK,
        num_gpu_blocks_override=NUM_BLOCKS,
        max_num_batched_tokens=32,
        max_num_seqs=4,
        max_model_len=256,
        enable_prefix_caching=False,
        gpu_memory_utilization=0.6,
        seed=0,
    )
    engine = llm.llm_engine
    scheduler = engine.engine_core.engine_core.scheduler
    manager = scheduler.kv_cache_manager
    num_groups = len(scheduler.kv_cache_config.kv_cache_groups)
    lookahead = scheduler.num_lookahead_tokens
    decode_step = 1 + scheduler.num_spec_tokens
    params = SamplingParams(temperature=0, max_tokens=40, ignore_eos=True)

    # Two decodes, each with most of its last block still free so that they do
    # not take a new block while `v` prefills.
    for n in range(0 if os.environ.get("ONLY_V") else 2):
        prompt = [1000 + 97 * n + i for i in range(BLOCK + 1)]
        engine.add_request(f"d{n}", TokensPrompt(prompt_token_ids=prompt), params)
    outputs = {}
    while any(
        r.num_computed_tokens < r.num_prompt_tokens
        for r in scheduler.requests.values()
    ):
        for out in engine.step():
            outputs[out.request_id] = out

    free = manager.block_pool.get_num_free_blocks()
    v_blocks = int(os.environ.get("V_BLOCKS", "0")) or free // num_groups
    v_len = v_blocks * BLOCK - (lookahead if MODE == "decode_step" else 0)
    setup = {
        "variant": os.environ.get("VARIANT", ""),
        "mode": MODE,
        "num_groups": num_groups,
        "free_blocks": free,
        "v_blocks": v_blocks,
        "v_len": v_len,
        "lookahead": lookahead,
        "decode_step_tokens": decode_step,
        "last_chunk_slots": v_len + lookahead,
        "first_decode_step_slots": v_len + decode_step + lookahead,
    }
    if v_blocks < 1 or free % num_groups or v_len < 1:
        print(json.dumps({"error": "pool does not split into v's blocks", **setup}))
        sys.exit(2)

    v_prompt = [5000 + i for i in range(v_len)]
    engine.add_request(
        "v",
        TokensPrompt(prompt_token_ids=v_prompt),
        SamplingParams(temperature=0, max_tokens=8, ignore_eos=True),
    )
    steps = 0
    while engine.has_unfinished_requests():
        for out in engine.step():
            outputs[out.request_id] = out
        steps += 1

    v_preempted = [p for p in preemptions if p["id"].split("-")[0] == "v"]
    print(
        json.dumps(
            {
                **setup,
                "steps_after_v": steps,
                "v_preempted_in_prefill": sum(
                    p["phase"] == "prefill" for p in v_preempted
                ),
                "v_preempted_on_decode_step": sum(
                    p["phase"] == "decode_step" for p in v_preempted
                ),
                "v_prefill_tokens_discarded": sum(p["computed"] for p in v_preempted),
                "all_preemptions": preemptions,
                "outputs": {
                    rid: list(out.outputs[0].token_ids)
                    for rid, out in sorted(outputs.items())
                },
            }
        )
    )


if __name__ == "__main__":
    main()

Details

Root cause. The gate passed request.num_tokens both as num_tokens, which should include the lookahead, and as num_tokens_main_model, unlike the allocation path. It now passes the sequence through the first decode step (num_decode_step_tokens = 1 + num_spec_tokens, the margin _request_remaining_blocks already adds) and the lookahead on top, both capped at max_model_len.

Tradeoff. A request whose sequence ends within 1 + num_spec_tokens + num_lookahead_tokens of a block boundary needs one more free block to be admitted; with the pool full it would be preempted on that step anyway. Without speculative decoding the gate is unchanged, so a block-aligned prompt's first decode token stays unreserved as before.

How it was found. In vllm-rbln, whose scheduler never mixes prefill and decode, this mismatch is a livelock: the self-preempted prefill evicts the decode batch on every re-admission, so no block is ever freed (RBLN-SW/vllm-rbln#1330). On upstream's scheduler it shows as recomputed prefills instead. vllm-rbln patches the gate today; this fix lets it drop the patch.

Not a duplicate. #57104, #59880 and #63136 handle the spec decode margin for async KV loads (reserved_blocks); this PR covers every other admission. #57658 aligns Mamba admission with dynamic K; #52520 / #52530 concern Mamba align pool sizing.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀