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