编写一个 KV Cache 显存估算工具,针对给定的业务场景进行容量规划,并设计最优的 KV Cache 分层存储方案。
下次课前 (一周)
编写一个 Python 脚本 kv_cache_estimator.py,接收模型配置和业务参数,输出 KV Cache 显存估算:
# kv_cache_estimator.py
def estimate_kv_cache(model_config: dict, seq_len: int, batch_size: int,
dtype: str = "fp16") -> dict:
"""
model_config: {
"n_layers": int,
"n_kv_heads": int,
"d_head": int,
"model_params_gb": float # 模型权重大小
}
返回: {
"kv_cache_gb": float,
"model_gb": float,
"total_gb": float,
"kv_ratio": float # KV Cache 占总显存的比例
}
"""
pass
要求:
使用该工具完成 Qwen2.5-7B/14B/72B 和 DeepSeek-V3 的显存估算表。
假设你需要在以下三种场景中部署 LLM 推理服务:
| 场景 | 模型 | 平均 Prompt | 平均输出 | 并发 |
|---|---|---|---|---|
| A: 客服系统 | Qwen2.5-7B | 1500 tokens | 200 tokens | 50 |
| B: 代码助手 | Qwen2.5-14B | 8000 tokens | 500 tokens | 20 |
| C: 长文档分析 | Qwen2.5-72B | 30000 tokens | 1000 tokens | 5 |
对每个场景:
编写脚本模拟 prefix caching 的命中率:
# 假设场景:
# - 系统 Prompt 固定 2000 tokens
# - 用户问题随机 50-200 tokens (每次不同)
# - 100 个并发用户
# 计算:
# 1. 无 Prefix Caching: 每个请求的 Prefill 计算量
# 2. 有 Prefix Caching: 命中系统 Prompt 时的 Prefill 计算量
# 3. 节省的计算量比例
在实际环境中部署 LMCache + vLLM,完成:
# 1. 安装并配置 LMCache
pip install lmcache
# 2. 配置多层存储
export LMCACHE_CONFIG='{
"l1": {"backend": "cpu", "max_size_gb": 32},
"l2": {"backend": "disk", "path": "/tmp/lmcache", "max_size_gb": 128}
}'
# 3. 启动 vLLM + LMCache 服务
# 4. 运行多轮对话测试,对比:
# - 无缓存
# - L1 only (CPU)
# - L1 + L2 (CPU + Disk)
# 的 TTFT and throughput
从 AI-fundamentals 中选择一个 KV Cache 系统源码分析:
09_inference_system/kv_cache/02_systems/lmcache/09_inference_system/kv_cache/02_systems/kvbm/分析存储后端的核心实现逻辑 (如 LocalCPUBackend 的缓存管理)。
kv_cache_estimator.py 及运行结果| 维度 | 权重 | 要求 |
|---|---|---|
| 任务 1-3 完成度 | 60% | 估算器正确、场景分析完整 |
| 分析质量 | 25% | 推荐方案有依据、有数据支撑 |
| 进阶任务 | 15% | 完成至少一项进阶任务 |