多模态 DiT 的推理成本，真正跑过的人都知道，瓶颈往往不在参数规模，而在注意力那一步的显存和带宽。我最近在做一个图文联合生成的模型加速，模型结构是典型的单流 DiT，文本 token 和图像 patch 拼成一条序列送进 Transformer block。序列长度一上去，标准 Attention 的 O(N²) 就开始吃人：显存爆、延迟高、batch 上不去。试过直接砍 token、降分辨率，效果掉得厉害；也试过换更快的 kernel，收益有限。最后落到块稀疏 Attention（Block Sparse Attention，BSA）这条路上，才算把问题拆开看清楚了。这篇就把我在多模态 DiT 上做块稀疏 Attention 的完整思路、选型理由、实现细节和踩过的坑摊开讲一遍，适合正在做多模态生成加速、或者准备把 DiT 推到更高分辨率/更长序列的同行参考。

1. 多模态 DiT 里 Attention 到底贵在哪

1.1 单流 DiT 的序列构成与注意力形态

先把这个场景说清楚。所谓单流 DiT，指的是文本和图像不走两套独立的编码器再融合，而是把文本 token 和图像 latent patch 直接拼成一条序列，共享同一组 Transformer block。这跟双流结构（文本一路、图像一路，中间靠 cross-attention 交互）是两种思路。单流的优势是模态交互更充分、结构更统一，代价就是序列长度直接叠加。

举个具体的量级：一张 1024×1024 的图，patch size 取 2，图像 token 就是 (1024/2)² = 262144 个，这还没算文本。实际工程里当然不会这么干，通常会先过 VAE 把图压到 latent 空间，比如 128×128 的 latent，patch size 取 2，那就是 64×64 = 4096 个图像 token，再加几十到几百个文本 token。序列长度 N 落在 4000 到 8000 这个区间是很常见的。

标准自注意力的计算量是 O(N²·d)，显存占用也是 O(N²)（注意力矩阵本身）。N=4096 时，单头注意力矩阵就是 4096×4096，按 fp16 算一个头 32MB，多头叠加、多层叠加，很快就顶到显存天花板。更关键的是，这个 O(N²) 里绝大部分权重其实很小，真正起作用的注意力连接是稀疏的——这就是块稀疏能成立的前提。

1.2 为什么是"块"稀疏而不是元素级稀疏

很多人第一反应是元素级稀疏：把注意力矩阵里小的元素直接置零。理论上很美好，实际在 GPU 上几乎跑不动。原因是 GPU 的算力来自矩阵乘的规整性，元素级稀疏会打乱内存访问模式，非零元素散落各处，访存效率极低，最后省下的 FLOPs 全被访存开销吃回去了。

块稀疏的思路是把注意力矩阵切成固定大小的块（比如 64×64 或 128×128），以块为单位决定"算还是不算"。这样做的好处是：被选中的块内部依然是稠密矩阵乘，可以走高效的 tensor core 路径；被跳过的块整块不加载、不计算，显存和算力都省。对 GPU 来说，规整的块结构才是能真正兑现收益的稀疏形式。这也是我在这个项目里坚持用块稀疏而不是元素稀疏的核心原因。

1.3 多模态场景下稀疏模式的特殊性

纯图像 DiT 的注意力稀疏模式相对好找，因为图像 patch 之间有很强的局部性——邻近 patch 相关性高，远处 patch 相关性低，天然适合局部窗口加少量全局连接。

但多模态场景不一样。文本 token 数量少但语义密度极高，而且文本和图像之间存在跨模态的强关联：描述"一只红色的猫"的文本 token，跟图像里猫所在区域的 patch 关联极强，跟背景天空的 patch 关联很弱。这种关联是跨距离的、内容驱动的，不是简单的局部性。所以多模态 DiT 的块稀疏不能照搬图像那套固定局部窗口，必须考虑跨模态的块选择策略。这是整个项目里最需要动脑子的地方，后面会专门展开。

2. 块稀疏 Attention 的选块逻辑：从固定模式到内容驱动

2.1 固定稀疏模式：局部窗口 + 全局 token

最省事的做法是固定模式。常见的有两种：

• 局部窗口（Local Window）：每个 token 只跟前后各 w 个 token 做注意力。实现简单，稀疏率可控，对图像 patch 的局部性很友好。
• 全局 token（Global Token）：选一部分 token（比如文本 token、或者图像里均匀采样的 anchor patch）作为全局节点，所有 token 都跟它们做注意力，它们也跟所有 token 做注意力。

