1. 从一次显存告急说起为什么要盯住注意力分布做多模态推理优化的人大概都经历过这种时刻一张图加一段长 prompt 丢进去模型还没开始吐字显存先红了。你去看显存占用权重只占了一部分剩下的大头全在 KV cache 上。尤其是多模态场景图像经过视觉编码器之后动辄产生几百上千个 visual token再叠上文本 token序列长度轻松翻倍。这时候如果还按文本模型那套每层都缓存、每个 token 都保留的思路去做显存和带宽都会被吃干抹净。UniCache 这个系列要解决的就是这件事。它的核心命题很直接多模态理解与生成模型里注意力到底是怎么分布的哪些 KV 是真正被反复读到的哪些其实可以丢。只有先把注意力分布摸清楚后面的缓存淘汰、复用、压缩策略才有依据。这篇是系列的第一篇定位是纵览也就是把理解模型和生成模型这两大类多模态架构的注意力行为先摊开来看一遍建立一个统一的观察框架。我写这篇的出发点很实际。过去一年我在本地和集群上跑过不少多模态模型从早期的 CLIP 类双塔到后来的 LLaVA 式拼接结构再到 DiT 这类生成架构一个共同的感受是理解模型和生成模型的注意力分布差异极大但很多缓存优化方案却把它们当成一回事。结果就是理解任务上效果还行一换到生成任务就崩或者反过来。所以这篇不急着讲具体怎么压缩而是先把分布长什么样这件事讲透。适合谁看如果你正在做多模态推理加速、KV cache 管理、长上下文优化或者只是单纯好奇为什么多模态模型这么吃显存这篇都能给你一个可复现的观察视角。我会尽量用能跑起来的代码和实测数据说话而不是停留在概念层面。2. 多模态注意力分布的基本盘先搞清楚在看什么2.1 理解模型与生成模型的注意力结构差异要谈分布先得把结构对齐。多模态理解模型的典型代表是 LLaVA 系列视觉编码器通常是 ViT把图像切成 patch输出一批 visual token然后通过一个投影层MLP 或 Q-Former映射到语言模型的 embedding 空间和文本 token 拼在一起送进 decoder。这里的注意力是因果的每个 token 只能看到自己之前的 tokenvisual token 一般放在序列前部文本 token 在后。生成模型就完全是另一套逻辑。以 DiTDiffusion Transformer和各类多模态生成架构为例图像 token 和文本 token 之间往往是双向注意力或者至少在同一模态内部是双向的。扩散过程还要反复迭代去噪同一个序列会被前向计算几十次。这意味着生成模型的注意力分布不是一次成型而是随去噪步数动态变化的。这个结构差异直接决定了 KV cache 的行为。理解模型里KV 是自回归逐 token 增长的缓存是追加式的生成模型里如果做缓存更多是跨步复用同一位置的 KV 在不同去噪步之间被反复读取。两者对哪些 KV 重要的答案完全不同。2.2 注意力分布到底在观察什么指标很多人一说注意力分布就只想到 softmax 之后的注意力权重矩阵。但真正做缓存优化光看权重是不够的。我一般会同时观察四个维度注意力熵某一层某个 head 的注意力分布有多集中。熵低说明这个 head 只盯着少数几个 token这些 token 就是缓存的重点。KV 命中频次在自回归解码过程中某个位置的 KV 被后续 token 读取了多少次。频次高的位置淘汰代价大。层间差异浅层和深层的注意力模式往往截然不同。浅层偏向局部和位置深层偏向语义聚合。模态间注意力流向文本 token 对 visual token 的注意力占比以及反向的流向。这决定了跨模态缓存能不能分开管理。这四个指标里KV 命中频次是最贴近缓存决策的因为它直接对应这个 KV 被读了几次。注意力熵则更像是一个先验帮你判断哪些 head 值得重点保护。实际做的时候我建议先把命中频次统计出来再回头看熵两者交叉验证。2.3 为什么不能直接套用文本模型的结论文本模型上有一套很成熟的观察attention sink前几个 token 吸收大量注意力、局部性近期 token 更重要、层间稀疏性差异。这些结论搬到多模态上一部分成立一部分会翻车。最典型的翻车点是 attention sink。文本模型里序列开头的 BOS token 经常成为 sink几乎所有 head 都会分一部分注意力给它。但在多模态理解模型里序列开头是 visual token情况就复杂了有些 head 确实会盯着某几个锚点visual token但另一些 head 的注意力会均匀铺在整片图像区域上。如果你按文本模型的经验粗暴地只保留前几个 token 的 KV图像信息会丢得很惨。另一个差异是位置编码的干扰。多模态模型里 visual token 的位置编码方式和文本不同有的是二维位置编码有的是插值后的一维编码。这会导致注意力分布在 visual 区域内部呈现出和文本完全不同的几何结构。所以我的态度很明确文本模型的缓存经验可以参考但必须重新验证不能直接照搬。3. 理解模型的注意力分布visual token 到底被谁在读3.1 LLaVA 式拼接结构的实测观察我拿一个标准的 LLaVA 结构做了实测图像分辨率 336ViT 输出 576 个 visual token加上一段约 200 token 的文本 prompt总序列长度接近 800。用一张信息量中等的图片比如一张有多个物体的街景跑一次完整前向把每层的注意力权重和 KV 读取情况都记录下来。第一个发现是visual token 内部的注意力高度不均匀。576 个 visual token 里真正被文本 token 大量读取的集中在少数几十个位置上。这些位置往往对应图像里的显著物体或高对比度区域。换句话说图像被看的时候模型并不是均匀地扫描而是有选择地聚焦。第二个发现是层间分化非常明显。浅层前 1/3的注意力基本是局部的visual token 主要看自己附近的邻居文本 token 也主要看邻近的 visual token。到了中层跨模态注意力开始活跃文本 token 大量读取视觉信息。深层则反过来visual token 开始聚合全局语义注意力铺得更开。这个分层特征对缓存策略的启示很直接浅层和深层的 KV 重要性判断标准应该不一样。浅层可以更激进地按局部性淘汰深层则要保守一些因为深层 KV 承载的是聚合后的语义。3.2 视觉 token 的注意力锚点现象在实测里我注意到一个很有意思的现象我把它叫做视觉锚点。在每一层里总有那么几个 visual token它们的 KV 被读取的频次远高于其他 token而且这些锚点在不同层之间有一定的延续性——浅层的锚点位置往往在中层还是锚点。这跟文本模型的 attention sink 有点像但成因不同。文本 sink 更多是位置编码和 softmax 归一化导致的无处可去的注意力倾倒视觉锚点则更像是模型真的在用它作为图像信息的索引。我猜测这和 ViT 的 patch 聚合方式有关某些 patch 在编码阶段就吸收了更多信息到了语言模型里自然成为被反复查询的对象。实操上的价值是如果能稳定地识别出这些锚点就可以给它们开小灶优先保留。识别方法也不复杂跑一批样本统计每个 visual token 位置在各层的平均读取频次取 top-k 作为锚点候选。我在几个不同图片上验证过锚点位置有一定稳定性但不是完全固定所以最好按任务类型分别统计。3.3 文本到视觉的注意力流向分析跨模态注意力流向是理解模型缓存优化的关键。我统计了文本 token 对 visual token 的注意力占比发现一个规律不是所有文本 token 都同等程度地看图像。具体来说实体词、名词短语、形容词这些承载视觉语义的 token对 visual token 的注意力占比很高而功能词、标点、连接词注意力基本留在文本内部。这意味着如果你要做跨模态的缓存压缩可以按 token 类型区别对待视觉相关的文本 token 需要保留对视觉 KV 的访问能力纯功能性 token 则可以更激进地压缩。还有一个细节值得注意文本对视觉的注意力在生成回复的不同阶段是变化的。当模型开始描述图像内容时对 visual token 的读取会突然增强当它在做纯语言推理时读取会减弱。这个动态特征说明静态的缓存策略可能不是最优的理想情况下应该根据解码阶段动态调整。4. 生成模型的注意力分布跨步复用才是重点4.1 扩散架构下的注意力动态生成模型这边我主要观察了 DiT 类架构。和自回归理解模型最大的不同是扩散模型的注意力是跨去噪步的。同一个序列在 50 步去噪里会被前向计算 50 次每次的注意力分布都不一样。实测下来去噪早期高噪声阶段的注意力比较弥散模型还在探索整体结构去噪后期低噪声阶段注意力变得集中模型在打磨细节。这个变化对缓存的意义是早期步的 KV 复用价值低因为分布还在剧烈变化后期步的 KV 复用价值高因为分布趋于稳定。我试过一个简单的策略只在去噪后 1/3 的步数里启用 KV 缓存复用前面的步数每步重算。结果在几乎不损失生成质量的前提下省下了可观的计算量。这个思路和很多缓存方案全程复用的默认假设是相反的但实测更稳。4.2 图像 token 与文本 token 的双向注意力特征生成模型里图像 token 和文本 token 之间通常是双向注意力。我观察到的一个特征是文本 token 对图像 token 的注意力在去噪过程中相对稳定而图像 token 对文本 token 的注意力波动较大。这个不对称性很有意思。文本作为条件它的语义是固定的所以图像去查询文本的模式比较稳定而图像本身在去噪中不断变化它查询文本的方式也就跟着变。对缓存的启示是文本侧的 KV 可以更放心地跨步复用图像侧的 KV 要更谨慎。另外生成模型里图像 token 数量通常远大于文本 token比如 1024 个图像 token 对 77 个文本 token所以缓存压力主要在图像侧。这就更需要精细地判断哪些图像 token 的 KV 值得留。4.3 跨步 KV 复用的可行性边界跨步复用听起来很美但边界在哪我做了几组对比实验结论是复用能不能成立取决于相邻去噪步之间注意力分布的相似度。我用余弦相似度衡量相邻步的注意力分布发现在高噪声阶段相邻步相似度可能只有 0.3 到 0.5这时候复用基本等于瞎猜到了低噪声阶段相似度能到 0.8 以上复用就很安全。所以一个可操作的判据是设定一个相似度阈值比如 0.7只在超过阈值的步之间复用 KV。这个阈值不是拍脑袋定的需要按模型和任务调。我一般会先跑一批样本画出相似度随步数变化的曲线找到那个拐点把阈值设在拐点之后。这样既保证了质量又最大化了复用收益。5. 统一观察框架把两类模型的分布放在一起看5.1 理解与生成的分布对比表把前面两部分的观察整理成一张表能更清楚地看到差异观察维度多模态理解模型多模态生成模型注意力方向因果单向双向或模态内双向KV 增长方式自回归追加跨去噪步复用关键 token视觉锚点 视觉相关文本稳定文本条件 后期图像 token分布稳定性随解码位置变化随去噪步数变化缓存优化重点淘汰低频 visual token判断跨步复用边界浅层特征局部、位置主导结构探索、弥散深层特征全局语义聚合细节打磨、集中这张表是我做后续缓存策略设计的基础。它说明一件事不存在一套通用的多模态缓存方案必须按模型类型分而治之。UniCache 这个系列后面要讲的统一缓存框架也是建立在这个分而治之再统一抽象的思路上的。5.2 一个可复用的注意力统计脚本光看结论不够得能自己跑。我整理了一个最小可用的统计脚本基于 PyTorch hook 抓取每层注意力统计 KV 读取频次和注意力熵。核心逻辑如下import torch from collections import defaultdict class AttentionStats: def __init__(self): self.read_count defaultdict(lambda: None) self.entropy defaultdict(list) def hook_fn(self, layer_idx): def fn(module, inputs, outputs): # outputs 里通常包含 attention weights attn outputs[1] if isinstance(outputs, tuple) else None if attn is None: return # attn: [batch, heads, q_len, kv_len] with torch.no_grad(): # 统计每个 kv 位置被读取的总次数 read attn.sum(dim(0, 1, 2)) # [kv_len] if self.read_count[layer_idx] is None: self.read_count[layer_idx] read.cpu() else: self.read_count[layer_idx] read.cpu() # 计算每个 head 的注意力熵 p attn.clamp(min1e-9) ent -(p * p.log()).sum(dim-1).mean().item() self.entropy[layer_idx].append(ent) return fn stats AttentionStats() handles [] for idx, layer in enumerate(model.layers): handles.append(layer.self_attn.register_forward_hook(stats.hook_fn(idx)))跑完之后stats.read_count里就是每层的 KV 读取频次stats.entropy是每层的平均注意力熵。这两个量结合起来就能画出分布图判断哪些位置值得保留。注意不同框架的 attention 模块输出格式不一样有的返回 attention weights有的默认不返回需要开启output_attentionsTrue。抓之前先确认一下否则 hook 拿到的可能是 None。5.3 统计结果怎么读拿到数据之后怎么读很关键。我一般会做三件事第一画 KV 读取频次的热力图横轴是 token 位置纵轴是层。这样一眼就能看出哪些位置在哪些层被重点读取。视觉锚点会表现为几条竖直的高亮带。第二画注意力熵的层间曲线。熵低的层说明注意力集中这些层的 KV 更专一淘汰风险高熵高的层注意力分散KV 冗余度大可以更激进地压缩。第三做模态分组统计。把 visual token 和文本 token 分开算平均读取频次看跨模态注意力的整体强度。如果文本对视觉的读取占比很低说明这个任务其实不太依赖视觉缓存可以大幅压缩视觉侧。这三步做完你对一个模型的注意力分布就有了量化认识后面做任何缓存决策都有据可依而不是凭感觉。6. 实操中的坑与排查经验6.1 常见问题速查表问题现象可能原因排查方向hook 抓不到 attention框架默认不返回权重开启 output_attentions读取频次全为 0hook 挂错模块确认挂在 self_attn 而非整个 block熵值异常低softmax 前未缩放检查是否用了原始 logits视觉锚点不稳定样本差异大分任务类型分别统计跨步复用后质量崩阈值设太低提高相似度阈值显存没降缓存没真正释放检查是否只标记未释放这张表是我踩坑之后总结的基本覆盖了统计阶段最常见的几类问题。尤其是 hook 挂错模块这一条我一开始挂在 transformer block 上拿到的输出里根本没有 attention weights白白折腾了半天。6.2 几个容易忽略的细节第一个细节是batch 维度。统计的时候如果 batch size 大于 1不同样本的注意力分布会混在一起。做锚点分析时最好 batch size 设为 1或者按样本分别统计再聚合否则锚点位置会被平均掉。第二个细节是精度。用 fp16 跑统计注意力权重可能会有数值误差熵的计算会偏。建议统计阶段用 fp32虽然慢一点但数据可靠。真要省时间至少把 softmax 那一步放到 fp32 里算。第三个细节是序列长度对齐。不同样本的 visual token 数量可能不同如果做了动态分辨率统计时要么固定分辨率要么按比例归一化位置。我一般固定分辨率省得后面解释麻烦。6.3 我个人的几条经验做了这么多轮统计有几条经验我觉得比技术细节更重要。别一上来就追求全模型统计。先挑几层代表性的比如第 1、中间层、最后一层看建立直觉再决定要不要全量跑。全量统计又慢又占内存很多时候没必要。分布是会随 prompt 变的。同一个模型换一种提问方式注意力分布可能就不一样。所以统计要覆盖多样化的 prompt不能只用一两个样本就下结论。我一般至少准备 20 到 30 个不同类型的样本。把统计和实际缓存实验结合起来。光看分布不够得真的按分布去淘汰 KV看效果掉多少。分布告诉你哪些可能重要实验告诉你淘汰了到底行不行。两者结合结论才靠谱。7. 从分布到策略下一步该往哪走把理解模型和生成模型的注意力分布摊开看之后一个自然的下一步就是怎么把这些观察转化成可执行的缓存策略。这篇作为纵览先把地基打好具体的淘汰算法、复用调度、跨模态缓存分区会在后续展开。但有一个方向我现在就可以说统一缓存框架的关键不在于找一套万能规则而在于建立一个能自适应判断当前该用哪种规则的机制。理解模型和生成模型的分布差异这么大硬套一套规则必然顾此失彼。更现实的做法是先识别当前处于什么模式理解还是生成、高噪声还是低噪声、浅层还是深层再切换到对应的缓存策略。我在实际项目里试过一个简化版的自适应逻辑用一个轻量的相似度检测判断当前步的注意力分布是否稳定稳定就复用不稳定就重算。这个逻辑在理解模型和生成模型上都能跑虽然还不够精细但至少证明了自适应这条路是通的。最后分享一个我常用的调试技巧把注意力分布可视化出来存成图片按层和步数排成网格。很多时候数值看不出的模式一眼图就看出来了。尤其是跨模态注意力的流向画出来比看表格直观得多。这个习惯帮我省了很多瞎猜的时间。