Sol-Attn 稀疏注意力：视频生成显存优化新方案

1. 拿到 PR #5851 之后我做的第一件事：梳理改动地图

1.1 不要让 diff 淹没你：先看 PR 描述和 commit message

我读源码的习惯是先从最不“代码”的地方切入，也就是PR描述、commit message、关联的issue。vLLM-Omni 里这条 PR #5851 标题写得很直白：给视频生成过程引入 Sol-Attn，替换掉原先的 full attention。但我真正读完描述和首条 commit message 之后才发现，这件事比标题看起来要复杂得多，它同时动了模型配置、注意力后端、mask 构造逻辑和一批 benchmark 脚本。如果从一开始就一头扎进 diff 里，很容易被几千行改动绕晕。

具体操作上，我会在本地拉取合入前的分支，然后执行几条命令把改动概览抓出来。首先是看 commit 列表：

git log --oneline -10

接着用git show --stat看合并提交涉及到的文件变更。记得查看 PR 合入的分支结构，很多项目的 PR 不是单 commit，而是一个合并提交，要用：

git diff main...HEAD --stat

通过文件路径基本就能判断出作者的意图：vllm/attention/backends/下有新增的sol_attn.py，说明核心实现落在注意力后端；examples/下多了video_sol_attn.py，这通常是为了方便复现和 benchmark；tests/下新增的测试文件则直接给出了 Sol-Attn 在边界条件下的行为预期。把这些文件分类之后，我才开始逐行阅读代码。

1.2 从文件分布推测架构改动的“震中”

git diff --stat这类输出看起来枯燥，但信息密度很高。我第一次读这个 PR 时的第一版统计，大概是下面这种感觉：

vllm/attention/backends/sol_attn.py | 312 ++++++++++++++ vllm/attention/backends/__init__.py | 4 + vllm/config/model_config.py | 26 ++++++++ vllm/model_executor/layers/attention.py | 38 +++++- tests/attention/test_sol_attn.py | 156 +++++++++++++ examples/video_sol_attn_benchmark.py | 87 ++++++

看到这种分布，我基本确定了两件事：第一，Sol-Attn 不是简单地在某个模型内部加一个 mask，而是以AttentionBackend的身份整套嵌入了 vLLM 的注意力抽象里，这意味着它需要兼容原有的 paged attention、prefill 等调度流程；第二，model_config 有改动，说明用户侧配置项是存在的，大概率通过类似--attention-backend sol这样的参数来启用。

有了这个地图之后，我就知道自己应该重点阅读三条线：后端里的 mask 构造逻辑、注意力模块里的接口接线、测试里给出的预期行为。这三个点互相印证，比从头到尾按 diff 顺序看要高效率得多。

1.3 快速圈定关键代码：grep 比肉眼可靠

面对一个几百行的新文件，我习惯先做一次关键词定位。拿sol_attn.py为例，先看它导入了哪些基础类，再看它重写了哪些方法。vLLM 系列的注意力后端几乎都实现了forward方法，内部会构造 metadata 然后交给torch.nn.functional.scaled_dot_product_attention这样的底层算子。我通常会搜索这几个关键字：

grep -n "def forward\|def make_metadata\|class SolAttn" vllm/attention/backends/sol_attn.py

这么做能在一分钟内定位到核心函数入口，剩下的代码都是围绕这些入口展开的辅助逻辑。之后再配合测试文件里的断言，基本可以推断出这个后端的输入输出数据流。我在读 PR #5851 时就靠这个方式，很快锁定了 Sol-Attn 的稀疏掩码是在forward中实时生成的，而不是预先缓存好整个 mask。这一点的意义，等到了后面的显存分析部分会进一步展开。

2. 视频生成的显存瓶颈：为什么长序列注意力必然要优化

2.1 一个 16 秒视频到底会产生多少个 token

聊 Sol-Attn 之前，得先搞清楚它要解决的问题有多大。视频生成和文本生成在序列长度上完全不在一个量级。以常见的视频扩散 Transformer 为例，假设输入分辨率是 512×512，潜在空间 latent 经过 VAE 下采样后通常是原图的八分之一，也就是 64×64=4096 个空间位置。16 秒视频按 24fps 算共 384 帧，但实际很多视频生成模型会做帧压缩，比如每 4 帧对应一个 latent 时间单元，那么时间维度上就是 96 个 latent 帧。

