From 86941a344901f259b190d6db5c64a41962a03072 Mon Sep 17 00:00:00 2001 From: Itay Etelis Date: Sun, 12 Jul 2026 14:59:53 +0300 Subject: [PATCH 1/2] Add per-layer KV head region schema for TP-agnostic offload Signed-off-by: Itay Etelis --- .../kv_connector/v1/offloading/worker.py | 56 +++++++++++++++++++ vllm/v1/kv_offload/base.py | 17 ++++++ 2 files changed, 73 insertions(+) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 29914e7388e5..a023218f65b6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -15,6 +15,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, MambaSpec, + MLAAttentionSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.kv_offload.base import ( @@ -22,6 +23,7 @@ CanonicalKVCaches, CanonicalKVCacheTensor, GPULoadStoreSpec, + KVHeadRegion, LoadStoreSpec, OffloadingSpec, OffloadingWorker, @@ -30,6 +32,34 @@ logger = init_logger(__name__) +def _attention_head_regions( + kv_cache: torch.Tensor, spec: AttentionSpec, num_blocks: int +) -> tuple[KVHeadRegion, ...] | None: + """Derive K/V head regions from a (num_blocks, 2, block_size, num_heads, + head_size) attention cache. Returns None when the physical layout cannot + be established unambiguously (fail closed).""" + bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size + if tuple(kv_cache.shape) != (num_blocks, 2, bs, heads, head_size): + return None + elem = kv_cache.element_size() + if 2 * bs * heads * head_size * elem != spec.real_page_size_bytes: + return None + s = kv_cache.stride() + if s[4] != 1 or s[1] != bs * heads * head_size: + return None + k_size = bs * heads * head_size * elem + if s[2] == heads * head_size and s[3] == head_size: # NHD + fragment_size, num_fragments = heads * head_size * elem, bs + elif s[3] == bs * head_size and s[2] == head_size: # HND + fragment_size, num_fragments = k_size, 1 + else: + return None + return ( + KVHeadRegion(0, fragment_size, num_fragments, heads), + KVHeadRegion(k_size, fragment_size, num_fragments, heads), + ) + + class OffloadingConnectorWorker: """Implementation of Worker side methods""" @@ -52,6 +82,13 @@ def register_kv_caches( ): kv_cache_config = self.spec.kv_cache_config num_blocks = kv_cache_config.num_blocks + parallel_config = self.spec.vllm_config.parallel_config + total_num_kv_heads = self.spec.vllm_config.model_config.get_total_num_kv_heads() + tp_size = parallel_config.tensor_parallel_size + cp_size = ( + parallel_config.decode_context_parallel_size + * parallel_config.prefill_context_parallel_size + ) # 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). @@ -69,6 +106,8 @@ def register_kv_caches( unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} + head_regions: dict[str, tuple[KVHeadRegion, ...]] = {} + replicated_layers: set[str] = set() for kv_cache_group in kv_cache_config.kv_cache_groups: group_layer_names = kv_cache_group.layer_names group_kv_cache_spec = kv_cache_group.kv_cache_spec @@ -108,6 +147,21 @@ def register_kv_caches( unpadded_page_size_bytes[layer_name] = ( layer_kv_cache_spec.real_page_size_bytes ) + if isinstance(layer_kv_cache_spec, MLAAttentionSpec): + # Replicated latent, unless CP shards tokens + if cp_size == 1: + replicated_layers.add(layer_name) + elif ( + cp_size == 1 + and layer_kv_cache_spec.num_kv_heads * tp_size + == total_num_kv_heads + and not layer_kv_cache_spec.kv_quant_mode.is_per_token_head + ): + regions = _attention_head_regions( + layer_kv_cache, layer_kv_cache_spec, num_blocks + ) + if regions is not None: + head_regions[layer_name] = regions elif isinstance(layer_kv_cache_spec, MambaSpec): state_tensors = kv_caches[layer_name] @@ -204,6 +258,8 @@ def register_kv_caches( CanonicalKVCacheRef( tensor_idx=curr_tensor_idx, page_size_bytes=(unpadded_page_size_bytes[layer_name]), + head_regions=head_regions.get(layer_name), + replicated=layer_name in replicated_layers, ) ) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 5a2e3c184d39..ee8a4c296c78 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -404,6 +404,18 @@ class CanonicalKVCacheTensor: page_size_bytes: int +@dataclass(frozen=True) +class KVHeadRegion: + """One region of this worker's page as a run of equally-sized fragments, + each holding num_kv_heads heads (the unit of one copy op). Describes this + worker's physical page only — in-process use, never serialized.""" + + offset: int + fragment_size: int + num_fragments: int + num_kv_heads: int + + @dataclass class CanonicalKVCacheRef: """ @@ -415,6 +427,11 @@ class CanonicalKVCacheRef: tensor_idx: int # The un-padded page size per block in bytes page_size_bytes: int + # Per-head source decomposition of the page; None = no head slicing + head_regions: tuple[KVHeadRegion, ...] | None = None + # When head_regions is None: True = page is identical on every TP rank + # (MLA latent), False = rank-specific shards (Mamba states, packed) + replicated: bool = False @dataclass From 5d96fe507f3b4d74b843e1dac739e2b671a4832d Mon Sep 17 00:00:00 2001 From: Itay Etelis Date: Sat, 18 Jul 2026 20:38:57 +0300 Subject: [PATCH 2/2] Replace KV head regions with canonical page mappings All parallelism reasoning moves to vllm/v1/kv_offload/sharding.py, covering TP/DCP/PCP, packed and split KV layouts, with single-writer election. Signed-off-by: Itay Etelis --- tests/v1/kv_offload/test_kv_sharding.py | 457 ++++++++++++++++++ .../kv_connector/v1/offloading/worker.py | 69 +-- vllm/v1/kv_offload/base.py | 43 +- vllm/v1/kv_offload/sharding.py | 384 +++++++++++++++ 4 files changed, 886 insertions(+), 67 deletions(-) create mode 100644 tests/v1/kv_offload/test_kv_sharding.py create mode 100644 vllm/v1/kv_offload/sharding.py diff --git a/tests/v1/kv_offload/test_kv_sharding.py b/tests/v1/kv_offload/test_kv_sharding.py new file mode 100644 index 000000000000..3a2cdaceaf4a --- /dev/null +++ b/tests/v1/kv_offload/test_kv_sharding.py @@ -0,0 +1,457 @@ +# 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.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + KVCacheSpec, + KVQuantMode, + MLAAttentionSpec, +) +from vllm.v1.kv_offload.base import CanonicalPageMapping, MappedRun +from vllm.v1.kv_offload.sharding import ( + _layer_mapping, + _rank_private_mapping, + _RankContext, + _verify_mappings, + derive_canonical_mappings, +) + +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[MappedRun, ...]) -> 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.parallel_invariant + k_dst = [256, 768, 1280, 1792] + assert _triples(mapping.store_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], + ) + ] + assert mapping.store_runs == mapping.load_runs + + +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.store_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.store_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.store_runs) == [(0, 0, 1024)] + + +# --------------------------------------------------------------------------- +# Replication and writer election +# --------------------------------------------------------------------------- + + +def test_gqa_replicated_heads_elect_single_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 + writer = _mapping(spec, cache, ctx(2)) + replica = _mapping(spec, cache, ctx(3)) + assert writer.canonical_page_size_bytes == 2 * 512 + # K region: head shard 1 at 64B offsets within 128B token rows + assert _triples(writer.store_runs)[:4] == [ + (0, 64, 64), + (64, 192, 64), + (128, 320, 64), + (192, 448, 64), + ] + assert replica.store_runs == () + assert replica.load_runs == writer.load_runs + _verify_mappings("gqa", [_mapping(spec, cache, ctx(r)) for r in range(4)]) + + +def test_mla_single_writer_tp_only(): + spec = _mla_spec() + writer = _mapping(spec, None, _ctx(rank=0, tp=2)) + reader = _mapping(spec, None, _ctx(rank=1, tp=2)) + # Latent pages are stored once, not once per rank + assert writer.canonical_page_size_bytes == 256 + assert _triples(writer.store_runs) == [(0, 0, 256)] + assert reader.store_runs == () + assert _triples(reader.load_runs) == [(0, 0, 256)] + + +# --------------------------------------------------------------------------- +# 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].parallel_invariant + # rank 2 = head shard 1, cp rank 0: K tokens 0,2,4,6 at head offset 64 + assert _triples(per_rank[2].store_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.store_runs for m in per_rank) + _verify_mappings("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.store_runs) == [(o, 2 * o, 64) for o in (0, 64, 128, 192)] + assert _triples(rank1.store_runs) == [ + (o, 2 * o + 64, 64) for o in (0, 64, 128, 192) + ] + _verify_mappings("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.store_runs for m in per_rank) + _verify_mappings("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.store_runs) == [(0, 0, 128), (128, 256, 128)] + _verify_mappings( + "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_mappings("layer", per_rank) + + +# --------------------------------------------------------------------------- +# Byte-level round trips +# --------------------------------------------------------------------------- + + +def _store_all(mappings, pages, size: int) -> bytes: + buf = bytearray(size) + for mapping, page in zip(mappings, pages): + for local, canonical, n in _triples(mapping.store_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.load_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 + + +# --------------------------------------------------------------------------- +# 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_rank_private_places_page_whole(): + mapping = _rank_private_mapping(1024, 4, 2) + assert mapping.canonical_page_size_bytes == 4096 + assert not mapping.parallel_invariant + assert _triples(mapping.store_runs) == [(0, 2048, 1024)] + assert mapping.store_runs == mapping.load_runs + _verify_mappings("opaque", [_rank_private_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"].parallel_invariant + assert mappings["attn"].store_runs + # Uncertifiable layers degrade to rank-private, never disappear + assert not mappings["quant"].parallel_invariant + assert mappings["quant"].canonical_page_size_bytes == 4 * 512 + + +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 + ) + == {} + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index a023218f65b6..bb099a55690e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -15,7 +15,6 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, MambaSpec, - MLAAttentionSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.kv_offload.base import ( @@ -23,43 +22,15 @@ CanonicalKVCaches, CanonicalKVCacheTensor, GPULoadStoreSpec, - KVHeadRegion, LoadStoreSpec, OffloadingSpec, OffloadingWorker, ) +from vllm.v1.kv_offload.sharding import derive_canonical_mappings logger = init_logger(__name__) -def _attention_head_regions( - kv_cache: torch.Tensor, spec: AttentionSpec, num_blocks: int -) -> tuple[KVHeadRegion, ...] | None: - """Derive K/V head regions from a (num_blocks, 2, block_size, num_heads, - head_size) attention cache. Returns None when the physical layout cannot - be established unambiguously (fail closed).""" - bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size - if tuple(kv_cache.shape) != (num_blocks, 2, bs, heads, head_size): - return None - elem = kv_cache.element_size() - if 2 * bs * heads * head_size * elem != spec.real_page_size_bytes: - return None - s = kv_cache.stride() - if s[4] != 1 or s[1] != bs * heads * head_size: - return None - k_size = bs * heads * head_size * elem - if s[2] == heads * head_size and s[3] == head_size: # NHD - fragment_size, num_fragments = heads * head_size * elem, bs - elif s[3] == bs * head_size and s[2] == head_size: # HND - fragment_size, num_fragments = k_size, 1 - else: - return None - return ( - KVHeadRegion(0, fragment_size, num_fragments, heads), - KVHeadRegion(k_size, fragment_size, num_fragments, heads), - ) - - class OffloadingConnectorWorker: """Implementation of Worker side methods""" @@ -82,12 +53,8 @@ def register_kv_caches( ): kv_cache_config = self.spec.kv_cache_config num_blocks = kv_cache_config.num_blocks - parallel_config = self.spec.vllm_config.parallel_config - total_num_kv_heads = self.spec.vllm_config.model_config.get_total_num_kv_heads() - tp_size = parallel_config.tensor_parallel_size - cp_size = ( - parallel_config.decode_context_parallel_size - * parallel_config.prefill_context_parallel_size + mappings = derive_canonical_mappings( + self.spec.vllm_config, kv_cache_config, kv_caches ) # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use @@ -106,8 +73,6 @@ def register_kv_caches( unpadded_page_size_bytes: dict[str, int] = {} # layer_name -> size of page in bytes page_size_bytes: dict[str, int] = {} - head_regions: dict[str, tuple[KVHeadRegion, ...]] = {} - replicated_layers: set[str] = set() for kv_cache_group in kv_cache_config.kv_cache_groups: group_layer_names = kv_cache_group.layer_names group_kv_cache_spec = kv_cache_group.kv_cache_spec @@ -147,21 +112,6 @@ def register_kv_caches( unpadded_page_size_bytes[layer_name] = ( layer_kv_cache_spec.real_page_size_bytes ) - if isinstance(layer_kv_cache_spec, MLAAttentionSpec): - # Replicated latent, unless CP shards tokens - if cp_size == 1: - replicated_layers.add(layer_name) - elif ( - cp_size == 1 - and layer_kv_cache_spec.num_kv_heads * tp_size - == total_num_kv_heads - and not layer_kv_cache_spec.kv_quant_mode.is_per_token_head - ): - regions = _attention_head_regions( - layer_kv_cache, layer_kv_cache_spec, num_blocks - ) - if regions is not None: - head_regions[layer_name] = regions elif isinstance(layer_kv_cache_spec, MambaSpec): state_tensors = kv_caches[layer_name] @@ -254,12 +204,21 @@ def register_kv_caches( 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]), - head_regions=head_regions.get(layer_name), - replicated=layer_name in replicated_layers, + mapping=mapping, ) ) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index ee8a4c296c78..7d230a15a504 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -405,15 +405,37 @@ class CanonicalKVCacheTensor: @dataclass(frozen=True) -class KVHeadRegion: - """One region of this worker's page as a run of equally-sized fragments, - each holding num_kv_heads heads (the unit of one copy op). Describes this - worker's physical page only — in-process use, never serialized.""" - - offset: int +class MappedRun: + """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 - num_kv_heads: 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. store_runs may be empty when another + worker contributes the same bytes; load_runs cover the full local page. + """ + + # 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 + # Bytes this worker contributes when writing a canonical page + store_runs: tuple[MappedRun, ...] + # Bytes this worker reads back from a canonical page + load_runs: tuple[MappedRun, ...] + # Canonical bytes identical under any parallel config with this block span + parallel_invariant: bool @dataclass @@ -427,11 +449,8 @@ class CanonicalKVCacheRef: tensor_idx: int # The un-padded page size per block in bytes page_size_bytes: int - # Per-head source decomposition of the page; None = no head slicing - head_regions: tuple[KVHeadRegion, ...] | None = None - # When head_regions is None: True = page is identical on every TP rank - # (MLA latent), False = rank-specific shards (Mamba states, packed) - replicated: bool = False + # How this worker's page maps into a canonical page; None = uncertified + mapping: CanonicalPageMapping | None = None @dataclass diff --git a/vllm/v1/kv_offload/sharding.py b/vllm/v1/kv_offload/sharding.py new file mode 100644 index 000000000000..48bd52c79997 --- /dev/null +++ b/vllm/v1/kv_offload/sharding.py @@ -0,0 +1,384 @@ +# 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 a rank-private 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, MappedRun + +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 + + +def _coalesce_runs(runs: list[MappedRun]) -> tuple[MappedRun, ...]: + """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[MappedRun] = [] + 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 = MappedRun(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] = MappedRun( + prev.local_offset, prev.canonical_offset, size, 1, size, size + ) + else: + out.append(run) + return tuple(out) + + +def _chunk_runs( + channels: list[tuple[int, int, int, int]], + num_tokens: int, + ctx: _RankContext, +) -> tuple[MappedRun, ...]: + """Place each channel's num_tokens rows: local token l of CP rank c is + canonical token ((l // I) * cp + c) * I + l % I. A channel is + (local_base, canonical_base, local_row, canonical_row).""" + runs: list[MappedRun] = [] + interleave, cp = ctx.interleave, ctx.cp_size + for local_base, canonical_base, local_row, canonical_row in channels: + if cp == 1: + runs.append( + MappedRun( + local_base, + canonical_base, + local_row, + num_tokens, + local_row, + canonical_row, + ) + ) + continue + for chunk in range(num_tokens // interleave): + canonical_token = (chunk * cp + ctx.total_cp_rank) * interleave + runs.append( + MappedRun( + local_base + chunk * interleave * local_row, + canonical_base + canonical_token * canonical_row, + local_row, + interleave, + local_row, + canonical_row, + ) + ) + return _coalesce_runs(runs) + + +def _attention_channels( + kv_cache: torch.Tensor, + spec: AttentionSpec, + num_blocks: int, + head_shard: int, + num_head_shards: int, + cp_size: int, +) -> list[tuple[int, int, int, int]] | None: + """Byte channels of an attention page, given this rank's head shard. + + Recognizes packed KV (num_blocks, heads, block_size, 2 * head_size) and + split KV (num_blocks, 2, block_size, heads, head_size), in NHD or HND + stride order. None when the layout is ambiguous (fail closed).""" + bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size + elem = kv_cache.element_size() + page = spec.real_page_size_bytes + content = 2 * head_size + span = bs * cp_size # canonical tokens per offloaded block + + if tuple(kv_cache.shape) == (num_blocks, heads, bs, content): + if heads * bs * content * elem != page: + return None + s = kv_cache.stride() + if s[3] != 1: + return None + row = heads * content * elem + if s[1] == content and s[2] == heads * content: # NHD: token-major + return [(0, head_shard * row, row, num_head_shards * row)] + if s[1] == bs * content and s[2] == content: # HND: head-major + g = content * elem + return [ + ( + j * bs * g, + (head_shard * heads + j) * span * g, + g, + g, + ) + for j in range(heads) + ] + return None + + if tuple(kv_cache.shape) != (num_blocks, 2, bs, heads, head_size): + return None + if 2 * bs * heads * head_size * elem != page: + return None + s = kv_cache.stride() + if s[4] != 1 or s[1] != bs * heads * head_size: + return None + k_size = bs * heads * head_size * elem + if s[2] == heads * head_size and s[3] == head_size: # NHD + row = heads * head_size * elem + return [ + ( + region * bs * row, + region * span * num_head_shards * row + head_shard * row, + row, + num_head_shards * row, + ) + for region in range(2) # K, then V + ] + if s[3] == bs * head_size and s[2] == head_size: # HND + f = head_size * elem + total_heads = num_head_shards * heads + return [ + ( + region * k_size + j * bs * f, + (region * total_heads + head_shard * heads + j) * span * f, + f, + f, + ) + for region in range(2) + for j in range(heads) + ] + 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; first DCP group writes + if spec.compress_ratio != 1 or page % bs: + return None + row = page // bs + runs = _chunk_runs([(0, 0, row, row)], bs, ctx) + return CanonicalPageMapping( + canonical_page_size_bytes=ctx.cp_size * page, + local_page_size_bytes=page, + store_runs=runs if ctx.tp_rank < ctx.dcp_size else (), + load_runs=runs, + parallel_invariant=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 + channels = _attention_channels( + kv_cache, spec, num_blocks, head_shard, num_head_shards, ctx.cp_size + ) + if channels is None: + return None + runs = _chunk_runs(channels, bs, ctx) + # One writer among ranks holding identical bytes + contributor = (ctx.tp_rank % replication) // ctx.dcp_size == 0 + return CanonicalPageMapping( + canonical_page_size_bytes=ctx.cp_size * num_head_shards * page, + local_page_size_bytes=page, + store_runs=runs if contributor else (), + load_runs=runs, + parallel_invariant=ctx.cp_size == 1, + ) + + +def _rank_private_mapping( + page_size_bytes: int, num_ranks: int, rank: int +) -> CanonicalPageMapping: + """Fallback: place the worker's page whole at a worker-exclusive offset.""" + run = MappedRun( + 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, + store_runs=(run,), + load_runs=(run,), + parallel_invariant=False, + ) + + +def _run_intervals( + runs: tuple[MappedRun, ...], 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_mappings(layer_name: str, per_rank: list[CanonicalPageMapping]) -> None: + """All ranks' store runs must tile the canonical page exactly once, and + each rank's load runs must cover exactly its local page.""" + size = per_rank[0].canonical_page_size_bytes + store_intervals: list[tuple[int, int]] = [] + for mapping in per_rank: + assert mapping.canonical_page_size_bytes == size + store_intervals += _run_intervals(mapping.store_runs, canonical=True) + local = _run_intervals(mapping.load_runs, canonical=False) + assert _is_exact_partition(local, mapping.local_page_size_bytes), ( + f"load runs do not cover the local page of layer {layer_name}" + ) + store_intervals.sort() + assert _is_exact_partition(store_intervals, size), ( + f"store runs do not tile the canonical page of layer {layer_name}" + ) + + +def _unpadded_page_size(spec: KVCacheSpec) -> int | None: + if isinstance(spec, AttentionSpec): + return spec.real_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 = [ + _rank_private_mapping(page, group_size, rank) + for rank in range(group_size) + ] + _verify_mappings(layer_name, per_rank) + mappings[layer_name] = per_rank[my_rank] + return mappings