[Scheduler] Fuse fixed-depth decode loops into one submission, keeping the per-loop prefill share

tests/test_async_pipeline.py

@@ -198,10 +198,13 @@ def execute_with_delay(batch, prepared=None):
 
     def inspect_prepare(batch, **kwargs):
         result = prepare(batch, **kwargs)
-        if injected and not observations and batch.operation == Operation.RECURRENT:
+        # The submission after the delayed core (its CODA once loops are fused,
+        # or the next loop when they are not) is prepared while the core runs.
+        if injected and not observations:
             assert not delayed.query(), "prepare waited for the preceding forward"
-            assert result.kv.position_ids.device.type == "cpu"
-            assert result.kv.position_ids.is_pinned()
+            if result.kv is not None:
+                assert result.kv.position_ids.device.type == "cpu"
+                assert result.kv.position_ids.is_pinned()
             observations.append(True)
         return result
 

tests/test_engine.py

@@ -22,7 +22,7 @@ def drain(engine, limit=1000):
             return outputs, trace
         emitted = engine.step()
         batch = engine.last_schedule
-        trace.append((batch.phase, tuple(i.request.request_id for i in batch.requests)))
+        trace.append((batch.phase, tuple(i.request.request_id for i in batch.requests), batch))
         for output in emitted:
             if output.finished:
                 outputs[output.request_id] = output
@@ -80,16 +80,48 @@ def test_refill_recycles_exited_token_before_slow_token_coda(mode):
     engine.add_request("fast", [1], SamplingParams(max_tokens=3, exit_threshold=0, ignore_eos=True))
     engine.add_request("slow", [2], SamplingParams(max_tokens=3, exit_threshold=1, ignore_eos=True))
     outputs, trace = drain(engine)
-    recurrent = [ids for stage, ids in trace if stage == RequestPhase.RECURRENT]
+    recurrent = [ids for stage, ids, _ in trace if stage == RequestPhase.RECURRENT]
     # In refill mode fast returns with its next token while slow is on its first.
+    # Alone in a batch, slow's remaining fixed-depth loops run as one submission.
     assert recurrent[:4] == (
         [("fast", "slow")] * 2 + [("slow", "fast")] * 2
         if mode == "refill"
-        else [("fast", "slow")] * 2 + [("slow",)] * 2
+        else [("fast", "slow")] * 2 + [("slow",), ("fast", "slow")]
     )
     assert set(outputs) == {"fast", "slow"}
 
 
+@pytest.mark.parametrize("mode", ["refill", "mixed"])
+def test_fixed_depth_decode_fuses_remaining_loops_without_changing_outputs(mode):
+    # Fixed-depth decode rows need no gate readback, so their loops share one submission.
+    model = tiny_model()
+    params = SamplingParams(max_tokens=4, exit_threshold=1, ignore_eos=True)
+    shallow = SamplingParams(max_tokens=4, max_loops=3, ignore_eos=True)
+    fused = LLMEngine(model, scheduler_config=SchedulerConfig(mode=mode))
+    fused.add_request("a", [1, 2], params)
+    fused.add_request("b", [3], shallow)
+    outputs, trace = drain(fused)
+    counts = [batch.loop_count for stage, _, batch in trace if stage == RequestPhase.RECURRENT]
+    # A batch advances by the smallest remaining depth among its rows, never one by one.
+    assert counts and max(counts) > 1 and all(1 <= c <= 4 for c in counts), counts
+    assert outputs["a"].exit_depths == [4, 4, 4, 4]
+    assert outputs["b"].exit_depths == [4, 3, 3, 3]
+    serial = LLMEngine(model, scheduler_config=SchedulerConfig(mode=mode))
+    serial.scheduler._fused_decode_loops = lambda scheduled_requests: 1
+    serial.add_request("a", [1, 2], params)
+    serial.add_request("b", [3], shallow)
+    reference, reference_trace = drain(serial)
+    assert all(
+        batch.loop_count == 1
+        for stage, _, batch in reference_trace
+        if stage == RequestPhase.RECURRENT
+    )
+    for rid in ("a", "b"):
+        assert reference[rid].token_ids == outputs[rid].token_ids
+        assert reference[rid].exit_depths == outputs[rid].exit_depths
+    assert len(reference_trace) > len(trace)
+
+
 def test_memory_pressure_waits_then_reuses_pages_and_abort_reclaims():
     engine = LLMEngine(tiny_model(), cache_config=CacheConfig(8, 2))
     params = SamplingParams(max_tokens=3, ignore_eos=True)
