[KV Connector] Support heterogeneous TP sharing in Mooncake Store Connector

[z-zanez]

Purpose

This PR allows vLLM instances whose TP sizes divide a common Store TP size to
share the same KV cache through Mooncake Store. Sharing vLLM instances must use
the same PP size. Both local HND and NHD KV cache layouts are supported, as are
vLLM instances whose ranks hold different numbers of KV heads because they use
different TP sizes.

Previously, Mooncake Store keys and physical data layouts were tied to the
local TP rank. Even when Prefill and Decode instances used the same model,
different TP sizes produced different Store shard counts, KV-head ranges, and
key namespaces, so the instances could not hit the same KV cache entries.

This PR introduces a fixed Store TP layout that decouples Store objects from
the local TP size of a vLLM instance:

• store_tp_size defines a stable number of logical Store shards;
• for regular MHA/GQA sharing, each Store shard contains a fixed range of
  global KV heads;
• one local TP rank maps to one or more consecutive Store shards;
• vLLM instances with different local TP sizes can read and write the same key
  namespace when they select the same Store TP size;
• a Decode consumer with save_decode_cache enabled writes newly completed
  decode KV blocks back in the same Store TP layout.

For example, with Store TP size 4:

• each TP4 rank maps to one Store shard;
• each TP2 rank maps to two Store shards;
• KV written by a TP4 Prefill instance can be read by a TP2 Decode instance;
• KV written back by the TP2 Decode instance can be reused by compatible TP4
  or TP2 instances.

The Store uses canonical HND as its stable interchange format. Each Store
shard serializes its assigned KV heads layer by layer, with each layer stored
in head-major HND order. Because this format does not depend on a vLLM
instance's local TP size or local KV cache layout, every compatible reader and
writer interprets a Store object in the same way.

• A local HND KV cache uses Mooncake multi-buffer I/O directly.
• A local CUDA NHD KV cache uses a reusable GPU staging arena and Triton
  kernels to convert between local NHD and canonical HND.
• Each staged Store object is exposed to Mooncake as one contiguous segment,
  avoiding a segment for every head/token cell.
• A staging GET unpacks only objects that were read successfully and preserves
  their original staging slots.

Deployments with multiple Prefill TP sizes can explicitly select a common
Store TP size:

{
    "kv_connector_extra_config": {
        "enable_store_tp_lcm": true,
        "prefill_tp_sizes": [4, 2]
    }
}

When enable_store_tp_lcm is enabled, the connector uses the least common
multiple of prefill_tp_sizes as the common Store TP size. For example, the
least common multiple of 4 and 2 is 4, so the configuration above selects
Store TP size 4. This selection does not change the runtime TP size of any
vLLM instance.

All Prefill and Decode instances sharing the namespace calculate the same
Store TP size from the same configuration. Decode does not inspect which
Prefill TP size produced an individual request and does not change its write
layout per request. With Prefill TP sizes 4 and 2 and Decode TP size 2:

• each TP4 Prefill rank writes one Store shard;
• each TP2 Prefill rank writes two Store shards;
• each TP2 Decode rank reads and writes two Store shards.

Decode therefore writes according to the common Store TP size, rather than its
own local TP size or the TP size of a particular Prefill instance. Its local KV
heads are mapped back to the same global Store shards, so decode KV written by
one compatible topology remains reusable by the others.

When automatic common Store TP size selection is not enabled, all vLLM
instances use the explicitly configured store_tp_size and follow the same
mapping rules.

The vLLM instances must already satisfy the existing Mooncake Store
compatibility requirements (same model/KV format, block and hash configuration,
and cache namespace). Heterogeneous-TP sharing additionally requires:

• Store TP size is greater than or equal to local TP size;
• Store TP size is divisible by local TP size;
• the model's total number of KV heads is divisible by Store TP size, or the
  model uses MQA (Hkv == 1), for which the connector uses a replicated
  shared layout;
• all sharing vLLM instances use the same PP size;
• the KV cache contains one Full Attention group;
• PCP, DCP, and cross-layer KV blocks are not used.

