From 81840a172f2886f4c59dafd9dce9917bb789280c Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:44:54 +0300 Subject: [PATCH] [KV Connector] Canonical CPU layout for parallelism-agnostic KV offload (#48414) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis --- .../unit/offloading_connector/test_config.py | 58 +++ .../kv_offload/cpu/test_canonical_layout.py | 274 +++++++++++++ tests/v1/kv_offload/cpu/test_gpu_worker.py | 4 +- .../cpu/test_shared_offload_region.py | 94 ++--- tests/v1/kv_offload/test_file_mapper.py | 17 + .../v1/offloading/canonical_mapping.py | 12 + .../kv_connector/v1/offloading/config.py | 44 +- vllm/v1/kv_offload/config.py | 10 +- vllm/v1/kv_offload/cpu/gpu_worker.py | 377 +++++++++++++++--- .../kv_offload/cpu/shared_offload_region.py | 47 ++- vllm/v1/kv_offload/file_mapper.py | 11 + vllm/v1/kv_offload/tiering/spec.py | 46 +++ 12 files changed, 891 insertions(+), 103 deletions(-) create mode 100644 tests/v1/kv_offload/cpu/test_canonical_layout.py diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 5a66b463e3..cc58aaf7a7 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -47,6 +47,9 @@ def _make_vllm_config( config.cache_config.cache_dtype = torch.float16 config.model_config.model = "test-model" config.model_config.use_mla = False + # _full_attention_spec's heads at tp=1: the parallelism-agnostic gate + # requires the head shard to cover the model's KV heads exactly + config.model_config.get_total_num_kv_heads.return_value = 4 world_size = ( tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size ) @@ -603,6 +606,61 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): assert not _parallelism_agnostic(kv_cache_groups) +def test_canonical_layout_widens_parallelism_agnostic_to_mla(): + """The canonical layout dedups the TP-replicated MLA latent into one + portable copy, so the gate admits MLA — but only when canonical_layout + is requested.""" + mla_groups = [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))] + assert not _parallelism_agnostic(mla_groups) + + config = _make_vllm_config(extra_config={"canonical_layout": True}) + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=mla_groups, + ) + offloading_config = build_offloading_config(config, kv_cache_config) + assert offloading_config.parallel.is_parallelism_agnostic + assert offloading_config.canonical_layout + + # hybrid groupings stay out: their non-full-attention layers can only + # derive opaque (exact-topology) mappings + hybrid_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["l0"], _full_attention_spec()), + KVCacheGroupSpec(["l1"], _full_attention_spec()), + ], + ) + assert not build_offloading_config( + config, hybrid_config + ).parallel.is_parallelism_agnostic + + +def test_canonical_layout_certifies_v2_model_runner(): + """Canonical bytes are certified per layer against live tensor strides at + registration, so the static gate must not depend on the model-runner + version — the v2 runner is the case the canonical layout exists for.""" + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())], + ) + + config = _make_vllm_config() + config.use_v2_model_runner = True + assert not build_offloading_config( + config, kv_cache_config + ).parallel.is_parallelism_agnostic + + config = _make_vllm_config(extra_config={"canonical_layout": True}) + config.use_v2_model_runner = True + assert build_offloading_config( + config, kv_cache_config + ).parallel.is_parallelism_agnostic + + def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() config.use_v2_model_runner = True diff --git a/tests/v1/kv_offload/cpu/test_canonical_layout.py b/tests/v1/kv_offload/cpu/test_canonical_layout.py new file mode 100644 index 0000000000..f10e534cc7 --- /dev/null +++ b/tests/v1/kv_offload/cpu/test_canonical_layout.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time +import uuid + +import numpy as np +import pytest +import torch + +from vllm.v1.kv_offload.base import ( + CanonicalKVCacheRef, + CanonicalPageMapping, + CopyRun, + GPULoadStoreSpec, +) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec +from vllm.v1.kv_offload.cpu.gpu_worker import ( + SingleDirectionOffloadingHandler, + _build_copy_plan, + _canonical_block_sizes, + _canonical_page_ids, + pin_mmap_region, +) +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion + + +def _ref(mapping: CanonicalPageMapping, tensor_idx: int = 0) -> CanonicalKVCacheRef: + return CanonicalKVCacheRef( + tensor_idx=tensor_idx, + page_size_bytes=mapping.local_page_size_bytes, + mapping=mapping, + ) + + +def _nhd_mapping() -> CanonicalPageMapping: + # 4-token NHD page at tp=4, rank 2: K and V runs of 4x128B fragments + runs = ( + CopyRun(0, 256, 128, 4, 128, 512), + CopyRun(512, 2304, 128, 4, 128, 512), + ) + return CanonicalPageMapping(4096, 1024, runs, 1, 0, True) + + +def test_copy_plan_unrolls_runs(): + plan = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=True) + k_dst = [256, 768, 1280, 1792] + assert plan.frag_offsets_src.tolist() == [0, 128, 256, 384, 512, 640, 768, 896] + assert plan.frag_offsets_dst.tolist() == k_dst + [2048 + o for o in k_dst] + assert plan.frag_sizes.tolist() == [128] * 8 + assert plan.num_frags == 8 + + +def test_load_direction_swaps_offsets(): + store = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=True) + load = _build_copy_plan(_ref(_nhd_mapping()), gpu_to_cpu=False) + assert np.array_equal(store.frag_offsets_src, load.frag_offsets_dst) + assert np.array_equal(store.frag_offsets_dst, load.frag_offsets_src) + assert np.array_equal(store.frag_sizes, load.frag_sizes) + + +def test_writer_rotation_matches_is_writer(): + # Replicas take turns writing shared canonical pages, keyed by the + # CPU-side canonical page id; the enumeration must agree with is_writer + identity = CopyRun(0, 0, 256, 1, 256, 256) + mapping = CanonicalPageMapping(256, 256, (identity,), 2, 1, True) + ids = _canonical_page_ids( + np.array([3, 7, 9]), blocks_per_chunk=4, count=10, skip_count=2 + ) + assert ids.tolist() == [14, 15, 28, 29, 30, 31, 36, 37, 38, 39] + mask = ids % mapping.num_writers == mapping.writer_index + assert mask.tolist() == [mapping.is_writer(int(b)) for b in ids] + + +def test_canonical_block_sizes_take_max_per_tensor(): + identity = CopyRun(0, 0, 512, 1, 512, 512) + small = CanonicalPageMapping(2048, 512, (identity,), 1, 0, False) + refs = [[_ref(_nhd_mapping(), 0), _ref(small, 0)], [_ref(small, 1)]] + assert _canonical_block_sizes(refs, 2) == [4096, 2048] + + +def _tp2_rank_mapping(rank: int) -> CanonicalPageMapping: + # 4-token NHD page, 4 total heads of 64B, tp=2: rank holds 2 heads, + # so K and V each scatter as 4 per-token 128B fragments + runs = ( + CopyRun(0, rank * 128, 128, 4, 128, 256), + CopyRun(512, 1024 + rank * 128, 128, 4, 128, 256), + ) + return CanonicalPageMapping(2048, 1024, runs, 1, 0, True) + + +def _whole_page_mapping() -> CanonicalPageMapping: + identity = CopyRun(0, 0, 2048, 1, 2048, 2048) + return CanonicalPageMapping(2048, 2048, (identity,), 1, 0, True) + + +def _transfer(handler, num_blocks: int, gpu_to_cpu: bool) -> None: + block_ids = list(range(num_blocks)) + gpu_spec = GPULoadStoreSpec( + block_ids, group_sizes=(num_blocks,), block_indices=(0,) + ) + cpu_spec = CPULoadStoreSpec(block_ids) + src, dst = (gpu_spec, cpu_spec) if gpu_to_cpu else (cpu_spec, gpu_spec) + assert handler.transfer_async(0, src, dst) + deadline = time.time() + 30 + while time.time() < deadline: + if handler.get_finished(): + return + time.sleep(0.001) + raise TimeoutError("transfer did not complete") + + +def _canonical_handler(gpu_tensor, cpu_tensor, mapping, gpu_to_cpu): + page = mapping.local_page_size_bytes + refs = [[CanonicalKVCacheRef(tensor_idx=0, page_size_bytes=page, mapping=mapping)]] + return SingleDirectionOffloadingHandler( + gpu_tensors=[gpu_tensor], + cpu_tensors=[cpu_tensor], + blocks_per_chunk=1, + layer_refs_per_group=refs, + gpu_to_cpu=gpu_to_cpu, + canonical_layout=True, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_gpu_roundtrip_assembles_canonical_page_across_ranks(): + """Two TP2-style rank handlers must scatter into one shared canonical + CPU page such that each rank reloads bit-exact and a whole-page (TP1) + reader sees both shards — the cross-topology contract.""" + torch.manual_seed(0) + num_blocks = 4 + gpu_rank = [ + torch.randint(-128, 128, (num_blocks, 1024), dtype=torch.int8, device="cuda") + for _ in range(2) + ] + cpu_canonical = torch.zeros(num_blocks, 2048, dtype=torch.int8, pin_memory=True) + + for rank in (0, 1): + store = _canonical_handler( + gpu_rank[rank], cpu_canonical, _tp2_rank_mapping(rank), gpu_to_cpu=True + ) + _transfer(store, num_blocks, gpu_to_cpu=True) + torch.accelerator.synchronize() + + # independent oracle: replay each rank's runs in numpy + expected = np.zeros((num_blocks, 2048), dtype=np.int8) + for rank in (0, 1): + local = gpu_rank[rank].cpu().numpy() + for run in _tp2_rank_mapping(rank).runs: + for i in range(run.num_fragments): + lo = run.local_offset + i * run.local_stride + co = run.canonical_offset + i * run.canonical_stride + expected[:, co : co + run.fragment_size] = local[ + :, lo : lo + run.fragment_size + ] + assert np.array_equal(cpu_canonical.numpy(), expected) + + # each rank reloads its shard bit-exact + gpu_back = torch.zeros(num_blocks, 1024, dtype=torch.int8, device="cuda") + load = _canonical_handler( + gpu_back, cpu_canonical, _tp2_rank_mapping(0), gpu_to_cpu=False + ) + _transfer(load, num_blocks, gpu_to_cpu=False) + torch.accelerator.synchronize() + assert torch.equal(gpu_back, gpu_rank[0]) + + # a whole-page reader (TP1 topology) sees the assembled page + gpu_full = torch.zeros(num_blocks, 2048, dtype=torch.int8, device="cuda") + load_full = _canonical_handler( + gpu_full, cpu_canonical, _whole_page_mapping(), gpu_to_cpu=False + ) + _transfer(load_full, num_blocks, gpu_to_cpu=False) + torch.accelerator.synchronize() + assert torch.equal(gpu_full.cpu(), cpu_canonical) + + +# 4-token NHD page, 4 total heads of 64 bytes: canonical page holds +# [K: token x head][V: token x head] = 2048 bytes per block +_TOTAL_HEADS = 4 +_HEAD_BYTES = 64 +_BLOCK_TOKENS = 4 +_CANONICAL_PAGE = 2 * _BLOCK_TOKENS * _TOTAL_HEADS * _HEAD_BYTES + + +def _nhd_shard_mapping(tp: int, rank: int) -> CanonicalPageMapping: + """Rank's head-shard mapping into the canonical NHD page at the given tp.""" + local_heads = _TOTAL_HEADS // tp + frag = local_heads * _HEAD_BYTES + canonical_row = _TOTAL_HEADS * _HEAD_BYTES + k_run = CopyRun(0, rank * frag, frag, _BLOCK_TOKENS, frag, canonical_row) + v_run = CopyRun( + _BLOCK_TOKENS * frag, + _BLOCK_TOKENS * canonical_row + rank * frag, + frag, + _BLOCK_TOKENS, + frag, + canonical_row, + ) + return CanonicalPageMapping( + _CANONICAL_PAGE, 2 * _BLOCK_TOKENS * frag, (k_run, v_run), 1, 0, True + ) + + +def _head_shard(full_kv: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """This rank's local page rows out of the (blocks, 2, tokens, heads, + head_bytes) ground truth.""" + local_heads = _TOTAL_HEADS // tp + shard = full_kv[:, :, :, rank * local_heads : (rank + 1) * local_heads, :] + return shard.reshape(full_kv.shape[0], -1).contiguous() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("writer_tp,reader_tp", [(2, 4), (4, 2), (2, 1), (4, 4)]) +def test_cross_topology_roundtrip(writer_tp: int, reader_tp: int): + """KV written at one tp must be readable at another: writer ranks scatter + head shards into a shared canonical region, reader ranks gather their own + shards, and every reader must see the writers' ground-truth bytes.""" + torch.manual_seed(0) + num_blocks = 3 + full_kv = torch.randint( + -128, + 128, + (num_blocks, 2, _BLOCK_TOKENS, _TOTAL_HEADS, _HEAD_BYTES), + dtype=torch.int8, + ) + + engine_id = str(uuid.uuid4()) + row_stride = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + assert row_stride >= _CANONICAL_PAGE + regions: list[SharedOffloadRegion] = [] + + def canonical_view(rank: int, world_size: int) -> torch.Tensor: + region = SharedOffloadRegion( + engine_id=engine_id, + num_blocks=num_blocks, + rank=rank, + kv_bytes_per_block=row_stride, + cpu_page_size=row_stride // world_size, + ) + regions.append(region) + # The Triton load path dereferences CPU pointers on the GPU, which is + # only legal on pinned memory; production pins via CPUOffloadingWorker + pin_mmap_region(region) + return region.create_next_canonical_view(_CANONICAL_PAGE) + + try: + for rank in range(writer_tp): + store = _canonical_handler( + _head_shard(full_kv, writer_tp, rank).cuda(), + canonical_view(rank, writer_tp), + _nhd_shard_mapping(writer_tp, rank), + gpu_to_cpu=True, + ) + _transfer(store, num_blocks, gpu_to_cpu=True) + torch.accelerator.synchronize() + + for rank in range(reader_tp): + expected = _head_shard(full_kv, reader_tp, rank) + gpu_out = torch.zeros_like(expected, device="cuda") + load = _canonical_handler( + gpu_out, + canonical_view(rank, reader_tp), + _nhd_shard_mapping(reader_tp, rank), + gpu_to_cpu=False, + ) + _transfer(load, num_blocks, gpu_to_cpu=False) + torch.accelerator.synchronize() + assert torch.equal(gpu_out.cpu(), expected), ( + f"reader tp={reader_tp} rank={rank} bytes diverge from the " + f"tp={writer_tp} writers' ground truth" + ) + finally: + for region in regions: + region.cleanup() diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index bfbb18251c..694c33712e 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -195,7 +195,7 @@ def test_transfer( assert finished[0].success assert finished[0].transfer_size == ( len(gpu_blocks) - * sum([x.page_size_bytes for x in handler.kv_cache_groups_data_refs[0]]) + * sum([x.page_size_bytes for x in handler.layer_refs_per_group[0]]) ) assert finished[0].transfer_time > 0 assert finished[0].transfer_time < (time.time() - start_time) @@ -409,7 +409,7 @@ def test_transfer_multi_group( expected_bytes = sum( group_size * sum([x.page_size_bytes for x in data_refs]) for group_size, data_refs in zip( - group_sizes, handler.kv_cache_groups_data_refs + group_sizes, handler.layer_refs_per_group ) ) assert finished[0].transfer_size == expected_bytes diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index 7b5867ddb7..1cdb589ef1 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -173,7 +173,7 @@ def _mp_race_construct_and_write( kv_bytes_per_block=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, ) - t = region.create_next_view(cpu_page_size) + t = region.create_next_worker_view(cpu_page_size) t[:, :] = fill_value done_queue.put({"rank": rank, "error": None}) cleanup_queue.get() # wait for parent's verification to finish @@ -190,110 +190,112 @@ def iid(): # --------------------------------------------------------------------------- -# create_next_view — shape, stride and storage offset +# create_next_worker_view — shape, stride and storage offset # --------------------------------------------------------------------------- -def test_create_next_view_shape_and_stride(iid): +def test_create_next_worker_view_shape_and_stride(iid): """Returned tensor must have shape (num_blocks, tensor_page_size) and stride (row_stride, 1) where row_stride = cpu_page_size * num_workers.""" with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r: - t = r.create_next_view(PAGE_SIZE) + t = r.create_next_worker_view(PAGE_SIZE) assert t.shape == (4, PAGE_SIZE) # num_workers=1 → row_stride = cpu_page_size assert t.stride() == (2 * PAGE_SIZE, 1) del t -def test_create_next_view_storage_offset_rank0(iid): +def test_create_next_worker_view_storage_offset_rank0(iid): """rank=0 worker's first tensor must start at byte 0 of the mmap.""" with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r: - t = r.create_next_view(PAGE_SIZE) + t = r.create_next_worker_view(PAGE_SIZE) assert t.data_ptr() == r._base.data_ptr() # storage_offset == 0 del t -def test_create_next_view_storage_offset_rank1(iid): +def test_create_next_worker_view_storage_offset_rank1(iid): """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap.""" with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1): - t1 = r1.create_next_view(PAGE_SIZE) + t1 = r1.create_next_worker_view(PAGE_SIZE) assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE del t1 -def test_create_next_view_row_stride_with_multiple_workers(iid): +def test_create_next_worker_view_row_stride_with_multiple_workers(iid): """With num_workers=4, row_stride must be 4 * cpu_page_size.""" with _region(iid, num_blocks=2, num_workers=4) as r: - t = r.create_next_view(PAGE_SIZE) + t = r.create_next_worker_view(PAGE_SIZE) assert t.stride(0) == 4 * PAGE_SIZE del t # --------------------------------------------------------------------------- -# create_next_view — cursor advancement +# create_next_worker_view — cursor advancement # --------------------------------------------------------------------------- -def test_create_next_view_cursor_advances(iid): - """Each call to create_next_view must advance _worker_offset by tensor_page_size.""" +def test_create_next_worker_view_cursor_advances(iid): + """Each create_next_worker_view call must advance _worker_offset by + tensor_page_size.""" with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r: assert r._worker_offset == 0 - r.create_next_view(PAGE_SIZE) + r.create_next_worker_view(PAGE_SIZE) assert r._worker_offset == PAGE_SIZE - r.create_next_view(PAGE_SIZE) + r.create_next_worker_view(PAGE_SIZE) assert r._worker_offset == 2 * PAGE_SIZE - r.create_next_view(PAGE_SIZE) + r.create_next_worker_view(PAGE_SIZE) assert r._worker_offset == 3 * PAGE_SIZE # exactly at area end -def test_create_next_view_exact_fill_succeeds(iid): +def test_create_next_worker_view_exact_fill_succeeds(iid): """Allocations whose total exactly equals cpu_page_size must all succeed.""" with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r: - r.create_next_view(PAGE_SIZE) # first half - r.create_next_view(PAGE_SIZE) # fills to area end — must not raise + r.create_next_worker_view(PAGE_SIZE) # first half + r.create_next_worker_view(PAGE_SIZE) # fills to area end — must not raise # --------------------------------------------------------------------------- -# create_next_view — overflow guard +# create_next_worker_view — overflow guard # --------------------------------------------------------------------------- -def test_create_next_view_single_overflow_raises(iid): +def test_create_next_worker_view_single_overflow_raises(iid): """A single allocation larger than cpu_page_size must raise AssertionError.""" with ( _region(iid) as r, pytest.raises(AssertionError, match="exceeds worker area end"), ): - r.create_next_view(PAGE_SIZE + 1) + r.create_next_worker_view(PAGE_SIZE + 1) -def test_create_next_view_cumulative_overflow_raises(iid): +def test_create_next_worker_view_cumulative_overflow_raises(iid): """Successive allocations that cumulatively exceed cpu_page_size must raise.""" with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r: - r.create_next_view(PAGE_SIZE) # ok — half used - r.create_next_view(PAGE_SIZE) # ok — full + r.create_next_worker_view(PAGE_SIZE) # ok — half used + r.create_next_worker_view(PAGE_SIZE) # ok — full with pytest.raises(AssertionError, match="exceeds worker area end"): - r.create_next_view(1) # one byte too many + r.create_next_worker_view(1) # one byte too many -def test_create_next_view_overflow_does_not_mutate_cursor(iid): - """A failed create_next_view must leave _worker_offset unchanged.""" +def test_create_next_worker_view_overflow_does_not_mutate_cursor(iid): + """A failed create_next_worker_view must leave _worker_offset unchanged.""" with _region(iid) as r: offset_before = r._worker_offset with pytest.raises(AssertionError): - r.create_next_view(PAGE_SIZE + 1) + r.create_next_worker_view(PAGE_SIZE + 1) assert r._worker_offset == offset_before # --------------------------------------------------------------------------- -# create_next_view — data correctness and layout +# create_next_worker_view — data correctness and layout # --------------------------------------------------------------------------- -def test_create_next_view_write_visible_in_raw_mmap(iid): - """Writes into a create_next_view view must appear at the correct raw mmap offset""" +def test_create_next_worker_view_write_visible_in_raw_mmap(iid): + """Writes into a create_next_worker_view view must appear at the correct + raw mmap offset""" with _region(iid, num_blocks=4) as r: - t = r.create_next_view(PAGE_SIZE) + t = r.create_next_worker_view(PAGE_SIZE) t[2, :] = 42 # write to block row 2 raw = memoryview(r.mmap_obj) @@ -303,11 +305,11 @@ def test_create_next_view_write_visible_in_raw_mmap(iid): del raw, t -def test_create_next_view_multi_tensor_layout(iid): +def test_create_next_worker_view_multi_tensor_layout(iid): """Two tensors from the same worker land at consecutive byte offsets per row.""" with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r: - ta = r.create_next_view(PAGE_SIZE) - tb = r.create_next_view(PAGE_SIZE) + ta = r.create_next_worker_view(PAGE_SIZE) + tb = r.create_next_worker_view(PAGE_SIZE) ta[:, :] = 1 tb[:, :] = 2 @@ -322,8 +324,8 @@ def test_create_next_view_multi_tensor_layout(iid): del raw, ta, tb -def test_create_next_view_multiprocess_slots(iid): - """Each worker process calls create_next_view and writes distinct data; +def test_create_next_worker_view_multiprocess_slots(iid): + """Each worker process calls create_next_worker_view and writes distinct data; the parent verifies each slot lands at the correct interleaved offset.""" num_workers = 2 num_blocks = 4 @@ -356,7 +358,7 @@ def test_create_next_view_multiprocess_slots(iid): ) child.start() - t0 = region.create_next_view(PAGE_SIZE) + t0 = region.create_next_worker_view(PAGE_SIZE) t0[:, :] = 11 result = done_queue.get(timeout=30) @@ -379,13 +381,13 @@ def test_create_next_view_multiprocess_slots(iid): _cleanup_file(region.mmap_path) -def test_create_next_view_worker_isolation(iid): +def test_create_next_worker_view_worker_isolation(iid): """Writes by worker 0 must not affect worker 1's slot and vice versa.""" num_workers = 2 num_blocks = 4 with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions: - t0 = regions[0].create_next_view(PAGE_SIZE) - t1 = regions[1].create_next_view(PAGE_SIZE) + t0 = regions[0].create_next_worker_view(PAGE_SIZE) + t1 = regions[1].create_next_worker_view(PAGE_SIZE) t0[:, :] = 11 t1[:, :] = 22 @@ -712,14 +714,14 @@ def test_cleanup_idempotent(iid): r.cleanup() # must be a no-op -def test_cleanup_after_create_next_view_releases_mmap(iid): - """cleanup() must close the mmap even after create_next_view was called. - create_next_view returns a view that shares storage with _base; both must be +def test_cleanup_after_create_next_worker_view_releases_mmap(iid): + """cleanup() must close the mmap even after create_next_worker_view was called. + create_next_worker_view returns a view that shares storage with _base; both must be released before mmap.close() can succeed.""" r = _make_region(iid) mmap_obj = r.mmap_obj - t = r.create_next_view(PAGE_SIZE) + t = r.create_next_worker_view(PAGE_SIZE) del t r.cleanup() diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 027f01133b..0c8416b0c0 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -52,6 +52,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False), ), replicated_layout=kwargs.get("replicated_layout", False), + canonical_layout=kwargs.get("canonical_layout", False), ) spec = MagicMock(spec=OffloadingSpec) spec.config = config @@ -208,6 +209,22 @@ def test_parallel_agnostic_separates_persistent_layouts(): assert specific.fields["parallel_agnostic"] is False +def test_canonical_layout_changes_storage_namespace(): + # Canonical bytes are not interchangeable with the direct layout, so the + # format id must fork the storage namespace. + from vllm.v1.attention.backends.utils import set_kv_cache_layout + + set_kv_cache_layout("NHD") + try: + direct = make_mapper_from_offloading_spec() + canonical = make_mapper_from_offloading_spec(canonical_layout=True) + finally: + set_kv_cache_layout(None) + assert "canonical_format" not in direct.fields + assert canonical.fields["canonical_format"] == "v1-nhd" + assert direct.base_path != canonical.base_path + + # --------------------------------------------------------------------------- # replicated_layout: OR'd into parallel-agnostic identity for compact rows # --------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py index 581f7fdbc8..209616bcc9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py @@ -27,6 +27,18 @@ from vllm.v1.kv_offload.base import CanonicalPageMapping, CopyRun if TYPE_CHECKING: from vllm.config import VllmConfig +# Version of the canonical byte format; bump on any layout change +CANONICAL_FORMAT_VERSION = 1 + + +def canonical_format_id() -> str: + """Identity of the canonical byte format, for namespacing persisted KV. + Canonical pages keep the worker's KV layout family, so the id couples the + format version with that family; consumers must match it exactly.""" + from vllm.v1.attention.backends.utils import get_kv_cache_layout + + return f"v{CANONICAL_FORMAT_VERSION}-{get_kv_cache_layout().lower()}" + @dataclass(frozen=True) class _RankContext: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index b9837bcb4b..2bf140f1e3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -135,15 +135,52 @@ def build_offloading_config( and parallel_config.nnodes_within_dp == 1 ) - # Only a single non-MLA full-attention group is parallelism-invariant: - # MLA latent KV is replicated per rank (never head-sharded), and the V2 - # model runner's KV layout is not known to be parallelism-invariant. + canonical_layout = bool(extra_config.get("canonical_layout", False)) + + # Only a single non-MLA full-attention group with genuinely head-sharded + # pages is parallelism-invariant: replicated latent or GQA heads, + # per-token-head scales, CP token sharding, and the V2 model runner's + # layout are all excluded. is_parallelism_agnostic = ( not vllm_config.use_v2_model_runner and single_group_spec is not None and isinstance(single_group_spec, FullAttentionSpec) and not isinstance(single_group_spec, MLAAttentionSpec) + and single_group_spec.num_kv_heads * parallel_config.tensor_parallel_size + == vllm_config.model_config.get_total_num_kv_heads() + and not single_group_spec.kv_quant_mode.is_per_token_head + and parallel_config.decode_context_parallel_size == 1 + and parallel_config.prefill_context_parallel_size == 1 ) + # Canonical pages are topology-free by construction, so the canonical + # layout widens the gate to every config whose mappings derive portable: + # exactly-sharded or replicated GQA heads (writer rotation) and the + # TP-replicated MLA latent (one canonical copy for all replicas). The + # model-runner version is irrelevant here — certification happens per + # layer against live tensor strides at registration, and create_worker + # fails closed against this flag if any layer cannot be certified. + if canonical_layout and not is_parallelism_agnostic: + tp_size = parallel_config.tensor_parallel_size + if isinstance(single_group_spec, FullAttentionSpec): + if type(single_group_spec) is MLAAttentionSpec: + spec_certifiable = ( + single_group_spec.compress_ratio == 1 + and single_group_spec.real_page_size_bytes + % single_group_spec.block_size + == 0 + ) + else: + total_kv_heads = vllm_config.model_config.get_total_num_kv_heads() + spec_certifiable = ( + total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0 + ) + is_parallelism_agnostic = ( + spec_certifiable + and not single_group_spec.kv_quant_mode.is_per_token_head + and parallel_config.decode_context_parallel_size == 1 + and parallel_config.prefill_context_parallel_size == 1 + and parallel_config.world_size == tp_size + ) kv_events_config = vllm_config.kv_events_config cache_dtype = ( @@ -179,4 +216,5 @@ def build_offloading_config( is_parallelism_agnostic=is_parallelism_agnostic, ), replicated_layout=replicated_layout, + canonical_layout=canonical_layout, ) diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index 46ca489cc2..fdbce491e0 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -48,8 +48,11 @@ class OffloadingParallelConfig: dcp_size: int # Data parallel replica index of this engine. data_parallel_index: int - # True when concatenating a block's data across all workers yields - # the same result regardless of the parallelism configuration. + # True when the bytes that will be persisted for a block are portable + # across parallelism configurations: for the direct layout, concatenating + # the block's data across all workers in rank order yields the same bytes + # under any topology; for the canonical layout, the canonical page itself + # is topology-free. is_parallelism_agnostic: bool @@ -74,3 +77,6 @@ class OffloadingConfig: # support it. Aggregate layout decision; per-layer replication metadata # is planned for CanonicalKVCacheRef (#48408). replicated_layout: bool = False + # True when the canonical per-layer host byte layout was requested via + # kv_connector_extra_config; certified per-layer at worker registration. + canonical_layout: bool = False diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 40b57ac867..c7d75b636c 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -3,7 +3,9 @@ import functools import time from collections import deque +from collections.abc import Sequence from dataclasses import dataclass +from typing import NamedTuple import numpy as np import torch @@ -18,6 +20,7 @@ from vllm.v1.kv_offload.base import ( BlockIDsLoadStoreSpec, CanonicalKVCacheRef, CanonicalKVCaches, + CanonicalPageMapping, GPULoadStoreSpec, LoadStoreSpec, OffloadingWorker, @@ -33,7 +36,7 @@ logger = init_logger(__name__) def _select_swap_blocks_fn( - kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], + layer_refs_per_group: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, ): """Resolve the swap_blocks function for a handler at init time.""" @@ -46,7 +49,7 @@ def _select_swap_blocks_fn( # so the Triton kernel's tl.load(cpu_ptr) is invalid on XPU). if not HAS_TRITON or current_platform.is_xpu() or current_platform.is_rocm(): return ops.swap_blocks_batch - page_sizes = [r.page_size_bytes for g in kv_cache_groups_data_refs for r in g] + page_sizes = [r.page_size_bytes for g in layer_refs_per_group for r in g] # Triton wins only on small, 8-byte-aligned payloads. if ( not page_sizes @@ -120,6 +123,73 @@ def compute_sub_block_ptrs( output[:] = flat[skip_count : skip_count + num_sub_blocks] +class CopyPlan(NamedTuple): + """Precomputed fragment-copy template for one data ref under the canonical + CPU layout, unrolled from the ref's mapped runs. Offsets are relative to + the per-block base pointers on each side.""" + + frag_offsets_src: np.ndarray + frag_offsets_dst: np.ndarray + frag_sizes: np.ndarray + total_bytes: int + + @property + def num_frags(self) -> int: + return len(self.frag_sizes) + + +def _build_copy_plan(ref: CanonicalKVCacheRef, gpu_to_cpu: bool) -> CopyPlan: + """Unroll one data ref's mapped runs into a per-fragment CopyPlan.""" + mapping = ref.mapping + assert mapping is not None + local: list[int] = [] + canonical: list[int] = [] + sizes: list[int] = [] + for run in mapping.runs: + for i in range(run.num_fragments): + local.append(run.local_offset + i * run.local_stride) + canonical.append(run.canonical_offset + i * run.canonical_stride) + sizes.append(run.fragment_size) + src, dst = (local, canonical) if gpu_to_cpu else (canonical, local) + return CopyPlan( + frag_offsets_src=np.asarray(src, dtype=np.uint64), + frag_offsets_dst=np.asarray(dst, dtype=np.uint64), + frag_sizes=np.asarray(sizes, dtype=np.int64), + total_bytes=sum(sizes), + ) + + +def _canonical_page_ids( + block_ids: np.ndarray, blocks_per_chunk: int, count: int, skip_count: int +) -> np.ndarray: + """Global canonical page ids matching compute_sub_block_ptrs' enumeration. + These identify canonical pages consistently across ranks, so they key + CanonicalPageMapping.is_writer rotation.""" + if blocks_per_chunk == 1: + return block_ids[:count] + flat = ( + block_ids[:, np.newaxis] * blocks_per_chunk + np.arange(blocks_per_chunk) + ).ravel() + return flat[skip_count : skip_count + count] + + +def _canonical_block_sizes( + layer_refs_per_group: list[list[CanonicalKVCacheRef]], num_tensors: int +) -> list[int]: + """Canonical CPU bytes per GPU block for each tensor, taken from the refs' + mappings. Requires every ref to carry a mapping.""" + canonical_bytes_per_block = [0] * num_tensors + for layer_refs in layer_refs_per_group: + for ref in layer_refs: + assert ref.mapping is not None + canonical_bytes_per_block[ref.tensor_idx] = max( + canonical_bytes_per_block[ref.tensor_idx], + ref.mapping.canonical_page_size_bytes, + ) + assert all(size > 0 for size in canonical_bytes_per_block) + return canonical_bytes_per_block + + def pin_mmap_region(region: SharedOffloadRegion) -> None: """Register the entire mmap as CUDA pinned memory via cudaHostRegister.""" if not current_platform.is_cuda_alike(): @@ -176,9 +246,10 @@ class SingleDirectionOffloadingHandler: gpu_tensors: list[torch.Tensor], cpu_tensors: list[torch.Tensor], blocks_per_chunk: int, - kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], + layer_refs_per_group: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, mmap_region: SharedOffloadRegion | None = None, + canonical_layout: bool = False, ): """ Initialize a SingleDirectionOffloadingHandler. @@ -189,14 +260,22 @@ class SingleDirectionOffloadingHandler: cpu_tensors: list of CPU KV cache tensors. Each of shape (num_cpu_blocks, cpu_page_size_bytes) with dtype int8. Order should match gpu_tensors. - kv_cache_groups_data_refs: list of CanonicalKVCacheRef per group. + layer_refs_per_group: list of CanonicalKVCacheRef per group. gpu_to_cpu: if True, transfer from GPU to CPU; otherwise CPU to GPU. + canonical_layout: if True, CPU pages use the canonical layout + described by the refs' mappings. """ assert len(gpu_tensors) == len(cpu_tensors) assert len(gpu_tensors) > 0 + canonical_bytes_per_block = ( + _canonical_block_sizes(layer_refs_per_group, len(gpu_tensors)) + if canonical_layout + else None + ) + # assert input tensors are as expected - for gpu_tensor, cpu_tensor in zip(gpu_tensors, cpu_tensors): + for t_idx, (gpu_tensor, cpu_tensor) in enumerate(zip(gpu_tensors, cpu_tensors)): assert gpu_tensor.dtype == torch.int8 assert gpu_tensor.ndim == 2 assert gpu_tensor.is_cuda or gpu_tensor.is_xpu @@ -205,7 +284,12 @@ class SingleDirectionOffloadingHandler: assert cpu_tensor.device.type == "cpu" _, gpu_page_size = gpu_tensor.shape _, cpu_page_size = cpu_tensor.shape - assert cpu_page_size == gpu_page_size * blocks_per_chunk + if canonical_bytes_per_block is not None: + assert ( + cpu_page_size == canonical_bytes_per_block[t_idx] * blocks_per_chunk + ) + else: + assert cpu_page_size == gpu_page_size * blocks_per_chunk self.src_tensors: list[torch.Tensor] = ( gpu_tensors if gpu_to_cpu else cpu_tensors @@ -214,9 +298,9 @@ class SingleDirectionOffloadingHandler: cpu_tensors if gpu_to_cpu else gpu_tensors ) self.gpu_to_cpu: bool = gpu_to_cpu - self.kv_cache_groups_data_refs = kv_cache_groups_data_refs + self.layer_refs_per_group = layer_refs_per_group self._swap_blocks_batch = _select_swap_blocks_fn( - kv_cache_groups_data_refs, gpu_to_cpu + layer_refs_per_group, gpu_to_cpu ) # GPU blocks may be smaller @@ -224,6 +308,24 @@ class SingleDirectionOffloadingHandler: self.src_blocks_per_chunk = 1 if self.gpu_to_cpu else blocks_per_chunk self.dst_blocks_per_chunk = blocks_per_chunk if self.gpu_to_cpu else 1 + # Per (group, ref) static copy plans for the canonical layout + self._canonical_copy_plans: list[list[CopyPlan]] | None = ( + [ + [_build_copy_plan(ref, gpu_to_cpu) for ref in layer_refs] + for layer_refs in layer_refs_per_group + ] + if canonical_layout + else None + ) + self._fill_group_ops = ( + self._fill_canonical_ops if canonical_layout else self._fill_direct_ops + ) + # Reusable per-block base-pointer scratch for the canonical fill, + # sized to the largest possible group (grown on demand) + num_scratch_blocks = gpu_tensors[0].shape[0] if canonical_layout else 0 + self._scratch_bases_src = np.empty(num_scratch_blocks, dtype=np.uint64) + self._scratch_bases_dst = np.empty(num_scratch_blocks, dtype=np.uint64) + # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it) self._mmap_region = mmap_region # job_id -> event @@ -237,6 +339,179 @@ class SingleDirectionOffloadingHandler: # list of pinned descriptor buffer sets available for re-use self._buffer_pool: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + def _estimate_max_copy_ops(self, group_sizes: Sequence[int]) -> int: + """Upper bound on the number of copy descriptors for a transfer. + + Exact for the direct layout. The canonical path may fill fewer: + writer rotation later drops the blocks this rank does not write.""" + num_copy_ops = 0 + for g_idx, (group_size, layer_refs) in enumerate( + zip(group_sizes, self.layer_refs_per_group) + ): + if self._canonical_copy_plans is None: + num_copy_ops += group_size * len(layer_refs) + else: + num_copy_ops += group_size * sum( + plan.num_frags for plan in self._canonical_copy_plans[g_idx] + ) + return num_copy_ops + + def _fill_direct_ops( + self, + g_idx: int, + group_src: np.ndarray, + group_dst: np.ndarray, + group_size: int, + src_skip_count: int, + dst_skip_count: int, + all_src: np.ndarray, + all_dst: np.ndarray, + all_sizes: np.ndarray, + op_idx: int, + ) -> tuple[int, int]: + """Fill one group's copy descriptors for the direct (worker-private) + layout: one whole-page copy per (block, ref). + + Returns (op_idx past the filled descriptors, bytes added).""" + num_bytes = 0 + for data_ref in self.layer_refs_per_group[g_idx]: + t_idx = data_ref.tensor_idx + end_idx = op_idx + group_size + + compute_sub_block_ptrs( + group_src, + self.src_blocks_per_chunk, + all_src[op_idx:end_idx], + self.src_tensors[t_idx], + skip_count=src_skip_count, + ) + compute_sub_block_ptrs( + group_dst, + self.dst_blocks_per_chunk, + all_dst[op_idx:end_idx], + self.dst_tensors[t_idx], + skip_count=dst_skip_count, + ) + + all_sizes[op_idx:end_idx] = data_ref.page_size_bytes + num_bytes += group_size * data_ref.page_size_bytes + op_idx = end_idx + return op_idx, num_bytes + + def _fill_canonical_ops( + self, + g_idx: int, + group_src: np.ndarray, + group_dst: np.ndarray, + group_size: int, + src_skip_count: int, + dst_skip_count: int, + all_src: np.ndarray, + all_dst: np.ndarray, + all_sizes: np.ndarray, + op_idx: int, + ) -> tuple[int, int]: + """Fill one group's copy descriptors for the canonical layout: + scatter each block through the ref's precomputed CopyPlan, keeping + only the blocks this rank writes. + + Returns (op_idx past the filled descriptors, bytes added).""" + assert self._canonical_copy_plans is not None + # Zero-copy reinterpretation for pointer arithmetic: uint64 and the + # buffers' int64 are bit-equivalent for addresses + all_src_u64 = all_src.view(np.uint64) + all_dst_u64 = all_dst.view(np.uint64) + if group_size > len(self._scratch_bases_src): + self._scratch_bases_src = np.empty(group_size, dtype=np.uint64) + self._scratch_bases_dst = np.empty(group_size, dtype=np.uint64) + + num_bytes = 0 + for plan, data_ref in zip( + self._canonical_copy_plans[g_idx], self.layer_refs_per_group[g_idx] + ): + if plan.num_frags == 0: + continue + t_idx = data_ref.tensor_idx + + # 1. Base byte pointer of every block on each side + block_bases_src = self._scratch_bases_src[:group_size] + block_bases_dst = self._scratch_bases_dst[:group_size] + compute_sub_block_ptrs( + group_src, + self.src_blocks_per_chunk, + block_bases_src, + self.src_tensors[t_idx], + skip_count=src_skip_count, + ) + compute_sub_block_ptrs( + group_dst, + self.dst_blocks_per_chunk, + block_bases_dst, + self.dst_tensors[t_idx], + skip_count=dst_skip_count, + ) + + # 2. On store, keep only the blocks this rank is elected to write + mapping = data_ref.mapping + assert mapping is not None + if self.gpu_to_cpu and mapping.num_writers > 1: + block_bases_src, block_bases_dst = self._filter_writer_blocks( + block_bases_src, + block_bases_dst, + mapping, + group_dst, + group_size, + dst_skip_count, + ) + num_active_blocks = len(block_bases_src) + + # 3. Expand (block base + fragment offset) into one descriptor + # per (block, fragment), writing straight into the descriptor + # buffers: reshaping a contiguous 1D slice is a view, so the + # broadcasts below allocate nothing + end_idx = op_idx + num_active_blocks * plan.num_frags + np.add( + block_bases_src[:, None], + plan.frag_offsets_src[None, :], + out=all_src_u64[op_idx:end_idx].reshape( + num_active_blocks, plan.num_frags + ), + ) + np.add( + block_bases_dst[:, None], + plan.frag_offsets_dst[None, :], + out=all_dst_u64[op_idx:end_idx].reshape( + num_active_blocks, plan.num_frags + ), + ) + all_sizes[op_idx:end_idx].reshape(num_active_blocks, plan.num_frags)[:] = ( + plan.frag_sizes + ) + num_bytes += num_active_blocks * plan.total_bytes + op_idx = end_idx + return op_idx, num_bytes + + def _filter_writer_blocks( + self, + block_bases_src: np.ndarray, + block_bases_dst: np.ndarray, + mapping: CanonicalPageMapping, + group_dst: np.ndarray, + group_size: int, + dst_skip_count: int, + ) -> tuple[np.ndarray, np.ndarray]: + """Keep only the blocks this rank writes: replicated ranks take turns + writing shared canonical pages, keyed by the rank-consistent CPU-side + canonical page id.""" + cpu_page_ids = _canonical_page_ids( + group_dst, + self.dst_blocks_per_chunk, + group_size, + dst_skip_count, + ) + writer_mask = cpu_page_ids % mapping.num_writers == mapping.writer_index + return block_bases_src[writer_mask], block_bases_dst[writer_mask] + def transfer_async( self, job_id: int, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec ) -> bool: @@ -274,17 +549,13 @@ class SingleDirectionOffloadingHandler: gpu_spec = src_spec if self.gpu_to_cpu else dst_spec assert isinstance(gpu_spec, GPULoadStoreSpec) group_sizes = gpu_spec.group_sizes - assert len(group_sizes) == len(self.kv_cache_groups_data_refs) + assert len(group_sizes) == len(self.layer_refs_per_group) # extract block indices from the GPU spec block_indices = gpu_spec.block_indices - assert len(block_indices) == len(self.kv_cache_groups_data_refs) + assert len(block_indices) == len(self.layer_refs_per_group) - num_copy_ops = 0 - for group_size, group_data_refs in zip( - group_sizes, self.kv_cache_groups_data_refs - ): - num_copy_ops += group_size * len(group_data_refs) + num_copy_ops = self._estimate_max_copy_ops(group_sizes) # reuse a pooled buffer set, growing it if this transfer needs more room batch_src, batch_dst, batch_sizes = ( @@ -307,8 +578,8 @@ class SingleDirectionOffloadingHandler: op_idx = 0 # count total number of bytes copied num_transfer_bytes = 0 - for group_size, block_idx, group_data_refs in zip( - group_sizes, block_indices, self.kv_cache_groups_data_refs + for g_idx, (group_size, block_idx) in enumerate( + zip(group_sizes, block_indices) ): if group_size == 0: continue @@ -326,38 +597,31 @@ class SingleDirectionOffloadingHandler: src_end_offset = src_offset + src_blocks_count assert src_end_offset <= num_src_blocks - group_src = src_blocks[src_offset:src_end_offset] - group_dst = dst_blocks[dst_offset:dst_end_offset] - - for data_ref in group_data_refs: - t_idx = data_ref.tensor_idx - end_idx = op_idx + group_size - - compute_sub_block_ptrs( - group_src, - self.src_blocks_per_chunk, - all_src[op_idx:end_idx], - self.src_tensors[t_idx], - skip_count=src_logical_blocks_to_skip, - ) - compute_sub_block_ptrs( - group_dst, - self.dst_blocks_per_chunk, - all_dst[op_idx:end_idx], - self.dst_tensors[t_idx], - skip_count=dst_logical_blocks_to_skip, - ) - - all_sizes[op_idx:end_idx] = data_ref.page_size_bytes - num_transfer_bytes += group_size * data_ref.page_size_bytes - op_idx = end_idx + op_idx, group_bytes = self._fill_group_ops( + g_idx, + group_src=src_blocks[src_offset:src_end_offset], + group_dst=dst_blocks[dst_offset:dst_end_offset], + group_size=group_size, + src_skip_count=src_logical_blocks_to_skip, + dst_skip_count=dst_logical_blocks_to_skip, + all_src=all_src, + all_dst=all_dst, + all_sizes=all_sizes, + op_idx=op_idx, + ) + num_transfer_bytes += group_bytes src_offset = src_end_offset dst_offset = dst_end_offset assert src_offset == num_src_blocks assert dst_offset == num_dst_blocks - assert op_idx == num_copy_ops + # Writer rotation may skip non-writer blocks, leaving op_idx below + # the sized upper bound + assert op_idx <= num_copy_ops + src = src[:op_idx] + dst = dst[:op_idx] + sizes = sizes[:op_idx] stream = ( self._stream_pool.pop() if self._stream_pool else current_platform.Stream() @@ -390,7 +654,7 @@ class SingleDirectionOffloadingHandler: is_src_access_order_any = not self.gpu_to_cpu with current_platform.stream(stream): start_event.record(stream) - if num_copy_ops > 0: + if op_idx > 0: self._swap_blocks_batch( src, dst, @@ -475,23 +739,36 @@ class CPUOffloadingWorker(OffloadingWorker): blocks_per_chunk: int, num_cpu_blocks: int, mmap_region: SharedOffloadRegion | None = None, + canonical_layout: bool = False, ): + assert not canonical_layout or mmap_region is not None pin_memory = PIN_MEMORY logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors)) if mmap_region is not None and pin_memory: pin_mmap_region(mmap_region) + canonical_bytes_per_block = ( + _canonical_block_sizes(kv_caches.group_data_refs, len(kv_caches.tensors)) + if canonical_layout + else None + ) + gpu_tensors: list[torch.Tensor] = [] cpu_tensors: list[torch.Tensor] = [] - for kv_cache_tensor in kv_caches.tensors: + for t_idx, kv_cache_tensor in enumerate(kv_caches.tensors): gpu_page_size_bytes = kv_cache_tensor.page_size_bytes gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view( (-1, gpu_page_size_bytes) ) cpu_page_size_bytes = gpu_page_size_bytes * blocks_per_chunk - if mmap_region is not None: - cpu_tensor = mmap_region.create_next_view(cpu_page_size_bytes) + if canonical_bytes_per_block is not None: + assert mmap_region is not None + cpu_tensor = mmap_region.create_next_canonical_view( + canonical_bytes_per_block[t_idx] * blocks_per_chunk + ) + elif mmap_region is not None: + cpu_tensor = mmap_region.create_next_worker_view(cpu_page_size_bytes) else: t0 = time.monotonic() cpu_tensor = torch.zeros( @@ -515,17 +792,19 @@ class CPUOffloadingWorker(OffloadingWorker): gpu_tensors=gpu_tensors, cpu_tensors=cpu_tensors, blocks_per_chunk=blocks_per_chunk, - kv_cache_groups_data_refs=kv_caches.group_data_refs, + layer_refs_per_group=kv_caches.group_data_refs, gpu_to_cpu=True, mmap_region=mmap_region, + canonical_layout=canonical_layout, ) self._load_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, cpu_tensors=cpu_tensors, blocks_per_chunk=blocks_per_chunk, - kv_cache_groups_data_refs=kv_caches.group_data_refs, + layer_refs_per_group=kv_caches.group_data_refs, gpu_to_cpu=False, + canonical_layout=canonical_layout, ) def submit_store( diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 69cec67bff..75bea4e2e0 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -167,9 +167,10 @@ class SharedOffloadRegion: self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8) self._views: list[torch.Tensor] = [] + self._canonical_offset = 0 self.is_pinned: bool = False - def create_next_view(self, tensor_page_size: int) -> torch.Tensor: + def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: """Allocate a strided int8 view for this worker, one canonical tensor. Must be called once per canonical tensor. The full mmap layout is: @@ -208,6 +209,50 @@ class SharedOffloadRegion: self._views.append(worker_layer_view) return worker_layer_view + def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor: + """Allocate a strided int8 view shared by all workers for one + canonical tensor (canonical layout). + + Must be called once per canonical tensor, instead of + create_next_worker_view. The full mmap layout is: + + |<-------- canonical area ------->|<-------- unused ------->| + | all workers share this area | | + | | | + | [ canonical_t0 | canonical_t1 ] | | + | [ canonical_t0 | canonical_t1 ] | | + | [ canonical_t0 | canonical_t1 ] | | + ^ ^ + _canonical_offset=0, then advances by each tensor's size + + Each canonical_t{i} cell is that tensor's canonical page for the + block. Canonical areas are carved consecutively from the start of + each block row; consecutive rows are separated by row_stride. Every + worker gets the identical byte ranges and writes only its disjoint + bytes within them, as described by its canonical mappings — unlike + create_next_worker_view, which gives each worker a private + cpu_page_size slot per row. + + The trailing unused bytes exist only when the canonical pages sum to + less than row_stride: page-alignment padding of the row, or + deduplication of KV replicated across workers (e.g. the MLA latent), + where one canonical copy replaces world_size worker copies. + + Args: + tensor_page_size: Canonical bytes per block for this tensor. + """ + new_offset = self._canonical_offset + tensor_page_size + assert new_offset <= self._row_stride + view = torch.as_strided( + self._base, + size=(self.num_blocks, tensor_page_size), + stride=(self._row_stride, 1), + storage_offset=self._canonical_offset, + ) + self._canonical_offset = new_offset + self._views.append(view) + return view + def create_kv_memoryview(self) -> memoryview: """Return a zero-copy memoryview over the entire KV buffer. diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index 4b12dba913..01c31551bc 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -4,6 +4,9 @@ import hashlib import json +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import ( + canonical_format_id, +) from vllm.v1.kv_offload.base import ( OffloadingSpec, OffloadKey, @@ -36,6 +39,7 @@ class FileMapper: inference_engine: str = "vllm", parallel_agnostic: bool = False, replicated_layout: bool = False, + canonical_format: str | None = None, ): """ Initialize the file mapper. Each worker constructs its own, but @@ -65,6 +69,11 @@ class FileMapper: # unchanged (False is the historical default and must not appear). if replicated_layout: self.fields["replicated_layout"] = True + # The canonical byte format is not interchangeable with the direct + # layout (or with other canonical format versions/families), so its + # identity participates in the storage namespace. + if canonical_format is not None: + self.fields["canonical_format"] = canonical_format self.base_path: str = self._compute_base_path(root_dir, self.fields) @classmethod @@ -85,6 +94,7 @@ class FileMapper: for group in config.groups ] parallel = config.parallel + canonical_format = canonical_format_id() if config.canonical_layout else None return cls( root_dir=root_dir, model_name=config.model.name, @@ -102,6 +112,7 @@ class FileMapper: and (parallel.is_parallelism_agnostic or config.replicated_layout) ), replicated_layout=(parallel_agnostic and config.replicated_layout), + canonical_format=canonical_format, ) def get_file_name(self, key: OffloadKey) -> str: diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index d1d84c356e..bb0325ca28 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -264,6 +264,11 @@ class TieringOffloadingSpec(CPUOffloadingSpec): # Scheduler-side mmap (rank=None); kept for cleanup self._scheduler_mmap: SharedOffloadRegion | None = None + # Set by create_worker when canonical_layout is enabled: True when + # every layer's canonical bytes are parallelism-agnostic (portable), + # False when some layers use the opaque fallback (exact-topology only) + self.all_layers_portable: bool | None = None + # engine_id is unique per DP replica (suffixed with _dp{rank} in both # the Ray and multiprocessing paths), so it names a per-replica offload # region. @@ -394,12 +399,53 @@ class TieringOffloadingSpec(CPUOffloadingSpec): cpu_page_size=self.cpu_page_size_per_worker, ) try: + if self.config.canonical_layout: + self._validate_canonical_refs(kv_caches) return CPUOffloadingWorker( kv_caches=kv_caches, blocks_per_chunk=self.blocks_per_chunk, num_cpu_blocks=self.num_blocks, mmap_region=worker_mmap, + canonical_layout=self.config.canonical_layout, ) except Exception: worker_mmap.cleanup() raise + + def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None: + """Require a mapping on every ref and record layer portability. + + Fails loudly rather than persist direct-layout bytes under a + canonical format identity.""" + all_refs = [ + ref for group_refs in kv_caches.group_data_refs for ref in group_refs + ] + if any(ref.mapping is None for ref in all_refs): + raise RuntimeError( + "canonical_layout was requested but the KV cache layout " + "could not be certified for canonical offload (offload " + "workers must be exactly the TP group, and packed / " + "cross-layer KV layouts are not supported). Remove " + "canonical_layout from kv_connector_extra_config." + ) + self.all_layers_portable = all( + ref.mapping is not None and ref.mapping.parallelism_agnostic + for ref in all_refs + ) + if self.config.parallel.is_parallelism_agnostic and not ( + self.all_layers_portable + ): + # The scheduler-side storage namespace was already collapsed on + # the static portability claim; opaque per-topology bytes must + # not land in it. + raise RuntimeError( + "canonical_layout could not certify every layer as " + "parallelism-agnostic, but the storage namespace is shared " + "across topologies. Remove canonical_layout from " + "kv_connector_extra_config or disable parallel-agnostic " + "secondary tiers." + ) + logger.info( + "Canonical KV layout enabled (all_layers_portable=%s)", + self.all_layers_portable, + )