forked from Karylab-cklius/vllm
716 lines
27 KiB
Python
716 lines
27 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""GPU integration tests for the extensible KV cache allocation paths.
|
|
|
|
Drives `GPUModelRunner._allocate_kv_cache_tensors` / `_reshape_kv_cache_tensors`
|
|
/ `extend_kv_cache` directly with fake attention backends, covering the buffer
|
|
layouts the extensible flow supports: block-major (one committed prefix),
|
|
K/V-split (one prefix per half), Mamba (block-major per layer), and hybrid
|
|
attention + Mamba (attention re-strided to block-major). Buffer sizes exceed
|
|
the CUDA VMM allocation granularity so touching a block that the commit logic
|
|
missed would fault instead of silently passing.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.v1.attention.backend import AttentionBackend
|
|
from vllm.v1.kv_cache_interface import (
|
|
FullAttentionSpec,
|
|
KVCacheConfig,
|
|
KVCacheGroupSpec,
|
|
KVCacheTensor,
|
|
MambaSpec,
|
|
)
|
|
from vllm.v1.worker.gpu.attn_utils import (
|
|
_allocate_extensible_kv_cache,
|
|
_kv_cache_num_segments_by_layer,
|
|
_reshape_kv_cache,
|
|
narrow_kv_caches_to_num_blocks,
|
|
)
|
|
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
|
from vllm.v1.worker.gpu_worker import Worker
|
|
from vllm.v1.worker.utils import AttentionGroup
|
|
|
|
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
|
|
|
BLOCK_SIZE = 16
|
|
NUM_BLOCKS = 256
|
|
|
|
|
|
class _SplitKVBackend(AttentionBackend):
|
|
"""Fake backend with a K/V-split layout, like FlashAttention."""
|
|
|
|
@staticmethod
|
|
def get_kv_cache_shape(
|
|
num_blocks: int,
|
|
block_size: int,
|
|
num_kv_heads: int,
|
|
head_size: int,
|
|
cache_dtype_str: str = "auto",
|
|
) -> tuple[int, ...]:
|
|
return (2, num_blocks, block_size, num_kv_heads, head_size)
|
|
|
|
|
|
class _BlockMajorBackend(AttentionBackend):
|
|
"""Fake backend with a num-blocks-first layout, like FlashInfer."""
|
|
|
|
@staticmethod
|
|
def get_kv_cache_shape(
|
|
num_blocks: int,
|
|
block_size: int,
|
|
num_kv_heads: int,
|
|
head_size: int,
|
|
cache_dtype_str: str = "auto",
|
|
) -> tuple[int, ...]:
|
|
return (num_blocks, 2, block_size, num_kv_heads, head_size)
|
|
|
|
|
|
class _StrideOrderBackend(AttentionBackend):
|
|
"""Fake backend whose stride order makes a kv-first shape block-major."""
|
|
|
|
@staticmethod
|
|
def get_kv_cache_shape(
|
|
num_blocks: int,
|
|
block_size: int,
|
|
num_kv_heads: int,
|
|
head_size: int,
|
|
cache_dtype_str: str = "auto",
|
|
) -> tuple[int, ...]:
|
|
return (2, num_blocks, block_size, num_kv_heads, head_size)
|
|
|
|
@staticmethod
|
|
def get_kv_cache_stride_order(
|
|
include_num_layers_dimension: bool = False,
|
|
) -> tuple[int, ...]:
|
|
assert not include_num_layers_dimension
|
|
return (1, 0, 2, 3, 4)
|
|
|
|
|
|
def _full_attention_spec() -> FullAttentionSpec:
|
|
# page_size_bytes = 2 (K+V) * 16 * 8 * 128 * 2 bytes = 64 KiB; 256 blocks
|
|
# = 16 MiB, several VMM granules per buffer.
|
|
return FullAttentionSpec(
|
|
block_size=BLOCK_SIZE,
|
|
num_kv_heads=8,
|
|
head_size=128,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
|
|
def _mamba_spec() -> MambaSpec:
|
|
# page_size_bytes = (8*128 + 16*64) * 4 bytes = 8 KiB per block per layer.
|
|
return MambaSpec(
|
|
block_size=BLOCK_SIZE,
|
|
shapes=((8, 128), (16, 64)),
|
|
dtypes=(torch.float32, torch.float32),
|
|
)
|
|
|
|
|
|
def _make_runner(kv_cache_config: KVCacheConfig, attn_groups) -> GPUModelRunner:
|
|
runner = object.__new__(GPUModelRunner)
|
|
runner.device = torch.device("cuda:0")
|
|
runner.kv_cache_config = kv_cache_config
|
|
runner.attn_groups = attn_groups
|
|
runner.runner_only_attn_layers = set()
|
|
runner.cache_config = SimpleNamespace(cache_dtype="auto")
|
|
return runner
|
|
|
|
|
|
def _attention_config(spec: FullAttentionSpec, backend) -> tuple[KVCacheConfig, list]:
|
|
kv_cache_config = KVCacheConfig(
|
|
num_blocks=NUM_BLOCKS,
|
|
kv_cache_tensors=[
|
|
KVCacheTensor(size=NUM_BLOCKS * spec.page_size_bytes, shared_by=["layer.0"])
|
|
],
|
|
kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)],
|
|
)
|
|
attn_groups = [
|
|
[
|
|
AttentionGroup(
|
|
backend=backend,
|
|
layer_names=["layer.0"],
|
|
kv_cache_spec=spec,
|
|
kv_cache_group_id=0,
|
|
)
|
|
]
|
|
]
|
|
return kv_cache_config, attn_groups
|
|
|
|
|
|
def _free_buffers(runner: GPUModelRunner) -> None:
|
|
buffers = getattr(runner, "extensible_kv_buffers", None)
|
|
if buffers is not None:
|
|
buffers.free()
|
|
|
|
|
|
def test_kv_cache_num_segments_by_layer() -> None:
|
|
"""Segment counts follow the physical layout of each layer's backend."""
|
|
spec = _full_attention_spec()
|
|
for backend, expected in (
|
|
(_SplitKVBackend, 2),
|
|
(_BlockMajorBackend, 1),
|
|
# kv-first logical shape but block-major physical order -> 1 segment.
|
|
(_StrideOrderBackend, 1),
|
|
):
|
|
kv_cache_config, attn_groups = _attention_config(spec, backend)
|
|
runner = _make_runner(kv_cache_config, attn_groups)
|
|
assert runner._kv_cache_num_segments_by_layer() == {"layer.0": expected}
|
|
|
|
|
|
def test_extensible_split_layout_grows_both_halves() -> None:
|
|
"""A K/V-split layer keeps its natural layout and both halves grow in
|
|
lockstep."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
|
|
runner = _make_runner(kv_cache_config, attn_groups)
|
|
try:
|
|
raw_tensors = runner._allocate_kv_cache_tensors(
|
|
kv_cache_config, extensible=True
|
|
)
|
|
kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
|
|
kv_cache = kv_caches["layer.0"]
|
|
assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128)
|
|
[(buffer, bytes_per_block_per_segment)] = runner.extensible_kv_buffers.buffers
|
|
assert buffer.num_segments == 2
|
|
assert bytes_per_block_per_segment == spec.page_size_bytes // 2
|
|
|
|
# Only block 0 is committed -- in each half.
|
|
kv_cache[0, 0].fill_(1) # K, block 0
|
|
kv_cache[1, 0].fill_(2) # V, block 0
|
|
torch.accelerator.synchronize()
|
|
|
|
runner.extend_kv_cache(NUM_BLOCKS)
|
|
# Old data survives the grow; new blocks are usable in both halves and
|
|
# zeroed.
|
|
assert torch.all(kv_cache[0, 0] == 1)
|
|
assert torch.all(kv_cache[1, 0] == 2)
|
|
kv_cache[0, NUM_BLOCKS - 1].fill_(3)
|
|
kv_cache[1, NUM_BLOCKS - 1].fill_(4)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 3)
|
|
assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4)
|
|
assert torch.count_nonzero(kv_cache[:, 1 : NUM_BLOCKS - 1]) == 0
|
|
finally:
|
|
_free_buffers(runner)
|
|
|
|
|
|
def test_extensible_block_major_layout() -> None:
|
|
"""A layer whose physical layout is block-major uses a single segment."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _BlockMajorBackend)
|
|
runner = _make_runner(kv_cache_config, attn_groups)
|
|
try:
|
|
raw_tensors = runner._allocate_kv_cache_tensors(
|
|
kv_cache_config, extensible=True
|
|
)
|
|
kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
|
|
kv_cache = kv_caches["layer.0"]
|
|
assert kv_cache.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128)
|
|
[(buffer, bytes_per_block_per_segment)] = runner.extensible_kv_buffers.buffers
|
|
assert buffer.num_segments == 1
|
|
assert bytes_per_block_per_segment == spec.page_size_bytes
|
|
|
|
kv_cache[0].fill_(1)
|
|
runner.extend_kv_cache(NUM_BLOCKS)
|
|
kv_cache[NUM_BLOCKS - 1].fill_(2)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(kv_cache[0] == 1)
|
|
assert torch.all(kv_cache[NUM_BLOCKS - 1] == 2)
|
|
assert torch.count_nonzero(kv_cache[1 : NUM_BLOCKS - 1]) == 0
|
|
finally:
|
|
_free_buffers(runner)
|
|
|
|
|
|
def test_legacy_split_layout_commits_everything() -> None:
|
|
"""Without `extensible`, the full buffer is committed up front."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
|
|
runner = _make_runner(kv_cache_config, attn_groups)
|
|
raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config, extensible=False)
|
|
kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
|
|
kv_cache = kv_caches["layer.0"]
|
|
kv_cache[0, NUM_BLOCKS - 1].fill_(1)
|
|
kv_cache[1, NUM_BLOCKS - 1].fill_(2)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(kv_cache[0, NUM_BLOCKS - 1] == 1)
|
|
assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 2)
|
|
with pytest.raises(RuntimeError, match="extensible"):
|
|
runner.extend_kv_cache(NUM_BLOCKS)
|
|
|
|
|
|
def test_extensible_mamba_grows_per_layer() -> None:
|
|
"""Mamba per-layer buffers are block-major and grow with the KV cache."""
|
|
spec = _mamba_spec()
|
|
num_blocks = 512
|
|
layer_names = ["mamba.0", "mamba.1"]
|
|
kv_cache_config = KVCacheConfig(
|
|
num_blocks=num_blocks,
|
|
kv_cache_tensors=[
|
|
KVCacheTensor(size=num_blocks * spec.page_size_bytes, shared_by=[name])
|
|
for name in layer_names
|
|
],
|
|
kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
|
|
)
|
|
attn_groups = [
|
|
[
|
|
AttentionGroup(
|
|
backend=_BlockMajorBackend,
|
|
layer_names=layer_names,
|
|
kv_cache_spec=spec,
|
|
kv_cache_group_id=0,
|
|
)
|
|
]
|
|
]
|
|
runner = _make_runner(kv_cache_config, attn_groups)
|
|
try:
|
|
raw_tensors = runner._allocate_kv_cache_tensors(
|
|
kv_cache_config, extensible=True
|
|
)
|
|
kv_caches = runner._reshape_kv_cache_tensors(raw_tensors, [BLOCK_SIZE])
|
|
assert set(kv_caches) == set(layer_names)
|
|
assert len(runner.extensible_kv_buffers.buffers) == len(layer_names)
|
|
for buffer, bytes_per_block_per_segment in runner.extensible_kv_buffers.buffers:
|
|
assert buffer.num_segments == 1
|
|
assert bytes_per_block_per_segment == spec.page_size_bytes
|
|
|
|
# Write block 0 of every state of every layer (the committed
|
|
# prefixes), then grow.
|
|
for name in layer_names:
|
|
for state_tensor in kv_caches[name]:
|
|
state_tensor[0].fill_(1)
|
|
torch.accelerator.synchronize()
|
|
runner.extend_kv_cache(num_blocks)
|
|
for name in layer_names:
|
|
for state_tensor in kv_caches[name]:
|
|
state_tensor[num_blocks - 1].fill_(2)
|
|
torch.accelerator.synchronize()
|
|
for name in layer_names:
|
|
for state_tensor in kv_caches[name]:
|
|
assert torch.all(state_tensor[0] == 1)
|
|
assert torch.all(state_tensor[num_blocks - 1] == 2)
|
|
assert torch.count_nonzero(state_tensor[1 : num_blocks - 1]) == 0
|
|
finally:
|
|
_free_buffers(runner)
|
|
|
|
|
|
def test_extensible_hybrid_attention_mamba() -> None:
|
|
"""In hybrid models the attention cache is re-strided to block-major, so
|
|
its buffer must use a single segment."""
|
|
attn_spec = _full_attention_spec()
|
|
mamba_spec = _mamba_spec()
|
|
kv_cache_config = KVCacheConfig(
|
|
num_blocks=NUM_BLOCKS,
|
|
kv_cache_tensors=[
|
|
KVCacheTensor(
|
|
size=NUM_BLOCKS * attn_spec.page_size_bytes, shared_by=["attn.0"]
|
|
),
|
|
KVCacheTensor(
|
|
size=NUM_BLOCKS * mamba_spec.page_size_bytes, shared_by=["mamba.0"]
|
|
),
|
|
],
|
|
kv_cache_groups=[
|
|
KVCacheGroupSpec(layer_names=["attn.0"], kv_cache_spec=attn_spec),
|
|
KVCacheGroupSpec(layer_names=["mamba.0"], kv_cache_spec=mamba_spec),
|
|
],
|
|
)
|
|
attn_groups = [
|
|
[
|
|
AttentionGroup(
|
|
backend=_SplitKVBackend,
|
|
layer_names=["attn.0"],
|
|
kv_cache_spec=attn_spec,
|
|
kv_cache_group_id=0,
|
|
)
|
|
],
|
|
[
|
|
AttentionGroup(
|
|
backend=_BlockMajorBackend,
|
|
layer_names=["mamba.0"],
|
|
kv_cache_spec=mamba_spec,
|
|
kv_cache_group_id=1,
|
|
)
|
|
],
|
|
]
|
|
runner = _make_runner(kv_cache_config, attn_groups)
|
|
try:
|
|
# The K/V-split attention layer is forced to one segment by the hybrid
|
|
# block-major re-stride.
|
|
assert runner._kv_cache_num_segments_by_layer() == {"attn.0": 1, "mamba.0": 1}
|
|
|
|
raw_tensors = runner._allocate_kv_cache_tensors(
|
|
kv_cache_config, extensible=True
|
|
)
|
|
kv_caches = runner._reshape_kv_cache_tensors(
|
|
raw_tensors, [BLOCK_SIZE, BLOCK_SIZE]
|
|
)
|
|
attn_cache = kv_caches["attn.0"]
|
|
# `_update_hybrid_attention_mamba_layout` re-strides to interleave K/V
|
|
# per block: block b spans one contiguous page.
|
|
hidden_size = attn_cache.shape[2:].numel()
|
|
assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size)
|
|
|
|
attn_cache[0, 0].fill_(1) # K, block 0
|
|
attn_cache[1, 0].fill_(2) # V, block 0
|
|
for state_tensor in kv_caches["mamba.0"]:
|
|
state_tensor[0].fill_(3)
|
|
torch.accelerator.synchronize()
|
|
|
|
runner.extend_kv_cache(NUM_BLOCKS)
|
|
attn_cache[0, NUM_BLOCKS - 1].fill_(4)
|
|
attn_cache[1, NUM_BLOCKS - 1].fill_(5)
|
|
for state_tensor in kv_caches["mamba.0"]:
|
|
state_tensor[NUM_BLOCKS - 1].fill_(6)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(attn_cache[0, 0] == 1)
|
|
assert torch.all(attn_cache[1, 0] == 2)
|
|
assert torch.all(attn_cache[0, NUM_BLOCKS - 1] == 4)
|
|
assert torch.all(attn_cache[1, NUM_BLOCKS - 1] == 5)
|
|
assert torch.count_nonzero(attn_cache[:, 1 : NUM_BLOCKS - 1]) == 0
|
|
for state_tensor in kv_caches["mamba.0"]:
|
|
assert torch.all(state_tensor[0] == 3)
|
|
assert torch.all(state_tensor[NUM_BLOCKS - 1] == 6)
|
|
assert torch.count_nonzero(state_tensor[1 : NUM_BLOCKS - 1]) == 0
|
|
finally:
|
|
_free_buffers(runner)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V2 model runner (vllm.v1.worker.gpu) extensible allocation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _v2_allocate(kv_cache_config, attn_groups, kernel_block_sizes):
|
|
flat_groups = [g for groups in attn_groups for g in groups]
|
|
raw_tensors, buffers = _allocate_extensible_kv_cache(
|
|
kv_cache_config,
|
|
{},
|
|
torch.device("cuda:0"),
|
|
flat_groups,
|
|
kernel_block_sizes,
|
|
"auto",
|
|
)
|
|
kv_caches = _reshape_kv_cache(
|
|
attn_groups=flat_groups,
|
|
kv_cache_raw_tensors=raw_tensors,
|
|
cache_dtype="auto",
|
|
kernel_block_sizes=kernel_block_sizes,
|
|
shared_kv_cache_layers={},
|
|
kv_cache_config=kv_cache_config,
|
|
)
|
|
return kv_caches, buffers
|
|
|
|
|
|
def test_v2_num_segments_by_layer() -> None:
|
|
"""V2 segment counts follow the layer's physical layout, and hybrid
|
|
models force block-major (one segment)."""
|
|
spec = _full_attention_spec()
|
|
for backend, expected in (
|
|
(_SplitKVBackend, 2),
|
|
(_BlockMajorBackend, 1),
|
|
(_StrideOrderBackend, 1),
|
|
):
|
|
_, attn_groups = _attention_config(spec, backend)
|
|
flat_groups = [g for groups in attn_groups for g in groups]
|
|
assert _kv_cache_num_segments_by_layer(
|
|
flat_groups, [BLOCK_SIZE], "auto", has_mamba=False
|
|
) == {"layer.0": expected}
|
|
assert _kv_cache_num_segments_by_layer(
|
|
flat_groups, [BLOCK_SIZE], "auto", has_mamba=True
|
|
) == {"layer.0": 1}
|
|
|
|
|
|
def test_v2_extensible_split_layout_grows_incrementally() -> None:
|
|
"""A K/V-split layer grows both halves in lockstep through the staged
|
|
commits the V2 flow performs (init -> warmup prefix -> final size)."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
|
|
kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
|
|
try:
|
|
kv_cache = kv_caches["layer.0"]
|
|
assert kv_cache.shape == (2, NUM_BLOCKS, BLOCK_SIZE, 8, 128)
|
|
assert buffers.num_blocks_committed == 1
|
|
|
|
kv_cache[0, 0].fill_(1) # K, block 0
|
|
kv_cache[1, 0].fill_(2) # V, block 0
|
|
torch.accelerator.synchronize()
|
|
|
|
# Warmup-style prefix commit, then the final post-warmup commit.
|
|
buffers.commit(8)
|
|
kv_cache[0, 7].fill_(3)
|
|
torch.accelerator.synchronize()
|
|
buffers.commit(NUM_BLOCKS)
|
|
# Shrink requests are ignored.
|
|
buffers.commit(1)
|
|
assert buffers.num_blocks_committed == NUM_BLOCKS
|
|
|
|
kv_cache[1, NUM_BLOCKS - 1].fill_(4)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(kv_cache[0, 0] == 1)
|
|
assert torch.all(kv_cache[1, 0] == 2)
|
|
assert torch.all(kv_cache[0, 7] == 3)
|
|
assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 4)
|
|
assert torch.count_nonzero(kv_cache[:, 1:7]) == 0
|
|
assert torch.count_nonzero(kv_cache[:, 8 : NUM_BLOCKS - 1]) == 0
|
|
assert buffers.physical_bytes >= NUM_BLOCKS * spec.page_size_bytes
|
|
finally:
|
|
buffers.free()
|
|
|
|
|
|
def test_v2_extensible_hybrid_attention_mamba() -> None:
|
|
"""V2 hybrid models re-stride attention to block-major; both the
|
|
attention and Mamba buffers grow as single-segment prefixes."""
|
|
attn_spec = _full_attention_spec()
|
|
mamba_spec = _mamba_spec()
|
|
kv_cache_config = KVCacheConfig(
|
|
num_blocks=NUM_BLOCKS,
|
|
kv_cache_tensors=[
|
|
KVCacheTensor(
|
|
size=NUM_BLOCKS * attn_spec.page_size_bytes, shared_by=["attn.0"]
|
|
),
|
|
KVCacheTensor(
|
|
size=NUM_BLOCKS * mamba_spec.page_size_bytes, shared_by=["mamba.0"]
|
|
),
|
|
],
|
|
kv_cache_groups=[
|
|
KVCacheGroupSpec(layer_names=["attn.0"], kv_cache_spec=attn_spec),
|
|
KVCacheGroupSpec(layer_names=["mamba.0"], kv_cache_spec=mamba_spec),
|
|
],
|
|
)
|
|
attn_groups = [
|
|
[
|
|
AttentionGroup(
|
|
backend=_SplitKVBackend,
|
|
layer_names=["attn.0"],
|
|
kv_cache_spec=attn_spec,
|
|
kv_cache_group_id=0,
|
|
)
|
|
],
|
|
[
|
|
AttentionGroup(
|
|
backend=_BlockMajorBackend,
|
|
layer_names=["mamba.0"],
|
|
kv_cache_spec=mamba_spec,
|
|
kv_cache_group_id=1,
|
|
)
|
|
],
|
|
]
|
|
kv_caches, buffers = _v2_allocate(
|
|
kv_cache_config, attn_groups, [BLOCK_SIZE, BLOCK_SIZE]
|
|
)
|
|
try:
|
|
attn_cache = kv_caches["attn.0"]
|
|
# Re-strided to interleave K/V per block: block b spans one page.
|
|
hidden_size = attn_cache.shape[2:].numel()
|
|
assert attn_cache.stride()[:2] == (hidden_size, 2 * hidden_size)
|
|
|
|
attn_cache[0, 0].fill_(1)
|
|
attn_cache[1, 0].fill_(2)
|
|
for state_tensor in kv_caches["mamba.0"]:
|
|
state_tensor[0].fill_(3)
|
|
torch.accelerator.synchronize()
|
|
|
|
buffers.commit(NUM_BLOCKS)
|
|
attn_cache[0, NUM_BLOCKS - 1].fill_(4)
|
|
attn_cache[1, NUM_BLOCKS - 1].fill_(5)
|
|
for state_tensor in kv_caches["mamba.0"]:
|
|
state_tensor[NUM_BLOCKS - 1].fill_(6)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(attn_cache[0, 0] == 1)
|
|
assert torch.all(attn_cache[1, 0] == 2)
|
|
assert torch.all(attn_cache[0, NUM_BLOCKS - 1] == 4)
|
|
assert torch.all(attn_cache[1, NUM_BLOCKS - 1] == 5)
|
|
assert torch.count_nonzero(attn_cache[:, 1 : NUM_BLOCKS - 1]) == 0
|
|
for state_tensor in kv_caches["mamba.0"]:
|
|
assert torch.all(state_tensor[0] == 3)
|
|
assert torch.all(state_tensor[NUM_BLOCKS - 1] == 6)
|
|
assert torch.count_nonzero(state_tensor[1 : NUM_BLOCKS - 1]) == 0
|
|
finally:
|
|
buffers.free()
|
|
|
|
|
|
def test_v2_extensible_packed_layout() -> None:
|
|
"""A packed (block_stride) layout uses one shared block-major buffer;
|
|
per-layer pages within a block stay isolated across commits."""
|
|
spec = _full_attention_spec()
|
|
page_bytes = spec.page_size_bytes
|
|
block_stride = 2 * page_bytes # two layers packed per block
|
|
layer_names = ["packed.0", "packed.1"]
|
|
kv_cache_config = KVCacheConfig(
|
|
num_blocks=NUM_BLOCKS,
|
|
kv_cache_tensors=[
|
|
KVCacheTensor(
|
|
size=NUM_BLOCKS * block_stride,
|
|
shared_by=[name],
|
|
offset=i * page_bytes,
|
|
block_stride=block_stride,
|
|
)
|
|
for i, name in enumerate(layer_names)
|
|
],
|
|
kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
|
|
)
|
|
attn_groups = [
|
|
[
|
|
AttentionGroup(
|
|
backend=_BlockMajorBackend,
|
|
layer_names=layer_names,
|
|
kv_cache_spec=spec,
|
|
kv_cache_group_id=0,
|
|
)
|
|
]
|
|
]
|
|
kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
|
|
try:
|
|
assert len(buffers.buffers) == 1
|
|
[(buffer, bytes_per_block)] = buffers.buffers
|
|
assert buffer.num_segments == 1
|
|
assert bytes_per_block == block_stride
|
|
|
|
cache0, cache1 = kv_caches["packed.0"], kv_caches["packed.1"]
|
|
assert cache0.shape == (NUM_BLOCKS, 2, BLOCK_SIZE, 8, 128)
|
|
|
|
cache0[0].fill_(1)
|
|
cache1[0].fill_(2)
|
|
torch.accelerator.synchronize()
|
|
buffers.commit(NUM_BLOCKS)
|
|
cache0[NUM_BLOCKS - 1].fill_(3)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(cache0[0] == 1)
|
|
assert torch.all(cache1[0] == 2)
|
|
assert torch.all(cache0[NUM_BLOCKS - 1] == 3)
|
|
# The other layer's page of the same block is untouched, and all
|
|
# middle blocks were zeroed on commit.
|
|
assert torch.count_nonzero(cache1[1:]) == 0
|
|
assert torch.count_nonzero(cache0[1 : NUM_BLOCKS - 1]) == 0
|
|
|
|
committed = NUM_BLOCKS // 2
|
|
narrowed = narrow_kv_caches_to_num_blocks(
|
|
kv_caches,
|
|
[g for groups in attn_groups for g in groups],
|
|
[BLOCK_SIZE],
|
|
"auto",
|
|
committed,
|
|
kv_cache_config,
|
|
)
|
|
narrowed0 = narrowed["packed.0"]
|
|
narrowed1 = narrowed["packed.1"]
|
|
assert narrowed0.untyped_storage().data_ptr() == buffer.base_ptr
|
|
assert (
|
|
narrowed0.untyped_storage().data_ptr()
|
|
== narrowed1.untyped_storage().data_ptr()
|
|
)
|
|
assert narrowed0.untyped_storage().nbytes() == committed * block_stride
|
|
assert narrowed0.stride() == cache0.stride()
|
|
assert narrowed1.stride() == cache1.stride()
|
|
finally:
|
|
buffers.free()
|
|
|
|
|
|
def test_v2_extensible_release_and_recommit() -> None:
|
|
"""Sleep/wake cycle: release_physical discards data but keeps VA and
|
|
views valid; recommit restores the committed size with zeroed pages."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
|
|
kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
|
|
try:
|
|
kv_cache = kv_caches["layer.0"]
|
|
base_ptr = buffers.buffers[0][0].base_ptr
|
|
buffers.commit(NUM_BLOCKS)
|
|
kv_cache.fill_(7)
|
|
torch.accelerator.synchronize()
|
|
assert buffers.physical_bytes > 0
|
|
|
|
buffers.release_physical()
|
|
assert buffers.physical_bytes == 0
|
|
assert buffers.num_blocks_committed == 0
|
|
|
|
buffers.recommit()
|
|
assert buffers.num_blocks_committed == NUM_BLOCKS
|
|
assert buffers.buffers[0][0].base_ptr == base_ptr
|
|
torch.accelerator.synchronize()
|
|
# Data was discarded; fresh pages are zeroed and writable through
|
|
# the original views.
|
|
assert torch.count_nonzero(kv_cache) == 0
|
|
kv_cache[1, NUM_BLOCKS - 1].fill_(9)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 9)
|
|
finally:
|
|
buffers.free()
|
|
|
|
|
|
def test_v2_extensible_connector_sleep_fails_before_remapping() -> None:
|
|
"""Connector registrations must not survive physical-page replacement."""
|
|
worker = object.__new__(Worker)
|
|
worker.model_runner = SimpleNamespace(extensible_kv_buffers=object())
|
|
worker.vllm_config = SimpleNamespace(kv_transfer_config=object())
|
|
|
|
with pytest.raises(RuntimeError, match="invalidates.*memory registration"):
|
|
worker.sleep()
|
|
|
|
|
|
def test_v2_narrow_kv_caches_to_num_blocks() -> None:
|
|
"""Connector-registration views are trimmed to the committed block count
|
|
along each layout's block dim, keeping base pointers and strides (so the
|
|
K and V segment prefixes are addressed exactly)."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
|
|
kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
|
|
try:
|
|
committed = 16
|
|
buffers.commit(committed)
|
|
narrowed = narrow_kv_caches_to_num_blocks(
|
|
kv_caches,
|
|
[g for groups in attn_groups for g in groups],
|
|
[BLOCK_SIZE],
|
|
"auto",
|
|
committed,
|
|
kv_cache_config,
|
|
)
|
|
full = kv_caches["layer.0"]
|
|
trimmed = narrowed["layer.0"]
|
|
assert trimmed.shape == (2, committed, BLOCK_SIZE, 8, 128)
|
|
assert trimmed.stride() == full.stride()
|
|
# K prefix starts at the buffer base; V prefix at the segment offset.
|
|
assert trimmed[0].data_ptr() == full[0].data_ptr()
|
|
assert trimmed[1].data_ptr() == full[1].data_ptr()
|
|
# The narrowed views cover only committed memory.
|
|
trimmed[0, committed - 1].fill_(1)
|
|
trimmed[1, committed - 1].fill_(2)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(full[0, committed - 1] == 1)
|
|
assert torch.all(full[1, committed - 1] == 2)
|
|
finally:
|
|
buffers.free()
|
|
|
|
|
|
def test_v2_extensible_defragment_on_commit() -> None:
|
|
"""commit(defragment=True) re-maps each segment prefix as ONE physical
|
|
chunk (required for KV-transfer registration), discarding prior data."""
|
|
spec = _full_attention_spec()
|
|
kv_cache_config, attn_groups = _attention_config(spec, _SplitKVBackend)
|
|
kv_caches, buffers = _v2_allocate(kv_cache_config, attn_groups, [BLOCK_SIZE])
|
|
try:
|
|
kv_cache = kv_caches["layer.0"]
|
|
# Staged commits spanning multiple VMM granules -> multiple physical
|
|
# chunks per segment.
|
|
buffers.commit(8)
|
|
buffers.commit(NUM_BLOCKS // 2)
|
|
kv_cache[0, 0].fill_(1)
|
|
torch.accelerator.synchronize()
|
|
[(buffer, _)] = buffers.buffers
|
|
assert len(buffer._buffer._handles) > 2
|
|
|
|
buffers.commit(NUM_BLOCKS, defragment=True)
|
|
# One chunk per segment; data discarded (zeroed); views still work.
|
|
assert len(buffer._buffer._handles) == 2
|
|
assert buffers.num_blocks_committed == NUM_BLOCKS
|
|
torch.accelerator.synchronize()
|
|
assert torch.count_nonzero(kv_cache) == 0
|
|
kv_cache[1, NUM_BLOCKS - 1].fill_(3)
|
|
torch.accelerator.synchronize()
|
|
assert torch.all(kv_cache[1, NUM_BLOCKS - 1] == 3)
|
|
finally:
|
|
buffers.free()
|