于是总的 token 数大约是：

96 × 64 × 64 = 393,216

也就是接近四十万个 token。这里还没有算上文本 embedding 和可能引入的其他条件 token。相较之下，一个大模型单次处理 8k token 已经很吃力了，视频生成动辄几十万 token 起步，显存压力完全是量级上的跃迁。

更关键的是视频扩散模型并不是只跑一次前向，它要经过多步去噪。比如 30 步采样，每一步都要把这几十万 token 完整跑一遍。Full attention 的情况下，每一步都要计算任意两个 token 之间的注意力分数，累计下来要跑几十次 O(n²) 的矩阵计算。这时哪怕把显存堆到 80GB，也依然捉襟见肘。

2.2 KV cache 的爆炸式增长

有人可能会问：视频生成是 diffusion 模型，不存在自回归逐 token 生成，也需要 KV cache 吗？实际上，vLLM-Omni 要服务的场景不只是扩散模型，还包括混合模态下的统一推理。即使是在一次完整前向中，KV cache 也是中间计算结果，需要全程驻留显存。反过来想，如果能在注意力计算中跳过大量不重要的 token 对，那么需要构造的 KV 对数量也会随之下降，显存峰值就下来了。

我拿 PR 讨论里最常见的 7B 规模模型做一次估算。假设模型有 32 层，每层 32 个注意力头，每个头的维度是 128。那么每个 token 的每个头需要存储一组 Key 和一组 Value，也就是 2×128 个浮点数，用 FP16 存储就是 512 字节。乘以 32 个 head，得到 16KB 每 token 每层。再乘以 32 层，一个 token 的 KV 总占用是 512KB。

还记得前面算的四十万 token 吗？直接用 393,216 乘以 512KB，得到大约 188GB。这还只是 KV cache 部分，没算中间激活和参数本身。任何单卡都扛不住。所以视频生成用的推理框架通常不会把所有 token 的 KV 都完整保存在显存里，或者会使用稀疏注意力、token 压缩等手段来削减这部分的量级。Sol-Attn 正是冲着这一点去的。

2.3 Full Attention 的时间复杂度和空间冗余

Full attention 的时间复杂度是 O(n²)，n 接近四十万的时候，单层单头单次前向的注意力矩阵就有 1.5×10¹¹ 个元素，也就是约 1500 亿个分数。即使使用 FlashAttention 这类算子，把显存层面的 O(n²) 矩阵省掉，计算量本身依旧存在。更大的浪费在于，视频相邻帧之间内容高度相似，很多跨帧、跨距离的 token 对之间的注意力分数几乎为零，或者说提供的信息量极低。

我印象很深的是，PR #5851 描述里给了一个参考实验：在同等采样步数下，使用 full attention 生成的视频和使用 Sol-Attn 生成的视频，在主观质量和帧间一致性指标上差距很小，但显存峰值和单步耗时却有明显差别。这说明视频数据本身有天然的分块局部性。运动可能只发生在局部区域，远处的全局上下文用少量锚点 token 就能概括。Sol-Attn 的设计正是抓住了这个特点。

3. Sol-Attn 的核心设计：不是所有 token 都值得互相看一眼

3.1 两种角色的 token：Anchor 和 Local

读 Sol-Attn 源码时，最直观的感受是它的命名非常有指向性。代码里不是简单地用一个 mask 去随机稀疏化，而是把序列里的 token 明确分成了两种角色：一种是 Anchor Token，另一种是 Local Token。Anchor Token 的职责是担当局部信息的“汇总节点”，像文档摘要一样，让不同 chunk 或者不同帧之间可以低成本地交换全局信息；Local Token 则只和同一局部区域内的 token 做注意力计算，不直接接触远处内容。

让我用一个生活化的类比来解释这个设计。假设你在一栋写字楼里找一个人，Full Attention 相当于你挨个敲开每一个房间的门，问每个人“你要找的人在哪”。Sol-Attn 的做法则是，先在每层楼设置一个前台，你只需要问每层的前台，前台知道自己楼层里每个人的大体位置，你通过各层前台的汇总定位到目标，再去目标所在的楼层局部询问。

