CANN-昇腾NPU显存优化-大模型推理怎么把64GB用出128GB的感觉

Atlas 800I A2 的 64GB 显存跑 Llama2-70B 不够跑 Llama2-7B 又太浪费。这篇把昇腾NPU上的显存优化手段从易到难排列逐个讲解。显存都花在哪了Llama2-7B 推理的显存分布权重14GBfp16 KV Cache2GB × batch_size 运行时 buffer1GB 编译缓存0.5GB可释放 安全余量1GB 可用 KV Cache 空间 64 - 14 - 1 - 1 48GB48GB 全给 KV Cache单 token KV 约 512KBLlama2-7B GQA最多支持 48GB/512KB ≈ 96000 token 并发。batch32 时 max_seq_len3000。优化 1KV Cache Paged 分配已在前面的文章讲过。显存利用率从 40-60% 提升到 90%。等效多出 30-50% 的可用显存。优化 2KV Cache fp8 压缩KV Cache 默认用 fp16 存储。fp8 可以减半fromatbimportLLM modelLLM(model_id,devicenpu:0,kv_cache_dtypefp8)# KV Cache 用 fp8fp8 的 KV Cache 精度损失Attention 输出的最大误差约 0.5-1%对生成质量几乎无影响。因为 KV Cache 存的是 Attention 的 key/value不是最终输出中间的精度损失在后续计算中被稀释。fp8 KV Cache 让可用 token 数翻倍96000 → 192000。优化 3权重 W8A16 量化权重从 fp16 压到 int8显存减半modelLLM(model_id,devicenpu:0,quantizew8a16)Llama2-7B 权重从 14GB 降到 7GB。可用 KV Cache 空间从 48GB 增加到 55GB。W8A16 的精度损失约 0.05-0.2%比 KV Cache fp8 还小。优化 4权重分片加载不常用的层权重放在 CPU 内存用到时才搬到 NPU。类似 CPU 的虚拟内存modelLLM(model_id,devicenpu:0,offload_ratio0.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.weightmodel.model.embed_tokens.weightLlama2-7B 的 Embedding 约 0.5GB共享后省 0.5GB。优化组合效果优化组合权重 (GB)KV Cache 可用 (GB)最大并发 token无优化144896K Paged KV1448 (利用率 90%)86K KV fp81448 (利用率 90%)172K W8A16755 (利用率 90%)194K 共享权重6.555.5197K从 96K 到 197K token 并发等效显存翻倍。实际部署建议7B 模型64GB 显存Paged KV KV fp8 足够不需要量化权重13B 模型64GB 显存Paged KV KV fp8 W8A1670B 模型4×64GB 显存Paged KV KV fp8 TP470B 模型2×64GB 显存Paged KV KV fp8 W8A16 TP2显存优化的核心思路先确保 KV Cache 的分配效率Paged再压缩 KV Cachefp8最后压缩权重W8A16。每一步都有精度代价但都在可接受范围内。仓库在这里https://atomgit.com/cann/ATB