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

Fork 0

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

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

Synced Nov 2, 20:45 UTC

[KV Connector] Propagate MooncakeStore load failures (#42788)

Signed-off-by: Dao Le <Dao007forever@gmail.com>

Dao007forever committed May 25, 2026 at 22:12 UTC c2a4005c70e8155624e9687de5e83b395f2a2c83

4 files changed +143 -15

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -103,14 +103,18 @@ def test_worker_methods_delegate_to_store_worker():

103

104 worker = mock_worker_cls.return_value

105 worker.get_finished.return_value = ({"req-1"}, {"req-2"})

106 + worker.get_block_ids_with_load_errors.return_value = {3, 4}

107 connector.bind_connector_metadata(metadata)

108

109 connector.register_kv_caches(kv_caches)

110 result = connector.get_finished(finished_req_ids)

111 + invalid_block_ids = connector.get_block_ids_with_load_errors()

112

113 worker.register_kv_caches.assert_called_once_with(kv_caches)

114 worker.get_finished.assert_called_once_with(finished_req_ids, metadata)

115 assert result == ({"req-1"}, {"req-2"})

116 + worker.get_block_ids_with_load_errors.assert_called_once_with()

117 + assert invalid_block_ids == {3, 4}

118

119

120 def test_get_kv_connector_kv_cache_events_returns_none_when_empty():

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+71 -3

@@ -77,6 +77,7 @@ def _make_store_sending_thread(

77 def _make_store_recving_thread(

78 store: MagicMock,

79 *,

80 + tp_rank: int = 0,

81 disk_offload_buffer_budget_bytes: int | None = None,

82 ) -> mooncake_store_worker.KVCacheStoreRecvingThread:

83 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec

@@ -96,7 +97,7 @@ def _make_store_recving_thread(

97 store=store,

98 token_databases=[token_database],

99 block_size=16,

99 - tp_rank=0,

100 + tp_rank=tp_rank,

101 ready_event=threading.Event(),

102 coord=coord,

103 disk_offload_buffer_budget_bytes=disk_offload_buffer_budget_bytes,

@@ -460,6 +461,67 @@ def test_store_sending_thread_only_skips_on_no_available_handle():

461 assert store.batch_put_from_multi_buffers.call_count == 2

462

463

464 +def test_store_recving_thread_reports_failed_block_ids():

465 + store = MagicMock()

466 + store.batch_get_into_multi_buffers.return_value = [256, -5, -7]

467 + thread = _make_store_recving_thread(store)

468 +

469 + thread._handle_request(

470 + _make_load_req(

471 + "req-a",

472 + [b"a0", b"a1", b"a2"],

473 + token_len=48,

474 + )

475 + )

476 +

477 + assert thread.get_and_clear_finished_requests() == {"req-a"}

478 + assert thread.get_and_clear_block_ids_with_load_errors() == {1, 2}

479 + assert thread.get_and_clear_block_ids_with_load_errors() == set()

480 +

481 +

482 +def test_store_recving_thread_reports_failed_block_ids_after_rotation():

483 + store = MagicMock()

484 + store.batch_get_into_multi_buffers.return_value = [256, -5, 256]

485 + thread = _make_store_recving_thread(store, tp_rank=1)

486 +

487 + thread._handle_request(

488 + _make_load_req(

489 + "req-a",

490 + [b"a0", b"a1", b"a2"],

491 + token_len=48,

492 + )

493 + )

494 +

495 + assert thread.get_and_clear_block_ids_with_load_errors() == {2}

496 +

497 +

498 +def test_store_recving_thread_reports_all_attempted_blocks_on_exception():

499 + store = MagicMock()

500 + store.batch_get_into_multi_buffers.side_effect = RuntimeError("boom")

501 + thread = _make_store_recving_thread(store)

502 +

503 + thread._handle_request(

504 + _make_load_req(

505 + "req-a",

506 + [b"a0", b"a1", b"a2"],

507 + token_len=48,

508 + )

509 + )

510 +

511 + assert thread.get_and_clear_finished_requests() == {"req-a"}

512 + assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}

513 +

514 +

515 +def test_store_worker_get_block_ids_with_load_errors_delegates_to_recv_thread():

516 + recv_thread = MagicMock()

517 + recv_thread.get_and_clear_block_ids_with_load_errors.return_value = {3, 4}

518 + w = _make_bare_worker()

519 + w.kv_recv_thread = recv_thread

520 +

521 + assert w.get_block_ids_with_load_errors() == {3, 4}

522 + recv_thread.get_and_clear_block_ids_with_load_errors.assert_called_once_with()

523 +

524 +

525 def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set():

526 store = MagicMock()

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

@@ -682,6 +744,7 @@ def test_recv_thread_stops_after_first_failing_disk_offload_sub_batch():

744 thread._handle_request(req)

745

746 assert store.batch_get_into_multi_buffers.call_count == 1

747 + assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1}

748

749

750 def test_recv_thread_skips_split_when_budget_holds_all_keys():

@@ -711,21 +774,26 @@ def test_recv_thread_skips_split_when_budget_holds_all_keys():

774

775

776 def test_recv_thread_reports_unsplittable_key_larger_than_budget():

777 + # Non-zero tp_rank exercises the rotation: the oversized key may not sit

778 + # at the original request's first block. Every block must still be marked

779 + # invalid, since none are loaded.

780 store = MagicMock()

781 thread = _make_store_recving_thread(

782 store,

783 + tp_rank=2,

784 disk_offload_buffer_budget_bytes=_DISK_OFFLOAD_BUDGET_TOO_SMALL,

785 )

786

787 req = _make_load_req(

788 "req-a",

722 - [b"a0"],

723 - token_len=16,

789 + [b"a0", b"a1", b"a2"],

790 + token_len=48,

791 )

792

793 thread._handle_request(req)

794

795 assert store.batch_get_into_multi_buffers.call_count == 0

796 + assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}

