[Feature]: Deduplicate replicated MLA KV across TP ranks in native offloading

[Change72]

🚀 The feature, motivation and pitch

Problem

For MLA models, the latent KV cache is logically replicated across TP ranks — and expected to be byte-identical in homogeneous TP setups without context parallelism. MLA hardcodes num_kv_heads=1 (vllm/model_executor/layers/attention/mla_attention.py:385) and the latent projection is not TP-sharded (ReplicatedLinear, or disable_tp=True in the fused-projection paths). The codebase states this directly:

"MLA is always replicated as the hidden dim can't be split" — vllm/distributed/kv_transfer/kv_connector/utils.py:562

The native KV offloading path (OffloadingConnector + vllm/v1/kv_offload/) does not exploit this. In the pure-MLA, DCP=PCP=1 baseline configuration (refs against main @ 7bd154375):

• Native CPU/tiering offload materializes one copy of the logically replicated latent KV per TP worker. CPUOffloadingSpec prices each logical block at per-worker bytes × world size (vllm/v1/kv_offload/cpu/spec.py:86), so at TP=N the unique-content capacity is roughly 1/N of the configured --kv-offloading-size budget.
• Default CPUOffloadingSpec: each rank allocates its own private CPU buffer (pinned when supported) (vllm/v1/kv_offload/cpu/gpu_worker.py:501).
• TieringOffloadingSpec: on single-node TP, the copies live in different per-worker slots of the same shared-mmap row (worker0_block0 | worker1_block0 | ..., vllm/v1/kv_offload/cpu/shared_offload_region.py:123); the FS secondary tier then writes the entire row as one layout-specific block file per key, so the duplicate copies sit inside that file as distinct worker slots (vllm/v1/kv_offload/tiering/fs/manager.py:159). These observations are verified for single-node TP; multi-node tiering — where both /dev/shm and the scheduler's memory view are node-local — is treated as unsupported/unverified here.
• The fs tier's parallel-agnostic mode acknowledges the replication but excludes MLA from it: "MLA is excluded: its latent KV is replicated per rank, never head-sharded" (vllm/v1/kv_offload/file_mapper.py:85, from [KV offload] Parallel-agnostic fs-tier cache for single full-attention group #44733). Note that relaxing this predicate alone would not be enough — because the file payload is the full multi-worker row, files produced under different TP sizes differ in length and layout; a compact, TP-independent payload is a prerequisite.

Net effect: --kv-offloading-size 100 (GiB) with TP=8 on a pure-MLA model yields only ~12.5 GiB of unique cached KV. Neither the kv_offloading_size docstring ("When TP > 1, this is the total buffer size summed across all TP ranks", vllm/config/cache.py:182) nor docs/features/kv_offloading_usage.md mentions this caveat.

MLA is the best case for deduplication (replicated content, no reshard needed), yet it is the one case excluded from the existing single-copy machinery.

Scope and correctness notes

• In general the duplication factor for pure-MLA is about TP / DCP, not always TP: DCP shards tokens across context-parallel ranks, but each DCP shard is still replicated across the TP ranks that share it. PCP is outer to TP, so TP-dedup can in principle be applied within each PCP namespace independently.
• world_size also includes PP, and PP stages hold different layers — that dimension is real sharding, not replication, and must remain untouched by any dedup.
• Proposed initial scope: pure-MLA models with PP=PCP=DCP=1 — a safe MVP boundary where the replica factor is exactly TP; this is an implementation simplification, not a theoretical requirement. Extensions (DCP-aware factor, per-PCP-namespace dedup, PP composition) can follow.
• Hybrid MLA + full-attention models may place both kinds of layers in the same KV cache group, so gating purely at group level is insufficient; generalization needs per-layer/region (canonical tensor) granularity. The initial version would simply not engage for hybrid models.
• Scheduler-side block accounting (num_cpu_blocks, medium/events) must stay consistent with the single-copy worker-side layout.

Prior art