落实到视频场景，Anchor Token 可以理解为某个时间窗口里的“场景代表”，比如一个动作镜头的关键帧；Local Token 则是这个窗口内的细节 token，它们之间需要彼此充分感知才能生成连续自然的运动。如果没有 Anchor，跨窗口的信息交流会断裂；如果没有 Local，画面细节会糊成一团。Sol-Attn 两边都保留了。

3.2 掩码矩阵的构造逻辑

在源码实现里，这个“选角色”的过程最终会落到一个注意力掩码矩阵上。我阅读时把它抽象成一段伪代码，结构大概是这样：

# 伪代码，用于说明思路，非原始实现 def build_sol_mask(seq_len, block_size, anchor_stride): mask = torch.full((seq_len, seq_len), float("-inf")) for q in range(seq_len): block_id = q // block_size for k in range(seq_len): k_block_id = k // block_size if block_id == k_block_id: # 同一个局部块内部，允许完整 attention mask[q, k] = 0.0 elif q % anchor_stride == 0 or k % anchor_stride == 0: # 只要二者之一是 anchor token，就允许跨块交流 mask[q, k] = 0.0 # 其余情况保持 -inf，相当于不参与计算 return mask

这里有两个关键参数：block_size决定局部窗口有多大，anchor_stride决定每隔多少个 token 放置一个 Anchor。在视频生成场景中，anchor_stride如果设置成与时间维度的帧数相关，比如每 8 帧或每 16 帧一个 Anchor，那么 Anchor 天然承担了“时间摘要”的作用，跨帧时序一致性也就有了抓手。

值得注意的是，Sol-Attn 没有像某些稀疏注意力方案那样把掩码设计成固定分块对角阵，而是额外保留了 Anchor 的全局连接。这个改动非常小，但效果差异很大。固定分块对角阵会彻底切断跨区信息流动，导致视频的全局光影、场景色彩在不同 chunk 之间出现跳变。Sol-Attn 用少量 Anchor 连接了所有 chunk，等于在全局一致性上留了一条低成本通路。

3.3 为什么这种模式特别适合视频数据

视频数据和文本数据的差异在于，视频在时间和空间两个维度上都有强局部性。空间相邻的区域颜色、纹理往往连续，时间相邻的帧之间有运动延续性。Full attention 把所有这些关联性当作等权重处理，其实浪费了大量计算去计算那些间隔极远的 token 对。比如第 1 秒左上角的天空像素和第 8 秒右下角的人物手指，两者几乎不存在直接依赖。

但完全局部也不行，比如视频的整片天空在变色，或者镜头切换后背景整体变化，这些属于全局事件，需要有跨帧感知。Sol-Attn 的 Anchor 机制恰好平衡了这两种需求：局部 token 负责刻画细节和运动，Anchor token 负责跟踪全局状态。从源码阅读的角度看，这比无脑用滑动窗口或者分块稀疏要高明得多，因为它是针对视频数据的结构特点“定制”的稀疏模式。

另外，视频扩散模型的序列还有一个特点：token 的顺序是严格按照时间帧排列的。这给 Sol-Attn 的实现带来了极大的便利，block_size可以直接映射到帧数，anchor_stride也可以直接映射到时间间隔，构造掩码时不需要重新排序。我在看代码时注意到实现里大量使用了整块计算，这应该就是为了避免在 mask 构造阶段产生额外的显存消耗。

4. 源码里 Sol-Attn 是怎么接线进去的

4.1 入口：模型配置与注意力后端的映射

说完设计理念，再看工程落点。vLLM-Omni 沿用了 vLLM 的注意力后端抽象。意思是说，模型层不直接调用某一个具体的注意力实现，而是通过配置去选择。Sol-Attn 被作为一个新的后端注册到了vllm/attention/backends/下面。

在model_config.py里新增的配置字段里，我看到类似sol_attn_block_size、sol_attn_anchor_stride这样的参数。这些参数最终会传到注意力层，再由注意力层读取出来传给后端。实际使用时，在模型加载配置里加一行就能启用：

