[KV Offload] Enable single-copy MLA layout for CPUOffloadingSpec

[Change72]

Purpose

Enable the single-copy (replicated) MLA layout for the default KV-offload backend CPUOffloadingSpec, closing out the plan in #47929. After #50094 moved the default backend's worker CPU buffer onto the shared SharedOffloadRegion (on CUDA/ROCm), the default backend can now deduplicate the per-TP-rank-replicated MLA KV the same way TieringOffloadingSpec already does — but only on deployments that actually allocate on the shared region.

MLA latent KV is logically replicated across TP ranks (expected to be byte-identical in homogeneous TP without context parallelism), so storing one copy instead of world_size copies is safe under the existing replicated_layout gate (pure-MLA single group, TP>1, PP=PCP=DCP=1, world==tp, single-node mp).

Maintainer context (upstream reviewers do not see the issue thread):

• The two-PR plan (allocation swap in [KV Offload] Move CPUOffloadingSpec onto SharedOffloadRegion #50094, then this dedup), proposed and +1'd by @orozery: [Feature]: Deduplicate replicated MLA KV across TP ranks in native offloading #47929 (comment)
• @orozery confirmed an independent PR is fine: [Feature]: Deduplicate replicated MLA KV across TP ranks in native offloading #47929 (comment)
• Prerequisite (merged): [KV Offload] Move CPUOffloadingSpec onto SharedOffloadRegion #50094 (CPUOffloadingSpec → SharedOffloadRegion).

What changes

• replicated_layout is gated by a new per-spec hook _uses_shared_region() — current_platform.is_cuda_alike() for the default spec, True for TieringOffloadingSpec. It is the single source of truth used by both the replicated-layout sizing (num_copies) and the create_worker allocation path, so the two cannot drift. This replaces the static SUPPORTS_REPLICATED_LAYOUT capability flag, which is removed: now that the default spec can deduplicate on the shared region, that flag was True for every in-tree spec and carried no information.
• CPUOffloadingSpec.create_worker now maps every rank to slot 0 under replicated layout (single MLA copy) and otherwise folds the physical device index into [0, world_size) — mirroring TieringOffloadingSpec.create_worker.
• TieringOffloadingSpec overrides _uses_shared_region() to True (it always allocates on the shared region, on every platform), so its behavior is unchanged.

Why the platform gate is load-bearing (fail-closed)

The replicated_layout config gate has no platform clause, and #50094 intentionally kept non-CUDA-alike platforms (currently XPU) on a per-rank private pinned tensor. Enabling dedup there would be a correctness bug, not a perf trade-off: the connector's rank-0 writer gate (offloading/worker.py, not spec.replicated_layout or rank == 0) would ack rank>0 stores without performing the D2H copy, leaving those private buffers empty and corrupting later loads. Gating the flip on _uses_shared_region() makes non-shared-region deployments fail closed to the existing private, per-rank layout; a dedicated test pins this.

Behavior changes

• Replicated-MLA deployments on the default backend: one host copy instead of world_size — cache capacity scales ×TP for the same cpu_bytes_to_use, and D2H store traffic drops to 1/TP (only rank 0 writes). Store-throughput metrics become writer-only, same as the tiering spec's existing disclosure.
• Non-CUDA-alike platforms: unchanged (private per-rank tensor, per-rank sizing).
• TieringOffloadingSpec: unchanged.
• Deduplication rests on the existing gate's premise that the per-rank MLA KV is logically replicated (expected byte-identical in homogeneous TP without context parallelism) — the same assumption TieringOffloadingSpec already relies on; this PR stores one copy instead of world_size. The default backend has no persistent tier, so there is no cross-version on-disk compatibility concern.

Not in this PR

• No change to the replicated_layout config gate predicate, the transfer protocol, or the tiering spec.
• XPU pinning is not addressed ([KV Offload] Move CPUOffloadingSpec onto SharedOffloadRegion #50094 already scoped around it; this PR inherits that boundary).
• Multi-group / non-MLA uniform-layout extensions remain future work (tracked with [KV Connector] Add per-layer canonical KV page mappings for parallelism-agnostic offload #48408).

Duplicate-work check

Searched open PRs (SUPPORTS_REPLICATED_LAYOUT, CPUOffloadingSpec replicated, 47929 in:body) and the touched-path history — no open PR enables dedup for the default CPU offload backend. Nearest neighbor: #48414 (@Etelis) works the parallelism-agnostic canonical allocation layout (the #50094 area); it does not touch the replicated-layout gate, so there is no overlap with this PR.

Test Plan

# Unit (offloading spec gating + sizing + create_worker)
.venv/bin/python -m pytest tests/v1/kv_offload/test_factory.py -q
pre-commit run --files vllm/v1/kv_offload/cpu/spec.py vllm/v1/kv_offload/tiering/spec.py tests/v1/kv_offload/test_factory.py

A100 TP=2 end-to-end (MLA model to exercise replicated_layout): serve DeepSeek-V2-Lite with -tp 2 and the OffloadingConnector, before (no dedup) vs after (dedup), and compare the connector's D2H store bytes and greedy output.

Test Result

Unit (NVIDIA L4, torch 2.13.0+cu130): tests/v1/kv_offload/test_factory.py — 38 passed. Covers the platform × config truth matrix, shared-region single-copy sizing (num_copies=1), the non-shared-region data-loss guard (replicated stays off, per-rank sizing preserved), and the replicated create_worker (rank 0). pre-commit (ruff, ruff-format, mypy, SPDX, torch.cuda-API guard) all pass.

A100 TP=2 E2E (2× A100-SXM4-80GB, DeepSeek-V2-Lite, --load-format dummy --seed 0 -tp 2, cpu_bytes_to_use=4 GiB, image vllm/vllm-openai:nightly-6f91edf96d whose v1/kv_offload/ is byte-identical to this PR's base — only the two changed spec files were overlaid for the "after" run):

	connector D2H store_bytes_total	greedy output
before (no dedup, world_size copies)	24.8832 MB	reference
after (dedup, single copy)	12.4416 MB	identical

The store bytes drop to exactly 1/TP — only rank 0 performs the D2H copy under the writer gate — confirming replicated_layout engages for the default backend on an MLA model. Greedy generation is unchanged for this deterministic dummy-weights run. (This run exercises the D2H store side; it does not force a CPU→GPU reload and is not itself a per-rank KV byte-equality proof.) The capacity ×TP effect is the num_copies=1 sizing, covered by the unit tests; the byte-exact GPU↔shared-region transfer is covered by #50094's test_transfer[use_shared_memory=True] and the tiering suite.

---

This PR was developed with AI assistance (implementation and tests). Unit tests and lint were run on the submitter's NVIDIA L4 workstation; the A100 TP=2 E2E was run on 2× A100-80GB.

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

[orozery]

Can we minimize this comment?

[Change72]

Done in 3161201.

[orozery]

Claude suggests this name instead test_cpu_spec_replicated_disabled_without_shared_region

[Change72]

Done in 3d396f3.

[orozery]

Claude claims this test can be made a bit less fragile:

@pytest.mark.parametrize(
    "replicated_layout, device_index, world_size, expected_rank",
    [
        (True, 5, 4, 0),       # replicated: always slot 0
        (True, 0, 4, 0),       # replicated: slot 0 regardless of device
        (False, 5, 4, 1),      # non-replicated: 5 % 4 = 1
        (False, 7, 4, 3),      # non-replicated: 7 % 4 = 3
    ],
)
def test_cpu_spec_create_worker_rank_assignment(
    monkeypatch, replicated_layout, device_index, world_size, expected_rank
):
    import vllm.v1.kv_offload.cpu.spec as cpu_spec_module

    monkeypatch.setattr(cpu_spec_module.current_platform, "is_cuda_alike", lambda: True)
    worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
    spec = _create_spec(
        cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
        worker_kv_bytes_per_block=worker_kv_bytes_per_block,
        world_size=world_size,
        replicated_layout=replicated_layout,
    )

    region_calls: list[dict[str, Any]] = []

    monkeypatch.setattr(
        cpu_spec_module, "SharedOffloadRegion", lambda **kw: region_calls.append(kw) or MagicMock()
    )
    monkeypatch.setattr(cpu_spec_module, "CPUOffloadingWorker", lambda **kw: MagicMock())
    monkeypatch.setattr(
        cpu_spec_module.torch.accelerator, "current_device_index", lambda: device_index
    )

    spec.create_worker(MagicMock())

    assert region_calls[0]["rank"] == expected_rank

What this fixes:

1. Asserts only the behavioral delta — rank is the only assertion. No kv_bytes_per_block (tested elsewhere), no mmap_region is region (trivial wiring).
2. Tests both branches — parametrizes replicated vs non-replicated, proving the if/else works correctly. This closes the gap I identified earlier (the "else" branch was untested after the refactor).
3. Descriptive parametrize IDs — when this test fails, the parametrize label tells you immediately whether it's the replicated path or the non-replicated path that broke.
4. Single assertion per concern — the test proves exactly one thing: "given (replicated_layout, device_index, world_size), the rank passed to SharedOffloadRegion is X." One assertion, one failure mode.

[Change72]

Done.

[orozery]

Thanks @Change72 !