CANN-昇腾NPU显存优化-大模型推理怎么把64GB用出128GB的感觉
Atlas 800I A2 的 64GB 显存，跑 Llama2-70B 不够，跑 Llama2-7B 又太浪费。这篇把昇腾NPU上的显存优化手段从易到难排列，逐个讲解。
显存都花在哪了
Llama2-7B 推理的显存分布：
权重：14GB（fp16） KV Cache：2GB × batch_size 运行时 buffer：1GB 编译缓存：0.5GB（可释放） 安全余量：1GB 可用 KV Cache 空间 = 64 - 14 - 1 - 1 = 48GB48GB 全给 KV Cache，单 token KV 约 512KB（Llama2-7B GQA），最多支持 48GB/512KB ≈ 96000 token 并发。batch=32 时 max_seq_len=3000。
优化 1：KV Cache Paged 分配
已在前面的文章讲过。显存利用率从 40-60% 提升到 90%。等效多出 30-50% 的可用显存。
优化 2：KV Cache fp8 压缩
KV Cache 默认用 fp16 存储。fp8 可以减半：
fromatbimportLLM model=LLM("model_id",device="npu:0",kv_cache_dtype="fp8")# KV Cache 用 fp8fp8 的 KV Cache 精度损失：Attention 输出的最大误差约 0.5-1%，对生成质量几乎无影响。因为 KV Cache 存的是 Attention 的 key/value，不是最终输出，中间的精度损失在后续计算中被稀释。
fp8 KV Cache 让可用 token 数翻倍：96000 → 192000。
优化 3：权重 W8A16 量化
权重从 fp16 压到 int8，显存减半：
model=LLM("model_id",device="npu:0",quantize="w8a16")Llama2-7B 权重从 14GB 降到 7GB。可用 KV Cache 空间从 48GB 增加到 55GB。
W8A16 的精度损失约 0.05-0.2%，比 KV Cache fp8 还小。
优化 4：权重分片加载
不常用的层权重放在 CPU 内存，用到时才搬到 NPU。类似 CPU 的虚拟内存：
model=LLM("model_id",device="npu:0",offload_ratio=0.3)# 30% 的层权重 offload 到 CPU30% 的权重在 CPU，每次用到这些层时通过 PCIe 搬到 NPU。PCIe 带宽约 32GB/s，搬一层约 0.5ms。
代价：每步推理多 10-15 次 PCIe 搬运，约 5-8ms 的额外延迟。decode 速度从 3200 tokens/s 降到 1800 tokens/s。
适用场景：显存绝对不够（比如 7B 模型在 32GB 显存的卡上），愿意用速度换容量。
优化 5：共享权重
Embedding 和 LM Head 的权重通常相同（tied weights）。确保模型加载时没有重复存储：
# 检查权重是否 tiedprint(model.config.tie_word_embeddings)# True 说明共享了# 如果没共享，手动共享model.lm_head.weight=model.model.embed_tokens.weightLlama2-7B 的 Embedding 约 0.5GB，共享后省 0.5GB。
优化组合效果
| 优化组合 | 权重 (GB) | KV Cache 可用 (GB) | 最大并发 token |
|---|---|---|---|
| 无优化 | 14 | 48 | 96K |
| + Paged KV | 14 | 48 (利用率 90%) | 86K |
| + KV fp8 | 14 | 48 (利用率 90%) | 172K |
| + W8A16 | 7 | 55 (利用率 90%) | 194K |
| + 共享权重 | 6.5 | 55.5 | 197K |
从 96K 到 197K token 并发，等效显存翻倍。
实际部署建议
- 7B 模型，64GB 显存：Paged KV + KV fp8 足够，不需要量化权重
- 13B 模型，64GB 显存：Paged KV + KV fp8 + W8A16
- 70B 模型，4×64GB 显存：Paged KV + KV fp8 + TP=4
- 70B 模型，2×64GB 显存：Paged KV + KV fp8 + W8A16 + TP=2
显存优化的核心思路：先确保 KV Cache 的分配效率（Paged），再压缩 KV Cache（fp8），最后压缩权重（W8A16）。每一步都有精度代价，但都在可接受范围内。仓库在这里：
https://atomgit.com/cann/ATB