@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7cb9a88f50

[Bugfix][Mooncake] Offload producer partial tails on request finish (#52832)

Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: OpenAI Codex <noreply@openai.com>

Dao007forever committed Sep 1, 2026 at 13:23 UTC 7cb9a88f50088453bbcf86704c376103dcd6b699
10 files changed +375 -10
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+111
@@ -12,6 +12,7 @@ import pytest
12 import torch
13
14 from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request
15 +from vllm.distributed.kv_transfer.kv_connector.v1.base import SupportsHMA
16 from vllm.utils.hashing import sha256
17 from vllm.v1.core.kv_cache_utils import (
18 KVCacheBlockCopy,
@@ -735,11 +736,121 @@ def test_boundary_state_offloads_returns_cow_target():
736 assert block_id != source_block_id
737 # Draining clears it.
738 assert manager.take_boundary_state_offloads() == {}
739 + # CoW consumed the producer marker, so finishing cannot emit the same
740 + # boundary a second time.
741 + assert manager.finalize_partial_tail_offloads(req0) == []
742
743 manager.block_pool.free_blocks(retained)
744 manager.free(req0)
745
746
747 +def test_finished_partial_tail_uses_table_source_once():
748 + """A finish without another allocation hands off the unchanged table block."""
749 + hash_block_size = 2
750 + block_size = 2 * hash_block_size
751 + kv_cache_config = KVCacheConfig(
752 + num_blocks=24,
753 + kv_cache_tensors=[],
754 + kv_cache_groups=[
755 + KVCacheGroupSpec(
756 + ["full"],
757 + FullAttentionSpec(
758 + block_size=hash_block_size,
759 + num_kv_heads=1,
760 + head_size=1,
761 + dtype=torch.float32,
762 + ),
763 + ),
764 + KVCacheGroupSpec(
765 + ["mamba"],
766 + MambaSpec(
767 + block_size=block_size,
768 + shapes=(1, 1),
769 + dtypes=(torch.float32,),
770 + mamba_cache_mode="align",
771 + ),
772 + ),
773 + ],
774 + )
775 + manager = make_kv_cache_manager(
776 + kv_cache_config=kv_cache_config,
777 + max_model_len=8192,
778 + enable_caching=True,
779 + hash_block_size=hash_block_size,
780 + )
781 +
782 + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
783 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
784 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
785 + source_block_id = manager.get_blocks("0").get_block_ids()[1][1]
786 +
787 + req0.num_computed_tokens = 6
788 + offloads = manager.finalize_partial_tail_offloads(req0)
789 +
790 + assert offloads == [(1, source_block_id, 6)]
791 + assert manager.finalize_partial_tail_offloads(req0) == []
792 + assert manager.take_boundary_state_offloads() == {}
793 +
794 + # A table state that advanced beyond the boundary is unsafe and is
795 + # consumed without producing a handoff.
796 + advanced_manager = make_kv_cache_manager(
797 + kv_cache_config=kv_cache_config,
798 + max_model_len=8192,
799 + enable_caching=True,
800 + hash_block_size=hash_block_size,
801 + )
802 + req1 = make_request("1", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
803 + computed_blocks, num_computed, _ = advanced_manager.get_computed_blocks(req1)
804 + assert (
805 + advanced_manager.allocate_slots(req1, 6, num_computed, computed_blocks)
806 + is not None
807 + )
808 + req1.num_computed_tokens = 7
809 + assert advanced_manager.finalize_partial_tail_offloads(req1) == []
810 + assert advanced_manager.finalize_partial_tail_offloads(req1) == []
811 +
812 +
813 +def test_connector_finish_registers_partial_tail_before_cleanup():
814 + class _Connector(SupportsHMA):
815 + def register_finished_partial_tail(self, request, block_ids, offloads):
816 + self.registered = (request, block_ids, offloads)
817 + return True
818 +
819 + def request_finished_all_groups(self, request, block_ids):
820 + return False, None
821 +
822 + scheduler = object.__new__(Scheduler)
823 + scheduler.connector = _Connector()
824 + scheduler.vllm_config = SimpleNamespace(
825 + kv_transfer_config=SimpleNamespace(is_kv_producer=True)
826 + )
827 + scheduler.kv_cache_manager = MagicMock()
828 + scheduler.kv_cache_manager.finalize_partial_tail_offloads.return_value = [
829 + (1, 9, 12)
830 + ]
831 + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.return_value = (
832 + [3],
833 + [9],
834 + )
835 + request = SimpleNamespace(
836 + request_id="0",
837 + num_computed_tokens=12,
838 + num_in_flight_tokens=0,
839 + num_prompt_tokens=12,
840 + )
841 +
842 + delay_free, kv_xfer_params = scheduler._connector_finished(request)
843 +
844 + assert delay_free is True
845 + assert kv_xfer_params is None
846 + assert scheduler.connector.registered == (
847 + request,
848 + ([3], [9]),
849 + [(1, 9, 12)],
850 + )
851 + scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once()
852 +
853 +
854 def test_block_pool_touch_pins_released_cow_target():
855 """The connector can rescue an offered CoW target after its step-scoped
856 retention is released by using the bound BlockPool's touch method."""
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+60
@@ -42,6 +42,7 @@ def _make_bare_scheduler(
42 scheduler._unfinished_request_ids = {"req-0"}
43 scheduler._unfinished_requests = {}
44 scheduler._request_trackers = {}
45 + scheduler._finished_partial_tail_metas = {}
46 scheduler._gpu_block_pool = BlockPool(
47 num_gpu_blocks=64, enable_caching=True, hash_block_size=hash_block_size
48 )
@@ -971,6 +972,65 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():
972 assert tracker.has_pending_offload is True
973
974
975 +def test_finished_partial_tail_is_pre_pinned_as_store_job():
976 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
977 + scheduler.client = SimpleNamespace(discard=lambda *_: None)
978 + request = SimpleNamespace(
979 + request_id="req-0",
980 + block_hashes=[b"h0", b"h1", b"h2"],
981 + )
982 + scheduler._request_trackers["req-0"] = RequestTracker(
983 + req_id="req-0",
984 + token_len=12,
985 + allocated_block_ids=([3], [9]),
986 + token_ids=list(range(12)),
987 + prefill_end_tokens=12,
988 + )
989 + block_ids = ([3], [9])
990 +
991 + delay_free = scheduler.register_finished_partial_tail(
992 + request,
993 + block_ids,
994 + [(1, 9, 12)],
995 + )
996 +
997 + # The exact source is pinned immediately, so the request can be freed
998 + # before the next connector metadata build.
999 + assert delay_free is False
1000 + assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1
1001 + assert scheduler._gpu_block_pool.blocks[3].ref_cnt == 0
1002 +
1003 + out = SimpleNamespace(
1004 + finished_req_ids={"req-0"},
1005 + preempted_req_ids=set(),
1006 + scheduled_new_reqs=[],
1007 + scheduled_cached_reqs=SimpleNamespace(
1008 + req_ids=[],
1009 + new_block_ids=[],
1010 + num_computed_tokens=[],
1011 + resumed_req_ids=set(),
1012 + ),
1013 + num_scheduled_tokens={},
1014 + scheduled_spec_decode_tokens={},
1015 + kv_connector_block_state=_make_connector_block_state(),
1016 + )
1017 + meta = scheduler.build_connector_meta(out)
1018 +
1019 + assert len(meta.requests) == 1
1020 + req_meta = meta.requests[0]
1021 + assert req_meta.req_id == "req-0"
1022 + assert req_meta.token_len_chunk == 0
1023 + assert req_meta.block_ids == block_ids
1024 + assert req_meta.block_hashes == request.block_hashes
1025 + assert req_meta.boundary_state_offloads == [(1, 9, 12)]
1026 + assert scheduler._pinned_saves[req_meta.store_job_id][0] == [9]
1027 + assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1
1028 + assert scheduler._finished_partial_tail_metas == {}
1029 +
1030 + scheduler.update_connector_output(_make_worker_output({req_meta.store_job_id: 1}))
1031 + assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 0
1032 +
1033 +
1034 def test_decode_boundary_state_offload_dropped_unclaimed():
1035 # A hand-off past the prefill end can never complete a joint hybrid hit
1036 # (every other group stops saving there), so it is neither transferred nor
tests/v1/kv_connector/unit/test_multi_connector.py
+20
@@ -148,6 +148,26 @@ KVConnectorFactory.register_connector(
148 )
149
150
151 +def test_register_finished_partial_tail_notifies_every_connector():
152 + connector = object.__new__(MultiConnector)
153 + first = MagicMock(spec_set=KVConnectorBase_V1)
154 + second = MagicMock(spec_set=KVConnectorBase_V1)
155 + first.register_finished_partial_tail.return_value = True
156 + second.register_finished_partial_tail.return_value = False
157 + connector._connectors = [first, second]
158 + request = MagicMock()
159 + block_ids = ([1], [2])
160 + offloads = [(1, 2, 12)]
161 +
162 + assert connector.register_finished_partial_tail(request, block_ids, offloads)
163 + first.register_finished_partial_tail.assert_called_once_with(
164 + request, block_ids, offloads
165 + )
166 + second.register_finished_partial_tail.assert_called_once_with(
167 + request, block_ids, offloads
168 + )
169 +
170 +
171 @pytest.fixture
172 def mc() -> MultiConnector:
173 """MultiConnector using two mocked connectors"""
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+13
@@ -561,6 +561,19 @@ class KVConnectorBase_V1(ABC):
561 """
562 return False, None
563
564 + def register_finished_partial_tail(
565 + self,
566 + request: "Request",
567 + block_ids: tuple[list[int], ...],
568 + partial_tail_offloads: list[tuple[int, int, int]],
569 + ) -> bool:
570 + """Register finish-time partial-tail sources before block cleanup.
571 +
572 + Returns True when the connector accepts responsibility for the sources
573 + and the request's blocks must remain alive until ``get_finished()``.
574 + """
575 + return False
576 +
577 def take_events(self) -> Iterable["KVCacheEvent"]:
578 """
579 Take the KV cache events from the connector.
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+11
@@ -265,6 +265,17 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
265 # so a finishing request never has to defer freeing them.
266 return False, None
267
268 + def register_finished_partial_tail(
269 + self,
270 + request: Request,
271 + block_ids: tuple[list[int], ...],
272 + partial_tail_offloads: list[tuple[int, int, int]],
273 + ) -> bool:
274 + assert self.connector_scheduler is not None
275 + return self.connector_scheduler.register_finished_partial_tail(
276 + request, block_ids, partial_tail_offloads
277 + )
278 +
279 def reset_cache(self) -> bool | None:
280 """Reset the external Mooncake store on prefix-cache reset.
281
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+66
@@ -102,6 +102,7 @@ class MooncakeStoreScheduler:
102 self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests
103 self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}
104 self._unfinished_request_ids: set[str] = set()
105 + self._finished_partial_tail_metas: dict[str, ReqMeta] = {}
106
107 def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
108 self._gpu_block_pool = gpu_block_pool
@@ -402,6 +403,14 @@ class MooncakeStoreScheduler:
403 )
404
405 self._apply_current_save_block_ids(meta, scheduler_output)
406 +
407 + # Finish-time handoffs arrive after the producing step's metadata was
408 + # built. Their exact blocks were pinned when they were registered, so
409 + # they remain valid after request cleanup and need no current snapshot.
410 + for req_meta in self._finished_partial_tail_metas.values():
411 + meta.add_request(req_meta)
412 + self._finished_partial_tail_metas.clear()
413 +
414 self._reference_save_blocks(meta)
415 return meta
416
@@ -440,6 +449,9 @@ class MooncakeStoreScheduler:
449 assert pool is not None, (
450 "GPU block pool must be bound before any store job is emitted"
451 )
452 + if req_meta.store_job_id is not None:
453 + assert req_meta.store_job_id in self._pinned_saves
454 + continue
455 req_meta.store_job_id = store_job_id = self._next_store_job_id
456 self._next_store_job_id += 1
457 block_ids: list[int] = []
@@ -470,6 +482,60 @@ class MooncakeStoreScheduler:
482 self._pinned_saves[store_job_id] = (block_ids, self._num_workers)
483 pool.touch([pool.blocks[block_id] for block_id in block_ids])
484
485 + def register_finished_partial_tail(
486 + self,
487 + request: Request,
488 + block_ids: tuple[list[int], ...],
489 + partial_tail_offloads: list[tuple[int, int, int]],
490 + ) -> bool:
491 + """Queue and pin a finish-time tail for the next connector step."""
492 + if self.kv_role == "kv_consumer" or not partial_tail_offloads:
493 + return False
494 + tracker = self._request_trackers.get(request.request_id)
495 + if tracker is None or not any(block_ids):
496 + return False
497 + boundaries = {boundary for _, _, boundary in partial_tail_offloads}
498 + if len(boundaries) != 1:
499 + raise ValueError(
500 + "Partial-tail offloads for one request must share a boundary"
501 + )
502 + boundary_tokens = next(iter(boundaries))
503 + if boundary_tokens > tracker.prefill_end_tokens:
504 + return False
505 +
506 + pinned_block_ids: list[int] = []
507 + for group_id, block_id, _ in partial_tail_offloads:
508 + if group_id not in self._boundary_state_group_ids:
509 + return False
510 + if block_id == NULL_BLOCK_ID:
511 + return False
512 + pinned_block_ids.append(block_id)
513 + pinned_block_ids = list(dict.fromkeys(pinned_block_ids))
514 +
515 + pool = self._gpu_block_pool
516 + assert pool is not None, (
517 + "GPU block pool must be bound before a finish-time handoff"
518 + )
519 + assert request.request_id not in self._finished_partial_tail_metas
520 + store_job_id = self._next_store_job_id
521 + self._next_store_job_id += 1
522 + self._pinned_saves[store_job_id] = (pinned_block_ids, self._num_workers)
523 + pool.touch([pool.blocks[block_id] for block_id in pinned_block_ids])
524 +
525 + self._finished_partial_tail_metas[request.request_id] = ReqMeta(
526 + req_id=request.request_id,
527 + token_len_chunk=0,
528 + block_ids=tuple(group.copy() for group in block_ids),
529 + block_hashes=list(request.block_hashes),
530 + can_save=True,
531 + num_prompt_tokens=tracker.prefill_end_tokens,
532 + store_job_id=store_job_id,
533 + boundary_state_offloads=partial_tail_offloads,
534 + )
535 + tracker.has_pending_offload = True
536 + # The store job owns exact block refs, so request cleanup need not wait.
537 + return False
538 +
539 def _handle_boundary_state_offloads(
540 self,
541 offloads: dict[str, list[tuple[int, int, int]]],
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+12
@@ -523,6 +523,18 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
523 lambda c: c.request_finished(request, blocks),
524 )
525
526 + def register_finished_partial_tail(
527 + self,
528 + request: "Request",
529 + block_ids: tuple[list[int], ...],
530 + partial_tail_offloads: list[tuple[int, int, int]],
531 + ) -> bool:
532 + accepted = [
533 + c.register_finished_partial_tail(request, block_ids, partial_tail_offloads)
534 + for c in self._connectors
535 + ]
536 + return any(accepted)
537 +
538 def request_finished_all_groups(
539 self,
540 request: "Request",
vllm/v1/core/kv_cache_manager.py
+23
@@ -874,6 +874,29 @@ class KVCacheManager:
874 )
875 return offloads
876
877 + def finalize_partial_tail_offloads(
878 + self, request: Request
879 + ) -> list[tuple[int, int, int]]:
880 + """Consume safe producer partial tails when a request finishes.
881 +
882 + A mamba align table block is a valid boundary source only if no later
883 + token was forwarded. The connector pins and queues the exact table
884 + block before request cleanup, then releases the pin when every worker
885 + reports the store job complete.
886 + """
887 + offloads: list[tuple[int, int, int]] = []
888 + for mgr in self.coordinator.single_type_managers:
889 + finalized = mgr.finalize_partial_tail_offload(
890 + request.request_id,
891 + request.num_computed_tokens,
892 + request.num_in_flight_tokens,
893 + )
894 + if finalized is None:
895 + continue
896 + group_id, block, boundary_tokens = finalized
897 + offloads.append((group_id, block.block_id, boundary_tokens))
898 + return offloads
899 +
900 def new_step_starts(self) -> None:
901 """Notify the coordinator that a new step is starting."""
902 self.coordinator.new_step_starts()
vllm/v1/core/sched/scheduler.py
+22 -3
@@ -2775,6 +2775,13 @@ class Scheduler(SchedulerInterface):
2775 if self.connector is None:
2776 return False, None
2777
2778 + finished_partial_tails: list[tuple[int, int, int]] = []
2779 + kv_transfer_config = self.vllm_config.kv_transfer_config
2780 + if kv_transfer_config is not None and kv_transfer_config.is_kv_producer:
2781 + finished_partial_tails = (
2782 + self.kv_cache_manager.finalize_partial_tail_offloads(request)
2783 + )
2784 +
2785 # Free any out-of-window prefix blocks before we hand the block table to
2786 # the connector, on the processed-token basis (see `allocate_slots`).
2787 self.kv_cache_manager.remove_skipped_blocks(
@@ -2789,6 +2796,13 @@ class Scheduler(SchedulerInterface):
2796 request_id=request.request_id,
2797 num_computed_tokens=request.num_computed_tokens,
2798 )
2799 + partial_tail_delay = False
2800 + if finished_partial_tails:
2801 + partial_tail_delay = self.connector.register_finished_partial_tail(
2802 + request,
2803 + block_ids,
2804 + finished_partial_tails,
2805 + )
2806
2807 if not isinstance(self.connector, SupportsHMA):
2808 # NOTE(Kuntai): We should deprecate this code path after we enforce
@@ -2796,9 +2810,14 @@ class Scheduler(SchedulerInterface):
2810 # Hybrid memory allocator should be already turned off for this
2811 # code path, but let's double-check here.
2812 assert len(self.kv_cache_config.kv_cache_groups) == 1
2799 - return self.connector.request_finished(request, block_ids[0])
2800 -
2801 - return self.connector.request_finished_all_groups(request, block_ids)
2813 + delay_free, kv_xfer_params = self.connector.request_finished(
2814 + request, block_ids[0]
2815 + )
2816 + else:
2817 + delay_free, kv_xfer_params = self.connector.request_finished_all_groups(
2818 + request, block_ids
2819 + )
2820 + return delay_free or partial_tail_delay, kv_xfer_params
2821
2822 def _request_remaining_blocks(self, request: Request) -> int:
2823 """Blocks `request` still needs to allocate to hold its full sequence."""
vllm/v1/core/single_type_kv_cache_manager.py
+37 -7
@@ -407,6 +407,15 @@ class SingleTypeKVCacheManager(ABC):
407 self._pending_boundary_state_offloads = []
408 return pending
409
410 + def finalize_partial_tail_offload(
411 + self,
412 + request_id: str,
413 + num_computed_tokens: int,
414 + num_in_flight_tokens: int,
415 + ) -> tuple[int, KVCacheBlock, int] | None:
416 + """Finalize a producer partial tail when its request finishes."""
417 + return None
418 +
419 def _apply_cow(
420 self,
421 request_id: str,
@@ -1383,10 +1392,10 @@ class MambaManager(SingleTypeKVCacheManager):
1392 # current allocation.
1393 self._num_checkpoint_blocks: dict[str, int] = {}
1394 # Requests that registered their own last-prompt-boundary partial
1386 - # tail (producers). On the next step's CoW the boundary state moves
1387 - # into a private cow_block; we record that block for connector
1388 - # offload (see _pending_boundary_state_offloads).
1389 - self._producer_partial_tail_reqs: dict[str, int] = {}
1395 + # tail (producers). A later CoW hands its private copy to the
1396 + # connector; a request that finishes first hands off this table
1397 + # source directly.
1398 + self._producer_partial_tail_reqs: dict[str, tuple[KVCacheBlock, int]] = {}
1399
1400 @classmethod
1401 def find_longest_cache_hit(
@@ -1749,10 +1758,12 @@ class MambaManager(SingleTypeKVCacheManager):
1758 self.block_pool.move_block_hashes(source_block, cow_block)
1759 self._pending_cow_copies.append((source_block, cow_block))
1760 source_block.ref_cnt += 1
1752 - boundary_tokens = self._producer_partial_tail_reqs.pop(
1761 + producer_tail = self._producer_partial_tail_reqs.pop(
1762 request_id, None
1763 )
1755 - if boundary_tokens is not None:
1764 + if producer_tail is not None:
1765 + marker_block, boundary_tokens = producer_tail
1766 + assert marker_block is source_block
1767 # This CoW preserved a producer's own boundary
1768 # state in cow_block; hand it to the connector for
1769 # partial-tail offload once the copy has run.
@@ -1788,6 +1799,22 @@ class MambaManager(SingleTypeKVCacheManager):
1799 req_blocks.append(block)
1800 req_blocks[block_idx] = self._null_block
1801
1802 + def finalize_partial_tail_offload(
1803 + self,
1804 + request_id: str,
1805 + num_computed_tokens: int,
1806 + num_in_flight_tokens: int,
1807 + ) -> tuple[int, KVCacheBlock, int] | None:
1808 + if self.mamba_cache_mode != "align":
1809 + return None
1810 + producer_tail = self._producer_partial_tail_reqs.pop(request_id, None)
1811 + if producer_tail is None:
1812 + return None
1813 + source_block, boundary_tokens = producer_tail
1814 + if num_in_flight_tokens != 0 or num_computed_tokens != boundary_tokens:
1815 + return None
1816 + return self.kv_cache_group_id, source_block, boundary_tokens
1817 +
1818 def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
1819 if self.mamba_cache_mode == "align":
1820 self._allocated_block_reqs.discard(request_id)
@@ -1894,7 +1921,10 @@ class MambaManager(SingleTypeKVCacheManager):
1921 # in ``source_block`` but the next step's forward overwrites it. The
1922 # upcoming CoW copies it into a durable cow_block; record the req so
1923 # allocate_new_blocks hands that block to the connector for offload.
1897 - self._producer_partial_tail_reqs[request.request_id] = num_tokens
1924 + self._producer_partial_tail_reqs[request.request_id] = (
1925 + source_block,
1926 + num_tokens,
1927 + )
1928 return partial_hash
1929
1930