LMCache - Redis for LLMs 无限高速 KV 缓存系统

docker pull lmcache/vllm-openai:latest-nightly

pip install -r requirements.txtpip install -e .

pip install -r requirements/lint.txtpre-commit install

LMCACHE_CONFIG_FILE=config.yaml python -m lmcache_vllm.vllm.entrypoints.openai.api_server --model mistralai/Mistral-7B-Instruct-v0.2

python3 multi-round-qa.py \    --num-users 10 \    --num-rounds 5 \    --qps 0.5 \    --shared-system-prompt 1000 \    --user-history-prompt 2000 \    --answer-len 100 \    --model mistralai/Mistral-7B-Instruct-v0.2 \    --base-url http://localhost:8000/v1

import openai
response = openai.Completion.create(    model="mistralai/Mistral-7B-Instruct-v0.2",    prompt="The capital of France is\nA. Berlin\nB. Madrid\nC. Paris\nD. Rome\nAnswer:",    temperature=0,    max_tokens=3)

def init_lmcache_engine(    model_config: ModelConfig,    tp_size: int,    rank: int,    world_size: int,) -> Optional[LMCacheEngine]:    config = lmcache_get_config()    kv_dtype = get_kv_cache_torch_dtype(model_config.dtype)        # 构造KV形状(用于内存池)    num_layer = model_config.num_hidden_layers    chunk_size = config.chunk_size    num_kv_head = model_config.get_num_kv_heads(tp_size)    head_dim = model_config.head_dim    kv_shape = (num_layer, 2, chunk_size, num_kv_head, head_dim)
    metadata = LMCacheEngineMetadata(        model_config.model_path,        world_size,        rank,        "sgl",        kv_dtype,        kv_shape,    )        engine = LMCacheEngineBuilder.get_or_create(        ENGINE_NAME, config, metadata, gpu_connector    )    return engine

def retrieve(    self,    token_ids: torch.Tensor,    mask: Optional[torch.Tensor] = None,    kvcaches: Optional[List[torch.Tensor]] = None,    slot_mapping: Optional[torch.Tensor] = None,    offset: int = 0,) -> torch.Tensor:    # 处理输入验证    assert isinstance(token_ids, torch.Tensor)    if mask is None:        mask = torch.ones_like(token_ids, dtype=torch.bool)        # 从缓存中检索KV    ret_token_mask = torch.zeros_like(token_ids, dtype=torch.bool)    for start, end, key in self.token_db.process_tokens(token_ids, mask):        cached_data = self.engine.get(key)        if cached_data is not None:            # 将缓存数据加载到GPU缓冲区            self.gpu_connector.from_gpu(cached_data, start, end, kvcaches, slot_mapping)            ret_token_mask[start:end] = True        return ret_token_mask

def process_qkv(    self,    q: torch.Tensor,    k: torch.Tensor,    v: torch.Tensor,    residual: torch.Tensor,    layer_id: int,    attn_output: Optional[torch.Tensor] = None,) -> torch.Tensor:    # 检查是否需要在此层执行混合计算    if layer_id not in self.common_metadata.check_layers:        return attn_output        # 从缓存中检索相关KV    retrieved_k, retrieved_v = self.cache_engine.retrieve_kv(q)        # 执行混合注意力计算    blended_output = self.blender.blend(        q, k, v,         retrieved_k, retrieved_v,        self.common_metadata.recomp_ratios[0]    )        # 残差连接    if residual is not None:        blended_output += residual            return blended_output