[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:
co-authored by
Itay Etelis
parent
31cd109f18
commit
81840a172f
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user