把两者结合，就是"局部窗口 + 全局 token"的经典结构。文本 token 天然适合当全局 token，因为它们数量少、语义密度高，让所有图像 patch 都能"看到"文本，跨模态对齐就有了通路。

这个方案我在第一版里用了，跑得通，但效果有上限。问题在于：全局 token 是固定的，不管内容是什么，每个图像 patch 都去关注全部文本 token。当文本很长（比如一段详细描述）时，全局连接的开销又上来了，而且很多文本 token 跟某个具体 patch 其实没关系，属于无效计算。

2.2 内容驱动的块选择：让注意力自己决定看哪里

第二版我换成了内容驱动的块选择。核心思路是：先用一个轻量的打分机制估计每个 query 块对每个 key 块的重要性，然后只保留 top-k 个 key 块参与真正的注意力计算。

打分机制有几种常见做法：

我最后用的是均值池化打分加一个小的可学习投影。具体来说，对每个块内的 token 特征做均值池化得到一个块级向量，然后用一个小的线性层投影到打分空间，算 query 块和 key 块的相似度，取 top-k。这个投影层参数量很小，可以在训练时一起学，让打分机制适配多模态的数据分布。

提示：打分机制一定要轻。如果打分本身的开销接近省下来的注意力开销，那整个块稀疏就没意义了。我实测下来，打分部分的开销控制在总注意力的 5% 以内是比较健康的。

2.3 跨模态块选择的特殊处理

多模态的关键点在这里。如果对文本块和图像块用同一套打分逻辑，会出现一个问题：文本 token 数量少，池化后信息损失大，打分不准，导致跨模态连接被误砍。

我的处理是分而治之：

• 对图像到图像的注意力，用标准的块打分 + top-k，充分利用图像局部性。
• 对图像到文本的注意力，不做块稀疏，或者只做很轻的稀疏。因为文本 token 本来就少，这部分开销可控，而且跨模态对齐对生成质量影响极大，砍不得。
• 对文本到图像的注意力，同样保持较稠密，保证文本能充分"指挥"图像生成。

换句话说，稀疏主要施加在占大头的图像自注意力上，跨模态那部分谨慎处理。这个策略听起来朴素，但实测效果比"一刀切稀疏"好很多，尤其是文本遵循度（prompt following）这个指标上差距明显。

3. 在 DiT Block 里落地块稀疏的工程细节

3.1 与 FlashAttention 的关系和取舍

这里要说清楚一个容易混淆的点：块稀疏 Attention 和 FlashAttention 不是一回事，但可以结合。

FlashAttention 解决的是稠密注意力的 IO 效率问题——它通过分块计算和在线 softmax，避免把完整的 N×N 注意力矩阵写回显存，大幅降低显存占用和访存开销。但它计算的还是完整的稠密注意力，FLOPs 没变。

块稀疏解决的是计算量问题——它直接跳过一部分块，FLOPs 真的降了。

两者结合的逻辑是：先用块稀疏决定哪些块要算，然后对这些被选中的块用 FlashAttention 式的分块计算。这样既省了 FLOPs，又省了 IO。我在实现时，被选中的块集合是不规则的，所以没法直接调用现成的稠密 FlashAttention kernel，需要自己写一个支持块索引的变体，或者用支持块稀疏的注意力库。

注意：如果你的稀疏模式是固定且规整的（比如纯局部窗口），很多框架已经有现成的滑动窗口注意力实现，直接调就行，别自己造轮子。只有内容驱动的动态稀疏才需要自己写 kernel。

3.2 块大小的选择：64 还是 128

块大小是个需要权衡的参数。我做过一组对比实验，序列长度 N=4096，头维度 64：

块太小（32），稀疏粒度细但 kernel 调度开销大，加速比上不去；块太大（256），稀疏粒度粗，容易误砍有用连接，质量掉得明显。64 到 128 是比较舒服的区间。我最终选了 64，因为它在质量和速度之间平衡得最好，而且 64 对齐到常见的 tensor core tile 尺寸，硬件利用率高。

3.3 稀疏率的动态调整

固定稀疏率在不同层、不同去噪步上未必最优。DiT 是迭代去噪的，早期步（噪声大）和晚期步（接近收敛）对注意力的需求不一样。我的观察是：早期步全局结构还没成型，需要更稠密的注意力来建立整体布局；晚期步细节已经定了，稀疏一点影响不大。