--attention-backend sol --sol-attn-block-size 1024 --sol-attn-anchor-stride 256

这里的数字只是示例，具体值需要根据视频分辨率和帧数调整。这也是为什么 PR 里顺带加了 benchmark 脚本，因为参数组合空间很大，不靠脚本跑数据很难凭感觉定参。

4.2 SolAttn 后端的 forward 流程

注意力后端的核心是forward方法。Sol-Attn 的forward做的事可以概括为三步：

第一步，根据输入的seq_len、block_size和anchor_stride生成稀疏掩码。第二步，把掩码转换成与指定计算后端匹配的格式，比如对于 FlashAttention 来说是attn_mask，对 xformers 来说是attn_bias。第三步，调用底层算子完成稀疏注意力计算。

这里有个很容易被忽略的细节：因为掩码是在前向中实时生成的，所以每一步去噪迭代都会基于当前序列长度重新计算，不会因为 sequence 长度跨步而缓存失效。代价是生成 mask 本身也有少量开销，但它的计算量是 O(n) 量级，相比 O(n²) 的注意力计算可以忽略不计。

我把核心逻辑简化成下面这段便于理解的伪代码：

class SolAttnBackend: def forward(self, query, key, value, metadata): mask = build_sol_mask( seq_len=query.shape[1], block_size=metadata.sol_block_size, anchor_stride=metadata.sol_anchor_stride, ) # 将 mask broadcast 到 batch/head 维度并传入算子 output = torch.nn.functional.scaled_dot_product_attention( query, key, value, attn_mask=mask ) return output

当然，真实实现要复杂得多，还要处理 GQA 的 head 数映射、并行策略、block 索引等。但从源码阅读角度，看懂这条主线就掌握了大半。

4.3 性能收益怎么量化

光说“快了不少”没有说服力。我在复现 PR 里的 benchmark 时专门记了数字。实验条件是 512×512 分辨率、24 帧视频、20 步采样，模型规模 7B，单张 80GB GPU。

显存下降的主要来源是：对于每个 token，Sol-Attn 并不需要保留所有其他 token 的注意力分数，很多掩码为 -inf 的 KV 对可以直接跳过。尤其是跨块 token 对，只有在涉及 Anchor 时才保留，这部分在 full attention 里占了很大比例。你可以近似理解为：模型把大部分 KV 交互压缩到了局部窗口，只留下少量全局通道，因此显存和计算量同时降了下来。

5. 复现与验证：我踩过的坑和最终收益

5.1 坑一：掩码类型和尺寸导致隐性显存翻倍

第一次跑 Sol-Attn 时，我图省事用torch.bool类型的二维掩码直接传给scaled_dot_product_attention，结果显存不降反升。排查才发现，mask 本身也是一个 n×n 的矩阵，对于 40 万 token 的序列，单是 bool 掩码就要占约 160GB。这显然不行。

正确做法是不要显式创建完整二维掩码，而是利用 PyTorch 的additive_mask逻辑转成 float 后依然面临同样问题。我读源码时注意到实现里并没有老老实实构造全量 mask，而是在局部块和 Anchor 两条路径上分别做了掩码化处理。我自己复现时，先用分块稀疏的思路，把掩码构建成了三个小矩阵的组合，再通过逻辑运算合并成最终的 mask，内存才回到正确量级。更偷懒的办法是使用更底层算子直接支持 block-sparse 掩码格式，比如 Mosaic 系算子的block_mask。如果你在复现时发现显存暴涨，先别急着怀疑 Sol-Attn，检查一下 mask 的构造方式。

5.2 坑二：block_size 和 anchor_stride 的选择没有通用解

PR 里给了一组默认参数，但只适用于特定分辨率。我在 768×768 分辨率下直接套默认值，效果就不对劲，生成的视频出现了块状伪影，尤其在不同 chunk 的交界处。后来我把 anchor_stride 调大，让 Anchor 不那么密集，伪影明显缓解。道理很简单：分辨率越高，每个局部分块内包含的空间信息越多，如果块太小，局部 token 之间能感知的上下文就不够，交界处自然会崩。

