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

lmcache/v1/compute/blend/metadata.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass
+from typing import List, Optional
+
+# Third Party
+import torch
+
+
+@dataclass
+class LMCBlendCommonMetadata:
+    """
+    CommonMetadata (fixed hyperparams) for blending operations in LMCache.
+    """
+
+    check_layers: List[int]
+    recomp_ratios: Optional[List[float]] = None
+    thresholds: Optional[List[float]] = None
+
+
+@dataclass
+class LMCBlendMetadata:
+    """
+    Metadata (determined during runtime) for blending operations in LMCache.
+    """
+
+    imp_indices: Optional[torch.Tensor] = None
+    attn_mask: Optional[torch.Tensor] = None
+    positions: Optional[torch.Tensor] = None
+
+    def clean(self):
+        self.imp_indices = None
+        self.attn_mask = None
+        self.positions = None

lmcache/v1/compute/blend/utils.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING, Dict
+
+# Third Party
+from torch import nn
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.compute.blend.blender import LMCBlender
+from lmcache.v1.compute.models.utils import VLLMModelTracker
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.cache_engine import LMCacheEngine
+    from lmcache.v1.config import LMCacheEngineConfig
+    from lmcache.v1.gpu_connector import GPUConnectorInterface
+
+logger = init_logger(__name__)
+
+
+class LMCBlenderBuilder:
+    _blenders: Dict[str, LMCBlender] = {}
+
+    @classmethod
+    def get_or_create(
+        cls,
+        instance_id: str,
+        cache_engine: "LMCacheEngine",
+        gpu_connector: "GPUConnectorInterface",
+        config: "LMCacheEngineConfig",
+    ):
+        """
+        Get or create a blender for the given instance_id.
+        """
+
+        if instance_id not in cls._blenders:
+            logger.info(f"Creating blender for {instance_id}")
+            vllm_model = VLLMModelTracker.get_model(instance_id)
+            blender = LMCBlender(
+                cache_engine=cache_engine,
+                gpu_connector=gpu_connector,
+                vllm_model=vllm_model,
+                config=config,
+            )
+            cls._blenders[instance_id] = blender
+        else:
+            logger.info(
+                f"Blender for {instance_id} already exists, returning the original one."
+            )
+        return cls._blenders[instance_id]
+
+    @classmethod
+    def get(
+        cls,
+        instance_id: str,
+    ) -> nn.Module:
+        """
+        Get the blender by instance_id.
+        """
+        if instance_id not in cls._blenders:
+            raise ValueError(f"Blender for {instance_id} not found.")
+        return cls._blenders[instance_id]

