diff --git a/tests/test_async_pipeline.py b/tests/test_async_pipeline.py index 76c5dc7c..03776b87 100644 --- a/tests/test_async_pipeline.py +++ b/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 diff --git a/tests/test_engine.py b/tests/test_engine.py index 1ad00979..1a5f0815 100644 --- a/tests/test_engine.py +++ b/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) diff --git a/tests/test_input_preparation.py b/tests/test_input_preparation.py index 6f580135..189b6d22 100644 --- a/tests/test_input_preparation.py +++ b/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) diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 8872d018..c1cf8d0a 100644 --- a/vllm_rlt/config.py +++ b/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"): diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index f443d742..492d600c 100644 --- a/vllm_rlt/core/scheduler.py +++ b/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" diff --git a/vllm_rlt/core/scheduling_policy.py b/vllm_rlt/core/scheduling_policy.py index 39b405ec..7c0b568b 100644 --- a/vllm_rlt/core/scheduling_policy.py +++ b/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): diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..cb8fa69f 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/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 < (