wehub-resource-sync
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache import torch_device_type
|
|
from lmcache.storage_backend.serde.cachegen_basics import CacheGenEncoderOutput
|
|
from lmcache.storage_backend.serde.cachegen_decoder import CacheGenDeserializer
|
|
from lmcache.storage_backend.serde.cachegen_encoder import CacheGenSerializer
|
|
from lmcache.v1.config import LMCacheEngineConfig
|
|
from lmcache.v1.metadata import LMCacheMetadata
|
|
|
|
# CacheGen has hand-written CUDA and SYCL kernels only; gate the test on
|
|
# those backends explicitly (whitelist), so HPU / CPU-only CI / future
|
|
# backends without a CacheGen kernel are skipped automatically.
|
|
# torch_device_type is set to "cuda"/"xpu" only after is_available() passes
|
|
# in lmcache.__init__, so no extra availability check is needed here.
|
|
_GPU_AVAILABLE = torch_device_type == "cuda" or torch_device_type == "xpu"
|
|
|
|
|
|
def generate_kv_cache(num_tokens, device):
|
|
ret = []
|
|
num_layers = 32
|
|
num_heads = 8
|
|
head_size = 128
|
|
shape = [num_tokens, num_heads, head_size]
|
|
dtype = torch.bfloat16
|
|
|
|
for i in range(num_layers):
|
|
k = torch.rand(shape, dtype=dtype, device=device)
|
|
v = torch.rand(shape, dtype=dtype, device=device)
|
|
ret.append((k, v))
|
|
|
|
return tuple(ret)
|
|
|
|
|
|
def to_blob(kv_tuples):
|
|
return torch.stack(
|
|
[torch.stack(inner_tuple, dim=0) for inner_tuple in kv_tuples], dim=0
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("chunk_size", [16, 128, 256])
|
|
@pytest.mark.skipif(
|
|
not _GPU_AVAILABLE,
|
|
reason="No GPU backend (CUDA or XPU) available",
|
|
)
|
|
def test_cachegen_encoder(chunk_size):
|
|
config = LMCacheEngineConfig.from_defaults(chunk_size=chunk_size)
|
|
metadata = LMCacheMetadata(
|
|
model_name="mistralai/Mistral-7B-Instruct-v0.2",
|
|
world_size=1,
|
|
local_world_size=1,
|
|
worker_id=0,
|
|
local_worker_id=0,
|
|
kv_dtype=torch.bfloat16,
|
|
kv_shape=None,
|
|
)
|
|
# metadata2 = LMCacheMetadata(
|
|
# model_name="mistralai/Mistral-7B-Instruct-v0.2",
|
|
# world_size=1,
|
|
# local_world_size=1,
|
|
# worker_id=0,
|
|
# local_worker_id=0,
|
|
# kv_dtype=torch.bfloat16,
|
|
# kv_shape=None,
|
|
# )
|
|
serializer = CacheGenSerializer(config, metadata)
|
|
# serializer2 = CacheGenSerializer(config, metadata2)
|
|
|
|
kv = to_blob(generate_kv_cache(chunk_size, torch_device_type))
|
|
output = serializer.to_bytes(kv)
|
|
# huggingface:
|
|
# kv2 = kv.permute([0, 1, 3, 2, 4])
|
|
# output2 = serializer2.to_bytes(kv2)
|
|
|
|
# assert abs(len(output) - len(output2)) < 10
|
|
output_dict = CacheGenEncoderOutput.from_bytes(output)
|
|
assert output_dict.num_heads == 8
|
|
assert output_dict.head_size == 128
|
|
|
|
|
|
@pytest.mark.parametrize("chunk_size", [16, 128, 256])
|
|
@pytest.mark.skipif(
|
|
not _GPU_AVAILABLE,
|
|
reason="No GPU backend (CUDA or XPU) available",
|
|
)
|
|
def test_cachegen_decoder(chunk_size):
|
|
config = LMCacheEngineConfig.from_defaults(chunk_size=chunk_size)
|
|
metadata = LMCacheMetadata(
|
|
model_name="mistralai/Mistral-7B-Instruct-v0.2",
|
|
world_size=1,
|
|
local_world_size=1,
|
|
worker_id=0,
|
|
local_worker_id=0,
|
|
kv_dtype=torch.bfloat16,
|
|
kv_shape=None,
|
|
)
|
|
serializer = CacheGenSerializer(config, metadata)
|
|
deserializer = CacheGenDeserializer(config, metadata, torch.bfloat16)
|
|
|
|
kv = to_blob(generate_kv_cache(chunk_size, torch_device_type))
|
|
output = serializer.to_bytes(kv)
|
|
|
|
decoded_kv = deserializer.from_bytes(output)
|
|
assert decoded_kv.shape == kv.shape
|
|
assert decoded_kv.mean() != 0
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _GPU_AVAILABLE,
|
|
reason="No GPU backend (CUDA or XPU) available",
|
|
)
|
|
def test_cachegen_unmatched_size():
|
|
chunk_size = 256
|
|
config = LMCacheEngineConfig.from_defaults(chunk_size=chunk_size)
|
|
metadata = LMCacheMetadata(
|
|
model_name="mistralai/Mistral-7B-Instruct-v0.2",
|
|
world_size=1,
|
|
local_world_size=1,
|
|
worker_id=0,
|
|
local_worker_id=0,
|
|
kv_dtype=torch.bfloat16,
|
|
kv_shape=None,
|
|
)
|
|
serializer = CacheGenSerializer(config, metadata)
|
|
deserializer = CacheGenDeserializer(config, metadata, torch.bfloat16)
|
|
|
|
kv = to_blob(generate_kv_cache(chunk_size - 20, torch_device_type))
|
|
output = serializer.to_bytes(kv)
|
|
|
|
decoded_kv = deserializer.from_bytes(output)
|
|
assert decoded_kv.shape == kv.shape
|
|
assert decoded_kv.mean() != 0
|