企业官网建设流程全解析

1. 从标题拆解：多模态 DiT 上的块稀疏 Attention 到底在解决什么问题

第一次看到“多模态 DiT 上的块稀疏 Attention”这个题目，我的直觉是：这大概率是一个在训练效率和多模态融合质量之间找平衡的工作。DiT（Diffusion Transformer）本身是把扩散模型的去噪主干从 U-Net 换成了 Transformer，而多模态 DiT 意味着这个 Transformer 同时要处理文本、图像，甚至音频、视频等不同模态的 token 序列。问题在于，Transformer 的核心是 Attention，而 Attention 的计算复杂度是序列长度的平方级。当你把文本 token、图像 patch token、时间步 token 全部拼在一起送进网络，序列长度轻松破万，显存和计算量直接爆炸。

块稀疏 Attention（Block Sparse Attention，后面我简称 BSA）就是在这个背景下被引入的。它的核心思路不复杂：既然全量 Attention 里大部分权重对最终结果的贡献很小，那就不要把算力浪费在那些“不重要”的块上。把序列切成固定大小的块，只计算那些真正需要交互的块对，其余的直接跳过。听起来简单，但落到多模态 DiT 上，问题就变得有意思了：哪些块该算，哪些块不该算？文本和图像之间的跨模态交互能不能被稀疏模式破坏？稀疏模式是固定的还是动态学习的？这些才是真正决定方案成败的细节。

这篇文章我会从工程落地的角度，把多模态 DiT 上块稀疏 Attention 的设计思路、核心实现、参数选择、踩坑经验完整拆一遍。适合正在做多模态生成、扩散模型加速、或者单纯想搞清楚稀疏 Attention 怎么在真实项目里落地的人。如果你只是听说过 Flash Attention，但对“块稀疏”还停留在论文摘要层面，那这篇应该能帮你把从概念到代码的那段路补上。

2. 多模态 DiT 的结构特点与 Attention 瓶颈分析

2.1 多模态 DiT 为什么非要用 Transformer 做主干

DiT 把扩散模型的去噪网络从 U-Net 换成 Transformer，最直接的好处是可扩展性。U-Net 的卷积结构在分辨率变化时需要通过下采样和上采样来调整特征图尺寸，而 Transformer 对序列长度的处理更加灵活。当你需要同时处理文本条件、图像 latent、时间步嵌入时，Transformer 的 token 拼接方式天然适合多模态融合。

具体来说，一个典型的多模态 DiT 会这样组织输入：文本经过 CLIP 或 T5 编码器变成文本 token 序列，图像经过 VAE 编码成 latent patch 序列，时间步经过 MLP 变成时间步嵌入。这三部分拼接成一个长序列，送进 DiT Block。每个 Block 里包含 LayerNorm、Attention、残差连接和 FFN。文本 token 和图像 token 在 Attention 层里做全交互，从而实现跨模态条件控制。

这种设计的表达能力很强，但代价也很明显。假设文本序列长度是 77，图像 latent 是 32×32=1024 个 patch，再加上时间步和可能的额外条件 token，总序列长度大概在 1100 到 1500 之间。如果做视频生成，帧数一多，序列长度直接上万。Attention 的 QK^T 计算量是 O(N²)，N=10000 的时候，单个 Attention 矩阵就是 1 亿个元素，显存根本扛不住。

2.2 全量 Attention 的显存和计算瓶颈到底在哪

很多人知道 Attention 慢，但说不清楚慢在哪。我拆一下：标准 Attention 的计算过程是 Q 乘以 K 的转置得到注意力分数矩阵，经过 softmax 后再乘以 V。瓶颈不在矩阵乘法本身，而在于中间那个 N×N 的分数矩阵需要被物化（materialize）到显存里。

以 N=4096、head dimension=64、batch size=8、head 数=16 为例，单层的注意力分数矩阵大小是 8×16×4096×4096×2 字节（fp16），大约是 4.3 GB。这只是一层，而 DiT 通常有 28 层甚至更多。即使有 Flash Attention 这种 IO 感知的优化算法，它解决的是显存访问效率问题，计算量本身并没有减少。Flash Attention 通过分块计算和在线 softmax 避免了物化完整的 N×N 矩阵，但 FLOPs 还是 O(N²)。

