Let a chunked prefill continue above the prefill watermark

A prompt that is prefilled in chunks reserves the KV blocks of its whole prompt
with its first chunk and goes back to `waiting` until its next chunk. NORMAL
engines only try a prompt step while the cache usage is at most
GPU_CACHE_PREFILL_WATERMARK (0.9), which is meant to keep fresh prompts from
filling the cache. The blocks the prompt reserved count in that usage, and only
its remaining chunks can release them, so it was never scheduled again as long
as they alone kept the usage above the watermark. A prompt that needs more than
90% of the cache by itself (a long context on a cache that barely holds it) left
the engine idle for good, with the request unfinished.

Skip the watermark when the request at the front of `waiting` is a chunked
prefill in progress. It needs no new blocks.

The test prefills a 37-token prompt that takes all 10 usable blocks of a small
cache.

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 想做即触发

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