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这通常是为了方便复现和 benchmarktests/下新增的测试文件则直接给出了 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×644096 个空间位置。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 两种角色的 tokenAnchor 和 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_lenquery.shape[1], block_sizemetadata.sol_block_size, anchor_stridemetadata.sol_anchor_stride, ) # 将 mask broadcast 到 batch/head 维度并传入算子 output torch.nn.functional.scaled_dot_product_attention( query, key, value, attn_maskmask ) return output当然真实实现要复杂得多还要处理 GQA 的 head 数映射、并行策略、block 索引等。但从源码阅读角度看懂这条主线就掌握了大半。4.3 性能收益怎么量化光说“快了不少”没有说服力。我在复现 PR 里的 benchmark 时专门记了数字。实验条件是 512×512 分辨率、24 帧视频、20 步采样模型规模 7B单张 80GB GPU。指标Full AttentionSol-Attn变化峰值显存占用约 38.6 GB约 21.4 GB下降约 44%单步平均耗时约 4.2 秒约 3.1 秒下降约 26%20 步总耗时约 84 秒约 62 秒节省约 22 秒显存下降的主要来源是对于每个 tokenSol-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 构造逻辑画一遍再看实现收获会大得多。