块稀疏Attention实战：破解多模态DiT长序列计算瓶颈

1. 为什么DiT里的Attention这么“贵”

1.1 扩散模型换骨架之后，计算瓶颈转移了

早年做扩散模型，绝大多数人脑子里默认的骨架是U-Net。U-Net里最核心的算子是卷积，参数量大，FLOPs也高，但好在卷积是“局部操作”，它的计算复杂度跟输入分辨率基本是线性关系，GPU跑起来非常舒服。到了DiT（Diffusion Transformer）出来之后，情况变了：图像被切成patch后拉成token序列，整个扩散模型变成了一个纯Transformer堆栈。Transformer里最核心的算子就从卷积变成了Attention。而Attention的复杂度是序列长度N的平方，也就是O(N²d)。这意味着，只要输入分辨率上调一档，或者视频帧数多几帧，Attention部分的燃烧速度比卷积恐怖得多。

我在实际调试DiT模型时有个很直观的感受：同样的总FLOPs预算下，U-Net是均匀分布的“体力活”，DiT是集中爆发的“烧卡活”。你观察profiling结果会发现，当序列长度超过几千之后，Attention算子的占比能冲到60%甚至70%以上，线性层反而退居其次。这时候你再去调MLP的hidden size，或者换GELU的近似方式，收益都已经很小了，真正的优化空间在Attention上。这就是块稀疏Attention被关注的直接原因：Attention贵，但我们有办法让它“少算一点”。

更扎心的是，多模态模型把这个瓶颈进一步放大了。图像token本身就已经不少了，再往序列里塞文本token、视频时空token，序列长度动不动就上万甚至上十万。在这个量级上，全量Attention不只是贵，而是物理上跑不动。我试过直接在一个长序列的多模态DiT上跑标准Attention，显存和算力双双爆掉，根本轮不到谈精度。所以做起多模态DiT，块稀疏Attention不是锦上添花的优化技巧，而是不得不面对的核心工程问题。

1.2 多模态让序列长度雪上加霜

多模态DiT的输入不是单一模态。视觉token、文本token、视频token、音频token……不同的模态会被分别编码，然后拼接成一条长序列进入Transformer。举个例子：一张512×512的图像，patch size为8×8，产生的token数是64×64=4096。如果再加一段16帧的视频，每帧同样512×512，视频token就是4096×16=65536。再加一段文本描述，几百个token。最后拼接起来的序列长度可能直逼七万。

标准Attention在这个长度上做一次矩阵乘法，score矩阵的规模是七万乘七万，光中间结果的显存就要几十个GB，这个数字在单卡上就是灾难。所以做多模态DiT的人，要么狠心降低分辨率，要么限制视频帧数，要么把模态拆开来做稀疏交互。这三种方案里，前两种是牺牲效果换显存，只有第三种是在想办法保留效果的同时压低计算量。块稀疏Attention就是第三种思路的技术化实现。

而且多模态序列一个天然的优势在于，它本身就带有明显的“块结构”。图像token按空间位置排布，视频token按时间-空间排布，文本token按句法语义聚合。不同模态内部有强局部相关性，跨模态之间的联系集中在语义摘要层面。这跟语言模型里那种长距离依赖为主的模式很不一样：多模态序列的信息结构天生适合做“块级别”的稀疏化，而不是均匀采样式的稀疏化。这也是我把块稀疏Attention和多模态DiT放在一起讨论的原因——它们非常契合。

1.3 块稀疏Attention到底要解决什么

一句话概括：在不显著损失模型表达能力的前提下，把Attention从O(N²)的计算量降下来，让长序列多模态DiT能够真正跑起来。这里的“块稀疏”指的是，你不去计算完整的N×N注意力矩阵，而是预先定义一批需要保留的注意力块，只有这些块参与QK矩阵乘法和softmax计算，其余位置直接跳过，用0填充。

