lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/storage_backend/serde/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

lmcache/storage_backend/serde/cachegen_basics.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass
+from typing import List
+import io
+import pickle
+
+# Third Party
+from transformers import AutoConfig
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.utils import _lmcache_nvtx_annotate
+
+logger = init_logger(__name__)
+
+CACHEGEN_GPU_MAX_TOKENS_PER_CHUNK = 256
+
+
+@dataclass
+class QuantizationSpec:
+    start_layer: int
+    end_layer: int
+    bins: int
+
+    def __getitem__(self, key: str) -> int:
+        return getattr(self, key)
+
+
+@dataclass
+class CacheGenConfig:
+    # TODO: move this class to another file like "cachegen_basics.py"
+    nlayers: int
+    kspecs: List[QuantizationSpec]
+    vspecs: List[QuantizationSpec]
+
+    def __getitem__(self, key: str) -> int:
+        return getattr(self, key)
+
+    @staticmethod
+    def from_model_name(model_name: str) -> "CacheGenConfig":
+        family_7b = [
+            "mistralai/Mistral-7B-Instruct-v0.2",
+            "lmsys/longchat-7b-16k",
+            "Qwen/Qwen-7B",
+        ]
+        family_8b = ["meta-llama/Llama-3.1-8B-Instruct"]
+        family_9b = ["THUDM/glm-4-9b-chat"]
+        if model_name in family_7b:
+            return CacheGenConfig(
+                nlayers=32,
+                kspecs=[
+                    QuantizationSpec(start_layer=0, end_layer=10, bins=32),
+                    QuantizationSpec(start_layer=10, end_layer=32, bins=16),
+                ],
+                vspecs=[
+                    QuantizationSpec(start_layer=0, end_layer=2, bins=32),
+                    QuantizationSpec(start_layer=2, end_layer=32, bins=16),
+                ],
+            )
+        elif model_name in family_8b:
+            return CacheGenConfig(
+                nlayers=32,
+                kspecs=[
+                    QuantizationSpec(start_layer=0, end_layer=10, bins=32),
+                    QuantizationSpec(start_layer=10, end_layer=32, bins=16),
+                ],
+                vspecs=[
+                    QuantizationSpec(start_layer=0, end_layer=2, bins=32),
+                    QuantizationSpec(start_layer=2, end_layer=32, bins=16),
+                ],
+            )
+        # TODO(Jiayi): needs tuning for better quality
+        elif model_name in family_9b:
+            return CacheGenConfig(
+                nlayers=40,
+                kspecs=[
+                    QuantizationSpec(start_layer=0, end_layer=10, bins=32),
+                    QuantizationSpec(start_layer=10, end_layer=40, bins=16),
+                ],
+                vspecs=[
+                    QuantizationSpec(start_layer=0, end_layer=2, bins=32),
+                    QuantizationSpec(start_layer=2, end_layer=40, bins=16),
+                ],
+            )
+        else:
+            try:
+                config = AutoConfig.from_pretrained(model_name)
+                # Default name caught by num_hidden_layers
+                if config.num_hidden_layers is None:
+                    raise ValueError(
+                        f"num_hidden_layers is None for model {model_name}"
+                    )
+                if config.num_hidden_layers < 10:
+                    return CacheGenConfig(
+                        nlayers=config.num_hidden_layers,
+                        kspecs=[
+                            QuantizationSpec(
+                                start_layer=0,
+                                end_layer=config.num_hidden_layers,
+                                bins=32,
+                            ),
+                        ],
+                        vspecs=[
+                            QuantizationSpec(
+                                start_layer=0,
+                                end_layer=config.num_hidden_layers,
+                                bins=32,
+                            ),
+                        ],
+                    )
+                else:
+                    return CacheGenConfig(
+                        nlayers=config.num_hidden_layers,
+                        kspecs=[
+                            QuantizationSpec(start_layer=0, end_layer=10, bins=32),
+                            QuantizationSpec(
+                                start_layer=10,
+                                end_layer=config.num_hidden_layers,
+                                bins=16,
+                            ),
+                        ],
+                        vspecs=[
+                            QuantizationSpec(start_layer=0, end_layer=2, bins=32),
+                            QuantizationSpec(
+                                start_layer=2,
+                                end_layer=config.num_hidden_layers,
+                                bins=16,
+                            ),
+                        ],
+                    )
+            except Exception as e:
+                raise ValueError(
+                    f"Model {model_name} not supported by CacheGenConfig"
+                ) from e
+
+
+@dataclass
+class CacheGenEncoderOutput:
+    # TODO: maybe use numpy array so that we can directly tobytes() and
+    # frombuffer() to have a better performance
+    bytestream: bytes
+    start_indices: torch.Tensor
+    cdf: torch.Tensor
+    max_tensors_key: torch.Tensor
+    max_tensors_value: torch.Tensor
+    num_heads: int
+    head_size: int
+
+    def __getitem__(self, key: str) -> int:
+        return getattr(self, key)
+
+    def to_bytes(self) -> bytes:
+        """Save the output to a file"""
+        with io.BytesIO() as f:
+            # torch.save(self, f)
+            pickle.dump(self, f)
+            return f.getvalue()
+
+    @staticmethod
+    def from_bytes(bs: bytes) -> "CacheGenEncoderOutput":
+        with io.BytesIO(bs) as f:
+            return pickle.load(f)
+
+
+@dataclass
+class CacheGenGPUBytestream:
+    bytestream: torch.Tensor
+    bytestream_lengths: torch.Tensor  # [nlayers, nchannels, bytestream_length]
+    ntokens: int
+
+    def __getitem__(self, key: str) -> int:
+        return getattr(self, key)
+
+
+@dataclass
+class CacheGenGPUEncoderOutput:
+    data_chunks: List[CacheGenGPUBytestream]
+    cdf: torch.Tensor
+    max_tensors_key: torch.Tensor
+    max_tensors_value: torch.Tensor
+    num_heads: int
+    head_size: int
+
+    def __getitem__(self, key: str) -> int:
+        return getattr(self, key)
+
+    @_lmcache_nvtx_annotate
+    def to_bytes(self) -> bytes:
+        """Save the output to a file"""
+        with io.BytesIO() as f:
+            pickle.dump(self, f)
+            return f.getvalue()
+
+    @staticmethod
+    @_lmcache_nvtx_annotate
+    def from_bytes(bs: bytes) -> "CacheGenGPUEncoderOutput":
+        with io.BytesIO(bs) as f:
+            return pickle.load(f)
+
+    def debug_print_device(self):
+        logger.debug(f"bytestream device: {self.data_chunks[0].bytestream.device}")
+        logger.debug(
+            f"bytestream_lengths device: "
+            f"{self.data_chunks[0].bytestream_lengths.device}"
+        )
+        logger.debug(f"cdf device: {self.cdf.device}")
+        logger.debug(f"max_tensors_key device: {self.max_tensors_key.device}")
+        logger.debug(f"max_tensors_value device: {self.max_tensors_value.device}")

