wehub-resource-sync
45 lines
1.5 KiB
Python
45 lines
1.5 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
|
|
|
|
|
|
@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 ID-buffer
|
|
# lifetime without constructing model attention groups.
|
|
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
|
|
zeroer.device = device
|
|
zeroer.pin_memory = True
|
|
zeroer.max_concurrency = 2
|
|
zeroer._id_cap = 8
|
|
zeroer._allocate_id_buffers()
|
|
zeroer._meta = (
|
|
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
|
|
page_size_el,
|
|
page_size_el,
|
|
1,
|
|
)
|
|
|
|
stream = torch.cuda.Stream()
|
|
with torch.cuda.stream(stream):
|
|
# Keep the first nonblocking H2D copy pending while the host submits the
|
|
# second call. A single shared pinned source would be overwritten here.
|
|
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)
|