| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for deferred block freeing under async scheduling. |
| 4 | |
| 5 | With async scheduling, a finished/preempted request's blocks may still be |
| 6 | written by a speculatively over-scheduled in-flight GPU step (mamba/GDN |
| 7 | layers rewrite the whole state block every step). If such a block is |
| 8 | reallocated to a request arriving via PD disaggregation, the NIC/RDMA write |
| 9 | of the received state races with the in-flight stale write. The scheduler |
| 10 | closes the race by deferring the return of blocks to the block pool until |
| 11 | the newest scheduled step's output has been processed. |
| 12 | """ |
| 13 | |
| 14 | import os |
| 15 | import time |
| 16 | from unittest.mock import PropertyMock, patch |
| 17 | |
| 18 | import pytest |
| 19 | |
| 20 | from vllm.config import VllmConfig |
| 21 | from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy |
| 22 | from vllm.v1.core.sched.output import SchedulerOutput |
| 23 | from vllm.v1.outputs import ModelRunnerOutput |
| 24 | from vllm.v1.request import RequestStatus |
| 25 | |
| 26 | from .utils import create_requests, create_scheduler, mock_kv |
| 27 | |
| 28 | pytestmark = pytest.mark.cpu_test |
| 29 | |
| 30 | # Allow overriding the model with a local path for offline environments. |
| 31 | MODEL = os.environ.get("VLLM_TEST_DEFER_FREE_MODEL", "facebook/opt-125m") |
| 32 | STOP_TOKEN_ID = 42 |
| 33 | NUM_PROMPT_TOKENS = 33 # 3 blocks with block_size=16 |
| 34 | |
| 35 | |
| 36 | def _make_model_runner_output( |
| 37 | scheduler_output: SchedulerOutput, |
| 38 | token_id: int = 0, |
| 39 | ) -> ModelRunnerOutput: |
| 40 | req_ids = list(scheduler_output.num_scheduled_tokens.keys()) |
| 41 | return ModelRunnerOutput( |
| 42 | req_ids=req_ids, |
| 43 | req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, |
| 44 | sampled_token_ids=[[token_id] for _ in req_ids], |
| 45 | logprobs=None, |
| 46 | prompt_logprobs_dict={}, |
| 47 | pooler_output=[], |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def _create_deferring_scheduler(scheduling_policy="fcfs"): |
| 52 | """Async scheduler with deferred block freeing forced on. |
| 53 | |
| 54 | The production gate additionally requires a PD KV-consumer connector; |
| 55 | the mechanism itself is independent of it. |
| 56 | """ |
| 57 | scheduler = create_scheduler( |
| 58 | model=MODEL, |
| 59 | async_scheduling=True, |
| 60 | scheduling_policy=scheduling_policy, |
| 61 | ) |
| 62 | scheduler.defer_block_free = True |
| 63 | return scheduler |
| 64 | |
| 65 | |
| 66 | def _setup_request_with_inflight_step(scheduler, max_tokens: int = 5): |
| 67 | """Schedule a request's prefill (step 1) and one speculatively |
| 68 | over-scheduled decode (step 2), mimicking async scheduling depth 1. |
| 69 | |
| 70 | Returns (request, out0, out1). |
| 71 | """ |
| 72 | request = create_requests( |
| 73 | num_requests=1, |
| 74 | num_tokens=NUM_PROMPT_TOKENS, |
| 75 | max_tokens=max_tokens, |
| 76 | stop_token_ids=[STOP_TOKEN_ID], |
| 77 | )[0] |
| 78 | scheduler.add_request(request) |
| 79 | out0 = scheduler.schedule() |
| 80 | assert out0.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS |
| 81 | out1 = scheduler.schedule() |
| 82 | assert out1.num_scheduled_tokens[request.request_id] == 1 |
| 83 | return request, out0, out1 |
| 84 | |
| 85 | |
| 86 | def _fail_one_allocation(scheduler, call_number: int): |
| 87 | real_allocate = scheduler.kv_cache_manager.allocate_slots |
| 88 | calls = 0 |
| 89 | |
| 90 | def allocate(*args, **kwargs): |
| 91 | nonlocal calls |
| 92 | calls += 1 |
| 93 | if calls == call_number: |
| 94 | return None |
| 95 | return real_allocate(*args, **kwargs) |
| 96 | |
| 97 | return patch.object( |
| 98 | scheduler.kv_cache_manager, |
| 99 | "allocate_slots", |
| 100 | side_effect=allocate, |
| 101 | ) |
| 102 | |
| 103 | |
| 104 | def test_gate_enabled_for_async_consumer(): |
| 105 | # Overlapping batches + consumer-side connector enables the gate. Async |
| 106 | # scheduling (which would give >1 concurrent batches) is force-disabled on |
| 107 | # CPU, where this test runs, and PP can't be built without GPUs, so force |
| 108 | # max_concurrent_batches to exercise the enabled path on any platform. |
| 109 | with patch.object( |
| 110 | VllmConfig, |
| 111 | "max_concurrent_batches", |
| 112 | new_callable=PropertyMock, |
| 113 | return_value=2, |
| 114 | ): |
| 115 | scheduler = create_scheduler( |
| 116 | model=MODEL, |
| 117 | async_scheduling=True, |
| 118 | use_kv_connector=mock_kv(matched_tokens=0, is_async=False), |
| 119 | ) |
| 120 | assert scheduler.defer_block_free |
| 121 | |
| 122 | |
| 123 | def test_gate_disabled_without_connector(): |
| 124 | # Async scheduling alone (no PD connector): the gate must stay off |
| 125 | # and freeing must remain immediate. |
| 126 | scheduler = create_scheduler(model=MODEL, async_scheduling=True) |
| 127 | assert not scheduler.defer_block_free |
| 128 | |
| 129 | pool = scheduler.kv_cache_manager.block_pool |
| 130 | num_free_initially = pool.get_num_free_blocks() |
| 131 | |
| 132 | request, out0, out1 = _setup_request_with_inflight_step(scheduler) |
| 133 | assert pool.get_num_free_blocks() < num_free_initially |
| 134 | |
| 135 | # Request stops early while step 2 is in flight: blocks are freed |
| 136 | # immediately because deferral is disabled. |
| 137 | scheduler.update_from_output( |
| 138 | out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) |
| 139 | ) |
| 140 | assert request.is_finished() |
| 141 | assert not scheduler.deferred_frees |
| 142 | assert pool.get_num_free_blocks() == num_free_initially |
| 143 | |
| 144 | |
| 145 | def test_finish_defers_free_until_inflight_step_done(): |
| 146 | scheduler = _create_deferring_scheduler() |
| 147 | pool = scheduler.kv_cache_manager.block_pool |
| 148 | num_free_initially = pool.get_num_free_blocks() |
| 149 | |
| 150 | request, out0, out1 = _setup_request_with_inflight_step(scheduler) |
| 151 | num_free_running = pool.get_num_free_blocks() |
| 152 | assert num_free_running < num_free_initially |
| 153 | |
| 154 | # The request stops early (stop token) while the over-scheduled step 2 |
| 155 | # is still in flight: its blocks must NOT return to the pool yet. |
| 156 | scheduler.update_from_output( |
| 157 | out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) |
| 158 | ) |
| 159 | assert request.is_finished() |
| 160 | assert len(scheduler.deferred_frees) == 1 |
| 161 | assert pool.get_num_free_blocks() == num_free_running |
| 162 | |
| 163 | # Step 2's output is processed: every GPU write of step 2 has |
| 164 | # completed, so the blocks can now be returned to the pool. |
| 165 | scheduler.update_from_output(out1, _make_model_runner_output(out1)) |
| 166 | assert not scheduler.deferred_frees |
| 167 | assert pool.get_num_free_blocks() == num_free_initially |
| 168 | |
| 169 | |
| 170 | def test_finish_frees_immediately_when_no_inflight_step(): |
| 171 | scheduler = _create_deferring_scheduler() |
| 172 | pool = scheduler.kv_cache_manager.block_pool |
| 173 | num_free_initially = pool.get_num_free_blocks() |
| 174 | |
| 175 | request = create_requests( |
| 176 | num_requests=1, |
| 177 | num_tokens=NUM_PROMPT_TOKENS, |
| 178 | max_tokens=5, |
| 179 | stop_token_ids=[STOP_TOKEN_ID], |
| 180 | )[0] |
| 181 | scheduler.add_request(request) |
| 182 | out0 = scheduler.schedule() |
| 183 | |
| 184 | # Synchronous-like flow: out0 is the newest scheduled step and its |
| 185 | # output is being processed, so no other step can still write the |
| 186 | # blocks and the free happens immediately. |
| 187 | scheduler.update_from_output( |
| 188 | out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) |
| 189 | ) |
| 190 | assert request.is_finished() |
| 191 | assert not scheduler.deferred_frees |
| 192 | assert pool.get_num_free_blocks() == num_free_initially |
| 193 | |
| 194 | |
| 195 | def test_abort_defers_free(): |
| 196 | scheduler = _create_deferring_scheduler() |
| 197 | pool = scheduler.kv_cache_manager.block_pool |
| 198 | num_free_initially = pool.get_num_free_blocks() |
| 199 | |
| 200 | request, out0, out1 = _setup_request_with_inflight_step(scheduler) |
| 201 | num_free_running = pool.get_num_free_blocks() |
| 202 | |
| 203 | # External abort arrives while steps 1 and 2 are both in flight. |
| 204 | scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) |
| 205 | assert len(scheduler.deferred_frees) == 1 |
| 206 | assert pool.get_num_free_blocks() == num_free_running |
| 207 | |
| 208 | # Step 1's output: step 2 is still in flight, keep holding the blocks. |
| 209 | scheduler.update_from_output(out0, _make_model_runner_output(out0)) |
| 210 | assert len(scheduler.deferred_frees) == 1 |
| 211 | assert pool.get_num_free_blocks() == num_free_running |
| 212 | |
| 213 | # Step 2's output: now the blocks can be freed. |
| 214 | scheduler.update_from_output(out1, _make_model_runner_output(out1)) |
| 215 | assert not scheduler.deferred_frees |
| 216 | assert pool.get_num_free_blocks() == num_free_initially |
| 217 | |
| 218 | |
| 219 | def test_preempt_defers_free_and_clears_bookkeeping(): |
| 220 | scheduler = _create_deferring_scheduler() |
| 221 | pool = scheduler.kv_cache_manager.block_pool |
| 222 | num_free_initially = pool.get_num_free_blocks() |
| 223 | |
| 224 | request, out0, out1 = _setup_request_with_inflight_step(scheduler) |
| 225 | num_free_running = pool.get_num_free_blocks() |
| 226 | |
| 227 | # Preempt the request while steps are in flight (mirrors the |
| 228 | # preemption path inside schedule()). |
| 229 | scheduler.running.remove(request) |
| 230 | scheduler._preempt_request(request, time.monotonic()) |
| 231 | assert request.status == RequestStatus.PREEMPTED |
| 232 | |
| 233 | # Blocks are withheld from the pool, but the manager bookkeeping is |
| 234 | # cleared immediately so the request can be rescheduled safely. |
| 235 | assert len(scheduler.deferred_frees) == 1 |
| 236 | assert pool.get_num_free_blocks() == num_free_running |
| 237 | for manager in scheduler.kv_cache_manager.coordinator.single_type_managers: |
| 238 | assert request.request_id not in manager.req_to_blocks |
| 239 | |
| 240 | # Outputs of both in-flight steps are processed: blocks return to the |
| 241 | # pool only after the newest one. |
| 242 | scheduler.update_from_output(out0, _make_model_runner_output(out0)) |
| 243 | assert len(scheduler.deferred_frees) == 1 |
| 244 | scheduler.update_from_output(out1, _make_model_runner_output(out1)) |
| 245 | assert not scheduler.deferred_frees |
| 246 | assert pool.get_num_free_blocks() == num_free_initially |
| 247 | |
| 248 | |
| 249 | @pytest.mark.parametrize( |
| 250 | ("policy", "failed_call"), |
| 251 | [ |
| 252 | ("fcfs", 1), |
| 253 | ("priority", 2), |
| 254 | ], |
| 255 | ) |
| 256 | def test_allocation_retry_waits_for_fence_then_succeeds(policy, failed_call): |
| 257 | scheduler = _create_deferring_scheduler(policy) |
| 258 | requests = create_requests( |
| 259 | num_requests=3, |
| 260 | num_tokens=NUM_PROMPT_TOKENS, |
| 261 | max_tokens=10, |
| 262 | stop_token_ids=[STOP_TOKEN_ID], |
| 263 | ) |
| 264 | for request in requests: |
| 265 | scheduler.add_request(request) |
| 266 | out0 = scheduler.schedule() |
| 267 | out1 = scheduler.schedule() |
| 268 | |
| 269 | worst, trigger, tail = requests |
| 270 | worst.priority, trigger.priority, tail.priority = 9, 0, 1 |
| 271 | |
| 272 | with _fail_one_allocation(scheduler, failed_call) as allocate: |
| 273 | blocked = scheduler.schedule() |
| 274 | assert allocate.call_count == failed_call |
| 275 | assert not blocked.preempted_req_ids |
| 276 | retry_trigger = requests[failed_call - 1] |
| 277 | assert retry_trigger.request_id not in blocked.num_scheduled_tokens |
| 278 | |
| 279 | for output in (out0, out1, blocked): |
| 280 | if output.total_num_scheduled_tokens: |
| 281 | scheduler.update_from_output( |
| 282 | output, |
| 283 | _make_model_runner_output(output), |
| 284 | ) |
| 285 | |
| 286 | expected_victim = worst if policy == "priority" else tail |
| 287 | with _fail_one_allocation(scheduler, failed_call) as allocate: |
| 288 | resumed = scheduler.schedule() |
| 289 | |
| 290 | assert allocate.call_count > failed_call |
| 291 | assert expected_victim.request_id in resumed.preempted_req_ids |
| 292 | assert retry_trigger.request_id in resumed.num_scheduled_tokens |
| 293 | if policy == "priority": |
| 294 | assert tail.request_id in resumed.num_scheduled_tokens |
| 295 | |
| 296 | |
| 297 | def test_multiple_deferred_frees_drain_in_order(): |
| 298 | scheduler = _create_deferring_scheduler() |
| 299 | pool = scheduler.kv_cache_manager.block_pool |
| 300 | num_free_initially = pool.get_num_free_blocks() |
| 301 | |
| 302 | requests = create_requests( |
| 303 | num_requests=2, |
| 304 | num_tokens=NUM_PROMPT_TOKENS, |
| 305 | max_tokens=5, |
| 306 | stop_token_ids=[STOP_TOKEN_ID], |
| 307 | ) |
| 308 | for request in requests: |
| 309 | scheduler.add_request(request) |
| 310 | out0 = scheduler.schedule() |
| 311 | out1 = scheduler.schedule() |
| 312 | |
| 313 | # Both requests stop early at step 1's output while step 2 is in |
| 314 | # flight: two deferred entries with the same fence. |
| 315 | scheduler.update_from_output( |
| 316 | out0, _make_model_runner_output(out0, token_id=STOP_TOKEN_ID) |
| 317 | ) |
| 318 | assert len(scheduler.deferred_frees) == 2 |
| 319 | assert pool.get_num_free_blocks() < num_free_initially |
| 320 | |
| 321 | scheduler.update_from_output(out1, _make_model_runner_output(out1)) |
| 322 | assert not scheduler.deferred_frees |
| 323 | assert pool.get_num_free_blocks() == num_free_initially |
| 324 | |
| 325 | |
| 326 | def test_fence_held_across_multiple_inflight_steps(): |
| 327 | """Pipeline-parallel / deep async: with several steps scheduled ahead, |
| 328 | a freed request's blocks must stay held until the *newest* in-flight |
| 329 | step's output is processed, not the first. |
| 330 | |
| 331 | Depth-1 tests only check a single intervening update; with PP the |
| 332 | scheduler can dispatch up to pp_size steps ahead, so the fence must |
| 333 | survive multiple intervening update_from_output calls. |
| 334 | """ |
| 335 | scheduler = _create_deferring_scheduler() |
| 336 | pool = scheduler.kv_cache_manager.block_pool |
| 337 | num_free_initially = pool.get_num_free_blocks() |
| 338 | |
| 339 | request = create_requests( |
| 340 | num_requests=1, |
| 341 | num_tokens=NUM_PROMPT_TOKENS, |
| 342 | max_tokens=10, |
| 343 | )[0] |
| 344 | scheduler.add_request(request) |
| 345 | |
| 346 | # Schedule three steps ahead without processing any output: a prefill |
| 347 | # plus two speculatively over-scheduled decodes, all in flight at once. |
| 348 | outs = [scheduler.schedule() for _ in range(3)] |
| 349 | assert outs[0].num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS |
| 350 | assert outs[1].num_scheduled_tokens[request.request_id] == 1 |
| 351 | assert outs[2].num_scheduled_tokens[request.request_id] == 1 |
| 352 | assert scheduler.sched_step_seq == 3 |
| 353 | num_free_running = pool.get_num_free_blocks() |
| 354 | assert num_free_running < num_free_initially |
| 355 | |
| 356 | # Abort while all three steps are in flight: the fence is the newest |
| 357 | # scheduled step (3), since any of them may still write the blocks. |
| 358 | scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) |
| 359 | assert len(scheduler.deferred_frees) == 1 |
| 360 | assert scheduler.deferred_frees[0][0] == 3 |
| 361 | assert pool.get_num_free_blocks() == num_free_running |
| 362 | |
| 363 | # Draining the two earlier in-flight steps must NOT release the blocks: |
| 364 | # their outputs don't fence the still-pending newest write. |
| 365 | for out in (outs[0], outs[1]): |
| 366 | scheduler.update_from_output(out, _make_model_runner_output(out)) |
| 367 | assert len(scheduler.deferred_frees) == 1 |
| 368 | assert pool.get_num_free_blocks() == num_free_running |
| 369 | |
| 370 | # Only once the newest scheduled step's output is processed do the |
| 371 | # blocks return to the pool. |
| 372 | scheduler.update_from_output(outs[2], _make_model_runner_output(outs[2])) |
| 373 | assert not scheduler.deferred_frees |
| 374 | assert pool.get_num_free_blocks() == num_free_initially |
| 375 | |
| 376 | |
| 377 | def test_max_tokens_finish_frees_immediately_with_other_inflight(): |
| 378 | """A request finishing by reaching max_tokens is never over-scheduled past |
| 379 | its final-token step, so no in-flight step writes its blocks: it is freed |
| 380 | immediately even while another request's step is still in flight. |
| 381 | """ |
| 382 | scheduler = _create_deferring_scheduler() |
| 383 | pool = scheduler.kv_cache_manager.block_pool |
| 384 | |
| 385 | # Short request finishes at max_tokens=1; long request keeps running. |
| 386 | short = create_requests( |
| 387 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=1, req_ids=["short"] |
| 388 | )[0] |
| 389 | long = create_requests( |
| 390 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=100, req_ids=["long"] |
| 391 | )[0] |
| 392 | scheduler.add_request(short) |
| 393 | scheduler.add_request(long) |
| 394 | |
| 395 | out0 = scheduler.schedule() # prefill both |
| 396 | out1 = scheduler.schedule() # short is skipped (at max_tokens); long decodes |
| 397 | assert "short" not in out1.num_scheduled_tokens |
| 398 | assert "long" in out1.num_scheduled_tokens |
| 399 | |
| 400 | free_before = pool.get_num_free_blocks() |
| 401 | # Process step 0: `short` reaches max_tokens and finishes while step 1 |
| 402 | # (which scheduled `long`, not `short`) is still in flight. |
| 403 | scheduler.update_from_output(out0, _make_model_runner_output(out0)) |
| 404 | |
| 405 | assert short.is_finished() |
| 406 | # A step IS globally in flight (the old global fence would have deferred), |
| 407 | # but the per-request gate frees `short` immediately since nothing writes |
| 408 | # its blocks anymore. |
| 409 | assert scheduler.sched_step_seq > scheduler.processed_step_seq |
| 410 | assert not scheduler.deferred_frees |
| 411 | assert pool.get_num_free_blocks() > free_before # short's blocks returned |
| 412 | |
| 413 | |
| 414 | def test_abort_mid_prefill_defers_free(): |
| 415 | """Intermediate prefill chunks don't allocate output placeholders, so the |
| 416 | deferral must key off is_prefill_chunk: aborting a request whose prefill |
| 417 | chunk is still in flight must withhold its blocks. |
| 418 | """ |
| 419 | scheduler = create_scheduler( |
| 420 | model=MODEL, async_scheduling=True, max_num_batched_tokens=16 |
| 421 | ) |
| 422 | scheduler.defer_block_free = True |
| 423 | pool = scheduler.kv_cache_manager.block_pool |
| 424 | num_free_initially = pool.get_num_free_blocks() |
| 425 | |
| 426 | request = create_requests( |
| 427 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5 |
| 428 | )[0] |
| 429 | scheduler.add_request(request) |
| 430 | |
| 431 | out0 = scheduler.schedule() |
| 432 | # Partial prefill: a chunk is in flight, with no output placeholders yet. |
| 433 | assert out0.num_scheduled_tokens[request.request_id] == 16 |
| 434 | assert request.num_output_placeholders == 0 |
| 435 | assert request.is_prefill_chunk |
| 436 | num_free_running = pool.get_num_free_blocks() |
| 437 | assert num_free_running < num_free_initially |
| 438 | |
| 439 | # Abort while the prefill chunk is in flight: blocks must be withheld |
| 440 | # (keyed off is_prefill_chunk, since there are no placeholders). |
| 441 | scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) |
| 442 | assert len(scheduler.deferred_frees) == 1 |
| 443 | assert pool.get_num_free_blocks() == num_free_running |
| 444 | |
| 445 | # Once the in-flight prefill step's output is processed, blocks return. |
| 446 | scheduler.update_from_output(out0, _make_model_runner_output(out0)) |
| 447 | assert not scheduler.deferred_frees |
| 448 | assert pool.get_num_free_blocks() == num_free_initially |
| 449 | |
| 450 | |
| 451 | def test_non_async_abort_defers_via_last_sched_seq(): |
| 452 | """Without async (e.g. PP filling the pipeline) there are no placeholders |
| 453 | and a full prefill isn't a partial chunk, yet an abort with a step in flight |
| 454 | must defer. Only the last-scheduled-step fence catches this. |
| 455 | |
| 456 | PP=2 can't be built on a single-GPU host, so force the flag and exercise the |
| 457 | mechanism; the gate itself is covered by test_gate_enabled_for_async_consumer. |
| 458 | """ |
| 459 | scheduler = create_scheduler(model=MODEL, async_scheduling=False) |
| 460 | scheduler.defer_block_free = True |
| 461 | pool = scheduler.kv_cache_manager.block_pool |
| 462 | num_free_initially = pool.get_num_free_blocks() |
| 463 | |
| 464 | request = create_requests( |
| 465 | num_requests=1, num_tokens=NUM_PROMPT_TOKENS, max_tokens=5 |
| 466 | )[0] |
| 467 | scheduler.add_request(request) |
| 468 | |
| 469 | out0 = scheduler.schedule() |
| 470 | # Neither async-only signal marks this request as in flight. |
| 471 | assert request.num_output_placeholders == 0 |
| 472 | assert not request.is_prefill_chunk |
| 473 | # Only the last-scheduled-step fence does. |
| 474 | assert request.last_sched_seq > scheduler.processed_step_seq |
| 475 | num_free_running = pool.get_num_free_blocks() |
| 476 | assert num_free_running < num_free_initially |
| 477 | |
| 478 | # Abort while out0 is in flight: blocks must be withheld. |
| 479 | scheduler.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) |
| 480 | assert len(scheduler.deferred_frees) == 1 |
| 481 | assert pool.get_num_free_blocks() == num_free_running |
| 482 | |
| 483 | scheduler.update_from_output(out0, _make_model_runner_output(out0)) |
| 484 | assert not scheduler.deferred_frees |
| 485 | assert pool.get_num_free_blocks() == num_free_initially |
| 486 | |
| 487 | |
| 488 | def test_cow_retentions_deferred_until_copy_step_processed(): |
| 489 | """The endpoints of a queued KV block copy must stay out of the free |
| 490 | pool until the step that runs the copy has been processed. Freed |
| 491 | earlier, an endpoint can be reallocated (e.g. as a PD KV-load |
| 492 | destination) and overwritten by a transfer that is not ordered against |
| 493 | the copy still pending in the in-flight step. |
| 494 | """ |
| 495 | scheduler = _create_deferring_scheduler() |
| 496 | pool = scheduler.kv_cache_manager.block_pool |
| 497 | manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] |
| 498 | |
| 499 | request = create_requests( |
| 500 | num_requests=1, |
| 501 | num_tokens=NUM_PROMPT_TOKENS, |
| 502 | max_tokens=5, |
| 503 | stop_token_ids=[STOP_TOKEN_ID], |
| 504 | )[0] |
| 505 | scheduler.add_request(request) |
| 506 | |
| 507 | # Simulate a partial-hit CoW performed while scheduling step 1, whose |
| 508 | # hitting request was freed within the same step: the copy rides out0 |
| 509 | # and each endpoint stays alive only through its copy retention. |
| 510 | src_block, dst_block = pool.get_new_blocks(2) |
| 511 | block_copy = KVCacheBlockCopy( |
| 512 | src_block_id=src_block.block_id, dst_block_id=dst_block.block_id |
| 513 | ) |
| 514 | manager._pending_cow_copies.append((src_block, dst_block)) |
| 515 | out0 = scheduler.schedule() |
| 516 | assert out0.kv_cache_block_copies == [block_copy] |
| 517 | |
| 518 | # Exhaust the rest of the pool so the copy endpoints are the only blocks |
| 519 | # a new request could receive, then add one that fits exactly in them. |
| 520 | pool.get_new_blocks(pool.get_num_free_blocks()) |
| 521 | late_request = create_requests( |
| 522 | num_requests=1, |
| 523 | num_tokens=2 * scheduler.block_size, |
| 524 | max_tokens=5, |
| 525 | req_ids=["late"], |
| 526 | )[0] |
| 527 | scheduler.add_request(late_request) |
| 528 | |
| 529 | # Step 2 is scheduled while step 1 (which runs the copy) is still in |
| 530 | # flight: the retentions are released against the copy's fence, so the |
| 531 | # endpoints must not reach the free pool -- the late request must not be |
| 532 | # scheduled onto them. |
| 533 | out1 = scheduler.schedule() |
| 534 | assert src_block.ref_cnt == 1 |
| 535 | assert dst_block.ref_cnt == 1 |
| 536 | assert scheduler.deferred_frees |
| 537 | assert not out1.scheduled_new_reqs |
| 538 | |
| 539 | # Step 1's output is processed: the copy has run, endpoints return to |
| 540 | # the pool and the late request can be scheduled onto them safely. |
| 541 | scheduler.update_from_output(out0, _make_model_runner_output(out0)) |
| 542 | assert src_block.ref_cnt == 0 |
| 543 | assert dst_block.ref_cnt == 0 |
| 544 | assert not scheduler.deferred_frees |
| 545 | out2 = scheduler.schedule() |
| 546 | assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"] |