[KV Connector] Add per-layer canonical KV page mappings for parallelism-agnostic offload (#48408)

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
Co-authored-by: Itay Etelis <Itay.etelis@gmail.com>

tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py

@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import (
+    _layer_mapping,
+    _opaque_fallback_mapping,
+    _RankContext,
+    _verify_tiling,
+    derive_canonical_mappings,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    KVQuantMode,
+    MLAAttentionSpec,
+)
+from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun
+
+NUM_BLOCKS = 3
+
+
+def _ctx(rank, tp=1, dcp=1, pcp=1, interleave=1, total=None, heads=2):
+    return _RankContext(
+        tp_size=tp,
+        dcp_size=dcp,
+        pcp_size=pcp,
+        interleave=interleave,
+        total_kv_heads=heads * tp if total is None else total,
+        rank=rank,
+    )
+
+
+def _full_spec(num_kv_heads: int = 2, **kwargs) -> FullAttentionSpec:
+    # block_size=4, head_dim=64, int8
+    return FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=num_kv_heads,
+        head_size=64,
+        dtype=torch.int8,
+        **kwargs,
+    )
+
+
+def _mla_spec(**kwargs) -> MLAAttentionSpec:
+    # page = 4 * 64 = 256B, one 64B latent row per token
+    return MLAAttentionSpec(
+        block_size=4, num_kv_heads=1, head_size=64, dtype=torch.int8, **kwargs
+    )
+
+
+def _split_nhd_cache(spec) -> torch.Tensor:
+    return torch.zeros(
+        NUM_BLOCKS,
+        2,
+        spec.block_size,
+        spec.num_kv_heads,
+        spec.head_size,
+        dtype=torch.int8,
+    )
+
+
+def _split_hnd_cache(spec) -> torch.Tensor:
+    return torch.zeros(
+        NUM_BLOCKS,
+        2,
+        spec.num_kv_heads,
+        spec.block_size,
+        spec.head_size,
+        dtype=torch.int8,
+    ).permute(0, 1, 3, 2, 4)
+
+
+def _packed_nhd_cache(spec) -> torch.Tensor:
+    """Logical (num_blocks, heads, block_size, 2 * head_size) over an NHD
+    physical layout — the FlashAttention/FlashInfer/Triton/Flex form."""
+    return torch.zeros(
+        NUM_BLOCKS,
+        spec.block_size,
+        spec.num_kv_heads,
+        2 * spec.head_size,
+        dtype=torch.int8,
+    ).permute(0, 2, 1, 3)
+
+
+def _packed_hnd_cache(spec) -> torch.Tensor:
+    return torch.zeros(
+        NUM_BLOCKS,
+        spec.num_kv_heads,
+        spec.block_size,
+        2 * spec.head_size,
+        dtype=torch.int8,
+    )
+
+
+CACHE_BUILDERS = {
+    "split_nhd": _split_nhd_cache,
+    "split_hnd": _split_hnd_cache,
+    "packed_nhd": _packed_nhd_cache,
+    "packed_hnd": _packed_hnd_cache,
+}
+
+
+def _try_mapping(spec, kv_cache, ctx) -> CanonicalPageMapping | None:
+    return _layer_mapping(spec, kv_cache, NUM_BLOCKS, ctx)
+
+
+def _mapping(spec, kv_cache, ctx) -> CanonicalPageMapping:
+    mapping = _try_mapping(spec, kv_cache, ctx)
+    assert mapping is not None
+    return mapping
+
+
+def _triples(runs: tuple[CopyRun, ...]) -> list[tuple[int, int, int]]:
+    """Expand runs to explicit (local_offset, canonical_offset, size) copies."""
+    out = []
+    for run in runs:
+        for i in range(run.num_fragments):
+            out.append(
+                (
+                    run.local_offset + i * run.local_stride,
+                    run.canonical_offset + i * run.canonical_stride,
+                    run.fragment_size,
+                )
+            )
+    return out
+
+
+# ---------------------------------------------------------------------------
+# TP-only placement (byte-compatible with the uniform interleave layout)
+# ---------------------------------------------------------------------------
+
+
+def test_split_nhd_placement_rank2_of_4():
+    spec = _full_spec()
+    mapping = _mapping(spec, _split_nhd_cache(spec), _ctx(rank=2, tp=4))
+    assert mapping.canonical_page_size_bytes == 4 * 1024
+    assert mapping.parallelism_agnostic
+    k_dst = [256, 768, 1280, 1792]
+    assert _triples(mapping.runs) == [
+        (local, canonical, 128)
+        for local, canonical in zip(
+            [0, 128, 256, 384, 512, 640, 768, 896],
+            k_dst + [2048 + o for o in k_dst],
+        )
+    ]
+    # Heads are sharded, not replicated: this rank writes every block
+    assert mapping.num_writers == 1
+
+
+def test_packed_nhd_placement_rank2_of_4():
+    spec = _full_spec()
+    mapping = _mapping(spec, _packed_nhd_cache(spec), _ctx(rank=2, tp=4))
+    assert _triples(mapping.runs) == [
+        (0, 512, 256),
+        (256, 1536, 256),
+        (512, 2560, 256),
+        (768, 3584, 256),
+    ]
+
+
+def test_packed_hnd_placement_rank1_of_4():
+    spec = _full_spec()
+    mapping = _mapping(spec, _packed_hnd_cache(spec), _ctx(rank=1, tp=4))
+    # Two heads, each a contiguous canonical head region of 4 tokens x 128B
+    assert _triples(mapping.runs) == [(0, 1024, 1024)]
+
+
+@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS))
+def test_single_rank_coalesces_to_one_run(form):
+    spec = _full_spec()
+    mapping = _mapping(spec, CACHE_BUILDERS[form](spec), _ctx(rank=0))
+    assert mapping.canonical_page_size_bytes == 1024
+    assert _triples(mapping.runs) == [(0, 0, 1024)]
+
+
+# ---------------------------------------------------------------------------
+# Replication and writer election
+# ---------------------------------------------------------------------------
+
+
+def test_gqa_replicated_heads_rotate_writer():
+    # total 2 KV heads on tp=4: replication factor 2, head shard = rank // 2
+    spec = _full_spec(num_kv_heads=1)
+    cache = _split_nhd_cache(spec)
+    ctx = lambda rank: _ctx(rank, tp=4, total=2)  # noqa: E731
+    rank2 = _mapping(spec, cache, ctx(2))
+    rank3 = _mapping(spec, cache, ctx(3))
+    assert rank2.canonical_page_size_bytes == 2 * 512
+    # K region: head shard 1 at 64B offsets within 128B token rows
+    assert _triples(rank2.runs)[:4] == [
+        (0, 64, 64),
+        (64, 192, 64),
+        (128, 320, 64),
+        (192, 448, 64),
+    ]
+    # Same head shard, so identical bytes: the two take alternate blocks
+    assert rank3.runs == rank2.runs
+    assert (rank2.is_writer(0), rank3.is_writer(0)) == (True, False)
+    assert (rank2.is_writer(1), rank3.is_writer(1)) == (False, True)
+    _verify_tiling("gqa", [_mapping(spec, cache, ctx(r)) for r in range(4)])
+
+
+def test_mla_replicas_rotate_writer():
+    spec = _mla_spec()
+    rank0 = _mapping(spec, None, _ctx(rank=0, tp=2))
+    rank1 = _mapping(spec, None, _ctx(rank=1, tp=2))
+    # Latent pages are stored once per block, not once per rank
+    assert rank0.canonical_page_size_bytes == 256
+    assert _triples(rank0.runs) == [(0, 0, 256)]
+    assert rank1.runs == rank0.runs
+    assert (rank0.is_writer(0), rank1.is_writer(0)) == (True, False)
+    assert (rank0.is_writer(1), rank1.is_writer(1)) == (False, True)
+
+
+# ---------------------------------------------------------------------------
+# DCP / PCP token sharding
+# ---------------------------------------------------------------------------
+
+
+def test_dcp_interleaves_tokens_within_replicas():
+    # tp=4, dcp=2, total 2 KV heads: head shard = rank // 2, cp rank = rank % 2
+    spec = _full_spec(num_kv_heads=1)
+    cache = _split_nhd_cache(spec)
+    ctx = lambda rank: _ctx(rank, tp=4, dcp=2, total=2)  # noqa: E731
+    per_rank = [_mapping(spec, cache, ctx(rank)) for rank in range(4)]
+    assert all(m is not None for m in per_rank)
+    # 8 canonical tokens x 2 heads x 64B per region
+    assert per_rank[0].canonical_page_size_bytes == 2048
+    assert not per_rank[0].parallelism_agnostic
+    # rank 2 = head shard 1, cp rank 0: K tokens 0,2,4,6 at head offset 64
+    assert _triples(per_rank[2].runs)[:4] == [
+        (0, 64, 64),
+        (64, 320, 64),
+        (128, 576, 64),
+        (192, 832, 64),
+    ]
+    # every rank contributes (dcp == replication: no residual replicas)
+    assert all(m.num_writers == 1 for m in per_rank)
+    _verify_tiling("dcp", per_rank)
+
+
+def test_mla_dcp_shards_latent_tokens():
+    spec = _mla_spec()
+    ctx = lambda rank: _ctx(rank, tp=2, dcp=2)  # noqa: E731
+    rank0 = _mapping(spec, None, ctx(0))
+    rank1 = _mapping(spec, None, ctx(1))
+    assert rank0.canonical_page_size_bytes == 512
+    assert _triples(rank0.runs) == [(o, 2 * o, 64) for o in (0, 64, 128, 192)]
+    assert _triples(rank1.runs) == [(o, 2 * o + 64, 64) for o in (0, 64, 128, 192)]
+    _verify_tiling("mla-dcp", [rank0, rank1])
+
+
+def test_pcp_tokens_and_tp_heads_compose():
+    # tp=2 x pcp=2: rank = pcp_rank * 2 + tp_rank; 4 workers tile the page
+    spec = _full_spec(num_kv_heads=1)
+    cache = _packed_nhd_cache(spec)
+    per_rank = [
+        _mapping(spec, cache, _ctx(rank, tp=2, pcp=2, total=2)) for rank in range(4)
+    ]
+    assert all(m is not None and m.runs for m in per_rank)
+    _verify_tiling("pcp", per_rank)
+
+
+def test_interleave_chunks_stay_contiguous():
+    # interleave=2: chunks of 2 tokens alternate between the 2 cp ranks and
+    # coalesce into one contiguous fragment per chunk
+    spec = _mla_spec()
+    mapping = _mapping(spec, None, _ctx(rank=0, tp=2, dcp=2, interleave=2))
+    assert _triples(mapping.runs) == [(0, 0, 128), (128, 256, 128)]
+    _verify_tiling(
+        "interleave",
+        [
+            _mapping(spec, None, _ctx(rank, tp=2, dcp=2, interleave=2))
+            for rank in range(2)
+        ],
+    )
+
+
+@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS))
+def test_all_ranks_tile_canonical_page(form):
+    spec = _full_spec()
+    per_rank = [
+        _mapping(spec, CACHE_BUILDERS[form](spec), _ctx(rank, tp=4))
+        for rank in range(4)
+    ]
+    _verify_tiling("layer", per_rank)
+
+
+# ---------------------------------------------------------------------------
+# Byte-level round trips
+# ---------------------------------------------------------------------------
+
+
+def _store_all(mappings, pages, size: int, block_id: int = 0) -> bytes:
+    buf = bytearray(size)
+    for mapping, page in zip(mappings, pages):
+        if not mapping.is_writer(block_id):
+            continue
+        for local, canonical, n in _triples(mapping.runs):
+            buf[canonical : canonical + n] = page[local : local + n]
+    return bytes(buf)
+
+
+def _load_one(mapping, canonical_bytes: bytes) -> bytes:
+    page = bytearray(mapping.local_page_size_bytes)
+    for local, canonical, n in _triples(mapping.runs):
+        page[local : local + n] = canonical_bytes[canonical : canonical + n]
+    return bytes(page)
+
+
+@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS))
+def test_cross_tp_store_load(form):
+    """Bytes stored under one TP size are the bytes another TP size loads."""
+    total_heads, canonical_size = 8, 4096
+    reference = bytes((7 + 31 * i) % 256 for i in range(canonical_size))
+
+    def mappings_at(tp: int):
+        spec = _full_spec(num_kv_heads=total_heads // tp)
+        cache = CACHE_BUILDERS[form](spec)
+        return [
+            _mapping(spec, cache, _ctx(rank, tp=tp, total=total_heads))
+            for rank in range(tp)
+        ]
+
+    for tp in (4, 2, 1):
+        mappings = mappings_at(tp)
+        assert all(m is not None for m in mappings)
+        pages = [_load_one(m, reference) for m in mappings]
+        assert _store_all(mappings, pages, canonical_size) == reference
+
+
+def test_cp_round_trip():
+    # tp=4 / dcp=2 / 2 KV heads: 4 workers jointly hold one canonical page
+    spec = _full_spec(num_kv_heads=1)
+    cache = _split_nhd_cache(spec)
+    mappings = [
+        _mapping(spec, cache, _ctx(rank, tp=4, dcp=2, total=2)) for rank in range(4)
+    ]
+    reference = bytes((3 + 17 * i) % 256 for i in range(2048))
+    pages = [_load_one(m, reference) for m in mappings]
+    assert _store_all(mappings, pages, 2048) == reference
+
+
+def test_replica_rotation_round_trip():
+    """Whichever replica a block elects reproduces the same canonical page."""
+    spec = _mla_spec()
+    mappings = [_mapping(spec, None, _ctx(rank, tp=2)) for rank in range(2)]
+    reference = bytes((5 + 11 * i) % 256 for i in range(256))
+    pages = [_load_one(m, reference) for m in mappings]
+    for block_id in (0, 1):
+        assert _store_all(mappings, pages, 256, block_id) == reference
+
+
+# ---------------------------------------------------------------------------
+# Fail-closed gates
+# ---------------------------------------------------------------------------
+
+
+def test_fail_closed_cases():
+    spec = _full_spec()
+    nhd = _split_nhd_cache(spec)
+    # Spec heads inconsistent with total heads / tp
+    assert _try_mapping(spec, nhd, _ctx(0, tp=4, total=2)) is None
+    # tp not divisible by total KV heads
+    one_head = _full_spec(num_kv_heads=1)
+    assert (
+        _try_mapping(one_head, _split_nhd_cache(one_head), _ctx(0, tp=3, total=2))
+        is None
+    )
+    # DCP wider than the KV replication factor (tokens would shard across
+    # ranks holding different heads)
+    assert _try_mapping(spec, nhd, _ctx(0, tp=4, dcp=2)) is None
+    # Interleave must divide the block size
+    assert _try_mapping(spec, nhd, _ctx(0, tp=2, dcp=2, interleave=3, total=2)) is None
+    # Per-token-head scales are packed with the data
+    quant_spec = _full_spec(kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD)
+    assert _try_mapping(quant_spec, _split_nhd_cache(quant_spec), _ctx(0, tp=4)) is None
+    # Compressed MLA slots are not 1:1 with tokens
+    assert _try_mapping(_mla_spec(compress_ratio=2), None, _ctx(0, tp=2, dcp=2)) is None
+    # Unrecognized physical layouts
+    swapped = torch.zeros(
+        NUM_BLOCKS,
+        spec.block_size,
+        2,
+        spec.num_kv_heads,
+        spec.head_size,
+        dtype=torch.int8,
+    ).permute(0, 2, 1, 3, 4)
+    assert _try_mapping(spec, swapped, _ctx(0, tp=4)) is None
+    # Non-attention specs
+    assert _try_mapping(KVCacheSpec(block_size=4), None, _ctx(0, tp=4, total=8)) is None
+
+
+def test_opaque_fallback_places_page_whole():
+    mapping = _opaque_fallback_mapping(1024, 4, 2)
+    assert mapping.canonical_page_size_bytes == 4096
+    assert not mapping.parallelism_agnostic
+    assert _triples(mapping.runs) == [(0, 2048, 1024)]
+    assert mapping.num_writers == 1
+    _verify_tiling("opaque", [_opaque_fallback_mapping(1024, 4, r) for r in range(4)])
+
+
+# ---------------------------------------------------------------------------
+# derive_canonical_mappings end to end
+# ---------------------------------------------------------------------------
+
+
+def _vllm_config(tp=1, dcp=1, pcp=1, pp=1, interleave=1, total_kv_heads=2):
+    config = MagicMock()
+    config.parallel_config.tensor_parallel_size = tp
+    config.parallel_config.decode_context_parallel_size = dcp
+    config.parallel_config.prefill_context_parallel_size = pcp
+    config.parallel_config.cp_kv_cache_interleave_size = interleave
+    config.parallel_config.world_size = pp * tp * pcp
+    config.parallel_config.rank = 0
+    config.model_config.get_total_num_kv_heads.return_value = total_kv_heads
+    return config
+
+
+def _kv_cache_config(groups):
+    config = MagicMock()
+    config.kv_cache_groups = groups
+    config.num_blocks = NUM_BLOCKS
+    return config
+
+
+def test_derive_mixed_model_with_dcp():
+    attn_spec = _full_spec(num_kv_heads=1)
+    mla_spec = _mla_spec()
+    quant_spec = _full_spec(
+        num_kv_heads=1, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD
+    )
+    kv_cache_config = _kv_cache_config(
+        [
+            KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=attn_spec),
+            KVCacheGroupSpec(layer_names=["mla"], kv_cache_spec=mla_spec),
+            KVCacheGroupSpec(layer_names=["quant"], kv_cache_spec=quant_spec),
+        ]
+    )
+    kv_caches = {
+        "attn": _split_nhd_cache(attn_spec),
+        "quant": _split_nhd_cache(quant_spec),
+    }
+    mappings = derive_canonical_mappings(
+        _vllm_config(tp=4, dcp=2, total_kv_heads=2), kv_cache_config, kv_caches
+    )
+    assert set(mappings) == {"attn", "mla", "quant"}
+    assert not mappings["attn"].parallelism_agnostic
+    assert mappings["attn"].runs
+    # Uncertifiable layers degrade to an opaque page, never disappear
+    assert not mappings["quant"].parallelism_agnostic
+    assert (
+        mappings["quant"].canonical_page_size_bytes
+        == 4 * quant_spec.unpadded_page_size_bytes
+    )
+
+
+def test_derive_refuses_foreign_worker_groups():
+    attn_spec = _full_spec()
+    kv_cache_config = _kv_cache_config(
+        [KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=attn_spec)]
+    )
+    kv_caches = {"attn": _split_nhd_cache(attn_spec)}
+    assert (
+        derive_canonical_mappings(
+            _vllm_config(tp=2, pp=2, total_kv_heads=4), kv_cache_config, kv_caches
+        )
+        == {}
+    )

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

