@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70052fb924

[Bugfix][KV Connector][Mooncake] Keep TP-sharded Mamba state out of the KV-head dedup (#49499)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com>

Yifan Qiao committed Jul 24, 2026 at 21:02 UTC 70052fb924fbd633580308562d43b9df838ddf6d

3 files changed +311 -93

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

+1 -2

@@ -173,7 +173,6 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

173 worker = _build_worker_with_dict_store(vllm_config, cfg, store)

174 worker.tp_size = 1

175 worker.pp_size = 1

176 - worker.put_step = 1

176 worker.num_kv_head = 8

177

178 # Register kv_caches using mocked thread classes so register_kv_caches

@@ -215,7 +214,7 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

214 block_size=worker.block_size,

215 coord=worker.coord,

216 tp_rank=worker.tp_rank,

218 - put_step=worker.put_step,

217 + group_put_steps=worker._group_tp_replication_factors,

218 kv_role=worker.kv_role,

219 ready_event=ready,

220 enable_kv_event=False,

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+213 -13

@@ -90,7 +90,7 @@ def _make_store_sending_thread(

90 block_size=block_size,

91 coord=coord,

92 tp_rank=tp_rank,

93 - put_step=put_step,

93 + group_put_steps=[put_step] * len(token_databases),

94 kv_role="kv_producer",

95 ready_event=threading.Event(),

96 replicate_config=replicate_config,

@@ -527,6 +527,29 @@ def test_store_sending_thread_delta_strides_with_local_phase():

527 assert store.batch_put_from_multi_buffers.call_args.args[0] == keys

528

529

530 +def test_tp_sharded_group_saves_every_block_on_every_rank():

531 + """Sharded ranks must write every block because peers hold different bytes."""

532 + store = MagicMock()

533 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

534 + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)

535 + thread = _make_store_sending_thread(store, tp_rank=0, put_step=2)

536 + thread.group_put_steps = [1]

537 +

538 + thread.add_stored_request("req-a")

539 + thread._handle_request(

540 + ReqMeta(

541 + req_id="req-a",

542 + token_len_chunk=64,

543 + block_ids=([0, 1, 2, 3],),

544 + block_hashes=[b"a0", b"a1", b"a2", b"a3"],

545 + can_save=True,

546 + )

547 + )

548 +

549 + keys = store.batch_is_exist.call_args.args[0]

550 + assert len(keys) == 4

551 +

552 +

553 def test_store_sending_thread_retries_skipped_range_after_pressure():

554 store = MagicMock()

555 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

@@ -1274,7 +1297,8 @@ def test_worker_put_striding_covers_every_rank_get_namespace(

1297 ]

1298 assert len(keys) == len(block_hashes)

1299 # PUT side: mirrors KVCacheStoreSendingThread's striding slice.

1277 - put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step])

1300 + put_step = w._group_tp_replication_factors[0]

1301 + put_keys.update(keys[w.tp_rank % put_step :: put_step])

1302 # GET side: KVCacheStoreRecvingThread fetches every key.

1303 get_keys_per_rank[tp_rank] = set(keys)

1304

