Implementation:LMCache LMCache KV Layer Groups

Knowledge Sources	LMCache
Domains	KV Cache Management, Model Architecture
Last Updated	2026-02-09 00:00 GMT

Overview

Manages grouping of transformer layers by their KV cache tensor shape and dtype, enabling efficient handling of models with heterogeneous layer structures.

Description

This module provides KVLayerGroupInfo and KVLayerGroupsManager for organizing transformer layers into groups based on their KV cache structure. KVLayerGroupInfo is a dataclass describing a group of layers sharing the same tensor shape and dtype, with fast O(1) membership checking via internal sets. It supports both MHA layouts (5D: [2, num_blocks, block_size, num_heads, head_size]) and MLA layouts (3D: [num_blocks, block_size, head_size]), and provides hidden_dim_size as a computed property. KVLayerGroupsManager holds a list of groups and provides lookup methods by layer index or name, shape/dtype queries, and a build_kv_layer_groups method that analyzes a dictionary of KV cache tensors to automatically discover and organize the group structure. Groups are sorted by their first layer index to maintain order.

Usage

Use KVLayerGroupsManager during model initialization to discover and organize layer groups based on the actual KV cache tensors. Query it throughout the system when needing to know the shape, dtype, or group membership of specific layers, particularly for models with mixed attention types where different layer groups have different KV cache structures.

Code Reference

Source Location

• Repository: LMCache
• File: lmcache/v1/kv_layer_groups.py
• Lines: 1-211

Signature

@dataclass
class KVLayerGroupInfo:
    layer_names: list[str]
    layer_indices: list[int]
    shape: torch.Size
    dtype: torch.dtype

    @property
    def num_layers(self) -> int: ...
    @property
    def hidden_dim_size(self) -> int: ...
    def contains_layer(self, layer_idx: int) -> bool: ...
    def contains_layer_name(self, layer_name: str) -> bool: ...

@dataclass
class KVLayerGroupsManager:
    kv_layer_groups: list[KVLayerGroupInfo]

    @property
    def num_groups(self) -> int: ...
    def get_group_by_layer_idx(self, layer_idx: int) -> Optional[KVLayerGroupInfo]: ...
    def get_group_by_layer_name(self, layer_name: str) -> Optional[KVLayerGroupInfo]: ...
    def get_layer_shape(self, layer_idx: int) -> Optional[torch.Size]: ...
    def get_layer_dtype(self, layer_idx: int) -> Optional[torch.dtype]: ...
    def build_kv_layer_groups(self, kv_caches: dict[str, torch.Tensor]) -> None: ...

Import

from lmcache.v1.kv_layer_groups import KVLayerGroupInfo, KVLayerGroupsManager

I/O Contract

Inputs

Name	Type	Required	Description
kv_caches	dict[str, torch.Tensor]	Yes	Dictionary mapping layer names to their KV cache tensors (for build_kv_layer_groups)
layer_idx	int	Yes	0-based layer index for lookup methods
layer_name	str	Yes	Layer name string for lookup methods

Outputs

Name	Type	Description
KVLayerGroupInfo	KVLayerGroupInfo or None	Group info for the queried layer, or None if not found
shape	torch.Size or None	KV cache tensor shape for the queried layer
dtype	torch.dtype or None	KV cache tensor dtype for the queried layer
num_groups	int	Total number of distinct layer groups

Usage Examples

from lmcache.v1.kv_layer_groups import KVLayerGroupsManager

# Build groups from actual KV caches
manager = KVLayerGroupsManager()
manager.build_kv_layer_groups(kv_caches=model.kv_caches)

# Query group structure
print(f"Number of groups: {manager.num_groups}")

# Look up a specific layer
group = manager.get_group_by_layer_idx(5)
if group is not None:
    print(f"Layer 5 shape: {group.shape}, dtype: {group.dtype}")
    print(f"Hidden dim: {group.hidden_dim_size}")
    print(f"Group has {group.num_layers} layers")

# Get shape/dtype directly
shape = manager.get_layer_shape(10)
dtype = manager.get_layer_dtype(10)