为什么是“块”而不是单个“点”？背后的工程原因很现实：GPU对密集矩阵乘法做了极致优化，但对随机稀疏的单个点访问非常不友好。一个warp要读取的数据如果散落在内存各处，访存效率就会迅速恶化，计算核心反而在等待数据。如果改成块级别稀疏，例如8×8、16×16的连续小块，每个块内部的访存是连续的，GPU可以像处理密集矩阵一样高效处理这些块，只是跳过一部分块而已。实际做下来，块稀疏Attention在A100/H100上能把Attention的FLOPs有效降一半以上，而端到端速度相比全量Attention只损失很小一部分。

另外，块稀疏Attention也天然契合FlashAttention这类kernel融合方案。FlashAttention本身已经用tiling的方式避免了完整score矩阵的物化，块稀疏Attention只需要在tiling循环里增加一个“当前块是否需要计算”的判断，就能在几乎不增加kernel复杂度的情况下实现稀疏化。这一点在后面介绍实现时会详细展开。

2. 块稀疏Attention的核心设计思路

2.1 从均匀采样到结构稀疏：Mask定生死

实现块稀疏Attention的第一步，是确定一个注意力Mask。这个Mask是一个与Attention score矩阵形状完全一致的0/1矩阵，只不过是把逐点的0/1换成了“块级别”的0/1。每个块被mask为0，意味着这个块不参与任何计算；mask为1时，这个块的QK乘法、softmax和加权求和都会正常执行。

Mask的确定是整个方案的灵魂。如果Mask选得不好，计算结果和全量Attention差距过大，精度崩掉，稀疏化就没有意义；如果Mask选得太保守，全是1，稀疏化等于没做。常见的结构稀疏Mask有以下几种：

• 局部窗口Mask（Local Window）：每个query token只跟它前后W个token的key做注意力。适合图像和视频，因为空间上相邻的token相关性最强。
• 分块对角Mask（Block Diagonal）：序列按模态或者语义分成几个大段，段内全连接，段间仅保留少量连接。适合多模态拼接场景，减少跨模态冗余计算。
• 跨步采样Mask（Strided）：每隔S个token采样一个key，保证一定的全局感知能力，计算量线性下降。
• 全局+局部混合Mask：一部分token（如图像中的CLS token、文本中的summary token）作为“锚点”，这些锚点与所有token都计算注意力，其余token只在局部窗口内计算。这是实践中效果与效率最平衡的方案。

我个人的经验是，不要指望靠某一种Mask打天下。多模态DiT里不同模态的统计特性差异很大，图像token的局部性明显，文本token的长距离依赖更突出，最理想的Mask往往是几种基本结构的组合。先跑一版全部用全局Attention的小模型，分析Attention矩阵的能量分布，再根据热力图去设计稀疏Mask，比上来就套一个固定稀疏模板要靠谱得多。

画一个粗糙的类比：全量Attention相当于全班同学互相认识，每个人都跟其他人打招呼；块稀疏Attention相当于先按座位分组，组内熟络，组间只找几个代表交流。绝大多数情况下，你不需要全班每个人彼此认识，只需要把信息传递路径保留下来，效果就差不多。

2.2 为什么稀疏粒度必须是“块”

严格来说，稀疏Attention可以做到任意细粒度：每个单独的(query, key)位置都能单独控制是否计算。这种细粒度稀疏在理论上最优，但工程实现上几乎行不通。原因有三点：

第一，GPU的访存是按cache line和寄存器批量进行的，一个block处理16×16的连续数据比处理16个分散数据快得多。随机点稀疏的kernel，访存可能比计算还贵，稀疏率上去了速度反而下降。这不是GPU不行，而是访存体系决定的物理规律。

第二，实现复杂度。逐点稀疏需要为每个元素存储坐标索引，索引本身的存储和计算开销不可忽略。块稀疏只需要记录“块坐标”，索引量减少两个数量级，mask矩阵也可以直接用位图压缩，实现简单很多。