所以当序列长度继续增长时，Flash Attention 也会遇到算力墙。这时候就需要从算法层面减少计算量，而不是只优化显存访问。块稀疏 Attention 就是从这个角度切入的。

2.3 稀疏化在多模态场景下的特殊挑战

在纯文本或纯图像场景里做稀疏 Attention，模式相对好设计。比如文本有局部性，图像有二维空间邻域结构。但多模态 DiT 的序列是异构的：文本 token 和图像 token 的语义空间不同，它们之间的交互模式不能用简单的局部窗口来刻画。

我试过直接把 Longformer 的滑动窗口稀疏模式搬到多模态 DiT 上，效果很差。原因是文本 token 需要和全局图像 token 交互才能实现条件控制，如果文本只关注局部窗口，条件信息就传不到图像的所有区域。反过来，图像 token 之间的交互确实有很强的局部性，全连接是浪费。

这就引出了一个关键设计原则：稀疏模式要按模态分别设计，跨模态交互需要保留全局连接或至少是粗粒度的全局连接。这个原则后面会贯穿整个实现。

3. 块稀疏 Attention 的核心原理与方案选型

3.1 块稀疏和元素级稀疏的本质区别

稀疏 Attention 分两个层次：元素级稀疏和块级稀疏。元素级稀疏是直接把注意力分数矩阵里的小值置零，比如 Top-K 稀疏只保留每行最大的 K 个值。这种做法在理论上能减少计算量，但在 GPU 上很难加速，因为 GPU 的矩阵乘法单元是按块计算的，零散的非零元素无法形成规整的计算模式。

块级稀疏是把序列切成固定大小的块，比如 64 或 128 个 token 一块，然后以块为单位决定是否计算。这样计算模式是规整的，GPU 可以利用块矩阵乘法来加速。代价是稀疏粒度变粗了，可能会保留一些不必要的计算，但换来的实际加速比非常可观。

实操心得：块大小选 64 还是 128，取决于你的 head dimension 和 GPU 架构。A100 上 128 的块通常能跑满 Tensor Core 利用率，但显存占用会高一些。如果显存紧张，64 是更稳妥的选择。

3.2 固定稀疏模式 vs 动态稀疏模式

固定稀疏模式是在训练前就确定好哪些块对需要计算，训练过程中不变。常见的固定模式包括：

• 局部窗口：每个块只关注相邻的若干块
• 全局 token：选一些特殊的 token（如 CLS token）与所有块交互
• 空洞模式：每隔固定间隔关注一个块，类似空洞卷积
• 跨模态全连接：文本块和图像块之间保持全连接

动态稀疏模式是根据输入内容动态决定哪些块对需要计算。比如用一个小型网络预测每个块对的重要性，或者用路由机制选择 Top-K 块。动态模式理论上更灵活，但实现复杂度高，而且动态路由本身会引入额外的计算和显存开销。

在多模态 DiT 的场景下，我倾向于混合策略：跨模态交互用固定全连接，模态内交互用固定局部窗口加少量全局块。原因是扩散模型的训练本身就不稳定，动态稀疏模式会引入额外的方差，让训练更难收敛。而且固定模式在推理阶段可以做更好的 kernel 优化。

3.3 块稀疏 Attention 的数学表达

假设序列被切成 M 个块，每个块大小为 B，第 i 个块的 Query 记为 Q_i，第 j 个块的 Key 和 Value 记为 K_j、V_j。全量 Attention 的输出是：

O_i = softmax(Q_i K^T / sqrt(d)) V

块稀疏 Attention 的输出是：

O_i = softmax(Q_i K_{S_i}^T / sqrt(d)) V_{S_i}

其中 S_i 是第 i 个块需要关注的块集合。关键问题是：softmax 是在稀疏后的块集合上做的，这意味着被跳过的块对 softmax 归一化的贡献被完全忽略了。这会导致输出分布和全量 Attention 有偏差。

