Implementation:LMCache LMCache GPU Connectors

Knowledge Sources	LMCache
Domains	GPU Computing, KV Cache Transfer
Last Updated	2026-02-09 00:00 GMT

Overview

This module defines the GPU connector interface and multiple implementations for transferring KV cache data between CPU memory objects and GPU paged memory, supporting vLLM and SGLang backends.

Description

The gpu_connectors.py module provides an abstract GPUConnectorInterface and six concrete implementations for efficiently moving KV cache data between CPU-side MemoryObj instances and GPU-side paged KV caches. The connectors handle the complexities of paged memory layouts, slot mappings, CUDA stream management, and optional GPU intermediate buffers. Key implementations include:

• VLLMPagedMemGPUConnectorV2: The standard vLLM connector that transfers all layers at once using pointer-based multi-layer CUDA kernels, producing/consuming KV_2LTD format memory objects.
• VLLMPagedMemGPUConnectorV3: An enhanced vLLM connector supporting heterogeneous KV layer groups with per-group pointer management.
• VLLMBufferLayerwiseGPUConnector: A layer-by-layer vLLM connector using ping-pong GPU buffers, supporting fused rotary embedding re-application and gap zeroing for blended chunks. Uses a generator pattern for pipeline-style processing.
• VLLMPagedMemLayerwiseGPUConnector: A layer-by-layer vLLM connector without rotary embedding support, using single-layer CUDA kernel transfers.
• SGLangGPUConnector: Connector for SGLang's separated key/value tensor layout using unilateral transfer kernels.
• SGLangLayerwiseGPUConnector: Layer-by-layer variant of the SGLang connector with lazy GPU buffer initialization.

All connectors support MLA (Multi-Latent Attention) format alongside standard multi-head attention formats.

Usage

GPU connectors are instantiated during LMCache engine initialization, typically via the from_metadata factory method. The appropriate connector variant is selected based on the serving backend (vLLM or SGLang) and whether layerwise processing is enabled. They are used by the cache engine's store and load paths.

Code Reference

Source Location

• Repository: LMCache
• File: lmcache/v1/gpu_connector/gpu_connectors.py
• Lines: 1-1820

Signature