所以我做了个简单的分层分步稀疏率调度：浅层和早期去噪步用较高密度（比如保留 30% 的块），深层和晚期步用较低密度（保留 10% 到 15%）。这个调度不需要训练，纯推理时控制，实现成本低，收益还挺明显——整体加速比能再提 10% 到 15%，质量几乎无损。

4. 实测中暴露的问题与排查过程

4.1 生成图出现块状伪影：定位到块边界处理

第一版跑通后，生成的图上有明显的网格状伪影，规律性很强，间距正好对应块大小。这个现象很典型，我一开始怀疑是打分机制的问题，排查了一圈才发现根因在块边界的处理。

具体是这样：块稀疏是按块决定算不算，但块与块之间的边界 token，其注意力需求可能跨越多个块。如果某个边界 token 真正需要的 key 恰好落在被跳过的块里，它的信息就断了，表现出来就是块边界处的不连续，累积成网格伪影。

解决办法有两个：一是对块边界做重叠处理（overlap），让相邻块共享一部分 token；二是在打分时对边界 token 给更高的保留权重。我用了第二种，改动小，效果够用。伪影基本消失。

4.2 文本遵循度下降：跨模态块被误砍

第二个问题是文本遵循度变差。给"一只戴帽子的猫"，生成的猫经常没帽子，或者帽子位置乱。这个问题的排查链路是这样的：

1. 先确认不是模型本身的问题——用稠密注意力跑同样的 prompt，帽子正常。说明是稀疏引入的。
2. 可视化注意力块的选择情况，发现描述"帽子"的文本 token 对应的图像区域 patch，在若干层里没有被选中参与注意力。
3. 根因清楚了：文本 token 少，池化打分时"帽子"这个 token 的信号被同块内其他 token 稀释，打分偏低，被 top-k 砍掉了。

修复方案就是前面 2.3 说的，跨模态那部分不做块稀疏，或者给文本相关的块一个保底保留名额。改完之后文本遵循度恢复到接近稠密水平。

4.3 加速比不及预期：kernel 启动开销

理论上稀疏率 85% 应该带来接近 6 倍的 FLOPs 下降，但实测端到端只快了 2 倍出头。这个落差一开始让我很困惑。

用 profiler 一查就明白了：被选中的块集合是动态的、不规则的，每个 batch、每个头、每一层的块索引都不一样，导致 kernel 启动频繁、调度开销大，而且不规则的内存访问让 tensor core 利用率下降。省下的 FLOPs 有一部分被这些开销吃掉了。

优化方向有几个：把块索引按 batch 内对齐（同一 batch 内不同样本用相近的稀疏模式，减少 kernel 种类）；把稀疏模式在若干去噪步之间复用（相邻步的注意力分布变化不大，没必要每步都重算块选择）；用更粗的粒度做 kernel 调度。我做了前两个，端到端加速比从 2.1x 提到了 2.8x。

提示：块稀疏的收益永远达不到理论 FLOPs 下降的比例，因为动态稀疏有调度成本。心里要有个预期，能到理论值的 40% 到 60% 就算不错了。

5. 一些值得记下来的经验

块稀疏 Attention 用下来，最大的体会是：稀疏模式的设计比 kernel 实现更决定成败。kernel 写得再好，如果选块逻辑把有用的连接砍了，质量就是上不去；反过来，选块逻辑合理，哪怕 kernel 朴素一点，整体也是赚的。

另外几点实操心得：

• 先做稠密基线，再上稀疏。没有稠密基线，你根本不知道质量掉了多少、加速了多少。我见过有人直接上稀疏，结果质量崩了都不知道是稀疏的锅还是模型本身的锅。
• 可视化是排查稀疏问题的第一工具。把每层选中的块画出来，很多问题一眼就能看出来，比盲猜快得多。
• 跨模态连接要保守。图像自注意力可以大胆稀疏，跨模态那部分能稠密就稠密，文本 token 本来就少，省不了多少，但砍了代价很大。
• 稀疏率调度是免费的午餐。不改模型、不重训，纯推理时控制，收益稳定，值得做。

后续如果要把这套东西推到视频 DiT（时序维度再叠一层），块稀疏的选块逻辑会更复杂，时序上的局部性和跨帧关联需要单独设计。这个我还在试，等有稳定结论再单独写一篇。