wehub-resource-sync
346 lines
12 KiB
Python
346 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Abstract base class for platform cache contexts.
|
|
|
|
Defines the common interface shared by :class:`GPUCacheContext` and
|
|
:class:`CPUCacheContext`. Concrete subclasses provide
|
|
device-specific implementations of stream / buffer / copy primitives
|
|
while the base class owns layout-agnostic helpers (shape calculation,
|
|
status reporting, block-ID staging).
|
|
"""
|
|
|
|
# Future
|
|
from __future__ import annotations
|
|
|
|
# Standard
|
|
from abc import ABC, abstractmethod
|
|
from typing import TYPE_CHECKING, Any, ClassVar
|
|
import array
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import (
|
|
get_attention_backend,
|
|
get_concrete_engine_kv_shape_from_shape_desc,
|
|
get_engine_kv_shape_description,
|
|
is_mla,
|
|
)
|
|
from lmcache.v1.kv_layer_groups import KVLayerGroupsManager
|
|
|
|
if TYPE_CHECKING:
|
|
# First Party
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
|
|
class BaseCacheContext(ABC):
|
|
"""Abstract base for GPU and CPU cache contexts.
|
|
|
|
Subclasses call :meth:`__init__` after computing the common
|
|
layout parameters and before setting up device-specific state.
|
|
All keyword arguments are required so the contract is explicit.
|
|
|
|
Concrete subclasses MUST set :attr:`device_type` to the
|
|
``torch.device.type`` string they handle (``"cuda"``, ``"cpu"``,
|
|
...). The platform-agnostic :func:`create_cache_context` factory
|
|
uses this attribute (via the platform registry) to pick the right
|
|
subclass without any ``isinstance`` / ``if-elif`` chain.
|
|
"""
|
|
|
|
#: ``torch.device.type`` string the subclass handles. Concrete
|
|
#: subclasses MUST override this.
|
|
device_type: ClassVar[str] = ""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
kv_caches: list[torch.Tensor],
|
|
device: torch.device,
|
|
num_layers: int,
|
|
kv_layer_groups_manager: KVLayerGroupsManager,
|
|
block_ids_buffer: torch.Tensor,
|
|
lmcache_tokens_per_chunk: int,
|
|
) -> None:
|
|
self.kv_caches_ = kv_caches
|
|
self.device_ = device
|
|
self.num_layers_ = num_layers
|
|
self.kv_layer_groups_manager_ = kv_layer_groups_manager
|
|
self.block_ids_buffer_ = block_ids_buffer
|
|
self.lmcache_tokens_per_chunk = lmcache_tokens_per_chunk
|
|
|
|
# ------------------------------------------------------------------
|
|
# Abstract -- subclasses MUST implement
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
@abstractmethod
|
|
def stream(self) -> Any:
|
|
"""Returns the device-specific stream for async operations."""
|
|
...
|
|
|
|
@property
|
|
@abstractmethod
|
|
def cupy_stream(self) -> Any:
|
|
"""Returns the cupy ExternalStream wrapping *stream*."""
|
|
...
|
|
|
|
@property
|
|
@abstractmethod
|
|
def max_batch_size(self) -> int:
|
|
"""Returns the maximum number of concurrent batches."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def close(self) -> None:
|
|
"""Release device-specific resources (GDS staging buffers, etc.)."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_kernel_group_kv_pointers(self, kernel_group_idx: int) -> torch.Tensor:
|
|
"""Returns the KV-cache pointer tensor for *kernel_group_idx*."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_temp_kernel_group_buffer(
|
|
self, batch_idx: int, kernel_group_idx: int
|
|
) -> torch.Tensor:
|
|
"""Returns a typed temp-buffer view for a (batch, kernel-group)
|
|
pair."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_temp_object_group_buffer(
|
|
self, batch_idx: int, object_group_idx: int
|
|
) -> torch.Tensor:
|
|
"""Returns a flat uint8 temp-buffer view for a (batch, object-group)
|
|
pair."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def get_kernel_group_shape_dtype(
|
|
self,
|
|
num_tokens: int,
|
|
kernel_group_idx: int,
|
|
) -> tuple[torch.Size, torch.dtype]:
|
|
"""Returns ``(shape, dtype)`` for *kernel_group_idx*."""
|
|
...
|
|
|
|
@abstractmethod
|
|
def cache_size_per_token(self) -> int:
|
|
"""Returns cache size per logical token in bytes (all groups)."""
|
|
...
|
|
|
|
# ------------------------------------------------------------------
|
|
# Concrete -- shared implementations
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
"""Returns the device where KV-cache tensors live."""
|
|
return self.device_
|
|
|
|
@property
|
|
def kv_tensors(self) -> list[torch.Tensor]:
|
|
"""Returns the list of per-layer KV cache tensors."""
|
|
return self.kv_caches_
|
|
|
|
@property
|
|
def num_layers(self) -> int:
|
|
"""Returns the number of layers in the model."""
|
|
return self.num_layers_
|
|
|
|
@property
|
|
def num_blocks(self) -> int:
|
|
"""Returns the number of blocks in the KV cache.
|
|
|
|
Sourced from the kernel groups (one shared block-id space), not a
|
|
representative-format computation.
|
|
"""
|
|
return self.kv_layer_groups_manager_.num_blocks
|
|
|
|
@property
|
|
def hidden_dim_sizes(self) -> list[int]:
|
|
"""Returns hidden dimension sizes per KV layer group."""
|
|
return [
|
|
group.hidden_dim_size
|
|
for group in self.kv_layer_groups_manager.kv_layer_groups
|
|
]
|
|
|
|
@property
|
|
def kv_layer_groups_manager(self) -> KVLayerGroupsManager:
|
|
"""Returns the KV layer groups manager."""
|
|
return self.kv_layer_groups_manager_
|
|
|
|
def calculate_num_blocks(self, num_tokens: int, kernel_group_idx: int) -> int:
|
|
"""Calculate the number of blocks for *num_tokens* in a kernel
|
|
group."""
|
|
return self.kv_layer_groups_manager.calculate_num_blocks(
|
|
kernel_group_idx, num_tokens
|
|
)
|
|
|
|
def get_shape_desc(self, group_idx: int) -> "lmc_ops.PageBufferShapeDesc":
|
|
"""Returns the PageBufferShapeDesc for *group_idx*."""
|
|
return self.kv_layer_groups_manager.get_shape_desc(group_idx)
|
|
|
|
def get_engine_kv_format(self, kernel_group_idx: int) -> "lmc_ops.EngineKVFormat":
|
|
"""Returns the Engine KV format of kernel *kernel_group_idx*.
|
|
|
|
Raises:
|
|
ValueError: If the group has no format (a bookkeeping group built by
|
|
``parse_kvcache_shape_spec`` should never reach the transfer
|
|
path; detection-built groups always carry one).
|
|
"""
|
|
groups = self.kv_layer_groups_manager.kv_layer_groups
|
|
engine_kv_format = groups[kernel_group_idx].engine_kv_format
|
|
if engine_kv_format is None:
|
|
raise ValueError(
|
|
f"kernel group {kernel_group_idx} has no engine_kv_format; a "
|
|
"formatless bookkeeping group reached the transfer path"
|
|
)
|
|
return engine_kv_format
|
|
|
|
def engine_kv_formats(self) -> list["lmc_ops.EngineKVFormat"]:
|
|
"""Returns the Engine KV format of each kernel group, in group order."""
|
|
num_groups = len(self.kv_layer_groups_manager.kernel_groups)
|
|
return [self.get_engine_kv_format(idx) for idx in range(num_groups)]
|
|
|
|
def engine_kv_format_per_layer(self) -> list["lmc_ops.EngineKVFormat | None"]:
|
|
"""Returns each layer's Engine KV format, indexed by layer index.
|
|
|
|
Formats differ across layers for a mixed-format model. ``None`` marks a
|
|
layer in no kernel group (a cross-layer KV-sharing layer).
|
|
"""
|
|
formats: list["lmc_ops.EngineKVFormat | None"] = [None] * len(self.kv_caches_)
|
|
for kernel_group_idx, group in enumerate(
|
|
self.kv_layer_groups_manager.kernel_groups
|
|
):
|
|
fmt = self.get_engine_kv_format(kernel_group_idx)
|
|
for layer_idx in group.layer_indices:
|
|
formats[layer_idx] = fmt
|
|
return formats
|
|
|
|
def get_slots_per_chunk_in_sw(self, kernel_group_idx: int) -> int:
|
|
"""Returns the number of slots per lmcache chunk for D/H
|
|
transfer."""
|
|
return self.kv_layer_groups_manager.get_slots_per_chunk_in_sw(kernel_group_idx)
|
|
|
|
def get_kv_buffer_shape(
|
|
self, logical_num_tokens: int, group_idx: int = 0
|
|
) -> torch.Size:
|
|
"""Returns the KV buffer shape for *logical_num_tokens*."""
|
|
group = self.kv_layer_groups_manager.kv_layer_groups[group_idx]
|
|
compress_ratio = group.tokens_per_block // group.slots_per_block
|
|
if logical_num_tokens % compress_ratio != 0:
|
|
raise ValueError(
|
|
"logical_num_tokens (%d) is not a multiple of "
|
|
"compress_ratio (%d) for group %d"
|
|
% (logical_num_tokens, compress_ratio, group_idx)
|
|
)
|
|
num_slots = logical_num_tokens // compress_ratio
|
|
sd = group.shape_desc
|
|
return torch.Size(
|
|
(sd.kv_size, group.num_layers, num_slots, group.hidden_dim_size)
|
|
)
|
|
|
|
def stage_block_ids(
|
|
self, block_ids_per_group: list[list[int]]
|
|
) -> list[torch.Tensor]:
|
|
"""Stage per-group block IDs into the shared staging buffer.
|
|
|
|
Returns one non-overlapping view per LMCache group.
|
|
"""
|
|
offsets = [0]
|
|
flat: array.array = array.array("q")
|
|
for view_block_ids in block_ids_per_group:
|
|
flat.extend(view_block_ids)
|
|
offsets.append(len(flat))
|
|
|
|
total = offsets[-1]
|
|
if total > self.block_ids_buffer_.shape[0]:
|
|
raise ValueError(
|
|
"block ID total %d exceeds the pre-allocated buffer "
|
|
"size %d" % (total, self.block_ids_buffer_.shape[0])
|
|
)
|
|
if total:
|
|
cpu_tensor = torch.frombuffer(flat, dtype=torch.long)
|
|
self.block_ids_buffer_[:total].copy_(cpu_tensor, non_blocking=True)
|
|
|
|
return [
|
|
self.block_ids_buffer_[offsets[i] : offsets[i + 1]]
|
|
for i in range(len(block_ids_per_group))
|
|
]
|
|
|
|
# ------------------------------------------------------------------
|
|
# Derived properties (pure helpers)
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def concrete_engine_kv_shape(self) -> str:
|
|
"""Returns the engine KV shape with actual numeric values."""
|
|
group = self.kv_layer_groups_manager.kv_layer_groups[0]
|
|
return get_concrete_engine_kv_shape_from_shape_desc(
|
|
group.shape_desc, group.engine_kv_format
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Shared report_status
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_group_report_map(self) -> dict[int, int]:
|
|
"""Map each kernel-group index to its owning object-group index."""
|
|
return {
|
|
kg_idx: og_idx
|
|
for og_idx, og in enumerate(self.kv_layer_groups_manager.object_groups)
|
|
for kg_idx in og.kernel_group_indices
|
|
}
|
|
|
|
def _build_single_group_report(
|
|
self,
|
|
kernel_group_idx: int,
|
|
group: Any,
|
|
group_map: dict[int, int],
|
|
) -> dict:
|
|
"""Build a status dict for a single kernel group.
|
|
|
|
Override this in subclasses to inject extra per-group fields
|
|
without duplicating the whole :meth:`report_status` method.
|
|
"""
|
|
engine_kv_format = self.get_engine_kv_format(kernel_group_idx)
|
|
return {
|
|
"kernel_group_idx": kernel_group_idx,
|
|
"engine_group_idx": group.engine_group_idx,
|
|
"object_group_idx": group_map.get(kernel_group_idx, 0),
|
|
"num_layers": group.num_layers,
|
|
"layer_indices": list(group.layer_indices),
|
|
"tokens_per_block": group.tokens_per_block,
|
|
"slots_per_block": group.slots_per_block,
|
|
"dtype": str(group.dtype),
|
|
"engine_kv_concrete_shape": (
|
|
get_concrete_engine_kv_shape_from_shape_desc(
|
|
group.shape_desc, engine_kv_format
|
|
)
|
|
),
|
|
"is_mla": is_mla(engine_kv_format),
|
|
"engine_kv_format": engine_kv_format.name,
|
|
"engine_kv_shape": get_engine_kv_shape_description(engine_kv_format),
|
|
"attention_backend": get_attention_backend(engine_kv_format),
|
|
}
|
|
|
|
def report_status(self) -> dict:
|
|
"""Return this context's KV cache layout metadata."""
|
|
manager = self.kv_layer_groups_manager
|
|
kernel_groups = manager.kernel_groups
|
|
group_map = self._build_group_report_map()
|
|
|
|
group_reports = [
|
|
self._build_single_group_report(kernel_group_idx, group, group_map)
|
|
for kernel_group_idx, group in enumerate(kernel_groups)
|
|
]
|
|
|
|
return {
|
|
"num_layers": self.num_layers,
|
|
"num_blocks": self.num_blocks,
|
|
"cache_size_per_token": self.cache_size_per_token(),
|
|
"kernel_groups": group_reports,
|
|
}
|