🧠 vLLM System Prompt Cache —— 系统提示词缓存优化模块

⚡ 基于 vLLM 的系统提示词（System Prompt）缓存与复用机制，提升多轮对话场景下的推理性能。

---

📌 项目简介

本项目基于 vLLM 进行扩展，实现了 系统提示词永不释放 + 缓存复用机制，特别适用于以下场景：

• 多用户共享相同系统提示并且系统提示词长度特别大，预填充阶段耗时多
• 原生Vllm的block管理策略（LUR机制）在没有请求时释放cache，新来的请求又会重新计算
• 高频请求中重复使用相同 system prompt
• 需要节省显存并提高吞吐量的部署环境

通过该模块，你可以：

✅ 避免重复分配 KV Cache
✅ 显著减少内存占用与推理延迟
✅ 无缝兼容 OpenAI API 接口

---

🔍 功能亮点

---

🛠️ 安装与部署

1. 克隆项目

git clone https://github.com/finnchen11/VLLM_PromptCache.git
cd VLLM_PromptCache

2. 安装依赖

建议使用 Python 3.8+ 和 PyTorch 2.0+

pip install -e .

或者从原始 vLLM 安装后手动替换文件。

3. 启动服务

python -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --port 8000 --model your_model_name

---

📦 使用方式

1. 使用 curl 调用

curl -X POST http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "llama3",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "讲个笑话"}
        ]
    }'

2. 使用 OpenAI SDK

import openai

openai.api_base = "http://localhost:8000/v1"
openai.api_key = "EMPTY"

completion = openai.ChatCompletion.create(
    model="llama3",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "讲个笑话"},
    ]
)
print(completion.choices[0].message.content)

---

🧩 核心修改说明

以下是项目中主要修改或新增的文件及功能：

---

📊 性能对比

---

🔄 系统提示词缓存机制工作流程

用户请求
    ↓
[vLLM API Server] api_server.py
    → 提取消息中的 "role": "system" 内容
    → 调用 engine.add_system_prompt(prompt)

    ↓
[LLM Engine] llm_engine.py
    → 检查 system_prompt_cache 是否已缓存相同提示词
    → 若存在则直接复用，若不存在则创建新 Sequence 并标记为 is_system_prompt=True
    → 通过 CacheEngine 分配 KV Cache 块

    ↓
[Cache Engine] cache_engine.py
    → 在 allocate() 中为系统提示词分配缓存块
    → 所有后续调用 free(seq) 时，检查 seq.is_system_prompt
    → 若为 True，则跳过释放，实现“永不释放”

    ↓
[Worker] worker.py
    → 所有 CacheEngine 实例共享同一个 block_manager
    → 保证缓存块在不同 pipeline stage 间一致且可复用

    ↓
[缓存复用]
    → 当再次出现相同 system prompt 时
    → 直接复用已缓存的 Sequence 和其 KV Cache 块
    → 避免重复计算和内存分配，显著提升性能

    ↓
返回结果给用户 ✅

---