[KV Offload] Deduplicate replicated MLA KV in the shared CPU region

[Change72]

Purpose

This implements the first shared CPU-region replica-reduction path from #47929.

For pure MLA tensor parallelism, each TP rank holds a replica of the latent KV payload. This duplicate host-copy path is present in both V1 and V2: both runners hand per-rank canonical MLA tensors to the same offloading connector, which previously reserved and stored one host slot per rank. The runners organize their GPU tensors differently, but replica ownership is the same.

This PR:

• derives a conservative replicated_layout decision in the offloading connector from KVCacheSpec and the parallel configuration;
• stores one copy per block in the shared CPU mmap and removes the TP replication factor from capacity sizing;
• lets TP rank 0 perform the D2H store while non-writers still acknowledge the dispatched job;
• lets every TP rank load the same shared host slot into its local GPU cache; and
• adds a DeepSeek-V2-Lite TP=2 shared CPU offload evaluation to the existing KV offload CI job.

At TP=N, this lets the configured CPU capacity hold approximately N times as many replicated MLA blocks and reduces D2H store traffic by a factor of N. Scheduler dispatch and completion accounting remain unchanged.

Scope

The replicated layout is enabled only for:

• a single bare MLAAttentionSpec group with matching per-layer page accounting;
• TP > 1 with PP = PCP = DCP = 1 and world_size == TP; and
• the single-node mp executor.

Both V1 and V2 model runners are supported under this gate. Mixed or wrapped cache layouts, sliding-window MLA, multi-node execution, and other parallel-axis combinations fail closed to the existing per-rank layout.

The default CPUOffloadingSpec remains unchanged because its buffers are process-private. This PR also does not add XPU shared-region pinning, change writer selection beyond rank 0, or change scheduler routing and completion accounting.

Shared CPU layout

SharedOffloadRegion remains unchanged. TieringOffloadingSpec selects a one-copy row size and maps every replicated worker to slot 0. Normal layouts still use one slot per worker. The existing mmap create, open, size-wait, and cleanup behavior is unchanged, and no metadata sidecar is added.

Compatibility notes

This PR changes replica ownership and sizing only in the shared CPU region. It does not make FS/OBJ/P2P ownership, naming, or cache identity parallel-agnostic. Persistent entries should be cleared when testing across the old TP-wide and new single-copy layouts. Secondary-storage identity and cleanup remain follow-up work.

V1 and V2 are supported independently under the same pure-MLA gate. This PR does not make persistent cache entries interchangeable between the two runners. #49440 addresses runner-specific persistent-cache namespacing separately.

Duplicate-work check

I checked #47929 and searched open PRs by issue number and MLA/offloading replication keywords. No open PR implements this shared CPU-region replica-reduction path.

#46954 was closed without merging in favor of #48408. That PR adds per-layer canonical KV page mappings for parallelism-agnostic offload. This PR instead needs an aggregate layout decision before canonical KV refs are registered, and the scheduler cannot see those refs. The changes are complementary; this PR does not depend on or modify #48408.

#49440 namespaces persistent cache entries by model runner. It does not deduplicate TP replicas or change shared-region ownership, so it is also separate from this PR.

Test plan

Workstation unit and GPU tests

.venv/bin/python -m pytest tests/v1/kv_offload/test_factory.py tests/v1/kv_offload/cpu/test_shared_offload_region.py tests/v1/kv_connector/unit/offloading_connector/test_config.py tests/v1/kv_connector/unit/offloading_connector/test_worker.py tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py -q

.venv/bin/python -m pytest tests/v1/kv_offload/cpu/test_gpu_worker.py -q

.venv/bin/pre-commit run --from-ref 6bcda970f --to-ref HEAD
.venv/bin/pre-commit run mypy-3.12 --all-files --hook-stage manual
git diff --check 6bcda970f...HEAD

TP=2 A100 model evaluation

