[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>
This commit is contained in:
Itay Etelis
2026-08-10 10:44:54 +03:00
committed by GitHub
co-authored by Itay Etelis
parent 31cd109f18
commit 81840a172f
12 changed files with 891 additions and 103 deletions
@@ -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
@@ -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()
+2 -2
View File
@@ -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
@@ -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()
+17
View File
@@ -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
# ---------------------------------------------------------------------------
@@ -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:
@@ -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,
)
+8 -2
View File
@@ -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
+328 -49
View File
@@ -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(
@@ -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.
+11
View File
@@ -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:
+46
View File
@@ -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,
)