@@ -1613,6 +1637,15 @@ def _register_with_mocked_threads(

1637 worker.register_kv_caches(kv_caches)

1638

1639

1640 +def _refresh_group_tp_replication_factors(

1641 + worker: mooncake_store_worker.MooncakeStoreWorker,

1642 +) -> None:

1643 + worker._group_tp_replication_factors = (

1644 + worker._compute_group_tp_replication_factors()

1645 + )

1646 + worker._init_lookup_key_prefixes()

1647 +

1648 +

1649 def _make_bare_worker(

1650 *,

1651 num_gpu_blocks: int = 10,

@@ -1630,11 +1663,9 @@ def _make_bare_worker(

1663 worker.cache_config.num_gpu_blocks = num_gpu_blocks

1664 worker.store = MagicMock()

1665 worker.store.register_buffer.return_value = 0

1633 - worker.use_mla = False

1666 worker.kv_role = kv_role

1667 worker.block_size = block_size

1668 worker.tp_rank = 0

1637 - worker.put_step = 1

1669 worker.enable_kv_events = False

1670 worker.kv_send_thread = None

1671 worker.kv_recv_threads = []

@@ -1664,7 +1695,6 @@ def _make_bare_worker(

1695 worker.pcp_size = 1

1696 worker.dcp_size = 1

1697 worker.hash_block_size = block_size

1667 - worker.metadata = KeyMetadata("test-model", 0, 0, 0, 0)

1698 # Pre-build a single-group token_dbs so lookup-only tests don't have to

1699 # call register_kv_caches.

1700 worker.token_dbs = [

@@ -1679,7 +1709,7 @@ def _make_bare_worker(

1709 scheduler_block_size=block_size,

1710 hash_block_size=block_size,

1711 )

1682 - worker._init_lookup_key_prefixes()

1712 + _refresh_group_tp_replication_factors(worker)

1713 return worker

1714

1715

@@ -1688,9 +1718,8 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces():

1718 worker.tp_size = 4

1719 worker.num_kv_head = 1

1720 worker.dcp_size = 4

1691 - worker._init_lookup_key_prefixes()

1721 + _refresh_group_tp_replication_factors(worker)

1722

1693 - assert worker._lookup_expected_per_key == 4

1723 assert worker._lookup_key_prefixes[0] == (

1724 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",

1725 "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0",

@@ -1705,21 +1734,192 @@ def test_lookup_key_prefixes_cover_pcp_rank_namespaces():

1734 worker.num_kv_head = 1

1735 worker.pcp_size = 2

1736 worker.dcp_size = 1

1708 - worker._init_lookup_key_prefixes()

1737 + _refresh_group_tp_replication_factors(worker)

1738

1710 - assert worker._lookup_expected_per_key == 2

1739 assert worker._lookup_key_prefixes[0] == (

1740 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",

1741 "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0",

1742 )

1743

1744

1745 +def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank():

1746 + """Replicated attention needs one namespace; sharded Mamba needs every rank."""

1747 + from vllm.v1.kv_cache_interface import (

1748 + FullAttentionSpec,

1749 + KVCacheGroupSpec,

1750 + MambaSpec,

1751 + )

1752 +

1753 + worker = _make_bare_worker(block_size=16)

1754 + worker.tp_size = 2

1755 + worker.num_kv_head = 1

1756 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

1757 + mamba = MambaSpec(

1758 + block_size=16,

1759 + shapes=((1, 1),),

1760 + dtypes=(torch.float32,),

1761 + mamba_cache_mode="align",

1762 + )

1763 + worker._kv_cache_groups = [

1764 + KVCacheGroupSpec(["l0"], fa),

1765 + KVCacheGroupSpec(["l1"], mamba),

1766 + ]

1767 + worker.token_dbs = [

1768 + ChunkedTokenDatabase(

1769 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16

1770 + ),

1771 + ChunkedTokenDatabase(

1772 + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16

1773 + ),

1774 + ]

1775 + _refresh_group_tp_replication_factors(worker)

1776 +

1777 + assert worker._lookup_key_prefixes[0] == (

1778 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",

1779 + )

1780 + assert worker._lookup_key_prefixes[1] == (

1781 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1",

1782 + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1",

1783 + )

1784 +

1785 +

1786 +def test_group_tp_replication_factors_mixed_mla_gqa_mamba():

1787 + from vllm.v1.kv_cache_interface import (

1788 + FullAttentionSpec,

1789 + KVCacheGroupSpec,

1790 + MambaSpec,

1791 + MLAAttentionSpec,

1792 + )

1793 +

1794 + worker = _make_bare_worker(block_size=16)

1795 + worker.tp_size = 4

1796 + worker.num_kv_head = 2

1797 + mla = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=None)

1798 + gqa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

1799 + mamba = MambaSpec(

1800 + block_size=16,

1801 + shapes=((1, 1),),

1802 + dtypes=(torch.float32,),

1803 + mamba_cache_mode="align",

1804 + )

1805 + worker._kv_cache_groups = [

1806 + KVCacheGroupSpec(["l0"], mla),

1807 + KVCacheGroupSpec(["l1"], gqa),

1808 + KVCacheGroupSpec(["l2"], mamba),

1809 + ]

1810 + worker.token_dbs = [

1811 + ChunkedTokenDatabase(

1812 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=g_idx), block_size=16

1813 + )

1814 + for g_idx in range(3)

1815 + ]

1816 +

1817 + _refresh_group_tp_replication_factors(worker)

1818 + assert worker._group_tp_replication_factors == (4, 2, 1)

1819 + assert worker._lookup_key_prefixes[0] == (

1820 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",

1821 + )

1822 + assert worker._lookup_key_prefixes[1] == (

1823 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1",

1824 + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1",

1825 + )

1826 + assert worker._lookup_key_prefixes[2] == (

1827 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:2",

1828 + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:2",

1829 + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:2",

1830 + "test-model@tp_rank:3@pcp0@dcp0@pp_rank:0@group:2",

1831 + )

1832 +

1833 +

1834 +@pytest.mark.parametrize("spec_order", [("mla", "gqa"), ("gqa", "mla")])

1835 +def test_uniform_group_uses_common_inner_replication_factor(spec_order):

1836 + from vllm.v1.kv_cache_interface import (

1837 + FullAttentionSpec,

1838 + KVCacheGroupSpec,

1839 + MLAAttentionSpec,

1840 + UniformTypeKVCacheSpecs,

1841 + )

1842 +

1843 + worker = _make_bare_worker(block_size=16)

1844 + worker.tp_size = 4

1845 + worker.num_kv_head = 2

1846 + specs_by_name = {

1847 + "mla": MLAAttentionSpec(

1848 + block_size=16, num_kv_heads=1, head_size=64, dtype=None

1849 + ),

1850 + "gqa": FullAttentionSpec(

1851 + block_size=16, num_kv_heads=1, head_size=64, dtype=None

1852 + ),

1853 + }

1854 + inner_specs = {name: specs_by_name[name] for name in spec_order}

1855 + uniform_spec = UniformTypeKVCacheSpecs(

1856 + block_size=16,

1857 + kv_cache_specs=inner_specs,

1858 + )

1859 + worker._kv_cache_groups = [

1860 + KVCacheGroupSpec(list(inner_specs), uniform_spec),

1861 + ]

1862 +

1863 + _refresh_group_tp_replication_factors(worker)

1864 +

1865 + assert worker._group_tp_replication_factors == (2,)

1866 + assert worker._lookup_key_prefixes[0] == (

1867 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0",

1868 + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:0",

1869 + )

1870 +

1871 +

1872 +def test_lookup_rejects_boundary_missing_one_mamba_shard():

1873 + from vllm.v1.kv_cache_interface import (

1874 + FullAttentionSpec,

1875 + KVCacheGroupSpec,

1876 + MambaSpec,

1877 + )

1878 +

1879 + worker = _make_bare_worker(block_size=16)

1880 + worker.tp_size = 2

1881 + worker.num_kv_head = 1

1882 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

1883 + mamba = MambaSpec(

1884 + block_size=16,

1885 + shapes=((1, 1),),

1886 + dtypes=(torch.float32,),

1887 + mamba_cache_mode="align",

1888 + )

1889 + worker._kv_cache_groups = [

1890 + KVCacheGroupSpec(["l0"], fa),

1891 + KVCacheGroupSpec(["l1"], mamba),

1892 + ]

1893 + worker.token_dbs = [

1894 + ChunkedTokenDatabase(

1895 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16

1896 + ),

1897 + ChunkedTokenDatabase(

1898 + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16

1899 + ),

1900 + ]

1901 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(

1902 + worker._kv_cache_groups,

1903 + scheduler_block_size=16,

1904 + hash_block_size=16,

1905 + )

1906 + _refresh_group_tp_replication_factors(worker)

1907 +

1908 + worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys)

1909 + assert worker.lookup(32, [b"h0", b"h1"]) == 32

1910 +

1911 + worker.store.batch_is_exist.side_effect = lambda keys: [

1912 + 0 if "tp_rank:1" in k and "group:1" in k else 1 for k in keys

1913 + ]

1914 + assert worker.lookup(32, [b"h0", b"h1"]) == 0

1915 +

1916 +

1917 def test_lookup_requires_all_dcp_rank_namespaces():

1918 worker = _make_bare_worker(block_size=16)

1919 worker.tp_size = 4

1920 worker.num_kv_head = 1

1921 worker.dcp_size = 4

1722 - worker._init_lookup_key_prefixes()

1922 + _refresh_group_tp_replication_factors(worker)

1923 worker.store.batch_is_exist.return_value = [1, 1, 0, 1]

1924

1925 assert worker.lookup(16, [b"a0"]) == 0

@@ -1850,7 +2050,7 @@ def test_lookup_checks_all_potential_swa_hit_boundaries():

2050 hash_block_size=8,

2051 retention_interval=0,

2052 )

1853 - worker._init_lookup_key_prefixes()

2053 + _refresh_group_tp_replication_factors(worker)

2054 # Candidate order: 3 full-attention chunks, then SWA chunks 3, 7, 11.

2055 # Only the first full chunk and the SWA chunk ending at token 32 exist, so

2056 # lookup should recover a 32-token external prefix hit. A sparse

@@ -1911,7 +2111,7 @@ def test_lookup_applies_swa_mask_before_accessing_hashes():

2111 hash_block_size=8,

2112 retention_interval=0,

2113 )

1914 - worker._init_lookup_key_prefixes()

2114 + _refresh_group_tp_replication_factors(worker)

2115

2116 block_hashes = _RecordingBlockHashes([f"h{i}".encode() for i in range(12)])

2117 accessed_before_rpc: list[int] = []

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+97 -78

@@ -66,7 +66,15 @@ from vllm.v1.core.kv_cache_utils import (

66 maybe_convert_block_hash,

67 resolve_kv_cache_block_sizes,

68 )

69 -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec

69 +from vllm.v1.kv_cache_interface import (

70 + KVCacheConfig,

71 + KVCacheGroupSpec,

72 + KVCacheSpec,

73 + MambaSpec,

74 + MLAAttentionSpec,

75 + SlidingWindowMLASpec,

76 + UniformTypeKVCacheSpecs,

77 +)

79 from .metrics import MooncakeStoreConnectorStats

@@ -448,7 +456,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

456 token_databases: list[ChunkedTokenDatabase],

457 block_size: int,

458 tp_rank: int,

451 - put_step: int,

459 + group_put_steps: Sequence[int],

460 kv_role: str,

461 ready_event: threading.Event,

462 enable_kv_event: bool = False,

@@ -464,7 +472,8 @@ class KVCacheStoreSendingThread(KVTransferThread):

472 name="KVCacheStoreSendingThread",

473 record_operation=record_operation,

474 )

