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

[ivanium]

Purpose

Fix silent Mamba/linear-attention state corruption in the Mooncake store connector when serving hybrid models with TP > 1.

With MLA (or GQA with num_kv_head < tp_size) all TP ranks hold identical attention KV, so the store dedups: ranks share one tp_rank key namespace and stripe PUTs across blocks (put_step). That assumption is wrong for hybrid models' Mamba/linear-attention KV-cache groups, whose state is head/dim-sharded across TP ranks: 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 TP=2 this corrupted every external mamba-state load — gsm8k warm 0.72 vs cold 0.87, with every wrong answer a plausible paraphrase; per-layer probing showed rank 1's layer-0 initial state wrong while the transported bytes matched the stored bytes exactly, i.e. the transfer was faithful but the stored shard belonged to the other rank.

Fix: per-group TP replication factor

The root cause is that replication is a per-group property but the connector computed it model-wide (one num_kv_head/put_step/head_or_tp_rank triple applied to every KV-cache group). This PR replaces that triple with a per-group TP replication factor — the number of TP ranks holding byte-identical cache bytes for that group:

Group	Factor
MLA (latent KV, replicated)	tp_size
GQA	max(1, tp_size // num_kv_head)
Mamba / linear-attention state (head/dim-sharded)	1
Any group under DCP (sequence-sharded)	1

Everything else is derived from the factor uniformly: each group's key namespace is its shard id (tp_rank // factor), PUTs stripe only across the ranks within a replication set, and lookup requires one key per distinct shard before reporting a boundary (the exists accounting is per-group instead of a single global _lookup_expected_per_key).

This produces byte-identical store keys for every currently-supported model — replicated attention groups keep the dedup exactly as before; Mamba groups simply stop sharing a namespace and striping (the bug fix). It also gives future mixed MLA + GQA + Mamba hybrids the correct per-group treatment instead of a model-wide guess, and unwraps UniformTypeKVCacheSpecs before dispatch (matching the coordinator) so wrapped Mamba groups are classified correctly.

Why this is not duplicating an existing PR

Searched open PRs/issues for mooncake+mamba, mooncake+dedup, KV-head striping, and hybrid-model KV-transfer accuracy. #48361 (EAGLE cache-peek gating for Mamba lookups), #48970 (e2e hybrid-Mamba regression tests for #43559), and #48216 (group-aware KV load-failure recovery) touch nearby code but none address the TP-shard vs KV-head-dedup namespace bug; no open PR changes the put_step/namespace logic for sharded groups.

Test Plan

Unit tests:

.venv/bin/python -m pytest tests/v1/kv_connector/unit/test_mooncake_store_worker.py tests/v1/kv_connector/unit/test_mooncake_store_connector.py tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py -v

New tests: test_tp_sharded_group_saves_every_block_on_every_rank and test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank fail on current main (reproducing the bug) and pass with this change; test_group_tp_replication_factors_mixed_mla_gqa_mamba covers the per-group factors on a mixed MLA+GQA+Mamba worker (factors 4/2/1 at TP=4); test_lookup_rejects_boundary_missing_one_mamba_shard locks down the per-group exists accounting in lookup() end-to-end. All pre-existing tests pass unchanged (including the DCP/PCP namespace and PUT-striding coverage tests, which pin the key format).

E2E validation on Kimi-Linear-48B-A3B-Instruct, TP=2, Mooncake store as external KV cache, two-pass (cold then warm) runs at temp=0.

Test Result

• Unit: 108 passed across the three mooncake store suites.

E2E on this exact branch (GB200, TP=2, Mooncake store, temp=0, per-run nonce so every pass 1 is cold; hit/failure counters verified nonzero/zero from /metrics for every claim below):

Kimi-Linear-48B-A3B-Instruct TP2 (hybrid MLA + KDA; the bug scenario — without partial-hit support the aligned hit granularity is 2048 tokens, so prompts use a ~4.2k-token primer):

• Serial (conc=1) two-pass with 12 per-prompt-unique block chains: every request stored and reloaded its own mamba boundary state (load_get ok = 12 req × 2 ranks, external hits = 12 × 2048 tokens, save_put on both ranks — no striping for the sharded group, per this PR's design): 0/12 cold-vs-warm divergence, byte-identical outputs, zero KV load failures.
• Concurrency-8 two-pass (n=40): zero load failures, warm accuracy ≥ cold (0.45 vs 0.375), warm latency 8.5× lower (external hits served). 21/40 outputs differed, all splitting mid-generation (median ~200 chars in) — MoE batch-composition numerics at temp 0, not state corruption; the serial run above shows 0 divergence when batching noise is removed.

DeepSeek-V4-Flash TP2, fp8 KV (MLA num_kv_heads=1 → KV-head dedup + PUT striping active — the regression guard for the path this PR must not break):

• gsm8k 120 prompts, conc=16: zero load failures, external hits nonzero, warm accuracy ≥ cold (0.933 vs 0.892).
• Three-pass byte-identity control (cold → GPU-cache-warm → reset → external-store-warm, serial): GPU-warm vs external-warm outputs byte-identical 0/12 with verified external loads (12 req × 2 ranks) — the deduped/striped store path returns exactly the bytes the local prefix cache would have served. (Cold-vs-warm shows divergence with fp8 KV cache for any prefix-cache source, GPU or external — cache-read vs in-chunk prefill numerics — hence the GPU-vs-external control.)

Historical (internal tree, same fix, pre-cherry-pick): Kimi-Linear TP2 40-prompt repro 38/40 divergent before → 0/40 after; gsm8k two-pass strict-match cold 0.8673 / warm 0.8726 (warm was 0.7225 before the fix).

AI assistance disclosure

This change was developed with AI assistance (Claude). The submitting human (@ivanium) has reviewed every changed line, ran the tests above, and can defend the change end-to-end. The commit was additionally reviewed by an independent code-review model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[GirasoleY]

Nice fix :) ! Minor nits, nothing blocking.

[Dao007forever]

Use Sequence[int]?

[ivanium]

I realized that tuple can indicate the sequence is immutable; so maybe we will keep tuple.

[Dao007forever]

Should we update group_put_steps to tuple as well?

[ivanium]

some discussion with agents and basically the tl;dr is for input arg type we can type loose to be more general; while for return type we can be more concrete. I think it makes sense to me, so we should still keep group_put_steps Sequence?

[Dao007forever]

Sure.

[Dao007forever]

Suggested change

	if isinstance(spec, UniformTypeKVCacheSpecs):
	inner_factors = tuple(
	self._spec_tp_replication_factor(inner_spec)
	for inner_spec in spec.kv_cache_specs.values()
	)
	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)
	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)
	inner_specs = (
	tuple(spec.kv_cache_specs.values())
	if isinstance(spec, UniformTypeKVCacheSpecs)
	else (spec,)
	)
	# Any rank-specific state makes the complete packed value rank-specific.
	if any(isinstance(inner, MambaSpec) for inner in inner_specs):
	return 1
	# A pure MLA packed value is replicated on every TP rank.
	if all(isinstance(inner, MLA_SPECS) for inner in inner_specs):
	return self.tp_size
	# All ordinary attention layers use the same model-wide KV-head count.
	return max(1, self.tp_size // self.num_kv_head)

If we assume UniformTypeKVCacheSpecs has only 1 level (which I think is a good assumption)

[ivanium]

done