class GPUConnectorInterface(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def to_gpu(self, memory_obj: MemoryObj, start: int, end: int, **kwargs): ...
    @abc.abstractmethod
    def from_gpu(self, memory_obj: MemoryObj, start: int, end: int, **kwargs): ...
    @abc.abstractmethod
    def batched_from_gpu(self, memory_objs, starts, ends, **kwargs): ...
    @abc.abstractmethod
    def batched_to_gpu(self, memory_objs=None, starts=None, ends=None, **kwargs): ...
    @abc.abstractmethod
    def get_shape(self, num_tokens: int) -> torch.Size: ...
    def initialize_kvcaches_ptr(self, **kwargs): ...

class VLLMPagedMemGPUConnectorV2(GPUConnectorInterface):
    def __init__(self, hidden_dim_size: int, num_layers: int, use_gpu: bool = False, **kwargs): ...
    @classmethod
    def from_metadata(cls, metadata: LMCacheMetadata, use_gpu=False, device=None) -> "VLLMPagedMemGPUConnectorV2": ...

class VLLMPagedMemGPUConnectorV3(GPUConnectorInterface):
    def __init__(self, metadata: LMCacheMetadata, device: torch.device, use_gpu: bool = False): ...
    @classmethod
    def from_metadata(cls, metadata: LMCacheMetadata, use_gpu=False, device=None) -> "VLLMPagedMemGPUConnectorV3": ...

class VLLMBufferLayerwiseGPUConnector(GPUConnectorInterface):
    def __init__(self, hidden_dim_size: int, num_layers: int, use_gpu: bool = False, use_double_buffer: bool = True, **kwargs): ...
    @classmethod
    def from_metadata(cls, metadata: LMCacheMetadata, use_gpu=False, device=None) -> "VLLMBufferLayerwiseGPUConnector": ...
    def get_kv(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]: ...

class VLLMPagedMemLayerwiseGPUConnector(GPUConnectorInterface):
    def __init__(self, hidden_dim_size: int, num_layers: int, use_gpu: bool = False, **kwargs): ...
    @classmethod
    def from_metadata(cls, metadata: LMCacheMetadata, use_gpu=False, device=None) -> "VLLMPagedMemLayerwiseGPUConnector": ...

class SGLangGPUConnector(GPUConnectorInterface):
    def __init__(self, hidden_dim_size: int, num_layers: int, use_gpu: bool = False, **kwargs): ...

class SGLangLayerwiseGPUConnector(GPUConnectorInterface):
    def __init__(self, hidden_dim_size: int, num_layers: int, use_gpu: bool = False, **kwargs): ...

Import

from lmcache.v1.gpu_connector.gpu_connectors import (
    GPUConnectorInterface,
    VLLMPagedMemGPUConnectorV2,
    VLLMPagedMemGPUConnectorV3,
    VLLMBufferLayerwiseGPUConnector,
    VLLMPagedMemLayerwiseGPUConnector,
    SGLangGPUConnector,
    SGLangLayerwiseGPUConnector,
)

I/O Contract

Inputs

Name	Type	Required	Description
memory_obj	MemoryObj	Yes	Memory object containing the KV cache tensor to transfer (must have non-None tensor)
start	int	Yes	Starting token index in the full token sequence for this chunk
end	int	Yes	Ending token index in the full token sequence for this chunk
kvcaches	List[torch.Tensor]	Yes (via kwargs)	GPU-side paged KV cache tensors (one per layer)
slot_mapping	torch.Tensor	Yes (via kwargs)	Mapping from token positions to physical page slots
hidden_dim_size	int	Yes (constructor)	Hidden dimension size (num_kv_heads * head_size)
num_layers	int	Yes (constructor)	Number of transformer layers
use_gpu	bool	No	Whether to allocate a GPU intermediate buffer for transfers (default False)
metadata	LMCacheMetadata	Yes (from_metadata)	Engine metadata containing model shape and dtype information
sync	bool	Yes (layerwise kwargs)	Whether to synchronize CUDA streams between layers

Outputs

Name	Type	Description
(via to_gpu)	None	Data is written directly into the GPU paged KV cache tensors
(via from_gpu)	None	Data is written into the memory_obj tensor, metadata.fmt is set appropriately
(via batched_to_gpu generator)	generator yields	Layerwise connectors yield control points for pipeline processing
(via batched_from_gpu generator)	generator yields	Layerwise connectors yield after each layer is transferred
torch.Size	torch.Size	Shape of the KV tensor for a given number of tokens

Usage Examples

from lmcache.v1.gpu_connector.gpu_connectors import VLLMPagedMemGPUConnectorV2
from lmcache.v1.metadata import LMCacheMetadata

# Create connector from metadata
connector = VLLMPagedMemGPUConnectorV2.from_metadata(
    metadata=metadata,
    use_gpu=True,
    device=torch.device("cuda:0"),
)

# Transfer from CPU memory object to GPU paged cache
connector.to_gpu(
    memory_obj,
    start=0,
    end=256,
    kvcaches=gpu_kv_caches,
    slot_mapping=slot_mapping,
)

# Transfer from GPU paged cache to CPU memory object
connector.from_gpu(
    memory_obj,
    start=0,
    end=256,
    kvcaches=gpu_kv_caches,
    slot_mapping=slot_mapping,
)

# Batched transfer
connector.batched_to_gpu(
    memory_objs=[mem_obj_1, mem_obj_2],
    starts=[0, 256],
    ends=[256, 512],
    kvcaches=gpu_kv_caches,
    slot_mapping=slot_mapping,
)