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