第三，优化生态。PyTorch的torch.nn.functional.scaled_dot_product_attention已经内置了FlashAttention和Memory-Efficient Attention的实现，这些kernel天然就是按块做tiling的。你自己实现块稀疏Attention，只需要在tiling循环里做个判断，就能复用FlashAttention成熟的高效kernel骨架。这个路径比从零写一个逐点稀疏Attention要被验证得多，踩坑成本低很多。

所以“块”这个选择不是理论最优，而是工程最优。它牺牲了一点理论上的稀疏灵活性，换来了实现简单、访存友好、能与现有kernel库无缝结合的巨大好处。

2.3 Mask生成与稀疏率估算

生成Mask的核心流程是：计算总块数，标记需要保留的块，存储成稀疏索引或位图。假设序列长度N=16384，块大小B=64，那么总的块数是：

num_blocks = N / B = 16384 / 64 = 256

Mask矩阵的维度是256×256个块，每个块标记为0或1。如果用位图存储，一个bit代表一个块，整个mask只需要256 * 256 / 8 = 8192字节，约8KB，小到可以放进常量内存甚至直接嵌入kernel代码。如果用坐标列表存储，也只需要记录非零块的(row_idx, col_idx)，在稀疏率较高时存储开销更小。

稀疏率（sparsity ratio）定义为：

sparsity = 1 - nonzero_blocks / total_blocks

比如256×256的块矩阵里只保留32768个块（约50%保留率），稀疏率就是50%。对应的Attention FLOPs大约降到全量的50%。这个数据可以直接用来预估训练速度提升和显存节省。我一般会在设计Mask的时候先跑一个公式估算：

estimated_FLOPs_saved = (1 - sparsity) * 2 * N^2 * d_head * num_heads

这个数字虽然粗糙，但能帮助我们快速判断一个Mask方案有没有价值。如果算出来的FLOPs节省不到20%，那说明这个Mask太密了，不值得付出稀疏化的实现成本；如果节省超过80%，就要警惕信息丢失，先跑小规模实验验证精度。

3. 多模态DiT中的稀疏Mask设计实践

3.1 多模态序列的混合排布策略

多模态DiT里，序列不是简单把各模态token拼在一起就完事，排布方式直接决定了Mask设计。

我常用的一种排布方式如下：文本token放在最前面，作为全局摘要；图像token按行优先顺序排列在中间；视频token按时间-空间展开放在最后。这样的排布有几个好处：文本在序列最前面，它的token能够作为所有后续token的“全局引导”；图像和视频本身内部相关性按空间和时间组织，容易配局部窗口Mask；跨模态交互主要发生在文本token与视觉token之间，用锚点Mask就可以覆盖。

跨模态交互还有另一种做法：在序列中周期性插入模态摘要token。比如每16个图像token插入一个CLS token，一段视频每8帧插入一个CLS token，这些CLS token负责与全局交互，其他token只跟相邻token和CLS token交互。这种“锚点+局部”的排布方式在视频多模态任务中效果特别好，因为视频的时间维度天然存在关键帧概念，关键帧token承担起和文本交叉注意力的职责，计算量大幅下降。

这里要特别提醒一点：Mask的块边界必须与token排布对齐。如果你的块大小是16，但图像的一个patch行有20个token，那块的边界就会割裂局部窗口的连续性。我在第一次做多模态稀疏attention时没注意这个问题，Mask热力图看起来完全正常，但局部窗口的覆盖率莫名下降，精度从0.82掉到0.77。后来仔细排查才发现是token排布和块大小错位导致部分关联被切断。解决方法很简单，在设计token排布时就把块大小考虑进去，让图像行宽、视频单帧token数都能被块大小整除。

3.2 一个可复用的Mask生成流程

基于上面的经验，我整理了一套Mask生成流程。以多模态DiT为例，假设序列包含文本段（T_len个token）、图像段（I_len个token）、视频段（V_len个token），块大小B=16：

第一步，确定各段的块区间。文本段对应的块范围是[0, T_len/B)，图像段是[T_len/B, (T_len+I_len)/B)，视频段是[(T_len+I_len)/B, total/B)。

