Implementation:LMCache LMCache CacheGen Decoder

Knowledge Sources	LMCache
Domains	Compression, Serialization, CUDA
Last Updated	2026-02-09 00:00 GMT

Overview

Implements the CacheGen GPU-accelerated decoder for deserializing compressed KV cache bytestreams back into tensor format.

Description

This module provides the CacheGenDeserializer class and supporting functions for decoding KV cache data that was compressed using the CacheGen entropy coding scheme. The decoder operates on GPU using custom CUDA kernels via lmc_ops.decode_fast_prefsum. The decoding pipeline involves: parsing the compressed bytestream into CacheGenGPUEncoderOutput structures, performing entropy decoding using cumulative distribution functions (CDFs), and then dequantizing the decoded integer values back to floating-point tensors using stored per-layer maximum values and bin configurations. The final output is reshaped into the standard LMCache KV format of [nlayers, 2, ntokens, num_heads, head_size].

Usage

Use CacheGenDeserializer as the deserialization backend when LMCache is configured to use CacheGen compression. It is instantiated with the engine config, model metadata, and target dtype, and its from_bytes method converts compressed bytes to KV tensors.

Code Reference

Source Location

• Repository: LMCache
• File: lmcache/storage_backend/serde/cachegen_decoder.py
• Lines: 1-211

Signature

def quant(bins: int, xq: torch.Tensor, max1: float) -> torch.Tensor: ...
def do_dequantize(t: torch.Tensor, bins: torch.Tensor, maxtensors: torch.Tensor) -> torch.Tensor: ...
def recombine_bytes(bytes_tensor, output_lengths) -> torch.Tensor: ...
def decode_chunk(cdf: torch.Tensor, data_chunk: CacheGenGPUBytestream, target_buffer: torch.Tensor) -> None: ...
def decode_function_gpu(cdf: torch.Tensor, data_chunks: List[CacheGenGPUBytestream], layers_in_key: int, chunk_size: int, output: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: ...

class CacheGenDeserializer(Deserializer):
    def __init__(self, config: LMCacheEngineConfig, metadata: LMCacheMetadata, dtype): ...
    def make_key_bins(self, config: CacheGenConfig) -> torch.Tensor: ...
    def make_value_bins(self, config: CacheGenConfig) -> torch.Tensor: ...
    def get_output_buffer(self, nlayers: int, nchannels: int, ntokens: int) -> torch.Tensor: ...
    def from_bytes(self, bs: bytes) -> torch.Tensor: ...

Import

from lmcache.storage_backend.serde.cachegen_decoder import CacheGenDeserializer

I/O Contract

Inputs

Name	Type	Required	Description
config	LMCacheEngineConfig	Yes	Engine configuration with chunk_size and other settings
metadata	LMCacheMetadata	Yes	Model metadata including model_name for CacheGen config lookup
dtype	torch.dtype	Yes	Target data type for the output tensors
bs	bytes	Yes (from_bytes)	Compressed CacheGen bytestream
cdf	torch.Tensor	Yes (decode functions)	Cumulative distribution function tensor of shape [2*nlayers, nchannels, bins+1]
data_chunks	List[CacheGenGPUBytestream]	Yes (decode_function_gpu)	List of compressed bytestream chunks

Outputs

Name	Type	Description
tensor	torch.Tensor	Decoded KV cache tensor of shape [nlayers, 2, ntokens, num_heads, head_size]
key, value	Tuple[torch.Tensor, torch.Tensor]	Separate key and value tensors of shape [nlayers, ntokens, nchannels] (from decode_function_gpu)

Usage Examples

from lmcache.storage_backend.serde.cachegen_decoder import CacheGenDeserializer
import torch

deserializer = CacheGenDeserializer(
    config=engine_config,
    metadata=model_metadata,
    dtype=torch.float16,
)

# Decode compressed bytes to KV tensor
kv_tensor = deserializer.from_bytes(compressed_bytes)
# kv_tensor.shape: [nlayers, 2, ntokens, num_heads, head_size]