.venv/bin/python -m pytest tests/evals/gsm8k/test_gsm8k_offloading.py -k deepseek-v2-lite -s -v

The committed evaluation uses TieringOffloadingSpec, TP=2, and the default V2 model runner. The V1 run used the same case with VLLM_USE_V2_MODEL_RUNNER=0.

Test results

The branch was rebased onto upstream main at 6bcda970f.

• focused factory, shared-region, connector config, connector worker, and scheduler tests on the workstation: 219 passed;
• L4 GPU transfer tests: 32 passed;
• all changed-file pre-commit hooks passed;
• full Python 3.12 mypy passed; and
• git diff --check passed.

The forced-V1 model evaluation ran on 2x NVIDIA A100-SXM4-80GB and passed: 1 passed, 8 deselected in 398.14s.

• run 1: accuracy 0.3600, invalid rate 0.005, latency 34.6s;
• run 2 after prefix-cache reset: accuracy 0.3450, invalid rate 0.005, latency 25.2s;
• both TP workers mapped rank 0 of the same shared mmap; and
• the configured 8 GiB CPU budget produced one 8,589,549,568-byte region with 17189 blocks.

The default-V2 model evaluation also ran on 2x NVIDIA A100-SXM4-80GB. The logs confirmed Using V2 Model Runner, and the probe-free PR tree passed: 1 passed, 8 deselected in 379.99s.

• run 1 accuracy: 0.3250;
• run 2 after prefix-cache reset: 0.3500;
• replay recorded 438,441,984 D2H store bytes and 876,883,968 H2D load bytes, matching one shared store and one load per TP rank; and
• validation-only instrumentation sampled the same canonical block on both TP ranks and produced the same SHA-256: 1b361a83278d97779bd65318aee1008c30480e50e5176c28244c1fd72a2dadac.

The A100 model evaluations ran before the metadata sidecar was removed. The sidecar did not control replica sizing, store ownership, or loading. The latest head was validated by the workstation unit and GPU suites above after restoring the original mmap lifecycle.

Source inspection and these runner-specific A100 runs confirm that the TP replica-ownership problem exists in both V1 and V2 and that the single-copy shared-region path works for each runner independently.

AI assistance

AI assistance was used for design exploration, implementation, test creation, and adversarial review. The human submitter reviewed every changed line and owns the change.

[coltonottley]

Thanks for taking this on. We independently built and qualified the same rank-0-store / all-rank-load ownership model on two independent TP=2 hosts, each with 2x RTX PRO 6000 Blackwell GPUs, using a downstream compact CPU representation for a hybrid DeepSeek V4 layout. This is not a test of this PR's exact commit; it is production evidence for the architecture and for a likely follow-on generalization beyond the pure-MLA MVP.

Real TP payload identity

We checksummed logical CPU payload bytes after GPU→CPU CUDA completion and before transfer-resource recycling. Across one real TP=2 request, both ranks emitted 126 exactly comparable offload keys spanning all five model-produced KV groups:

Group	Comparable keys	Matches	Mismatches	Logical bytes/key/rank
0	103	103	0	985,664
1	2	2	0	822,272
2	3	3	0	822,272
3	2	2	0	860,160
4	16	16	0	655,360
Total	126	126	0	—

The hashes covered logical payload bytes only and excluded page padding. Every observed full-history MLA, sliding-window/tail, compressor-state, and indexer payload was byte-identical across TP ranks. That supports this PR's conservative pure-MLA MVP and suggests the eventual canonical-mapping consumer can generalize replica reduction per layer/group once #48408 supplies authoritative mappings.

End-to-end ownership proof

Our downstream implementation kept scheduler job dispatch/completion accounting unchanged:

• rank 0 performed the physical D2H store;
• rank 1 returned a zero-byte success vote;
• the scheduler did not mark the key ready until all worker votes arrived;
• both ranks loaded directly from the same shared CPU bytes.