lmcache/v1/compute/models/base.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from abc import ABC, abstractmethod
+
+# Third Party
+from torch import nn
+import torch
+
+# First Party
+from lmcache.v1.compute.attention.utils import infer_attn_backend_from_vllm
+from lmcache.v1.compute.positional_encoding import get_fused_rope
+
+# TODO(Jiayi): A few things need to be tested/supported:
+# TP, PP, Multimodal
+
+
+class LMCBaseModel(nn.Module, ABC):
+    def __init__(
+        self,
+        vllm_model,
+        blender,
+        enable_sparse: bool = False,
+    ):
+        super().__init__()
+        self.vllm_model = vllm_model
+
+        self.num_layers = len(vllm_model.model.layers)
+
+        self.vllm_attn_layers = []
+        self.lmc_attn_layers = []
+        for i in range(self.num_layers):
+            vllm_attn = vllm_model.model.layers[i].self_attn.attn
+            self.vllm_attn_layers.append(vllm_attn)
+
+            self.lmc_attn_layers.append(
+                infer_attn_backend_from_vllm(vllm_attn, enable_sparse)
+            )
+
+        # NOTE(Jiayi): better not to pass the blender in init
+        # if we want to make this LMCModel more general.
+        self.blender = blender
+
+        # remove hard code
+        rotary_emb = vllm_model.model.layers[0].self_attn.rotary_emb
+        head_dim = rotary_emb.head_size
+        max_position_embeddings = rotary_emb.max_position_embeddings
+        rope_scaling = None
+        base = rotary_emb.base
+        is_neox_style = rotary_emb.is_neox_style
+        dtype = rotary_emb.dtype
+        self.fused_rotary_emb = get_fused_rope(
+            head_dim,
+            rotary_dim=head_dim,
+            max_position=max_position_embeddings,
+            base=base,
+            rope_scaling=rope_scaling,
+            is_neox_style=is_neox_style,
+            dtype=dtype,
+        )
+
+    @abstractmethod
+    def _process_qkv(self, q, k, v, layer):
+        """Process QKV tensors. Model-specific implementation."""
+        pass
+
+    @torch.compile
+    def compute_layer(
+        self,
+        input_ids: torch.Tensor,
+    ):
+        input_ids = input_ids.cuda()
+        hidden_states = self.vllm_model.get_input_embeddings(input_ids)
+        residual = None
+
+        attn_output = None
+
+        # TODO(Jiayi): Need to build `attn_metadata` more elegantly.
+        attn_metadata = self.lmc_attn_layers[0].init_attn_metadata(
+            input_ids=input_ids,
+        )
+
+        for idx, layer in enumerate(
+            self.vllm_model.model.layers[
+                self.vllm_model.model.start_layer : self.vllm_model.model.end_layer
+            ]
+        ):
+            # TODO(Jiayi) The last layer doesn't have to be computed
+            # hidden_states, residual = layer(positions, hidden_states, residual)
+
+            # Self Attention
+            if residual is None:
+                residual = hidden_states
+                hidden_states = layer.input_layernorm(hidden_states)
+            else:
+                hidden_states, residual = layer.input_layernorm(hidden_states, residual)
+            # hidden_states = self.self_attn(positions=positions,
+            #                            hidden_states=hidden_states)
+
+            qkv, _ = layer.self_attn.qkv_proj(hidden_states)
+            q, k, v = qkv.split(
+                [
+                    layer.self_attn.q_size,
+                    layer.self_attn.kv_size,
+                    layer.self_attn.kv_size,
+                ],
+                dim=-1,
+            )
+
+            # Model-specific QKV processing
+            q, k, v = self._process_qkv(q, k, v, layer)
+
+            q, k, v, residual, attn_output, attn_metadata = self.blender.process_qkv(
+                q, k, v, residual, idx, attn_output, attn_metadata
+            )
+
+            num_heads = self.vllm_attn_layers[idx].num_heads
+            num_kv_heads = self.vllm_attn_layers[idx].num_kv_heads
+            head_size = self.vllm_attn_layers[idx].head_size
+
+            q = q.view(-1, num_heads, head_size)
+            k = k.view(-1, num_kv_heads, head_size)
+            v = v.view(-1, num_kv_heads, head_size)
+            attn_output = attn_output.view(-1, num_heads, head_size)
+
+            attn_output = self.lmc_attn_layers[idx].forward_contiguous(
+                q, k, v, attn_output, attn_metadata
+            )
+
+            attn_output = attn_output.view(-1, num_heads * head_size)
+            k = k.view(-1, num_kv_heads * head_size)
+            v = v.view(-1, num_kv_heads * head_size)
+
+            hidden_states, _ = layer.self_attn.o_proj(attn_output)
+
+            # Fully Connected
+            hidden_states, residual = layer.post_attention_layernorm(
+                hidden_states, residual
+            )
+            hidden_states = layer.mlp(hidden_states)
+
+            yield

lmcache/v1/compute/models/llama.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.compute.models.base import LMCBaseModel
+
+
+class LMCLlamaModel(LMCBaseModel):
+    def _process_qkv(self, q, k, v, layer):
+        """Process QKV tensors for LLaMa model (no additional processing)."""
+        return q, k, v

lmcache/v1/compute/models/qwen3.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.compute.models.base import LMCBaseModel
+
+
+class LMCQwen3Model(LMCBaseModel):
+    def _process_qkv(self, q, k, v, layer):
+        """Process QKV tensors for Qwen3 model with q_norm and k_norm layers."""
+        # Qwen3 has q_norm and k_norm layers
+        q_by_head = q.view(
+            *q.shape[:-1],
+            q.shape[-1] // layer.self_attn.head_dim,
+            layer.self_attn.head_dim,
+        )
+        q_by_head = layer.self_attn.q_norm(q_by_head)
+        q = q_by_head.view(q.shape)
+        k_by_head = k.view(
+            *k.shape[:-1],
+            k.shape[-1] // layer.self_attn.head_dim,
+            layer.self_attn.head_dim,
+        )
+        k_by_head = layer.self_attn.k_norm(k_by_head)
+        k = k_by_head.view(k.shape)
+        return q, k, v