467 - self.put_step = put_step

475 + # Only ranks with identical group bytes may stripe PUTs (e.g., MLA).

476 + self.group_put_steps = group_put_steps

477 self.coord = coord

478 self.kv_role = kv_role

479 self.stored_requests: defaultdict[str, int] = defaultdict(int)

@@ -568,13 +577,14 @@ class KVCacheStoreSendingThread(KVTransferThread):

577 group_indices: list[int] = []

578 for g_idx, db in enumerate(self.token_databases):

579 # Rotate the stride phase per group to balance load across ranks.

571 - put_step_rank = (self.tp_rank + g_idx) % self.put_step

580 + put_step = self.group_put_steps[g_idx]

581 + put_step_rank = (self.tp_rank + g_idx) % put_step

582 for start, end, block_hash in db.process_tokens(

583 token_len,

584 req_meta.block_hashes,

585 mask_num=save_start,

586 chunk_mask=store_masks[g_idx],

577 - put_step=self.put_step,

587 + put_step=put_step,

588 put_step_rank=put_step_rank,

589 ):

590 starts.append(start)

@@ -1002,44 +1012,7 @@ class MooncakeStoreWorker:

1012 )

1013 self.num_layers = model_config.get_num_layers(parallel_config)

1014

1005 - self.use_mla = False

