mirror of
https://github.com/vllm-project/vllm.git
synced 2026-06-06 00:16:14 +00:00
Yan Ru Pei
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
|
|
from vllm.distributed.kv_events import BlockRemoved, BlockStored
|
|
|
|
# Minimal ExternalBlockHash for testing (bytes are a valid ExternalBlockHash).
|
|
_FAKE_HASH: bytes = b"\xab" * 32
|
|
|
|
|
|
def _make_block_stored(
|
|
group_idx: int | None = None,
|
|
kv_cache_spec_sliding_window: int | None = None,
|
|
) -> BlockStored:
|
|
return BlockStored(
|
|
block_hashes=[_FAKE_HASH],
|
|
parent_block_hash=None,
|
|
token_ids=[1, 2, 3, 4],
|
|
block_size=4,
|
|
lora_id=None,
|
|
medium="GPU",
|
|
lora_name=None,
|
|
group_idx=group_idx,
|
|
kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
|
|
)
|
|
|
|
|
|
def _make_block_removed(
|
|
group_idx: int | None = None,
|
|
) -> BlockRemoved:
|
|
return BlockRemoved(
|
|
block_hashes=[_FAKE_HASH],
|
|
medium="GPU",
|
|
group_idx=group_idx,
|
|
)
|
|
|
|
|
|
def test_block_stored_default_group_idx_is_none():
|
|
"""group_idx defaults to None when not provided."""
|
|
event = _make_block_stored()
|
|
assert event.group_idx is None
|
|
|
|
|
|
def test_block_removed_default_group_idx_is_none():
|
|
"""group_idx defaults to None when not provided."""
|
|
event = _make_block_removed()
|
|
assert event.group_idx is None
|
|
|
|
|
|
@pytest.mark.parametrize("group_idx", [1, 2, 3])
|
|
def test_block_stored_hash_differs_by_group_idx(group_idx: int):
|
|
"""BlockStored events that differ only in group_idx must hash differently."""
|
|
other_group_idx = group_idx + 1
|
|
event_a = _make_block_stored(group_idx=group_idx)
|
|
event_b = _make_block_stored(group_idx=other_group_idx)
|
|
assert hash(event_a) != hash(event_b)
|
|
|
|
|
|
def test_block_stored_hash_same_for_equal_group_idx():
|
|
"""Two BlockStored events with identical fields produce the same hash."""
|
|
event_a = _make_block_stored(group_idx=1)
|
|
event_b = _make_block_stored(group_idx=1)
|
|
assert hash(event_a) == hash(event_b)
|
|
|
|
|
|
@pytest.mark.parametrize("group_idx", [1, 2, 3])
|
|
def test_block_removed_hash_differs_by_group_idx(group_idx: int):
|
|
"""BlockRemoved events that differ only in group_idx must hash differently."""
|
|
other_group_idx = group_idx + 1
|
|
event_a = _make_block_removed(group_idx=group_idx)
|
|
event_b = _make_block_removed(group_idx=other_group_idx)
|
|
assert hash(event_a) != hash(event_b)
|
|
|
|
|
|
def test_block_removed_hash_same_for_equal_group_idx():
|
|
"""Two BlockRemoved events with identical fields produce the same hash."""
|
|
event_a = _make_block_removed(group_idx=1)
|
|
event_b = _make_block_removed(group_idx=1)
|
|
assert hash(event_a) == hash(event_b)
|
|
|
|
|
|
def test_block_stored_hash_differs_by_sliding_window():
|
|
event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)
|
|
event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)
|
|
assert hash(event_a) != hash(event_b)
|