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

[KV Connector][Mooncake] Add store group semantics (#44956)

Signed-off-by: Schatten <czhengt@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Schatten committed Aug 5, 2026 at 16:25 UTC fd859d2e0689ce254bfb0e67483fba76405ca9df
2 files changed +531 -3
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+463 -2
@@ -76,6 +76,8 @@ def _make_store_sending_thread(
76 tp_rank: int = 0,
77 put_step: int = 1,
78 replicate_config: object | None = None,
79 + enable_group_semantics: bool = False,
80 + supports_group_ids: bool = False,
81 ) -> mooncake_store_worker.KVCacheStoreSendingThread:
82 if coord is None:
83 coord = _default_send_coord()
@@ -94,6 +96,8 @@ def _make_store_sending_thread(
96 kv_role="kv_producer",
97 ready_event=threading.Event(),
98 replicate_config=replicate_config,
99 + enable_group_semantics=enable_group_semantics,
100 + supports_group_ids=supports_group_ids,
101 )
102 thread.request_queue.task_done = MagicMock()
103 return thread
@@ -162,6 +166,16 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
166 )
167
168
169 +def _make_multi_group_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
170 + return ReqMeta(
171 + req_id=req_id,
172 + token_len_chunk=32,
173 + block_ids=([0, 1], [2, 3]),
174 + block_hashes=block_hashes,
175 + can_save=True,
176 + )
177 +
178 +
179 _DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256])
180 _DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9
181 _DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES
@@ -276,6 +290,7 @@ def _patch_worker_runtime(
290 monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
291 monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)
292 monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
293 + monkeypatch.setattr(worker, "LookupKeyServer", MagicMock())
294
295
296 def test_pool_key_to_string_without_prefix_is_unchanged():
@@ -599,7 +614,13 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():
614 assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
615
616
602 -def _make_partial_tail_send_thread(store):
617 +def _make_partial_tail_send_thread(
618 + store,
619 + *,
620 + replicate_config=None,
621 + enable_group_semantics=False,
622 + supports_group_ids=False,
623 +):
624 coord = SimpleNamespace(
625 enable_partial_hash_hits=True,
626 hash_block_size=4,
@@ -616,6 +637,9 @@ def _make_partial_tail_send_thread(store):
637 store,
638 coord=coord,
639 token_databases=[db],
640 + replicate_config=replicate_config,
641 + enable_group_semantics=enable_group_semantics,
642 + supports_group_ids=supports_group_ids,
643 )
644
645
@@ -648,6 +672,32 @@ def test_partial_tail_offload_skips_null_source_blocks():
672 assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]]
673
674
675 +def test_partial_tail_offload_replaces_stale_group_ids_after_filtering():
676 + store = MagicMock()
677 + store.batch_is_exist.return_value = [1, 0, 0]
678 + store.batch_put_from_multi_buffers.return_value = [256, 256]
679 + replicate_config = SimpleNamespace(group_ids=["stale"])
680 + thread = _make_partial_tail_send_thread(
681 + store,
682 + replicate_config=replicate_config,
683 + enable_group_semantics=True,
684 + supports_group_ids=True,
685 + )
686 +
687 + assert thread._maybe_offload_partial_tail(_make_partial_tail_req([1, 2, 3]))
688 +
689 + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
690 + assert keys == [
691 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
692 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
693 + ]
694 + assert config is replicate_config
695 + assert config.group_ids == [
696 + "vllm-mooncake-store:test-model@6131",
697 + "vllm-mooncake-store:test-model@6132",
698 + ]
699 +
700 +
701 def test_partial_tail_offload_honors_active_pressure_gate():
702 store = MagicMock()
703 thread = _make_partial_tail_send_thread(store)
@@ -969,6 +1019,291 @@ def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_
1019 assert call_args[3] is replicate_config
1020
1021
1022 +def test_group_id_detection_uses_class_attribute_without_instantiating():
1023 + class ReplicateConfigWithClassGroupIdsAndRequiredInit:
1024 + group_ids = None
1025 +
1026 + def __init__(self, required):
1027 + self.group_ids = required
1028 +
1029 + assert worker._replicate_config_supports_group_ids(
1030 + ReplicateConfigWithClassGroupIdsAndRequiredInit,
1031 + object(),
1032 + )
1033 +
1034 +
1035 +def test_group_id_detection_uses_instance_attribute_when_class_lacks_attribute():
1036 + class ReplicateConfigWithInstanceGroupIds:
1037 + def __init__(self) -> None:
1038 + self.group_ids = None
1039 +
1040 + replicate_config = ReplicateConfigWithInstanceGroupIds()
1041 +
1042 + assert worker._replicate_config_supports_group_ids(
1043 + ReplicateConfigWithInstanceGroupIds,
1044 + replicate_config,
1045 + )
1046 +
1047 +
1048 +def test_group_id_detection_rejects_old_replicate_config():
1049 + class ReplicateConfigWithoutGroupIds:
1050 + __slots__ = ()
1051 +
1052 + assert not worker._replicate_config_supports_group_ids(
1053 + ReplicateConfigWithoutGroupIds,
1054 + ReplicateConfigWithoutGroupIds(),
1055 + )
1056 +
1057 +
1058 +def test_store_sending_thread_sets_group_ids_when_enabled():
1059 + store = MagicMock()
1060 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1061 + store.batch_put_from_multi_buffers.return_value = [256, 256]
1062 + replicate_config = SimpleNamespace(group_ids=None)
1063 + thread = _make_store_sending_thread(
1064 + store,
1065 + replicate_config=replicate_config,
1066 + enable_group_semantics=True,
1067 + supports_group_ids=True,
1068 + )
1069 +
1070 + thread.add_stored_request("req-a")
1071 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
1072 +
1073 + assert store.batch_put_from_multi_buffers.call_count == 1
1074 + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
1075 + assert config is replicate_config
1076 + assert config.group_ids == [
1077 + "vllm-mooncake-store:test-model@6130",
1078 + "vllm-mooncake-store:test-model@6131",
1079 + ]
1080 + assert len(config.group_ids) == len(keys)
1081 +
1082 +
1083 +def test_store_sending_thread_leaves_group_ids_unchanged_when_flag_disabled():
1084 + store = MagicMock()
1085 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1086 + store.batch_put_from_multi_buffers.return_value = [256, 256]
1087 + replicate_config = SimpleNamespace(group_ids=["existing"])
1088 + thread = _make_store_sending_thread(
1089 + store,
1090 + replicate_config=replicate_config,
1091 + enable_group_semantics=False,
1092 + supports_group_ids=True,
1093 + )
1094 +
1095 + thread.add_stored_request("req-a")
1096 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
1097 +
1098 + assert store.batch_put_from_multi_buffers.call_count == 1
1099 + assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config
1100 + assert replicate_config.group_ids == ["existing"]
1101 +
1102 +
1103 +def test_store_sending_thread_leaves_group_ids_unchanged_when_unsupported():
1104 + store = MagicMock()
1105 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1106 + store.batch_put_from_multi_buffers.return_value = [256, 256]
1107 + replicate_config = SimpleNamespace(group_ids=["existing"])
1108 + thread = _make_store_sending_thread(
1109 + store,
1110 + replicate_config=replicate_config,
1111 + enable_group_semantics=True,
1112 + supports_group_ids=False,
1113 + )
1114 +
1115 + thread.add_stored_request("req-a")
1116 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
1117 +
1118 + assert store.batch_put_from_multi_buffers.call_count == 1
1119 + assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config
1120 + assert replicate_config.group_ids == ["existing"]
1121 +
1122 +
1123 +def test_mooncake_group_id_uses_logical_chunk_and_excludes_physical_ranks():
1124 + group_id = worker._make_mooncake_group_id(
1125 + KeyMetadata(
1126 + "test-model",
1127 + tp_rank=3,
1128 + pcp_rank=5,
1129 + dcp_rank=7,
1130 + pp_rank=11,
1131 + group_id=13,
1132 + ),
1133 + "abcdef",
1134 + )
1135 +
1136 + assert group_id == "vllm-mooncake-store:test-model@abcdef"
1137 + assert "tp_rank" not in group_id
1138 + assert "pp_rank" not in group_id
1139 + assert "pcp" not in group_id
1140 + assert "dcp" not in group_id
1141 + assert "group:" not in group_id
1142 +
1143 +
1144 +def test_mooncake_group_id_uses_cache_prefix_namespace():
1145 + metadata_a = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="deployment-a")
1146 + metadata_b = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="deployment-b")
1147 +
1148 + group_id_a = worker._make_mooncake_group_id(metadata_a, "abcdef")
1149 + group_id_b = worker._make_mooncake_group_id(metadata_b, "abcdef")
1150 +
1151 + assert group_id_a == "vllm-mooncake-store:deployment-a@test-model@abcdef"
1152 + assert group_id_b == "vllm-mooncake-store:deployment-b@test-model@abcdef"
1153 + assert group_id_a != group_id_b
1154 +
1155 +
1156 +def test_store_sending_thread_group_id_excludes_physical_sharding():
1157 + store = MagicMock()
1158 + store.batch_is_exist.return_value = [0, 0]
1159 + store.batch_put_from_multi_buffers.return_value = [256, 256]
1160 + replicate_config = SimpleNamespace(group_ids=None)
1161 + db = ChunkedTokenDatabase(
1162 + KeyMetadata(
1163 + "test-model",
1164 + tp_rank=2,
1165 + pcp_rank=0,
1166 + dcp_rank=0,
1167 + pp_rank=0,
1168 + group_id=0,
1169 + ),
1170 + block_size=16,
1171 + )
1172 + db.set_kv_caches_base_addr([0x1000])
1173 + db.set_block_len([256])
1174 + thread = _make_store_sending_thread(
1175 + store,
1176 + token_databases=[db],
1177 + tp_rank=2,
1178 + replicate_config=replicate_config,
1179 + enable_group_semantics=True,
1180 + supports_group_ids=True,
1181 + )
1182 +
1183 + thread.add_stored_request("req-a")
1184 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
1185 +
1186 + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
1187 + assert keys == [
1188 + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:0@6130",
1189 + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:0@6131",
1190 + ]
1191 + assert config.group_ids == [
1192 + "vllm-mooncake-store:test-model@6130",
1193 + "vllm-mooncake-store:test-model@6131",
1194 + ]
1195 +
1196 +
1197 +def test_store_sending_thread_multiple_segments_share_logical_group_id():
1198 + store = MagicMock()
1199 + store.batch_is_exist.return_value = [0, 0]
1200 + store.batch_put_from_multi_buffers.return_value = [512, 512]
1201 + replicate_config = SimpleNamespace(group_ids=None)
1202 + db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
1203 + db.set_kv_caches_base_addr([0x1000, 0x2000])
1204 + db.set_block_len([256, 256])
1205 + thread = _make_store_sending_thread(
1206 + store,
1207 + token_databases=[db],
1208 + replicate_config=replicate_config,
1209 + enable_group_semantics=True,
1210 + supports_group_ids=True,
1211 + )
1212 +
1213 + thread.add_stored_request("req-a")
1214 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
1215 +
1216 + keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args
1217 + assert keys == [
1218 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
1219 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
1220 + ]
1221 + assert addrs == [[0x1000, 0x2000], [0x1100, 0x2100]]
1222 + assert sizes == [[256, 256], [256, 256]]
1223 + assert config.group_ids == [
1224 + "vllm-mooncake-store:test-model@6130",
1225 + "vllm-mooncake-store:test-model@6131",
1226 + ]
1227 +
1228 +
1229 +def test_store_sending_thread_group_ids_share_across_kv_cache_groups():
1230 + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
1231 +
1232 + store = MagicMock()
1233 + store.batch_is_exist.return_value = [0, 0, 0, 0]
1234 + store.batch_put_from_multi_buffers.return_value = [256, 256, 256, 256]
1235 + replicate_config = SimpleNamespace(group_ids=None)
1236 + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
1237 + coord = mooncake_store_worker.MooncakeStoreCoordinator(
1238 + [
1239 + KVCacheGroupSpec(["layer0"], spec),
1240 + KVCacheGroupSpec(["layer1"], spec),
1241 + ],
1242 + scheduler_block_size=16,
1243 + hash_block_size=16,
1244 + )
1245 + token_databases = []
1246 + for group_id, base_addr in enumerate([0x1000, 0x3000]):
1247 + db = ChunkedTokenDatabase(
1248 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
1249 + block_size=16,
1250 + )
1251 + db.set_kv_caches_base_addr([base_addr])
1252 + db.set_block_len([256])
1253 + token_databases.append(db)
1254 + thread = _make_store_sending_thread(
1255 + store,
1256 + coord=coord,
1257 + token_databases=token_databases,
1258 + replicate_config=replicate_config,
1259 + enable_group_semantics=True,
1260 + supports_group_ids=True,
1261 + )
1262 +
1263 + thread.add_stored_request("req-a")
1264 + thread._handle_request(_make_multi_group_store_req("req-a", [b"a0", b"a1"]))
1265 +
1266 + keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args
1267 + assert keys == [
1268 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
1269 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
1270 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6130",
1271 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6131",
1272 + ]
1273 + assert addrs == [[0x1000], [0x1100], [0x3200], [0x3300]]
1274 + assert sizes == [[256], [256], [256], [256]]
1275 + # Different vLLM KV cache groups for the same prefix chunk share the
1276 + # same Mooncake lifecycle group id.
1277 + assert config.group_ids == [
1278 + "vllm-mooncake-store:test-model@6130",
1279 + "vllm-mooncake-store:test-model@6131",
1280 + "vllm-mooncake-store:test-model@6130",
1281 + "vllm-mooncake-store:test-model@6131",
1282 + ]
1283 + assert config.group_ids[0] == config.group_ids[2]
1284 + assert config.group_ids[1] == config.group_ids[3]
1285 +
1286 +
1287 +def test_store_sending_thread_group_ids_follow_missing_key_filter():
1288 + store = MagicMock()
1289 + store.batch_is_exist.return_value = [1, 0]
1290 + store.batch_put_from_multi_buffers.return_value = [256]
1291 + replicate_config = SimpleNamespace(group_ids=None)
1292 + thread = _make_store_sending_thread(
1293 + store,
1294 + replicate_config=replicate_config,
1295 + enable_group_semantics=True,
1296 + supports_group_ids=True,
1297 + )
1298 +
1299 + thread.add_stored_request("req-a")
1300 + thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
1301 +
1302 + keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
1303 + assert keys == ["test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131"]
1304 + assert config.group_ids == ["vllm-mooncake-store:test-model@6131"]
1305 +
1306 +
1307 def test_estimate_disk_offload_staging_bytes_sums_multi_segment_sizes():
1308 assert worker._estimate_disk_offload_staging_bytes([256, 512]) == 12288
1309
@@ -1392,6 +1727,130 @@ def test_worker_put_striding_covers_every_rank_get_namespace(
1727 )
1728
1729
1730 +def test_requester_worker_group_semantics_falls_back_without_group_ids(
1731 + tmp_path,
1732 + monkeypatch,
1733 +):
1734 + store = MagicMock()
1735 + store.setup.return_value = 0
1736 + _install_fake_mooncake(monkeypatch, store)
1737 + _patch_worker_runtime(monkeypatch)
1738 + warning = MagicMock()
1739 + monkeypatch.setattr(worker.logger, "warning", warning)
1740 + monkeypatch.setenv(
1741 + "MOONCAKE_CONFIG_PATH",
1742 + _write_mooncake_config(
1743 + tmp_path,
1744 + {
1745 + "metadata_server": "http://metadata/endpoint",
1746 + "protocol": "tcp",
1747 + "device_name": "",
1748 + "master_server_address": "10.0.0.7:50051",
1749 + },
1750 + ),
1751 + )
1752 +
1753 + w = worker.MooncakeStoreWorker(
1754 + _make_vllm_config(
1755 + extra_config={
1756 + "enable_group_semantics": True,
1757 + }
1758 + ),
1759 + _make_kv_cache_config(),
1760 + )
1761 +
1762 + assert w.enable_group_semantics is True
1763 + assert w._supports_group_ids is False
1764 + assert any(
1765 + "does not support ReplicateConfig.group_ids" in call.args[0]
1766 + for call in warning.call_args_list
1767 + )
1768 +
1769 +
1770 +def test_requester_worker_group_semantics_string_false_does_not_enable(
1771 + tmp_path,
1772 + monkeypatch,
1773 +):
1774 + store = MagicMock()
1775 + store.setup.return_value = 0
1776 + _install_fake_mooncake(monkeypatch, store)
1777 + _patch_worker_runtime(monkeypatch)
1778 + warning = MagicMock()
1779 + monkeypatch.setattr(worker.logger, "warning", warning)
1780 + monkeypatch.setenv(
1781 + "MOONCAKE_CONFIG_PATH",
1782 + _write_mooncake_config(
1783 + tmp_path,
1784 + {
1785 + "metadata_server": "http://metadata/endpoint",
1786 + "protocol": "tcp",
1787 + "device_name": "",
1788 + "master_server_address": "10.0.0.7:50051",
1789 + },
1790 + ),
1791 + )
1792 +
1793 + w = worker.MooncakeStoreWorker(
1794 + _make_vllm_config(
1795 + extra_config={
1796 + "enable_group_semantics": "false",
1797 + }
1798 + ),
1799 + _make_kv_cache_config(),
1800 + )
1801 +
1802 + assert w.enable_group_semantics is False
1803 + assert not any(
1804 + "does not support ReplicateConfig.group_ids" in call.args[0]
1805 + for call in warning.call_args_list
1806 + )
1807 +
1808 +
1809 +def test_requester_worker_group_semantics_string_true_enables(
1810 + tmp_path,
1811 + monkeypatch,
1812 +):
1813 + store = MagicMock()
1814 + store.setup.return_value = 0
1815 +
1816 + class FakeReplicateConfig:
1817 + def __init__(self) -> None:
1818 + self.group_ids = None
1819 +
1820 + fake_store_module = types.ModuleType("mooncake.store")
1821 + fake_store_module.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined]
1822 + fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined]
1823 + fake_mooncake_module = types.ModuleType("mooncake")
1824 + fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined]
1825 + monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
1826 + monkeypatch.setitem(sys.modules, "mooncake.store", fake_store_module)
1827 + _patch_worker_runtime(monkeypatch)
1828 + monkeypatch.setenv(
1829 + "MOONCAKE_CONFIG_PATH",
1830 + _write_mooncake_config(
1831 + tmp_path,
1832 + {
1833 + "metadata_server": "http://metadata/endpoint",
1834 + "protocol": "tcp",
1835 + "device_name": "",
1836 + "master_server_address": "10.0.0.7:50051",
1837 + },
1838 + ),
1839 + )
1840 +
1841 + w = worker.MooncakeStoreWorker(
1842 + _make_vllm_config(
1843 + extra_config={
1844 + "enable_group_semantics": "true",
1845 + }
1846 + ),
1847 + _make_kv_cache_config(),
1848 + )
1849 +
1850 + assert w.enable_group_semantics is True
1851 + assert w._supports_group_ids is True
1852 +
1853 +
1854 # ---------------------------------------------------------------------------
1855 # Helpers for register_kv_caches tests
1856 # ---------------------------------------------------------------------------
@@ -1481,7 +1940,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
1940 store = MagicMock()
1941 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
1942 store.batch_put_from_multi_buffers.side_effect = (
1484 - lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
1943 + lambda keys, addrs, sizes, *_args: [256] * len(keys)
1944 )
1945
1946 full_spec = FullAttentionSpec(
@@ -1767,6 +2226,8 @@ def _make_bare_worker(
2226
2227 worker.disk_offload_buffer_budget_bytes = None
2228 worker.store_replicate_config = SimpleNamespace()
2229 + worker.enable_group_semantics = False
2230 + worker._supports_group_ids = False
2231 worker._kv_connector_stats_lock = threading.Lock()
2232 worker.kv_connector_stats = MooncakeStoreConnectorStats()
2233
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+68 -1
@@ -89,6 +89,22 @@ def _rotate_list(values: list[_T], offset: int) -> list[_T]:
89 return values[offset:] + values[:offset]
90
91
92 +def _replicate_config_supports_group_ids(
93 + replicate_config_cls: type[Any],
94 + replicate_config: Any,
95 +) -> bool:
96 + if hasattr(replicate_config_cls, "group_ids"):
97 + return True
98 + return hasattr(replicate_config, "group_ids")
99 +
100 +
101 +def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str:
102 + # Mooncake group ids describe the lifecycle unit. For vLLM, that unit is
103 + # a prefix chunk, so shard dimensions stay only in the object key.
104 + prefix = f"{metadata.cache_prefix}@" if metadata.cache_prefix else ""
105 + return f"vllm-mooncake-store:{prefix}{metadata.model_name}@{chunk_hash}"
106 +
107 +
108 # Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.
109 DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024
110
@@ -459,6 +475,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
475 ready_event: threading.Event,
476 enable_kv_event: bool = False,
477 replicate_config: Any = None,
478 + enable_group_semantics: bool = False,
479 + supports_group_ids: bool = False,
480 record_operation: Callable[..., None] | None = None,
481 ):
482 super().__init__(
@@ -479,6 +497,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
497 # Caller always passes a non-None ReplicateConfig — see
498 # MooncakeStoreWorker.__init__ where store_replicate_config is built.
499 self.replicate_config = replicate_config
500 + self.enable_group_semantics = enable_group_semantics
501 + self.supports_group_ids = supports_group_ids
502
503 # Pause store requests when CPU/disk offloading is under pressure.
504 self._store_pressure_active = False
@@ -567,6 +587,9 @@ class KVCacheStoreSendingThread(KVTransferThread):
587 keys: list[str] = []
588 addrs: list[list[int]] = []
589 sizes: list[list[int]] = []
590 + group_ids: list[str] | None = (
591 + [] if self.enable_group_semantics and self.supports_group_ids else None
592 + )
593 saved = self._saved_offset.get(req_meta.req_id, 0)
594 for g_idx, db in enumerate(self.token_databases):
595 group_blocks = req_meta.block_ids[g_idx]
@@ -603,9 +626,17 @@ class KVCacheStoreSendingThread(KVTransferThread):
626 )
627 continue
628 addr, size = db.prepare_value_for_block(block_id)
606 - keys.append(db.key_for(key_hash))
629 + key = db.key_for(key_hash)
630 + keys.append(key)
631 addrs.append(addr)
632 sizes.append(size)
633 + if group_ids is not None:
634 + group_ids.append(
635 + _make_mooncake_group_id(
636 + db.metadata,
637 + key.rsplit("@", 1)[-1],
638 + )
639 + )
640
641 if not keys:
642 return True
@@ -633,9 +664,14 @@ class KVCacheStoreSendingThread(KVTransferThread):
664 keys = [keys[i] for i in missing]
665 addrs = [addrs[i] for i in missing]
666 sizes = [sizes[i] for i in missing]
667 + if group_ids is not None:
668 + group_ids = [group_ids[i] for i in missing]
669 if req_meta.current_event is not None:
670 # Fence the CoW block copy enqueued earlier this step.
671 req_meta.current_event.synchronize()
672 + if group_ids is not None:
673 + assert len(group_ids) == len(keys)
674 + self.replicate_config.group_ids = group_ids
675 batch_bytes = _sum_batch_bytes(sizes)
676 put_start = time.perf_counter()
677 try:
@@ -797,6 +833,18 @@ class KVCacheStoreSendingThread(KVTransferThread):
833 ]
834 group_indices = [group_indices[i] for i in missing_indices]
835
836 + group_ids = (
837 + [
838 + _make_mooncake_group_id(
839 + self.token_databases[g_idx].metadata,
840 + key.rsplit("@", 1)[-1],
841 + )
842 + for key, g_idx in zip(keys, group_indices, strict=True)
843 + ]
844 + if self.enable_group_semantics and self.supports_group_ids
845 + else None
846 + )
847 +
848 logger.debug(
849 "Storing KV cache for %d blocks (groups=%s) for request %s",
850 len(keys),
@@ -855,6 +903,10 @@ class KVCacheStoreSendingThread(KVTransferThread):
903 if current_event is not None:
904 current_event.synchronize()
905
906 + if group_ids is not None:
907 + assert len(group_ids) == len(keys)
908 + self.replicate_config.group_ids = group_ids
909 +
910 batch_bytes = _sum_batch_bytes(sizes)
911 put_start = time.perf_counter()
912 try:
@@ -1210,6 +1262,19 @@ class MooncakeStoreWorker:
1262 preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config)
1263 self.preferred_segment = preferred_segment
1264 self.store_replicate_config = ReplicateConfig()
1265 + self.enable_group_semantics = (
1266 + str(extra_config.get("enable_group_semantics", "False")).strip().lower()
1267 + == "true"
1268 + )
1269 + self._supports_group_ids = _replicate_config_supports_group_ids(
1270 + ReplicateConfig, self.store_replicate_config
1271 + )
1272 + if self.enable_group_semantics and not self._supports_group_ids:
1273 + logger.warning(
1274 + "Mooncake group semantics is enabled, but the installed "
1275 + "Mooncake package does not support ReplicateConfig.group_ids. "
1276 + "Falling back to the existing batch_put_from_multi_buffers path."
1277 + )
1278 if preferred_segment is not None:
1279 self.store_replicate_config.preferred_segment = preferred_segment
1280
@@ -1503,6 +1568,8 @@ class MooncakeStoreWorker:
1568 ready_event_sending,
1569 self.enable_kv_events,
1570 self.store_replicate_config,
1571 + enable_group_semantics=self.enable_group_semantics,
1572 + supports_group_ids=self._supports_group_ids,
1573 record_operation=self._record_kv_connector_operation,
1574 )
1575 self.kv_send_thread.start()