@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dc9ae4b8ac

[Bugfix][Mooncake] Reference GPU blocks for in-flight store jobs and key the store ledger by store_job_id (#52372)

Signed-off-by: Guanyi Chen <939416532@qq.com>

Guanyi Chen committed Aug 17, 2026 at 07:09 UTC dc9ae4b8ac2331991ad7091812ef82ece4f8fdc2

7 files changed +474 -250

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

+7 -8

@@ -228,12 +228,11 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

228 block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),

229 block_hashes=hs,

230 can_save=True,

231 + store_job_id=1,

232 )

232 - send_thread.add_stored_request("r0")

233 - # Put the request in the queue so task_done() doesn't underflow.

234 - send_thread.request_queue.put(save_req)

235 - req = send_thread.request_queue.get()

236 - send_thread._handle_request(req)

233 + # add_request also queues the job, so task_done() doesn't underflow.

234 + send_thread.add_request(save_req)

235 + send_thread._handle_request(send_thread.request_queue.get())

236

237 # Point worker.store at the dict store (the worker constructor captured

238 # the MagicMock; replace with the real dict store for lookup).

@@ -499,15 +498,15 @@ def test_offload_syncs_event_before_put():

498 can_save=True,

499 num_prompt_tokens=12,

500 partial_tail_offloads=[(1, 7, 12)],

501 + store_job_id=1,

502 )

503 req.current_event = event

504 - send.add_stored_request("r1")

504

506 - send.request_queue.put(req)

505 + send.add_request(req)

506 send._handle_request(send.request_queue.get())

507 assert send.request_queue.qsize() == 0

508 assert store._data

510 - assert send.stored_requests["r1"] == 0

509 + assert send.stored_requests["r1"] == set()

510 event.synchronize.assert_called_once()

511

512

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

+125 -54

@@ -5,12 +5,14 @@ from types import SimpleNamespace

6 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

7 LoadSpec,

8 + MooncakeStoreWorkerMetadata,

9 ReqMeta,

10 RequestTracker,

11 )

12 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import (

13 MooncakeStoreScheduler,

14 )

15 +from vllm.v1.core.block_pool import BlockPool

