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