biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

OffloadingConnector: Support kernel_block_size != block_size (#30692)

Browse Source 
Signed-off-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Or Ozeri
2026-01-22 14:30:04 +02:00
committed by GitHub
parent 841d53aaa8
commit 421012b63a
7 changed files with 158 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
tests/v1/kv_offload/test_cpu_gpu.py
View File

@@ -25,8 +25,9 @@ if not current_platform.is_rocm():
NUM_GPU_BLOCKS = [64]
NUM_CPU_BLOCKS = [256]
GPU_BLOCK_SIZES = [16]
GPU_BLOCKS_PER_CPU_BLOCK = [1, 3]
KERNEL_BLOCK_SIZES = [16]
LOGICAL_BLOCK_SIZES = [16, 32]
LOGICAL_BLOCKS_PER_CPU_BLOCK = [1, 3]
HEAD_SIZES = [64]
NUM_HEADS = [8]
NUM_LAYERS = [4]
@@ -40,8 +41,9 @@ NUM_MAPPINGS = [3]
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("gpu_block_size", GPU_BLOCK_SIZES)
@pytest.mark.parametrize("gpu_blocks_per_cpu_block", GPU_BLOCKS_PER_CPU_BLOCK)
@pytest.mark.parametrize("kernel_block_size", KERNEL_BLOCK_SIZES)
@pytest.mark.parametrize("logical_block_size", LOGICAL_BLOCK_SIZES)
@pytest.mark.parametrize("logical_blocks_per_cpu_block", LOGICAL_BLOCKS_PER_CPU_BLOCK)
@pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
@pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
@pytest.mark.parametrize("num_layers", NUM_LAYERS)
@@ -55,8 +57,9 @@ def test_transfer(
num_mappings: int,
head_size: int,
num_heads: int,
gpu_block_size: int,
gpu_blocks_per_cpu_block: int,
kernel_block_size: int,
logical_block_size: int,
logical_blocks_per_cpu_block: int,
num_gpu_blocks: int,
num_cpu_blocks: int,
num_layers: int,
@@ -69,6 +72,10 @@ def test_transfer(
# create per-layer GPU KV caches based on available attn_backends
attn_backends_list = BACKENDS_TO_TEST
assert logical_block_size % kernel_block_size == 0
kernel_blocks_per_gpu_block = logical_block_size // kernel_block_size
num_gpu_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
gpu_caches = {}
attn_backends = {}
for i in range(num_layers):
@@ -78,15 +85,16 @@ def test_transfer(
attn_backends[layer_name] = attn_backend
gpu_cache_shape = attn_backend.get_kv_cache_shape(
num_gpu_blocks, gpu_block_size, num_heads, head_size
num_gpu_kernel_blocks, kernel_block_size, num_heads, head_size
)
gpu_caches[layer_name] = torch.rand(gpu_cache_shape, dtype=dtype, device=device)
# create handler
cpu_block_size = gpu_blocks_per_cpu_block * gpu_block_size
cpu_block_size = logical_blocks_per_cpu_block * logical_block_size
kernel_blocks_per_cpu_block = cpu_block_size // kernel_block_size
handlers = CpuGpuOffloadingHandlers(
attn_backends=attn_backends,
gpu_block_size=gpu_block_size,
gpu_block_size=logical_block_size,
cpu_block_size=cpu_block_size,
num_cpu_blocks=num_cpu_blocks,
gpu_caches=gpu_caches,
@@ -94,22 +102,34 @@ def test_transfer(
# select block mappings
gpu_blocks = random.sample(
range(num_gpu_blocks), num_mappings * gpu_blocks_per_cpu_block
range(num_gpu_blocks), num_mappings * logical_blocks_per_cpu_block
)
cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
# convert cpu blocks to gpu block size
cpu_blocks_in_gpu_block_size = []
for cpu_block in cpu_blocks:
base_block_id = cpu_block * gpu_blocks_per_cpu_block
for i in range(gpu_blocks_per_cpu_block):
cpu_blocks_in_gpu_block_size.append(i + base_block_id)
# convert gpu blocks to kernel block size
gpu_blocks_in_kernel_block_size = []
for gpu_block in gpu_blocks:
base_block_id = gpu_block * kernel_blocks_per_gpu_block
for i in range(kernel_blocks_per_gpu_block):
gpu_blocks_in_kernel_block_size.append(i + base_block_id)
# maybe skip a GPU block to test reading from the middle of a CPU block
# convert cpu blocks to gpu block size
cpu_blocks_in_kernel_block_size = []
for cpu_block in cpu_blocks:
base_block_id = cpu_block * kernel_blocks_per_cpu_block
for i in range(kernel_blocks_per_cpu_block):
cpu_blocks_in_kernel_block_size.append(i + base_block_id)
# maybe skip some GPU block to test reading from the middle of a CPU block
if not gpu_to_cpu:
gpu_blocks = gpu_blocks[gpu_blocks_per_cpu_block - 1 :]
cpu_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size[
gpu_blocks_per_cpu_block - 1 :
gpu_blocks_to_skip = logical_blocks_per_cpu_block - 1
gpu_blocks = gpu_blocks[gpu_blocks_to_skip:]
kernel_blocks_to_skip = gpu_blocks_to_skip * kernel_blocks_per_gpu_block
gpu_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size[
kernel_blocks_to_skip:
]
cpu_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size[
kernel_blocks_to_skip:
]
# set transfer direction
@@ -119,23 +139,23 @@ def test_transfer(
dst_spec_class = CPULoadStoreSpec
src_blocks = gpu_blocks
dst_blocks = cpu_blocks
src_blocks_in_gpu_block_size = gpu_blocks
dst_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size
dst_size_in_gpu_blocks = num_cpu_blocks * gpu_blocks_per_cpu_block
src_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
dst_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
dst_size_in_kernel_blocks = num_cpu_blocks * kernel_blocks_per_cpu_block
else:
handler = handlers.cpu_to_gpu_handler
src_spec_class = CPULoadStoreSpec
dst_spec_class = GPULoadStoreSpec
src_blocks = cpu_blocks
dst_blocks = gpu_blocks
src_blocks_in_gpu_block_size = cpu_blocks_in_gpu_block_size
dst_blocks_in_gpu_block_size = gpu_blocks
dst_size_in_gpu_blocks = num_gpu_blocks
src_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
dst_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
dst_size_in_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
# build dst -> src mapping
dst_to_src = {}
for src_block, dst_block in zip(
src_blocks_in_gpu_block_size, dst_blocks_in_gpu_block_size
src_blocks_in_kernel_block_size, dst_blocks_in_kernel_block_size
):
dst_to_src[dst_block] = src_block
@@ -165,29 +185,15 @@ def test_transfer(
assert torch.equal(orig_tensor, tensor)
# verify dst tensors
for dst_block in range(dst_size_in_gpu_blocks):
for dst_block in range(dst_size_in_kernel_blocks):
src_block_candidate = dst_to_src.get(dst_block)
for src_cache, dst_cache, orig_dst_cache, kv_dim in zip(
for src_cache, dst_cache, orig_dst_cache in zip(
handler.src_tensors,
handler.dst_tensors,
orig_dst_caches,
handler.kv_dim_before_num_blocks,
):
if kv_dim:
# iterate over key, value
for i in range(2):
if src_block_candidate is not None:
expected_value = src_cache[i][src_block_candidate]
else:
expected_value = orig_dst_cache[i][dst_block]
torch.testing.assert_close(
dst_cache[i][dst_block].cpu(), expected_value.cpu()
)
if src_block_candidate is not None:
expected_value = src_cache[src_block_candidate]
else:
if src_block_candidate is not None:
expected_value = src_cache[src_block_candidate]
else:
expected_value = orig_dst_cache[dst_block]
torch.testing.assert_close(
dst_cache[dst_block].cpu(), expected_value.cpu()
)
expected_value = orig_dst_cache[dst_block]
torch.testing.assert_close(dst_cache[dst_block].cpu(), expected_value.cpu())