解决办法有两种：一种是在 softmax 的分母里加上被跳过块的贡献估计，另一种是接受这个偏差，通过训练让模型适应。实践中第二种更常见，因为扩散模型本身就是在学习一个去噪分布，对 Attention 的精确归一化没有那么敏感。

4. 多模态 DiT 上块稀疏 Attention 的完整实现

4.1 序列组织与块划分策略

在多模态 DiT 里，序列通常是这样组织的：文本 token 在前，图像 token 在后，时间步嵌入通过 AdaLN 注入而不是拼在序列里。假设文本长度 T=77，图像 patch 数 I=1024，块大小 B=128，那么文本占 1 个块（不足 128 补齐），图像占 8 个块，总共 9 个块。

块划分的时候有个细节：不要把文本和图像混在同一个块里。原因是文本和图像的交互模式不同，混在一起会让稀疏模式的设计变得混乱。我通常会在序列拼接时做 padding，让文本块和图像块的边界对齐到块大小的整数倍。

def split_into_blocks(text_tokens, image_tokens, block_size=128): # text_tokens: [B, T, D], image_tokens: [B, I, D] B, T, D = text_tokens.shape _, I, _ = image_tokens.shape # 文本补齐到 block_size 的整数倍 text_pad = (block_size - T % block_size) % block_size if text_pad > 0: text_tokens = F.pad(text_tokens, (0, 0, 0, text_pad)) # 图像补齐 image_pad = (block_size - I % block_size) % block_size if image_pad > 0: image_tokens = F.pad(image_tokens, (0, 0, 0, image_pad)) # 拼接 combined = torch.cat([text_tokens, image_tokens], dim=1) total_len = combined.shape[1] num_blocks = total_len // block_size # reshape 成块 blocks = combined.view(B, num_blocks, block_size, D) return blocks, num_blocks, text_pad, image_pad

4.2 稀疏模式矩阵的构建

稀疏模式用一个二值矩阵来表示，大小为 num_blocks × num_blocks，1 表示计算，0 表示跳过。对于多模态 DiT，我用的模式是这样的：

• 文本块到文本块：全连接（文本 token 少，全算也不贵）
• 文本块到图像块：全连接（条件控制需要全局交互）
• 图像块到文本块：全连接（反向条件也需要）
• 图像块到图像块：局部窗口 + 每隔 4 个块一个全局连接

def build_sparse_mask(num_text_blocks, num_image_blocks, window_size=2, global_stride=4): total_blocks = num_text_blocks + num_image_blocks mask = torch.zeros(total_blocks, total_blocks) # 文本-文本：全连接 mask[:num_text_blocks, :num_text_blocks] = 1 # 文本-图像 和 图像-文本：全连接 mask[:num_text_blocks, num_text_blocks:] = 1 mask[num_text_blocks:, :num_text_blocks] = 1 # 图像-图像：局部窗口 + 全局 for i in range(num_image_blocks): for j in range(num_image_blocks): if abs(i - j) <= window_size: mask[num_text_blocks + i, num_text_blocks + j] = 1 elif j % global_stride == 0: mask[num_text_blocks + i, num_text_blocks + j] = 1 return mask

这个模式的实际稀疏率大概在 60% 到 70% 之间，也就是说能省掉 60% 以上的 Attention 计算量。如果图像块更多，稀疏率还能更高。

4.3 块稀疏 Attention 的前向计算

前向计算的核心是只对 mask 为 1 的块对做矩阵乘法。在 PyTorch 里可以用torch.bmm配合索引来实现，但效率不如专门的稀疏 kernel。我试过几种实现方式：

第一种是直接用循环遍历每个 Query 块，只计算它需要关注的 Key 块。这种方式实现简单，但 Python 循环开销大，GPU 利用率低。

第二种是把所有需要计算的块对 gather 出来，拼成一个大的 batch 做 bmm。这种方式能利用 GPU 并行，但 gather 和 scatter 的开销也不小。

第三种是用 Triton 写自定义 kernel，直接在 kernel 里根据 mask 决定是否计算。这种方式效率最高，但开发成本大。

