@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9ce74042d3

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 19:00 UTC

[Bugfix][SimpleCPUOffloadBackend] Dedup in-flight CPU offload stores across scheduler steps (#41289)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Yifan Qiao committed May 13, 2026 at 01:53 UTC 9ce74042d368d6f985741d048ad53d6d66cace6d

3 files changed +126 -21

.buildkite/test_areas/misc.yaml

@@ -79,6 +79,7 @@ steps:

79 - tests/v1/core

80 - tests/v1/executor

81 - tests/v1/kv_offload

82 + - tests/v1/simple_kv_offload

83 - tests/v1/worker

84 - tests/v1/kv_connector/unit

85 - tests/v1/metrics

@@ -90,6 +91,7 @@ steps:

91 - pytest -v -s -m 'not cpu_test' v1/core

92 - pytest -v -s v1/executor

93 - pytest -v -s v1/kv_offload

94 + - pytest -v -s v1/simple_kv_offload

95 - pytest -v -s v1/worker

96 - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit

97 - pytest -v -s -m 'not cpu_test' v1/metrics

tests/v1/simple_kv_offload/test_scheduler.py

+116 -16

@@ -179,6 +179,7 @@ _req_counter = 0

179 def make_request(

180 num_blocks: int = 2,

181 request_id: str | None = None,

182 + extra_tokens: int = 1,

183 ) -> Request:

184 """Create a Request with deterministic block hashes."""

185 global _req_counter

@@ -186,13 +187,7 @@ def make_request(

187 if request_id is None:

188 request_id = f"req-{_req_counter}"

189

189 - # Add one extra token beyond the last full block so that

190 - # ``max_cache_hit_length = num_tokens - 1`` (see

191 - # KVCacheManager.get_computed_blocks) does not truncate the final

192 - # full block: ``find_longest_cache_hit`` uses

193 - # ``max_length // block_size`` and would otherwise drop one block

194 - # when the prompt is an exact multiple of block_size.

195 - num_tokens = num_blocks * BLOCK_SIZE + 1

190 + num_tokens = num_blocks * BLOCK_SIZE + extra_tokens

191 start = _req_counter * 10000

192 prompt_token_ids = list(range(start, start + num_tokens))

193 sampling_params = SamplingParams(max_tokens=1)

@@ -386,6 +381,8 @@ def test_eager_store_and_load_roundtrip() -> None:

381 block_hasher=req._block_hasher,

382 )

383 hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

384 + # make_request pads num_tokens by +1 beyond the last full block, so the

385 + # manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.

386 assert hit_tokens == num_blocks * BLOCK_SIZE

387 assert is_async is True

388

@@ -405,7 +402,44 @@ def test_eager_store_and_load_roundtrip() -> None:

402

403

404 # ---------------------------------------------------------------------------

408 -# Test 1b: Lazy store-and-load roundtrip

405 +# Test 1b: Boundary — max_hit_len cap drops the last full block when the

406 +# prompt is an exact multiple of BLOCK_SIZE.

407 +# ---------------------------------------------------------------------------

408 +def test_max_hit_len_cap_drops_last_full_block() -> None:

409 + """When num_tokens is an exact multiple of BLOCK_SIZE, the manager's

410 + ``max_hit_len = num_tokens - 1`` cap forces ``find_longest_cache_hit`` to

411 + drop the final block (since ``max_length // block_size`` rounds down).

412 + """

413 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)

414 + sched = fix.scheduler

415 +

416 + num_blocks = 2

417 + req = make_request(num_blocks=num_blocks, extra_tokens=0)

418 + assert req.num_tokens == num_blocks * BLOCK_SIZE

419 +

420 + kv_blocks = _alloc_and_register(fix, req, num_blocks)

421 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

422 + sched_out = make_scheduler_output(

423 + {req.request_id: num_blocks * BLOCK_SIZE},

424 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

425 + )

426 + meta = sched.build_connector_meta(sched_out)

427 + simulate_store_completion(sched, meta.store_event)

428 +