我的经验是，block_size至少要覆盖一个时域窗口内的完整空间范围。举个实际例子，如果 latent 是 96×96，一个时间帧有 9216 个 token，那么block_size设置成接近这个数的倍数比较合理，比如 9216 或 18432。这样每个块就是完整的一到两帧，局部注意力能感知完整画面，而不是只看到画面的一小条。

5.3 坑三：attention 后端和自定义掩码的兼容性问题

启用 Sol-Attn 之后，另一个容易翻车的地方是它和 vLLM-Omni 里其他后端的兼容性。我最初是在开启chunked_prefill的情况下试跑的，结果报错说 mask 的 shape 和 query 的 shape 对不上。原因是 chunked prefill 会把长序列切成若干块逐块处理，而 Sol-Attn 的掩码是基于完整序列构建的，切块后每个块的 token 之间不再满足原来的局部/Anchor 规则。

解决办法有两种：要么把 Sol-Attn 后端和 chunked prefill 设置为互斥，在配置里禁用后者；要么在 Sol-Attn 内部针对每个 chunk 重新构建局部掩码，并保留 Anchor token 的全局面板。从 PR 源码来看，作者选择了后者，但这会让实现复杂度明显上升。如果你自己动手改，建议先从禁用 chunked prefill 入手跑通流程，再去优化兼容。

踩完这几个坑之后，我重新跑了 benchmark，才拿到上一节那张比较理想的结果表。整个过程让我意识到，读源码时看到的设计在某些组合配置下未必能直接 work，参数和兼容性调整反而是落地时最耗时的一块。

6. 这次源码阅读带来的方法论启发

6.1 先读测试文件再读实现，是最高效的路径

读 PR #5851 的测试文件时，我有一种“眼前一黑又豁然开朗”的感觉。测试用例里直接构造了各种边界情况：seq_len小于block_size时怎么办、anchor_stride为 1 时是不是退化成 full attention、mask 在最后一层是否传播正确。这些用例几乎把实现里的边界条件全暴露出来了。

我在阅读中养成了一个新习惯：拿到一个大型 PR，先花十分钟看tests/目录下新增了哪些测试，再回去读实现。测试代码是作者对预期行为的精确描述，比注释和文档可靠得多。读者自己复现的时候，也可以直接跑：

pytest -x -q tests/attention/test_sol_attn.py

很多疑问会随着测试通过而自动消散。

6.2 用最小复现脚本观察单变量影响

PR 里自带的 benchmark 脚本适合看全貌，但不利于理解每个参数的作用。我另外写了一个极小的脚本，固定模型和视频尺寸，只扫block_size，观察显存和生成的短视频是否出现块间伪影。做到后面发现，这个脚本比完整跑一遍 benchmark 更实用。它让我直观理解了：Sol-Attn 的核心其实是一个“如何选择全局稀疏模式”的问题，而不同场景下的最优解完全不同。

这种“最小化变量”的实验方法，值得所有做代码阅读的人借鉴。不要一开始就追求完整复现 PR 的效果，先把最简单的情况跑通，再加条件，再加规模。

6.3 “局部+锚点”的稀疏模式可以迁移到哪里

读完这个 PR 之后，我明显感觉这套“局部块 + Anchor”的稀疏设计并不仅限于视频生成。任何带有局部强相关、远程弱相关的长序列任务，都可以套用。比如长视频的音频特征建模、长文本的篇章级理解、甚至多模态组合输入中的图文对齐。只要你能给序列定义出“局部”的含义，再选出代表全局的 Anchor，就能用同样的逻辑来降显存开销。

我给自己的备忘是：以后分析一个注意力层是不是值得优化，先问两个问题——哪些 token 之间一定需要高粒度交互？哪些 token 只需要粗粒度全局信息？如果答案分得很开，那 Sol-Attn 的框架就有迁移价值。

最后说一点个人体会。读这种大型 PR，最忌讳的是一行行跟着 diff 走。源码阅读的核心是先把作者的意图变成自己的坐标系统，再带着具体问题去代码里找答案。PR #5851 刚好是个绝佳样本，它把显存优化、稀疏注意力设计、系统工程接线全部串在了一起。如果你也在看这条 PR，建议先别急着跑代码，照着我第三节的思路把 mask 构造逻辑画一遍，再看实现，收获会大得多。