Different local TP sizes give each rank a different number of KV heads. This
implementation first divides the model's global KV heads into fixed Store
shards and then maps each local rank's KV heads to one or more consecutive
Store shards. The number of KV heads held by a local rank may therefore differ
between vLLM instances, as long as the total number of KV heads is divisible by
the common Store TP size. MQA is handled separately: because every TP rank
holds the same single KV head, it uses a replicated shared layout instead of
regular Store-head partitioning.

vLLM instances with different PP sizes do not share a Store namespace. The PP
size is encoded in the namespace, preventing incompatible PP topologies from
hitting each other's KV cache entries.

If the sharing conditions are not satisfied, the connector falls back to an
isolated rank-local namespace that includes the local TP, PP, and KV layout.
Existing valid configurations therefore remain valid, while incompatible
topologies and layouts cannot read each other's entries.

This PR also updates the Mooncake Store Connector documentation and adds
regression coverage for:

• Store-shard mapping across different TP sizes;
• NHD-to-NHD, NHD-to-HND, HND-to-NHD, and HND-to-HND transfers;
• vLLM instances whose ranks hold different numbers of KV heads;
• the canonical Store key namespace;
• common Store TP size selection for multiple Prefill TP sizes;
• PP namespace isolation;
• rank-local fallback for malformed configurations and incompatible
  topologies;
• the shared MQA namespace;
• staging PUT batching;
• partial staging GET failures and preservation of original staging slots;
• CUDA Triton staging pack/unpack round trips.

Related Issues and PRs

• Issue [RFC]: [Roadmap] Mooncake Store Connector Feature Enrichment Roadmap #45036 tracks the Mooncake Store Connector roadmap and lists
  "Decode-Phase KV Cache Put with GQA TP Asymmetry" as an open item. This PR
  covers the Full Attention case where TP sizes are divisible, PP sizes match,
  and KV heads can be partitioned by a common Store TP size.
• PR [KVTransfer][Mooncake] Add heterogeneous TP support for disaggregated P/D in MooncakeConnector #36869 added heterogeneous-TP support to MooncakeConnector for direct
  P/D KV transfer. In contrast, this PR targets MooncakeStoreConnector, where
  KV is persisted using a stable Store-shard layout and can be reused and
  updated by multiple Prefill and Decode instances.
• PR [KV Connector] Add decode offloading to Mooncake Store consumers #52466 added Decode KV offloading to Mooncake Store consumers. This PR
  extends that path so Prefill and Decode instances with different TP sizes can
  read and write the same KV cache using a common Store TP size.
• PR [Bugfix] Fix Mooncake heterogeneous TP with replicated GQA heads #52516 fixes heterogeneous-TP mapping in MooncakeConnector when TP size
  exceeds the number of KV heads and the heads are replicated. This PR uses a
  separate Store-shard mapping and limits regular sharing to topologies where
  the total number of KV heads is divisible by Store TP size. Unsupported
  topologies fall back to an isolated rank-local namespace.

Test Plan

Unit and static checks

pre-commit run --files \
  docs/features/mooncake_store_connector_usage.md \
  tests/v1/kv_connector/unit/test_mooncake_store_connector.py \
  tests/v1/kv_connector/unit/test_mooncake_store_worker.py \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/tp_layout.py \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

python -m pytest -q \
  tests/v1/kv_connector/unit/test_mooncake_store_worker.py \
  tests/v1/kv_connector/unit/test_mooncake_store_connector.py

RDMA cluster tests

Reusable E2E test code:

• Heterogeneous-TP cluster test directory
• Single-node RDMA layout matrix
• Multiple-Prefill-TP correctness runner

The layout matrix used the following configuration:

Item	Configuration
Model	Qwen3-32B-FP8
Hardware	One node with 8 NVIDIA H20 GPUs and NVLink
Store transport	Mooncake standalone Store over RDMA using mlx5_bond_0
TP topologies	Prefill TP4 -> Decode TP2 and Prefill TP2 -> Decode TP1
KV layouts	HND -> HND, HND -> NHD, NHD -> HND, and NHD -> NHD for each TP topology
Functional workload	One 976-token prompt, 960 expected cached tokens, and 64 deterministic output tokens
Functional criteria	Store hit covers the block-aligned prefix, generated token IDs match a same-TP no-connector reference, and no Store key fails
Performance workload	64 requests, concurrency 16, 64 output tokens per request, 2 warmup requests, and cold-before-cached ordering
Performance criteria	Every cached request hits 1,041 tokens on average, cold requests have zero cached tokens, and Mooncake reports zero failed keys

