[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>
This commit is contained in:
Schatten
2026-08-05 08:25:38 +00:00
committed by GitHub
co-authored by OpenAI Codex Yifan Qiao
parent f5cd862dbd
commit fd859d2e06
2 changed files with 531 additions and 3 deletions
@@ -76,6 +76,8 @@ def _make_store_sending_thread(
tp_rank: int = 0,
put_step: int = 1,
replicate_config: object | None = None,
enable_group_semantics: bool = False,
supports_group_ids: bool = False,
) -> mooncake_store_worker.KVCacheStoreSendingThread:
if coord is None:
coord = _default_send_coord()
@@ -94,6 +96,8 @@ def _make_store_sending_thread(
kv_role="kv_producer",
ready_event=threading.Event(),
replicate_config=replicate_config,
enable_group_semantics=enable_group_semantics,
supports_group_ids=supports_group_ids,
)
thread.request_queue.task_done = MagicMock()
return thread
@@ -162,6 +166,16 @@ def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
)
def _make_multi_group_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
return ReqMeta(
req_id=req_id,
token_len_chunk=32,
block_ids=([0, 1], [2, 3]),
block_hashes=block_hashes,
can_save=True,
)
_DISK_OFFLOAD_SINGLE_KEY_BYTES = worker._estimate_disk_offload_staging_bytes([256])
_DISK_OFFLOAD_USABLE_BUDGET_RATIO = 0.9
_DISK_OFFLOAD_BUDGET_FOR_THREE_KEYS = 4 * _DISK_OFFLOAD_SINGLE_KEY_BYTES
@@ -276,6 +290,7 @@ def _patch_worker_runtime(
monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)
monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
monkeypatch.setattr(worker, "LookupKeyServer", MagicMock())
def test_pool_key_to_string_without_prefix_is_unchanged():
@@ -599,7 +614,13 @@ def test_store_sending_thread_retries_skipped_range_after_pressure():
assert store.batch_put_from_multi_buffers.call_args.args[0] == keys
def _make_partial_tail_send_thread(store):
def _make_partial_tail_send_thread(
store,
*,
replicate_config=None,
enable_group_semantics=False,
supports_group_ids=False,
):
coord = SimpleNamespace(
enable_partial_hash_hits=True,
hash_block_size=4,
@@ -616,6 +637,9 @@ def _make_partial_tail_send_thread(store):
store,
coord=coord,
token_databases=[db],
replicate_config=replicate_config,
enable_group_semantics=enable_group_semantics,
supports_group_ids=supports_group_ids,
)
@@ -648,6 +672,32 @@ def test_partial_tail_offload_skips_null_source_blocks():
assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]]
def test_partial_tail_offload_replaces_stale_group_ids_after_filtering():
store = MagicMock()
store.batch_is_exist.return_value = [1, 0, 0]
store.batch_put_from_multi_buffers.return_value = [256, 256]
replicate_config = SimpleNamespace(group_ids=["stale"])
thread = _make_partial_tail_send_thread(
store,
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=True,
)
assert thread._maybe_offload_partial_tail(_make_partial_tail_req([1, 2, 3]))
keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
assert keys == [
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",
]
assert config is replicate_config
assert config.group_ids == [
"vllm-mooncake-store:test-model@6131",
"vllm-mooncake-store:test-model@6132",
]
def test_partial_tail_offload_honors_active_pressure_gate():
store = MagicMock()
thread = _make_partial_tail_send_thread(store)
@@ -969,6 +1019,291 @@ def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_
assert call_args[3] is replicate_config
def test_group_id_detection_uses_class_attribute_without_instantiating():
class ReplicateConfigWithClassGroupIdsAndRequiredInit:
group_ids = None
def __init__(self, required):
self.group_ids = required
assert worker._replicate_config_supports_group_ids(
ReplicateConfigWithClassGroupIdsAndRequiredInit,
object(),
)
def test_group_id_detection_uses_instance_attribute_when_class_lacks_attribute():
class ReplicateConfigWithInstanceGroupIds:
def __init__(self) -> None:
self.group_ids = None
replicate_config = ReplicateConfigWithInstanceGroupIds()
assert worker._replicate_config_supports_group_ids(
ReplicateConfigWithInstanceGroupIds,
replicate_config,
)
def test_group_id_detection_rejects_old_replicate_config():
class ReplicateConfigWithoutGroupIds:
__slots__ = ()
assert not worker._replicate_config_supports_group_ids(
ReplicateConfigWithoutGroupIds,
ReplicateConfigWithoutGroupIds(),
)
def test_store_sending_thread_sets_group_ids_when_enabled():
store = MagicMock()
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
store.batch_put_from_multi_buffers.return_value = [256, 256]
replicate_config = SimpleNamespace(group_ids=None)
thread = _make_store_sending_thread(
store,
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=True,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
assert store.batch_put_from_multi_buffers.call_count == 1
keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
assert config is replicate_config
assert config.group_ids == [
"vllm-mooncake-store:test-model@6130",
"vllm-mooncake-store:test-model@6131",
]
assert len(config.group_ids) == len(keys)
def test_store_sending_thread_leaves_group_ids_unchanged_when_flag_disabled():
store = MagicMock()
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
store.batch_put_from_multi_buffers.return_value = [256, 256]
replicate_config = SimpleNamespace(group_ids=["existing"])
thread = _make_store_sending_thread(
store,
replicate_config=replicate_config,
enable_group_semantics=False,
supports_group_ids=True,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
assert store.batch_put_from_multi_buffers.call_count == 1
assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config
assert replicate_config.group_ids == ["existing"]
def test_store_sending_thread_leaves_group_ids_unchanged_when_unsupported():
store = MagicMock()
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
store.batch_put_from_multi_buffers.return_value = [256, 256]
replicate_config = SimpleNamespace(group_ids=["existing"])
thread = _make_store_sending_thread(
store,
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=False,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
assert store.batch_put_from_multi_buffers.call_count == 1
assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config
assert replicate_config.group_ids == ["existing"]
def test_mooncake_group_id_uses_logical_chunk_and_excludes_physical_ranks():
group_id = worker._make_mooncake_group_id(
KeyMetadata(
"test-model",
tp_rank=3,
pcp_rank=5,
dcp_rank=7,
pp_rank=11,
group_id=13,
),
"abcdef",
)
assert group_id == "vllm-mooncake-store:test-model@abcdef"
assert "tp_rank" not in group_id
assert "pp_rank" not in group_id
assert "pcp" not in group_id
assert "dcp" not in group_id
assert "group:" not in group_id
def test_mooncake_group_id_uses_cache_prefix_namespace():
metadata_a = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="deployment-a")
metadata_b = KeyMetadata("test-model", 0, 0, 0, 0, cache_prefix="deployment-b")
group_id_a = worker._make_mooncake_group_id(metadata_a, "abcdef")
group_id_b = worker._make_mooncake_group_id(metadata_b, "abcdef")
assert group_id_a == "vllm-mooncake-store:deployment-a@test-model@abcdef"
assert group_id_b == "vllm-mooncake-store:deployment-b@test-model@abcdef"
assert group_id_a != group_id_b
def test_store_sending_thread_group_id_excludes_physical_sharding():
store = MagicMock()
store.batch_is_exist.return_value = [0, 0]
store.batch_put_from_multi_buffers.return_value = [256, 256]
replicate_config = SimpleNamespace(group_ids=None)
db = ChunkedTokenDatabase(
KeyMetadata(
"test-model",
tp_rank=2,
pcp_rank=0,
dcp_rank=0,
pp_rank=0,
group_id=0,
),
block_size=16,
)
db.set_kv_caches_base_addr([0x1000])
db.set_block_len([256])
thread = _make_store_sending_thread(
store,
token_databases=[db],
tp_rank=2,
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=True,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
assert keys == [
"test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:0@6130",
"test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:0@6131",
]
assert config.group_ids == [
"vllm-mooncake-store:test-model@6130",
"vllm-mooncake-store:test-model@6131",
]
def test_store_sending_thread_multiple_segments_share_logical_group_id():
store = MagicMock()
store.batch_is_exist.return_value = [0, 0]
store.batch_put_from_multi_buffers.return_value = [512, 512]
replicate_config = SimpleNamespace(group_ids=None)
db = ChunkedTokenDatabase(KeyMetadata("test-model", 0, 0, 0, 0), block_size=16)
db.set_kv_caches_base_addr([0x1000, 0x2000])
db.set_block_len([256, 256])
thread = _make_store_sending_thread(
store,
token_databases=[db],
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=True,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args
assert keys == [
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
]
assert addrs == [[0x1000, 0x2000], [0x1100, 0x2100]]
assert sizes == [[256, 256], [256, 256]]
assert config.group_ids == [
"vllm-mooncake-store:test-model@6130",
"vllm-mooncake-store:test-model@6131",
]
def test_store_sending_thread_group_ids_share_across_kv_cache_groups():
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec
store = MagicMock()
store.batch_is_exist.return_value = [0, 0, 0, 0]
store.batch_put_from_multi_buffers.return_value = [256, 256, 256, 256]
replicate_config = SimpleNamespace(group_ids=None)
spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
coord = mooncake_store_worker.MooncakeStoreCoordinator(
[
KVCacheGroupSpec(["layer0"], spec),
KVCacheGroupSpec(["layer1"], spec),
],
scheduler_block_size=16,
hash_block_size=16,
)
token_databases = []
for group_id, base_addr in enumerate([0x1000, 0x3000]):
db = ChunkedTokenDatabase(
KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),
block_size=16,
)
db.set_kv_caches_base_addr([base_addr])
db.set_block_len([256])
token_databases.append(db)
thread = _make_store_sending_thread(
store,
coord=coord,
token_databases=token_databases,
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=True,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_multi_group_store_req("req-a", [b"a0", b"a1"]))
keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args
assert keys == [
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6130",
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6130",
"test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6131",
]
assert addrs == [[0x1000], [0x1100], [0x3200], [0x3300]]
assert sizes == [[256], [256], [256], [256]]
# Different vLLM KV cache groups for the same prefix chunk share the
# same Mooncake lifecycle group id.
assert config.group_ids == [
"vllm-mooncake-store:test-model@6130",
"vllm-mooncake-store:test-model@6131",
"vllm-mooncake-store:test-model@6130",
"vllm-mooncake-store:test-model@6131",
]
assert config.group_ids[0] == config.group_ids[2]
assert config.group_ids[1] == config.group_ids[3]
def test_store_sending_thread_group_ids_follow_missing_key_filter():
store = MagicMock()
store.batch_is_exist.return_value = [1, 0]
store.batch_put_from_multi_buffers.return_value = [256]
replicate_config = SimpleNamespace(group_ids=None)
thread = _make_store_sending_thread(
store,
replicate_config=replicate_config,
enable_group_semantics=True,
supports_group_ids=True,
)
thread.add_stored_request("req-a")
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args
assert keys == ["test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131"]
assert config.group_ids == ["vllm-mooncake-store:test-model@6131"]
def test_estimate_disk_offload_staging_bytes_sums_multi_segment_sizes():
assert worker._estimate_disk_offload_staging_bytes([256, 512]) == 12288
@@ -1392,6 +1727,130 @@ def test_worker_put_striding_covers_every_rank_get_namespace(
)
def test_requester_worker_group_semantics_falls_back_without_group_ids(
tmp_path,
monkeypatch,
):
store = MagicMock()
store.setup.return_value = 0
_install_fake_mooncake(monkeypatch, store)
_patch_worker_runtime(monkeypatch)
warning = MagicMock()
monkeypatch.setattr(worker.logger, "warning", warning)
monkeypatch.setenv(
"MOONCAKE_CONFIG_PATH",
_write_mooncake_config(
tmp_path,
{
"metadata_server": "http://metadata/endpoint",
"protocol": "tcp",
"device_name": "",
"master_server_address": "10.0.0.7:50051",
},
),
)
w = worker.MooncakeStoreWorker(
_make_vllm_config(
extra_config={
"enable_group_semantics": True,
}
),
_make_kv_cache_config(),
)
assert w.enable_group_semantics is True
assert w._supports_group_ids is False
assert any(
"does not support ReplicateConfig.group_ids" in call.args[0]
for call in warning.call_args_list
)
def test_requester_worker_group_semantics_string_false_does_not_enable(
tmp_path,
monkeypatch,
):
store = MagicMock()
store.setup.return_value = 0
_install_fake_mooncake(monkeypatch, store)
_patch_worker_runtime(monkeypatch)
warning = MagicMock()
monkeypatch.setattr(worker.logger, "warning", warning)
monkeypatch.setenv(
"MOONCAKE_CONFIG_PATH",
_write_mooncake_config(
tmp_path,
{
"metadata_server": "http://metadata/endpoint",
"protocol": "tcp",
"device_name": "",
"master_server_address": "10.0.0.7:50051",
},
),
)
w = worker.MooncakeStoreWorker(
_make_vllm_config(
extra_config={
"enable_group_semantics": "false",
}
),
_make_kv_cache_config(),
)
assert w.enable_group_semantics is False
assert not any(
"does not support ReplicateConfig.group_ids" in call.args[0]
for call in warning.call_args_list
)
def test_requester_worker_group_semantics_string_true_enables(
tmp_path,
monkeypatch,
):
store = MagicMock()
store.setup.return_value = 0
class FakeReplicateConfig:
def __init__(self) -> None:
self.group_ids = None
fake_store_module = types.ModuleType("mooncake.store")
fake_store_module.MooncakeDistributedStore = lambda: store # type: ignore[attr-defined]
fake_store_module.ReplicateConfig = FakeReplicateConfig # type: ignore[attr-defined]
fake_mooncake_module = types.ModuleType("mooncake")
fake_mooncake_module.store = fake_store_module # type: ignore[attr-defined]
monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake_module)
monkeypatch.setitem(sys.modules, "mooncake.store", fake_store_module)
_patch_worker_runtime(monkeypatch)
monkeypatch.setenv(
"MOONCAKE_CONFIG_PATH",
_write_mooncake_config(
tmp_path,
{
"metadata_server": "http://metadata/endpoint",
"protocol": "tcp",
"device_name": "",
"master_server_address": "10.0.0.7:50051",
},
),
)
w = worker.MooncakeStoreWorker(
_make_vllm_config(
extra_config={
"enable_group_semantics": "true",
}
),
_make_kv_cache_config(),
)
assert w.enable_group_semantics is True
assert w._supports_group_ids is True
# ---------------------------------------------------------------------------
# Helpers for register_kv_caches tests
# ---------------------------------------------------------------------------
@@ -1481,7 +1940,7 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window():
store = MagicMock()
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
store.batch_put_from_multi_buffers.side_effect = (
lambda keys, addrs, sizes, replicate_config: [256] * len(keys)
lambda keys, addrs, sizes, *_args: [256] * len(keys)
)
full_spec = FullAttentionSpec(
@@ -1767,6 +2226,8 @@ def _make_bare_worker(
worker.disk_offload_buffer_budget_bytes = None
worker.store_replicate_config = SimpleNamespace()
worker.enable_group_semantics = False
worker._supports_group_ids = False
worker._kv_connector_stats_lock = threading.Lock()
worker.kv_connector_stats = MooncakeStoreConnectorStats()
@@ -89,6 +89,22 @@ def _rotate_list(values: list[_T], offset: int) -> list[_T]:
return values[offset:] + values[:offset]
def _replicate_config_supports_group_ids(
replicate_config_cls: type[Any],
replicate_config: Any,
) -> bool:
if hasattr(replicate_config_cls, "group_ids"):
return True
return hasattr(replicate_config, "group_ids")
def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str:
# Mooncake group ids describe the lifecycle unit. For vLLM, that unit is
# a prefix chunk, so shard dimensions stay only in the object key.
prefix = f"{metadata.cache_prefix}@" if metadata.cache_prefix else ""
return f"vllm-mooncake-store:{prefix}{metadata.model_name}@{chunk_hash}"
# Mirrors FileStorageConfig::local_buffer_size in Mooncake C++.
DEFAULT_MOONCAKE_DISK_STAGING_BUFFER_BYTES = 1280 * 1024 * 1024
@@ -459,6 +475,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
ready_event: threading.Event,
enable_kv_event: bool = False,
replicate_config: Any = None,
enable_group_semantics: bool = False,
supports_group_ids: bool = False,
record_operation: Callable[..., None] | None = None,
):
super().__init__(
@@ -479,6 +497,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
# Caller always passes a non-None ReplicateConfig — see
# MooncakeStoreWorker.__init__ where store_replicate_config is built.
self.replicate_config = replicate_config
self.enable_group_semantics = enable_group_semantics
self.supports_group_ids = supports_group_ids
# Pause store requests when CPU/disk offloading is under pressure.
self._store_pressure_active = False
@@ -567,6 +587,9 @@ class KVCacheStoreSendingThread(KVTransferThread):
keys: list[str] = []
addrs: list[list[int]] = []
sizes: list[list[int]] = []
group_ids: list[str] | None = (
[] if self.enable_group_semantics and self.supports_group_ids else None
)
saved = self._saved_offset.get(req_meta.req_id, 0)
for g_idx, db in enumerate(self.token_databases):
group_blocks = req_meta.block_ids[g_idx]
@@ -603,9 +626,17 @@ class KVCacheStoreSendingThread(KVTransferThread):
)
continue
addr, size = db.prepare_value_for_block(block_id)
keys.append(db.key_for(key_hash))
key = db.key_for(key_hash)
keys.append(key)
addrs.append(addr)
sizes.append(size)
if group_ids is not None:
group_ids.append(
_make_mooncake_group_id(
db.metadata,
key.rsplit("@", 1)[-1],
)
)
if not keys:
return True
@@ -633,9 +664,14 @@ class KVCacheStoreSendingThread(KVTransferThread):
keys = [keys[i] for i in missing]
addrs = [addrs[i] for i in missing]
sizes = [sizes[i] for i in missing]
if group_ids is not None:
group_ids = [group_ids[i] for i in missing]
if req_meta.current_event is not None:
# Fence the CoW block copy enqueued earlier this step.
req_meta.current_event.synchronize()
if group_ids is not None:
assert len(group_ids) == len(keys)
self.replicate_config.group_ids = group_ids
batch_bytes = _sum_batch_bytes(sizes)
put_start = time.perf_counter()
try:
@@ -797,6 +833,18 @@ class KVCacheStoreSendingThread(KVTransferThread):
]
group_indices = [group_indices[i] for i in missing_indices]
group_ids = (
[
_make_mooncake_group_id(
self.token_databases[g_idx].metadata,
key.rsplit("@", 1)[-1],
)
for key, g_idx in zip(keys, group_indices, strict=True)
]
if self.enable_group_semantics and self.supports_group_ids
else None
)
logger.debug(
"Storing KV cache for %d blocks (groups=%s) for request %s",
len(keys),
@@ -855,6 +903,10 @@ class KVCacheStoreSendingThread(KVTransferThread):
if current_event is not None:
current_event.synchronize()
if group_ids is not None:
assert len(group_ids) == len(keys)
self.replicate_config.group_ids = group_ids
batch_bytes = _sum_batch_bytes(sizes)
put_start = time.perf_counter()
try:
@@ -1210,6 +1262,19 @@ class MooncakeStoreWorker:
preferred_segment = rdma_utils.get_configured_preferred_segment(extra_config)
self.preferred_segment = preferred_segment
self.store_replicate_config = ReplicateConfig()
self.enable_group_semantics = (
str(extra_config.get("enable_group_semantics", "False")).strip().lower()
== "true"
)
self._supports_group_ids = _replicate_config_supports_group_ids(
ReplicateConfig, self.store_replicate_config
)
if self.enable_group_semantics and not self._supports_group_ids:
logger.warning(
"Mooncake group semantics is enabled, but the installed "
"Mooncake package does not support ReplicateConfig.group_ids. "
"Falling back to the existing batch_put_from_multi_buffers path."
)
if preferred_segment is not None:
self.store_replicate_config.preferred_segment = preferred_segment
@@ -1503,6 +1568,8 @@ class MooncakeStoreWorker:
ready_event_sending,
self.enable_kv_events,
self.store_replicate_config,
enable_group_semantics=self.enable_group_semantics,
supports_group_ids=self._supports_group_ids,
record_operation=self._record_kv_connector_operation,
)
self.kv_send_thread.start()