18 def _make_bare_scheduler(

@@ -27,6 +29,12 @@ def _make_bare_scheduler(

29 scheduler._unfinished_request_ids = {"req-0"}

30 scheduler._unfinished_requests = {}

31 scheduler._request_trackers = {}

32 + scheduler._gpu_block_pool = BlockPool(

33 + num_gpu_blocks=64, enable_caching=True, hash_block_size=hash_block_size

34 + )

35 + scheduler._num_workers = 1

36 + scheduler._next_store_job_id = 0

37 + scheduler._pinned_saves = {}

38 return scheduler

@@ -64,6 +72,14 @@ def _make_preemption_scheduler_output():

72 )

75 +def _make_worker_output(completed_saves: dict[int, int]) -> SimpleNamespace:

76 + return SimpleNamespace(

77 + kv_connector_worker_meta=MooncakeStoreWorkerMetadata(

78 + completed_saves=completed_saves

79 + )

80 + )

81 +

82 +

83 def _add_unfinished_request(

84 scheduler: MooncakeStoreScheduler,

85 *,

@@ -133,7 +149,7 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap():

149 assert tracker.num_saved_tokens == 48

150

151

136 -def test_preemption_resets_tracker_before_request_finished():

152 +def test_preemption_resets_tracker():

153 scheduler = _make_bare_scheduler()

154 _add_unfinished_request(

155 scheduler,

@@ -152,8 +168,6 @@ def test_preemption_resets_tracker_before_request_finished():

168 assert tracker.token_ids is None

169 assert tracker.has_pending_offload is False

170 assert tracker.prefill_end_tokens == 0

155 - request = SimpleNamespace(request_id="req-0")

156 - assert scheduler.request_finished(request, ([0, 1],)) == (False, None)

171

172

173 def test_preemption_clears_stale_load_state():

@@ -215,7 +229,7 @@ def test_pending_load_does_not_co_queue_save():

229 # enqueue a save in the same scheduling step. Co-queuing both produces a

230 # recv+send pair for the same req_id, and the scheduler's

231 # _update_from_kv_xfer_finished then trips `assert req_id in self.requests`

218 - # when both completions land for the delay-freed request.

232 + # when a completion lands for a request it has already dropped.

233 scheduler = _make_bare_scheduler()

234 _make_pending_load_unfinished_request(

235 scheduler,

@@ -238,9 +252,7 @@ def test_pending_load_does_not_co_queue_save():

252 # Load is still issued as planned.

253 assert req_meta.load_spec is not None

254 assert req_meta.load_spec.can_load is True

241 - # And the tracker's saved-tokens watermark stays at 0 so request_finished

242 - # later sees `num_saved_tokens <= 0` and frees immediately rather than

243 - # waiting for a finished_sending that will never come.

255 + # And the save watermark does not advance for a save that was never queued.

256 tracker = scheduler._request_trackers["req-0"]

257 assert tracker.num_saved_tokens == 0

258

@@ -661,30 +673,34 @@ def test_disabled_lookup_reports_no_hit_without_querying_client():

673 assert scheduler.load_specs == {}

674

675

664 -def test_pending_partial_tail_emits_offload_only_reqmeta():

665 - # A sub-block prompt never produces a block-aligned save, so the partial-

666 - # tail offload arriving this step is emitted as an offload-only ReqMeta

667 - # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so

668 - # the worker skips the normal save). Pending-offload state delays the free

669 - # without advancing the normal-save watermark before the put succeeds.

670 - scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

676 +def _add_pending_partial_tail_request(

677 + scheduler: MooncakeStoreScheduler,

678 + *,

679 + num_tokens: int,

680 + block_hashes: list[bytes],

681 + block_ids: tuple[list[int], ...],

682 +) -> SimpleNamespace:

683 + """Register a sub-block request and return the step that offloads its tail.

684 +

685 + The CoW block holding the tail is block 7, which the core deliberately keeps

686 + out of the request's block table.

687 + """

688 request = SimpleNamespace(

672 - all_token_ids=list(range(12)),

673 - block_hashes=[b"h0", b"h1", b"h2"],

689 + all_token_ids=list(range(num_tokens)),

690 + block_hashes=block_hashes,

691 num_output_placeholders=0,

692 num_prompt_tokens=12,

693 )

677 - scheduler._unfinished_requests["req-0"] = (request, ([0],))

694 + scheduler._unfinished_requests["req-0"] = (request, block_ids)

695 scheduler._request_trackers["req-0"] = RequestTracker(

696 req_id="req-0",

680 - token_len=12,

681 - allocated_block_ids=([0],),

697 + token_len=num_tokens,

698 + allocated_block_ids=block_ids,

699 num_saved_tokens=0,

683 - token_ids=list(range(12)),

684 - prefill_end_tokens=12,

700 + token_ids=list(range(num_tokens)),

701 + prefill_end_tokens=num_tokens,

702 )

686 -

687 - out = SimpleNamespace(

703 + return SimpleNamespace(

704 finished_req_ids=set(),

705 preempted_req_ids=set(),

706 scheduled_new_reqs=[],

@@ -699,6 +715,21 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():

715 partial_tail_offloads={"req-0": [(1, 7, 12)]},

716 )

717

718 +

719 +def test_pending_partial_tail_emits_offload_only_reqmeta():

720 + # A sub-block prompt never produces a block-aligned save, so the partial-

721 + # tail offload arriving this step is emitted as an offload-only ReqMeta

722 + # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so

723 + # the worker skips the normal save), without advancing the normal-save

724 + # watermark before the put succeeds.

725 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

726 + out = _add_pending_partial_tail_request(

727 + scheduler,

728 + num_tokens=12,

729 + block_hashes=[b"h0", b"h1", b"h2"],

730 + block_ids=([0],),

731 + )

732 +

733 meta = scheduler.build_connector_meta(out)

734

735 assert len(meta.requests) == 1

@@ -712,42 +743,16 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():

743 tracker = scheduler._request_trackers["req-0"]

744 assert tracker.num_saved_tokens == 0

745 assert tracker.has_pending_offload is True

715 - request = SimpleNamespace(request_id="req-0")

716 - assert scheduler.request_finished(request, ([0],)) == (True, None)

746

747

748 def test_resumed_partial_tail_uses_handoff_boundary():

749 scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

721 - request = SimpleNamespace(

722 - all_token_ids=list(range(20)),

750 + # Resumption replays prompt + previously generated tokens.

751 + out = _add_pending_partial_tail_request(

752 + scheduler,

753 + num_tokens=20,

754 block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4"],

724 - num_output_placeholders=0,

725 - num_prompt_tokens=12,

726 - )

727 - scheduler._unfinished_requests["req-0"] = (request, ([0, 1],))

728 - scheduler._request_trackers["req-0"] = RequestTracker(

729 - req_id="req-0",

730 - token_len=20,

731 - allocated_block_ids=([0, 1],),

732 - num_saved_tokens=0,

733 - token_ids=list(range(20)),

734 - # Resumption replays prompt + previously generated tokens.

735 - prefill_end_tokens=20,

736 - )

737 -

738 - out = SimpleNamespace(

739 - finished_req_ids=set(),

740 - preempted_req_ids=set(),

741 - scheduled_new_reqs=[],

742 - scheduled_cached_reqs=SimpleNamespace(

743 - req_ids=[],

744 - new_block_ids=[],

745 - num_computed_tokens=[],

746 - resumed_req_ids=set(),

747 - ),

748 - num_scheduled_tokens={},

749 - scheduled_spec_decode_tokens={},

750 - partial_tail_offloads={"req-0": [(1, 7, 12)]},

755 + block_ids=([0, 1],),

756 )

757

758 meta = scheduler.build_connector_meta(out)

@@ -791,3 +796,69 @@ def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary():

796 tracker = scheduler._request_trackers["req-0"]

797 assert tracker.num_saved_tokens == 48

798 assert tracker.has_pending_offload is True

799 +

800 +

801 +def test_partial_tail_cow_block_is_referenced_for_the_job():

802 + # The CoW block a partial-tail offload reads is deliberately kept out of the

803 + # request block table, so it is absent from ReqMeta.block_ids. The worker

804 + # DMAs out of it just as asynchronously, so it needs its own reference.

805 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

806 + out = _add_pending_partial_tail_request(

807 + scheduler,

808 + num_tokens=12,

809 + block_hashes=[b"h0", b"h1", b"h2"],

810 + block_ids=([0],),

811 + )

812 + pool = scheduler._gpu_block_pool

813 +

814 + meta = scheduler.build_connector_meta(out)

815 +

816 + store_job_id = meta.requests[0].store_job_id

817 + # It leads the list, as in `pop_blocks_for_free`, so that the reversed free

818 + # puts it last in eviction priority.

819 + assert scheduler._pinned_saves[store_job_id][0] == [7, 0]

820 + assert pool.blocks[7].ref_cnt == 1

821 +

822 + scheduler.update_connector_output(_make_worker_output({store_job_id: 1}))

823 + assert pool.blocks[7].ref_cnt == 0

824 +

825 +

826 +def test_store_job_blocks_are_released_once_every_rank_reports():

827 + # Every rank DMAs the job's blocks on its own, so the reference can only be

828 + # dropped once the last of them reports. Until then the engine has to keep

829 + # stepping: a completion only reaches the scheduler as worker metadata

830 + # attached to a step, and a finishing request no longer defers its own free.

831 + scheduler = _make_bare_scheduler()

832 + scheduler._num_workers = 2

833 + _add_unfinished_request(

834 + scheduler,

835 + token_ids=list(range(48)),

836 + block_hashes=[b"h0", b"h1", b"h2"],

837 + prefill_end_tokens=48,

838 + )

839 + pool = scheduler._gpu_block_pool

840 + assert scheduler.has_pending_push_work() is False

841 +

842 + meta = scheduler.build_connector_meta(

843 + _make_scheduler_output(scheduled_spec_tokens=None)

844 + )

845 + store_job_id = meta.requests[0].store_job_id

846 + assert pool.blocks[2].ref_cnt == 1

847 + assert scheduler.has_pending_push_work() is True

848 +

849 + scheduler.update_connector_output(_make_worker_output({store_job_id: 1}))

850 + assert pool.blocks[2].ref_cnt == 1

851 + assert scheduler.has_pending_push_work() is True

852 +

853 + scheduler.update_connector_output(_make_worker_output({store_job_id: 1}))

854 + assert pool.blocks[2].ref_cnt == 0

855 + assert scheduler.has_pending_push_work() is False

856 +

857 +

858 +def test_worker_metadata_aggregates_completions_across_ranks():

859 + # The engine merges each rank's metadata before the scheduler sees it, so a

860 + # job that every rank finished in one step arrives as a single count.

861 + merged = MooncakeStoreWorkerMetadata(completed_saves={1: 1}).aggregate(

862 + MooncakeStoreWorkerMetadata(completed_saves={1: 1, 2: 1})

863 + )

864 + assert merged.completed_saves == {1: 2, 2: 1}

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+117 -99

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import contextlib

5 +import itertools

6 import json

7 import logging

8 import math

@@ -156,6 +158,17 @@ def _make_load_req(

158 )

159

160

161 +_TEST_SAVE_SEQ = itertools.count(1)

162 +

163 +

164 +def _run_store_req(thread, req_meta: ReqMeta) -> None:

165 + """Register, enqueue and run a store job the way the worker does."""

166 + if req_meta.store_job_id is None:

167 + req_meta.store_job_id = next(_TEST_SAVE_SEQ)

168 + thread.add_request(req_meta)

169 + thread._handle_request(thread.request_queue.get())

170 +

171 +

172 def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:

173 return ReqMeta(

174 req_id=req_id,

@@ -385,27 +398,23 @@ def test_store_sending_thread_skips_request_during_cpu_pressure():

398 ]

399 thread = _make_store_sending_thread(store)

400

388 - thread.add_stored_request("req-a")

389 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

401 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

402

403 assert thread._store_pressure_active is True

404 assert "req-a" in thread._skip_store_requests

405 assert store.batch_put_from_multi_buffers.call_count == 1

406

395 - thread.add_stored_request("req-a")

396 - thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))

407 + _run_store_req(thread, _make_store_req("req-a", [b"a2", b"a3"]))

408

409 assert store.batch_put_from_multi_buffers.call_count == 1

410

400 - thread.add_stored_request("req-b")

401 - thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"]))

411 + _run_store_req(thread, _make_store_req("req-b", [b"b0", b"b1"]))

412

413 assert thread._store_pressure_active is False

414 assert "req-a" not in thread._skip_store_requests

415 assert store.batch_put_from_multi_buffers.call_count == 2

416

407 - thread.add_stored_request("req-a")

408 - thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"]))

417 + _run_store_req(thread, _make_store_req("req-a", [b"a4", b"a5"]))

418

419 assert store.batch_put_from_multi_buffers.call_count == 3

420

@@ -418,8 +427,7 @@ def test_store_sending_thread_records_mooncake_metrics():

427 stats = MooncakeStoreConnectorStats()

428 thread._record_operation_cb = stats.record_operation

429

421 - thread.add_stored_request("req-a")

422 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

430 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

431

432 assert len(stats.data["save_exists"]) == 1

433 assert stats.data["save_exists"][0]["num_keys"] == 2

@@ -497,16 +505,16 @@ def test_store_sending_thread_delta_saves_only_new_full_attention_chunks():

505 store.batch_put_from_multi_buffers.return_value = [256, 256]

506 thread = _make_store_sending_thread(store)

507

500 - thread.add_stored_request("req-a")

508 thread._saved_offset["req-a"] = 32

502 - thread._handle_request(

509 + _run_store_req(

510 + thread,

511 ReqMeta(

512 req_id="req-a",

513 token_len_chunk=64,

514 block_ids=([0, 1, 2, 3],),

515 block_hashes=[b"a0", b"a1", b"a2", b"a3"],

516 can_save=True,

509 - )

517 + ),

518 )

519

520 keys = store.batch_is_exist.call_args.args[0]

@@ -523,16 +531,16 @@ def test_store_sending_thread_delta_strides_with_local_phase():

531 store.batch_put_from_multi_buffers.return_value = [256]

532 thread = _make_store_sending_thread(store, tp_rank=0, put_step=2)

533

526 - thread.add_stored_request("req-a")

534 thread._saved_offset["req-a"] = 16

528 - thread._handle_request(

535 + _run_store_req(

536 + thread,

537 ReqMeta(

538 req_id="req-a",

539 token_len_chunk=64,

540 block_ids=([0, 1, 2, 3],),

541 block_hashes=[b"a0", b"a1", b"a2", b"a3"],

542 can_save=True,

535 - )

543 + ),

544 )

545

546 keys = store.batch_is_exist.call_args.args[0]

@@ -550,15 +558,15 @@ def test_tp_sharded_group_saves_every_block_on_every_rank():

558 thread = _make_store_sending_thread(store, tp_rank=0, put_step=2)

559 thread.group_put_steps = [1]

560

553 - thread.add_stored_request("req-a")

554 - thread._handle_request(

561 + _run_store_req(

562 + thread,

563 ReqMeta(

564 req_id="req-a",

565 token_len_chunk=64,

566 block_ids=([0, 1, 2, 3],),

567 block_hashes=[b"a0", b"a1", b"a2", b"a3"],

568 can_save=True,

561 - )

569 + ),

570 )

571

572 keys = store.batch_is_exist.call_args.args[0]

@@ -576,15 +584,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():

584 thread._store_pressure_active = True

585 thread._skip_store_requests.add("req-a")

586

579 - thread.add_stored_request("req-a")

580 - thread._handle_request(

587 + _run_store_req(

588 + thread,

589 ReqMeta(

590 req_id="req-a",

591 token_len_chunk=16,

592 block_ids=([0],),

593 block_hashes=[b"a0"],

594 can_save=True,

587 - )

595 + ),

596 )

597

598 store.batch_is_exist.assert_not_called()

@@ -595,15 +603,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():

603

604 # The next batch resumes from offset 0, re-covering the chunk skipped under

605 # pressure (chunk 0) rather than losing it.

598 - thread.add_stored_request("req-a")

599 - thread._handle_request(

606 + _run_store_req(

607 + thread,

608 ReqMeta(

609 req_id="req-a",

610 token_len_chunk=64,

611 block_ids=([0, 1, 2, 3],),

612 block_hashes=[b"a0", b"a1", b"a2", b"a3"],

613 can_save=True,

606 - )

614 + ),

615 )

616

617 keys = store.batch_is_exist.call_args.args[0]

@@ -703,13 +711,11 @@ def test_partial_tail_offload_honors_active_pressure_gate():

711 thread = _make_partial_tail_send_thread(store)

712 thread._store_pressure_active = True

713 thread._skip_store_requests.add("req-a")

706 - thread.add_stored_request("req-a")

714

708 - thread._handle_request(_make_partial_tail_req([1, 2, 3]))

715 + _run_store_req(thread, _make_partial_tail_req([1, 2, 3]))

716

717 store.batch_is_exist.assert_not_called()

718 store.batch_put_from_multi_buffers.assert_not_called()

712 - assert thread.stored_requests["req-a"] == 0

719

720

721 def test_partial_tail_put_failure_activates_pressure_gate():

@@ -717,17 +723,14 @@ def test_partial_tail_put_failure_activates_pressure_gate():

723 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

724 store.batch_put_from_multi_buffers.return_value = [256, -200, 256]

725 thread = _make_partial_tail_send_thread(store)

720 - thread.add_stored_request("req-a")

726

722 - thread._handle_request(_make_partial_tail_req([1, 2, 3]))

727 + _run_store_req(thread, _make_partial_tail_req([1, 2, 3]))

728

729 assert thread._store_pressure_active is True

730 assert thread._skip_store_requests == {"req-a"}

731 assert thread._saved_offset.get("req-a", 0) == 0

727 - assert thread.stored_requests["req-a"] == 0

732

729 - thread.add_stored_request("req-a")

730 - thread._handle_request(_make_partial_tail_req([1, 2, 3]))

733 + _run_store_req(thread, _make_partial_tail_req([1, 2, 3]))

734 assert store.batch_put_from_multi_buffers.call_count == 1

735

736

@@ -737,16 +740,16 @@ def test_store_sending_thread_delta_start_rank_saves_second_local_chunk():

740 store.batch_put_from_multi_buffers.return_value = [256, 256]

741 thread = _make_store_sending_thread(store, tp_rank=1, put_step=2)

742

740 - thread.add_stored_request("req-a")

743 thread._saved_offset["req-a"] = 16

742 - thread._handle_request(

744 + _run_store_req(

745 + thread,

746 ReqMeta(

747 req_id="req-a",

748 token_len_chunk=64,

749 block_ids=([0, 1, 2, 3],),

750 block_hashes=[b"a0", b"a1", b"a2", b"a3"],

751 can_save=True,

749 - )

752 + ),

753 )

754

755 keys = store.batch_is_exist.call_args.args[0]

@@ -790,16 +793,16 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks():

793 token_databases=[db_full, db_masked],

794 )

795

793 - thread.add_stored_request("req-a")

796 thread._saved_offset["req-a"] = 32

795 - thread._handle_request(

797 + _run_store_req(

798 + thread,

799 ReqMeta(

800 req_id="req-a",

801 token_len_chunk=64,

802 block_ids=([0, 1, 2, 3], [0, 1, 2, 3]),

803 block_hashes=[b"a0", b"a1", b"a2", b"a3"],

804 can_save=True,

802 - )

805 + ),

806 )

807

808 keys = store.batch_is_exist.call_args.args[0]

@@ -845,15 +848,15 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group():

848 coord=coord,

849 token_databases=[db0, db1],

850 )

848 - thread.add_stored_request("req-a")

849 - thread._handle_request(

851 + _run_store_req(

852 + thread,

853 ReqMeta(

854 req_id="req-a",

855 token_len_chunk=48,

856 block_ids=([0, 1, 2], [2, 1, 0]),

857 block_hashes=[b"a0", b"a1", b"a2"],

858 can_save=True,

856 - )

859 + ),

860 )

861

862 db0.prepare_value.assert_not_called()

@@ -881,46 +884,71 @@ def test_store_sending_thread_only_skips_on_no_available_handle():

884 ]

885 thread = _make_store_sending_thread(store)

886

884 - thread.add_stored_request("req-a")

885 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

887 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

888

889 assert thread._store_pressure_active is False

890 assert "req-a" not in thread._skip_store_requests

891 assert store.batch_put_from_multi_buffers.call_count == 1

892

891 - thread.add_stored_request("req-a")

892 - thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))