The multiple-Prefill-TP correctness run used Prefill TP sizes 4 and 2, one
Decode TP2 instance, and common Store TP size lcm(4, 2) = 4. Two distinct
prefixes were written to the same Store pool, one by the TP4 Prefill and one by
the TP2 Prefill. The same Decode instance reused both prefixes without knowing
the source Prefill TP of each request. Decode writeback was also checked after
each hit. This run was correctness-only; no performance claim is made from it.

Test Result

Unit and static checks

All pre-commit hooks passed.
182 passed

RDMA cluster tests

All eight topology/layout combinations passed functional validation:

Prefill -> Decode	Prefill layout -> Decode layout	Cached / prompt tokens	Output tokens	Reference token IDs	Failed keys
TP4 -> TP2	HND -> HND	960 / 976	64	Match	0
TP4 -> TP2	HND -> NHD	960 / 976	64	Match	0
TP4 -> TP2	NHD -> HND	960 / 976	64	Match	0
TP4 -> TP2	NHD -> NHD	960 / 976	64	Match	0
TP2 -> TP1	HND -> HND	960 / 976	64	Match	0
TP2 -> TP1	HND -> NHD	960 / 976	64	Match	0
TP2 -> TP1	NHD -> HND	960 / 976	64	Match	0
TP2 -> TP1	NHD -> NHD	960 / 976	64	Match	0

Performance results are end-to-end server measurements. Cached req/s and
Cold req/s include scheduling, Store reads, KV placement, and generation.
RDMA RPC GiB/s measures successful Mooncake load_get RPC time, while
phase GiB/s divides the same bytes by the complete cached-phase wall time.

Prefill -> Decode	Layout	Cached req/s	Cold req/s	Throughput speedup	Cached / cold mean latency (s)	RDMA RPC GiB/s	Phase GiB/s
TP4 -> TP2	HND -> HND	6.276	3.351	1.873x	2.529 / 4.763	11.734	1.595
TP4 -> TP2	HND -> NHD	6.214	3.376	1.841x	2.554 / 4.727	18.607	1.579
TP4 -> TP2	NHD -> HND	6.411	3.385	1.894x	2.476 / 4.714	11.759	1.629
TP4 -> TP2	NHD -> NHD	6.280	3.392	1.851x	2.527 / 4.705	19.416	1.596
TP2 -> TP1	HND -> HND	4.769	1.961	2.432x	3.342 / 8.149	24.667	1.212
TP2 -> TP1	HND -> NHD	4.726	1.964	2.406x	3.366 / 8.139	27.987	1.201
TP2 -> TP1	NHD -> HND	4.778	1.965	2.431x	3.337 / 8.135	24.761	1.214
TP2 -> TP1	NHD -> NHD	4.758	1.960	2.427x	3.338 / 8.155	24.552	1.209

Across the matrix, cached throughput was 1.841x-2.432x cold throughput. Each
cached phase transferred 17,465,081,856 bytes from Mooncake Store, with no
cache misses, unexpected cold hits, or failed Store keys.

Multiple-Prefill-TP correctness results:

One TP2/NHD Decode instance was connected to a single shared Store pool. The
pool contained distinct prefixes produced with both supported Prefill TP sizes:

Prefix	KV producer in the shared pool	Decode consumer	Decode cached / prompt tokens	Decode output tokens	Decode-writeback verification	Result
A	TP4/HND Prefill	Same TP2/NHD Decode	1,056 / 1,057	64	TP2/NHD Prefill subsequently hit 1,104 tokens	Passed
B	TP2/NHD Prefill	Same TP2/NHD Decode	1,056 / 1,057	64	TP4/HND Prefill subsequently hit 1,104 tokens	Passed

This demonstrates that one Decode service is not tied to a single Prefill TP
size: across requests, it can reuse KV produced by both TP4 and TP2 through the
same common Store TP size 4 namespace, without per-request source-TP metadata.
It does not mean that one request splices together KV fragments from different
Prefill executions. Mooncake reported zero failed keys and zero operation
errors.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new 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.

[mergify]

Documentation preview: https://vllm--53129.org.readthedocs.build/en/53129/

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