wehub-resource-sync
472 lines
14 KiB
Python
472 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# Standard
|
|
import random
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
pytest.importorskip(
|
|
"lmcache.c_ops",
|
|
reason="Requires CUDA extension lmcache.c_ops",
|
|
)
|
|
|
|
# First Party
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
# Skip all tests if cuda is unavailable
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available() or torch.cuda.device_count() == 0,
|
|
reason="No CUDA GPU present",
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tensor factories (ported from kernel harness)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _create_random_tensor(
|
|
shape: list, dtype: torch.dtype, device: torch.device
|
|
) -> torch.Tensor:
|
|
if dtype == torch.float8_e4m3fn:
|
|
return torch.rand(shape, dtype=torch.bfloat16, device=device).to(dtype)
|
|
return torch.rand(shape, dtype=dtype, device=device)
|
|
|
|
|
|
def _create_zero_tensor(
|
|
shape: list, dtype: torch.dtype, device: torch.device
|
|
) -> torch.Tensor:
|
|
if dtype == torch.float8_e4m3fn:
|
|
return torch.zeros(shape, dtype=torch.bfloat16, device=device).to(dtype)
|
|
return torch.zeros(shape, dtype=dtype, device=device)
|
|
|
|
|
|
# Format enum values from c_ops
|
|
FMT_NORMAL = lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
|
|
FMT_CROSS_LAYER = lmc_ops.EngineKVFormat.NB_NL_TWO_BS_NH_HS
|
|
FMT_FLASH_INFER = lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS
|
|
FMT_MLA = lmc_ops.EngineKVFormat.NL_X_NB_BS_HS
|
|
FMT_SGLANG_MHA = lmc_ops.EngineKVFormat.TWO_X_NL_X_NBBS_NH_HS
|
|
FMT_SGLANG_MLA = lmc_ops.EngineKVFormat.NL_X_NBBS_ONE_HS
|
|
FMT_NORMAL_HND = lmc_ops.EngineKVFormat.NL_X_TWO_NB_NH_BS_HS
|
|
FMT_FLASH_INFER_HND = lmc_ops.EngineKVFormat.NL_X_NB_TWO_NH_BS_HS
|
|
|
|
# Format parameters: (engine_kv_format, num_layers, num_heads, head_size, is_mla)
|
|
# Use small layer counts to keep GPU memory usage low in CI
|
|
FORMAT_PARAMS = [
|
|
(FMT_NORMAL, 4, 8, 128, False),
|
|
(FMT_CROSS_LAYER, 4, 8, 128, False),
|
|
(FMT_FLASH_INFER, 4, 8, 128, False),
|
|
(FMT_MLA, 4, 1, 576, True),
|
|
(FMT_SGLANG_MHA, 4, 8, 128, False),
|
|
(FMT_SGLANG_MLA, 4, 1, 576, True),
|
|
(FMT_NORMAL_HND, 4, 8, 128, False),
|
|
(FMT_FLASH_INFER_HND, 4, 8, 128, False),
|
|
]
|
|
|
|
|
|
def create_vllm_tensors(
|
|
engine_kv_format,
|
|
nl: int,
|
|
nb: int,
|
|
bs: int,
|
|
nh: int,
|
|
hs: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> list[torch.Tensor]:
|
|
nbbs = nb * bs
|
|
if engine_kv_format == FMT_NORMAL:
|
|
shape = [2, nb, bs, nh, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_NORMAL_HND:
|
|
shape = [2, nb, nh, bs, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_CROSS_LAYER:
|
|
shape = [nb, nl, 2, bs, nh, hs]
|
|
return [_create_random_tensor(shape, dtype, device)]
|
|
elif engine_kv_format == FMT_FLASH_INFER:
|
|
shape = [nb, 2, bs, nh, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_FLASH_INFER_HND:
|
|
shape = [nb, 2, nh, bs, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_MLA:
|
|
shape = [nb, bs, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_SGLANG_MHA:
|
|
shape = [nbbs, nh, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(2 * nl)]
|
|
elif engine_kv_format == FMT_SGLANG_MLA:
|
|
shape = [nbbs, 1, hs]
|
|
return [_create_random_tensor(shape, dtype, device) for _ in range(nl)]
|
|
raise ValueError(f"Unknown format: {engine_kv_format}")
|
|
|
|
|
|
def create_zero_vllm_tensors(
|
|
engine_kv_format,
|
|
nl: int,
|
|
nb: int,
|
|
bs: int,
|
|
nh: int,
|
|
hs: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> list[torch.Tensor]:
|
|
nbbs = nb * bs
|
|
if engine_kv_format == FMT_NORMAL:
|
|
shape = [2, nb, bs, nh, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_NORMAL_HND:
|
|
shape = [2, nb, nh, bs, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_CROSS_LAYER:
|
|
shape = [nb, nl, 2, bs, nh, hs]
|
|
return [_create_zero_tensor(shape, dtype, device)]
|
|
elif engine_kv_format == FMT_FLASH_INFER:
|
|
shape = [nb, 2, bs, nh, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_FLASH_INFER_HND:
|
|
shape = [nb, 2, nh, bs, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_MLA:
|
|
shape = [nb, bs, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(nl)]
|
|
elif engine_kv_format == FMT_SGLANG_MHA:
|
|
shape = [nbbs, nh, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(2 * nl)]
|
|
elif engine_kv_format == FMT_SGLANG_MLA:
|
|
shape = [nbbs, 1, hs]
|
|
return [_create_zero_tensor(shape, dtype, device) for _ in range(nl)]
|
|
raise ValueError(f"Unknown format: {engine_kv_format}")
|
|
|
|
|
|
def create_memory_objects(
|
|
kv_dim: int,
|
|
nl: int,
|
|
tokens_per_object: int,
|
|
hidden_dim: int,
|
|
num_objects: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> list[torch.Tensor]:
|
|
shape = [kv_dim, nl, tokens_per_object, hidden_dim]
|
|
objects = []
|
|
for _ in range(num_objects):
|
|
if dtype == torch.float8_e4m3fn:
|
|
t = torch.zeros(shape, dtype=torch.bfloat16, device=device)
|
|
t = t.to(dtype)
|
|
else:
|
|
t = torch.zeros(shape, dtype=dtype, device=device)
|
|
if device.type == "cpu":
|
|
t = t.pin_memory()
|
|
objects.append(t)
|
|
return objects
|
|
|
|
|
|
def get_block_data(
|
|
vllm_tensors: list[torch.Tensor],
|
|
engine_kv_format,
|
|
nl: int,
|
|
bs: int,
|
|
nh: int,
|
|
block_idx: int,
|
|
) -> list[torch.Tensor]:
|
|
"""Extract all layer data for a given block."""
|
|
results = []
|
|
for layer_idx in range(nl):
|
|
if engine_kv_format == FMT_NORMAL:
|
|
results.append(vllm_tensors[layer_idx][:, block_idx, :, :, :].clone())
|
|
elif engine_kv_format == FMT_NORMAL_HND:
|
|
results.append(vllm_tensors[layer_idx][:, block_idx, :, :, :].clone())
|
|
elif engine_kv_format == FMT_CROSS_LAYER:
|
|
results.append(vllm_tensors[0][block_idx, layer_idx, :, :, :, :].clone())
|
|
elif engine_kv_format == FMT_FLASH_INFER:
|
|
results.append(vllm_tensors[layer_idx][block_idx, :, :, :, :].clone())
|
|
elif engine_kv_format == FMT_FLASH_INFER_HND:
|
|
results.append(vllm_tensors[layer_idx][block_idx, :, :, :, :].clone())
|
|
elif engine_kv_format == FMT_MLA:
|
|
results.append(vllm_tensors[layer_idx][block_idx, :, :].clone())
|
|
elif engine_kv_format == FMT_SGLANG_MHA:
|
|
ts, ed = block_idx * bs, (block_idx + 1) * bs
|
|
k = vllm_tensors[layer_idx][ts:ed, :, :].clone()
|
|
v = vllm_tensors[nl + layer_idx][ts:ed, :, :].clone()
|
|
results.append(torch.stack([k, v], dim=0))
|
|
elif engine_kv_format == FMT_SGLANG_MLA:
|
|
ts, ed = block_idx * bs, (block_idx + 1) * bs
|
|
results.append(vllm_tensors[layer_idx][ts:ed, 0, :].clone())
|
|
return results
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Kernel call helper
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def call_block_kernel(
|
|
vllm_tensors: list[torch.Tensor],
|
|
mem_objects: list[torch.Tensor],
|
|
block_ids: list[int],
|
|
engine_kv_format,
|
|
direction,
|
|
nl: int,
|
|
nb: int,
|
|
bs: int,
|
|
nh: int,
|
|
hs: int,
|
|
is_mla: bool,
|
|
tokens_per_object: int,
|
|
skip_prefix_n_blocks: int = 0,
|
|
) -> None:
|
|
device = vllm_tensors[0].device
|
|
|
|
shape_desc = lmc_ops.PageBufferShapeDesc()
|
|
shape_desc.kv_size = 1 if is_mla else 2
|
|
shape_desc.nl = nl
|
|
shape_desc.nb = nb
|
|
shape_desc.bs = bs
|
|
shape_desc.nh = nh
|
|
shape_desc.hs = hs
|
|
shape_desc.element_size = vllm_tensors[0].element_size()
|
|
|
|
ptrs = [t.data_ptr() for t in vllm_tensors]
|
|
paged_buffer_ptrs_tensor = torch.tensor(ptrs, dtype=torch.int64, device=device)
|
|
lmcache_objects_ptrs = [m.data_ptr() for m in mem_objects]
|
|
|
|
block_ids_gpu = torch.tensor(block_ids, dtype=torch.int64, device=device)
|
|
lmc_ops.multi_layer_block_kv_transfer(
|
|
paged_buffer_ptrs_tensor,
|
|
lmcache_objects_ptrs,
|
|
block_ids_gpu,
|
|
device,
|
|
direction,
|
|
shape_desc,
|
|
tokens_per_object,
|
|
engine_kv_format,
|
|
skip_prefix_n_blocks,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NB = 1000
|
|
BS = 16
|
|
NUM_MEMORY_OBJECTS = 4
|
|
TOKENS_PER_OBJECT = 256
|
|
BLOCKS_PER_OBJECT = TOKENS_PER_OBJECT // BS # 16
|
|
TOTAL_BLOCKS = NUM_MEMORY_OBJECTS * BLOCKS_PER_OBJECT # 64
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"engine_kv_format,nl,nh,hs,is_mla",
|
|
FORMAT_PARAMS,
|
|
ids=[
|
|
"normal",
|
|
"cross_layer",
|
|
"flash_infer",
|
|
"mla",
|
|
"sglang_mha",
|
|
"sglang_mla",
|
|
"normal_hnd",
|
|
"flash_infer_hnd",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]
|
|
)
|
|
@pytest.mark.parametrize("mem_device", ["cuda", "cpu"], ids=["mem_gpu", "mem_cpu"])
|
|
def test_block_transfer_roundtrip(
|
|
engine_kv_format, nl, nh, hs, is_mla, dtype, mem_device
|
|
):
|
|
"""
|
|
D2H -> H2D roundtrip with different block IDs proves data flows through
|
|
memory objects.
|
|
"""
|
|
device = torch.device("cuda")
|
|
mem_dev = torch.device(mem_device)
|
|
kv_dim = 1 if is_mla else 2
|
|
hidden_dim = nh * hs
|
|
|
|
# Create tensors
|
|
source_vllm = create_vllm_tensors(
|
|
engine_kv_format, nl, NB, BS, nh, hs, dtype, device
|
|
)
|
|
target_vllm = create_zero_vllm_tensors(
|
|
engine_kv_format, nl, NB, BS, nh, hs, dtype, device
|
|
)
|
|
mem_objects = create_memory_objects(
|
|
kv_dim,
|
|
nl,
|
|
TOKENS_PER_OBJECT,
|
|
hidden_dim,
|
|
NUM_MEMORY_OBJECTS,
|
|
dtype,
|
|
mem_dev,
|
|
)
|
|
|
|
# Disjoint block IDs for D2H and H2D
|
|
rng_d2h = random.Random(42)
|
|
block_ids_d2h = rng_d2h.sample(range(NB), TOTAL_BLOCKS)
|
|
excluded = set(block_ids_d2h)
|
|
available = [i for i in range(NB) if i not in excluded]
|
|
rng_h2d = random.Random(123)
|
|
block_ids_h2d = rng_h2d.sample(available, TOTAL_BLOCKS)
|
|
|
|
# D2H: source -> mem_objects
|
|
call_block_kernel(
|
|
source_vllm,
|
|
mem_objects,
|
|
block_ids_d2h,
|
|
engine_kv_format,
|
|
lmc_ops.TransferDirection.D2H,
|
|
nl,
|
|
NB,
|
|
BS,
|
|
nh,
|
|
hs,
|
|
is_mla,
|
|
TOKENS_PER_OBJECT,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# H2D: mem_objects -> target
|
|
call_block_kernel(
|
|
target_vllm,
|
|
mem_objects,
|
|
block_ids_h2d,
|
|
engine_kv_format,
|
|
lmc_ops.TransferDirection.H2D,
|
|
nl,
|
|
NB,
|
|
BS,
|
|
nh,
|
|
hs,
|
|
is_mla,
|
|
TOKENS_PER_OBJECT,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# Verify: target[h2d_block] == source[d2h_block]
|
|
for i in range(TOTAL_BLOCKS):
|
|
src_data = get_block_data(
|
|
source_vllm, engine_kv_format, nl, BS, nh, block_ids_d2h[i]
|
|
)
|
|
tgt_data = get_block_data(
|
|
target_vllm, engine_kv_format, nl, BS, nh, block_ids_h2d[i]
|
|
)
|
|
for layer_idx in range(nl):
|
|
assert torch.equal(src_data[layer_idx], tgt_data[layer_idx]), (
|
|
f"Mismatch at block index {i}, layer {layer_idx}"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"engine_kv_format,nl,nh,hs,is_mla",
|
|
FORMAT_PARAMS,
|
|
ids=[
|
|
"normal",
|
|
"cross_layer",
|
|
"flash_infer",
|
|
"mla",
|
|
"sglang_mha",
|
|
"sglang_mla",
|
|
"normal_hnd",
|
|
"flash_infer_hnd",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"])
|
|
def test_block_transfer_skip_prefix(engine_kv_format, nl, nh, hs, is_mla, dtype):
|
|
"""Verify skip_prefix_n_blocks=4 skips the first 4 blocks globally."""
|
|
device = torch.device("cuda")
|
|
kv_dim = 1 if is_mla else 2
|
|
hidden_dim = nh * hs
|
|
skip = 4
|
|
|
|
source_vllm = create_vllm_tensors(
|
|
engine_kv_format, nl, NB, BS, nh, hs, dtype, device
|
|
)
|
|
target_vllm = create_zero_vllm_tensors(
|
|
engine_kv_format, nl, NB, BS, nh, hs, dtype, device
|
|
)
|
|
mem_objects = create_memory_objects(
|
|
kv_dim,
|
|
nl,
|
|
TOKENS_PER_OBJECT,
|
|
hidden_dim,
|
|
NUM_MEMORY_OBJECTS,
|
|
dtype,
|
|
device,
|
|
)
|
|
|
|
rng_d2h = random.Random(42)
|
|
block_ids_d2h = rng_d2h.sample(range(NB), TOTAL_BLOCKS)
|
|
excluded = set(block_ids_d2h)
|
|
available = [i for i in range(NB) if i not in excluded]
|
|
rng_h2d = random.Random(123)
|
|
block_ids_h2d = rng_h2d.sample(available, TOTAL_BLOCKS)
|
|
|
|
# D2H with skip
|
|
call_block_kernel(
|
|
source_vllm,
|
|
mem_objects,
|
|
block_ids_d2h,
|
|
engine_kv_format,
|
|
lmc_ops.TransferDirection.D2H,
|
|
nl,
|
|
NB,
|
|
BS,
|
|
nh,
|
|
hs,
|
|
is_mla,
|
|
TOKENS_PER_OBJECT,
|
|
skip_prefix_n_blocks=skip,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# H2D with skip
|
|
call_block_kernel(
|
|
target_vllm,
|
|
mem_objects,
|
|
block_ids_h2d,
|
|
engine_kv_format,
|
|
lmc_ops.TransferDirection.H2D,
|
|
nl,
|
|
NB,
|
|
BS,
|
|
nh,
|
|
hs,
|
|
is_mla,
|
|
TOKENS_PER_OBJECT,
|
|
skip_prefix_n_blocks=skip,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# Non-skipped blocks should match
|
|
for i in range(skip, TOTAL_BLOCKS):
|
|
src_data = get_block_data(
|
|
source_vllm, engine_kv_format, nl, BS, nh, block_ids_d2h[i]
|
|
)
|
|
tgt_data = get_block_data(
|
|
target_vllm, engine_kv_format, nl, BS, nh, block_ids_h2d[i]
|
|
)
|
|
for layer_idx in range(nl):
|
|
assert torch.equal(src_data[layer_idx], tgt_data[layer_idx]), (
|
|
f"Mismatch at block index {i}, layer {layer_idx}"
|
|
)
|
|
|
|
# Skipped blocks in target should remain zero
|
|
for i in range(skip):
|
|
tgt_data = get_block_data(
|
|
target_vllm, engine_kv_format, nl, BS, nh, block_ids_h2d[i]
|
|
)
|
|
for layer_idx in range(nl):
|
|
block = tgt_data[layer_idx]
|
|
if block.dtype == torch.float8_e4m3fn:
|
|
block = block.to(torch.float32)
|
|
assert block.abs().sum().item() == 0, (
|
|
f"Skipped block {i}, layer {layer_idx} is not zero"
|
|
)
|