From b6ddc56a0c811ea26b561bce43edba87f067f9fb Mon Sep 17 00:00:00 2001 From: Jahidul Hossain Mekat Date: Tue, 20 Oct 2026 14:49:40 +0000 Subject: [PATCH 1/3] Release the KV blocks and state slot of requests aborted while waiting A request that is aborted (or fails multimodal processing / structured output setup) while it sits in `waiting` is popped by the scheduler and handed back as an ignored sequence. Nothing released what it still held. A long prompt that is prefilled in chunks is moved back to `waiting` after each chunk with the KV blocks of the whole prompt (and its state slot) already allocated, so aborting it between two chunks leaked those blocks for the lifetime of the engine, and so did a fork that pinned its parent's blocks. Every such abort shrinks the cache until the engine can no longer admit new requests. _schedule_prefill and the waiting sweep in _schedule_decode now release the sequence's resources when they drop it. free_seq_nolock also reset nothing after freeing the state slot, so releasing a sequence a second time (a preempted request that is aborted had already given its slot back) freed the slot twice and tripped the block manager's "ref_count corrupted" check. It now clears state_id, which makes the release idempotent like the block table already was. Tests drive the real scheduler through schedule / free / process_outputs on a CPU process group. A decode step in the middle of a prefill is forced through decode_step, which is what a data-parallel rank sees when the other ranks do not agree to a prompt step. --- csrc/engine/scheduler.hpp | 5 + tests/engine/test_scheduler_resources.py | 215 +++++++++++++++++++++++ 2 files changed, 220 insertions(+) create mode 100644 tests/engine/test_scheduler_resources.py diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index ed6e897..ac73888 100644 --- a/csrc/engine/scheduler.hpp +++ b/csrc/engine/scheduler.hpp @@ -362,6 +362,8 @@ struct Scheduler { // 入队前已判死的 seq 不跑 prefill，进 ignored_seqs 带 finish_reason 回传 if (seq->schedule.status == SequenceStatus::FINISHED_ABORTED || seq->schedule.status == SequenceStatus::FINISHED_VE_FAILED || seq->schedule.status == SequenceStatus::FINISHED_SO_FAILED) { + // A chunked prefill that was aborted between chunks (or a fork) still holds KV blocks / a state slot. + free_seq_nolock(*seq); ignored_seqs.push_back(seq); waiting.pop_front(); continue; @@ -569,6 +571,7 @@ struct Scheduler { std::vector ignored_seqs; for (auto it = waiting.begin(); it != waiting.end();) { if ((*it)->is_finished()) { + free_seq_nolock(**it); // see _schedule_prefill: an aborted waiting seq may still hold KV blocks / a state slot ignored_seqs.push_back(std::move(*it)); it = waiting.erase(it); } else { @@ -735,10 +738,12 @@ struct Scheduler { return py::make_tuple(std::move(seq_metadata_list), std::move(p_scheduler_outputs)); } + /// Release everything the sequence holds. Safe to call again: nothing is left to release the second time. void free_seq_nolock(Sequence &seq) { release_mtp_info_slot(seq); if (seq.cache.state_id != -1) { state_manager->free_blocks_nolock({seq.cache.state_id}); + seq.cache.state_id = -1; } block_manager->free_blocks_nolock(seq.cache.block_table); seq.cache.block_table.clear(); diff --git a/tests/engine/test_scheduler_resources.py b/tests/engine/test_scheduler_resources.py new file mode 100644 index 0000000..caa6693 --- /dev/null +++ b/tests/engine/test_scheduler_resources.py @@ -0,0 +1,215 @@ +"""KV blocks and state slots must always find their way back to the managers. + +Drives the real C++ scheduler through the part of LLMEngine.step() that touches it: schedule, free the finished +sequences, feed the (fake) model outputs back. No GPU or model is needed, only the compiled extension and a CPU +process group. +""" + +from __future__ import annotations + +from collections.abc import Generator +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch # noqa: F401 Load torch libraries before the C++ extension. +import torch.distributed as dist + +from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool, process_MTPInfo +from deepinfer.sampling_params import SamplingParams + +pytestmark = [pytest.mark.ci, pytest.mark.ci_npu] + +EOS = 1 +VOCAB_SIZE = 1000 +GENERATED_TOKEN = 5 # never EOS, so requests end by max_tokens + + +@pytest.fixture(scope="module") +def cpu_group(tmp_path_factory: pytest.TempPathFactory) -> Generator[dist.ProcessGroup]: + initialized_here = not dist.is_initialized() + if initialized_here: + rendezvous = Path(tmp_path_factory.mktemp("rendezvous")) / "store" + dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=0, world_size=1) + yield dist.new_group(ranks=[0], backend="gloo") + if initialized_here: + dist.destroy_process_group() + + +class _Engine: + """schedule -> free finished sequences -> process outputs, without speculative decoding or zero-overhead scheduling.""" + + def __init__( + self, + group: dist.ProcessGroup, + monkeypatch: pytest.MonkeyPatch, + *, + num_blocks: int, + block_size: int, + max_num_batched_tokens: int, + num_states: int = 0, + max_num_seqs: int = 2, + prefill_attempt_interval: int = 2, + ) -> None: + monkeypatch.setenv("DEEPINFER_USE_NEXTN_DECODE", "0") # read by the C++ scheduler when it is built + self.num_blocks = num_blocks + self.num_states = num_states + self.block_size = block_size + self.block_manager = BlockManager(num_blocks) + self.state_manager = BlockManager(num_states) if num_states else None + max_model_len = 64 + config = SimpleNamespace( + max_model_len=max_model_len, + runtime_max_model_len=max_model_len + 16, + max_num_batched_tokens=max_num_batched_tokens, + max_num_seqs=max_num_seqs, + stop_preempted_request=False, + chunked_prefill=True, + engine_type=0, # EngineType.NORMAL + prefill_attempt_interval=prefill_attempt_interval, + prefill_any_attempt_interval=prefill_attempt_interval + 1, + ) + self.scheduler = Scheduler( + block_size, self.block_manager, self.state_manager, EOS, group, config, 0, 0, 2**30, None, [], 1, VOCAB_SIZE, False + ) + self.pool = SequenceMetadataPool() + self.sequences: dict[str, Sequence] = {} + + def add_request(self, prompt_len: int, max_tokens: int) -> Sequence: + request_id = str(len(self.sequences)) + prompt = np.full(prompt_len, 7, dtype=np.int32) + seq = Sequence( + seq_id=len(self.sequences), prompt_token_ids=prompt, compressed_prompt_token_ids=prompt, multi_modal_data=None, + block_size=self.block_size, request_id=request_id, user_id="user", + sampling_params=SamplingParams(max_tokens=max_tokens, eos_id=EOS), arrival_time=0.0, golden_token_ids=None, + ) # fmt: skip + self.sequences[request_id] = seq + self.scheduler.add_sequence(seq) + return seq + + def step(self, *, decode_only: bool = False) -> None: + if decode_only: + # In a data-parallel engine every rank has to agree to a prompt step, so a rank can run decode steps between two chunks + # of a prompt it is prefilling. A decode_step that is no attempt point gives the same schedule on a single rank. + self.scheduler.decode_step = 1 + seq_metadata_list, outputs = self.scheduler.schedule(self.pool) + for seq in outputs.finished_seqs: + self.scheduler.free_seq(seq) + n = len(outputs.scheduled_seqs) + process_MTPInfo(outputs, False, torch.zeros((n, 0), dtype=torch.int32), None) + self.scheduler.process_outputs( + outputs, torch.empty(0, dtype=torch.int32), torch.full((n,), GENERATED_TOKEN, dtype=torch.int32), None + ) + self.pool.release(seq_metadata_list) + + def run(self, max_steps: int = 300) -> None: + for _ in range(max_steps): + if not self.scheduler.has_unfinished_seqs(): + return + self.step() + raise AssertionError(f"no progress after {max_steps} steps: {self.scheduler.get_num_unfinished_seqs()} requests unfinished") + + def assert_all_resources_returned(self) -> None: + # Block 0 (and state slot 0) is reserved by the managers and never handed out. + assert self.scheduler.get_num_free_gpu_blocks() == self.num_blocks - 1 + if self.state_manager is not None: + assert self.state_manager.get_num_free_blocks() == self.num_states - 1 + + +@pytest.mark.parametrize("num_states", [0, 8], ids=["no_state", "with_state"]) +def test_abort_between_prefill_chunks_releases_the_sequence_resources( + cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, num_states: int +) -> None: + engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=num_states) + long_prompt = engine.add_request(prompt_len=30, max_tokens=3) # 4 chunks of at most 8 tokens + engine.add_request(prompt_len=6, max_tokens=3) + + engine.step() + # The first chunk is done: the request waits for its next chunk and already owns the KV blocks of the whole prompt. + assert long_prompt.status.name == "WAITING" + assert long_prompt.num_cached_tokens == 8 + assert engine.scheduler.get_num_free_gpu_blocks() < engine.num_blocks - 1 + + engine.scheduler.abort_sequence([long_prompt.request_id]) + engine.run() + + assert long_prompt.status.name == "FINISHED_ABORTED" + engine.assert_all_resources_returned() + + +@pytest.mark.parametrize("num_states", [0, 8], ids=["no_state", "with_state"]) +def test_abort_during_a_decode_step_releases_the_sequence_resources( + cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, num_states: int +) -> None: + engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=num_states) + engine.add_request(prompt_len=6, max_tokens=12) + engine.step() + long_prompt = engine.add_request(prompt_len=30, max_tokens=3) # 4 chunks of at most 8 tokens + engine.step() + assert long_prompt.status.name == "WAITING" + assert long_prompt.num_cached_tokens == 8 + + # The aborted prompt is dropped from the queue by a decode step this time, not by a prompt step. + engine.scheduler.abort_sequence([long_prompt.request_id]) + engine.step(decode_only=True) + assert long_prompt.status.name == "FINISHED_ABORTED" + assert not any(seq.request_id == long_prompt.request_id for seq in engine.scheduler.waiting) + + engine.run() + engine.assert_all_resources_returned() + + +def test_abort_before_the_first_chunk_leaves_nothing_to_release(cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch) -> None: + engine = _Engine(cpu_group, monkeypatch, num_blocks=17, block_size=4, max_num_batched_tokens=8, num_states=8) + request = engine.add_request(prompt_len=30, max_tokens=3) + engine.scheduler.abort_sequence([request.request_id]) + + engine.run() + + assert request.status.name == "FINISHED_ABORTED" + engine.assert_all_resources_returned() + + +def _preempt_during_chunked_prefill( + group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch, *, num_states: int = 0 +) -> tuple[_Engine, list[Sequence], Sequence]: + """Two requests decode while a long prompt is prefilled in chunks, and the cache is too small for all of them to keep growing. + + The cache has 6 usable blocks of 4 tokens. The prompt (12 tokens, chunks of 4) reserves 3 of them with its first chunk, the + decoding requests hold one each, and the last block goes to one of them for its next token. The other one needs a new block too, + is preempted and has to be recomputed, which takes 2 blocks while only the 1 it just released is free. + """ + engine = _Engine(group, monkeypatch, num_blocks=7, block_size=4, max_num_batched_tokens=4, max_num_seqs=4, num_states=num_states) + decoding = [engine.add_request(prompt_len=4, max_tokens=12), engine.add_request(prompt_len=4, max_tokens=12)] + long_prompt = engine.add_request(prompt_len=12, max_tokens=2) + for _ in range(10): + if long_prompt.num_cached_tokens: + break + engine.step() + assert long_prompt.status.name == "WAITING" + assert long_prompt.num_cached_tokens == 4 + assert engine.scheduler.get_num_free_gpu_blocks() == 1 + + for _ in range(10): + if engine.scheduler.get_num_preempted_seqs(): + break + engine.step(decode_only=True) + assert engine.scheduler.get_num_preempted_seqs() == 1 + assert sorted(seq.status.name for seq in decoding) == ["RUNNING", "WAITING"] + assert long_prompt.status.name == "WAITING" + return engine, decoding, long_prompt + + +def test_abort_of_a_preempted_request_does_not_free_its_state_slot_twice( + cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch +) -> None: + engine, decoding, _ = _preempt_during_chunked_prefill(cpu_group, monkeypatch, num_states=8) + preempted = next(seq for seq in decoding if seq.status.name == "WAITING") + # Preemption already gave its state slot back; releasing the aborted sequence again must not release it a second time. + engine.scheduler.abort_sequence([preempted.request_id]) + + engine.run() + + assert preempted.status.name == "FINISHED_ABORTED" + engine.assert_all_resources_returned() From e25a404e2c0da62fef03145a4e6580453bfc89e9 Mon Sep 17 00:00:00 2001 From: Jahidul Hossain Mekat Date: Tue, 20 Oct 2026 14:49:40 +0000 Subject: [PATCH 2/3] Don't let a preempted request block a chunked prefill that holds blocks A request that is preempted for recompute was pushed to the front of `waiting`. _schedule_prefill takes `waiting` in order and stops at the first sequence it cannot allocate, and a preempted request usually does not fit right away: it was preempted because it needed a block more than it had, and recomputing it takes that block as well. If a chunked prefill is in progress at that moment (it is back at the front of the queue holding the KV blocks of its whole prompt), the preempted request lands in front of it and ends the scan before the prefill is reached. The prefill never gets its next chunk, so it never finishes and never releases the blocks the preempted request is waiting for, and the engine stops making progress (or keeps recomputing the requests that are left, which are preempted again for the same reason). Decode steps between two chunks of a prompt are all it takes. They happen when the cache usage is above the prefill watermark, and on a data-parallel rank whenever the other ranks do not agree to a prompt step. Queue the preempted request ahead of the fresh requests, as before, but behind the sequences that already hold blocks. They only need compute to finish. A sequence that cannot be scheduled yet (a multimodal request that is still being preprocessed, but already owns its blocks as in examples/benchmark.py) is not one of them: _schedule_prefill stops at it, so the preempted request stays in front of it. --- csrc/engine/scheduler.hpp | 10 ++++- tests/engine/test_scheduler_resources.py | 49 ++++++++++++++++++++++-- 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index ac73888..b24fe19 100644 --- a/csrc/engine/scheduler.hpp +++ b/csrc/engine/scheduler.hpp @@ -787,7 +787,15 @@ struct Scheduler { LOG_WARN("Request {} is preempted and will be recomputed.", seq->request.request_id); seq->reset_for_recompute(); // 复位 seq 自身的滚动状态（cache.num_tokens、inflight 计数、MTP 状态），清单见 // Sequence::reset_for_recompute - waiting.push_front(seq); + // Queue it ahead of the fresh requests but behind the chunked prefills in progress. Those already hold the KV blocks of + // their whole prompt and only need compute to finish. _schedule_prefill stops at the first sequence it cannot allocate, and + // a preempted request usually does not fit right away (it needs a block more than it released): in front of them it would + // keep them from ever being scheduled, and the blocks they hold from ever being released. A sequence that cannot be + // scheduled yet (its multimodal inputs are still being processed) does not count as one of them, it would hold this + // request up instead. + waiting.insert(std::find_if(waiting.begin(), waiting.end(), + [](const SequencePtr &s) { return s->cache.block_table.empty() || !s->schedule.schedulable; }), + seq); seq->schedule.status = SequenceStatus::WAITING; free_seq_nolock(*seq); } diff --git a/tests/engine/test_scheduler_resources.py b/tests/engine/test_scheduler_resources.py index caa6693..02a1b16 100644 --- a/tests/engine/test_scheduler_resources.py +++ b/tests/engine/test_scheduler_resources.py @@ -1,4 +1,4 @@ -"""KV blocks and state slots must always find their way back to the managers. +"""KV blocks and state slots must always find their way back to the managers, and the requests holding them must always be able to finish. Drives the real C++ scheduler through the part of LLMEngine.step() that touches it: schedule, free the finished sequences, feed the (fake) model outputs back. No GPU or model is needed, only the compiled extension and a CPU @@ -76,15 +76,19 @@ class _Engine: self.pool = SequenceMetadataPool() self.sequences: dict[str, Sequence] = {} - def add_request(self, prompt_len: int, max_tokens: int) -> Sequence: + def add_request( + self, prompt_len: int, max_tokens: int, *, multi_modal_data: dict | None = None, allocate_blocks: bool = False + ) -> Sequence: request_id = str(len(self.sequences)) prompt = np.full(prompt_len, 7, dtype=np.int32) seq = Sequence( - seq_id=len(self.sequences), prompt_token_ids=prompt, compressed_prompt_token_ids=prompt, multi_modal_data=None, + seq_id=len(self.sequences), prompt_token_ids=prompt, compressed_prompt_token_ids=prompt, multi_modal_data=multi_modal_data, block_size=self.block_size, request_id=request_id, user_id="user", sampling_params=SamplingParams(max_tokens=max_tokens, eos_id=EOS), arrival_time=0.0, golden_token_ids=None, ) # fmt: skip self.sequences[request_id] = seq + if allocate_blocks: + assert self.scheduler.try_allocate_seq(seq) self.scheduler.add_sequence(seq) return seq @@ -213,3 +217,42 @@ def test_abort_of_a_preempted_request_does_not_free_its_state_slot_twice( assert preempted.status.name == "FINISHED_ABORTED" engine.assert_all_resources_returned() + + +def test_preempted_request_does_not_block_a_chunked_prefill_that_holds_blocks( + cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch +) -> None: + engine, decoding, long_prompt = _preempt_during_chunked_prefill(cpu_group, monkeypatch) + + # The preempted request cannot be recomputed before the long prompt has finished its prefill and released its blocks. + engine.run() + + assert [seq.status.name for seq in decoding] == ["FINISHED_LENGTH_CAPPED"] * 2 + assert long_prompt.status.name == "FINISHED_LENGTH_CAPPED" + engine.assert_all_resources_returned() + + +def test_preempted_request_is_not_queued_behind_a_sequence_that_cannot_be_scheduled( + cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch +) -> None: + # 7 usable blocks of 4 tokens. The two decoding requests need 3 blocks each in the end, the multimodal request holds 3. + engine = _Engine(cpu_group, monkeypatch, num_blocks=8, block_size=4, max_num_batched_tokens=64, max_num_seqs=3) + decoding = [engine.add_request(prompt_len=4, max_tokens=8), engine.add_request(prompt_len=4, max_tokens=8)] + # It is not schedulable until its images are processed, but it owns its KV blocks (examples/benchmark.py allocates them up front). + multimodal = engine.add_request(prompt_len=12, max_tokens=3, multi_modal_data={"image": 1}, allocate_blocks=True) + assert not multimodal.schedulable + + for _ in range(40): + if all(seq.status.name == "FINISHED_LENGTH_CAPPED" for seq in decoding): + break + engine.step() + + # The cache is too small for both, so one of them was preempted and recomputed. Queued behind the multimodal request, which + # cannot be scheduled yet, it would wait for it forever. + assert engine.scheduler.get_num_preempted_seqs() > 0 + assert [seq.status.name for seq in decoding] == ["FINISHED_LENGTH_CAPPED"] * 2 + + engine.scheduler.set_schedulable(multimodal.request_id) + engine.run() + assert multimodal.status.name == "FINISHED_LENGTH_CAPPED" + engine.assert_all_resources_returned() From 6a934611d538bcce54434c73a886459bca36b021 Mon Sep 17 00:00:00 2001 From: Jahidul Hossain Mekat Date: Tue, 20 Oct 2026 14:49:40 +0000 Subject: [PATCH 3/3] Let a chunked prefill continue above the prefill watermark A prompt that is prefilled in chunks reserves the KV blocks of its whole prompt with its first chunk and goes back to `waiting` until its next chunk. NORMAL engines only try a prompt step while the cache usage is at most GPU_CACHE_PREFILL_WATERMARK (0.9), which is meant to keep fresh prompts from filling the cache. The blocks the prompt reserved count in that usage, and only its remaining chunks can release them, so it was never scheduled again as long as they alone kept the usage above the watermark. A prompt that needs more than 90% of the cache by itself (a long context on a cache that barely holds it) left the engine idle for good, with the request unfinished. Skip the watermark when the request at the front of `waiting` is a chunked prefill in progress. It needs no new blocks. The test prefills a 37-token prompt that takes all 10 usable blocks of a small cache. --- csrc/engine/scheduler.hpp | 6 +++++- tests/engine/test_scheduler_resources.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index b24fe19..e82cbb7 100644 --- a/csrc/engine/scheduler.hpp +++ b/csrc/engine/scheduler.hpp @@ -310,7 +310,11 @@ struct Scheduler { if (!is_any_point && !is_all_point) { return _schedule_decode(acquire); } - bool want = (get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK) && _has_prefill_waiting(); + // A prompt prefilled in chunks reserves the blocks of its whole prompt with its first chunk, and only its remaining chunks can + // give them back. Holding it to the watermark would stall it for as long as its own blocks keep the usage above it, which is + // for good if it needs more than that share of the cache by itself. + const bool continues_chunked_prefill = !waiting.empty() && waiting.front()->schedule.chunked_rest_len > 0; + bool want = _has_prefill_waiting() && (continues_chunked_prefill || get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK); bool do_prefill; if (is_any_point) { do_prefill = (_dp_allreduce(want ? 0 : 1, c10d::ReduceOp::MIN) == 0); // ANY：任一 rank 想做即触发 diff --git a/tests/engine/test_scheduler_resources.py b/tests/engine/test_scheduler_resources.py index 02a1b16..b2fe054 100644 --- a/tests/engine/test_scheduler_resources.py +++ b/tests/engine/test_scheduler_resources.py @@ -256,3 +256,20 @@ def test_preempted_request_is_not_queued_behind_a_sequence_that_cannot_be_schedu engine.run() assert multimodal.status.name == "FINISHED_LENGTH_CAPPED" engine.assert_all_resources_returned() + + +def test_prompt_that_takes_nearly_the_whole_cache_is_prefilled_to_the_end( + cpu_group: dist.ProcessGroup, monkeypatch: pytest.MonkeyPatch +) -> None: + # 10 usable blocks of 4 tokens. The first chunk reserves all of them, far above the usage at which new prompts are admitted. + engine = _Engine(cpu_group, monkeypatch, num_blocks=11, block_size=4, max_num_batched_tokens=8) + prompt = engine.add_request(prompt_len=37, max_tokens=2) # 5 chunks of at most 8 tokens + + engine.step() + assert prompt.status.name == "WAITING" + assert engine.scheduler.get_num_free_gpu_blocks() == 0 + + engine.run() + + assert prompt.status.name == "FINISHED_LENGTH_CAPPED" + engine.assert_all_resources_returned()