lmcache/storage_backend/serde/cachegen_decoder.py

@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import List, Optional
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.storage_backend.serde.cachegen_basics import (
+    CacheGenConfig,
+    CacheGenGPUBytestream,
+    CacheGenGPUEncoderOutput,
+)
+from lmcache.storage_backend.serde.serde import Deserializer
+from lmcache.utils import _lmcache_nvtx_annotate
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.metadata import LMCacheMetadata
+import lmcache.c_ops as lmc_ops
+import lmcache.storage_backend.serde.cachegen_basics as CGBasics
+
+logger = init_logger(__name__)
+
+
+@_lmcache_nvtx_annotate
+def quant(bins: int, xq: torch.Tensor, max1: float):
+    C = bins // 2 - 1
+    x = xq / C * max1
+    return x
+
+
+def do_dequantize(t: torch.Tensor, bins: torch.Tensor, maxtensors: torch.Tensor):
+    """
+    t: [nlayers, ntokens, nchannels]
+    bins: [nlayers]
+    maxtensors: [nlayers, ntokens, 1]
+    """
+    C = (bins // 2 - 1)[:, None, None]
+    t = t - C
+    t = t / C
+    t = t * maxtensors
+    return t
+
+
+@_lmcache_nvtx_annotate
+def recombine_bytes(bytes_tensor, output_lengths) -> torch.Tensor:
+    output_buffer_size = CGBasics.CACHEGEN_GPU_MAX_TOKENS_PER_CHUNK
+    offsets = output_lengths.flatten().cumsum(0).roll(1).reshape(output_lengths.shape)
+    offsets[0][0] = 0
+    indexes = torch.arange(output_buffer_size, device=offsets.device).tile(
+        (output_lengths.shape[0], output_lengths.shape[1], 1)
+    )
+    final_indexes = (indexes + offsets[:, :, None]).clamp(max=len(bytes_tensor) - 1)
+    return bytes_tensor[final_indexes]
+
+
+@_lmcache_nvtx_annotate
+def decode_chunk(
+    cdf: torch.Tensor,
+    data_chunk: CacheGenGPUBytestream,
+    target_buffer: torch.Tensor,
+) -> None:
+    """
+    Write the decode output in target_buffer
+    Expected shape: [nlayers (kv in total), ntokens, nchannels]
+    """
+    bytes_tensor = data_chunk.bytestream
+    length_prefsum = (
+        data_chunk.bytestream_lengths.flatten()
+        .cumsum(0)
+        .reshape(data_chunk.bytestream_lengths.shape)
+    )
+    lmc_ops.decode_fast_prefsum(cdf, bytes_tensor, length_prefsum, target_buffer)
+
+
+@_lmcache_nvtx_annotate
+def decode_function_gpu(
+    cdf: torch.Tensor,
+    data_chunks: List[CacheGenGPUBytestream],
+    layers_in_key: int,
+    chunk_size: int,
+    output: torch.Tensor,
+):
+    # TODO: dtype and shape -- still have 128 and 8
+    """
+    Given the path to the encoded KV bytestream, decode the KV cache
+
+    Inputs:
+        cdf: the cdf tensor, in shape [2 * nlayers, nchannels, bins + 1]
+        data_chunks: the data_chunks in the encoder's output
+        layers_in_key: number of layers in K (or V)
+        (K/V should have the same number of layers)
+        chunk_size: the chunk_size
+        output: output buffer, in shape [ntokens, 2 * nlayers * nchannels]
+
+    Outputs:
+        key: the decoded key tensor in the shape of (layers, tokens, nchannels)
+        value: the decoded value tensor in the shape of
+        (layers, tokens, nchannels)
+    """
+    nlayers, nchannels, _ = cdf.shape
+    output = output.reshape((nlayers, chunk_size, nchannels))
+
+    start = 0
+    for data_chunk in data_chunks:
+        end = start + data_chunk.ntokens
+        decode_chunk(cdf, data_chunk, output[:, start:end, :])
+        start = end
+
+    out = output.reshape((2, layers_in_key, chunk_size, nchannels))
+    key, value = out.float()
+
+    return key, value
+
+
+class CacheGenDeserializer(Deserializer):
+    def __init__(
+        self,
+        config: LMCacheEngineConfig,
+        metadata: LMCacheMetadata,
+        dtype,
+    ):
+        self.dtype = dtype
+        self.cachegen_config = CacheGenConfig.from_model_name(metadata.model_name)
+        self.chunk_size = config.chunk_size
+        self.output_buffer: Optional[torch.Tensor] = None
+        self.key_bins = self.make_key_bins(self.cachegen_config)
+        self.value_bins = self.make_value_bins(self.cachegen_config)
+
+    def make_key_bins(self, config: CacheGenConfig) -> torch.Tensor:
+        ret = torch.zeros(config.nlayers)
+        for spec in config.kspecs:
+            ret[spec.start_layer : spec.end_layer] = spec.bins
+        return ret.cuda()
+
+    def make_value_bins(self, config: CacheGenConfig) -> torch.Tensor:
+        ret = torch.zeros(config.nlayers)
+        for spec in config.vspecs:
+            ret[spec.start_layer : spec.end_layer] = spec.bins
+        return ret.cuda()
+
+    def get_output_buffer(self, nlayers: int, nchannels: int, ntokens: int):
+        if (
+            self.output_buffer is None
+            or self.output_buffer.shape[1] != 2 * nlayers * nchannels
+        ):
+            self.output_buffer = torch.zeros(
+                (self.chunk_size, 2 * nlayers * nchannels), dtype=torch.uint8
+            ).cuda()
+        return self.output_buffer[:ntokens, :]
+
+    @_lmcache_nvtx_annotate
+    def from_bytes(self, bs: bytes) -> torch.Tensor:
+        encoder_output = CacheGenGPUEncoderOutput.from_bytes(bs)
+        encoder_output.max_tensors_key = encoder_output.max_tensors_key.cuda()
+        encoder_output.max_tensors_value = encoder_output.max_tensors_value.cuda()
+
+        ntokens = encoder_output.max_tensors_key.shape[1]
+        layers_in_key = encoder_output.max_tensors_key.shape[0]
+        key, value = decode_function_gpu(
+            encoder_output.cdf,
+            encoder_output.data_chunks,
+            layers_in_key,
+            ntokens,
+            self.get_output_buffer(
+                encoder_output.cdf.shape[0] // 2,
+                encoder_output.cdf.shape[1],
+                ntokens,
+            ),
+        )
+
+        # Temporary fix for #83: change the device of key_bins and value_bins
+        # to the device of key and value
+        # This requires a long-term fix in the future. Currently,
+        # CacheGenGPUEncoderOutput has implicit device in itself.
+        # More specifically, if the encoder encodes the tensor on GPU0, the
+        # from_bytes will also return a tensor on GPU0
+        # We may want to dynamically configure the device based on config and
+        # metadata in the future
+        if self.key_bins.device != key.device:
+            self.key_bins = self.key_bins.to(key.device)
+
+        if self.value_bins.device != value.device:
+            self.value_bins = self.value_bins.cuda()
+
+        key = do_dequantize(key, self.key_bins, encoder_output.max_tensors_key)
+        value = do_dequantize(value, self.value_bins, encoder_output.max_tensors_value)
+        """ merge key and value back and reshape """
+        nlayers, ntokens, nchannels = key.shape
+        blob = torch.stack([key, value])  # [2, nlayers, ntokens, nchannels]
+        blob = blob.reshape(
+            (
+                2,
+                nlayers,
+                ntokens,
+                encoder_output.num_heads,
+                encoder_output.head_size,
+            )
+        )
+
+        return blob.permute((1, 0, 2, 3, 4)).to(
+            self.dtype
+        )  # [nlayers, 2, ntokens, num_heads, head_size]
+        # huggingface
+        # return blob.permute((1, 0, 3, 2, 4)).to(
+        #     self.dtype
+        # )  # [nlayers, 2, num_heads, ntokens, head_size]