From 2675643ced61d02d0c5514bc0fc9dfba9cf294c9 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Wed, 22 Jul 2026 23:38:03 +0000 Subject: [PATCH 1/3] Keep TP-sharded mamba state out of the KV-head dedup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With MLA (or GQA with num_kv_head < tp_size) all TP ranks hold the same attention KV, so the store dedups: ranks share one tp_rank namespace and stripe PUTs across blocks. That assumption is wrong for the hybrid models' Mamba/linear-attention groups, whose state is head/dim-sharded: striping persists each boundary block's state from only one rank, and on a warm hit every other rank silently loads a foreign shard. On Kimi-Linear TP2 this corrupted every external mamba-state load (gsm8k warm 0.72 vs cold 0.87, every answer a plausible paraphrase; per-layer probing shows rank 1's layer-0 initial state wrong while the transported bytes match the stored bytes exactly). Replace the model-wide dedup triple (num_kv_head, put_step, head_or_tp_rank) with a per-group TP replication factor: the number of ranks holding byte-identical bytes for that group (MLA latent KV: tp_size; GQA: tp_size // num_kv_head ranks per shared KV head; Mamba state: 1; any group under DCP: 1). Each group's key namespace is its shard id (tp_rank // factor), PUTs stripe across the ranks within a replication set, and lookup requires one key per distinct shard before reporting a boundary. This produces identical store keys for every currently-supported model — Mamba groups simply stop sharing a namespace and striping — and gives future mixed MLA+GQA+Mamba hybrids the correct per-group treatment instead of a model-wide guess. Cold/warm divergence on a 40-prompt temp=0 repro drops from 38/40 to 0/40; full gsm8k two-pass: cold strict 0.8673, warm strict 0.8726 (was 0.7225), zero KV load failures. Co-authored-by: Claude Signed-off-by: Yifan Qiao --- .../unit/test_mooncake_store_hma_e2e.py | 3 +- .../unit/test_mooncake_store_worker.py | 181 +++++++++++++++++- .../kv_connector/v1/mooncake/store/worker.py | 176 +++++++++-------- 3 files changed, 273 insertions(+), 87 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index f09e0a24729b..40d72b3b5cf1 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -173,7 +173,6 @@ def test_e2e_swa_plus_full_save_then_lookup_hits(): worker = _build_worker_with_dict_store(vllm_config, cfg, store) worker.tp_size = 1 worker.pp_size = 1 - worker.put_step = 1 worker.num_kv_head = 8 # Register kv_caches using mocked thread classes so register_kv_caches @@ -215,7 +214,7 @@ def _fake_thread_init(*args, **kwargs): block_size=worker.block_size, coord=worker.coord, tp_rank=worker.tp_rank, - put_step=worker.put_step, + group_put_steps=worker._group_tp_replication_factors(), kv_role=worker.kv_role, ready_event=ready, enable_kv_event=False, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index f8f43662f26a..d4fe518a69c8 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -90,7 +90,7 @@ def _make_store_sending_thread( block_size=block_size, coord=coord, tp_rank=tp_rank, - put_step=put_step, + group_put_steps=[put_step] * len(token_databases), kv_role="kv_producer", ready_event=threading.Event(), replicate_config=replicate_config, @@ -527,6 +527,33 @@ def test_store_sending_thread_delta_strides_with_local_phase(): assert store.batch_put_from_multi_buffers.call_args.args[0] == keys +def test_tp_sharded_group_saves_every_block_on_every_rank(): + # Mamba state is TP-sharded (different bytes per rank), so the cross-rank + # PUT striping that dedups replicated MLA/GQA KV must not apply: a rank + # skipping a block would leave its shard unwritten, and consumers would + # load another rank's shard (silent state corruption on TP>1 warm hits). + store = MagicMock() + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) + thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) + thread.group_put_steps = [1] + + thread.add_stored_request("req-a") + thread._handle_request( + ReqMeta( + req_id="req-a", + token_len_chunk=64, + block_ids=([0, 1, 2, 3],), + block_hashes=[b"a0", b"a1", b"a2", b"a3"], + can_save=True, + ) + ) + + keys = store.batch_is_exist.call_args.args[0] + # All four blocks written, not the strided half. + assert len(keys) == 4 + + def test_store_sending_thread_retries_skipped_range_after_pressure(): store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) @@ -1212,7 +1239,8 @@ def test_worker_put_striding_covers_every_rank_get_namespace( ] assert len(keys) == len(block_hashes) # PUT side: mirrors KVCacheStoreSendingThread's striding slice. - put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step]) + put_step = w._group_tp_replication_factors()[0] + put_keys.update(keys[w.tp_rank % put_step :: put_step]) # GET side: KVCacheStoreRecvingThread fetches every key. get_keys_per_rank[tp_rank] = set(keys) @@ -1568,11 +1596,9 @@ def _make_bare_worker( worker.cache_config.num_gpu_blocks = num_gpu_blocks worker.store = MagicMock() worker.store.register_buffer.return_value = 0 - worker.use_mla = False worker.kv_role = kv_role worker.block_size = block_size worker.tp_rank = 0 - worker.put_step = 1 worker.enable_kv_events = False worker.kv_send_thread = None worker.kv_recv_threads = [] @@ -1602,7 +1628,6 @@ def _make_bare_worker( worker.pcp_size = 1 worker.dcp_size = 1 worker.hash_block_size = block_size - worker.metadata = KeyMetadata("test-model", 0, 0, 0, 0) # Pre-build a single-group token_dbs so lookup-only tests don't have to # call register_kv_caches. worker.token_dbs = [ @@ -1628,7 +1653,6 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker.dcp_size = 4 worker._init_lookup_key_prefixes() - assert worker._lookup_expected_per_key == 4 assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", "test-model@tp_rank:1@pcp0@dcp1@pp_rank:0@group:0", @@ -1645,13 +1669,156 @@ def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): worker.dcp_size = 1 worker._init_lookup_key_prefixes() - assert worker._lookup_expected_per_key == 2 assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", "test-model@tp_rank:0@pcp1@dcp0@pp_rank:0@group:0", ) +def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank(): + # Under MLA KV-head dedup the replicated attention group probes a single + # shared namespace, but a TP-sharded Mamba group has one shard per rank: + # a boundary is only usable when every rank's shard exists. + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 2 + worker.num_kv_head = 1 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], fa), + KVCacheGroupSpec(["l1"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16 + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 + ), + ] + worker._init_lookup_key_prefixes() + + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + ) + assert worker._lookup_key_prefixes[1] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1", + ) + + +def test_group_tp_replication_factors_mixed_mla_gqa_mamba(): + # Replication is a per-group property: MLA latent KV is replicated on + # every rank, each GQA KV head on tp_size // num_kv_head ranks, Mamba + # state on none. Namespace count per group is tp_size // factor. + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + MLAAttentionSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 2 + mla = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=None) + gqa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], mla), + KVCacheGroupSpec(["l1"], gqa), + KVCacheGroupSpec(["l2"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=g_idx), block_size=16 + ) + for g_idx in range(3) + ] + + assert worker._group_tp_replication_factors() == [4, 2, 1] + + worker._init_lookup_key_prefixes() + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + ) + assert worker._lookup_key_prefixes[1] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:1", + ) + assert worker._lookup_key_prefixes[2] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:2@pcp0@dcp0@pp_rank:0@group:2", + "test-model@tp_rank:3@pcp0@dcp0@pp_rank:0@group:2", + ) + + +def test_lookup_rejects_boundary_missing_one_mamba_shard(): + # Mixed model: the deduped attention group probes 1 namespace per hash, + # the TP-sharded Mamba group probes tp_size. Exercises the per-group + # exists accounting in lookup(): a boundary whose Mamba shard is absent + # on one rank (what PUT striping used to produce) must not be a hit. + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MambaSpec, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 2 + worker.num_kv_head = 1 + fa = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None) + mamba = MambaSpec( + block_size=16, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(["l0"], fa), + KVCacheGroupSpec(["l1"], mamba), + ] + worker.token_dbs = [ + ChunkedTokenDatabase( + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16 + ), + ChunkedTokenDatabase( + KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 + ), + ] + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator( + worker._kv_cache_groups, + scheduler_block_size=16, + hash_block_size=16, + ) + worker._init_lookup_key_prefixes() + + worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys) + assert worker.lookup(32, [b"h0", b"h1"]) == 32 + + worker.store.batch_is_exist.side_effect = lambda keys: [ + 0 if "tp_rank:1" in k and "group:1" in k else 1 for k in keys + ] + assert worker.lookup(32, [b"h0", b"h1"]) == 0 + + def test_lookup_requires_all_dcp_rank_namespaces(): worker = _make_bare_worker(block_size=16) worker.tp_size = 4 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 38f4cb0c3a1e..425a54bb3343 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -43,6 +43,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 ExternalCachedBlockPool, MooncakeStoreCoordinator, + _unwrap_spec, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 BlobBlockHashes, @@ -66,7 +67,13 @@ maybe_convert_block_hash, resolve_kv_cache_block_sizes, ) -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) from .metrics import MooncakeStoreConnectorStats @@ -448,7 +455,7 @@ def __init__( token_databases: list[ChunkedTokenDatabase], block_size: int, tp_rank: int, - put_step: int, + group_put_steps: list[int], kv_role: str, ready_event: threading.Event, enable_kv_event: bool = False, @@ -464,7 +471,10 @@ def __init__( name="KVCacheStoreSendingThread", record_operation=record_operation, ) - self.put_step = put_step + # Per-group PUT stride: the ranks holding byte-identical bytes for a + # group stripe its blocks across themselves. TP-sharded groups (Mamba + # state) have stride 1 — every rank writes its own shard. + self.group_put_steps = group_put_steps self.coord = coord self.kv_role = kv_role self.stored_requests: defaultdict[str, int] = defaultdict(int) @@ -568,13 +578,14 @@ def _handle_request(self, req_meta: ReqMeta): group_indices: list[int] = [] for g_idx, db in enumerate(self.token_databases): # Rotate the stride phase per group to balance load across ranks. - put_step_rank = (self.tp_rank + g_idx) % self.put_step + put_step = self.group_put_steps[g_idx] + put_step_rank = (self.tp_rank + g_idx) % put_step for start, end, block_hash in db.process_tokens( token_len, req_meta.block_hashes, mask_num=save_start, chunk_mask=store_masks[g_idx], - put_step=self.put_step, + put_step=put_step, put_step_rank=put_step_rank, ): starts.append(start) @@ -989,44 +1000,7 @@ def __init__( ) self.num_layers = model_config.get_num_layers(parallel_config) - self.use_mla = False - if ( - hasattr(model_config, "use_mla") - and isinstance(model_config.use_mla, bool) - and model_config.use_mla - ): - self.use_mla = True - - if self.use_mla: - self.num_kv_head = 1 - else: - self.num_kv_head = model_config.get_total_num_kv_heads() - - if self.num_kv_head < self.tp_size and self.dcp_size <= 1: - # Dedup: TP ranks holding the same KV heads stripe PUTs across - # one shared key namespace. DCP splits the TP group, so with - # DCP>1 those ranks have different `@dcpN` namespaces and - # striping would leave keys unwritten (OBJECT_NOT_FOUND on - # GET). PCP is outer to TP (pcp_rank is constant within a TP - # group), so it needs no guard. - self.put_step = self.tp_size // self.num_kv_head - self.head_or_tp_rank = self.tp_rank // self.put_step - else: - self.head_or_tp_rank = self.tp_rank - self.put_step = 1 - - self.metadata = KeyMetadata( - model_name=model_config.model.rstrip("/").split("/")[-1], - tp_rank=self.head_or_tp_rank, - pcp_rank=self.pcp_rank, - dcp_rank=self.dcp_rank, - pp_rank=self.pp_rank, - cache_prefix=str( - vllm_config.kv_transfer_config.kv_connector_extra_config.get( - "cache_prefix", "" - ) - ), - ) + self.num_kv_head = model_config.get_total_num_kv_heads() # Initialize MooncakeDistributedStore with its own TransferEngine store_config = MooncakeStoreConfig.load_from_config() @@ -1143,10 +1117,30 @@ def __init__( retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL, ) # One ChunkedTokenDatabase per group; addresses populated in - # register_kv_caches once the kv-cache layout is known. + # 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. + metadata = KeyMetadata( + model_name=model_config.model.rstrip("/").split("/")[-1], + tp_rank=self.tp_rank, + pcp_rank=self.pcp_rank, + dcp_rank=self.dcp_rank, + pp_rank=self.pp_rank, + cache_prefix=str( + vllm_config.kv_transfer_config.kv_connector_extra_config.get( + "cache_prefix", "" + ) + ), + ) + factors = self._group_tp_replication_factors() self.token_dbs: list[ChunkedTokenDatabase] = [ ChunkedTokenDatabase( - dataclasses.replace(self.metadata, group_id=g_idx), + dataclasses.replace( + metadata, + group_id=g_idx, + tp_rank=self.tp_rank // factors[g_idx], + ), g.kv_cache_spec.block_size, hash_block_size=self.hash_block_size, ) @@ -1154,27 +1148,52 @@ def __init__( ] self._init_lookup_key_prefixes() + def _group_tp_replication_factors(self) -> list[int]: + """Per-group count of TP ranks holding byte-identical cache bytes. + + The ranks in a replication set share one key namespace (their shard + id, ``tp_rank // factor``) and stripe PUTs across themselves; lookup + probes one namespace per distinct shard. MLA latent KV is replicated + on every rank; GQA replicates each KV head across + ``tp_size // num_kv_head`` ranks; Mamba/linear-attention state is + head/dim-sharded, so nothing is replicated. DCP splits the TP group + along the sequence dim, so nothing is replicated there either (and + striping across `@dcpN` namespaces would leave keys unwritten). + """ + factors: list[int] = [] + for group in self._kv_cache_groups: + spec = _unwrap_spec(group.kv_cache_spec) + if isinstance(spec, MambaSpec) or self.dcp_size > 1: + factors.append(1) + elif isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)): + factors.append(self.tp_size) + else: + factors.append(max(1, self.tp_size // self.num_kv_head)) + return factors + def _init_lookup_key_prefixes(self) -> None: - """Prepare per-group key prefixes across parallel rank namespaces.""" - # (tp_rank, pcp_rank, dcp_rank, pp_rank) namespaces - if self.dcp_size > 1: - # DCP reuses the TP workers and splits each TP group into - # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size. - # Store/load paths do not apply KV-head dedup under DCP - rank_namespaces = tuple( - (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) - for pcp_rank in range(self.pcp_size) - for tp_rank in range(self.tp_size) - for pp_rank in range(self.pp_size) - ) - else: - # Without DCP, TP ranks that share a KV head write identical KV, so - # lookup only needs one TP namespace per unique KV head. - tp_count = min(self.tp_size, self.num_kv_head) - rank_namespaces = tuple( - (tp_rank, pcp_rank, 0, pp_rank) + """Prepare per-group key prefixes across parallel rank namespaces. + + A boundary is usable only when every namespace the load path will + read has the key: one (tp, pcp, dcp, pp) namespace per distinct TP + shard of each group. + """ + factors = self._group_tp_replication_factors() + + def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]: + if self.dcp_size > 1: + # DCP reuses the TP workers and splits each TP group into + # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size. + return tuple( + (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) + for pcp_rank in range(self.pcp_size) + for tp_rank in range(self.tp_size) + for pp_rank in range(self.pp_size) + ) + return tuple( + (shard_rank, pcp_rank, 0, pp_rank) for pcp_rank in range(self.pcp_size) - for tp_rank in range(tp_count) + for shard_rank in range(self.tp_size // factor) for pp_rank in range(self.pp_size) ) @@ -1187,11 +1206,12 @@ def _init_lookup_key_prefixes(self) -> None: dcp_rank=dcp_rank, pp_rank=pp_rank, ) - for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces + for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces( + factors[g_idx] + ) ) - for db in self.token_dbs + for g_idx, db in enumerate(self.token_dbs) ) - self._lookup_expected_per_key = len(rank_namespaces) def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: """Register a cross-layers KV cache tensor. @@ -1284,7 +1304,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: self.token_dbs, self.block_size, self.tp_rank, - self.put_step, + self._group_tp_replication_factors(), self.kv_role, ready_event_sending, self.enable_kv_events, @@ -1511,16 +1531,16 @@ def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int: logger.error("Remote connection failed in lookup: %s", e) return 0 - # A (group, hash) is "present" only when every TP*PP rank has it. - ranks_per_candidate = self._lookup_expected_per_key - exists_set = { - (g_idx, hash_bytes) - for i, (g_idx, hash_bytes) in enumerate(candidate_meta) - if all( - res[i * ranks_per_candidate + j] == 1 - for j in range(ranks_per_candidate) - ) - } + # A (group, hash) is "present" only when every namespace that will be + # loaded has it (per-group count: sharded groups need every rank's + # shard, replicated groups one namespace per unique KV head). + exists_set = set() + pos = 0 + for g_idx, hash_bytes in candidate_meta: + count = len(self._lookup_key_prefixes[g_idx]) + if all(res[pos + j] == 1 for j in range(count)): + exists_set.add((g_idx, hash_bytes)) + pos += count _masks, hit_length = self.coord.find_longest_cache_hit( block_hashes, From e1c3576d8d01baf40ceb375baff9f42c411ee7ce Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 23 Jul 2026 07:42:32 +0000 Subject: [PATCH 2/3] Cache and harden per-group TP replication factors Signed-off-by: Yifan Qiao --- .../unit/test_mooncake_store_hma_e2e.py | 2 +- .../unit/test_mooncake_store_worker.py | 87 +++++++++++++------ .../kv_connector/v1/mooncake/store/worker.py | 76 ++++++++-------- 3 files changed, 97 insertions(+), 68 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 40d72b3b5cf1..ce1cc48d4e05 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -214,7 +214,7 @@ def _fake_thread_init(*args, **kwargs): block_size=worker.block_size, coord=worker.coord, tp_rank=worker.tp_rank, - group_put_steps=worker._group_tp_replication_factors(), + group_put_steps=worker._group_tp_replication_factors, kv_role=worker.kv_role, ready_event=ready, enable_kv_event=False, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index d4fe518a69c8..4ba67e53dc54 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -528,10 +528,7 @@ def test_store_sending_thread_delta_strides_with_local_phase(): def test_tp_sharded_group_saves_every_block_on_every_rank(): - # Mamba state is TP-sharded (different bytes per rank), so the cross-rank - # PUT striping that dedups replicated MLA/GQA KV must not apply: a rank - # skipping a block would leave its shard unwritten, and consumers would - # load another rank's shard (silent state corruption on TP>1 warm hits). + """Sharded ranks must write every block because peers hold different bytes.""" store = MagicMock() store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys) @@ -550,7 +547,6 @@ def test_tp_sharded_group_saves_every_block_on_every_rank(): ) keys = store.batch_is_exist.call_args.args[0] - # All four blocks written, not the strided half. assert len(keys) == 4 @@ -1239,7 +1235,7 @@ def test_worker_put_striding_covers_every_rank_get_namespace( ] assert len(keys) == len(block_hashes) # PUT side: mirrors KVCacheStoreSendingThread's striding slice. - put_step = w._group_tp_replication_factors()[0] + put_step = w._group_tp_replication_factors[0] put_keys.update(keys[w.tp_rank % put_step :: put_step]) # GET side: KVCacheStoreRecvingThread fetches every key. get_keys_per_rank[tp_rank] = set(keys) @@ -1579,6 +1575,15 @@ def _register_with_mocked_threads( worker.register_kv_caches(kv_caches) +def _refresh_group_tp_replication_factors( + worker: mooncake_store_worker.MooncakeStoreWorker, +) -> None: + worker._group_tp_replication_factors = ( + worker._compute_group_tp_replication_factors() + ) + worker._init_lookup_key_prefixes() + + def _make_bare_worker( *, num_gpu_blocks: int = 10, @@ -1642,7 +1647,7 @@ def _make_bare_worker( scheduler_block_size=block_size, hash_block_size=block_size, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) return worker @@ -1651,7 +1656,7 @@ def test_lookup_key_prefixes_cover_dcp_rank_namespaces(): worker.tp_size = 4 worker.num_kv_head = 1 worker.dcp_size = 4 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", @@ -1667,7 +1672,7 @@ def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): worker.num_kv_head = 1 worker.pcp_size = 2 worker.dcp_size = 1 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", @@ -1676,9 +1681,7 @@ def test_lookup_key_prefixes_cover_pcp_rank_namespaces(): def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank(): - # Under MLA KV-head dedup the replicated attention group probes a single - # shared namespace, but a TP-sharded Mamba group has one shard per rank: - # a boundary is only usable when every rank's shard exists. + """Replicated attention needs one namespace; sharded Mamba needs every rank.""" from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1707,7 +1710,7 @@ def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank(): KeyMetadata("test-model", 1, 0, 0, 0, group_id=1), block_size=16 ), ] - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", @@ -1719,9 +1722,6 @@ def test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank(): def test_group_tp_replication_factors_mixed_mla_gqa_mamba(): - # Replication is a per-group property: MLA latent KV is replicated on - # every rank, each GQA KV head on tp_size // num_kv_head ranks, Mamba - # state on none. Namespace count per group is tp_size // factor. from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1752,9 +1752,8 @@ def test_group_tp_replication_factors_mixed_mla_gqa_mamba(): for g_idx in range(3) ] - assert worker._group_tp_replication_factors() == [4, 2, 1] - - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) + assert worker._group_tp_replication_factors == (4, 2, 1) assert worker._lookup_key_prefixes[0] == ( "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", ) @@ -1770,11 +1769,45 @@ def test_group_tp_replication_factors_mixed_mla_gqa_mamba(): ) +@pytest.mark.parametrize("spec_order", [("mla", "gqa"), ("gqa", "mla")]) +def test_uniform_group_uses_common_inner_replication_factor(spec_order): + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + ) + + worker = _make_bare_worker(block_size=16) + worker.tp_size = 4 + worker.num_kv_head = 2 + specs_by_name = { + "mla": MLAAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=None + ), + "gqa": FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=None + ), + } + inner_specs = {name: specs_by_name[name] for name in spec_order} + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs=inner_specs, + ) + worker._kv_cache_groups = [ + KVCacheGroupSpec(list(inner_specs), uniform_spec), + ] + + _refresh_group_tp_replication_factors(worker) + + assert worker._group_tp_replication_factors == (2,) + assert worker._lookup_key_prefixes[0] == ( + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0", + "test-model@tp_rank:1@pcp0@dcp0@pp_rank:0@group:0", + ) + + def test_lookup_rejects_boundary_missing_one_mamba_shard(): - # Mixed model: the deduped attention group probes 1 namespace per hash, - # the TP-sharded Mamba group probes tp_size. Exercises the per-group - # exists accounting in lookup(): a boundary whose Mamba shard is absent - # on one rank (what PUT striping used to produce) must not be a hit. from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -1808,7 +1841,7 @@ def test_lookup_rejects_boundary_missing_one_mamba_shard(): scheduler_block_size=16, hash_block_size=16, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.side_effect = lambda keys: [1] * len(keys) assert worker.lookup(32, [b"h0", b"h1"]) == 32 @@ -1824,7 +1857,7 @@ def test_lookup_requires_all_dcp_rank_namespaces(): worker.tp_size = 4 worker.num_kv_head = 1 worker.dcp_size = 4 - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) worker.store.batch_is_exist.return_value = [1, 1, 0, 1] assert worker.lookup(16, [b"a0"]) == 0 @@ -1900,7 +1933,7 @@ def test_lookup_checks_all_potential_swa_hit_boundaries(): hash_block_size=8, retention_interval=0, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) # Candidate order: 3 full-attention chunks, then SWA chunks 3, 7, 11. # Only the first full chunk and the SWA chunk ending at token 32 exist, so # lookup should recover a 32-token external prefix hit. A sparse @@ -1961,7 +1994,7 @@ def test_lookup_applies_swa_mask_before_accessing_hashes(): hash_block_size=8, retention_interval=0, ) - worker._init_lookup_key_prefixes() + _refresh_group_tp_replication_factors(worker) block_hashes = _RecordingBlockHashes([f"h{i}".encode() for i in range(12)]) accessed_before_rpc: list[int] = [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 425a54bb3343..186362d71f90 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -21,6 +21,7 @@ from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass +from math import gcd from typing import Any, Literal, TypeVar import regex as re @@ -43,7 +44,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import ( # noqa: E501 ExternalCachedBlockPool, MooncakeStoreCoordinator, - _unwrap_spec, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 BlobBlockHashes, @@ -70,9 +70,11 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, + UniformTypeKVCacheSpecs, ) from .metrics import MooncakeStoreConnectorStats @@ -455,7 +457,7 @@ def __init__( token_databases: list[ChunkedTokenDatabase], block_size: int, tp_rank: int, - group_put_steps: list[int], + group_put_steps: Sequence[int], kv_role: str, ready_event: threading.Event, enable_kv_event: bool = False, @@ -471,9 +473,7 @@ def __init__( name="KVCacheStoreSendingThread", record_operation=record_operation, ) - # Per-group PUT stride: the ranks holding byte-identical bytes for a - # group stripe its blocks across themselves. TP-sharded groups (Mamba - # state) have stride 1 — every rank writes its own shard. + # Only ranks with identical group bytes may stripe PUTs (e.g., MLA). self.group_put_steps = group_put_steps self.coord = coord self.kv_role = kv_role @@ -1133,13 +1133,15 @@ def __init__( ) ), ) - factors = self._group_tp_replication_factors() + self._group_tp_replication_factors: tuple[int, ...] = ( + self._compute_group_tp_replication_factors() + ) self.token_dbs: list[ChunkedTokenDatabase] = [ ChunkedTokenDatabase( dataclasses.replace( metadata, group_id=g_idx, - tp_rank=self.tp_rank // factors[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, @@ -1148,42 +1150,36 @@ def __init__( ] self._init_lookup_key_prefixes() - def _group_tp_replication_factors(self) -> list[int]: - """Per-group count of TP ranks holding byte-identical cache bytes. - - The ranks in a replication set share one key namespace (their shard - id, ``tp_rank // factor``) and stripe PUTs across themselves; lookup - probes one namespace per distinct shard. MLA latent KV is replicated - on every rank; GQA replicates each KV head across - ``tp_size // num_kv_head`` ranks; Mamba/linear-attention state is - head/dim-sharded, so nothing is replicated. DCP splits the TP group - along the sequence dim, so nothing is replicated there either (and - striping across `@dcpN` namespaces would leave keys unwritten). + def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int: + if self.dcp_size > 1: + return 1 + if isinstance(spec, UniformTypeKVCacheSpecs): + inner_factors = tuple( + self._spec_tp_replication_factor(inner_spec) + for inner_spec in spec.kv_cache_specs.values() + ) + return gcd(*inner_factors) if inner_factors else 1 + if isinstance(spec, MambaSpec): + return 1 + if isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)): + return self.tp_size + return max(1, self.tp_size // self.num_kv_head) + + def _compute_group_tp_replication_factors(self) -> tuple[int, ...]: + """Return the number of byte-identical TP replicas per cache group. + + DCP and Mamba use 1; MLA uses ``tp_size``; GQA uses + ``tp_size // num_kv_head``. Packed groups use the GCD of inner factors. """ - factors: list[int] = [] - for group in self._kv_cache_groups: - spec = _unwrap_spec(group.kv_cache_spec) - if isinstance(spec, MambaSpec) or self.dcp_size > 1: - factors.append(1) - elif isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)): - factors.append(self.tp_size) - else: - factors.append(max(1, self.tp_size // self.num_kv_head)) - return factors + return tuple( + self._spec_tp_replication_factor(group.kv_cache_spec) + for group in self._kv_cache_groups + ) def _init_lookup_key_prefixes(self) -> None: - """Prepare per-group key prefixes across parallel rank namespaces. - - A boundary is usable only when every namespace the load path will - read has the key: one (tp, pcp, dcp, pp) namespace per distinct TP - shard of each group. - """ - factors = self._group_tp_replication_factors() - def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]: if self.dcp_size > 1: - # DCP reuses the TP workers and splits each TP group into - # contiguous DCP groups, so dcp_rank == tp_rank % dcp_size. + # DCP is a TP subdivision: dcp_rank == tp_rank % dcp_size. return tuple( (tp_rank, pcp_rank, tp_rank % self.dcp_size, pp_rank) for pcp_rank in range(self.pcp_size) @@ -1207,7 +1203,7 @@ def rank_namespaces(factor: int) -> tuple[tuple[int, int, int, int], ...]: pp_rank=pp_rank, ) for tp_rank, pcp_rank, dcp_rank, pp_rank in rank_namespaces( - factors[g_idx] + self._group_tp_replication_factors[g_idx] ) ) for g_idx, db in enumerate(self.token_dbs) @@ -1304,7 +1300,7 @@ def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor: self.token_dbs, self.block_size, self.tp_rank, - self._group_tp_replication_factors(), + self._group_tp_replication_factors, self.kv_role, ready_event_sending, self.enable_kv_events, From eb7d24a76f00f41f375c6c7e7756c3588b2bfc82 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 23 Jul 2026 18:45:59 +0000 Subject: [PATCH 3/3] Flatten TP replication factor logic per review UniformTypeKVCacheSpecs only ever holds raw per-layer specs, so replace the recursive gcd over inner factors with a flat any-Mamba / all-MLA / GQA check. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GzJ49XLtcNtpm8r7iLemaT Signed-off-by: Yifan Qiao --- .../kv_connector/v1/mooncake/store/worker.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 186362d71f90..0aa92a9e46d6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -21,7 +21,6 @@ from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass -from math import gcd from typing import Any, Literal, TypeVar import regex as re @@ -1153,15 +1152,19 @@ def __init__( def _spec_tp_replication_factor(self, spec: KVCacheSpec) -> int: if self.dcp_size > 1: return 1 - if isinstance(spec, UniformTypeKVCacheSpecs): - inner_factors = tuple( - self._spec_tp_replication_factor(inner_spec) - for inner_spec in spec.kv_cache_specs.values() - ) - return gcd(*inner_factors) if inner_factors else 1 - if isinstance(spec, MambaSpec): + inner_specs = ( + tuple(spec.kv_cache_specs.values()) + if isinstance(spec, UniformTypeKVCacheSpecs) + else (spec,) + ) + # Any rank-specific state makes the whole packed value rank-specific. + if any(isinstance(inner, MambaSpec) for inner in inner_specs): return 1 - if isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)): + # A pure MLA packed value is replicated on every TP rank. + if all( + isinstance(inner, (MLAAttentionSpec, SlidingWindowMLASpec)) + for inner in inner_specs + ): return self.tp_size return max(1, self.tp_size // self.num_kv_head) @@ -1169,7 +1172,7 @@ def _compute_group_tp_replication_factors(self) -> tuple[int, ...]: """Return the number of byte-identical TP replicas per cache group. DCP and Mamba use 1; MLA uses ``tp_size``; GQA uses - ``tp_size // num_kv_head``. Packed groups use the GCD of inner factors. + ``tp_size // num_kv_head``. """ return tuple( self._spec_tp_replication_factor(group.kv_cache_spec)