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/lmcache/python_ops_fallback.py
T
wehub-resource-sync f213ec8976 chore: import upstream snapshot with attribution
2026-07-13 12:24:33 +08:00

2765 lines
102 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
#
# This file contains Python non-CUDA fallback implementations for
# CUDA-specific operations.
#
# Standard
from concurrent.futures import ThreadPoolExecutor
from enum import IntEnum
from multiprocessing import shared_memory
from typing import Any, Optional, Tuple
import ctypes
import ctypes.util
import os
import threading
import warnings
# Third Party
from numba import njit
import numpy as np
import torch
# First Party
from lmcache import torch_dev, torch_device_type
from lmcache.logging import init_logger
from lmcache.v1.platform import current_device_spec
# Store the tensor objects in memory so that they can be accessed
# outside the scope of this file
_tensor_registry: dict[int, torch.Tensor] = {}
_shm_registry: dict[int, shared_memory.SharedMemory] = {}
_buf_registry: dict[int, ctypes.Array] = {}
_pinned_ptr_registry: dict[int, int] = {} # ptr -> size, for cudaHostUnregister
# Cached copy library for lmcache_memcpy_async (lazy-initialized)
_copy_lib_NOT_LOADED = object()
_copy_lib: Optional[ctypes.CDLL] = _copy_lib_NOT_LOADED # type: ignore
def _get_copy_lib() -> Optional[ctypes.CDLL]:
"""Lazily load and cache the CUDA/ROCm runtime library, or None for CPU fallback."""
global _copy_lib
if _copy_lib is _copy_lib_NOT_LOADED:
# Try to load GPU runtime libraries in priority order: CUDA first, then ROCm
# TODO: ROCm path to be validated on real device
for name, fallback in [
("cudart", "libcudart.so"), # NVIDIA CUDA Runtime
("amdhip64", "libamdhip64.so"), # AMD ROCm HIP Runtime
]:
try:
path = ctypes.util.find_library(name)
if path:
_copy_lib = ctypes.CDLL(path)
else:
_copy_lib = ctypes.CDLL(fallback)
break # Successfully loaded, stop trying
except OSError:
continue # Current library not available, try next
else:
# All GPU libraries failed to load, fall back to CPU
_copy_lib = None
return _copy_lib
def _tensor_from_ptr(
ptr: int,
shape: tuple[int, ...],
dtype: torch.dtype,
device: torch.device | str | None = None,
) -> torch.Tensor:
"""
Create a tensor view over a raw pointer (zero-copy where possible).
Supports both CPU (pinned or regular) and CUDA device pointers.
Args:
ptr: Raw memory pointer as int (must be non-zero).
shape: Desired tensor shape.
dtype: Desired tensor dtype, must match the memory layout.
device: Where the pointer lives.
- None / "cpu" / torch.device("cpu") → CPU pointer
- "cuda" / "cuda:N" / torch.device("cuda", N) → CUDA pointer
If None and ptr looks like a CUDA ptr, pass device explicitly.
Returns:
A tensor that shares memory with the original pointer.
For CPU: always zero-copy via ctypes + torch.frombuffer.
For CUDA: zero-copy via torch._C._construct_storage_from_data_pointer
(PyTorch >= 2.0) or __cuda_array_interface__, with a
cudaMemcpy D2D fallback.
Raises:
ValueError: if ptr is 0.
Warning:
The caller is responsible for keeping the underlying memory alive
for the entire lifetime of the returned tensor.
"""
if ptr == 0:
raise ValueError("Pointer must be non-zero")
# ------------------------------------------------------------------ #
# Normalise device #
# ------------------------------------------------------------------ #
if device is None:
device = torch.device("cpu")
elif not isinstance(device, torch.device):
device = torch.device(device)
assert isinstance(device, torch.device)
# ------------------------------------------------------------------ #
# Compute size #
# ------------------------------------------------------------------ #
numel = 1
for dim in shape:
numel *= int(dim)
element_size = torch.empty((), dtype=dtype).element_size()
total_bytes = numel * element_size
# ------------------------------------------------------------------ #
# CPU path #
# ------------------------------------------------------------------ #
if device.type == "cpu":
return _tensor_from_cpu_ptr(ptr, shape, dtype, numel, total_bytes)
# ------------------------------------------------------------------ #
# CUDA path #
# ------------------------------------------------------------------ #
if device.type == "cuda":
return _tensor_from_cuda_ptr(ptr, shape, dtype, device, numel, total_bytes)
raise ValueError(
f"Unsupported device type: {device.type!r}. Expected 'cpu' or 'cuda'."
)
# ====================================================================== #
# CPU implementation #
# ====================================================================== #
def _tensor_from_cpu_ptr(
ptr: int,
shape: tuple[int, ...],
dtype: torch.dtype,
numel: int,
total_bytes: int,
) -> torch.Tensor:
"""
Zero-copy CPU tensor from a raw host pointer via ctypes + torch.frombuffer.
"""
buffer_type = ctypes.c_uint8 * total_bytes
buf = buffer_type.from_address(ptr)
# torch.frombuffer is zero-copy for contiguous byte buffers on CPU.
return torch.frombuffer(buf, dtype=dtype).view(*shape)
# ====================================================================== #
# CUDA implementation #
# ====================================================================== #
def _tensor_from_cuda_ptr(
ptr: int,
shape: tuple[int, ...],
dtype: torch.dtype,
device: torch.device,
numel: int,
total_bytes: int,
) -> torch.Tensor:
"""Zero-copy CUDA tensor from a raw device pointer."""
try:
_DTYPE_TO_TYPESTR = {
torch.float16: "<f2",
torch.float32: "<f4",
torch.float64: "<f8",
torch.int8: "|i1",
torch.int16: "<i2",
torch.int32: "<i4",
torch.int64: "<i8",
torch.uint8: "|u1",
torch.bool: "|b1",
}
is_bf16 = dtype == torch.bfloat16
# Determine the correct typestr, smuggle bfloat16 as int16
typestr = "<i2" if is_bf16 else _DTYPE_TO_TYPESTR.get(dtype, "|u1")
class _CudaArrayWrapper:
def __init__(self, ptr_int: int, shape_tuple: tuple, type_str: str):
self.__cuda_array_interface__ = {
"data": (ptr_int, False),
"shape": shape_tuple,
"typestr": type_str,
"version": 3,
}
t = torch.as_tensor(_CudaArrayWrapper(ptr, (numel,), typestr), device=device)
if is_bf16:
t = t.view(torch.bfloat16)
return t.view(*shape)
except Exception:
pass
# Strategy 2: cudaMemcpy Device-to-Device (Fallback)
libcudart = _get_copy_lib()
if libcudart is None:
raise RuntimeError("Failed to load libcudart/libamdhip")
cudaMemcpy = libcudart.cudaMemcpy
cudaMemcpy.restype = ctypes.c_int
cudaMemcpy.argtypes = [
ctypes.c_void_p,
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_int,
]
_MEMCPY_D2D = 3
dst = torch.empty(numel, dtype=dtype, device=device)
err = cudaMemcpy(
ctypes.c_void_p(dst.data_ptr()),
ctypes.c_void_p(ptr),
ctypes.c_size_t(total_bytes),
ctypes.c_int(_MEMCPY_D2D),
)
if err != 0:
raise RuntimeError(f"cudaMemcpy D2D failed with error code {err}.")
return dst.view(*shape)
def _copy_bytes_with_tensor(dst: int, src: int, num_bytes: int) -> None:
"""Copy raw bytes between pointers using torch tensor semantics.
Note: This function only works for CPU-accessible memory. For device
memory (CUDA/XPU), use lmcache_memcpy_async with the appropriate runtime
library or PyTorch's tensor copy operations.
"""
if num_bytes <= 0:
return
buffer_type = ctypes.c_uint8 * num_bytes
dst_tensor = torch.frombuffer(buffer_type.from_address(dst), dtype=torch.uint8)
src_tensor = torch.frombuffer(buffer_type.from_address(src), dtype=torch.uint8)
dst_tensor.copy_(src_tensor)
class TransferDirection(IntEnum):
"""Specifies the direction of a memory transfer.
Inherits from IntEnum so that members compare equal to plain ints
and to native pybind11 enum members with the same integer value.
Several call sites (and the fallback ops themselves) use
``int(direction)`` to compare across backend / fallback boundaries.
"""
H2D = 0
D2H = 1
class EngineKVFormat(IntEnum):
"""Enumeration of different engine KV cache memory layouts."""
# used by: vLLM CROSS_LAYER mode
NB_NL_TWO_BS_NH_HS = 0
# used by: vLLM non-MLA flash attention
NL_X_TWO_NB_BS_NH_HS = 1
# used by: vLLM non-MLA flash infer
NL_X_NB_TWO_BS_NH_HS = 2
# used by: vLLM MLA
NL_X_NB_BS_HS = 3
# used by: SGLang MHA (flash attention and flash infer)
TWO_X_NL_X_NBBS_NH_HS = 4
# used by: SGLang MLA
NL_X_NBBS_ONE_HS = 5
# used by: vLLM non-MLA flash attention (HND layout)
NL_X_TWO_NB_NH_BS_HS = 6
# used by: vLLM non-MLA flash infer (HND layout)
NL_X_NB_TWO_NH_BS_HS = 7
# used by: TRT-LLM cross-layer (HND layout)
NB_NL_TWO_NH_BS_HS = 8
# used by: SGLang MHA via the MP daemon path
TWO_X_NL_X_NB_BS_NH_HS = 9
# used by: vLLM non-MLA blocks-first attention with K/V fused into the
# trailing dim. Per-layer physical shape
# [num_blocks, num_heads, block_size, 2, head_size] -- the K/V "2" axis is
# second-to-last, recovered by splitting the fused [..., 2 * head_size].
NL_X_NB_NH_BS_TWO_HS = 10
# Backward-compat alias
GPUKVFormat = EngineKVFormat
class PageBufferShapeDesc:
"""Python stand-in for the C++ ``PageBufferShapeDesc`` struct.
Mirrors the pybind ``def_readwrite`` attributes in ``csrc/pybind.cpp``
so non-CUDA code paths can construct and inspect shape descriptors
without the compiled extension.
``block_stride_elems`` captures the *physical* per-block step in
element units (= ``tensor.stride(0)``). For a tightly-packed paged
buffer it equals ``bs * kv_size * nh * hs`` (non-MLA) /
``bs * nh * hs`` (MLA); for a vLLM KV pool where a group's row is
padded to the pool's maximum row width (e.g. DeepSeek V4 compressor
/ indexer caches), it is strictly larger. Downstream kernels must
use this value instead of recomputing a "tight" stride from the
logical shape, otherwise they'll skip into the next block's padding
region and read/write the wrong slots.
"""
__slots__ = (
"kv_size",
"nl",
"nb",
"bs",
"nh",
"hs",
"element_size",
"block_stride_elems",
"dtype",
)
def __init__(self) -> None:
self.kv_size: int = 0
self.nl: int = 0
self.nb: int = 0
self.bs: int = 0
self.nh: int = 0
self.hs: int = 0
self.element_size: int = 0
# 0 means "unset — fall back to tight stride"; any downstream
# consumer that needs exact addressing must check this.
self.block_stride_elems: int = 0
self.dtype: torch.dtype | None = None
class _NativePlanType:
"""Base for object-group transfer plan types that only exist natively.
The plan value structs (see ``csrc/mp_mem_kernels.cuh``) are built on the
Python side and consumed by the native ``execute_object_group_transfer``.
They have no pure-Python fallback, so constructing one without the compiled
``c_ops`` extension is unsupported. Subclasses exist only so the CPU-only
build exposes the same names as ``c_ops``.
"""
def __init__(self, *args: object, **kwargs: object) -> None:
raise NotImplementedError(
f"{type(self).__name__} requires the c_ops native extension; "
"no pure-Python fallback exists."
)
class StagingCopy(_NativePlanType):
"""Fallback stub for the native ``StagingCopy`` plan type."""
class LaunchVar(_NativePlanType):
"""Fallback stub for the native ``LaunchVar`` plan type."""
class BatchStep(_NativePlanType):
"""Fallback stub for the native ``BatchStep`` plan type."""
class KernelGroupSpec(_NativePlanType):
"""Fallback stub for the native ``KernelGroupSpec`` plan type."""
def set_shape_desc_dtype(shape_desc: Any, dtype: torch.dtype) -> None:
"""Best-effort ``shape_desc.dtype = dtype``.
The pure-Python ``PageBufferShapeDesc`` exposes a ``dtype`` slot so
the CPU fallback kernel can disambiguate float16 vs bfloat16 (both
have ``element_size == 2``). The pybind C++ struct in
``csrc/pybind.cpp`` has no such field; assignment raises
``AttributeError`` and is silently swallowed here so call sites
don't need to branch on the active backend.
Args:
shape_desc: A ``PageBufferShapeDesc`` instance (either the
pure-Python fallback or the C++ pybind struct).
dtype: The torch dtype to assign.
"""
try:
shape_desc.dtype = dtype
except AttributeError:
pass
# Cuda path goes through func cudaHostAlloc, which is
# already page aligned by CUDA spec. This fallback shim mirrors that
# guarantee so consumers that require page-aligned host buffers, in
# particular the Rust raw-block backend when O_DIRECT is enabled, which
# requires page-aligned buffer pointer
try:
_PAGE_SIZE = os.sysconf("SC_PAGESIZE")
except (AttributeError, ValueError, OSError):
_PAGE_SIZE = 4096
logger = init_logger(__name__)
# Cached one-shot decision: pin host buffers only when an accelerator is
# present. Probed lazily on first allocation; if a pinned allocation ever
# fails at runtime we flip this to False permanently and fall back to
# pageable memory for all subsequent allocations.
_use_pinned: Optional[bool] = None
def _alloc_page_aligned_pinned_view(size: int) -> Tuple[torch.Tensor, int]:
"""
Allocate a pinned CPU buffer whose first usable byte is page-aligned,
and return a torch view of ``size`` bytes plus its base pointer.
Internally over-allocates one extra page on a backing tensor, then
slices the aligned region out. The slice shares storage with the
backing tensor, so keeping the slice alive keeps the underlying
memory alive (no need to track the backing tensor separately).
"""
# Pin the host buffer when an accelerator is present (probed once).
# StubCPUDevice.is_available returns False on CPU-only hosts.
global _use_pinned
if _use_pinned is None:
_use_pinned = torch_dev.is_available()
try:
backing = torch.empty(
size + _PAGE_SIZE, dtype=torch.uint8, pin_memory=_use_pinned
)
except RuntimeError:
if not _use_pinned:
# Pure host allocation failed (e.g. OOM); nothing to fall back to.
raise
logger.warning(
"Pinned host allocation failed on device '%s'; falling back to "
"unpinned allocation from now on.",
torch_device_type,
)
_use_pinned = False
backing = torch.empty(size + _PAGE_SIZE, dtype=torch.uint8, pin_memory=False)
# First-touch initialization on the entire backing region
backing.fill_(0)
base = backing.data_ptr()
# Distance from `base` to the next page boundary (0..PAGE_SIZE-1).
offset = (-base) % _PAGE_SIZE
aligned_view = backing[offset : offset + size]
return aligned_view, aligned_view.data_ptr()
def alloc_pinned_numa_ptr(size: int, numa_id: int = 0) -> int:
"""Non-CUDA equivalent of allocating pinned memory with NUMA awareness.
On XPU, uses pin_memory=True (SYCL USM host allocation) for fast transfers.
Note: NUMA node selection is not supported on non-CUDA."""
view, aligned_ptr = _alloc_page_aligned_pinned_view(size)
# view shares storage with its over-allocated backing tensor;
# holding the view in the registry transitively keeps the underlying
# memory alive.
_tensor_registry[aligned_ptr] = view
return aligned_ptr
def free_pinned_numa_ptr(ptr: int, size: int | None = None) -> None:
"""Non-CUDA equivalent of freeing a previously allocated NUMA pointer."""
# Release the tensor object for that pointer reference
_tensor_registry.pop(ptr, None)
def alloc_pinned_ptr(size: int, device_id: int = 0) -> int:
"""Non-CUDA equivalent of allocating pinned memory and returning pointer
to it. On XPU, uses pin_memory=True (SYCL USM host allocation) for
fast DMA transfers. On other non-CUDA platforms, pinning is not supported."""
view, aligned_ptr = _alloc_page_aligned_pinned_view(size)
_tensor_registry[aligned_ptr] = view
return aligned_ptr
def free_pinned_ptr(ptr: int) -> None:
"""Non-CUDA equivalent of freeing a previously allocated pinned pointer."""
# Release the tensor object for that pointer reference
_tensor_registry.pop(ptr, None)
def batched_memcpy(src_ptrs: list[int], dst_ptrs: list[int], sizes: list[int]) -> None:
"""Non-CUDA equivalent of the native batched memcpy helper."""
if len(src_ptrs) != len(dst_ptrs) or len(src_ptrs) != len(sizes):
raise ValueError(
"batched_memcpy expects equally sized src_ptrs, dst_ptrs, and sizes"
)
for src_ptr, dst_ptr, size in zip(src_ptrs, dst_ptrs, sizes, strict=True):
if size <= 0:
continue
ctypes.memmove(
ctypes.c_void_p(dst_ptr),
ctypes.c_void_p(src_ptr),
size,
)
def alloc_shm_pinned_ptr(size: int, shm_name: str = "") -> int:
"""Non-CUDA equivalent of allocating shared memory pinned pointer.
Uses multiprocessing.shared_memory for cross-platform POSIX shm.
Attempts to pin the buffer via cudaHostRegister for async D2H;
if pinning fails, continues without pinning."""
# Strip leading '/' for SharedMemory name
name = shm_name.lstrip("/") if shm_name else None
# Clean up stale shm segment if it exists
if name:
try:
stale = shared_memory.SharedMemory(name=name, create=False)
stale.close()
stale.unlink()
except FileNotFoundError:
pass
shm = shared_memory.SharedMemory(name=name, create=True, size=size)
array_type = ctypes.c_uint8 * size
buf = array_type.from_buffer(shm.buf)
ptr = ctypes.addressof(buf)
# Store references to keep them alive
tensor = torch.frombuffer(buf, dtype=torch.uint8)
_tensor_registry[ptr] = tensor
_buf_registry[ptr] = buf
_shm_registry[ptr] = shm
# Try to pin the SHM buffer for async D2H copies
if current_device_spec.pin_memory(ptr, size):
_pinned_ptr_registry[ptr] = size
return ptr
def free_shm_pinned_ptr(ptr: int, size: int = 0, shm_name: str = "") -> None:
"""Non-CUDA equivalent of freeing a shared memory
pinned pointer. Unregisters pinned memory if it was pinned."""
# Unpin if previously registered
if ptr in _pinned_ptr_registry:
current_device_spec.unpin_memory(ptr)
_pinned_ptr_registry.pop(ptr, None)
# Release in order: tensor -> ctypes buf -> shm
_tensor_registry.pop(ptr, None)
_buf_registry.pop(ptr, None)
shm = _shm_registry.pop(ptr, None)
if shm is not None:
shm.close()
shm.unlink()
# Hugepage variants: non-CUDA platforms do not support hugepages, so these
# fall back to the same regular pinned allocation.
def alloc_hugepage_pinned_ptr(size: int, device_id: int = 0) -> int:
"""Non-CUDA fallback for alloc_hugepage_pinned_ptr (no hugepage support)."""
warnings.warn(
"Hugepages requested but not available on non-CUDA platforms; "
"falling back to regular allocation.",
RuntimeWarning,
stacklevel=2,
)
return alloc_pinned_ptr(size, device_id)
def free_hugepage_pinned_ptr(ptr: int, size: int = 0) -> None:
"""Non-CUDA fallback for free_hugepage_pinned_ptr (no hugepage support)."""
free_pinned_ptr(ptr)
def alloc_hugepage_pinned_numa_ptr(size: int, numa_id: int = 0) -> int:
"""Non-CUDA fallback for alloc_hugepage_pinned_numa_ptr (no hugepage support)."""
warnings.warn(
"Hugepages requested but not available on non-CUDA platforms; "
"falling back to regular allocation.",
RuntimeWarning,
stacklevel=2,
)
return alloc_pinned_numa_ptr(size, numa_id)
def free_hugepage_pinned_numa_ptr(ptr: int, size: int = 0) -> None:
"""Non-CUDA fallback for free_hugepage_pinned_numa_ptr (no hugepage support)."""
free_pinned_numa_ptr(ptr, size)
def alloc_numa_ptr(size: int, numa_id: int = 0) -> int:
"""Non-CUDA equivalent of allocating numa memory and returning pointer
to it. Note: Numa memory is not supported on non-CUDA."""
return alloc_pinned_numa_ptr(size, numa_id)
def free_numa_ptr(ptr: int, size: int | None = None) -> None:
"""Non-CUDA equivalent of freeing a previously allocated NUMA pointer."""
return free_pinned_numa_ptr(ptr, size)
def multi_layer_kv_transfer(
key_value: torch.Tensor,
key_value_ptrs: torch.Tensor | list[torch.Tensor],
slot_mapping: torch.Tensor,
paged_memory_device: torch.device,
page_buffer_size: int,
direction: TransferDirection,
engine_kv_format: EngineKVFormat,
block_size: int = 0,
head_size: int = 0,
skip_prefix_n_tokens: int = 0,
):
"""
Fully vectorized Python fallback for multi_layer_kv_transfer.
Eliminates ALL token- and KV-level Python loops.
"""
if not isinstance(key_value_ptrs, (torch.Tensor, list)):
raise TypeError(
f"Expected torch.Tensor or list, but got {type(key_value_ptrs).__name__}"
)
# TODO: Implement head_size support for HND layouts (NL_X_TWO_NB_NH_BS_HS,
# NL_X_NB_TWO_NH_BS_HS) as next step.
if int(engine_kv_format) in (
int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS),
int(EngineKVFormat.NL_X_NB_TWO_NH_BS_HS),
):
raise NotImplementedError(
"HND layouts (NL_X_TWO_NB_NH_BS_HS, NL_X_NB_TWO_NH_BS_HS) "
"are not supported in the non-CUDA fallback. "
"head_size parameter is required but not implemented in this path."
)
# 1. Filter out invalid slots.
# valid_mask_kv: on key_value.device, used to index key_value
# valid_slots: on paged_memory_device, used to index paged_tensor
kv_device = key_value.device
slots_kv = slot_mapping.to(dtype=torch.long).to(kv_device)
valid_mask_kv = slots_kv >= 0
# Skip the first skip_prefix_n_tokens tokens from transfer.
# This matches the CUDA kernel semantics where the grid starts at
# token_id=0 but indexes key_value/slot_mapping at
# kv_token_id = token_id + skip_prefix_n_tokens.
# By masking them as invalid, the vectorized indexing via valid_mask_kv
# naturally skips them while keeping key_value indices aligned.
if skip_prefix_n_tokens > 0:
valid_mask_kv[:skip_prefix_n_tokens] = False
if not valid_mask_kv.any():
return
valid_slots = slots_kv[valid_mask_kv].to(paged_memory_device)
# 2. Determine architecture variant and tensor dimensions.
is_mla = int(engine_kv_format) in (
int(EngineKVFormat.NL_X_NB_BS_HS),
int(EngineKVFormat.NL_X_NBBS_ONE_HS),
)
is_flash_infer = int(engine_kv_format) == int(EngineKVFormat.NL_X_NB_TWO_BS_NH_HS)
num_layers = key_value.size(1)
hidden_size = key_value.size(3)
# For the flash_infer interleaved layout, pre-compute block-level indices.
if is_flash_infer:
block_indices = valid_slots // block_size
block_offsets = valid_slots % block_size
# Determine the physical shape of the underlying paged tensor
# (used when wrapping a raw pointer).
layer_shape: Tuple[int, ...]
if is_mla:
layer_shape = (page_buffer_size, hidden_size)
elif is_flash_infer:
num_blocks = page_buffer_size // block_size
layer_shape = (num_blocks, 2, block_size, hidden_size)
else:
layer_shape = (2, page_buffer_size, hidden_size)
# 3. Iterate over layers — the only remaining Python-level loop.
for layer_id in range(num_layers):
# --- A. Obtain the physical device-memory view for this layer. ---
if isinstance(key_value_ptrs, list):
paged_tensor = key_value_ptrs[layer_id]
else:
ptr = int(key_value_ptrs[layer_id].item())
# Convert a raw device pointer into a PyTorch tensor view.
paged_tensor = _tensor_from_ptr(
ptr, layer_shape, key_value.dtype, paged_memory_device
)
# --- B. Vectorized bulk data transfer. ---
if is_mla:
# Paged layout : [page_buffer_size, hidden_size]
# key_value layout: [1, num_layers, num_tokens, hidden_size]
if int(direction) == int(TransferDirection.H2D):
lmc_valid = key_value[0, layer_id, valid_mask_kv, :]
paged_tensor.index_copy_(
0, valid_slots, lmc_valid.to(paged_tensor.device)
)
else:
gathered = paged_tensor.index_select(0, valid_slots)
key_value[0, layer_id, valid_mask_kv, :] = gathered.to(
kv_device, non_blocking=False
)
elif is_flash_infer:
# Paged layout : [num_blocks, 2, block_size, hidden_size]
# key_value layout: [2, num_layers, num_tokens, hidden_size]
if int(direction) == int(TransferDirection.H2D):
lmc_valid = key_value[:, layer_id, valid_mask_kv, :]
src_data = lmc_valid.transpose(0, 1).to(paged_memory_device)
# src_data: [num_valid, 2, hidden_size]
paged_tensor[block_indices, :, block_offsets, :] = src_data
else:
gathered = paged_tensor[block_indices, :, block_offsets, :]
# gathered: [num_valid, 2, hidden_size]
key_value[:, layer_id, valid_mask_kv, :] = gathered.to(
kv_device, non_blocking=False
).transpose(0, 1)
else:
# Paged layout : [2, page_buffer_size, hidden_size]
# key_value layout: [2, num_layers, num_tokens, hidden_size]
if int(direction) == int(TransferDirection.H2D):
lmc_valid = key_value[:, layer_id, valid_mask_kv, :]
paged_tensor.index_copy_(
1, valid_slots, lmc_valid.to(paged_memory_device)
)
else:
gathered = paged_tensor.index_select(1, valid_slots)
key_value[:, layer_id, valid_mask_kv, :] = gathered.to(
kv_device, non_blocking=False
)
def multi_layer_kv_transfer_unilateral(
key_value: torch.Tensor,
key_value_ptrs: torch.Tensor | list[torch.Tensor],
slot_mapping: torch.Tensor,
paged_memory_device: torch.device,
page_buffer_size: int,
direction: TransferDirection,
engine_kv_format: EngineKVFormat,
):
"""
Python fallback for multi_layer_kv_transfer_unilateral
Handles SGLang MHA format where K and V paged buffers are stored separately:
ptrs = [K_layer0, K_layer1, ..., V_layer0, V_layer1, ...]
each buffer shape: [page_buffer_size, hidden_size]
For MLA, delegates to multi_layer_kv_transfer (same as C++ implementation).
key_value_ptrs:
- If torch.Tensor: int64 tensor containing raw memory pointers.
- If list[torch.Tensor]: list of tensor objects.
key_value layout:
- Standard: [2, num_layers, num_tokens, hidden_size]
- MLA: [1, num_layers, num_tokens, hidden_size]
direction:
H2D = LMCache -> PagedBuffer
D2H = PagedBuffer -> LMCache
"""
is_mla = int(engine_kv_format) in (
int(EngineKVFormat.NL_X_NB_BS_HS),
int(EngineKVFormat.NL_X_NBBS_ONE_HS),
)
# MLA case collapses back to multi_layer_kv_transfer
# (vLLM and SGLang indexing are compatible)
if is_mla:
return multi_layer_kv_transfer(
key_value,
key_value_ptrs,
slot_mapping,
paged_memory_device,
page_buffer_size,
direction,
engine_kv_format,
0, # block_size unused for MLA formats
)
# ── Non-MLA path: unilateral (separate K/V buffers per layer) ──
num_layers = key_value.size(1)
hidden_size = key_value.size(3)
layer_shape = (page_buffer_size, hidden_size)
kv_device = key_value.device
slots_kv = slot_mapping.to(dtype=torch.long).to(kv_device)
valid_mask_kv = slots_kv >= 0
if not valid_mask_kv.any():
return
valid_slots = slots_kv[valid_mask_kv].to(paged_memory_device)
for layer_id in range(num_layers):
for kv_idx in range(2): # 0 = K, 1 = V
buffer_idx = layer_id + kv_idx * num_layers
if isinstance(key_value_ptrs, list):
paged_tensor = key_value_ptrs[buffer_idx]
else:
ptr = int(key_value_ptrs[buffer_idx].item())
paged_tensor = _tensor_from_ptr(
ptr, layer_shape, key_value.dtype, paged_memory_device
)
if int(direction) == int(TransferDirection.H2D):
lmc_valid = key_value[kv_idx, layer_id, valid_mask_kv, :]
paged_tensor.index_copy_(
0, valid_slots, lmc_valid.to(paged_memory_device)
)
else:
gathered = paged_tensor.index_select(0, valid_slots)
key_value[kv_idx, layer_id, valid_mask_kv, :] = gathered.to(kv_device)
# Pure-Python mirrors of the c_ops predicates (csrc/engine_kv_format.h); the
# parity test pins these to the same names and signatures.
def is_cross_layer(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when all layers live in one fused tensor."""
return int(engine_kv_format) in (
int(EngineKVFormat.NB_NL_TWO_BS_NH_HS),
int(EngineKVFormat.NB_NL_TWO_NH_BS_HS),
)
def is_kv_list(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when keys and values are two separate top-level lists."""
return int(engine_kv_format) in (
int(EngineKVFormat.TWO_X_NL_X_NBBS_NH_HS),
int(EngineKVFormat.TWO_X_NL_X_NB_BS_NH_HS),
)
def is_layer_list(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when the structure is one list entry per layer."""
return int(engine_kv_format) in (
int(EngineKVFormat.NL_X_TWO_NB_BS_NH_HS),
int(EngineKVFormat.NL_X_NB_TWO_BS_NH_HS),
int(EngineKVFormat.NL_X_NB_BS_HS),
int(EngineKVFormat.NL_X_NBBS_ONE_HS),
int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS),
int(EngineKVFormat.NL_X_NB_TWO_NH_BS_HS),
int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS),
)
def is_mla(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when a KV format uses MLA paged layout."""
return int(engine_kv_format) in (
int(EngineKVFormat.NL_X_NB_BS_HS),
int(EngineKVFormat.NL_X_NBBS_ONE_HS),
)
def _is_cross_layer_format(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when a KV format uses a single cross-layer tensor."""
return is_cross_layer(engine_kv_format)
def _is_sglang_mha_format(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when a KV format uses SGLang MHA layout (2*NL tensors)."""
return is_kv_list(engine_kv_format)
def _is_hnd_format(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when a per-layer KV format stores heads before block tokens (HND)."""
return int(engine_kv_format) in (
int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS),
int(EngineKVFormat.NL_X_NB_TWO_NH_BS_HS),
int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS),
)
def _is_mla_format(engine_kv_format: EngineKVFormat) -> bool:
"""Return True when a KV format uses MLA paged layout."""
return is_mla(engine_kv_format)
_ELEMENT_SIZE_TO_DTYPE: dict[int, torch.dtype] = {
# Maps the byte width of a KV-cache element to a representative torch dtype.
# Only widths that commonly appear in KV caches are listed; 1-byte entries
# are treated as uint8 (raw bytes), 2-byte as float16, 4-byte as float32.
# Note: bfloat16 also has element_size == 2 but cannot be distinguished here;
# callers that need exact dtype should supply it explicitly.
1: torch.uint8,
2: torch.float16,
4: torch.float32,
}
def _is_ptr_tensor(x: object) -> bool:
"""Return True when *x* is a 1-D pointer tensor (int64 or uint64)."""
return (
isinstance(x, torch.Tensor)
and x.dtype in (torch.int64, torch.uint64)
and x.ndim == 1
)
def _per_layer_paged_shape(
engine_kv_format: EngineKVFormat,
nb: int,
bs: int,
nh: int,
hs: int,
) -> tuple[int, ...]:
"""Return the logical shape of a single per-layer paged buffer tensor.
Args:
engine_kv_format: The format enum that describes how K/V tokens are laid out.
nb: Number of blocks in the paged buffer (``shape_desc.nb``).
bs: Tokens per block / block size (``shape_desc.bs``).
nh: Number of attention heads (``shape_desc.nh``).
hs: Per-head hidden size (``shape_desc.hs``).
Returns:
A tuple representing the shape needed to reconstruct one layer's tensor
from a raw pointer via :func:`_tensor_from_ptr`.
"""
fmt = int(engine_kv_format)
if fmt == int(EngineKVFormat.NL_X_NBBS_ONE_HS):
return (nb * bs, 1, hs)
if fmt == int(EngineKVFormat.NL_X_NB_BS_HS):
return (nb, bs, hs)
if fmt == int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS):
return (2, nb, nh, bs, hs)
if fmt == int(EngineKVFormat.NL_X_NB_TWO_NH_BS_HS):
return (nb, 2, nh, bs, hs)
if fmt == int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS):
# vLLM CPU blocks-first fused KV: K and V interleaved at the
# second-to-last dim so each layer is [NB, NH, BS, 2, HS].
return (nb, nh, bs, 2, hs)
if fmt == int(EngineKVFormat.NL_X_TWO_NB_BS_NH_HS):
return (2, nb, bs, nh, hs)
if fmt == int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS):
return (nb, nh, bs, 2, hs)
# Covers NL_X_NB_TWO_BS_NH_HS and any future NHD variants.
return (nb, 2, bs, nh, hs)
def _infer_kv_dtype(
paged_buffer_ptrs_tensor: object,
lmcache_objects_ptrs: object,
shape_desc: "PageBufferShapeDesc",
) -> torch.dtype:
"""Infer the KV element dtype from whichever inputs carry it.
Inference order (first match wins):
1. ``shape_desc.dtype`` — authoritative when set (requires the
``set_shape_desc_dtype`` helper from PR #3514; correctly distinguishes
float16 vs bfloat16 which share ``element_size == 2``).
2. ``paged_buffer_ptrs_tensor`` — if it is a non-pointer tensor or a list
of tensors (including nested SGLang MHA lists), the dtype of the first
tensor is used.
3. ``lmcache_objects_ptrs`` — if it is a list of tensors, the dtype of the
first chunk tensor is used.
4. ``shape_desc.element_size`` — looked up in :data:`_ELEMENT_SIZE_TO_DTYPE`
(ambiguous for 2-byte types; kept only as last-resort fallback).
5. ``torch.bfloat16`` — silent default when no other source is available.
"""
# Prefer shape_desc.dtype — it is exact and avoids the element_size ambiguity.
if shape_desc is not None:
sd_dtype = getattr(shape_desc, "dtype", None)
if sd_dtype is not None:
return sd_dtype
if isinstance(paged_buffer_ptrs_tensor, list) and paged_buffer_ptrs_tensor:
first = paged_buffer_ptrs_tensor[0]
if isinstance(first, list) and first and isinstance(first[0], torch.Tensor):
return first[0].dtype
if isinstance(first, torch.Tensor):
return first.dtype
if isinstance(paged_buffer_ptrs_tensor, torch.Tensor) and not _is_ptr_tensor(
paged_buffer_ptrs_tensor
):
return paged_buffer_ptrs_tensor.dtype
if isinstance(lmcache_objects_ptrs, list) and lmcache_objects_ptrs:
if isinstance(lmcache_objects_ptrs[0], torch.Tensor):
return lmcache_objects_ptrs[0].dtype
if shape_desc is not None and shape_desc.element_size > 0:
dtype = _ELEMENT_SIZE_TO_DTYPE.get(shape_desc.element_size)
if dtype is None:
raise ValueError(
f"Unsupported element_size {shape_desc.element_size!r} in "
"shape_desc; cannot infer KV dtype. "
f"Supported sizes: {sorted(_ELEMENT_SIZE_TO_DTYPE)}"
)
return dtype
return torch.bfloat16
def _normalize_paged_layers(
paged_buffer_ptrs_tensor: "torch.Tensor | list",
engine_kv_format: EngineKVFormat,
shape_desc: "PageBufferShapeDesc | None" = None,
device: "torch.device | str | None" = None,
dtype: "torch.dtype | None" = None,
) -> "torch.Tensor | list[torch.Tensor] | list[list[torch.Tensor]]":
"""Normalize paged buffer input based on GPU KV format.
Accepts either tensor-form inputs (list / Tensor) or a 1-D pointer tensor
(int64 / uint64). When a pointer tensor is provided *shape_desc*, *device*,
and *dtype* must be supplied so the tensors can be reconstructed via
:func:`_tensor_from_ptr`.
Returns:
- Single ``torch.Tensor`` for cross-layer formats.
- ``list[list[torch.Tensor]]`` (2 x NL) for SGLang MHA formats.
- ``list[torch.Tensor]`` (per-layer) for all other formats.
"""
if _is_cross_layer_format(engine_kv_format):
if isinstance(paged_buffer_ptrs_tensor, torch.Tensor):
if _is_ptr_tensor(paged_buffer_ptrs_tensor):
# 1-D pointer tensor with a single entry → reconstruct full tensor.
if shape_desc is None or device is None or dtype is None:
raise ValueError(
"_normalize_paged_layers: shape_desc, device, and dtype are "
"required when paged_buffer_ptrs_tensor is a pointer tensor"
)
nb = int(shape_desc.nb)
nl = int(shape_desc.nl)
bs = int(shape_desc.bs)
nh = int(shape_desc.nh)
hs = int(shape_desc.hs)
if int(engine_kv_format) == int(EngineKVFormat.NB_NL_TWO_NH_BS_HS):
shape: tuple[int, ...] = (nb, nl, 2, nh, bs, hs)
else:
shape = (nb, nl, 2, bs, nh, hs)
ptr = int(paged_buffer_ptrs_tensor[0].item())
return _tensor_from_ptr(ptr, shape, dtype, device)
return paged_buffer_ptrs_tensor
raise TypeError(
"Cross-layer formats require a single torch.Tensor input; "
"got: " + type(paged_buffer_ptrs_tensor).__name__
)
if _is_sglang_mha_format(engine_kv_format):
if _is_ptr_tensor(paged_buffer_ptrs_tensor):
# 1-D pointer tensor [K_L0,...,K_LN-1, V_L0,...,V_LN-1] → nested list.
if shape_desc is None or device is None or dtype is None:
raise ValueError(
"_normalize_paged_layers: shape_desc, device, and dtype are "
"required when paged_buffer_ptrs_tensor is a pointer tensor"
)
nb = int(shape_desc.nb)
nl = int(shape_desc.nl)
bs = int(shape_desc.bs)
nh = int(shape_desc.nh)
hs = int(shape_desc.hs)
is_flat = int(engine_kv_format) == int(EngineKVFormat.TWO_X_NL_X_NBBS_NH_HS)
per_layer_shape: tuple[int, ...] = (
(nb * bs, nh, hs) if is_flat else (nb, bs, nh, hs)
)
ptrs = [int(p.item()) for p in paged_buffer_ptrs_tensor]
k_tensors = [
_tensor_from_ptr(ptrs[i], per_layer_shape, dtype, device)
for i in range(nl)
]
v_tensors = [
_tensor_from_ptr(ptrs[nl + i], per_layer_shape, dtype, device)
for i in range(nl)
]
return [k_tensors, v_tensors]
if isinstance(paged_buffer_ptrs_tensor, list):
# Already nested [[K tensors], [V tensors]]
if (
len(paged_buffer_ptrs_tensor) == 2
and isinstance(paged_buffer_ptrs_tensor[0], list)
and all(
isinstance(t, torch.Tensor)
for group in paged_buffer_ptrs_tensor
for t in group
)
):
return paged_buffer_ptrs_tensor
# Flat list [K_L0, ..., K_LN-1, V_L0, ..., V_LN-1]
if all(isinstance(t, torch.Tensor) for t in paged_buffer_ptrs_tensor):
if len(paged_buffer_ptrs_tensor) % 2 != 0:
raise ValueError(
"Flat SGLang MHA list must have even length (2*NL)"
)
half = len(paged_buffer_ptrs_tensor) // 2
return [
paged_buffer_ptrs_tensor[:half],
paged_buffer_ptrs_tensor[half:],
]
raise TypeError(
"SGLang MHA formats require a list[list[torch.Tensor]], a flat "
"list[torch.Tensor] (2*NL entries), or a 1-D pointer tensor; "
"got: " + type(paged_buffer_ptrs_tensor).__name__
)
# Per-layer formats
if _is_ptr_tensor(paged_buffer_ptrs_tensor):
# 1-D pointer tensor [ptr_L0, ..., ptr_LN-1] → list of per-layer tensors.
if shape_desc is None or device is None or dtype is None:
raise ValueError(
"_normalize_paged_layers: shape_desc, device, and dtype are "
"required when paged_buffer_ptrs_tensor is a pointer tensor"
)
nb = int(shape_desc.nb)
bs = int(shape_desc.bs)
nh = int(shape_desc.nh)
hs = int(shape_desc.hs)
per_shape = _per_layer_paged_shape(engine_kv_format, nb, bs, nh, hs)
return [
_tensor_from_ptr(int(p.item()), per_shape, dtype, device)
for p in paged_buffer_ptrs_tensor
]
if isinstance(paged_buffer_ptrs_tensor, list):
if not all(isinstance(t, torch.Tensor) for t in paged_buffer_ptrs_tensor):
raise TypeError(
"paged_buffer_ptrs_tensor list must contain torch.Tensor entries"
)
return paged_buffer_ptrs_tensor
raise TypeError(
"paged_buffer_ptrs_tensor must be a list[torch.Tensor] or 1-D pointer tensor; "
"got: " + type(paged_buffer_ptrs_tensor).__name__
)
def _normalize_lmcache_objects(
lmcache_objects_ptrs: "list[int] | list[torch.Tensor]",
shape_desc: "PageBufferShapeDesc | None" = None,
lmcache_chunk_size: "int | None" = None,
engine_kv_format: "EngineKVFormat | None" = None,
dtype: "torch.dtype | None" = None,
) -> list[torch.Tensor]:
"""Normalize LMCache object inputs to chunk tensors.
Accepts either a list of chunk tensors or a ``list[int]`` of raw CPU pointers.
When a pointer list is provided *shape_desc*, *lmcache_chunk_size*,
*engine_kv_format*, and *dtype* must be supplied so the tensors can be
reconstructed via :func:`_tensor_from_ptr` on the CPU.
"""
if not isinstance(lmcache_objects_ptrs, list):
raise TypeError(
"lmcache_objects_ptrs must be a list[torch.Tensor] or list[int]; "
"got: " + type(lmcache_objects_ptrs).__name__
)
if not lmcache_objects_ptrs:
return []
if isinstance(lmcache_objects_ptrs[0], torch.Tensor):
return lmcache_objects_ptrs # type: ignore[return-value]
if isinstance(lmcache_objects_ptrs[0], int):
# Pointer mode: reconstruct chunk tensors (always on CPU).
if (
shape_desc is None
or lmcache_chunk_size is None
or engine_kv_format is None
or dtype is None
):
raise ValueError(
"_normalize_lmcache_objects: shape_desc, lmcache_chunk_size, "
"engine_kv_format, and dtype are required when lmcache_objects_ptrs "
"contains raw int pointers"
)
nl = int(shape_desc.nl)
nh = int(shape_desc.nh)
hs = int(shape_desc.hs)
chunk_tokens = lmcache_chunk_size
if _is_mla_format(engine_kv_format):
chunk_shape: tuple[int, ...] = (nl, chunk_tokens, hs)
else:
chunk_shape = (2, nl, chunk_tokens, nh * hs)
return [
_tensor_from_ptr(ptr, chunk_shape, dtype, "cpu")
for ptr in lmcache_objects_ptrs
]
raise TypeError(
"lmcache_objects_ptrs must be a list[torch.Tensor] or list[int]; "
"got list containing: " + type(lmcache_objects_ptrs[0]).__name__
)
def _to_block_id_list(block_ids: torch.Tensor | list[int]) -> list[int]:
"""Convert block IDs from tensor/list form into a Python ``list[int]``."""
if isinstance(block_ids, torch.Tensor):
return [int(x) for x in block_ids.to(dtype=torch.int64).cpu().tolist()]
if isinstance(block_ids, list):
return [int(x) for x in block_ids]
raise TypeError("block_ids must be a torch.Tensor or list[int]")
def multi_layer_block_kv_transfer(
paged_buffer_ptrs_tensor: "torch.Tensor | list",
lmcache_objects_ptrs: list[int] | list[torch.Tensor],
block_ids: torch.Tensor | list[int],
device: torch.device | str,
direction: TransferDirection,
shape_desc: PageBufferShapeDesc,
lmcache_chunk_size: int,
engine_kv_format: EngineKVFormat,
skip_prefix_n_blocks: int,
) -> None:
"""Python fallback implementation of block-based multi-layer KV transfer.
Signature intentionally mirrors the C++ binding so callers can invoke
``lmcache.c_ops.multi_layer_block_kv_transfer`` uniformly on native and
fallback backends.
Args:
paged_buffer_ptrs_tensor: Paged buffer pointers or tensors.
lmcache_objects_ptrs: LMCache object pointers or chunk tensors.
block_ids: Ordered engine block IDs for the transfer.
device: Target device for the transfer.
direction: Transfer direction (H2D or D2H).
shape_desc: Shape descriptor of the page buffer.
lmcache_chunk_size: Chunk size of LMCache objects.
engine_kv_format: GPU KV cache format.
skip_prefix_n_blocks: Number of leading blocks to skip.
Returns:
None
Raises:
ValueError: If chunk size is invalid, or transfer direction is unsupported.
TypeError: If input types do not match expected types.
"""
if lmcache_chunk_size <= 0:
raise ValueError("lmcache_chunk_size must be positive")
if int(shape_desc.bs) <= 0 or lmcache_chunk_size % int(shape_desc.bs) != 0:
raise ValueError(
"lmcache_chunk_size must be a positive multiple of shape_desc.bs"
)
if skip_prefix_n_blocks < 0:
raise ValueError("skip_prefix_n_blocks must be >= 0")
is_d2h = int(direction) == int(TransferDirection.D2H)
is_h2d = int(direction) == int(TransferDirection.H2D)
if not (is_d2h or is_h2d):
raise ValueError(f"Unsupported transfer direction: {direction!r}")
kv_dtype = _infer_kv_dtype(
paged_buffer_ptrs_tensor, lmcache_objects_ptrs, shape_desc
)
normalized = _normalize_paged_layers(
paged_buffer_ptrs_tensor,
engine_kv_format,
shape_desc=shape_desc,
device=device,
dtype=kv_dtype,
)
object_tensors = _normalize_lmcache_objects(
lmcache_objects_ptrs,
shape_desc=shape_desc,
lmcache_chunk_size=lmcache_chunk_size,
engine_kv_format=engine_kv_format,
dtype=kv_dtype,
)
n_block_ids = (
int(block_ids.numel())
if isinstance(block_ids, torch.Tensor)
else len(block_ids)
)
blocks_per_object = lmcache_chunk_size // int(shape_desc.bs)
block_size = int(shape_desc.bs)
if _is_cross_layer_format(engine_kv_format):
_transfer_cross_layer(
normalized,
object_tensors,
block_ids,
n_block_ids,
blocks_per_object,
block_size,
engine_kv_format,
is_d2h,
skip_prefix_n_blocks,
)
elif _is_sglang_mha_format(engine_kv_format):
_transfer_sglang_mha(
normalized,
object_tensors,
block_ids,
n_block_ids,
blocks_per_object,
block_size,
engine_kv_format,
is_d2h,
skip_prefix_n_blocks,
)
elif _is_mla_format(engine_kv_format):
_transfer_per_layer_mla(
normalized,
object_tensors,
block_ids,
n_block_ids,
blocks_per_object,
block_size,
engine_kv_format,
is_d2h,
skip_prefix_n_blocks,
)
elif _is_hnd_format(engine_kv_format):
_transfer_per_layer_hnd(
normalized,
object_tensors,
block_ids,
n_block_ids,
blocks_per_object,
block_size,
engine_kv_format,
is_d2h,
skip_prefix_n_blocks,
)
else:
_transfer_per_layer_nhd(
normalized,
object_tensors,
block_ids,
n_block_ids,
blocks_per_object,
block_size,
engine_kv_format,
is_d2h,
skip_prefix_n_blocks,
)
def execute_object_group_transfer(
direction: TransferDirection,
device: torch.device | str,
host_buffer_alignment: int,
kernel_group_specs: list,
batch_steps: list,
) -> None:
"""Python fallback for the native object-group transfer plan executor.
The planned/batched object-group transfer (see ``csrc/mp_mem_kernels.cuh``
and ``execute_object_group_transfer``) is only implemented in the compiled
``c_ops`` extension. The signature mirrors the C++ binding so callers can
dispatch uniformly, but there is no pure-Python equivalent.
Args:
direction: Transfer direction (H2D or D2H).
device: CUDA device of the transfer.
host_buffer_alignment: Host buffer alignment for staging copies.
kernel_group_specs: Per-kernel-group invariants (native ``KernelGroupSpec``).
batch_steps: Ordered per-batch staging + launch work (native ``BatchStep``).
Raises:
NotImplementedError: Always; requires the c_ops native extension.
"""
raise NotImplementedError(
"execute_object_group_transfer requires the c_ops native extension; "
"no pure-Python fallback exists."
)
def _valid_block_range(
object_idx: int,
block_id_list: list[int],
blocks_per_object: int,
block_size: int,
skip_prefix_n_blocks: int,
) -> tuple[list[int], int] | None:
"""Return valid engine block IDs and their LMCache object token offset.
Args:
object_idx: Index of the LMCache object/chunk being processed.
block_id_list: Full ordered engine block ids for the transfer.
blocks_per_object: Number of blocks represented by one LMCache object.
block_size: Number of tokens per block.
skip_prefix_n_blocks: Number of leading flat block positions to skip.
Returns:
``None`` if this object has no valid blocks after skip handling.
Otherwise, a tuple of valid engine block ids and the token offset
within this LMCache object where those blocks start.
"""
object_flat_start = object_idx * blocks_per_object
valid_flat_start = max(object_flat_start, skip_prefix_n_blocks)
valid_flat_end = min(object_flat_start + blocks_per_object, len(block_id_list))
if valid_flat_start >= valid_flat_end:
return None
offset_in_object = (valid_flat_start - object_flat_start) * block_size
return block_id_list[valid_flat_start:valid_flat_end], offset_in_object
def _valid_block_range_indices(
object_idx: int,
n_block_ids: int,
blocks_per_object: int,
block_size: int,
skip_prefix_n_blocks: int,
) -> tuple[int, int, int] | None:
"""Return valid [start, end) range over flat block IDs and object token offset."""
object_flat_start = object_idx * blocks_per_object
valid_flat_start = max(object_flat_start, skip_prefix_n_blocks)
valid_flat_end = min(object_flat_start + blocks_per_object, n_block_ids)
if valid_flat_start >= valid_flat_end:
return None
offset_in_object = (valid_flat_start - object_flat_start) * block_size
return valid_flat_start, valid_flat_end, offset_in_object
def _transfer_cross_layer(
paged_tensor: torch.Tensor,
object_tensors: list[torch.Tensor],
block_ids: torch.Tensor | list[int],
n_block_ids: int,
blocks_per_object: int,
block_size: int,
engine_kv_format: EngineKVFormat,
is_d2h: bool,
skip_prefix_n_blocks: int,
) -> None:
"""Handle cross-layer formats: single tensor [NB, NL, 2, ...]."""
# NHD: [NB, NL, 2, BS, NH, HS] HND: [NB, NL, 2, NH, BS, HS]
is_hnd = int(engine_kv_format) == int(EngineKVFormat.NB_NL_TWO_NH_BS_HS)
num_layers = paged_tensor.shape[1]
if is_hnd:
# [NB, NL, 2, NH, BS, HS]
nh = paged_tensor.shape[3]
hs = paged_tensor.shape[5]
else:
# [NB, NL, 2, BS, NH, HS]
nh = paged_tensor.shape[4]
hs = paged_tensor.shape[5]
# H2D: pre-transfer objects to paged device
if not is_d2h and object_tensors:
objs_on_device = [obj.to(paged_tensor.device) for obj in object_tensors]
block_ids_dev = torch.as_tensor(
block_ids, dtype=torch.long, device=paged_tensor.device
)
for object_idx, obj in enumerate(object_tensors):
valid = _valid_block_range_indices(
object_idx,
n_block_ids,
blocks_per_object,
block_size,
skip_prefix_n_blocks,
)
if valid is None:
continue
idx_start, idx_end, offset_in_object = valid
n_valid = idx_end - idx_start
token_end = offset_in_object + n_valid * block_size
eff_idx = block_ids_dev[idx_start:idx_end]
if is_d2h:
selected = paged_tensor.index_select(0, eff_idx)
for layer_idx in range(num_layers):
for kv in range(2):
if is_d2h:
slice_t = selected[:, layer_idx, kv]
if is_hnd:
# N=n_valid, BS=block_size:
# [N, NH, BS, HS] -> [N, BS, NH, HS] -> [N*BS, NH*HS]
flat = slice_t.permute(0, 2, 1, 3).reshape(
n_valid * block_size, nh * hs
)
else:
# [N, BS, NH, HS] → [N*BS, NH*HS]
flat = slice_t.reshape(n_valid * block_size, nh * hs)
obj[kv, layer_idx, offset_in_object:token_end].copy_(
flat, non_blocking=True
)
else:
obj_device = objs_on_device[object_idx]
src = obj_device[kv, layer_idx, offset_in_object:token_end]
if is_hnd:
# N=n_valid, BS=block_size:
# [N*BS, NH*HS] -> [N, BS, NH, HS] -> [N, NH, BS, HS]
src_blocks = src.reshape(n_valid, block_size, nh, hs).permute(
0, 2, 1, 3
)
else:
# N=n_valid, BS=block_size:
# [N*BS, NH*HS] -> [N, BS, NH, HS]
src_blocks = src.reshape(n_valid, block_size, nh, hs)
paged_tensor[:, layer_idx, kv].index_copy_(0, eff_idx, src_blocks)
def _transfer_sglang_mha(
paged_tensors: list[list[torch.Tensor]],
object_tensors: list[torch.Tensor],
block_ids: torch.Tensor | list[int],
n_block_ids: int,
blocks_per_object: int,
block_size: int,
engine_kv_format: EngineKVFormat,
is_d2h: bool,
skip_prefix_n_blocks: int,
) -> None:
"""Handle SGLang MHA formats: 2*NL tensors (list[list[Tensor]])."""
# TWO_X_NL_X_NBBS_NH_HS: each tensor [NB*BS, NH, HS]
# TWO_X_NL_X_NB_BS_NH_HS: each tensor [NB, BS, NH, HS]
is_flat = int(engine_kv_format) == int(EngineKVFormat.TWO_X_NL_X_NBBS_NH_HS)
num_layers = len(paged_tensors[0])
# Determine target device from first tensor
target_device = paged_tensors[0][0].device
# H2D: pre-transfer objects
if not is_d2h and object_tensors:
objs_on_device = [obj.to(target_device) for obj in object_tensors]
block_ids_dev = torch.as_tensor(block_ids, dtype=torch.long, device=target_device)
for object_idx, obj in enumerate(object_tensors):
valid = _valid_block_range_indices(
object_idx,
n_block_ids,
blocks_per_object,
block_size,
skip_prefix_n_blocks,
)
if valid is None:
continue
idx_start, idx_end, offset_in_object = valid
n_valid = idx_end - idx_start
token_end = offset_in_object + n_valid * block_size
eff_idx = block_ids_dev[idx_start:idx_end]
if is_flat:
# Flat token positions for all valid blocks:
# block_id * block_size + token offset. Reused across layer/KV pairs.
token_indices = (
eff_idx[:, None] * block_size
+ torch.arange(block_size, dtype=torch.long, device=target_device)
).reshape(-1)
for layer_idx in range(num_layers):
for kv in range(2):
layer_t = paged_tensors[kv][layer_idx]
nh = layer_t.shape[-2]
hs = layer_t.shape[-1]
if is_d2h:
if is_flat:
# [NB*BS, NH, HS]
gathered = layer_t.index_select(0, token_indices)
else:
# [NB, BS, NH, HS]
gathered = layer_t.index_select(0, eff_idx).reshape(
n_valid * block_size, nh, hs
)
flat = gathered.reshape(n_valid * block_size, nh * hs)
obj[kv, layer_idx, offset_in_object:token_end].copy_(
flat, non_blocking=True
)
else:
obj_device = objs_on_device[object_idx]
src = obj_device[kv, layer_idx, offset_in_object:token_end]
src_shaped = src.reshape(n_valid * block_size, nh, hs)
if is_flat:
# scatter into [NB*BS, NH, HS]
layer_t.index_copy_(0, token_indices, src_shaped)
else:
# N=n_valid, BS=block_size:
# [N*BS, NH, HS] -> [N, BS, NH, HS]
src_blocks = src_shaped.reshape(n_valid, block_size, nh, hs)
layer_t.index_copy_(0, eff_idx, src_blocks)
def _transfer_per_layer_mla(
layer_tensors: list[torch.Tensor],
object_tensors: list[torch.Tensor],
block_ids: torch.Tensor | list[int],
n_block_ids: int,
blocks_per_object: int,
block_size: int,
engine_kv_format: EngineKVFormat,
is_d2h: bool,
skip_prefix_n_blocks: int,
) -> None:
"""Handle MLA per-layer formats: [NB, BS, HS]."""
if not layer_tensors or not object_tensors:
return
is_flat = int(engine_kv_format) == int(EngineKVFormat.NL_X_NBBS_ONE_HS)
target_device = layer_tensors[0].device
if is_flat:
token_offsets = torch.arange(block_size, dtype=torch.long, device=target_device)
block_ids_dev = torch.as_tensor(block_ids, dtype=torch.long, device=target_device)
for object_idx, obj in enumerate(object_tensors):
valid = _valid_block_range_indices(
object_idx,
n_block_ids,
blocks_per_object,
block_size,
skip_prefix_n_blocks,
)
if valid is None:
continue
idx_start, idx_end, offset_in_object = valid
n_valid = idx_end - idx_start
token_end = offset_in_object + n_valid * block_size
eff_idx = block_ids_dev[idx_start:idx_end]
if is_flat:
token_indices = (
eff_idx[:, None] * block_size + token_offsets[None, :]
).reshape(-1)
if is_d2h:
hidden_size = layer_tensors[0].shape[-1]
chunk_gpu = torch.empty(
len(layer_tensors),
n_valid * block_size,
hidden_size,
dtype=layer_tensors[0].dtype,
device=target_device,
)
for layer_idx, layer in enumerate(layer_tensors):
if is_flat:
dst = chunk_gpu[layer_idx].view(
n_valid * block_size, 1, hidden_size
)
torch.index_select(layer, 0, token_indices, out=dst)
else:
dst = chunk_gpu[layer_idx].view(n_valid, block_size, hidden_size)
torch.index_select(layer, 0, eff_idx, out=dst)
obj[:, offset_in_object:token_end].copy_(chunk_gpu, non_blocking=True)
else:
chunk_gpu = obj[:, offset_in_object:token_end].to(
target_device, non_blocking=True
)
for layer_idx, layer in enumerate(layer_tensors):
src = chunk_gpu[layer_idx]
hidden_size = layer.shape[-1]
if is_flat:
src_tokens = src.reshape(n_valid * block_size, 1, hidden_size)
layer.index_copy_(0, token_indices, src_tokens)
else:
src_blocks = src.reshape(n_valid, block_size, hidden_size)
layer.index_copy_(0, eff_idx, src_blocks)
def _transfer_per_layer_hnd(
layer_tensors: list[torch.Tensor],
object_tensors: list[torch.Tensor],
block_ids: torch.Tensor | list[int],
n_block_ids: int,
blocks_per_object: int,
block_size: int,
engine_kv_format: EngineKVFormat,
is_d2h: bool,
skip_prefix_n_blocks: int,
) -> None:
"""Handle per-layer HND formats: heads before block tokens."""
if not layer_tensors or not object_tensors:
return
target_device = layer_tensors[0].device
block_ids_dev = torch.as_tensor(block_ids, dtype=torch.long, device=target_device)
first_layer = layer_tensors[0]
if int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS):
first_k = first_layer[0]
elif int(engine_kv_format) == int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS):
first_k = first_layer[:, :, :, 0]
else:
first_k = first_layer[:, 0]
_nb0, nh0, _bs0, hs0 = first_k.shape
for object_idx, obj in enumerate(object_tensors):
valid = _valid_block_range_indices(
object_idx,
n_block_ids,
blocks_per_object,
block_size,
skip_prefix_n_blocks,
)
if valid is None:
continue
idx_start, idx_end, offset_in_object = valid
n_valid = idx_end - idx_start
token_end = offset_in_object + n_valid * block_size
eff_idx = block_ids_dev[idx_start:idx_end]
if is_d2h:
chunk_gpu = torch.empty(
2,
len(layer_tensors),
n_valid * block_size,
nh0 * hs0,
dtype=first_k.dtype,
device=target_device,
)
scratch = torch.empty(
n_valid,
nh0,
block_size,
hs0,
dtype=first_k.dtype,
device=target_device,
)
for layer_idx, layer in enumerate(layer_tensors):
if int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS):
k_t, v_t = layer[0], layer[1]
torch.index_select(k_t, 0, eff_idx, out=scratch)
chunk_gpu[0, layer_idx].view(n_valid, block_size, nh0, hs0).copy_(
scratch.permute(0, 2, 1, 3)
)
torch.index_select(v_t, 0, eff_idx, out=scratch)
chunk_gpu[1, layer_idx].view(n_valid, block_size, nh0, hs0).copy_(
scratch.permute(0, 2, 1, 3)
)
elif int(engine_kv_format) == int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS):
k_t, v_t = layer[:, :, :, 0], layer[:, :, :, 1]
torch.index_select(k_t, 0, eff_idx, out=scratch)
chunk_gpu[0, layer_idx].view(n_valid, block_size, nh0, hs0).copy_(
scratch.permute(0, 2, 1, 3)
)
torch.index_select(v_t, 0, eff_idx, out=scratch)
chunk_gpu[1, layer_idx].view(n_valid, block_size, nh0, hs0).copy_(
scratch.permute(0, 2, 1, 3)
)
else:
# FlashInfer HND stores KV as [NB, 2, NH, BS, HS].
# Gather on dim=0 first so reads stay contiguous in memory;
# index_select on layer[:, 0]/layer[:, 1] non-contiguous views
# triggers slower element-wise gather reads.
selected = layer.index_select(0, eff_idx)
chunk_gpu[0, layer_idx].view(n_valid, block_size, nh0, hs0).copy_(
selected[:, 0].permute(0, 2, 1, 3)
)
chunk_gpu[1, layer_idx].view(n_valid, block_size, nh0, hs0).copy_(
selected[:, 1].permute(0, 2, 1, 3)
)
obj[:, :, offset_in_object:token_end].copy_(chunk_gpu, non_blocking=True)
else:
chunk_gpu = obj[:, :, offset_in_object:token_end].to(
target_device, non_blocking=True
)
for layer_idx, layer in enumerate(layer_tensors):
if int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_NH_BS_HS):
k_t, v_t = layer[0], layer[1]
elif int(engine_kv_format) == int(EngineKVFormat.NL_X_NB_NH_BS_TWO_HS):
k_t, v_t = layer[:, :, :, 0], layer[:, :, :, 1]
else:
k_t, v_t = layer[:, 0], layer[:, 1]
_nb, nh, _bs, hs = k_t.shape
k_blocks = (
chunk_gpu[0, layer_idx]
.reshape(n_valid, block_size, nh, hs)
.permute(0, 2, 1, 3)
)
v_blocks = (
chunk_gpu[1, layer_idx]
.reshape(n_valid, block_size, nh, hs)
.permute(0, 2, 1, 3)
)
if int(engine_kv_format) == int(EngineKVFormat.NL_X_NB_TWO_NH_BS_HS):
layer.index_copy_(
0, eff_idx, torch.stack([k_blocks, v_blocks], dim=1)
)
else:
k_t.index_copy_(0, eff_idx, k_blocks)
v_t.index_copy_(0, eff_idx, v_blocks)
def _transfer_per_layer_nhd(
layer_tensors: list[torch.Tensor],
object_tensors: list[torch.Tensor],
block_ids: torch.Tensor | list[int],
n_block_ids: int,
blocks_per_object: int,
block_size: int,
engine_kv_format: EngineKVFormat,
is_d2h: bool,
skip_prefix_n_blocks: int,
) -> None:
"""Handle per-layer NHD formats: block tokens before heads."""
if not layer_tensors or not object_tensors:
return
target_device = layer_tensors[0].device
block_ids_dev = torch.as_tensor(block_ids, dtype=torch.long, device=target_device)
first_layer = layer_tensors[0]
if int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_BS_NH_HS):
first_k = first_layer[0]
else:
first_k = first_layer[:, 0]
_nb0, _bs0, nh0, hs0 = first_k.shape
for object_idx, obj in enumerate(object_tensors):
valid = _valid_block_range_indices(
object_idx,
n_block_ids,
blocks_per_object,
block_size,
skip_prefix_n_blocks,
)
if valid is None:
continue
idx_start, idx_end, offset_in_object = valid
n_valid = idx_end - idx_start
token_end = offset_in_object + n_valid * block_size
eff_idx = block_ids_dev[idx_start:idx_end]
if is_d2h:
chunk_gpu = torch.empty(
2,
len(layer_tensors),
n_valid * block_size,
nh0 * hs0,
dtype=first_k.dtype,
device=target_device,
)
for layer_idx, layer in enumerate(layer_tensors):
if int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_BS_NH_HS):
k_t, v_t = layer[0], layer[1]
torch.index_select(
k_t,
0,
eff_idx,
out=chunk_gpu[0, layer_idx].view(n_valid, block_size, nh0, hs0),
)
torch.index_select(
v_t,
0,
eff_idx,
out=chunk_gpu[1, layer_idx].view(n_valid, block_size, nh0, hs0),
)
else:
# FlashInfer NHD stores KV as [NB, 2, BS, NH, HS].
# Gather on dim=0 first to avoid index_select from
# non-contiguous layer[:, 0]/layer[:, 1] views, which
# trigger slower element-wise gather reads.
selected = layer.index_select(0, eff_idx)
chunk_gpu[0, layer_idx].copy_(
selected[:, 0].reshape(n_valid * block_size, nh0 * hs0)
)
chunk_gpu[1, layer_idx].copy_(
selected[:, 1].reshape(n_valid * block_size, nh0 * hs0)
)
obj[:, :, offset_in_object:token_end].copy_(chunk_gpu, non_blocking=True)
else:
chunk_gpu = obj[:, :, offset_in_object:token_end].to(
target_device, non_blocking=True
)
for layer_idx, layer in enumerate(layer_tensors):
if int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_BS_NH_HS):
k_t, v_t = layer[0], layer[1]
k_t.index_copy_(
0,
eff_idx,
chunk_gpu[0, layer_idx].reshape(n_valid, block_size, nh0, hs0),
)
v_t.index_copy_(
0,
eff_idx,
chunk_gpu[1, layer_idx].reshape(n_valid, block_size, nh0, hs0),
)
else:
k_blocks = chunk_gpu[0, layer_idx].reshape(
n_valid, block_size, nh0, hs0
)
v_blocks = chunk_gpu[1, layer_idx].reshape(
n_valid, block_size, nh0, hs0
)
layer.index_copy_(
0, eff_idx, torch.stack([k_blocks, v_blocks], dim=1)
)
def single_layer_kv_transfer(
lmc_key_value_cache: torch.Tensor,
vllm_key_value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
direction: TransferDirection,
engine_kv_format: EngineKVFormat,
token_major: bool = False,
):
"""
Vectorized Python fallback for single_layer_kv_transfer
(eliminates per-token loops).
Transfers KV data between LMCache buffer
and a single vLLM paged KV cache layer.
lmc_key_value_cache layout:
- MLA: [num_tokens, aligned_head_size]
- token_major=True: [num_tokens, 2, num_heads * head_size]
- token_major=False: [2, num_tokens, num_heads * head_size]
vllm_key_value_cache layout:
- NL_X_TWO_NB_BS_NH_HS (flash attn):
[2, num_blocks, block_size, num_heads, head_size]
- NL_X_NB_TWO_BS_NH_HS (flash infer):
[num_blocks, 2, block_size, num_heads, head_size]
- NL_X_NB_BS_HS (vLLM MLA):
[num_blocks, block_size, head_size]
direction:
H2D = LMCache -> vLLM GPU
D2H = vLLM GPU -> LMCache
"""
kv_device = lmc_key_value_cache.device
paged_memory_device = vllm_key_value_cache.device
slots_kv = slot_mapping.to(dtype=torch.long).to(kv_device)
valid_mask_kv = slots_kv >= 0
if not valid_mask_kv.any():
return
valid_token_indices = torch.nonzero(valid_mask_kv, as_tuple=True)[0]
valid_slots = slots_kv[valid_mask_kv].to(paged_memory_device)
is_mla = int(engine_kv_format) in (
int(EngineKVFormat.NL_X_NB_BS_HS),
int(EngineKVFormat.NL_X_NBBS_ONE_HS),
)
if is_mla:
# ── MLA format ──
# vllm: [num_blocks, block_size, head_size]
# lmc: [num_tokens, aligned_head_size]
block_size = vllm_key_value_cache.size(1)
block_indices = valid_slots // block_size
block_offsets = valid_slots % block_size
if int(direction) == int(TransferDirection.D2H):
# vLLM -> LMCache
lmc_key_value_cache[valid_token_indices] = vllm_key_value_cache[
block_indices, block_offsets
].to(lmc_key_value_cache.device)
else:
# LMCache -> vLLM
vllm_key_value_cache[block_indices, block_offsets] = lmc_key_value_cache[
valid_token_indices
].to(paged_memory_device)
else:
# ── Non-MLA format ──
# Determine vLLM layout and block_size
is_two_major = int(engine_kv_format) == int(EngineKVFormat.NL_X_TWO_NB_BS_NH_HS)
# flash attn:
# [2, num_blocks, block_size, num_heads, head_size]
# -> dim2 = block_size
# flash infer:
# [num_blocks, 2, block_size, num_heads, head_size]
# -> dim2 = block_size
block_size = vllm_key_value_cache.size(2)
num_heads = vllm_key_value_cache.size(3)
head_size = vllm_key_value_cache.size(4)
block_indices = valid_slots // block_size
block_offsets = valid_slots % block_size
for kv in range(2):
if int(direction) == int(TransferDirection.D2H):
if is_two_major:
gathered = vllm_key_value_cache[kv, block_indices, block_offsets]
else:
gathered = vllm_key_value_cache[block_indices, kv, block_offsets]
gathered_flat = gathered.reshape(-1, num_heads * head_size).to(
lmc_key_value_cache.device
)
if token_major:
lmc_key_value_cache[valid_token_indices, kv] = gathered_flat
else:
lmc_key_value_cache[kv, valid_token_indices] = gathered_flat
else:
if token_major:
lmc_src = lmc_key_value_cache[valid_token_indices, kv]
else:
lmc_src = lmc_key_value_cache[kv, valid_token_indices]
lmc_reshaped = lmc_src.reshape(-1, num_heads, head_size).to(
vllm_key_value_cache.device
)
if is_two_major:
vllm_key_value_cache[kv, block_indices, block_offsets] = (
lmc_reshaped
)
else:
vllm_key_value_cache[block_indices, kv, block_offsets] = (
lmc_reshaped
)
def single_layer_kv_transfer_sgl(
lmc_key_value_cache: torch.Tensor,
sgl_key_cache: torch.Tensor,
sgl_value_cache: torch.Tensor,
slot_mapping: torch.Tensor,
direction: TransferDirection,
token_major: bool = False,
):
"""
Python fallback implementation of single_layer_kv_transfer_sgl.
Args:
lmc_key_value_cache:
[num_tokens, 2, num_heads*head_size] or
[2, num_tokens, num_heads*head_size]
sgl_key_cache: [num_blocks, block_size, num_heads, head_size]
sgl_value_cache: [num_blocks, block_size, num_heads, head_size]
slot_mapping: [num_tokens] - maps each token to a global slot index
direction: False for LMCache -> SGLang, True for SGLang -> LMCache
token_major: Boolean to determine the layout of lmc_key_value_cache
"""
kv_device = lmc_key_value_cache.device
paged_memory_device = sgl_key_cache.device
slots_kv = slot_mapping.to(dtype=torch.long).to(kv_device)
valid_mask_kv = slots_kv >= 0
if not valid_mask_kv.any():
return
# 1. Get basic dimensions
block_size = sgl_key_cache.size(1)
num_heads = sgl_key_cache.size(2)
head_size = sgl_key_cache.size(3)
# 2. Calculate block indices and offsets within the blocks from slot_mapping
# In SGLang/vLLM, slot_idx = block_idx * block_size + block_offset
valid_slots = slots_kv[valid_mask_kv].to(paged_memory_device)
block_indices = valid_slots // block_size
block_offsets = valid_slots % block_size
# 3. Prepare LMCache views for K and V
if token_major:
# Layout: [num_tokens, 2, hidden_size]
lmc_k = lmc_key_value_cache[:, 0, :]
lmc_v = lmc_key_value_cache[:, 1, :]
else:
# Layout: [2, num_tokens, hidden_size]
lmc_k = lmc_key_value_cache[0, :, :]
lmc_v = lmc_key_value_cache[1, :, :]
# 4. Perform the transfer
if int(direction) == int(TransferDirection.H2D):
# --- Direction: LMCache to SGLang (Paged Buffer) ---
# Reshape LMC flat tensors to match SGL [num_heads, head_size]
src_k_reshaped = (
lmc_k[valid_mask_kv]
.reshape(-1, num_heads, head_size)
.to(paged_memory_device)
)
src_v_reshaped = (
lmc_v[valid_mask_kv]
.reshape(-1, num_heads, head_size)
.to(paged_memory_device)
)
# Advanced indexing: update specific slots in the paged cache
sgl_key_cache[block_indices, block_offsets] = src_k_reshaped
sgl_value_cache[block_indices, block_offsets] = src_v_reshaped
else:
# --- Direction: SGLang (Paged Buffer) to LMCache ---
# Gather tensors from paged cache based on mapping
sampled_k = sgl_key_cache[block_indices, block_offsets].to(kv_device)
sampled_v = sgl_value_cache[block_indices, block_offsets].to(kv_device)
# Flatten the head dimensions and copy into LMC tensors
lmc_k[valid_mask_kv] = sampled_k.reshape(-1, num_heads * head_size)
lmc_v[valid_mask_kv] = sampled_v.reshape(-1, num_heads * head_size)
def load_and_reshape_flash(
key_value: torch.Tensor,
# Destination (Dst): Pinned CPU Tensor [2, L, T, H]
key_cache: torch.Tensor,
# Source (Src): GPU Cache [Blocks, BlockSize, NumHeads, HeadSize]
value_cache: torch.Tensor, # Source (Src): GPU Cache
slot_mapping: torch.Tensor, # Mapping indices [num_tokens]
layer_idx: int,
):
"""
Python equivalent of load_and_reshape_flash.
Note: In the context of 'test_extract_and_load_back', this function performs
an EXTRACT operation (Reads from GPU Cache and writes to Pinned CPU memory).
"""
# 1. Prepare indices on the target device
# Mapping must be on the same GPU as the cache to perform indexing
device = key_cache.device
slot_mapping = slot_mapping.to(device=device, dtype=torch.long)
block_size = key_cache.size(1)
# Calculate physical locations within the paged cache
block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor")
block_offsets = slot_mapping % block_size
# 2. Extract data from Cache (Gather operation)
# The result k_out/v_out will be on the GPU
# Shape: [num_tokens, num_heads, head_size]
k_out = key_cache[block_indices, block_offsets]
v_out = value_cache[block_indices, block_offsets]
# 3. Write to the destination tensor (CPU Copy)
# Target shape: [2, num_layers, num_tokens, hidden_dim]
# Flatten heads into the hidden dimension: [T, NumHeads, HeadSize] -> [T, HiddenDim]
hidden_dim = k_out.shape[1] * k_out.shape[2]
# Assignment automatically handles the Device-to-Host (D2H) transfer
key_value[0, layer_idx] = k_out.view(-1, hidden_dim)
key_value[1, layer_idx] = v_out.view(-1, hidden_dim)
def reshape_and_cache_back_flash(
key_value: torch.Tensor,
# Source: [2, num_layer, num_tokens, num_heads * head_size]
# (Can be on CPU/Pinned Memory or GPU)
key_cache: torch.Tensor,
# Destination: [num_blocks, block_size, num_heads, head_size]
# (Must be on GPU)
value_cache: torch.Tensor, # Destination: (Must be on GPU)
slot_mapping: torch.Tensor, # Indices: [num_tokens]
layer_idx: int,
):
"""
Python implementation of reshape_and_cache_back_flash.
Operation:
Flat Tensor (Source) -> Paged Attention Cache (Destination)
Logic:
1. Extract the specific layer's data from key_value.
2. Move it to the GPU (if it's on CPU).
3. Reshape it to match the cache's head structure.
4. Scatter (write) it into the non-contiguous cache blocks using slot_mapping.
"""
# 1. Setup Device & Dimensions
# The cache is on the GPU, so all indices and source data must eventually be there.
device = key_cache.device
block_size = key_cache.size(1)
num_heads = key_cache.size(2)
head_size = key_cache.size(3)
# 2. Prepare Indices
# slot_mapping might be on CPU, must move to GPU for indexing.
slot_mapping = slot_mapping.to(device=device, dtype=torch.long)
# Calculate physical block indices and offsets
block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor")
block_offsets = slot_mapping % block_size
# 3. Process Source Data (Key)
# Step A: Slice the specific layer from the source tensor
# Source shape: [2, num_layers, num_tokens, hidden_dim] -> [num_tokens, hidden_dim]
k_src_flat = key_value[0, layer_idx]
v_src_flat = key_value[1, layer_idx]
# Step B: Reshape & Move to GPU
# .to(device) handles the CPU -> GPU transfer if key_value is in pinned memory.
# View shape: [num_tokens, num_heads, head_size]
k_src = k_src_flat.to(device).view(-1, num_heads, head_size)
v_src = v_src_flat.to(device).view(-1, num_heads, head_size)
# 4. Write to Cache (Scatter)
# Using Advanced Indexing to write data into specific blocks/offsets
key_cache[block_indices, block_offsets] = k_src
value_cache[block_indices, block_offsets] = v_src
def lmcache_memcpy_async(
dest: int | torch.Tensor,
src: int | torch.Tensor,
nbytes: int,
direction: TransferDirection,
host_buffer_offset: int,
host_buffer_alignments: int,
):
"""
Python fallback for lmcache_memcpy_async.
- Tensor mode (non-CUDA devices like HPU): uses .to(device) + copy_()
- Pointer mode with libcudart: uses synchronous cudaMemcpy (cudaMemcpyDefault)
- Pointer mode without libcudart: uses CPU tensor copy
Unlike the C++ version (which uses cudaMemcpyAsync and must split copies
at cudaHostRegister boundaries), this Python fallback does NOT need
alignment-based chunking because:
- cudaMemcpy (synchronous) handles cross-cudaHostRegister boundaries
internally via staging buffers
- CPU tensor copy has no alignment constraints
- Tensor mode bypasses raw pointers entirely
dest:
- If int: raw memory pointer (used for CUDA/CPU devices where we
work with pointers).
- If torch.Tensor: tensor object (used for non-CUDA/CPU devices
where we operate on tensor objects directly).
src:
- If int: raw memory pointer (used for CUDA/CPU devices where we
work with pointers).
- If torch.Tensor: tensor object (used for non-CUDA/CPU devices
where we operate on tensor objects directly).
"""
# 1. Power of two check (kept for API compatibility)
if host_buffer_alignments <= 0 or (
host_buffer_alignments & (host_buffer_alignments - 1) != 0
):
raise ValueError("host_buffer_alignments must be power of two")
# 2. Validate direction
if int(direction) not in (int(TransferDirection.H2D), int(TransferDirection.D2H)):
raise ValueError(f"Unsupported direction: {direction}")
# 3. Tensor-backed mode.
# Mixed pointer/tensor are not allowed
if isinstance(dest, torch.Tensor) or isinstance(src, torch.Tensor):
if not (isinstance(dest, torch.Tensor) and isinstance(src, torch.Tensor)):
raise TypeError(
"Mixed types are not allowed: both dest and src must be torch.Tensor "
"if either of them is a tensor."
)
if nbytes % dest.element_size() != 0:
raise ValueError("nbytes must align with tensor element size")
num_elements = nbytes // dest.element_size()
dest_slice = dest.flatten()[:num_elements]
src_slice = src.flatten()[:num_elements]
copied = src_slice.to(dest_slice.device)
dest_slice.copy_(copied)
return
# 4. Pointer mode
if not isinstance(dest, int) or not isinstance(src, int):
raise TypeError(
"dest and src must be both int (pointer mode) "
"or both torch.Tensor (tensor mode)"
)
libcudart = _get_copy_lib()
if libcudart is not None and hasattr(libcudart, "cudaMemcpy"):
try:
# Synchronous cudaMemcpy handles cross-cudaHostRegister boundaries
# internally — no manual alignment splitting needed.
ret = libcudart.cudaMemcpy(
ctypes.c_void_p(dest),
ctypes.c_void_p(src),
ctypes.c_size_t(nbytes),
ctypes.c_int(4), # cudaMemcpyDefault
)
if ret != 0:
raise RuntimeError(f"cudaMemcpy failed with error code {ret}")
except AttributeError:
raise
else:
# Pure CPU copy — no alignment constraints.
_copy_bytes_with_tensor(dest, src, nbytes)
@njit(cache=True)
def _encode_single_channel(
cdf_layer_c, # np.uint32 [lp]
sym_channel, # np.uint8 [n_tokens]
out_buf_lc, # np.uint8 [buffer_size]
):
"""Core arithmetic encoding for a single (layer, channel).
Returns number of bytes written."""
MASK32 = 0xFFFFFFFF
precision = 16
max_symbol = len(cdf_layer_c) - 2
n_tokens = len(sym_channel)
low, high = 0, MASK32
pending_bits = 0
output_reg, output_reg_len = 0, 0
ptr = 0
buf_size = len(out_buf_lc)
# Inline flush_bit to avoid closure (numba does not support nonlocal)
for token_idx in range(n_tokens):
sym = int(sym_channel[token_idx])
c_low = int(cdf_layer_c[sym])
c_high = 0x10000 if sym == max_symbol else int(cdf_layer_c[sym + 1])
span = (high - low + 1) & MASK32
if span == 0:
span = 0x100000000
high = (low + ((span * c_high) >> precision) - 1) & MASK32
low = (low + ((span * c_low) >> precision)) & MASK32
while True:
if (high & 0x80000000) == (low & 0x80000000):
# flush_bit(bit)
bit = (high >> 31) & 1
output_reg = (output_reg << 1) | bit
output_reg_len += 1
if output_reg_len == 8:
if ptr < buf_size:
out_buf_lc[ptr] = output_reg & 0xFF
ptr += 1
output_reg, output_reg_len = 0, 0
# flush pending bits
for _ in range(pending_bits):
output_reg = (output_reg << 1) | (1 - bit)
output_reg_len += 1
if output_reg_len == 8:
if ptr < buf_size:
out_buf_lc[ptr] = output_reg & 0xFF
ptr += 1
output_reg, output_reg_len = 0, 0
pending_bits = 0
low = (low << 1) & MASK32
high = ((high << 1) | 1) & MASK32
elif (low & 0x40000000) != 0 and (high & 0x40000000) == 0:
pending_bits += 1
low = (low << 1) & 0x7FFFFFFF
high = ((high << 1) | 0x80000001) & MASK32
else:
break
# Final flushing sequence
pending_bits += 1
bit = 1 if (low & 0x40000000) != 0 else 0
output_reg = (output_reg << 1) | bit
output_reg_len += 1
if output_reg_len == 8:
if ptr < buf_size:
out_buf_lc[ptr] = output_reg & 0xFF
ptr += 1
output_reg, output_reg_len = 0, 0
for _ in range(pending_bits):
output_reg = (output_reg << 1) | (1 - bit)
output_reg_len += 1
if output_reg_len == 8:
if ptr < buf_size:
out_buf_lc[ptr] = output_reg & 0xFF
ptr += 1
output_reg, output_reg_len = 0, 0
pending_bits = 0 # noqa: F841
if output_reg_len > 0:
if ptr < buf_size:
out_buf_lc[ptr] = (output_reg << (8 - output_reg_len)) & 0xFF
ptr += 1
return ptr
def encode_fast_new(cdf, input_sym, output_buffer, output_lengths):
"""
Python equivalent of C++ Arithmetic Encoder.
Strictly emulates 32-bit unsigned overflow for high/low.
"""
cdf_np = cdf.cpu().numpy().view(np.uint16).astype(np.uint32)
sym_np = input_sym.cpu().numpy().astype(np.uint8)
n_layers, n_tokens, n_channels = sym_np.shape
out_buf_np = np.zeros(output_buffer.shape, dtype=np.uint8)
out_len_np = np.zeros(output_lengths.shape, dtype=np.int32)
def encode_one(args):
layer_idx, c = args
length = _encode_single_channel(
cdf_np[layer_idx, c],
sym_np[layer_idx, :, c],
out_buf_np[layer_idx, c],
)
out_len_np[layer_idx, c] = length
tasks = [(layer_idx, c) for layer_idx in range(n_layers) for c in range(n_channels)]
with ThreadPoolExecutor() as executor:
list(executor.map(encode_one, tasks))
output_buffer.copy_(torch.from_numpy(out_buf_np))
output_lengths.copy_(torch.from_numpy(out_len_np))
@njit(cache=True)
def _decode_single_channel(
cdf_layer_c,
bs_np,
start_off,
end_off,
n_tokens,
out_layer_c,
):
MASK32 = 0xFFFFFFFF
precision = 16
max_symbol = len(cdf_layer_c) - 2
v_val = 0
if start_off + 4 <= len(bs_np):
v_val = (
(int(bs_np[start_off]) << 24)
| (int(bs_np[start_off + 1]) << 16)
| (int(bs_np[start_off + 2]) << 8)
| int(bs_np[start_off + 3])
) & MASK32
low, high = 0, MASK32
byte_buffer_offset = start_off + 4
bit_idx = 1
byte_buffer = int(bs_np[byte_buffer_offset]) if byte_buffer_offset < end_off else 0
for i in range(n_tokens):
span = (high - low + 1) & MASK32
if span == 0:
span = 0x100000000
v_minus_l = (v_val - low) & MASK32
count = ((v_minus_l + 1) * 0x10000 - 1) // span
count = count & 0xFFFF
left = 0
right = max_symbol + 1
while left + 1 < right:
m = (left + right) // 2
if int(cdf_layer_c[m]) < count:
left = m
elif int(cdf_layer_c[m]) > count:
right = m
else:
left = m
break
out_layer_c[i] = left
if i == n_tokens - 1:
break
sym_i = left
c_low = int(cdf_layer_c[sym_i])
c_high = 0x10000 if sym_i == max_symbol else int(cdf_layer_c[sym_i + 1])
high = (low + ((span * c_high) >> precision) - 1) & MASK32
low = (low + ((span * c_low) >> precision)) & MASK32
while True:
if low >= 0x80000000 or high < 0x80000000:
v_val = ((v_val << 1) | ((byte_buffer >> (8 - bit_idx)) & 1)) & MASK32
low = (low << 1) & MASK32
high = ((high << 1) | 1) & MASK32
bit_idx += 1
elif low >= 0x40000000 and high < 0xC0000000:
v_val = (v_val - 0x40000000) & MASK32
v_val = ((v_val << 1) | ((byte_buffer >> (8 - bit_idx)) & 1)) & MASK32
low = (low << 1) & 0x7FFFFFFF
high = ((high << 1) | 0x80000001) & MASK32
bit_idx += 1
else:
break
if bit_idx == 9:
bit_idx = 1
byte_buffer_offset += 1
byte_buffer = (
int(bs_np[byte_buffer_offset])
if byte_buffer_offset < end_off
else 0
)
# Standard
def decode_fast_new(cdf, bytestreams, lengths, output):
"""
Python implementation of Arithmetic Decoding.
Strictly aligned with CUDA decode_with_accessor_kernel.
bytestreams shape: [nlayers, nchannels, buffer_size]
"""
cdf_np = cdf.cpu().numpy().view(np.uint16).astype(np.uint32)
bs_np = bytestreams.cpu().numpy().astype(np.uint8)
len_np = lengths.cpu().numpy().astype(np.int32)
n_layers, n_tokens, n_channels = output.shape
out_np = np.zeros(output.shape, dtype=np.uint8)
def decode_one(args):
layer_idx, c = args
curr_len = int(len_np[layer_idx, c])
# For decode_fast_new, each channel has its own contiguous buffer,
# so start_off=0 and end_off=curr_len within channel_bs
channel_bs = bs_np[layer_idx, c] # shape [buffer_size]
_decode_single_channel(
cdf_np[layer_idx, c],
channel_bs,
0,
curr_len,
n_tokens,
out_np[layer_idx, :, c],
)
tasks = [(layer_idx, c) for layer_idx in range(n_layers) for c in range(n_channels)]
with ThreadPoolExecutor() as executor:
list(executor.map(decode_one, tasks))
if output is not None:
output.copy_(torch.from_numpy(out_np))
def decode_fast_prefsum(cdf, bytestreams, lengths_prefsum, output):
"""
Python equivalent of C++ decode_fast_prefsum.
bytestreams shape: [total_bytes] (1D, all channels packed)
"""
cdf_np = cdf.cpu().numpy().view(np.uint16).astype(np.uint32)
pref_np = lengths_prefsum.cpu().numpy().astype(np.int64).flatten()
# WA: CUDA kernel reads out-of-bound in two ways:
# 1. max(prefsum) may equal len(bytestreams) (off-by-one on exclusive-end)
# 2. v_val init reads 4 bytes starting at start_off, may exceed bytestreams
# Pad with zeros to make all reads safe.
max_prefsum = int(pref_np.max())
pad_size = max(0, max_prefsum + 4 - bytestreams.shape[0])
if pad_size > 0:
bytestreams = torch.nn.functional.pad(bytestreams, (0, pad_size), value=0)
bs_np = bytestreams.cpu().numpy().astype(np.uint8) # must be after padding
n_layers, n_tokens, n_channels = output.shape
out_np = np.zeros(output.shape, dtype=np.uint8)
def decode_one(args):
layer_idx, c = args
cid = layer_idx * n_channels + c
start_off = 0 if cid == 0 else int(pref_np[cid - 1])
end_off = int(pref_np[cid])
_decode_single_channel(
cdf_np[layer_idx, c],
bs_np,
start_off,
end_off,
n_tokens,
out_np[layer_idx, :, c],
)
tasks = [(layer_idx, c) for layer_idx in range(n_layers) for c in range(n_channels)]
with ThreadPoolExecutor() as executor:
list(executor.map(decode_one, tasks))
output.copy_(torch.from_numpy(out_np))
def calculate_cdf(input_tensor: torch.Tensor, num_bins: int) -> torch.Tensor:
"""Equivalent to CUDA calculate_cdf.
Calculates the CDF across tokens for each (layer, channel) pair.
Args:
input_tensor: 3D tensor with shape [nlayers, ntokens, nchannels].
num_bins: Maximum number of bins (i.e., Lp - 1).
Returns:
int16 tensor with shape [nlayers, nchannels, num_bins + 1]
containing normalized CDF values.
"""
nlayers, ntokens, nchannels = input_tensor.shape
device = input_tensor.device
# Compute per-(layer, channel) histogram via scatter_add.
# Permute to [nlayers, nchannels, ntokens] then flatten first two dims.
input_perm = input_tensor.permute(0, 2, 1).reshape(-1, ntokens).long()
src = torch.ones_like(input_perm)
counts = torch.zeros(nlayers * nchannels, num_bins, dtype=torch.long, device=device)
counts.scatter_add_(1, input_perm.clamp(0, num_bins - 1), src)
counts = counts.reshape(nlayers, nchannels, num_bins)
# Build CDF: cdf[..., 0] = 0, cdf[..., i] = sum(counts[..., 0:i])
cdf = torch.zeros(nlayers, nchannels, num_bins + 1, dtype=torch.long, device=device)
cdf[:, :, 1:] = torch.cumsum(counts, dim=2)
# Total count per (layer, channel)
total = cdf[:, :, -1:] # [nlayers, nchannels, 1]
# Normalize: (0xFFFF - num_bins) * cdf / total + bin_index
max_uint16_value = 0xFFFF - num_bins
bin_offsets = torch.arange(num_bins + 1, dtype=torch.long, device=device)
safe_total = total.clamp(min=1)
normalized = (max_uint16_value * cdf) // safe_total + bin_offsets
# Where total is 0, use just the bin offsets
normalized = torch.where(
total > 0, normalized, bin_offsets.unsqueeze(0).unsqueeze(0)
)
return normalized.to(torch.int16)
def rotary_embedding_k_fused(
old_positions: torch.Tensor,
new_positions: torch.Tensor,
key: torch.Tensor,
head_size: int,
cos_sin_cache: torch.Tensor,
is_neox: bool,
) -> None:
"""Apply fused rotary embedding undo/redo to key tensor in-place.
Reverses the rotary embedding at old_positions and applies the rotary
embedding at new_positions. head_size is unused but kept for API
compatibility with the CUDA equivalent.
Args:
old_positions: Token positions whose rotary embedding to reverse.
new_positions: Token positions whose rotary embedding to apply.
key: Key tensor to update in-place.
head_size: Head size (unused; kept for API compatibility).
cos_sin_cache: Precomputed cosine/sine cache indexed by position.
is_neox: If True, uses NeoX-style rotary (contiguous halves);
otherwise uses GPT-J-style (interleaved).
"""
rot_dim = cos_sin_cache.shape[1]
half_rot = rot_dim // 2
old_cs = cos_sin_cache[old_positions]
new_cs = cos_sin_cache[new_positions]
oc, os = old_cs[:, :half_rot].unsqueeze(1), old_cs[:, half_rot:].unsqueeze(1)
nc, ns = new_cs[:, :half_rot].unsqueeze(1), new_cs[:, half_rot:].unsqueeze(1)
if is_neox:
x = key[..., :half_rot]
y = key[..., half_rot:rot_dim]
else:
x = key[..., :rot_dim:2]
y = key[..., 1:rot_dim:2]
x_rev = x * oc + y * os
y_rev = y * oc - x * os
x_out = x_rev * nc - y_rev * ns
y_out = y_rev * nc + x_rev * ns
if is_neox:
key[..., :half_rot] = x_out
key[..., half_rot:rot_dim] = y_out
else:
key[..., :rot_dim:2] = x_out
key[..., 1:rot_dim:2] = y_out
def get_gpu_pci_bus_id(device_id: int = 0) -> str | None:
"""
Get the PCI bus ID via CUDA/ROCm runtime.
Other backends return None.
Args:
device_id (int): CUDA/ROCm device index.
Returns:
str | None: PCI bus ID (e.g., "0000:29:00.0") or None if unavailable.
"""
try:
if torch_dev.is_available() and device_id < torch_dev.device_count():
props = torch_dev.get_device_properties(device_id)
# PCI function number is always 0 for GPUs
bus_id = (
f"{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:"
f"{props.pci_device_id:02x}.0"
)
return bus_id.upper()
except Exception:
pass
return None
# ---------------------------------------------------------------------------
# Completion recorder fallback (no CUDA stream ordering; enqueue immediately)
# ---------------------------------------------------------------------------
_completion_lock = threading.Lock()
_completion_buffer: list[tuple[str, bytes]] = []
def record_completion_on_stream(
cuda_stream_ptr: int, kind: str, payload: bytes
) -> None:
"""Fallback: immediately enqueue the completion without stream ordering.
Args:
cuda_stream_ptr: Ignored on non-CUDA path.
kind: Dispatch key identifying the handler (e.g. "finish_write").
payload: Opaque msgpack-encoded bytes forwarded to the handler.
"""
with _completion_lock:
_completion_buffer.append((kind, payload))
def drain_recorded_completions() -> list[tuple[str, bytes]]:
"""Fallback: atomically drain and return all pending completions.
Returns:
List of (kind, payload) pairs recorded since the last drain.
"""
with _completion_lock:
items = list(_completion_buffer)
_completion_buffer.clear()
return items
# ---------------------------------------------------------------------------
# Event recorder fallback (no CUDA stream ordering; timestamp immediately)
# ---------------------------------------------------------------------------
_event_lock = threading.Lock()
_event_buffer: list[tuple[str, str, float, dict[str, str], dict[str, int]]] = []
def record_event_on_stream(
cuda_stream_ptr: int,
event_type_name: str,
session_id: str,
str_metadata: dict[str, str],
int_metadata: dict[str, int],
) -> None:
"""Fallback: immediately record the event without CUDA stream ordering.
The wall-clock timestamp is captured at call time (no host-callback).
Args:
cuda_stream_ptr: Ignored on non-CUDA path.
event_type_name: Event type identifier (e.g. "mp.store.start").
session_id: Session identifier for the event.
str_metadata: String-valued metadata dict.
int_metadata: Integer-valued metadata dict.
"""
# Standard
import time
ts = time.time()
with _event_lock:
_event_buffer.append(
(event_type_name, session_id, ts, dict(str_metadata), dict(int_metadata))
)
def drain_recorded_events() -> list[
tuple[str, str, float, dict[str, str], dict[str, int]]
]:
"""Fallback: atomically drain and return all pending events.
Returns:
List of (event_type_name, session_id, timestamp, str_metadata,
int_metadata) tuples recorded since the last drain.
"""
with _event_lock:
items = list(_event_buffer)
_event_buffer.clear()
return items
Reference in New Issue View Git Blame Copy Permalink