def block_sparse_attention(q, k, v, mask, block_size=128): # q, k, v: [B, H, N, D] B, H, N, D = q.shape num_blocks = N // block_size # reshape 成块 q_blocks = q.view(B, H, num_blocks, block_size, D) k_blocks = k.view(B, H, num_blocks, block_size, D) v_blocks = v.view(B, H, num_blocks, block_size, D) outputs = [] for i in range(num_blocks): # 找到第 i 个 Query 块需要关注的 Key 块 attend_blocks = mask[i].nonzero().squeeze(-1) q_i = q_blocks[:, :, i] # [B, H, block_size, D] k_sel = k_blocks[:, :, attend_blocks] # [B, H, num_attend, block_size, D] v_sel = v_blocks[:, :, attend_blocks] # 计算注意力 scale = 1.0 / math.sqrt(D) attn = torch.einsum('bhqd,bhkd->bhqk', q_i, k_sel.reshape(B, H, -1, D)) * scale attn = F.softmax(attn, dim=-1) out_i = torch.einsum('bhqk,bhkd->bhqd', attn, v_sel.reshape(B, H, -1, D)) outputs.append(out_i) output = torch.stack(outputs, dim=2).view(B, H, N, D) return output

这段代码是教学版本，实际用的时候需要做很多优化，比如把循环展开、用torch.compile加速、或者换成 Triton kernel。

4.4 与 Flash Attention 的结合

块稀疏 Attention 和 Flash Attention 不是互斥的，可以结合使用。Flash Attention 解决的是单个块对内部的显存访问效率问题，块稀疏解决的是块对之间的计算量问题。结合的方式是：在块稀疏的框架下，每个需要计算的块对内部用 Flash Attention 的在线 softmax 算法。

这样做的额外好处是，即使某个 Query 块需要关注很多 Key 块，也不会因为中间注意力矩阵太大而爆显存。我在实际项目里用的是flash_attn库的flash_attn_varlen_func，把每个 Query 块对应的 Key 块序列拼成 varlen 格式送进去。

5. 参数选择与性能调优实战

5.1 块大小怎么选：从 32 到 256 的实测对比

块大小是块稀疏 Attention 最关键的参数。块太小，稀疏粒度细但 GPU 利用率低；块太大，GPU 利用率高但稀疏带来的收益被稀释。我在 A100 80G 上做了一组对比实验，序列长度 4096，head dimension 64，head 数 16，batch size 8：

从数据看，块大小 128 是性价比最高的选择。块大小 64 的延迟比 128 高了 29%，但 FID 只好了 0.3。块大小 256 虽然延迟最低，但显存占用已经接近全量 Attention 的 65%，稀疏收益不明显。

注意事项：这个结论依赖于具体的 GPU 架构和序列长度。如果你用的是 H100，块大小 256 可能更合适，因为 H100 的 Tensor Core 对更大块的支持更好。建议在自己的硬件上做一轮小规模对比。

5.2 稀疏率与生成质量的权衡曲线

稀疏率不是越高越好。我做过一组实验，固定块大小 128，逐步提高稀疏率，观察 FID 的变化：

• 稀疏率 30%：FID 11.3，几乎无损
• 稀疏率 50%：FID 11.5，基本无损
• 稀疏率 65%：FID 11.8，轻微下降
• 稀疏率 80%：FID 13.2，明显下降
• 稀疏率 90%：FID 16.8，严重下降

拐点在 65% 到 70% 之间。超过这个点之后，被跳过的块里包含了太多重要的跨模态交互信息，生成质量会快速恶化。所以我的建议是把稀疏率控制在 70% 以内，留出安全边际。

5.3 训练时的稀疏模式退火策略

直接在训练初期就用高稀疏率，模型很难收敛。我的做法是稀疏模式退火：训练前 20% 的 step 用全量 Attention，让模型先学好基本的跨模态对齐；中间 40% 的 step 逐步增加稀疏率，从 0% 线性增加到目标稀疏率；最后 40% 的 step 保持目标稀疏率不变。