On the final immutable image, a deterministic 52,173-token semantic test produced two byte-identical cold answers. After GPU-only reset, replay restored 51,968 source-attributed external tokens with zero local GPU-cache hits, 206 recomputed tokens, 431,169,920 aggregate H2D bytes across both ranks, and exactly identical six-answer output in 0.872 s. The cold store wrote 216,407,232 bytes—one physical TP copy rather than two.

We also exercised two active ~171K-token requests while repeatedly attempting cache reset. Resets were refused while transfers were live, then succeeded after drain; a fresh post-reset request completed with the same process and no stale transfer-job references.

Shared-region lifecycle finding

We hit one operational failure worth carrying into review: after an ungraceful EngineCore exit, the shared mmap remained and blocked the next startup. A generic hint to rm /dev/shm/vllm_offload_<engine>.* is unsafe when multiple endpoint owners or generations can coexist.

The durable downstream fix uses an owner-scoped 0600 receipt containing exact version, owner key, fingerprint, byte size, and mmap path. Recovery takes flock, validates the complete schema/path/regular-file/size/fingerprint contract, and unlinks only when the host launcher proves that exact named endpoint owner is stopped or absent. Malformed receipts fail closed; receipt-write failure aborts creation; joiners never delete the creator receipt; already-absent exact mmaps clean up idempotently. Container PID was specifically rejected as host liveness evidence because it is in the wrong PID namespace and can be reused.

That lifecycle work is separable from replica dedup, but I strongly recommend replacing wildcard/manual cleanup before treating the shared primary as unattended-production safe.

The capacity result matched the model: canonical TP1 storage nearly halved resident bytes and D2H write traffic versus the replicated TP2 layout. Our measured compact density projected about 16.7M unique tokens per 64 GiB shared CPU region versus 8.35M with two private copies.

Happy to share focused receipt/recovery tests or review a follow-on against #48408. I am deliberately not opening a competing PR because #48906 owns this path.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Change72.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[coltonottley]

@Change72 — thank you for carrying the shared single-copy CPU-primary design forward and for tightening the stale-region cleanup after our production note.

I want to make one important boundary in our earlier evidence explicit, because our approximately 40× end-to-end result has two separate components and your current work covers one of them:

1. TP-rank replica deduplication: approximately 2× for TP=2. This is the broadly useful replicated-MLA improvement your current work implements: rank 0 stores one physical CPU copy and every rank loads from it.

2. Compact heterogeneous-group storage: approximately 19× on DeepSeek V4 Flash. The stock packed CPU row reserved roughly 187–195 KiB of server RAM per unique logical token while carrying only about 10–11 KiB of useful replay payload per TP rank. We replaced that coarse packed row with fixed-page scatter pools that store each hybrid group’s actual logical bytes rather than its padded alignment footprint.

Combined, compact storage and TP=2 replica deduplication reduced measured server-wide density to approximately 4.1 KiB per unique logical token—roughly 38×, reasonably summarized as about 40×, versus the original replicated packed layout.

This distinction matters because the approximately 19× component is not a general claim for ordinary uniform transformer KV caches. It comes from DeepSeek V4’s heterogeneous hybrid layout: full-history MLA, sliding-window groups, MTP/EAGLE state, compressor state, and indexer state with very different physical payload shapes forced into one coarse representation. The generic fixed-page mechanism may benefit future heterogeneous models, but the measured amplification is specific to this geometry.

Production results from the complete design:

• approximately 16.7 million unique tokens projected per 64 GiB shared CPU region;
• DeepSeek V4 Flash with MTP2, FP8 KV, and full 1M-token admission preserved;
• approximately 190 output tok/s decode on 2x RTX PRO 6000 Blackwell GPUs;
• exact 100K-token histories restored from CPU RAM in roughly 0.4–0.6 seconds after GPU eviction;
• verified on two independent TP=2 hosts, each with 2x RTX PRO 6000 Blackwell GPUs, plus a TP=4 qualification;
• natural production traffic is now showing 29.7% external-cache token hits on one endpoint and 8.7% on another.