第二步，为每种交互关系分配稀疏策略：

• 文本内部：全量Attention，保留所有块。文本token数量通常不多，全量计算成本可接受。
• 图像内部：局部窗口Attention，每个query块只与附近的key块交互，窗口半径W=8个块。
• 视频内部：局部窗口Attention，窗口按时间维度展开，保持时序一致性。
• 文本到图像/视频：只让文本块的“summary block”与视觉块全连接，其他文本块只与summary block交互。
• 视觉到文本：与文本到视觉对称，也做稀疏处理。

第三步，将上述策略落到块矩阵上，生成0/1矩阵，再用torch.sparse_coo_tensor或者位图形式存储。

整个流程可以用一个简单的伪码表示：

def build_multimodal_block_mask(segments, block_size, window_size): total_tokens = sum(seg.values()) num_blocks = total_tokens // block_size mask = torch.zeros(num_blocks, num_blocks, dtype=torch.bool) # 遍历每个段的块区间 # 根据段类型决定稀疏规则 # 填充mask return mask

当然，实际代码要比这个复杂，涉及段边界处理和窗口边界截断，但整体思路就是这么个思路。

3.3 与FlashAttention的结合：kernel内部怎么吃Mask

块稀疏Attention最终的执行依赖kernel。最简单的实现是在PyTorch里用torch.masked_fill把mask加进去，但这样score矩阵还是会被完整计算出来，只是结果被置零了，计算量一点没省。正确做法是让kernel在tiling循环里直接跳过全零块。

FlashAttention的tiling循环大致是这样运作的：把Q按block_M分块，把K按block_N分块，逐个加载Q_block和K_block，计算score_block，做online softmax，再与V_block加权求和。要加稀疏Mask，只需要在加载K_block之前检查mask[row_block, col_block]是否为0，如果是0，整个block直接跳过，不加载K和V的数据，也不做任何矩阵乘法。

在Triton里实现这个逻辑非常简单：

@triton.jit def flash_attn_sparse_kernel(Q, K, V, Mask, ...): pid_m = tl.program_id(0) pid_n = tl.program_id(1) if tl.load(Mask + pid_m * num_col_blocks + pid_n) == 0: return # 当前块被mask掉，跳过整个block # 正常加载Q_block、K_block、V_block，做flash attention计算

注意上面代码里的return：Triton kernel里提前返回，程序对这个block就不做任何操作，GPU资源立刻释放给其他block。这是块稀疏Attention在kernel层能真正带来加速的关键所在。如果只是在PyTorch层面用masked_fill，那只是掩耳盗铃，FLOPs没有减少，显存也没有省。

我在实际项目里用的是Triton + FlashAttention的思路。FlashAttention原版kernel本身就是按块计算的，我在它的基础上加一个Mask参数，替换原来无条件加载K/V的流程。这个改造很轻量，但带来的收益非常直接，跑多模态DiT的forward阶段，Attention部分的时间从原来的42%降到21%，端到端训练吞吐提升接近1.6倍。这个数字在不同模型上会有波动，但整体趋势是一致的。

另外，如果不想手写Triton kernel，可以先用xformers的BlockSparseAttention顶上。它的API设计直观，接受coo格式的稀疏索引，内部已经实现了mask-aware的block sparse kernel。不过实测下来，xformers的块稀疏kernel在自定义mask切换上灵活度不如自己直接改FlashAttention，如果你的Mask方案经常调整，我还是建议用Triton自己改。

4. 实操记录：在多模态DiT上替换Attention算子

4.1 基础替换实验设计

前面讲了不少原理和设计，实际操作时第一步永远是“先能跑通，再谈效果”。我建议按下面这个顺序推进：

