@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2c4d348848

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

Itay Etelis committed Jul 31, 2026 at 14:42 UTC 2c4d3488483a4ce90b44661aefdb0cd31c6fbffa
6 files changed +1016 -7
tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py new
+474
@@ -0,0 +1,474 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from unittest.mock import MagicMock
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import (
9 + _layer_mapping,
10 + _opaque_fallback_mapping,
11 + _RankContext,
12 + _verify_tiling,
13 + derive_canonical_mappings,
14 +)
15 +from vllm.v1.kv_cache_interface import (
16 + FullAttentionSpec,
17 + KVCacheGroupSpec,
18 + KVCacheSpec,
19 + KVQuantMode,
20 + MLAAttentionSpec,
21 +)
22 +from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun
23 +
24 +NUM_BLOCKS = 3
25 +
26 +
27 +def _ctx(rank, tp=1, dcp=1, pcp=1, interleave=1, total=None, heads=2):
28 + return _RankContext(
29 + tp_size=tp,
30 + dcp_size=dcp,
31 + pcp_size=pcp,
32 + interleave=interleave,
33 + total_kv_heads=heads * tp if total is None else total,
34 + rank=rank,
35 + )
36 +
37 +
38 +def _full_spec(num_kv_heads: int = 2, **kwargs) -> FullAttentionSpec:
39 + # block_size=4, head_dim=64, int8
40 + return FullAttentionSpec(
41 + block_size=4,
42 + num_kv_heads=num_kv_heads,
43 + head_size=64,
44 + dtype=torch.int8,
45 + **kwargs,
46 + )
47 +
48 +
49 +def _mla_spec(**kwargs) -> MLAAttentionSpec:
50 + # page = 4 * 64 = 256B, one 64B latent row per token
51 + return MLAAttentionSpec(
52 + block_size=4, num_kv_heads=1, head_size=64, dtype=torch.int8, **kwargs
53 + )
54 +
55 +
56 +def _split_nhd_cache(spec) -> torch.Tensor:
57 + return torch.zeros(
58 + NUM_BLOCKS,
59 + 2,
60 + spec.block_size,
61 + spec.num_kv_heads,
62 + spec.head_size,
63 + dtype=torch.int8,
64 + )
65 +
66 +
67 +def _split_hnd_cache(spec) -> torch.Tensor:
68 + return torch.zeros(
69 + NUM_BLOCKS,
70 + 2,
71 + spec.num_kv_heads,
72 + spec.block_size,
73 + spec.head_size,
74 + dtype=torch.int8,
75 + ).permute(0, 1, 3, 2, 4)
76 +
77 +
78 +def _packed_nhd_cache(spec) -> torch.Tensor:
79 + """Logical (num_blocks, heads, block_size, 2 * head_size) over an NHD
80 + physical layout — the FlashAttention/FlashInfer/Triton/Flex form."""
81 + return torch.zeros(
82 + NUM_BLOCKS,
83 + spec.block_size,
84 + spec.num_kv_heads,
85 + 2 * spec.head_size,
86 + dtype=torch.int8,
87 + ).permute(0, 2, 1, 3)
88 +
89 +
90 +def _packed_hnd_cache(spec) -> torch.Tensor:
91 + return torch.zeros(
92 + NUM_BLOCKS,
93 + spec.num_kv_heads,
94 + spec.block_size,
95 + 2 * spec.head_size,
96 + dtype=torch.int8,
97 + )
98 +
99 +
100 +CACHE_BUILDERS = {
101 + "split_nhd": _split_nhd_cache,
102 + "split_hnd": _split_hnd_cache,
103 + "packed_nhd": _packed_nhd_cache,
104 + "packed_hnd": _packed_hnd_cache,
105 +}
106 +
107 +
108 +def _try_mapping(spec, kv_cache, ctx) -> CanonicalPageMapping | None:
109 + return _layer_mapping(spec, kv_cache, NUM_BLOCKS, ctx)
110 +
111 +
112 +def _mapping(spec, kv_cache, ctx) -> CanonicalPageMapping:
113 + mapping = _try_mapping(spec, kv_cache, ctx)
114 + assert mapping is not None
115 + return mapping
116 +
117 +
118 +def _triples(runs: tuple[CopyRun, ...]) -> list[tuple[int, int, int]]:
119 + """Expand runs to explicit (local_offset, canonical_offset, size) copies."""
120 + out = []
121 + for run in runs:
122 + for i in range(run.num_fragments):
123 + out.append(
124 + (
125 + run.local_offset + i * run.local_stride,
126 + run.canonical_offset + i * run.canonical_stride,
127 + run.fragment_size,
128 + )
129 + )
130 + return out
131 +
132 +
133 +# ---------------------------------------------------------------------------
134 +# TP-only placement (byte-compatible with the uniform interleave layout)
135 +# ---------------------------------------------------------------------------
136 +
137 +
138 +def test_split_nhd_placement_rank2_of_4():
139 + spec = _full_spec()
140 + mapping = _mapping(spec, _split_nhd_cache(spec), _ctx(rank=2, tp=4))
141 + assert mapping.canonical_page_size_bytes == 4 * 1024
142 + assert mapping.parallelism_agnostic
143 + k_dst = [256, 768, 1280, 1792]
144 + assert _triples(mapping.runs) == [
145 + (local, canonical, 128)
146 + for local, canonical in zip(
147 + [0, 128, 256, 384, 512, 640, 768, 896],
148 + k_dst + [2048 + o for o in k_dst],
149 + )
150 + ]
151 + # Heads are sharded, not replicated: this rank writes every block
152 + assert mapping.num_writers == 1
153 +
154 +
155 +def test_packed_nhd_placement_rank2_of_4():
156 + spec = _full_spec()
157 + mapping = _mapping(spec, _packed_nhd_cache(spec), _ctx(rank=2, tp=4))
158 + assert _triples(mapping.runs) == [
159 + (0, 512, 256),
160 + (256, 1536, 256),
161 + (512, 2560, 256),
162 + (768, 3584, 256),
163 + ]
164 +
165 +
166 +def test_packed_hnd_placement_rank1_of_4():
167 + spec = _full_spec()
168 + mapping = _mapping(spec, _packed_hnd_cache(spec), _ctx(rank=1, tp=4))
169 + # Two heads, each a contiguous canonical head region of 4 tokens x 128B
170 + assert _triples(mapping.runs) == [(0, 1024, 1024)]
171 +
172 +
173 +@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS))
174 +def test_single_rank_coalesces_to_one_run(form):
175 + spec = _full_spec()
176 + mapping = _mapping(spec, CACHE_BUILDERS[form](spec), _ctx(rank=0))
177 + assert mapping.canonical_page_size_bytes == 1024
178 + assert _triples(mapping.runs) == [(0, 0, 1024)]
179 +
180 +
181 +# ---------------------------------------------------------------------------
182 +# Replication and writer election
183 +# ---------------------------------------------------------------------------
184 +
185 +
186 +def test_gqa_replicated_heads_rotate_writer():
187 + # total 2 KV heads on tp=4: replication factor 2, head shard = rank // 2
188 + spec = _full_spec(num_kv_heads=1)
189 + cache = _split_nhd_cache(spec)
190 + ctx = lambda rank: _ctx(rank, tp=4, total=2) # noqa: E731
191 + rank2 = _mapping(spec, cache, ctx(2))
192 + rank3 = _mapping(spec, cache, ctx(3))
193 + assert rank2.canonical_page_size_bytes == 2 * 512
194 + # K region: head shard 1 at 64B offsets within 128B token rows
195 + assert _triples(rank2.runs)[:4] == [
196 + (0, 64, 64),
197 + (64, 192, 64),
198 + (128, 320, 64),
199 + (192, 448, 64),
200 + ]
201 + # Same head shard, so identical bytes: the two take alternate blocks
202 + assert rank3.runs == rank2.runs
203 + assert (rank2.is_writer(0), rank3.is_writer(0)) == (True, False)
204 + assert (rank2.is_writer(1), rank3.is_writer(1)) == (False, True)
205 + _verify_tiling("gqa", [_mapping(spec, cache, ctx(r)) for r in range(4)])
206 +
207 +
208 +def test_mla_replicas_rotate_writer():
209 + spec = _mla_spec()
210 + rank0 = _mapping(spec, None, _ctx(rank=0, tp=2))
211 + rank1 = _mapping(spec, None, _ctx(rank=1, tp=2))
212 + # Latent pages are stored once per block, not once per rank
213 + assert rank0.canonical_page_size_bytes == 256
214 + assert _triples(rank0.runs) == [(0, 0, 256)]
215 + assert rank1.runs == rank0.runs
216 + assert (rank0.is_writer(0), rank1.is_writer(0)) == (True, False)
217 + assert (rank0.is_writer(1), rank1.is_writer(1)) == (False, True)
218 +
219 +
220 +# ---------------------------------------------------------------------------
221 +# DCP / PCP token sharding
222 +# ---------------------------------------------------------------------------
223 +
224 +
225 +def test_dcp_interleaves_tokens_within_replicas():
226 + # tp=4, dcp=2, total 2 KV heads: head shard = rank // 2, cp rank = rank % 2
227 + spec = _full_spec(num_kv_heads=1)
228 + cache = _split_nhd_cache(spec)
229 + ctx = lambda rank: _ctx(rank, tp=4, dcp=2, total=2) # noqa: E731
230 + per_rank = [_mapping(spec, cache, ctx(rank)) for rank in range(4)]
231 + assert all(m is not None for m in per_rank)
232 + # 8 canonical tokens x 2 heads x 64B per region
233 + assert per_rank[0].canonical_page_size_bytes == 2048
234 + assert not per_rank[0].parallelism_agnostic
235 + # rank 2 = head shard 1, cp rank 0: K tokens 0,2,4,6 at head offset 64
236 + assert _triples(per_rank[2].runs)[:4] == [
237 + (0, 64, 64),
238 + (64, 320, 64),
239 + (128, 576, 64),
240 + (192, 832, 64),
241 + ]
242 + # every rank contributes (dcp == replication: no residual replicas)
243 + assert all(m.num_writers == 1 for m in per_rank)
244 + _verify_tiling("dcp", per_rank)
245 +
246 +
247 +def test_mla_dcp_shards_latent_tokens():
248 + spec = _mla_spec()
249 + ctx = lambda rank: _ctx(rank, tp=2, dcp=2) # noqa: E731
250 + rank0 = _mapping(spec, None, ctx(0))
251 + rank1 = _mapping(spec, None, ctx(1))
252 + assert rank0.canonical_page_size_bytes == 512
253 + assert _triples(rank0.runs) == [(o, 2 * o, 64) for o in (0, 64, 128, 192)]
254 + assert _triples(rank1.runs) == [(o, 2 * o + 64, 64) for o in (0, 64, 128, 192)]
255 + _verify_tiling("mla-dcp", [rank0, rank1])
256 +
257 +
258 +def test_pcp_tokens_and_tp_heads_compose():
259 + # tp=2 x pcp=2: rank = pcp_rank * 2 + tp_rank; 4 workers tile the page
260 + spec = _full_spec(num_kv_heads=1)
261 + cache = _packed_nhd_cache(spec)
262 + per_rank = [
263 + _mapping(spec, cache, _ctx(rank, tp=2, pcp=2, total=2)) for rank in range(4)
264 + ]
265 + assert all(m is not None and m.runs for m in per_rank)
266 + _verify_tiling("pcp", per_rank)
267 +
268 +
269 +def test_interleave_chunks_stay_contiguous():
270 + # interleave=2: chunks of 2 tokens alternate between the 2 cp ranks and
271 + # coalesce into one contiguous fragment per chunk
272 + spec = _mla_spec()
273 + mapping = _mapping(spec, None, _ctx(rank=0, tp=2, dcp=2, interleave=2))
274 + assert _triples(mapping.runs) == [(0, 0, 128), (128, 256, 128)]
275 + _verify_tiling(
276 + "interleave",
277 + [
278 + _mapping(spec, None, _ctx(rank, tp=2, dcp=2, interleave=2))
279 + for rank in range(2)
280 + ],
281 + )
282 +
283 +
284 +@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS))
285 +def test_all_ranks_tile_canonical_page(form):
286 + spec = _full_spec()
287 + per_rank = [
288 + _mapping(spec, CACHE_BUILDERS[form](spec), _ctx(rank, tp=4))
289 + for rank in range(4)
290 + ]
291 + _verify_tiling("layer", per_rank)
292 +
293 +
294 +# ---------------------------------------------------------------------------
295 +# Byte-level round trips
296 +# ---------------------------------------------------------------------------
297 +
298 +
299 +def _store_all(mappings, pages, size: int, block_id: int = 0) -> bytes:
300 + buf = bytearray(size)
301 + for mapping, page in zip(mappings, pages):
302 + if not mapping.is_writer(block_id):
303 + continue
304 + for local, canonical, n in _triples(mapping.runs):
305 + buf[canonical : canonical + n] = page[local : local + n]
306 + return bytes(buf)
307 +
308 +
309 +def _load_one(mapping, canonical_bytes: bytes) -> bytes:
310 + page = bytearray(mapping.local_page_size_bytes)
311 + for local, canonical, n in _triples(mapping.runs):
312 + page[local : local + n] = canonical_bytes[canonical : canonical + n]
313 + return bytes(page)
314 +
315 +
316 +@pytest.mark.parametrize("form", sorted(CACHE_BUILDERS))
317 +def test_cross_tp_store_load(form):
318 + """Bytes stored under one TP size are the bytes another TP size loads."""
319 + total_heads, canonical_size = 8, 4096
320 + reference = bytes((7 + 31 * i) % 256 for i in range(canonical_size))
321 +
322 + def mappings_at(tp: int):
323 + spec = _full_spec(num_kv_heads=total_heads // tp)
324 + cache = CACHE_BUILDERS[form](spec)
325 + return [
326 + _mapping(spec, cache, _ctx(rank, tp=tp, total=total_heads))
327 + for rank in range(tp)
328 + ]
329 +
330 + for tp in (4, 2, 1):
331 + mappings = mappings_at(tp)
332 + assert all(m is not None for m in mappings)
333 + pages = [_load_one(m, reference) for m in mappings]
334 + assert _store_all(mappings, pages, canonical_size) == reference
335 +
336 +
337 +def test_cp_round_trip():
338 + # tp=4 / dcp=2 / 2 KV heads: 4 workers jointly hold one canonical page
339 + spec = _full_spec(num_kv_heads=1)
340 + cache = _split_nhd_cache(spec)
341 + mappings = [
342 + _mapping(spec, cache, _ctx(rank, tp=4, dcp=2, total=2)) for rank in range(4)
343 + ]
344 + reference = bytes((3 + 17 * i) % 256 for i in range(2048))
345 + pages = [_load_one(m, reference) for m in mappings]
346 + assert _store_all(mappings, pages, 2048) == reference
347 +
348 +
349 +def test_replica_rotation_round_trip():
350 + """Whichever replica a block elects reproduces the same canonical page."""
351 + spec = _mla_spec()
352 + mappings = [_mapping(spec, None, _ctx(rank, tp=2)) for rank in range(2)]
353 + reference = bytes((5 + 11 * i) % 256 for i in range(256))
354 + pages = [_load_one(m, reference) for m in mappings]
355 + for block_id in (0, 1):
356 + assert _store_all(mappings, pages, 256, block_id) == reference
357 +
358 +
359 +# ---------------------------------------------------------------------------
360 +# Fail-closed gates
361 +# ---------------------------------------------------------------------------
362 +
363 +
364 +def test_fail_closed_cases():
365 + spec = _full_spec()
366 + nhd = _split_nhd_cache(spec)
367 + # Spec heads inconsistent with total heads / tp
368 + assert _try_mapping(spec, nhd, _ctx(0, tp=4, total=2)) is None
369 + # tp not divisible by total KV heads
370 + one_head = _full_spec(num_kv_heads=1)
371 + assert (
372 + _try_mapping(one_head, _split_nhd_cache(one_head), _ctx(0, tp=3, total=2))
373 + is None
374 + )
375 + # DCP wider than the KV replication factor (tokens would shard across
376 + # ranks holding different heads)
377 + assert _try_mapping(spec, nhd, _ctx(0, tp=4, dcp=2)) is None
378 + # Interleave must divide the block size
379 + assert _try_mapping(spec, nhd, _ctx(0, tp=2, dcp=2, interleave=3, total=2)) is None
380 + # Per-token-head scales are packed with the data
381 + quant_spec = _full_spec(kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD)
382 + assert _try_mapping(quant_spec, _split_nhd_cache(quant_spec), _ctx(0, tp=4)) is None
383 + # Compressed MLA slots are not 1:1 with tokens
384 + assert _try_mapping(_mla_spec(compress_ratio=2), None, _ctx(0, tp=2, dcp=2)) is None
385 + # Unrecognized physical layouts
386 + swapped = torch.zeros(
387 + NUM_BLOCKS,
388 + spec.block_size,
389 + 2,
390 + spec.num_kv_heads,
391 + spec.head_size,
392 + dtype=torch.int8,
393 + ).permute(0, 2, 1, 3, 4)
394 + assert _try_mapping(spec, swapped, _ctx(0, tp=4)) is None
395 + # Non-attention specs
396 + assert _try_mapping(KVCacheSpec(block_size=4), None, _ctx(0, tp=4, total=8)) is None
397 +
398 +
399 +def test_opaque_fallback_places_page_whole():
400 + mapping = _opaque_fallback_mapping(1024, 4, 2)
401 + assert mapping.canonical_page_size_bytes == 4096
402 + assert not mapping.parallelism_agnostic
403 + assert _triples(mapping.runs) == [(0, 2048, 1024)]
404 + assert mapping.num_writers == 1
405 + _verify_tiling("opaque", [_opaque_fallback_mapping(1024, 4, r) for r in range(4)])
406 +
407 +
408 +# ---------------------------------------------------------------------------
409 +# derive_canonical_mappings end to end
410 +# ---------------------------------------------------------------------------
411 +
412 +
413 +def _vllm_config(tp=1, dcp=1, pcp=1, pp=1, interleave=1, total_kv_heads=2):
414 + config = MagicMock()
415 + config.parallel_config.tensor_parallel_size = tp
416 + config.parallel_config.decode_context_parallel_size = dcp
417 + config.parallel_config.prefill_context_parallel_size = pcp
418 + config.parallel_config.cp_kv_cache_interleave_size = interleave
419 + config.parallel_config.world_size = pp * tp * pcp
420 + config.parallel_config.rank = 0
421 + config.model_config.get_total_num_kv_heads.return_value = total_kv_heads
422 + return config
423 +
424 +
425 +def _kv_cache_config(groups):
426 + config = MagicMock()
427 + config.kv_cache_groups = groups
428 + config.num_blocks = NUM_BLOCKS
429 + return config
430 +
431 +
432 +def test_derive_mixed_model_with_dcp():
433 + attn_spec = _full_spec(num_kv_heads=1)
434 + mla_spec = _mla_spec()
435 + quant_spec = _full_spec(
436 + num_kv_heads=1, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD
437 + )
438 + kv_cache_config = _kv_cache_config(
439 + [
440 + KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=attn_spec),
441 + KVCacheGroupSpec(layer_names=["mla"], kv_cache_spec=mla_spec),
442 + KVCacheGroupSpec(layer_names=["quant"], kv_cache_spec=quant_spec),
443 + ]
444 + )
445 + kv_caches = {
446 + "attn": _split_nhd_cache(attn_spec),
447 + "quant": _split_nhd_cache(quant_spec),
448 + }
449 + mappings = derive_canonical_mappings(
450 + _vllm_config(tp=4, dcp=2, total_kv_heads=2), kv_cache_config, kv_caches
451 + )
452 + assert set(mappings) == {"attn", "mla", "quant"}
453 + assert not mappings["attn"].parallelism_agnostic
454 + assert mappings["attn"].runs
455 + # Uncertifiable layers degrade to an opaque page, never disappear
456 + assert not mappings["quant"].parallelism_agnostic
457 + assert (
458 + mappings["quant"].canonical_page_size_bytes
459 + == 4 * quant_spec.unpadded_page_size_bytes
460 + )
461 +
462 +
463 +def test_derive_refuses_foreign_worker_groups():
464 + attn_spec = _full_spec()
465 + kv_cache_config = _kv_cache_config(
466 + [KVCacheGroupSpec(layer_names=["attn"], kv_cache_spec=attn_spec)]
467 + )
468 + kv_caches = {"attn": _split_nhd_cache(attn_spec)}
469 + assert (
470 + derive_canonical_mappings(
471 + _vllm_config(tp=2, pp=2, total_kv_heads=4), kv_cache_config, kv_caches
472 + )
473 + == {}
474 + )
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+31 -6
@@ -101,6 +101,21 @@ def _allocate_and_reshape_kv_caches(
101 set_kv_cache_layout(None)
102
103
104 +def _single_rank_vllm_config(total_kv_heads: int):
105 + """A one-rank (TP=1) parallel config, as canonical mappings are derived
106 + from it."""
107 + vllm_config = MagicMock()
108 + parallel_config = vllm_config.parallel_config
109 + parallel_config.tensor_parallel_size = 1
110 + parallel_config.decode_context_parallel_size = 1
111 + parallel_config.prefill_context_parallel_size = 1
112 + parallel_config.cp_kv_cache_interleave_size = 1
113 + parallel_config.world_size = 1
114 + parallel_config.rank = 0
115 + vllm_config.model_config.get_total_num_kv_heads.return_value = total_kv_heads
116 + return vllm_config
117 +
118 +
119 def _make_worker(
120 kv_cache_config: KVCacheConfig,
121 replicated_layout: bool = False,
@@ -121,6 +136,7 @@ def _make_worker(
136
137 worker = OffloadingConnectorWorker(
138 spec=spec,
139 + vllm_config=_single_rank_vllm_config(NUM_KV_HEADS),
140 kv_cache_config=kv_cache_config,
141 )
142 worker.worker = MagicMock()
@@ -287,6 +303,7 @@ def test_offloading_connector_worker_accepts_plugin_spec_default_layout():
303
304 OffloadingConnectorWorker(
305 spec=spec,
306 + vllm_config=_single_rank_vllm_config(NUM_KV_HEADS),
307 kv_cache_config=KVCacheConfig(
308 num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
309 ),
@@ -522,6 +539,8 @@ def test_register_kv_caches(backend):
539 for actual, expected in zip(actual_refs, exp_refs):
540 assert actual.tensor_idx == expected.tensor_idx
541 assert actual.page_size_bytes == expected.page_size_bytes
542 + # Every layer gets a canonical mapping, certified or opaque
543 + assert actual.mapping is not None
544
545
546 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
@@ -622,9 +641,15 @@ def test_register_kv_caches_uniform_type(backend):
641 assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
642 assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
643
625 - assert group_refs[0] == CanonicalKVCacheRef(
626 - tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
627 - )
628 - assert group_refs[1] == CanonicalKVCacheRef(
629 - tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
630 - )
644 + for ref, expected_tensor_idx, expected_spec in (
645 + (group_refs[0], 0, spec_a),
646 + (group_refs[1], 1, spec_b),
647 + ):
648 + assert ref.tensor_idx == expected_tensor_idx
649 + assert ref.page_size_bytes == expected_spec.page_size_bytes
650 + assert ref.mapping is not None
651 +
652 + # Only layer_a matches the model's total KV head count, so layer_b gets an
653 + # opaque mapping rather than a certified, parallelism-agnostic one
654 + assert group_refs[0].mapping.parallelism_agnostic
655 + assert not group_refs[1].mapping.parallelism_agnostic
vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py new
+444
@@ -0,0 +1,444 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Derivation of canonical page mappings for KV offloading.
4 +
5 +The only place in the offloading stack that reasons about parallelism
6 +(TP/DCP/PCP); everything downstream consumes byte mappings. The canonical
7 +page of a layer is the full offloaded block without parallelism: all KV
8 +heads, all block_size * dcp * pcp tokens, in the worker's page encoding.
9 +Uncertifiable layers get an opaque fallback mapping (fail closed).
10 +"""
11 +
12 +from dataclasses import dataclass, replace
13 +from typing import TYPE_CHECKING
14 +
15 +import torch
16 +
17 +from vllm.v1.kv_cache_interface import (
18 + AttentionSpec,
19 + KVCacheConfig,
20 + KVCacheSpec,
21 + MambaSpec,
22 + MLAAttentionSpec,
23 + UniformTypeKVCacheSpecs,
24 +)
25 +from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun
26 +
27 +if TYPE_CHECKING:
28 + from vllm.config import VllmConfig
29 +
30 +
31 +@dataclass(frozen=True)
32 +class _RankContext:
33 + """Sharding parameters of one worker rank within the offload group."""
34 +
35 + tp_size: int
36 + dcp_size: int
37 + pcp_size: int
38 + interleave: int
39 + total_kv_heads: int
40 + rank: int
41 +
42 + @property
43 + def cp_size(self) -> int:
44 + return self.dcp_size * self.pcp_size
45 +
46 + @property
47 + def tp_rank(self) -> int:
48 + return self.rank % self.tp_size
49 +
50 + @property
51 + def total_cp_rank(self) -> int:
52 + pcp_rank = self.rank // self.tp_size
53 + return pcp_rank * self.dcp_size + self.tp_rank % self.dcp_size
54 +
55 +
56 +@dataclass(frozen=True)
57 +class ByteRegion:
58 + """A byte region within a page that repeats once per token."""
59 +
60 + local_offset: int
61 + canonical_offset: int
62 + bytes_per_token: int
63 + canonical_token_stride: int
64 +
65 +
66 +def _coalesce_runs(runs: list[CopyRun]) -> tuple[CopyRun, ...]:
67 + """Collapse contiguous fragments within and across runs to minimize the
68 + number of copy ops (e.g. a single-rank mapping becomes one whole-page run).
69 + """
70 + out: list[CopyRun] = []
71 + for run in runs:
72 + if (
73 + run.num_fragments > 1
74 + and run.local_stride == run.fragment_size
75 + and run.canonical_stride == run.fragment_size
76 + ):
77 + size = run.fragment_size * run.num_fragments
78 + run = CopyRun(run.local_offset, run.canonical_offset, size, 1, size, size)
79 + prev = out[-1] if out else None
80 + if (
81 + prev is not None
82 + and prev.num_fragments == 1
83 + and run.num_fragments == 1
84 + and prev.local_offset + prev.fragment_size == run.local_offset
85 + and prev.canonical_offset + prev.fragment_size == run.canonical_offset
86 + ):
87 + size = prev.fragment_size + run.fragment_size
88 + out[-1] = CopyRun(
89 + prev.local_offset, prev.canonical_offset, size, 1, size, size
90 + )
91 + else:
92 + out.append(run)
93 + return tuple(out)
94 +
95 +
96 +def _local_to_canonical_token(local_idx: int, ctx: _RankContext) -> int:
97 + """Canonical position of one of this rank's local token indices."""
98 + chunk, pos_in_chunk = divmod(local_idx, ctx.interleave)
99 + return (chunk * ctx.cp_size + ctx.total_cp_rank) * ctx.interleave + pos_in_chunk
100 +
101 +
102 +def _interleave_cp_tokens(
103 + regions: list[ByteRegion],
104 + num_tokens: int,
105 + ctx: _RankContext,
106 +) -> tuple[CopyRun, ...]:
107 + """Place each region's num_tokens rows at their canonical token positions,
108 + one run per chunk of interleaved tokens."""
109 + runs: list[CopyRun] = []
110 + for region in regions:
111 + if ctx.cp_size == 1:
112 + runs.append(
113 + CopyRun(
114 + region.local_offset,
115 + region.canonical_offset,
116 + region.bytes_per_token,
117 + num_tokens,
118 + region.bytes_per_token,
119 + region.canonical_token_stride,
120 + )
121 + )
122 + continue
123 + for chunk_start in range(0, num_tokens, ctx.interleave):
124 + canonical_token = _local_to_canonical_token(chunk_start, ctx)
125 + runs.append(
126 + CopyRun(
127 + region.local_offset + chunk_start * region.bytes_per_token,
128 + region.canonical_offset
129 + + canonical_token * region.canonical_token_stride,
130 + region.bytes_per_token,
131 + ctx.interleave,
132 + region.bytes_per_token,
133 + region.canonical_token_stride,
134 + )
135 + )
136 + return _coalesce_runs(runs)
137 +
138 +
139 +def _packed_kv_regions(
140 + kv_cache: torch.Tensor,
141 + spec: AttentionSpec,
142 + head_shard: int,
143 + num_head_shards: int,
144 + cp_size: int,
145 +) -> list[ByteRegion] | None:
146 + """K and V adjacent per (token, head), in NHD or HND stride order."""
147 + bs, heads = spec.block_size, spec.num_kv_heads
148 + elem = kv_cache.element_size()
149 + head_elems = 2 * spec.head_size
150 + if heads * bs * head_elems * elem != spec.real_page_size_bytes:
151 + return None
152 + _, head_stride, token_stride, inner_stride = kv_cache.stride()
153 + if inner_stride != 1:
154 + return None
155 + head_bytes = head_elems * elem
156 + token_row_bytes = heads * head_bytes
157 +
158 + if head_stride == head_elems and token_stride == heads * head_elems: # NHD
159 + return [
160 + ByteRegion(
161 + local_offset=0,
162 + canonical_offset=head_shard * token_row_bytes,
163 + bytes_per_token=token_row_bytes,
164 + canonical_token_stride=num_head_shards * token_row_bytes,
165 + )
166 + ]
167 +
168 + if head_stride == bs * head_elems and token_stride == head_elems: # HND
169 + canonical_span = bs * cp_size # canonical tokens per offloaded block
170 + return [
171 + ByteRegion(
172 + local_offset=head * bs * head_bytes,
173 + canonical_offset=(head_shard * heads + head)
174 + * canonical_span
175 + * head_bytes,
176 + bytes_per_token=head_bytes,
177 + canonical_token_stride=head_bytes,
178 + )
179 + for head in range(heads)
180 + ]
181 + return None
182 +
183 +
184 +def _split_kv_regions(
185 + kv_cache: torch.Tensor,
186 + spec: AttentionSpec,
187 + head_shard: int,
188 + num_head_shards: int,
189 + cp_size: int,
190 +) -> list[ByteRegion] | None:
191 + """K and V in separate page halves, in NHD or HND stride order."""
192 + bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size
193 + elem = kv_cache.element_size()
194 + if 2 * bs * heads * head_size * elem != spec.real_page_size_bytes:
195 + return None
196 + _, half_stride, token_stride, head_stride, inner_stride = kv_cache.stride()
197 + if inner_stride != 1 or half_stride != bs * heads * head_size:
198 + return None
199 + head_bytes = head_size * elem
200 + token_row_bytes = heads * head_bytes
201 + canonical_span = bs * cp_size # canonical tokens per offloaded block
202 +
203 + if token_stride == heads * head_size and head_stride == head_size: # NHD
204 + canonical_half_bytes = canonical_span * num_head_shards * token_row_bytes
205 + return [
206 + ByteRegion(
207 + local_offset=half * bs * token_row_bytes,
208 + canonical_offset=half * canonical_half_bytes
209 + + head_shard * token_row_bytes,
210 + bytes_per_token=token_row_bytes,
211 + canonical_token_stride=num_head_shards * token_row_bytes,
212 + )
213 + for half in range(2) # K, then V
214 + ]
215 +
216 + if head_stride == bs * head_size and token_stride == head_size: # HND
217 + local_half_bytes = bs * heads * head_bytes
218 + total_heads = num_head_shards * heads
219 + return [
220 + ByteRegion(
221 + local_offset=half * local_half_bytes + head * bs * head_bytes,
222 + canonical_offset=(half * total_heads + head_shard * heads + head)
223 + * canonical_span
224 + * head_bytes,
225 + bytes_per_token=head_bytes,
226 + canonical_token_stride=head_bytes,
227 + )
228 + for half in range(2)
229 + for head in range(heads)
230 + ]
231 + return None
232 +
233 +
234 +def _attention_byte_regions(
235 + kv_cache: torch.Tensor,
236 + spec: AttentionSpec,
237 + num_blocks: int,
238 + head_shard: int,
239 + num_head_shards: int,
240 + cp_size: int,
241 +) -> list[ByteRegion] | None:
242 + """Byte regions of an attention page, given this rank's head shard.
243 + None when the physical layout is not recognized (fail closed)."""
244 + bs, heads, head_size = spec.block_size, spec.num_kv_heads, spec.head_size
245 + if tuple(kv_cache.shape) == (num_blocks, heads, bs, 2 * head_size):
246 + return _packed_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size)
247 + if tuple(kv_cache.shape) == (num_blocks, 2, bs, heads, head_size):
248 + return _split_kv_regions(kv_cache, spec, head_shard, num_head_shards, cp_size)
249 + return None
250 +
251 +
252 +def _layer_mapping(
253 + spec: KVCacheSpec,
254 + kv_cache: torch.Tensor | list[torch.Tensor] | None,
255 + num_blocks: int,
256 + ctx: _RankContext,
257 +) -> CanonicalPageMapping | None:
258 + """Certified mapping for one layer at one rank, or None (fail closed)."""
259 + if not isinstance(spec, AttentionSpec):
260 + return None
261 + bs = spec.block_size
262 + page = spec.real_page_size_bytes
263 + if ctx.cp_size > 1 and (ctx.interleave > bs or bs % ctx.interleave):
264 + return None
265 +
266 + if isinstance(spec, MLAAttentionSpec):
267 + # TP-replicated latent; CP shards its tokens across the DCP groups
268 + if (
269 + spec.compress_ratio != 1
270 + or page % bs
271 + or ctx.tp_size % ctx.dcp_size
272 + or spec.kv_quant_mode.is_per_token_head
273 + ):
274 + return None
275 + row = page // bs
276 + return CanonicalPageMapping(
277 + canonical_page_size_bytes=ctx.cp_size * page,
278 + local_page_size_bytes=page,
279 + runs=_interleave_cp_tokens([ByteRegion(0, 0, row, row)], bs, ctx),
280 + num_writers=ctx.tp_size // ctx.dcp_size,
281 + writer_index=ctx.tp_rank // ctx.dcp_size,
282 + parallelism_agnostic=ctx.cp_size == 1,
283 + )
284 +
285 + if spec.kv_quant_mode.is_per_token_head or not isinstance(kv_cache, torch.Tensor):
286 + return None
287 + total, tp = ctx.total_kv_heads, ctx.tp_size
288 + if spec.num_kv_heads != max(1, total // tp):
289 + return None
290 + if total >= tp:
291 + if total % tp:
292 + return None
293 + num_head_shards, replication = tp, 1
294 + else:
295 + if tp % total:
296 + return None
297 + num_head_shards, replication = total, tp // total
298 + # DCP shards tokens across ranks holding replicated KV
299 + if replication % ctx.dcp_size:
300 + return None
301 +
302 + head_shard = ctx.tp_rank // replication
303 + regions = _attention_byte_regions(
304 + kv_cache, spec, num_blocks, head_shard, num_head_shards, ctx.cp_size
305 + )
306 + if regions is None:
307 + return None
308 + return CanonicalPageMapping(
309 + canonical_page_size_bytes=ctx.cp_size * num_head_shards * page,
310 + local_page_size_bytes=page,
311 + runs=_interleave_cp_tokens(regions, bs, ctx),
312 + num_writers=replication // ctx.dcp_size,
313 + writer_index=(ctx.tp_rank % replication) // ctx.dcp_size,
314 + parallelism_agnostic=ctx.cp_size == 1,
315 + )
316 +
317 +
318 +def _opaque_fallback_mapping(
319 + page_size_bytes: int, num_ranks: int, rank: int
320 +) -> CanonicalPageMapping:
321 + """Fallback: place the worker's page whole at a worker-exclusive offset."""
322 + run = CopyRun(
323 + 0, rank * page_size_bytes, page_size_bytes, 1, page_size_bytes, page_size_bytes
324 + )
325 + return CanonicalPageMapping(
326 + canonical_page_size_bytes=num_ranks * page_size_bytes,
327 + local_page_size_bytes=page_size_bytes,
328 + runs=(run,),
329 + num_writers=1,
330 + writer_index=0,
331 + parallelism_agnostic=False,
332 + )
333 +
334 +
335 +def _run_intervals(runs: tuple[CopyRun, ...], canonical: bool) -> list[tuple[int, int]]:
336 + intervals = []
337 + for run in runs:
338 + offset = run.canonical_offset if canonical else run.local_offset
339 + stride = run.canonical_stride if canonical else run.local_stride
340 + for i in range(run.num_fragments):
341 + start = offset + i * stride
342 + intervals.append((start, start + run.fragment_size))
343 + return sorted(intervals)
344 +
345 +
346 +def _is_exact_partition(intervals: list[tuple[int, int]], size: int) -> bool:
347 + return (
348 + bool(intervals)
349 + and intervals[0][0] == 0
350 + and intervals[-1][1] == size
351 + and all(a[1] == b[0] for a, b in zip(intervals, intervals[1:]))
352 + )
353 +
354 +
355 +def _verify_tiling(layer_name: str, per_rank: list[CanonicalPageMapping]) -> None:
356 + """Whichever ranks a block elects as writers must tile the canonical page
357 + exactly once, and each rank's runs must cover exactly its local page."""
358 + size = per_rank[0].canonical_page_size_bytes
359 + num_writers = per_rank[0].num_writers
360 + for mapping in per_rank:
361 + assert mapping.canonical_page_size_bytes == size
362 + assert mapping.num_writers == num_writers
363 + local = _run_intervals(mapping.runs, canonical=False)
364 + assert _is_exact_partition(local, mapping.local_page_size_bytes), (
365 + f"runs do not cover the local page of layer {layer_name}"
366 + )
367 + for block_id in range(num_writers):
368 + stored: list[tuple[int, int]] = []
369 + for mapping in per_rank:
370 + if mapping.is_writer(block_id):
371 + stored += _run_intervals(mapping.runs, canonical=True)
372 + stored.sort()
373 + assert _is_exact_partition(stored, size), (
374 + f"writers of block {block_id} do not tile the canonical page "
375 + f"of layer {layer_name}"
376 + )
377 +
378 +
379 +def _unpadded_page_size(spec: KVCacheSpec) -> int | None:
380 + if isinstance(spec, AttentionSpec):
381 + return spec.unpadded_page_size_bytes
382 + if isinstance(spec, MambaSpec):
383 + return replace(spec, page_size_padded=None).page_size_bytes
384 + return None
385 +
386 +
387 +def derive_canonical_mappings(
388 + vllm_config: "VllmConfig",
389 + kv_cache_config: KVCacheConfig,
390 + kv_caches: dict[str, torch.Tensor | list[torch.Tensor]],
391 +) -> dict[str, CanonicalPageMapping]:
392 + """Per-layer canonical page mappings for this worker.
393 +
394 + Empty when the worker group is not exactly the TP x PCP grid; layers
395 + absent from the result have no canonical representation.
396 + """
397 + parallel_config = vllm_config.parallel_config
398 + tp_size = parallel_config.tensor_parallel_size
399 + pcp_size = parallel_config.prefill_context_parallel_size
400 + group_size = tp_size * pcp_size
401 + if parallel_config.world_size != group_size:
402 + return {}
403 +
404 + def ctx(rank: int) -> _RankContext:
405 + return _RankContext(
406 + tp_size=tp_size,
407 + dcp_size=parallel_config.decode_context_parallel_size,
408 + pcp_size=pcp_size,
409 + interleave=parallel_config.cp_kv_cache_interleave_size,
410 + total_kv_heads=vllm_config.model_config.get_total_num_kv_heads(),
411 + rank=rank,
412 + )
413 +
414 + my_rank = parallel_config.rank
415 + num_blocks = kv_cache_config.num_blocks
416 +
417 + mappings: dict[str, CanonicalPageMapping] = {}
418 + for kv_cache_group in kv_cache_config.kv_cache_groups:
419 + group_kv_cache_spec = kv_cache_group.kv_cache_spec
420 + if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
421 + per_layer_specs = group_kv_cache_spec.kv_cache_specs
422 + else:
423 + per_layer_specs = {}
424 + for layer_name in kv_cache_group.layer_names:
425 + spec = per_layer_specs.get(layer_name, group_kv_cache_spec)
426 + per_rank: list[CanonicalPageMapping] = []
427 + for rank in range(group_size):
428 + mapping = _layer_mapping(
429 + spec, kv_caches.get(layer_name), num_blocks, ctx(rank)
430 + )
431 + if mapping is None:
432 + break
433 + per_rank.append(mapping)
434 + if len(per_rank) != group_size:
435 + page = _unpadded_page_size(spec)
436 + if page is None:
437 + continue
438 + per_rank = [
439 + _opaque_fallback_mapping(page, group_size, rank)
440 + for rank in range(group_size)
441 + ]
442 + _verify_tiling(layer_name, per_rank)
443 + mappings[layer_name] = per_rank[my_rank]
444 + return mappings
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+20
@@ -5,6 +5,10 @@ from dataclasses import replace
5
6 import torch
7
8 +from vllm.config import VllmConfig
9 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import (
10 + derive_canonical_mappings,
11 +)
12 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
13 OffloadingConnectorMetadata,
14 OffloadingWorkerMetadata,
@@ -40,9 +44,11 @@ class OffloadingConnectorWorker:
44 def __init__(
45 self,
46 spec: OffloadingSpec,
47 + vllm_config: "VllmConfig",
48 kv_cache_config: KVCacheConfig,
49 ):
50 self.spec = spec
51 + self.vllm_config = vllm_config
52 self.kv_cache_config = kv_cache_config
53 self.worker: OffloadingWorker | None = None
54 # Non-writers still ack: pending_count waits for world_size per job.
@@ -63,6 +69,9 @@ class OffloadingConnectorWorker:
69 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
70 kv_cache_config = self.kv_cache_config
71 num_blocks = kv_cache_config.num_blocks
72 + mappings = derive_canonical_mappings(
73 + self.vllm_config, kv_cache_config, kv_caches
74 + )
75
76 # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
77 # stride(0) as the manager-block stride (equals total_num_bytes_per_block).
@@ -201,10 +210,21 @@ class OffloadingConnectorWorker:
210
211 curr_tensor_idx = len(block_tensors) - 1
212 for layer_name in tensor_layer_names:
213 + mapping = (
214 + mappings.get(layer_name)
215 + if len(tensors_per_block[first_layer_name]) == 1
216 + else None
217 + )
218 + assert (
219 + mapping is None
220 + or mapping.local_page_size_bytes
221 + == unpadded_page_size_bytes[layer_name]
222 + )
223 block_data_refs[layer_name].append(
224 CanonicalKVCacheRef(
225 tensor_idx=curr_tensor_idx,
226 page_size_bytes=(unpadded_page_size_bytes[layer_name]),
227 + mapping=mapping,
228 )
229 )
230
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+3 -1
@@ -75,7 +75,9 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
75 spec, vllm_config, kv_cache_config
76 )
77 elif role == KVConnectorRole.WORKER:
78 - self.connector_worker = OffloadingConnectorWorker(spec, kv_cache_config)
78 + self.connector_worker = OffloadingConnectorWorker(
79 + spec, vllm_config, kv_cache_config
80 + )
81
82 def shutdown(self) -> None:
83 if self.connector_worker is not None:
vllm/v1/kv_offload/base.py
+44
@@ -457,6 +457,48 @@ class CanonicalKVCacheTensor:
457 page_size_bytes: int
458
459
460 +@dataclass(frozen=True)
461 +class CopyRun:
462 + """A strided byte correspondence between this worker's physical page and
463 + a canonical page: for i in range(num_fragments), fragment i spans
464 + [local_offset + i * local_stride, +fragment_size) in the worker's page and
465 + [canonical_offset + i * canonical_stride, +fragment_size) canonically."""
466 +
467 + local_offset: int
468 + canonical_offset: int
469 + fragment_size: int
470 + num_fragments: int
471 + local_stride: int
472 + canonical_stride: int
473 +
474 +
475 +@dataclass(frozen=True)
476 +class CanonicalPageMapping:
477 + """How this worker's page maps into a canonical (parallelism-free) page.
478 + In-process only, never serialized. Runs cover the full local page in both
479 + directions; ranks holding identical bytes take turns writing them.
480 + """
481 +
482 + # Size of the canonical page in bytes
483 + canonical_page_size_bytes: int
484 + # Size of this worker's (un-padded) page in bytes
485 + local_page_size_bytes: int
486 + # Byte correspondences between this worker's page and a canonical page
487 + runs: tuple[CopyRun, ...]
488 + # Number of ranks holding these exact bytes
489 + num_writers: int
490 + # This worker's index among those ranks
491 + writer_index: int
492 + # Canonical bytes identical under any parallel config with this block span
493 + parallelism_agnostic: bool
494 +
495 + def is_writer(self, block_id: int) -> bool:
496 + """Whether this worker stores the canonical page of the given block.
497 + Rotating by block spreads writes across ranks holding identical bytes.
498 + """
499 + return block_id % self.num_writers == self.writer_index
500 +
501 +
502 @dataclass
503 class CanonicalKVCacheRef:
504 """
@@ -468,6 +510,8 @@ class CanonicalKVCacheRef:
510 tensor_idx: int
511 # The un-padded page size per block in bytes
512 page_size_bytes: int
513 + # How this worker's page maps into a canonical page; None = uncertified
514 + mapping: CanonicalPageMapping | None = None
515
516
517 @dataclass