893 + _run_store_req(thread, _make_store_req("req-a", [b"a2", b"a3"]))

894

895 assert store.batch_put_from_multi_buffers.call_count == 2

896

897

897 -def test_store_sending_thread_releases_pin_on_batch_is_exist_failure():

898 - # `batch_is_exist` raising must still decrement `stored_requests` so the

899 - # scheduler can drop `delay_free_blocks` and release the pinned GPU blocks.

898 +@pytest.mark.parametrize(

899 + "failing_call", ["batch_is_exist", "batch_put_from_multi_buffers"]

900 +)

901 +def test_store_sending_thread_reports_job_when_store_raises(failing_call):

902 + # A store that blows up must still report its job, or the scheduler keeps

903 + # the job's GPU block references for the rest of the run.

904 store = MagicMock()

901 - store.batch_is_exist.side_effect = RuntimeError("mooncake down")

905 + store.batch_is_exist.return_value = [0, 0]

906 + getattr(store, failing_call).side_effect = RuntimeError("mooncake down")

907 thread = _make_store_sending_thread(store)

908 + req = _make_store_req("req-a", [b"a0", b"a1"])

909

904 - thread.add_stored_request("req-a")

905 - with pytest.raises(RuntimeError):

906 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

910 + with contextlib.suppress(RuntimeError):

