wehub/lmcache--lmcache
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
f213ec89765936f8d36a71674c92a605834810fc
lmcache--lmcache/tests/v1/test_python_ops_fallback.py
T
wehub-resource-sync f213ec8976 chore: import upstream snapshot with attribution
2026-07-13 12:24:33 +08:00

2868 lines
98 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# Standard
from typing import Any, Union
import ctypes
import os
import time
import unittest.mock
# Third Party
import pytest
import torch
# First Party
from lmcache.v1.multiprocess.native_completion import (
DeviceHostFuncDispatcher,
submit_callback_to_stream,
)
import lmcache.python_ops_fallback as _py_ops
# ==========================================
# 0. utils functions.
# ==========================================
# Device utility functions
def device_sync(device: str) -> None:
"""Synchronize device operations.
Args:
device: Device string ("cuda", "xpu", or "cpu")
This function synchronizes operations for GPU devices:
- For CUDA devices, calls torch.cuda.synchronize()
- For XPU devices, calls torch.xpu.synchronize()
- For CPU, no synchronization is needed (returns immediately)
"""
if device == "cuda":
torch.cuda.synchronize()
elif device == "xpu":
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.synchronize()
else:
# TODO: add more device here
pass
# CPU requires no synchronization
# ==========================================
# 1. Backend Configuration
# ==========================================
def _build_backend_params() -> list:
"""Build pytest parameter list for the backend fixture.
Returns one entry per available backend configuration:
- cuda_c_ops: uses lmcache.c_ops (requires CUDA and the CUDA extension)
- cuda_py_ops: uses lmcache.python_ops_fallback with GPU visible
- cpy_py_ops: uses lmcache.python_ops_fallback with GPU mocked away
- xpu_sycl_ops: uses lmcache.xpu_ops (requires XPU and the SYCL extension)
- xpu_py_ops: uses lmcache.python_ops_fallback with XPU visible
"""
params = []
cuda_available = torch.cuda.is_available()
params.append(pytest.param(("cpu_py_ops", _py_ops, "cpu"), id="cpu_py_ops"))
if cuda_available:
try:
# First Party
import lmcache.c_ops as cuda_c_ops
params.append(
pytest.param(("cuda_c_ops", cuda_c_ops, "cuda"), id="cuda_c_ops")
)
except ImportError:
pass
if _py_ops._get_copy_lib() is not None:
params.append(
pytest.param(("cuda_py_ops", _py_ops, "cuda"), id="cuda_cuda_py_ops")
)
else:
params.append(
pytest.param(("cuda_py_ops", _py_ops, "cpu"), id="cuda_cpu_py_ops")
)
if hasattr(torch, "xpu") and torch.xpu.is_available():
try:
# First Party
import lmcache.c_ops as xpu_sycl_ops
params.append(
pytest.param(("xpu_sycl_ops", xpu_sycl_ops, "xpu"), id="xpu_sycl_ops")
)
except ImportError:
pass
return params
_BACKEND_PARAMS = _build_backend_params()
# Module-level storage: (scenario_name, backend_id) -> {result_key: tensor}
_results: dict[tuple[str, str], dict[str, Any]] = {}
@pytest.fixture(scope="module", params=_BACKEND_PARAMS)
def backend(request: pytest.FixtureRequest) -> Any:
"""Yield (backend_id, ops_module, device) for each backend config.
For the cpu_py_ops variant, torch.cuda.is_available is patched to
return False so the scenario code behaves as if no GPU is present.
"""
backend_id, ops, device = request.param
if backend_id == "cpu_py_ops":
with unittest.mock.patch("torch.cuda.is_available", return_value=False):
yield backend_id, ops, device
else:
yield backend_id, ops, device
# ==========================================
# 2. Scenario functions
# ==========================================
def scenario_get_gpu_pci_bus_id(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test get_gpu_pci_bus_id returns a valid string on CUDA backends."""
res = ops.get_gpu_pci_bus_id(0)
is_valid = isinstance(res, str) and len(res) > 0
if torch.cuda.is_available() is False:
# for non cuda device, not expecting a valid return
# but crash should not happen
# and we mock a valid return here
is_valid = True
# 1 = PASS (call succeeded without crash)
# 0 = FAIL
return {
"get_gpu_pci_bus_id": torch.tensor([1 if is_valid else 0], dtype=torch.int32)
}
def scenario_calculate_cdf(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test calculate_cdf for multiple bin counts."""
num_bins_list = [1, 2, 5, 11, 15, 31, 32, 63]
results: dict[str, torch.Tensor] = {}
# all bins should be smaller than 64 -> align with C ops
assert all(n < 64 for n in num_bins_list), (
f"All num_bins must be < 64, got {[n for n in num_bins_list if n >= 64]}"
)
for num_bins in num_bins_list:
torch.manual_seed(42)
# Create on CPU first for consistent RNG across backends
input_tensor = torch.randint(0, num_bins, (1, 1000, 1), dtype=torch.int8).to(
device
)
raw_output = ops.calculate_cdf(input_tensor, num_bins)
# Both CUDA and non-CUDA return int16 with shape
# [nlayers, nchannels, num_bins + 1]
nlayers, _, nchannels = input_tensor.shape
assert raw_output.shape == (nlayers, nchannels, num_bins + 1), (
f"Expected shape ({nlayers}, {nchannels}, {num_bins + 1}), "
f"got {raw_output.shape}"
)
out_cpu = raw_output.flatten().cpu()
out_int32 = out_cpu.to(torch.int32)
out_uint16 = torch.where(out_int32 < 0, out_int32 + 65536, out_int32)
final_result = out_uint16.float() / 65536.0
results[f"calculate_cdf_bins{num_bins}"] = final_result
return results
def scenario_rotary_embedding_k_fused(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test rotary_embedding_k_fused for both NeoX and GPT-J rotation styles."""
torch.manual_seed(42)
# 1. Setup Dimensions
num_tokens = 128
num_kv_heads = 32
head_size = 128
max_position = 2048
rotary_dim = head_size
# 2. Generate Inputs on CPU first for consistent RNG across backends
old_positions = torch.randint(0, 1000, (num_tokens,), dtype=torch.long).to(device)
new_positions = old_positions + 1
cos_sin_cache = torch.randn(max_position, rotary_dim, dtype=torch.float32).to(
device
)
# Test both is_neox=True (NeoX-style, contiguous halves) and
# is_neox=False (GPT-J-style, interleaved)
results: dict[str, torch.Tensor] = {}
for is_neox in [True, False]:
# Reset seed for consistent key tensor across both tests
torch.manual_seed(42)
key = torch.randn(num_tokens, num_kv_heads, head_size, dtype=torch.float32).to(
device
)
# 3. Execute (in-place update on key)
ops.rotary_embedding_k_fused(
old_positions,
new_positions,
key,
head_size,
cos_sin_cache,
is_neox,
)
# 4. Collect with is_neox suffix to distinguish the two test cases
neox_suffix = "neox" if is_neox else "gptj"
results[f"rotary_embedding_k_fused_{neox_suffix}"] = key.cpu()
return results
def scenario_lmcache_memcpy_async(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test lmcache_memcpy_async for H2D and D2H memory transfers.
Tests both pointer mode (int pointers) and tensor mode (torch.Tensor objects).
Uses pointer mode for CPU/CUDA devices and tensor mode for other devices.
Exercises multiple boundary conditions to verify correct behaviour for
both the CUDA c_ops backend (which chunks at alignment boundaries via
cudaMemcpyAsync) and the Python fallback backend (which issues a single
synchronous copy):
- copy spanning exactly one aligned block
- copy entirely within one aligned block (no boundary crossing)
- copy crossing a single alignment boundary
- copy crossing multiple alignment boundaries
- copy starting exactly at an alignment boundary
- copy with an unaligned start offset crossing many boundaries
Both H2D and D2H directions are verified for each boundary condition.
"""
torch.manual_seed(42)
# Buffer large enough for all boundary test cases (max offset+nbytes = 16+200 = 216)
buf_size = 256
alignment = 64 # boundary interval in bytes (must be a power of two)
src_host = torch.randint(0, 256, (buf_size,), dtype=torch.uint8)
gpu_buffer = torch.zeros(buf_size, dtype=torch.uint8, device=device)
dst_host = torch.zeros(buf_size, dtype=torch.uint8)
if device in ("cuda", "xpu"):
dst_host = dst_host.pin_memory()
h2d_dir = ops.TransferDirection.H2D
d2h_dir = ops.TransferDirection.D2H
# Decide mode based on the running device.
# The native CUDA/XPU backend only accepts a tensor of uint64 pointers;
# only the Python fallback supports list[Tensor].
use_tensor_mode = device not in ("cpu", "cuda", "xpu")
# (host_buffer_offset, nbytes) boundary test cases:
# (0, 64): exactly one aligned block from the start
# (32, 32): entirely within one block, no boundary crossing
# (32, 64): crosses one boundary — [32..64) then [64..96)
# (0, 192): three full aligned blocks (boundaries at 64 and 128)
# (64, 128): starts at alignment boundary, spans two full blocks
# (16, 200): unaligned start, crosses multiple boundaries
test_cases = [
(0, 64),
(32, 32),
(32, 64),
(0, 192),
(64, 128),
(16, 200),
]
all_results: list[torch.Tensor] = []
for offset, nbytes in test_cases:
gpu_buffer.zero_()
dst_host.zero_()
expected = src_host[offset : offset + nbytes].clone()
device_sync(device)
if use_tensor_mode:
ops.lmcache_memcpy_async(
gpu_buffer[offset : offset + nbytes],
src_host[offset : offset + nbytes],
nbytes,
h2d_dir,
0,
alignment,
)
device_sync(device)
ops.lmcache_memcpy_async(
dst_host[offset : offset + nbytes],
gpu_buffer[offset : offset + nbytes],
nbytes,
d2h_dir,
0,
alignment,
)
else:
ops.lmcache_memcpy_async(
gpu_buffer.data_ptr() + offset,
src_host.data_ptr() + offset,
nbytes,
h2d_dir,
offset,
alignment,
)
device_sync(device)
ops.lmcache_memcpy_async(
dst_host.data_ptr() + offset,
gpu_buffer.data_ptr() + offset,
nbytes,
d2h_dir,
offset,
alignment,
)
device_sync(device)
result = dst_host[offset : offset + nbytes].cpu()
assert torch.equal(result, expected), (
f"Boundary test (offset={offset}, nbytes={nbytes}, alignment={alignment}): "
f"data corrupted, max diff = "
f"{(result.float() - expected.float()).abs().max().item()}"
)
all_results.append(result)
return {"lmcache_memcpy_async": torch.cat(all_results)}
def scenario_load_and_reshape_flash(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test load_and_reshape_flash extracts KV cache tokens into contiguous buffer."""
torch.manual_seed(42)
# 1. Standard Params
src_device = device
dst_device = "cpu"
num_blocks = 100
block_size = 16
num_heads = 8
head_size = 128
num_layers = 32
num_tokens = 256
chunk_size = 256
dtype = torch.bfloat16
# 2. Setup Data (Deterministic Pattern)
total_elements = num_blocks * block_size * num_heads * head_size
kv_cache_cpu = []
for i in range(num_layers):
base_tensor = torch.linspace(i, i + 1, total_elements, dtype=torch.float32)
base_tensor = base_tensor.reshape(
num_blocks, block_size, num_heads, head_size
).to(dtype)
k = base_tensor
v = base_tensor + 0.5
kv_cache_cpu.append([k, v])
kv_cache = [
[layer[0].to(src_device), layer[1].to(src_device)] for layer in kv_cache_cpu
]
# Slot mapping: deterministic strided selection
step = (num_blocks * block_size) // num_tokens
slot_indices = list(range(0, num_blocks * block_size, step))[:num_tokens]
slot_mapping = torch.tensor(slot_indices, device=src_device, dtype=torch.int64)
slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
# 3. Extract (to CPU pinned)
extracted_chunks = []
for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
mem_obj_shape = (2, num_layers, len(slot_mapping_temp), num_heads * head_size)
mem_obj_tensor = torch.zeros(mem_obj_shape, dtype=dtype, device=dst_device)
if device in ("cuda", "xpu"):
mem_obj_tensor = mem_obj_tensor.pin_memory()
for layer_id in range(num_layers):
ops.load_and_reshape_flash(
mem_obj_tensor,
kv_cache[layer_id][0],
kv_cache[layer_id][1],
slot_mapping_temp,
layer_id,
)
extracted_chunks.append(mem_obj_tensor)
device_sync(device)
# 4. Verify: compare extracted data against original kv_cache
# mem_obj_tensor layout:
# [2, num_layers, num_tokens_in_chunk, num_heads * head_size]
# dim 0: K=0, V=1
# Original kv_cache layout: [num_blocks, block_size, num_heads, head_size]
# slot_mapping tells us which (block, offset) each token comes from
for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
slots = slot_mapping_temp.cpu()
extracted = extracted_chunks[chunk_id].cpu()
for layer_id in range(num_layers):
orig_k = kv_cache_cpu[layer_id][
0
] # [num_blocks, block_size, num_heads, head_size]
orig_v = kv_cache_cpu[layer_id][1]
for tok_idx, slot in enumerate(slots):
block_idx = slot.item() // block_size
offset = slot.item() % block_size
# Expected: flattened [num_heads * head_size]
expected_k = orig_k[block_idx, offset].reshape(-1)
expected_v = orig_v[block_idx, offset].reshape(-1)
# Extracted
got_k = extracted[0, layer_id, tok_idx]
got_v = extracted[1, layer_id, tok_idx]
k_diff = (got_k.float() - expected_k.float()).abs().max().item()
assert torch.equal(got_k, expected_k), (
f"K mismatch layer={layer_id}, slot={slot.item()}, "
f"max diff={k_diff}"
)
v_diff = (got_v.float() - expected_v.float()).abs().max().item()
assert torch.equal(got_v, expected_v), (
f"V mismatch layer={layer_id}, slot={slot.item()}, "
f"max diff={v_diff}"
)
# 5. Return ALL extracted chunks concatenated for cross-backend comparison
return {
"load_and_reshape_flash": torch.cat([c.cpu() for c in extracted_chunks], dim=2)
}
def scenario_reshape_and_cache_back_flash(
ops: Any, device: str
) -> dict[str, torch.Tensor]:
"""Test reshape_and_cache_back_flash writes tokens back into paged KV cache."""
torch.manual_seed(42)
# 1. Environment Setup
src_device = "cpu"
dst_device = device
num_blocks = 100
block_size = 16
num_heads = 8
head_size = 128
num_layers = 32
num_tokens = 256
chunk_size = 256
dtype = torch.bfloat16
# 2. Prepare Source Data (CPU Buffer)
# Shape: [2, num_layers, num_tokens, num_heads * head_size]
mem_obj_shape = (2, num_layers, num_tokens, num_heads * head_size)
src_buffer = torch.zeros(mem_obj_shape, dtype=dtype, device=src_device)
# Data Pattern: Odd numbers (1.0, 3.0, 5.0, ...)
for i in range(num_tokens):
val = 1.0 + (i * 2.0)
src_buffer[0, :, i, :] = val # Key
src_buffer[1, :, i, :] = val + 0.5 # Value
if device in ("cuda", "xpu"):
src_buffer = src_buffer.pin_memory()
# 3. Prepare Destination (Empty Cache)
kv_cache = [
[
torch.zeros(
num_blocks,
block_size,
num_heads,
head_size,
device=dst_device,
dtype=dtype,
),
torch.zeros(
num_blocks,
block_size,
num_heads,
head_size,
device=dst_device,
dtype=dtype,
),
]
for _ in range(num_layers)
]
# 4. Slot Mapping (Continuous: Token 0 → Slot 0, Token 1 → Slot 1, ...)
slot_indices = list(range(num_tokens))
slot_mapping = torch.tensor(slot_indices, device=dst_device, dtype=torch.int64)
slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
# 5. Execute Operator (Load Back)
current_token_offset = 0
for chunk_id, slot_chunk in enumerate(slot_mapping_chunked):
chunk_len = len(slot_chunk)
buffer_chunk = src_buffer[
:, :, current_token_offset : current_token_offset + chunk_len, :
]
if not buffer_chunk.is_contiguous():
buffer_chunk = buffer_chunk.contiguous()
for layer_id in range(num_layers):
ops.reshape_and_cache_back_flash(
buffer_chunk,
kv_cache[layer_id][0],
kv_cache[layer_id][1],
slot_chunk,
layer_id,
)
current_token_offset += chunk_len
device_sync(device)
# 6. Verify: check written values against source pattern
for layer_id in range(num_layers):
k_cache = kv_cache[layer_id][
0
].cpu() # [num_blocks, block_size, num_heads, head_size]
v_cache = kv_cache[layer_id][1].cpu()
for tok_idx, slot in enumerate(slot_indices):
block_idx = slot // block_size
offset = slot % block_size
expected_k_val = 1.0 + (tok_idx * 2.0)
expected_v_val = expected_k_val + 0.5
got_k = k_cache[block_idx, offset]
got_v = v_cache[block_idx, offset]
expected_k = torch.full_like(got_k, expected_k_val)
expected_v = torch.full_like(got_v, expected_v_val)
assert torch.allclose(got_k.float(), expected_k.float(), atol=0.1), (
f"K mismatch at layer={layer_id}, slot={slot}, "
f"expected={expected_k_val}, got={got_k[0, 0].item()}"
)
assert torch.allclose(got_v.float(), expected_v.float(), atol=0.1), (
f"V mismatch at layer={layer_id}, slot={slot}, "
f"expected={expected_v_val}, got={got_v[0, 0].item()}"
)
# 7. Return first block of layer 0 key cache for cross-backend comparison
return {"reshape_and_cache_back_flash": kv_cache[0][0][0].cpu()}
def scenario_encode_fast_new(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test encode_fast_new produces valid, non-empty encoded output."""
torch.manual_seed(42)
# 1. Hyperparameters
nlayers = 2
nchannels = 4
ntokens = 128
alphabet_size = 16
max_buf_len = ntokens * 2
# 2. Construct Data on target device
# A. CDF: uniform distribution, strictly increasing
step = 100 // alphabet_size
base_cdf = torch.arange(0, 100, step, dtype=torch.int32)
base_cdf = base_cdf[:alphabet_size]
cdf_cpu = (
base_cdf.unsqueeze(0).unsqueeze(0).expand(nlayers, nchannels, -1).contiguous()
)
cdf = cdf_cpu.to(dtype=torch.int16, device=device)
# B. Input symbols: cycling 0..14
total_syms = nlayers * ntokens * nchannels
input_cpu = torch.arange(total_syms, dtype=torch.float32)
input_cpu = (input_cpu % (alphabet_size - 1)).to(torch.int8)
input_cpu = input_cpu.reshape(nlayers, ntokens, nchannels)
input_sym = input_cpu.to(device=device)
# 3. Prepare Outputs
output_buffer = torch.zeros(
(nlayers, nchannels, max_buf_len),
dtype=torch.uint8,
device=device,
)
output_lengths = torch.zeros(
(nlayers, nchannels),
dtype=torch.int32,
device=device,
)
# 4. Execute
ops.encode_fast_new(
cdf,
input_sym,
output_buffer,
output_lengths,
)
device_sync(device)
# 5. Verify
lengths_cpu = output_lengths.cpu()
assert (lengths_cpu > 0).all(), "Encoding produced zero-length output!"
assert (lengths_cpu <= max_buf_len).all(), "Buffer overflow detected!"
# 6. Return: first 200 bytes of layer 0, channel 0
valid_len = int(lengths_cpu[0, 0].item())
res = output_buffer[0, 0, : min(valid_len, 200)].cpu()
return {"encode_fast_new": res}
def scenario_decode_fast_new(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test decode_fast_new correctly decodes a round-tripped encoded stream."""
torch.manual_seed(42)
# 1. Config
nlayers = 2
nchannels = 4
ntokens = 128
alphabet_size = 16
max_buf_len = ntokens * 2
# 2. Data Generation on CPU first for consistent RNG across backends
cdf = torch.randint(
1,
100,
(nlayers, nchannels, alphabet_size),
dtype=torch.int32,
)
cdf = torch.cumsum(cdf, dim=-1).to(torch.int16).to(device)
input_sym = torch.randint(
0,
alphabet_size - 2,
(nlayers, ntokens, nchannels),
dtype=torch.int8,
).to(device)
# 3. Encode first (need encoded data to test decode)
encoded_buffer = torch.zeros(
(nlayers, nchannels, max_buf_len),
dtype=torch.uint8,
device=device,
)
encoded_lengths = torch.zeros(
(nlayers, nchannels),
dtype=torch.int32,
device=device,
)
ops.encode_fast_new(
cdf,
input_sym,
encoded_buffer,
encoded_lengths,
)
device_sync(device)
# 4. Decode
decoded_sym = torch.zeros_like(input_sym, dtype=torch.uint8)
ops.decode_fast_new(
cdf,
encoded_buffer,
encoded_lengths,
decoded_sym,
)
device_sync(device)
# 5. Verify: decoded must match original
input_uint8 = input_sym.to(torch.uint8)
mismatch = (input_uint8 != decoded_sym).sum().item()
if mismatch > 0:
mask = input_uint8 != decoded_sym
ly, t, c = mask.nonzero()[0].tolist()
pytest.fail(
f"Decode mismatch: {mismatch} errors. "
f"First diff at L{ly}T{t}C{c}: "
f"orig={input_uint8[ly, t, c].item()} "
f"decoded={decoded_sym[ly, t, c].item()}"
)
# 6. Return decoded slice for cross-backend comparison
return {"decode_fast_new": decoded_sym[0, :20, 0].cpu()}
def scenario_decode_fast_prefsum(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test decode_fast_prefsum correctly decodes with prefix-sum offsets."""
torch.manual_seed(42)
# 1. Configuration
nlayers = 2
nchannels = 4
ntokens = 128
alphabet_size = 16
max_buf_len = ntokens * 2
# 2. Data Generation (Normalized CDF) on CPU first for consistent RNG
cdf = torch.randint(
1,
100,
(nlayers, nchannels, alphabet_size),
dtype=torch.int32,
)
cdf = torch.cumsum(cdf, dim=-1).float()
cdf = (cdf / cdf[..., -1:] * 65536).to(torch.int32)
cdf[..., -1] = 65536
cdf = cdf.to(torch.int16).to(device).contiguous()
input_sym = torch.randint(
0,
alphabet_size - 2,
(nlayers, ntokens, nchannels),
dtype=torch.int8,
).to(device)
# 3. Encode to get variable lengths
tmp_buf = torch.zeros(
(nlayers, nchannels, max_buf_len),
dtype=torch.uint8,
device=device,
)
tmp_len = torch.zeros(
(nlayers, nchannels),
dtype=torch.int32,
device=device,
)
ops.encode_fast_new(cdf, input_sym, tmp_buf, tmp_len)
device_sync(device)
# 4. Pack into 1D dense bytestream
lens_flat = tmp_len.cpu().flatten().tolist()
bufs_flat = tmp_buf.cpu().reshape(-1, max_buf_len).numpy()
all_bytes = []
for i, length in enumerate(lens_flat):
all_bytes.extend(bufs_flat[i, :length].tolist())
bytestream_1d = torch.tensor(
all_bytes,
dtype=torch.uint8,
device=device,
).contiguous()
# 5. Offsets (end-position via cumsum)
lengths_prefsum = (
tmp_len.flatten().cumsum(0).reshape(tmp_len.shape).to(torch.int64).to(device)
).contiguous()
# 6. Decode
decoded_sym = (
torch.zeros_like(
input_sym,
dtype=torch.uint8,
)
.to(device)
.contiguous()
)
ops.decode_fast_prefsum(
cdf,
bytestream_1d,
lengths_prefsum,
decoded_sym,
)
device_sync(device)
# 7. Verify roundtrip
input_ref = input_sym.to(torch.uint8)
mismatch = (input_ref != decoded_sym).sum().item()
if mismatch > 0:
mask = input_ref != decoded_sym
ly, t, c = mask.nonzero()[0].tolist()
pytest.fail(
f"Prefsum mismatch: {mismatch} errors. "
f"First diff at L{ly}T{t}C{c}: "
f"orig={input_ref[ly, t, c].item()} "
f"decoded={decoded_sym[ly, t, c].item()}"
)
# 8. Return
return {
"decode_fast_prefsum": decoded_sym[0, :20, 0].cpu(),
}
def scenario_single_layer_kv_transfer(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test single_layer_kv_transfer for multiple KV formats and directions."""
torch.manual_seed(42)
num_tokens = 64
num_blocks = 256
block_size = 16
num_heads = 12
head_size = 64
hidden_size = num_heads * head_size
slot_mapping = torch.arange(
0,
num_tokens * 2,
2,
device=device,
).to(torch.int64)
# (engine_kv_format, is_mla, token_major, direction)
# direction: False = LMC→vLLM (H2D), True = vLLM→LMC (D2H)
test_cases = [
# flash attn: [2, NB, BS, NH, HS] — two_major
(ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS, False, True, False),
(ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS, False, False, False),
(ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS, False, True, True),
# flash infer: [NB, 2, BS, NH, HS]
(ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS, False, True, False),
(ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS, False, False, False),
(ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS, False, True, True),
# vLLM MLA: [NB, BS, HS]
(ops.EngineKVFormat.NL_X_NB_BS_HS, True, True, False),
(ops.EngineKVFormat.NL_X_NB_BS_HS, True, True, True),
]
for engine_kv_format, is_mla, token_major, direction in test_cases:
dir_tag = "v2l" if direction else "l2v"
is_two_major = engine_kv_format == ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
case_desc = (
f"fmt={engine_kv_format}, MLA={is_mla}, TM={token_major}, Dir={dir_tag}"
)
# ── 1. Setup Shapes ──
lmc_shape: tuple[int, ...] = ()
vllm_shape: tuple[int, ...] = ()
if is_mla:
lmc_shape = (num_tokens, hidden_size)
vllm_shape = (num_blocks, block_size, hidden_size)
else:
lmc_shape = (
(num_tokens, 2, hidden_size)
if token_major
else (2, num_tokens, hidden_size)
)
if is_two_major:
# flash attn: [2, num_blocks, block_size, num_heads, head_size]
vllm_shape = (2, num_blocks, block_size, num_heads, head_size)
else:
# flash infer: [num_blocks, 2, block_size, num_heads, head_size]
vllm_shape = (num_blocks, 2, block_size, num_heads, head_size)
# ── 2. Deterministic Data ──
lmc_size = 1
for s in lmc_shape:
lmc_size *= s
vllm_size = 1
for s in vllm_shape:
vllm_size *= s
lmc_tensor = (
(torch.arange(lmc_size, device=device) % 1000)
.to(torch.float16)
.reshape(lmc_shape)
)
vllm_tensor = (
(torch.arange(vllm_size, device=device) % 1000)
.to(torch.float16)
.reshape(vllm_shape)
)
# ── 3. Golden Reference ──
lmc_ref = lmc_tensor.clone()
vllm_ref = vllm_tensor.clone()
block_indices = slot_mapping // block_size
block_offsets = slot_mapping % block_size
if not direction: # LMC → vLLM
if is_mla:
vllm_ref[block_indices, block_offsets, :] = lmc_ref
else:
src = lmc_ref if token_major else lmc_ref.permute(1, 0, 2)
src = src.view(num_tokens, 2, num_heads, head_size)
if is_two_major:
# [2, NB, BS, NH, HS]
vllm_ref[0, block_indices, block_offsets] = src[:, 0, :, :]
vllm_ref[1, block_indices, block_offsets] = src[:, 1, :, :]
else:
# [NB, 2, BS, NH, HS]
vllm_ref[block_indices, 0, block_offsets] = src[:, 0, :, :]
vllm_ref[block_indices, 1, block_offsets] = src[:, 1, :, :]
else: # vLLM → LMC
if is_mla:
lmc_ref = vllm_ref[block_indices, block_offsets, :]
else:
if is_two_major:
k = vllm_ref[0, block_indices, block_offsets]
v = vllm_ref[1, block_indices, block_offsets]
else:
k = vllm_ref[block_indices, 0, block_offsets]
v = vllm_ref[block_indices, 1, block_offsets]
combined = torch.stack(
[k, v],
dim=1,
).view(num_tokens, 2, hidden_size)
lmc_ref = combined if token_major else combined.permute(1, 0, 2)
# ── 4. Execute ──
xfer_dir = ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D
ops.single_layer_kv_transfer(
lmc_tensor,
vllm_tensor,
slot_mapping,
xfer_dir,
engine_kv_format,
token_major,
)
device_sync(device)
# ── 5. Verify ──
if not direction:
torch.testing.assert_close(
vllm_tensor,
vllm_ref,
rtol=1e-3,
atol=1e-3,
msg=f"Mismatch in {case_desc}",
)
else:
torch.testing.assert_close(
lmc_tensor,
lmc_ref,
rtol=1e-3,
atol=1e-3,
msg=f"Mismatch in {case_desc}",
)
# ── 6. Collect canonical results for cross-backend comparison ──
# Use flash attn (two_major) format to match original file names
canonical_cases = [
(False, True, False), # l2v, non-MLA
(False, True, True), # v2l, non-MLA
(True, True, False), # l2v, MLA
(True, True, True), # v2l, MLA
]
results: dict[str, torch.Tensor] = {}
for is_mla, token_major, direction in canonical_cases:
dir_tag = "v2l" if direction else "l2v"
if is_mla:
lmc_shape = (num_tokens, hidden_size)
vllm_shape = (num_blocks, block_size, hidden_size)
fmt = ops.EngineKVFormat.NL_X_NB_BS_HS
else:
lmc_shape = (num_tokens, 2, hidden_size)
vllm_shape = (2, num_blocks, block_size, num_heads, head_size)
fmt = ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
lmc_size = 1
for s in lmc_shape:
lmc_size *= s
vllm_size = 1
for s in vllm_shape:
vllm_size *= s
lmc_tensor = (
(torch.arange(lmc_size, device=device) % 1000)
.to(torch.float16)
.reshape(lmc_shape)
)
vllm_tensor = (
(torch.arange(vllm_size, device=device) % 1000)
.to(torch.float16)
.reshape(vllm_shape)
)
xfer_dir = ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D
ops.single_layer_kv_transfer(
lmc_tensor,
vllm_tensor,
slot_mapping,
xfer_dir,
fmt,
token_major,
)
device_sync(device)
result = lmc_tensor.cpu() if direction else vllm_tensor.cpu()
results[f"single_layer_kv_transfer_{dir_tag}_mla_{is_mla}"] = result
return results
def scenario_single_layer_kv_transfer_sgl(
ops: Any, device: str
) -> dict[str, torch.Tensor]:
"""Test single_layer_kv_transfer_sgl for SGLang KV format."""
torch.manual_seed(42)
num_tokens = 32
num_blocks = 128
block_size = 16
num_heads = 8
head_size = 64
hidden_size = num_heads * head_size
slot_mapping = torch.arange(
0,
num_tokens * 3,
3,
device=device,
).to(torch.int64)
# (token_major, direction)
# direction: False = LMC→SGL, True = SGL→LMC
test_cases = [
(True, False),
(False, False),
(True, True),
(False, True),
]
results: dict[str, torch.Tensor] = {}
for token_major, direction in test_cases:
dir_tag = "s2l" if direction else "l2s"
# 1. Setup Shapes
lmc_shape = (
(num_tokens, 2, hidden_size)
if token_major
else (2, num_tokens, hidden_size)
)
sgl_shape = (
num_blocks,
block_size,
num_heads,
head_size,
)
# 2. Deterministic Data
lmc_size = 1
for s in lmc_shape:
lmc_size *= s
sgl_size = 1
for s in sgl_shape:
sgl_size *= s
lmc_tensor = (
(torch.arange(lmc_size, device=device) % 500)
.to(torch.float16)
.reshape(lmc_shape)
)
sgl_k_tensor = (
(torch.arange(sgl_size, device=device) % 500 + 500)
.to(torch.float16)
.reshape(sgl_shape)
)
sgl_v_tensor = (
(torch.arange(sgl_size, device=device) % 500 + 1000)
.to(torch.float16)
.reshape(sgl_shape)
)
# 3. Golden Reference
lmc_ref = lmc_tensor.clone()
sgl_k_ref = sgl_k_tensor.clone()
sgl_v_ref = sgl_v_tensor.clone()
block_indices = slot_mapping // block_size
block_offsets = slot_mapping % block_size
if not direction: # LMC → SGL
src = lmc_ref if token_major else lmc_ref.permute(1, 0, 2)
src_k = src[:, 0, :].view(
num_tokens,
num_heads,
head_size,
)
src_v = src[:, 1, :].view(
num_tokens,
num_heads,
head_size,
)
sgl_k_ref[block_indices, block_offsets] = src_k
sgl_v_ref[block_indices, block_offsets] = src_v
else: # SGL → LMC
k_data = sgl_k_ref[block_indices, block_offsets].reshape(
num_tokens, hidden_size
)
v_data = sgl_v_ref[block_indices, block_offsets].reshape(
num_tokens, hidden_size
)
combined = torch.stack(
[k_data, v_data],
dim=1,
) # [N, 2, H]
lmc_ref = combined if token_major else combined.permute(1, 0, 2)
# 4. Execute
ops.single_layer_kv_transfer_sgl(
lmc_tensor,
sgl_k_tensor,
sgl_v_tensor,
slot_mapping,
ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D,
token_major,
)
device_sync(device)
# 5. Verify
case_desc = f"TM={token_major}, Dir={dir_tag}"
if not direction:
torch.testing.assert_close(
sgl_k_tensor,
sgl_k_ref,
rtol=1e-3,
atol=1e-3,
msg=f"K mismatch in {case_desc}",
)
torch.testing.assert_close(
sgl_v_tensor,
sgl_v_ref,
rtol=1e-3,
atol=1e-3,
msg=f"V mismatch in {case_desc}",
)
else:
torch.testing.assert_close(
lmc_tensor,
lmc_ref,
rtol=1e-3,
atol=1e-3,
msg=f"Mismatch in {case_desc}",
)
# 6. Collect each case separately
result = lmc_tensor.cpu() if direction else sgl_k_tensor.cpu()
results[f"single_layer_kv_transfer_sgl_{dir_tag}_tm_{token_major}"] = result
return results
def scenario_multi_layer_kv_transfer(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test multi_layer_kv_transfer for multiple paged KV formats and directions.
Tests both pointer mode (torch.Tensor of int64 pointers) and tensor list mode
(list[torch.Tensor]) for key_value_ptrs.
"""
torch.manual_seed(42)
num_layers = 2
num_tokens = 4
head_size = 16
page_buffer_size = 10
block_size = 5
dtype = torch.float32
slot_mapping = torch.tensor(
[0, 2, 5, 9],
dtype=torch.int64,
device=device,
)
# ── Format-specific test cases ──
# Each: (engine_kv_format, is_mla, block_size_arg)
format_cases = [
(ops.EngineKVFormat.NB_NL_TWO_BS_NH_HS, False, 1), # vLLM cross layer
(ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS, False, 1), # flash attn
(ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS, False, block_size), # flash infer
(ops.EngineKVFormat.NL_X_NB_BS_HS, True, 1), # vLLM MLA
(ops.EngineKVFormat.NL_X_NBBS_ONE_HS, True, 1), # SGLang MLA
]
# Decide mode based on the running device.
# The native CUDA/XPU backend only accepts a tensor of uint64 pointers;
# only the Python fallback supports list[Tensor].
use_tensor_list = device not in ("cpu", "cuda", "xpu")
for engine_kv_format, is_mla, bs_arg in format_cases:
k_or_v_size = 1 if is_mla else 2
for direction in [True, False]:
dir_tag = "paged2lmc" if direction else "lmc2paged"
fmt_name = str(engine_kv_format).split(".")[-1]
# ── 1. LMCache Tensor ──
lmc_shape = (k_or_v_size, num_layers, num_tokens, head_size)
key_value = torch.zeros(lmc_shape, dtype=dtype, device="cpu")
if device in ("cuda", "xpu"):
key_value = key_value.pin_memory()
if not direction: # LMC → Paged
for kv in range(k_or_v_size):
for ly in range(num_layers):
for t in range(num_tokens):
val = (
kv * 5000
+ ly * 1000
+ t * 10
+ torch.arange(head_size, device=device)
).to(dtype)
key_value[kv, ly, t] = val
# ── 2. Paged Buffers (one per layer) ──
page_buffers = []
for ly in range(num_layers):
if engine_kv_format == ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS:
num_blocks = page_buffer_size // bs_arg
pb = torch.zeros(
(num_blocks, 2, bs_arg, head_size),
dtype=dtype,
device=device,
)
elif is_mla:
pb = torch.zeros(
(page_buffer_size, head_size),
dtype=dtype,
device=device,
)
else:
# Handles NB_NL_TWO_BS_NH_HS and NL_X_TWO_NB_BS_NH_HS
pb = torch.zeros(
(2, page_buffer_size, head_size),
dtype=dtype,
device=device,
)
if direction: # Paged → LMC
for s in range(page_buffer_size):
for kv in range(k_or_v_size):
val = (
kv * 7000
+ ly * 2000
+ s * 10
+ torch.arange(head_size, device=device)
).to(dtype)
if (
engine_kv_format
== ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS
):
blk_idx = s // bs_arg
blk_off = s % bs_arg
pb[blk_idx, kv, blk_off] = val
elif is_mla:
pb[s] = val
else:
# Handles NB_NL_TWO_BS_NH_HS and NL_X_TWO_NB_BS_NH_HS
pb[kv, s] = val
page_buffers.append(pb)
# ── 3. Prepare key_value_ptrs (pointer mode or tensor list mode) ──
key_value_ptrs: Union[list[torch.Tensor], torch.Tensor]
if use_tensor_list:
# Tensor list mode: pass the tensor objects directly
key_value_ptrs = page_buffers
else:
# Pointer mode: create tensor of uint64 pointers
key_value_ptrs = torch.tensor(
[pb.data_ptr() for pb in page_buffers],
dtype=torch.uint64,
device=device,
)
# ── 4. Execute ──
xfer_dir = (
ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D
)
ops.multi_layer_kv_transfer(
key_value,
key_value_ptrs,
slot_mapping,
torch.device(device),
page_buffer_size,
xfer_dir,
engine_kv_format,
bs_arg,
)
device_sync(device)
# ── 5. Verify (internal, per-format) ──
for t_id in range(num_tokens):
s_idx = int(slot_mapping[t_id].item())
for ly in range(num_layers):
for kv in range(k_or_v_size):
lmc_val = key_value[kv, ly, t_id]
if engine_kv_format == ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS:
blk_idx = s_idx // bs_arg
blk_off = s_idx % bs_arg
paged_val = page_buffers[ly][blk_idx, kv, blk_off]
elif is_mla:
paged_val = page_buffers[ly][s_idx]
else:
# Handles NB_NL_TWO_BS_NH_HS and NL_X_TWO_NB_BS_NH_HS
paged_val = page_buffers[ly][kv, s_idx]
torch.testing.assert_close(
lmc_val.to("cpu"),
paged_val.to("cpu"),
msg=(
f"Mismatch: {fmt_name} {dir_tag} "
f"(tensor_list={use_tensor_list}), "
f"kv={kv}, layer={ly}, token={t_id}"
),
)
# ── 6. Collect ONE canonical result for cross-backend comparison ──
# Use flash attn format (NL_X_TWO_NB_BS_NH_HS), re-run canonical cases
canonical_format = ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
results: dict[str, torch.Tensor] = {}
for direction in [True, False]:
dir_tag = "paged2lmc" if direction else "lmc2paged"
lmc_shape = (2, num_layers, num_tokens, head_size)
key_value = torch.zeros(lmc_shape, dtype=dtype, device=device)
if not direction:
for ly in range(num_layers):
for t in range(num_tokens):
val = (
ly * 1000 + t * 10 + torch.arange(head_size, device=device)
).to(dtype)
key_value[0, ly, t] = val
key_value[1, ly, t] = val + 500
page_buffers = []
for ly in range(num_layers):
pb = torch.zeros(
(2, page_buffer_size, head_size),
dtype=dtype,
device=device,
)
if direction:
for s in range(page_buffer_size):
val = (
ly * 2000 + s * 10 + torch.arange(head_size, device=device)
).to(dtype)
pb[0, s] = val
pb[1, s] = val + 700
page_buffers.append(pb)
if use_tensor_list:
key_value_ptrs = page_buffers
else:
key_value_ptrs = torch.tensor(
[pb.data_ptr() for pb in page_buffers],
dtype=torch.uint64,
device=device,
)
xfer_dir = ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D
ops.multi_layer_kv_transfer(
key_value,
key_value_ptrs,
slot_mapping,
torch.device(device),
page_buffer_size,
xfer_dir,
canonical_format,
1,
)
device_sync(device)
results[f"multi_layer_kv_transfer_{dir_tag}"] = key_value.cpu()
return results
def scenario_multi_layer_kv_transfer_unilateral(
ops: Any, device: str
) -> dict[str, torch.Tensor]:
"""Test multi_layer_kv_transfer_unilateral for non-interleaved K/V pointers.
Tests both pointer mode (torch.Tensor of int64 pointers) and tensor list mode
(list[torch.Tensor]) for key_value_ptrs.
"""
torch.manual_seed(42)
num_layers = 2
num_tokens = 4
head_size = 16
page_buffer_size = 10
dtype = torch.float32
slot_mapping = torch.tensor(
[1, 3, 4, 7],
dtype=torch.int64,
device=device,
)
# ── Test cases: (engine_kv_format, is_mla) ──
format_cases = [
(
ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS,
False,
), # SGLang MHA (unilateral path)
(
ops.EngineKVFormat.NL_X_NB_BS_HS,
True,
), # vLLM MLA (delegates to multi_layer_kv_transfer)
(
ops.EngineKVFormat.NL_X_NBBS_ONE_HS,
True,
), # SGLang MLA (delegates to multi_layer_kv_transfer)
]
# Decide mode based on the running device.
# The native CUDA/XPU backend only accepts a tensor of uint64 pointers;
# only the Python fallback supports list[Tensor].
use_tensor_list = device not in ("cpu", "cuda", "xpu")
for engine_kv_format, is_mla in format_cases:
k_or_v_size = 1 if is_mla else 2
for direction in [True, False]:
dir_tag = "p2l" if direction else "l2p"
# ── 1. LMCache Tensor ──
lmc_shape = (k_or_v_size, num_layers, num_tokens, head_size)
lmc_tensor = torch.zeros(lmc_shape, dtype=dtype, device="cpu")
if device in ("cuda", "xpu"):
lmc_tensor = lmc_tensor.pin_memory()
if not direction: # LMC → Paged
for kv in range(k_or_v_size):
for ly in range(num_layers):
for t in range(num_tokens):
val = (
kv * 5000
+ ly * 1000
+ t * 10
+ torch.arange(head_size, device=device)
).to(dtype)
lmc_tensor[kv, ly, t] = val
# ── 2. Paged Buffers ──
key_value_ptrs: Union[list[torch.Tensor], torch.Tensor]
if is_mla:
# MLA delegates to multi_layer_kv_transfer
# ptrs: [layer0, layer1, ...], each -> [page_buffer_size, head_size]
page_buffers = []
for ly in range(num_layers):
pb = torch.zeros(
(page_buffer_size, head_size),
dtype=dtype,
device=device,
)
if direction: # Paged → LMC
for s in range(page_buffer_size):
val = (
ly * 2000
+ s * 10
+ torch.arange(head_size, device=device)
).to(dtype)
pb[s] = val
page_buffers.append(pb)
if use_tensor_list:
key_value_ptrs = page_buffers
else:
key_value_ptrs = torch.tensor(
[pb.data_ptr() for pb in page_buffers],
dtype=torch.uint64,
device=device,
)
else:
# Non-MLA unilateral: separate K/V buffers
# ptrs: [K_l0, K_l1, ..., V_l0, V_l1, ...]
# each -> [page_buffer_size, head_size]
buffers = {}
for kv in range(2):
for ly in range(num_layers):
pb = torch.zeros(
(page_buffer_size, head_size),
dtype=dtype,
device=device,
)
if direction: # Paged → LMC
for s in range(page_buffer_size):
val = (
kv * 7000
+ ly * 2000
+ s * 10
+ torch.arange(head_size, device=device)
).to(dtype)
pb[s] = val
buffers[(kv, ly)] = pb
if use_tensor_list:
# Tensor list mode: [K_l0, K_l1, ..., V_l0, V_l1, ...]
tensor_list = []
for ly in range(num_layers):
tensor_list.append(buffers[(0, ly)])
for ly in range(num_layers):
tensor_list.append(buffers[(1, ly)])
key_value_ptrs = tensor_list
else:
# Pointer mode
ptr_list = []
for ly in range(num_layers):
ptr_list.append(buffers[(0, ly)].data_ptr())
for ly in range(num_layers):
ptr_list.append(buffers[(1, ly)].data_ptr())
key_value_ptrs = torch.tensor(
ptr_list,
dtype=torch.uint64,
device=device,
).contiguous()
# ── 3. Execute ──
xfer_dir = (
ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D
)
ops.multi_layer_kv_transfer_unilateral(
lmc_tensor,
key_value_ptrs,
slot_mapping,
torch.device(device),
page_buffer_size,
xfer_dir,
engine_kv_format,
)
device_sync(device)
# ── 4. Verify ──
for t_id in range(num_tokens):
s_idx = int(slot_mapping[t_id].item())
for ly in range(num_layers):
for kv in range(k_or_v_size):
lmc_val = lmc_tensor[kv, ly, t_id]
if is_mla:
paged_val = page_buffers[ly][s_idx]
else:
paged_val = buffers[(kv, ly)][s_idx]
torch.testing.assert_close(
lmc_val.to("cpu"),
paged_val.to("cpu"),
msg=(
f"Mismatch: {engine_kv_format} {dir_tag} "
f"(tensor_list={use_tensor_list}), "
f"KV={kv}, layer={ly}, "
f"token={t_id}, slot={s_idx}"
),
)
# ── 5. Collect canonical result for cross-backend comparison ──
# Use non-MLA unilateral (the primary use case of this function)
results: dict[str, torch.Tensor] = {}
for direction in [True, False]:
dir_tag = "p2l" if direction else "l2p"
lmc_shape = (2, num_layers, num_tokens, head_size)
lmc_tensor = torch.zeros(lmc_shape, dtype=dtype, device=device)
if not direction:
for kv in range(2):
for ly in range(num_layers):
for t in range(num_tokens):
val = (
kv * 5000
+ ly * 1000
+ t * 10
+ torch.arange(head_size, device=device)
).to(dtype)
lmc_tensor[kv, ly, t] = val
buffers = {}
for kv in range(2):
for ly in range(num_layers):
pb = torch.zeros(
(page_buffer_size, head_size),
dtype=dtype,
device=device,
)
if direction:
for s in range(page_buffer_size):
val = (
kv * 7000
+ ly * 2000
+ s * 10
+ torch.arange(head_size, device=device)
).to(dtype)
pb[s] = val
buffers[(kv, ly)] = pb
key_value_ptrs: Union[list[torch.Tensor], torch.Tensor] # type: ignore[no-redef]
if use_tensor_list:
# Tensor list mode: [K_l0, K_l1, ..., V_l0, V_l1, ...]
tensor_list = []
for ly in range(num_layers):
tensor_list.append(buffers[(0, ly)])
for ly in range(num_layers):
tensor_list.append(buffers[(1, ly)])
key_value_ptrs = tensor_list
else:
ptr_list = []
for ly in range(num_layers):
ptr_list.append(buffers[(0, ly)].data_ptr())
for ly in range(num_layers):
ptr_list.append(buffers[(1, ly)].data_ptr())
key_value_ptrs = torch.tensor(
ptr_list,
dtype=torch.uint64,
device=device,
).contiguous()
xfer_dir = ops.TransferDirection.D2H if direction else ops.TransferDirection.H2D
ops.multi_layer_kv_transfer_unilateral(
lmc_tensor,
key_value_ptrs,
slot_mapping,
torch.device(device),
page_buffer_size,
xfer_dir,
ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS,
)
device_sync(device)
results[f"multi_layer_kv_transfer_unilateral_{dir_tag}"] = lmc_tensor.cpu()
return results
def scenario_alloc_free_pinned_ptr(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test alloc_pinned_ptr and free_pinned_ptr round-trip."""
alloc_size = 4096
flags = 0 # cudaHostAllocDefault
# 1. Allocate
ptr = ops.alloc_pinned_ptr(alloc_size, flags)
assert isinstance(ptr, int), f"Expected int, got {type(ptr)}"
assert ptr != 0, "alloc_pinned_ptr returned null"
# 2. Free
ops.free_pinned_ptr(ptr)
return {"alloc_free_pinned_ptr": torch.tensor([1], dtype=torch.int32)}
def scenario_alloc_free_numa_ptr(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test alloc_numa_ptr and free_numa_ptr round-trip."""
alloc_size = 4096
node = 0 # NUMA node 0 (always exists)
# 1. Allocate
ptr = ops.alloc_numa_ptr(alloc_size, node)
assert isinstance(ptr, int), f"Expected int, got {type(ptr)}"
assert ptr != 0, "alloc_numa_ptr returned null"
# 2. Free (must pass same size as alloc)
ops.free_numa_ptr(ptr, alloc_size)
return {"alloc_free_numa_ptr": torch.tensor([1], dtype=torch.int32)}
def scenario_alloc_free_pinned_numa_ptr(
ops: Any, device: str
) -> dict[str, torch.Tensor]:
"""Test alloc_pinned_numa_ptr and free_pinned_numa_ptr round-trip."""
alloc_size = 4096
node = 0 # NUMA node 0
# 1. Allocate (NUMA + cudaHostRegister)
ptr = ops.alloc_pinned_numa_ptr(alloc_size, node)
assert isinstance(ptr, int), f"Expected int, got {type(ptr)}"
assert ptr != 0, "alloc_pinned_numa_ptr returned null"
# 2. Free (cudaHostUnregister + munmap)
ops.free_pinned_numa_ptr(ptr, alloc_size)
return {"alloc_free_pinned_numa_ptr": torch.tensor([1], dtype=torch.int32)}
def scenario_alloc_free_shm_pinned_ptr(
ops: Any, device: str
) -> dict[str, torch.Tensor]:
"""Test alloc_shm_pinned_ptr and free_shm_pinned_ptr round-trip."""
alloc_size = 4096
shm_name = "/test_lmcache_shm"
# 1. Allocate
ptr = ops.alloc_shm_pinned_ptr(alloc_size, shm_name)
assert isinstance(ptr, int), f"Expected int, got {type(ptr)}"
assert ptr != 0, "alloc_shm_pinned_ptr returned null"
# 2. Free
ops.free_shm_pinned_ptr(ptr, alloc_size, shm_name)
return {"alloc_free_shm_pinned_ptr": torch.tensor([1], dtype=torch.int32)}
def scenario_transfer_direction_enum(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test TransferDirection enum has distinct H2D and D2H members."""
# 1. Verify enum members exist
td = ops.TransferDirection
assert hasattr(td, "H2D"), "Missing TransferDirection.H2D"
assert hasattr(td, "D2H"), "Missing TransferDirection.D2H"
# 2. Verify values are distinct
assert td.H2D != td.D2H, "H2D and D2H should be distinct"
# 3. Extract int value (compatible with both
# pybind11 enum and Python enum)
h2d = td.H2D
d2h = td.D2H
h2d_val = h2d.value if hasattr(h2d, "value") else int(h2d)
d2h_val = d2h.value if hasattr(d2h, "value") else int(d2h)
return {
"transfer_direction_enum": torch.tensor(
[h2d_val, d2h_val],
dtype=torch.int32,
)
}
def scenario_record_drain_completion(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test record_completion_on_stream / drain_recorded_completions contracts.
Verified backend-agnostic: native c_ops uses cudaLaunchHostFunc on the
default stream (ptr=0), which fires synchronously after device_sync; the
fallback enqueues immediately. Both paths satisfy every assertion below.
"""
ops.drain_recorded_completions() # clear residual global state
assert ops.drain_recorded_completions() == []
ops.record_completion_on_stream(0, "kind-a", b"payload-a")
ops.record_completion_on_stream(0, "kind-b", b"payload-b")
device_sync(device)
result = ops.drain_recorded_completions()
assert result == [("kind-a", b"payload-a"), ("kind-b", b"payload-b")]
assert all(isinstance(k, str) and isinstance(p, bytes) for k, p in result)
assert ops.drain_recorded_completions() == []
# Multiple records on the default stream (ptr=0) must all enqueue.
# Note: ptr=0 is the only value safe on both backends — the fallback
# ignores the field, but the native path casts it to cudaStream_t, so
# arbitrary values like -1 / 2**32 would be invalid there.
for _ in range(3):
ops.record_completion_on_stream(0, "k", b"v")
device_sync(device)
assert len(ops.drain_recorded_completions()) == 3
return {"record_drain_completion": torch.tensor([1], dtype=torch.int32)}
def scenario_dispatcher_integration(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test submit_callback_to_stream -> DeviceHostFuncDispatcher -> handler.
Works on all backends: submit_callback_to_stream only reads stream.ptr, so
_FakeStream(ptr=0) is accepted; the native path routes through the CUDA
default stream which fires before the dispatcher's drain loop polls.
"""
# First Party
import lmcache.v1.multiprocess.native_completion as nc
original = nc._lmc_ops
nc._lmc_ops = ops
try:
ops.drain_recorded_completions()
class _FakeStream:
ptr: int = 0
dispatcher = DeviceHostFuncDispatcher(drain_interval_seconds=0.001)
received: list[list[bytes]] = []
dispatcher.register("finish_write", received.append, payload_type=list[bytes])
dispatcher.start()
try:
submit_callback_to_stream(_FakeStream(), "finish_write", [b"k0", b"k1"])
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline and not received:
time.sleep(0.01)
finally:
dispatcher.stop()
assert received == [[b"k0", b"k1"]]
finally:
nc._lmc_ops = original
return {"dispatcher_integration": torch.tensor([1], dtype=torch.int32)}
def scenario_multi_layer_block_kv_transfer(
ops: Any, device: str
) -> dict[str, torch.Tensor]:
"""Test multi_layer_block_kv_transfer across all GPU KV formats.
Exercises the block-based transfer path for:
- NHD per-layer format (D2H and H2D round-trip)
- FlashInfer NHD per-layer format (interleaved K/V)
- HND per-layer format (with permute)
- FlashInfer HND per-layer format (interleaved K/V)
- MLA per-layer format (no K/V split)
- SGLang MLA flat per-layer format
- Cross-layer NHD (single tensor [NB, NL, 2, BS, NH, HS])
- Cross-layer HND (single tensor [NB, NL, 2, NH, BS, HS])
- SGLang MHA flat (2*NL tensors [NB*BS, NH, HS])
- SGLang MHA block (2*NL tensors [NB, BS, NH, HS])
- non-sequential block_ids and list[int] block_ids input
- skip_prefix_n_blocks > 0
- num_blocks > blocks_per_chunk (multi-chunk)
"""
results = {}
# C++ bindings (cuda_c_ops, xpu_sycl_ops) expect uint64 pointer tensors for
# paged_buffer_ptrs_tensor and list[int] for lmcache_objects_ptrs.
# The Python fallback also supports both modes on cpu/cuda (pointer inputs
# are reconstructed internally via _tensor_from_ptr).
use_tensor_list = device not in ("cpu", "cuda")
def _alloc_chunks(shape: tuple[int, ...], count: int) -> list[torch.Tensor]:
chunks = [torch.zeros(shape, dtype=dtype) for _ in range(count)]
if device in ("cuda"):
chunks = [chunk.pin_memory() for chunk in chunks]
return chunks
# --- NHD per-layer ---
torch.manual_seed(123)
num_layers, num_blocks, block_size = 2, 8, 4
num_heads, head_size = 2, 8
blocks_per_chunk = 4
chunk_tokens = blocks_per_chunk * block_size
hidden_dim = num_heads * head_size
dtype = torch.float32
paged_layers = [
torch.randn(2, num_blocks, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
shape_desc = ops.PageBufferShapeDesc()
shape_desc.nl = num_layers
shape_desc.nb = num_blocks
shape_desc.bs = block_size
shape_desc.nh = num_heads
shape_desc.hs = head_size
shape_desc.element_size = dtype.itemsize
shape_desc.kv_size = 2
engine_kv_format = ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
num_chunks = num_blocks // blocks_per_chunk
d2h_chunks = _alloc_chunks((2, num_layers, chunk_tokens, hidden_dim), num_chunks)
block_ids = list(range(num_blocks))
ops.multi_layer_block_kv_transfer(
paged_layers
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers],
dtype=torch.uint64,
device=device,
),
d2h_chunks if use_tensor_list else [c.data_ptr() for c in d2h_chunks],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format,
0,
)
paged_h2d = [torch.zeros_like(layer) for layer in paged_layers]
ops.multi_layer_block_kv_transfer(
paged_h2d
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d], dtype=torch.uint64, device=device
),
d2h_chunks if use_tensor_list else [c.data_ptr() for c in d2h_chunks],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format,
0,
)
for i in range(num_layers):
orig = paged_layers[i].cpu()
recon = paged_h2d[i].cpu()
results[f"nhd_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"NHD Layer {i} round-trip mismatch"
)
# --- FlashInfer NHD per-layer ---
torch.manual_seed(234)
paged_layers_fi_nhd = [
torch.randn(num_blocks, 2, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
engine_kv_format_fi_nhd = ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS
d2h_chunks_fi_nhd = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_layers_fi_nhd
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_fi_nhd],
dtype=torch.uint64,
device=device,
),
d2h_chunks_fi_nhd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_fi_nhd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_fi_nhd,
0,
)
paged_h2d_fi_nhd = [torch.zeros_like(layer) for layer in paged_layers_fi_nhd]
ops.multi_layer_block_kv_transfer(
paged_h2d_fi_nhd
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_fi_nhd],
dtype=torch.uint64,
device=device,
),
d2h_chunks_fi_nhd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_fi_nhd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_fi_nhd,
0,
)
for i in range(num_layers):
orig = paged_layers_fi_nhd[i].cpu()
recon = paged_h2d_fi_nhd[i].cpu()
results[f"flashinfer_nhd_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"FlashInfer NHD Layer {i} round-trip mismatch"
)
# --- HND per-layer ---
torch.manual_seed(456)
paged_layers_hnd = [
torch.randn(2, num_blocks, num_heads, block_size, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
engine_kv_format_hnd = ops.EngineKVFormat.NL_X_TWO_NB_NH_BS_HS
d2h_chunks_hnd = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_layers_hnd
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_hnd],
dtype=torch.uint64,
device=device,
),
d2h_chunks_hnd if use_tensor_list else [c.data_ptr() for c in d2h_chunks_hnd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_hnd,
0,
)
paged_h2d_hnd = [torch.zeros_like(layer) for layer in paged_layers_hnd]
ops.multi_layer_block_kv_transfer(
paged_h2d_hnd
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_hnd],
dtype=torch.uint64,
device=device,
),
d2h_chunks_hnd if use_tensor_list else [c.data_ptr() for c in d2h_chunks_hnd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_hnd,
0,
)
for i in range(num_layers):
orig = paged_layers_hnd[i].cpu()
recon = paged_h2d_hnd[i].cpu()
results[f"hnd_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"HND Layer {i} round-trip mismatch"
)
# --- FlashInfer HND per-layer ---
torch.manual_seed(567)
paged_layers_fi_hnd = [
torch.randn(num_blocks, 2, num_heads, block_size, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
engine_kv_format_fi_hnd = ops.EngineKVFormat.NL_X_NB_TWO_NH_BS_HS
d2h_chunks_fi_hnd = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_layers_fi_hnd
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_fi_hnd],
dtype=torch.uint64,
device=device,
),
d2h_chunks_fi_hnd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_fi_hnd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_fi_hnd,
0,
)
paged_h2d_fi_hnd = [torch.zeros_like(layer) for layer in paged_layers_fi_hnd]
ops.multi_layer_block_kv_transfer(
paged_h2d_fi_hnd
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_fi_hnd],
dtype=torch.uint64,
device=device,
),
d2h_chunks_fi_hnd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_fi_hnd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_fi_hnd,
0,
)
for i in range(num_layers):
orig = paged_layers_fi_hnd[i].cpu()
recon = paged_h2d_fi_hnd[i].cpu()
results[f"flashinfer_hnd_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"FlashInfer HND Layer {i} round-trip mismatch"
)
# --- MLA per-layer ---
torch.manual_seed(789)
mla_hidden = 16
paged_layers_mla = [
torch.randn(num_blocks, block_size, mla_hidden, dtype=dtype).to(device)
for _ in range(num_layers)
]
shape_desc_mla = ops.PageBufferShapeDesc()
shape_desc_mla.nl = num_layers
shape_desc_mla.nb = num_blocks
shape_desc_mla.bs = block_size
shape_desc_mla.nh = 1
shape_desc_mla.hs = mla_hidden
shape_desc_mla.element_size = dtype.itemsize
shape_desc_mla.kv_size = 1
engine_kv_format_mla = ops.EngineKVFormat.NL_X_NB_BS_HS
d2h_chunks_mla = _alloc_chunks((num_layers, chunk_tokens, mla_hidden), num_chunks)
ops.multi_layer_block_kv_transfer(
paged_layers_mla
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_mla],
dtype=torch.uint64,
device=device,
),
d2h_chunks_mla if use_tensor_list else [c.data_ptr() for c in d2h_chunks_mla],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc_mla,
chunk_tokens,
engine_kv_format_mla,
0,
)
paged_h2d_mla = [torch.zeros_like(layer) for layer in paged_layers_mla]
ops.multi_layer_block_kv_transfer(
paged_h2d_mla
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_mla],
dtype=torch.uint64,
device=device,
),
d2h_chunks_mla if use_tensor_list else [c.data_ptr() for c in d2h_chunks_mla],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc_mla,
chunk_tokens,
engine_kv_format_mla,
0,
)
for i in range(num_layers):
orig = paged_layers_mla[i].cpu()
recon = paged_h2d_mla[i].cpu()
results[f"mla_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"MLA Layer {i} round-trip mismatch"
)
# --- SGLang MLA flat per-layer ---
torch.manual_seed(890)
paged_layers_sglang_mla = [
torch.randn(num_blocks * block_size, 1, mla_hidden, dtype=dtype).to(device)
for _ in range(num_layers)
]
engine_kv_format_sglang_mla = ops.EngineKVFormat.NL_X_NBBS_ONE_HS
d2h_chunks_sglang_mla = _alloc_chunks(
(num_layers, chunk_tokens, mla_hidden), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_layers_sglang_mla
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_sglang_mla],
dtype=torch.uint64,
device=device,
),
d2h_chunks_sglang_mla
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_sglang_mla],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc_mla,
chunk_tokens,
engine_kv_format_sglang_mla,
0,
)
paged_h2d_sglang_mla = [
torch.zeros_like(layer) for layer in paged_layers_sglang_mla
]
ops.multi_layer_block_kv_transfer(
paged_h2d_sglang_mla
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_sglang_mla],
dtype=torch.uint64,
device=device,
),
d2h_chunks_sglang_mla
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_sglang_mla],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc_mla,
chunk_tokens,
engine_kv_format_sglang_mla,
0,
)
for i in range(num_layers):
orig = paged_layers_sglang_mla[i].cpu()
recon = paged_h2d_sglang_mla[i].cpu()
results[f"sglang_mla_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"SGLang MLA Layer {i} round-trip mismatch"
)
# --- Cross-layer NHD ---
torch.manual_seed(101)
paged_cross_nhd = torch.randn(
num_blocks,
num_layers,
2,
block_size,
num_heads,
head_size,
dtype=dtype,
).to(device)
engine_kv_format_cross_nhd = ops.EngineKVFormat.NB_NL_TWO_BS_NH_HS
d2h_chunks_cross_nhd = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_cross_nhd
if use_tensor_list
else torch.tensor(
[paged_cross_nhd.data_ptr()], dtype=torch.uint64, device=device
),
d2h_chunks_cross_nhd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_cross_nhd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_cross_nhd,
0,
)
paged_h2d_cross_nhd = torch.zeros_like(paged_cross_nhd)
ops.multi_layer_block_kv_transfer(
paged_h2d_cross_nhd
if use_tensor_list
else torch.tensor(
[paged_h2d_cross_nhd.data_ptr()], dtype=torch.uint64, device=device
),
d2h_chunks_cross_nhd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_cross_nhd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_cross_nhd,
0,
)
orig = paged_cross_nhd.cpu()
recon = paged_h2d_cross_nhd.cpu()
results["cross_nhd"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), "Cross-layer NHD mismatch"
# --- Cross-layer HND ---
torch.manual_seed(202)
paged_cross_hnd = torch.randn(
num_blocks,
num_layers,
2,
num_heads,
block_size,
head_size,
dtype=dtype,
).to(device)
engine_kv_format_cross_hnd = ops.EngineKVFormat.NB_NL_TWO_NH_BS_HS
d2h_chunks_cross_hnd = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_cross_hnd
if use_tensor_list
else torch.tensor(
[paged_cross_hnd.data_ptr()], dtype=torch.uint64, device=device
),
d2h_chunks_cross_hnd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_cross_hnd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_cross_hnd,
0,
)
paged_h2d_cross_hnd = torch.zeros_like(paged_cross_hnd)
ops.multi_layer_block_kv_transfer(
paged_h2d_cross_hnd
if use_tensor_list
else torch.tensor(
[paged_h2d_cross_hnd.data_ptr()], dtype=torch.uint64, device=device
),
d2h_chunks_cross_hnd
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_cross_hnd],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_cross_hnd,
0,
)
orig = paged_cross_hnd.cpu()
recon = paged_h2d_cross_hnd.cpu()
results["cross_hnd"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), "Cross-layer HND mismatch"
# --- SGLang MHA flat (NBBS) ---
torch.manual_seed(303)
pbs = num_blocks * block_size
paged_sglang_nbbs = [
[
torch.randn(pbs, num_heads, head_size, dtype=dtype).to(device)
for _ in range(num_layers)
],
[
torch.randn(pbs, num_heads, head_size, dtype=dtype).to(device)
for _ in range(num_layers)
],
]
engine_kv_format_sglang_nbbs = ops.EngineKVFormat.TWO_X_NL_X_NBBS_NH_HS
d2h_chunks_sglang_nbbs = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_sglang_nbbs
if use_tensor_list
else torch.tensor(
[t.data_ptr() for t in paged_sglang_nbbs[0]]
+ [t.data_ptr() for t in paged_sglang_nbbs[1]],
dtype=torch.uint64,
device=device,
),
d2h_chunks_sglang_nbbs
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_sglang_nbbs],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_sglang_nbbs,
0,
)
paged_h2d_sglang_nbbs = [
[torch.zeros_like(t) for t in group] for group in paged_sglang_nbbs
]
ops.multi_layer_block_kv_transfer(
paged_h2d_sglang_nbbs
if use_tensor_list
else torch.tensor(
[t.data_ptr() for t in paged_h2d_sglang_nbbs[0]]
+ [t.data_ptr() for t in paged_h2d_sglang_nbbs[1]],
dtype=torch.uint64,
device=device,
),
d2h_chunks_sglang_nbbs
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_sglang_nbbs],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_sglang_nbbs,
0,
)
for kv in range(2):
for i in range(num_layers):
orig = paged_sglang_nbbs[kv][i].cpu()
recon = paged_h2d_sglang_nbbs[kv][i].cpu()
results[f"sglang_nbbs_kv{kv}_l{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"SGLang NBBS kv={kv} layer={i} mismatch"
)
# --- SGLang MHA block (NB_BS) ---
torch.manual_seed(404)
paged_sglang_nb = [
[
torch.randn(num_blocks, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
],
[
torch.randn(num_blocks, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
],
]
engine_kv_format_sglang_nb = ops.EngineKVFormat.TWO_X_NL_X_NB_BS_NH_HS
d2h_chunks_sglang_nb = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_sglang_nb
if use_tensor_list
else torch.tensor(
[t.data_ptr() for t in paged_sglang_nb[0]]
+ [t.data_ptr() for t in paged_sglang_nb[1]],
dtype=torch.uint64,
device=device,
),
d2h_chunks_sglang_nb
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_sglang_nb],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_sglang_nb,
0,
)
paged_h2d_sglang_nb = [
[torch.zeros_like(t) for t in group] for group in paged_sglang_nb
]
ops.multi_layer_block_kv_transfer(
paged_h2d_sglang_nb
if use_tensor_list
else torch.tensor(
[t.data_ptr() for t in paged_h2d_sglang_nb[0]]
+ [t.data_ptr() for t in paged_h2d_sglang_nb[1]],
dtype=torch.uint64,
device=device,
),
d2h_chunks_sglang_nb
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_sglang_nb],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_sglang_nb,
0,
)
for kv in range(2):
for i in range(num_layers):
orig = paged_sglang_nb[kv][i].cpu()
recon = paged_h2d_sglang_nb[kv][i].cpu()
results[f"sglang_nb_kv{kv}_l{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"SGLang NB kv={kv} layer={i} mismatch"
)
# --- skip_prefix_n_blocks > 0 ---
torch.manual_seed(505)
skip_n = 2
paged_layers_skip = [
torch.randn(2, num_blocks, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
engine_kv_format_nhd = ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
# With skip=2, effective blocks start at index 2.
# Object 0 occupies flat indices [0, blocks_per_chunk), skipping first 2.
# Object 1 occupies flat indices [blocks_per_chunk, 2*blocks_per_chunk).
d2h_chunks_skip = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_layers_skip
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_skip],
dtype=torch.uint64,
device=device,
),
d2h_chunks_skip if use_tensor_list else [c.data_ptr() for c in d2h_chunks_skip],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_nhd,
skip_n,
)
paged_h2d_skip = [torch.zeros_like(layer) for layer in paged_layers_skip]
ops.multi_layer_block_kv_transfer(
paged_h2d_skip
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_skip],
dtype=torch.uint64,
device=device,
),
d2h_chunks_skip if use_tensor_list else [c.data_ptr() for c in d2h_chunks_skip],
torch.tensor(block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_nhd,
skip_n,
)
# Blocks [skip_n:num_blocks] should round-trip at their original positions
for layer_idx in range(num_layers):
for kv in range(2):
orig_blocks = paged_layers_skip[layer_idx][kv, skip_n:num_blocks]
recon_blocks = paged_h2d_skip[layer_idx][kv, skip_n:num_blocks]
key = f"skip_l{layer_idx}_kv{kv}"
results[key] = torch.stack([orig_blocks.cpu(), recon_blocks.cpu()])
assert torch.allclose(orig_blocks, recon_blocks, atol=1e-6), (
f"Skip mismatch at layer={layer_idx} kv={kv}"
)
# Skipped blocks [0:skip_n] should remain zero
for layer_idx in range(num_layers):
for kv in range(2):
skipped = paged_h2d_skip[layer_idx][kv, :skip_n]
assert torch.all(skipped == 0), (
f"Skipped blocks not zero at layer={layer_idx} kv={kv}"
)
assert torch.all(d2h_chunks_skip[0][:, :, : skip_n * block_size] == 0), (
"Skipped D2H chunk region should remain zero"
)
# --- Non-sequential block_ids ---
torch.manual_seed(515)
paged_layers_permuted = [
torch.randn(2, num_blocks, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
generator = torch.Generator(device="cpu").manual_seed(616)
permuted_block_ids = torch.randperm(num_blocks, generator=generator).tolist()
d2h_chunks_permuted = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks
)
ops.multi_layer_block_kv_transfer(
paged_layers_permuted
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_permuted],
dtype=torch.uint64,
device=device,
),
d2h_chunks_permuted
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_permuted],
torch.tensor(permuted_block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc,
chunk_tokens,
engine_kv_format_nhd,
0,
)
paged_h2d_permuted = [torch.zeros_like(layer) for layer in paged_layers_permuted]
ops.multi_layer_block_kv_transfer(
paged_h2d_permuted
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_permuted],
dtype=torch.uint64,
device=device,
),
d2h_chunks_permuted
if use_tensor_list
else [c.data_ptr() for c in d2h_chunks_permuted],
torch.tensor(permuted_block_ids, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc,
chunk_tokens,
engine_kv_format_nhd,
0,
)
for i in range(num_layers):
orig = paged_layers_permuted[i].cpu()
recon = paged_h2d_permuted[i].cpu()
results[f"permuted_nhd_layer{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"Permuted NHD Layer {i} round-trip mismatch"
)
# --- Multi-chunk (num_blocks > blocks_per_chunk) ---
torch.manual_seed(606)
num_blocks_mc = 12
paged_layers_mc = [
torch.randn(2, num_blocks_mc, block_size, num_heads, head_size, dtype=dtype).to(
device
)
for _ in range(num_layers)
]
shape_desc_mc = ops.PageBufferShapeDesc()
shape_desc_mc.nl = num_layers
shape_desc_mc.nb = num_blocks_mc
shape_desc_mc.bs = block_size
shape_desc_mc.nh = num_heads
shape_desc_mc.hs = head_size
shape_desc_mc.element_size = dtype.itemsize
shape_desc_mc.kv_size = 2
num_chunks_mc = num_blocks_mc // blocks_per_chunk # 3 chunks
d2h_chunks_mc = _alloc_chunks(
(2, num_layers, chunk_tokens, hidden_dim), num_chunks_mc
)
block_ids_mc = list(range(num_blocks_mc))
ops.multi_layer_block_kv_transfer(
paged_layers_mc
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_layers_mc],
dtype=torch.uint64,
device=device,
),
d2h_chunks_mc if use_tensor_list else [c.data_ptr() for c in d2h_chunks_mc],
torch.tensor(block_ids_mc, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.D2H,
shape_desc_mc,
chunk_tokens,
engine_kv_format_nhd,
0,
)
paged_h2d_mc = [torch.zeros_like(layer) for layer in paged_layers_mc]
ops.multi_layer_block_kv_transfer(
paged_h2d_mc
if use_tensor_list
else torch.tensor(
[layer.data_ptr() for layer in paged_h2d_mc],
dtype=torch.uint64,
device=device,
),
d2h_chunks_mc if use_tensor_list else [c.data_ptr() for c in d2h_chunks_mc],
torch.tensor(block_ids_mc, dtype=torch.int64, device=device),
torch.device(device),
ops.TransferDirection.H2D,
shape_desc_mc,
chunk_tokens,
engine_kv_format_nhd,
0,
)
for i in range(num_layers):
orig = paged_layers_mc[i].cpu()
recon = paged_h2d_mc[i].cpu()
results[f"multi_chunk_l{i}"] = torch.stack([orig, recon])
assert torch.allclose(orig, recon, atol=1e-6), (
f"Multi-chunk Layer {i} round-trip mismatch"
)
return results
def scenario_record_drain_event(ops: Any, device: str) -> dict[str, torch.Tensor]:
"""Test record_event_on_stream / drain_recorded_events contracts.
Verified backend-agnostic: native c_ops uses cudaLaunchHostFunc on the
default stream (ptr=0), which fires synchronously after device_sync; the
fallback enqueues immediately with time.time(). Both paths satisfy every
assertion below.
"""
ops.drain_recorded_events() # clear residual global state
assert ops.drain_recorded_events() == []
ops.record_event_on_stream(
0, "mp.store.start", "sess-1", {"device": "cuda:0"}, {"count": 10}
)
ops.record_event_on_stream(
0, "mp.store.end", "sess-1", {"device": "cuda:0"}, {"count": 10}
)
device_sync(device)
result = ops.drain_recorded_events()
assert len(result) == 2
# Each element is (event_type_name, session_id, timestamp, str_meta, int_meta)
name0, sid0, ts0, str_meta0, int_meta0 = result[0]
name1, sid1, ts1, str_meta1, int_meta1 = result[1]
assert name0 == "mp.store.start"
assert name1 == "mp.store.end"
assert sid0 == "sess-1"
assert sid1 == "sess-1"
assert ts0 > 0.0
assert ts1 >= ts0 # timestamps must be monotonically non-decreasing
assert str_meta0 == {"device": "cuda:0"}
assert int_meta0 == {"count": 10}
assert str_meta1 == {"device": "cuda:0"}
assert int_meta1 == {"count": 10}
# Drain clears the buffer
assert ops.drain_recorded_events() == []
# Multiple records on the default stream (ptr=0) must all enqueue.
for i in range(5):
ops.record_event_on_stream(0, f"mp.test.{i}", f"sess-{i}", {}, {"idx": i})
device_sync(device)
events = ops.drain_recorded_events()
assert len(events) == 5
for i, (name, sid, ts, str_meta, int_meta) in enumerate(events):
assert name == f"mp.test.{i}"
assert sid == f"sess-{i}"
assert int_meta == {"idx": i}
assert ts > 0.0
return {"record_drain_event": torch.tensor([1], dtype=torch.int32)}
# ==========================================
# 3. Registry
# ==========================================
# cover pybind list in csrc/pybind.cpp
SCENARIO_REGISTRY = {
"transfer_direction_enum": scenario_transfer_direction_enum,
"multi_layer_kv_transfer": scenario_multi_layer_kv_transfer,
"multi_layer_kv_transfer_unilateral": scenario_multi_layer_kv_transfer_unilateral,
"multi_layer_block_kv_transfer": scenario_multi_layer_block_kv_transfer,
"single_layer_kv_transfer": scenario_single_layer_kv_transfer,
"single_layer_kv_transfer_sgl": scenario_single_layer_kv_transfer_sgl,
"load_and_reshape_flash": scenario_load_and_reshape_flash,
"reshape_and_cache_back_flash": scenario_reshape_and_cache_back_flash,
"lmcache_memcpy_async": scenario_lmcache_memcpy_async,
"encode_fast_new": scenario_encode_fast_new,
"decode_fast_new": scenario_decode_fast_new,
"decode_fast_prefsum": scenario_decode_fast_prefsum,
"calculate_cdf": scenario_calculate_cdf,
"rotary_embedding_k_fused": scenario_rotary_embedding_k_fused,
"alloc_free_pinned_ptr": scenario_alloc_free_pinned_ptr,
"alloc_free_pinned_numa_ptr": scenario_alloc_free_pinned_numa_ptr,
"alloc_free_numa_ptr": scenario_alloc_free_numa_ptr,
"alloc_free_shm_pinned_ptr": scenario_alloc_free_shm_pinned_ptr,
"get_gpu_pci_bus_id": scenario_get_gpu_pci_bus_id,
"record_drain_completion": scenario_record_drain_completion,
"dispatcher_integration": scenario_dispatcher_integration,
"record_drain_event": scenario_record_drain_event,
}
# ==========================================
# 4. Test functions pytest sees
# ==========================================
class TestScenarios:
"""Test class to ensure test_scenario runs before test_compare.
**Execution Order Guarantee:**
By grouping tests in a class and using alphabetically ordered method names
(test_1_scenario, test_2_compare), pytest will execute all scenario tests
before any compare tests, ensuring _results dict is fully populated.
"""
@pytest.mark.parametrize("name,fn", list(SCENARIO_REGISTRY.items()))
def test_1_scenario(
self,
backend: tuple,
name: str,
fn: Any,
) -> None:
"""Run a single scenario with a specific backend configuration.
Each (scenario, backend) pair is a separate pytest test case, giving
per-scenario per-backend failure reporting without subprocess indirection.
Results are stored in the module-level _results dict for later comparison
by test_2_compare.
"""
backend_id, ops, device = backend
result = fn(ops, device)
if result is not None:
_results[(name, backend_id)] = result
@pytest.mark.parametrize("name", list(SCENARIO_REGISTRY.keys()))
def test_2_compare(self, name: str) -> None:
"""Compare results across backends for a single scenario.
When multiple backends ran (CUDA available), asserts that python
fallback equivalents produce numerically identical results to cuda_ops.
When only one backend ran (no CUDA), simply verifies results were stored.
This test runs after test_1_scenario due to alphabetical ordering of
method names within the TestScenarios class.
"""
available = [p.values[0][0] for p in _BACKEND_PARAMS] # list of backend_ids
backend_results = {
bid: _results[(name, bid)] for bid in available if (name, bid) in _results
}
assert len(backend_results) >= 1, (
f"{name}: no results collected — were all test_1_scenario tests skipped?"
)
if len(backend_results) < 2:
# Only one backend available (no CUDA); just verify results exist
return
# Collect all result keys across all backends for this scenario
all_keys: set[str] = set()
for res in backend_results.values():
all_keys.update(res.keys())
base_bid = next(iter(backend_results)) # first available backend as reference
for key in sorted(all_keys):
base_val = backend_results[base_bid].get(key)
if base_val is None:
continue
for bid, res in backend_results.items():
if bid == base_bid:
continue
val = res.get(key)
if val is None:
pytest.fail(
f"{name}/{key}: backend '{bid}' has no result "
f"(reference backend '{base_bid}' does)"
)
continue
if isinstance(val, torch.Tensor):
v_current = val.detach().cpu().float()
v_base = base_val.detach().cpu().float()
if not torch.allclose(v_current, v_base, rtol=1e-4, atol=1e-4):
max_diff = (v_current - v_base).abs().max().item()
pytest.fail(
f"{name}/{key}: '{bid}' vs '{base_bid}' mismatch, "
f"max diff = {max_diff:.2e}"
)
else:
if val != base_val:
pytest.fail(
f"{name}/{key}: '{bid}'={val} != '{base_bid}'={base_val}"
)
# ==========================================
# Allocation page alignment
# ==========================================
#
# Rust raw-block backend with O_DIRECT requires page-aligned buffer
# pointers; CUDA path gets this for free via cudaHostAlloc, and the
# non-CUDA fallback in lmcache.non_cuda_equivalents shall mirror the same
# guarantee.
_PINNED_ALLOC_SIZES = [1, 4095, 4096, 8192, 1024 * 1024]
@pytest.mark.parametrize("size", _PINNED_ALLOC_SIZES)
def test_alloc_pinned_ptr_is_page_aligned(size: int) -> None:
page_size = os.sysconf("SC_PAGESIZE")
ptr = _py_ops.alloc_pinned_ptr(size)
try:
assert ptr != 0
if ptr % page_size != 0:
raise AssertionError(
f"alloc_pinned_ptr({size}) returned non-page-aligned ptr "
f"{hex(ptr)} (page size {page_size})"
)
# Touch every byte in the requested region through the raw pointer
# to confirm the registered view covers the full requested size
# (an undersized view would corrupt adjacent memory or segfault).
buf = (ctypes.c_uint8 * size).from_address(ptr)
for i in range(size):
buf[i] = (i & 0xFF) ^ 0xA5
for i in range(size):
assert buf[i] == ((i & 0xFF) ^ 0xA5)
finally:
_py_ops.free_pinned_ptr(ptr)
Reference in New Issue View Git Blame Copy Permalink