@@ -101,6 +101,21 @@ def _allocate_and_reshape_kv_caches(
         set_kv_cache_layout(None)
 
 
+def _single_rank_vllm_config(total_kv_heads: int):
+    """A one-rank (TP=1) parallel config, as canonical mappings are derived
+    from it."""
+    vllm_config = MagicMock()
+    parallel_config = vllm_config.parallel_config
+    parallel_config.tensor_parallel_size = 1
+    parallel_config.decode_context_parallel_size = 1
+    parallel_config.prefill_context_parallel_size = 1
+    parallel_config.cp_kv_cache_interleave_size = 1
+    parallel_config.world_size = 1
+    parallel_config.rank = 0
+    vllm_config.model_config.get_total_num_kv_heads.return_value = total_kv_heads
+    return vllm_config
+
+
 def _make_worker(
     kv_cache_config: KVCacheConfig,
     replicated_layout: bool = False,
@@ -121,6 +136,7 @@ def _make_worker(
 
     worker = OffloadingConnectorWorker(
         spec=spec,
+        vllm_config=_single_rank_vllm_config(NUM_KV_HEADS),
         kv_cache_config=kv_cache_config,
     )
     worker.worker = MagicMock()
@@ -287,6 +303,7 @@ def test_offloading_connector_worker_accepts_plugin_spec_default_layout():
 
     OffloadingConnectorWorker(
         spec=spec,
+        vllm_config=_single_rank_vllm_config(NUM_KV_HEADS),
         kv_cache_config=KVCacheConfig(
             num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
         ),
@@ -522,6 +539,8 @@ def test_register_kv_caches(backend):
         for actual, expected in zip(actual_refs, exp_refs):
             assert actual.tensor_idx == expected.tensor_idx
             assert actual.page_size_bytes == expected.page_size_bytes
+            # Every layer gets a canonical mapping, certified or opaque
+            assert actual.mapping is not None
 
 
 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
@@ -622,9 +641,15 @@ def test_register_kv_caches_uniform_type(backend):
     assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
     assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
 
-    assert group_refs[0] == CanonicalKVCacheRef(
-        tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
-    )
-    assert group_refs[1] == CanonicalKVCacheRef(
-        tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
-    )
+    for ref, expected_tensor_idx, expected_spec in (
+        (group_refs[0], 0, spec_a),
+        (group_refs[1], 1, spec_b),
+    ):
+        assert ref.tensor_idx == expected_tensor_idx
+        assert ref.page_size_bytes == expected_spec.page_size_bytes
+        assert ref.mapping is not None
+
+    # Only layer_a matches the model's total KV head count, so layer_b gets an
+    # opaque mapping rather than a certified, parallelism-agnostic one
+    assert group_refs[0].mapping.parallelism_agnostic
+    assert not group_refs[1].mapping.parallelism_agnostic

vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py

@@ -0,0 +1,444 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Derivation of canonical page mappings for KV offloading.
+
+The only place in the offloading stack that reasons about parallelism
+(TP/DCP/PCP); everything downstream consumes byte mappings. The canonical
+page of a layer is the full offloaded block without parallelism: all KV
+heads, all block_size * dcp * pcp tokens, in the worker's page encoding.
+Uncertifiable layers get an opaque fallback mapping (fail closed).
+"""
+
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheConfig,
+    KVCacheSpec,
+    MambaSpec,
+    MLAAttentionSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+
+@dataclass(frozen=True)
+class _RankContext:
+    """Sharding parameters of one worker rank within the offload group."""
+
+    tp_size: int
+    dcp_size: int
+    pcp_size: int
+    interleave: int
+    total_kv_heads: int
+    rank: int
+
+    @property
+    def cp_size(self) -> int:
+        return self.dcp_size * self.pcp_size
+
+    @property
+    def tp_rank(self) -> int:
+        return self.rank % self.tp_size
+
+    @property
+    def total_cp_rank(self) -> int:
+        pcp_rank = self.rank // self.tp_size
+        return pcp_rank * self.dcp_size + self.tp_rank % self.dcp_size
+
+
+@dataclass(frozen=True)
+class ByteRegion:
+    """A byte region within a page that repeats once per token."""
+
+    local_offset: int
+    canonical_offset: int
+    bytes_per_token: int
+    canonical_token_stride: int
+
+
+def _coalesce_runs(runs: list[CopyRun]) -> tuple[CopyRun, ...]:
+    """Collapse contiguous fragments within and across runs to minimize the
+    number of copy ops (e.g. a single-rank mapping becomes one whole-page run).
+    """
+    out: list[CopyRun] = []
+    for run in runs:
+        if (
+            run.num_fragments > 1
+            and run.local_stride == run.fragment_size
+            and run.canonical_stride == run.fragment_size
+        ):
+            size = run.fragment_size * run.num_fragments
+            run = CopyRun(run.local_offset, run.canonical_offset, size, 1, size, size)
+        prev = out[-1] if out else None
+        if (
+            prev is not None
+            and prev.num_fragments == 1
+            and run.num_fragments == 1
+            and prev.local_offset + prev.fragment_size == run.local_offset
+            and prev.canonical_offset + prev.fragment_size == run.canonical_offset
+        ):
+            size = prev.fragment_size + run.fragment_size
+            out[-1] = CopyRun(
+                prev.local_offset, prev.canonical_offset, size, 1, size, size
+            )
+        else:
+            out.append(run)
+    return tuple(out)
+
+
+def _local_to_canonical_token(local_idx: int, ctx: _RankContext) -> int:
+    """Canonical position of one of this rank's local token indices."""
+    chunk, pos_in_chunk = divmod(local_idx, ctx.interleave)
+    return (chunk * ctx.cp_size + ctx.total_cp_rank) * ctx.interleave + pos_in_chunk
+
+
+def _interleave_cp_tokens(
+    regions: list[ByteRegion],
+    num_tokens: int,
+    ctx: _RankContext,
+) -> tuple[CopyRun, ...]:
+    """Place each region's num_tokens rows at their canonical token positions,
+    one run per chunk of interleaved tokens."""
+    runs: list[CopyRun] = []
+    for region in regions:
+        if ctx.cp_size == 1:
+            runs.append(
+                CopyRun(
+                    region.local_offset,
+                    region.canonical_offset,
+                    region.bytes_per_token,
+                    num_tokens,
+                    region.bytes_per_token,
+                    region.canonical_token_stride,
+                )
+            )
+            continue
+        for chunk_start in range(0, num_tokens, ctx.interleave):
+            canonical_token = _local_to_canonical_token(chunk_start, ctx)
+            runs.append(
+                CopyRun(
+                    region.local_offset + chunk_start * region.bytes_per_token,
+                    region.canonical_offset
+                    + canonical_token * region.canonical_token_stride,
+                    region.bytes_per_token,
+                    ctx.interleave,
+                    region.bytes_per_token,
+                    region.canonical_token_stride,
+                )
+            )
+    return _coalesce_runs(runs)
+
+
+def _packed_kv_regions(
+    kv_cache: torch.Tensor,
+    spec: AttentionSpec,
+    head_shard: int,
+    num_head_shards: int,
+    cp_size: int,
+) -> list[ByteRegion] | None:
+    """K and V adjacent per (token, head), in NHD or HND stride order."""
+    bs, heads = spec.block_size, spec.num_kv_heads
+    elem = kv_cache.element_size()
+    head_elems = 2 * spec.head_size
+    if heads * bs * head_elems * elem != spec.real_page_size_bytes:
+        return None
+    _, head_stride, token_stride, inner_stride = kv_cache.stride()
+    if inner_stride != 1:
+        return None
+    head_bytes = head_elems * elem
+    token_row_bytes = heads * head_bytes
+
+    if head_stride == head_elems and token_stride == heads * head_elems:  # NHD
+        return [
+            ByteRegion(
+                local_offset=0,
+                canonical_offset=head_shard * token_row_bytes,
+                bytes_per_token=token_row_bytes,
+                canonical_token_stride=num_head_shards * token_row_bytes,
+            )
+        ]
+
+    if head_stride == bs * head_elems and token_stride == head_elems:  # HND
+        canonical_span = bs * cp_size  # canonical tokens per offloaded block
+        return [
+            ByteRegion(
+                local_offset=head * bs * head_bytes,
+                canonical_offset=(head_shard * heads + head)
+                * canonical_span
+                * head_bytes,
+                bytes_per_token=head_bytes,
+                canonical_token_stride=head_bytes,
+            )
+            for head in range(heads)
+        ]
+    return None
+
+
+def _split_kv_regions(
+    kv_cache: torch.Tensor,
+    spec: AttentionSpec,
+    head_shard: int,
+    num_head_shards: int,
+    cp_size: int,
+) -> list[ByteRegion] | None:
+    """K and V in separate page halves, in NHD or HND stride order."""
+    bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size
+    elem = kv_cache.element_size()
+    if 2 * bs * heads * head_size * elem != spec.real_page_size_bytes:
+        return None
+    _, half_stride, token_stride, head_stride, inner_stride = kv_cache.stride()
+    if inner_stride != 1 or half_stride != bs * heads * head_size:
+        return None
+    head_bytes = head_size * elem
+    token_row_bytes = heads * head_bytes
+    canonical_span = bs * cp_size  # canonical tokens per offloaded block
+
+    if token_stride == heads * head_size and head_stride == head_size:  # NHD
+        canonical_half_bytes = canonical_span * num_head_shards * token_row_bytes
+        return [
+            ByteRegion(
+                local_offset=half * bs * token_row_bytes,
+                canonical_offset=half * canonical_half_bytes
+                + head_shard * token_row_bytes,
+                bytes_per_token=token_row_bytes,
+                canonical_token_stride=num_head_shards * token_row_bytes,
+            )
+            for half in range(2)  # K, then V
+        ]
+
+    if head_stride == bs * head_size and token_stride == head_size:  # HND
+        local_half_bytes = bs * heads * head_bytes
+        total_heads = num_head_shards * heads
+        return [
+            ByteRegion(
+                local_offset=half * local_half_bytes + head * bs * head_bytes,
+                canonical_offset=(half * total_heads + head_shard * heads + head)
+                * canonical_span
+                * head_bytes,
+                bytes_per_token=head_bytes,
+                canonical_token_stride=head_bytes,
+            )
+            for half in range(2)
+            for head in range(heads)
+        ]
+    return None
+
+
+def _attention_byte_regions(
+    kv_cache: torch.Tensor,
+    spec: AttentionSpec,
+    num_blocks: int,
+    head_shard: int,
+    num_head_shards: int,
+    cp_size: int,
+) -> list[ByteRegion] | None:
+    """Byte regions of an attention page, given this rank's head shard.
+    None when the physical layout is not recognized (fail closed)."""
+    bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size
+    if tuple(kv_cache.shape) == (num_blocks, heads, bs, 2 * head_size):
+        return _packed_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size)
+    if tuple(kv_cache.shape) == (num_blocks, 2, bs, heads, head_size):
+        return _split_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size)
+    return None
+
+
+def _layer_mapping(
+    spec: KVCacheSpec,
+    kv_cache: torch.Tensor | list[torch.Tensor] | None,
+    num_blocks: int,
+    ctx: _RankContext,
+) -> CanonicalPageMapping | None:
+    """Certified mapping for one layer at one rank, or None (fail closed)."""
+    if not isinstance(spec, AttentionSpec):
+        return None
+    bs = spec.block_size
+    page = spec.real_page_size_bytes
+    if ctx.cp_size > 1 and (ctx.interleave > bs or bs % ctx.interleave):
+        return None
+
+    if isinstance(spec, MLAAttentionSpec):
+        # TP-replicated latent; CP shards its tokens across the DCP groups
+        if (
+            spec.compress_ratio != 1
+            or page % bs
+            or ctx.tp_size % ctx.dcp_size
+            or spec.kv_quant_mode.is_per_token_head
+        ):
+            return None
+        row = page // bs
+        return CanonicalPageMapping(
+            canonical_page_size_bytes=ctx.cp_size * page,
+            local_page_size_bytes=page,
+            runs=_interleave_cp_tokens([ByteRegion(0, 0, row, row)], bs, ctx),
+            num_writers=ctx.tp_size // ctx.dcp_size,
+            writer_index=ctx.tp_rank // ctx.dcp_size,
+            parallelism_agnostic=ctx.cp_size == 1,
+        )
+
+    if spec.kv_quant_mode.is_per_token_head or not isinstance(kv_cache, torch.Tensor):
+        return None
+    total, tp = ctx.total_kv_heads, ctx.tp_size
+    if spec.num_kv_heads != max(1, total // tp):
+        return None
+    if total >= tp:
+        if total % tp:
+            return None
+        num_head_shards, replication = tp, 1
+    else:
+        if tp % total:
+            return None
+        num_head_shards, replication = total, tp // total
+    # DCP shards tokens across ranks holding replicated KV
+    if replication % ctx.dcp_size:
+        return None
+
+    head_shard = ctx.tp_rank // replication
+    regions = _attention_byte_regions(
+        kv_cache, spec, num_blocks, head_shard, num_head_shards, ctx.cp_size
+    )
+    if regions is None:
+        return None
+    return CanonicalPageMapping(
+        canonical_page_size_bytes=ctx.cp_size * num_head_shards * page,
+        local_page_size_bytes=page,
+        runs=_interleave_cp_tokens(regions, bs, ctx),
+        num_writers=replication // ctx.dcp_size,
+        writer_index=(ctx.tp_rank % replication) // ctx.dcp_size,
+        parallelism_agnostic=ctx.cp_size == 1,
+    )
+
+
+def _opaque_fallback_mapping(
+    page_size_bytes: int, num_ranks: int, rank: int
+) -> CanonicalPageMapping:
+    """Fallback: place the worker's page whole at a worker-exclusive offset."""
+    run = CopyRun(
+        0, rank * page_size_bytes, page_size_bytes, 1, page_size_bytes, page_size_bytes
+    )
+    return CanonicalPageMapping(
+        canonical_page_size_bytes=num_ranks * page_size_bytes,
+        local_page_size_bytes=page_size_bytes,
+        runs=(run,),
+        num_writers=1,
+        writer_index=0,
+        parallelism_agnostic=False,
+    )
+
+
+def _run_intervals(runs: tuple[CopyRun, ...], canonical: bool) -> list[tuple[int, int]]:
+    intervals = []
+    for run in runs:
+        offset = run.canonical_offset if canonical else run.local_offset
+        stride = run.canonical_stride if canonical else run.local_stride
+        for i in range(run.num_fragments):
+            start = offset + i * stride
+            intervals.append((start, start + run.fragment_size))
+    return sorted(intervals)
+
+
+def _is_exact_partition(intervals: list[tuple[int, int]], size: int) -> bool:
+    return (
+        bool(intervals)
+        and intervals[0][0] == 0
+        and intervals[-1][1] == size
+        and all(a[1] == b[0] for a, b in zip(intervals, intervals[1:]))
+    )
+
+
+def _verify_tiling(layer_name: str, per_rank: list[CanonicalPageMapping]) -> None:
+    """Whichever ranks a block elects as writers must tile the canonical page
+    exactly once, and each rank's runs must cover exactly its local page."""
+    size = per_rank[0].canonical_page_size_bytes
+    num_writers = per_rank[0].num_writers
+    for mapping in per_rank:
+        assert mapping.canonical_page_size_bytes == size
+        assert mapping.num_writers == num_writers
+        local = _run_intervals(mapping.runs, canonical=False)
+        assert _is_exact_partition(local, mapping.local_page_size_bytes), (
+            f"runs do not cover the local page of layer {layer_name}"
+        )
+    for block_id in range(num_writers):
+        stored: list[tuple[int, int]] = []
+        for mapping in per_rank:
+            if mapping.is_writer(block_id):
+                stored += _run_intervals(mapping.runs, canonical=True)
+        stored.sort()
+        assert _is_exact_partition(stored, size), (
+            f"writers of block {block_id} do not tile the canonical page "
+            f"of layer {layer_name}"
+        )
+
+
+def _unpadded_page_size(spec: KVCacheSpec) -> int | None:
+    if isinstance(spec, AttentionSpec):
+        return spec.unpadded_page_size_bytes
+    if isinstance(spec, MambaSpec):
+        return replace(spec, page_size_padded=None).page_size_bytes
+    return None
+
+
+def derive_canonical_mappings(
+    vllm_config: "VllmConfig",
+    kv_cache_config: KVCacheConfig,
+    kv_caches: dict[str, torch.Tensor | list[torch.Tensor]],
+) -> dict[str, CanonicalPageMapping]:
+    """Per-layer canonical page mappings for this worker.
+
+    Empty when the worker group is not exactly the TP x PCP grid; layers
+    absent from the result have no canonical representation.
+    """
+    parallel_config = vllm_config.parallel_config
+    tp_size = parallel_config.tensor_parallel_size
+    pcp_size = parallel_config.prefill_context_parallel_size
+    group_size = tp_size * pcp_size
+    if parallel_config.world_size != group_size:
+        return {}
+
+    def ctx(rank: int) -> _RankContext:
+        return _RankContext(
+            tp_size=tp_size,
+            dcp_size=parallel_config.decode_context_parallel_size,
+            pcp_size=pcp_size,
+            interleave=parallel_config.cp_kv_cache_interleave_size,
+            total_kv_heads=vllm_config.model_config.get_total_num_kv_heads(),
+            rank=rank,
+        )
+
+    my_rank = parallel_config.rank
+    num_blocks = kv_cache_config.num_blocks
+
+    mappings: dict[str, CanonicalPageMapping] = {}
+    for kv_cache_group in kv_cache_config.kv_cache_groups:
+        group_kv_cache_spec = kv_cache_group.kv_cache_spec
+        if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
+            per_layer_specs = group_kv_cache_spec.kv_cache_specs
+        else:
+            per_layer_specs = {}
+        for layer_name in kv_cache_group.layer_names:
+            spec = per_layer_specs.get(layer_name, group_kv_cache_spec)
+            per_rank: list[CanonicalPageMapping] = []
+            for rank in range(group_size):
+                mapping = _layer_mapping(
+                    spec, kv_caches.get(layer_name), num_blocks, ctx(rank)
+                )
+                if mapping is None:
+                    break
+                per_rank.append(mapping)
+            if len(per_rank) != group_size:
+                page = _unpadded_page_size(spec)
+                if page is None:
+                    continue
+                per_rank = [
+                    _opaque_fallback_mapping(page, group_size, rank)
+                    for rank in range(group_size)
+                ]
+            _verify_tiling(layer_name, per_rank)
+            mappings[layer_name] = per_rank[my_rank]
+    return mappings

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -5,6 +5,10 @@ from dataclasses import replace
 
 import torch
 
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import (
+    derive_canonical_mappings,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
     OffloadingConnectorMetadata,
     OffloadingWorkerMetadata,
@@ -40,9 +44,11 @@ class OffloadingConnectorWorker:
     def __init__(
         self,
         spec: OffloadingSpec,
+        vllm_config: "VllmConfig",
         kv_cache_config: KVCacheConfig,
     ):
         self.spec = spec
+        self.vllm_config = vllm_config
         self.kv_cache_config = kv_cache_config
         self.worker: OffloadingWorker | None = None
         # Non-writers still ack: pending_count waits for world_size per job.
@@ -63,6 +69,9 @@ class OffloadingConnectorWorker:
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         kv_cache_config = self.kv_cache_config
         num_blocks = kv_cache_config.num_blocks
+        mappings = derive_canonical_mappings(
+            self.vllm_config, kv_cache_config, kv_caches
+        )
 
         # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
         # stride(0) as the manager-block stride (equals total_num_bytes_per_block).
@@ -201,10 +210,21 @@ class OffloadingConnectorWorker:
 
                 curr_tensor_idx = len(block_tensors) - 1
                 for layer_name in tensor_layer_names:
+                    mapping = (
+                        mappings.get(layer_name)
+                        if len(tensors_per_block[first_layer_name]) == 1
+                        else None
+                    )
+                    assert (
+                        mapping is None
+                        or mapping.local_page_size_bytes
+                        == unpadded_page_size_bytes[layer_name]
+                    )
                     block_data_refs[layer_name].append(
                         CanonicalKVCacheRef(
                             tensor_idx=curr_tensor_idx,
                             page_size_bytes=(unpadded_page_size_bytes[layer_name]),
+                            mapping=mapping,
                         )
                     )
 

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -75,7 +75,9 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
                 spec, vllm_config, kv_cache_config
             )
         elif role == KVConnectorRole.WORKER:
-            self.connector_worker = OffloadingConnectorWorker(spec, kv_cache_config)
+            self.connector_worker = OffloadingConnectorWorker(
+                spec, vllm_config, kv_cache_config
+            )
 
     def shutdown(self) -> None:
         if self.connector_worker is not None:

vllm/v1/kv_offload/base.py

@@ -457,6 +457,48 @@ class CanonicalKVCacheTensor:
     page_size_bytes: int
 
 
+@dataclass(frozen=True)
+class CopyRun:
+    """A strided byte correspondence between this worker's physical page and
+    a canonical page: for i in range(num_fragments), fragment i spans
+    [local_offset + i * local_stride, +fragment_size) in the worker's page and
+    [canonical_offset + i * canonical_stride, +fragment_size) canonically."""
+
+    local_offset: int
+    canonical_offset: int
+    fragment_size: int
+    num_fragments: int
+    local_stride: int
+    canonical_stride: int
+
+
+@dataclass(frozen=True)
+class CanonicalPageMapping:
+    """How this worker's page maps into a canonical (parallelism-free) page.
+    In-process only, never serialized. Runs cover the full local page in both
+    directions; ranks holding identical bytes take turns writing them.
+    """
+
+    # Size of the canonical page in bytes
+    canonical_page_size_bytes: int
+    # Size of this worker's (un-padded) page in bytes
+    local_page_size_bytes: int
+    # Byte correspondences between this worker's page and a canonical page
+    runs: tuple[CopyRun, ...]
+    # Number of ranks holding these exact bytes
+    num_writers: int
+    # This worker's index among those ranks
+    writer_index: int
+    # Canonical bytes identical under any parallel config with this block span
+    parallelism_agnostic: bool
+
+    def is_writer(self, block_id: int) -> bool:
+        """Whether this worker stores the canonical page of the given block.
+        Rotating by block spreads writes across ranks holding identical bytes.
+        """
+        return block_id % self.num_writers == self.writer_index
+
+
 @dataclass
 class CanonicalKVCacheRef:
     """
@@ -468,6 +510,8 @@ class CanonicalKVCacheRef:
     tensor_idx: int
     # The un-padded page size per block in bytes
     page_size_bytes: int
+    # How this worker's page maps into a canonical page; None = uncertified
+    mapping: CanonicalPageMapping | None = None
 
 
 @dataclass