@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 12b9573c98

[Bugfix][KV Offloading] Fix eager SimpleCPUOffload cache registration and final flush (#53532)

Signed-off-by: maithilijoshi20 <maithilijoshi2003@gmail.com> Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Co-authored-by: maithilijoshi20 <maithilijoshi2003@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

maithilijoshi20 committed Sep 2, 2026 at 12:23 UTC 12b9573c98c32932e83fdb22e22c152b248c8c88

3 files changed +354 -94

tests/v1/simple_kv_offload/test_scheduler.py

+248 -6

@@ -18,6 +18,9 @@ from vllm.config import (

18 SchedulerConfig,

19 VllmConfig,

20 )

21 +from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (

22 + SimpleCPUOffloadConnector,

23 +)

24 from vllm.utils.hashing import sha256

25 from vllm.v1.core.block_pool import BlockPool

26 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

@@ -39,12 +42,16 @@ from vllm.v1.kv_cache_interface import (

42 KVCacheConfig,

43 KVCacheGroupSpec,

44 KVCacheTensor,

45 + SlidingWindowSpec,

46 )

47 from vllm.v1.outputs import KVConnectorOutput

48 from vllm.v1.request import Request

49 from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler

50 from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadWorkerMetadata

52 +pytestmark = pytest.mark.skip_global_cleanup

53 +

54 +

55 # ---------------------------------------------------------------------------

56 # Constants

57 # ---------------------------------------------------------------------------

@@ -77,15 +84,26 @@ def _make_kv_cache_config(

84 ) # Ensure specs are registered for tests

85 for g in range(num_groups):

86 layer_names = [f"layer_{g}"]

87 + spec = (

88 + FullAttentionSpec(

89 + block_size=BLOCK_SIZE,

90 + num_kv_heads=NUM_KV_HEADS,

91 + head_size=HEAD_SIZE,

92 + dtype=DTYPE,

93 + )

94 + if g == 0

95 + else SlidingWindowSpec(

96 + block_size=BLOCK_SIZE,

97 + num_kv_heads=NUM_KV_HEADS,

98 + head_size=HEAD_SIZE,

99 + dtype=DTYPE,

100 + sliding_window=BLOCK_SIZE * 4,

101 + )

102 + )

103 groups.append(

104 KVCacheGroupSpec(

105 layer_names,

83 - FullAttentionSpec(

84 - block_size=BLOCK_SIZE,

85 - num_kv_heads=NUM_KV_HEADS,

86 - head_size=HEAD_SIZE,

87 - dtype=DTYPE,

88 - ),

106 + spec,

107 )

108 )

109 tensors.append(

@@ -412,6 +430,230 @@ def test_eager_store_and_load_roundtrip() -> None:

430 assert len(meta2.load_cpu_blocks) == len(meta2.load_gpu_blocks)

431

432

433 +def test_eager_store_preserves_secondary_block_hashes() -> None:

434 + """CPU copies retain fine-grained hashes owned by the GPU block."""

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

436 + sched = fix.scheduler

437 + req = make_request(num_blocks=1)

438 + kv_blocks = _alloc_and_register(fix, req, num_blocks=1)

439 + gpu_block = kv_blocks.blocks[0][0]

440 + primary_hash = gpu_block.block_hash

441 + primary_num_tokens = gpu_block.block_hash_num_tokens

442 + assert primary_hash is not None

443 +

444 + fine_grained_req = Request(

445 + request_id="req-fine-grained-hash",

446 + prompt_token_ids=req.prompt_token_ids,

447 + sampling_params=req.sampling_params,

448 + pooling_params=None,

449 + mm_features=None,

450 + block_hasher=get_request_block_hasher(BLOCK_SIZE // 2, sha256),

451 + )

452 + secondary_hash = make_block_hash_with_group_id(fine_grained_req.block_hashes[0], 0)

453 + fix.gpu_block_pool._insert_block_hash(

454 + secondary_hash, gpu_block, num_tokens=BLOCK_SIZE // 2

455 + )

456 + assert (

457 + secondary_hash

458 + in fix.gpu_block_pool.cached_block_hashes_by_block[gpu_block.block_id]

459 + )

460 +

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

462 + block_ids = kv_blocks.get_block_ids()

463 + meta = sched.build_connector_meta(

464 + make_scheduler_output(

465 + {req.request_id: BLOCK_SIZE},

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

467 + )

468 + )

469 + simulate_store_completion(sched, meta.store_event)

470 +

471 + cpu_block = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(

472 + secondary_hash

473 + )

474 + assert cpu_block is not None

475 + assert cpu_block.block_hash == primary_hash

476 + assert cpu_block.block_hash_num_tokens == primary_num_tokens

477 + assert (

478 + secondary_hash

479 + in sched.cpu_block_pool.cached_block_hashes_by_block[cpu_block.block_id]

480 + )

481 +

482 +

483 +# ---------------------------------------------------------------------------

484 +# Test 1b: Finished eager stores

485 +# ---------------------------------------------------------------------------

486 +def test_finished_eager_stores_are_batched() -> None:

487 + """Finished requests in one step retain, complete, and cache every store."""

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

489 + sched = fix.scheduler

490 +

491 + requests = [make_request(num_blocks=2) for _ in range(2)]

492 + block_ids = []

493 + for request in requests:

494 + kv_blocks = _alloc_and_register(fix, request, num_blocks=2)

495 + sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0)

496 + block_ids.append(kv_blocks.get_block_ids())

497 +

498 + for request, request_block_ids in zip(requests, block_ids):

499 + sched.request_finished_all_groups(request, request_block_ids)

500 +

501 + meta = sched.build_connector_meta(make_scheduler_output({}))

502 + assert len(meta.store_gpu_blocks) == 4

503 + assert len(meta.store_cpu_blocks) == 4

504 + simulate_store_completion(sched, meta.store_event)

505 +

506 + for request in requests:

507 + for block_hash in request.block_hashes[:2]:

508 + cached = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(

509 + make_block_hash_with_group_id(block_hash, 0)

510 + )

511 + assert cached is not None

512 +

513 + request = requests[0]

514 + cached_request = Request(

515 + request_id="req-finished-eager-load",

516 + prompt_token_ids=request.prompt_token_ids,

517 + sampling_params=request.sampling_params,

518 + pooling_params=None,

519 + mm_features=None,

520 + block_hasher=request._block_hasher,

521 + )

522 + hit_tokens, is_async = sched.get_num_new_matched_tokens(

523 + cached_request, num_computed_tokens=0

524 + )

525 + assert hit_tokens == 2 * BLOCK_SIZE

526 + assert is_async is True

527 +

528 + cached_gpu_blocks = fix.gpu_block_pool.get_new_blocks(2)

529 + sched.update_state_after_alloc(

530 + cached_request,

531 + KVCacheBlocks(blocks=(cached_gpu_blocks,)),

532 + num_external_tokens=hit_tokens,

533 + )

534 + load_meta = sched.build_connector_meta(

535 + make_scheduler_output(

536 + {cached_request.request_id: 1},

537 + new_reqs={

538 + cached_request.request_id: (

539 + [block.block_id for block in cached_gpu_blocks],

540 + )

541 + },

542 + )

543 + )

544 + assert load_meta.load_event >= 0

545 + assert len(load_meta.load_gpu_blocks) == 2

546 + assert len(load_meta.load_cpu_blocks) == 2

547 +

548 +

549 +def test_finished_eager_store_skips_existing_cpu_cache_entries() -> None:

550 + """A final flush must not consume CPU capacity for an existing hash."""

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

552 + sched = fix.scheduler

553 +

554 + first = make_request(num_blocks=2)

555 + first_blocks = _alloc_and_register(fix, first, num_blocks=2)

556 + sched.update_state_after_alloc(first, first_blocks, num_external_tokens=0)

557 + sched.request_finished_all_groups(first, first_blocks.get_block_ids())

558 + first_meta = sched.build_connector_meta(make_scheduler_output({}))

559 + simulate_store_completion(sched, first_meta.store_event)

560 + free_blocks = get_cpu_free_blocks(sched)

561 +

562 + duplicate = Request(

563 + request_id="req-finished-eager-duplicate",

564 + prompt_token_ids=first.prompt_token_ids,

565 + sampling_params=first.sampling_params,

566 + pooling_params=None,

567 + mm_features=None,

568 + block_hasher=first._block_hasher,

569 + )

570 + duplicate_blocks = _alloc_and_register(fix, duplicate, num_blocks=2)

571 + sched.update_state_after_alloc(duplicate, duplicate_blocks, num_external_tokens=0)

572 + sched.request_finished_all_groups(duplicate, duplicate_blocks.get_block_ids())

573 +

574 + duplicate_meta = sched.build_connector_meta(make_scheduler_output({}))

575 + assert duplicate_meta.store_event == -1

576 + assert get_cpu_free_blocks(sched) == free_blocks

577 +

578 +

579 +def test_finished_eager_store_is_reported_pending_before_metadata_build() -> None:

580 + """Queued finished stores must keep the connector pending until submitted."""

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

582 + sched = fix.scheduler

583 + request = make_request(num_blocks=2)

584 + kv_blocks = _alloc_and_register(fix, request, num_blocks=2)

585 + connector = SimpleCPUOffloadConnector.__new__(SimpleCPUOffloadConnector)

586 + connector.scheduler_manager = sched

587 + sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0)

588 +

589 + sched.request_finished_all_groups(request, kv_blocks.get_block_ids())

590 +

591 + assert sched.has_pending_stores()

592 + assert connector.has_pending_push_work()

593 + meta = sched.build_connector_meta(make_scheduler_output({}))

594 + assert meta.store_event >= 0

595 + assert not sched._pending_finished_stores

596 + assert connector.has_pending_push_work()

597 + simulate_store_completion(sched, meta.store_event)

598 + assert not sched.has_pending_stores()

599 + assert not connector.has_pending_push_work()

600 +

601 +

602 +def test_finished_eager_store_caches_all_groups() -> None:

603 + """The final flush retains completed blocks from every KV cache group."""

604 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=2, lazy=False)