@@ -141,7 +173,7 @@ def test_eos_finishes_without_extra_decode_and_invalid_requests_do_not_poison_qu
     outputs, trace = drain(engine)
     assert outputs["eos"].token_ids == [0]
     assert outputs["eos"].finish_reason == "stop"
-    assert [stage for stage, _ in trace] == [RequestPhase.PREFILL, RequestPhase.CODA]
+    assert [stage for stage, *_ in trace] == [RequestPhase.PREFILL, RequestPhase.CODA]
     assert engine.cache_manager.num_used_blocks == 0
 
 
@@ -333,6 +365,8 @@ def test_repeated_short_arrivals_cannot_starve_existing_decode(mode):
     while not engine.step():
         pass
     # Leave the request at PRELUDE, also covering no_refill's fill phase.
+    # Count decode progress in loops: a fused submission advances several loops and
+    # leaves the same per-loop prefill share, so batches alone undercount decode.
     recurrent_steps = 0
     for index in range(12):
         request_id = f"short-{index}"
@@ -344,7 +378,7 @@ def test_repeated_short_arrivals_cannot_starve_existing_decode(mode):
                 scheduled_request.request.request_id == "long"
                 for scheduled_request in batch.requests
             ):
-                recurrent_steps += 1
+                recurrent_steps += batch.loop_count or 1
             if request_id not in engine.scheduler.requests:
                 break
         else:
@@ -415,3 +449,39 @@ def test_finish_reason_preserves_public_string_values(reason):
     output = RequestOutput.from_request(request)
     assert output.finished and type(output.finish_reason) is str
     assert output.finish_reason == reason.value
+
+
+def test_fused_decode_keeps_the_per_loop_prefill_share():
+    """A fused L-loop decode step leaves L prefill batches, so late prompts are admitted as fast
+    as with one-loop steps, without changing any output."""
+
+    def run(share):
+        model = tiny_model()
+        engine = LLMEngine(
+            model,
+            scheduler_config=SchedulerConfig(
+                max_num_seqs=4,
+                max_num_batched_tokens=4,
+                prefill_chunk_size=4,
+                prefill_share_per_loop=share,
+            ),
+        )
+        params = SamplingParams(max_tokens=12, exit_threshold=1, ignore_eos=True)
+        engine.add_request("early", [1, 2, 3], params)
+        for _ in range(4):  # early is decoding with fused loops
+            engine.step()
+        engine.add_request("late", list(range(1, 25)), params)  # six prefill chunks
+        outputs, first_late_token, step = {}, None, 0
+        while engine.has_unfinished_requests():
+            step += 1
+            for output in engine.step():
+                if output.request_id == "late" and output.token_ids and first_late_token is None:
+                    first_late_token = step
+                if output.finished:
+                    outputs[output.request_id] = (list(output.token_ids), list(output.exit_depths))
+        return outputs, first_late_token
+
+    shared, late_shared = run(True)
+    single, late_single = run(False)
+    assert shared == single  # scheduling order only; tokens and exit depths unchanged
+    assert late_shared < late_single, (late_shared, late_single)

tests/test_input_preparation.py

@@ -162,11 +162,15 @@ def reject(*args, **kwargs):
         pytest.fail("graph selection, capture and input staging belong to preparation")
 
     def checked_execute(scheduled, prepared):
+        # Graph selection and capture belong to preparation. Input staging of a
+        # fused later loop necessarily follows its producer loop inside execution;
+        # the first loop's inputs are still staged in preparation.
         with monkeypatch.context() as patch:
             patch.setattr(runner.graphs, "select", reject)
-            patch.setattr(runner.graphs, "prepare_inputs", reject)
-            patch.setattr(cuda_graph, "_copy_graph_inputs", reject)
             patch.setattr(cuda_graph, "_capture", reject)
+            if prepared.batch.loop_count <= 1:
+                patch.setattr(runner.graphs, "prepare_inputs", reject)
+                patch.setattr(cuda_graph, "_copy_graph_inputs", reject)
             return execute(scheduled, prepared)
 
     monkeypatch.setattr(runner, "_execute", checked_execute)

vllm_rlt/config.py

@@ -61,6 +61,9 @@ class SchedulerConfig:
     policy: str = "fcfs"
     enable_preemption: bool = False
     wavefront_chunked_prefill: bool = False
