36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import itertools
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.jit_kernel.kvcache import store_cache
|
|
|
|
BS_LIST = [2**n for n in range(0, 15)]
|
|
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
|
|
HIDDEN_DIMS = [64, 128, 256, 512, 1024, 96, 98, 100]
|
|
CACHE_SIZE = 1024 * 1024
|
|
DTYPE = torch.bfloat16
|
|
DEVICE = "cuda"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"batch_size,element_dim",
|
|
list(itertools.product(BS_LIST, HIDDEN_DIMS)),
|
|
)
|
|
def test_store_cache(batch_size: int, element_dim: int) -> None:
|
|
k = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE)
|
|
v = torch.randn((batch_size, element_dim), dtype=DTYPE, device=DEVICE)
|
|
k_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE)
|
|
v_cache = torch.randn((CACHE_SIZE, element_dim), dtype=DTYPE, device=DEVICE)
|
|
indices = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size]
|
|
|
|
# AOT store cache
|
|
store_cache(k, v, k_cache, v_cache, indices)
|
|
|
|
assert torch.all(k_cache[indices] == k)
|
|
assert torch.all(v_cache[indices] == v)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|