• MooncakeStoreConnector ([KV Transfer] Add MooncakeStoreConnector for KV cache offloading via Mooncake distributed store #40900) implements exactly this storage dedup: with num_kv_head < tp_size (MLA ⇒ 1) and dcp_size <= 1, TP ranks collapse to one shared key namespace and stripe PUTs across ranks so each block is stored once (vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py:1004). The DCP guard was added in [Bugfix][KV Connector] Disable Mooncake TP put-striding when DCP > 1 #45371; PCP needs no guard since it is outer to TP.
• LMCache (LMCacheConnectorV1): for MLA, save_only_first_rank defaults on — non-leader ranks neither store nor read; the leader retrieves and performs a TP-group device broadcast (v0.5.1 cache_engine.py L113-117). Cache keys collapse world_size to 1 (TP-size-agnostic) since LMCache#2098 (first released in v0.3.10). A separate remote-only scheme, remote_enable_mla_worker_id_as0, serves as an alternative when save_only_first_rank is disabled (each rank reads the single rank-0 copy from remote storage directly). The built-in single-server LMCacheMPConnector fallback collapses only the TP dimension and suppresses MLA stores on ranks where rank % tp_size != 0 ("Tensor parallel does not change the KV caches for MLA models", vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py); in a TP-only deployment this yields one stored copy while every TP rank retrieves directly. The external connector preferred by current vLLM has separate multi-server writer semantics.
• (The NIXL connector's single-source MLA handling is a transfer optimization for replicated KV, not storage dedup — noted only to delimit scope.)

The native CPU/tiering layout currently lacks an equivalent MLA replica-reduction path.

Proposed direction (candidates — feedback requested)

1. Sizing and single-copy layout must land together. The per-block cost in cpu/spec.py should drop only the confirmed TP replica factor, leaving PP (different layers per stage) and any context-parallel sharding intact. Adjusting the sizing alone is unsafe: with per-worker private buffers still in place, simply growing num_cpu_blocks would multiply actual memory use by roughly world_size beyond the configured budget — the single-copy physical layout is a prerequisite for the sizing change.
2. Candidate A — single-node tiering shared-mmap single slot: store one worker slot per block row; the TP leader performs the d2h write, all node-local ranks read the same slot. Open questions to resolve: store-completion synchronization across ranks, pinning and ownership of the shared region, row-stride/layout change, and the fact that /dev/shm is node-local — multi-node TP cannot share the slice without communication.
3. Candidate B — leader store + TP-group broadcast on load (LMCache-style): works uniformly across single- and multi-node TP; adds TP-group broadcast traffic on load.
4. FS tier: produce a compact, TP-independent payload (single logical copy) and then admit MLA into the parallel-agnostic naming of file_mapper.py. Note that this naming mode erases all parallel dimensions (TP/PP/DCP/PCP) and rank, not just TP ("tp/pp/pcp/dcp are forced to 1 and rank to 0", vllm/v1/kv_offload/file_mapper.py:43), so the admission stays gated to the PP=PCP=DCP=1 MVP; extending beyond that requires TP-only naming or canonicalizing the remaining axes first. This aligns naturally with the canonical KV layout direction in [KV Connector] Add canonical KV layout fields for TP-agnostic offload #46954.

Happy to converge on whichever direction fits the multi-tier design; a hybrid (A for node-local tiers, B for the default spec) is also plausible.

Coordination

This slots under RFC #38260 (its "canonical TP1 form" principle implies exactly one MLA copy in storage) and relates to #46954 (canonical KV layout fields — MLA latent modeled head-less) and #44865 (transfer data model refactor touching the same paths).

Not a duplicate of #38395 ("Fix/mla kv offloading", open draft): that patch keeps the world_size multiplication and per-TP-rank copies — it addresses MLA/HMA correctness and early disk-tier experimentation, not cross-rank replica dedup.

cc @orozery @Etelis @dannyharnik

Alternatives

• Document the caveat only (docs + kv_offloading_size docstring): cheap, but leaves a ~TPx capacity loss for exactly the models (DeepSeek-class) where KV offloading matters most.
• Point users to kv_offloading_backend="lmcache", which already dedups (MLA TP collapsing added in [DeepSeek + LMCache Multiprocess] handle MLA for deepseek model + LMCache Multiprocess connector #29039; the backend was wired to the LMCacheMPConnector in [KV Connector] Update lmcache kv_offloading_backend to use LMCacheMPConnector #42865): a valid workaround, but the native backend should not silently spend (TP-1)/TP of its budget on replicas.

Additional context

Duplicate-work check: searched open/closed issues and PRs (keyword sweeps, vllm/v1/kv_offload/ git history, and the comment threads of #19854 / #38260 / #33689). No existing issue or PR implements native MLA TP replica dedup. Closest items: #44733 (explicitly excludes MLA from parallel-agnostic caching), #46954 (layout schema only, no write-side dedup semantics), #44865 (overlapping-path refactor), #38395 (see above — retains per-rank copies).

I intend to follow up with a PR implementing the initial scope above, pending feedback on the preferred direction.

[orozery]

Thanks for taking on this @Change72 !

I like Candidate A.
We should tag replicated: bool = False at CanonicalKVCacheTensor to support that.
On the scheduler-side, since we don't have access to the canonical tensors metadata, we should deduce this from the KVCacheSpec.
Actually, it would be good if we can move all KVCacheSpec parsing logic outside from /v1/kv_offload to offloading_connector/ (e.g. the kv_bytes_per_block calculation in cpu/spec.py).

Regarding store completion: For start, I think we should still disperse store jobs to all workers to keep things simpler on the scheduler level. We can consider changing that as a follow-up.
For writer selection, we can start off with tp_rank=0 worker, and perhaps tweak that as a follow up.

[Change72]

Thanks @orozery — Candidate A makes sense.

The replicated flag also fits the p2p tier I am working on: once the row is compact, p2p transfers one slot instead of the full TP row. Workers can use the flag directly, while the scheduler derives the same property from KVCacheSpec.

One scope question before I start:

• Candidate A naturally applies to TieringOffloadingSpec, which already has a shared mmap.
• The default CPUOffloadingSpec still uses private per-worker buffers, so it would remain duplicated.
• Since CPUOffloadingSpec is the default native backend, should this be a Tiering-first MVP, or should the default backend also move to a shared region as part of this work?

If Tiering-first is the intended scope, I suggest this split:

1. Layout refactor: move the KVCacheSpec/KVCacheConfig-derived layout parsing from /v1/kv_offload into connector-owned config. Preserve the existing OffloadingSpec plugin contract. No intended runtime behavior change.
2. Tiering primary dedup: add the replicated tag and scheduler-side inference, then implement the single-slot shared-CPU row, matching sizing, and tp_rank=0 writer. Keep store jobs on all workers for now; non-writers skip D2H but still report completion. Gate this to pure MLA with PP=PCP=DCP=1.
3. Secondary tiers: reuse the compact row for fs and p2p, with TP-agnostic fs naming and p2p peer/config identity under the same gate.

Writer-only job dispatch/completion accounting and improved writer selection can remain follow-ups.

Does this scope and split match what you had in mind?

[orozery]

• or should the default backend also move to a shared region as part of this work?

We can open a separate independent PR to move the CPU spec to use SharedOffloadRegion.

Does this scope and split match what you had in mind?

Sounds great! Thanks!

[orozery]

BTW currently I think that SharedOffloadRegion will not work for XPU as it only supports CUDA pinning.
@jikunshang @chaojun-zhang what are your thoughts? Can XPU work well without pinning? can we make SharedOffloadRegion work with XPU pinning?

[orozery]

2. Tiering primary dedup: add the replicated tag

@Etelis should be adding replicated: bool at the CanonicalKVCacheRef level (not the tensor level) in #46954.

[Change72]

I'd like to take the CPUOffloadingSpec → SharedOffloadRegion follow-up @orozery suggested above, as two PRs:

1. Allocation swap. CPUOffloadingSpec workers open a SharedOffloadRegion instead of allocating private pin_memory=True tensors (CPUOffloadingWorker already takes an optional mmap_region). Per-worker copy semantics are unchanged; block accounting only picks up the region's page-size alignment. The region is an exactly-sized mmap, registered with cudaHostRegister where pinning is available, so the caching host allocator's power-of-2 rounding no longer applies — superseding the fix in [KV Offload] Bypass power-of-2 rounding in KV offload CPU pinned allocation #48436 (thanks @procr1337 for the diagnosis; I'll reproduce the RSS verification in the PR). Non-CUDA-alike platforms keep the current tensor path, so this doesn't depend on the XPU pinning question above.

2. Enable dedup for the default backend. Flip CPUOffloadingSpec.SUPPORTS_REPLICATED_LAYOUT under the existing gate from [KV Offload] Deduplicate replicated MLA KV in the shared CPU region #48906, so the default backend gets the same single-copy MLA row as tiering.

@Etelis: PR 1 touches the same allocation path as #48414; I'll keep the diff small to avoid rebase pain.

[Etelis]

I'll take a look today :)