The compact implementation also replaces contiguous variable-size extents with fixed 64 KiB pages. That removed the fatal spatial-fragmentation behavior we observed when long histories filled the CPU tier. Compact scatter transfers use per-copy asynchronous CUDA operations because the batched-copy API reproducibly crashed both TP workers above roughly 32,000 descriptors.

Our final compact/shared implementation currently lives in a downstream branch based on the pre-current-main scheduler lineage, so its commit hashes are not yet a useful public handoff. I do not want to dump an opaque downstream patch onto your work or open a competing implementation.

If you’re interested, we can publish a clean source branch or patch series, write a precise old-to-current scheduler mapping, and collaborate on porting the compact per-group page storage into the CPU-primary architecture you already own. We can carry the DeepSeek V4 TP=2 and TP=4 hardware qualification and provide the fragmentation, transfer-descriptor, lifecycle, reset, and semantic-replay regressions.

The clean ownership split seems to be:

• your current work remains the canonical shared CPU-primary and replica-deduplication path;
• compact per-group page storage becomes an extension for heterogeneous hybrid cache layouts;
• canonical per-layer/group mappings determine when compact storage and replica ownership are valid;
• unsupported layouts fail closed to the existing representation.

Would that direction be useful to you? If so, we’ll prepare the public branch and implementation map rather than making you reconstruct it from scattered benchmark comments.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Change72.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Change72]

Thanks @coltonottley. The 2x vs. 19x split is clear. Replica dedup is general; the 19x result is specific to the V4 heterogeneous layout.

Please go ahead. I suggest starting with a short RFC or a comment on #38260 covering the fixed-page pool, fail-closed fallback, and scheduler mapping. Please cc @orozery and include a clean reference branch. The implementation should build on #48408 and fall back to the current layout when those mappings are unavailable.

Please also file the crash above 32K copy descriptors as a separate bug.

My next work is FS/P2P reuse of the shared single-copy rows, as scoped on #47929 , so it should not overlap with per-group compaction.

[Change72]

@orozery One design choice I’d like your input on:

The existing size check cannot detect different layouts with the same total size. For example:

• normal: 2 blocks × 24 KiB = 48 KiB
• replicated: 4 blocks × 12 KiB = 48 KiB

I added one JSON sidecar per engine in /dev/shm. The creator writes the layout and mmap inode once; joiners validate it before mapping. The non-replicated path does not read or write this file.

Would you keep this validation, or avoid the extra file?

[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]

Hi @Change72, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[orozery]

Thanks @Change72 !

[coltonottley]

@Change72 — the requested package is published:

• short fixed-page / fail-closed / scheduler-mapping comment for [RFC]: Multi-tier KV offloading via the vLLM offloading connector #38260, CC @orozery: [RFC]: Multi-tier KV offloading via the vLLM offloading connector #38260 (comment)
• clean reference branch composed directly on current [KV Offload] Deduplicate replicated MLA KV in the shared CPU region #48906 head 47de1bd and [KV Connector] Add per-layer canonical KV page mappings for parallelism-agnostic offload #48408 head 7793ccf: https://github.com/coltonottley/vllm/tree/feature/native-compact-cpu-kv-current-main
• separate large/repeated batch-copy failure: [Bug]: cuMemcpyBatchAsync segfaults or hangs on large/repeated compact scatter submissions #49276

The reference branch inherits your writer gate and shared-primary ownership directly; it does not duplicate them. It consumes #48408 mappings, permanently falls back to the existing CPU representation when compact geometry cannot be certified, and keeps secondary-tier/P2P orchestration out of scope.

The full reference is intentionally larger than a comfortable single PR. The cover note maps it into three owner-aligned review units: fixed-page manager, mapping/scatter consumer, and scheduler admission/lifecycle. The exact immutable image passed full TP=2 production acceptance and fresh TP=4 qualification; refreshed-source validation and human approval are complete.