上周帮团队排查一个部署项目凌晨两点对着torch profiler一顿看发现一个7B模型推理占满A100可batch size明明没调大显存却比模型权重多出好几个G。查来查去问题出在KV Cache的容量规划上——sequence长度一旦拉长KV Cache的增长比权重本身还猛。这个事情让我感触挺深很多人对大模型训练和推理的理解还停留在“显存大头是模型权重”的阶段但实际跑起来训练有激活值推理有KV Cache每一项都可能反客为主。所以这一讲就把话题完全聚焦在显存账本上。训练阶段权重、梯度、优化器状态、激活值四个户头各吃得下多少推理阶段权重固定之后KV Cache和激活值又是怎么抢占卡上空间的。同时把FLOPs手算一遍——不是看别人结论而是自己拿笔在白板上推从attention到FFN每个矩阵乘法的计算量都验一遍。最后给一套可以照着做的KV Cache容量规划方法和避坑清单。这篇内容适合正在做模型训练调优、推理服务部署、以及刚入行想做容量预算的工程师。看完之后你至少能把“我这模型到底该买几张卡”“服务上线前显存够不够”这两个问题从拍脑袋变成拉表计算。1. 训练显存全景账本钱到底都花在哪几个户头上训练阶段显存占用不像推理那么直观。很多人一开始以为模型多大显存就多大比如7B参数用bf16就是14GB。但真跑起来你会看到一张80GB的卡有时还不够训一个7B模型甚至加上各种优化手段才勉强塞下。原因在于训练时显存里有四笔开销同时在跑模型权重、梯度、优化器状态、激活值。每一笔都不是小数目而且优化器状态常常比权重本身还大。1.1 模型状态三件套权重、梯度、优化器状态的真面目先说“模型状态”这个概念我常用它统称权重、梯度和优化器相关的显存开销。用混合精度训练比如bf16混合精度跑AdamW优化器时经典经验公式是单参数大约占16字节。这16个字节怎么来的我拆给你看。首先是bf16权重副本2字节反向传播时算出来的bf16梯度2字节这两份是计算过程中真正参与前向和反向的。但AdamW优化器还需要维护一份fp32的主权重用于数值稳定性4字节以及fp32的一阶动量m和二阶动量v各4字节。把这些加起来权重2 梯度2 主权重4 动量m 4 动量v 4正好16字节。你可能会问为什么不直接用fp32训练省得复制一份主权重理论上可以但显存翻倍不说很多现代GPU对fp16/bf16的矩阵计算有专门优化速度更快。所以工程上普遍保留一份fp32主权重作为“慢速账本”每步更新时用它算完再转回bf16参与前向误差可控显存多花一点但换来稳定收敛和更高吞吐。如果不用AdamW改用SGDMomentum情况会好很多bf16权重2字节bf16梯度2字节主权重4字节动量4字节合计12字节。但如果用Adagrad这类需要累积梯度平方的状态单参数开销比AdamW还高。所以很多大模型训练框架默认推AdamW不是因为它省显存而是因为收敛稳定你得为这种稳定付显存成本。1.2 激活值训练时最容易被忽略的隐形吞噬者权重、梯度、优化器状态这三笔还算好算因为有明确的参数量可以乘。激活值就不一样了它跟batch size、序列长度、隐藏层维度、层数、视野大小全部挂钩呈乘法关系爆发。举个具体例子感受一下。假设一个7B规模模型隐藏层维度4096层数32序列长度4096batch size设为1。每个Transformer层前向要保存Q、K、V、注意力矩阵输出、FFN中间激活等。如果每个激活都保留一个token在单层的激活存储就要几十KB到几百KB不等再乘上序列长度、层数、batch size轻松突破几十GB。这也是为什么很多人训长文本时固定显存没变但OOM频发多半就是激活值在悄悄暴涨。缓解办法最常见的是激活检查点activation checkpointing也叫梯度检查点。原理简单说前向时不全存激活只存每层输入和少数关键张量反向时重新算一遍前向来补齐需要的激活。显然这是用计算换显存一次重计算大约是原计算量的1.3到1.5倍但显存可以从几十GB级别直接压到个位数GB。我在实际项目里激活检查点几乎是长序列训练必开的选项不开它序列长度稍微拉长就寸步难行。1.3 四种分布式并行策略怎么影响显存账本训练不是单卡就能搞定的事分布式并行策略直接影响上面四笔开销的分配。数据并行Data Parallel是把模型复制到每张卡上各跑各的batch梯度同步。它的特点是权重、优化器状态每张卡都有一份显存占用随卡数线性复制。显存紧张的场景下数据并行是最浪费的。张量并行Tensor Parallel把单个Transformer层的矩阵按列或按行切分到多卡上每张卡只持有部分权重和对应梯度适合单层特别大、单卡装不下的情况。但它在每层前向和反向都会有通信开销。流水线并行Pipeline Parallel把不同层切到不同卡上每张卡只持有自己那几层的权重和梯度这样单卡显存压力小很多但存在流水线气泡问题。Zero优化则把优化器状态、梯度和权重做切分典型做法是每张卡只负责更新一部分参数的状态其他参数状态通过通信临时聚合。它的好处是显存随卡数近似线性下降通信成本也在可控范围。实际做显存规划时要结合这几种并行方式搞清楚“每卡账本”。我的经验是先把每卡理论显存算一遍再按20%缓冲留余量比用Nz放缩法靠谱。2. FLOPs白板手算从attention到FFN一次算给每个人看显存搞清楚了接下来该弄清楚计算量。FLOPs算的不是显存但显存规划经常需要结合计算量来定并行方案和batch size而且面试、方案评审、论文复现时也经常被问“这个模型推理一次大概多少计算量”。如果你只会背结论“forward 2倍参数量训练是6倍”那遇到具体结构还是会打磕绊。我建议你像我一样在白板上从矩阵乘法开始推一遍。2.1 Attention部分QK^T和AV的计算量推导一个标准的自注意力模块输入是一个形状为[batch, seq_len, hidden]的序列。第一大步是生成Q、K、V通过三个线性变换完成。每个线性变换本质上是矩阵乘把hidden维映射到hidden维那么单个变换的计算量就是2 × batch × seq_len × hidden × hidden。为什么有个2因为一次矩阵乘法[M, N]乘以[N, P]需要M×N×P次乘法、M×(N-1)×P次加法近似记作2×M×N×P。以这个标准计算Q、K、V三个变换加起来就是6 × batch × seq_len × hidden × hidden接着是注意力分数Q和K做点积。在多头设定下每个头的维度是hidden/h所有头并行计算。Q乘K的转置计算量是2 × batch × seq_len × seq_len × hidden。注意这里跟heads数量没有直接关系因为分头后hidden总维度保持不变最终乘出来还是2×batch×seq_len×seq_len×hidden。然后是注意力分数乘V同样形状是[batch, heads, seq_len, seq_len]乘以[batch, heads, seq_len, head_dim]合并之后也是2 × batch × seq_len × seq_len × hidden。所以attention内部跟序列长度平方相关的项合计是4 × batch × seq_len^2 × hidden。这就是为什么长序列的注意力计算量会爆炸seq_len一拉长平方项直接支配整个FLOPs预算。2.2 FFN部分为什么12倍hidden值得重视Transformer层里attention之后还有一个FFN一般结构是先把hidden映射到4倍hidden再映射回来。两个线性变换计算量分别是2 × batch × seq_len × hidden × 4hidden以及 2 × batch × seq_len × 4hidden × hidden两项相加就是 16 × batch × seq_len × hidden × hidden。对比attention部分的QKV投影6倍hidden^2FFN的计算量明显更大。再算上attention内部与hidden^2相关的部分实际上single layer的FLOPs大头是FFN约占三分之二以上。合并一下单个Transformer层的FLOPs近似为24 × batch × seq_len × hidden × hidden 4 × batch × seq_len^2 × hidden如果层数为L总FLOPs还要乘以L。这个公式非常实用你只需要知道模型的hidden size、层数和序列长度就能在一个数量级内估算出前向的计算量。2.3 推理与训练的FLOPs对比2倍与6倍的由来大模型场景下有个快速估算法则每token每参数的FLOPs约等于2乘以参数量。为什么因为一次矩阵乘法的FLOPs约等于2倍的参数量×token数。把所有参数都作为矩阵乘法的一部分计算单个token前向推理的FLOPs就约等于2×N×TN是参数量T是token数。比如7B模型推理1个token计算量大约14GFLOPs。这个估算在量化前是相当准的误差通常在10%以内。训练阶段的FLOPs要更高。反向传播需要计算梯度梯度的计算在数学上相当于再跑一遍前向的矩阵乘法而且还要额外计算权重梯度。业界公认的经验值训练一个token的总FLOPs大约是前向的3倍也就是6×N×T。所以你在对比模型训练和推理成本时可以这样拉账如果同样的tokens和模型训练的总计算量是推理的3倍左右。这种算账方式不需要任何profiler口算个数量级就能判断一个训练任务大约需要多少卡时。3. KV Cache是什么为什么推理时必须规划它推理阶段跟训练阶段不一样模型权重固定后基本不再变化显存大头就变成两类固定的模型权重以及随序列长度动态变化的KV Cache。很多人刚接触推理服务时只按权重定卡结果上线一测并发发现KV Cache把显存吃了大半。3.1 从自回归解码过程理解KV Cache的产生机制大模型推理是自回归的也就是每次只生成一个token把新token拼到已有序列后面再作为输入继续生成下一个token。如果每次生成都重新对所有token算一遍Q乘K代价极高。好在这里有一个关键观察在生成第t1个token时前面t个token的K和V在注意力计算中是固定的其实可以直接缓存下来下一轮直接取出复用。这就是KV Cache的由来K和V在生成过程中不断累积每一层都缓存一份。但Q不用缓存因为Q只跟当前token的当前位置有关。PagedAttention这类技术也是围绕KV Cache的显存管理做文章把KV Cache分块存储按需分页。3.2 KV Cache容量规划公式一次拉通所有变量KV Cache的显存占用可以写成下面这个公式KV Cache显存 2 × layers × seq_len × head_dim × kv_heads × batch_size × bytes_per_element这里2代表K和V各一份layers代表层数每一层都有独立的K和V缓存seq_len是目前上下文长度head_dim是每个注意力头的维度kv_heads在每个头单独存KV时等于hidden/head_dim在GQA或MQA架构下则是更少的组batch_size是同时处理的序列数bytes_per_element是精度字节数fp16/bf16是2fp8是1int4是0.5左右。以7B模型为例32层hidden4096GQA设为8个KV头head_dim128bf16精度seq_len4096batch12 × 32 × 4096 × 128 × 8 × 1 × 2 536,870,912字节即0.5GB如果seq_len拉到32K这个数字直接变成4GB。如果batch涨到32那KV Cache就是128GB直接爆掉一张A100。3.3 GQA、MQA如何在KV Cache规划上省钱GQAGrouped Query Attention和MQAMulti-Query Attention的核心目标都很一致减少KV头从而减少KV Cache的显存和带宽压力。在标准多头注意力中Q、K、V各有相同数量的头。GQA的做法是让Q仍有多个头但K和V共享成更少的头比如8个KV头对应32个Q头也就是4个Q头分成一组共用一个K和V。这样KV Cache的占用缩减为原来的1/4。MQA更极端所有Q头共用同一组K和VKV Cache直接变成原来的1/32。市面上很多模型的架构选型就考虑了这点。某公司在做长序列部署时把原版多头注意力改成GQA后同一条推理链路里的KV Cache显存直接下降四倍。这个改动对于推理服务来说几乎是纯赚训练效果损失也极小。4. 推理显存全景拿一个实际部署例子做完整拉账训练讲了KV Cache也讲了最后要把推理侧完整的显存账本放在一起用一个实际例子做演算。这样等下你拿自己的模型套公式时心里就有底了。4.1 预算一个7B模型的推理显存假设我有一个7B模型bf16精度模型权重占14GB。再假设32层hidden4096kv_heads8head_dim128上下文长度2048batch8。KV Cache按公式算2 × 32 × 2048 × 128 × 8 × 8 × 2 ≈ 4.3GB激活值部分虽然推理不像训练那样保存大量中间梯度但每个token在每层也有激活缓存估算很小约100MB到几百MB。加上框架本身的一些CUDA context和临时buffer约占1GB左右。总显存预算是14 4.3 0.5 1 ≈ 19.8GB。也就是说一张24GB的卡刚好能跑但只剩不到4GB余量。如果上下文拉长到8192kv_cache变成17GB整卡直接爆掉。这个例子很能说明问题在长上下文场景下KV Cache的增长跟序列长度线性相关权重是固定成本KV Cache才是动态成本很多人忽略这个动态成本所以总是在上线前一刻发现卡不够用。4.2 batch size与KV Cache的线性关系以及如何取舍KV Cache另一个容易被低估的特征是它跟batch_size成正比。推理服务为了提高吞吐往往一个batch拼进多个请求同时处理。batch从1涨到16KV Cache直接乘16。训练里你可以靠梯度累积来减小batch size对显存的影响但推理服务中batch是吞吐的基石只能硬扛。所以实践中要在延迟和吞吐之间找平衡。一般做法是先根据延迟要求确定单请求的最大并发数再套KV Cache公式估算总显存。比如对响应时间很敏感的在线服务batch不会拉太高离线批量打分场景batch可以尽量拉高只要显存不爆。4.3 常用显存估算速查表我把常用的容量估算浓缩成一张参考表方便平时拿来快速估算。表里的值全是近似够用就好真到规划阶段再精确计算。项目估算方式备注模型权重显存bf16参数量 × 2字节7B约14GB训练单参数模型状态AdamW混合精度参数量 × 16字节7B约112GB激活值训练开销无法简单乘算需按层与序列估算开启activation checkpoint可大幅降低推理KV Cache显存2 × layers × seq_len × head_dim × kv_heads × batch × 2动态增长最快的一项推理前向FLOPs/token2 × 参数量7B约14GFLOPs训练FLOPs/token6 × 参数量对应一个token的forwardbackward这张表你背下来日常方案评审里基本能应对90%的显存估算问题。5. KV Cache容量规划实战步骤、技巧、坑位全记录前面算是把原理和公式都理清了接下来讲我在实际部署项目里的一些操作方法和踩坑经验这部分的含金量不在推导之下。5.1 KV Cache容量规划的实操步骤第一步确定上下文长度上限。这个决定KV Cache能涨到多大。注意一个容易踩的坑很多项目声称支持8K上下文但只在短输入时测试过一旦用户真的输入8K的文本KV Cache瞬间涨到最大如果没预留空间系统就开始频繁触发缓存淘汰或直接OOM。第二步确定batch size目标。在线推理batch_size取决于延迟上限和请求到达率离线批处理batch_size只受显存限制。我一般先把batch设为目标值套公式算总显存再看有没有余量给激活值和临时buffer。第三步确定KV Cache策略。如果是自研推理框架优先考虑PagedAttention式的分块管理如果是用成熟框架确认一下是否支持KV Cache的量化。KV Cache量化成8bit或4bit在长序列场景里非常有效显存直接减半或减四分之一实测精度损失通常可以忽略。第四步设置安全线和告警。显存预留不要低于总显存的20%。一个简单的经验把公式算出的数值乘以1.2再作为规划值。如果规划值接近卡的上限就要考虑换更大的卡或者减少batch或者缩短最大上下文长度。5.2 我踩过的几个坑第一个坑是只看权重不看KV Cache。曾经部署一个13B模型权重占26GB一张40GB的卡看起来绰绰有余。结果跑并发时发现模型服务总显存冲到36GB以上一查才知道是KV Cache没算进去batch4时KV Cache已经吃掉10GB。第二个坑是不区分GQA和MHA。同一个模型如果用了GQAKV Cache的显存跟MHA版本差好几倍。做方案对比时一定要确认对方的attention架构否则估算会严重失真。第三个坑是忽略了上下文长度对新版本模型的加成。支持长上下文的模型在推理时KV Cache的增长率跟长度是线性关系看起来12K和32K只是2.7倍但加上batch就是成倍的差距。上线长文本功能前建议专门做一次最坏情况压测——最大长度、最大并发、最大batch一起上看显存究竟能不能扛住。5.3 一些额外的优化经验如果KV Cache实在压不下来我会按这个优先级来解决问题。先在架构上选或改造成GQA/MQA这是最彻底的方案直接砍头数然后做KV Cache量化8bit起步cuda加速下部署影响不大再考虑窗口注意力或滑动注意力只保留最近一段长度的KV Cache虽然结论是旧信息被丢弃但在很多生成任务中效果仍然够用最后考虑把KV Cache换到CPU和GPU混合方案但这需要框架支持且带宽可能成为新瓶颈。训练侧也有反向影响KV Cache的手段。在训练时就采用较短的上下文窗口或在训练数据里限制长文本比例能让模型对远距离依赖不那么敏感。但这个方法会影响效果不建议作为首选。6. FLOPs手算与显存算账怎么配合使用FLOPs算计算量显存算存储量看起来是两码事但在规划集群和推理资源时经常必须放在一起考虑。6.1 用FLOPs估算训练时长和集群规模训练时长估算公式可以这样写训练总FLOPs ÷ 集群总算力 ÷ 实际利用率。实际利用率也就是MFU业界通常只能做到40%到60%再高就要花很大精力调通信和计算重叠。举个例子假设要处理100B tokens模型7B训练FLOPs按6×N×T算也就是6 × 7B × 100B 4.2e21FLOPs。如果用8张A100每张理论算力约312TFLOPs峰值总算力2.5e15FLOPs。如果MFU50%实际有效算力约1.25e15FLOPs训练时间是4.2e21÷1.25e15 ≈ 3.36e6秒约39天。这个估算能帮你在开会讨论方案时快速判断“这个规模是否合理”。6.2 推理场景用FLOPs估算单请求消耗推理侧单token的FLOPs约2×N再乘以输出token数就是一次请求的总计算量。比如7B模型回答200个token单请求约2×7B×200 2.8TFLOPs。如果单卡峰值算力312TFLOPs理论延迟大概是9ms量级但实际还要算上显存带宽、attention计算、框架开销等给个10倍左右的折减就是现实延迟。这样估算出来的值虽然粗糙但比拍脑袋强很多。6.3 显存和FLOPs的交叉决策案例有一次我给某服务做容量规划目标QPS是50平均输出长度500 token模型7B。先算KV Cachebatch要50并发seq_len按500假设输入200那就是约700代入公式全套跑一遍KV Cache约12GB模型权重14GB加起来超24GB单卡A10G不够转用A100 40GB稳了。再算FLOPs看延迟单token计算量14GFLOPs500token总FLOPs约7TFLOPsA100算力下延迟理论极低但实际受显存带宽影响。最后判断瓶颈优先在KV Cache容量而非算力。这类交叉判断在方案设计里几乎每天都在用。7. 个人实操中的一些体会最后分享几个我经手多个项目后的经验总结。第一个体会任何显存估算都要做压力测试验证不要让公式成为唯一的依据。公式给的是理想账本真实环境里有CUDA context、框架缓存、通信buffer等各种隐形开销。我一般先把估算值乘以1.2到1.3作为初始分配再跑一个最长序列、最大batch的压测脚本拿实际的torch profiler或者nvidia-smi数据回来校准。第二个体会最简单的显存优化往往藏在架构选择里而不是花式优化技巧里。比如模型本身用GQA和不用GQAKV Cache差距可以高达4倍甚至32倍这比任何缓存淘汰算法都有效。我遇到过很多人一上来就研究各种复杂的KV Cache调度策略结果发现模型架构本身就不是省显存的设计换一个GQA变体重训一下或者微调一下问题迎刃而解。第三个体会训练和推理的显存规划虽然公式不同但思维模式一致都是先列全所有开销项再做叠加最后留出安全边际。我见过不少方案只看模型权重不看KV Cache或者只看模型状态不看激活值最终上线阶段都付出过代价。按这篇的账本方式先把所有项列出来后面每一步决策都会轻松很多。第三个体会顺带补充一个小技巧日常做显存估算时建议把常用公式做成一张随手可查的表包含权重、梯度、优化器状态、激活值、KV Cache、FLOPs这几类。时间久了这些估算就能变成条件反射开会时不用开电脑也能心算大概成本非常省事。