605 + sched = fix.scheduler

606 + request = make_request(num_blocks=2)

607 + group_blocks = tuple(

608 + _allocate_gpu_blocks(fix.gpu_block_pool, request, 2, group_id=group_id)

609 + for group_id in range(2)

610 + )

611 + request.num_computed_tokens = 2 * BLOCK_SIZE

612 + kv_blocks = KVCacheBlocks(blocks=group_blocks)

613 + sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0)

614 +

615 + sched.request_finished_all_groups(request, kv_blocks.get_block_ids())

616 + meta = sched.build_connector_meta(make_scheduler_output({}))

617 + assert len(meta.store_gpu_blocks) == 4

618 + simulate_store_completion(sched, meta.store_event)

619 +

620 + for group_id in range(2):

621 + for block_hash in request.block_hashes[:2]:

622 + cached = sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(

623 + make_block_hash_with_group_id(block_hash, group_id)

624 + )

625 + assert cached is not None

626 +

627 +

628 +def test_finished_eager_store_respects_cpu_capacity() -> None:

629 + """A finished request skips excess blocks when the CPU cache is full."""

630 + fix = make_scheduler(num_cpu_blocks=3, num_gpu_blocks=16, lazy=False)

631 + sched = fix.scheduler

632 + request = make_request(num_blocks=3)

