wehub-resource-sync
378 lines
12 KiB
Python
378 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Unit tests for the LMCache-side TRT-LLM integration components.
|
|
|
|
Covers ``EngineType.TRTLLM``, the new ``NB_NL_TWO_NH_BS_HS`` GPU KV
|
|
format, format-aware accessors in ``gpu_connector/utils.py``, the
|
|
``TRTLLMGPUConnector`` construction path, and the
|
|
``RawCudaIPCWrapper`` subclass relationship — without requiring
|
|
TensorRT-LLM to be installed.
|
|
"""
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.utils import EngineType
|
|
|
|
|
|
def _has_lmc_ops() -> bool:
|
|
try:
|
|
# First Party
|
|
import lmcache.c_ops # noqa: F401
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def _has_cuda() -> bool:
|
|
return torch.cuda.is_available()
|
|
|
|
|
|
class TestEngineType:
|
|
def test_trtllm_exists(self) -> None:
|
|
assert hasattr(EngineType, "TRTLLM")
|
|
assert EngineType.TRTLLM.value == "trtllm"
|
|
|
|
def test_all_engine_types(self) -> None:
|
|
expected = {"vllm", "sglang", "trtllm", "mock"}
|
|
actual = {e.value for e in EngineType}
|
|
assert expected == actual
|
|
|
|
|
|
class TestAssertContiguous:
|
|
def test_contiguous_tensor_passes(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import assert_contiguous
|
|
|
|
t = torch.zeros(2, 3, 4)
|
|
assert_contiguous(t) # no raise
|
|
|
|
def test_non_contiguous_raises(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import assert_contiguous
|
|
|
|
t = torch.zeros(2, 3, 4).transpose(0, 1)
|
|
with pytest.raises(ValueError, match="not contiguous"):
|
|
assert_contiguous(t)
|
|
|
|
def test_nonzero_storage_offset_raises(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import assert_contiguous
|
|
|
|
base = torch.zeros(8)
|
|
view = base[2:6] # storage_offset = 2
|
|
with pytest.raises(ValueError, match="storage_offset"):
|
|
assert_contiguous(view)
|
|
|
|
|
|
@pytest.mark.skipif(not _has_lmc_ops(), reason="lmcache C ops not built")
|
|
class TestGPUKVFormatEnum:
|
|
def test_nb_nl_two_nh_bs_hs_exists(self) -> None:
|
|
# First Party
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
assert hasattr(lmc_ops.EngineKVFormat, "NB_NL_TWO_NH_BS_HS")
|
|
|
|
|
|
@pytest.mark.skipif(not _has_lmc_ops(), reason="lmcache C ops not built")
|
|
class TestNormalizeTRTLLM:
|
|
"""``normalize_kv_and_discover_format`` for ``EngineType.TRTLLM``."""
|
|
|
|
def test_4d_tensor_reshape_to_6d(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import (
|
|
LayoutHints,
|
|
normalize_kv_and_discover_format,
|
|
)
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
nb, nl, kv = 4, 2, 2
|
|
nh, bs, hs = 8, 16, 64
|
|
flat = nh * bs * hs
|
|
|
|
t = torch.zeros(nb, nl, kv, flat, dtype=torch.bfloat16)
|
|
layout_hints: LayoutHints = {
|
|
"kv_layout": "HND",
|
|
"num_kv_heads": nh,
|
|
"tokens_per_block": bs,
|
|
"head_dim": hs,
|
|
}
|
|
|
|
fmt, normalized = normalize_kv_and_discover_format(
|
|
t, EngineType.TRTLLM, layout_hints=layout_hints
|
|
)
|
|
|
|
assert fmt == lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS
|
|
assert isinstance(normalized, torch.Tensor)
|
|
assert tuple(normalized.shape) == (nb, nl, kv, nh, bs, hs)
|
|
|
|
def test_collapses_one_element_list(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import (
|
|
LayoutHints,
|
|
normalize_kv_and_discover_format,
|
|
)
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
nb, nl, kv, nh, bs, hs = 2, 2, 2, 4, 8, 32
|
|
t = torch.zeros(nb, nl, kv, nh * bs * hs)
|
|
layout_hints: LayoutHints = {
|
|
"kv_layout": "HND",
|
|
"num_kv_heads": nh,
|
|
"tokens_per_block": bs,
|
|
"head_dim": hs,
|
|
}
|
|
fmt, normalized = normalize_kv_and_discover_format(
|
|
[t], EngineType.TRTLLM, layout_hints=layout_hints
|
|
)
|
|
assert fmt == lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS
|
|
assert isinstance(normalized, torch.Tensor)
|
|
assert normalized.shape == (nb, nl, kv, nh, bs, hs)
|
|
|
|
def test_missing_layout_hints_raises(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import normalize_kv_and_discover_format
|
|
|
|
t = torch.zeros(2, 2, 2, 16)
|
|
with pytest.raises(ValueError, match="num_kv_heads"):
|
|
normalize_kv_and_discover_format(t, EngineType.TRTLLM, layout_hints={})
|
|
|
|
def test_flat_dim_mismatch_raises(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import (
|
|
LayoutHints,
|
|
normalize_kv_and_discover_format,
|
|
)
|
|
|
|
t = torch.zeros(2, 2, 2, 17) # not divisible by 4*2*2 = 16
|
|
layout_hints: LayoutHints = {
|
|
"kv_layout": "HND",
|
|
"num_kv_heads": 4,
|
|
"tokens_per_block": 2,
|
|
"head_dim": 2,
|
|
}
|
|
with pytest.raises(ValueError, match="flat dim"):
|
|
normalize_kv_and_discover_format(
|
|
t, EngineType.TRTLLM, layout_hints=layout_hints
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not _has_lmc_ops(), reason="lmcache C ops not built")
|
|
class TestAccessorsTRTLLM:
|
|
"""Format accessors for ``NB_NL_TWO_NH_BS_HS``."""
|
|
|
|
def _tensor(
|
|
self,
|
|
nb: int = 4,
|
|
nl: int = 3,
|
|
kv: int = 2,
|
|
nh: int = 8,
|
|
bs: int = 16,
|
|
hs: int = 64,
|
|
) -> torch.Tensor:
|
|
return torch.zeros(nb, nl, kv, nh, bs, hs, dtype=torch.bfloat16)
|
|
|
|
def test_get_num_layers(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_num_layers
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert get_num_layers(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 3
|
|
|
|
def test_get_num_blocks(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_num_blocks
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert get_num_blocks(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 4
|
|
|
|
def test_get_block_size(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_block_size
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert get_block_size(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 16
|
|
|
|
def test_get_num_heads(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_num_heads
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert get_num_heads(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 8
|
|
|
|
def test_get_head_size(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_head_size
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert get_head_size(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 64
|
|
|
|
def test_get_hidden_dim_size(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_hidden_dim_size
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert (
|
|
get_hidden_dim_size(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 8 * 64
|
|
)
|
|
|
|
def test_get_page_buffer_size(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_page_buffer_size
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert (
|
|
get_page_buffer_size(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == 4 * 16
|
|
)
|
|
|
|
def test_get_dtype(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_dtype
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
assert get_dtype(t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS) == torch.bfloat16
|
|
|
|
def test_get_group_data_ptrs_returns_single_base_pointer(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import get_group_data_ptrs
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
t = self._tensor()
|
|
ptrs = get_group_data_ptrs(
|
|
t, lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS, list(range(3))
|
|
)
|
|
assert len(ptrs) == 1
|
|
assert ptrs[0] == t.data_ptr()
|
|
|
|
def test_shape_description_strings(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import (
|
|
get_attention_backend,
|
|
get_concrete_engine_kv_shape,
|
|
get_engine_kv_shape_description,
|
|
)
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
fmt = lmc_ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS
|
|
assert get_engine_kv_shape_description(fmt) == "[NB, NL, 2, NH, BS, HS]"
|
|
assert "TRT-LLM" in get_attention_backend(fmt)
|
|
assert get_concrete_engine_kv_shape(self._tensor(), fmt) == (
|
|
"[4, 3, 2, 8, 16, 64]"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not _has_cuda(), reason="CUDA required for connector init")
|
|
@pytest.mark.skipif(not _has_lmc_ops(), reason="lmcache C ops not built")
|
|
class TestTRTLLMGPUConnector:
|
|
def test_construct(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import TRTLLMGPUConnector
|
|
|
|
device = torch.device("cuda:0")
|
|
c = TRTLLMGPUConnector(
|
|
num_kv_heads=2,
|
|
head_dim=64,
|
|
hidden_dim_size=128,
|
|
num_layers=4,
|
|
chunk_size=256,
|
|
dtype=torch.bfloat16,
|
|
device=device,
|
|
)
|
|
assert c.num_kv_heads == 2
|
|
assert c.head_dim == 64
|
|
assert c.kv_cache_tensor is None # not registered yet
|
|
|
|
def test_get_shape(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import TRTLLMGPUConnector
|
|
|
|
device = torch.device("cuda:0")
|
|
c = TRTLLMGPUConnector(
|
|
num_kv_heads=2,
|
|
head_dim=64,
|
|
hidden_dim_size=128,
|
|
num_layers=4,
|
|
chunk_size=256,
|
|
dtype=torch.bfloat16,
|
|
device=device,
|
|
)
|
|
# Memory-obj layout: [2, num_layers, num_tokens, hidden_dim]
|
|
assert tuple(c.get_shape(256)) == (2, 4, 256, 128)
|
|
|
|
def test_register_kv_caches_with_4d_tensor(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.gpu_connectors import TRTLLMGPUConnector
|
|
|
|
device = torch.device("cuda:0")
|
|
nh, bs, hs = 2, 16, 64
|
|
nb, nl, kv = 4, 4, 2
|
|
flat = nh * bs * hs
|
|
c = TRTLLMGPUConnector(
|
|
num_kv_heads=nh,
|
|
head_dim=hs,
|
|
hidden_dim_size=nh * hs,
|
|
num_layers=nl,
|
|
chunk_size=256,
|
|
dtype=torch.bfloat16,
|
|
device=device,
|
|
)
|
|
t = torch.zeros(nb, nl, kv, flat, dtype=torch.bfloat16, device=device)
|
|
c.register_kv_caches(t)
|
|
assert c.kv_cache_tensor is not None
|
|
assert tuple(c.kv_cache_tensor.shape) == (nb, nl, kv, nh, bs, hs)
|
|
assert c.tokens_per_block == bs
|
|
assert c.blocks_per_chunk == 256 // bs
|
|
assert c.shape_desc is not None
|
|
assert c.shape_desc.nl == nl
|
|
assert c.shape_desc.nb == nb
|
|
assert c.shape_desc.bs == bs
|
|
assert c.shape_desc.nh == nh
|
|
assert c.shape_desc.hs == hs
|
|
|
|
|
|
class TestRawCudaIPCWrapperType:
|
|
"""Subclass relationship — without invoking CUDA IPC syscalls."""
|
|
|
|
def test_is_subclass(self) -> None:
|
|
# First Party
|
|
from lmcache.v1.multiprocess.custom_types import (
|
|
DeviceIPCWrapper,
|
|
)
|
|
from lmcache.v1.platform.cuda.ipc_wrapper import RawCudaIPCWrapper
|
|
|
|
assert issubclass(RawCudaIPCWrapper, DeviceIPCWrapper)
|
|
|
|
def test_kvcache_typing_unchanged(self) -> None:
|
|
"""``KVCache = list[DeviceIPCWrapper]`` should accept concrete wrappers
|
|
— load-bearing for msgspec ext-code reuse over ZMQ.
|
|
"""
|
|
# First Party
|
|
from lmcache.v1.multiprocess.custom_types import (
|
|
DeviceIPCWrapper,
|
|
KVCache,
|
|
)
|
|
from lmcache.v1.platform.cuda.ipc_wrapper import RawCudaIPCWrapper
|
|
|
|
assert KVCache == list[DeviceIPCWrapper]
|
|
# Static check substitute: a concrete wrapper fits the list type.
|
|
assert issubclass(RawCudaIPCWrapper, DeviceIPCWrapper)
|
|
|
|
def test_serde_uses_shared_ext_code(self) -> None:
|
|
"""Ext code 1 dispatches via ``DeviceIPCWrapper`` for all wrappers."""
|
|
# First Party
|
|
from lmcache.v1.multiprocess import custom_types
|
|
|
|
registry = custom_types._CUSTOMERIZED_SERIALIZERS # noqa: PLC2701
|
|
assert custom_types.DeviceIPCWrapper in registry
|
|
assert registry[custom_types.DeviceIPCWrapper].code == 1
|