lmcache/v1/compute/models/utils.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Dict
+
+# Third Party
+from torch import nn
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+def infer_model_from_vllm(vllm_model, blender, enable_sparse: bool = False):
+    model_name = type(vllm_model).__name__
+    if model_name == "LlamaForCausalLM":
+        # First Party
+        from lmcache.v1.compute.models.llama import LMCLlamaModel
+
+        return LMCLlamaModel(vllm_model, blender, enable_sparse)
+    elif model_name == "Qwen2ForCausalLM":
+        # First Party
+        from lmcache.v1.compute.models.llama import LMCLlamaModel
+
+        return LMCLlamaModel(vllm_model, blender, enable_sparse)
+    elif model_name == "Qwen3ForCausalLM":
+        # First Party
+        from lmcache.v1.compute.models.qwen3 import LMCQwen3Model
+
+        return LMCQwen3Model(vllm_model, blender, enable_sparse)
+    else:
+        # TODO(Jiayi): Add support for more models
+        raise NotImplementedError(
+            f"Model type {model_name} is not supported in LMCache."
+        )
+
+
+class VLLMModelTracker:
+    _vllm_models: Dict[str, nn.Module] = {}
+
+    @classmethod
+    def register_model(
+        cls,
+        instance_id: str,
+        vllm_model: nn.Module,
+    ):
+        """
+        Register a vllm model by instance_id.
+        """
+        logger.info(f"Registering vllm model for {instance_id}")
+        if instance_id not in cls._vllm_models:
+            cls._vllm_models[instance_id] = vllm_model
+        else:
+            logger.warning(
+                f"vllm model for {instance_id} already registered, doing nothing."
+            )
+
+    @classmethod
+    def get_model(
+        cls,
+        instance_id: str,
+    ) -> nn.Module:
+        """
+        Get the vllm model by instance_id.
+        """
+        if instance_id not in cls._vllm_models:
+            raise ValueError(f"vllm model for {instance_id} not found.")
+        return cls._vllm_models[instance_id]