911 + _run_store_req(thread, req)

912

908 - assert thread.stored_requests["req-a"] == 0

909 - store.batch_put_from_multi_buffers.assert_not_called()

913 + assert thread.take_completed_saves() == {req.store_job_id: 1}

914

915

912 -def test_store_sending_thread_releases_pin_on_batch_put_failure():

913 - # `batch_put_from_multi_buffers` raising is logged (not re-raised), and the

914 - # pin must still be released through the finally block.

915 - store = MagicMock()

916 - store.batch_is_exist.return_value = [0, 0]

917 - store.batch_put_from_multi_buffers.side_effect = RuntimeError("rdma error")

918 - thread = _make_store_sending_thread(store)

916 +def test_store_sending_thread_reports_job_when_the_preamble_raises():

917 + # The report has to survive a failure before the first store call too, so

918 + # every dequeue leaves through the same exit.

919 + thread = _make_store_sending_thread(MagicMock())

920 + req = _make_store_req("req-a", [b"a0", b"a1"])

921 + req.token_len_chunk = None # type: ignore[assignment]

922 +

923 + with contextlib.suppress(TypeError):

924 + _run_store_req(thread, req)

925

920 - thread.add_stored_request("req-a")

921 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

926 + assert thread.take_completed_saves() == {req.store_job_id: 1}

927 + thread.request_queue.task_done.assert_called_once()

928

923 - assert thread.stored_requests["req-a"] == 0

929 +

930 +def test_stale_store_job_cannot_touch_a_reused_request_id():

931 + # A preempted request resumes under its original id, so a job left over from

932 + # the retired generation carries a req_id that now belongs to a live one.

933 + thread = _make_store_sending_thread(MagicMock())

934 + stale = _make_store_req("req-a", [b"a0", b"a1"])

935 + stale.store_job_id = 1

936 + thread.add_request(stale)

937 + thread._record_saved(stale, 32)

938 + thread.delete_finished_stored_request("req-a")

939 +

940 + live = _make_store_req("req-a", [b"a0", b"a1"])

941 + live.store_job_id = 2

942 + thread.add_request(live)

943 +

944 + thread.finish_store_job(stale)

945 + thread._record_saved(stale, 64)

946 + thread._mark_request_skipped_for_pressure(stale)

947 +

948 + assert thread.is_live_store_job(live)

949 + assert not thread.is_live_store_job(stale)

950 + assert thread._saved_offset.get("req-a") is None

951 + assert "req-a" not in thread._skip_store_requests

952

953

954 def test_store_recving_thread_reports_failed_block_ids():

@@ -991,8 +1019,7 @@ def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set

1019 replicate_config = SimpleNamespace(preferred_segment="10.0.0.7:50053")

1020 thread = _make_store_sending_thread(store, replicate_config=replicate_config)

1021

994 - thread.add_stored_request("req-a")

995 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1022 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1023

1024 assert store.batch_put_from_multi_buffers.call_count == 1

1025 call_args = store.batch_put_from_multi_buffers.call_args.args

@@ -1010,8 +1037,7 @@ def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_

1037 replicate_config = SimpleNamespace()

1038 thread = _make_store_sending_thread(store, replicate_config=replicate_config)

1039

1013 - thread.add_stored_request("req-a")

1014 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1040 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1041

1042 assert store.batch_put_from_multi_buffers.call_count == 1

1043 call_args = store.batch_put_from_multi_buffers.call_args.args

@@ -1067,8 +1093,7 @@ def test_store_sending_thread_sets_group_ids_when_enabled():

1093 supports_group_ids=True,

1094 )

1095

1070 - thread.add_stored_request("req-a")

1071 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1096 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1097

1098 assert store.batch_put_from_multi_buffers.call_count == 1

