diff --git a/docs/features/mooncake_store_connector_usage.md b/docs/features/mooncake_store_connector_usage.md index a35a709793..79e200bf93 100644 --- a/docs/features/mooncake_store_connector_usage.md +++ b/docs/features/mooncake_store_connector_usage.md @@ -137,15 +137,12 @@ vllm serve meta-llama/Llama-3.1-8B-Instruct \ ``` To also offload newly completed decode KV blocks, add the following extra -configuration to the decoder's `MooncakeStoreConnector` entry. This currently -requires the prefill and decode instances to use the same tensor-parallel size -and compatible KV-cache topology. +configuration to the decoder's `MooncakeStoreConnector` entry. When decode processing starts, the consumer checks the block-aligned prompt prefix and fills any blocks missing from the Store. Subsequent saves append newly completed decode blocks. This keeps a complete, reusable prefix in the -Store and also covers deployments where prompt KV is delivered directly by -`MooncakeConnector` instead of through the Store. +Store. It also covers prompt KV delivered directly by `MooncakeConnector`. ```json { @@ -155,9 +152,41 @@ Store and also covers deployments where prompt KV is delivered directly by } ``` +#### Sharing one Store across multiple Prefill TP sizes + +Heterogeneous-TP sharing normally uses a fixed `store_tp_size`. When several +prefillers use different TP sizes, opt in to a common Store TP derived from +their least common multiple: + +```json +{ + "kv_connector_extra_config": { + "enable_store_tp_lcm": true, + "prefill_tp_sizes": [4, 2] + } +} +``` + +Every prefiller and decoder that shares these entries must use the same list. +The example selects Store TP 4: a TP4 endpoint maps each rank to one Store +shard, while TP2 endpoints map each rank to two Store shards. Runtime TP sizes +remain unchanged. A decoder configured with `"save_decode_cache": true` uses +the same Store TP for decode KV from every prefiller. + +The list may contain positive integer TP sizes. Sharing requires a Store TP that +is at least the local TP and divisible by it, an LBHNC or LBNHC local KV cache, +and the existing topology and KV-head constraints. The Store namespace includes +the attention backend's selected layout. Different layouts use separate Store +entries. Malformed lists and unsupported endpoints use an isolated rank-local +key layout. When +`enable_store_tp_lcm` is absent or false, `prefill_tp_sizes` has no effect and +the existing `store_tp_size` behavior is unchanged. + **Proxy:** -A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details. +A disaggregation proxy routes requests between prefiller and decoder nodes. +When `MooncakeConnector` is also used for direct P2P transfer, refer to its +[usage guide](mooncake_connector_usage.md) for proxy setup details. ### Disk Offloading @@ -249,9 +278,26 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`. - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format). - `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`. +- `store_tp_size` (int): Common Store TP for endpoints with different local TP sizes. It supports LBHNC and LBNHC local KV caches, with `store_tp_size >= local_tp_size` and `store_tp_size % local_tp_size == 0`. The current topology is one full-attention cache group, PCP/DCP disabled, and cross-layer blocks disabled. For GQA and MHA, the total KV-head count must be divisible by `store_tp_size`. Store shards contain fixed global KV-head ranges in the local layout. Shared endpoints use the same KV cache layout, pipeline-parallel size, and Store TP. The Store namespace includes the layout and PP size. Unsupported configurations use a topology-specific rank-local namespace. -Decode offloading uses the existing TP-rank key namespace. Cross-TP sharing is -not yet supported. +LBHNC/HND is strongly recommended for TP-sharded Store when supported. LBNHC/NHD +creates many transfer segments and may significantly reduce PUT/GET performance. + +For example, with prefill TP 4, decode TP 2, and eight KV heads, set +`store_tp_size` to 4 on both instances. Each decode rank reads and writes two +of the four Store shards. + +MQA with one total KV head uses a replicated-head layout. For the supported +prefill TP 4 to decode TP 2 case, every rank uses the +same rank-0 key namespace. The four prefill replicas stripe block PUTs so each +object is stored once, while both decode ranks GET every block into their local +KV replica. `store_tp_size` does not appear in MQA keys, so identical MQA +objects written at different store TP sizes share the same pool entry when PP +sizes match. + +Tensor-parallel collectives and low-precision arithmetic are not bitwise +invariant across TP sizes, so heterogeneous-TP reuse does not guarantee the +same greedy output as recomputing the prefix at the decode TP size. ## Notes diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index bc3432154d..96315b026d 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py @@ -48,6 +48,31 @@ def _make_vllm_config(): ) +@pytest.mark.parametrize( + "extra_config", + [ + {"store_tp_size": 4}, + {"enable_store_tp_lcm": True, "prefill_tp_sizes": [6, 4]}, + {"prefill_tp_sizes": [4, 2]}, + {"enable_store_tp_lcm": True, "prefill_tp_sizes": [4, 0]}, + None, + ], +) +def test_store_tp_does_not_override_backend_kv_cache_layout(extra_config): + vllm_config = create_vllm_config( + kv_connector="MooncakeStoreConnector", + kv_role="kv_both", + kv_connector_extra_config=extra_config, + ) + + assert ( + mooncake_store_connector.MooncakeStoreConnector.get_required_kvcache_layout( + vllm_config + ) + is None + ) + + def _make_kv_cache_config() -> KVCacheConfig: """Single-group full-attention KVCacheConfig — enough for the connector constructor's validate() pass.""" @@ -90,15 +115,18 @@ def test_scheduler_requires_align_mode_for_mamba(): scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config) -def _make_block_stored() -> BlockStored: +def _make_block_stored( + block_hash: bytes = b"hash", group_idx: int | None = None +) -> BlockStored: return BlockStored( - block_hashes=[b"hash"], + block_hashes=[block_hash], parent_block_hash=None, token_ids=[1, 2, 3], block_size=16, lora_id=None, medium="cpu", lora_name=None, + group_idx=group_idx, ) @@ -169,7 +197,7 @@ def test_worker_methods_delegate_to_store_worker(): assert invalid_block_ids == {3, 4} -def test_get_kv_connector_kv_cache_events_returns_none_when_empty(): +def test_get_kv_connector_kv_cache_events_returns_none_when_disabled(): vllm_config = _make_vllm_config() kv_cache_config = _make_kv_cache_config() @@ -184,7 +212,7 @@ def test_get_kv_connector_kv_cache_events_returns_none_when_empty(): vllm_config, KVConnectorRole.WORKER, kv_cache_config ) - mock_worker_cls.return_value.get_kv_events.return_value = [] + mock_worker_cls.return_value.enable_kv_events = False assert connector.get_kv_connector_kv_cache_events() is None @@ -248,6 +276,9 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events(): ) mock_worker_cls.return_value.get_kv_events.return_value = [event] + mock_worker_cls.return_value.enable_kv_events = True + mock_worker_cls.return_value.kv_send_thread = MagicMock() + mock_worker_cls.return_value.group_tp_replication_factors = (1,) kv_events = connector.get_kv_connector_kv_cache_events() assert isinstance(kv_events, mooncake_store_connector.MooncakeStoreKVEvents) @@ -255,6 +286,119 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events(): assert kv_events.get_all_events() == [event] +def test_get_kv_connector_kv_cache_events_keeps_empty_worker_contribution(): + vllm_config = _make_vllm_config() + kv_cache_config = _make_kv_cache_config() + + with ( + set_current_vllm_config(vllm_config), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store." + "connector.MooncakeStoreWorker" + ) as mock_worker_cls, + ): + connector = mooncake_store_connector.MooncakeStoreConnector( + vllm_config, KVConnectorRole.WORKER, kv_cache_config + ) + + worker = mock_worker_cls.return_value + worker.enable_kv_events = True + worker.kv_send_thread = MagicMock() + worker.group_tp_replication_factors = (2,) + worker.get_kv_events.return_value = [] + + kv_events = connector.get_kv_connector_kv_cache_events() + + assert isinstance(kv_events, mooncake_store_connector.MooncakeStoreKVEvents) + assert kv_events.get_number_of_workers() == 1 + assert kv_events.get_all_events() == [] + + +def _aggregate_store_events( + worker_events: list[list[BlockStored]], + replication_factors: tuple[int, ...], +) -> list[BlockStored]: + combined = mooncake_store_connector.MooncakeStoreKVEvents( + num_workers=1, + group_tp_replication_factors=replication_factors, + ) + combined.add_events(worker_events[0]) + for events in worker_events[1:]: + combined.add_events(events) + combined.increment_workers() + combined.aggregate() + return combined.get_all_events() + + +def test_gqa_store_events_follow_tp_replica_striping(): + block_a = _make_block_stored(b"a", group_idx=0) + block_b = _make_block_stored(b"b", group_idx=0) + + events = _aggregate_store_events( + [[block_a], [block_b], [block_a], [block_b]], + replication_factors=(2,), + ) + + assert events == [block_a, block_b] + + +def test_mqa_store_events_require_one_replica(): + blocks = [_make_block_stored(bytes([index]), group_idx=0) for index in range(4)] + + events = _aggregate_store_events( + [[block] for block in blocks], + replication_factors=(4,), + ) + + assert events == blocks + + +def test_store_events_require_every_non_replicated_worker(): + complete = _make_block_stored(b"complete", group_idx=0) + incomplete = _make_block_stored(b"incomplete", group_idx=0) + + events = _aggregate_store_events( + [ + [complete, incomplete], + [complete], + [complete], + [complete], + ], + replication_factors=(1,), + ) + + assert events == [complete] + + +def test_store_events_reject_incomplete_replica_coverage(): + complete = _make_block_stored(b"complete", group_idx=0) + incomplete = _make_block_stored(b"incomplete", group_idx=0) + + events = _aggregate_store_events( + [[complete, incomplete], [complete], [], []], + replication_factors=(2,), + ) + + assert events == [complete] + + +def test_store_events_use_each_groups_replication_factor(): + full_attention = _make_block_stored(b"full", group_idx=0) + replicated_attention = _make_block_stored(b"replicated", group_idx=1) + + events = _aggregate_store_events( + [ + [full_attention, replicated_attention], + [full_attention], + [full_attention, replicated_attention], + [full_attention], + ], + replication_factors=(1, 2), + ) + + assert events == [full_attention, replicated_attention] + + def test_update_connector_output_and_take_events(): vllm_config = _make_vllm_config() kv_cache_config = _make_kv_cache_config() @@ -284,6 +428,28 @@ def test_update_connector_output_and_take_events(): assert connector._kv_cache_events is None +def test_store_events_persist_across_polls(): + connector = object.__new__(mooncake_store_connector.MooncakeStoreConnector) + connector.connector_scheduler = MagicMock() + connector._kv_cache_events = None + complete = _make_block_stored(b"complete", group_idx=0) + pending = _make_block_stored(b"pending", group_idx=0) + + def poll(worker_events: list[BlockStored]) -> list[BlockStored]: + events = mooncake_store_connector.MooncakeStoreKVEvents( + num_workers=4, + group_tp_replication_factors=(2,), + ) + events.add_events(worker_events) + connector.update_connector_output(KVConnectorOutput(kv_cache_events=events)) + return list(connector.take_events()) + + assert poll([complete, pending]) == [] + assert poll([complete]) == [complete] + assert poll([pending, pending]) == [pending] + assert connector._kv_cache_events is None + + # ============================================================ # reset_cache() — RL hard-reset path via typed LookupKey protocol # ============================================================ diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_layout.py b/tests/v1/kv_connector/unit/test_mooncake_store_layout.py new file mode 100644 index 0000000000..00ff3eea4d --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_store_layout.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for Mooncake Store payload layouts.""" + +import ctypes +import random + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( + ChunkedTokenDatabase, + KeyMetadata, + LBHNCStoreLayout, + LBNHCStoreLayout, + RankLocalStoreLayout, +) +from vllm.utils.math_utils import cdiv + +BLOCK_SIZE = 128 + + +def _rank_local_layout(num_regions: int, num_block_lens: int) -> RankLocalStoreLayout: + metadata = KeyMetadata("test-model", 1, 0, 0, 0) + layout = RankLocalStoreLayout(metadata, BLOCK_SIZE, BLOCK_SIZE) + layout.set_kv_caches_base_addr( + [0x7F00_0000_0000 + i * (1 << 30) for i in range(num_regions)] + ) + layout.set_block_len([30_208 + 512 * i for i in range(num_block_lens)]) + return layout + + +def _rank_local_reference( + layout: RankLocalStoreLayout, start: int, end: int, block_ids: list[int] +) -> tuple[list[int], list[int], int]: + block_id = block_ids[start // layout.block_size] + length = len(layout.block_len) + addrs = [ + base_addr + block_id * layout.block_len[index % length] + for index, base_addr in enumerate(layout.kv_caches_base_addr) + ] + sizes = [ + layout.block_len[index % length] * cdiv(end - start, layout.block_size) + for index in range(len(layout.kv_caches_base_addr)) + ] + return addrs, sizes, block_id + + +@pytest.mark.parametrize("num_regions,num_block_lens", [(96, 96), (96, 2), (1, 1)]) +def test_rank_local_descriptors_match_reference(num_regions: int, num_block_lens: int): + layout = _rank_local_layout(num_regions, num_block_lens) + rng = random.Random(0) + block_ids = [rng.randrange(0, 1 << 20) for _ in range(300)] + chunks = [] + block = 0 + while block < len(block_ids) - 4: + span = rng.choice([1, 1, 1, 2, 4]) + chunks.append((block * BLOCK_SIZE, (block + span) * BLOCK_SIZE)) + block += span + rng.choice([0, 1]) + + addrs, sizes, selected_blocks = layout.prepare_values( + chunks, block_ids, [0] * len(chunks) + ) + + for chunk, chunk_addrs, chunk_sizes, block_id in zip( + chunks, addrs, sizes, selected_blocks, strict=True + ): + assert (chunk_addrs, chunk_sizes, block_id) == _rank_local_reference( + layout, *chunk, block_ids + ) + assert all(type(addr) is int for addr in chunk_addrs) + assert type(block_id) is int + + +def test_rank_local_descriptors_handle_empty_and_invalid_chunks(): + layout = _rank_local_layout(4, 4) + assert layout.prepare_values([], [1, 2, 3], []) == ([], [], []) + with pytest.raises(AssertionError): + layout.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1], [0]) + + +@pytest.mark.parametrize( + ("layout_cls", "store_format"), + [ + (LBHNCStoreLayout, "tp_shared_lbhnc"), + (LBNHCStoreLayout, "tp_shared_lbnhc"), + ], +) +def test_tp_shared_layout_owns_store_namespace(layout_cls, store_format): + assert ( + layout_cls.shared_namespace(4, 2) + == f"@store_tp:4@store_pp:2@store_format:{store_format}" + ) + + +def test_rank_local_database_api_rejects_tp_shared_layout(): + metadata = KeyMetadata("test-model", 0, 0, 0, 0) + layout = LBHNCStoreLayout( + metadata, + block_size=16, + hash_block_size=16, + local_tp_size=2, + store_tp_size=4, + tp_rank=0, + num_kv_heads=8, + ) + database = ChunkedTokenDatabase(metadata, 16, store_layout=layout) + + with pytest.raises(RuntimeError, match="rank-local"): + database.prepare_value_for_block(0) + + +@pytest.mark.parametrize( + ("layout_cls", "producer_strides", "consumer_strides"), + [ + (LBHNCStoreLayout, (128, 64, 4, 1), (256, 64, 4, 1)), + (LBNHCStoreLayout, (128, 4, 8, 1), (256, 4, 16, 1)), + ], +) +def test_tp_shared_layout_round_trip_across_tp_sizes( + layout_cls, producer_strides, consumer_strides +): + block_size = 16 + shape = (1, 2, block_size, 4) + stored: dict[int, bytes] = {} + producer_tensors = [] + + for tp_rank in range(4): + tensor = torch.empty_strided(shape, producer_strides, dtype=torch.float16) + tensor.copy_( + torch.arange(128, dtype=torch.float16).view(shape) + tp_rank * 1000 + ) + producer_tensors.append(tensor) + metadata = KeyMetadata("test-model", tp_rank, 0, 0, 0) + layout = layout_cls( + metadata, + block_size, + block_size, + local_tp_size=4, + store_tp_size=4, + tp_rank=tp_rank, + num_kv_heads=8, + ) + layout.register_kv_caches([tensor], 1) + addrs, sizes, _ = layout.prepare_values([(0, block_size)], [0], [tp_rank]) + stored[tp_rank] = b"".join( + ctypes.string_at(addr, size) + for addr, size in zip(addrs[0], sizes[0], strict=True) + ) + + for tp_rank in range(2): + tensor = torch.empty_strided( + (1, 4, block_size, 4), consumer_strides, dtype=torch.float16 + ) + tensor.zero_() + metadata = KeyMetadata("test-model", tp_rank, 0, 0, 0) + layout = layout_cls( + metadata, + block_size, + block_size, + local_tp_size=2, + store_tp_size=4, + tp_rank=tp_rank, + num_kv_heads=8, + ) + layout.register_kv_caches([tensor], 1) + shard_ids = layout.local_shard_ids + addrs, sizes, _ = layout.prepare_values( + [(0, block_size)] * len(shard_ids), [0], shard_ids + ) + for shard_id, shard_addrs, shard_sizes in zip( + shard_ids, addrs, sizes, strict=True + ): + offset = 0 + for addr, size in zip(shard_addrs, shard_sizes, strict=True): + ctypes.memmove(addr, stored[shard_id][offset : offset + size], size) + offset += size + + expected = torch.cat(producer_tensors[tp_rank * 2 : tp_rank * 2 + 2], dim=1) + torch.testing.assert_close(tensor, expected) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py b/tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py deleted file mode 100644 index 77cae42ffc..0000000000 --- a/tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py +++ /dev/null @@ -1,90 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for ChunkedTokenDatabase.prepare_values.""" - -import random - -import pytest - -from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( - ChunkedTokenDatabase, - KeyMetadata, -) -from vllm.utils.math_utils import cdiv - -BLOCK_SIZE = 128 - - -def _reference_prepare_value( - db: ChunkedTokenDatabase, start: int, end: int, block_ids: list[int] -) -> tuple[list[int], list[int], int]: - """Compute a token range with the original scalar implementation.""" - addr_list = [] - size_list = [] - block_id = block_ids[start // db.block_size] - length = len(db.block_len) - for index, base_addr in enumerate(db.kv_caches_base_addr): - addr = base_addr + block_id * db.block_len[index % length] - assert (end - start) % db.block_size == 0 - size = db.block_len[index % length] * cdiv(end - start, db.block_size) - addr_list.append(addr) - size_list.append(size) - return addr_list, size_list, block_id - - -def _make_db(num_regions: int, num_block_lens: int) -> ChunkedTokenDatabase: - md = KeyMetadata(model_name="t", tp_rank=1, pcp_rank=0, dcp_rank=0, pp_rank=0) - db = ChunkedTokenDatabase(md, BLOCK_SIZE) - db.set_kv_caches_base_addr( - [0x7F00_0000_0000 + i * (1 << 30) for i in range(num_regions)] - ) - # Exercise repeated block lengths when there are more cache regions. - db.set_block_len([30_208 + 512 * i for i in range(num_block_lens)]) - return db - - -@pytest.mark.parametrize("num_regions,num_block_lens", [(96, 96), (96, 2), (1, 1)]) -def test_prepare_values_matches_reference(num_regions: int, num_block_lens: int): - db = _make_db(num_regions, num_block_lens) - rng = random.Random(0) - n_blocks = 300 - block_ids = [rng.randrange(0, 1 << 20) for _ in range(n_blocks)] - chunks = [] - b = 0 - while b < n_blocks - 4: - span = rng.choice([1, 1, 1, 2, 4]) - chunks.append((b * BLOCK_SIZE, (b + span) * BLOCK_SIZE)) - b += span + rng.choice([0, 1]) - - addrs, sizes, bids = db.prepare_values(chunks, block_ids) - assert len(addrs) == len(sizes) == len(bids) == len(chunks) - for (start, end), addr, size, bid in zip(chunks, addrs, sizes, bids): - ref_addr, ref_size, ref_bid = _reference_prepare_value( - db, start, end, block_ids - ) - assert addr == ref_addr - assert size == ref_size - assert bid == ref_bid - # Native bindings require Python ints rather than numpy scalars. - assert all(type(a) is int for a in addr) - assert type(bid) is int - - -def test_prepare_value_single_matches_reference(): - db = _make_db(8, 8) - block_ids = list(range(64)) - got = db.prepare_value(5 * BLOCK_SIZE, 7 * BLOCK_SIZE, block_ids) - assert got == _reference_prepare_value( - db, 5 * BLOCK_SIZE, 7 * BLOCK_SIZE, block_ids - ) - - -def test_prepare_values_empty(): - db = _make_db(4, 4) - assert db.prepare_values([], [1, 2, 3]) == ([], [], []) - - -def test_prepare_values_rejects_unaligned_chunk(): - db = _make_db(4, 4) - with pytest.raises(AssertionError): - db.prepare_values([(0, BLOCK_SIZE + 1)], [0, 1]) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 7bfc666667..52df7531e7 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -17,6 +17,7 @@ import pytest import torch from tests.v1.attention.utils import dense_kv_cache_views +from vllm.distributed.kv_events import KVEventAggregator from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import ( rdma_utils, ) @@ -30,9 +31,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( BlobBlockHashes, ChunkedTokenDatabase, KeyMetadata, + LBHNCStoreLayout, + LBNHCStoreLayout, LoadSpec, MooncakeLookupResult, PoolKey, + RankLocalStoreLayout, ReqMeta, TailKeyBoundary, ) @@ -44,8 +48,19 @@ from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, - KVCacheLayout, + RSWASpec, ) +from vllm.v1.kv_cache_layout import KVCacheLayout + +_TP_SHARED_NAMESPACE = "@store_tp:4@store_pp:1@store_format:tp_shared_lbhnc" +_TP_SHARED_LBNHC_NAMESPACE = "@store_tp:4@store_pp:1@store_format:tp_shared_lbnhc" + + +def _tp_shared_prefix(tp_rank: int) -> str: + return ( + f"test-model{_TP_SHARED_NAMESPACE}@tp_rank:{tp_rank}" + "@pcp0@dcp0@pp_rank:0@group:0" + ) class _RecordingBlockHashes: @@ -78,6 +93,46 @@ def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator: ) +def _make_tp_shared_db( + tp_rank: int = 0, +) -> tuple[ChunkedTokenDatabase, torch.Tensor]: + num_blocks = 2 + local_heads = 4 + block_size = 16 + content_size = 4 + strides = ( + local_heads * block_size * content_size, + block_size * content_size, + content_size, + 1, + ) + tensor = torch.empty_strided( + (num_blocks, local_heads, block_size, content_size), + strides, + dtype=torch.float16, + ) + metadata = KeyMetadata( + "test-model", + tp_rank, + 0, + 0, + 0, + store_namespace=_TP_SHARED_NAMESPACE, + ) + layout = LBHNCStoreLayout( + metadata, + block_size, + block_size, + local_tp_size=2, + store_tp_size=4, + tp_rank=tp_rank, + num_kv_heads=8, + ) + db = ChunkedTokenDatabase(metadata, block_size, store_layout=layout) + layout.register_kv_caches([tensor], num_blocks) + return db, tensor + + def _make_store_sending_thread( store: MagicMock, *, @@ -243,11 +298,15 @@ def _make_vllm_config( rank: int = 0, decode_context_parallel_size: int = 1, kv_role: str = "kv_both", + pipeline_parallel_size: int = 1, + kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC, ) -> SimpleNamespace: + cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10) + cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout return SimpleNamespace( model_config=_FakeModelConfig(), parallel_config=SimpleNamespace( - pipeline_parallel_size=1, + pipeline_parallel_size=pipeline_parallel_size, rank=rank, data_parallel_index=0, decode_context_parallel_size=decode_context_parallel_size, @@ -256,7 +315,7 @@ def _make_vllm_config( kv_transfer_config=_FakeKVTransferConfig( kv_role=kv_role, extra_config=extra_config ), - cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=10), + cache_config=cache_config, kv_events_config=SimpleNamespace(enable_kv_cache_events=False), speculative_config=None, ) @@ -334,6 +393,138 @@ def test_pool_key_to_string_without_prefix_is_unchanged(): ) +def test_pool_key_store_tp_namespace(): + key = PoolKey( + KeyMetadata( + "test-model", + 2, + 0, + 0, + 0, + store_namespace=_TP_SHARED_NAMESPACE, + ), + "deadbeef", + ) + + assert key.to_string() == f"{_tp_shared_prefix(2)}@deadbeef" + + +def test_tp_shared_sending_writes_each_store_shard(): + store = MagicMock() + store.batch_is_exist.return_value = [0, 0, 0, 0] + store.batch_put_from_multi_buffers.return_value = [1, 1, 1, 1] + db, _ = _make_tp_shared_db() + thread = _make_store_sending_thread( + store, + token_databases=[db], + replicate_config=SimpleNamespace(), + ) + req = _make_store_req("req", [b"h0", b"h1"]) + _run_store_req(thread, req) + + keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args + assert keys == [ + f"{_tp_shared_prefix(0)}@6830", + f"{_tp_shared_prefix(1)}@6830", + f"{_tp_shared_prefix(0)}@6831", + f"{_tp_shared_prefix(1)}@6831", + ] + assert len(addrs) == len(sizes) == 4 + + +def test_tp_shared_kv_events_aggregate_once_per_logical_block(): + aggregator = KVEventAggregator(num_workers=2) + + for tp_rank in range(2): + store = MagicMock() + store.batch_is_exist.return_value = [0, 0] + store.batch_put_from_multi_buffers.return_value = [256, 256] + db, _ = _make_tp_shared_db(tp_rank) + thread = _make_store_sending_thread( + store, + token_databases=[db], + tp_rank=tp_rank, + replicate_config=SimpleNamespace(), + ) + thread.enable_kv_event = True + + _run_store_req(thread, _make_event_store_req(16)) + + events = thread.get_kv_events() + assert len(events) == 1 + aggregator.add_events(events) + + assert len(aggregator.get_common_events()) == 1 + + +def test_tp_shared_kv_event_waits_for_all_missing_shards(): + store = MagicMock() + store.batch_is_exist.return_value = [0, 0] + store.batch_put_from_multi_buffers.return_value = [256, -1] + db, _ = _make_tp_shared_db() + thread = _make_store_sending_thread( + store, + token_databases=[db], + replicate_config=SimpleNamespace(), + ) + thread.enable_kv_event = True + + _run_store_req(thread, _make_event_store_req(16)) + + assert thread.get_kv_events() == [] + + +def test_tp_shared_kv_event_counts_existing_shards_as_satisfied(): + store = MagicMock() + store.batch_is_exist.return_value = [1, 0] + store.batch_put_from_multi_buffers.return_value = [256] + db, _ = _make_tp_shared_db() + thread = _make_store_sending_thread( + store, + token_databases=[db], + replicate_config=SimpleNamespace(), + ) + thread.enable_kv_event = True + + _run_store_req(thread, _make_event_store_req(16)) + + assert len(thread.get_kv_events()) == 1 + + +def test_tp_shared_receiving_reads_each_local_store_shard(): + from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec + + store = MagicMock() + store.batch_get_into_multi_buffers.return_value = [1, 1, 1, 1] + db, _ = _make_tp_shared_db() + spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + coord = mooncake_store_worker.MooncakeStoreCoordinator( + [KVCacheGroupSpec(["layer0"], spec)], + scheduler_block_size=16, + hash_block_size=16, + ) + thread = mooncake_store_worker.KVCacheStoreRecvingThread( + store=store, + token_databases=[db], + block_size=16, + tp_rank=0, + ready_event=threading.Event(), + coord=coord, + ) + thread.request_queue.task_done = MagicMock() + + thread._handle_request(_make_load_req("req", [b"h0", b"h1"], token_len=32)) + + keys, addrs, sizes = store.batch_get_into_multi_buffers.call_args.args + assert keys == [ + f"{_tp_shared_prefix(0)}@6830", + f"{_tp_shared_prefix(1)}@6830", + f"{_tp_shared_prefix(0)}@6831", + f"{_tp_shared_prefix(1)}@6831", + ] + assert len(addrs) == len(sizes) == 4 + + def test_pool_key_cache_prefix_namespaces_and_disambiguates(): """A non-empty cache_prefix is prepended, and two instances with different prefixes never collide on identical block hashes.""" @@ -1155,8 +1346,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): ) db0.set_kv_caches_base_addr([0x1000]) db0.set_block_len([256]) - db0.prepare_values = MagicMock(wraps=db0.prepare_values) - db0.prepare_value = MagicMock(side_effect=AssertionError("scalar path called")) + db0.store_layout.prepare_values = MagicMock(wraps=db0.store_layout.prepare_values) db1 = ChunkedTokenDatabase( KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), @@ -1164,8 +1354,7 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): ) db1.set_kv_caches_base_addr([0x2000]) db1.set_block_len([512]) - db1.prepare_values = MagicMock(wraps=db1.prepare_values) - db1.prepare_value = MagicMock(side_effect=AssertionError("scalar path called")) + db1.store_layout.prepare_values = MagicMock(wraps=db1.store_layout.prepare_values) thread = _make_store_sending_thread( store, @@ -1183,10 +1372,12 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): ), ) - db0.prepare_value.assert_not_called() - db1.prepare_value.assert_not_called() - db0.prepare_values.assert_called_once_with([(0, 16), (32, 48)], [1, 2, 3]) - db1.prepare_values.assert_called_once_with([(16, 32), (32, 48)], [3, 2, 1]) + db0.store_layout.prepare_values.assert_called_once_with( + [(0, 16), (32, 48)], [1, 2, 3], [0, 0] + ) + db1.store_layout.prepare_values.assert_called_once_with( + [(16, 32), (32, 48)], [3, 2, 1], [0, 0] + ) keys, addrs, sizes, _ = store.batch_put_from_multi_buffers.call_args.args assert [key.rsplit("@", 1)[-1] for key in keys] == [ @@ -1836,8 +2027,10 @@ def test_recv_thread_splits_disk_offload_loads_by_budget(): assert second_keys == [ "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132", ] - base_addr = thread.token_databases[0].kv_caches_base_addr[0] - block_len = thread.token_databases[0].block_len[0] + layout = thread.token_databases[0].store_layout + assert isinstance(layout, RankLocalStoreLayout) + base_addr = layout.kv_caches_base_addr[0] + block_len = layout.block_len[0] assert first_addrs == [[base_addr], [base_addr + block_len]] assert second_addrs == [[base_addr + 2 * block_len]] expected_size = block_len @@ -2004,6 +2197,294 @@ def test_requester_worker_init_skips_disk_budget_when_offload_disabled( assert w.disk_offload_buffer_budget_bytes is None +def test_worker_enables_lbhnc_store_tp_layout(tmp_path, monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch, tp_size=2) + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + w = worker.MooncakeStoreWorker( + _make_vllm_config(extra_config={"store_tp_size": 4}), + _make_kv_cache_config(), + ) + + assert w.store_tp_size == 4 + assert w.token_dbs[0].metadata.store_namespace == _TP_SHARED_NAMESPACE + assert isinstance(w.token_dbs[0].store_layout, LBHNCStoreLayout) + assert w.token_dbs[0].store_layout.local_shard_ids == (0, 1) + assert len(w._lookup_key_prefixes[0]) == 4 + + +def test_worker_enables_lbnhc_store_tp_layout(tmp_path, monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch, tp_size=2) + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + with patch.object(worker.logger, "warning_once") as warning_once: + w = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={"store_tp_size": 4}, + kv_cache_layout=KVCacheLayout.LBNHC, + ), + _make_kv_cache_config(), + ) + + warning_once.assert_called_once() + assert w.store_tp_size == 4 + assert w.token_dbs[0].metadata.store_namespace == _TP_SHARED_LBNHC_NAMESPACE + assert isinstance(w.token_dbs[0].store_layout, LBNHCStoreLayout) + assert w.token_dbs[0].store_layout.local_shard_ids == (0, 1) + assert len(w._lookup_key_prefixes[0]) == 4 + + +def test_tp_sharded_layout_rejects_full_attention_subclasses(): + spec = RSWASpec( + block_size=16, + num_kv_heads=8, + head_size=64, + dtype=None, + rswa_window=32, + ) + store_worker = object.__new__(worker.MooncakeStoreWorker) + store_worker.pcp_size = 1 + store_worker.dcp_size = 1 + store_worker._kv_cache_groups = [SimpleNamespace(kv_cache_spec=spec)] + + assert not store_worker._supports_tp_sharded_store_layout(LBHNCStoreLayout, {}) + + +@pytest.mark.parametrize( + ("extra_config", "expected"), + [ + ({"store_tp_size": 4}, 4), + ( + { + "enable_store_tp_lcm": True, + "prefill_tp_sizes": [3, 4, 6], + }, + 12, + ), + ( + { + "enable_store_tp_lcm": False, + "prefill_tp_sizes": [4, 2], + }, + None, + ), + ({"enable_store_tp_lcm": True, "prefill_tp_sizes": []}, None), + ({"enable_store_tp_lcm": True, "prefill_tp_sizes": [4, 0]}, None), + ({"enable_store_tp_lcm": True, "prefill_tp_sizes": [4, 2.0]}, None), + ({"store_tp_size": 0}, None), + ({"store_tp_size": -1}, None), + ({"store_tp_size": True}, None), + ({"store_tp_size": 4.0}, None), + ], +) +def test_resolve_store_tp_size(extra_config: dict[str, object], expected: int | None): + assert worker.resolve_store_tp_size(extra_config) == expected + + +def test_lcm_store_tp_gives_prefill_and_decode_common_namespace(tmp_path, monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + lcm_config = { + "enable_store_tp_lcm": True, + "prefill_tp_sizes": [4, 2], + } + + workers = [] + topologies = [ + (4, 2, "kv_both", (2,)), + (2, 1, "kv_both", (2, 3)), + (2, 1, "kv_consumer", (2, 3)), + ] + for tp_size, tp_rank, kv_role, expected_shards in topologies: + _patch_worker_runtime( + monkeypatch, + tp_size=tp_size, + tp_rank=tp_rank, + ) + extra_config = dict(lcm_config) + if kv_role == "kv_consumer": + extra_config["save_decode_cache"] = True + store_worker = worker.MooncakeStoreWorker( + _make_vllm_config(extra_config=extra_config, kv_role=kv_role), + _make_kv_cache_config(), + ) + assert store_worker.store_tp_size == 4 + assert isinstance(store_worker.token_dbs[0].store_layout, LBHNCStoreLayout) + assert store_worker.token_dbs[0].store_layout.local_shard_ids == expected_shards + workers.append(store_worker) + + shard_keys = [ + store_worker.token_dbs[0].store_layout.key_for(2, BlockHash(b"h")) + for store_worker in workers + ] + assert len(set(shard_keys)) == 1 + assert workers[-1].can_put + + +@pytest.mark.parametrize( + ("tp_size", "prefill_tp_sizes"), + [ + (3, [4, 2]), + (2, [4, 3]), + (2, []), + (2, [4, 0]), + ], +) +def test_lcm_store_tp_falls_back_when_topology_is_incompatible( + tmp_path, + monkeypatch, + tp_size: int, + prefill_tp_sizes: list[int], +): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch, tp_size=tp_size) + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + store_worker = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={ + "enable_store_tp_lcm": True, + "prefill_tp_sizes": prefill_tp_sizes, + } + ), + _make_kv_cache_config(), + ) + + assert store_worker.store_tp_size is None + assert type(store_worker.token_dbs[0]) is ChunkedTokenDatabase + assert "@store_format:rank_local_tp" in store_worker.token_dbs[0].key_for( + BlockHash(b"h") + ) + + +@pytest.mark.parametrize("store_tp_size", [1, 3, 6]) +def test_worker_keeps_rank_local_layout_for_unsupported_store_tp( + tmp_path, + monkeypatch, + store_tp_size: int, +): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch, tp_size=2) + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + w = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={"store_tp_size": store_tp_size}, + kv_cache_layout=KVCacheLayout.LBNHC, + ), + _make_kv_cache_config(), + ) + + assert w.store_tp_size is None + assert type(w.token_dbs[0]) is ChunkedTokenDatabase + key = w.token_dbs[0].key_for(BlockHash(b"h")) + assert "@store_tp:" not in key + assert "@store_pp:1" in key + assert "@store_format:rank_local_tp2_layout_LBNHC" in key + + +def test_store_tp_opt_in_isolates_different_pp_sizes(tmp_path, monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + _patch_worker_runtime(monkeypatch, tp_size=2) + monkeypatch.setattr(_FakeModelConfig, "get_total_num_kv_heads", lambda _self: 8) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + keys = [] + for pp_size in (1, 2): + w = worker.MooncakeStoreWorker( + _make_vllm_config( + extra_config={"store_tp_size": 4}, + pipeline_parallel_size=pp_size, + ), + _make_kv_cache_config(), + ) + layout = w.token_dbs[0].store_layout + keys.append(layout.key_for(layout.local_shard_ids[0], BlockHash(b"h"))) + + assert "@store_pp:1" in keys[0] + assert "@store_pp:2" in keys[1] + assert keys[0] != keys[1] + + def test_save_decode_cache_keeps_transfer_path_enabled(tmp_path, monkeypatch): store = MagicMock() store.setup.return_value = 0 @@ -2185,6 +2666,58 @@ def test_worker_put_striding_covers_every_rank_get_namespace( ) +def test_mqa_p4_to_d2_uses_shared_rank_zero_namespace(tmp_path, monkeypatch): + store = MagicMock() + store.setup.return_value = 0 + _install_fake_mooncake(monkeypatch, store) + monkeypatch.setenv( + "MOONCAKE_CONFIG_PATH", + _write_mooncake_config( + tmp_path, + { + "metadata_server": "http://metadata/endpoint", + "protocol": "tcp", + "master_server_address": "10.0.0.7:50051", + }, + ), + ) + + block_hash = BlockHash(b"hash") + put_keys: set[str] = set() + get_keys: set[str] = set() + for tp_size, target in ((4, put_keys), (2, get_keys)): + for tp_rank in range(tp_size): + _patch_worker_runtime(monkeypatch, tp_rank=tp_rank, tp_size=tp_size) + w = worker.MooncakeStoreWorker( + _make_vllm_config( + rank=tp_rank, + extra_config={"store_tp_size": 4}, + ), + _make_kv_cache_config(), + ) + + assert w.store_tp_size is None + assert w.token_dbs[0].metadata.tp_rank == 0 + key = w.token_dbs[0].key_for(block_hash) + if tp_size == 4: + put_step = w._group_tp_replication_factors[0] + if tp_rank % put_step == 0: + target.add(key) + else: + target.add(key) + + assert ( + put_keys + == get_keys + == { + ( + "test-model@store_pp:1@store_format:tp_shared_mqa" + "@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@68617368" + ) + } + ) + + def test_requester_worker_group_semantics_falls_back_without_group_ids( tmp_path, monkeypatch, @@ -2396,9 +2929,11 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) - store.batch_put_from_multi_buffers.side_effect = ( - lambda keys, addrs, sizes, *_args: ([256] * len(keys)) - ) + + def successful_put(keys, addrs, sizes, *_args): + return [256] * len(keys) + + store.batch_put_from_multi_buffers.side_effect = successful_put full_spec = FullAttentionSpec( block_size=32, num_kv_heads=8, head_size=64, dtype=None @@ -2754,6 +3289,7 @@ def _make_bare_worker( worker.recv_request_queue = queue.Queue() worker.finished_store_req = set() worker.tp_size = 1 + worker.store_tp_size = None worker.num_kv_head = 1 worker.pp_size = 1 # Minimal single-full-attention-group config so the coordinator-based @@ -2812,6 +3348,40 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): ) +def test_lookup_key_prefixes_cover_store_tp_shards(): + worker = _make_bare_worker() + worker.store_tp_size = 4 + metadata = KeyMetadata( + "test-model", + 0, + 0, + 0, + 0, + group_id=0, + store_namespace=_TP_SHARED_NAMESPACE, + ) + worker.token_dbs = [ + ChunkedTokenDatabase( + metadata, + block_size=16, + store_layout=LBHNCStoreLayout( + metadata, + block_size=16, + hash_block_size=16, + local_tp_size=2, + store_tp_size=4, + tp_rank=0, + num_kv_heads=8, + ), + ) + ] + worker._init_lookup_key_prefixes() + + assert worker._lookup_key_prefixes[0] == ( + *(_tp_shared_prefix(tp_rank) for tp_rank in range(4)), + ) + + def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): worker = _make_bare_worker() worker.tp_size = 4 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index b6128b28db..a73c6dea43 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -11,16 +11,17 @@ and consumer instances read/write KV to/from the store independently, enabling prefix caching via hash-based deduplication. """ -from collections.abc import Iterable +from collections import Counter +from collections.abc import Iterable, Sequence from typing import Any import torch from vllm.config import VllmConfig from vllm.distributed.kv_events import ( + BlockStored, KVCacheEvent, KVConnectorKVEvents, - KVEventAggregator, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, @@ -56,31 +57,66 @@ logger = init_logger(__name__) class MooncakeStoreKVEvents(KVConnectorKVEvents): """KV event aggregation for MooncakeStoreConnector.""" - def __init__(self, num_workers: int) -> None: - self._aggregator = KVEventAggregator(num_workers) + def __init__( + self, + num_workers: int, + group_tp_replication_factors: Sequence[int] = (1,), + ) -> None: + if num_workers <= 0: + raise ValueError("num_workers must be greater than zero.") + if any(factor <= 0 for factor in group_tp_replication_factors): + raise ValueError("TP replication factors must be greater than zero.") + self._event_counter: Counter[KVCacheEvent] = Counter() + self._num_workers = num_workers + self._group_tp_replication_factors = tuple(group_tp_replication_factors) def add_events(self, events: list[KVCacheEvent]) -> None: - self._aggregator.add_events(events) + if not isinstance(events, list): + raise TypeError("events must be a list of KVCacheEvent.") + self._event_counter.update(events) + + def _replication_factor(self, event: KVCacheEvent) -> int: + if not isinstance(event, BlockStored) or event.group_idx is None: + return 1 + return self._group_tp_replication_factors[event.group_idx] + + def _is_common_event(self, event: KVCacheEvent, count: int) -> bool: + return count * self._replication_factor(event) >= self._num_workers def aggregate(self) -> "MooncakeStoreKVEvents": - common_events = self._aggregator.get_common_events() - self._aggregator.clear_events() - self._aggregator.add_events(common_events) - self._aggregator.reset_workers() + common_events = self.pop_common_events() + self._event_counter.clear() + self._event_counter.update(common_events) + self._num_workers = 1 return self + def pop_common_events(self) -> list[KVCacheEvent]: + common_events = [ + event + for event, count in self._event_counter.items() + if self._is_common_event(event, count) + ] + for event in common_events: + del self._event_counter[event] + return common_events + + def has_events(self) -> bool: + return bool(self._event_counter) + def increment_workers(self, count: int = 1) -> None: - self._aggregator.increment_workers(count) + if count <= 0: + raise ValueError("count must be positive.") + self._num_workers += count def get_all_events(self) -> list[KVCacheEvent]: - return self._aggregator.get_all_events() + return list(self._event_counter.elements()) def get_number_of_workers(self) -> int: - return self._aggregator.get_number_of_workers() + return self._num_workers def clear_events(self) -> None: - self._aggregator.clear_events() - self._aggregator.reset_workers() + self._event_counter.clear() + self._num_workers = 1 def __repr__(self) -> str: return f"" @@ -260,16 +296,13 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): self._kv_cache_events = kv_cache_events else: self._kv_cache_events.add_events(kv_cache_events.get_all_events()) - self._kv_cache_events.increment_workers( - kv_cache_events.get_number_of_workers() - ) def take_events(self) -> Iterable[KVCacheEvent]: if self._kv_cache_events is not None: - self._kv_cache_events.aggregate() - yield from self._kv_cache_events.get_all_events() - self._kv_cache_events.clear_events() - self._kv_cache_events = None + events = self._kv_cache_events.pop_common_events() + if not self._kv_cache_events.has_events(): + self._kv_cache_events = None + yield from events # ============================================================ # Worker-side methods @@ -321,11 +354,19 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA): self, ) -> MooncakeStoreKVEvents | None: assert self.connector_worker is not None - events = self.connector_worker.get_kv_events() - if not events: + if ( + not self.connector_worker.enable_kv_events + or self.connector_worker.kv_send_thread is None + ): return None - - kv_events = MooncakeStoreKVEvents(num_workers=1) + events = self.connector_worker.get_kv_events() + # Empty containers still count this worker toward the poll's quorum. + kv_events = MooncakeStoreKVEvents( + num_workers=1, + group_tp_replication_factors=( + self.connector_worker.group_tp_replication_factors + ), + ) kv_events.add_events(events) return kv_events diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 34b7ac565d..b9df9648d5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( ) from vllm.logger import init_logger from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import is_non_overlapping_and_dense from vllm.v1.core.kv_cache_utils import ( BlockHash, BlockHashListWithBlockSize, @@ -111,6 +112,9 @@ class KeyMetadata: # share one Mooncake master without colliding on identical block hashes. # Empty (the default) keeps keys byte-identical to the unprefixed format. cache_prefix: str = "" + # Complete namespace for an opt-in Store payload format. Empty keeps + # historical keys byte-identical. + store_namespace: str = "" @dataclass(order=True) @@ -125,6 +129,7 @@ class PoolKey: ( self.key_metadata.cache_prefix, self.key_metadata.model_name, + self.key_metadata.store_namespace, self.key_metadata.tp_rank, self.key_metadata.pcp_rank, self.key_metadata.dcp_rank, @@ -148,6 +153,7 @@ class PoolKey: return ( f"{prefix}" f"{key_metadata.model_name}" + f"{key_metadata.store_namespace}" f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}" f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}" f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}" @@ -165,61 +171,149 @@ class PoolKey: ) -class ChunkedTokenDatabase: - """Maps token positions to store keys and GPU memory addresses.""" +StoreShardId = int +# Physical ranks encoded in a rank-local key: (TP, PCP, DCP, PP). +RankNamespace = tuple[int, int, int, int] + + +class StoreLayout: + """Store payload layout for one KV-cache group on one local TP rank.""" def __init__( self, metadata: KeyMetadata, block_size: int, - hash_block_size: int | None = None, - ): + hash_block_size: int, + ) -> None: self.metadata = metadata self.block_size = block_size - self.hash_block_size = hash_block_size or block_size - if self.block_size % self.hash_block_size != 0: - raise ValueError( - f"block_size ({self.block_size}) must be a multiple of " - f"hash_block_size ({self.hash_block_size})" - ) - self.kv_caches_base_addr: list[int] = [] - self.block_len: list[int] = [] - self._key_prefix = PoolKey.build_prefix(metadata) + self.hash_block_size = hash_block_size - def key_for(self, chunk_hash: BlockHash) -> str: - return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex()) + @property + def local_shard_ids(self) -> tuple[StoreShardId, ...]: + """Store objects contributed or consumed by this local rank.""" + raise NotImplementedError - def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]): - self.kv_caches_base_addr = kv_caches_base_addr + def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str: + raise NotImplementedError - def set_block_len(self, block_len: list[int]): - self.block_len = block_len + def lookup_key_prefixes( + self, + rank_namespaces: Sequence[RankNamespace], + ) -> tuple[str, ...]: + """Prefixes that must exist for a logical block to be reusable.""" + raise NotImplementedError - def prepare_value( - self, start: int, end: int, block_ids: list[int] - ) -> tuple[list[int], list[int], int]: - """Compute memory addresses and sizes for a single token range. - - Returns: - (addr_list, size_list, block_id) - """ - addr_lists, size_lists, chunk_block_ids = self.prepare_values( - ((start, end),), block_ids - ) - return addr_lists[0], size_lists[0], chunk_block_ids[0] + def register_kv_caches( + self, + kv_caches: Sequence[torch.Tensor], + num_blocks: int, + ) -> None: + """Build descriptors from the local KV cache tensors.""" + raise NotImplementedError def prepare_values( self, chunks: Sequence[tuple[int, int]], block_ids: list[int], + shard_ids: Sequence[StoreShardId], ) -> tuple[list[list[int]], list[list[int]], list[int]]: - """Compute memory addresses and sizes for multiple token ranges. + """Build Store multi-buffer descriptors for logical chunks.""" + raise NotImplementedError - Returns: - (addr_lists, size_lists, chunk_block_ids), one entry per chunk. - """ + +class RankLocalStoreLayout(StoreLayout): + """Historical rank-local Store payload layout.""" + + def __init__( + self, + metadata: KeyMetadata, + block_size: int, + hash_block_size: int, + ) -> None: + super().__init__(metadata, block_size, hash_block_size) + self._key_prefix = PoolKey.build_prefix(metadata) + self.kv_caches_base_addr: list[int] = [] + self.block_len: list[int] = [] + + @property + def local_shard_ids(self) -> tuple[StoreShardId, ...]: + return (0,) + + def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str: + assert shard_id == 0 + return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex()) + + def lookup_key_prefixes( + self, + rank_namespaces: Sequence[RankNamespace], + ) -> tuple[str, ...]: + return tuple( + PoolKey.build_prefix( + self.metadata, + tp_rank=tp_rank, + pcp_rank=pcp_rank, + dcp_rank=dcp_rank, + pp_rank=pp_rank, + ) + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces + ) + + def set_kv_caches_base_addr(self, base_addrs: list[int]) -> None: + self.kv_caches_base_addr = base_addrs + + def set_block_len(self, block_lens: list[int]) -> None: + self.block_len = block_lens + + def register_kv_caches( + self, + kv_caches: Sequence[torch.Tensor], + num_blocks: int, + ) -> None: + seen_region_ptrs: set[int] = set() + base_addrs: list[int] = [] + block_lens: list[int] = [] + for cache in kv_caches: + storage = cache.untyped_storage() + storage_addr = storage.data_ptr() + storage_len = storage.nbytes() + if not is_non_overlapping_and_dense(cache[0]): + for head_idx in range(cache.shape[1]): + head_cache = cache[:, head_idx] + assert is_non_overlapping_and_dense(head_cache[0]) + region_addr = head_cache.data_ptr() + if region_addr in seen_region_ptrs: + continue + seen_region_ptrs.add(region_addr) + base_addrs.append(region_addr) + block_lens.append(head_cache.stride(0) * head_cache.element_size()) + elif cache.stride(0) * cache.element_size() * num_blocks == storage_len: + if storage_addr in seen_region_ptrs: + continue + seen_region_ptrs.add(storage_addr) + base_addrs.append(storage_addr) + block_lens.append(storage_len // num_blocks) + else: + region_addr = cache.data_ptr() + if region_addr in seen_region_ptrs: + continue + seen_region_ptrs.add(region_addr) + base_addrs.append(region_addr) + block_lens.append(cache.stride(0) * cache.element_size()) + self.kv_caches_base_addr = base_addrs + self.block_len = block_lens + + def prepare_values( + self, + chunks: Sequence[tuple[int, int]], + block_ids: list[int], + shard_ids: Sequence[StoreShardId], + ) -> tuple[list[list[int]], list[list[int]], list[int]]: if not chunks: return [], [], [] + assert len(chunks) == len(shard_ids) and all( + shard_id == 0 for shard_id in shard_ids + ) base = np.asarray(self.kv_caches_base_addr, dtype=np.int64) length = len(self.block_len) blen = np.asarray( @@ -229,7 +323,7 @@ class ChunkedTokenDatabase: n = len(chunks) starts = np.fromiter((c[0] for c in chunks), dtype=np.int64, count=n) spans = np.fromiter((c[1] for c in chunks), dtype=np.int64, count=n) - starts - assert not (spans % self.hash_block_size).any() + assert np.all(spans % self.hash_block_size == 0) bids = np.fromiter( (block_ids[i] for i in (starts // self.block_size).tolist()), dtype=np.int64, @@ -241,15 +335,272 @@ class ChunkedTokenDatabase: return addrs.tolist(), sizes.tolist(), bids.tolist() def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]: - """Return addresses and sizes for one physical block slot.""" - addr_list = [] - size_list = [] length = len(self.block_len) - for index, base_addr in enumerate(self.kv_caches_base_addr): - addr = base_addr + block_id * self.block_len[index % length] - addr_list.append(addr) - size_list.append(self.block_len[index % length]) - return addr_list, size_list + return ( + [ + base_addr + block_id * self.block_len[index % length] + for index, base_addr in enumerate(self.kv_caches_base_addr) + ], + [ + self.block_len[index % length] + for index in range(len(self.kv_caches_base_addr)) + ], + ) + + +class TPShardedStoreLayout(StoreLayout): + """Store layout shared by divisible TP sizes.""" + + store_format: str + + @classmethod + def shared_namespace(cls, store_tp_size: int, pp_size: int) -> str: + return ( + f"@store_tp:{store_tp_size}@store_pp:{pp_size}" + f"@store_format:{cls.store_format}" + ) + + def __init__( + self, + metadata: KeyMetadata, + block_size: int, + hash_block_size: int, + local_tp_size: int, + store_tp_size: int, + tp_rank: int, + num_kv_heads: int, + ) -> None: + super().__init__(metadata, block_size, hash_block_size) + self.shards_per_rank = store_tp_size // local_tp_size + first_shard = tp_rank * self.shards_per_rank + self.store_shard_ids = tuple( + range(first_shard, first_shard + self.shards_per_rank) + ) + self.store_tp_size = store_tp_size + self.heads_per_store_shard = num_kv_heads // store_tp_size + self.local_num_kv_heads = num_kv_heads // local_tp_size + self._shard_key_prefixes = { + shard_id: PoolKey.build_prefix(metadata, tp_rank=shard_id) + for shard_id in self.store_shard_ids + } + self._shard_addr_bases = np.empty((self.shards_per_rank, 0), dtype=np.uint64) + self._shard_block_strides = np.empty((self.shards_per_rank, 0), dtype=np.uint64) + self._shard_sizes = np.empty((self.shards_per_rank, 0), dtype=np.uint64) + + @property + def local_shard_ids(self) -> tuple[StoreShardId, ...]: + return self.store_shard_ids + + def key_for(self, shard_id: StoreShardId, chunk_hash: BlockHash) -> str: + return PoolKey.build_key_string( + self._shard_key_prefixes[shard_id], chunk_hash.hex() + ) + + def lookup_key_prefixes( + self, + rank_namespaces: Sequence[RankNamespace], + ) -> tuple[str, ...]: + pp_ranks = sorted({namespace[3] for namespace in rank_namespaces}) + return tuple( + PoolKey.build_prefix(self.metadata, tp_rank=shard_id, pp_rank=pp_rank) + for shard_id in range(self.store_tp_size) + for pp_rank in pp_ranks + ) + + def _set_segment_templates( + self, templates: Sequence[tuple[list[int], list[int], list[int]]] + ) -> None: + self._shard_addr_bases = np.asarray( + [template[0] for template in templates], dtype=np.uint64 + ) + self._shard_block_strides = np.asarray( + [template[1] for template in templates], dtype=np.uint64 + ) + self._shard_sizes = np.asarray( + [template[2] for template in templates], dtype=np.uint64 + ) + + def _has_packed_strides( + self, head_stride: int, token_stride: int, content_bytes: int + ) -> bool: + raise NotImplementedError + + def _shard_segments( + self, head_start: int, head_stride: int, token_stride: int, content_bytes: int + ) -> Iterable[tuple[int, int]]: + raise NotImplementedError + + def register_kv_caches( + self, + kv_caches: Sequence[torch.Tensor], + num_blocks: int, + ) -> None: + templates: list[tuple[list[int], list[int], list[int]]] = [ + ([], [], []) for _ in range(self.shards_per_rank) + ] + for cache in kv_caches: + if cache.ndim != 4 or tuple(cache.shape[:3]) != ( + num_blocks, + self.local_num_kv_heads, + self.block_size, + ): + raise ValueError( + "TP-shared Mooncake store requires packed KV caches with " + "logical shape (num_blocks, local_kv_heads, block_size, content)" + ) + + element_size = cache.element_size() + block_stride, head_stride, token_stride, content_stride = ( + stride * element_size for stride in cache.stride() + ) + content_bytes = cache.shape[3] * element_size + if content_stride != element_size or not self._has_packed_strides( + head_stride, token_stride, content_bytes + ): + raise ValueError( + "TP-shared Mooncake store requires packed " + f"{self.store_format} KV layout" + ) + + for local_shard, (addr_bases, block_strides, sizes) in enumerate(templates): + head_start = local_shard * self.heads_per_store_shard + for offset, size in self._shard_segments( + head_start, head_stride, token_stride, content_bytes + ): + addr_bases.append(cache.data_ptr() + offset) + block_strides.append(block_stride) + sizes.append(size) + + self._set_segment_templates(templates) + + def prepare_values( + self, + chunks: Sequence[tuple[int, int]], + block_ids: list[int], + shard_ids: Sequence[StoreShardId], + ) -> tuple[list[list[int]], list[list[int]], list[int]]: + if not chunks: + return [], [], [] + if len(chunks) != len(shard_ids): + raise ValueError("Each Store chunk must have one shard ID") + + count = len(chunks) + starts = np.fromiter( + (chunk[0] for chunk in chunks), dtype=np.int64, count=count + ) + ends = np.fromiter((chunk[1] for chunk in chunks), dtype=np.int64, count=count) + if np.any(ends - starts != self.block_size): + raise ValueError("TP-shared Mooncake store requires full KV blocks") + + first_shard = self.store_shard_ids[0] + local_shards = np.fromiter( + (shard_id - first_shard for shard_id in shard_ids), + dtype=np.int64, + count=count, + ) + if np.any(local_shards < 0) or np.any(local_shards >= self.shards_per_rank): + raise ValueError("Store shard is not owned by this TP rank") + chunk_block_ids = np.fromiter( + (block_ids[index] for index in (starts // self.block_size).tolist()), + dtype=np.uint64, + count=count, + ) + addr_bases = self._shard_addr_bases[local_shards] + block_strides = self._shard_block_strides[local_shards] + addrs = addr_bases + chunk_block_ids[:, None] * block_strides + sizes = self._shard_sizes[local_shards] + return addrs.tolist(), sizes.tolist(), chunk_block_ids.tolist() + + +class LBHNCStoreLayout(TPShardedStoreLayout): + """Native head-major layout shared by divisible TP sizes.""" + + store_format = "tp_shared_lbhnc" + + def _has_packed_strides( + self, head_stride: int, token_stride: int, content_bytes: int + ) -> bool: + return ( + token_stride == content_bytes + and head_stride == self.block_size * content_bytes + ) + + def _shard_segments( + self, head_start: int, head_stride: int, token_stride: int, content_bytes: int + ) -> Iterable[tuple[int, int]]: + yield head_start * head_stride, self.heads_per_store_shard * head_stride + + +class LBNHCStoreLayout(TPShardedStoreLayout): + """Native token-major layout shared by divisible TP sizes.""" + + store_format = "tp_shared_lbnhc" + + def _has_packed_strides( + self, head_stride: int, token_stride: int, content_bytes: int + ) -> bool: + return ( + head_stride == content_bytes + and token_stride == self.local_num_kv_heads * content_bytes + ) + + def _shard_segments( + self, head_start: int, head_stride: int, token_stride: int, content_bytes: int + ) -> Iterable[tuple[int, int]]: + for token_idx in range(self.block_size): + yield ( + token_idx * token_stride + head_start * head_stride, + self.heads_per_store_shard * content_bytes, + ) + + +class ChunkedTokenDatabase: + """Enumerates logical token chunks and their hashes.""" + + def __init__( + self, + metadata: KeyMetadata, + block_size: int, + hash_block_size: int | None = None, + store_layout: StoreLayout | None = None, + ): + self.metadata = metadata + self.block_size = block_size + self.hash_block_size = hash_block_size or block_size + if self.block_size % self.hash_block_size != 0: + raise ValueError( + f"block_size ({self.block_size}) must be a multiple of " + f"hash_block_size ({self.hash_block_size})" + ) + self.store_layout = store_layout or RankLocalStoreLayout( + metadata, block_size, self.hash_block_size + ) + + @property + def kv_caches_base_addr(self) -> list[int]: + return self._rank_local_layout().kv_caches_base_addr + + @property + def block_len(self) -> list[int]: + return self._rank_local_layout().block_len + + def _rank_local_layout(self) -> RankLocalStoreLayout: + if not isinstance(self.store_layout, RankLocalStoreLayout): + raise RuntimeError("This operation requires a rank-local Store layout") + return self.store_layout + + def key_for(self, chunk_hash: BlockHash) -> str: + return self._rank_local_layout().key_for(0, chunk_hash) + + def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]): + self._rank_local_layout().set_kv_caches_base_addr(kv_caches_base_addr) + + def set_block_len(self, block_len: list[int]): + self._rank_local_layout().set_block_len(block_len) + + def prepare_value_for_block(self, block_id: int) -> tuple[list[int], list[int]]: + """Return addresses and sizes for one physical block slot.""" + return self._rank_local_layout().prepare_value_for_block(block_id) def process_tokens( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 95c892c06f..89d5315aa9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -12,6 +12,7 @@ and MooncakeDistributedStore integration. import dataclasses import json +import math import os import queue import socket @@ -44,12 +45,16 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( BlobBlockHashes, ChunkedTokenDatabase, KeyMetadata, + LBHNCStoreLayout, + LBNHCStoreLayout, MooncakeLookupResult, MooncakeStoreConnectorMetadata, MooncakeStoreWorkerMetadata, PoolKey, ReqMeta, + StoreShardId, TailKeyBoundary, + TPShardedStoreLayout, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import ( # noqa: E501 LOOKUP_MSG, @@ -62,7 +67,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip, make_zmq_socket -from vllm.utils.torch_utils import is_non_overlapping_and_dense from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.core.kv_cache_utils import ( BlockHash, @@ -71,6 +75,7 @@ from vllm.v1.core.kv_cache_utils import ( resolve_kv_cache_block_sizes, ) from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, KVCacheConfig, KVCacheSpec, MambaSpec, @@ -79,6 +84,7 @@ from vllm.v1.kv_cache_interface import ( UniformTypeKVCacheSpecs, group_kernel_blocks, ) +from vllm.v1.kv_cache_layout import KVCacheLayout from .metrics import MooncakeStoreConnectorStats @@ -92,6 +98,22 @@ MOONCAKE_NO_AVAILABLE_HANDLE = -200 _T = TypeVar("_T") +def resolve_store_tp_size(extra_config: dict[str, Any]) -> int | None: + """Resolve the common Store TP requested by connector config.""" + if extra_config.get("enable_store_tp_lcm") is True: + prefill_tp_sizes = extra_config.get("prefill_tp_sizes") + if not isinstance(prefill_tp_sizes, list) or not prefill_tp_sizes: + return None + if any( + type(tp_size) is not int or tp_size <= 0 for tp_size in prefill_tp_sizes + ): + return None + return math.lcm(*prefill_tp_sizes) + + store_tp_size = extra_config.get("store_tp_size") + return store_tp_size if type(store_tp_size) is int and store_tp_size > 0 else None + + def _rotate_list(values: list[_T], offset: int) -> list[_T]: return values[offset:] + values[:offset] @@ -109,7 +131,10 @@ def _make_mooncake_group_id(metadata: KeyMetadata, chunk_hash: str) -> str: # Mooncake group ids describe the lifecycle unit. For vLLM, that unit is # a prefix chunk, so shard dimensions stay only in the object key. prefix = f"{metadata.cache_prefix}@" if metadata.cache_prefix else "" - return f"vllm-mooncake-store:{prefix}{metadata.model_name}@{chunk_hash}" + return ( + f"vllm-mooncake-store:{prefix}{metadata.model_name}" + f"{metadata.store_namespace}@{chunk_hash}" + ) # Mirrors FileStorageConfig::local_buffer_size in Mooncake C++. @@ -952,8 +977,11 @@ class KVCacheStoreSendingThread(KVTransferThread): starts: list[int] = [] ends: list[int] = [] keys: list[str] = [] - kv_event_block_hashes: list[BlockHash] = [] + event_specs: list[tuple[int, int, int, BlockHash]] | None = ( + [] if self.enable_kv_event else None + ) group_indices: list[int] = [] + store_shard_ids: list[StoreShardId] = [] for g_idx, db in enumerate(self.token_databases): # Rotate the stride phase per group to balance load across ranks. put_step = self.group_put_steps[g_idx] @@ -980,12 +1008,14 @@ class KVCacheStoreSendingThread(KVTransferThread): block_idx, ) continue - starts.append(start) - ends.append(end) - keys.append(db.key_for(block_hash)) - if self.enable_kv_event: - kv_event_block_hashes.append(block_hash) - group_indices.append(g_idx) + for store_shard_id in db.store_layout.local_shard_ids: + starts.append(start) + ends.append(end) + keys.append(db.store_layout.key_for(store_shard_id, block_hash)) + group_indices.append(g_idx) + store_shard_ids.append(store_shard_id) + if event_specs is not None: + event_specs.append((start, end, g_idx, block_hash)) if not keys: self._record_saved(req_meta, token_len) @@ -1023,11 +1053,10 @@ class KVCacheStoreSendingThread(KVTransferThread): starts = [starts[i] for i in missing_indices] ends = [ends[i] for i in missing_indices] keys = [keys[i] for i in missing_indices] - if self.enable_kv_event: - kv_event_block_hashes = [ - kv_event_block_hashes[i] for i in missing_indices - ] + if event_specs is not None: + event_specs = [event_specs[i] for i in missing_indices] group_indices = [group_indices[i] for i in missing_indices] + store_shard_ids = [store_shard_ids[i] for i in missing_indices] group_ids = ( [ @@ -1050,61 +1079,33 @@ class KVCacheStoreSendingThread(KVTransferThread): addrs: list[list[int]] = [] sizes: list[list[int]] = [] - stored_events: list[BlockStored] = [] chunks_per_group: list[list[tuple[int, int]]] = [ [] for _ in self.token_databases ] - for start, end, g_idx in zip(starts, ends, group_indices, strict=True): + shards_per_group: list[list[StoreShardId]] = [ + [] for _ in self.token_databases + ] + for start, end, g_idx, store_shard_id in zip( + starts, + ends, + group_indices, + store_shard_ids, + strict=True, + ): chunks_per_group[g_idx].append((start, end)) + shards_per_group[g_idx].append(store_shard_id) for g_idx, chunks in enumerate(chunks_per_group): if not chunks: continue db = self.token_databases[g_idx] - group_addrs, group_sizes, _ = db.prepare_values( - chunks, block_ids_per_group[g_idx] + group_addrs, group_sizes, _ = db.store_layout.prepare_values( + chunks, + block_ids_per_group[g_idx], + shards_per_group[g_idx], ) addrs.extend(group_addrs) sizes.extend(group_sizes) - if self.enable_kv_event: - new_block_hashes = [ - maybe_convert_block_hash(bh) for bh in kv_event_block_hashes - ] - token_ids_end = token_ids_start + len(event_token_ids or ()) - - for idx, (s, e, g_idx) in enumerate( - zip(starts, ends, group_indices, strict=True) - ): - db = self.token_databases[g_idx] - if self.enable_kv_event: - token_ids = ( - event_token_ids[s - token_ids_start : e - token_ids_start] - if event_token_ids is not None - and token_ids_start <= s - and e <= token_ids_end - else [] - ) - stored_event = BlockStored( - block_hashes=[new_block_hashes[idx]], - # Derive the direct predecessor from the unfiltered - # request chain. Adjacent PUTs need not be adjacent in - # that chain after Store dedup, masks, or TP striding. - parent_block_hash=( - maybe_convert_block_hash( - req_meta.block_hashes[s // db.hash_block_size - 1] - ) - if s > 0 - else None - ), - token_ids=token_ids, - block_size=db.block_size, - lora_id=None, - medium="cpu", - lora_name=None, - group_idx=g_idx, - ) - stored_events.append(stored_event) - if current_event is not None: current_event.synchronize() @@ -1112,6 +1113,8 @@ class KVCacheStoreSendingThread(KVTransferThread): assert len(group_ids) == len(keys) self.replicate_config.group_ids = group_ids + failed_indices: set[int] = set() + put_had_exception = False batch_bytes = _sum_batch_bytes(sizes) put_start = time.perf_counter() try: @@ -1121,54 +1124,6 @@ class KVCacheStoreSendingThread(KVTransferThread): sizes, self.replicate_config, ) - failed = [i for i, v in enumerate(res) if v < 0] - self._record_operation( - "save_put", - put_start, - len(keys), - num_bytes=batch_bytes, - status="partial_failure" if failed else "ok", - num_failed_keys=len(failed), - ) - if failed: - failed_codes = set(res[i] for i in failed) - if self.enable_kv_event: - failed_indices = set(failed) - stored_events = [ - event - for i, event in enumerate(stored_events) - if i not in failed_indices - ] - logger.warning( - "batch_put failed: %d/%d keys failed " - "(codes=%s, batch_bytes=%d, num_keys=%d), " - "first_key=%s", - len(failed), - len(keys), - failed_codes, - batch_bytes, - len(keys), - keys[0] if keys else "N/A", - ) - if ( - MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes - and not self._mark_request_skipped_for_pressure(req_meta) - ): - logger.warning( - "Detected Mooncake CPU/disk offloading pressure " - "(NO_AVAILABLE_HANDLE); skipping future store " - "batches for request %s until a later store " - "batch succeeds", - req_id, - ) - else: - self._record_saved(req_meta, token_len) - save_completed = True - if self._clear_store_pressure(): - logger.info( - "Mooncake CPU/disk offloading pressure cleared " - "after a successful store batch" - ) except Exception as e: self._record_operation( "save_put", @@ -1179,7 +1134,96 @@ class KVCacheStoreSendingThread(KVTransferThread): num_failed_keys=len(keys), ) logger.error("Failed to put key %s, error: %s", keys, e) - stored_events.clear() + put_had_exception = True + else: + failed_indices = {i for i, value in enumerate(res) if value < 0} + self._record_operation( + "save_put", + put_start, + len(keys), + num_bytes=batch_bytes, + status="partial_failure" if failed_indices else "ok", + num_failed_keys=len(failed_indices), + ) + failed_codes = {res[i] for i in failed_indices} + if failed_indices: + logger.warning( + "batch_put failed: %d/%d keys failed " + "(codes=%s, batch_bytes=%d), first_key=%s", + len(failed_indices), + len(keys), + failed_codes, + batch_bytes, + keys[0], + ) + if ( + MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes + and not self._mark_request_skipped_for_pressure(req_meta) + ): + logger.warning( + "Detected Mooncake CPU/disk offloading pressure " + "(NO_AVAILABLE_HANDLE); skipping future store " + "batches for request %s until a later store batch succeeds", + req_id, + ) + + if not put_had_exception and not failed_indices: + self._record_saved(req_meta, token_len) + save_completed = True + if self._clear_store_pressure(): + logger.info( + "Mooncake CPU/disk offloading pressure cleared " + "after a successful store batch" + ) + + stored_events: list[BlockStored] = [] + if self.enable_kv_event and not put_had_exception: + assert event_specs is not None + # BlockStored is a logical-block event, while one block may map + # to several Store shards. Emit once only after every missing + # shard for that block succeeded. Shards that already existed + # were removed before this mapping and are already satisfied. + indices_by_event: dict[tuple[int, int, int, BlockHash], list[int]] = {} + for index, event_spec in enumerate(event_specs): + indices_by_event.setdefault(event_spec, []).append(index) + + token_ids_end = token_ids_start + len(event_token_ids or ()) + for ( + s, + end, + g_idx, + block_hash, + ), event_indices in indices_by_event.items(): + if any(index in failed_indices for index in event_indices): + continue + db = self.token_databases[g_idx] + token_ids = ( + event_token_ids[s - token_ids_start : end - token_ids_start] + if event_token_ids is not None + and token_ids_start <= s + and end <= token_ids_end + else [] + ) + stored_events.append( + BlockStored( + block_hashes=[maybe_convert_block_hash(block_hash)], + # Store filtering can separate adjacent request + # blocks, so derive the predecessor from the request. + parent_block_hash=( + maybe_convert_block_hash( + req_meta.block_hashes[s // db.hash_block_size - 1] + ) + if s > 0 + else None + ), + token_ids=token_ids, + block_size=db.block_size, + lora_id=None, + medium="cpu", + lora_name=None, + group_idx=g_idx, + ) + ) if self.enable_kv_event and stored_events: self.update_kv_event(stored_events) @@ -1269,6 +1313,7 @@ class KVCacheStoreRecvingThread(KVTransferThread): for g_idx, db in enumerate(self.token_databases): mask = load_mask_per_group[g_idx] chunks: list[tuple[int, int]] = [] + store_shard_ids: list[StoreShardId] = [] for start, end, block_hash in db.process_tokens( token_len, req_meta.block_hashes, mask_num ): @@ -1282,10 +1327,14 @@ class KVCacheStoreRecvingThread(KVTransferThread): block_hash = req_meta.block_hashes[ boundary_tokens // db.hash_block_size - 1 ] - key_list.append(db.key_for(block_hash)) - chunks.append((start, end)) - g_addrs, g_sizes, g_block_ids = db.prepare_values( - chunks, req_meta.block_ids[g_idx] + for store_shard_id in db.store_layout.local_shard_ids: + key_list.append(db.store_layout.key_for(store_shard_id, block_hash)) + chunks.append((start, end)) + store_shard_ids.append(store_shard_id) + g_addrs, g_sizes, g_block_ids = db.store_layout.prepare_values( + chunks, + req_meta.block_ids[g_idx], + store_shard_ids, ) addr_list.extend(g_addrs) size_list.extend(g_sizes) @@ -1298,7 +1347,14 @@ class KVCacheStoreRecvingThread(KVTransferThread): size_list_c = _rotate_list(size_list, rotation) block_id_list_c = _rotate_list(block_id_list, rotation) - load_batches = [(key_list_c, addr_list_c, size_list_c, block_id_list_c)] + load_batches = [ + ( + key_list_c, + addr_list_c, + size_list_c, + block_id_list_c, + ) + ] if self.usable_disk_offload_buffer_budget_bytes is not None: total_staging_bytes = sum( _estimate_disk_offload_staging_bytes(size) for size in size_list_c @@ -1340,7 +1396,12 @@ class KVCacheStoreRecvingThread(KVTransferThread): block_id_offset:next_block_id_offset ] load_batches.append( - (batch_keys, batch_addrs, batch_sizes, batch_block_ids) + ( + batch_keys, + batch_addrs, + batch_sizes, + batch_block_ids, + ) ) block_id_offset = next_block_id_offset @@ -1348,7 +1409,12 @@ class KVCacheStoreRecvingThread(KVTransferThread): current_batch_block_ids: list[int] = block_id_list_c batch_bytes = 0 try: - for batch_keys, batch_addrs, batch_sizes, batch_block_ids in load_batches: + for ( + batch_keys, + batch_addrs, + batch_sizes, + batch_block_ids, + ) in load_batches: current_batch_keys = batch_keys current_batch_block_ids = batch_block_ids batch_bytes = _sum_batch_bytes(batch_sizes) @@ -1607,11 +1673,9 @@ class MooncakeStoreWorker: retention_interval=kv_cache_config.prefix_cache_retention_interval, dcp_world_size=self.dcp_size, ) - # One ChunkedTokenDatabase per group; addresses populated in - # register_kv_caches once the kv-cache layout is known. Each group's - # key namespace is its TP shard id: ranks holding identical bytes - # (MLA / shared GQA KV heads) share a namespace, TP-sharded Mamba - # state gets one namespace per rank. + self.store_tp_size, store_namespace, store_layout_cls = ( + self._select_store_layout(extra_config) + ) metadata = KeyMetadata( model_name=model_config.model.rstrip("/").split("/")[-1], tp_rank=self.tp_rank, @@ -1623,24 +1687,140 @@ class MooncakeStoreWorker: "cache_prefix", "" ) ), + store_namespace=store_namespace, ) self._group_tp_replication_factors: tuple[int, ...] = ( self._compute_group_tp_replication_factors() ) - self.token_dbs = [ - ChunkedTokenDatabase( - dataclasses.replace( - metadata, - group_id=g_idx, - tp_rank=self.tp_rank // self._group_tp_replication_factors[g_idx], - ), - g.kv_cache_spec.block_size, - hash_block_size=self.hash_block_size, - ) - for g_idx, g in enumerate(self._kv_cache_groups) - ] + self.token_dbs = self._build_token_databases(metadata, store_layout_cls) self._init_lookup_key_prefixes() + def _supports_tp_sharded_store_layout( + self, + layout_cls: type[TPShardedStoreLayout] | None, + extra_config: dict[str, Any], + ) -> bool: + if ( + layout_cls is None + or self.pcp_size != 1 + or self.dcp_size != 1 + or len(self._kv_cache_groups) != 1 + ): + return False + + # Subclasses may use different sharding or cache-lifetime semantics. + return ( + type(self._kv_cache_groups[0].kv_cache_spec) is FullAttentionSpec + and str(extra_config.get("enable_cross_layers_blocks", "False")).lower() + != "true" + ) + + def _select_store_layout( + self, extra_config: dict[str, Any] + ) -> tuple[int | None, str, type[TPShardedStoreLayout] | None]: + """Select the opt-in TP layout and its Store namespace.""" + lcm_store_tp_enabled = extra_config.get("enable_store_tp_lcm") is True + store_tp_requested = ( + lcm_store_tp_enabled or extra_config.get("store_tp_size") is not None + ) + if not store_tp_requested: + return None, "", None + + requested_store_tp_size = resolve_store_tp_size(extra_config) + cache_layout = self.cache_config.get_resolved_kv_cache_layout() + layout_cls: type[TPShardedStoreLayout] | None = { + KVCacheLayout.LBHNC: LBHNCStoreLayout, + KVCacheLayout.LBNHC: LBNHCStoreLayout, + }.get(cache_layout) + + if ( + requested_store_tp_size is not None + and requested_store_tp_size >= self.tp_size + and requested_store_tp_size % self.tp_size == 0 + and self._supports_tp_sharded_store_layout(layout_cls, extra_config) + ): + assert layout_cls is not None + if self.num_kv_head % requested_store_tp_size == 0: + if layout_cls is LBNHCStoreLayout: + logger.warning_once( + "Mooncake Store TP sharding is using the LBNHC (NHC) " + "KV cache layout, which creates many transfer segments " + "and may significantly reduce PUT/GET performance. Use " + "LBHNC (HNC) when supported by the attention backend." + ) + return ( + requested_store_tp_size, + layout_cls.shared_namespace(requested_store_tp_size, self.pp_size), + layout_cls, + ) + if self.num_kv_head == 1: + logger.info( + "Mooncake heterogeneous-TP store sharing uses the replicated " + "MQA layout for store_tp_size=%d", + requested_store_tp_size, + ) + return ( + None, + f"@store_pp:{self.pp_size}@store_format:tp_shared_mqa", + None, + ) + + store_namespace = ( + f"@store_pp:{self.pp_size}@store_format:" + f"rank_local_tp{self.tp_size}_layout_{cache_layout.name}" + ) + requested_topology = ( + extra_config.get("prefill_tp_sizes") + if lcm_store_tp_enabled + else extra_config.get("store_tp_size") + ) + logger.warning( + "Mooncake heterogeneous-TP store sharing is disabled for " + "Store TP configuration %r with KV layout %s; using a " + "compatibility-namespaced rank-local store layout", + requested_topology, + cache_layout, + ) + return None, store_namespace, None + + def _build_token_databases( + self, + metadata: KeyMetadata, + layout_cls: type[TPShardedStoreLayout] | None, + ) -> list[ChunkedTokenDatabase]: + """Construct token databases and their Store layouts.""" + token_dbs: list[ChunkedTokenDatabase] = [] + for group_idx, group in enumerate(self._kv_cache_groups): + group_tp_rank = self.tp_rank + if layout_cls is None: + group_tp_rank //= self._group_tp_replication_factors[group_idx] + group_metadata = dataclasses.replace( + metadata, + group_id=group_idx, + tp_rank=group_tp_rank, + ) + store_layout: TPShardedStoreLayout | None = None + if layout_cls is not None: + assert self.store_tp_size is not None + store_layout = layout_cls( + group_metadata, + group.kv_cache_spec.block_size, + self.hash_block_size, + local_tp_size=self.tp_size, + store_tp_size=self.store_tp_size, + tp_rank=self.tp_rank, + num_kv_heads=self.num_kv_head, + ) + token_dbs.append( + ChunkedTokenDatabase( + group_metadata, + group.kv_cache_spec.block_size, + hash_block_size=self.hash_block_size, + store_layout=store_layout, + ) + ) + return token_dbs + def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int: if self.dcp_size > 1: return 1 @@ -1689,17 +1869,8 @@ class MooncakeStoreWorker: ) self._lookup_key_prefixes = tuple( - tuple( - PoolKey.build_prefix( - db.metadata, - tp_rank=tp_rank, - pcp_rank=pcp_rank, - dcp_rank=dcp_rank, - pp_rank=pp_rank, - ) - for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces( - self._group_tp_replication_factors[g_idx] - ) + db.store_layout.lookup_key_prefixes( + rank_namespaces(self._group_tp_replication_factors[g_idx]) ) for g_idx, db in enumerate(self.token_dbs) ) @@ -1719,12 +1890,11 @@ class MooncakeStoreWorker: self.num_blocks = self.cache_config.num_gpu_blocks seen_storage_ptrs: set[int] = set() - seen_region_ptrs: set[int] = set() - addrs: list[int] = [] - block_lens: list[int] = [] + cache_tensors: list[torch.Tensor] = [] for cache in kv_caches.values(): cache = group_kernel_blocks(cache, self.num_blocks) + cache_tensors.append(cache) cache_storage = cache.untyped_storage() base_addr = cache_storage.data_ptr() region_len = cache_storage.nbytes() @@ -1740,48 +1910,15 @@ class MooncakeStoreWorker: ret, ) - if not is_non_overlapping_and_dense(cache[0]): - # A block is scattered across per-head regions; each region's - # blocks are contiguous. - for head_idx in range(cache.shape[1]): - head_cache = cache[:, head_idx] - assert is_non_overlapping_and_dense(head_cache[0]) - region_addr = head_cache.data_ptr() - if region_addr in seen_region_ptrs: - continue - seen_region_ptrs.add(region_addr) - addrs.append(region_addr) - block_lens.append(head_cache.stride(0) * head_cache.element_size()) - elif cache.stride(0) * cache.element_size() * self.num_blocks == region_len: - # The block stride spans the whole per-block window - # (block-outermost and packed layouts, and single-layer - # tensors), which may hold other layers' pages at higher - # offsets. Register the storage once as one whole-window - # region; per-layer regions would copy overlapping windows and - # run past the storage for offset layers. - if base_addr in seen_region_ptrs: - continue - seen_region_ptrs.add(base_addr) - addrs.append(base_addr) - block_lens.append(region_len // self.num_blocks) - else: - region_addr = cache.data_ptr() - if region_addr in seen_region_ptrs: - continue - seen_region_ptrs.add(region_addr) - addrs.append(region_addr) - block_lens.append(cache.stride(0) * cache.element_size()) - logger.info( - "Registered KV caches: num_groups=%d, num_segments=%d, num_blocks=%d", + "Registered KV caches: num_groups=%d, num_tensors=%d, num_blocks=%d", len(self.token_dbs), - len(addrs), + len(cache_tensors), self.num_blocks, ) for db in self.token_dbs: - db.set_kv_caches_base_addr(addrs) - db.set_block_len(block_lens) + db.store_layout.register_kv_caches(cache_tensors, self.num_blocks) # Start transfer threads if self.can_put: @@ -2128,6 +2265,10 @@ class MooncakeStoreWorker: return self.kv_send_thread.get_kv_events() return [] + @property + def group_tp_replication_factors(self) -> tuple[int, ...]: + return self._group_tp_replication_factors + def close(self) -> None: """Release the MooncakeDistributedStore handle on teardown.