1006 - if (

1007 - hasattr(model_config, "use_mla")

1008 - and isinstance(model_config.use_mla, bool)

1009 - and model_config.use_mla

1010 - ):

1011 - self.use_mla = True

1012 -

1013 - if self.use_mla:

1014 - self.num_kv_head = 1

1015 - else:

1016 - self.num_kv_head = model_config.get_total_num_kv_heads()

1017 -

1018 - if self.num_kv_head < self.tp_size and self.dcp_size <= 1:

1019 - # Dedup: TP ranks holding the same KV heads stripe PUTs across

1020 - # one shared key namespace. DCP splits the TP group, so with

1021 - # DCP>1 those ranks have different `@dcpN` namespaces and

1022 - # striping would leave keys unwritten (OBJECT_NOT_FOUND on

1023 - # GET). PCP is outer to TP (pcp_rank is constant within a TP

1024 - # group), so it needs no guard.

1025 - self.put_step = self.tp_size // self.num_kv_head

1026 - self.head_or_tp_rank = self.tp_rank // self.put_step

1027 - else:

1028 - self.head_or_tp_rank = self.tp_rank

1029 - self.put_step = 1

1030 -

1031 - self.metadata = KeyMetadata(

1032 - model_name=model_config.model.rstrip("/").split("/")[-1],

1033 - tp_rank=self.head_or_tp_rank,

1034 - pcp_rank=self.pcp_rank,

1035 - dcp_rank=self.dcp_rank,

1036 - pp_rank=self.pp_rank,

1037 - cache_prefix=str(

1038 - vllm_config.kv_transfer_config.kv_connector_extra_config.get(

1039 - "cache_prefix", ""

1040 - )

1041 - ),

1042 - )

