From cd823daab1525e61a0628021daaf8d0ded1b7093 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=81=AA=E6=98=8E=E4=BC=81=E9=B9=85=5C70733?= <707334817@qq.com> Date: Mon, 6 Jul 2026 19:47:15 +0800 Subject: [PATCH] Archive vLLM EPD connector changes for CCF --- ccf-vllm-epd-archive/README.md | 28 + .../test_mooncake_store_hidden_connector.py | 419 +++ .../test_mooncake_store_hidden_protocol.py | 143 + .../unit/test_mooncake_store_hidden_worker.py | 850 +++++ .../ec_transfer/ec_connector/factory.py | 93 + .../mooncake_store_hidden/__init__.py | 9 + .../mooncake_store_hidden/connector.py | 377 +++ .../mooncake_store_hidden/data.py | 202 ++ .../mooncake_store_hidden/keys.py | 22 + .../mooncake_store_hidden/store_client.py | 567 ++++ .../mooncake_store_hidden/worker.py | 643 ++++ ...-ec-connector-b4482f0a1-full-feature.patch | 3004 +++++++++++++++++ ...onnector-b4482f0a1-latest-refinement.patch | 2018 +++++++++++ ...epd-hidden-ec-connector-review-fixes.patch | Bin 0 -> 33054 bytes 14 files changed, 8375 insertions(+) create mode 100644 ccf-vllm-epd-archive/README.md create mode 100644 ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py create mode 100644 ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py create mode 100644 ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/factory.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py create mode 100644 ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py create mode 100644 ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-full-feature.patch create mode 100644 ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-latest-refinement.patch create mode 100644 ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-review-fixes.patch diff --git a/ccf-vllm-epd-archive/README.md b/ccf-vllm-epd-archive/README.md new file mode 100644 index 00000000..a814e6df --- /dev/null +++ b/ccf-vllm-epd-archive/README.md @@ -0,0 +1,28 @@ +# CCF vLLM EPD 代码归档 + +本目录用于 CCF Mooncake 赛题提交归档。由于本赛题 GitLink 侧只提供 +Mooncake 仓库，而本部分工作实现于 vLLM，因此将相关 vLLM 改动源码快照 +和 patch 一并放入 Mooncake 仓库根目录，便于比赛评审在 GitLink 上查看。 + +对应的 vLLM 实现分支为： + +- 仓库：https://github.com/kanceler/vllm +- 分支：`mooncake-store-ec-hidden` +- 归档基准提交：`b4482f0a10a25ddbf34b687129c344594fab4613` + +## 目录内容 + +- `code/`：vLLM EPD Hidden State connector 相关源码快照，包括 connector + 实现、factory 注册入口和单元测试。 +- `patches/vllm-epd-hidden-ec-connector-b4482f0a1-full-feature.patch`：归档 + vLLM 实现的完整功能 patch。 +- `patches/vllm-epd-hidden-ec-connector-b4482f0a1-latest-refinement.patch`： + 归档实现过程中的 refinement patch。 +- `patches/vllm-epd-hidden-ec-connector-review-fixes.patch`：初赛提交前代码 + 审查阶段补充的增量修正，主要包括异步保存失败可观测性和 connector + 资源释放逻辑。 + +## 归档范围 + +本目录不作为 Mooncake 主工程的可合入代码，仅用于满足比赛要求中“上游软件 +托管至 GitLink”的提交要求。实际 vLLM 实现请以上述 vLLM 仓库和分支为准。 diff --git a/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py b/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py new file mode 100644 index 00000000..8759fa4e --- /dev/null +++ b/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden import ( + connector as connector_module, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.connector import ( + MooncakeStoreECConnector, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenPoolKey, + LoadSpec, + MMMeta, + MooncakeStoreConnectorMetadata, +) +from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange + + +class FakeWorker: + def __init__(self): + self.requests = [] + self.key_metadata = HiddenKeyMetadata( + cache_prefix="", + kind="encoder_output", + model_name="qwen", + encoder="encoder-config-a", + storage="replicated_object", + parallel="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", + tensor_layout=HIDDEN_TENSOR_LAYOUT, + ) + + def make_pool_key(self, identifier: str) -> HiddenPoolKey: + return HiddenPoolKey(self.key_metadata, identifier) + + def enqueue_save(self, request): + self.requests.append(request) + + def get_finished_sending(self): + return set() + + def get_failed_sending(self): + return {} + + +def make_connector(*, soft_pin_video_hidden: bool = False): + connector = MooncakeStoreECConnector.__new__(MooncakeStoreECConnector) + connector._is_producer = True + connector._is_consumer = False + connector.lookup_client = None + connector.lookup_async = True + connector.worker = FakeWorker() + connector._connector_metadata = None + connector.soft_pin_video_hidden = soft_pin_video_hidden + connector.load_specs = {} + connector.lookup_result_cache = {} + connector.identifier_waiters = {} + connector._candidate_consumes = {} + connector._candidate_loads = {} + connector._candidate_saves = {} + connector._load_modalities = {} + connector._save_modalities = {} + return connector + + +class FakeLookupClient: + def __init__(self, results): + self.results = list(results) + self.calls = [] + self.discarded = [] + + def lookup_batch(self, identifiers, non_block=True): + self.calls.append((tuple(identifiers), non_block)) + return self.results.pop(0) + + def discard(self, identifier): + self.discarded.append(identifier) + + +def make_request(request_id, features): + mm_features = [ + MultiModalFeatureSpec( + data=None, + modality=modality, + identifier=identifier, + mm_position=PlaceholderRange(offset=offset, length=length), + ) + for identifier, offset, length, modality in features + ] + return SimpleNamespace( + request_id=request_id, + mm_features=mm_features, + num_tokens=1000, + ) + + +def make_scheduler_output(*, finished_req_ids=None, preempted_req_ids=None): + return SimpleNamespace( + finished_req_ids=finished_req_ids or set(), + preempted_req_ids=preempted_req_ids, + ) + + +def test_build_hidden_key_metadata_uses_structured_key_fields(monkeypatch): + monkeypatch.setattr( + connector_module, + "get_tensor_model_parallel_world_size", + lambda: 4, + ) + monkeypatch.setattr( + connector_module, + "get_pcp_group", + lambda: SimpleNamespace(world_size=1), + ) + monkeypatch.setattr( + connector_module, + "get_dcp_group", + lambda: SimpleNamespace(world_size=1), + ) + multimodal_config = SimpleNamespace( + compute_hash=lambda: "encoder-config-a", + mm_encoder_tp_mode="data", + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + model="/models/qwen", + multimodal_config=multimodal_config, + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=2), + ec_transfer_config=SimpleNamespace( + ec_connector_extra_config={ + "cache_prefix": "shared-prefix", + "hidden_cache_prefix": "hidden-prefix", + } + ), + ) + + metadata = connector_module.build_hidden_key_metadata(vllm_config) + + assert metadata.cache_prefix == "hidden-prefix" + assert metadata.kind == "encoder_output" + assert metadata.model_name == "qwen" + assert metadata.encoder == "encoder-config-a" + assert metadata.storage == "replicated_object" + assert metadata.parallel == "tp:4@pp:2@pcp:1@dcp:1@mm_tp:data" + assert "storage" not in metadata.parallel + assert metadata.tensor_layout == "tensor" + + +def test_ensure_cache_available_defers_pending_batch_lookup(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_client = FakeLookupClient([None]) + request = make_request( + "req-1", + [ + ("image-1", 20, 60, "image"), + ("image-2", 500, 60, "image"), + ], + ) + + assert not connector.ensure_cache_available(request, num_computed_tokens=0) + + assert connector.lookup_client.calls == [ + (("image-1", "image-2"), True), + ] + assert connector.identifier_waiters == { + "image-1": {"req-1"}, + "image-2": {"req-1"}, + } + + +def test_ensure_cache_available_deduplicates_request_waiters_and_lookup_results(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_client = FakeLookupClient( + [ + {"image-1": True, "image-2": False}, + ] + ) + request = make_request( + "req-1", + [ + ("image-1", 20, 60, "image"), + ("image-2", 500, 60, "image"), + ], + ) + + assert connector.ensure_cache_available(request, num_computed_tokens=0) + assert connector.ensure_cache_available(request, num_computed_tokens=0) + + assert connector.lookup_client.calls == [ + (("image-1", "image-2"), True), + ] + assert connector.identifier_waiters == { + "image-1": {"req-1"}, + "image-2": {"req-1"}, + } + assert connector.lookup_result_cache == { + "image-1": True, + "image-2": False, + } + + +def test_has_cache_item_is_local_only(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_client = SimpleNamespace(lookup=lambda identifier: True) + connector.lookup_result_cache = {"image-1": True, "image-2": False} + + assert connector.has_cache_item("image-1") + assert not connector.has_cache_item("image-2") + assert not connector.has_cache_item("unknown") + + +def test_build_connector_meta_commits_waiter_consumes_and_keeps_unreached_image(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_result_cache = {"image-1": True, "image-2": True} + connector.identifier_waiters = { + "image-1": {"req-1"}, + "image-2": {"req-1"}, + } + connector.load_specs["image-1"] = LoadSpec(can_load=False) + request = make_request( + "req-1", + [ + ("image-1", 20, 60, "image"), + ("image-2", 500, 60, "image"), + ], + ) + + connector.update_state_after_alloc(request, 0) + meta = connector.build_connector_meta(make_scheduler_output()) + + assert [item.identifier for item in meta.items] == ["image-1"] + assert "image-1" not in connector.identifier_waiters + assert "image-1" not in connector.lookup_result_cache + assert connector.identifier_waiters == {"image-2": {"req-1"}} + assert connector.lookup_result_cache == {"image-2": True} + + +def test_build_connector_meta_rolls_back_preempted_candidate_state(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_result_cache = {"image-1": True} + connector.identifier_waiters = {"image-1": {"req-1"}} + connector.load_specs["image-1"] = LoadSpec(can_load=False) + request = make_request("req-1", [("image-1", 20, 60, "image")]) + + connector.update_state_after_alloc(request, 0) + meta = connector.build_connector_meta( + make_scheduler_output(preempted_req_ids={"req-1"}) + ) + + assert meta.items == [] + assert connector.identifier_waiters == {"image-1": {"req-1"}} + assert connector.lookup_result_cache == {"image-1": True} + assert "image-1" in connector.load_specs + + +def test_build_connector_meta_cleans_finished_waiters(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_client = FakeLookupClient([]) + connector.lookup_result_cache = {"image-1": True, "image-2": True} + connector.identifier_waiters = { + "image-1": {"req-1"}, + "image-2": {"req-1", "req-2"}, + } + + connector.build_connector_meta(make_scheduler_output(finished_req_ids={"req-1"})) + + assert "image-1" not in connector.identifier_waiters + assert "image-1" not in connector.lookup_result_cache + assert connector.identifier_waiters == {"image-2": {"req-2"}} + assert connector.lookup_result_cache == {"image-2": True} + assert connector.lookup_client.discarded == ["image-1"] + + +def test_cleanup_lookup_results_discards_inflight_lookup_without_waiters(): + connector = make_connector() + connector._is_consumer = True + connector._is_producer = False + connector.lookup_client = FakeLookupClient([]) + connector.identifier_waiters = {"image-1": set()} + connector.lookup_result_cache = {"image-1": True} + connector.load_specs["image-1"] = LoadSpec(can_load=False) + + connector._cleanup_lookup_results_without_waiters() + + assert connector.identifier_waiters == {} + assert connector.lookup_result_cache == {} + assert connector.load_specs == {} + assert connector.lookup_client.discarded == ["image-1"] + + +def test_build_connector_meta_merges_load_and_save_item_by_identifier(): + connector = make_connector() + connector._is_consumer = True + connector.load_specs["video-hash"] = LoadSpec(can_load=False) + connector.lookup_result_cache["video-hash"] = True + connector.identifier_waiters["video-hash"] = {"req-1"} + request = SimpleNamespace( + request_id="req-1", + mm_features=[ + SimpleNamespace( + identifier="video-hash", + modality="video", + ) + ], + ) + + connector.update_state_after_alloc(request, 0) + meta = connector.build_connector_meta(make_scheduler_output()) + + assert len(meta.items) == 1 + item = meta.items[0] + assert item.identifier == "video-hash" + assert item.modality == "video" + assert item.can_save + assert item.load_spec is not None + assert item.load_spec.can_load + assert connector.load_specs == {} + + +def test_save_caches_skips_items_without_save_plan(): + connector = make_connector() + connector.bind_connector_metadata( + MooncakeStoreConnectorMetadata( + items=[ + MMMeta( + identifier="image-hash", + modality="image", + can_save=False, + ) + ] + ) + ) + + connector.save_caches({"image-hash": torch.zeros((1, 2))}, "image-hash") + + assert connector.worker.requests == [] + + +def test_save_caches_enqueues_video_hidden_with_soft_pin(): + connector = make_connector(soft_pin_video_hidden=True) + tensor = torch.zeros((1, 2)) + connector.bind_connector_metadata( + MooncakeStoreConnectorMetadata( + items=[ + MMMeta( + identifier="video-hash", + modality="video", + can_save=True, + load_spec=LoadSpec(can_load=False), + ) + ] + ) + ) + + connector.save_caches({"video-hash": tensor}, "video-hash") + + assert len(connector.worker.requests) == 1 + request = connector.worker.requests[0] + assert request.identifier == "video-hash" + assert request.tensor is tensor + assert request.with_soft_pin + + +def test_save_caches_does_not_soft_pin_image_hidden(): + connector = make_connector(soft_pin_video_hidden=True) + connector.bind_connector_metadata( + MooncakeStoreConnectorMetadata( + items=[ + MMMeta( + identifier="image-hash", + modality="image", + can_save=True, + ) + ] + ) + ) + + connector.save_caches({"image-hash": torch.zeros((1, 2))}, "image-hash") + + assert len(connector.worker.requests) == 1 + assert not connector.worker.requests[0].with_soft_pin + + +def test_get_finished_logs_failed_hidden_saves(caplog): + class FailedWorker(FakeWorker): + def get_finished_sending(self): + return {"image-ok"} + + def get_failed_sending(self): + return {"image-failed": "batch put failed"} + + connector = make_connector() + connector.worker = FailedWorker() + + finished_sending, finished_recving = connector.get_finished({"req-1"}) + + assert finished_sending == {"image-ok"} + assert finished_recving is None + assert "hidden_store_save_failed" in caplog.text + assert "image-failed" in caplog.text + assert "batch put failed" in caplog.text diff --git a/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py b/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py new file mode 100644 index 00000000..d370014a --- /dev/null +++ b/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenPoolKey, + LoadSpec, + MMMeta, + MooncakeStoreConnectorMetadata, + build_tensor_meta, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( + make_hidden_data_key, +) + + +def make_pool_key( + identifier: str = "image-hash", + *, + cache_prefix: str = "", + kind: str = "encoder_output", + model_name: str = "qwen", + encoder: str = "encoder-config-a", + storage: str = "replicated_object", + parallel: str = "tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", + tensor_layout: str = HIDDEN_TENSOR_LAYOUT, +) -> HiddenPoolKey: + return HiddenPoolKey( + key_metadata=HiddenKeyMetadata( + cache_prefix=cache_prefix, + kind=kind, + model_name=model_name, + encoder=encoder, + storage=storage, + parallel=parallel, + tensor_layout=tensor_layout, + ), + identifier=identifier, + ) + + +def test_hidden_pool_key_is_the_single_tensor_object_key(): + pool_key = make_pool_key() + + data_key = make_hidden_data_key(pool_key) + + assert data_key == pool_key.to_string() + assert data_key.startswith("hidden@") + assert "kind:encoder_output" in data_key + assert "model:qwen" in data_key + assert "encoder:encoder-config-a" in data_key + assert "storage:replicated_object" in data_key + assert ( + "parallel:tp%3A1%40pp%3A1%40pcp%3A1%40dcp%3A1%40mm_tp%3Aweights" + in data_key + ) + assert "tensor_layout:tensor" in data_key + assert "storage%3Areplicated" not in data_key + assert "writer" not in data_key + assert "adapter:" not in data_key + assert "modality:" not in data_key + assert "image-hash" in data_key + + +def test_same_identifier_with_different_encoder_config_uses_different_keys(): + pool_key_a = make_pool_key(encoder="encoder-config-a") + pool_key_b = make_pool_key(encoder="encoder-config-b") + + assert make_hidden_data_key(pool_key_a) != make_hidden_data_key(pool_key_b) + + +def test_cache_prefix_namespaces_hidden_pool_key(): + pool_key_a = make_pool_key(cache_prefix="deployment-a") + pool_key_b = make_pool_key(cache_prefix="deployment-b") + + data_key_a = make_hidden_data_key(pool_key_a) + data_key_b = make_hidden_data_key(pool_key_b) + + assert data_key_a.startswith("deployment-a@hidden@") + assert data_key_b.startswith("deployment-b@hidden@") + assert data_key_a != data_key_b + + +def test_request_id_and_modality_are_not_part_of_hidden_pool_key(): + pool_key = make_pool_key(identifier="image-hash") + + assert "req-1" not in make_hidden_data_key(pool_key) + assert "request" not in make_hidden_data_key(pool_key) + assert "image@" not in make_hidden_data_key(pool_key) + assert "modality" not in make_hidden_data_key(pool_key) + + +def test_mm_meta_carries_hidden_item_plan(): + item = MMMeta( + identifier="image-hash", + modality="video", + can_save=True, + load_spec=LoadSpec(can_load=True), + ) + meta = MooncakeStoreConnectorMetadata(items=[item]) + + assert meta.items == [item] + assert meta.items[0].identifier == "image-hash" + assert meta.items[0].modality == "video" + assert meta.items[0].can_save + assert meta.items[0].load_spec is not None + assert meta.items[0].load_spec.can_load + + +def test_tensor_meta_describes_canonical_contiguous_tensor(): + pool_key = make_pool_key() + source = torch.zeros((4, 8), dtype=torch.float16).t() + stored = source.contiguous() + tensor_meta = build_tensor_meta(pool_key, stored) + + assert tensor_meta.pool_key == pool_key + assert tensor_meta.layout == HIDDEN_TENSOR_LAYOUT + assert tensor_meta.shape == tuple(stored.shape) + assert tensor_meta.dtype == "torch.float16" + assert tensor_meta.nbytes == stored.numel() * stored.element_size() + + +def test_tensor_meta_rejects_non_contiguous_tensor(): + pool_key = make_pool_key() + source = torch.zeros((4, 8), dtype=torch.float16).t() + + try: + build_tensor_meta(pool_key, source) + except ValueError as exc: + assert "contiguous" in str(exc) + else: + raise AssertionError("non-contiguous tensor descriptor should fail") + + +def test_pool_key_namespace_carries_reuse_compatibility(): + pool_key_a = make_pool_key(encoder="encoder-config-a") + pool_key_b = make_pool_key(encoder="encoder-config-b") + + assert pool_key_a != pool_key_b + assert make_hidden_data_key(pool_key_a) != make_hidden_data_key(pool_key_b) diff --git a/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py b/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py new file mode 100644 index 00000000..5bd1fe1f --- /dev/null +++ b/ccf-vllm-epd-archive/code/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py @@ -0,0 +1,850 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ctypes +import sys +import struct +import types +from concurrent.futures import Future + +import torch + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenPoolKey, + HiddenSaveRequest, + HiddenTensorDatabase, + LoadSpec, + MMMeta, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( + make_hidden_data_key, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( + HiddenStoreError, + HiddenStoreLoadError, + HiddenStoreSaveError, + MooncakeHiddenStoreClient, + _get_hidden_state_object_data_type, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.worker import ( + HiddenLookupClient, + HiddenLookupServer, + HiddenStoreSendingThread, + HiddenStoreWorker, +) + +TENSOR_METADATA_SIZE = 304 +TENSOR_OBJECT_MAGIC = 0x4D4F4F4E +TENSOR_OBJECT_VERSION = 1 +TORCH_DTYPE_TO_MOONCAKE_DTYPE = { + torch.float32: 0, + torch.float16: 11, + torch.bfloat16: 12, +} + + +class FakeStore: + def __init__(self): + self.objects = {} + self.batch_is_exist_calls = [] + self.registered = [] + self.unregistered = [] + self.pub_tensors = [] + self.range_gets = [] + self.fail_register_addrs = set() + self.raise_on_batch_put = False + self.batch_put_results = [0] + + def batch_is_exist(self, keys): + self.batch_is_exist_calls.append(list(keys)) + return [1 if key in self.objects else 0 for key in keys] + + def register_buffer(self, addr, size): + if addr in self.fail_register_addrs: + return -1 + self.registered.append((addr, size)) + return 0 + + def unregister_buffer(self, addr): + self.unregistered.append(addr) + return 0 + + def pub_tensor(self, key, tensor, replicate_config=None): + self.pub_tensors.append((key, tensor, replicate_config)) + self.objects[key] = _serialize_tensor_object(tensor) + return 0 + + def put_tensor(self, key, tensor): + return self.pub_tensor(key, tensor) + + def get_into_ranges( + self, + buffer_ptrs, + all_keys, + all_dst_offsets, + all_src_offsets, + all_sizes, + ): + self.range_gets.append( + (buffer_ptrs, all_keys, all_dst_offsets, all_src_offsets, all_sizes) + ) + results = [] + for buffer_ptr, keys, dst_offsets, src_offsets, sizes in zip( + buffer_ptrs, + all_keys, + all_dst_offsets, + all_src_offsets, + all_sizes, + strict=True, + ): + key_results = [] + for key, key_dst_offsets, key_src_offsets, key_sizes in zip( + keys, + dst_offsets, + src_offsets, + sizes, + strict=True, + ): + payload = self.objects.get(key) + fragment_results = [] + for dst_offset, src_offset, size in zip( + key_dst_offsets, + key_src_offsets, + key_sizes, + strict=True, + ): + if payload is None or src_offset + size > len(payload): + fragment_results.append(-1) + continue + ctypes.memmove( + buffer_ptr + dst_offset, + payload[src_offset : src_offset + size], + size, + ) + fragment_results.append(size) + key_results.append(fragment_results) + results.append(key_results) + return results + + +class FakeClosableStore(FakeStore): + def __init__(self): + super().__init__() + self.closed = False + + def close(self): + self.closed = True + + +class FakeStoreWithTeardown(FakeStore): + def __init__(self): + super().__init__() + self.teardown_called = False + + def teardown(self): + self.teardown_called = True + + +class FakeSocket: + def __init__(self): + self.closed = False + self.linger = None + + def close(self, linger=0): + self.closed = True + self.linger = linger + + +class FakeContext: + def __init__(self): + self.destroy_called = False + self.term_called = False + + def destroy(self, linger=0): + self.destroy_called = True + + def term(self): + self.term_called = True + + +class FakeThread: + def __init__(self): + self.join_called = False + self.timeout = None + + def join(self, timeout=None): + self.join_called = True + self.timeout = timeout + + def is_alive(self): + return False + + +class FakeBufferStore(FakeStore): + def batch_put_from_multi_buffers( + self, + keys, + buffer_ptrs, + buffer_sizes, + replicate_config=None, + ): + if self.raise_on_batch_put: + raise RuntimeError("batch put failed") + self.objects[keys[0]] = b"tensor-object" + return self.batch_put_results + + +class FakeReplicateConfig: + def __init__(self): + self.replica_num = 1 + self.nof_replica_num = 0 + self.with_soft_pin = False + self.with_hard_pin = False + self.preferred_segments = [] + self.preferred_nof_segments = [] + self.preferred_segment = "" + self.prefer_alloc_in_same_node = False + self.data_type = None + self.group_ids = None + + +class FakeReplicateConfigWithoutGroups: + def __init__(self): + self.replica_num = 1 + + +class FakeObjectDataTypeWithHidden: + HIDDEN_STATE = 10 + TENSOR = 2 + + +class FakeObjectDataTypeOnlyTensor: + TENSOR = 2 + + +class FakeObjectDataTypeNoTensor: + UNKNOWN = 0 + + +def make_pool_key(identifier: str = "image-hash") -> HiddenPoolKey: + return HiddenPoolKey( + key_metadata=HiddenKeyMetadata( + cache_prefix="", + kind="encoder_output", + model_name="qwen", + encoder="encoder-config-a", + storage="replicated_object", + parallel="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", + tensor_layout=HIDDEN_TENSOR_LAYOUT, + ), + identifier=identifier, + ) + + +def test_hidden_tensor_database_prepares_data_key_addrs_and_sizes(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + + key, addrs, sizes = HiddenTensorDatabase().prepare_value(pool_key, tensor) + + assert key == make_hidden_data_key(pool_key) + assert addrs == [tensor.data_ptr()] + assert sizes == [tensor.numel() * tensor.element_size()] + + +def test_store_client_checks_single_tensor_object_exists(): + pool_key = make_pool_key() + store = FakeStore() + client = MooncakeHiddenStoreClient(store) + + assert not client.exists(pool_key) + + store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" + assert client.exists(pool_key) + + +def test_worker_lookup_checks_existence_without_reading_tensor_metadata(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) + worker.save_tensor(pool_key, tensor) + + assert worker.lookup(pool_key.identifier) + assert not worker.lookup("missing-image-hash") + assert store.range_gets == [] + + +def test_worker_batch_lookup_checks_existence_in_one_store_call(): + pool_key_a = make_pool_key("image-a") + pool_key_b = make_pool_key("image-b") + store = FakeBufferStore() + store.objects[make_hidden_data_key(pool_key_a)] = b"tensor-object" + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key_a.key_metadata, + ) + + results = worker.lookup_batch(["image-a", "image-b"]) + + assert results == {"image-a": True, "image-b": False} + assert store.batch_is_exist_calls == [ + [make_hidden_data_key(pool_key_a), make_hidden_data_key(pool_key_b)] + ] + assert store.range_gets == [] + + +def test_lookup_client_discard_removes_identifier_future_mapping(): + client = HiddenLookupClient.__new__(HiddenLookupClient) + future: Future[dict[str, bool]] = Future() + client.futures = { + "image-a": future, + "image-b": future, + } + + client.discard("image-a") + + assert "image-a" not in client.futures + assert client.futures == {"image-b": future} + assert not future.cancelled() + + client.discard("image-b") + + assert client.futures == {} + assert future.cancelled() + + +def test_worker_lookup_records_minimal_operation_stats(): + pool_key = make_pool_key() + store = FakeBufferStore() + store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) + + assert worker.lookup_batch(["image-hash", "missing-image-hash"]) == { + "image-hash": True, + "missing-image-hash": False, + } + + stats = worker.get_operation_stats() + records = stats.data["lookup_exists"] + assert len(records) == 1 + assert records[0]["num_keys"] == 2 + assert records[0]["num_bytes"] == 0 + assert records[0]["status"] == "miss" + assert records[0]["num_failed_keys"] == 1 + assert worker.get_operation_stats() is None + + +def test_worker_save_stores_hidden_as_single_tensor_object(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient( + store, + replicate_config=FakeReplicateConfig(), + ), + tensor_database=HiddenTensorDatabase(), + ) + + worker.save_tensor(pool_key, tensor) + + assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) + assert store.pub_tensors[0][2] is not None + assert make_hidden_data_key(pool_key) in store.objects + + +def test_worker_save_rejects_dtype_that_load_cannot_decode(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float64) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + + try: + worker.save_tensor(pool_key, tensor) + except HiddenStoreSaveError as exc: + assert "unsupported hidden tensor dtype" in str(exc) + else: + raise AssertionError("unsupported hidden dtype should fail before store put") + + assert store.pub_tensors == [] + + +def test_buffer_put_unregisters_payload_and_metadata_buffers(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeBufferStore() + client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) + + client.put_tensor(pool_key, tensor) + + payload_addr = tensor.data_ptr() + metadata_addr = next( + addr for addr, size in store.registered if size == TENSOR_METADATA_SIZE + ) + assert payload_addr in store.unregistered + assert metadata_addr in store.unregistered + assert store.unregistered[-2:] == [metadata_addr, payload_addr] + + +def test_buffer_put_unregisters_payload_and_metadata_when_put_raises(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeBufferStore() + store.raise_on_batch_put = True + client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) + + try: + client.put_tensor(pool_key, tensor) + except RuntimeError as exc: + assert "batch put failed" in str(exc) + else: + raise AssertionError("batch put exception should propagate") + + payload_addr = tensor.data_ptr() + metadata_addr = next( + addr for addr, size in store.registered if size == TENSOR_METADATA_SIZE + ) + assert payload_addr in store.unregistered + assert metadata_addr in store.unregistered + + +def test_buffer_put_unregisters_payload_when_metadata_registration_fails(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeBufferStore() + original_register = store.register_buffer + + def register_buffer(addr, size): + if size == TENSOR_METADATA_SIZE: + store.fail_register_addrs.add(addr) + return original_register(addr, size) + + store.register_buffer = register_buffer + client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) + + try: + client.put_tensor(pool_key, tensor) + except HiddenStoreError: + pass + else: + raise AssertionError("metadata registration failure should raise") + + assert tensor.data_ptr() in store.unregistered + + +def test_worker_save_marks_hidden_state_data_type(monkeypatch): + fake_mooncake = types.ModuleType("mooncake") + fake_store = types.ModuleType("mooncake.store") + fake_store.ObjectDataType = FakeObjectDataTypeWithHidden + monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake) + monkeypatch.setitem(sys.modules, "mooncake.store", fake_store) + + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + replicate_config = FakeReplicateConfig() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient( + store, + replicate_config=replicate_config, + ), + tensor_database=HiddenTensorDatabase(), + ) + + worker.save_tensor(pool_key, tensor) + + used_config = store.pub_tensors[0][2] + assert used_config is not replicate_config + assert int(used_config.data_type) == 10 + + +def test_hidden_state_data_type_falls_back_to_tensor(monkeypatch): + fake_mooncake = types.ModuleType("mooncake") + fake_store = types.ModuleType("mooncake.store") + fake_store.ObjectDataType = FakeObjectDataTypeOnlyTensor + monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake) + monkeypatch.setitem(sys.modules, "mooncake.store", fake_store) + + assert _get_hidden_state_object_data_type() == FakeObjectDataTypeOnlyTensor.TENSOR + + +def test_hidden_state_data_type_missing_type_returns_none(monkeypatch): + fake_mooncake = types.ModuleType("mooncake") + fake_store = types.ModuleType("mooncake.store") + fake_store.ObjectDataType = FakeObjectDataTypeNoTensor + monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake) + monkeypatch.setitem(sys.modules, "mooncake.store", fake_store) + + assert _get_hidden_state_object_data_type() is None + + +def test_worker_save_does_not_require_mooncake_object_group_support(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient( + store, + replicate_config=FakeReplicateConfigWithoutGroups(), + ), + tensor_database=HiddenTensorDatabase(), + ) + + worker.save_tensor(pool_key, tensor) + + assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) + + +def test_worker_save_skips_existing_tensor_object(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + + worker.save_tensor(pool_key, tensor) + worker.save_tensor(pool_key, tensor) + + assert len(store.pub_tensors) == 1 + + +def test_worker_save_records_exists_and_put_operation_stats(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + + worker.save_tensor(pool_key, tensor) + + stats = worker.get_operation_stats() + assert stats.data["save_exists"][0]["status"] == "miss" + assert stats.data["save_exists"][0]["num_keys"] == 1 + assert stats.data["save_put"][0]["status"] == "ok" + assert stats.data["save_put"][0]["num_keys"] == 1 + assert stats.data["save_put"][0]["num_bytes"] == ( + tensor.numel() * tensor.element_size() + ) + + +def test_worker_save_existing_records_only_save_exists(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + + worker.save_tensor(pool_key, tensor) + + stats = worker.get_operation_stats() + assert stats.data["save_exists"][0]["status"] == "ok" + assert "save_put" not in stats.data + + +def test_sending_thread_stores_hidden_tensor_asynchronously(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + sending_thread = HiddenStoreSendingThread(worker) + sending_thread.start() + + sending_thread.add_request( + HiddenSaveRequest(pool_key=pool_key, tensor=tensor) + ) + sending_thread.request_queue.join() + + assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) + assert sending_thread.get_and_clear_finished_identifiers() == {pool_key.identifier} + sending_thread.close() + + +def test_sending_thread_records_failed_identifier_without_finishing(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeBufferStore() + store.raise_on_batch_put = True + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + sending_thread = HiddenStoreSendingThread(worker) + sending_thread.start() + + sending_thread.add_request( + HiddenSaveRequest(pool_key=pool_key, tensor=tensor) + ) + sending_thread.request_queue.join() + + assert sending_thread.get_and_clear_finished_identifiers() == set() + assert sending_thread.get_and_clear_failed_identifiers() == {pool_key.identifier} + assert pool_key.identifier in sending_thread.failure_reasons + assert worker.get_operation_stats().data["save_put"][0]["status"] == "error" + sending_thread.close() + + +def test_worker_drains_failed_sending_reasons(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeBufferStore() + store.raise_on_batch_put = True + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + worker.start_sending_thread() + + assert worker.sending_thread is not None + worker.enqueue_save(HiddenSaveRequest(pool_key=pool_key, tensor=tensor)) + worker.sending_thread.request_queue.join() + + failed = worker.get_failed_sending() + + assert set(failed) == {pool_key.identifier} + assert "batch put failed" in failed[pool_key.identifier] + assert worker.get_failed_sending() == {} + worker.shutdown() + + +def test_sending_thread_close_joins_worker_thread(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + sending_thread = HiddenStoreSendingThread(worker) + sending_thread.start() + sending_thread.add_request(HiddenSaveRequest(pool_key=pool_key, tensor=tensor)) + sending_thread.request_queue.join() + + sending_thread.close() + + assert not sending_thread.is_alive() + + +def test_worker_shutdown_closes_store_client(): + store = FakeClosableStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + ) + + worker.shutdown() + + assert store.closed + + +def test_store_client_close_uses_fallback_close_method(): + store = FakeStoreWithTeardown() + client = MooncakeHiddenStoreClient(store) + + client.close() + + assert store.teardown_called + + +def test_lookup_server_close_joins_thread_and_closes_context(tmp_path): + socket = FakeSocket() + ctx = FakeContext() + thread = FakeThread() + ipc_path = tmp_path / "hidden_lookup.ipc" + ipc_path.write_text("socket") + server = HiddenLookupServer.__new__(HiddenLookupServer) + server.running = True + server.socket = socket + server.ctx = ctx + server.thread = thread + server._ipc_path = str(ipc_path) + + server.close() + + assert not server.running + assert socket.closed + assert socket.linger == 0 + assert thread.join_called + assert ctx.destroy_called or ctx.term_called + assert not ipc_path.exists() + + +def test_lookup_client_close_shuts_down_executor_socket_and_context(): + socket = FakeSocket() + ctx = FakeContext() + executor = types.SimpleNamespace( + shutdown_called=False, + shutdown=lambda wait=False, cancel_futures=True: setattr( + executor, "shutdown_called", True + ), + ) + client = HiddenLookupClient.__new__(HiddenLookupClient) + client.executor = executor + client.futures = {"image-hash": Future()} + client.socket = socket + client.ctx = ctx + + client.close() + + assert executor.shutdown_called + assert client.futures == {} + assert socket.closed + assert socket.linger == 0 + assert ctx.destroy_called or ctx.term_called + + +def test_worker_load_gets_tensor_data_into_encoder_cache_before_returning(): + pool_key = make_pool_key() + stored = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) + worker.save_tensor(pool_key, stored) + + encoder_cache = {} + worker.load( + [MMMeta(identifier=pool_key.identifier, load_spec=LoadSpec(can_load=True))], + encoder_cache, + device="cpu", + ) + + assert pool_key.identifier in encoder_cache + assert tuple(encoder_cache[pool_key.identifier].shape) == tuple(stored.shape) + assert str(encoder_cache[pool_key.identifier].dtype) == str(stored.dtype) + assert store.range_gets[0][1] == [[make_hidden_data_key(pool_key)]] + assert store.range_gets[0][3] == [[[0]]] + assert store.range_gets[0][4] == [[[TENSOR_METADATA_SIZE]]] + assert store.range_gets[-1][1] == [[make_hidden_data_key(pool_key)]] + assert store.range_gets[-1][3] == [[[TENSOR_METADATA_SIZE]]] + + stats = worker.get_operation_stats() + assert stats.data["load_get"][0]["status"] == "ok" + assert stats.data["load_get"][0]["num_keys"] == 1 + assert stats.data["load_get"][0]["num_bytes"] == ( + stored.numel() * stored.element_size() + ) + + +def test_worker_load_records_error_without_writing_encoder_cache(): + pool_key = make_pool_key() + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) + encoder_cache = {} + + try: + worker.load( + [MMMeta(identifier=pool_key.identifier, load_spec=LoadSpec(can_load=True))], + encoder_cache, + device="cpu", + ) + except HiddenStoreLoadError: + pass + else: + raise AssertionError("missing hidden tensor should fail fast") + + assert pool_key.identifier not in encoder_cache + stats = worker.get_operation_stats() + assert stats.data["load_get"][0]["status"] == "error" + assert stats.data["load_get"][0]["num_failed_keys"] == 1 + + +def test_get_tensor_payload_unregisters_target_buffer_after_success(): + pool_key = make_pool_key() + stored = torch.zeros((2, 4), dtype=torch.float16) + store = FakeStore() + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) + worker.save_tensor(pool_key, stored) + target = torch.empty_like(stored) + + worker.store_client.get_tensor_payload( + pool_key, + target.data_ptr(), + target.numel() * target.element_size(), + TENSOR_METADATA_SIZE, + ) + + assert target.data_ptr() in store.unregistered + + +def test_get_tensor_payload_unregisters_target_buffer_after_load_error(): + pool_key = make_pool_key() + target = torch.empty((2, 4), dtype=torch.float16) + store = FakeStore() + client = MooncakeHiddenStoreClient(store) + + try: + client.get_tensor_payload( + pool_key, + target.data_ptr(), + target.numel() * target.element_size(), + TENSOR_METADATA_SIZE, + ) + except HiddenStoreLoadError: + pass + else: + raise AssertionError("missing payload should raise") + + assert target.data_ptr() in store.unregistered + + +def _serialize_tensor_object(tensor: torch.Tensor) -> bytes: + tensor = tensor.detach().cpu().contiguous() + nbytes = tensor.numel() * tensor.element_size() + header = struct.pack( + " bytes: + dims = list(shape) + [-1] * (8 - len(shape)) + return struct.pack("<8q", *dims) diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/factory.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/factory.py new file mode 100644 index 00000000..7fedd29f --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/factory.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from collections.abc import Callable +from typing import TYPE_CHECKING + +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorRole, +) +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import ECTransferConfig, VllmConfig + +logger = init_logger(__name__) + + +class ECConnectorFactory: + _registry: dict[str, Callable[[], type[ECConnectorBase]]] = {} + + @classmethod + def register_connector(cls, name: str, module_path: str, class_name: str) -> None: + """Register a connector with a lazy-loading module and class name.""" + if name in cls._registry: + raise ValueError(f"Connector '{name}' is already registered.") + + def loader() -> type[ECConnectorBase]: + module = importlib.import_module(module_path) + return getattr(module, class_name) + + cls._registry[name] = loader + + @classmethod + def create_connector( + cls, + config: "VllmConfig", + role: ECConnectorRole, + ) -> ECConnectorBase: + ec_transfer_config = config.ec_transfer_config + if ec_transfer_config is None: + raise ValueError("ec_transfer_config must be set to create a connector") + connector_cls = cls.get_connector_class(ec_transfer_config) + logger.info( + "Creating connector with name: %s and engine_id: %s", + connector_cls.__name__, + ec_transfer_config.engine_id, + ) + # Connector is explicitly separated into two roles. + # Scheduler connector: + # - Co-locate with scheduler process + # - Should only be used inside the Scheduler class + # Worker connector: + # - Co-locate with worker process + return connector_cls(config, role) + + @classmethod + def get_connector_class( + cls, ec_transfer_config: "ECTransferConfig" + ) -> type[ECConnectorBase]: + """Get the connector class by name.""" + connector_name = ec_transfer_config.ec_connector + if connector_name is None: + raise ValueError("EC connect must not be None") + connector_module_path = ec_transfer_config.ec_connector_module_path + if connector_module_path is not None and not connector_module_path: + raise ValueError("ec_connector_module_path cannot be an empty string.") + if connector_module_path: + connector_module = importlib.import_module(connector_module_path) + connector_cls = getattr(connector_module, connector_name) + elif connector_name in cls._registry: + connector_cls = cls._registry[connector_name]() + else: + raise ValueError(f"Unsupported connector type: {connector_name}") + return connector_cls + + +# Register various connectors here. +# The registration should not be done in each individual file, as we want to +# only load the files corresponding to the current connector. + +ECConnectorFactory.register_connector( + "ECExampleConnector", + "vllm.distributed.ec_transfer.ec_connector.example_connector", + "ECExampleConnector", +) + +ECConnectorFactory.register_connector( + "MooncakeStoreECConnector", + "vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden", + "MooncakeStoreECConnector", +) diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py new file mode 100644 index 00000000..77f80e01 --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hidden-state Mooncake Store EC connector support.""" + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.connector import ( + MooncakeStoreECConnector, +) + +__all__ = ["MooncakeStoreECConnector"] diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py new file mode 100644 index 00000000..d9dc005c --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EC connector backed by Mooncake Store for hidden-state tensors.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from vllm.distributed import ( + get_dcp_group, + get_pcp_group, + get_tensor_model_parallel_world_size, +) +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, + ECConnectorRole, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HIDDEN_OBJECT_KIND, + HIDDEN_STORAGE_LAYOUT, + HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenSaveRequest, + LoadSpec, + MMMeta, + MooncakeStoreConnectorMetadata, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( + MooncakeHiddenStoreClient, + create_mooncake_hidden_store_client, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.worker import ( + HiddenLookupClient, + HiddenLookupServer, + HiddenStoreWorker, +) +from vllm.logger import init_logger +from vllm.multimodal.utils import get_mm_features_in_window +from vllm.v1.core.sched.output import SchedulerOutput + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class MooncakeStoreECConnector(ECConnectorBase): + """Hidden-state EC connector that stores tensors in Mooncake Store.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: ECConnectorRole, + store_client: MooncakeHiddenStoreClient | None = None, + ): + super().__init__(vllm_config=vllm_config, role=role) + self.lookup_client: HiddenLookupClient | None = None + self.lookup_server: HiddenLookupServer | None = None + self.store_client: MooncakeHiddenStoreClient | None = None + self.worker: HiddenStoreWorker | None = None + assert vllm_config.ec_transfer_config is not None + extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config + self.soft_pin_video_hidden = bool( + extra_config.get("soft_pin_video_hidden", False) + ) + self.lookup_async = bool(extra_config.get("lookup_async", True)) + + if role == ECConnectorRole.SCHEDULER: + if self.is_consumer: + self.lookup_client = HiddenLookupClient(vllm_config) + else: + if not (self.is_producer or self.is_consumer): + return + hidden_key_metadata = build_hidden_key_metadata(vllm_config) + self.store_client = store_client or create_mooncake_hidden_store_client() + self.worker = HiddenStoreWorker( + store_client=self.store_client, + key_metadata=hidden_key_metadata, + ) + if self.is_producer: + self.worker.start_sending_thread() + if self.is_consumer and vllm_config.parallel_config.rank == 0: + self.lookup_server = HiddenLookupServer(self.worker, vllm_config) + + self.load_specs: dict[str, LoadSpec] = {} + self.lookup_result_cache: dict[str, bool] = {} + self.identifier_waiters: dict[str, set[str]] = {} + self._candidate_consumes: dict[str, set[str]] = {} + self._candidate_loads: dict[str, set[str]] = {} + self._candidate_saves: dict[str, set[str]] = {} + self._load_modalities: dict[str, str | None] = {} + self._save_modalities: dict[str, str | None] = {} + + def shutdown(self) -> None: + if self.lookup_client is not None: + self.lookup_client.close() + if self.lookup_server is not None: + self.lookup_server.close() + if self.worker is not None: + self.worker.shutdown() + + def has_cache_item(self, identifier: str) -> bool: + if not self.is_consumer: + return False + + if not self.lookup_result_cache.get(identifier, False): + self.load_specs.pop(identifier, None) + logger.info( + "hidden_store_scheduler_miss identifier=%s " + "reason=local_lookup_result_miss", + identifier, + ) + return False + + self.load_specs.setdefault(identifier, LoadSpec(can_load=False)) + logger.info( + "hidden_store_scheduler_hit identifier=%s", + identifier, + ) + return True + + def ensure_cache_available( + self, + request: Request, + num_computed_tokens: int, + ) -> bool: + if not self.is_consumer: + return True + if not request.mm_features: + return True + assert self.lookup_client is not None + + start = num_computed_tokens + end = request.num_tokens + lo, hi = get_mm_features_in_window(request.mm_features, start, end) + identifiers = list( + dict.fromkeys( + feature.identifier for feature in request.mm_features[lo:hi] + ) + ) + if not identifiers: + return True + + request_id = request.request_id + for identifier in identifiers: + self.identifier_waiters.setdefault(identifier, set()).add(request_id) + + unknown_identifiers = [ + identifier + for identifier in identifiers + if identifier not in self.lookup_result_cache + ] + if not unknown_identifiers: + return True + + lookup_results = self.lookup_client.lookup_batch( + unknown_identifiers, + non_block=self.lookup_async, + ) + if lookup_results is None: + return False + + for identifier in unknown_identifiers: + self.lookup_result_cache[identifier] = lookup_results.get( + identifier, + False, + ) + return True + + def update_state_after_alloc(self, request: Request, index: int) -> None: + mm_feature = request.mm_features[index] + identifier = mm_feature.identifier + modality = mm_feature.modality + request_id = request.request_id + + self._candidate_consumes.setdefault(request_id, set()).add(identifier) + + if self.is_consumer and identifier in self.load_specs: + self._candidate_loads.setdefault(request_id, set()).add(identifier) + self._load_modalities[identifier] = modality + + if self.is_producer: + self._save_modalities[identifier] = modality + self._candidate_saves.setdefault(request_id, set()).add(identifier) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> ECConnectorMetadata: + items_by_identifier: dict[str, MMMeta] = {} + preempted_ids = getattr(scheduler_output, "preempted_req_ids", None) or set() + + for request_id, identifiers in self._candidate_consumes.items(): + if request_id in preempted_ids: + continue + for identifier in identifiers: + waiters = self.identifier_waiters.get(identifier) + if waiters is not None: + waiters.discard(request_id) + + for request_id, identifiers in self._candidate_loads.items(): + if request_id in preempted_ids: + continue + for identifier in identifiers: + load_spec = self.load_specs.pop(identifier, None) + if load_spec is None: + continue + load_spec.can_load = True + items_by_identifier[identifier] = MMMeta( + identifier=identifier, + modality=self._load_modalities.get(identifier), + load_spec=load_spec, + ) + + for request_id, identifiers in self._candidate_saves.items(): + if request_id in preempted_ids: + continue + for identifier in identifiers: + item = items_by_identifier.get(identifier) + if item is None: + item = MMMeta( + identifier=identifier, + modality=self._save_modalities.get(identifier), + ) + items_by_identifier[identifier] = item + item.can_save = True + if item.modality is None: + item.modality = self._save_modalities.get(identifier) + + finished_req_ids = getattr(scheduler_output, "finished_req_ids", set()) + for finished_req_id in finished_req_ids: + for waiters in self.identifier_waiters.values(): + waiters.discard(finished_req_id) + + self._cleanup_lookup_results_without_waiters() + + metadata = MooncakeStoreConnectorMetadata( + items=list(items_by_identifier.values()), + ) + + self._candidate_consumes.clear() + self._candidate_loads.clear() + self._candidate_saves.clear() + self._load_modalities.clear() + self._save_modalities.clear() + return metadata + + def _cleanup_lookup_results_without_waiters(self) -> None: + for identifier, waiters in list(self.identifier_waiters.items()): + if waiters: + continue + del self.identifier_waiters[identifier] + self.lookup_result_cache.pop(identifier, None) + self.load_specs.pop(identifier, None) + if self.lookup_client is not None: + self.lookup_client.discard(identifier) + + def start_load_caches( + self, + encoder_cache: dict[str, torch.Tensor], + **kwargs, + ) -> None: + metadata = self._get_connector_metadata() + assert isinstance(metadata, MooncakeStoreConnectorMetadata) + assert self.worker is not None + self.worker.load( + metadata.items, + encoder_cache, + device=kwargs.get("device"), + ) + + def save_caches( + self, + encoder_cache: dict[str, torch.Tensor], + mm_hash: str, + **kwargs, + ) -> None: + if not self.is_producer: + return + assert self.worker is not None + identifier = mm_hash + if identifier not in encoder_cache: + logger.warning( + "Skip hidden store save; identifier %s is missing", + identifier, + ) + return + item = self._find_metadata_item(identifier) + if item is None or not item.can_save: + logger.debug( + "Skip hidden store save; identifier %s has no save plan", + identifier, + ) + return + pool_key = self.worker.make_pool_key(identifier) + self.worker.enqueue_save( + HiddenSaveRequest( + pool_key=pool_key, + tensor=encoder_cache[identifier], + with_soft_pin=self._should_soft_pin(item), + ) + ) + + def get_finished( + self, finished_req_ids: set[str] + ) -> tuple[set[str] | None, set[str] | None]: + if self.worker is None or not self.is_producer: + return None, None + finished_sending = self.worker.get_finished_sending() + failed_sending = self.worker.get_failed_sending() + for identifier, reason in failed_sending.items(): + logger.error( + "hidden_store_save_failed identifier=%s reason=%s", + identifier, + reason, + ) + return finished_sending or None, None + + def _find_metadata_item(self, identifier: str) -> MMMeta | None: + metadata = self._get_connector_metadata() + assert isinstance(metadata, MooncakeStoreConnectorMetadata) + for item in metadata.items: + if item.identifier == identifier: + return item + return None + + def _should_soft_pin(self, item: MMMeta) -> bool: + return self.soft_pin_video_hidden and item.modality == "video" + + +def build_hidden_key_metadata(vllm_config: VllmConfig) -> HiddenKeyMetadata: + model_config = vllm_config.model_config + parallel_config = vllm_config.parallel_config + assert vllm_config.ec_transfer_config is not None + extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config + + multimodal_config = getattr(model_config, "multimodal_config", None) + compute_hash = getattr(multimodal_config, "compute_hash", None) + mm_encoder_config_hash = ( + compute_hash() if callable(compute_hash) else "encoder:default" + ) + + tp_size = get_tensor_model_parallel_world_size() + pp_size = parallel_config.pipeline_parallel_size + pcp_size = get_pcp_group().world_size + dcp_size = get_dcp_group().world_size + mm_encoder_tp_mode = getattr( + multimodal_config, + "mm_encoder_tp_mode", + "unknown", + ) + parallel = ( + f"tp:{tp_size}" + f"@pp:{pp_size}" + f"@pcp:{pcp_size}" + f"@dcp:{dcp_size}" + f"@mm_tp:{mm_encoder_tp_mode}" + ) + + return HiddenKeyMetadata( + cache_prefix=str( + extra_config.get( + "hidden_cache_prefix", + extra_config.get("cache_prefix", ""), + ) + ), + kind=HIDDEN_OBJECT_KIND, + model_name=model_config.model.rstrip("/").split("/")[-1], + encoder=str(mm_encoder_config_hash), + storage=HIDDEN_STORAGE_LAYOUT, + parallel=parallel, + tensor_layout=HIDDEN_TENSOR_LAYOUT, + ) diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py new file mode 100644 index 00000000..f3fa09ef --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Data classes for the hidden-state Mooncake Store EC connector.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( + escape_key_part, + make_hidden_data_key, +) + +HIDDEN_OBJECT_KIND = "encoder_output" +HIDDEN_STORAGE_LAYOUT = "replicated_object" +HIDDEN_TENSOR_LAYOUT = "tensor" +HIDDEN_PROTOCOL_VERSION = "v1" +MOONCAKE_TENSOR_METADATA_NBYTES = 304 + + +@dataclass(frozen=True) +class HiddenKeyMetadata: + """Metadata that defines the semantic namespace for hidden reuse.""" + + cache_prefix: str + kind: str + model_name: str + encoder: str + storage: str + parallel: str + tensor_layout: str + + +@dataclass(frozen=True, order=True) +class HiddenPoolKey: + """Key for addressing one hidden tensor in the distributed store.""" + + key_metadata: HiddenKeyMetadata + identifier: str + + def to_string(self) -> str: + meta = self.key_metadata + prefix = ( + f"{escape_key_part(meta.cache_prefix)}@" if meta.cache_prefix else "" + ) + return ( + f"{prefix}hidden" + f"@kind:{escape_key_part(meta.kind)}" + f"@model:{escape_key_part(meta.model_name)}" + f"@encoder:{escape_key_part(meta.encoder)}" + f"@storage:{escape_key_part(meta.storage)}" + f"@parallel:{escape_key_part(meta.parallel)}" + f"@tensor_layout:{escape_key_part(meta.tensor_layout)}" + f"@id:{escape_key_part(self.identifier)}" + ) + + +@dataclass +class MMMeta: + """Per hidden object metadata passed from scheduler to worker.""" + + identifier: str + modality: str | None = None + can_save: bool = False + load_spec: LoadSpec | None = None + + +@dataclass(frozen=True) +class TensorMeta: + """Canonical contiguous tensor descriptor for one hidden store object.""" + + pool_key: HiddenPoolKey + protocol_version: str + layout: str + shape: tuple[int, ...] + dtype: str + nbytes: int + device_type: str + data_offset: int = MOONCAKE_TENSOR_METADATA_NBYTES + producer_stage: str = "encoder" + + +@dataclass +class LoadSpec: + """Specification for loading a hidden tensor from external store.""" + + can_load: bool = False + + +@dataclass +class HiddenSaveRequest: + """Specification for asynchronously storing one hidden tensor.""" + + pool_key: HiddenPoolKey + tensor: torch.Tensor + with_soft_pin: bool = False + + @property + def identifier(self) -> str: + return self.pool_key.identifier + + +@dataclass +class MooncakeStoreConnectorMetadata(ECConnectorMetadata): + """Metadata passed from scheduler to worker for hidden store operations.""" + + items: list[MMMeta] = field(default_factory=list) + + def add_item(self, item: MMMeta) -> None: + self.items.append(item) + + +@dataclass +class HiddenStoreOperationStats: + """Minimal per-operation telemetry aligned with Mooncake KV store stats.""" + + data: dict[str, list[dict[str, int | float | str]]] = field(default_factory=dict) + + def is_empty(self) -> bool: + return not self.data + + def record_operation( + self, + operation: str, + duration_seconds: float, + num_keys: int, + *, + num_bytes: int = 0, + status: str = "ok", + num_failed_keys: int = 0, + ) -> None: + self.data.setdefault(operation, []).append( + { + "duration_seconds": duration_seconds, + "num_keys": num_keys, + "num_bytes": num_bytes, + "status": status, + "num_failed_keys": num_failed_keys, + } + ) + + +class HiddenTensorDatabase: + """Maps hidden tensors to store keys and GPU memory descriptors.""" + + def prepare_value( + self, + pool_key: HiddenPoolKey, + tensor: torch.Tensor, + ) -> tuple[str, list[int], list[int]]: + return ( + make_hidden_data_key(pool_key), + [tensor.data_ptr()], + [tensor.numel() * tensor.element_size()], + ) + + +def build_tensor_meta( + pool_key: HiddenPoolKey, + tensor: torch.Tensor, +) -> TensorMeta: + """Build metadata for the canonical stored hidden tensor layout.""" + if not tensor.is_contiguous(): + raise ValueError("Hidden tensor descriptor requires a contiguous tensor") + + return TensorMeta( + pool_key=pool_key, + protocol_version=HIDDEN_PROTOCOL_VERSION, + layout=HIDDEN_TENSOR_LAYOUT, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype), + nbytes=tensor.numel() * tensor.element_size(), + device_type=tensor.device.type, + data_offset=MOONCAKE_TENSOR_METADATA_NBYTES, + ) + + +def validate_loaded_tensor(tensor: torch.Tensor, meta: TensorMeta) -> None: + if tuple(tensor.shape) != tuple(meta.shape): + raise ValueError( + "Hidden tensor shape mismatch: " + f"actual={tuple(tensor.shape)} expected={meta.shape}" + ) + + if str(tensor.dtype) != meta.dtype: + raise ValueError( + "Hidden tensor dtype mismatch: " + f"actual={tensor.dtype} expected={meta.dtype}" + ) + + actual_nbytes = tensor.numel() * tensor.element_size() + if actual_nbytes != meta.nbytes: + raise ValueError( + "Hidden tensor nbytes mismatch: " + f"actual={actual_nbytes} expected={meta.nbytes}" + ) + + if meta.layout != HIDDEN_TENSOR_LAYOUT: + raise ValueError(f"Unsupported hidden tensor layout: {meta.layout}") diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py new file mode 100644 index 00000000..b05c99b2 --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Store key helpers for the hidden-state Mooncake connector.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from urllib.parse import quote + +if TYPE_CHECKING: + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HiddenPoolKey, + ) + + +def escape_key_part(value: str) -> str: + """Escape one key component while keeping simple values readable.""" + return quote(str(value), safe="-_.~") + + +def make_hidden_data_key(pool_key: "HiddenPoolKey") -> str: + return pool_key.to_string() diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py new file mode 100644 index 00000000..21206ff6 --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py @@ -0,0 +1,567 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Thin Mooncake Store client for hidden-state objects.""" + +from __future__ import annotations + +import copy +import ctypes +import json +import os +import re +import struct +from dataclasses import dataclass +from typing import Any + +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HIDDEN_PROTOCOL_VERSION, + HIDDEN_TENSOR_LAYOUT, + MOONCAKE_TENSOR_METADATA_NBYTES, + HiddenPoolKey, + TensorMeta, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( + make_hidden_data_key, +) +from vllm.logger import init_logger +from vllm.utils.network_utils import get_ip + +logger = init_logger(__name__) + +DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 +DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 +_MOONCAKE_TENSOR_OBJECT_MAGIC = 0x4D4F4F4E +_MOONCAKE_TENSOR_OBJECT_VERSION = 1 +_MOONCAKE_TENSOR_HEADER_FORMAT = " MooncakeHiddenStoreConfig: + with open(file_path, encoding="utf-8") as file: + config = json.load(file) + mode = config.get("mode", "embedded") + return MooncakeHiddenStoreConfig( + metadata_server=config.get("metadata_server", ""), + master_server_address=config.get("master_server_address", ""), + protocol=config.get("protocol", "rdma"), + device_name=config.get("device_name", ""), + mode=mode, + global_segment_size=_parse_size( + config.get( + "global_segment_size", + 0 if mode == "standalone-store" else DEFAULT_GLOBAL_SEGMENT_SIZE, + ) + ), + local_buffer_size=_parse_size( + config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) + ), + ) + + @staticmethod + def load_from_env() -> MooncakeHiddenStoreConfig: + config_path = os.getenv("MOONCAKE_CONFIG_PATH") + if not config_path: + raise ValueError( + "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." + ) + return MooncakeHiddenStoreConfig.from_file(config_path) + + +def _parse_size(value: Any) -> int: + if isinstance(value, int): + return value + if not isinstance(value, str): + return int(value) + + cleaned = value.strip().lower() + match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) + if not match: + raise ValueError(f"Invalid size format: {value!r}") + + multipliers = { + "gb": 1024**3, + "mb": 1024**2, + "kb": 1024, + "b": 1, + None: 1, + } + return int(float(match.group(1)) * multipliers[match.group(2)]) + + +def create_mooncake_hidden_store_client() -> MooncakeHiddenStoreClient: + try: + from mooncake.store import ( # type: ignore + MooncakeDistributedStore, + ReplicateConfig, + ) + except ImportError as e: + raise ImportError( + "Please install mooncake to run vLLM with " "MooncakeStoreECConnector." + ) from e + + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils + + config = MooncakeHiddenStoreConfig.load_from_env() + config.device_name = rdma_utils.get_configured_worker_rnic( + protocol=config.protocol, + configured_device=config.device_name, + ) + + store = MooncakeDistributedStore() + local_ip = get_ip() + local_hostname = rdma_utils.get_requester_local_hostname(local_ip) + ret = store.setup( + local_hostname, + config.metadata_server, + config.global_segment_size, + config.local_buffer_size, + config.protocol, + config.device_name, + config.master_server_address, + ) + if ret != 0: + raise RuntimeError("Initialize MooncakeDistributedStore failed.") + + logger.info( + "Initialized hidden Mooncake store mode=%s global_segment_size=%d " + "local_buffer_size=%d", + config.mode, + config.global_segment_size, + config.local_buffer_size, + ) + return MooncakeHiddenStoreClient(store, replicate_config=ReplicateConfig()) + + +class HiddenStoreError(RuntimeError): + pass + + +class HiddenStoreLoadError(HiddenStoreError): + pass + + +class HiddenStoreSaveError(HiddenStoreError): + pass + + +class MooncakeHiddenStoreClient: + """Wraps Mooncake object and buffer APIs used by hidden transfer.""" + + def __init__(self, store: Any, replicate_config: Any | None = None): + self.store = store + self.replicate_config = replicate_config + + def close(self) -> None: + """Best-effort shutdown for Mooncake store implementations.""" + for method_name in ("close", "teardown", "disconnect", "finalize"): + close_fn = getattr(self.store, method_name, None) + if close_fn is None: + continue + try: + close_fn() + except Exception: + logger.warning( + "failed to close hidden Mooncake store with %s()", + method_name, + exc_info=True, + ) + return + + def exists(self, pool_key: HiddenPoolKey) -> bool: + data_key = make_hidden_data_key(pool_key) + states = self.store.batch_is_exist([data_key]) + return len(states) == 1 and states[0] == 1 + + def batch_exists(self, pool_keys: list[HiddenPoolKey]) -> list[bool]: + if not pool_keys: + return [] + + keys = [make_hidden_data_key(pool_key) for pool_key in pool_keys] + states = self.store.batch_is_exist(keys) + return [state == 1 for state in states] + + def get_tensor_meta(self, pool_key: HiddenPoolKey) -> TensorMeta | None: + metadata = self._read_range( + pool_key, + src_offset=0, + size=MOONCAKE_TENSOR_METADATA_NBYTES, + ) + if metadata is None: + return None + try: + return _decode_mooncake_tensor_metadata(pool_key, metadata) + except HiddenStoreLoadError: + logger.exception( + "failed to decode hidden Mooncake tensor metadata for %s", + pool_key.to_string(), + ) + return None + + def put_tensor( + self, + pool_key: HiddenPoolKey, + tensor: Any, + *, + with_soft_pin: bool = False, + ) -> None: + _validate_supported_hidden_tensor_dtype(tensor) + key = make_hidden_data_key(pool_key) + replicate_config = _make_hidden_replicate_config( + self.replicate_config, + with_soft_pin=with_soft_pin, + ) + batch_put_from_multi_buffers = getattr( + self.store, + "batch_put_from_multi_buffers", + None, + ) + if batch_put_from_multi_buffers is not None: + self._put_tensor_from_buffers( + pool_key, + tensor, + replicate_config=replicate_config, + ) + return + + if replicate_config is None: + put_fn = getattr(self.store, "put_tensor", None) + if put_fn is None: + raise HiddenStoreSaveError( + "Mooncake Hidden Store requires put_tensor or pub_tensor " + "support for single-object hidden tensors." + ) + ret = put_fn(key, tensor) + else: + put_fn = getattr(self.store, "pub_tensor", None) + if put_fn is None: + raise HiddenStoreSaveError( + "Mooncake Hidden Store requires pub_tensor support when " + "a ReplicateConfig is configured." + ) + ret = put_fn(key, tensor, replicate_config) + if ret != 0: + raise HiddenStoreSaveError( + f"failed to put hidden tensor for {pool_key.to_string()}: {ret}" + ) + + def _put_tensor_from_buffers( + self, + pool_key: HiddenPoolKey, + tensor: Any, + *, + replicate_config: Any | None, + ) -> None: + if not tensor.is_contiguous(): + raise HiddenStoreSaveError( + "hidden tensor must be contiguous before batch buffer put" + ) + data_size = tensor.numel() * tensor.element_size() + metadata = _encode_mooncake_tensor_metadata(tensor) + metadata_buffer = (ctypes.c_ubyte * len(metadata)).from_buffer_copy(metadata) + metadata_ptr = ctypes.addressof(metadata_buffer) + payload_ptr = tensor.data_ptr() + registered_addrs: list[int] = [] + try: + self.register_tensor(payload_ptr, data_size) + registered_addrs.append(payload_ptr) + self.register_tensor(metadata_ptr, len(metadata)) + registered_addrs.append(metadata_ptr) + + key = make_hidden_data_key(pool_key) + results = self.store.batch_put_from_multi_buffers( + [key], + [[metadata_ptr, payload_ptr]], + [[len(metadata), data_size]], + replicate_config, + ) + failed = [result for result in results if result < 0] + if failed: + raise HiddenStoreSaveError( + "failed to put hidden tensor for " + f"{pool_key.to_string()}: {failed}" + ) + finally: + for addr in reversed(registered_addrs): + self.unregister_tensor(addr) + + def register_tensor(self, addr: int, size: int) -> None: + ret = self.store.register_buffer(addr, size) + if ret != 0: + raise HiddenStoreError( + f"failed to register hidden buffer addr={addr:#x} size={size}: {ret}" + ) + + def unregister_tensor(self, addr: int) -> None: + unregister_fn = getattr(self.store, "unregister_buffer", None) + if unregister_fn is None: + return + try: + ret = unregister_fn(addr) + except Exception: + logger.warning( + "failed to unregister hidden buffer addr=%#x", + addr, + exc_info=True, + ) + return + if ret != 0: + logger.warning( + "unregister hidden buffer failed addr=%#x ret=%s", + addr, + ret, + ) + + def get_tensor_payload( + self, + pool_key: HiddenPoolKey, + addr: int, + size: int, + src_offset: int, + ) -> int: + self.register_tensor(addr, size) + try: + key = make_hidden_data_key(pool_key) + results = self.store.get_into_ranges( + [addr], + [[key]], + [[[0]]], + [[[src_offset]]], + [[[size]]], + ) + result = _single_range_result(results) + if result != size: + raise HiddenStoreLoadError( + "failed to get hidden tensor payload for " + f"{pool_key.to_string()}: {result}" + ) + return result + finally: + self.unregister_tensor(addr) + + def _read_range( + self, + pool_key: HiddenPoolKey, + *, + src_offset: int, + size: int, + ) -> bytes | None: + buffer = (ctypes.c_ubyte * size)() + buffer_ptr = ctypes.addressof(buffer) + self.register_tensor(buffer_ptr, size) + key = make_hidden_data_key(pool_key) + try: + results = self.store.get_into_ranges( + [buffer_ptr], + [[key]], + [[[0]]], + [[[src_offset]]], + [[[size]]], + ) + finally: + self.unregister_tensor(buffer_ptr) + if _single_range_result(results) != size: + return None + return bytes(buffer) + + +def _single_range_result(results: Any) -> int: + try: + return int(results[0][0][0]) + except Exception: + return -1 + + +def _decode_mooncake_tensor_metadata( + pool_key: HiddenPoolKey, + metadata: bytes, +) -> TensorMeta: + if len(metadata) < MOONCAKE_TENSOR_METADATA_NBYTES: + raise HiddenStoreLoadError( + f"hidden tensor metadata is too small: {len(metadata)}" + ) + ( + magic, + version, + header_size, + dtype, + ndim, + _layout_kind, + _reserved_flags, + data_offset, + data_bytes, + ) = struct.unpack_from(_MOONCAKE_TENSOR_HEADER_FORMAT, metadata, 0) + if ( + magic != _MOONCAKE_TENSOR_OBJECT_MAGIC + or version != _MOONCAKE_TENSOR_OBJECT_VERSION + or header_size != MOONCAKE_TENSOR_METADATA_NBYTES + ): + raise HiddenStoreLoadError( + "invalid Mooncake tensor metadata header for " f"{pool_key.to_string()}" + ) + if ndim < 0 or ndim > 8: + raise HiddenStoreLoadError( + f"invalid hidden tensor ndim for {pool_key.to_string()}: {ndim}" + ) + if dtype not in _MOONCAKE_DTYPE_TO_TORCH_DTYPE: + raise HiddenStoreLoadError( + f"unsupported Mooncake tensor dtype for {pool_key.to_string()}: {dtype}" + ) + local_shape = struct.unpack_from( + "<8q", + metadata, + _MOONCAKE_TENSOR_LOCAL_SHAPE_OFFSET, + ) + shape = tuple(int(dim) for dim in local_shape[:ndim]) + if any(dim < 0 for dim in shape): + raise HiddenStoreLoadError( + f"invalid hidden tensor shape for {pool_key.to_string()}: {shape}" + ) + return TensorMeta( + pool_key=pool_key, + protocol_version=HIDDEN_PROTOCOL_VERSION, + layout=HIDDEN_TENSOR_LAYOUT, + shape=shape, + dtype=_MOONCAKE_DTYPE_TO_TORCH_DTYPE[dtype], + nbytes=int(data_bytes), + device_type="cpu", + data_offset=int(data_offset), + ) + + +def _encode_mooncake_tensor_metadata(tensor: Any) -> bytes: + dtype = str(tensor.dtype) + _validate_supported_hidden_tensor_dtype(tensor) + shape = tuple(int(dim) for dim in tensor.shape) + if len(shape) > 8: + raise HiddenStoreSaveError( + f"hidden tensor has too many dimensions: {len(shape)}" + ) + nbytes = tensor.numel() * tensor.element_size() + header = struct.pack( + _MOONCAKE_TENSOR_HEADER_FORMAT, + _MOONCAKE_TENSOR_OBJECT_MAGIC, + _MOONCAKE_TENSOR_OBJECT_VERSION, + MOONCAKE_TENSOR_METADATA_NBYTES, + _TORCH_DTYPE_TO_MOONCAKE_DTYPE[dtype], + len(shape), + 0, + 0, + MOONCAKE_TENSOR_METADATA_NBYTES, + nbytes, + ) + dims = shape + (-1,) * (8 - len(shape)) + tensor_shape = struct.pack("<8q", *dims) + axes = b"\0" * (32 * 4) + metadata = header + tensor_shape + tensor_shape + struct.pack(" None: + dtype = str(tensor.dtype) + if dtype not in _SUPPORTED_HIDDEN_TORCH_DTYPES: + raise HiddenStoreSaveError(f"unsupported hidden tensor dtype: {dtype}") + + +def _make_hidden_replicate_config( + replicate_config: Any | None, + *, + with_soft_pin: bool, +) -> Any | None: + if replicate_config is None: + return None + + config = _clone_replicate_config(replicate_config) + hidden_state_data_type = _get_hidden_state_object_data_type() + if hidden_state_data_type is not None and hasattr(config, "data_type"): + config.data_type = hidden_state_data_type + if hasattr(config, "with_soft_pin"): + config.with_soft_pin = bool(config.with_soft_pin) or with_soft_pin + return config + + +def _clone_replicate_config(replicate_config: Any) -> Any: + try: + return copy.copy(replicate_config) + except Exception: + config = type(replicate_config)() + for attr in ( + "replica_num", + "nof_replica_num", + "with_soft_pin", + "with_hard_pin", + "preferred_segments", + "preferred_nof_segments", + "preferred_segment", + "prefer_alloc_in_same_node", + "data_type", + "group_ids", + ): + if hasattr(replicate_config, attr) and hasattr(config, attr): + setattr(config, attr, getattr(replicate_config, attr)) + return config + + +def _get_hidden_state_object_data_type() -> Any | None: + try: + from mooncake.store import ObjectDataType # type: ignore + except Exception: + return None + hidden_state_type = getattr(ObjectDataType, "HIDDEN_STATE", None) + if hidden_state_type is not None: + return hidden_state_type + return getattr(ObjectDataType, "TENSOR", None) diff --git a/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py new file mode 100644 index 00000000..89fe5a22 --- /dev/null +++ b/ccf-vllm-epd-archive/code/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker-side hidden-state load/save logic for Mooncake Store.""" + +from __future__ import annotations + +import os +import queue +import socket +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor + +import torch +import zmq + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( + HiddenKeyMetadata, + HiddenPoolKey, + HiddenSaveRequest, + HiddenStoreOperationStats, + HiddenTensorDatabase, + MMMeta, + build_tensor_meta, + validate_loaded_tensor, +) +from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( + HiddenStoreLoadError, + MooncakeHiddenStoreClient, +) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( + get_mooncake_dp_engine_index, +) +from vllm.logger import init_logger +from vllm.utils.network_utils import make_zmq_socket + +logger = init_logger(__name__) + +LOOKUP_MSG = b"LOOKUP" +BATCH_LOOKUP_MSG = b"BATCH_LOOKUP" +RESP_BATCH = b"BATCH" +RESP_HIT = b"HIT" +RESP_MISS = b"MISS" +RESP_ERR = b"ERR" +THREAD_JOIN_TIMEOUT_SECONDS = 5.0 + + +class HiddenStoreWorker: + """Synchronous hidden tensor load/save path used by the EC connector.""" + + def __init__( + self, + store_client: MooncakeHiddenStoreClient, + tensor_database: HiddenTensorDatabase | None = None, + key_metadata: HiddenKeyMetadata | None = None, + ): + self.store_client = store_client + self.tensor_database = tensor_database or HiddenTensorDatabase() + self.key_metadata = key_metadata + self.sending_thread: HiddenStoreSendingThread | None = None + self._operation_stats_lock = threading.Lock() + self._operation_stats = HiddenStoreOperationStats() + + def make_pool_key(self, identifier: str) -> HiddenPoolKey: + assert self.key_metadata is not None + return HiddenPoolKey( + key_metadata=self.key_metadata, + identifier=identifier, + ) + + def start_sending_thread(self) -> None: + if self.sending_thread is not None: + return + self.sending_thread = HiddenStoreSendingThread(self) + self.sending_thread.start() + + def enqueue_save(self, request: HiddenSaveRequest) -> None: + if self.sending_thread is None: + self.save_tensor( + request.pool_key, + request.tensor, + with_soft_pin=request.with_soft_pin, + ) + return + self.sending_thread.add_request(request) + + def get_finished_sending(self) -> set[str]: + if self.sending_thread is None: + return set() + return self.sending_thread.get_and_clear_finished_identifiers() + + def get_failed_sending(self) -> dict[str, str]: + if self.sending_thread is None: + return {} + return self.sending_thread.get_and_clear_failure_reasons() + + def get_operation_stats(self) -> HiddenStoreOperationStats | None: + with self._operation_stats_lock: + if self._operation_stats.is_empty(): + return None + stats = self._operation_stats + self._operation_stats = HiddenStoreOperationStats() + return stats + + def _record_operation( + self, + operation: str, + duration_seconds: float, + num_keys: int, + *, + num_bytes: int = 0, + status: str = "ok", + num_failed_keys: int = 0, + ) -> None: + with self._operation_stats_lock: + self._operation_stats.record_operation( + operation=operation, + duration_seconds=duration_seconds, + num_keys=num_keys, + num_bytes=num_bytes, + status=status, + num_failed_keys=num_failed_keys, + ) + + def shutdown(self) -> None: + if self.sending_thread is not None: + self.sending_thread.close() + self.sending_thread = None + close_fn = getattr(self.store_client, "close", None) + if close_fn is not None: + close_fn() + + def lookup(self, identifier: str) -> bool: + """Return whether the hidden object exists in Mooncake Store.""" + return self.lookup_batch([identifier]).get(identifier, False) + + def lookup_batch(self, identifiers: list[str]) -> dict[str, bool]: + """Return whether hidden objects exist in Mooncake Store.""" + pool_keys = [self.make_pool_key(identifier) for identifier in identifiers] + started = time.perf_counter() + try: + exists = self.store_client.batch_exists(pool_keys) + except Exception: + self._record_operation( + "lookup_exists", + time.perf_counter() - started, + len(pool_keys), + status="error", + num_failed_keys=len(pool_keys), + ) + raise + + failed_keys = sum(1 for hit in exists if not hit) + self._record_operation( + "lookup_exists", + time.perf_counter() - started, + len(pool_keys), + status="miss" if failed_keys else "ok", + num_failed_keys=failed_keys, + ) + results = dict(zip(identifiers, exists, strict=True)) + for pool_key, hit in zip(pool_keys, exists, strict=True): + if hit: + logger.info( + "hidden_store_lookup_hit identifier=%s hidden_pool_key=%s", + pool_key.identifier, + pool_key.to_string(), + ) + else: + logger.info( + "hidden_store_lookup_miss identifier=%s hidden_pool_key=%s " + "reason=missing_object", + pool_key.identifier, + pool_key.to_string(), + ) + return results + + def save_tensor( + self, + pool_key: HiddenPoolKey, + tensor: torch.Tensor, + with_soft_pin: bool = False, + ) -> None: + exists_started = time.perf_counter() + try: + exists = self.store_client.exists(pool_key) + except Exception: + self._record_operation( + "save_exists", + time.perf_counter() - exists_started, + 1, + status="error", + num_failed_keys=1, + ) + raise + + self._record_operation( + "save_exists", + time.perf_counter() - exists_started, + 1, + status="ok" if exists else "miss", + ) + if exists: + logger.info( + "hidden_store_save_skip identifier=%s hidden_pool_key=%s " + "reason=exists", + pool_key.identifier, + pool_key.to_string(), + ) + return + + started = time.perf_counter() + stored_tensor = tensor if tensor.is_contiguous() else tensor.contiguous() + used_staging = stored_tensor is not tensor + tensor_meta = build_tensor_meta(pool_key, stored_tensor) + try: + self.store_client.put_tensor( + pool_key, + stored_tensor, + with_soft_pin=with_soft_pin, + ) + except Exception: + self._record_operation( + "save_put", + time.perf_counter() - started, + 1, + num_bytes=tensor_meta.nbytes, + status="error", + num_failed_keys=1, + ) + raise + self._record_operation( + "save_put", + time.perf_counter() - started, + 1, + num_bytes=tensor_meta.nbytes, + status="ok", + ) + logger.info( + "hidden_store_put identifier=%s hidden_pool_key=%s nbytes=%d " + "used_staging=%s hidden_store_put_ms=%.3f", + pool_key.identifier, + pool_key.to_string(), + tensor_meta.nbytes, + used_staging, + (time.perf_counter() - started) * 1000.0, + ) + + def load( + self, + items: list[MMMeta], + encoder_cache: dict[str, torch.Tensor], + *, + device: torch.device | str | None = None, + ) -> None: + for item in items: + load_spec = item.load_spec + if load_spec is None or not load_spec.can_load: + continue + if item.identifier in encoder_cache: + logger.debug( + "hidden_store_load_skip identifier=%s " + "reason=local_encoder_cache", + item.identifier, + ) + continue + + started = time.perf_counter() + pool_key = self.make_pool_key(item.identifier) + tensor_meta = None + load_stage = "metadata" + try: + tensor_meta = self.store_client.get_tensor_meta(pool_key) + if tensor_meta is None: + raise HiddenStoreLoadError( + "failed to load hidden tensor metadata for " + f"{pool_key.to_string()}" + ) + + load_stage = "allocate" + target_device = device + if target_device is None: + target_device = "cuda" if torch.cuda.is_available() else None + target = torch.empty( + tensor_meta.shape, + dtype=_resolve_torch_dtype(tensor_meta.dtype), + device=target_device, + ) + _data_key, addrs, sizes = self.tensor_database.prepare_value( + pool_key, + target, + ) + load_stage = "payload" + self.store_client.get_tensor_payload( + pool_key, + addrs[0], + sizes[0], + tensor_meta.data_offset, + ) + load_stage = "validate" + validate_loaded_tensor(target, tensor_meta) + except Exception as e: + self._record_operation( + "load_get", + time.perf_counter() - started, + 1, + num_bytes=tensor_meta.nbytes if tensor_meta is not None else 0, + status="error", + num_failed_keys=1, + ) + logger.exception( + "hidden_store_load_failed identifier=%s hidden_pool_key=%s " + "stage=%s shape=%s dtype=%s nbytes=%s error=%s", + item.identifier, + pool_key.to_string(), + load_stage, + tensor_meta.shape if tensor_meta is not None else None, + tensor_meta.dtype if tensor_meta is not None else None, + tensor_meta.nbytes if tensor_meta is not None else 0, + e, + ) + raise + encoder_cache[item.identifier] = target + self._record_operation( + "load_get", + time.perf_counter() - started, + 1, + num_bytes=tensor_meta.nbytes, + status="ok", + ) + logger.info( + "hidden_store_get identifier=%s hidden_pool_key=%s nbytes=%d " + "hidden_store_get_ms=%.3f", + item.identifier, + pool_key.to_string(), + tensor_meta.nbytes, + (time.perf_counter() - started) * 1000.0, + ) + + +def _resolve_torch_dtype(dtype: str) -> torch.dtype: + if dtype == "torch.float16": + return torch.float16 + if dtype == "torch.bfloat16": + return torch.bfloat16 + if dtype == "torch.float32": + return torch.float32 + raise HiddenStoreLoadError(f"unsupported hidden tensor dtype: {dtype}") + + +class HiddenStoreSendingThread(threading.Thread): + """Background thread for storing hidden tensors to the store.""" + + def __init__(self, store_worker: HiddenStoreWorker): + super().__init__(daemon=True, name="HiddenStoreSendingThread") + self.store_worker = store_worker + self.request_queue: queue.Queue[HiddenSaveRequest | None] = queue.Queue() + self.done_task_lock = threading.Lock() + self.finished_identifiers: set[str] = set() + self.failed_identifiers: set[str] = set() + self.failure_reasons: dict[str, str] = {} + self._closed = threading.Event() + + def add_request(self, request: HiddenSaveRequest) -> None: + self.request_queue.put(request) + + def get_and_clear_finished_identifiers(self) -> set[str]: + with self.done_task_lock: + finished = self.finished_identifiers.copy() + self.finished_identifiers.clear() + return finished + + def get_and_clear_failed_identifiers(self) -> set[str]: + with self.done_task_lock: + failed = self.failed_identifiers.copy() + self.failed_identifiers.clear() + return failed + + def get_and_clear_failure_reasons(self) -> dict[str, str]: + with self.done_task_lock: + failures = { + identifier: self.failure_reasons.get(identifier, "") + for identifier in self.failed_identifiers + } + for identifier in self.failed_identifiers: + self.failure_reasons.pop(identifier, None) + self.failed_identifiers.clear() + return failures + + def set_finished_identifier(self, identifier: str) -> None: + with self.done_task_lock: + self.finished_identifiers.add(identifier) + + def set_failed_identifier(self, identifier: str, error: Exception) -> None: + with self.done_task_lock: + self.failed_identifiers.add(identifier) + self.failure_reasons[identifier] = str(error) + + def run(self) -> None: + while True: + request = self.request_queue.get() + try: + if request is None: + return + self.store_worker.save_tensor( + request.pool_key, + request.tensor, + with_soft_pin=request.with_soft_pin, + ) + self.set_finished_identifier(request.identifier) + except Exception as e: + if request is not None: + self.set_failed_identifier(request.identifier, e) + logger.error("Error in %s: %s", self.name, e) + finally: + self.request_queue.task_done() + + def close(self) -> None: + if self._closed.is_set(): + return + self._closed.set() + self.request_queue.put(None) + if threading.current_thread() is not self: + self.join(timeout=THREAD_JOIN_TIMEOUT_SECONDS) + if self.is_alive(): + logger.warning( + "%s did not exit within %.1f seconds", + self.name, + THREAD_JOIN_TIMEOUT_SECONDS, + ) + + +class HiddenLookupServer: + """Worker rank-0 admin channel for scheduler-side hidden lookups.""" + + def __init__( + self, + store_worker: HiddenStoreWorker, + vllm_config: VllmConfig, + ): + self.ctx = zmq.Context() # type: ignore[attr-defined] + socket_path = get_zmq_rpc_path_hidden_lookup(vllm_config) + self._ipc_path = socket_path.removeprefix("ipc://") + if os.path.exists(self._ipc_path): + os.unlink(self._ipc_path) + self.socket = make_zmq_socket( + self.ctx, + socket_path, + zmq.REP, # type: ignore[attr-defined] + bind=True, + ) + + self.store_worker = store_worker + self.running = True + + def process_request(): + while self.running: + try: + all_frames = self.socket.recv_multipart(copy=False) + except zmq.error.ZMQError: + if not self.running: + return + logger.exception("HiddenLookupServer recv failed") + continue + msg_type = bytes(all_frames[0]) + + if msg_type == LOOKUP_MSG: + try: + identifier = bytes(all_frames[1]).decode("utf-8") + exists = self.store_worker.lookup(identifier) + if not exists: + self.socket.send_multipart([RESP_MISS]) + else: + self.socket.send_multipart([RESP_HIT]) + except Exception: + logger.exception("HiddenLookupServer lookup failed") + self.socket.send_multipart([RESP_ERR]) + elif msg_type == BATCH_LOOKUP_MSG: + try: + identifiers = [ + bytes(frame).decode("utf-8") for frame in all_frames[1:] + ] + exists = self.store_worker.lookup_batch(identifiers) + frames = [ + RESP_HIT if exists.get(identifier, False) else RESP_MISS + for identifier in identifiers + ] + self.socket.send_multipart([RESP_BATCH, *frames]) + except Exception: + logger.exception("HiddenLookupServer batch lookup failed") + self.socket.send_multipart([RESP_ERR]) + else: + logger.warning( + "HiddenLookupServer received unknown msg_type: %r", + msg_type, + ) + self.socket.send_multipart([RESP_ERR]) + + self.thread = threading.Thread(target=process_request, daemon=True) + self.thread.start() + + def close(self): + self.running = False + self.socket.close(linger=0) + self.thread.join(timeout=THREAD_JOIN_TIMEOUT_SECONDS) + if self.thread.is_alive(): + logger.warning( + "HiddenLookupServer thread did not exit within %.1f seconds", + THREAD_JOIN_TIMEOUT_SECONDS, + ) + _close_zmq_context(self.ctx) + if os.path.exists(self._ipc_path): + os.unlink(self._ipc_path) + + +class HiddenLookupClient: + """Scheduler-side client for worker rank-0 hidden lookup queries.""" + + def __init__(self, vllm_config: VllmConfig): + self.ctx = zmq.Context() # type: ignore[attr-defined] + socket_path = get_zmq_rpc_path_hidden_lookup(vllm_config) + self.socket = make_zmq_socket( + self.ctx, + socket_path, + zmq.REQ, # type: ignore[attr-defined] + bind=False, + ) + self.executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="HiddenLookupClient", + ) + self.futures: dict[str, Future[dict[str, bool]]] = {} + + def lookup(self, identifier: str) -> bool: + result = self.lookup_batch([identifier], non_block=False) + assert result is not None + return result.get(identifier, False) + + def _lookup_batch(self, identifiers: list[str]) -> dict[str, bool]: + self.socket.send_multipart( + [ + BATCH_LOOKUP_MSG, + *(identifier.encode("utf-8") for identifier in identifiers), + ] + ) + resp = self.socket.recv_multipart() + msg_type = bytes(resp[0]) + if msg_type == RESP_BATCH: + states = [bytes(frame) == RESP_HIT for frame in resp[1:]] + if len(states) != len(identifiers): + logger.warning( + "HiddenLookupClient received malformed batch response: " + "identifiers=%d states=%d", + len(identifiers), + len(states), + ) + return {identifier: False for identifier in identifiers} + return dict(zip(identifiers, states, strict=True)) + if msg_type == RESP_ERR: + return {identifier: False for identifier in identifiers} + logger.warning("HiddenLookupClient received unknown response: %r", msg_type) + return {identifier: False for identifier in identifiers} + + def lookup_batch( + self, + identifiers: list[str], + non_block: bool = False, + ) -> dict[str, bool] | None: + identifiers = list(dict.fromkeys(identifiers)) + if not identifiers: + return {} + + new_identifiers = [ + identifier for identifier in identifiers if identifier not in self.futures + ] + if new_identifiers: + future = self.executor.submit(self._lookup_batch, new_identifiers) + for identifier in new_identifiers: + self.futures[identifier] = future + + if non_block and any( + not self.futures[identifier].done() for identifier in identifiers + ): + return None + + results: dict[str, bool] = {} + for identifier in identifiers: + future = self.futures[identifier] + try: + batch_results = future.result() + results[identifier] = batch_results.get(identifier, False) + except Exception as e: + logger.error("Async hidden lookup failed for %s: %s", identifier, e) + results[identifier] = False + finally: + self.futures.pop(identifier, None) + return results + + def discard(self, identifier: str) -> None: + future = self.futures.pop(identifier, None) + if future is None: + return + if not any(existing is future for existing in self.futures.values()): + future.cancel() + + def close(self): + self.executor.shutdown(wait=False, cancel_futures=True) + self.futures.clear() + self.socket.close(linger=0) + _close_zmq_context(self.ctx) + + +def get_zmq_rpc_path_hidden_lookup(vllm_config: VllmConfig) -> str: + """Construct IPC path for Hidden Store lookup socket.""" + assert vllm_config.ec_transfer_config is not None + dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config) + base_url = envs.VLLM_RPC_BASE_PATH + hostname = socket.gethostname() + extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config + rpc_port = extra_config.get( + "hidden_lookup_rpc_port", + extra_config.get("lookup_rpc_port", 0), + ) + logger.debug("Hidden lookup Base URL: %s, RPC Port: %s", base_url, rpc_port) + return ( + f"ipc://{base_url}/hidden_lookup_rpc_port_{rpc_port}_host_{hostname}" + f"_dp_rank{dp_rank}" + ) + + +def _close_zmq_context(ctx) -> None: + try: + destroy = getattr(ctx, "destroy", None) + if destroy is not None: + destroy(linger=0) + return + term = getattr(ctx, "term", None) + if term is not None: + term() + except Exception: + logger.warning("failed to close hidden lookup ZMQ context", exc_info=True) diff --git a/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-full-feature.patch b/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-full-feature.patch new file mode 100644 index 00000000..248c7cef --- /dev/null +++ b/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-full-feature.patch @@ -0,0 +1,3004 @@ +diff --git a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py +new file mode 100644 +index 000000000..df8baf237 +--- /dev/null ++++ b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py +@@ -0,0 +1,393 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++from types import SimpleNamespace ++ ++import torch ++ ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden import ( ++ connector as connector_module, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.connector import ( ++ MooncakeStoreECConnector, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HIDDEN_TENSOR_LAYOUT, ++ HiddenKeyMetadata, ++ HiddenPoolKey, ++ LoadSpec, ++ MMMeta, ++ MooncakeStoreConnectorMetadata, ++) ++from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange ++ ++ ++class FakeWorker: ++ def __init__(self): ++ self.requests = [] ++ self.key_metadata = HiddenKeyMetadata( ++ cache_prefix="", ++ kind="encoder_output", ++ model_name="qwen", ++ encoder="encoder-config-a", ++ storage="replicated_object", ++ parallel="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", ++ tensor_layout=HIDDEN_TENSOR_LAYOUT, ++ ) ++ ++ def make_pool_key(self, identifier: str) -> HiddenPoolKey: ++ return HiddenPoolKey(self.key_metadata, identifier) ++ ++ def enqueue_save(self, request): ++ self.requests.append(request) ++ ++ ++def make_connector(*, soft_pin_video_hidden: bool = False): ++ connector = MooncakeStoreECConnector.__new__(MooncakeStoreECConnector) ++ connector._is_producer = True ++ connector._is_consumer = False ++ connector.lookup_client = None ++ connector.lookup_async = True ++ connector.worker = FakeWorker() ++ connector._connector_metadata = None ++ connector.soft_pin_video_hidden = soft_pin_video_hidden ++ connector.load_specs = {} ++ connector.lookup_result_cache = {} ++ connector.identifier_waiters = {} ++ connector._candidate_consumes = {} ++ connector._candidate_loads = {} ++ connector._candidate_saves = {} ++ connector._load_modalities = {} ++ connector._save_modalities = {} ++ return connector ++ ++ ++class FakeLookupClient: ++ def __init__(self, results): ++ self.results = list(results) ++ self.calls = [] ++ self.discarded = [] ++ ++ def lookup_batch(self, identifiers, non_block=True): ++ self.calls.append((tuple(identifiers), non_block)) ++ return self.results.pop(0) ++ ++ def discard(self, identifier): ++ self.discarded.append(identifier) ++ ++ ++def make_request(request_id, features): ++ mm_features = [ ++ MultiModalFeatureSpec( ++ data=None, ++ modality=modality, ++ identifier=identifier, ++ mm_position=PlaceholderRange(offset=offset, length=length), ++ ) ++ for identifier, offset, length, modality in features ++ ] ++ return SimpleNamespace( ++ request_id=request_id, ++ mm_features=mm_features, ++ num_tokens=1000, ++ ) ++ ++ ++def make_scheduler_output(*, finished_req_ids=None, preempted_req_ids=None): ++ return SimpleNamespace( ++ finished_req_ids=finished_req_ids or set(), ++ preempted_req_ids=preempted_req_ids, ++ ) ++ ++ ++def test_build_hidden_key_metadata_uses_structured_key_fields(monkeypatch): ++ monkeypatch.setattr( ++ connector_module, ++ "get_tensor_model_parallel_world_size", ++ lambda: 4, ++ ) ++ monkeypatch.setattr( ++ connector_module, ++ "get_pcp_group", ++ lambda: SimpleNamespace(world_size=1), ++ ) ++ monkeypatch.setattr( ++ connector_module, ++ "get_dcp_group", ++ lambda: SimpleNamespace(world_size=1), ++ ) ++ multimodal_config = SimpleNamespace( ++ compute_hash=lambda: "encoder-config-a", ++ mm_encoder_tp_mode="data", ++ ) ++ vllm_config = SimpleNamespace( ++ model_config=SimpleNamespace( ++ model="/models/qwen", ++ multimodal_config=multimodal_config, ++ ), ++ parallel_config=SimpleNamespace(pipeline_parallel_size=2), ++ ec_transfer_config=SimpleNamespace( ++ ec_connector_extra_config={ ++ "cache_prefix": "shared-prefix", ++ "hidden_cache_prefix": "hidden-prefix", ++ } ++ ), ++ ) ++ ++ metadata = connector_module.build_hidden_key_metadata(vllm_config) ++ ++ assert metadata.cache_prefix == "hidden-prefix" ++ assert metadata.kind == "encoder_output" ++ assert metadata.model_name == "qwen" ++ assert metadata.encoder == "encoder-config-a" ++ assert metadata.storage == "replicated_object" ++ assert metadata.parallel == "tp:4@pp:2@pcp:1@dcp:1@mm_tp:data" ++ assert "storage" not in metadata.parallel ++ assert metadata.tensor_layout == "tensor" ++ ++ ++def test_ensure_cache_available_defers_pending_batch_lookup(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient([None]) ++ request = make_request( ++ "req-1", ++ [ ++ ("image-1", 20, 60, "image"), ++ ("image-2", 500, 60, "image"), ++ ], ++ ) ++ ++ assert not connector.ensure_cache_available(request, num_computed_tokens=0) ++ ++ assert connector.lookup_client.calls == [ ++ (("image-1", "image-2"), True), ++ ] ++ assert connector.identifier_waiters == { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1"}, ++ } ++ ++ ++def test_ensure_cache_available_deduplicates_request_waiters_and_lookup_results(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient( ++ [ ++ {"image-1": True, "image-2": False}, ++ ] ++ ) ++ request = make_request( ++ "req-1", ++ [ ++ ("image-1", 20, 60, "image"), ++ ("image-2", 500, 60, "image"), ++ ], ++ ) ++ ++ assert connector.ensure_cache_available(request, num_computed_tokens=0) ++ assert connector.ensure_cache_available(request, num_computed_tokens=0) ++ ++ assert connector.lookup_client.calls == [ ++ (("image-1", "image-2"), True), ++ ] ++ assert connector.identifier_waiters == { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1"}, ++ } ++ assert connector.lookup_result_cache == { ++ "image-1": True, ++ "image-2": False, ++ } ++ ++ ++def test_has_cache_item_is_local_only(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = SimpleNamespace(lookup=lambda identifier: True) ++ connector.lookup_result_cache = {"image-1": True, "image-2": False} ++ ++ assert connector.has_cache_item("image-1") ++ assert not connector.has_cache_item("image-2") ++ assert not connector.has_cache_item("unknown") ++ ++ ++def test_build_connector_meta_commits_waiter_consumes_and_keeps_unreached_image(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_result_cache = {"image-1": True, "image-2": True} ++ connector.identifier_waiters = { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1"}, ++ } ++ connector.load_specs["image-1"] = LoadSpec(can_load=False) ++ request = make_request( ++ "req-1", ++ [ ++ ("image-1", 20, 60, "image"), ++ ("image-2", 500, 60, "image"), ++ ], ++ ) ++ ++ connector.update_state_after_alloc(request, 0) ++ meta = connector.build_connector_meta(make_scheduler_output()) ++ ++ assert [item.identifier for item in meta.items] == ["image-1"] ++ assert "image-1" not in connector.identifier_waiters ++ assert "image-1" not in connector.lookup_result_cache ++ assert connector.identifier_waiters == {"image-2": {"req-1"}} ++ assert connector.lookup_result_cache == {"image-2": True} ++ ++ ++def test_build_connector_meta_rolls_back_preempted_candidate_state(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_result_cache = {"image-1": True} ++ connector.identifier_waiters = {"image-1": {"req-1"}} ++ connector.load_specs["image-1"] = LoadSpec(can_load=False) ++ request = make_request("req-1", [("image-1", 20, 60, "image")]) ++ ++ connector.update_state_after_alloc(request, 0) ++ meta = connector.build_connector_meta( ++ make_scheduler_output(preempted_req_ids={"req-1"}) ++ ) ++ ++ assert meta.items == [] ++ assert connector.identifier_waiters == {"image-1": {"req-1"}} ++ assert connector.lookup_result_cache == {"image-1": True} ++ assert "image-1" in connector.load_specs ++ ++ ++def test_build_connector_meta_cleans_finished_waiters(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient([]) ++ connector.lookup_result_cache = {"image-1": True, "image-2": True} ++ connector.identifier_waiters = { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1", "req-2"}, ++ } ++ ++ connector.build_connector_meta(make_scheduler_output(finished_req_ids={"req-1"})) ++ ++ assert "image-1" not in connector.identifier_waiters ++ assert "image-1" not in connector.lookup_result_cache ++ assert connector.identifier_waiters == {"image-2": {"req-2"}} ++ assert connector.lookup_result_cache == {"image-2": True} ++ assert connector.lookup_client.discarded == ["image-1"] ++ ++ ++def test_cleanup_lookup_results_discards_inflight_lookup_without_waiters(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient([]) ++ connector.identifier_waiters = {"image-1": set()} ++ connector.lookup_result_cache = {"image-1": True} ++ connector.load_specs["image-1"] = LoadSpec(can_load=False) ++ ++ connector._cleanup_lookup_results_without_waiters() ++ ++ assert connector.identifier_waiters == {} ++ assert connector.lookup_result_cache == {} ++ assert connector.load_specs == {} ++ assert connector.lookup_client.discarded == ["image-1"] ++ ++ ++def test_build_connector_meta_merges_load_and_save_item_by_identifier(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector.load_specs["video-hash"] = LoadSpec(can_load=False) ++ connector.lookup_result_cache["video-hash"] = True ++ connector.identifier_waiters["video-hash"] = {"req-1"} ++ request = SimpleNamespace( ++ request_id="req-1", ++ mm_features=[ ++ SimpleNamespace( ++ identifier="video-hash", ++ modality="video", ++ ) ++ ], ++ ) ++ ++ connector.update_state_after_alloc(request, 0) ++ meta = connector.build_connector_meta(make_scheduler_output()) ++ ++ assert len(meta.items) == 1 ++ item = meta.items[0] ++ assert item.identifier == "video-hash" ++ assert item.modality == "video" ++ assert item.can_save ++ assert item.load_spec is not None ++ assert item.load_spec.can_load ++ assert connector.load_specs == {} ++ ++ ++def test_save_caches_skips_items_without_save_plan(): ++ connector = make_connector() ++ connector.bind_connector_metadata( ++ MooncakeStoreConnectorMetadata( ++ items=[ ++ MMMeta( ++ identifier="image-hash", ++ modality="image", ++ can_save=False, ++ ) ++ ] ++ ) ++ ) ++ ++ connector.save_caches({"image-hash": torch.zeros((1, 2))}, "image-hash") ++ ++ assert connector.worker.requests == [] ++ ++ ++def test_save_caches_enqueues_video_hidden_with_soft_pin(): ++ connector = make_connector(soft_pin_video_hidden=True) ++ tensor = torch.zeros((1, 2)) ++ connector.bind_connector_metadata( ++ MooncakeStoreConnectorMetadata( ++ items=[ ++ MMMeta( ++ identifier="video-hash", ++ modality="video", ++ can_save=True, ++ load_spec=LoadSpec(can_load=False), ++ ) ++ ] ++ ) ++ ) ++ ++ connector.save_caches({"video-hash": tensor}, "video-hash") ++ ++ assert len(connector.worker.requests) == 1 ++ request = connector.worker.requests[0] ++ assert request.identifier == "video-hash" ++ assert request.tensor is tensor ++ assert request.with_soft_pin ++ ++ ++def test_save_caches_does_not_soft_pin_image_hidden(): ++ connector = make_connector(soft_pin_video_hidden=True) ++ connector.bind_connector_metadata( ++ MooncakeStoreConnectorMetadata( ++ items=[ ++ MMMeta( ++ identifier="image-hash", ++ modality="image", ++ can_save=True, ++ ) ++ ] ++ ) ++ ) ++ ++ connector.save_caches({"image-hash": torch.zeros((1, 2))}, "image-hash") ++ ++ assert len(connector.worker.requests) == 1 ++ assert not connector.worker.requests[0].with_soft_pin +diff --git a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py +new file mode 100644 +index 000000000..d370014ab +--- /dev/null ++++ b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py +@@ -0,0 +1,143 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import torch ++ ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HIDDEN_TENSOR_LAYOUT, ++ HiddenKeyMetadata, ++ HiddenPoolKey, ++ LoadSpec, ++ MMMeta, ++ MooncakeStoreConnectorMetadata, ++ build_tensor_meta, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( ++ make_hidden_data_key, ++) ++ ++ ++def make_pool_key( ++ identifier: str = "image-hash", ++ *, ++ cache_prefix: str = "", ++ kind: str = "encoder_output", ++ model_name: str = "qwen", ++ encoder: str = "encoder-config-a", ++ storage: str = "replicated_object", ++ parallel: str = "tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", ++ tensor_layout: str = HIDDEN_TENSOR_LAYOUT, ++) -> HiddenPoolKey: ++ return HiddenPoolKey( ++ key_metadata=HiddenKeyMetadata( ++ cache_prefix=cache_prefix, ++ kind=kind, ++ model_name=model_name, ++ encoder=encoder, ++ storage=storage, ++ parallel=parallel, ++ tensor_layout=tensor_layout, ++ ), ++ identifier=identifier, ++ ) ++ ++ ++def test_hidden_pool_key_is_the_single_tensor_object_key(): ++ pool_key = make_pool_key() ++ ++ data_key = make_hidden_data_key(pool_key) ++ ++ assert data_key == pool_key.to_string() ++ assert data_key.startswith("hidden@") ++ assert "kind:encoder_output" in data_key ++ assert "model:qwen" in data_key ++ assert "encoder:encoder-config-a" in data_key ++ assert "storage:replicated_object" in data_key ++ assert ( ++ "parallel:tp%3A1%40pp%3A1%40pcp%3A1%40dcp%3A1%40mm_tp%3Aweights" ++ in data_key ++ ) ++ assert "tensor_layout:tensor" in data_key ++ assert "storage%3Areplicated" not in data_key ++ assert "writer" not in data_key ++ assert "adapter:" not in data_key ++ assert "modality:" not in data_key ++ assert "image-hash" in data_key ++ ++ ++def test_same_identifier_with_different_encoder_config_uses_different_keys(): ++ pool_key_a = make_pool_key(encoder="encoder-config-a") ++ pool_key_b = make_pool_key(encoder="encoder-config-b") ++ ++ assert make_hidden_data_key(pool_key_a) != make_hidden_data_key(pool_key_b) ++ ++ ++def test_cache_prefix_namespaces_hidden_pool_key(): ++ pool_key_a = make_pool_key(cache_prefix="deployment-a") ++ pool_key_b = make_pool_key(cache_prefix="deployment-b") ++ ++ data_key_a = make_hidden_data_key(pool_key_a) ++ data_key_b = make_hidden_data_key(pool_key_b) ++ ++ assert data_key_a.startswith("deployment-a@hidden@") ++ assert data_key_b.startswith("deployment-b@hidden@") ++ assert data_key_a != data_key_b ++ ++ ++def test_request_id_and_modality_are_not_part_of_hidden_pool_key(): ++ pool_key = make_pool_key(identifier="image-hash") ++ ++ assert "req-1" not in make_hidden_data_key(pool_key) ++ assert "request" not in make_hidden_data_key(pool_key) ++ assert "image@" not in make_hidden_data_key(pool_key) ++ assert "modality" not in make_hidden_data_key(pool_key) ++ ++ ++def test_mm_meta_carries_hidden_item_plan(): ++ item = MMMeta( ++ identifier="image-hash", ++ modality="video", ++ can_save=True, ++ load_spec=LoadSpec(can_load=True), ++ ) ++ meta = MooncakeStoreConnectorMetadata(items=[item]) ++ ++ assert meta.items == [item] ++ assert meta.items[0].identifier == "image-hash" ++ assert meta.items[0].modality == "video" ++ assert meta.items[0].can_save ++ assert meta.items[0].load_spec is not None ++ assert meta.items[0].load_spec.can_load ++ ++ ++def test_tensor_meta_describes_canonical_contiguous_tensor(): ++ pool_key = make_pool_key() ++ source = torch.zeros((4, 8), dtype=torch.float16).t() ++ stored = source.contiguous() ++ tensor_meta = build_tensor_meta(pool_key, stored) ++ ++ assert tensor_meta.pool_key == pool_key ++ assert tensor_meta.layout == HIDDEN_TENSOR_LAYOUT ++ assert tensor_meta.shape == tuple(stored.shape) ++ assert tensor_meta.dtype == "torch.float16" ++ assert tensor_meta.nbytes == stored.numel() * stored.element_size() ++ ++ ++def test_tensor_meta_rejects_non_contiguous_tensor(): ++ pool_key = make_pool_key() ++ source = torch.zeros((4, 8), dtype=torch.float16).t() ++ ++ try: ++ build_tensor_meta(pool_key, source) ++ except ValueError as exc: ++ assert "contiguous" in str(exc) ++ else: ++ raise AssertionError("non-contiguous tensor descriptor should fail") ++ ++ ++def test_pool_key_namespace_carries_reuse_compatibility(): ++ pool_key_a = make_pool_key(encoder="encoder-config-a") ++ pool_key_b = make_pool_key(encoder="encoder-config-b") ++ ++ assert pool_key_a != pool_key_b ++ assert make_hidden_data_key(pool_key_a) != make_hidden_data_key(pool_key_b) +diff --git a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py +new file mode 100644 +index 000000000..c1fd96b18 +--- /dev/null ++++ b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py +@@ -0,0 +1,687 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import ctypes ++import sys ++import struct ++import types ++from concurrent.futures import Future ++ ++import torch ++ ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HIDDEN_TENSOR_LAYOUT, ++ HiddenKeyMetadata, ++ HiddenPoolKey, ++ HiddenSaveRequest, ++ HiddenTensorDatabase, ++ LoadSpec, ++ MMMeta, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( ++ make_hidden_data_key, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( ++ HiddenStoreError, ++ HiddenStoreLoadError, ++ HiddenStoreSaveError, ++ MooncakeHiddenStoreClient, ++ _get_hidden_state_object_data_type, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.worker import ( ++ HiddenLookupClient, ++ HiddenStoreSendingThread, ++ HiddenStoreWorker, ++) ++ ++TENSOR_METADATA_SIZE = 304 ++TENSOR_OBJECT_MAGIC = 0x4D4F4F4E ++TENSOR_OBJECT_VERSION = 1 ++TORCH_DTYPE_TO_MOONCAKE_DTYPE = { ++ torch.float32: 0, ++ torch.float16: 11, ++ torch.bfloat16: 12, ++} ++ ++ ++class FakeStore: ++ def __init__(self): ++ self.objects = {} ++ self.batch_is_exist_calls = [] ++ self.registered = [] ++ self.unregistered = [] ++ self.pub_tensors = [] ++ self.range_gets = [] ++ self.fail_register_addrs = set() ++ self.raise_on_batch_put = False ++ self.batch_put_results = [0] ++ ++ def batch_is_exist(self, keys): ++ self.batch_is_exist_calls.append(list(keys)) ++ return [1 if key in self.objects else 0 for key in keys] ++ ++ def register_buffer(self, addr, size): ++ if addr in self.fail_register_addrs: ++ return -1 ++ self.registered.append((addr, size)) ++ return 0 ++ ++ def unregister_buffer(self, addr): ++ self.unregistered.append(addr) ++ return 0 ++ ++ def pub_tensor(self, key, tensor, replicate_config=None): ++ self.pub_tensors.append((key, tensor, replicate_config)) ++ self.objects[key] = _serialize_tensor_object(tensor) ++ return 0 ++ ++ def put_tensor(self, key, tensor): ++ return self.pub_tensor(key, tensor) ++ ++ def get_into_ranges( ++ self, ++ buffer_ptrs, ++ all_keys, ++ all_dst_offsets, ++ all_src_offsets, ++ all_sizes, ++ ): ++ self.range_gets.append( ++ (buffer_ptrs, all_keys, all_dst_offsets, all_src_offsets, all_sizes) ++ ) ++ results = [] ++ for buffer_ptr, keys, dst_offsets, src_offsets, sizes in zip( ++ buffer_ptrs, ++ all_keys, ++ all_dst_offsets, ++ all_src_offsets, ++ all_sizes, ++ strict=True, ++ ): ++ key_results = [] ++ for key, key_dst_offsets, key_src_offsets, key_sizes in zip( ++ keys, ++ dst_offsets, ++ src_offsets, ++ sizes, ++ strict=True, ++ ): ++ payload = self.objects.get(key) ++ fragment_results = [] ++ for dst_offset, src_offset, size in zip( ++ key_dst_offsets, ++ key_src_offsets, ++ key_sizes, ++ strict=True, ++ ): ++ if payload is None or src_offset + size > len(payload): ++ fragment_results.append(-1) ++ continue ++ ctypes.memmove( ++ buffer_ptr + dst_offset, ++ payload[src_offset : src_offset + size], ++ size, ++ ) ++ fragment_results.append(size) ++ key_results.append(fragment_results) ++ results.append(key_results) ++ return results ++ ++ ++class FakeBufferStore(FakeStore): ++ def batch_put_from_multi_buffers( ++ self, ++ keys, ++ buffer_ptrs, ++ buffer_sizes, ++ replicate_config=None, ++ ): ++ if self.raise_on_batch_put: ++ raise RuntimeError("batch put failed") ++ self.objects[keys[0]] = b"tensor-object" ++ return self.batch_put_results ++ ++ ++class FakeReplicateConfig: ++ def __init__(self): ++ self.replica_num = 1 ++ self.nof_replica_num = 0 ++ self.with_soft_pin = False ++ self.with_hard_pin = False ++ self.preferred_segments = [] ++ self.preferred_nof_segments = [] ++ self.preferred_segment = "" ++ self.prefer_alloc_in_same_node = False ++ self.data_type = None ++ self.group_ids = None ++ ++ ++class FakeReplicateConfigWithoutGroups: ++ def __init__(self): ++ self.replica_num = 1 ++ ++ ++class FakeObjectDataTypeWithHidden: ++ HIDDEN_STATE = 10 ++ TENSOR = 2 ++ ++ ++class FakeObjectDataTypeOnlyTensor: ++ TENSOR = 2 ++ ++ ++class FakeObjectDataTypeNoTensor: ++ UNKNOWN = 0 ++ ++ ++def make_pool_key(identifier: str = "image-hash") -> HiddenPoolKey: ++ return HiddenPoolKey( ++ key_metadata=HiddenKeyMetadata( ++ cache_prefix="", ++ kind="encoder_output", ++ model_name="qwen", ++ encoder="encoder-config-a", ++ storage="replicated_object", ++ parallel="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", ++ tensor_layout=HIDDEN_TENSOR_LAYOUT, ++ ), ++ identifier=identifier, ++ ) ++ ++ ++def test_hidden_tensor_database_prepares_data_key_addrs_and_sizes(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ ++ key, addrs, sizes = HiddenTensorDatabase().prepare_value(pool_key, tensor) ++ ++ assert key == make_hidden_data_key(pool_key) ++ assert addrs == [tensor.data_ptr()] ++ assert sizes == [tensor.numel() * tensor.element_size()] ++ ++ ++def test_store_client_checks_single_tensor_object_exists(): ++ pool_key = make_pool_key() ++ store = FakeStore() ++ client = MooncakeHiddenStoreClient(store) ++ ++ assert not client.exists(pool_key) ++ ++ store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" ++ assert client.exists(pool_key) ++ ++ ++def test_worker_lookup_checks_existence_without_reading_tensor_metadata(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ worker.save_tensor(pool_key, tensor) ++ ++ assert worker.lookup(pool_key.identifier) ++ assert not worker.lookup("missing-image-hash") ++ assert store.range_gets == [] ++ ++ ++def test_worker_batch_lookup_checks_existence_in_one_store_call(): ++ pool_key_a = make_pool_key("image-a") ++ pool_key_b = make_pool_key("image-b") ++ store = FakeBufferStore() ++ store.objects[make_hidden_data_key(pool_key_a)] = b"tensor-object" ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key_a.key_metadata, ++ ) ++ ++ results = worker.lookup_batch(["image-a", "image-b"]) ++ ++ assert results == {"image-a": True, "image-b": False} ++ assert store.batch_is_exist_calls == [ ++ [make_hidden_data_key(pool_key_a), make_hidden_data_key(pool_key_b)] ++ ] ++ assert store.range_gets == [] ++ ++ ++def test_lookup_client_discard_removes_identifier_future_mapping(): ++ client = HiddenLookupClient.__new__(HiddenLookupClient) ++ future: Future[dict[str, bool]] = Future() ++ client.futures = { ++ "image-a": future, ++ "image-b": future, ++ } ++ ++ client.discard("image-a") ++ ++ assert "image-a" not in client.futures ++ assert client.futures == {"image-b": future} ++ assert not future.cancelled() ++ ++ client.discard("image-b") ++ ++ assert client.futures == {} ++ assert future.cancelled() ++ ++ ++def test_worker_lookup_records_minimal_operation_stats(): ++ pool_key = make_pool_key() ++ store = FakeBufferStore() ++ store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ ++ assert worker.lookup_batch(["image-hash", "missing-image-hash"]) == { ++ "image-hash": True, ++ "missing-image-hash": False, ++ } ++ ++ stats = worker.get_operation_stats() ++ records = stats.data["lookup_exists"] ++ assert len(records) == 1 ++ assert records[0]["num_keys"] == 2 ++ assert records[0]["num_bytes"] == 0 ++ assert records[0]["status"] == "miss" ++ assert records[0]["num_failed_keys"] == 1 ++ assert worker.get_operation_stats() is None ++ ++ ++def test_worker_save_stores_hidden_as_single_tensor_object(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient( ++ store, ++ replicate_config=FakeReplicateConfig(), ++ ), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) ++ assert store.pub_tensors[0][2] is not None ++ assert make_hidden_data_key(pool_key) in store.objects ++ ++ ++def test_worker_save_rejects_dtype_that_load_cannot_decode(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float64) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ try: ++ worker.save_tensor(pool_key, tensor) ++ except HiddenStoreSaveError as exc: ++ assert "unsupported hidden tensor dtype" in str(exc) ++ else: ++ raise AssertionError("unsupported hidden dtype should fail before store put") ++ ++ assert store.pub_tensors == [] ++ ++ ++def test_buffer_put_unregisters_payload_and_metadata_buffers(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) ++ ++ client.put_tensor(pool_key, tensor) ++ ++ payload_addr = tensor.data_ptr() ++ metadata_addr = next( ++ addr for addr, size in store.registered if size == TENSOR_METADATA_SIZE ++ ) ++ assert payload_addr in store.unregistered ++ assert metadata_addr in store.unregistered ++ assert store.unregistered[-2:] == [metadata_addr, payload_addr] ++ ++ ++def test_buffer_put_unregisters_payload_and_metadata_when_put_raises(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ store.raise_on_batch_put = True ++ client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) ++ ++ try: ++ client.put_tensor(pool_key, tensor) ++ except RuntimeError as exc: ++ assert "batch put failed" in str(exc) ++ else: ++ raise AssertionError("batch put exception should propagate") ++ ++ payload_addr = tensor.data_ptr() ++ metadata_addr = next( ++ addr for addr, size in store.registered if size == TENSOR_METADATA_SIZE ++ ) ++ assert payload_addr in store.unregistered ++ assert metadata_addr in store.unregistered ++ ++ ++def test_buffer_put_unregisters_payload_when_metadata_registration_fails(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ original_register = store.register_buffer ++ ++ def register_buffer(addr, size): ++ if size == TENSOR_METADATA_SIZE: ++ store.fail_register_addrs.add(addr) ++ return original_register(addr, size) ++ ++ store.register_buffer = register_buffer ++ client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) ++ ++ try: ++ client.put_tensor(pool_key, tensor) ++ except HiddenStoreError: ++ pass ++ else: ++ raise AssertionError("metadata registration failure should raise") ++ ++ assert tensor.data_ptr() in store.unregistered ++ ++ ++def test_worker_save_marks_hidden_state_data_type(monkeypatch): ++ fake_mooncake = types.ModuleType("mooncake") ++ fake_store = types.ModuleType("mooncake.store") ++ fake_store.ObjectDataType = FakeObjectDataTypeWithHidden ++ monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake) ++ monkeypatch.setitem(sys.modules, "mooncake.store", fake_store) ++ ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ replicate_config = FakeReplicateConfig() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient( ++ store, ++ replicate_config=replicate_config, ++ ), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ used_config = store.pub_tensors[0][2] ++ assert used_config is not replicate_config ++ assert int(used_config.data_type) == 10 ++ ++ ++def test_hidden_state_data_type_falls_back_to_tensor(monkeypatch): ++ fake_mooncake = types.ModuleType("mooncake") ++ fake_store = types.ModuleType("mooncake.store") ++ fake_store.ObjectDataType = FakeObjectDataTypeOnlyTensor ++ monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake) ++ monkeypatch.setitem(sys.modules, "mooncake.store", fake_store) ++ ++ assert _get_hidden_state_object_data_type() == FakeObjectDataTypeOnlyTensor.TENSOR ++ ++ ++def test_hidden_state_data_type_missing_type_returns_none(monkeypatch): ++ fake_mooncake = types.ModuleType("mooncake") ++ fake_store = types.ModuleType("mooncake.store") ++ fake_store.ObjectDataType = FakeObjectDataTypeNoTensor ++ monkeypatch.setitem(sys.modules, "mooncake", fake_mooncake) ++ monkeypatch.setitem(sys.modules, "mooncake.store", fake_store) ++ ++ assert _get_hidden_state_object_data_type() is None ++ ++ ++def test_worker_save_does_not_require_mooncake_object_group_support(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient( ++ store, ++ replicate_config=FakeReplicateConfigWithoutGroups(), ++ ), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) ++ ++ ++def test_worker_save_skips_existing_tensor_object(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ worker.save_tensor(pool_key, tensor) ++ ++ assert len(store.pub_tensors) == 1 ++ ++ ++def test_worker_save_records_exists_and_put_operation_stats(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ stats = worker.get_operation_stats() ++ assert stats.data["save_exists"][0]["status"] == "miss" ++ assert stats.data["save_exists"][0]["num_keys"] == 1 ++ assert stats.data["save_put"][0]["status"] == "ok" ++ assert stats.data["save_put"][0]["num_keys"] == 1 ++ assert stats.data["save_put"][0]["num_bytes"] == ( ++ tensor.numel() * tensor.element_size() ++ ) ++ ++ ++def test_worker_save_existing_records_only_save_exists(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ stats = worker.get_operation_stats() ++ assert stats.data["save_exists"][0]["status"] == "ok" ++ assert "save_put" not in stats.data ++ ++ ++def test_sending_thread_stores_hidden_tensor_asynchronously(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ sending_thread = HiddenStoreSendingThread(worker) ++ sending_thread.start() ++ ++ sending_thread.add_request( ++ HiddenSaveRequest(pool_key=pool_key, tensor=tensor) ++ ) ++ sending_thread.request_queue.join() ++ ++ assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) ++ assert sending_thread.get_and_clear_finished_identifiers() == {pool_key.identifier} ++ sending_thread.close() ++ ++ ++def test_sending_thread_records_failed_identifier_without_finishing(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ store.raise_on_batch_put = True ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ sending_thread = HiddenStoreSendingThread(worker) ++ sending_thread.start() ++ ++ sending_thread.add_request( ++ HiddenSaveRequest(pool_key=pool_key, tensor=tensor) ++ ) ++ sending_thread.request_queue.join() ++ ++ assert sending_thread.get_and_clear_finished_identifiers() == set() ++ assert sending_thread.get_and_clear_failed_identifiers() == {pool_key.identifier} ++ assert pool_key.identifier in sending_thread.failure_reasons ++ assert worker.get_operation_stats().data["save_put"][0]["status"] == "error" ++ sending_thread.close() ++ ++ ++def test_worker_load_gets_tensor_data_into_encoder_cache_before_returning(): ++ pool_key = make_pool_key() ++ stored = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ worker.save_tensor(pool_key, stored) ++ ++ encoder_cache = {} ++ worker.load( ++ [MMMeta(identifier=pool_key.identifier, load_spec=LoadSpec(can_load=True))], ++ encoder_cache, ++ device="cpu", ++ ) ++ ++ assert pool_key.identifier in encoder_cache ++ assert tuple(encoder_cache[pool_key.identifier].shape) == tuple(stored.shape) ++ assert str(encoder_cache[pool_key.identifier].dtype) == str(stored.dtype) ++ assert store.range_gets[0][1] == [[make_hidden_data_key(pool_key)]] ++ assert store.range_gets[0][3] == [[[0]]] ++ assert store.range_gets[0][4] == [[[TENSOR_METADATA_SIZE]]] ++ assert store.range_gets[-1][1] == [[make_hidden_data_key(pool_key)]] ++ assert store.range_gets[-1][3] == [[[TENSOR_METADATA_SIZE]]] ++ ++ stats = worker.get_operation_stats() ++ assert stats.data["load_get"][0]["status"] == "ok" ++ assert stats.data["load_get"][0]["num_keys"] == 1 ++ assert stats.data["load_get"][0]["num_bytes"] == ( ++ stored.numel() * stored.element_size() ++ ) ++ ++ ++def test_worker_load_records_error_without_writing_encoder_cache(): ++ pool_key = make_pool_key() ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ encoder_cache = {} ++ ++ try: ++ worker.load( ++ [MMMeta(identifier=pool_key.identifier, load_spec=LoadSpec(can_load=True))], ++ encoder_cache, ++ device="cpu", ++ ) ++ except HiddenStoreLoadError: ++ pass ++ else: ++ raise AssertionError("missing hidden tensor should fail fast") ++ ++ assert pool_key.identifier not in encoder_cache ++ stats = worker.get_operation_stats() ++ assert stats.data["load_get"][0]["status"] == "error" ++ assert stats.data["load_get"][0]["num_failed_keys"] == 1 ++ ++ ++def test_get_tensor_payload_unregisters_target_buffer_after_success(): ++ pool_key = make_pool_key() ++ stored = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ worker.save_tensor(pool_key, stored) ++ target = torch.empty_like(stored) ++ ++ worker.store_client.get_tensor_payload( ++ pool_key, ++ target.data_ptr(), ++ target.numel() * target.element_size(), ++ TENSOR_METADATA_SIZE, ++ ) ++ ++ assert target.data_ptr() in store.unregistered ++ ++ ++def test_get_tensor_payload_unregisters_target_buffer_after_load_error(): ++ pool_key = make_pool_key() ++ target = torch.empty((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ client = MooncakeHiddenStoreClient(store) ++ ++ try: ++ client.get_tensor_payload( ++ pool_key, ++ target.data_ptr(), ++ target.numel() * target.element_size(), ++ TENSOR_METADATA_SIZE, ++ ) ++ except HiddenStoreLoadError: ++ pass ++ else: ++ raise AssertionError("missing payload should raise") ++ ++ assert target.data_ptr() in store.unregistered ++ ++ ++def _serialize_tensor_object(tensor: torch.Tensor) -> bytes: ++ tensor = tensor.detach().cpu().contiguous() ++ nbytes = tensor.numel() * tensor.element_size() ++ header = struct.pack( ++ " bytes: ++ dims = list(shape) + [-1] * (8 - len(shape)) ++ return struct.pack("<8q", *dims) +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py +new file mode 100644 +index 000000000..77f80e012 +--- /dev/null ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/__init__.py +@@ -0,0 +1,9 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Hidden-state Mooncake Store EC connector support.""" ++ ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.connector import ( ++ MooncakeStoreECConnector, ++) ++ ++__all__ = ["MooncakeStoreECConnector"] +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py +new file mode 100644 +index 000000000..7213e1cbf +--- /dev/null ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py +@@ -0,0 +1,370 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""EC connector backed by Mooncake Store for hidden-state tensors.""" ++ ++from __future__ import annotations ++ ++from typing import TYPE_CHECKING ++ ++import torch ++ ++from vllm.distributed import ( ++ get_dcp_group, ++ get_pcp_group, ++ get_tensor_model_parallel_world_size, ++) ++from vllm.distributed.ec_transfer.ec_connector.base import ( ++ ECConnectorBase, ++ ECConnectorMetadata, ++ ECConnectorRole, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HIDDEN_OBJECT_KIND, ++ HIDDEN_STORAGE_LAYOUT, ++ HIDDEN_TENSOR_LAYOUT, ++ HiddenKeyMetadata, ++ HiddenSaveRequest, ++ LoadSpec, ++ MMMeta, ++ MooncakeStoreConnectorMetadata, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( ++ MooncakeHiddenStoreClient, ++ create_mooncake_hidden_store_client, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.worker import ( ++ HiddenLookupClient, ++ HiddenLookupServer, ++ HiddenStoreWorker, ++) ++from vllm.logger import init_logger ++from vllm.multimodal.utils import get_mm_features_in_window ++from vllm.v1.core.sched.output import SchedulerOutput ++ ++if TYPE_CHECKING: ++ from vllm.config import VllmConfig ++ from vllm.v1.request import Request ++ ++logger = init_logger(__name__) ++ ++ ++class MooncakeStoreECConnector(ECConnectorBase): ++ """Hidden-state EC connector that stores tensors in Mooncake Store.""" ++ ++ def __init__( ++ self, ++ vllm_config: VllmConfig, ++ role: ECConnectorRole, ++ store_client: MooncakeHiddenStoreClient | None = None, ++ ): ++ super().__init__(vllm_config=vllm_config, role=role) ++ self.lookup_client: HiddenLookupClient | None = None ++ self.lookup_server: HiddenLookupServer | None = None ++ self.store_client: MooncakeHiddenStoreClient | None = None ++ self.worker: HiddenStoreWorker | None = None ++ assert vllm_config.ec_transfer_config is not None ++ extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config ++ self.soft_pin_video_hidden = bool( ++ extra_config.get("soft_pin_video_hidden", False) ++ ) ++ self.lookup_async = bool(extra_config.get("lookup_async", True)) ++ ++ if role == ECConnectorRole.SCHEDULER: ++ if self.is_consumer: ++ self.lookup_client = HiddenLookupClient(vllm_config) ++ else: ++ if not (self.is_producer or self.is_consumer): ++ return ++ hidden_key_metadata = build_hidden_key_metadata(vllm_config) ++ self.store_client = store_client or create_mooncake_hidden_store_client() ++ self.worker = HiddenStoreWorker( ++ store_client=self.store_client, ++ key_metadata=hidden_key_metadata, ++ ) ++ if self.is_producer: ++ self.worker.start_sending_thread() ++ if self.is_consumer and vllm_config.parallel_config.rank == 0: ++ self.lookup_server = HiddenLookupServer(self.worker, vllm_config) ++ ++ self.load_specs: dict[str, LoadSpec] = {} ++ self.lookup_result_cache: dict[str, bool] = {} ++ self.identifier_waiters: dict[str, set[str]] = {} ++ self._candidate_consumes: dict[str, set[str]] = {} ++ self._candidate_loads: dict[str, set[str]] = {} ++ self._candidate_saves: dict[str, set[str]] = {} ++ self._load_modalities: dict[str, str | None] = {} ++ self._save_modalities: dict[str, str | None] = {} ++ ++ def shutdown(self) -> None: ++ if self.lookup_client is not None: ++ self.lookup_client.close() ++ if self.lookup_server is not None: ++ self.lookup_server.close() ++ if self.worker is not None: ++ self.worker.shutdown() ++ ++ def has_cache_item(self, identifier: str) -> bool: ++ if not self.is_consumer: ++ return False ++ ++ if not self.lookup_result_cache.get(identifier, False): ++ self.load_specs.pop(identifier, None) ++ logger.info( ++ "hidden_store_scheduler_miss identifier=%s " ++ "reason=local_lookup_result_miss", ++ identifier, ++ ) ++ return False ++ ++ self.load_specs.setdefault(identifier, LoadSpec(can_load=False)) ++ logger.info( ++ "hidden_store_scheduler_hit identifier=%s", ++ identifier, ++ ) ++ return True ++ ++ def ensure_cache_available( ++ self, ++ request: Request, ++ num_computed_tokens: int, ++ ) -> bool: ++ if not self.is_consumer: ++ return True ++ if not request.mm_features: ++ return True ++ assert self.lookup_client is not None ++ ++ start = num_computed_tokens ++ end = request.num_tokens ++ lo, hi = get_mm_features_in_window(request.mm_features, start, end) ++ identifiers = list( ++ dict.fromkeys( ++ feature.identifier for feature in request.mm_features[lo:hi] ++ ) ++ ) ++ if not identifiers: ++ return True ++ ++ request_id = request.request_id ++ for identifier in identifiers: ++ self.identifier_waiters.setdefault(identifier, set()).add(request_id) ++ ++ unknown_identifiers = [ ++ identifier ++ for identifier in identifiers ++ if identifier not in self.lookup_result_cache ++ ] ++ if not unknown_identifiers: ++ return True ++ ++ lookup_results = self.lookup_client.lookup_batch( ++ unknown_identifiers, ++ non_block=self.lookup_async, ++ ) ++ if lookup_results is None: ++ return False ++ ++ for identifier in unknown_identifiers: ++ self.lookup_result_cache[identifier] = lookup_results.get( ++ identifier, ++ False, ++ ) ++ return True ++ ++ def update_state_after_alloc(self, request: Request, index: int) -> None: ++ mm_feature = request.mm_features[index] ++ identifier = mm_feature.identifier ++ modality = mm_feature.modality ++ request_id = request.request_id ++ ++ self._candidate_consumes.setdefault(request_id, set()).add(identifier) ++ ++ if self.is_consumer and identifier in self.load_specs: ++ self._candidate_loads.setdefault(request_id, set()).add(identifier) ++ self._load_modalities[identifier] = modality ++ ++ if self.is_producer: ++ self._save_modalities[identifier] = modality ++ self._candidate_saves.setdefault(request_id, set()).add(identifier) ++ ++ def build_connector_meta( ++ self, ++ scheduler_output: SchedulerOutput, ++ ) -> ECConnectorMetadata: ++ items_by_identifier: dict[str, MMMeta] = {} ++ preempted_ids = getattr(scheduler_output, "preempted_req_ids", None) or set() ++ ++ for request_id, identifiers in self._candidate_consumes.items(): ++ if request_id in preempted_ids: ++ continue ++ for identifier in identifiers: ++ waiters = self.identifier_waiters.get(identifier) ++ if waiters is not None: ++ waiters.discard(request_id) ++ ++ for request_id, identifiers in self._candidate_loads.items(): ++ if request_id in preempted_ids: ++ continue ++ for identifier in identifiers: ++ load_spec = self.load_specs.pop(identifier, None) ++ if load_spec is None: ++ continue ++ load_spec.can_load = True ++ items_by_identifier[identifier] = MMMeta( ++ identifier=identifier, ++ modality=self._load_modalities.get(identifier), ++ load_spec=load_spec, ++ ) ++ ++ for request_id, identifiers in self._candidate_saves.items(): ++ if request_id in preempted_ids: ++ continue ++ for identifier in identifiers: ++ item = items_by_identifier.get(identifier) ++ if item is None: ++ item = MMMeta( ++ identifier=identifier, ++ modality=self._save_modalities.get(identifier), ++ ) ++ items_by_identifier[identifier] = item ++ item.can_save = True ++ if item.modality is None: ++ item.modality = self._save_modalities.get(identifier) ++ ++ finished_req_ids = getattr(scheduler_output, "finished_req_ids", set()) ++ for finished_req_id in finished_req_ids: ++ for waiters in self.identifier_waiters.values(): ++ waiters.discard(finished_req_id) ++ ++ self._cleanup_lookup_results_without_waiters() ++ ++ metadata = MooncakeStoreConnectorMetadata( ++ items=list(items_by_identifier.values()), ++ ) ++ ++ self._candidate_consumes.clear() ++ self._candidate_loads.clear() ++ self._candidate_saves.clear() ++ self._load_modalities.clear() ++ self._save_modalities.clear() ++ return metadata ++ ++ def _cleanup_lookup_results_without_waiters(self) -> None: ++ for identifier, waiters in list(self.identifier_waiters.items()): ++ if waiters: ++ continue ++ del self.identifier_waiters[identifier] ++ self.lookup_result_cache.pop(identifier, None) ++ self.load_specs.pop(identifier, None) ++ if self.lookup_client is not None: ++ self.lookup_client.discard(identifier) ++ ++ def start_load_caches( ++ self, ++ encoder_cache: dict[str, torch.Tensor], ++ **kwargs, ++ ) -> None: ++ metadata = self._get_connector_metadata() ++ assert isinstance(metadata, MooncakeStoreConnectorMetadata) ++ assert self.worker is not None ++ self.worker.load( ++ metadata.items, ++ encoder_cache, ++ device=kwargs.get("device"), ++ ) ++ ++ def save_caches( ++ self, ++ encoder_cache: dict[str, torch.Tensor], ++ mm_hash: str, ++ **kwargs, ++ ) -> None: ++ if not self.is_producer: ++ return ++ assert self.worker is not None ++ identifier = mm_hash ++ if identifier not in encoder_cache: ++ logger.warning( ++ "Skip hidden store save; identifier %s is missing", ++ identifier, ++ ) ++ return ++ item = self._find_metadata_item(identifier) ++ if item is None or not item.can_save: ++ logger.debug( ++ "Skip hidden store save; identifier %s has no save plan", ++ identifier, ++ ) ++ return ++ pool_key = self.worker.make_pool_key(identifier) ++ self.worker.enqueue_save( ++ HiddenSaveRequest( ++ pool_key=pool_key, ++ tensor=encoder_cache[identifier], ++ with_soft_pin=self._should_soft_pin(item), ++ ) ++ ) ++ ++ def get_finished( ++ self, finished_req_ids: set[str] ++ ) -> tuple[set[str] | None, set[str] | None]: ++ if self.worker is None or not self.is_producer: ++ return None, None ++ finished_sending = self.worker.get_finished_sending() ++ return finished_sending or None, None ++ ++ def _find_metadata_item(self, identifier: str) -> MMMeta | None: ++ metadata = self._get_connector_metadata() ++ assert isinstance(metadata, MooncakeStoreConnectorMetadata) ++ for item in metadata.items: ++ if item.identifier == identifier: ++ return item ++ return None ++ ++ def _should_soft_pin(self, item: MMMeta) -> bool: ++ return self.soft_pin_video_hidden and item.modality == "video" ++ ++ ++def build_hidden_key_metadata(vllm_config: VllmConfig) -> HiddenKeyMetadata: ++ model_config = vllm_config.model_config ++ parallel_config = vllm_config.parallel_config ++ assert vllm_config.ec_transfer_config is not None ++ extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config ++ ++ multimodal_config = getattr(model_config, "multimodal_config", None) ++ compute_hash = getattr(multimodal_config, "compute_hash", None) ++ mm_encoder_config_hash = ( ++ compute_hash() if callable(compute_hash) else "encoder:default" ++ ) ++ ++ tp_size = get_tensor_model_parallel_world_size() ++ pp_size = parallel_config.pipeline_parallel_size ++ pcp_size = get_pcp_group().world_size ++ dcp_size = get_dcp_group().world_size ++ mm_encoder_tp_mode = getattr( ++ multimodal_config, ++ "mm_encoder_tp_mode", ++ "unknown", ++ ) ++ parallel = ( ++ f"tp:{tp_size}" ++ f"@pp:{pp_size}" ++ f"@pcp:{pcp_size}" ++ f"@dcp:{dcp_size}" ++ f"@mm_tp:{mm_encoder_tp_mode}" ++ ) ++ ++ return HiddenKeyMetadata( ++ cache_prefix=str( ++ extra_config.get( ++ "hidden_cache_prefix", ++ extra_config.get("cache_prefix", ""), ++ ) ++ ), ++ kind=HIDDEN_OBJECT_KIND, ++ model_name=model_config.model.rstrip("/").split("/")[-1], ++ encoder=str(mm_encoder_config_hash), ++ storage=HIDDEN_STORAGE_LAYOUT, ++ parallel=parallel, ++ tensor_layout=HIDDEN_TENSOR_LAYOUT, ++ ) +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py +new file mode 100644 +index 000000000..f3fa09ef1 +--- /dev/null ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py +@@ -0,0 +1,202 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Data classes for the hidden-state Mooncake Store EC connector.""" ++ ++from __future__ import annotations ++ ++from dataclasses import dataclass, field ++ ++import torch ++ ++from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( ++ escape_key_part, ++ make_hidden_data_key, ++) ++ ++HIDDEN_OBJECT_KIND = "encoder_output" ++HIDDEN_STORAGE_LAYOUT = "replicated_object" ++HIDDEN_TENSOR_LAYOUT = "tensor" ++HIDDEN_PROTOCOL_VERSION = "v1" ++MOONCAKE_TENSOR_METADATA_NBYTES = 304 ++ ++ ++@dataclass(frozen=True) ++class HiddenKeyMetadata: ++ """Metadata that defines the semantic namespace for hidden reuse.""" ++ ++ cache_prefix: str ++ kind: str ++ model_name: str ++ encoder: str ++ storage: str ++ parallel: str ++ tensor_layout: str ++ ++ ++@dataclass(frozen=True, order=True) ++class HiddenPoolKey: ++ """Key for addressing one hidden tensor in the distributed store.""" ++ ++ key_metadata: HiddenKeyMetadata ++ identifier: str ++ ++ def to_string(self) -> str: ++ meta = self.key_metadata ++ prefix = ( ++ f"{escape_key_part(meta.cache_prefix)}@" if meta.cache_prefix else "" ++ ) ++ return ( ++ f"{prefix}hidden" ++ f"@kind:{escape_key_part(meta.kind)}" ++ f"@model:{escape_key_part(meta.model_name)}" ++ f"@encoder:{escape_key_part(meta.encoder)}" ++ f"@storage:{escape_key_part(meta.storage)}" ++ f"@parallel:{escape_key_part(meta.parallel)}" ++ f"@tensor_layout:{escape_key_part(meta.tensor_layout)}" ++ f"@id:{escape_key_part(self.identifier)}" ++ ) ++ ++ ++@dataclass ++class MMMeta: ++ """Per hidden object metadata passed from scheduler to worker.""" ++ ++ identifier: str ++ modality: str | None = None ++ can_save: bool = False ++ load_spec: LoadSpec | None = None ++ ++ ++@dataclass(frozen=True) ++class TensorMeta: ++ """Canonical contiguous tensor descriptor for one hidden store object.""" ++ ++ pool_key: HiddenPoolKey ++ protocol_version: str ++ layout: str ++ shape: tuple[int, ...] ++ dtype: str ++ nbytes: int ++ device_type: str ++ data_offset: int = MOONCAKE_TENSOR_METADATA_NBYTES ++ producer_stage: str = "encoder" ++ ++ ++@dataclass ++class LoadSpec: ++ """Specification for loading a hidden tensor from external store.""" ++ ++ can_load: bool = False ++ ++ ++@dataclass ++class HiddenSaveRequest: ++ """Specification for asynchronously storing one hidden tensor.""" ++ ++ pool_key: HiddenPoolKey ++ tensor: torch.Tensor ++ with_soft_pin: bool = False ++ ++ @property ++ def identifier(self) -> str: ++ return self.pool_key.identifier ++ ++ ++@dataclass ++class MooncakeStoreConnectorMetadata(ECConnectorMetadata): ++ """Metadata passed from scheduler to worker for hidden store operations.""" ++ ++ items: list[MMMeta] = field(default_factory=list) ++ ++ def add_item(self, item: MMMeta) -> None: ++ self.items.append(item) ++ ++ ++@dataclass ++class HiddenStoreOperationStats: ++ """Minimal per-operation telemetry aligned with Mooncake KV store stats.""" ++ ++ data: dict[str, list[dict[str, int | float | str]]] = field(default_factory=dict) ++ ++ def is_empty(self) -> bool: ++ return not self.data ++ ++ def record_operation( ++ self, ++ operation: str, ++ duration_seconds: float, ++ num_keys: int, ++ *, ++ num_bytes: int = 0, ++ status: str = "ok", ++ num_failed_keys: int = 0, ++ ) -> None: ++ self.data.setdefault(operation, []).append( ++ { ++ "duration_seconds": duration_seconds, ++ "num_keys": num_keys, ++ "num_bytes": num_bytes, ++ "status": status, ++ "num_failed_keys": num_failed_keys, ++ } ++ ) ++ ++ ++class HiddenTensorDatabase: ++ """Maps hidden tensors to store keys and GPU memory descriptors.""" ++ ++ def prepare_value( ++ self, ++ pool_key: HiddenPoolKey, ++ tensor: torch.Tensor, ++ ) -> tuple[str, list[int], list[int]]: ++ return ( ++ make_hidden_data_key(pool_key), ++ [tensor.data_ptr()], ++ [tensor.numel() * tensor.element_size()], ++ ) ++ ++ ++def build_tensor_meta( ++ pool_key: HiddenPoolKey, ++ tensor: torch.Tensor, ++) -> TensorMeta: ++ """Build metadata for the canonical stored hidden tensor layout.""" ++ if not tensor.is_contiguous(): ++ raise ValueError("Hidden tensor descriptor requires a contiguous tensor") ++ ++ return TensorMeta( ++ pool_key=pool_key, ++ protocol_version=HIDDEN_PROTOCOL_VERSION, ++ layout=HIDDEN_TENSOR_LAYOUT, ++ shape=tuple(tensor.shape), ++ dtype=str(tensor.dtype), ++ nbytes=tensor.numel() * tensor.element_size(), ++ device_type=tensor.device.type, ++ data_offset=MOONCAKE_TENSOR_METADATA_NBYTES, ++ ) ++ ++ ++def validate_loaded_tensor(tensor: torch.Tensor, meta: TensorMeta) -> None: ++ if tuple(tensor.shape) != tuple(meta.shape): ++ raise ValueError( ++ "Hidden tensor shape mismatch: " ++ f"actual={tuple(tensor.shape)} expected={meta.shape}" ++ ) ++ ++ if str(tensor.dtype) != meta.dtype: ++ raise ValueError( ++ "Hidden tensor dtype mismatch: " ++ f"actual={tensor.dtype} expected={meta.dtype}" ++ ) ++ ++ actual_nbytes = tensor.numel() * tensor.element_size() ++ if actual_nbytes != meta.nbytes: ++ raise ValueError( ++ "Hidden tensor nbytes mismatch: " ++ f"actual={actual_nbytes} expected={meta.nbytes}" ++ ) ++ ++ if meta.layout != HIDDEN_TENSOR_LAYOUT: ++ raise ValueError(f"Unsupported hidden tensor layout: {meta.layout}") +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py +new file mode 100644 +index 000000000..b05c99b20 +--- /dev/null ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/keys.py +@@ -0,0 +1,22 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Store key helpers for the hidden-state Mooncake connector.""" ++ ++from __future__ import annotations ++ ++from typing import TYPE_CHECKING ++from urllib.parse import quote ++ ++if TYPE_CHECKING: ++ from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HiddenPoolKey, ++ ) ++ ++ ++def escape_key_part(value: str) -> str: ++ """Escape one key component while keeping simple values readable.""" ++ return quote(str(value), safe="-_.~") ++ ++ ++def make_hidden_data_key(pool_key: "HiddenPoolKey") -> str: ++ return pool_key.to_string() +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py +new file mode 100644 +index 000000000..0852ee1db +--- /dev/null ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py +@@ -0,0 +1,551 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Thin Mooncake Store client for hidden-state objects.""" ++ ++from __future__ import annotations ++ ++import copy ++import ctypes ++import json ++import os ++import re ++import struct ++from dataclasses import dataclass ++from typing import Any ++ ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HIDDEN_PROTOCOL_VERSION, ++ HIDDEN_TENSOR_LAYOUT, ++ MOONCAKE_TENSOR_METADATA_NBYTES, ++ HiddenPoolKey, ++ TensorMeta, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import ( ++ make_hidden_data_key, ++) ++from vllm.logger import init_logger ++from vllm.utils.network_utils import get_ip ++ ++logger = init_logger(__name__) ++ ++DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 ++DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024 ++_MOONCAKE_TENSOR_OBJECT_MAGIC = 0x4D4F4F4E ++_MOONCAKE_TENSOR_OBJECT_VERSION = 1 ++_MOONCAKE_TENSOR_HEADER_FORMAT = " MooncakeHiddenStoreConfig: ++ with open(file_path, encoding="utf-8") as file: ++ config = json.load(file) ++ mode = config.get("mode", "embedded") ++ return MooncakeHiddenStoreConfig( ++ metadata_server=config.get("metadata_server", ""), ++ master_server_address=config.get("master_server_address", ""), ++ protocol=config.get("protocol", "rdma"), ++ device_name=config.get("device_name", ""), ++ mode=mode, ++ global_segment_size=_parse_size( ++ config.get( ++ "global_segment_size", ++ 0 if mode == "standalone-store" else DEFAULT_GLOBAL_SEGMENT_SIZE, ++ ) ++ ), ++ local_buffer_size=_parse_size( ++ config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE) ++ ), ++ ) ++ ++ @staticmethod ++ def load_from_env() -> MooncakeHiddenStoreConfig: ++ config_path = os.getenv("MOONCAKE_CONFIG_PATH") ++ if not config_path: ++ raise ValueError( ++ "The environment variable 'MOONCAKE_CONFIG_PATH' is not set." ++ ) ++ return MooncakeHiddenStoreConfig.from_file(config_path) ++ ++ ++def _parse_size(value: Any) -> int: ++ if isinstance(value, int): ++ return value ++ if not isinstance(value, str): ++ return int(value) ++ ++ cleaned = value.strip().lower() ++ match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned) ++ if not match: ++ raise ValueError(f"Invalid size format: {value!r}") ++ ++ multipliers = { ++ "gb": 1024**3, ++ "mb": 1024**2, ++ "kb": 1024, ++ "b": 1, ++ None: 1, ++ } ++ return int(float(match.group(1)) * multipliers[match.group(2)]) ++ ++ ++def create_mooncake_hidden_store_client() -> MooncakeHiddenStoreClient: ++ try: ++ from mooncake.store import ( # type: ignore ++ MooncakeDistributedStore, ++ ReplicateConfig, ++ ) ++ except ImportError as e: ++ raise ImportError( ++ "Please install mooncake to run vLLM with " "MooncakeStoreECConnector." ++ ) from e ++ ++ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import rdma_utils ++ ++ config = MooncakeHiddenStoreConfig.load_from_env() ++ config.device_name = rdma_utils.get_configured_worker_rnic( ++ protocol=config.protocol, ++ configured_device=config.device_name, ++ ) ++ ++ store = MooncakeDistributedStore() ++ local_ip = get_ip() ++ local_hostname = rdma_utils.get_requester_local_hostname(local_ip) ++ ret = store.setup( ++ local_hostname, ++ config.metadata_server, ++ config.global_segment_size, ++ config.local_buffer_size, ++ config.protocol, ++ config.device_name, ++ config.master_server_address, ++ ) ++ if ret != 0: ++ raise RuntimeError("Initialize MooncakeDistributedStore failed.") ++ ++ logger.info( ++ "Initialized hidden Mooncake store mode=%s global_segment_size=%d " ++ "local_buffer_size=%d", ++ config.mode, ++ config.global_segment_size, ++ config.local_buffer_size, ++ ) ++ return MooncakeHiddenStoreClient(store, replicate_config=ReplicateConfig()) ++ ++ ++class HiddenStoreError(RuntimeError): ++ pass ++ ++ ++class HiddenStoreLoadError(HiddenStoreError): ++ pass ++ ++ ++class HiddenStoreSaveError(HiddenStoreError): ++ pass ++ ++ ++class MooncakeHiddenStoreClient: ++ """Wraps Mooncake object and buffer APIs used by hidden transfer.""" ++ ++ def __init__(self, store: Any, replicate_config: Any | None = None): ++ self.store = store ++ self.replicate_config = replicate_config ++ ++ def exists(self, pool_key: HiddenPoolKey) -> bool: ++ data_key = make_hidden_data_key(pool_key) ++ states = self.store.batch_is_exist([data_key]) ++ return len(states) == 1 and states[0] == 1 ++ ++ def batch_exists(self, pool_keys: list[HiddenPoolKey]) -> list[bool]: ++ if not pool_keys: ++ return [] ++ ++ keys = [make_hidden_data_key(pool_key) for pool_key in pool_keys] ++ states = self.store.batch_is_exist(keys) ++ return [state == 1 for state in states] ++ ++ def get_tensor_meta(self, pool_key: HiddenPoolKey) -> TensorMeta | None: ++ metadata = self._read_range( ++ pool_key, ++ src_offset=0, ++ size=MOONCAKE_TENSOR_METADATA_NBYTES, ++ ) ++ if metadata is None: ++ return None ++ try: ++ return _decode_mooncake_tensor_metadata(pool_key, metadata) ++ except HiddenStoreLoadError: ++ logger.exception( ++ "failed to decode hidden Mooncake tensor metadata for %s", ++ pool_key.to_string(), ++ ) ++ return None ++ ++ def put_tensor( ++ self, ++ pool_key: HiddenPoolKey, ++ tensor: Any, ++ *, ++ with_soft_pin: bool = False, ++ ) -> None: ++ _validate_supported_hidden_tensor_dtype(tensor) ++ key = make_hidden_data_key(pool_key) ++ replicate_config = _make_hidden_replicate_config( ++ self.replicate_config, ++ with_soft_pin=with_soft_pin, ++ ) ++ batch_put_from_multi_buffers = getattr( ++ self.store, ++ "batch_put_from_multi_buffers", ++ None, ++ ) ++ if batch_put_from_multi_buffers is not None: ++ self._put_tensor_from_buffers( ++ pool_key, ++ tensor, ++ replicate_config=replicate_config, ++ ) ++ return ++ ++ if replicate_config is None: ++ put_fn = getattr(self.store, "put_tensor", None) ++ if put_fn is None: ++ raise HiddenStoreSaveError( ++ "Mooncake Hidden Store requires put_tensor or pub_tensor " ++ "support for single-object hidden tensors." ++ ) ++ ret = put_fn(key, tensor) ++ else: ++ put_fn = getattr(self.store, "pub_tensor", None) ++ if put_fn is None: ++ raise HiddenStoreSaveError( ++ "Mooncake Hidden Store requires pub_tensor support when " ++ "a ReplicateConfig is configured." ++ ) ++ ret = put_fn(key, tensor, replicate_config) ++ if ret != 0: ++ raise HiddenStoreSaveError( ++ f"failed to put hidden tensor for {pool_key.to_string()}: {ret}" ++ ) ++ ++ def _put_tensor_from_buffers( ++ self, ++ pool_key: HiddenPoolKey, ++ tensor: Any, ++ *, ++ replicate_config: Any | None, ++ ) -> None: ++ if not tensor.is_contiguous(): ++ raise HiddenStoreSaveError( ++ "hidden tensor must be contiguous before batch buffer put" ++ ) ++ data_size = tensor.numel() * tensor.element_size() ++ metadata = _encode_mooncake_tensor_metadata(tensor) ++ metadata_buffer = (ctypes.c_ubyte * len(metadata)).from_buffer_copy(metadata) ++ metadata_ptr = ctypes.addressof(metadata_buffer) ++ payload_ptr = tensor.data_ptr() ++ registered_addrs: list[int] = [] ++ try: ++ self.register_tensor(payload_ptr, data_size) ++ registered_addrs.append(payload_ptr) ++ self.register_tensor(metadata_ptr, len(metadata)) ++ registered_addrs.append(metadata_ptr) ++ ++ key = make_hidden_data_key(pool_key) ++ results = self.store.batch_put_from_multi_buffers( ++ [key], ++ [[metadata_ptr, payload_ptr]], ++ [[len(metadata), data_size]], ++ replicate_config, ++ ) ++ failed = [result for result in results if result < 0] ++ if failed: ++ raise HiddenStoreSaveError( ++ "failed to put hidden tensor for " ++ f"{pool_key.to_string()}: {failed}" ++ ) ++ finally: ++ for addr in reversed(registered_addrs): ++ self.unregister_tensor(addr) ++ ++ def register_tensor(self, addr: int, size: int) -> None: ++ ret = self.store.register_buffer(addr, size) ++ if ret != 0: ++ raise HiddenStoreError( ++ f"failed to register hidden buffer addr={addr:#x} size={size}: {ret}" ++ ) ++ ++ def unregister_tensor(self, addr: int) -> None: ++ unregister_fn = getattr(self.store, "unregister_buffer", None) ++ if unregister_fn is None: ++ return ++ try: ++ ret = unregister_fn(addr) ++ except Exception: ++ logger.warning( ++ "failed to unregister hidden buffer addr=%#x", ++ addr, ++ exc_info=True, ++ ) ++ return ++ if ret != 0: ++ logger.warning( ++ "unregister hidden buffer failed addr=%#x ret=%s", ++ addr, ++ ret, ++ ) ++ ++ def get_tensor_payload( ++ self, ++ pool_key: HiddenPoolKey, ++ addr: int, ++ size: int, ++ src_offset: int, ++ ) -> int: ++ self.register_tensor(addr, size) ++ try: ++ key = make_hidden_data_key(pool_key) ++ results = self.store.get_into_ranges( ++ [addr], ++ [[key]], ++ [[[0]]], ++ [[[src_offset]]], ++ [[[size]]], ++ ) ++ result = _single_range_result(results) ++ if result != size: ++ raise HiddenStoreLoadError( ++ "failed to get hidden tensor payload for " ++ f"{pool_key.to_string()}: {result}" ++ ) ++ return result ++ finally: ++ self.unregister_tensor(addr) ++ ++ def _read_range( ++ self, ++ pool_key: HiddenPoolKey, ++ *, ++ src_offset: int, ++ size: int, ++ ) -> bytes | None: ++ buffer = (ctypes.c_ubyte * size)() ++ buffer_ptr = ctypes.addressof(buffer) ++ self.register_tensor(buffer_ptr, size) ++ key = make_hidden_data_key(pool_key) ++ try: ++ results = self.store.get_into_ranges( ++ [buffer_ptr], ++ [[key]], ++ [[[0]]], ++ [[[src_offset]]], ++ [[[size]]], ++ ) ++ finally: ++ self.unregister_tensor(buffer_ptr) ++ if _single_range_result(results) != size: ++ return None ++ return bytes(buffer) ++ ++ ++def _single_range_result(results: Any) -> int: ++ try: ++ return int(results[0][0][0]) ++ except Exception: ++ return -1 ++ ++ ++def _decode_mooncake_tensor_metadata( ++ pool_key: HiddenPoolKey, ++ metadata: bytes, ++) -> TensorMeta: ++ if len(metadata) < MOONCAKE_TENSOR_METADATA_NBYTES: ++ raise HiddenStoreLoadError( ++ f"hidden tensor metadata is too small: {len(metadata)}" ++ ) ++ ( ++ magic, ++ version, ++ header_size, ++ dtype, ++ ndim, ++ _layout_kind, ++ _reserved_flags, ++ data_offset, ++ data_bytes, ++ ) = struct.unpack_from(_MOONCAKE_TENSOR_HEADER_FORMAT, metadata, 0) ++ if ( ++ magic != _MOONCAKE_TENSOR_OBJECT_MAGIC ++ or version != _MOONCAKE_TENSOR_OBJECT_VERSION ++ or header_size != MOONCAKE_TENSOR_METADATA_NBYTES ++ ): ++ raise HiddenStoreLoadError( ++ "invalid Mooncake tensor metadata header for " f"{pool_key.to_string()}" ++ ) ++ if ndim < 0 or ndim > 8: ++ raise HiddenStoreLoadError( ++ f"invalid hidden tensor ndim for {pool_key.to_string()}: {ndim}" ++ ) ++ if dtype not in _MOONCAKE_DTYPE_TO_TORCH_DTYPE: ++ raise HiddenStoreLoadError( ++ f"unsupported Mooncake tensor dtype for {pool_key.to_string()}: {dtype}" ++ ) ++ local_shape = struct.unpack_from( ++ "<8q", ++ metadata, ++ _MOONCAKE_TENSOR_LOCAL_SHAPE_OFFSET, ++ ) ++ shape = tuple(int(dim) for dim in local_shape[:ndim]) ++ if any(dim < 0 for dim in shape): ++ raise HiddenStoreLoadError( ++ f"invalid hidden tensor shape for {pool_key.to_string()}: {shape}" ++ ) ++ return TensorMeta( ++ pool_key=pool_key, ++ protocol_version=HIDDEN_PROTOCOL_VERSION, ++ layout=HIDDEN_TENSOR_LAYOUT, ++ shape=shape, ++ dtype=_MOONCAKE_DTYPE_TO_TORCH_DTYPE[dtype], ++ nbytes=int(data_bytes), ++ device_type="cpu", ++ data_offset=int(data_offset), ++ ) ++ ++ ++def _encode_mooncake_tensor_metadata(tensor: Any) -> bytes: ++ dtype = str(tensor.dtype) ++ _validate_supported_hidden_tensor_dtype(tensor) ++ shape = tuple(int(dim) for dim in tensor.shape) ++ if len(shape) > 8: ++ raise HiddenStoreSaveError( ++ f"hidden tensor has too many dimensions: {len(shape)}" ++ ) ++ nbytes = tensor.numel() * tensor.element_size() ++ header = struct.pack( ++ _MOONCAKE_TENSOR_HEADER_FORMAT, ++ _MOONCAKE_TENSOR_OBJECT_MAGIC, ++ _MOONCAKE_TENSOR_OBJECT_VERSION, ++ MOONCAKE_TENSOR_METADATA_NBYTES, ++ _TORCH_DTYPE_TO_MOONCAKE_DTYPE[dtype], ++ len(shape), ++ 0, ++ 0, ++ MOONCAKE_TENSOR_METADATA_NBYTES, ++ nbytes, ++ ) ++ dims = shape + (-1,) * (8 - len(shape)) ++ tensor_shape = struct.pack("<8q", *dims) ++ axes = b"\0" * (32 * 4) ++ metadata = header + tensor_shape + tensor_shape + struct.pack(" None: ++ dtype = str(tensor.dtype) ++ if dtype not in _SUPPORTED_HIDDEN_TORCH_DTYPES: ++ raise HiddenStoreSaveError(f"unsupported hidden tensor dtype: {dtype}") ++ ++ ++def _make_hidden_replicate_config( ++ replicate_config: Any | None, ++ *, ++ with_soft_pin: bool, ++) -> Any | None: ++ if replicate_config is None: ++ return None ++ ++ config = _clone_replicate_config(replicate_config) ++ hidden_state_data_type = _get_hidden_state_object_data_type() ++ if hidden_state_data_type is not None and hasattr(config, "data_type"): ++ config.data_type = hidden_state_data_type ++ if hasattr(config, "with_soft_pin"): ++ config.with_soft_pin = bool(config.with_soft_pin) or with_soft_pin ++ return config ++ ++ ++def _clone_replicate_config(replicate_config: Any) -> Any: ++ try: ++ return copy.copy(replicate_config) ++ except Exception: ++ config = type(replicate_config)() ++ for attr in ( ++ "replica_num", ++ "nof_replica_num", ++ "with_soft_pin", ++ "with_hard_pin", ++ "preferred_segments", ++ "preferred_nof_segments", ++ "preferred_segment", ++ "prefer_alloc_in_same_node", ++ "data_type", ++ "group_ids", ++ ): ++ if hasattr(replicate_config, attr) and hasattr(config, attr): ++ setattr(config, attr, getattr(replicate_config, attr)) ++ return config ++ ++ ++def _get_hidden_state_object_data_type() -> Any | None: ++ try: ++ from mooncake.store import ObjectDataType # type: ignore ++ except Exception: ++ return None ++ hidden_state_type = getattr(ObjectDataType, "HIDDEN_STATE", None) ++ if hidden_state_type is not None: ++ return hidden_state_type ++ return getattr(ObjectDataType, "TENSOR", None) +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py +new file mode 100644 +index 000000000..16829f34c +--- /dev/null ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py +@@ -0,0 +1,573 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Worker-side hidden-state load/save logic for Mooncake Store.""" ++ ++from __future__ import annotations ++ ++import os ++import queue ++import socket ++import threading ++import time ++from concurrent.futures import Future, ThreadPoolExecutor ++ ++import torch ++import zmq ++ ++import vllm.envs as envs ++from vllm.config import VllmConfig ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( ++ HiddenKeyMetadata, ++ HiddenPoolKey, ++ HiddenSaveRequest, ++ HiddenStoreOperationStats, ++ HiddenTensorDatabase, ++ MMMeta, ++ build_tensor_meta, ++ validate_loaded_tensor, ++) ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( ++ HiddenStoreLoadError, ++ MooncakeHiddenStoreClient, ++) ++from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( ++ get_mooncake_dp_engine_index, ++) ++from vllm.logger import init_logger ++from vllm.utils.network_utils import make_zmq_socket ++ ++logger = init_logger(__name__) ++ ++LOOKUP_MSG = b"LOOKUP" ++BATCH_LOOKUP_MSG = b"BATCH_LOOKUP" ++RESP_BATCH = b"BATCH" ++RESP_HIT = b"HIT" ++RESP_MISS = b"MISS" ++RESP_ERR = b"ERR" ++ ++ ++class HiddenStoreWorker: ++ """Synchronous hidden tensor load/save path used by the EC connector.""" ++ ++ def __init__( ++ self, ++ store_client: MooncakeHiddenStoreClient, ++ tensor_database: HiddenTensorDatabase | None = None, ++ key_metadata: HiddenKeyMetadata | None = None, ++ ): ++ self.store_client = store_client ++ self.tensor_database = tensor_database or HiddenTensorDatabase() ++ self.key_metadata = key_metadata ++ self.sending_thread: HiddenStoreSendingThread | None = None ++ self._operation_stats_lock = threading.Lock() ++ self._operation_stats = HiddenStoreOperationStats() ++ ++ def make_pool_key(self, identifier: str) -> HiddenPoolKey: ++ assert self.key_metadata is not None ++ return HiddenPoolKey( ++ key_metadata=self.key_metadata, ++ identifier=identifier, ++ ) ++ ++ def start_sending_thread(self) -> None: ++ if self.sending_thread is not None: ++ return ++ self.sending_thread = HiddenStoreSendingThread(self) ++ self.sending_thread.start() ++ ++ def enqueue_save(self, request: HiddenSaveRequest) -> None: ++ if self.sending_thread is None: ++ self.save_tensor( ++ request.pool_key, ++ request.tensor, ++ with_soft_pin=request.with_soft_pin, ++ ) ++ return ++ self.sending_thread.add_request(request) ++ ++ def get_finished_sending(self) -> set[str]: ++ if self.sending_thread is None: ++ return set() ++ return self.sending_thread.get_and_clear_finished_identifiers() ++ ++ def get_operation_stats(self) -> HiddenStoreOperationStats | None: ++ with self._operation_stats_lock: ++ if self._operation_stats.is_empty(): ++ return None ++ stats = self._operation_stats ++ self._operation_stats = HiddenStoreOperationStats() ++ return stats ++ ++ def _record_operation( ++ self, ++ operation: str, ++ duration_seconds: float, ++ num_keys: int, ++ *, ++ num_bytes: int = 0, ++ status: str = "ok", ++ num_failed_keys: int = 0, ++ ) -> None: ++ with self._operation_stats_lock: ++ self._operation_stats.record_operation( ++ operation=operation, ++ duration_seconds=duration_seconds, ++ num_keys=num_keys, ++ num_bytes=num_bytes, ++ status=status, ++ num_failed_keys=num_failed_keys, ++ ) ++ ++ def shutdown(self) -> None: ++ if self.sending_thread is not None: ++ self.sending_thread.close() ++ self.sending_thread = None ++ ++ def lookup(self, identifier: str) -> bool: ++ """Return whether the hidden object exists in Mooncake Store.""" ++ return self.lookup_batch([identifier]).get(identifier, False) ++ ++ def lookup_batch(self, identifiers: list[str]) -> dict[str, bool]: ++ """Return whether hidden objects exist in Mooncake Store.""" ++ pool_keys = [self.make_pool_key(identifier) for identifier in identifiers] ++ started = time.perf_counter() ++ try: ++ exists = self.store_client.batch_exists(pool_keys) ++ except Exception: ++ self._record_operation( ++ "lookup_exists", ++ time.perf_counter() - started, ++ len(pool_keys), ++ status="error", ++ num_failed_keys=len(pool_keys), ++ ) ++ raise ++ ++ failed_keys = sum(1 for hit in exists if not hit) ++ self._record_operation( ++ "lookup_exists", ++ time.perf_counter() - started, ++ len(pool_keys), ++ status="miss" if failed_keys else "ok", ++ num_failed_keys=failed_keys, ++ ) ++ results = dict(zip(identifiers, exists, strict=True)) ++ for pool_key, hit in zip(pool_keys, exists, strict=True): ++ if hit: ++ logger.info( ++ "hidden_store_lookup_hit identifier=%s hidden_pool_key=%s", ++ pool_key.identifier, ++ pool_key.to_string(), ++ ) ++ else: ++ logger.info( ++ "hidden_store_lookup_miss identifier=%s hidden_pool_key=%s " ++ "reason=missing_object", ++ pool_key.identifier, ++ pool_key.to_string(), ++ ) ++ return results ++ ++ def save_tensor( ++ self, ++ pool_key: HiddenPoolKey, ++ tensor: torch.Tensor, ++ with_soft_pin: bool = False, ++ ) -> None: ++ exists_started = time.perf_counter() ++ try: ++ exists = self.store_client.exists(pool_key) ++ except Exception: ++ self._record_operation( ++ "save_exists", ++ time.perf_counter() - exists_started, ++ 1, ++ status="error", ++ num_failed_keys=1, ++ ) ++ raise ++ ++ self._record_operation( ++ "save_exists", ++ time.perf_counter() - exists_started, ++ 1, ++ status="ok" if exists else "miss", ++ ) ++ if exists: ++ logger.info( ++ "hidden_store_save_skip identifier=%s hidden_pool_key=%s " ++ "reason=exists", ++ pool_key.identifier, ++ pool_key.to_string(), ++ ) ++ return ++ ++ started = time.perf_counter() ++ stored_tensor = tensor if tensor.is_contiguous() else tensor.contiguous() ++ used_staging = stored_tensor is not tensor ++ tensor_meta = build_tensor_meta(pool_key, stored_tensor) ++ try: ++ self.store_client.put_tensor( ++ pool_key, ++ stored_tensor, ++ with_soft_pin=with_soft_pin, ++ ) ++ except Exception: ++ self._record_operation( ++ "save_put", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes, ++ status="error", ++ num_failed_keys=1, ++ ) ++ raise ++ self._record_operation( ++ "save_put", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes, ++ status="ok", ++ ) ++ logger.info( ++ "hidden_store_put identifier=%s hidden_pool_key=%s nbytes=%d " ++ "used_staging=%s hidden_store_put_ms=%.3f", ++ pool_key.identifier, ++ pool_key.to_string(), ++ tensor_meta.nbytes, ++ used_staging, ++ (time.perf_counter() - started) * 1000.0, ++ ) ++ ++ def load( ++ self, ++ items: list[MMMeta], ++ encoder_cache: dict[str, torch.Tensor], ++ *, ++ device: torch.device | str | None = None, ++ ) -> None: ++ for item in items: ++ load_spec = item.load_spec ++ if load_spec is None or not load_spec.can_load: ++ continue ++ if item.identifier in encoder_cache: ++ logger.debug( ++ "hidden_store_load_skip identifier=%s " ++ "reason=local_encoder_cache", ++ item.identifier, ++ ) ++ continue ++ ++ started = time.perf_counter() ++ pool_key = self.make_pool_key(item.identifier) ++ tensor_meta = None ++ try: ++ tensor_meta = self.store_client.get_tensor_meta(pool_key) ++ if tensor_meta is None: ++ raise HiddenStoreLoadError( ++ "failed to load hidden tensor metadata for " ++ f"{pool_key.to_string()}" ++ ) ++ ++ target_device = device ++ if target_device is None: ++ target_device = "cuda" if torch.cuda.is_available() else None ++ target = torch.empty( ++ tensor_meta.shape, ++ dtype=_resolve_torch_dtype(tensor_meta.dtype), ++ device=target_device, ++ ) ++ _data_key, addrs, sizes = self.tensor_database.prepare_value( ++ pool_key, ++ target, ++ ) ++ self.store_client.get_tensor_payload( ++ pool_key, ++ addrs[0], ++ sizes[0], ++ tensor_meta.data_offset, ++ ) ++ validate_loaded_tensor(target, tensor_meta) ++ except Exception: ++ self._record_operation( ++ "load_get", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes if tensor_meta is not None else 0, ++ status="error", ++ num_failed_keys=1, ++ ) ++ raise ++ encoder_cache[item.identifier] = target ++ self._record_operation( ++ "load_get", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes, ++ status="ok", ++ ) ++ logger.info( ++ "hidden_store_get identifier=%s hidden_pool_key=%s nbytes=%d " ++ "hidden_store_get_ms=%.3f", ++ item.identifier, ++ pool_key.to_string(), ++ tensor_meta.nbytes, ++ (time.perf_counter() - started) * 1000.0, ++ ) ++ ++ ++def _resolve_torch_dtype(dtype: str) -> torch.dtype: ++ if dtype == "torch.float16": ++ return torch.float16 ++ if dtype == "torch.bfloat16": ++ return torch.bfloat16 ++ if dtype == "torch.float32": ++ return torch.float32 ++ raise HiddenStoreLoadError(f"unsupported hidden tensor dtype: {dtype}") ++ ++ ++class HiddenStoreSendingThread(threading.Thread): ++ """Background thread for storing hidden tensors to the store.""" ++ ++ def __init__(self, store_worker: HiddenStoreWorker): ++ super().__init__(daemon=True, name="HiddenStoreSendingThread") ++ self.store_worker = store_worker ++ self.request_queue: queue.Queue[HiddenSaveRequest | None] = queue.Queue() ++ self.done_task_lock = threading.Lock() ++ self.finished_identifiers: set[str] = set() ++ self.failed_identifiers: set[str] = set() ++ self.failure_reasons: dict[str, str] = {} ++ self._closed = threading.Event() ++ ++ def add_request(self, request: HiddenSaveRequest) -> None: ++ self.request_queue.put(request) ++ ++ def get_and_clear_finished_identifiers(self) -> set[str]: ++ with self.done_task_lock: ++ finished = self.finished_identifiers.copy() ++ self.finished_identifiers.clear() ++ return finished ++ ++ def get_and_clear_failed_identifiers(self) -> set[str]: ++ with self.done_task_lock: ++ failed = self.failed_identifiers.copy() ++ self.failed_identifiers.clear() ++ return failed ++ ++ def set_finished_identifier(self, identifier: str) -> None: ++ with self.done_task_lock: ++ self.finished_identifiers.add(identifier) ++ ++ def set_failed_identifier(self, identifier: str, error: Exception) -> None: ++ with self.done_task_lock: ++ self.failed_identifiers.add(identifier) ++ self.failure_reasons[identifier] = str(error) ++ ++ def run(self) -> None: ++ while True: ++ request = self.request_queue.get() ++ try: ++ if request is None: ++ return ++ self.store_worker.save_tensor( ++ request.pool_key, ++ request.tensor, ++ with_soft_pin=request.with_soft_pin, ++ ) ++ self.set_finished_identifier(request.identifier) ++ except Exception as e: ++ if request is not None: ++ self.set_failed_identifier(request.identifier, e) ++ logger.error("Error in %s: %s", self.name, e) ++ finally: ++ self.request_queue.task_done() ++ ++ def close(self) -> None: ++ if self._closed.is_set(): ++ return ++ self._closed.set() ++ self.request_queue.put(None) ++ ++ ++class HiddenLookupServer: ++ """Worker rank-0 admin channel for scheduler-side hidden lookups.""" ++ ++ def __init__( ++ self, ++ store_worker: HiddenStoreWorker, ++ vllm_config: VllmConfig, ++ ): ++ self.ctx = zmq.Context() # type: ignore[attr-defined] ++ socket_path = get_zmq_rpc_path_hidden_lookup(vllm_config) ++ self._ipc_path = socket_path.removeprefix("ipc://") ++ if os.path.exists(self._ipc_path): ++ os.unlink(self._ipc_path) ++ self.socket = make_zmq_socket( ++ self.ctx, ++ socket_path, ++ zmq.REP, # type: ignore[attr-defined] ++ bind=True, ++ ) ++ ++ self.store_worker = store_worker ++ self.running = True ++ ++ def process_request(): ++ while self.running: ++ all_frames = self.socket.recv_multipart(copy=False) ++ msg_type = bytes(all_frames[0]) ++ ++ if msg_type == LOOKUP_MSG: ++ try: ++ identifier = bytes(all_frames[1]).decode("utf-8") ++ exists = self.store_worker.lookup(identifier) ++ if not exists: ++ self.socket.send_multipart([RESP_MISS]) ++ else: ++ self.socket.send_multipart([RESP_HIT]) ++ except Exception: ++ logger.exception("HiddenLookupServer lookup failed") ++ self.socket.send_multipart([RESP_ERR]) ++ elif msg_type == BATCH_LOOKUP_MSG: ++ try: ++ identifiers = [ ++ bytes(frame).decode("utf-8") for frame in all_frames[1:] ++ ] ++ exists = self.store_worker.lookup_batch(identifiers) ++ frames = [ ++ RESP_HIT if exists.get(identifier, False) else RESP_MISS ++ for identifier in identifiers ++ ] ++ self.socket.send_multipart([RESP_BATCH, *frames]) ++ except Exception: ++ logger.exception("HiddenLookupServer batch lookup failed") ++ self.socket.send_multipart([RESP_ERR]) ++ else: ++ logger.warning( ++ "HiddenLookupServer received unknown msg_type: %r", ++ msg_type, ++ ) ++ self.socket.send_multipart([RESP_ERR]) ++ ++ self.thread = threading.Thread(target=process_request, daemon=True) ++ self.thread.start() ++ ++ def close(self): ++ self.running = False ++ self.socket.close(linger=0) ++ if os.path.exists(self._ipc_path): ++ os.unlink(self._ipc_path) ++ ++ ++class HiddenLookupClient: ++ """Scheduler-side client for worker rank-0 hidden lookup queries.""" ++ ++ def __init__(self, vllm_config: VllmConfig): ++ self.ctx = zmq.Context() # type: ignore[attr-defined] ++ socket_path = get_zmq_rpc_path_hidden_lookup(vllm_config) ++ self.socket = make_zmq_socket( ++ self.ctx, ++ socket_path, ++ zmq.REQ, # type: ignore[attr-defined] ++ bind=False, ++ ) ++ self.executor = ThreadPoolExecutor( ++ max_workers=1, ++ thread_name_prefix="HiddenLookupClient", ++ ) ++ self.futures: dict[str, Future[dict[str, bool]]] = {} ++ ++ def lookup(self, identifier: str) -> bool: ++ result = self.lookup_batch([identifier], non_block=False) ++ assert result is not None ++ return result.get(identifier, False) ++ ++ def _lookup_batch(self, identifiers: list[str]) -> dict[str, bool]: ++ self.socket.send_multipart( ++ [ ++ BATCH_LOOKUP_MSG, ++ *(identifier.encode("utf-8") for identifier in identifiers), ++ ] ++ ) ++ resp = self.socket.recv_multipart() ++ msg_type = bytes(resp[0]) ++ if msg_type == RESP_BATCH: ++ states = [bytes(frame) == RESP_HIT for frame in resp[1:]] ++ if len(states) != len(identifiers): ++ logger.warning( ++ "HiddenLookupClient received malformed batch response: " ++ "identifiers=%d states=%d", ++ len(identifiers), ++ len(states), ++ ) ++ return {identifier: False for identifier in identifiers} ++ return dict(zip(identifiers, states, strict=True)) ++ if msg_type == RESP_ERR: ++ return {identifier: False for identifier in identifiers} ++ logger.warning("HiddenLookupClient received unknown response: %r", msg_type) ++ return {identifier: False for identifier in identifiers} ++ ++ def lookup_batch( ++ self, ++ identifiers: list[str], ++ non_block: bool = False, ++ ) -> dict[str, bool] | None: ++ identifiers = list(dict.fromkeys(identifiers)) ++ if not identifiers: ++ return {} ++ ++ new_identifiers = [ ++ identifier for identifier in identifiers if identifier not in self.futures ++ ] ++ if new_identifiers: ++ future = self.executor.submit(self._lookup_batch, new_identifiers) ++ for identifier in new_identifiers: ++ self.futures[identifier] = future ++ ++ if non_block and any( ++ not self.futures[identifier].done() for identifier in identifiers ++ ): ++ return None ++ ++ results: dict[str, bool] = {} ++ for identifier in identifiers: ++ future = self.futures[identifier] ++ try: ++ batch_results = future.result() ++ results[identifier] = batch_results.get(identifier, False) ++ except Exception as e: ++ logger.error("Async hidden lookup failed for %s: %s", identifier, e) ++ results[identifier] = False ++ finally: ++ self.futures.pop(identifier, None) ++ return results ++ ++ def discard(self, identifier: str) -> None: ++ future = self.futures.pop(identifier, None) ++ if future is None: ++ return ++ if not any(existing is future for existing in self.futures.values()): ++ future.cancel() ++ ++ def close(self): ++ self.executor.shutdown(wait=False, cancel_futures=True) ++ self.futures.clear() ++ self.socket.close(linger=0) ++ ++ ++def get_zmq_rpc_path_hidden_lookup(vllm_config: VllmConfig) -> str: ++ """Construct IPC path for Hidden Store lookup socket.""" ++ assert vllm_config.ec_transfer_config is not None ++ dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config) ++ base_url = envs.VLLM_RPC_BASE_PATH ++ hostname = socket.gethostname() ++ extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config ++ rpc_port = extra_config.get( ++ "hidden_lookup_rpc_port", ++ extra_config.get("lookup_rpc_port", 0), ++ ) ++ logger.debug("Hidden lookup Base URL: %s, RPC Port: %s", base_url, rpc_port) ++ return ( ++ f"ipc://{base_url}/hidden_lookup_rpc_port_{rpc_port}_host_{hostname}" ++ f"_dp_rank{dp_rank}" ++ ) diff --git a/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-latest-refinement.patch b/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-latest-refinement.patch new file mode 100644 index 00000000..e0747edd --- /dev/null +++ b/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-b4482f0a1-latest-refinement.patch @@ -0,0 +1,2018 @@ +commit b4482f0a10a25ddbf34b687129c344594fab4613 +Author: 聪明企鹅\70733 <707334817@qq.com> +AuthorDate: Sun Jul 5 19:11:14 2026 +0800 +Commit: 聪明企鹅\70733 <707334817@qq.com> +CommitDate: Sun Jul 5 19:11:14 2026 +0800 + + Refine Mooncake hidden EC connector lifecycle + +diff --git a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py +index 6264bd5f0..df8baf237 100644 +--- a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py ++++ b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_connector.py +@@ -5,29 +5,34 @@ from types import SimpleNamespace + + import torch + ++from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden import ( ++ connector as connector_module, ++) + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.connector import ( + MooncakeStoreECConnector, + ) + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( +- HIDDEN_LAYOUT_VERSION, ++ HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenPoolKey, + LoadSpec, + MMMeta, + MooncakeStoreConnectorMetadata, + ) ++from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange + + + class FakeWorker: + def __init__(self): + self.requests = [] + self.key_metadata = HiddenKeyMetadata( ++ cache_prefix="", ++ kind="encoder_output", + model_name="qwen", +- mm_encoder_config_hash="encoder-config-a", +- hidden_parallel_key=( +- "tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights@storage:replicated" +- ), +- layout=HIDDEN_LAYOUT_VERSION, ++ encoder="encoder-config-a", ++ storage="replicated_object", ++ parallel="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", ++ tensor_layout=HIDDEN_TENSOR_LAYOUT, + ) + + def make_pool_key(self, identifier: str) -> HiddenPoolKey: +@@ -41,22 +46,268 @@ def make_connector(*, soft_pin_video_hidden: bool = False): + connector = MooncakeStoreECConnector.__new__(MooncakeStoreECConnector) + connector._is_producer = True + connector._is_consumer = False ++ connector.lookup_client = None ++ connector.lookup_async = True + connector.worker = FakeWorker() + connector._connector_metadata = None + connector.soft_pin_video_hidden = soft_pin_video_hidden + connector.load_specs = {} +- connector._load_identifiers_to_schedule = set() +- connector._save_identifiers_to_schedule = set() ++ connector.lookup_result_cache = {} ++ connector.identifier_waiters = {} ++ connector._candidate_consumes = {} ++ connector._candidate_loads = {} ++ connector._candidate_saves = {} + connector._load_modalities = {} + connector._save_modalities = {} + return connector + + ++class FakeLookupClient: ++ def __init__(self, results): ++ self.results = list(results) ++ self.calls = [] ++ self.discarded = [] ++ ++ def lookup_batch(self, identifiers, non_block=True): ++ self.calls.append((tuple(identifiers), non_block)) ++ return self.results.pop(0) ++ ++ def discard(self, identifier): ++ self.discarded.append(identifier) ++ ++ ++def make_request(request_id, features): ++ mm_features = [ ++ MultiModalFeatureSpec( ++ data=None, ++ modality=modality, ++ identifier=identifier, ++ mm_position=PlaceholderRange(offset=offset, length=length), ++ ) ++ for identifier, offset, length, modality in features ++ ] ++ return SimpleNamespace( ++ request_id=request_id, ++ mm_features=mm_features, ++ num_tokens=1000, ++ ) ++ ++ ++def make_scheduler_output(*, finished_req_ids=None, preempted_req_ids=None): ++ return SimpleNamespace( ++ finished_req_ids=finished_req_ids or set(), ++ preempted_req_ids=preempted_req_ids, ++ ) ++ ++ ++def test_build_hidden_key_metadata_uses_structured_key_fields(monkeypatch): ++ monkeypatch.setattr( ++ connector_module, ++ "get_tensor_model_parallel_world_size", ++ lambda: 4, ++ ) ++ monkeypatch.setattr( ++ connector_module, ++ "get_pcp_group", ++ lambda: SimpleNamespace(world_size=1), ++ ) ++ monkeypatch.setattr( ++ connector_module, ++ "get_dcp_group", ++ lambda: SimpleNamespace(world_size=1), ++ ) ++ multimodal_config = SimpleNamespace( ++ compute_hash=lambda: "encoder-config-a", ++ mm_encoder_tp_mode="data", ++ ) ++ vllm_config = SimpleNamespace( ++ model_config=SimpleNamespace( ++ model="/models/qwen", ++ multimodal_config=multimodal_config, ++ ), ++ parallel_config=SimpleNamespace(pipeline_parallel_size=2), ++ ec_transfer_config=SimpleNamespace( ++ ec_connector_extra_config={ ++ "cache_prefix": "shared-prefix", ++ "hidden_cache_prefix": "hidden-prefix", ++ } ++ ), ++ ) ++ ++ metadata = connector_module.build_hidden_key_metadata(vllm_config) ++ ++ assert metadata.cache_prefix == "hidden-prefix" ++ assert metadata.kind == "encoder_output" ++ assert metadata.model_name == "qwen" ++ assert metadata.encoder == "encoder-config-a" ++ assert metadata.storage == "replicated_object" ++ assert metadata.parallel == "tp:4@pp:2@pcp:1@dcp:1@mm_tp:data" ++ assert "storage" not in metadata.parallel ++ assert metadata.tensor_layout == "tensor" ++ ++ ++def test_ensure_cache_available_defers_pending_batch_lookup(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient([None]) ++ request = make_request( ++ "req-1", ++ [ ++ ("image-1", 20, 60, "image"), ++ ("image-2", 500, 60, "image"), ++ ], ++ ) ++ ++ assert not connector.ensure_cache_available(request, num_computed_tokens=0) ++ ++ assert connector.lookup_client.calls == [ ++ (("image-1", "image-2"), True), ++ ] ++ assert connector.identifier_waiters == { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1"}, ++ } ++ ++ ++def test_ensure_cache_available_deduplicates_request_waiters_and_lookup_results(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient( ++ [ ++ {"image-1": True, "image-2": False}, ++ ] ++ ) ++ request = make_request( ++ "req-1", ++ [ ++ ("image-1", 20, 60, "image"), ++ ("image-2", 500, 60, "image"), ++ ], ++ ) ++ ++ assert connector.ensure_cache_available(request, num_computed_tokens=0) ++ assert connector.ensure_cache_available(request, num_computed_tokens=0) ++ ++ assert connector.lookup_client.calls == [ ++ (("image-1", "image-2"), True), ++ ] ++ assert connector.identifier_waiters == { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1"}, ++ } ++ assert connector.lookup_result_cache == { ++ "image-1": True, ++ "image-2": False, ++ } ++ ++ ++def test_has_cache_item_is_local_only(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = SimpleNamespace(lookup=lambda identifier: True) ++ connector.lookup_result_cache = {"image-1": True, "image-2": False} ++ ++ assert connector.has_cache_item("image-1") ++ assert not connector.has_cache_item("image-2") ++ assert not connector.has_cache_item("unknown") ++ ++ ++def test_build_connector_meta_commits_waiter_consumes_and_keeps_unreached_image(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_result_cache = {"image-1": True, "image-2": True} ++ connector.identifier_waiters = { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1"}, ++ } ++ connector.load_specs["image-1"] = LoadSpec(can_load=False) ++ request = make_request( ++ "req-1", ++ [ ++ ("image-1", 20, 60, "image"), ++ ("image-2", 500, 60, "image"), ++ ], ++ ) ++ ++ connector.update_state_after_alloc(request, 0) ++ meta = connector.build_connector_meta(make_scheduler_output()) ++ ++ assert [item.identifier for item in meta.items] == ["image-1"] ++ assert "image-1" not in connector.identifier_waiters ++ assert "image-1" not in connector.lookup_result_cache ++ assert connector.identifier_waiters == {"image-2": {"req-1"}} ++ assert connector.lookup_result_cache == {"image-2": True} ++ ++ ++def test_build_connector_meta_rolls_back_preempted_candidate_state(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_result_cache = {"image-1": True} ++ connector.identifier_waiters = {"image-1": {"req-1"}} ++ connector.load_specs["image-1"] = LoadSpec(can_load=False) ++ request = make_request("req-1", [("image-1", 20, 60, "image")]) ++ ++ connector.update_state_after_alloc(request, 0) ++ meta = connector.build_connector_meta( ++ make_scheduler_output(preempted_req_ids={"req-1"}) ++ ) ++ ++ assert meta.items == [] ++ assert connector.identifier_waiters == {"image-1": {"req-1"}} ++ assert connector.lookup_result_cache == {"image-1": True} ++ assert "image-1" in connector.load_specs ++ ++ ++def test_build_connector_meta_cleans_finished_waiters(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient([]) ++ connector.lookup_result_cache = {"image-1": True, "image-2": True} ++ connector.identifier_waiters = { ++ "image-1": {"req-1"}, ++ "image-2": {"req-1", "req-2"}, ++ } ++ ++ connector.build_connector_meta(make_scheduler_output(finished_req_ids={"req-1"})) ++ ++ assert "image-1" not in connector.identifier_waiters ++ assert "image-1" not in connector.lookup_result_cache ++ assert connector.identifier_waiters == {"image-2": {"req-2"}} ++ assert connector.lookup_result_cache == {"image-2": True} ++ assert connector.lookup_client.discarded == ["image-1"] ++ ++ ++def test_cleanup_lookup_results_discards_inflight_lookup_without_waiters(): ++ connector = make_connector() ++ connector._is_consumer = True ++ connector._is_producer = False ++ connector.lookup_client = FakeLookupClient([]) ++ connector.identifier_waiters = {"image-1": set()} ++ connector.lookup_result_cache = {"image-1": True} ++ connector.load_specs["image-1"] = LoadSpec(can_load=False) ++ ++ connector._cleanup_lookup_results_without_waiters() ++ ++ assert connector.identifier_waiters == {} ++ assert connector.lookup_result_cache == {} ++ assert connector.load_specs == {} ++ assert connector.lookup_client.discarded == ["image-1"] ++ ++ + def test_build_connector_meta_merges_load_and_save_item_by_identifier(): + connector = make_connector() + connector._is_consumer = True + connector.load_specs["video-hash"] = LoadSpec(can_load=False) ++ connector.lookup_result_cache["video-hash"] = True ++ connector.identifier_waiters["video-hash"] = {"req-1"} + request = SimpleNamespace( ++ request_id="req-1", + mm_features=[ + SimpleNamespace( + identifier="video-hash", +@@ -66,7 +317,7 @@ def test_build_connector_meta_merges_load_and_save_item_by_identifier(): + ) + + connector.update_state_after_alloc(request, 0) +- meta = connector.build_connector_meta(SimpleNamespace(finished_req_ids=set())) ++ meta = connector.build_connector_meta(make_scheduler_output()) + + assert len(meta.items) == 1 + item = meta.items[0] +@@ -75,7 +326,6 @@ def test_build_connector_meta_merges_load_and_save_item_by_identifier(): + assert item.can_save + assert item.load_spec is not None + assert item.load_spec.can_load +- assert meta.unfinished_identifiers == {"video-hash"} + assert connector.load_specs == {} + + +diff --git a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py +index a03ed88b6..d370014ab 100644 +--- a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py ++++ b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_protocol.py +@@ -4,7 +4,7 @@ + import torch + + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( +- HIDDEN_LAYOUT_VERSION, ++ HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenPoolKey, + LoadSpec, +@@ -20,17 +20,23 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import + def make_pool_key( + identifier: str = "image-hash", + *, ++ cache_prefix: str = "", ++ kind: str = "encoder_output", + model_name: str = "qwen", +- mm_encoder_config_hash: str = "encoder-config-a", +- hidden_parallel_key: str = "tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights@storage:replicated", +- layout: str = HIDDEN_LAYOUT_VERSION, ++ encoder: str = "encoder-config-a", ++ storage: str = "replicated_object", ++ parallel: str = "tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", ++ tensor_layout: str = HIDDEN_TENSOR_LAYOUT, + ) -> HiddenPoolKey: + return HiddenPoolKey( + key_metadata=HiddenKeyMetadata( ++ cache_prefix=cache_prefix, ++ kind=kind, + model_name=model_name, +- mm_encoder_config_hash=mm_encoder_config_hash, +- hidden_parallel_key=hidden_parallel_key, +- layout=layout, ++ encoder=encoder, ++ storage=storage, ++ parallel=parallel, ++ tensor_layout=tensor_layout, + ), + identifier=identifier, + ) +@@ -42,25 +48,42 @@ def test_hidden_pool_key_is_the_single_tensor_object_key(): + data_key = make_hidden_data_key(pool_key) + + assert data_key == pool_key.to_string() ++ assert data_key.startswith("hidden@") ++ assert "kind:encoder_output" in data_key + assert "model:qwen" in data_key +- assert "mm_encoder:encoder-config-a" in data_key ++ assert "encoder:encoder-config-a" in data_key ++ assert "storage:replicated_object" in data_key + assert ( +- "parallel:tp%3A1%40pp%3A1%40pcp%3A1%40dcp%3A1%40mm_tp%3Aweights%40storage%3Areplicated" ++ "parallel:tp%3A1%40pp%3A1%40pcp%3A1%40dcp%3A1%40mm_tp%3Aweights" + in data_key + ) +- assert "layout:vllm-encoder-cache-tensor-v1" in data_key ++ assert "tensor_layout:tensor" in data_key ++ assert "storage%3Areplicated" not in data_key ++ assert "writer" not in data_key + assert "adapter:" not in data_key + assert "modality:" not in data_key + assert "image-hash" in data_key + + + def test_same_identifier_with_different_encoder_config_uses_different_keys(): +- pool_key_a = make_pool_key(mm_encoder_config_hash="encoder-config-a") +- pool_key_b = make_pool_key(mm_encoder_config_hash="encoder-config-b") ++ pool_key_a = make_pool_key(encoder="encoder-config-a") ++ pool_key_b = make_pool_key(encoder="encoder-config-b") + + assert make_hidden_data_key(pool_key_a) != make_hidden_data_key(pool_key_b) + + ++def test_cache_prefix_namespaces_hidden_pool_key(): ++ pool_key_a = make_pool_key(cache_prefix="deployment-a") ++ pool_key_b = make_pool_key(cache_prefix="deployment-b") ++ ++ data_key_a = make_hidden_data_key(pool_key_a) ++ data_key_b = make_hidden_data_key(pool_key_b) ++ ++ assert data_key_a.startswith("deployment-a@hidden@") ++ assert data_key_b.startswith("deployment-b@hidden@") ++ assert data_key_a != data_key_b ++ ++ + def test_request_id_and_modality_are_not_part_of_hidden_pool_key(): + pool_key = make_pool_key(identifier="image-hash") + +@@ -94,6 +117,7 @@ def test_tensor_meta_describes_canonical_contiguous_tensor(): + tensor_meta = build_tensor_meta(pool_key, stored) + + assert tensor_meta.pool_key == pool_key ++ assert tensor_meta.layout == HIDDEN_TENSOR_LAYOUT + assert tensor_meta.shape == tuple(stored.shape) + assert tensor_meta.dtype == "torch.float16" + assert tensor_meta.nbytes == stored.numel() * stored.element_size() +@@ -112,8 +136,8 @@ def test_tensor_meta_rejects_non_contiguous_tensor(): + + + def test_pool_key_namespace_carries_reuse_compatibility(): +- pool_key_a = make_pool_key(mm_encoder_config_hash="encoder-config-a") +- pool_key_b = make_pool_key(mm_encoder_config_hash="encoder-config-b") ++ pool_key_a = make_pool_key(encoder="encoder-config-a") ++ pool_key_b = make_pool_key(encoder="encoder-config-b") + + assert pool_key_a != pool_key_b + assert make_hidden_data_key(pool_key_a) != make_hidden_data_key(pool_key_b) +diff --git a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py +index f159c2ea8..c1fd96b18 100644 +--- a/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py ++++ b/tests/v1/ec_connector/unit/test_mooncake_store_hidden_worker.py +@@ -5,11 +5,12 @@ import ctypes + import sys + import struct + import types ++from concurrent.futures import Future + + import torch + + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( +- HIDDEN_LAYOUT_VERSION, ++ HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenPoolKey, + HiddenSaveRequest, +@@ -21,10 +22,14 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import + make_hidden_data_key, + ) + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.store_client import ( ++ HiddenStoreError, ++ HiddenStoreLoadError, ++ HiddenStoreSaveError, + MooncakeHiddenStoreClient, + _get_hidden_state_object_data_type, + ) + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.worker import ( ++ HiddenLookupClient, + HiddenStoreSendingThread, + HiddenStoreWorker, + ) +@@ -42,18 +47,27 @@ TORCH_DTYPE_TO_MOONCAKE_DTYPE = { + class FakeStore: + def __init__(self): + self.objects = {} ++ self.batch_is_exist_calls = [] + self.registered = [] ++ self.unregistered = [] + self.pub_tensors = [] + self.range_gets = [] ++ self.fail_register_addrs = set() ++ self.raise_on_batch_put = False ++ self.batch_put_results = [0] + + def batch_is_exist(self, keys): ++ self.batch_is_exist_calls.append(list(keys)) + return [1 if key in self.objects else 0 for key in keys] + + def register_buffer(self, addr, size): ++ if addr in self.fail_register_addrs: ++ return -1 + self.registered.append((addr, size)) + return 0 + + def unregister_buffer(self, addr): ++ self.unregistered.append(addr) + return 0 + + def pub_tensor(self, key, tensor, replicate_config=None): +@@ -114,6 +128,20 @@ class FakeStore: + return results + + ++class FakeBufferStore(FakeStore): ++ def batch_put_from_multi_buffers( ++ self, ++ keys, ++ buffer_ptrs, ++ buffer_sizes, ++ replicate_config=None, ++ ): ++ if self.raise_on_batch_put: ++ raise RuntimeError("batch put failed") ++ self.objects[keys[0]] = b"tensor-object" ++ return self.batch_put_results ++ ++ + class FakeReplicateConfig: + def __init__(self): + self.replica_num = 1 +@@ -149,10 +177,13 @@ class FakeObjectDataTypeNoTensor: + def make_pool_key(identifier: str = "image-hash") -> HiddenPoolKey: + return HiddenPoolKey( + key_metadata=HiddenKeyMetadata( ++ cache_prefix="", ++ kind="encoder_output", + model_name="qwen", +- mm_encoder_config_hash="encoder-config-a", +- hidden_parallel_key="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights@storage:replicated", +- layout=HIDDEN_LAYOUT_VERSION, ++ encoder="encoder-config-a", ++ storage="replicated_object", ++ parallel="tp:1@pp:1@pcp:1@dcp:1@mm_tp:weights", ++ tensor_layout=HIDDEN_TENSOR_LAYOUT, + ), + identifier=identifier, + ) +@@ -189,13 +220,78 @@ def test_worker_lookup_checks_existence_without_reading_tensor_metadata(): + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) +- worker.save_tensor(pool_key, tensor, now_ms=1234) ++ worker.save_tensor(pool_key, tensor) + + assert worker.lookup(pool_key.identifier) + assert not worker.lookup("missing-image-hash") + assert store.range_gets == [] + + ++def test_worker_batch_lookup_checks_existence_in_one_store_call(): ++ pool_key_a = make_pool_key("image-a") ++ pool_key_b = make_pool_key("image-b") ++ store = FakeBufferStore() ++ store.objects[make_hidden_data_key(pool_key_a)] = b"tensor-object" ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key_a.key_metadata, ++ ) ++ ++ results = worker.lookup_batch(["image-a", "image-b"]) ++ ++ assert results == {"image-a": True, "image-b": False} ++ assert store.batch_is_exist_calls == [ ++ [make_hidden_data_key(pool_key_a), make_hidden_data_key(pool_key_b)] ++ ] ++ assert store.range_gets == [] ++ ++ ++def test_lookup_client_discard_removes_identifier_future_mapping(): ++ client = HiddenLookupClient.__new__(HiddenLookupClient) ++ future: Future[dict[str, bool]] = Future() ++ client.futures = { ++ "image-a": future, ++ "image-b": future, ++ } ++ ++ client.discard("image-a") ++ ++ assert "image-a" not in client.futures ++ assert client.futures == {"image-b": future} ++ assert not future.cancelled() ++ ++ client.discard("image-b") ++ ++ assert client.futures == {} ++ assert future.cancelled() ++ ++ ++def test_worker_lookup_records_minimal_operation_stats(): ++ pool_key = make_pool_key() ++ store = FakeBufferStore() ++ store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ ++ assert worker.lookup_batch(["image-hash", "missing-image-hash"]) == { ++ "image-hash": True, ++ "missing-image-hash": False, ++ } ++ ++ stats = worker.get_operation_stats() ++ records = stats.data["lookup_exists"] ++ assert len(records) == 1 ++ assert records[0]["num_keys"] == 2 ++ assert records[0]["num_bytes"] == 0 ++ assert records[0]["status"] == "miss" ++ assert records[0]["num_failed_keys"] == 1 ++ assert worker.get_operation_stats() is None ++ ++ + def test_worker_save_stores_hidden_as_single_tensor_object(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) +@@ -208,13 +304,95 @@ def test_worker_save_stores_hidden_as_single_tensor_object(): + tensor_database=HiddenTensorDatabase(), + ) + +- worker.save_tensor(pool_key, tensor, now_ms=1234) ++ worker.save_tensor(pool_key, tensor) + + assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) + assert store.pub_tensors[0][2] is not None + assert make_hidden_data_key(pool_key) in store.objects + + ++def test_worker_save_rejects_dtype_that_load_cannot_decode(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float64) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ try: ++ worker.save_tensor(pool_key, tensor) ++ except HiddenStoreSaveError as exc: ++ assert "unsupported hidden tensor dtype" in str(exc) ++ else: ++ raise AssertionError("unsupported hidden dtype should fail before store put") ++ ++ assert store.pub_tensors == [] ++ ++ ++def test_buffer_put_unregisters_payload_and_metadata_buffers(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) ++ ++ client.put_tensor(pool_key, tensor) ++ ++ payload_addr = tensor.data_ptr() ++ metadata_addr = next( ++ addr for addr, size in store.registered if size == TENSOR_METADATA_SIZE ++ ) ++ assert payload_addr in store.unregistered ++ assert metadata_addr in store.unregistered ++ assert store.unregistered[-2:] == [metadata_addr, payload_addr] ++ ++ ++def test_buffer_put_unregisters_payload_and_metadata_when_put_raises(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ store.raise_on_batch_put = True ++ client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) ++ ++ try: ++ client.put_tensor(pool_key, tensor) ++ except RuntimeError as exc: ++ assert "batch put failed" in str(exc) ++ else: ++ raise AssertionError("batch put exception should propagate") ++ ++ payload_addr = tensor.data_ptr() ++ metadata_addr = next( ++ addr for addr, size in store.registered if size == TENSOR_METADATA_SIZE ++ ) ++ assert payload_addr in store.unregistered ++ assert metadata_addr in store.unregistered ++ ++ ++def test_buffer_put_unregisters_payload_when_metadata_registration_fails(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ original_register = store.register_buffer ++ ++ def register_buffer(addr, size): ++ if size == TENSOR_METADATA_SIZE: ++ store.fail_register_addrs.add(addr) ++ return original_register(addr, size) ++ ++ store.register_buffer = register_buffer ++ client = MooncakeHiddenStoreClient(store, replicate_config=FakeReplicateConfig()) ++ ++ try: ++ client.put_tensor(pool_key, tensor) ++ except HiddenStoreError: ++ pass ++ else: ++ raise AssertionError("metadata registration failure should raise") ++ ++ assert tensor.data_ptr() in store.unregistered ++ ++ + def test_worker_save_marks_hidden_state_data_type(monkeypatch): + fake_mooncake = types.ModuleType("mooncake") + fake_store = types.ModuleType("mooncake.store") +@@ -234,7 +412,7 @@ def test_worker_save_marks_hidden_state_data_type(monkeypatch): + tensor_database=HiddenTensorDatabase(), + ) + +- worker.save_tensor(pool_key, tensor, now_ms=1234) ++ worker.save_tensor(pool_key, tensor) + + used_config = store.pub_tensors[0][2] + assert used_config is not replicate_config +@@ -273,7 +451,7 @@ def test_worker_save_does_not_require_mooncake_object_group_support(): + tensor_database=HiddenTensorDatabase(), + ) + +- worker.save_tensor(pool_key, tensor, now_ms=1234) ++ worker.save_tensor(pool_key, tensor) + + assert store.pub_tensors[0][0] == make_hidden_data_key(pool_key) + +@@ -287,12 +465,50 @@ def test_worker_save_skips_existing_tensor_object(): + tensor_database=HiddenTensorDatabase(), + ) + +- worker.save_tensor(pool_key, tensor, now_ms=1234) +- worker.save_tensor(pool_key, tensor, now_ms=1235) ++ worker.save_tensor(pool_key, tensor) ++ worker.save_tensor(pool_key, tensor) + + assert len(store.pub_tensors) == 1 + + ++def test_worker_save_records_exists_and_put_operation_stats(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ stats = worker.get_operation_stats() ++ assert stats.data["save_exists"][0]["status"] == "miss" ++ assert stats.data["save_exists"][0]["num_keys"] == 1 ++ assert stats.data["save_put"][0]["status"] == "ok" ++ assert stats.data["save_put"][0]["num_keys"] == 1 ++ assert stats.data["save_put"][0]["num_bytes"] == ( ++ tensor.numel() * tensor.element_size() ++ ) ++ ++ ++def test_worker_save_existing_records_only_save_exists(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ store.objects[make_hidden_data_key(pool_key)] = b"tensor-object" ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ ++ worker.save_tensor(pool_key, tensor) ++ ++ stats = worker.get_operation_stats() ++ assert stats.data["save_exists"][0]["status"] == "ok" ++ assert "save_put" not in stats.data ++ ++ + def test_sending_thread_stores_hidden_tensor_asynchronously(): + pool_key = make_pool_key() + tensor = torch.zeros((2, 4), dtype=torch.float16) +@@ -300,13 +516,12 @@ def test_sending_thread_stores_hidden_tensor_asynchronously(): + worker = HiddenStoreWorker( + store_client=MooncakeHiddenStoreClient(store), + tensor_database=HiddenTensorDatabase(), +- producer_engine_id="encoder-1", + ) + sending_thread = HiddenStoreSendingThread(worker) + sending_thread.start() + + sending_thread.add_request( +- HiddenSaveRequest(pool_key=pool_key, tensor=tensor, now_ms=1234) ++ HiddenSaveRequest(pool_key=pool_key, tensor=tensor) + ) + sending_thread.request_queue.join() + +@@ -315,6 +530,30 @@ def test_sending_thread_stores_hidden_tensor_asynchronously(): + sending_thread.close() + + ++def test_sending_thread_records_failed_identifier_without_finishing(): ++ pool_key = make_pool_key() ++ tensor = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeBufferStore() ++ store.raise_on_batch_put = True ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ ) ++ sending_thread = HiddenStoreSendingThread(worker) ++ sending_thread.start() ++ ++ sending_thread.add_request( ++ HiddenSaveRequest(pool_key=pool_key, tensor=tensor) ++ ) ++ sending_thread.request_queue.join() ++ ++ assert sending_thread.get_and_clear_finished_identifiers() == set() ++ assert sending_thread.get_and_clear_failed_identifiers() == {pool_key.identifier} ++ assert pool_key.identifier in sending_thread.failure_reasons ++ assert worker.get_operation_stats().data["save_put"][0]["status"] == "error" ++ sending_thread.close() ++ ++ + def test_worker_load_gets_tensor_data_into_encoder_cache_before_returning(): + pool_key = make_pool_key() + stored = torch.zeros((2, 4), dtype=torch.float16) +@@ -324,7 +563,7 @@ def test_worker_load_gets_tensor_data_into_encoder_cache_before_returning(): + tensor_database=HiddenTensorDatabase(), + key_metadata=pool_key.key_metadata, + ) +- worker.save_tensor(pool_key, stored, now_ms=1234) ++ worker.save_tensor(pool_key, stored) + + encoder_cache = {} + worker.load( +@@ -342,6 +581,83 @@ def test_worker_load_gets_tensor_data_into_encoder_cache_before_returning(): + assert store.range_gets[-1][1] == [[make_hidden_data_key(pool_key)]] + assert store.range_gets[-1][3] == [[[TENSOR_METADATA_SIZE]]] + ++ stats = worker.get_operation_stats() ++ assert stats.data["load_get"][0]["status"] == "ok" ++ assert stats.data["load_get"][0]["num_keys"] == 1 ++ assert stats.data["load_get"][0]["num_bytes"] == ( ++ stored.numel() * stored.element_size() ++ ) ++ ++ ++def test_worker_load_records_error_without_writing_encoder_cache(): ++ pool_key = make_pool_key() ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ encoder_cache = {} ++ ++ try: ++ worker.load( ++ [MMMeta(identifier=pool_key.identifier, load_spec=LoadSpec(can_load=True))], ++ encoder_cache, ++ device="cpu", ++ ) ++ except HiddenStoreLoadError: ++ pass ++ else: ++ raise AssertionError("missing hidden tensor should fail fast") ++ ++ assert pool_key.identifier not in encoder_cache ++ stats = worker.get_operation_stats() ++ assert stats.data["load_get"][0]["status"] == "error" ++ assert stats.data["load_get"][0]["num_failed_keys"] == 1 ++ ++ ++def test_get_tensor_payload_unregisters_target_buffer_after_success(): ++ pool_key = make_pool_key() ++ stored = torch.zeros((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ worker = HiddenStoreWorker( ++ store_client=MooncakeHiddenStoreClient(store), ++ tensor_database=HiddenTensorDatabase(), ++ key_metadata=pool_key.key_metadata, ++ ) ++ worker.save_tensor(pool_key, stored) ++ target = torch.empty_like(stored) ++ ++ worker.store_client.get_tensor_payload( ++ pool_key, ++ target.data_ptr(), ++ target.numel() * target.element_size(), ++ TENSOR_METADATA_SIZE, ++ ) ++ ++ assert target.data_ptr() in store.unregistered ++ ++ ++def test_get_tensor_payload_unregisters_target_buffer_after_load_error(): ++ pool_key = make_pool_key() ++ target = torch.empty((2, 4), dtype=torch.float16) ++ store = FakeStore() ++ client = MooncakeHiddenStoreClient(store) ++ ++ try: ++ client.get_tensor_payload( ++ pool_key, ++ target.data_ptr(), ++ target.numel() * target.element_size(), ++ TENSOR_METADATA_SIZE, ++ ) ++ except HiddenStoreLoadError: ++ pass ++ else: ++ raise AssertionError("missing payload should raise") ++ ++ assert target.data_ptr() in store.unregistered ++ + + def _serialize_tensor_object(tensor: torch.Tensor) -> bytes: + tensor = tensor.detach().cpu().contiguous() +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py +index b14b22b7c..7213e1cbf 100644 +--- a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/connector.py +@@ -4,7 +4,6 @@ + + from __future__ import annotations + +-import time + from typing import TYPE_CHECKING + + import torch +@@ -20,7 +19,9 @@ from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorRole, + ) + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( +- HIDDEN_LAYOUT_VERSION, ++ HIDDEN_OBJECT_KIND, ++ HIDDEN_STORAGE_LAYOUT, ++ HIDDEN_TENSOR_LAYOUT, + HiddenKeyMetadata, + HiddenSaveRequest, + LoadSpec, +@@ -37,6 +38,7 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.worker impo + HiddenStoreWorker, + ) + from vllm.logger import init_logger ++from vllm.multimodal.utils import get_mm_features_in_window + from vllm.v1.core.sched.output import SchedulerOutput + + if TYPE_CHECKING: +@@ -65,6 +67,7 @@ class MooncakeStoreECConnector(ECConnectorBase): + self.soft_pin_video_hidden = bool( + extra_config.get("soft_pin_video_hidden", False) + ) ++ self.lookup_async = bool(extra_config.get("lookup_async", True)) + + if role == ECConnectorRole.SCHEDULER: + if self.is_consumer: +@@ -72,13 +75,10 @@ class MooncakeStoreECConnector(ECConnectorBase): + else: + if not (self.is_producer or self.is_consumer): + return +- engine_id = vllm_config.ec_transfer_config.engine_id +- + hidden_key_metadata = build_hidden_key_metadata(vllm_config) + self.store_client = store_client or create_mooncake_hidden_store_client() + self.worker = HiddenStoreWorker( + store_client=self.store_client, +- producer_engine_id=engine_id, + key_metadata=hidden_key_metadata, + ) + if self.is_producer: +@@ -87,8 +87,11 @@ class MooncakeStoreECConnector(ECConnectorBase): + self.lookup_server = HiddenLookupServer(self.worker, vllm_config) + + self.load_specs: dict[str, LoadSpec] = {} +- self._load_identifiers_to_schedule: set[str] = set() +- self._save_identifiers_to_schedule: set[str] = set() ++ self.lookup_result_cache: dict[str, bool] = {} ++ self.identifier_waiters: dict[str, set[str]] = {} ++ self._candidate_consumes: dict[str, set[str]] = {} ++ self._candidate_loads: dict[str, set[str]] = {} ++ self._candidate_saves: dict[str, set[str]] = {} + self._load_modalities: dict[str, str | None] = {} + self._save_modalities: dict[str, str | None] = {} + +@@ -103,81 +106,159 @@ class MooncakeStoreECConnector(ECConnectorBase): + def has_cache_item(self, identifier: str) -> bool: + if not self.is_consumer: + return False +- assert self.lookup_client is not None + +- started = time.perf_counter() +- if not self.lookup_client.lookup(identifier): ++ if not self.lookup_result_cache.get(identifier, False): + self.load_specs.pop(identifier, None) + logger.info( + "hidden_store_scheduler_miss identifier=%s " +- "reason=worker_lookup_miss hidden_store_lookup_ms=%.3f", ++ "reason=local_lookup_result_miss", + identifier, +- (time.perf_counter() - started) * 1000.0, + ) + return False + +- self.load_specs[identifier] = LoadSpec(can_load=False) ++ self.load_specs.setdefault(identifier, LoadSpec(can_load=False)) + logger.info( +- "hidden_store_scheduler_hit identifier=%s " "hidden_store_lookup_ms=%.3f", ++ "hidden_store_scheduler_hit identifier=%s", + identifier, +- (time.perf_counter() - started) * 1000.0, + ) + return True + ++ def ensure_cache_available( ++ self, ++ request: Request, ++ num_computed_tokens: int, ++ ) -> bool: ++ if not self.is_consumer: ++ return True ++ if not request.mm_features: ++ return True ++ assert self.lookup_client is not None ++ ++ start = num_computed_tokens ++ end = request.num_tokens ++ lo, hi = get_mm_features_in_window(request.mm_features, start, end) ++ identifiers = list( ++ dict.fromkeys( ++ feature.identifier for feature in request.mm_features[lo:hi] ++ ) ++ ) ++ if not identifiers: ++ return True ++ ++ request_id = request.request_id ++ for identifier in identifiers: ++ self.identifier_waiters.setdefault(identifier, set()).add(request_id) ++ ++ unknown_identifiers = [ ++ identifier ++ for identifier in identifiers ++ if identifier not in self.lookup_result_cache ++ ] ++ if not unknown_identifiers: ++ return True ++ ++ lookup_results = self.lookup_client.lookup_batch( ++ unknown_identifiers, ++ non_block=self.lookup_async, ++ ) ++ if lookup_results is None: ++ return False ++ ++ for identifier in unknown_identifiers: ++ self.lookup_result_cache[identifier] = lookup_results.get( ++ identifier, ++ False, ++ ) ++ return True ++ + def update_state_after_alloc(self, request: Request, index: int) -> None: + mm_feature = request.mm_features[index] + identifier = mm_feature.identifier + modality = mm_feature.modality ++ request_id = request.request_id ++ ++ self._candidate_consumes.setdefault(request_id, set()).add(identifier) + + if self.is_consumer and identifier in self.load_specs: +- load_spec = self.load_specs[identifier] +- load_spec.can_load = True ++ self._candidate_loads.setdefault(request_id, set()).add(identifier) + self._load_modalities[identifier] = modality +- self._load_identifiers_to_schedule.add(identifier) + + if self.is_producer: + self._save_modalities[identifier] = modality +- self._save_identifiers_to_schedule.add(identifier) ++ self._candidate_saves.setdefault(request_id, set()).add(identifier) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> ECConnectorMetadata: + items_by_identifier: dict[str, MMMeta] = {} ++ preempted_ids = getattr(scheduler_output, "preempted_req_ids", None) or set() + +- for identifier in self._load_identifiers_to_schedule: +- load_spec = self.load_specs.pop(identifier, None) +- if load_spec is None: ++ for request_id, identifiers in self._candidate_consumes.items(): ++ if request_id in preempted_ids: + continue +- items_by_identifier[identifier] = MMMeta( +- identifier=identifier, +- modality=self._load_modalities.get(identifier), +- load_spec=load_spec, +- ) ++ for identifier in identifiers: ++ waiters = self.identifier_waiters.get(identifier) ++ if waiters is not None: ++ waiters.discard(request_id) + +- for identifier in self._save_identifiers_to_schedule: +- item = items_by_identifier.get(identifier) +- if item is None: +- item = MMMeta( ++ for request_id, identifiers in self._candidate_loads.items(): ++ if request_id in preempted_ids: ++ continue ++ for identifier in identifiers: ++ load_spec = self.load_specs.pop(identifier, None) ++ if load_spec is None: ++ continue ++ load_spec.can_load = True ++ items_by_identifier[identifier] = MMMeta( + identifier=identifier, +- modality=self._save_modalities.get(identifier), ++ modality=self._load_modalities.get(identifier), ++ load_spec=load_spec, + ) +- items_by_identifier[identifier] = item +- item.can_save = True +- if item.modality is None: +- item.modality = self._save_modalities.get(identifier) ++ ++ for request_id, identifiers in self._candidate_saves.items(): ++ if request_id in preempted_ids: ++ continue ++ for identifier in identifiers: ++ item = items_by_identifier.get(identifier) ++ if item is None: ++ item = MMMeta( ++ identifier=identifier, ++ modality=self._save_modalities.get(identifier), ++ ) ++ items_by_identifier[identifier] = item ++ item.can_save = True ++ if item.modality is None: ++ item.modality = self._save_modalities.get(identifier) ++ ++ finished_req_ids = getattr(scheduler_output, "finished_req_ids", set()) ++ for finished_req_id in finished_req_ids: ++ for waiters in self.identifier_waiters.values(): ++ waiters.discard(finished_req_id) ++ ++ self._cleanup_lookup_results_without_waiters() + + metadata = MooncakeStoreConnectorMetadata( + items=list(items_by_identifier.values()), +- unfinished_identifiers=self._save_identifiers_to_schedule.copy(), + ) + +- self._load_identifiers_to_schedule.clear() +- self._save_identifiers_to_schedule.clear() ++ self._candidate_consumes.clear() ++ self._candidate_loads.clear() ++ self._candidate_saves.clear() + self._load_modalities.clear() + self._save_modalities.clear() + return metadata + ++ def _cleanup_lookup_results_without_waiters(self) -> None: ++ for identifier, waiters in list(self.identifier_waiters.items()): ++ if waiters: ++ continue ++ del self.identifier_waiters[identifier] ++ self.lookup_result_cache.pop(identifier, None) ++ self.load_specs.pop(identifier, None) ++ if self.lookup_client is not None: ++ self.lookup_client.discard(identifier) ++ + def start_load_caches( + self, + encoder_cache: dict[str, torch.Tensor], +@@ -220,7 +301,6 @@ class MooncakeStoreECConnector(ECConnectorBase): + HiddenSaveRequest( + pool_key=pool_key, + tensor=encoder_cache[identifier], +- now_ms=kwargs.get("now_ms"), + with_soft_pin=self._should_soft_pin(item), + ) + ) +@@ -248,11 +328,13 @@ class MooncakeStoreECConnector(ECConnectorBase): + def build_hidden_key_metadata(vllm_config: VllmConfig) -> HiddenKeyMetadata: + model_config = vllm_config.model_config + parallel_config = vllm_config.parallel_config ++ assert vllm_config.ec_transfer_config is not None ++ extra_config = vllm_config.ec_transfer_config.ec_connector_extra_config + + multimodal_config = getattr(model_config, "multimodal_config", None) + compute_hash = getattr(multimodal_config, "compute_hash", None) + mm_encoder_config_hash = ( +- compute_hash() if callable(compute_hash) else "mm_encoder:default" ++ compute_hash() if callable(compute_hash) else "encoder:default" + ) + + tp_size = get_tensor_model_parallel_world_size() +@@ -264,18 +346,25 @@ def build_hidden_key_metadata(vllm_config: VllmConfig) -> HiddenKeyMetadata: + "mm_encoder_tp_mode", + "unknown", + ) +- hidden_parallel_key = ( ++ parallel = ( + f"tp:{tp_size}" + f"@pp:{pp_size}" + f"@pcp:{pcp_size}" + f"@dcp:{dcp_size}" + f"@mm_tp:{mm_encoder_tp_mode}" +- "@storage:replicated" + ) + + return HiddenKeyMetadata( ++ cache_prefix=str( ++ extra_config.get( ++ "hidden_cache_prefix", ++ extra_config.get("cache_prefix", ""), ++ ) ++ ), ++ kind=HIDDEN_OBJECT_KIND, + model_name=model_config.model.rstrip("/").split("/")[-1], +- mm_encoder_config_hash=str(mm_encoder_config_hash), +- hidden_parallel_key=hidden_parallel_key, +- layout=HIDDEN_LAYOUT_VERSION, ++ encoder=str(mm_encoder_config_hash), ++ storage=HIDDEN_STORAGE_LAYOUT, ++ parallel=parallel, ++ tensor_layout=HIDDEN_TENSOR_LAYOUT, + ) +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py +index 2ba125943..f3fa09ef1 100644 +--- a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/data.py +@@ -14,7 +14,9 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.keys import + make_hidden_data_key, + ) + +-HIDDEN_LAYOUT_VERSION = "vllm-encoder-cache-tensor-v1" ++HIDDEN_OBJECT_KIND = "encoder_output" ++HIDDEN_STORAGE_LAYOUT = "replicated_object" ++HIDDEN_TENSOR_LAYOUT = "tensor" + HIDDEN_PROTOCOL_VERSION = "v1" + MOONCAKE_TENSOR_METADATA_NBYTES = 304 + +@@ -23,10 +25,13 @@ MOONCAKE_TENSOR_METADATA_NBYTES = 304 + class HiddenKeyMetadata: + """Metadata that defines the semantic namespace for hidden reuse.""" + ++ cache_prefix: str ++ kind: str + model_name: str +- mm_encoder_config_hash: str +- hidden_parallel_key: str +- layout: str ++ encoder: str ++ storage: str ++ parallel: str ++ tensor_layout: str + + + @dataclass(frozen=True, order=True) +@@ -38,13 +43,18 @@ class HiddenPoolKey: + + def to_string(self) -> str: + meta = self.key_metadata ++ prefix = ( ++ f"{escape_key_part(meta.cache_prefix)}@" if meta.cache_prefix else "" ++ ) + return ( +- "hidden" ++ f"{prefix}hidden" ++ f"@kind:{escape_key_part(meta.kind)}" + f"@model:{escape_key_part(meta.model_name)}" +- f"@mm_encoder:{escape_key_part(meta.mm_encoder_config_hash)}" +- f"@parallel:{escape_key_part(meta.hidden_parallel_key)}" +- f"@layout:{escape_key_part(meta.layout)}" +- f"@{escape_key_part(self.identifier)}" ++ f"@encoder:{escape_key_part(meta.encoder)}" ++ f"@storage:{escape_key_part(meta.storage)}" ++ f"@parallel:{escape_key_part(meta.parallel)}" ++ f"@tensor_layout:{escape_key_part(meta.tensor_layout)}" ++ f"@id:{escape_key_part(self.identifier)}" + ) + + +@@ -86,7 +96,6 @@ class HiddenSaveRequest: + + pool_key: HiddenPoolKey + tensor: torch.Tensor +- now_ms: int | None = None + with_soft_pin: bool = False + + @property +@@ -99,12 +108,41 @@ class MooncakeStoreConnectorMetadata(ECConnectorMetadata): + """Metadata passed from scheduler to worker for hidden store operations.""" + + items: list[MMMeta] = field(default_factory=list) +- unfinished_identifiers: set[str] = field(default_factory=set) + + def add_item(self, item: MMMeta) -> None: + self.items.append(item) + + ++@dataclass ++class HiddenStoreOperationStats: ++ """Minimal per-operation telemetry aligned with Mooncake KV store stats.""" ++ ++ data: dict[str, list[dict[str, int | float | str]]] = field(default_factory=dict) ++ ++ def is_empty(self) -> bool: ++ return not self.data ++ ++ def record_operation( ++ self, ++ operation: str, ++ duration_seconds: float, ++ num_keys: int, ++ *, ++ num_bytes: int = 0, ++ status: str = "ok", ++ num_failed_keys: int = 0, ++ ) -> None: ++ self.data.setdefault(operation, []).append( ++ { ++ "duration_seconds": duration_seconds, ++ "num_keys": num_keys, ++ "num_bytes": num_bytes, ++ "status": status, ++ "num_failed_keys": num_failed_keys, ++ } ++ ) ++ ++ + class HiddenTensorDatabase: + """Maps hidden tensors to store keys and GPU memory descriptors.""" + +@@ -131,7 +169,7 @@ def build_tensor_meta( + return TensorMeta( + pool_key=pool_key, + protocol_version=HIDDEN_PROTOCOL_VERSION, +- layout=HIDDEN_LAYOUT_VERSION, ++ layout=HIDDEN_TENSOR_LAYOUT, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype), + nbytes=tensor.numel() * tensor.element_size(), +@@ -160,5 +198,5 @@ def validate_loaded_tensor(tensor: torch.Tensor, meta: TensorMeta) -> None: + f"actual={actual_nbytes} expected={meta.nbytes}" + ) + +- if meta.layout != HIDDEN_LAYOUT_VERSION: ++ if meta.layout != HIDDEN_TENSOR_LAYOUT: + raise ValueError(f"Unsupported hidden tensor layout: {meta.layout}") +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py +index f71159aad..0852ee1db 100644 +--- a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/store_client.py +@@ -14,8 +14,8 @@ from dataclasses import dataclass + from typing import Any + + from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import ( +- HIDDEN_LAYOUT_VERSION, + HIDDEN_PROTOCOL_VERSION, ++ HIDDEN_TENSOR_LAYOUT, + MOONCAKE_TENSOR_METADATA_NBYTES, + HiddenPoolKey, + TensorMeta, +@@ -69,6 +69,11 @@ _TORCH_DTYPE_TO_MOONCAKE_DTYPE = { + "torch.float8_e4m3fn": 13, + "torch.float8_e5m2": 14, + } ++_SUPPORTED_HIDDEN_TORCH_DTYPES = { ++ "torch.float16", ++ "torch.bfloat16", ++ "torch.float32", ++} + + + @dataclass +@@ -234,6 +239,7 @@ class MooncakeHiddenStoreClient: + *, + with_soft_pin: bool = False, + ) -> None: ++ _validate_supported_hidden_tensor_dtype(tensor) + key = make_hidden_data_key(pool_key) + replicate_config = _make_hidden_replicate_config( + self.replicate_config, +@@ -288,24 +294,30 @@ class MooncakeHiddenStoreClient: + metadata = _encode_mooncake_tensor_metadata(tensor) + metadata_buffer = (ctypes.c_ubyte * len(metadata)).from_buffer_copy(metadata) + metadata_ptr = ctypes.addressof(metadata_buffer) +- self.register_tensor(tensor.data_ptr(), data_size) +- self.register_tensor(metadata_ptr, len(metadata)) +- +- key = make_hidden_data_key(pool_key) +- results = self.store.batch_put_from_multi_buffers( +- [key], +- [[metadata_ptr, tensor.data_ptr()]], +- [[len(metadata), data_size]], +- replicate_config, +- ) +- unregister_fn = getattr(self.store, "unregister_buffer", None) +- if unregister_fn is not None: +- unregister_fn(metadata_ptr) +- failed = [result for result in results if result < 0] +- if failed: +- raise HiddenStoreSaveError( +- "failed to put hidden tensor for " f"{pool_key.to_string()}: {failed}" ++ payload_ptr = tensor.data_ptr() ++ registered_addrs: list[int] = [] ++ try: ++ self.register_tensor(payload_ptr, data_size) ++ registered_addrs.append(payload_ptr) ++ self.register_tensor(metadata_ptr, len(metadata)) ++ registered_addrs.append(metadata_ptr) ++ ++ key = make_hidden_data_key(pool_key) ++ results = self.store.batch_put_from_multi_buffers( ++ [key], ++ [[metadata_ptr, payload_ptr]], ++ [[len(metadata), data_size]], ++ replicate_config, + ) ++ failed = [result for result in results if result < 0] ++ if failed: ++ raise HiddenStoreSaveError( ++ "failed to put hidden tensor for " ++ f"{pool_key.to_string()}: {failed}" ++ ) ++ finally: ++ for addr in reversed(registered_addrs): ++ self.unregister_tensor(addr) + + def register_tensor(self, addr: int, size: int) -> None: + ret = self.store.register_buffer(addr, size) +@@ -314,6 +326,26 @@ class MooncakeHiddenStoreClient: + f"failed to register hidden buffer addr={addr:#x} size={size}: {ret}" + ) + ++ def unregister_tensor(self, addr: int) -> None: ++ unregister_fn = getattr(self.store, "unregister_buffer", None) ++ if unregister_fn is None: ++ return ++ try: ++ ret = unregister_fn(addr) ++ except Exception: ++ logger.warning( ++ "failed to unregister hidden buffer addr=%#x", ++ addr, ++ exc_info=True, ++ ) ++ return ++ if ret != 0: ++ logger.warning( ++ "unregister hidden buffer failed addr=%#x ret=%s", ++ addr, ++ ret, ++ ) ++ + def get_tensor_payload( + self, + pool_key: HiddenPoolKey, +@@ -322,21 +354,24 @@ class MooncakeHiddenStoreClient: + src_offset: int, + ) -> int: + self.register_tensor(addr, size) +- key = make_hidden_data_key(pool_key) +- results = self.store.get_into_ranges( +- [addr], +- [[key]], +- [[[0]]], +- [[[src_offset]]], +- [[[size]]], +- ) +- result = _single_range_result(results) +- if result != size: +- raise HiddenStoreLoadError( +- "failed to get hidden tensor payload for " +- f"{pool_key.to_string()}: {result}" ++ try: ++ key = make_hidden_data_key(pool_key) ++ results = self.store.get_into_ranges( ++ [addr], ++ [[key]], ++ [[[0]]], ++ [[[src_offset]]], ++ [[[size]]], + ) +- return result ++ result = _single_range_result(results) ++ if result != size: ++ raise HiddenStoreLoadError( ++ "failed to get hidden tensor payload for " ++ f"{pool_key.to_string()}: {result}" ++ ) ++ return result ++ finally: ++ self.unregister_tensor(addr) + + def _read_range( + self, +@@ -358,9 +393,7 @@ class MooncakeHiddenStoreClient: + [[[size]]], + ) + finally: +- unregister_fn = getattr(self.store, "unregister_buffer", None) +- if unregister_fn is not None: +- unregister_fn(buffer_ptr) ++ self.unregister_tensor(buffer_ptr) + if _single_range_result(results) != size: + return None + return bytes(buffer) +@@ -421,7 +454,7 @@ def _decode_mooncake_tensor_metadata( + return TensorMeta( + pool_key=pool_key, + protocol_version=HIDDEN_PROTOCOL_VERSION, +- layout=HIDDEN_LAYOUT_VERSION, ++ layout=HIDDEN_TENSOR_LAYOUT, + shape=shape, + dtype=_MOONCAKE_DTYPE_TO_TORCH_DTYPE[dtype], + nbytes=int(data_bytes), +@@ -432,8 +465,7 @@ def _decode_mooncake_tensor_metadata( + + def _encode_mooncake_tensor_metadata(tensor: Any) -> bytes: + dtype = str(tensor.dtype) +- if dtype not in _TORCH_DTYPE_TO_MOONCAKE_DTYPE: +- raise HiddenStoreSaveError(f"unsupported hidden tensor dtype: {dtype}") ++ _validate_supported_hidden_tensor_dtype(tensor) + shape = tuple(int(dim) for dim in tensor.shape) + if len(shape) > 8: + raise HiddenStoreSaveError( +@@ -463,6 +495,12 @@ def _encode_mooncake_tensor_metadata(tensor: Any) -> bytes: + return metadata + + ++def _validate_supported_hidden_tensor_dtype(tensor: Any) -> None: ++ dtype = str(tensor.dtype) ++ if dtype not in _SUPPORTED_HIDDEN_TORCH_DTYPES: ++ raise HiddenStoreSaveError(f"unsupported hidden tensor dtype: {dtype}") ++ ++ + def _make_hidden_replicate_config( + replicate_config: Any | None, + *, +diff --git a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py +index 7127ead58..16829f34c 100644 +--- a/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py ++++ b/vllm/distributed/ec_transfer/ec_connector/mooncake_store_hidden/worker.py +@@ -9,6 +9,7 @@ import queue + import socket + import threading + import time ++from concurrent.futures import Future, ThreadPoolExecutor + + import torch + import zmq +@@ -19,6 +20,7 @@ from vllm.distributed.ec_transfer.ec_connector.mooncake_store_hidden.data import + HiddenKeyMetadata, + HiddenPoolKey, + HiddenSaveRequest, ++ HiddenStoreOperationStats, + HiddenTensorDatabase, + MMMeta, + build_tensor_meta, +@@ -37,6 +39,8 @@ from vllm.utils.network_utils import make_zmq_socket + logger = init_logger(__name__) + + LOOKUP_MSG = b"LOOKUP" ++BATCH_LOOKUP_MSG = b"BATCH_LOOKUP" ++RESP_BATCH = b"BATCH" + RESP_HIT = b"HIT" + RESP_MISS = b"MISS" + RESP_ERR = b"ERR" +@@ -49,14 +53,14 @@ class HiddenStoreWorker: + self, + store_client: MooncakeHiddenStoreClient, + tensor_database: HiddenTensorDatabase | None = None, +- producer_engine_id: str | None = None, + key_metadata: HiddenKeyMetadata | None = None, + ): + self.store_client = store_client + self.tensor_database = tensor_database or HiddenTensorDatabase() +- self.producer_engine_id = producer_engine_id + self.key_metadata = key_metadata + self.sending_thread: HiddenStoreSendingThread | None = None ++ self._operation_stats_lock = threading.Lock() ++ self._operation_stats = HiddenStoreOperationStats() + + def make_pool_key(self, identifier: str) -> HiddenPoolKey: + assert self.key_metadata is not None +@@ -76,7 +80,6 @@ class HiddenStoreWorker: + self.save_tensor( + request.pool_key, + request.tensor, +- now_ms=request.now_ms, + with_soft_pin=request.with_soft_pin, + ) + return +@@ -87,6 +90,34 @@ class HiddenStoreWorker: + return set() + return self.sending_thread.get_and_clear_finished_identifiers() + ++ def get_operation_stats(self) -> HiddenStoreOperationStats | None: ++ with self._operation_stats_lock: ++ if self._operation_stats.is_empty(): ++ return None ++ stats = self._operation_stats ++ self._operation_stats = HiddenStoreOperationStats() ++ return stats ++ ++ def _record_operation( ++ self, ++ operation: str, ++ duration_seconds: float, ++ num_keys: int, ++ *, ++ num_bytes: int = 0, ++ status: str = "ok", ++ num_failed_keys: int = 0, ++ ) -> None: ++ with self._operation_stats_lock: ++ self._operation_stats.record_operation( ++ operation=operation, ++ duration_seconds=duration_seconds, ++ num_keys=num_keys, ++ num_bytes=num_bytes, ++ status=status, ++ num_failed_keys=num_failed_keys, ++ ) ++ + def shutdown(self) -> None: + if self.sending_thread is not None: + self.sending_thread.close() +@@ -94,32 +125,76 @@ class HiddenStoreWorker: + + def lookup(self, identifier: str) -> bool: + """Return whether the hidden object exists in Mooncake Store.""" +- pool_key = self.make_pool_key(identifier) +- if not self.store_client.exists(pool_key): +- logger.info( +- "hidden_store_lookup_miss identifier=%s hidden_pool_key=%s " +- "reason=missing_object", +- pool_key.identifier, +- pool_key.to_string(), +- ) +- return False ++ return self.lookup_batch([identifier]).get(identifier, False) + +- logger.info( +- "hidden_store_lookup_hit identifier=%s hidden_pool_key=%s", +- pool_key.identifier, +- pool_key.to_string(), ++ def lookup_batch(self, identifiers: list[str]) -> dict[str, bool]: ++ """Return whether hidden objects exist in Mooncake Store.""" ++ pool_keys = [self.make_pool_key(identifier) for identifier in identifiers] ++ started = time.perf_counter() ++ try: ++ exists = self.store_client.batch_exists(pool_keys) ++ except Exception: ++ self._record_operation( ++ "lookup_exists", ++ time.perf_counter() - started, ++ len(pool_keys), ++ status="error", ++ num_failed_keys=len(pool_keys), ++ ) ++ raise ++ ++ failed_keys = sum(1 for hit in exists if not hit) ++ self._record_operation( ++ "lookup_exists", ++ time.perf_counter() - started, ++ len(pool_keys), ++ status="miss" if failed_keys else "ok", ++ num_failed_keys=failed_keys, + ) +- return True ++ results = dict(zip(identifiers, exists, strict=True)) ++ for pool_key, hit in zip(pool_keys, exists, strict=True): ++ if hit: ++ logger.info( ++ "hidden_store_lookup_hit identifier=%s hidden_pool_key=%s", ++ pool_key.identifier, ++ pool_key.to_string(), ++ ) ++ else: ++ logger.info( ++ "hidden_store_lookup_miss identifier=%s hidden_pool_key=%s " ++ "reason=missing_object", ++ pool_key.identifier, ++ pool_key.to_string(), ++ ) ++ return results + + def save_tensor( + self, + pool_key: HiddenPoolKey, + tensor: torch.Tensor, +- now_ms: int | None = None, + with_soft_pin: bool = False, + ) -> None: +- if self.store_client.exists(pool_key): +- logger.debug( ++ exists_started = time.perf_counter() ++ try: ++ exists = self.store_client.exists(pool_key) ++ except Exception: ++ self._record_operation( ++ "save_exists", ++ time.perf_counter() - exists_started, ++ 1, ++ status="error", ++ num_failed_keys=1, ++ ) ++ raise ++ ++ self._record_operation( ++ "save_exists", ++ time.perf_counter() - exists_started, ++ 1, ++ status="ok" if exists else "miss", ++ ) ++ if exists: ++ logger.info( + "hidden_store_save_skip identifier=%s hidden_pool_key=%s " + "reason=exists", + pool_key.identifier, +@@ -131,10 +206,28 @@ class HiddenStoreWorker: + stored_tensor = tensor if tensor.is_contiguous() else tensor.contiguous() + used_staging = stored_tensor is not tensor + tensor_meta = build_tensor_meta(pool_key, stored_tensor) +- self.store_client.put_tensor( +- pool_key, +- stored_tensor, +- with_soft_pin=with_soft_pin, ++ try: ++ self.store_client.put_tensor( ++ pool_key, ++ stored_tensor, ++ with_soft_pin=with_soft_pin, ++ ) ++ except Exception: ++ self._record_operation( ++ "save_put", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes, ++ status="error", ++ num_failed_keys=1, ++ ) ++ raise ++ self._record_operation( ++ "save_put", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes, ++ status="ok", + ) + logger.info( + "hidden_store_put identifier=%s hidden_pool_key=%s nbytes=%d " +@@ -167,33 +260,52 @@ class HiddenStoreWorker: + + started = time.perf_counter() + pool_key = self.make_pool_key(item.identifier) +- tensor_meta = self.store_client.get_tensor_meta(pool_key) +- if tensor_meta is None: +- raise HiddenStoreLoadError( +- "failed to load hidden tensor metadata for " +- f"{pool_key.to_string()}" +- ) ++ tensor_meta = None ++ try: ++ tensor_meta = self.store_client.get_tensor_meta(pool_key) ++ if tensor_meta is None: ++ raise HiddenStoreLoadError( ++ "failed to load hidden tensor metadata for " ++ f"{pool_key.to_string()}" ++ ) + +- target_device = device +- if target_device is None: +- target_device = "cuda" if torch.cuda.is_available() else None +- target = torch.empty( +- tensor_meta.shape, +- dtype=_resolve_torch_dtype(tensor_meta.dtype), +- device=target_device, +- ) +- _data_key, addrs, sizes = self.tensor_database.prepare_value( +- pool_key, +- target, +- ) +- self.store_client.get_tensor_payload( +- pool_key, +- addrs[0], +- sizes[0], +- tensor_meta.data_offset, +- ) +- validate_loaded_tensor(target, tensor_meta) ++ target_device = device ++ if target_device is None: ++ target_device = "cuda" if torch.cuda.is_available() else None ++ target = torch.empty( ++ tensor_meta.shape, ++ dtype=_resolve_torch_dtype(tensor_meta.dtype), ++ device=target_device, ++ ) ++ _data_key, addrs, sizes = self.tensor_database.prepare_value( ++ pool_key, ++ target, ++ ) ++ self.store_client.get_tensor_payload( ++ pool_key, ++ addrs[0], ++ sizes[0], ++ tensor_meta.data_offset, ++ ) ++ validate_loaded_tensor(target, tensor_meta) ++ except Exception: ++ self._record_operation( ++ "load_get", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes if tensor_meta is not None else 0, ++ status="error", ++ num_failed_keys=1, ++ ) ++ raise + encoder_cache[item.identifier] = target ++ self._record_operation( ++ "load_get", ++ time.perf_counter() - started, ++ 1, ++ num_bytes=tensor_meta.nbytes, ++ status="ok", ++ ) + logger.info( + "hidden_store_get identifier=%s hidden_pool_key=%s nbytes=%d " + "hidden_store_get_ms=%.3f", +@@ -223,6 +335,8 @@ class HiddenStoreSendingThread(threading.Thread): + self.request_queue: queue.Queue[HiddenSaveRequest | None] = queue.Queue() + self.done_task_lock = threading.Lock() + self.finished_identifiers: set[str] = set() ++ self.failed_identifiers: set[str] = set() ++ self.failure_reasons: dict[str, str] = {} + self._closed = threading.Event() + + def add_request(self, request: HiddenSaveRequest) -> None: +@@ -234,10 +348,21 @@ class HiddenStoreSendingThread(threading.Thread): + self.finished_identifiers.clear() + return finished + ++ def get_and_clear_failed_identifiers(self) -> set[str]: ++ with self.done_task_lock: ++ failed = self.failed_identifiers.copy() ++ self.failed_identifiers.clear() ++ return failed ++ + def set_finished_identifier(self, identifier: str) -> None: + with self.done_task_lock: + self.finished_identifiers.add(identifier) + ++ def set_failed_identifier(self, identifier: str, error: Exception) -> None: ++ with self.done_task_lock: ++ self.failed_identifiers.add(identifier) ++ self.failure_reasons[identifier] = str(error) ++ + def run(self) -> None: + while True: + request = self.request_queue.get() +@@ -247,11 +372,12 @@ class HiddenStoreSendingThread(threading.Thread): + self.store_worker.save_tensor( + request.pool_key, + request.tensor, +- now_ms=request.now_ms, + with_soft_pin=request.with_soft_pin, + ) + self.set_finished_identifier(request.identifier) + except Exception as e: ++ if request is not None: ++ self.set_failed_identifier(request.identifier, e) + logger.error("Error in %s: %s", self.name, e) + finally: + self.request_queue.task_done() +@@ -302,6 +428,20 @@ class HiddenLookupServer: + except Exception: + logger.exception("HiddenLookupServer lookup failed") + self.socket.send_multipart([RESP_ERR]) ++ elif msg_type == BATCH_LOOKUP_MSG: ++ try: ++ identifiers = [ ++ bytes(frame).decode("utf-8") for frame in all_frames[1:] ++ ] ++ exists = self.store_worker.lookup_batch(identifiers) ++ frames = [ ++ RESP_HIT if exists.get(identifier, False) else RESP_MISS ++ for identifier in identifiers ++ ] ++ self.socket.send_multipart([RESP_BATCH, *frames]) ++ except Exception: ++ logger.exception("HiddenLookupServer batch lookup failed") ++ self.socket.send_multipart([RESP_ERR]) + else: + logger.warning( + "HiddenLookupServer received unknown msg_type: %r", +@@ -331,19 +471,87 @@ class HiddenLookupClient: + zmq.REQ, # type: ignore[attr-defined] + bind=False, + ) ++ self.executor = ThreadPoolExecutor( ++ max_workers=1, ++ thread_name_prefix="HiddenLookupClient", ++ ) ++ self.futures: dict[str, Future[dict[str, bool]]] = {} + + def lookup(self, identifier: str) -> bool: +- self.socket.send_multipart([LOOKUP_MSG, identifier.encode("utf-8")]) ++ result = self.lookup_batch([identifier], non_block=False) ++ assert result is not None ++ return result.get(identifier, False) ++ ++ def _lookup_batch(self, identifiers: list[str]) -> dict[str, bool]: ++ self.socket.send_multipart( ++ [ ++ BATCH_LOOKUP_MSG, ++ *(identifier.encode("utf-8") for identifier in identifiers), ++ ] ++ ) + resp = self.socket.recv_multipart() + msg_type = bytes(resp[0]) +- if msg_type == RESP_HIT: +- return True +- if msg_type in (RESP_MISS, RESP_ERR): +- return False ++ if msg_type == RESP_BATCH: ++ states = [bytes(frame) == RESP_HIT for frame in resp[1:]] ++ if len(states) != len(identifiers): ++ logger.warning( ++ "HiddenLookupClient received malformed batch response: " ++ "identifiers=%d states=%d", ++ len(identifiers), ++ len(states), ++ ) ++ return {identifier: False for identifier in identifiers} ++ return dict(zip(identifiers, states, strict=True)) ++ if msg_type == RESP_ERR: ++ return {identifier: False for identifier in identifiers} + logger.warning("HiddenLookupClient received unknown response: %r", msg_type) +- return False ++ return {identifier: False for identifier in identifiers} ++ ++ def lookup_batch( ++ self, ++ identifiers: list[str], ++ non_block: bool = False, ++ ) -> dict[str, bool] | None: ++ identifiers = list(dict.fromkeys(identifiers)) ++ if not identifiers: ++ return {} ++ ++ new_identifiers = [ ++ identifier for identifier in identifiers if identifier not in self.futures ++ ] ++ if new_identifiers: ++ future = self.executor.submit(self._lookup_batch, new_identifiers) ++ for identifier in new_identifiers: ++ self.futures[identifier] = future ++ ++ if non_block and any( ++ not self.futures[identifier].done() for identifier in identifiers ++ ): ++ return None ++ ++ results: dict[str, bool] = {} ++ for identifier in identifiers: ++ future = self.futures[identifier] ++ try: ++ batch_results = future.result() ++ results[identifier] = batch_results.get(identifier, False) ++ except Exception as e: ++ logger.error("Async hidden lookup failed for %s: %s", identifier, e) ++ results[identifier] = False ++ finally: ++ self.futures.pop(identifier, None) ++ return results ++ ++ def discard(self, identifier: str) -> None: ++ future = self.futures.pop(identifier, None) ++ if future is None: ++ return ++ if not any(existing is future for existing in self.futures.values()): ++ future.cancel() + + def close(self): ++ self.executor.shutdown(wait=False, cancel_futures=True) ++ self.futures.clear() + self.socket.close(linger=0) + + diff --git a/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-review-fixes.patch b/ccf-vllm-epd-archive/patches/vllm-epd-hidden-ec-connector-review-fixes.patch new file mode 100644 index 0000000000000000000000000000000000000000..9326cca8eee842915072edfc38effdb3ccccf015 GIT binary patch literal 33054 zcmeHQYj0e~k-eV-?0*;#2r!DRktv$CWPpprbiCe;WhEXZ0wM?lN*qdzB~l@&wH2-Z z`tF|7HHTH*x9?*p%9>yo1M)uRc6W8XtGaH_|NQr1bKD%^Gih$(^9)xPI6r7^;kazh z@&6p>|JiKg|FU`7Ebw>QoZ!Q&3!Ikv4*1LzKLyxq+3xzl{!>^66Ceb9W0 zC$^hU@c9hqpEhhw{NCIk+!Cgb(r2@<{M+WYz|u`%ZM*px_?rMmJKqMTC-|1%KW$zB z&*%6|ny*(J_AqfHdF`zEBH?_3Z|y@)BbSrk{tkX!ra$uJ0l0hFOmVhEpGO#Tl6cgf zq6FM%^`>6cW>&N~2R|HugI}fb4$(WgKK4lLPZEqDtso>%8H~?xyns|3C5}4=WY2Nu zFi8;Mr_7!}Rvl(R=zR%7j(dT*o}tAg&G7>5ZmdQ&9?-Mx@eeT4i)$MDZS!*&)#mPV zez~GyO&3g?-z2J>gT6-zGp`Z`7N8up$P)K+jq(NN`UH1R0sW-e!ykF%MVh~GrCF1g zN{wc1?xh}FgUvCv>95~6`}mZtW`JS=JTIFc)7i9Hw+Ew(Bwz>Hgfg%L_&>t`yU;Kn zLk4(l&xm5u)y~hv;{pE2Ip=`p5YXD$3C_q}ZEPMxXCHvCO8rcJkvtjS3U?-fls8%p za>@y~tBubYFhb0kY__H423qeWI&Bwnn~z|y*C;kL!XxP7$B@b&!7&%WN_+k$3-Ny& zyWNuc0`q&8FirjMQj^;`X`AzeC29z3C;YnnmikRSMjbb;YQA<~qL${-HgALaWz)hE#o;v@Gq7ja-ATi=(CuIl*k-P3SW6 zKpnnZ&fAdW^q02i=QCk+sPA%?w78SzJ8&&^d5*`tudTw5rMDmtSF%n!?wq`Zo@u?Q z!(9GnXt!*B!vFr!qD)<1ufzO44D)8h5xJH5FcCx>2GTUf4o@KoN3h6u;oUw-vBNHo zhw1JT|2HRgsL$m?MGbXm+*|qt-jT%%JFxa2!&{P{bOarA3VSqxE_;E`OJJQGOL;xR z@k!G1^s#42<1Nzt1IXz)C`JE9_iNg3e55-|BT7Ti#vh~ilaGDYv+qg&07Vfz` zfL^9nKLA8F=1}{GX!r$fZX+3(Bt!r?sOQBP673J?@6}fu6A?|5& zsyMPOA>v&gfqPN}8*_x}4K4(ID2HQ8LmpbioXk0q3O253n8eBO& z>riQ44$m0gQLC||+SoNc8m;vE;ex+``joXAtY0UNAeD3M)aXyohgn9;1u&mDMq)L( zE#^p}UE*10X8L7qD-`|j2PjJ%nV z)aplWdxV#~KlWB%=ZyH711F5KzsH}iPENtyxqR)l?wBQW>{1r|>PCzIbLdKmk2=sr74tF1 ztVJj{8{v2~a`eZAjOC+W4RgdgwXv`FYHoRdT-ByO z`YhK>rB;j@je5*Iq}<=CX)%A+cvLLQ zJKY)blG9@JV4U`@5k$n4FdavU*d3<{9J{t+-<>^=W7~P#lc5x9mCSt5v*e94hIbv6 z`!Uf|W*&wSdwiNSwKYom`eV()+Lm*hpkA0YJkOK2ZT|Km={sSXemHyDDtzCv ze^%XaZJ2nZ$0%%WL2evdTdhckk|jDCjM2J&i?)Jf9-eI14vuM6%1|7eaTM#axt{0M zJN1k}-Y(A1W)Sb>RLNViB3%4n=Y3DNIU(NZYst)4jTW4XDaE+!lY+1F_7TG#O;cz1`bZztMpaf%^sTQ@ z_cNPIZZJ*RUO(v@MYt{L$SRDVKwZ+=GK1Dj{6C~;qXmeXlN#8pHFc<>(nT7PIatQ3 zoEM*S`Hwm{u4<8o%BV87Gr5vSHV-{I_lX0dRH?;~mL;Ad)Ha?nm#j~p1CO)|T>(-Hf43@bI^o%wL06c z0|m+RWhN}|zb!lQP9J1ZdhFbu?r*bRgjiy#?fK_&l8>Eh=zQL7yw-DTc^?rsznpE6 zzP~Qp;B@kpJo;hO{$)mw^kQ~nk?JaF9n5Dt!I^1otrS>|6LmRtJ2eWU8(TT^cP6-| z!nbiOim{A|sR?yuIeLi0Vif>^ZH1=1xU>r>mzmF zVh?I3as#n<3%Z>?FEKszCg80^jr{Gx_$J_{LPt&|SCEhB|9*}&!)?5$!)o5A(Ae~s z?&9w$yt7ZzYH;7%ftx(bSo^!laB59{2Ndx42z5o;Xhoj-*0krOwPrcLl*alSU(=?3 z|I5_l^3n%%`PLcN>&;WK8pxP8(t4lP`6fw@eEpC9jkU6z<10uRb=Yh07xkF=of@fq zq|YQ>=9Mh;do&{Xx!rEeTNzz*nyvJ5UX4(HtFiP|1W8=svU~|8E=ISGnTXFg_Lx;a zji;ko0aAtXomX2BYnHR1o#A^OgoicU-6JsrK|A&m+P)SiXj~m>+n%lJ1?j7uC3+8g zwrXEyL)^v?i?#xs*JRY^y4U;~&*$^Cd-R8MpGISs$2tzjwT*Zl8qavB9Nl1d-XG>@ zS}>Q3wv6PGLjIs!i_7VcSPqrGiB&4l|1j$3bzwEkSTsXh%a!Gu&A;ORz4WmNXVUyv zfX?K8b8UO%yWGq5H+4R=P3eb`)Q)zZ6<)ir^-FxVlbxsCzYQDz8Tf#lu*r6Qc&?jd z?T4Y2%_sg@vt@5;=;}7^jj{PLnx0?U|NRhwciC6BJ`Srz!hW6W3+X3sNgc?n;nuA$s3J)tW_NH*qANQ%mxl zpfYPhJ7^lTT+4}(4y}zcQPh^Su$=ob+A@Bdr7_EVy&%#Ere|%fiun=CXvrQ=(TC66 z@;6%vzi|ym_bv9xV;OyCHiOyH<5VZe6@V^m2}U28mVy2Oqj;-D@;sHO!R&#eKw6yo z%%Uyo@|ZNfd7Eu0=s&u38=8fBGrW*KC-6DD>ZkQIt=F4D4C(0-#&$MTPwL2si+=Ab z;A{ea+^>v(#?C7DvY3VBo*Cv{4&@L9i;i2MvipCm`91|~KDSA>KGlvpa_LY#?n@7RlG)NPfpA`+@cOZIbU-hfw|=ci57*XXpOr_xG>}^X89`jxU?vVU=hf zy1=CBN%CWK&%BEXDxvlKH_gA}d4JZf+t1y@b8~$BN!@8b`+M~NeR$HY<=H#{Y#jf0 z%~yc;KbrgK`8%9F!0{gL-*5hm^99CY*6BXJ`|5JsU2vVP{Koe(@;7Rt7Mv%WL~Z*J z+EA5Sru&4~yXaAPwVjBVORM?NV|#XOTAoF{ZyL|{sVbi>qxh&F={+hZl-K!s7ARKD z*c-5l*QrOf*H$ea_cdp|k1g0_m&VnX3u5=X2{SY1j{pHZCR#E(d)(XIX&97+?sd(oaRMK)eXxKCCqcOwOw0zVSd(P0rv;%eu40R3O|6EGLIve>oUtq zDhg7kK%Iy0xB8*vuvcQUoqZ8r{OI1B#^efaKYUiP#EPsgZ$$liZnjvDVfpg4d5FqdoEx1x#w>WX%BtFgO0?*5ptj_LLsy5{0|tjj5*E!DxfG;$_2 zy<5wJ4*BbrsrVjRK=A}?OQ_vh14Z4yxQu6IGTeasyrJx23JG~@)WL0w&XA!H_q81{_%7#UL!nG@T^~W`)!wor9?jOXa^VD7HtM3E@pM0Tt6v?(lXcy#r*!P&FRQF%&7j5g zVG^cfU0+QF9zV4na&S^%^rd05!;CTp=;OxB3=+ zD=&p>ljaR-ZzIx@I@xvy^4ulo@7nfG(ze@(3%PSD<8t~ljMCpv^ipD`)RdMTE1&AV zj(8l~w#&|!qx5kqYwOxt`NW4J56JEAFqQCD0N)lFi4i8b!XYm;uzxDDm~(i`Cf|a zA$orCTetEyAMTu9|5#4?>iPen_S_AhvEFhp7886f5BdHbwb*Q3p5MBZPE|T;EkL6W zAIJHEtT(V61gS*YXfHk2v$i9(`2no0qF1{e_Fa6xo`89B^zK;!z6A(X8M8i4xRS*k zp5nTVa`)_=-upNkyN>+Ds-AM)YW+2*=SgDCyEYF<9;GyTKEHabr7EAiw^G;YDYFsX zabd}gtgNTp)Bjc<%(-zdT)z|dSh1R;6U1b15*LyhcF*Z;cSQB^(DX+oM)C2zl{lIK z-kLAEZof`F$KASp)O9QKbi%#+WA#|eJi{_CQ1ivE3pVG1&4VrShz0wEO*It#3(Xn|84B z=Bq+Gut~NT*E&(^9CPt}`U6M`^E>or=)3rpMVT|K&G+2%@*?7&u^XOMu$*ASX`FV`-ow`W;jBVvw|>nq*(S6*w#S8I=GaZDYrH5qNs zsxG7ARFlwLD@slY4c~lfS1s4zxms3TnCV(~x|*blg6(6(FnO6?M*j+|cCo2b*C^^% zfWOM6Oy#GmbuG&a*)zit${bjKe^wxJLr!(!$}wb_VXNKbPW8qQ-PXWca+I#UbnF*S zrkC%2bjgIWb#CKi_n5U~Ssv%TYh9s5KD4oXKTzJdtx<8cO7dxT;n(02{?>$92&>!V zKBDVKN{$Ov)~@TZUX4;?M=^?yt9|(v6w%t_99*m28}d~S{^lG#NP3L?1-jj&NxI$~ z`zBGuV(UgM2p->!UlXroL?L zf;NZo0jo%v3oEM?A8-AmtR7r9w#v+K}bf?OJlK2-aSk}fUP z#gm4~uM__G%#b(4Y5?t7l~?d^<$~+Jlq>s11oNr|w~h`;^jLm5Z=>`wawx9ibL3KX zHD6<_97brlKf^EOk-Q&2qRR$FpL>DTK5A^MoUytY+hNT)$@ajIq>IyZG>TJJ7^8BP zkGqYV_r~0AtW0E;V7dR%uwE;R@Dv1hf%aON!f zJ^N{ViwK4_0WlEwpRv4jzB}-{?F-aJ5ywrw(@HIt7Zv^Jt>0Z6y&I=F6@9N1<@Ln1 zXW+TGgI@ldx%n=6Q&?EqG5ggMX5#GqL*~K$vifEp8d&!||3^O79t6%&pQjbH9kl)k zRT%b5TCpBOw&cUZ?9|m6D~(JNzrb!Lw@`JY-R8NI<{7Z@Gw#yo_V<6Ni4|kehZp9# zBF7l{ZR_68xXZT}i9dd?bH!WT6- zJmdu?pmXWXA z1-<&vG4Evx>W}j_)PC59<{Aq3%Nxdpo?HhDYID{3tG)70PUq;=uk$8mzt^(vyOdG# aCaZXg@4--45q4T^!