SAFE-DiT：高分辨率扩散的语义感知快速路径执行 Transformer

高分辨率扩散 Transformer (DiT) 推理包含大量空间冗余，但许多空间自适应实现将区域计算编码为注意力掩码，这可能会无意中将缩放点积注意力 (SDPA) 移离 FlashAttention 快速路径。我们将这种可避免的系统瓶颈识别为掩模诱导调度税（MIDT），并表明它随着潜在序列长度的增长而增长。我们介绍 SAFE-DiT，这是一种 无需训练 语义感知快速路径执行框架，它将精确掩码省略与基于近似的空间调度分开。 SAFE-DiT 仅删除经过来源认证的图像自注意力掩码，这些掩码会引起注意力 logits 的行式恒定转移，保留带有语义的掩码（例如文本填充掩码），并通过提示条件词元分区、全局上下文的选择性状态更新和定期上下文刷新来实现空间适应。我们将这种仅加速配置称为 SAFE-Core，并将敏感度加权无分类器指导单独报告为 SAFE-DiT+SW。在评估的 PyTorch SDPA 堆栈上，冗余掩码使长序列注意力比无掩码路径慢 4.1 倍到 5.8 倍。在 Lumina-Next 上，SAFE-DiT 在 1024^2$ 分辨率下实现了 2.69\times$ 端到端加速，在 2560^2$ 分辨率下实现了 5.09\times$ 端到端加速，将 $2560^2$ 上的峰值内存从 94.1 GB 减少到 27.9 GB，并在密集推理内存不足时生成 3072²分辨率的图像。配对指标、组件消融和盲人研究支持 SAFE-Core 在视觉上不劣于密集的快速路径基线，而 SAFE-DiT+SW 提供了一个单独的提示对齐操作点，而无需重新引入空间自注意力掩模。代码可在 https://github.com/xuanhuayin/SAFE-DiT 获取。