1. 用一个小规模多模态DiT模型跑一版全量Attention的baseline，记录训练损失收敛曲线和端到端速度。
2. 实现一个最简单的局部窗口Mask（块大小16，窗口半径8），替换Attention，看loss曲线和baseline的偏离程度。
3. 如果局部窗口效果可以接受，再逐步加入文本-视觉间的锚点连接、视频时序建模等复杂结构。
4. 对每个Mask方案都记录“精度-速度-显存”三维指标，做成对照表。

这个顺序的关键在于，每一步改动都很小，出了问题容易定位。我见过不少同学一口气上了个非常复杂的多模态Mask，结果精度崩了，完全不知道是哪个模态的稀疏策略出了问题，最后只能全部回退。

以图像+文本二模态DiT为例，我分享一组实际替换时的数据。模型patch size=4，图像分辨率256×256，序列长度N=4096，batch size=32，8卡A100。全量Attention的GPU显存占用约38GB，单步耗时0.91秒。换成块稀疏Attention（块大小16，图像内部窗口半径8，文本内部全量，文本到图像锚点连接）之后，显存降到22GB，单步耗时0.58秒。精度方面，在相同训练步数下，FID指标从全量Attention的9.8降到10.3，损失非常小，但吞吐量提升约1.5倍，显存压力大幅缓解。

4.2 关键指标与统计方式

替换Attention之后，除了看loss和下游任务指标，还要监控两个容易被忽略的指标：实际有效FLOPs和Attention矩阵的熵。

实际有效FLOPs可以通过统计非零块的占比来计算。比如总的块数是256×256=65536，实际参与计算的块数是32768，那么非零占比就是50%。把这个数值和kernel实际运行时间做对比，可以判断kernel是不是真的按预期跳过了全零块。我遇到过一种很坑的情况：mask生成正确，但kernel实现时忘了检查mask，导致所有块都被计算，非零占比50%但实际速度没有任何提升。这个问题在profiling里一眼就能看出来——kernel运行时间和全量Attention几乎一致。

Attention矩阵的熵则用来衡量稀疏化是否破坏了信息多样性。如果某个块在Mask里被保留，但计算出来的attention weight几乎全是均匀分布（熵接近最大值），说明这个块贡献的有效信息很小，可以考虑从Mask里删掉；反过来，如果Mask里删掉了某个块，但训练过程中发现attention weight在相应区域出现了尖峰，说明这块信息很重要，应该加回来。这个方法特别适合用来“裁剪”Mask：先全量训练一个小模型，统计注意力权重的分布，砍掉那些持续低权重的块，保留高权重区域，让稀疏Mask逐步逼近最优。我管这个方法叫“注意力热力图裁剪法”，听起来挺玄乎，其实就是用数据驱动的方式设计Mask，而不是拍脑袋画。

4.3 训练稳定性与收敛速度的变化

块稀疏Attention对训练行为的影响，我观察到两个现象：

第一，前期收敛反而变快。这个听起来有点反直觉，但实际上是因为稀疏Mask相当于给Attention加了一个先验正则。全量Attention里大量低价值连接不仅浪费计算，还会引入噪声，让模型在早期被无关信息干扰。稀疏Mask把这些噪声连接提前剪掉，模型前端更专注在真正相关的token对上，loss下降在早期反而更陡。

第二，后期精度上限会有轻微损失。不管Mask设计得多好，丢掉的连接里总有一些是模型的能力支撑点。我的经验是，块稀疏Attention的精度损失通常在0.2-0.5个点以内（具体指标因任务而异），而且可以通过加一个小的“全局锚点块”来弥补。例如让序列最前端的几个token与所有token保持全量Attention，这个锚点块相当于一个“信息总览”，它能让丢失的远程信息有一个汇聚和重分发通道。成本是仅增加了少量FLOPs，但精度损失基本上能拉回来一多半。

我实际项目中最后采用的方案就是“锚点+局部窗口”的复合Mask：序列前8个token作为全局锚点，与所有token全量Attention；其余token在局部窗口内与附近8个块交互。这个方案在250M参数的多模态DiT上，端到端训练速度比全量快1.4倍，FID只增加了0.2，已经被我当成了默认配置。