429 + req2 = Request(

430 + request_id="req-cap-boundary",

431 + prompt_token_ids=req.prompt_token_ids,

432 + sampling_params=req.sampling_params,

433 + pooling_params=None,

434 + mm_features=None,

435 + block_hasher=req._block_hasher,

436 + )

437 + hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

438 + assert hit_tokens == (num_blocks - 1) * BLOCK_SIZE

439 +

440 +

441 +# ---------------------------------------------------------------------------

442 +# Test 1c: Lazy store-and-load roundtrip

443 # ---------------------------------------------------------------------------

444 def _flush_old_blocks_to_lru_head(

445 gpu_pool: BlockPool,

@@ -469,8 +503,11 @@ def test_lazy_store_and_load_roundtrip() -> None:

503 hit_tokens, is_async = sched.get_num_new_matched_tokens(

504 req_old2, num_computed_tokens=0

505 )

472 - assert hit_tokens == num_blocks * BLOCK_SIZE, (

473 - f"Expected {num_blocks * BLOCK_SIZE} hit tokens, got {hit_tokens}"

506 + # make_request pads num_tokens by +1 beyond the last full block, so the

507 + # manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.

508 + expected_hit = num_blocks * BLOCK_SIZE

509 + assert hit_tokens == expected_hit, (

510 + f"Expected {expected_hit} hit tokens, got {hit_tokens}"

511 )

512 assert is_async is True

513

@@ -537,7 +574,72 @@ def test_eager_duplicate_store_skipped() -> None:

574

575

576 # ---------------------------------------------------------------------------

540 -# Test 2b: Lazy duplicate store is skipped

577 +# Test 2b: Eager dedup of in-flight stores across consecutive steps

578 +# ---------------------------------------------------------------------------

579 +def test_eager_in_flight_store_dedup_across_steps() -> None:

580 + """Eager: a second request sharing a prefix with an in-flight store

581 + must not re-offload the same GPU blocks before completion lands.

582 +

583 + Simulates a GPU prefix-cache hit by reusing the first request's

584 + GPU block IDs in the second scheduler step, which is the path the

585 + real scheduler takes when two requests share a prefix.

586 + """

587 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)

588 + sched = fix.scheduler

589 +

590 + num_blocks = 2

591 + req = make_request(num_blocks=num_blocks)

592 +

593 + kv_blocks = _alloc_and_register(fix, req, num_blocks)

594 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

595 + block_ids = kv_blocks.get_block_ids()

596 + sched_out = make_scheduler_output(

597 + {req.request_id: num_blocks * BLOCK_SIZE},

598 + new_reqs={req.request_id: block_ids},

599 + )

600 +

601 + meta1 = sched.build_connector_meta(sched_out)

602 + assert meta1.store_event >= 0

603 + assert len(meta1.store_cpu_blocks) == num_blocks

604 + # In-flight set tracks the scheduled GPU blocks until completion.

605 + assert sched._in_flight_store_gpu_blocks == set(meta1.store_gpu_blocks)

606 + cpu_free_after_first = get_cpu_free_blocks(sched)

607 +

608 + # Second request shares the prefix and reuses the same GPU block IDs

609 + # (the real scheduler path: GPU prefix cache returns the same blocks).

610 + # Do NOT simulate completion — the first store is still in-flight.

611 + req2 = Request(

612 + request_id="req-dup-eager-inflight",

613 + prompt_token_ids=req.prompt_token_ids,

614 + sampling_params=req.sampling_params,

615 + pooling_params=None,

616 + mm_features=None,

617 + block_hasher=req._block_hasher,

618 + )

619 + req2.num_computed_tokens = num_blocks * BLOCK_SIZE

620 + sched.update_state_after_alloc(req2, kv_blocks, num_external_tokens=0)

621 + sched_out2 = make_scheduler_output(

622 + {req2.request_id: num_blocks * BLOCK_SIZE},

623 + new_reqs={req2.request_id: block_ids},

624 + )

625 +

626 + meta2 = sched.build_connector_meta(sched_out2)

627 + if meta2.store_event >= 0:

628 + assert len(meta2.store_cpu_blocks) == 0, (

629 + "Expected no new CPU blocks for in-flight duplicate hashes"

630 + )

631 + assert get_cpu_free_blocks(sched) == cpu_free_after_first, (

632 + "Second request should not consume CPU blocks while the first "

633 + "store is still in-flight"

634 + )

635 +

636 + # After completion, the in-flight set is cleared.

637 + simulate_store_completion(sched, meta1.store_event)

638 + assert sched._in_flight_store_gpu_blocks == set()

639 +

640 +

641 +# ---------------------------------------------------------------------------

642 +# Test 2c: Lazy duplicate store is skipped

643 # ---------------------------------------------------------------------------

644 def test_lazy_duplicate_store_skipped() -> None:

645 """Lazy: blocks already offloaded to CPU should not be offloaded again.

@@ -1092,11 +1194,9 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None:

1194 hit_tokens, is_async = sched.get_num_new_matched_tokens(

1195 req2, num_computed_tokens=gpu_local_computed

1196 )

1095 - # CPU should hit blocks 2,3 (not 4,5 — those are beyond the CPU range).

1096 - num_cpu_hit_blocks = 2

1097 - # Actually CPU has all 6 stored; it returns hits starting from position 2.

1098 - # The number of CPU hit blocks = min(remaining request blocks, CPU cached).

1099 - # Here remaining = 6 - 2 = 4 blocks are in CPU, so hit = 4 * BLOCK_SIZE.

1197 + # CPU has all 6 blocks stored. make_request pads num_tokens by +1, so

1198 + # the manager's num_tokens - 1 cap leaves all full blocks intact:

1199 + # remaining hashable range = 6 - 2 = 4 blocks, all hit.

1200 num_cpu_hit_blocks = 4

1201 assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, (

1202 f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}"

vllm/v1/simple_kv_offload/manager.py

+8 -5

@@ -146,6 +146,7 @@ class SimpleCPUOffloadScheduler:

146 # Eager mode only

147 self._reqs_to_store: dict[str, StoreRequestState] = {}

148 self._store_event_to_reqs: dict[int, list[str]] = {}

149 + self._in_flight_store_gpu_blocks: set[int] = set()

150

151 # Event counters

152 self._load_event_counter: int = 0

@@ -468,7 +469,8 @@ class SimpleCPUOffloadScheduler:

469 num_free = cpu_block_pool.get_num_free_blocks()

470 kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups

471 num_groups = len(kv_cache_groups)

471 - gpu_blocks_this_step: set[int] = set()

472 + # Dedup against blocks already scheduled.

473 + in_flight = self._in_flight_store_gpu_blocks

474

475 for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):

476 state = self._reqs_to_store.get(req_id)

@@ -523,10 +525,9 @@ class SimpleCPUOffloadScheduler:

525 if bhash_with_group is None:

526 break

527

526 - # Check if this group's data is already scheduled for store

527 - # in this step or already cached in CPU.

528 + # Skip if already scheduled for store or already cached in CPU.

529 if (

529 - gpu_block_id in gpu_blocks_this_step

530 + gpu_block_id in in_flight

531 or cpu_block_pool.cached_block_hash_to_block.get_one_block(

532 bhash_with_group

533 )

@@ -561,7 +562,7 @@ class SimpleCPUOffloadScheduler:

562 req_ids.append(req_id)

563 merged_gpu_block_ids.extend(gpu_block_ids)

564 merged_cpu_block_ids.extend(cpu_block_ids)

564 - gpu_blocks_this_step.update(gpu_block_ids)

565 + in_flight.update(gpu_block_ids)

566

567 # Touch GPU blocks to prevent freeing during async copy

568 gpu_block_pool.touch(

@@ -608,6 +609,8 @@ class SimpleCPUOffloadScheduler:

609 def _process_store_event(self, event_idx: int) -> None:

610 """Process a fully-completed store event."""

611 transfer = self._store_event_to_blocks.pop(event_idx)

612 + if not self._lazy_mode:

613 + self._in_flight_store_gpu_blocks.difference_update(transfer.gpu_block_ids)

614 self._process_store_completion(transfer.gpu_block_ids, transfer.cpu_block_ids)

615 logger.debug(

616 "Store event %d completed: cached %d blocks to CPU",