[KV Connector] Support heterogeneous TP sharing in Mooncake Store Connector (#53129)

Signed-off-by: Guanyi Chen <939416532@qq.com>
Signed-off-by: z-zanez <zhouzh93@mail2.sysu.edu.cn>
Signed-off-by: Zheng Zhou <zhouzh93@mail2.sysu.edu.cn>
Signed-off-by: zane <zhouzh93@mail2.sysu.edu.cn>
Co-authored-by: Guanyi Chen <939416532@qq.com>

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
 

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
 # ============================================================

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)

tests/v1/kv_connector/unit/test_mooncake_store_prepare_values.py

@@ -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])

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

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"<MooncakeStoreKVEvents events={self.get_all_events()}>"
@@ -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
 

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,

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.