633 + kv_blocks = _alloc_and_register(fix, request, num_blocks=3)

634 + sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0)

635 + available = get_cpu_free_blocks(sched)

636 +

637 + sched.request_finished_all_groups(request, kv_blocks.get_block_ids())

638 + meta = sched.build_connector_meta(make_scheduler_output({}))

639 +

640 + assert len(meta.store_gpu_blocks) == available

641 + assert len(meta.store_cpu_blocks) == available

642 +

643 +

644 +def test_reset_releases_unsubmitted_finished_stores() -> None:

645 + """Reset releases finished-store refs that were not sent to a worker."""

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

647 + sched = fix.scheduler

648 + request = make_request(num_blocks=2)

649 + kv_blocks = _alloc_and_register(fix, request, num_blocks=2)

650 + sched.update_state_after_alloc(request, kv_blocks, num_external_tokens=0)

651 +

652 + sched.request_finished_all_groups(request, kv_blocks.get_block_ids())

653 +

654 + assert sched.reset()

655 +

656 +

657 # ---------------------------------------------------------------------------

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

659 # prompt is an exact multiple of BLOCK_SIZE.

vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py

+1 -2

@@ -291,8 +291,7 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):

291 )

292 return False, None

293

294 - # NOTE: New API only for SimpleCPUOffloadConnector.

295 - def has_pending_transfers(self) -> bool:

294 + def has_pending_push_work(self) -> bool:

295 if self.scheduler_manager is not None:

296 return self.scheduler_manager.has_pending_stores()

297 return False

vllm/v1/simple_kv_offload/manager.py

+105 -86

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

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

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

176 self._in_flight_store_gpu_blocks: set[int] = set()

177 + self._pending_finished_stores: list[TransferMeta] = []

178 self._abandoned_reqs_to_load: dict[str, LoadRequestState] = {}

179

180 # Event counters

@@ -414,6 +415,11 @@ class SimpleCPUOffloadScheduler:

415 # --- Stores ---

416 store_event = -1

417 store_gpu, store_cpu, store_req_ids = self.prepare_store_specs(scheduler_output)

418 + for transfer in self._pending_finished_stores:

419 + store_gpu.extend(transfer.gpu_block_ids)

420 + store_cpu.extend(transfer.cpu_block_ids)

421 + self._pending_finished_stores.clear()

422 +

423 if store_gpu:

424 store_event = self._store_event_counter

425 self._store_event_counter += 1

@@ -492,7 +498,6 @@ class SimpleCPUOffloadScheduler:

498 self._cursor = None

499

500 gpu_ids: list[int] = []

495 - block_hashes: list[bytes] = []

501 last_visited = self._cursor

502

503 for covered, node in enumerate(free_queue.iter_blocks_after(self._cursor)):

@@ -508,16 +513,13 @@ class SimpleCPUOffloadScheduler:

513 and cpu_pool.cached_block_hash_to_block.get_one_block(bhash) is None

514 ):

515 gpu_ids.append(node.block_id)

511 - block_hashes.append(bhash)

516

517 self._cursor = last_visited

518

515 - # Batch-allocate CPU blocks and stamp hashes.

519 + # Batch-allocate CPU blocks.

520 if gpu_ids:

521 cpu_blocks = cpu_pool.get_new_blocks(len(gpu_ids))

522 cpu_ids = [blk.block_id for blk in cpu_blocks]

519 - for cpu_blk, bhash in zip(cpu_blocks, block_hashes): # type: ignore[assignment]

520 - cpu_blk._block_hash = bhash # type: ignore[assignment]

523 # Touch GPU blocks to prevent eviction during async copy.

524 gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])

525 else:

@@ -531,9 +533,8 @@ class SimpleCPUOffloadScheduler:

533 """Identify newly computed blocks to offload from scheduler requests.

534

535 Only considers blocks whose KV data has been **confirmed computed** by

534 - the GPU. This means blocks from the current step are NOT stored until the

535 - next step. If a request finishes in the same step as its last full block,

536 - that block may be missed. (TODO: flush on finish.)

536 + the GPU. Blocks from the current step are stored on a later step, or by

537 + the finish-time flush if the request completes first.

538

539 Returns:

540 (gpu_block_ids, cpu_block_ids, req_ids) for the store event.

@@ -547,7 +548,6 @@ class SimpleCPUOffloadScheduler:

548 if gpu_block_pool is None:

549 return [], [], []

550 cpu_block_pool = self.cpu_block_pool

550 - num_free = cpu_block_pool.get_num_free_blocks()

551 kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups

552 num_groups = len(kv_cache_groups)

553 # Dedup against blocks already scheduled.

@@ -575,73 +575,15 @@ class SimpleCPUOffloadScheduler:

575 if not block_ids_by_group:

576 continue

577

578 - # --- Phase 1: Scan blocks, classify as cached vs to-store ---

579 - gpu_block_ids: list[int] = []

580 - block_hashes_to_store: list[bytes] = []

581 - advanced_per_group: list[int] = [0] * num_groups

582 - out_of_space = False

583 - # Confirmed tokens: KV data written and visible to all streams.

584 - req = state.request

585 - confirmed_tokens = req.num_computed_tokens - req.num_output_placeholders

586 - # Cap to blocks with confirmed KV data.

587 - aligned_tokens = confirmed_tokens // self.block_size * self.block_size

588 -

589 - for g in range(num_groups):

590 - # FIXME (yifan): handle CPU cache eviction, where

591 - # num_stored_blocks can be stale and omit evicted blocks in

592 - # the middle of the request.

593 - already_stored_g = state.num_stored_blocks[g]

594 - group_gpu_ids = block_ids_by_group[g]

595 -

596 - g_block_size = (

597 - kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size

598 - )

599 - ready_blocks_g = aligned_tokens // g_block_size

600 - scannable = group_gpu_ids[already_stored_g:ready_blocks_g]

601 -

602 - for gpu_block_id in scannable:

603 - gpu_block = gpu_block_pool.blocks[gpu_block_id]

604 - if gpu_block.is_null:

605 - advanced_per_group[g] += 1

606 - continue

607 -

608 - bhash_with_group = gpu_block.block_hash

609 - if bhash_with_group is None:

610 - # Masked-out SWA position the coordinator chose not to

611 - # hash; it can never serve a prefix-cache hit, so skip.

612 - advanced_per_group[g] += 1

613 - continue

614 -

615 - # Skip if already scheduled for store or already cached in CPU.

616 - if (

617 - gpu_block_id in in_flight

618 - or cpu_block_pool.cached_block_hash_to_block.get_one_block(

619 - bhash_with_group

620 - )

621 - is not None

622 - ):

623 - advanced_per_group[g] += 1

624 - continue

625 -

626 - if num_free <= 0:

627 - out_of_space = True

628 - break

629 - num_free -= 1

630 -

631 - gpu_block_ids.append(gpu_block_id)

632 - block_hashes_to_store.append(bhash_with_group)

633 - advanced_per_group[g] += 1

634 -

635 - if out_of_space:

636 - break

578 + gpu_block_ids, advanced_per_group = self._select_eager_blocks_to_store(

579 + state, block_ids_by_group

580 + )

581

638 - # --- Phase 2: Batch allocate CPU blocks and stamp hashes ---

582 + # --- Phase 2: Batch allocate CPU blocks ---

583 n_to_alloc = len(gpu_block_ids)

584 if n_to_alloc > 0:

585 cpu_blocks_alloc = cpu_block_pool.get_new_blocks(n_to_alloc)

586 cpu_block_ids = [blk.block_id for blk in cpu_blocks_alloc]

643 - for cpu_blk, bhash in zip(cpu_blocks_alloc, block_hashes_to_store):

644 - cpu_blk._block_hash = bhash # type: ignore[assignment]

587 else:

588 cpu_block_ids = []

589

@@ -669,6 +611,53 @@ class SimpleCPUOffloadScheduler:

611

612 return merged_gpu_block_ids, merged_cpu_block_ids, req_ids

613

614 + def _select_eager_blocks_to_store(

615 + self,

616 + state: StoreRequestState,

617 + block_ids_by_group: tuple[list[int], ...],

618 + ) -> tuple[list[int], list[int]]:

619 + """Return confirmed, uncached eager blocks and per-group cursor advances."""

620 + assert self._gpu_block_pool is not None

621 + gpu_block_ids: list[int] = []

622 + advanced_per_group = [0] * len(self.cpu_kv_cache_config.kv_cache_groups)

623 + confirmed_tokens = (

624 + state.request.num_computed_tokens - state.request.num_output_placeholders

625 + )

626 + aligned_tokens = confirmed_tokens // self.block_size * self.block_size

627 + num_free = self.cpu_block_pool.get_num_free_blocks()

628 +

629 + for g, group_gpu_ids in enumerate(block_ids_by_group):

630 + if len(gpu_block_ids) >= num_free:

631 + break

632 + # FIXME (yifan): handle CPU cache eviction, where

633 + # num_stored_blocks can be stale and omit evicted blocks in

634 + # the middle of the request.

635 + group_size = (

636 + self.cpu_kv_cache_config.kv_cache_groups[g].kv_cache_spec.block_size

637 + * self.cp_world_size

638 + )

639 + ready = min(len(group_gpu_ids), aligned_tokens // group_size)

640 + for gpu_block_id in group_gpu_ids[state.num_stored_blocks[g] : ready]:

641 + gpu_block = self._gpu_block_pool.blocks[gpu_block_id]

642 + if gpu_block.is_null or gpu_block.block_hash is None:

643 + advanced_per_group[g] += 1

644 + continue

645 + if (

646 + gpu_block_id in self._in_flight_store_gpu_blocks

647 + or self.cpu_block_pool.cached_block_hash_to_block.get_one_block(

648 + gpu_block.block_hash

649 + )

650 + is not None

651 + ):

652 + advanced_per_group[g] += 1

653 + continue

654 + if len(gpu_block_ids) >= num_free:

655 + break

656 + gpu_block_ids.append(gpu_block_id)

657 + advanced_per_group[g] += 1

658 +

659 + return gpu_block_ids, advanced_per_group

660 +

661 def update_connector_output(self, connector_output: KVConnectorOutput) -> None:

662 """Handle async transfer completions from worker.

663

@@ -726,24 +715,31 @@ class SimpleCPUOffloadScheduler:

715 def _process_store_completion(

716 self, gpu_block_ids: list[int], cpu_block_ids: list[int]

717 ) -> None:

729 - """Cache CPU blocks per-group and release GPU refs.

730 -

731 - Block hashes were stamped on CPU blocks at allocation time (in

732 - ``_prepare_*_store_specs``). Here we just register them in the

733 - cache map so they become discoverable by the load path.

734 - """