+    # A decode submission that fuses L loops leaves L times the prefill-batch share
+    # (max_prefill_batches_before_decode per loop), as L unfused loop steps would.
+    prefill_share_per_loop: bool = True
 
     def __post_init__(self):
         for name in (
@@ -77,6 +80,8 @@ def __post_init__(self):
             raise ValueError("enable_preemption must be a boolean")
         if type(self.wavefront_chunked_prefill) is not bool:
             raise ValueError("wavefront_chunked_prefill must be a boolean")
+        if type(self.prefill_share_per_loop) is not bool:
+            raise ValueError("prefill_share_per_loop must be a boolean")
         if self.policy not in ("fcfs", "priority"):
             raise ValueError("policy must be fcfs or priority")
         if self.mode not in ("refill", "no_refill", "mixed"):

vllm_rlt/core/scheduler.py

@@ -771,8 +771,8 @@ def take_phases(self, phases, *, eligible=None):
             if prompt_phases
             else scheduled_requests[0].phase
         )
-        self.policy.record_batch(recorded)
         loop_count = self._planned_loop_count(scheduled_requests)
+        self.policy.record_batch(recorded, loop_count=loop_count)
         scheduled_requests = [
             replace(
                 item,
@@ -800,7 +800,11 @@ def _planned_loop_count(self, scheduled_requests):
         if RequestPhase.PRELUDE in phases:
             return 0
         if self.wavefront_chunked_prefill:
-            return 1 if RequestPhase.PREFILLING_RECURRENT in phases else 0
+            if RequestPhase.PREFILLING_RECURRENT in phases:
+                return 1
+            if phases == {RequestPhase.RECURRENT}:
+                return self._fused_decode_loops(scheduled_requests)
+            return 0
         if RequestPhase.PREFILLING_PRELUDE in phases:
             # In mixed mode, stage just the embeddings when recurrent decode work
             # is ready so the next depth can join that batch. Without a reason to
@@ -815,9 +819,40 @@ def _planned_loop_count(self, scheduled_requests):
                 return 1
             return min(self._max_contiguous_prefill_loops(item) for item in scheduled_requests)
         if RequestPhase.RECURRENT in phases:
-            return 1
+            return self._fused_decode_loops(scheduled_requests)
         raise ValueError(f"cannot derive recurrent boundary for phases: {phases}")
 
+    def _fused_decode_loops(self, scheduled_requests):
+        """Loops one decode RECURRENT submission may run without a host decision.
+
+        A row whose exit depth is already known (fixed depth, or a trace) needs no
+        gate readback between loops, so its remaining loops can run back to back
+        in one submission with per-loop KV metadata. Any row that still depends
+        on a hazard score keeps the one-loop boundary for the whole batch. The
+        batch advances every row by the same count, so the minimum remaining
+        depth bounds the fusion; rows with more loops left return to the
+        RECURRENT queue afterwards.
+        """
+        if self.cache_manager.layout != "last_exited":
+            # SHARED overwrites one plane through recurrence; keep its per-loop boundaries.
+            return 1
+        count = None
+        for item in scheduled_requests:
+            request = item.request
+            params = request.sampling_params
+            maximum = params.max_loops or self.total_ut_steps
+            if self.exit_mode == "trace":
+                target = request.exit_trace[request.num_scheduled_outputs]
+            elif params.exit_threshold >= 1.0:
+                target = maximum
+            else:
+                return 1
+            remaining = min(target, maximum) - request.loops_done
+            if remaining < 1:
+                return 1
+            count = remaining if count is None else min(count, remaining)
+        return count or 1
+
     def _batch_loop_count(self, batch):
         if batch.loop_count is not None:
             return batch.loop_count
@@ -1263,9 +1298,15 @@ def update_from_output(
                 request.pending_exit_depth = None
                 self.enqueue(request, RequestPhase.RECURRENT)
             elif phase == RequestPhase.RECURRENT:
-                request.loops_done += 1
+                # A zero count is the executor's one-loop default for RECURRENT.
+                request.loops_done += max(1, self._batch_loop_count(batch))
                 if self.exit_mode == "trace":
                     should_exit = self._trace_exit(request)
+                elif result.gate_scores is None:
+                    # Fused fixed-depth loops deliver no hazard scores; the exit
+                    # depth was decided before submission.
+                    maximum = request.sampling_params.max_loops or self.total_ut_steps
+                    should_exit = request.loops_done >= maximum
                 else:
                     score = result.gate_scores[row]
                     if self.exit_mode == "ouro":
@@ -1404,12 +1445,14 @@ def update_after_submission(
                 request.num_speculative_rounds += 1
                 self.enqueue(request, RequestPhase.SPECULATIVE)
             elif phase == RequestPhase.RECURRENT:
+                loop_count = max(1, self._batch_loop_count(batch))
                 consume_signal = (
-                    request.loops_done > 0
+                    loop_count == 1
+                    and request.loops_done > 0
                     and request.sampling_params.exit_threshold < 1
                     and self.needs_exit_signal(request)
                 )
-                request.loops_done += 1
+                request.loops_done += loop_count
                 should_exit = (
                     self._trace_exit(request)
                     if self.exit_mode == "trace"

vllm_rlt/core/scheduling_policy.py

@@ -23,17 +23,25 @@ class SchedulingPolicy:
     def __init__(self, config: SchedulerConfig):
         self.config = config
         self.prefill_batches_since_recurrent = 0
+        # Prefill batches allowed before the next decode step.
+        self.prefill_allowance = config.max_prefill_batches_before_decode
 
     @property
     def decode_due(self) -> bool:
-        return self.prefill_batches_since_recurrent >= self.config.max_prefill_batches_before_decode
+        return self.prefill_batches_since_recurrent >= self.prefill_allowance
 
-    def record_batch(self, phase: RequestPhase) -> None:
+    def record_batch(self, phase: RequestPhase, *, loop_count: int | None = None) -> None:
         """Count nonempty scheduled batches, not attempts or GPU completions."""
         if phase in self.PREFILL_PHASES:
             self.prefill_batches_since_recurrent += 1
         elif phase in (RequestPhase.RECURRENT, RequestPhase.SPECULATIVE):
             self.prefill_batches_since_recurrent = 0
+            # Unfused, each loop is its own decode step followed by a prefill turn.
+            # Fusing L loops into one submission would otherwise cut prompt work to
+            # 1/L of its share and delay admission (TTFT) under load.
+            loops = loop_count if phase != RequestPhase.SPECULATIVE and loop_count else 1
+            share = loops if self.config.prefill_share_per_loop else 1
+            self.prefill_allowance = self.config.max_prefill_batches_before_decode * max(1, share)
 
 
 class PromptOnlyPolicy(SchedulingPolicy):

vllm_rlt/worker/model_runner.py

@@ -366,6 +366,26 @@ def _size(self, count):
         self.last_effective_size, self.last_submitted_size = count, size
         return size
 
+    def _fused_decode(self, batch):
+        """A decode submission that runs several loops (or prelude + loops) at once."""
+        return (
+            batch.loop_count > 1
+            and not batch.has_chunks
+            or (
+                batch.operation == Operation.PRELUDE
+                and batch.loop_count > 0
+                and all(span.input_tokens is None for span in batch.spans)
+            )
+        )
+
+    def _row_size(self, batch):
+        """Submitted row count: padded for single-loop steps, exact for fused ones.
+
+        Fused loops stage one host KV view per depth from the real rows, so the
+        device view of the first loop and the hidden rows must use the same count.
+        """
+        return batch.num_rows if self._fused_decode(batch) else self._size(batch.num_rows)
+
     def _workspace(self, group):
         if not self.requestspaces:
             return None
@@ -926,7 +946,7 @@ def _prepare_cpu_inputs(self, batch: ExecutionBatch):
                 requests,
                 depths,
                 positions,
-                self._size(len(requests)),
+                self._row_size(batch),
                 recurrent=batch.operation == Operation.RECURRENT,
             )
             return replace(prepared, kv=kv, routing=routing)
@@ -980,7 +1000,7 @@ def prepare_on_gpu(self, prepared):
                 hidden = self._gather_prefill(prepared)
                 prepared = replace(prepared, batch_inputs=replace(metadata, hidden=hidden))
         elif batch.operation == Operation.PRELUDE and not batch.loop_count:
-            tokens = self._prepare_tokens(prepared, self._size(batch.num_rows))
+            tokens = self._prepare_tokens(prepared, self._row_size(batch))
             prepared = replace(prepared, token_ids=tokens)
         return self._prepare_forward(prepared)
 
@@ -989,7 +1009,7 @@ def _prepare_forward(self, prepared):
         if batch.operation == Operation.PRELUDE and batch.loop_count:
             return replace(prepared, forwards=self._prepare_loop_forwards(prepared))
         requests = [span.request for span in batch.spans]
-        size = batch.num_rows if batch.has_chunks else self._size(batch.num_rows)
+        size = batch.num_rows if batch.has_chunks else self._row_size(batch)
         inputs = prepared.forwards[0].forward
         metadata = inputs.metadata
         kv_metadata = inputs.prelude_coda_kv_metadata
@@ -1216,11 +1236,10 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu
                 # Keep graph selection and its captured compute unchanged.
                 if (
                     scheduled_batch is not None
-                    and self.execution_config.async_scheduling
                     and batch.operation == Operation.RECURRENT
                     and not any(
                         span.emit_gate
-                        and self.exit_config.mode in ("ouro_delayed", "random_lookahead")
+                        and self.exit_config.mode != "trace"
                         and span.request.sampling_params.exit_threshold < 1
                         and span.depth + batch.loop_count
                         < (