多模态 DiT 的推理成本真正跑过的人都知道瓶颈往往不在参数规模而在注意力那一步的显存和带宽。我最近在做一个图文联合生成的模型加速模型结构是典型的单流 DiT文本 token 和图像 patch 拼成一条序列送进 Transformer block。序列长度一上去标准 Attention 的 O(N²) 就开始吃人显存爆、延迟高、batch 上不去。试过直接砍 token、降分辨率效果掉得厉害也试过换更快的 kernel收益有限。最后落到块稀疏 AttentionBlock Sparse AttentionBSA这条路上才算把问题拆开看清楚了。这篇就把我在多模态 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 的 latentpatch size 取 2那就是 64×64 4096 个图像 token再加几十到几百个文本 token。序列长度 N 落在 4000 到 8000 这个区间是很常见的。标准自注意力的计算量是 O(N²·d)显存占用也是 O(N²)注意力矩阵本身。N4096 时单头注意力矩阵就是 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 的局部性很友好。全局 tokenGlobal Token选一部分 token比如文本 token、或者图像里均匀采样的 anchor patch作为全局节点所有 token 都跟它们做注意力它们也跟所有 token 做注意力。把两者结合就是局部窗口 全局 token的经典结构。文本 token 天然适合当全局 token因为它们数量少、语义密度高让所有图像 patch 都能看到文本跨模态对齐就有了通路。这个方案我在第一版里用了跑得通但效果有上限。问题在于全局 token 是固定的不管内容是什么每个图像 patch 都去关注全部文本 token。当文本很长比如一段详细描述时全局连接的开销又上来了而且很多文本 token 跟某个具体 patch 其实没关系属于无效计算。2.2 内容驱动的块选择让注意力自己决定看哪里第二版我换成了内容驱动的块选择。核心思路是先用一个轻量的打分机制估计每个 query 块对每个 key 块的重要性然后只保留 top-k 个 key 块参与真正的注意力计算。打分机制有几种常见做法方法原理优点缺点均值池化打分对 query/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块大小是个需要权衡的参数。我做过一组对比实验序列长度 N4096头维度 64块大小稀疏率相对稠密加速生成质量FID 相对变化32可到 90%1.6x0.8%64可到 85%2.1x0.3%128可到 75%2.4x1.5%256可到 60%2.2x4.2%块太小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 文本遵循度下降跨模态块被误砍第二个问题是文本遵循度变差。给一只戴帽子的猫生成的猫经常没帽子或者帽子位置乱。这个问题的排查链路是这样的先确认不是模型本身的问题——用稠密注意力跑同样的 prompt帽子正常。说明是稀疏引入的。可视化注意力块的选择情况发现描述帽子的文本 token 对应的图像区域 patch在若干层里没有被选中参与注意力。根因清楚了文本 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时序维度再叠一层块稀疏的选块逻辑会更复杂时序上的局部性和跨帧关联需要单独设计。这个我还在试等有稳定结论再单独写一篇。