1099 keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args

@@ -1092,8 +1117,7 @@ def test_store_sending_thread_leaves_group_ids_unchanged_when_flag_disabled():

1117 supports_group_ids=True,

1118 )

1119

1095 - thread.add_stored_request("req-a")

1096 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1120 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1121

1122 assert store.batch_put_from_multi_buffers.call_count == 1

1123 assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config

@@ -1112,8 +1136,7 @@ def test_store_sending_thread_leaves_group_ids_unchanged_when_unsupported():

1136 supports_group_ids=False,

1137 )

1138

1115 - thread.add_stored_request("req-a")

1116 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1139 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1140

1141 assert store.batch_put_from_multi_buffers.call_count == 1

1142 assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config

@@ -1180,8 +1203,7 @@ def test_store_sending_thread_group_id_excludes_physical_sharding():

1203 supports_group_ids=True,

1204 )

1205

1183 - thread.add_stored_request("req-a")

1184 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1206 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1207

1208 keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args

1209 assert keys == [

@@ -1210,8 +1232,7 @@ def test_store_sending_thread_multiple_segments_share_logical_group_id():

1232 supports_group_ids=True,

1233 )

1234

1213 - thread.add_stored_request("req-a")

1214 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1235 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1236

1237 keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args

1238 assert keys == [

@@ -1260,8 +1281,7 @@ def test_store_sending_thread_group_ids_share_across_kv_cache_groups():

1281 supports_group_ids=True,

1282 )

1283

1263 - thread.add_stored_request("req-a")

1264 - thread._handle_request(_make_multi_group_store_req("req-a", [b"a0", b"a1"]))

1284 + _run_store_req(thread, _make_multi_group_store_req("req-a", [b"a0", b"a1"]))

1285

1286 keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args

1287 assert keys == [

@@ -1296,8 +1316,7 @@ def test_store_sending_thread_group_ids_follow_missing_key_filter():

1316 supports_group_ids=True,

1317 )

1318

1299 - thread.add_stored_request("req-a")

1300 - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))

1319 + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"]))

1320

1321 keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args

1322 assert keys == ["test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131"]

@@ -1866,15 +1885,15 @@ def test_store_sending_thread_clamps_token_len_to_lcm():

1885 # token_len_chunk=33 clamps to 32 → 2 chunks (not 3 with a partial 1-token chunk).

1886 thread = _make_store_sending_thread(store)

1887

1869 - thread.add_stored_request("r0")

1870 - thread._handle_request(

1888 + _run_store_req(

1889 + thread,

1890 ReqMeta(

1891 req_id="r0",

1892 token_len_chunk=33,

1893 block_ids=([0, 1, 2],),

1894 block_hashes=[b"a0", b"a1", b"a2"],

1895 can_save=True,

1877 - )

1896 + ),

1897 )

1898

1899 keys = store.batch_put_from_multi_buffers.call_args.args[0]

@@ -1905,20 +1924,19 @@ def test_store_sending_thread_skips_when_token_len_below_lcm():

1924 store, coord=coord, token_databases=[db], block_size=64

1925 )

1926

1908 - thread.add_stored_request("r0")

1909 - thread._handle_request(

1927 + _run_store_req(

1928 + thread,

1929 ReqMeta(

1930 req_id="r0",

1931 token_len_chunk=32,

1932 block_ids=([0, 1],),

1933 block_hashes=[b"a0", b"a1"],

1934 can_save=True,

1916 - )

1935 + ),

1936 )

1937

1938 store.batch_is_exist.assert_not_called()

1939 store.batch_put_from_multi_buffers.assert_not_called()

1921 - assert thread.stored_requests["r0"] == 0

1940

1941

1942 def test_store_sending_thread_only_stores_swa_blocks_in_window():

@@ -1982,15 +2000,15 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():

2000 )

2001

2002 hs = [bytes([i + 1]) * 4 for i in range(8)]

1985 - thread.add_stored_request("r0")

1986 - thread._handle_request(

2003 + _run_store_req(

2004 + thread,

2005 ReqMeta(

2006 req_id="r0",

2007 token_len_chunk=64,

2008 block_ids=([0, 1], list(range(8))),

2009 block_hashes=hs,

2010 can_save=True,

1993 - )

2011 + ),

2012 )

2013

2014 keys = store.batch_put_from_multi_buffers.call_args.args[0]

@@ -2057,16 +2075,16 @@ def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks():

2075 )

2076

2077 hs = [bytes([i + 1]) * 4 for i in range(8)]

2060 - thread.add_stored_request("r0")

2078 thread._saved_offset["r0"] = 32

2062 - thread._handle_request(

2079 + _run_store_req(

2080 + thread,

2081 ReqMeta(

2082 req_id="r0",

2083 token_len_chunk=64,

2084 block_ids=([0, 1], list(range(8))),

2085 block_hashes=hs,

2086 can_save=True,

2069 - )

2087 + ),

2088 )

2089

2090 keys = store.batch_put_from_multi_buffers.call_args.args[0]

@@ -2128,8 +2146,8 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():

2146 thread.enable_kv_event = True

2147

2148 hs = [bytes([i + 1]) * 4 for i in range(4)]

2131 - thread.add_stored_request("r0")

2132 - thread._handle_request(

2149 + _run_store_req(

2150 + thread,

2151 ReqMeta(

2152 req_id="r0",

2153 token_len_chunk=32,

@@ -2137,7 +2155,7 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():

2155 block_hashes=hs,

2156 can_save=True,

2157 token_ids=list(range(32)),

2140 - )

2158 + ),

2159 )

2160

2161 full_event, swa_event = thread.get_kv_events()

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

+20 -2

@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

26 KVConnectorBase_V1,

27 KVConnectorMetadata,

28 KVConnectorRole,

29 + KVConnectorWorkerMetadata,

30 SupportsHMA,

31 )

32 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (

@@ -37,6 +38,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (

38 from vllm.forward_context import ForwardContext

39 from vllm.logger import init_logger

40 from vllm.v1.attention.backend import AttentionMetadata

41 +from vllm.v1.core.block_pool import BlockPool

42 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

43 from vllm.v1.core.sched.output import SchedulerOutput

44 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -201,6 +203,14 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

203 request, blocks, num_external_tokens

204 )

205

206 + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:

207 + assert self.connector_scheduler is not None

208 + self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool)

209 +

210 + def has_pending_push_work(self) -> bool:

211 + assert self.connector_scheduler is not None

212 + return self.connector_scheduler.has_pending_push_work()

213 +