lmcache/v1/compute/positional_encoding.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Any, Callable, Dict, Optional
+
+# Third Party
+from vllm.model_executor.layers.rotary_embedding import get_rope as vllm_get_rope
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+import lmcache.c_ops as lmc_ops
+
+logger = init_logger(__name__)
+
+# TODO(Jiayi): Add and test more types of rope
+# (e.g., rope scaling, (non-)neox style, dtype, etc.)
+
+
+class BasicReverseRope:
+    def __init__(self, rope, rotary_dim, is_neox_style):
+        self.rope = rope
+        self.rotary_dim = rotary_dim
+        self.is_neox_style = is_neox_style
+
+    def do_shuffle(self, t):
+        original_shape = t.shape
+        t = t.reshape(t.shape[0], -1, self.rotary_dim)
+
+        if self.is_neox_style:
+            o1, o2 = torch.chunk(t, 2, dim=-1)
+        else:
+            o1 = t[..., ::2]
+            o2 = t[..., 1::2]
+
+        if self.is_neox_style:
+            return torch.cat((o2, o1), dim=-1).reshape(original_shape)
+        else:
+            return torch.stack((o2, o1), dim=-1).reshape(original_shape)
+
+    def reverse_encode(self, positions, q, k):
+        sq = self.do_shuffle(q)
+        sk = self.do_shuffle(k)
+        nq, nk = self.rope(positions, sq, sk)
+        fq = self.do_shuffle(nq)
+        fk = self.do_shuffle(nk)
+        return fq, fk
+
+    def __call__(self, positions, q, k):
+        return self.reverse_encode(positions, q, k)
+
+
+class FusedRope:
+    """
+    Directly use the fused kernel to ratate K cache from
+    the old positions to the new positions.
+    """
+
+    def __init__(self, rope, is_neox_style):
+        self.rope = rope
+        self.is_neox_style = is_neox_style
+        self.head_size = rope.head_size
+        self.cos_sin_cache = rope.cos_sin_cache
+
+    def fused_encode(self, old_positions, new_positions, k):
+        num_tokens = k.shape[0]
+        k = k.view(num_tokens, -1, self.head_size)
+        lmc_ops.rotary_embedding_k_fused(
+            old_positions,
+            new_positions,
+            k,
+            self.head_size,
+            self.cos_sin_cache.to(k.device),
+            self.is_neox_style,
+        )
+        k = k.view(num_tokens, -1)
+        return k
+
+    def __call__(self, old_positions, new_positions, k):
+        return self.fused_encode(old_positions, new_positions, k)
+
+
+def validate_rope_params(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: float,
+    is_neox_style: bool = True,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    dtype: Optional[torch.dtype] = None,
+    partial_rotary_factor: float = 1.0,
+):
+    if rotary_dim != head_size:
+        logger.error("Currently KV blending only support rotary_dim == head_size.")
+        return False
+
+    if rope_scaling is not None:
+        logger.error("Currently KV blending do not support rope scaling.")
+        return False
+
+    if partial_rotary_factor != 1.0:
+        logger.error(
+            "Currently KV blending do not support rotary factor other than 1.0."
+        )
+        return False
+
+    return True
+
+
+def validate_reverse_correctness(rope, reverse_rope, fused_rope, head_size) -> bool:
+    hidden_dim = head_size * 8
+    num_tokens = 10
+
+    dumb_q = torch.rand((num_tokens, hidden_dim), device="cuda", dtype=torch.bfloat16)
+    dumb_k = torch.rand((num_tokens, hidden_dim), device="cuda", dtype=torch.bfloat16)
+    positions = torch.arange(num_tokens, device="cuda")
+
+    q1 = dumb_q.clone()
+    k1 = dumb_k.clone()
+    q1, k1 = rope(positions, q1, k1)
+    q1, k1 = reverse_rope(positions, q1, k1)
+
+    max_q_error = (dumb_q - q1).abs().max()
+    max_k_error = (dumb_k - k1).abs().max()
+
+    logger.info(f"Max Q error: {max_q_error.item()}")
+    logger.info(f"Max K error: {max_k_error.item()}")
+
+    q_no_pos = dumb_q.clone()
+    k_no_pos = dumb_k.clone()
+    positions2 = torch.arange(100, 100 + num_tokens, device="cuda")
+    _, k_pos2 = rope(positions2, q_no_pos, k_no_pos)
+
+    k_no_pos = dumb_k.clone()
+    _, k_pos1 = rope(positions, q_no_pos, k_no_pos)
+    k_pos2_fused = fused_rope(positions, positions2, k_pos1)
+
+    max_k_error_fused = (k_pos2 - k_pos2_fused).abs().max()
+
+    logger.info(f"Max K error (fused): {max_k_error.item()}")
+
+    return max_q_error < 0.1 and max_k_error < 0.1 and max_k_error_fused < 0.1
+
+
+# Main interface
+def get_fused_rope(
+    head_size: int,
+    rotary_dim: int,
+    max_position: int,
+    base: float,
+    is_neox_style: bool = True,
+    rope_scaling: Optional[Dict[str, Any]] = None,
+    dtype: Optional[torch.dtype] = None,
+    partial_rotary_factor: float = 1.0,
+) -> Optional[Callable[..., Any]]:
+    # Validate the ROPE parameters
+    if not validate_rope_params(
+        head_size,
+        rotary_dim,
+        max_position,
+        base,
+        is_neox_style,
+        rope_scaling,
+        dtype,
+        partial_rotary_factor,
+    ):
+        logger.warning(
+            "The rope parameters is not supported! Cannot use cacheblend in this case"
+        )
+        return None
+
+    new_rope_params = {
+        "rope_theta": base,
+        "partial_rotary_factor": partial_rotary_factor,
+    }
+    if rope_scaling is not None:
+        new_rope_params.update(rope_scaling)
+        if "type" in rope_scaling:
+            new_rope_params["rope_type"] = rope_scaling["type"]
+
+    rope = vllm_get_rope(
+        head_size=head_size,
+        max_position=max_position,
+        is_neox_style=is_neox_style,
+        rope_parameters=new_rope_params,
+        dtype=dtype,
+        dual_chunk_attention_config=None,
+    )
+
+    reverse_rope = BasicReverseRope(rope, rotary_dim, is_neox_style)
+    fused_rope = FusedRope(rope, is_neox_style)
+
+    correct = validate_reverse_correctness(rope, reverse_rope, fused_rope, head_size)
+    if not correct:
+        logger.error(
+            "Fused/reverse rotary encoding is not correct! Will disable blending!"
+        )
+        return None
+
+    return fused_rope