@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 81840a172f

[KV Connector] Canonical CPU layout for parallelism-agnostic KV offload (#48414)

Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>

Itay Etelis committed Aug 10, 2026 at 10:44 UTC 81840a172f2886f4c59dafd9dce9917bb789280c
12 files changed +891 -103
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+58
@@ -47,6 +47,9 @@ def _make_vllm_config(
47 config.cache_config.cache_dtype = torch.float16
48 config.model_config.model = "test-model"
49 config.model_config.use_mla = False
50 + # _full_attention_spec's heads at tp=1: the parallelism-agnostic gate
51 + # requires the head shard to cover the model's KV heads exactly
52 + config.model_config.get_total_num_kv_heads.return_value = 4
53 world_size = (
54 tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
55 )
@@ -603,6 +606,61 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
606 assert not _parallelism_agnostic(kv_cache_groups)
607
608
609 +def test_canonical_layout_widens_parallelism_agnostic_to_mla():
610 + """The canonical layout dedups the TP-replicated MLA latent into one
611 + portable copy, so the gate admits MLA — but only when canonical_layout
612 + is requested."""
613 + mla_groups = [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))]
614 + assert not _parallelism_agnostic(mla_groups)
615 +
616 + config = _make_vllm_config(extra_config={"canonical_layout": True})
617 + kv_cache_config = KVCacheConfig(
618 + num_blocks=0,
619 + kv_cache_tensors=[],
620 + kv_cache_groups=mla_groups,
621 + )
622 + offloading_config = build_offloading_config(config, kv_cache_config)
623 + assert offloading_config.parallel.is_parallelism_agnostic
624 + assert offloading_config.canonical_layout
625 +
626 + # hybrid groupings stay out: their non-full-attention layers can only
627 + # derive opaque (exact-topology) mappings
628 + hybrid_config = KVCacheConfig(
629 + num_blocks=0,
630 + kv_cache_tensors=[],
631 + kv_cache_groups=[
632 + KVCacheGroupSpec(["l0"], _full_attention_spec()),
633 + KVCacheGroupSpec(["l1"], _full_attention_spec()),
634 + ],
635 + )
636 + assert not build_offloading_config(
637 + config, hybrid_config
638 + ).parallel.is_parallelism_agnostic
639 +
640 +
641 +def test_canonical_layout_certifies_v2_model_runner():
642 + """Canonical bytes are certified per layer against live tensor strides at
643 + registration, so the static gate must not depend on the model-runner
644 + version — the v2 runner is the case the canonical layout exists for."""
645 + kv_cache_config = KVCacheConfig(
646 + num_blocks=0,
647 + kv_cache_tensors=[],
648 + kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
649 + )
650 +
651 + config = _make_vllm_config()
652 + config.use_v2_model_runner = True
653 + assert not build_offloading_config(
654 + config, kv_cache_config
655 + ).parallel.is_parallelism_agnostic
656 +
657 + config = _make_vllm_config(extra_config={"canonical_layout": True})
658 + config.use_v2_model_runner = True
659 + assert build_offloading_config(
660 + config, kv_cache_config
661 + ).parallel.is_parallelism_agnostic
662 +
663 +
664 def test_parallelism_agnostic_disabled_on_v2_model_runner():
665 config = _make_vllm_config()
666 config.use_v2_model_runner = True
tests/v1/kv_offload/cpu/test_canonical_layout.py new
+274
@@ -0,0 +1,274 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import time
4 +import uuid
5 +
6 +import numpy as np
7 +import pytest
8 +import torch
9 +
10 +from vllm.v1.kv_offload.base import (
11 + CanonicalKVCacheRef,
12 + CanonicalPageMapping,
13 + CopyRun,
14 + GPULoadStoreSpec,
15 +)
16 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
17 +from vllm.v1.kv_offload.cpu.gpu_worker import (
18 + SingleDirectionOffloadingHandler,
19 + _build_copy_plan,
20 + _canonical_block_sizes,
21 + _canonical_page_ids,
22 + pin_mmap_region,
23 +)
24 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
25 +
26 +
27 +def _ref(mapping: CanonicalPageMapping, tensor_idx: int = 0) -> CanonicalKVCacheRef:
28 + return CanonicalKVCacheRef(
29 + tensor_idx=tensor_idx,
30 + page_size_bytes=mapping.local_page_size_bytes,
31 + mapping=mapping,
32 + )
33 +
34 +
35 +def _nhd_mapping() -> CanonicalPageMapping:
36 + # 4-token NHD page at tp=4, rank 2: K and V runs of 4x128B fragments
37 + runs = (
38 + CopyRun(0, 256, 128, 4, 128, 512),
39 + CopyRun(512, 2304, 128, 4, 128, 512),
40 + )
41 + return CanonicalPageMapping(4096, 1024, runs, 1, 0, True)
42 +
43 +
44 +def test_copy_plan_unrolls_runs():
45 + plan = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=True)
46 + k_dst = [256, 768, 1280, 1792]
47 + assert plan.frag_offsets_src.tolist() == [0, 128, 256, 384, 512, 640, 768, 896]
48 + assert plan.frag_offsets_dst.tolist() == k_dst + [2048 + o for o in k_dst]
49 + assert plan.frag_sizes.tolist() == [128] * 8
50 + assert plan.num_frags == 8
51 +
52 +
53 +def test_load_direction_swaps_offsets():
54 + store = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=True)
55 + load = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=False)
56 + assert np.array_equal(store.frag_offsets_src, load.frag_offsets_dst)
57 + assert np.array_equal(store.frag_offsets_dst, load.frag_offsets_src)
58 + assert np.array_equal(store.frag_sizes, load.frag_sizes)
59 +
60 +
61 +def test_writer_rotation_matches_is_writer():
62 + # Replicas take turns writing shared canonical pages, keyed by the
63 + # CPU-side canonical page id; the enumeration must agree with is_writer
64 + identity = CopyRun(0, 0, 256, 1, 256, 256)
65 + mapping = CanonicalPageMapping(256, 256, (identity,), 2, 1, True)
66 + ids = _canonical_page_ids(
67 + np.array([3, 7, 9]), blocks_per_chunk=4, count=10, skip_count=2
68 + )
69 + assert ids.tolist() == [14, 15, 28, 29, 30, 31, 36, 37, 38, 39]
70 + mask = ids % mapping.num_writers == mapping.writer_index
71 + assert mask.tolist() == [mapping.is_writer(int(b)) for b in ids]
72 +
73 +
74 +def test_canonical_block_sizes_take_max_per_tensor():
75 + identity = CopyRun(0, 0, 512, 1, 512, 512)
76 + small = CanonicalPageMapping(2048, 512, (identity,), 1, 0, False)
77 + refs = [[_ref(_nhd_mapping(), 0), _ref(small, 0)], [_ref(small, 1)]]
78 + assert _canonical_block_sizes(refs, 2) == [4096, 2048]
79 +
80 +
81 +def _tp2_rank_mapping(rank: int) -> CanonicalPageMapping:
82 + # 4-token NHD page, 4 total heads of 64B, tp=2: rank holds 2 heads,
83 + # so K and V each scatter as 4 per-token 128B fragments
84 + runs = (
85 + CopyRun(0, rank * 128, 128, 4, 128, 256),
86 + CopyRun(512, 1024 + rank * 128, 128, 4, 128, 256),
87 + )
88 + return CanonicalPageMapping(2048, 1024, runs, 1, 0, True)
89 +
90 +
91 +def _whole_page_mapping() -> CanonicalPageMapping:
92 + identity = CopyRun(0, 0, 2048, 1, 2048, 2048)
93 + return CanonicalPageMapping(2048, 2048, (identity,), 1, 0, True)
94 +
95 +
96 +def _transfer(handler, num_blocks: int, gpu_to_cpu: bool) -> None:
97 + block_ids = list(range(num_blocks))
98 + gpu_spec = GPULoadStoreSpec(
99 + block_ids, group_sizes=(num_blocks,), block_indices=(0,)
100 + )
101 + cpu_spec = CPULoadStoreSpec(block_ids)
102 + src, dst = (gpu_spec, cpu_spec) if gpu_to_cpu else (cpu_spec, gpu_spec)
103 + assert handler.transfer_async(0, src, dst)
104 + deadline = time.time() + 30
105 + while time.time() < deadline:
106 + if handler.get_finished():
107 + return
108 + time.sleep(0.001)
109 + raise TimeoutError("transfer did not complete")
110 +
111 +
112 +def _canonical_handler(gpu_tensor, cpu_tensor, mapping, gpu_to_cpu):
113 + page = mapping.local_page_size_bytes
114 + refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=page, mapping=mapping)]]
115 + return SingleDirectionOffloadingHandler(
116 + gpu_tensors=[gpu_tensor],
117 + cpu_tensors=[cpu_tensor],
118 + blocks_per_chunk=1,
119 + layer_refs_per_group=refs,
120 + gpu_to_cpu=gpu_to_cpu,
121 + canonical_layout=True,
122 + )
123 +
124 +
125 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
126 +def test_gpu_roundtrip_assembles_canonical_page_across_ranks():
127 + """Two TP2-style rank handlers must scatter into one shared canonical
128 + CPU page such that each rank reloads bit-exact and a whole-page (TP1)
129 + reader sees both shards — the cross-topology contract."""
130 + torch.manual_seed(0)
131 + num_blocks = 4
132 + gpu_rank = [
133 + torch.randint(-128, 128, (num_blocks, 1024), dtype=torch.int8, device="cuda")
134 + for _ in range(2)
135 + ]
136 + cpu_canonical = torch.zeros(num_blocks, 2048, dtype=torch.int8, pin_memory=True)
137 +
138 + for rank in (0, 1):
139 + store = _canonical_handler(
140 + gpu_rank[rank], cpu_canonical, _tp2_rank_mapping(rank), gpu_to_cpu=True
141 + )
142 + _transfer(store, num_blocks, gpu_to_cpu=True)
143 + torch.accelerator.synchronize()
144 +
145 + # independent oracle: replay each rank's runs in numpy
146 + expected = np.zeros((num_blocks, 2048), dtype=np.int8)
147 + for rank in (0, 1):
148 + local = gpu_rank[rank].cpu().numpy()
149 + for run in _tp2_rank_mapping(rank).runs:
150 + for i in range(run.num_fragments):
151 + lo = run.local_offset + i * run.local_stride
152 + co = run.canonical_offset + i * run.canonical_stride
153 + expected[:, co : co + run.fragment_size] = local[
154 + :, lo : lo + run.fragment_size
155 + ]
156 + assert np.array_equal(cpu_canonical.numpy(), expected)
157 +
158 + # each rank reloads its shard bit-exact
159 + gpu_back = torch.zeros(num_blocks, 1024, dtype=torch.int8, device="cuda")
160 + load = _canonical_handler(
161 + gpu_back, cpu_canonical, _tp2_rank_mapping(0), gpu_to_cpu=False
162 + )
163 + _transfer(load, num_blocks, gpu_to_cpu=False)
164 + torch.accelerator.synchronize()
165 + assert torch.equal(gpu_back, gpu_rank[0])
166 +
167 + # a whole-page reader (TP1 topology) sees the assembled page
168 + gpu_full = torch.zeros(num_blocks, 2048, dtype=torch.int8, device="cuda")
169 + load_full = _canonical_handler(
170 + gpu_full, cpu_canonical, _whole_page_mapping(), gpu_to_cpu=False
171 + )
172 + _transfer(load_full, num_blocks, gpu_to_cpu=False)
173 + torch.accelerator.synchronize()
174 + assert torch.equal(gpu_full.cpu(), cpu_canonical)
175 +
176 +
177 +# 4-token NHD page, 4 total heads of 64 bytes: canonical page holds
178 +# [K: token x head][V: token x head] = 2048 bytes per block
179 +_TOTAL_HEADS = 4
180 +_HEAD_BYTES = 64
181 +_BLOCK_TOKENS = 4
182 +_CANONICAL_PAGE = 2 * _BLOCK_TOKENS * _TOTAL_HEADS * _HEAD_BYTES
183 +
184 +
185 +def _nhd_shard_mapping(tp: int, rank: int) -> CanonicalPageMapping:
186 + """Rank's head-shard mapping into the canonical NHD page at the given tp."""
187 + local_heads = _TOTAL_HEADS // tp
188 + frag = local_heads * _HEAD_BYTES
189 + canonical_row = _TOTAL_HEADS * _HEAD_BYTES
190 + k_run = CopyRun(0, rank * frag, frag, _BLOCK_TOKENS, frag, canonical_row)
191 + v_run = CopyRun(
192 + _BLOCK_TOKENS * frag,
193 + _BLOCK_TOKENS * canonical_row + rank * frag,
194 + frag,
195 + _BLOCK_TOKENS,
196 + frag,
197 + canonical_row,
198 + )
199 + return CanonicalPageMapping(
200 + _CANONICAL_PAGE, 2 * _BLOCK_TOKENS * frag, (k_run, v_run), 1, 0, True
201 + )
202 +
203 +
204 +def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor:
205 + """This rank's local page rows out of the (blocks, 2, tokens, heads,
206 + head_bytes) ground truth."""
207 + local_heads = _TOTAL_HEADS // tp
208 + shard = full_kv[:, :, :, rank * local_heads : (rank + 1) * local_heads, :]
209 + return shard.reshape(full_kv.shape[0], -1).contiguous()
210 +
211 +
212 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
213 +@pytest.mark.parametrize("writer_tp,reader_tp", [(2, 4), (4, 2), (2, 1), (4, 4)])
214 +def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int):
215 + """KV written at one tp must be readable at another: writer ranks scatter
216 + head shards into a shared canonical region, reader ranks gather their own
217 + shards, and every reader must see the writers' ground-truth bytes."""
218 + torch.manual_seed(0)
219 + num_blocks = 3
220 + full_kv = torch.randint(
221 + -128,
222 + 128,
223 + (num_blocks, 2, _BLOCK_TOKENS, _TOTAL_HEADS, _HEAD_BYTES),
224 + dtype=torch.int8,
225 + )
226 +
227 + engine_id = str(uuid.uuid4())
228 + row_stride = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
229 + assert row_stride >= _CANONICAL_PAGE
230 + regions: list[SharedOffloadRegion] = []
231 +
232 + def canonical_view(rank: int, world_size: int) -> torch.Tensor:
233 + region = SharedOffloadRegion(
234 + engine_id=engine_id,
235 + num_blocks=num_blocks,
236 + rank=rank,
237 + kv_bytes_per_block=row_stride,
238 + cpu_page_size=row_stride // world_size,
239 + )
240 + regions.append(region)
241 + # The Triton load path dereferences CPU pointers on the GPU, which is
242 + # only legal on pinned memory; production pins via CPUOffloadingWorker
243 + pin_mmap_region(region)
244 + return region.create_next_canonical_view(_CANONICAL_PAGE)
245 +
246 + try:
247 + for rank in range(writer_tp):
248 + store = _canonical_handler(
249 + _head_shard(full_kv, writer_tp, rank).cuda(),
250 + canonical_view(rank, writer_tp),
251 + _nhd_shard_mapping(writer_tp, rank),
252 + gpu_to_cpu=True,
253 + )
254 + _transfer(store, num_blocks, gpu_to_cpu=True)
255 + torch.accelerator.synchronize()
256 +
257 + for rank in range(reader_tp):
258 + expected = _head_shard(full_kv, reader_tp, rank)
259 + gpu_out = torch.zeros_like(expected, device="cuda")
260 + load = _canonical_handler(
261 + gpu_out,
262 + canonical_view(rank, reader_tp),
263 + _nhd_shard_mapping(reader_tp, rank),
264 + gpu_to_cpu=False,
265 + )
266 + _transfer(load, num_blocks, gpu_to_cpu=False)
267 + torch.accelerator.synchronize()
268 + assert torch.equal(gpu_out.cpu(), expected), (
269 + f"reader tp={reader_tp} rank={rank} bytes diverge from the "
270 + f"tp={writer_tp} writers' ground truth"
271 + )
272 + finally:
273 + for region in regions:
274 + region.cleanup()
tests/v1/kv_offload/cpu/test_gpu_worker.py
+2 -2
@@ -195,7 +195,7 @@ def test_transfer(
195 assert finished[0].success
196 assert finished[0].transfer_size == (
197 len(gpu_blocks)
198 - * sum([x.page_size_bytes for x in handler.kv_cache_groups_data_refs[0]])
198 + * sum([x.page_size_bytes for x in handler.layer_refs_per_group[0]])
199 )
200 assert finished[0].transfer_time > 0
201 assert finished[0].transfer_time < (time.time() - start_time)
@@ -409,7 +409,7 @@ def test_transfer_multi_group(
409 expected_bytes = sum(
410 group_size * sum([x.page_size_bytes for x in data_refs])
411 for group_size, data_refs in zip(
412 - group_sizes, handler.kv_cache_groups_data_refs
412 + group_sizes, handler.layer_refs_per_group
413 )
414 )
415 assert finished[0].transfer_size == expected_bytes
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+48 -46
@@ -173,7 +173,7 @@ def _mp_race_construct_and_write(
173 kv_bytes_per_block=num_workers * cpu_page_size,
174 cpu_page_size=cpu_page_size,
175 )
176 - t = region.create_next_view(cpu_page_size)
176 + t = region.create_next_worker_view(cpu_page_size)
177 t[:, :] = fill_value
178 done_queue.put({"rank": rank, "error": None})
179 cleanup_queue.get() # wait for parent's verification to finish
@@ -190,110 +190,112 @@ def iid():
190
191
192 # ---------------------------------------------------------------------------
193 -# create_next_view — shape, stride and storage offset
193 +# create_next_worker_view — shape, stride and storage offset
194 # ---------------------------------------------------------------------------
195
196
197 -def test_create_next_view_shape_and_stride(iid):
197 +def test_create_next_worker_view_shape_and_stride(iid):
198 """Returned tensor must have shape (num_blocks, tensor_page_size) and
199 stride (row_stride, 1) where row_stride = cpu_page_size * num_workers."""
200 with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r:
201 - t = r.create_next_view(PAGE_SIZE)
201 + t = r.create_next_worker_view(PAGE_SIZE)
202 assert t.shape == (4, PAGE_SIZE)
203 # num_workers=1 → row_stride = cpu_page_size
204 assert t.stride() == (2 * PAGE_SIZE, 1)
205 del t
206
207
208 -def test_create_next_view_storage_offset_rank0(iid):
208 +def test_create_next_worker_view_storage_offset_rank0(iid):
209 """rank=0 worker's first tensor must start at byte 0 of the mmap."""
210 with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r:
211 - t = r.create_next_view(PAGE_SIZE)
211 + t = r.create_next_worker_view(PAGE_SIZE)
212 assert t.data_ptr() == r._base.data_ptr() # storage_offset == 0
213 del t
214
215
216 -def test_create_next_view_storage_offset_rank1(iid):
216 +def test_create_next_worker_view_storage_offset_rank1(iid):
217 """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap."""
218 with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1):
219 - t1 = r1.create_next_view(PAGE_SIZE)
219 + t1 = r1.create_next_worker_view(PAGE_SIZE)
220 assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE
221 del t1
222
223
224 -def test_create_next_view_row_stride_with_multiple_workers(iid):
224 +def test_create_next_worker_view_row_stride_with_multiple_workers(iid):
225 """With num_workers=4, row_stride must be 4 * cpu_page_size."""
226 with _region(iid, num_blocks=2, num_workers=4) as r:
227 - t = r.create_next_view(PAGE_SIZE)
227 + t = r.create_next_worker_view(PAGE_SIZE)
228 assert t.stride(0) == 4 * PAGE_SIZE
229 del t
230
231
232 # ---------------------------------------------------------------------------
233 -# create_next_view — cursor advancement
233 +# create_next_worker_view — cursor advancement
234 # ---------------------------------------------------------------------------
235
236
237 -def test_create_next_view_cursor_advances(iid):
238 - """Each call to create_next_view must advance _worker_offset by tensor_page_size."""
237 +def test_create_next_worker_view_cursor_advances(iid):
238 + """Each create_next_worker_view call must advance _worker_offset by
239 + tensor_page_size."""
240 with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r:
241 assert r._worker_offset == 0
241 - r.create_next_view(PAGE_SIZE)
242 + r.create_next_worker_view(PAGE_SIZE)
243 assert r._worker_offset == PAGE_SIZE
243 - r.create_next_view(PAGE_SIZE)
244 + r.create_next_worker_view(PAGE_SIZE)
245 assert r._worker_offset == 2 * PAGE_SIZE
245 - r.create_next_view(PAGE_SIZE)
246 + r.create_next_worker_view(PAGE_SIZE)
247 assert r._worker_offset == 3 * PAGE_SIZE # exactly at area end
248
249
249 -def test_create_next_view_exact_fill_succeeds(iid):
250 +def test_create_next_worker_view_exact_fill_succeeds(iid):
251 """Allocations whose total exactly equals cpu_page_size must all succeed."""
252 with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
252 - r.create_next_view(PAGE_SIZE) # first half
253 - r.create_next_view(PAGE_SIZE) # fills to area end — must not raise
253 + r.create_next_worker_view(PAGE_SIZE) # first half
254 + r.create_next_worker_view(PAGE_SIZE) # fills to area end — must not raise
255
256
257 # ---------------------------------------------------------------------------
257 -# create_next_view — overflow guard
258 +# create_next_worker_view — overflow guard
259 # ---------------------------------------------------------------------------
260
261
261 -def test_create_next_view_single_overflow_raises(iid):
262 +def test_create_next_worker_view_single_overflow_raises(iid):
263 """A single allocation larger than cpu_page_size must raise AssertionError."""
264 with (
265 _region(iid) as r,
266 pytest.raises(AssertionError, match="exceeds worker area end"),
267 ):
267 - r.create_next_view(PAGE_SIZE + 1)
268 + r.create_next_worker_view(PAGE_SIZE + 1)
269
270
270 -def test_create_next_view_cumulative_overflow_raises(iid):
271 +def test_create_next_worker_view_cumulative_overflow_raises(iid):
272 """Successive allocations that cumulatively exceed cpu_page_size must raise."""
273 with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r:
273 - r.create_next_view(PAGE_SIZE) # ok — half used
274 - r.create_next_view(PAGE_SIZE) # ok — full
274 + r.create_next_worker_view(PAGE_SIZE) # ok — half used
275 + r.create_next_worker_view(PAGE_SIZE) # ok — full
276 with pytest.raises(AssertionError, match="exceeds worker area end"):
276 - r.create_next_view(1) # one byte too many
277 + r.create_next_worker_view(1) # one byte too many
278
279
279 -def test_create_next_view_overflow_does_not_mutate_cursor(iid):
280 - """A failed create_next_view must leave _worker_offset unchanged."""
280 +def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid):
281 + """A failed create_next_worker_view must leave _worker_offset unchanged."""
282 with _region(iid) as r:
283 offset_before = r._worker_offset
284 with pytest.raises(AssertionError):
284 - r.create_next_view(PAGE_SIZE + 1)
285 + r.create_next_worker_view(PAGE_SIZE + 1)
286 assert r._worker_offset == offset_before
287
288
289 # ---------------------------------------------------------------------------
289 -# create_next_view — data correctness and layout
290 +# create_next_worker_view — data correctness and layout
291 # ---------------------------------------------------------------------------
292
293
293 -def test_create_next_view_write_visible_in_raw_mmap(iid):
294 - """Writes into a create_next_view view must appear at the correct raw mmap offset"""
294 +def test_create_next_worker_view_write_visible_in_raw_mmap(iid):
295 + """Writes into a create_next_worker_view view must appear at the correct
296 + raw mmap offset"""
297 with _region(iid, num_blocks=4) as r:
296 - t = r.create_next_view(PAGE_SIZE)
298 + t = r.create_next_worker_view(PAGE_SIZE)
299 t[2, :] = 42 # write to block row 2
300
301 raw = memoryview(r.mmap_obj)
@@ -303,11 +305,11 @@ def test_create_next_view_write_visible_in_raw_mmap(iid):
305 del raw, t
306
307
306 -def test_create_next_view_multi_tensor_layout(iid):
308 +def test_create_next_worker_view_multi_tensor_layout(iid):
309 """Two tensors from the same worker land at consecutive byte offsets per row."""
310 with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r:
309 - ta = r.create_next_view(PAGE_SIZE)
310 - tb = r.create_next_view(PAGE_SIZE)
311 + ta = r.create_next_worker_view(PAGE_SIZE)
312 + tb = r.create_next_worker_view(PAGE_SIZE)
313
314 ta[:, :] = 1
315 tb[:, :] = 2
@@ -322,8 +324,8 @@ def test_create_next_view_multi_tensor_layout(iid):
324 del raw, ta, tb
325
326
325 -def test_create_next_view_multiprocess_slots(iid):
326 - """Each worker process calls create_next_view and writes distinct data;
327 +def test_create_next_worker_view_multiprocess_slots(iid):
328 + """Each worker process calls create_next_worker_view and writes distinct data;
329 the parent verifies each slot lands at the correct interleaved offset."""
330 num_workers = 2
331 num_blocks = 4
@@ -356,7 +358,7 @@ def test_create_next_view_multiprocess_slots(iid):
358 )
359 child.start()
360
359 - t0 = region.create_next_view(PAGE_SIZE)
361 + t0 = region.create_next_worker_view(PAGE_SIZE)
362 t0[:, :] = 11
363
364 result = done_queue.get(timeout=30)
@@ -379,13 +381,13 @@ def test_create_next_view_multiprocess_slots(iid):
381 _cleanup_file(region.mmap_path)
382
383
382 -def test_create_next_view_worker_isolation(iid):
384 +def test_create_next_worker_view_worker_isolation(iid):
385 """Writes by worker 0 must not affect worker 1's slot and vice versa."""
386 num_workers = 2
387 num_blocks = 4
388 with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions:
387 - t0 = regions[0].create_next_view(PAGE_SIZE)
388 - t1 = regions[1].create_next_view(PAGE_SIZE)
389 + t0 = regions[0].create_next_worker_view(PAGE_SIZE)
390 + t1 = regions[1].create_next_worker_view(PAGE_SIZE)
391
392 t0[:, :] = 11
393 t1[:, :] = 22
@@ -712,14 +714,14 @@ def test_cleanup_idempotent(iid):
714 r.cleanup() # must be a no-op
715
716
715 -def test_cleanup_after_create_next_view_releases_mmap(iid):
716 - """cleanup() must close the mmap even after create_next_view was called.
717 - create_next_view returns a view that shares storage with _base; both must be
717 +def test_cleanup_after_create_next_worker_view_releases_mmap(iid):
718 + """cleanup() must close the mmap even after create_next_worker_view was called.
719 + create_next_worker_view returns a view that shares storage with _base; both must be
720 released before mmap.close() can succeed."""
721 r = _make_region(iid)
722 mmap_obj = r.mmap_obj
723
722 - t = r.create_next_view(PAGE_SIZE)
724 + t = r.create_next_worker_view(PAGE_SIZE)
725 del t
726
727 r.cleanup()
tests/v1/kv_offload/test_file_mapper.py
+17
@@ -52,6 +52,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
52 is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False),
53 ),
54 replicated_layout=kwargs.get("replicated_layout", False),
55 + canonical_layout=kwargs.get("canonical_layout", False),
56 )
57 spec = MagicMock(spec=OffloadingSpec)
58 spec.config = config
@@ -208,6 +209,22 @@ def test_parallel_agnostic_separates_persistent_layouts():
209 assert specific.fields["parallel_agnostic"] is False
210
211
212 +def test_canonical_layout_changes_storage_namespace():
213 + # Canonical bytes are not interchangeable with the direct layout, so the
214 + # format id must fork the storage namespace.
215 + from vllm.v1.attention.backends.utils import set_kv_cache_layout
216 +
217 + set_kv_cache_layout("NHD")
218 + try:
219 + direct = make_mapper_from_offloading_spec()
220 + canonical = make_mapper_from_offloading_spec(canonical_layout=True)
221 + finally:
222 + set_kv_cache_layout(None)
223 + assert "canonical_format" not in direct.fields
224 + assert canonical.fields["canonical_format"] == "v1-nhd"
225 + assert direct.base_path != canonical.base_path
226 +
227 +
228 # ---------------------------------------------------------------------------
229 # replicated_layout: OR'd into parallel-agnostic identity for compact rows
230 # ---------------------------------------------------------------------------
vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py
+12
@@ -27,6 +27,18 @@ from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun
27 if TYPE_CHECKING:
28 from vllm.config import VllmConfig
29
30 +# Version of the canonical byte format; bump on any layout change
31 +CANONICAL_FORMAT_VERSION = 1
32 +
33 +
34 +def canonical_format_id() -> str:
35 + """Identity of the canonical byte format, for namespacing persisted KV.
36 + Canonical pages keep the worker's KV layout family, so the id couples the
37 + format version with that family; consumers must match it exactly."""
38 + from vllm.v1.attention.backends.utils import get_kv_cache_layout
39 +
40 + return f"v{CANONICAL_FORMAT_VERSION}-{get_kv_cache_layout().lower()}"
41 +
42
43 @dataclass(frozen=True)
44 class _RankContext:
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+41 -3
@@ -135,15 +135,52 @@ def build_offloading_config(
135 and parallel_config.nnodes_within_dp == 1
136 )
137
138 - # Only a single non-MLA full-attention group is parallelism-invariant:
139 - # MLA latent KV is replicated per rank (never head-sharded), and the V2
140 - # model runner's KV layout is not known to be parallelism-invariant.
138 + canonical_layout = bool(extra_config.get("canonical_layout", False))
139 +
140 + # Only a single non-MLA full-attention group with genuinely head-sharded
141 + # pages is parallelism-invariant: replicated latent or GQA heads,
142 + # per-token-head scales, CP token sharding, and the V2 model runner's
143 + # layout are all excluded.
144 is_parallelism_agnostic = (
145 not vllm_config.use_v2_model_runner
146 and single_group_spec is not None
147 and isinstance(single_group_spec, FullAttentionSpec)
148 and not isinstance(single_group_spec, MLAAttentionSpec)
149 + and single_group_spec.num_kv_heads * parallel_config.tensor_parallel_size
150 + == vllm_config.model_config.get_total_num_kv_heads()
151 + and not single_group_spec.kv_quant_mode.is_per_token_head
152 + and parallel_config.decode_context_parallel_size == 1
153 + and parallel_config.prefill_context_parallel_size == 1
154 )
155 + # Canonical pages are topology-free by construction, so the canonical
156 + # layout widens the gate to every config whose mappings derive portable:
157 + # exactly-sharded or replicated GQA heads (writer rotation) and the
158 + # TP-replicated MLA latent (one canonical copy for all replicas). The
159 + # model-runner version is irrelevant here — certification happens per
160 + # layer against live tensor strides at registration, and create_worker
161 + # fails closed against this flag if any layer cannot be certified.
162 + if canonical_layout and not is_parallelism_agnostic:
163 + tp_size = parallel_config.tensor_parallel_size
164 + if isinstance(single_group_spec, FullAttentionSpec):
165 + if type(single_group_spec) is MLAAttentionSpec:
166 + spec_certifiable = (
167 + single_group_spec.compress_ratio == 1
168 + and single_group_spec.real_page_size_bytes
169 + % single_group_spec.block_size
170 + == 0
171 + )
172 + else:
173 + total_kv_heads = vllm_config.model_config.get_total_num_kv_heads()
174 + spec_certifiable = (
175 + total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0
176 + )
177 + is_parallelism_agnostic = (
178 + spec_certifiable
179 + and not single_group_spec.kv_quant_mode.is_per_token_head
180 + and parallel_config.decode_context_parallel_size == 1
181 + and parallel_config.prefill_context_parallel_size == 1
182 + and parallel_config.world_size == tp_size
183 + )
184
185 kv_events_config = vllm_config.kv_events_config
186 cache_dtype = (
@@ -179,4 +216,5 @@ def build_offloading_config(
216 is_parallelism_agnostic=is_parallelism_agnostic,
217 ),
218 replicated_layout=replicated_layout,
219 + canonical_layout=canonical_layout,
220 )
vllm/v1/kv_offload/config.py
+8 -2
@@ -48,8 +48,11 @@ class OffloadingParallelConfig:
48 dcp_size: int
49 # Data parallel replica index of this engine.
50 data_parallel_index: int
51 - # True when concatenating a block's data across all workers yields
52 - # the same result regardless of the parallelism configuration.
51 + # True when the bytes that will be persisted for a block are portable
52 + # across parallelism configurations: for the direct layout, concatenating
53 + # the block's data across all workers in rank order yields the same bytes
54 + # under any topology; for the canonical layout, the canonical page itself
55 + # is topology-free.
56 is_parallelism_agnostic: bool
57
58
@@ -74,3 +77,6 @@ class OffloadingConfig:
77 # support it. Aggregate layout decision; per-layer replication metadata
78 # is planned for CanonicalKVCacheRef (#48408).
79 replicated_layout: bool = False
80 + # True when the canonical per-layer host byte layout was requested via
81 + # kv_connector_extra_config; certified per-layer at worker registration.
82 + canonical_layout: bool = False
vllm/v1/kv_offload/cpu/gpu_worker.py
+328 -49
@@ -3,7 +3,9 @@
3 import functools
4 import time
5 from collections import deque
6 +from collections.abc import Sequence
7 from dataclasses import dataclass
8 +from typing import NamedTuple
9
10 import numpy as np
11 import torch
@@ -18,6 +20,7 @@ from vllm.v1.kv_offload.base import (
20 BlockIDsLoadStoreSpec,
21 CanonicalKVCacheRef,
22 CanonicalKVCaches,
23 + CanonicalPageMapping,
24 GPULoadStoreSpec,
25 LoadStoreSpec,
26 OffloadingWorker,
@@ -33,7 +36,7 @@ logger = init_logger(__name__)
36
37
38 def _select_swap_blocks_fn(
36 - kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
39 + layer_refs_per_group: list[list[CanonicalKVCacheRef]],
40 gpu_to_cpu: bool,
41 ):
42 """Resolve the swap_blocks function for a handler at init time."""
@@ -46,7 +49,7 @@ def _select_swap_blocks_fn(
49 # so the Triton kernel's tl.load(cpu_ptr) is invalid on XPU).
50 if not HAS_TRITON or current_platform.is_xpu() or current_platform.is_rocm():
51 return ops.swap_blocks_batch
49 - page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g]
52 + page_sizes = [r.page_size_bytes for g in layer_refs_per_group for r in g]
53 # Triton wins only on small, 8-byte-aligned payloads.
54 if (
55 not page_sizes
@@ -120,6 +123,73 @@ def compute_sub_block_ptrs(
123 output[:] = flat[skip_count : skip_count + num_sub_blocks]
124
125
126 +class CopyPlan(NamedTuple):
127 + """Precomputed fragment-copy template for one data ref under the canonical
128 + CPU layout, unrolled from the ref's mapped runs. Offsets are relative to
129 + the per-block base pointers on each side."""
130 +
131 + frag_offsets_src: np.ndarray
132 + frag_offsets_dst: np.ndarray
133 + frag_sizes: np.ndarray
134 + total_bytes: int
135 +
136 + @property
137 + def num_frags(self) -> int:
138 + return len(self.frag_sizes)
139 +
140 +
141 +def _build_copy_plan(ref: CanonicalKVCacheRef, gpu_to_cpu: bool) -> CopyPlan:
142 + """Unroll one data ref's mapped runs into a per-fragment CopyPlan."""
143 + mapping = ref.mapping
144 + assert mapping is not None
145 + local: list[int] = []
146 + canonical: list[int] = []
147 + sizes: list[int] = []
148 + for run in mapping.runs:
149 + for i in range(run.num_fragments):
150 + local.append(run.local_offset + i * run.local_stride)
151 + canonical.append(run.canonical_offset + i * run.canonical_stride)
152 + sizes.append(run.fragment_size)
153 + src, dst = (local, canonical) if gpu_to_cpu else (canonical, local)
154 + return CopyPlan(
155 + frag_offsets_src=np.asarray(src, dtype=np.uint64),
156 + frag_offsets_dst=np.asarray(dst, dtype=np.uint64),
157 + frag_sizes=np.asarray(sizes, dtype=np.int64),
158 + total_bytes=sum(sizes),
159 + )
160 +
161 +
162 +def _canonical_page_ids(
163 + block_ids: np.ndarray, blocks_per_chunk: int, count: int, skip_count: int
164 +) -> np.ndarray:
165 + """Global canonical page ids matching compute_sub_block_ptrs' enumeration.
166 + These identify canonical pages consistently across ranks, so they key
167 + CanonicalPageMapping.is_writer rotation."""
168 + if blocks_per_chunk == 1:
169 + return block_ids[:count]
170 + flat = (
171 + block_ids[:, np.newaxis] * blocks_per_chunk + np.arange(blocks_per_chunk)
172 + ).ravel()
173 + return flat[skip_count : skip_count + count]
174 +
175 +
176 +def _canonical_block_sizes(
177 + layer_refs_per_group: list[list[CanonicalKVCacheRef]], num_tensors: int
178 +) -> list[int]:
179 + """Canonical CPU bytes per GPU block for each tensor, taken from the refs'
180 + mappings. Requires every ref to carry a mapping."""
181 + canonical_bytes_per_block = [0] * num_tensors
182 + for layer_refs in layer_refs_per_group:
183 + for ref in layer_refs:
184 + assert ref.mapping is not None
185 + canonical_bytes_per_block[ref.tensor_idx] = max(
186 + canonical_bytes_per_block[ref.tensor_idx],
187 + ref.mapping.canonical_page_size_bytes,
188 + )
189 + assert all(size > 0 for size in canonical_bytes_per_block)
190 + return canonical_bytes_per_block
191 +
192 +
193 def pin_mmap_region(region: SharedOffloadRegion) -> None:
194 """Register the entire mmap as CUDA pinned memory via cudaHostRegister."""
195 if not current_platform.is_cuda_alike():
@@ -176,9 +246,10 @@ class SingleDirectionOffloadingHandler:
246 gpu_tensors: list[torch.Tensor],
247 cpu_tensors: list[torch.Tensor],
248 blocks_per_chunk: int,
179 - kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]],
249 + layer_refs_per_group: list[list[CanonicalKVCacheRef]],
250 gpu_to_cpu: bool,
251 mmap_region: SharedOffloadRegion | None = None,
252 + canonical_layout: bool = False,
253 ):
254 """
255 Initialize a SingleDirectionOffloadingHandler.
@@ -189,14 +260,22 @@ class SingleDirectionOffloadingHandler:
260 cpu_tensors: list of CPU KV cache tensors.
261 Each of shape (num_cpu_blocks, cpu_page_size_bytes) with dtype int8.
262 Order should match gpu_tensors.
192 - kv_cache_groups_data_refs: list of CanonicalKVCacheRef per group.
263 + layer_refs_per_group: list of CanonicalKVCacheRef per group.
264 gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU.
265 + canonical_layout: if True, CPU pages use the canonical layout
266 + described by the refs' mappings.
267 """
268 assert len(gpu_tensors) == len(cpu_tensors)
269 assert len(gpu_tensors) > 0
270
271 + canonical_bytes_per_block = (
272 + _canonical_block_sizes(layer_refs_per_group, len(gpu_tensors))
273 + if canonical_layout
274 + else None
275 + )
276 +
277 # assert input tensors are as expected
199 - for gpu_tensor, cpu_tensor in zip(gpu_tensors, cpu_tensors):
278 + for t_idx, (gpu_tensor, cpu_tensor) in enumerate(zip(gpu_tensors, cpu_tensors)):
279 assert gpu_tensor.dtype == torch.int8
280 assert gpu_tensor.ndim == 2
281 assert gpu_tensor.is_cuda or gpu_tensor.is_xpu
@@ -205,7 +284,12 @@ class SingleDirectionOffloadingHandler:
284 assert cpu_tensor.device.type == "cpu"
285 _, gpu_page_size = gpu_tensor.shape
286 _, cpu_page_size = cpu_tensor.shape
208 - assert cpu_page_size == gpu_page_size * blocks_per_chunk
287 + if canonical_bytes_per_block is not None:
288 + assert (
289 + cpu_page_size == canonical_bytes_per_block[t_idx] * blocks_per_chunk
290 + )
291 + else:
292 + assert cpu_page_size == gpu_page_size * blocks_per_chunk
293
294 self.src_tensors: list[torch.Tensor] = (
295 gpu_tensors if gpu_to_cpu else cpu_tensors
@@ -214,9 +298,9 @@ class SingleDirectionOffloadingHandler:
298 cpu_tensors if gpu_to_cpu else gpu_tensors
299 )
300 self.gpu_to_cpu: bool = gpu_to_cpu
217 - self.kv_cache_groups_data_refs = kv_cache_groups_data_refs
301 + self.layer_refs_per_group = layer_refs_per_group
302 self._swap_blocks_batch = _select_swap_blocks_fn(
219 - kv_cache_groups_data_refs, gpu_to_cpu
303 + layer_refs_per_group, gpu_to_cpu
304 )
305
306 # GPU blocks may be smaller
@@ -224,6 +308,24 @@ class SingleDirectionOffloadingHandler:
308 self.src_blocks_per_chunk = 1 if self.gpu_to_cpu else blocks_per_chunk
309 self.dst_blocks_per_chunk = blocks_per_chunk if self.gpu_to_cpu else 1
310
311 + # Per (group, ref) static copy plans for the canonical layout
312 + self._canonical_copy_plans: list[list[CopyPlan]] | None = (
313 + [
314 + [_build_copy_plan(ref, gpu_to_cpu) for ref in layer_refs]
315 + for layer_refs in layer_refs_per_group
316 + ]
317 + if canonical_layout
318 + else None
319 + )
320 + self._fill_group_ops = (
321 + self._fill_canonical_ops if canonical_layout else self._fill_direct_ops
322 + )
323 + # Reusable per-block base-pointer scratch for the canonical fill,
324 + # sized to the largest possible group (grown on demand)
325 + num_scratch_blocks = gpu_tensors[0].shape[0] if canonical_layout else 0
326 + self._scratch_bases_src = np.empty(num_scratch_blocks, dtype=np.uint64)
327 + self._scratch_bases_dst = np.empty(num_scratch_blocks, dtype=np.uint64)
328 +
329 # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)
330 self._mmap_region = mmap_region
331 # job_id -> event
@@ -237,6 +339,179 @@ class SingleDirectionOffloadingHandler:
339 # list of pinned descriptor buffer sets available for re-use
340 self._buffer_pool: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
341
342 + def _estimate_max_copy_ops(self, group_sizes: Sequence[int]) -> int:
343 + """Upper bound on the number of copy descriptors for a transfer.
344 +
345 + Exact for the direct layout. The canonical path may fill fewer:
346 + writer rotation later drops the blocks this rank does not write."""
347 + num_copy_ops = 0
348 + for g_idx, (group_size, layer_refs) in enumerate(
349 + zip(group_sizes, self.layer_refs_per_group)
350 + ):
351 + if self._canonical_copy_plans is None:
352 + num_copy_ops += group_size * len(layer_refs)
353 + else:
354 + num_copy_ops += group_size * sum(
355 + plan.num_frags for plan in self._canonical_copy_plans[g_idx]
356 + )
357 + return num_copy_ops
358 +
359 + def _fill_direct_ops(
360 + self,
361 + g_idx: int,
362 + group_src: np.ndarray,
363 + group_dst: np.ndarray,
364 + group_size: int,
365 + src_skip_count: int,
366 + dst_skip_count: int,
367 + all_src: np.ndarray,
368 + all_dst: np.ndarray,
369 + all_sizes: np.ndarray,
370 + op_idx: int,
371 + ) -> tuple[int, int]:
372 + """Fill one group's copy descriptors for the direct (worker-private)
373 + layout: one whole-page copy per (block, ref).
374 +
375 + Returns (op_idx past the filled descriptors, bytes added)."""
376 + num_bytes = 0
377 + for data_ref in self.layer_refs_per_group[g_idx]:
378 + t_idx = data_ref.tensor_idx
379 + end_idx = op_idx + group_size
380 +
381 + compute_sub_block_ptrs(
382 + group_src,
383 + self.src_blocks_per_chunk,
384 + all_src[op_idx:end_idx],
385 + self.src_tensors[t_idx],
386 + skip_count=src_skip_count,
387 + )
388 + compute_sub_block_ptrs(
389 + group_dst,
390 + self.dst_blocks_per_chunk,
391 + all_dst[op_idx:end_idx],
392 + self.dst_tensors[t_idx],
393 + skip_count=dst_skip_count,
394 + )
395 +
396 + all_sizes[op_idx:end_idx] = data_ref.page_size_bytes
397 + num_bytes += group_size * data_ref.page_size_bytes
398 + op_idx = end_idx
399 + return op_idx, num_bytes
400 +
401 + def _fill_canonical_ops(
402 + self,
403 + g_idx: int,
404 + group_src: np.ndarray,
405 + group_dst: np.ndarray,
406 + group_size: int,
407 + src_skip_count: int,
408 + dst_skip_count: int,
409 + all_src: np.ndarray,
410 + all_dst: np.ndarray,
411 + all_sizes: np.ndarray,
412 + op_idx: int,
413 + ) -> tuple[int, int]:
414 + """Fill one group's copy descriptors for the canonical layout:
415 + scatter each block through the ref's precomputed CopyPlan, keeping
416 + only the blocks this rank writes.
417 +
418 + Returns (op_idx past the filled descriptors, bytes added)."""
419 + assert self._canonical_copy_plans is not None
420 + # Zero-copy reinterpretation for pointer arithmetic: uint64 and the
421 + # buffers' int64 are bit-equivalent for addresses
422 + all_src_u64 = all_src.view(np.uint64)
423 + all_dst_u64 = all_dst.view(np.uint64)
424 + if group_size > len(self._scratch_bases_src):
425 + self._scratch_bases_src = np.empty(group_size, dtype=np.uint64)
426 + self._scratch_bases_dst = np.empty(group_size, dtype=np.uint64)
427 +
428 + num_bytes = 0
429 + for plan, data_ref in zip(
430 + self._canonical_copy_plans[g_idx], self.layer_refs_per_group[g_idx]
431 + ):
432 + if plan.num_frags == 0:
433 + continue
434 + t_idx = data_ref.tensor_idx
435 +
436 + # 1. Base byte pointer of every block on each side
437 + block_bases_src = self._scratch_bases_src[:group_size]
438 + block_bases_dst = self._scratch_bases_dst[:group_size]
439 + compute_sub_block_ptrs(
440 + group_src,
441 + self.src_blocks_per_chunk,
442 + block_bases_src,
443 + self.src_tensors[t_idx],
444 + skip_count=src_skip_count,
445 + )
446 + compute_sub_block_ptrs(
447 + group_dst,
448 + self.dst_blocks_per_chunk,
449 + block_bases_dst,
450 + self.dst_tensors[t_idx],
451 + skip_count=dst_skip_count,
452 + )
453 +
454 + # 2. On store, keep only the blocks this rank is elected to write
455 + mapping = data_ref.mapping
456 + assert mapping is not None
457 + if self.gpu_to_cpu and mapping.num_writers > 1:
458 + block_bases_src, block_bases_dst = self._filter_writer_blocks(
459 + block_bases_src,
460 + block_bases_dst,
461 + mapping,
462 + group_dst,
463 + group_size,
464 + dst_skip_count,
465 + )
466 + num_active_blocks = len(block_bases_src)
467 +
468 + # 3. Expand (block base + fragment offset) into one descriptor
469 + # per (block, fragment), writing straight into the descriptor
470 + # buffers: reshaping a contiguous 1D slice is a view, so the
471 + # broadcasts below allocate nothing
472 + end_idx = op_idx + num_active_blocks * plan.num_frags
473 + np.add(
474 + block_bases_src[:, None],
475 + plan.frag_offsets_src[None, :],
476 + out=all_src_u64[op_idx:end_idx].reshape(
477 + num_active_blocks, plan.num_frags
478 + ),
479 + )
480 + np.add(
481 + block_bases_dst[:, None],
482 + plan.frag_offsets_dst[None, :],
483 + out=all_dst_u64[op_idx:end_idx].reshape(
484 + num_active_blocks, plan.num_frags
485 + ),
486 + )
487 + all_sizes[op_idx:end_idx].reshape(num_active_blocks, plan.num_frags)[:] = (
488 + plan.frag_sizes
489 + )
490 + num_bytes += num_active_blocks * plan.total_bytes
491 + op_idx = end_idx
492 + return op_idx, num_bytes
493 +
494 + def _filter_writer_blocks(
495 + self,
496 + block_bases_src: np.ndarray,
497 + block_bases_dst: np.ndarray,
498 + mapping: CanonicalPageMapping,
499 + group_dst: np.ndarray,
500 + group_size: int,
501 + dst_skip_count: int,
502 + ) -> tuple[np.ndarray, np.ndarray]:
503 + """Keep only the blocks this rank writes: replicated ranks take turns
504 + writing shared canonical pages, keyed by the rank-consistent CPU-side
505 + canonical page id."""
506 + cpu_page_ids = _canonical_page_ids(
507 + group_dst,
508 + self.dst_blocks_per_chunk,
509 + group_size,
510 + dst_skip_count,
511 + )
512 + writer_mask = cpu_page_ids % mapping.num_writers == mapping.writer_index
513 + return block_bases_src[writer_mask], block_bases_dst[writer_mask]
514 +
515 def transfer_async(
516 self, job_id: int, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec
517 ) -> bool:
@@ -274,17 +549,13 @@ class SingleDirectionOffloadingHandler:
549 gpu_spec = src_spec if self.gpu_to_cpu else dst_spec
550 assert isinstance(gpu_spec, GPULoadStoreSpec)
551 group_sizes = gpu_spec.group_sizes
277 - assert len(group_sizes) == len(self.kv_cache_groups_data_refs)
552 + assert len(group_sizes) == len(self.layer_refs_per_group)
553
554 # extract block indices from the GPU spec
555 block_indices = gpu_spec.block_indices
281 - assert len(block_indices) == len(self.kv_cache_groups_data_refs)
556 + assert len(block_indices) == len(self.layer_refs_per_group)
557
283 - num_copy_ops = 0
284 - for group_size, group_data_refs in zip(
285 - group_sizes, self.kv_cache_groups_data_refs
286 - ):
287 - num_copy_ops += group_size * len(group_data_refs)
558 + num_copy_ops = self._estimate_max_copy_ops(group_sizes)
559
560 # reuse a pooled buffer set, growing it if this transfer needs more room
561 batch_src, batch_dst, batch_sizes = (
@@ -307,8 +578,8 @@ class SingleDirectionOffloadingHandler:
578 op_idx = 0
579 # count total number of bytes copied
580 num_transfer_bytes = 0
310 - for group_size, block_idx, group_data_refs in zip(
311 - group_sizes, block_indices, self.kv_cache_groups_data_refs
581 + for g_idx, (group_size, block_idx) in enumerate(
582 + zip(group_sizes, block_indices)
583 ):
584 if group_size == 0:
585 continue
@@ -326,38 +597,31 @@ class SingleDirectionOffloadingHandler:
597 src_end_offset = src_offset + src_blocks_count
598 assert src_end_offset <= num_src_blocks
599
329 - group_src = src_blocks[src_offset:src_end_offset]
330 - group_dst = dst_blocks[dst_offset:dst_end_offset]
331 -
332 - for data_ref in group_data_refs:
333 - t_idx = data_ref.tensor_idx
334 - end_idx = op_idx + group_size
335 -
336 - compute_sub_block_ptrs(
337 - group_src,
338 - self.src_blocks_per_chunk,
339 - all_src[op_idx:end_idx],
340 - self.src_tensors[t_idx],
341 - skip_count=src_logical_blocks_to_skip,
342 - )
343 - compute_sub_block_ptrs(
344 - group_dst,
345 - self.dst_blocks_per_chunk,
346 - all_dst[op_idx:end_idx],
347 - self.dst_tensors[t_idx],
348 - skip_count=dst_logical_blocks_to_skip,
349 - )
350 -
351 - all_sizes[op_idx:end_idx] = data_ref.page_size_bytes
352 - num_transfer_bytes += group_size * data_ref.page_size_bytes
353 - op_idx = end_idx
600 + op_idx, group_bytes = self._fill_group_ops(
601 + g_idx,
602 + group_src=src_blocks[src_offset:src_end_offset],
603 + group_dst=dst_blocks[dst_offset:dst_end_offset],
604 + group_size=group_size,
605 + src_skip_count=src_logical_blocks_to_skip,
606 + dst_skip_count=dst_logical_blocks_to_skip,
607 + all_src=all_src,
608 + all_dst=all_dst,
609 + all_sizes=all_sizes,
610 + op_idx=op_idx,
611 + )
612 + num_transfer_bytes += group_bytes
613
614 src_offset = src_end_offset
615 dst_offset = dst_end_offset
616
617 assert src_offset == num_src_blocks
618 assert dst_offset == num_dst_blocks
360 - assert op_idx == num_copy_ops
619 + # Writer rotation may skip non-writer blocks, leaving op_idx below
620 + # the sized upper bound
621 + assert op_idx <= num_copy_ops
622 + src = src[:op_idx]
623 + dst = dst[:op_idx]
624 + sizes = sizes[:op_idx]
625
626 stream = (
627 self._stream_pool.pop() if self._stream_pool else current_platform.Stream()
@@ -390,7 +654,7 @@ class SingleDirectionOffloadingHandler:
654 is_src_access_order_any = not self.gpu_to_cpu
655 with current_platform.stream(stream):
656 start_event.record(stream)
393 - if num_copy_ops > 0:
657 + if op_idx > 0:
658 self._swap_blocks_batch(
659 src,
660 dst,
@@ -475,23 +739,36 @@ class CPUOffloadingWorker(OffloadingWorker):
739 blocks_per_chunk: int,
740 num_cpu_blocks: int,
741 mmap_region: SharedOffloadRegion | None = None,
742 + canonical_layout: bool = False,
743 ):
744 + assert not canonical_layout or mmap_region is not None
745 pin_memory = PIN_MEMORY
746 logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
747 if mmap_region is not None and pin_memory:
748 pin_mmap_region(mmap_region)
749
750 + canonical_bytes_per_block = (
751 + _canonical_block_sizes(kv_caches.group_data_refs, len(kv_caches.tensors))
752 + if canonical_layout
753 + else None
754 + )
755 +
756 gpu_tensors: list[torch.Tensor] = []
757 cpu_tensors: list[torch.Tensor] = []
486 - for kv_cache_tensor in kv_caches.tensors:
758 + for t_idx, kv_cache_tensor in enumerate(kv_caches.tensors):
759 gpu_page_size_bytes = kv_cache_tensor.page_size_bytes
760 gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view(
761 (-1, gpu_page_size_bytes)
762 )
763 cpu_page_size_bytes = gpu_page_size_bytes * blocks_per_chunk
764
493 - if mmap_region is not None:
494 - cpu_tensor = mmap_region.create_next_view(cpu_page_size_bytes)
765 + if canonical_bytes_per_block is not None:
766 + assert mmap_region is not None
767 + cpu_tensor = mmap_region.create_next_canonical_view(
768 + canonical_bytes_per_block[t_idx] * blocks_per_chunk
769 + )
770 + elif mmap_region is not None:
771 + cpu_tensor = mmap_region.create_next_worker_view(cpu_page_size_bytes)
772 else:
773 t0 = time.monotonic()
774 cpu_tensor = torch.zeros(
@@ -515,17 +792,19 @@ class CPUOffloadingWorker(OffloadingWorker):
792 gpu_tensors=gpu_tensors,
793 cpu_tensors=cpu_tensors,
794 blocks_per_chunk=blocks_per_chunk,
518 - kv_cache_groups_data_refs=kv_caches.group_data_refs,
795 + layer_refs_per_group=kv_caches.group_data_refs,
796 gpu_to_cpu=True,
797 mmap_region=mmap_region,
798 + canonical_layout=canonical_layout,
799 )
800
801 self._load_handler = SingleDirectionOffloadingHandler(
802 gpu_tensors=gpu_tensors,
803 cpu_tensors=cpu_tensors,
804 blocks_per_chunk=blocks_per_chunk,
527 - kv_cache_groups_data_refs=kv_caches.group_data_refs,
805 + layer_refs_per_group=kv_caches.group_data_refs,
806 gpu_to_cpu=False,
807 + canonical_layout=canonical_layout,
808 )
809
810 def submit_store(
vllm/v1/kv_offload/cpu/shared_offload_region.py
+46 -1
@@ -167,9 +167,10 @@ class SharedOffloadRegion:
167
168 self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8)
169 self._views: list[torch.Tensor] = []
170 + self._canonical_offset = 0
171 self.is_pinned: bool = False
172
172 - def create_next_view(self, tensor_page_size: int) -> torch.Tensor:
173 + def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor:
174 """Allocate a strided int8 view for this worker, one canonical tensor.
175
176 Must be called once per canonical tensor. The full mmap layout is:
@@ -208,6 +209,50 @@ class SharedOffloadRegion:
209 self._views.append(worker_layer_view)
210 return worker_layer_view
211
212 + def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor:
213 + """Allocate a strided int8 view shared by all workers for one
214 + canonical tensor (canonical layout).
215 +
216 + Must be called once per canonical tensor, instead of
217 + create_next_worker_view. The full mmap layout is:
218 +
219 + |<-------- canonical area ------->|<-------- unused ------->|
220 + | all workers share this area | |
221 + | | |
222 + | [ canonical_t0 | canonical_t1 ] | |
223 + | [ canonical_t0 | canonical_t1 ] | |
224 + | [ canonical_t0 | canonical_t1 ] | |
225 + ^ ^
226 + _canonical_offset=0, then advances by each tensor's size
227 +
228 + Each canonical_t{i} cell is that tensor's canonical page for the
229 + block. Canonical areas are carved consecutively from the start of
230 + each block row; consecutive rows are separated by row_stride. Every
231 + worker gets the identical byte ranges and writes only its disjoint
232 + bytes within them, as described by its canonical mappings — unlike
233 + create_next_worker_view, which gives each worker a private
234 + cpu_page_size slot per row.
235 +
236 + The trailing unused bytes exist only when the canonical pages sum to
237 + less than row_stride: page-alignment padding of the row, or
238 + deduplication of KV replicated across workers (e.g. the MLA latent),
239 + where one canonical copy replaces world_size worker copies.
240 +
241 + Args:
242 + tensor_page_size: Canonical bytes per block for this tensor.
243 + """
244 + new_offset = self._canonical_offset + tensor_page_size
245 + assert new_offset <= self._row_stride
246 + view = torch.as_strided(
247 + self._base,
248 + size=(self.num_blocks, tensor_page_size),
249 + stride=(self._row_stride, 1),
250 + storage_offset=self._canonical_offset,
251 + )
252 + self._canonical_offset = new_offset
253 + self._views.append(view)
254 + return view
255 +
256 def create_kv_memoryview(self) -> memoryview:
257 """Return a zero-copy memoryview over the entire KV buffer.
258
vllm/v1/kv_offload/file_mapper.py
+11
@@ -4,6 +4,9 @@
4 import hashlib
5 import json
6
7 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import (
8 + canonical_format_id,
9 +)
10 from vllm.v1.kv_offload.base import (
11 OffloadingSpec,
12 OffloadKey,
@@ -36,6 +39,7 @@ class FileMapper:
39 inference_engine: str = "vllm",
40 parallel_agnostic: bool = False,
41 replicated_layout: bool = False,
42 + canonical_format: str | None = None,
43 ):
44 """
45 Initialize the file mapper. Each worker constructs its own, but
@@ -65,6 +69,11 @@ class FileMapper:
69 # unchanged (False is the historical default and must not appear).
70 if replicated_layout:
71 self.fields["replicated_layout"] = True
72 + # The canonical byte format is not interchangeable with the direct
73 + # layout (or with other canonical format versions/families), so its
74 + # identity participates in the storage namespace.
75 + if canonical_format is not None:
76 + self.fields["canonical_format"] = canonical_format
77 self.base_path: str = self._compute_base_path(root_dir, self.fields)
78
79 @classmethod
@@ -85,6 +94,7 @@ class FileMapper:
94 for group in config.groups
95 ]
96 parallel = config.parallel
97 + canonical_format = canonical_format_id() if config.canonical_layout else None
98 return cls(
99 root_dir=root_dir,
100 model_name=config.model.name,
@@ -102,6 +112,7 @@ class FileMapper:
112 and (parallel.is_parallelism_agnostic or config.replicated_layout)
113 ),
114 replicated_layout=(parallel_agnostic and config.replicated_layout),
115 + canonical_format=canonical_format,
116 )
117
118 def get_file_name(self, key: OffloadKey) -> str:
vllm/v1/kv_offload/tiering/spec.py
+46
@@ -264,6 +264,11 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
264 # Scheduler-side mmap (rank=None); kept for cleanup
265 self._scheduler_mmap: SharedOffloadRegion | None = None
266
267 + # Set by create_worker when canonical_layout is enabled: True when
268 + # every layer's canonical bytes are parallelism-agnostic (portable),
269 + # False when some layers use the opaque fallback (exact-topology only)
270 + self.all_layers_portable: bool | None = None
271 +
272 # engine_id is unique per DP replica (suffixed with _dp{rank} in both
273 # the Ray and multiprocessing paths), so it names a per-replica offload
274 # region.
@@ -394,12 +399,53 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
399 cpu_page_size=self.cpu_page_size_per_worker,
400 )
401 try:
402 + if self.config.canonical_layout:
403 + self._validate_canonical_refs(kv_caches)
404 return CPUOffloadingWorker(
405 kv_caches=kv_caches,
406 blocks_per_chunk=self.blocks_per_chunk,
407 num_cpu_blocks=self.num_blocks,
408 mmap_region=worker_mmap,
409 + canonical_layout=self.config.canonical_layout,
410 )
411 except Exception:
412 worker_mmap.cleanup()
413 raise
414 +
415 + def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None:
416 + """Require a mapping on every ref and record layer portability.
417 +
418 + Fails loudly rather than persist direct-layout bytes under a
419 + canonical format identity."""
420 + all_refs = [
421 + ref for group_refs in kv_caches.group_data_refs for ref in group_refs
422 + ]
423 + if any(ref.mapping is None for ref in all_refs):
424 + raise RuntimeError(
425 + "canonical_layout was requested but the KV cache layout "
426 + "could not be certified for canonical offload (offload "
427 + "workers must be exactly the TP group, and packed / "
428 + "cross-layer KV layouts are not supported). Remove "
429 + "canonical_layout from kv_connector_extra_config."
430 + )
431 + self.all_layers_portable = all(
432 + ref.mapping is not None and ref.mapping.parallelism_agnostic
433 + for ref in all_refs
434 + )
435 + if self.config.parallel.is_parallelism_agnostic and not (
436 + self.all_layers_portable
437 + ):
438 + # The scheduler-side storage namespace was already collapsed on
439 + # the static portability claim; opaque per-topology bytes must
440 + # not land in it.
441 + raise RuntimeError(
442 + "canonical_layout could not certify every layer as "
443 + "parallelism-agnostic, but the storage namespace is shared "
444 + "across topologies. Remove canonical_layout from "
445 + "kv_connector_extra_config or disable parallel-agnostic "
446 + "secondary tiers."
447 + )
448 + logger.info(
449 + "Canonical KV layout enabled (all_layers_portable=%s)",
450 + self.all_layers_portable,
451 + )