5. 常见问题与排查技巧实录

5.1 问题速查表

块稀疏Attention在落地过程中，问题主要集中在Mask生成、kernel实现、效果调优三个层面。我整理了一张速查表：

这一节没有数学推导，全是我踩过的坑。每一行背后都有真实案例支撑，特别是“kernel没有真正跳过mask块”这个坑，我在几次优化中反复踩到，每次都是profiling才暴露。

5.2 三个容易忽略的实现细节

第一个细节：mask不该用Python循环逐块生成，要用向量化操作。我第一次用双重for循环构造mask，序列长度8192、块大小16，总块数512×512，Python循环跑了接近30秒。把循环改成布尔矩阵拼接后，耗时降到0.2秒以下。虽然这个mask生成只是整个训练流程的一小段，但如果你每步都要动态生成Mask，这个开销会快速累积。

第二个细节：softmax的计算必须只统计非掩码块。FlashAttention的online softmax本身就假设被跳过块对统计量没有贡献，但如果你在mask块对应的位置填充了一个极大的负数值然后参与softmax，很容易出现数值问题。正确做法是在kernel里直接跳过掩码块，让它完全不参与统计，而不是用一个特殊值代替。

第三个细节：多模态batch里不同样本的稀疏Mask可能不一样。文本长度不同、视频帧数不同，会导致batch中每个样本的序列长度和段边界不同。处理这种batch不整齐的情况，常见做法是按长度padding到batch内最大序列长度，再用key_padding_mask跳过padding部分。块稀疏Attention的mask也要跟padding方案配套，否则padding块会参与Attention计算，污染信息。

5.3 性能分析的几个建议

做完块稀疏Attention替换，别急着说“快了”，先用profiling工具量一下。我会按下面这个流程分析：

1. 用PyTorch Profiler分别记录全量Attention和块稀疏Attention的kernel耗时，确认稀疏kernel在总耗时中的占比确实下降了。
2. 对比FLOPs理论节省值和实际端到端提速，两者之间差距如果超过50%，说明kernel效率还不行，可能是分支判断太多导致GPU warp分支发散，或者访存模式不够连续。
3. 用Nsight Compute看GPU利用率（SM occupancy）和内存吞吐。块稀疏Attention要关注mask判断是否导致warp load不均衡，比如有些block相邻的mask全是0，有些全是1，不同block计算量差异大，GPU调度会拖慢整体。
4. 最后再看端到端收敛曲线，确保稀疏化后精度没有恶化。

这套流程看起来繁琐，但真能帮你避免“感觉快了但不知道为什么快”的盲目状态。优化Attention这种底层算子，一切以profiling数据为准，别靠直觉拍板。

6. 关于这套方案的一点个人体会

做完这个项目，我最深的体会是：块稀疏Attention的本质不是“减少信息”，而是“重新组织注意力计算的优先级”。全量Attention把计算资源平均分配给所有token对，但真实数据里绝大多数token对的关联性其实非常弱。稀疏Mask要做的就是把资源集中到真正有信息价值的连接上，这跟人类阅读文章时先扫标题、再看正文、最后盯着关键段落的逻辑是一样的。

如果你刚接触这个方向，我建议先别急着写kernel，先用全量Attention跑一个小模型，把注意力分布可视化一遍，你会非常直观地看到哪些区域存在大量近零权重。这些近零区域就是块稀疏Attention的原生利润空间。把这块空间榨干，你的多模态DiT速度和显存优化就已经成功了一大半。

另外，块稀疏Attention与FlashAttention、Triton这一类工具的组合已经非常成熟，我自己从零写到的第一版能跑通，只花了两到三天。真正花时间的不是实现kernel，而是设计一套适合你数据模态分布的Mask方案。这个方向很值得投入，因为多模态DiT序列越来越长是大趋势，块稀疏Attention是当前少数几种能在效果和效率之间取得平衡的实用方案。