718 + """Register copied blocks in the CPU prefix cache and release refs."""

719 assert len(cpu_block_ids) == len(gpu_block_ids)

720

721 cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in cpu_block_ids]

722

739 - for cpu_block in cpu_blocks:

740 - bhash = cpu_block.block_hash

741 - assert bhash is not None

742 - self.cpu_block_pool.cached_block_hash_to_block.insert(bhash, cpu_block)

723 + assert self._gpu_block_pool is not None

724 + for gpu_block_id, cpu_block in zip(gpu_block_ids, cpu_blocks):

725 + gpu_block = self._gpu_block_pool.blocks[gpu_block_id]

726 + primary_hash = gpu_block.block_hash

727 + assert primary_hash is not None

728 + self.cpu_block_pool._insert_block_hash(

729 + primary_hash,

730 + cpu_block,

731 + num_tokens=gpu_block.block_hash_num_tokens,

732 + )

733 + secondary_hashes = self._gpu_block_pool.cached_block_hashes_by_block.get(

734 + gpu_block_id, ()

735 + )

736 + for block_hash in secondary_hashes:

737 + self.cpu_block_pool._insert_block_hash(

738 + block_hash, cpu_block, num_tokens=None

739 + )

740

741 # Free CPU and GPU blocks' ref counts to turn them into prefix cache

742 self.cpu_block_pool.free_blocks(cpu_blocks)

746 - assert self._gpu_block_pool is not None

743 self._gpu_block_pool.free_blocks(

744 self._gpu_block_pool.blocks[bid] for bid in gpu_block_ids

745 )

@@ -751,8 +747,6 @@ class SimpleCPUOffloadScheduler:

747 def _release_transfer_refs(self, transfer: TransferMeta) -> None:

748 """Release transfer refs without making copied data cacheable."""

749 cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in transfer.cpu_block_ids]

754 - for cpu_block in cpu_blocks:

755 - cpu_block.reset_hash()

750 self.cpu_block_pool.free_blocks(cpu_blocks)

751 assert self._gpu_block_pool is not None

752 self._gpu_block_pool.free_blocks(

@@ -760,9 +754,11 @@ class SimpleCPUOffloadScheduler:

754 )

755

756 def has_pending_stores(self) -> bool:

763 - """Return True if there are in-flight store transfers."""

757 + """Return True if a store transfer is queued or in flight."""

758 return bool(

765 - self._store_event_to_blocks or self._abandoned_store_event_to_blocks

759 + self._pending_finished_stores

760 + or self._store_event_to_blocks

761 + or self._abandoned_store_event_to_blocks

762 )

763

764 def request_finished(

@@ -804,8 +800,27 @@ class SimpleCPUOffloadScheduler:

800 request: "Request",

801 block_ids: tuple[list[int], ...],

802 ) -> tuple[bool, dict[str, Any] | None]:

803 + self._queue_finished_eager_store(request, block_ids)

804 return self.request_finished(request, block_ids=[])

805

806 + def _queue_finished_eager_store(

807 + self, request: "Request", block_ids: tuple[list[int], ...]

808 + ) -> None:

809 + """Queue confirmed eager blocks omitted after a request enters decode."""

810 + state = self._reqs_to_store.get(request.request_id)

811 + gpu_pool = self._gpu_block_pool

812 + if state is None or gpu_pool is None:

813 + return

814 + gpu_ids, _ = self._select_eager_blocks_to_store(state, block_ids)

815 + if not gpu_ids:

816 + return

817 + cpu_blocks = self.cpu_block_pool.get_new_blocks(len(gpu_ids))

818 + self._pending_finished_stores.append(

819 + TransferMeta(gpu_ids, [b.block_id for b in cpu_blocks])

820 + )

821 + self._in_flight_store_gpu_blocks.update(gpu_ids)

822 + gpu_pool.touch([gpu_pool.blocks[bid] for bid in gpu_ids])

823 +

824 def _free_pending_cpu_hit(self, pending: tuple) -> None:

825 """Release the temporary CPU block pin taken in get_num_new_matched_tokens()."""

826 cpu_hit_blocks, _ = pending

@@ -879,6 +894,10 @@ class SimpleCPUOffloadScheduler:

894 """

895

896 self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks)

897 + for transfer in self._pending_finished_stores:

898 + self._release_transfer_refs(transfer)

899 + self._pending_finished_stores.clear()

900 +

901 self._store_event_to_blocks.clear()

902 self._in_flight_store_gpu_blocks.clear()

903