214 def build_connector_meta(

215 self,

216 scheduler_output: SchedulerOutput,

@@ -208,6 +218,10 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

218 assert self.connector_scheduler is not None

219 return self.connector_scheduler.build_connector_meta(scheduler_output)

220

221 + def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:

222 + assert self.connector_worker is not None

223 + return self.connector_worker.build_connector_worker_meta()

224 +

225 def request_finished(

226 self,

227 request: Request,

@@ -220,8 +234,9 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

234 request: Request,

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

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

223 - assert self.connector_scheduler is not None

224 - return self.connector_scheduler.request_finished(request, block_ids)

237 + # An in-flight store job holds its own reference on the blocks it reads,

238 + # so a finishing request never has to defer freeing them.

239 + return False, None

240

241 def reset_cache(self) -> bool | None:

242 """Reset the external Mooncake store on prefix-cache reset.

@@ -241,6 +256,9 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):

256 return None

257

258 def update_connector_output(self, connector_output: KVConnectorOutput):

259 + assert self.connector_scheduler is not None

260 + self.connector_scheduler.update_connector_output(connector_output)

261 +

262 kv_cache_events = connector_output.kv_cache_events

263 if not kv_cache_events or not isinstance(

264 kv_cache_events, MooncakeStoreKVEvents

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

+23 -1

@@ -6,7 +6,7 @@

6 """Data classes for MooncakeStoreConnector."""

8 from collections.abc import Iterable, Sequence

9 -from dataclasses import dataclass

9 +from dataclasses import dataclass, field

10 from typing import cast

12 import numpy as np

@@ -14,6 +14,7 @@ import torch

15 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

16 KVConnectorMetadata,

17 + KVConnectorWorkerMetadata,

18 )

19 from vllm.logger import init_logger

20 from vllm.utils.math_utils import cdiv

@@ -366,6 +367,10 @@ class ReqMeta:

367

368 token_ids: list[int] | None = None

369 num_prompt_tokens: int | None = None

370 + # Identifies this store job for the engine's lifetime. A request id cannot

371 + # serve that purpose: it is reused once a preempted request resumes, so it

372 + # would release the wrong job's blocks.

373 + store_job_id: int | None = None

374 # Core-provided per-mamba-group

375 # (group_id, cow_block_id, boundary_tokens) for this request's partial tail.

376 # Present only on the producer's CoW step; drives the connector's offload

@@ -434,6 +439,23 @@ class ReqMeta:

439 )

440

441

442 +@dataclass

443 +class MooncakeStoreWorkerMetadata(KVConnectorWorkerMetadata):

444 + """Maps ``ReqMeta.store_job_id`` to the number of ranks done with that job."""

445 +

446 + completed_saves: dict[int, int] = field(default_factory=dict)

447 +

448 + def aggregate(

449 + self, other: "KVConnectorWorkerMetadata"

450 + ) -> "MooncakeStoreWorkerMetadata":

451 + assert isinstance(other, MooncakeStoreWorkerMetadata)

452 + for store_job_id, count in other.completed_saves.items():

453 + self.completed_saves[store_job_id] = (

454 + self.completed_saves.get(store_job_id, 0) + count

455 + )

456 + return self

457 +

458 +

459 class MooncakeStoreConnectorMetadata(KVConnectorMetadata):

460 """Metadata passed from scheduler to worker."""

461

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

+78 -26

@@ -5,8 +5,6 @@

5 # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).

6 """Scheduler-side logic for MooncakeStoreConnector."""

8 -from typing import Any

9 -

8 from vllm.config import VllmConfig

9 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

10 KVConnectorMetadata,

@@ -17,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp

15 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501

16 LoadSpec,

17 MooncakeStoreConnectorMetadata,

18 + MooncakeStoreWorkerMetadata,

19 ReqMeta,

20 RequestTracker,

21 )

@@ -24,10 +23,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (

23 LookupKeyClient,

24 )

25 from vllm.logger import init_logger

26 +from vllm.v1.core.block_pool import BlockPool

27 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

28 from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes

29 from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput

30 from vllm.v1.kv_cache_interface import KVCacheConfig

31 +from vllm.v1.outputs import KVConnectorOutput

32 from vllm.v1.request import Request

34 logger = init_logger(__name__)

@@ -78,6 +79,15 @@ class MooncakeStoreScheduler:

79 self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}

80 self._unfinished_request_ids: set[str] = set()

82 + self._gpu_block_pool: BlockPool | None = None

83 + self._num_workers = vllm_config.parallel_config.world_size

84 + self._next_store_job_id = 0

85 + # store_job_id -> (referenced block ids, ranks yet to report completion)

86 + self._pinned_saves: dict[int, tuple[list[int], int]] = {}

87 +

88 + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:

89 + self._gpu_block_pool = gpu_block_pool

90 +

91 def get_num_new_matched_tokens(

92 self,

93 request: Request,

@@ -392,33 +402,75 @@ class MooncakeStoreScheduler:

402 )

403 )

404

405 + self._reference_save_blocks(meta)

406 return meta

407

397 - def request_finished(

398 - self,

399 - request: Request,

400 - block_ids: tuple[list[int], ...],

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

402 - """Determine whether to delay freeing blocks for async save."""

403 - if self.kv_role == "kv_consumer":

404 - return False, None

405 - tracker = self._request_trackers.get(request.request_id)

406 - # Missing tracker can happen when the request is aborted before the

407 - # connector observes the normal finished lifecycle or is preempted

408 - # before finishing.

409 - if tracker is None or (

410 - tracker.num_saved_tokens <= 0 and not tracker.has_pending_offload

411 - ):

412 - return False, None

413 - total_blocks = sum(len(g) for g in block_ids)

414 - delay_free_blocks = total_blocks > 0

415 - if delay_free_blocks:

416 - logger.debug(

417 - "Delaying free of %d blocks for request %s",

418 - total_blocks,

419 - request.request_id,

408 + def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None:

409 + """Take a GPU block reference for every store job this step emits.

410 +

411 + The worker DMAs out of these blocks after the step that scheduled them,

412 + so a reference keeps them out of the free queue even once the request

413 + itself is freed, until every rank reports the job done.

414 + """

415 + pool = self._gpu_block_pool

416 + for req_meta in meta.requests:

417 + if not req_meta.can_save:

418 + continue

419 + assert pool is not None, (

420 + "GPU block pool must be bound before any store job is emitted"

421 + )

422 + req_meta.store_job_id = store_job_id = self._next_store_job_id

423 + self._next_store_job_id += 1

424 + block_ids: list[int] = []

425 + if req_meta.partial_tail_offloads:

426 + # A partial-tail CoW block is deliberately kept out of the

427 + # request's block table, so it is absent from `block_ids` even

428 + # though the worker DMAs out of it just as asynchronously.

429 + # It leads the list, as in `pop_blocks_for_free`.

430 + block_ids += [bid for _, bid, _ in req_meta.partial_tail_offloads]

431 + # Every allocated block is referenced, not just the ones covering

432 + # this job's token range: a rank resumes from its own last

433 + # successful offset, which lags the scheduler's whenever a save was

434 + # skipped or failed, so it may read anywhere below the range.

435 + block_ids += [bid for group in req_meta.block_ids for bid in group]

436 + if not block_ids:

437 + continue

438 + self._pinned_saves[store_job_id] = (block_ids, self._num_workers)

439 + pool.touch([pool.blocks[bid] for bid in block_ids])

440 +

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

442 + """Drop the block references of store jobs every rank has finished."""

443 + meta = connector_output.kv_connector_worker_meta

444 + if not isinstance(meta, MooncakeStoreWorkerMetadata):

445 + return

446 + pool = self._gpu_block_pool

447 + assert pool is not None

448 + for store_job_id, count in meta.completed_saves.items():

449 + pinned = self._pinned_saves.get(store_job_id)

450 + if pinned is None:

451 + # The job referenced no blocks, so nothing was recorded for it.

452 + continue

453 + block_ids, remaining = pinned

454 + remaining -= count

455 + if remaining > 0:

456 + self._pinned_saves[store_job_id] = (block_ids, remaining)

457 + continue

458 + assert remaining == 0, (

459 + f"store job {store_job_id} reported by too many ranks"

460 )

421 - return delay_free_blocks, None

461 + del self._pinned_saves[store_job_id]

462 + # Tail-first, as elsewhere, so the shared prefix is evicted last.

463 + pool.free_blocks(pool.blocks[bid] for bid in reversed(block_ids))

464 +

465 + def has_pending_push_work(self) -> bool:

466 + """Keep the engine stepping while any store job still holds block refs.

467 +

468 + Completions only reach the scheduler as worker metadata on a step, so an

469 + engine that quiesced with jobs in flight would leave those references

470 + held indefinitely. Nothing else keeps it alive now that a finishing

471 + request no longer defers its own free.

472 + """

473 + return bool(self._pinned_saves)

474

475 def reset_store(self) -> bool:

476 """Trigger a global ``remove_all(force=True)`` on the Mooncake master.

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+104 -60

@@ -17,7 +17,6 @@ import queue

17 import socket

18 import threading

19 import time

20 -from collections import defaultdict

20 from collections.abc import Callable, Sequence

21 from concurrent.futures import Future, ThreadPoolExecutor

22 from dataclasses import dataclass

@@ -46,6 +45,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (

45 ChunkedTokenDatabase,

46 KeyMetadata,

47 MooncakeStoreConnectorMetadata,

48 + MooncakeStoreWorkerMetadata,

49 PoolKey,

50 ReqMeta,

51 )

@@ -506,7 +506,16 @@ class KVCacheStoreSendingThread(KVTransferThread):

506 self.group_put_steps = group_put_steps

507 self.coord = coord

508 self.kv_role = kv_role

509 - self.stored_requests: defaultdict[str, int] = defaultdict(int)

509 + # req_id -> ids of its store jobs that are still queued or running.

510 + # Keying by store_job_id, which never repeats for the engine's lifetime,

511 + # rather than counting jobs per request id makes the ledger immune to id

512 + # reuse across preemption: a job left over from a retired generation is

513 + # missing from the set its resumed generation builds, so it can no longer

514 + # retire that generation, rewind its resume offset, or mark it skipped.

515 + self.stored_requests: dict[str, set[int]] = {}

516 + # store_job_id -> times this rank finished with it, drained every step

517 + # so the scheduler can release the blocks it referenced for those jobs.

518 + self._completed_saves: dict[int, int] = {}

519 self.enable_kv_event = enable_kv_event

520 # Caller always passes a non-None ReplicateConfig — see

521 # MooncakeStoreWorker.__init__ where store_replicate_config is built.

@@ -522,14 +531,20 @@ class KVCacheStoreSendingThread(KVTransferThread):

531 # batch resumes here, so pressure-skipped or failed ranges are retried.

532 self._saved_offset: dict[str, int] = {}

533

525 - def add_stored_request(self, req_id: str):

534 + def add_request(self, request: ReqMeta) -> None:

535 + # Register before enqueueing so a job is never picked up unledgered.

536 + assert request.store_job_id is not None

537 with self.done_task_lock:

527 - self.stored_requests[req_id] += 1

538 + self.stored_requests.setdefault(request.req_id, set()).add(

539 + request.store_job_id

540 + )

541 + super().add_request(request)

542

529 - def dec_stored_request(self, req_id: str):

543 + def is_live_store_job(self, req_meta: ReqMeta) -> bool:

544 with self.done_task_lock:

531 - if req_id in self.stored_requests:

532 - self.stored_requests[req_id] -= 1

545 + return req_meta.store_job_id in self.stored_requests.get(

546 + req_meta.req_id, ()

547 + )

548

549 def delete_finished_stored_request(self, req_id: str):

550 with self.done_task_lock:

@@ -538,21 +553,51 @@ class KVCacheStoreSendingThread(KVTransferThread):

553 self._skip_store_requests.discard(req_id)

554 self._saved_offset.pop(req_id, None)

555

541 - def _record_saved(self, req_id: str, token_len: int) -> None:

542 - # Guard on liveness so a concurrent finish/preempt pop isn't recreated.

556 + def finish_store_job(self, req_meta: ReqMeta) -> None:

557 + """Retire a job from the ledger and report its blocks as no longer read.

558 +

559 + Every path out of a job must reach this, skips and failures included: a

560 + job that never reports leaves its blocks referenced for the rest of the

561 + run. The discard is a no-op for a job whose generation already retired.

562 + """

563 + store_job_id = req_meta.store_job_id

564 + assert store_job_id is not None, (

565 + "a queued store job always carries a store_job_id"

566 + )

567 with self.done_task_lock:

544 - if req_id in self.stored_requests:

545 - self._saved_offset[req_id] = token_len

568 + live = self.stored_requests.get(req_meta.req_id)

569 + if live is not None:

570 + live.discard(store_job_id)

571 + self._completed_saves[store_job_id] = (

572 + self._completed_saves.get(store_job_id, 0) + 1

573 + )

574 +

575 + def take_completed_saves(self) -> dict[int, int]:

576 + with self.done_task_lock:

577 + completed = self._completed_saves

578 + self._completed_saves = {}

579 + return completed

580 +

581 + def _record_saved(self, req_meta: ReqMeta, token_len: int) -> None:

582 + # Guard on job liveness so neither a concurrent finish/preempt pop nor a

583 + # stale job's offset is written back over the live generation's.

584 + with self.done_task_lock:

585 + if req_meta.store_job_id in self.stored_requests.get(req_meta.req_id, ()):

586 + self._saved_offset[req_meta.req_id] = token_len

587

588 def _should_skip_request(self, req_id: str) -> bool:

589 with self.done_task_lock:

590 return self._store_pressure_active and req_id in self._skip_store_requests

591

551 - def _mark_request_skipped_for_pressure(self, req_id: str) -> bool:

592 + def _mark_request_skipped_for_pressure(self, req_meta: ReqMeta) -> bool:

593 + req_id = req_meta.req_id

594 with self.done_task_lock:

595 already_skipped = req_id in self._skip_store_requests

596 self._store_pressure_active = True

555 - self._skip_store_requests.add(req_id)

597 + # The pressure itself is global, but only a live job may sentence its

598 + # own request to being skipped.

599 + if req_meta.store_job_id in self.stored_requests.get(req_id, ()):

600 + self._skip_store_requests.add(req_id)

601 return already_skipped

602

603 def _clear_store_pressure(self) -> bool:

@@ -727,7 +772,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

772 failed_codes,

773 )

774 if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes:

730 - self._mark_request_skipped_for_pressure(req_meta.req_id)

775 + self._mark_request_skipped_for_pressure(req_meta)

776 return False

777

778 if self._clear_store_pressure():

@@ -738,22 +783,20 @@ class KVCacheStoreSendingThread(KVTransferThread):

783 return True

784

785 def _handle_request(self, req_meta: ReqMeta):

741 - # Cache hits are always a multiple of ``lcm_block_size`` tokens, which

742 - # is also ``store_mask``'s precondition.

743 - lcm_block_size = self.coord.lcm_block_size

744 - token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size

745 - block_ids_per_group = req_meta.block_ids

746 - req_id = req_meta.req_id

747 - current_event = req_meta.current_event

748 -

749 - if req_id not in self.stored_requests:

750 - self.request_queue.task_done()

751 - return

752 -

753 - # Decrement the in-flight counter and signal task_done() in `finally`

754 - # so the scheduler can release the GPU blocks it pinned for this

755 - # request (via `delay_free_blocks`) even when the store path raises.

786 + # The single `finally` is the only way out, so the scheduler releases

787 + # this job's GPU block references however the job ends.

788 try:

789 + # Cache hits are always a multiple of ``lcm_block_size`` tokens,

790 + # which is also ``store_mask``'s precondition.

791 + lcm_block_size = self.coord.lcm_block_size

792 + token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size

793 + block_ids_per_group = req_meta.block_ids

794 + req_id = req_meta.req_id

795 + current_event = req_meta.current_event

796 +

797 + if not self.is_live_store_job(req_meta):

798 + return

799 +

800 if self._should_skip_request(req_id):

801 logger.debug(

802 "Skipping Mooncake store for request %s while CPU/disk "

@@ -808,7 +851,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

851 group_indices.append(g_idx)

852

853 if not keys:

811 - self._record_saved(req_id, token_len)

854 + self._record_saved(req_meta, token_len)

855 return

856

857 # Check which blocks already exist (dedup)

@@ -834,7 +877,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

877 ]

878

879 if not missing_indices:

837 - self._record_saved(req_id, token_len)

880 + self._record_saved(req_meta, token_len)

881 return

882

883 if len(missing_indices) != len(keys):

@@ -954,7 +997,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

997 )