797

798

799 def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):

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

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

270 assert isinstance(metadata, MooncakeStoreConnectorMetadata)

271 return self.connector_worker.get_finished(finished_req_ids, metadata)

272

273 + def get_block_ids_with_load_errors(self) -> set[int]:

274 + assert self.connector_worker is not None

275 + return self.connector_worker.get_block_ids_with_load_errors()

276 +

277 def get_kv_connector_kv_cache_events(

278 self,

279 ) -> MooncakeStoreKVEvents | None:

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

+64 -12

@@ -20,7 +20,7 @@ import time

20 from collections import defaultdict

21 from collections.abc import Callable

22 from dataclasses import dataclass

23 -from typing import Any, Literal

23 +from typing import Any, Literal, TypeVar

25 import regex as re

26 import torch

@@ -68,6 +68,12 @@ DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB

68 DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB

70 MOONCAKE_NO_AVAILABLE_HANDLE = -200

71 +_T = TypeVar("_T")

72 +

73 +

74 +def _rotate_list(values: list[_T], offset: int) -> list[_T]:

75 + return values[offset:] + values[:offset]

76 +

78 # Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.

79 DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024

@@ -711,6 +717,9 @@ class KVCacheStoreRecvingThread(KVTransferThread):

717 name="KVCacheStoreRecvingThread",

718 record_operation=record_operation,

719 )

720 + # _invalid_block_ids can be access by both the Worker and RecvingThread

721 + self._invalid_block_ids_lock = threading.Lock()

722 + self._invalid_block_ids: set[int] = set()

723 self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes

724 self.usable_disk_offload_buffer_budget_bytes = (

725 None

@@ -721,6 +730,16 @@ class KVCacheStoreRecvingThread(KVTransferThread):

730 )

731 self.coord = coord

732

733 + def _add_load_error_block_ids(self, block_ids: list[int]) -> None:

734 + with self._invalid_block_ids_lock:

735 + self._invalid_block_ids.update(block_ids)

736 +

737 + def get_and_clear_block_ids_with_load_errors(self) -> set[int]:

738 + with self._invalid_block_ids_lock:

739 + invalid_block_ids = self._invalid_block_ids.copy()

740 + self._invalid_block_ids.clear()

741 + return invalid_block_ids

742 +

743 def _handle_request(self, req_meta: ReqMeta):

744 token_len = req_meta.load_spec.token_len # type: ignore[union-attr]

745 req_id = req_meta.req_id

@@ -737,6 +756,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):

756 addr_list: list[list[int]] = []

757 size_list: list[list[int]] = []

758 key_list: list[str] = []

759 + block_id_list: list[int] = []

760 for g_idx, db in enumerate(self.token_databases):

761 mask = load_mask_per_group[g_idx]