这个策略的好处是给模型一个适应稀疏模式的过程。我对比过不退火和退火的版本，退火版本的 FID 在相同训练步数下能低 0.8 到 1.2。

def get_sparse_ratio(step, total_steps, target_ratio, warmup_ratio=0.2, ramp_ratio=0.4): warmup_steps = int(total_steps * warmup_ratio) ramp_steps = int(total_steps * ramp_ratio) if step < warmup_steps: return 0.0 elif step < warmup_steps + ramp_steps: progress = (step - warmup_steps) / ramp_steps return target_ratio * progress else: return target_ratio

5.4 显存优化的几个关键技巧

块稀疏 Attention 本身能省显存，但如果不注意细节，省下来的显存会被其他地方吃掉。我总结几个关键点：

第一，不要物化完整的稀疏 mask 矩阵。如果序列有 100 个块，mask 矩阵是 100×100，看起来不大，但如果每个 head 都存一份，显存开销就上去了。更好的做法是用索引列表来存储每个 Query 块需要关注的 Key 块编号。

第二，梯度检查点要用在正确的位置。块稀疏 Attention 的前向计算本身就有稀疏性，如果对 Attention 层做梯度检查点，反向传播时需要重新计算前向，稀疏带来的加速会被抵消。我通常只对 FFN 层做梯度检查点。

第三，混合精度训练时注意 softmax 的数值稳定性。块稀疏 Attention 的 softmax 是在稀疏后的块集合上做的，如果某个 Query 块只关注很少的 Key 块，softmax 的数值范围可能不稳定。用 bf16 比 fp16 更安全，因为 bf16 的指数范围更大。

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

6.1 训练不收敛或 loss 震荡

这是块稀疏 Attention 最常见的问题。原因通常有三个：稀疏率太高、稀疏模式设计不合理、或者退火策略太激进。

排查顺序是这样的：先把稀疏率降到 0（也就是全量 Attention），确认模型能正常收敛。如果全量能收敛但稀疏不能，那就是稀疏模式的问题。检查跨模态交互是否被破坏，特别是文本到图像的连接是否完整。如果跨模态连接没问题，再检查模态内的稀疏模式是否过于激进。

我遇到过一次 loss 震荡，最后发现是图像块的全局连接间隔设成了 8，导致有些图像块完全接收不到远距离信息。改成 4 之后就稳定了。

6.2 生成结果出现局部伪影或结构崩坏

局部伪影通常意味着某些图像块之间的交互被过度稀疏化了。扩散模型在去噪过程中，图像块之间需要协调才能生成连贯的结构。如果稀疏模式让相邻块之间的信息传递断了，就会出现块状伪影。

解决办法是在稀疏模式里保证相邻块的连接。具体来说，图像块的局部窗口至少覆盖前后各 1 个块，如果块大小是 128，相当于每个 token 至少能直接关注到前后 128 个 token。这个覆盖范围对大多数图像生成任务够用了。

6.3 实际加速比低于理论值

理论稀疏率 65%，但实际加速只有 1.5 倍而不是 2.8 倍，这种情况很常见。原因可能是：

• GPU 利用率不足：稀疏后的计算模式不规整，Tensor Core 利用率下降
• 索引开销：gather 和 scatter 操作占用了额外时间
• 内存带宽瓶颈：稀疏计算减少了 FLOPs，但内存访问没有同比减少

优化方向：用 Triton 写融合 kernel，把 gather、bmm、softmax、scatter 融合成一个 kernel；或者用torch.compile做图优化，让编译器自动融合操作。

6.4 常见问题速查表

实操心得：块稀疏 Attention 的调试最好从全量 Attention 的 checkpoint 开始做微调，而不是从头训练。这样模型已经学到了跨模态对齐，只需要适应稀疏模式，收敛速度快很多，通常 10% 到 20% 的额外训练步数就够了。

7. 多模态场景下的稀疏模式设计经验

7.1 文本 token 和图像 token 的交互不能省

这是我踩过最大的坑。一开始为了追求高稀疏率，我把文本到图像的连接也做了稀疏化，结果生成质量直接崩了。原因是文本条件控制是扩散模型的核心，如果文本 token 不能和所有图像 token 交互，条件信息就传不到图像的某些区域，生成结果会忽略文本描述。

