KV-Cache 全栈优化系统

本项目针对大语言模型（LLM）长上下文推理中 KV Cache 带来的显存瓶颈、访存开销和抢占恢复问题，提出了一套端到端的全栈优化方案。系统在 vLLM 框架基础上，集成了 AnDPro Value‑aware KV 剪枝、QServe A8W4 低比特线性算子与 KV8 FlashAttention，以及 前缀感知调度 与 CPU Offload 抢占恢复 三大模块，并在请求生命周期内实现三者的协同工作，从而在保持模型生成质量的前提下，显著提升单卡并发容量和系统吞吐。

主要特性

Value‑aware KV 剪枝
基于 AnDPro 评分思想，逐层评估 Token 对 Attention 输出的实际贡献，动态裁剪冗余 KV，减少显存占用，支持长上下文和高并发。
低比特线性与 KV‑cache 压缩
- 线性层采用 QServe A8W4 量化（4bit 权重 + 8bit 激活），配合 LayerNorm/SwiGLU 融合算子，降低模型存储与访存开销。
- 注意力层实现 8bit KV‑cache FlashAttention，将 K/V 按 token 动态量化为 INT8，压缩缓存容量，并在计算时按需恢复，提升固定显存下的有效吞吐。
前缀感知调度与 CPU Offload 抢占恢复
- 基于四分树和密度优先搜索，将前缀长度相近的请求组织为同质 batch，从源头降低抢占频率。
- 引入异步 CPU 备份机制，在请求被抢占前将 KV Block 预先搬运至主机内存，恢复时直接加载，避免重算开销；配合 Block 粒度的精确恢复，实现计算与传输重叠。
三模块协同集成
剪枝、低比特算子和调度/卸载模块共享同一份请求状态、Block Table 和缓存布局，通过显式反馈与两阶段提交保证资源一致性和安全回退，形成完整的推理优化链路。

系统架构概览

系统沿请求生命周期划分为四个层面：

层次	核心职责
控制面（Scheduler / Batch Policy）	KV 长度感知选批、Block 预算管理、抢占与 Commit 决策
模型执行面（GPUModelRunner）	Prefill/Decode 分流、逐层结果汇总、剪枝反馈生成
缓存算子面（Llama/Attention + QServe/KV8）	RoPE 后逐层评分、选中 K/V 直接写入紧凑布局、低比特量化与注意力计算
搬运恢复面（Offload/Recovery）	预防性 D2H 备份、两级命中（GPU/CPU）、原子恢复与回退

各层通过显式状态和反馈连接，任何关键条件不满足时均采用 fail‑closed 策略，回退至完整重算或标准路径，确保状态一致性。

已验证系统

项目	配置
操作系统	OpenEuler 24.03，Ubuntu 22.04
GPU	NVIDIA-L40S-48GB，NVIDIA GeForce RTX 4090
模型	Meta-Llama-3.1-8B-Instruct

部署指南

环境要求

操作系统：Linux x86_64
Python：3.12
CUDA：13.x（推荐）
PyTorch：2.11
GPU：NVIDIA Ampere/Ada 架构（测试于 L40S）
显存：建议 ≥ 24GB

1. 克隆仓库并创建虚拟环境

cd /path/to/acceleration_system
python3.12 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r modified_vllm/requirements/cuda.txt
pip install pytest ruff
pip install -r Experiment/longbench_e/requirements.txt

2. 编译 CUDA 扩展

项目包含两个核心扩展，需分别编译。

QServe W4A8 Linear 扩展

cd extensions/qserve_w4a8_ext
export CUDA_HOME=/opt/conda                 # 根据实际路径修改
export TORCH_CUDA_ARCH_LIST=8.9             # 适配 GPU 架构
export MAX_JOBS=2
export PYTHONPATH=/path/to/acceleration_system/modified_vllm
export PYTHONNOUSERSITE=1
python setup.py build_ext --inplace

KV8 FlashAttention 扩展

cd extensions/qserve_kv8_fa2_ext
cd vendor && ./prepare_vendor.sh && cd ..
python setup.py build_ext --inplace

3. 安装修改后的 vLLM

cd /path/to/acceleration_system/modified_vllm
pip install -e .

4. 配置运行环境

加载优化环境脚本：

source scripts/runtime_env.sh optimized

关键配置文件位于 config/ 目录：

**optimized.env**：控制 KV 剪枝、AnDPro 预算、物理回收、Inline 写入等开关。
**offload.json**：配置 CPU Offload 连接器（如 CPU 缓存大小、回收阈值等）。

示例 offload.json：

{
  "kv_connector": "SimpleCPUOffloadConnector",
  "kv_connector_module_path": "simple_kv_offload.simple_cpu_offload_connector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "cpu_bytes_to_use": 8589934592,
    "lazy_gpu_free_trigger": 0.35,
    "max_scan_reqs": 3
  }
}

5. 启动优化服务

使用统一启动脚本：

cd /path/to/acceleration_system
SERVER_HOST=0.0.0.0 PORT=8101 \
VLLM_ANDPRO_BUDGET_TOKENS=544 \
./scripts/start_optimized_server.sh

该脚本会自动加载环境、配置参数，并启动 vLLM OpenAI API 服务。可通过环境变量覆盖模型路径、上下文长度、批处理上限等：

MODEL=/path/to/QServe-model \
MAX_MODEL_LEN=32768 \
MAX_NUM_BATCHED_TOKENS=65536 \
MAX_NUM_SEQS=64 \
GPU_MEMORY_UTILIZATION=0.90 \
SERVER_HOST=0.0.0.0 PORT=8101 \
./scripts/start_optimized_server.sh

服务启动后可通过 OpenAI API 访问，例如：

curl http://127.0.0.1:8101/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"acceleration-system-optimized","prompt":"介绍 FlashAttention","max_tokens":128,"temperature":0}'

实验效果（摘要）

在 L40S 上的端到端测试（WikiText 继续生成，32 请求，多种输入/输出长度组合）显示：

端到端吞吐：优化版 vLLM 相对原始 vLLM 0.23 的输出 Token 吞吐加速 3.00× ~ 6.11×，平均 4.43×，且加速比随 Prompt 长度增加而扩大。
KV 剪枝容量收益：在关闭 Prefix Cache 的隔离测试中，AnDPro 剪枝使峰值 GPU KV Block 使用量降低 **52.18%**，Output TPS 提升 **14.92%**，无抢占稳定并发从 22 提升至 46。
低比特 KV‑cache：在 258K 上下文的相同显存预算下，KV8 FlashAttention 吞吐为原始 FA2 的 1.64×。
抢占恢复：CPU Offload 将 Block 丢失率降低 95% 以上，端到端延迟优化最高达 **9.8%**；前缀感知调度在高抢占场景下减少抢占次数 **33%~49%**。

详细实验数据和配置请参见项目技术报告。

引用

本项目基于以下关键技术：

AnDPro: NeurIPS 2025
QServe: MLSys 2025
FlashAttention‑2: arXiv:2307.08691
vLLM (PagedAttention): SOSP 2023
LongBench / RULER 基准

许可证

本代码仅供研究参考，具体许可请参见项目内 LICENSE 文件。

项目团队：欢迎通过 Issue 或邮件联系。