762 for start, end, key in db.process_tokens(

@@ -745,25 +765,29 @@ class KVCacheStoreRecvingThread(KVTransferThread):

765 chunk_idx = start // db.block_size

766 if chunk_idx >= len(mask) or not mask[chunk_idx]:

767 continue

748 - addr, size, _ = db.prepare_value(start, end, req_meta.block_ids[g_idx])

768 + addr, size, block_id = db.prepare_value(

769 + start, end, req_meta.block_ids[g_idx]

770 + )

771 key_list.append(key.to_string())

772 addr_list.append(addr)

773 size_list.append(size)

774 + block_id_list.append(block_id)

775

753 - # Rotate lists by tp_rank for load balancing

776 + # Rotate aligned lists by tp_rank for load balancing.

777 rotation = self.tp_rank % len(key_list)

755 - key_list_c = key_list[rotation:] + key_list[:rotation]

756 - addr_list_c = addr_list[rotation:] + addr_list[:rotation]

757 - size_list_c = size_list[rotation:] + size_list[:rotation]

778 + key_list_c = _rotate_list(key_list, rotation)

779 + addr_list_c = _rotate_list(addr_list, rotation)

780 + size_list_c = _rotate_list(size_list, rotation)

781 + block_id_list_c = _rotate_list(block_id_list, rotation)

782

759 - load_batches = [(key_list_c, addr_list_c, size_list_c)]

783 + load_batches = [(key_list_c, addr_list_c, size_list_c, block_id_list_c)]

784 if self.usable_disk_offload_buffer_budget_bytes is not None:

785 total_staging_bytes = sum(

786 _estimate_disk_offload_staging_bytes(size) for size in size_list_c

787 )

788 if total_staging_bytes > self.usable_disk_offload_buffer_budget_bytes:

789 assert self.disk_offload_buffer_budget_bytes is not None

766 - load_batches, oversized_key = _split_disk_offload_load_batches(

790 + split_batches, oversized_key = _split_disk_offload_load_batches(

791 key_list_c,

792 addr_list_c,

793 size_list_c,

@@ -772,6 +796,10 @@ class KVCacheStoreRecvingThread(KVTransferThread):

796 )

797 if oversized_key is not None:

798 oversized_key_index = key_list_c.index(oversized_key)

799 + # Mark every block: we skip the whole request, and the

800 + # tp_rank rotation means oversized_key isn't necessarily

801 + # the first block in the request's original order.

802 + self._add_load_error_block_ids(block_id_list_c)

803 oversized_key_bytes = _estimate_disk_offload_staging_bytes(

804 size_list_c[oversized_key_index]

805 )

@@ -786,12 +814,25 @@ class KVCacheStoreRecvingThread(KVTransferThread):

814 self.set_finished_request(req_id)

815 self.request_queue.task_done()

816 return

817 + load_batches = []

818 + block_id_offset = 0

819 + for batch_keys, batch_addrs, batch_sizes in split_batches:

820 + next_block_id_offset = block_id_offset + len(batch_keys)

821 + batch_block_ids = block_id_list_c[

822 + block_id_offset:next_block_id_offset

823 + ]

824 + load_batches.append(

825 + (batch_keys, batch_addrs, batch_sizes, batch_block_ids)

826 + )

827 + block_id_offset = next_block_id_offset

828

829 current_batch_keys: list[str] = key_list_c

830 + current_batch_block_ids: list[int] = block_id_list_c

831 batch_bytes = 0

832 try:

793 - for batch_keys, batch_addrs, batch_sizes in load_batches:

833 + for batch_keys, batch_addrs, batch_sizes, batch_block_ids in load_batches:

834 current_batch_keys = batch_keys

835 + current_batch_block_ids = batch_block_ids

836 batch_bytes = _sum_batch_bytes(batch_sizes)

837 tiers_by_key: dict[str, str] | None = None

838 if envs.VLLM_MOONCAKE_STORE_TIER_LOG:

@@ -806,8 +847,10 @@ class KVCacheStoreRecvingThread(KVTransferThread):

847 req_id, batch_keys, res, tiers_by_key

848 )

849 failed = [

809 - (key, value)

810 - for key, value in zip(batch_keys, res, strict=True)

850 + (key, value, block_id)

851 + for key, value, block_id in zip(

852 + batch_keys, res, batch_block_ids, strict=True

853 + )

854 if value < 0

855 ]

856 self._record_operation(

@@ -819,15 +862,19 @@ class KVCacheStoreRecvingThread(KVTransferThread):

862 num_failed_keys=len(failed),

863 )

864 if failed:

865 + self._add_load_error_block_ids(

866 + [block_id for _, _, block_id in failed]

867 + )

868 logger.warning(

869 "Failed to get %d Mooncake keys from sub-batch "

870 "(batch_keys=%d, first_failures=%s)",

871 len(failed),

872 len(batch_keys),

827 - failed[:3],

873 + [(key, value) for key, value, _ in failed[:3]],

874 )

875 break

876 except Exception as e:

877 + self._add_load_error_block_ids(current_batch_block_ids)

878 self._record_operation(

879 "load_get",

880 load_get_start,

@@ -1239,6 +1286,11 @@ class MooncakeStoreWorker:

1286 )

1287 return done_sending, done_recving

1288

1289 + def get_block_ids_with_load_errors(self) -> set[int]:

1290 + if self.kv_recv_thread is None:

1291 + return set()

1292 + return self.kv_recv_thread.get_and_clear_block_ids_with_load_errors()

1293 +

1294 def _record_kv_connector_operation(

1295 self,

1296 operation: str,