998 if (

999 MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes

957 - and not self._mark_request_skipped_for_pressure(req_id)

1000 + and not self._mark_request_skipped_for_pressure(req_meta)

1001 ):

1002 logger.warning(

1003 "Detected Mooncake CPU/disk offloading pressure "

@@ -964,7 +1007,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

1007 req_id,

1008 )

1009 else:

967 - self._record_saved(req_id, token_len)

1010 + self._record_saved(req_meta, token_len)

1011 if self._clear_store_pressure():

1012 logger.info(

1013 "Mooncake CPU/disk offloading pressure cleared "

@@ -984,7 +1027,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

1027 if self.enable_kv_event and stored_events:

1028 self.update_kv_event(stored_events)

1029 finally:

987 - self.dec_stored_request(req_id)

1030 + self.finish_store_job(req_meta)

1031 self.request_queue.task_done()

1032

1033

@@ -1670,16 +1713,15 @@ class MooncakeStoreWorker:

1713 continue

1714 request.current_event = current_event

1715 assert self.kv_send_thread is not None

1673 - self.kv_send_thread.add_stored_request(request.req_id)

1716 self.kv_send_thread.add_request(request)

1717

1676 - # Check completion of previously queued transfers

1677 - done_sending = (

1678 - self._get_and_clear_finished_sending(finished_req_ids, meta)

1679 - if self.kv_role in ["kv_producer", "kv_both"]

1680 - else set()

1681 - )

1718 + if self.kv_role in ["kv_producer", "kv_both"]:

1719 + self._close_ended_store_requests(finished_req_ids, meta)

1720

1721 + # Blocks read by a store job are released by the scheduler when the job

1722 + # reports back (see build_connector_worker_meta), so no request ever waits

1723 + # on a `finished_sending` signal to get its blocks back.

1724 + done_sending: set[str] = set()

1725 done_recving: set[str] = set()

1726 if self.load_async:

1727 for recv_thread in self.kv_recv_threads:

@@ -1727,35 +1769,37 @@ class MooncakeStoreWorker:

1769 self.kv_connector_stats = MooncakeStoreConnectorStats()

1770 return kv_connector_stats

1771

1730 - def _get_and_clear_finished_sending(

1772 + def _close_ended_store_requests(

1773 self,

1774 finished_req_ids: set[str],

1775 meta: MooncakeStoreConnectorMetadata,

1734 - ) -> set[str]:

1776 + ) -> None:

1777 + """Retire the ledger entries of requests that finished or were preempted.

1778 +

1779 + An entry may only go once its jobs have drained, because they still read

1780 + the resume offset it owns; a request that comes back after preemption

1781 + then saves from the start rather than from where the last attempt got to.

1782 + """

1783 assert self.kv_send_thread is not None

1736 - finished_sending: set[str] = set()

1784

1785 for req_id in meta.preempted_req_ids:

1786 self.kv_send_thread.delete_finished_stored_request(req_id)

1787

1741 - for req_id in self.kv_send_thread.stored_requests.copy():

1742 - if (

1743 - self.kv_send_thread.stored_requests[req_id] == 0

1744 - and req_id in self.finished_store_req

1745 - ):

1746 - self.finished_store_req.remove(req_id)

1747 - finished_sending.add(req_id)

1748 - self.kv_send_thread.delete_finished_stored_request(req_id)

1749 -

1750 - for req_id in finished_req_ids:

1751 - req_remain_jobs = self.kv_send_thread.stored_requests.get(req_id)

1752 - if req_remain_jobs == 0:

1753 - finished_sending.add(req_id)

1754 - self.kv_send_thread.delete_finished_stored_request(req_id)

1755 - elif req_remain_jobs is not None:

1788 + for req_id in finished_req_ids | self.finished_store_req:

1789 + if self.kv_send_thread.stored_requests.get(req_id):

1790 + # Queued jobs still need the resume offset; retire on a later step.

1791 self.finished_store_req.add(req_id)

1792 + else:

1793 + self.finished_store_req.discard(req_id)

1794 + self.kv_send_thread.delete_finished_stored_request(req_id)

1795

1758 - return finished_sending

1796 + def build_connector_worker_meta(self) -> MooncakeStoreWorkerMetadata | None:

1797 + if self.kv_send_thread is None:

1798 + return None

1799 + completed_saves = self.kv_send_thread.take_completed_saves()

1800 + if not completed_saves:

1801 + return None

1802 + return MooncakeStoreWorkerMetadata(completed_saves=completed_saves)

1803

1804 def lookup(self, num_tokens: int, block_hashes: Sequence[BlockHash]) -> int:

1805 """Check how many prefix tokens exist in the store.