diff --git a/csrc/engine/scheduler.hpp b/csrc/engine/scheduler.hpp index ed6e897..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 想做即触发 @@ -362,6 +366,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 +575,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 +742,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(); @@ -782,7 +791,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 new file mode 100644 index 0000000..b2fe054 --- /dev/null +++ b/tests/engine/test_scheduler_resources.py @@ -0,0 +1,275 @@ +"""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 +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, *, 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=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 + + 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() + + +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() + + +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()