后来我改成文本到图像全连接，图像到文本也全连接，只对图像到图像做稀疏。这样稀疏率虽然降到了 60% 左右，但生成质量几乎无损。这个取舍是值得的，因为文本 token 数量少（通常 77 个），全连接的计算开销很小。

7.2 不同模态用不同的块大小

文本 token 和图像 token 的语义密度不同。文本 token 每个都承载了丰富的语义信息，块大小可以小一些，比如 32 或 64。图像 token 的语义密度相对低，块大小可以大一些，比如 128 或 256。

但这样做会让块稀疏的实现变复杂，因为不同模态的块大小不同，块对的计算需要处理非均匀的块。我试过一段时间，实现复杂度上升了不少，但收益只有 5% 左右的延迟降低。后来还是统一用 128 的块大小，简单可靠。

7.3 时间步嵌入的处理

时间步嵌入在 DiT 里通常是通过 AdaLN 注入的，不占序列位置。但有些实现会把时间步嵌入也拼到序列里。如果拼进去，时间步 token 应该和所有块全连接，因为它影响整个去噪过程。

我建议不要把时间步拼进序列，用 AdaLN 更干净。如果非要拼，记得给时间步 token 单独一个块，并且这个块和所有其他块全连接。

8. 推理阶段的优化与部署建议

8.1 稀疏模式的预编译

推理阶段稀疏模式是固定的，可以预编译成高效的 kernel。具体做法是把每个 Query 块需要关注的 Key 块索引提前算好，存成一个静态的索引表。推理时直接根据索引表做 gather 和 bmm，省去了动态计算 mask 的开销。

如果用的是 TensorRT 或者 ONNX Runtime，可以把块稀疏 Attention 导出成自定义算子。TensorRT 对稀疏算子的支持还在完善中，目前比较稳妥的方案是用 Triton 写 kernel，然后在推理框架里调用。

8.2 batch 内不同样本的稀疏模式对齐

如果 batch 内不同样本的序列长度不同，稀疏模式也会不同。这会导致 batch 内的计算不规整，GPU 利用率下降。解决办法是把序列长度补齐到 batch 内最大长度，然后用统一的稀疏模式。补齐带来的额外计算通常比不规整计算的损失小。

8.3 与 KV Cache 的结合

扩散模型的推理是迭代去噪，每步都要重新计算 Attention。如果能把 KV Cache 利用起来，可以省掉重复的 Key 和 Value 计算。但块稀疏 Attention 的 KV Cache 管理和全量 Attention 不同，因为不是所有 Key 块都会被用到。

我的做法是只缓存那些被至少一个 Query 块关注的 Key 块。在稀疏模式固定的情况下，这个集合是确定的，可以提前算好。这样 KV Cache 的显存占用也能同比减少。

9. 一些个人体会和后续可以尝试的方向

块稀疏 Attention 在多模态 DiT 上的落地，核心不是算法有多复杂，而是稀疏模式的设计要和模态特性匹配。跨模态交互要保守，模态内交互可以激进。这个原则听起来简单，但实际调的时候很容易为了追求稀疏率而破坏跨模态连接，最后生成质量崩了还得回头改。

另一个体会是，块稀疏 Attention 的收益在序列长度越长时越明显。序列长度 1024 的时候，稀疏带来的加速可能只有 1.3 倍；序列长度 4096 的时候能到 2 倍以上；如果做视频生成，序列长度上万，加速比能到 3 倍甚至更多。所以这个技术最适合的场景是长序列多模态生成，比如视频扩散、高分辨率图像生成。

后续可以尝试的方向，一个是动态稀疏模式，用轻量级的路由网络根据输入内容决定哪些块对需要计算。另一个是跨层稀疏模式共享，相邻的 DiT Block 用相似的稀疏模式，减少模式切换的开销。还有一个是稀疏模式的可学习化，把稀疏模式作为可学习参数，通过 Gumbel Softmax 做端到端训练。这些方向我还在实验中，有结果了再分享。