Don't let a preempted request block a chunked prefill that holds blocks

A request that is preempted for recompute was pushed to the front of `waiting`.
_schedule_prefill takes `waiting` in order and stops at the first sequence it
cannot allocate, and a preempted request usually does not fit right away: it was
preempted because it needed a block more than it had, and recomputing it takes
that block as well. If a chunked prefill is in progress at that moment (it is
back at the front of the queue holding the KV blocks of its whole prompt), the
preempted request lands in front of it and ends the scan before the prefill is
reached. The prefill never gets its next chunk, so it never finishes and never
releases the blocks the preempted request is waiting for, and the engine stops
making progress (or keeps recomputing the requests that are left, which are
preempted again for the same reason).

Decode steps between two chunks of a prompt are all it takes. They happen when
the cache usage is above the prefill watermark, and on a data-parallel rank
whenever the other ranks do not agree to a prompt step.

Queue the preempted request ahead of the fresh requests, as before, but behind
the sequences that already hold blocks. They only need compute to finish. A
sequence that cannot be scheduled yet (a multimodal request that is still being
preprocessed, but already owns its blocks as in examples/benchmark.py) is not
one of them: _schedule_prefill stops at it, so the preempted request stays in
front of it.

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);
         }

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