[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:
co-authored by
OpenAI Codex
Yifan Qiao
parent
f5cd862dbd
commit
fd859d2e06
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user