153 lines
4.9 KiB
Python
153 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
|
|
device = torch.device("cuda")
|
|
num_blocks = 4
|
|
page_size_el = 4
|
|
storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device)
|
|
|
|
# Build the minimal zeroer state directly so the test can focus on the
|
|
# in-flight copy behavior without constructing model attention groups.
|
|
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
|
|
zeroer.device = device
|
|
zeroer._meta = (
|
|
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
|
|
torch.tensor([page_size_el], dtype=torch.int64, device=device),
|
|
page_size_el // page_size_el, # max_chunks = 1
|
|
page_size_el, # blk_size
|
|
1, # n_segs
|
|
)
|
|
|
|
stream = torch.cuda.Stream()
|
|
with torch.cuda.stream(stream):
|
|
# Keep the first nonblocking H2D copy pending while the host submits the
|
|
# second call. Each call must stage from its own pinned source so the
|
|
# first copy is not corrupted before it runs.
|
|
torch.cuda._sleep(10_000_000)
|
|
zeroer.zero_block_ids([1])
|
|
zeroer.zero_block_ids([2])
|
|
stream.synchronize()
|
|
|
|
assert torch.all(storage[0] == 1)
|
|
assert torch.all(storage[1] == 0)
|
|
assert torch.all(storage[2] == 0)
|
|
assert torch.all(storage[3] == 1)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
def test_non_uniform_page_sizes():
|
|
"""Two segments with different page sizes (e.g. MLA + DSA indexer)."""
|
|
device = torch.device("cuda")
|
|
num_blocks = 4
|
|
page_size_a = 10496 # int32 elements
|
|
page_size_b = 2112
|
|
|
|
storage_a = torch.ones((num_blocks, page_size_a), dtype=torch.int32, device=device)
|
|
storage_b = torch.ones((num_blocks, page_size_b), dtype=torch.int32, device=device)
|
|
|
|
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
|
|
zeroer.device = device
|
|
|
|
seg_page_sizes = [page_size_a, page_size_b]
|
|
max_ps = max(seg_page_sizes)
|
|
|
|
def largest_power_of_2_divisor(n):
|
|
return n & -n
|
|
|
|
blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024)
|
|
|
|
zeroer._meta = (
|
|
torch.tensor(
|
|
[storage_a.data_ptr(), storage_b.data_ptr()],
|
|
dtype=torch.uint64,
|
|
device=device,
|
|
),
|
|
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
|
|
max_ps // blk_size,
|
|
blk_size,
|
|
2,
|
|
)
|
|
|
|
stream = torch.cuda.Stream()
|
|
with torch.cuda.stream(stream):
|
|
zeroer.zero_block_ids([1, 2])
|
|
stream.synchronize()
|
|
|
|
for storage in (storage_a, storage_b):
|
|
assert torch.all(storage[0] == 1)
|
|
assert torch.all(storage[1] == 0)
|
|
assert torch.all(storage[2] == 0)
|
|
assert torch.all(storage[3] == 1)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
def test_warmup_compiles_every_n_blocks_specialization():
|
|
"""After warmup, no launch should trigger a first-request JIT compile.
|
|
|
|
``n_blocks`` is ``do_not_specialize``, so a single warmup launch must
|
|
cover every block count.
|
|
"""
|
|
device = torch.device("cuda")
|
|
num_blocks = 64
|
|
page_size_el = 4
|
|
storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device)
|
|
|
|
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
|
|
zeroer.device = device
|
|
zeroer._meta = (
|
|
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
|
|
torch.tensor([page_size_el], dtype=torch.int64, device=device),
|
|
1, # max_chunks
|
|
page_size_el, # blk_size
|
|
1, # n_segs
|
|
)
|
|
|
|
def compiled_variants() -> set:
|
|
return {
|
|
key
|
|
for caches in _zero_kv_blocks_kernel.device_caches.values()
|
|
for key in caches[0]
|
|
}
|
|
|
|
zeroer.warmup(num_blocks)
|
|
torch.accelerator.synchronize()
|
|
warmed = compiled_variants()
|
|
assert warmed
|
|
|
|
for n_blocks in (1, 2, 3, 16, 32):
|
|
zeroer.zero_block_ids(list(range(n_blocks)))
|
|
torch.accelerator.synchronize()
|
|
|
|
assert compiled_variants() == warmed
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
def test_warmup_respects_available_block_count():
|
|
"""An empty KV cache must not be warmed with out-of-range block IDs."""
|
|
device = torch.device("cuda")
|
|
page_size_el = 4
|
|
storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device)
|
|
|
|
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
|
|
zeroer.device = device
|
|
zeroer._meta = (
|
|
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
|
|
torch.tensor([page_size_el], dtype=torch.int64, device=device),
|
|
1,
|
|
page_size_el,
|
|
1,
|
|
)
|
|
|
|
zeroer.warmup(0)
|
|
torch.accelerator.synchronize()
|
|
|
|
assert torch.all(storage == 1)
|