1015 + self.num_kv_head = model_config.get_total_num_kv_heads()

1016

1017 # Initialize MooncakeDistributedStore with its own TransferEngine

1018 store_config = MooncakeStoreConfig.load_from_config()

@@ -1156,10 +1129,32 @@ class MooncakeStoreWorker:

1129 retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL,

1130 )

1131 # One ChunkedTokenDatabase per group; addresses populated in

1159 - # register_kv_caches once the kv-cache layout is known.

1132 + # register_kv_caches once the kv-cache layout is known. Each group's

1133 + # key namespace is its TP shard id: ranks holding identical bytes

1134 + # (MLA / shared GQA KV heads) share a namespace, TP-sharded Mamba

1135 + # state gets one namespace per rank.

1136 + metadata = KeyMetadata(

1137 + model_name=model_config.model.rstrip("/").split("/")[-1],

1138 + tp_rank=self.tp_rank,

1139 + pcp_rank=self.pcp_rank,

1140 + dcp_rank=self.dcp_rank,

1141 + pp_rank=self.pp_rank,

1142 + cache_prefix=str(

1143 + vllm_config.kv_transfer_config.kv_connector_extra_config.get(

1144 + "cache_prefix", ""

1145 + )

1146 + ),

1147 + )

1148 + self._group_tp_replication_factors: tuple[int, ...] = (

1149 + self._compute_group_tp_replication_factors()

1150 + )

1151 self.token_dbs: list[ChunkedTokenDatabase] = [

1152 ChunkedTokenDatabase(

1162 - dataclasses.replace(self.metadata, group_id=g_idx),

1153 + dataclasses.replace(

1154 + metadata,

1155 + group_id=g_idx,

1156 + tp_rank=self.tp_rank // self._group_tp_replication_factors[g_idx],

1157 + ),

1158 g.kv_cache_spec.block_size,

1159 hash_block_size=self.hash_block_size,

1160 )

@@ -1167,27 +1162,50 @@ class MooncakeStoreWorker:

1162 ]

1163 self._init_lookup_key_prefixes()

1164

1170 - def _init_lookup_key_prefixes(self) -> None:

1171 - """Prepare per-group key prefixes across parallel rank namespaces."""

1172 - # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces

1165 + def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int:

1166 if self.dcp_size > 1:

1174 - # DCP reuses the TP workers and splits each TP group into

1175 - # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size.

1176 - # Store/load paths do not apply KV-head dedup under DCP

1177 - rank_namespaces = tuple(

1178 - (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank)

1179 - for pcp_rank in range(self.pcp_size)

1180 - for tp_rank in range(self.tp_size)

1181 - for pp_rank in range(self.pp_size)

1182 - )

1183 - else:

1184 - # Without DCP, TP ranks that share a KV head write identical KV, so

1185 - # lookup only needs one TP namespace per unique KV head.

1186 - tp_count = min(self.tp_size, self.num_kv_head)

