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

@@ -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<SequencePtr> 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();

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()