| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections import defaultdict, deque |
| 4 | from collections.abc import Callable |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.v1.core.sched.async_scheduler import AsyncScheduler |
| 10 | from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput |
| 11 | from vllm.v1.outputs import ModelRunnerOutput |
| 12 | from vllm.v1.request import RequestStatus |
| 13 | from vllm.v1.structured_output import StructuredOutputGrammar |
| 14 | from vllm.v1.utils import ConstantList |
| 15 | |
| 16 | from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv |
| 17 | |
| 18 | pytestmark = pytest.mark.cpu_test |
| 19 | |
| 20 | |
| 21 | def _make_model_runner_output( |
| 22 | scheduler_output: SchedulerOutput, |
| 23 | ) -> ModelRunnerOutput: |
| 24 | req_ids = list(scheduler_output.num_scheduled_tokens.keys()) |
| 25 | return ModelRunnerOutput( |
| 26 | req_ids=req_ids, |
| 27 | req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, |
| 28 | sampled_token_ids=[[i] for i in range(len(req_ids))], |
| 29 | logprobs=None, |
| 30 | prompt_logprobs_dict={}, |
| 31 | pooler_output=[], |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | @pytest.mark.parametrize("max_tokens", [1, 2, 3, 5]) |
| 36 | def test_stop_by_max_tokens(max_tokens: int): |
| 37 | scheduler = create_scheduler(async_scheduling=True) |
| 38 | requests = create_requests(num_requests=2, max_tokens=max_tokens) |
| 39 | req0, req1 = requests |
| 40 | |
| 41 | expected_total_num_scheduled_tokens = 0 |
| 42 | sched_outputs: deque[SchedulerOutput] = deque() |
| 43 | scheduler.add_request(req0) |
| 44 | sched_outputs.append(scheduler.schedule()) |
| 45 | expected_total_num_scheduled_tokens += req0.num_prompt_tokens + max_tokens - 1 |
| 46 | |
| 47 | scheduler.add_request(req1) |
| 48 | sched_outputs.append(scheduler.schedule()) |
| 49 | expected_total_num_scheduled_tokens += req1.num_prompt_tokens + max_tokens - 1 |
| 50 | |
| 51 | total_num_scheduled_tokens = 0 |
| 52 | while sched_outputs: |
| 53 | sched_output = sched_outputs.popleft() |
| 54 | total_num_scheduled_tokens += sched_output.total_num_scheduled_tokens |
| 55 | model_runner_output = _make_model_runner_output(sched_output) |
| 56 | scheduler.update_from_output(sched_output, model_runner_output) |
| 57 | |
| 58 | sched_output = scheduler.schedule() |
| 59 | if sched_output.num_scheduled_tokens: |
| 60 | sched_outputs.append(sched_output) |
| 61 | |
| 62 | assert scheduler.get_num_unfinished_requests() == 0 |
| 63 | assert req0.num_output_tokens == max_tokens |
| 64 | assert req1.num_output_tokens == max_tokens |
| 65 | # Ensure we aren't scheduling more tokens than necessary. |
| 66 | assert total_num_scheduled_tokens == expected_total_num_scheduled_tokens |
| 67 | |
| 68 | |
| 69 | def test_no_spec_decode_padding_up_to_max_model_len(): |
| 70 | """Uniform spec-decode padding must leave room for the sampled token. |
| 71 | |
| 72 | Padding a single-token decode up to exactly max_model_len makes the |
| 73 | running-loop max_model_len cap negative on the next step (async |
| 74 | scheduling schedules it before the output that would stop the request), |
| 75 | and a negative num_scheduled_tokens crashes the model runner. The |
| 76 | request is scheduled un-padded instead. |
| 77 | |
| 78 | Modelled on a disaggregated decode worker, where every request arrives |
| 79 | with all but the last prompt token already computed. |
| 80 | """ |
| 81 | max_model_len = 32 |
| 82 | num_spec = 1 |
| 83 | scheduler = create_scheduler( |
| 84 | async_scheduling=True, |
| 85 | max_model_len=max_model_len, |
| 86 | num_speculative_tokens=num_spec, |
| 87 | speculative_method="ngram_gpu", |
| 88 | ) |
| 89 | |
| 90 | # A running decode, so padding's "something already scheduled and no |
| 91 | # prefill this step" precondition holds for the request below. |
| 92 | (filler,) = create_requests(num_requests=1, num_tokens=4, req_ids=["filler"]) |
| 93 | scheduler.add_request(filler) |
| 94 | sched_output = scheduler.schedule() |
| 95 | scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) |
| 96 | |
| 97 | (request,) = create_requests( |
| 98 | num_requests=1, num_tokens=max_model_len - num_spec, req_ids=["boundary"] |
| 99 | ) |
| 100 | request.num_computed_tokens = request.num_tokens - 1 |
| 101 | scheduler.add_request(request) |
| 102 | |
| 103 | sched_output = scheduler.schedule() |
| 104 | assert sched_output.num_scheduled_tokens["boundary"] == 1 |
| 105 | assert "boundary" not in sched_output.scheduled_spec_decode_tokens |
| 106 | assert request.num_computed_tokens < max_model_len |
| 107 | |
| 108 | scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) |
| 109 | assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED |
| 110 | |
| 111 | |
| 112 | def test_abort(): |
| 113 | scheduler = create_scheduler(async_scheduling=True) |
| 114 | requests = create_requests(num_requests=10, max_tokens=20) |
| 115 | |
| 116 | for req in requests: |
| 117 | scheduler.add_request(req) |
| 118 | |
| 119 | sched_outputs: deque[SchedulerOutput] = deque() |
| 120 | sched_outputs.append(scheduler.schedule()) |
| 121 | sched_outputs.append(scheduler.schedule()) |
| 122 | |
| 123 | abort_order = [0, 8, 3, 1, 6, 4, 2, 5, 7, 9] |
| 124 | abort_order_copy = abort_order.copy() |
| 125 | |
| 126 | def abort_request(): |
| 127 | if not abort_order: |
| 128 | return |
| 129 | req = requests[abort_order.pop(0)] |
| 130 | scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED) |
| 131 | |
| 132 | while sched_outputs: |
| 133 | # Abort a scheduled request. |
| 134 | abort_request() |
| 135 | sched_output = sched_outputs.popleft() |
| 136 | model_runner_output = _make_model_runner_output(sched_output) |
| 137 | scheduler.update_from_output(sched_output, model_runner_output) |
| 138 | |
| 139 | sched_output = scheduler.schedule() |
| 140 | if sched_output.num_scheduled_tokens: |
| 141 | sched_outputs.append(sched_output) |
| 142 | |
| 143 | for i, req in enumerate(requests): |
| 144 | assert req.status == RequestStatus.FINISHED_ABORTED |
| 145 | assert req.num_output_tokens == abort_order_copy.index(i) |
| 146 | |
| 147 | |
| 148 | def test_connector_metadata_precedes_async_placeholder_advance(monkeypatch): |
| 149 | """Mirroring must see earlier unresolved outputs, excluding the current step.""" |
| 150 | scheduler = create_scheduler( |
| 151 | async_scheduling=True, |
| 152 | use_kv_connector=mock_kv(matched_tokens=0, is_async=False), |
| 153 | ) |
| 154 | (request,) = create_requests(num_requests=1, num_tokens=4, max_tokens=4) |
| 155 | scheduler.add_request(request) |
| 156 | |
| 157 | assert scheduler.connector is not None |
| 158 | build_connector_meta = scheduler.connector.build_connector_meta |
| 159 | observed_placeholders = [] |
| 160 | |
| 161 | def record_placeholders(scheduler_output): |
| 162 | observed_placeholders.append(request.num_output_placeholders) |
| 163 | return build_connector_meta(scheduler_output) |
| 164 | |
| 165 | monkeypatch.setattr( |
| 166 | scheduler.connector, "build_connector_meta", record_placeholders |
| 167 | ) |
| 168 | scheduler.schedule() |
| 169 | scheduler.schedule() |
| 170 | |
| 171 | assert observed_placeholders == [0, 1] |
| 172 | |
| 173 | |
| 174 | def test_prefix_caching_for_prefill_dedup(): |
| 175 | CHUNK_SIZE = 1000 |
| 176 | BLOCK_SIZE = 16 |
| 177 | num_prompt_tokens = 100 |
| 178 | scheduler = create_scheduler( |
| 179 | async_scheduling=True, |
| 180 | max_num_batched_tokens=CHUNK_SIZE, |
| 181 | enable_prefix_caching=True, |
| 182 | block_size=BLOCK_SIZE, |
| 183 | ) |
| 184 | requests = create_requests( |
| 185 | num_requests=5, |
| 186 | num_tokens=num_prompt_tokens, |
| 187 | max_tokens=3, |
| 188 | same_prompt=True, |
| 189 | block_size=BLOCK_SIZE, |
| 190 | ) |
| 191 | |
| 192 | # Two requests with the same prompt. |
| 193 | req0 = requests.pop(0) |
| 194 | req1 = requests.pop(0) |
| 195 | scheduler.add_request(req0) |
| 196 | scheduler.add_request(req1) |
| 197 | |
| 198 | sched_outputs: deque[SchedulerOutput] = deque() |
| 199 | sched_output = scheduler.schedule() |
| 200 | sched_outputs.append(sched_output) |
| 201 | # Make sure prefix caching de-duplicates the prompts in the same step, |
| 202 | # so all the blocks except the last are shared between the two requests. |
| 203 | assert len(sched_output.num_scheduled_tokens) == 2 |
| 204 | assert sched_output.num_scheduled_tokens[req0.request_id] == num_prompt_tokens |
| 205 | assert ( |
| 206 | sched_output.num_scheduled_tokens[req1.request_id] |
| 207 | == num_prompt_tokens % BLOCK_SIZE |
| 208 | ) |
| 209 | |
| 210 | sched_outputs.append(scheduler.schedule()) |
| 211 | while sched_outputs: |
| 212 | added_req = None |
| 213 | if requests: |
| 214 | added_req = requests.pop(0) |
| 215 | scheduler.add_request(added_req) |
| 216 | sched_output = sched_outputs.popleft() |
| 217 | model_runner_output = _make_model_runner_output(sched_output) |
| 218 | scheduler.update_from_output(sched_output, model_runner_output) |
| 219 | sched_output = scheduler.schedule() |
| 220 | if sched_output.num_scheduled_tokens: |
| 221 | sched_outputs.append(sched_output) |
| 222 | if added_req: |
| 223 | assert ( |
| 224 | sched_output.num_scheduled_tokens[added_req.request_id] |
| 225 | == num_prompt_tokens % BLOCK_SIZE |
| 226 | ) |
| 227 | |
| 228 | assert scheduler.get_num_unfinished_requests() == 0 |
| 229 | |
| 230 | |
| 231 | def test_prefix_caching_for_multi_turn(): |
| 232 | CHUNK_SIZE = 1000 |
| 233 | BLOCK_SIZE = 16 |
| 234 | num_prompt_tokens = 100 |
| 235 | num_output_tokens = 200 |
| 236 | scheduler = create_scheduler( |
| 237 | async_scheduling=True, |
| 238 | max_num_batched_tokens=CHUNK_SIZE, |
| 239 | enable_prefix_caching=True, |
| 240 | block_size=BLOCK_SIZE, |
| 241 | ) |
| 242 | requests = create_requests( |
| 243 | num_requests=5, |
| 244 | num_tokens=num_prompt_tokens, |
| 245 | max_tokens=num_output_tokens, |
| 246 | block_size=BLOCK_SIZE, |
| 247 | ) |
| 248 | |
| 249 | for req in requests: |
| 250 | scheduler.add_request(req) |
| 251 | sched_outputs: deque[SchedulerOutput] = deque() |
| 252 | sched_outputs.append(scheduler.schedule()) |
| 253 | sched_outputs.append(scheduler.schedule()) |
| 254 | |
| 255 | # Process the requests. |
| 256 | while sched_outputs: |
| 257 | sched_output = sched_outputs.popleft() |
| 258 | model_runner_output = _make_model_runner_output(sched_output) |
| 259 | scheduler.update_from_output(sched_output, model_runner_output) |
| 260 | sched_output = scheduler.schedule() |
| 261 | if sched_output.num_scheduled_tokens: |
| 262 | sched_outputs.append(sched_output) |
| 263 | assert scheduler.get_num_unfinished_requests() == 0 |
| 264 | |
| 265 | # Create next-turn requests whose prompts are the full output of the |
| 266 | # previous turn. |
| 267 | next_turn_requests = create_requests( |
| 268 | num_requests=5, |
| 269 | num_tokens=num_prompt_tokens + num_output_tokens, |
| 270 | max_tokens=num_output_tokens, |
| 271 | block_size=BLOCK_SIZE, |
| 272 | ) |
| 273 | for i, req in enumerate(next_turn_requests): |
| 274 | req.prompt_token_ids = requests[i].prompt_token_ids + list( |
| 275 | requests[i].output_token_ids |
| 276 | ) |
| 277 | req._all_token_ids = req.prompt_token_ids.copy() |
| 278 | req.all_token_ids = ConstantList(req._all_token_ids) |
| 279 | req.block_hashes = [] |
| 280 | req.update_block_hashes() |
| 281 | |
| 282 | # Schedule the next-turn requests. |
| 283 | for req in next_turn_requests: |
| 284 | scheduler.add_request(req) |
| 285 | sched_output = scheduler.schedule() |
| 286 | sched_outputs.append(sched_output) |
| 287 | |
| 288 | # Make sure the next-turn requests get prefix cache hit by the previous |
| 289 | # requests. |
| 290 | for req in next_turn_requests: |
| 291 | assert sched_output.num_scheduled_tokens[req.request_id] == ( |
| 292 | req.num_prompt_tokens % BLOCK_SIZE |
| 293 | ) |
| 294 | |
| 295 | |
| 296 | def test_abort_request_when_structured_output_fsm_cannot_advance(): |
| 297 | scheduler = object.__new__(AsyncScheduler) |
| 298 | request = create_requests(num_requests=1, num_tokens=1)[0] |
| 299 | request.structured_output_request = Mock() |
| 300 | request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) |
| 301 | request.status = RequestStatus.RUNNING |
| 302 | request.num_computed_tokens = request.num_tokens |
| 303 | request.num_output_placeholders = 1 |
| 304 | |
| 305 | scheduler.perf_metrics = None |
| 306 | scheduler.connector = None |
| 307 | scheduler.ec_connector = None |
| 308 | scheduler.structured_output_manager = Mock() |
| 309 | scheduler.structured_output_manager.accept_tokens.return_value = False |
| 310 | scheduler.requests = {request.request_id: request} |
| 311 | scheduler.running = [request] |
| 312 | scheduler.waiting = Mock() |
| 313 | scheduler.kv_cache_manager = Mock() |
| 314 | scheduler.kv_cache_manager.take_events.return_value = None |
| 315 | scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0 |
| 316 | scheduler.kv_event_publisher = Mock() |
| 317 | scheduler.finished_req_ids = set() |
| 318 | scheduler.finished_req_ids_dict = None |
| 319 | scheduler.grammar_compile_error_reqs = set() |
| 320 | scheduler.encoder_cache_mismatch_reqs = set() |
| 321 | scheduler.vllm_config = Mock() |
| 322 | scheduler.aux_output_connector = None |
| 323 | scheduler.return_sampling_mask = False |
| 324 | scheduler.recompute_kv_load_failures = False |
| 325 | scheduler.defer_block_free = False |
| 326 | scheduler.make_stats = Mock(return_value=None) |
| 327 | scheduler.max_model_len = 128 |
| 328 | |
| 329 | def free_request(req, delay_free_blocks=False): |
| 330 | scheduler.finished_req_ids.add(req.request_id) |
| 331 | scheduler.requests.pop(req.request_id, None) |
| 332 | return None, None |
| 333 | |
| 334 | scheduler._free_request = Mock(side_effect=free_request) |
| 335 | |
| 336 | output = SchedulerOutput( |
| 337 | scheduled_new_reqs=[], |
| 338 | scheduled_cached_reqs=CachedRequestData.make_empty(), |
| 339 | num_scheduled_tokens={request.request_id: 1}, |
| 340 | total_num_scheduled_tokens=1, |
| 341 | scheduled_encoder_inputs={}, |
| 342 | scheduled_spec_decode_tokens={}, |
| 343 | num_common_prefix_blocks=[], |
| 344 | finished_req_ids=set(), |
| 345 | free_encoder_mm_hashes=[], |
| 346 | ) |
| 347 | model_runner_output = ModelRunnerOutput( |
| 348 | req_ids=[request.request_id], |
| 349 | req_id_to_index={request.request_id: 0}, |
| 350 | sampled_token_ids=[[123]], |
| 351 | logprobs=None, |
| 352 | prompt_logprobs_dict={}, |
| 353 | pooler_output=[], |
| 354 | ) |
| 355 | |
| 356 | scheduler.update_from_output(output, model_runner_output) |
| 357 | |
| 358 | assert request.resumable is False |
| 359 | assert request.status == RequestStatus.FINISHED_ERROR |
| 360 | assert request.request_id not in scheduler.requests |
| 361 | assert not scheduler.running |
| 362 | |
| 363 | |
| 364 | class PipelinedEngine: |
| 365 | """Drive a real AsyncScheduler like EngineCore.step_with_batch_queue: |
| 366 | schedule until the batch queue is full, then process the oldest step's |
| 367 | output. Async PP runs pp_size+1 concurrent batches, so up to pp_size |
| 368 | steps are in flight at each schedule() call -- the window in which |
| 369 | preemption must handle output that has not yet returned. (Single-GPU e2e |
| 370 | tests can never create this window: at PP=1, exactly one step is in |
| 371 | flight and it is processed before a preempted request can resume.) |
| 372 | |
| 373 | The model runner is emulated with the V2 runner's own bookkeeping, from |
| 374 | only what the scheduler serializes to it: slots flushed on |
| 375 | preempted_req_ids, resumed requests re-added from the NewRequestData |
| 376 | snapshot, sampling when a step reaches the end of the runner's own view |
| 377 | of the sequence. This makes preemption races observable: a stale token |
| 378 | delivered after a resume is scheduled extends the scheduler's sequence |
| 379 | but not the runner's. |
| 380 | |
| 381 | Every sample emits a globally unique token tagged with its sampled |
| 382 | position, so tests can assert exact delivery. |
| 383 | """ |
| 384 | |
| 385 | def __init__( |
| 386 | self, |
| 387 | scheduler: AsyncScheduler, |
| 388 | queue_size: int, |
| 389 | accept_drafts: Callable[[int, str, int], int] | None = None, |
| 390 | ): |
| 391 | self.scheduler = scheduler |
| 392 | self.queue_size = queue_size |
| 393 | self.accept_drafts = accept_drafts |
| 394 | # In-flight steps: (scheduler_output, new_reqs snapshot) in FIFO order. |
| 395 | self.queue: deque[tuple[SchedulerOutput, list[tuple[str, int, int]]]] = deque() |
| 396 | # Runner-side request state: req_id -> [seq_len, num_computed] as the |
| 397 | # runner sees them (its own sampled tokens, not the scheduler's). |
| 398 | self.runner_view: dict[str, list[int]] = {} |
| 399 | # All tokens the fake runner ever sampled, per request, in order. |
| 400 | self.emitted: dict[str, list[int]] = defaultdict(list) |
| 401 | # Sequence position each (globally unique) token was sampled for. |
| 402 | self.emitted_position: dict[int, int] = {} |
| 403 | self.step_idx = 0 |
| 404 | self._next_token = 1000 |
| 405 | |
| 406 | def _schedule(self) -> bool: |
| 407 | scheduler_output = self.scheduler.schedule() |
| 408 | self.step_idx += 1 |
| 409 | # Snapshot what NewRequestData serializes at schedule time (both new |
| 410 | # and resumed requests for the V2 runner). |
| 411 | new_reqs = [ |
| 412 | (r.req_id, len(r.prefill_token_ids), r.num_computed_tokens) |
| 413 | for r in scheduler_output.scheduled_new_reqs |
| 414 | ] |
| 415 | # Enqueue empty steps too (the engine executes them), so the runner |
| 416 | # still observes their preempted/finished request ids in step order. |
| 417 | self.queue.appendleft((scheduler_output, new_reqs)) |
| 418 | return True |
| 419 | |
| 420 | def _process_oldest_step(self) -> None: |
| 421 | scheduler_output, new_reqs = self.queue.pop() |
| 422 | # Worker-side state updates, in step order: flush preempted/finished |
| 423 | # slots, then (re-)add new/resumed requests. |
| 424 | for req_id in scheduler_output.preempted_req_ids or (): |
| 425 | self.runner_view.pop(req_id, None) |
| 426 | for req_id in scheduler_output.finished_req_ids or (): |
| 427 | self.runner_view.pop(req_id, None) |
| 428 | for req_id, seq_len, num_computed in new_reqs: |
| 429 | self.runner_view[req_id] = [seq_len, num_computed] |
| 430 | |
| 431 | req_ids = list(scheduler_output.num_scheduled_tokens.keys()) |
| 432 | sampled_token_ids: list[list[int]] = [] |
| 433 | for req_id in req_ids: |
| 434 | num_scheduled = scheduler_output.num_scheduled_tokens[req_id] |
| 435 | view = self.runner_view.get(req_id) |
| 436 | if view is None: |
| 437 | # Slot already flushed (request finished/aborted mid-flight). |
| 438 | sampled_token_ids.append([]) |
| 439 | continue |
| 440 | seq_len, num_computed = view |
| 441 | end = num_computed + num_scheduled |
| 442 | if end < seq_len: |
| 443 | # Partial prefill by the runner's own bookkeeping: no sample. |
| 444 | view[1] = end |
| 445 | sampled_token_ids.append([]) |
| 446 | continue |
| 447 | drafts = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()) |
| 448 | num_accepted = ( |
| 449 | min(self.accept_drafts(self.step_idx, req_id, len(drafts)), len(drafts)) |
| 450 | if drafts and self.accept_drafts |
| 451 | else 0 |
| 452 | ) |
| 453 | num_rejected = len(drafts) - num_accepted |
| 454 | tokens = list(range(self._next_token, self._next_token + 1 + num_accepted)) |
| 455 | self._next_token += 1 + num_accepted |
| 456 | self.emitted[req_id].extend(tokens) |
| 457 | sampled_token_ids.append(tokens) |
| 458 | # Rejected drafts roll back computed; the sampled tokens extend |
| 459 | # the runner's sequence. |
| 460 | view[1] = end - num_rejected |
| 461 | view[0] = view[1] + 1 |
| 462 | for offset, token in enumerate(tokens): |
| 463 | self.emitted_position[token] = view[0] - len(tokens) + offset |
| 464 | model_runner_output = ModelRunnerOutput( |
| 465 | req_ids=req_ids, |
| 466 | req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, |
| 467 | sampled_token_ids=sampled_token_ids, |
| 468 | logprobs=None, |
| 469 | prompt_logprobs_dict={}, |
| 470 | pooler_output=[], |
| 471 | ) |
| 472 | self.scheduler.update_from_output(scheduler_output, model_runner_output) |
| 473 | |
| 474 | def run( |
| 475 | self, |
| 476 | max_steps: int = 2000, |
| 477 | before_step: Callable[[int, "PipelinedEngine"], None] | None = None, |
| 478 | ) -> None: |
| 479 | for i in range(max_steps): |
| 480 | if not self.scheduler.has_requests() and not self.queue: |
| 481 | return |
| 482 | if before_step is not None: |
| 483 | before_step(i, self) |
| 484 | scheduled = ( |
| 485 | self.scheduler.has_requests() |
| 486 | and len(self.queue) < self.queue_size |
| 487 | and self._schedule() |
| 488 | ) |
| 489 | if scheduled and len(self.queue) < self.queue_size: |
| 490 | # Queue not yet full: the engine returns without blocking. |
| 491 | continue |
| 492 | if self.queue: |
| 493 | self._process_oldest_step() |
| 494 | raise AssertionError("engine loop did not converge") |
| 495 | |
| 496 | |
| 497 | def _create_async_pp_scheduler( |
| 498 | num_spec: int, pp_size: int = 3, num_blocks: int = 5 |
| 499 | ) -> AsyncScheduler: |
| 500 | scheduler = create_scheduler( |
| 501 | async_scheduling=True, |
| 502 | num_speculative_tokens=num_spec or None, |
| 503 | speculative_method="ngram_gpu" if num_spec else None, |
| 504 | use_v2_model_runner=True, |
| 505 | num_blocks=num_blocks, |
| 506 | block_size=16, |
| 507 | max_num_batched_tokens=512, |
| 508 | ) |
| 509 | # Emulate PP at the scheduler level; constructing with |
| 510 | # pipeline_parallel_size>1 requires that many visible GPUs. Drive with |
| 511 | # queue_size=pp_size+1 (V2 async PP runs pp_size+1 concurrent batches). |
| 512 | scheduler.pp_size = pp_size |
| 513 | scheduler.use_pp = pp_size > 1 |
| 514 | return scheduler |
| 515 | |
| 516 | |
| 517 | def _assert_ordered_subset(delivered: list[int], emitted: list[int]) -> None: |
| 518 | """Delivered tokens must be an order-preserving subset of the emitted |
| 519 | tokens with no duplicates (tokens are globally unique).""" |
| 520 | it = iter(emitted) |
| 521 | for token in delivered: |
| 522 | assert token in it, f"token {token} delivered out of order or twice" |
| 523 | |
| 524 | |
| 525 | def _assert_positions_consistent(req, engine: PipelinedEngine) -> None: |
| 526 | """The i-th delivered output token must be one the runner sampled for |
| 527 | exactly sequence position prompt_len + i: catches a preempted request's |
| 528 | stale output landing on a position the resumed request resampled (or |
| 529 | vice versa), which token-stream equality alone cannot see.""" |
| 530 | for i, token in enumerate(req.output_token_ids): |
| 531 | expected = req.num_prompt_tokens + i |
| 532 | actual = engine.emitted_position[token] |
| 533 | assert actual == expected, ( |
| 534 | f"output {i} of {req.request_id}: token sampled for position " |
| 535 | f"{actual}, delivered as position {expected}" |
| 536 | ) |
| 537 | |
| 538 | |
| 539 | @pytest.mark.parametrize("num_spec", [0, 3]) |
| 540 | def test_kv_pressure_preemption_with_inflight_output(num_spec: int): |
| 541 | """KV-pressure preemption of requests with in-flight async output. |
| 542 | |
| 543 | PP=3 + async scheduling (batch queue of 4), a block pool small enough |
| 544 | that decodes contend and preempt mid-flight, and staggered arrivals so |
| 545 | the batch queue actually pipelines. A preempted request's in-flight steps |
| 546 | still return: their tokens must be delivered exactly once, their stale |
| 547 | spec-rejection counts must not corrupt the rolled-back counters, and the |
| 548 | resume must not resample a position that output later delivers. |
| 549 | |
| 550 | Regression for the num_output_placeholders underflow EngineCore crash: |
| 551 | with the fix reverted, the num_spec=3 variant fails with exactly |
| 552 | ``assert request.num_output_placeholders >= 0`` when a stale spec output |
| 553 | returns after the preempted request was resumed and sampled. |
| 554 | """ |
| 555 | max_tokens = 24 |
| 556 | scheduler = _create_async_pp_scheduler(num_spec) |
| 557 | requests = create_requests( |
| 558 | num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True |
| 559 | ) |
| 560 | pending = list(requests) |
| 561 | for _ in range(2): |
| 562 | scheduler.add_request(pending.pop(0)) |
| 563 | |
| 564 | # Observe that the scenario under test actually occurs. |
| 565 | preempts_with_inflight_output = 0 |
| 566 | orig_preempt = scheduler._preempt_request |
| 567 | |
| 568 | def counting_preempt(request, timestamp, **kwargs): |
| 569 | nonlocal preempts_with_inflight_output |
| 570 | if request.num_in_flight_tokens > 0: |
| 571 | preempts_with_inflight_output += 1 |
| 572 | return orig_preempt(request, timestamp, **kwargs) |
| 573 | |
| 574 | scheduler._preempt_request = counting_preempt |
| 575 | |
| 576 | def add_requests(step: int, engine: PipelinedEngine): |
| 577 | if pending: |
| 578 | scheduler.add_request(pending.pop(0)) |
| 579 | |
| 580 | engine = PipelinedEngine( |
| 581 | scheduler, |
| 582 | queue_size=4, |
| 583 | # Deterministically vary spec acceptance so stale outputs carry |
| 584 | # nonzero rejection counts. |
| 585 | accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1), |
| 586 | ) |
| 587 | engine.run(before_step=add_requests) |
| 588 | |
| 589 | assert preempts_with_inflight_output > 0, ( |
| 590 | "test did not exercise preemption with in-flight output" |
| 591 | ) |
| 592 | for req in requests: |
| 593 | assert req.is_finished() |
| 594 | assert req.num_output_tokens == max_tokens |
| 595 | # Lossless: delivered tokens are exactly the sampled tokens, in order |
| 596 | # (the excluded tail was emitted after the request finished). |
| 597 | emitted = engine.emitted[req.request_id] |
| 598 | assert list(req.output_token_ids) == emitted[:max_tokens] |
| 599 | _assert_positions_consistent(req, engine) |
| 600 | |
| 601 | |
| 602 | @pytest.mark.parametrize("pp_size", [1, 3]) |
| 603 | def test_reset_prefix_cache_with_inflight_output_under_kv_pressure(pp_size: int): |
| 604 | """reset_prefix_cache(reset_running_requests=True) resumes requests in |
| 605 | the same step it preempts them, so in-flight output must be dropped (the |
| 606 | resume resamples those positions). |
| 607 | |
| 608 | pp_size=1: regression for the frame-based discard this fix replaces, |
| 609 | which with spec decode drained one *token* count per output frame and |
| 610 | over-discarded, corrupting the fresh frames after the resume. |
| 611 | pp_size=3: back-to-back resets, so the second re-preempts requests whose |
| 612 | dropped stale share is still in flight -- it must be recorded once (not |
| 613 | accumulated) and stay dropped. |
| 614 | """ |
| 615 | max_tokens = 24 |
| 616 | scheduler = _create_async_pp_scheduler(num_spec=3, pp_size=pp_size) |
| 617 | requests = create_requests( |
| 618 | num_requests=8, num_tokens=8, max_tokens=max_tokens, ignore_eos=True |
| 619 | ) |
| 620 | pending = list(requests) |
| 621 | for _ in range(2): |
| 622 | scheduler.add_request(pending.pop(0)) |
| 623 | |
| 624 | # Observe re-preemptions with an undrained stale share (the |
| 625 | # double-count hazard). |
| 626 | repreempts_with_stale = 0 |
| 627 | orig_preempt = scheduler._preempt_request |
| 628 | |
| 629 | def counting_preempt(request, timestamp, **kwargs): |
| 630 | nonlocal repreempts_with_stale |
| 631 | if getattr(request, "num_stale_output_tokens", 0) > 0: |
| 632 | repreempts_with_stale += 1 |
| 633 | return orig_preempt(request, timestamp, **kwargs) |
| 634 | |
| 635 | scheduler._preempt_request = counting_preempt |
| 636 | |
| 637 | resets = 0 |
| 638 | reset_steps = {6, 14} if pp_size == 1 else {6, 7, 18, 19} |
| 639 | |
| 640 | def before_step(step: int, engine: PipelinedEngine): |
| 641 | nonlocal resets |
| 642 | if pending: |
| 643 | scheduler.add_request(pending.pop(0)) |
| 644 | if step in reset_steps and (engine.queue or scheduler.running): |
| 645 | scheduler.reset_prefix_cache(reset_running_requests=True) |
| 646 | resets += 1 |
| 647 | |
| 648 | engine = PipelinedEngine( |
| 649 | scheduler, |
| 650 | queue_size=pp_size + 1, |
| 651 | accept_drafts=lambda step, req_id, n: (step + int(req_id)) % (n + 1), |
| 652 | ) |
| 653 | engine.run(before_step=before_step) |
| 654 | |
| 655 | assert resets > 0, "test did not exercise reset_prefix_cache" |
| 656 | if pp_size > 1: |
| 657 | # The re-preempt-while-stale-pending window needs pipeline depth. |
| 658 | assert repreempts_with_stale > 0, ( |
| 659 | "test did not exercise re-preemption with an undrained stale share" |
| 660 | ) |
| 661 | for req in requests: |
| 662 | assert req.is_finished() |
| 663 | assert req.num_output_tokens == max_tokens |
| 664 | # Dropped tokens are never delivered; order must be preserved with |
| 665 | # no duplicates. |
| 666 | _assert_ordered_subset( |
| 667 | list(req.output_token_ids), engine.emitted[req.request_id] |
| 668 | ) |
| 669 | _assert_positions_consistent(req, engine) |
| 670 | # All stale shares fully drained by the end. |
| 671 | assert getattr(req, "num_stale_output_tokens", 0) == 0 |
| 672 | |
| 673 | |
| 674 | def test_requires_kv_delivery_defaults_to_producer_role(): |
| 675 | # No connector: nothing is handed off, so keep the lossless deliver-stale |
| 676 | # path on preemption. |
| 677 | assert create_scheduler(async_scheduling=True).requires_kv_delivery is False |
| 678 | # Only a producer hands KV off when a request completes. |
| 679 | for role, expected in ( |
| 680 | ("kv_producer", True), |
| 681 | ("kv_both", True), |
| 682 | ("kv_consumer", False), |
| 683 | ): |
| 684 | scheduler = create_scheduler( |
| 685 | async_scheduling=True, use_kv_connector=True, kv_role=role |
| 686 | ) |
| 687 | assert scheduler.requires_kv_delivery is expected, role |
| 688 | |
| 689 | |
| 690 | @pytest.mark.parametrize( |
| 691 | ("kv_role", "defer_free"), |
| 692 | [ |
| 693 | ("kv_producer", False), |
| 694 | ("kv_consumer", False), |
| 695 | ("kv_consumer", True), |
| 696 | ], |
| 697 | ) |
| 698 | def test_kv_pressure_preempt_mid_handoff(kv_role: str, defer_free: bool): |
| 699 | """P/D race: KV pressure hits while the output of a request's final |
| 700 | prefill chunk -- the hand-off token that would finish it -- is in flight. |
| 701 | |
| 702 | When the victim's blocks free immediately, the request is preempted. On a |
| 703 | producer, that output must be dropped so the request recomputes; delivering |
| 704 | it would finish the request and hand off blocks the preemption already |
| 705 | freed, so the consumer pulls garbage. A consumer hands nothing off, so it |
| 706 | keeps the lossless deliver-stale path. |
| 707 | |
| 708 | A consumer with overlapping batches (async scheduling or PP) instead fences |
| 709 | the victim's free behind its in-flight output, so the allocation retry |
| 710 | stops instead of preempting; the request then finishes from that output |
| 711 | once it lands. The gate depends on the platform (async scheduling is |
| 712 | force-disabled on CPU), so force the flag to cover both paths everywhere. |
| 713 | """ |
| 714 | is_producer = kv_role == "kv_producer" |
| 715 | scheduler = create_scheduler( |
| 716 | async_scheduling=True, |
| 717 | use_kv_connector=True, |
| 718 | kv_role=kv_role, |
| 719 | num_blocks=5, |
| 720 | block_size=16, |
| 721 | max_num_batched_tokens=512, |
| 722 | ) |
| 723 | assert scheduler.requires_kv_delivery is is_producer |
| 724 | # The production gate requires overlapping batches, which async |
| 725 | # scheduling only provides off-CPU. |
| 726 | scheduler.defer_block_free = defer_free |
| 727 | |
| 728 | # 32-token prompts fill 2 blocks each, exhausting the usable pool, so the |
| 729 | # next decode allocation targets the tail of the running queue (the |
| 730 | # handoff request) while its prefill output is still in flight. |
| 731 | decoder = create_requests( |
| 732 | num_requests=1, num_tokens=32, max_tokens=8, req_ids=["decoder"] |
| 733 | )[0] |
| 734 | handoff = create_requests( |
| 735 | num_requests=1, num_tokens=32, max_tokens=1, req_ids=["handoff"] |
| 736 | )[0] |
| 737 | scheduler.add_request(decoder) |
| 738 | scheduler.add_request(handoff) |
| 739 | sched_output = scheduler.schedule() |
| 740 | assert handoff.status == RequestStatus.RUNNING |
| 741 | assert handoff.num_output_placeholders == 1 |
| 742 | |
| 743 | scheduler.schedule() |
| 744 | if defer_free: |
| 745 | # The victim's blocks are fenced behind its in-flight output, so |
| 746 | # preempting it could not satisfy the allocation; the retry stops |
| 747 | # instead of preempting. |
| 748 | assert handoff.status == RequestStatus.RUNNING |
| 749 | assert handoff.num_stale_output_tokens == 0 |
| 750 | else: |
| 751 | # Blocks free immediately, so the handoff request is preempted. |
| 752 | assert handoff.status == RequestStatus.PREEMPTED |
| 753 | assert handoff.num_stale_output_tokens == handoff.num_prompt_tokens |
| 754 | assert handoff.drop_stale_output is is_producer |
| 755 | |
| 756 | scheduler.update_from_output(sched_output, _make_model_runner_output(sched_output)) |
| 757 | |
| 758 | assert handoff.num_stale_output_tokens == 0 |
| 759 | if is_producer: |
| 760 | # Dropped: recomputed from the waiting queue, so the hand-off happens |
| 761 | # against real KV. |
| 762 | assert not handoff.is_finished() |
| 763 | assert handoff.status == RequestStatus.PREEMPTED |
| 764 | assert handoff.num_output_tokens == 0 |
| 765 | assert handoff.request_id in scheduler.requests |
| 766 | else: |
| 767 | assert handoff.is_finished() |
| 768 | assert handoff.num_output_tokens == 1 |
| 769 | |
| 770 | |
| 771 | def test_resumable_request_handoff(): |
| 772 | """New streaming input resumes with the correct context and stale results |
| 773 | are ignored. |
| 774 | """ |
| 775 | scheduler = create_scheduler( |
| 776 | async_scheduling=True, |
| 777 | use_v2_model_runner=True, |
| 778 | ) |
| 779 | engine = PipelinedEngine(scheduler, queue_size=3) |
| 780 | engine._next_token = EOS_TOKEN_ID |
| 781 | |
| 782 | old_prompt_len = 10 |
| 783 | continuation_len = 12 |
| 784 | request_id = "resumable-race" |
| 785 | |
| 786 | request = create_requests( |
| 787 | 1, num_tokens=old_prompt_len, max_tokens=4, req_ids=[request_id] |
| 788 | )[0] |
| 789 | request.resumable = True |
| 790 | continuation = create_requests( |
| 791 | 1, num_tokens=continuation_len, max_tokens=4, req_ids=[request_id] |
| 792 | )[0] |
| 793 | continuation.resumable = True |
| 794 | |
| 795 | scheduler.add_request(request) |
| 796 | |
| 797 | for _ in range(3): |
| 798 | assert engine._schedule() |
| 799 | |
| 800 | engine._process_oldest_step() |
| 801 | scheduler.add_request(continuation) |
| 802 | assert engine._schedule() |
| 803 | new_step, _ = engine.queue[0] |
| 804 | assert new_step.scheduled_new_reqs[0].num_computed_tokens == old_prompt_len |
| 805 | |
| 806 | all_token_ids_before_stale = list(request.all_token_ids) |
| 807 | for _ in range(2): |
| 808 | engine._process_oldest_step() |
| 809 | assert list(request.output_token_ids) == [] |
| 810 | assert list(request.all_token_ids) == all_token_ids_before_stale |
| 811 | |
| 812 | engine._process_oldest_step() |
| 813 | assert len(request.output_token_ids) == 1 |