1187 - rank_namespaces = tuple(

1188 - (tp_rank, pcp_rank, 0, pp_rank)

1167 + return 1

1168 + inner_specs = (

1169 + tuple(spec.kv_cache_specs.values())

1170 + if isinstance(spec, UniformTypeKVCacheSpecs)

1171 + else (spec,)

1172 + )

1173 + # Any rank-specific state makes the whole packed value rank-specific.

1174 + if any(isinstance(inner, MambaSpec) for inner in inner_specs):

1175 + return 1

1176 + # A pure MLA packed value is replicated on every TP rank.

1177 + if all(

1178 + isinstance(inner, (MLAAttentionSpec, SlidingWindowMLASpec))

1179 + for inner in inner_specs

1180 + ):

1181 + return self.tp_size

1182 + return max(1, self.tp_size // self.num_kv_head)

1183 +

1184 + def _compute_group_tp_replication_factors(self) -> tuple[int, ...]:

1185 + """Return the number of byte-identical TP replicas per cache group.

1186 +

1187 + DCP and Mamba use 1; MLA uses ``tp_size``; GQA uses

1188 + ``tp_size // num_kv_head``.

1189 + """

1190 + return tuple(

1191 + self._spec_tp_replication_factor(group.kv_cache_spec)

1192 + for group in self._kv_cache_groups

1193 + )

1194 +

1195 + def _init_lookup_key_prefixes(self) -> None:

1196 + def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]:

1197 + if self.dcp_size > 1:

1198 + # DCP is a TP subdivision: dcp_rank == tp_rank % dcp_size.

1199 + return tuple(

1200 + (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank)

1201 + for pcp_rank in range(self.pcp_size)

1202 + for tp_rank in range(self.tp_size)

1203 + for pp_rank in range(self.pp_size)

1204 + )

1205 + return tuple(

1206 + (shard_rank, pcp_rank, 0, pp_rank)

1207 for pcp_rank in range(self.pcp_size)

1190 - for tp_rank in range(tp_count)

1208 + for shard_rank in range(self.tp_size // factor)

1209 for pp_rank in range(self.pp_size)

1210 )

1211

@@ -1200,11 +1218,12 @@ class MooncakeStoreWorker:

1218 dcp_rank=dcp_rank,

1219 pp_rank=pp_rank,

1220 )

1203 - for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces

1221 + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces(

1222 + self._group_tp_replication_factors[g_idx]

1223 + )

1224 )

1205 - for db in self.token_dbs

1225 + for g_idx, db in enumerate(self.token_dbs)

1226 )

1207 - self._lookup_expected_per_key = len(rank_namespaces)

1227

1228 def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:

1229 """Register a cross-layers KV cache tensor.

@@ -1297,7 +1316,7 @@ class MooncakeStoreWorker:

1316 self.token_dbs,

1317 self.block_size,

1318 self.tp_rank,

1300 - self.put_step,

1319 + self._group_tp_replication_factors,

1320 self.kv_role,

1321 ready_event_sending,

1322 self.enable_kv_events,

@@ -1527,16 +1546,16 @@ class MooncakeStoreWorker:

1546 logger.error("Remote connection failed in lookup: %s", e)

1547 return 0

1548

1530 - # A (group, hash) is "present" only when every TP*PP rank has it.

1531 - ranks_per_candidate = self._lookup_expected_per_key

1532 - exists_set = {

1533 - (g_idx, hash_bytes)

1534 - for i, (g_idx, hash_bytes) in enumerate(candidate_meta)

1535 - if all(

1536 - res[i * ranks_per_candidate + j] == 1

1537 - for j in range(ranks_per_candidate)

1538 - )

1539 - }

1549 + # A (group, hash) is "present" only when every namespace that will be

1550 + # loaded has it (per-group count: sharded groups need every rank's

1551 + # shard, replicated groups one namespace per unique KV head).

1552 + exists_set = set()

1553 + pos = 0

1554 + for g_idx, hash_bytes in candidate_meta:

1555 + count = len(self._lookup_key_prefixes[g_idx])

1556 + if all(res[pos + j] == 1 for j in range(count)):

1557 + exists_set.add((g_idx, hash_bytes))

1558 + pos += count

1559

1560 cached_block_pool = ExternalCachedBlockPool(

1561 self.hash_block_size,