推理与部署

长文本生成时显存占用怎么优化?

91学AI·2026/7/26·12 阅读

考察点

长文本(几十 k 到上百 k token)场景下显存的第一矛盾是 KV Cache 随序列长度线性膨胀。面试官想看你能不能把手段分层讲清:模型架构层(GQA、滑窗)、存储层(量化、分页)、调度层(offload、chunked prefill)、以及输入侧(压缩、截断策略)。高分答案还会点破"长输入"和"长输出"的显存特征不同。追问常往具体数字估算和某个机制的取舍上走。

参考答案

先算账:显存到底花在哪

长文本推理的显存 = 模型权重 + KV Cache + 激活临时占用。权重是定值(7B fp16 约 14GB),激活临时占用很小,变量全在 KV Cache2 × 层数 × kv_heads × head_dim × 字节数 × 序列长度 × batch。以 LLaMA-2-13B(40 层、40 KV heads、head_dim 128)为例,单 token KV 约 1.6MB,一条 32k 的序列就是约 51GB——一张 A100 80G 都装不下几条。所以题目等价于:怎么治理 KV Cache。

注意区分两种"长":长输入的痛点是 prefill 阶段的算力和一次性 KV 写入;长输出的痛点是 decode 期间 KV 持续增长、长期占着槽位。两者手段有重叠也有侧重。

模型架构层:选型时就决定的事

  • GQA/MQA:KV head 数砍 4-8 倍,Cache 同比例缩。LLaMA-2/3、Qwen 全系都是 GQA,选模型时这已经是默认收益。
  • 滑窗注意力(SWA):Mistral 系的做法,每层只关注最近 W 个 token,KV Cache 从 O(N) 变 O(W),超长文本显存直接封顶;代价是窗口外的精确记忆丢失,靠信息逐层传播弥补。
  • MLA:DeepSeek 的多头潜注意力,把 KV 压成低秩潜向量缓存,Cache 比 GQA 再小一个量级,是长上下文模型的新方向。

存储层:不动模型能做的事

  • KV Cache 量化:int8 减半、int4 再减半,vLLM/TRT-LLM 均支持,是长上下文最常用的一手。注意 K 对量化更敏感,优先验证业务指标。
  • PagedAttention:按需按 block 分配,消除预留和碎片浪费,长文本场景长度方差大,收益比短文本更明显。
  • Prefix caching:长文档问答里同一份文档被反复问,文档部分的 KV 跨请求复用,既省显存又省 prefill。

调度层:以时间换空间

  • CPU offload:把不活跃请求的 KV 换出到 CPU 内存甚至 NVMe,需要时再换回(vLLM 的 swapping、Mooncake 的分布式 KV 池)。带宽换容量,适合多轮对话里"挂着不说话"的长会话。
  • Chunked prefill:超长 prompt 切块分批算,prefill 峰值显存(激活 + 临时 buffer)被摊平,避免长输入一把 OOM。
  • 限制上下文策略:业务上设 max_context,超限走摘要/截断,别把治理压力全甩给推理层。

输入侧:最便宜的手段往往是少喂点

RAG 场景的很多"长文本"是检索召回过多、prompt 里塞满低相关 chunk 造成的虚胖。重排后只留 top-k、对历史对话做滚动摘要、用 LLMLingua 这类 prompt 压缩技术删掉冗余 token,输入砍一半,KV Cache 和 prefill 成本都砍一半。先问"这些 token 真的都要进模型吗",再谈怎么省显存。

组合策略

生产上一个长文档问答服务的典型组合:GQA 模型 + KV int8 + PagedAttention + prefix caching + chunked prefill + RAG 召回控制在 4k 以内。这套下来 32k 上下文能在单卡 80G 上跑出不错的并发。真到 100k+,就得考虑 MLA 架构模型或 PD 分离 + 分布式 KV 池的系统方案了。

可能的追问

  • 长文本训练侧的显存问题和推理侧一样吗? 不一样。训练大头是激活值(反向要重算或保存,靠 FlashAttention 重计算、梯度检查点、序列并行来压);推理大头是 KV Cache。面试别混着答。
  • KV 量化到 int4 掉点怎么评估? 长文本任务对检索式记忆敏感,重点测"大海捞针"(needle in a haystack)准确率和业务端到端指标,困惑度参考价值有限。
  • 滑窗注意力为什么不会完全丢失远处信息? 信息逐层传播:第 l 层的每个位置聚合窗口 W,多层叠加后感受野是 层数 × W,远处信息以"接力"方式间接流入,但精确的逐 token 引用能力确实受损。

评论 (0)

暂无评论,快来抢沙发吧!

91学AI

© 2026 91学AI · 按岗位学 AI 与大数据. All rights reserved.