RT-Lynx：为扩散模型正确应用GEMM稀疏性

TL;DR RT-Lynx 将 N:M 半结构化稀疏从权重转向激活，发现 DiT 激活天然稀疏且更鲁棒，结合误差补偿与定制 CUDA 内核，在保持生成质量的同时实现线性层最高 1.55 倍推理加速。

问题

问题背景

Diffusion Transformers (DiT) 凭借全局建模能力在图像生成任务中取得领先质量，但高昂的推理延迟阻碍了实时应用与服务部署。在量化、蒸馏等加速方案之外，利用结构化稀疏性压缩 GEMM 操作已成为提升吞吐量的重要路径。

现有方法局限

主流稀疏化技术集中于 权重剪枝（weight sparsification），尤其是 N:M semi-structured sparsity，其声称可在 NVIDIA Tensor Core 上实现近理论 2 倍加速。然而，在 DiT 上直接对权重做 50% 剪枝会带来严重容量损失：

• 权重矩阵缺乏内在冗余，尤其是注意力投影和 FFN 中的关键映射通道，剪枝后细节纹理生成出现明显退化；
• 生成模型对输出分布的微小扰动敏感，稀疏化误差在去噪循环中累积，导致图像模糊或伪影；
• 现有方法依赖额外的微调或蒸馏来恢复精度，但训练代价高且不一定稳定，难以适配多样化的扩散模型变体。

为什么这个问题难且重要

DiT 的 激活值（activation） 在去噪过程中天然呈现稀疏性，原因在于 token 之间的信息竞争与叠加效应。将其转换为 N:M 稀疏模式，理论上可无损利用硬件加速指令。但挑战在于：

• 激活稀疏是动态、输入相关的，需要在线稀疏化，对延迟和精度控制要求严苛；
• 简单的 top-k 或阈值法会破坏数值分布，必须设计误差补偿机制（如 norm-compensated sparsification）来保持输出稳定性；
• 自定义 CUDA kernel 需将稀疏化、误差修正与 GEMM 高效融合，避免额外内存搬运。

从产业角度看，图像/视频生成模型被广泛应用于创意工具、虚拟试穿、交互式编辑等场景，推理加速直接关系到用户体验和运营成本。因此，实现 质量无损的激活稀疏化加速 既能突破权重剪枝的瓶颈，又能将半结构化稀疏从 LLM 领域拓展到视觉生成任务，具有很高的实践价值。

行业类比

类似于在 LLM 推理中通过 KV cache 稀疏化 利用注意力头的自然稀疏性来降低显存与计算量，RT-Lynx 将思路迁移至 DiT，挖掘激活值的稀疏结构来加速线性层，为生成模型的高效部署提供了新范式。

方法

RT-Lynx 方法核心是将 N:M 半结构化稀疏 从权重转向激活值，并配合误差补偿、微调与定制化 CUDA 内核，实现 DiT 推理加速。流程如下：

输入与稀疏化策略

输入为 DiT 线性层的激活张量（如形状 [B, N, C]）。RT-Lynx 在推理时对每个激活应用 2:4 稀疏模式（每 4 个值保留 2 个），而非传统权重剪枝。此举依赖观察：DiT 激活具有天然稀疏性且对剪枝更鲁棒，50% 激活去掉后仍能保持生成质量，而权重剪枝 50% 会严重损害模型容量。

核心模块

1. 范数补偿稀疏化 (Norm-Compensated Sparsification)
   直接丢弃元素会导致输出分布偏移。该方法计算被剪枝元素的 L1 范数，并将该能量按比例重新分配到保留的激活值上，从而补偿信息损失，维持层输出的统计一致性。
2. LoRA 适配与微调
   为恢复稀疏带来的细微精度损失，仅插入少量 低秩适配 (LoRA) 层进行微调，训练参数量极小，使模型快速适应激活稀疏化后的数据分布。
3. 选择性层跳过 (Selective Layer Skipping)
   在单流 DiT 架构中，并非所有层对稀疏同样敏感。RT-Lynx 识别并跳过某些关键层的稀疏化，或对不同层采用不同稀疏策略，以平衡速度与质量。
4. CUDA 核优化
   针对 2:4 稀疏模式编写高度优化的 CUDA 算子，充分利用 NVIDIA 稀疏 Tensor Core，在线性层中实现平均 1.55 倍加速。这些内核针对激活稀疏化定制，不同于权重稀疏 Kernel。

输出

保持与原始模型同等的图像生成质量（FID、CLIP Score 等指标相当），同时显著降低推理延迟。

与同类方法的差异：传统方法聚焦权重稀疏化，而 RT-Lynx 首次将 N:M 稀疏化系统性地应用于扩散模型的激活，并结合范数补偿与轻量 LoRA 微调，在不牺牲容量的前提下实现实际推理加速。

实验

实验设计

• 模型：多种 Diffusion Transformer 变体（DiT-XL/2、U-ViT 等），覆盖条件与无条件图像生成。
• 稀疏化方法：对激活应用 2:4 N:M 半结构化稀疏，采用 norm-compensated sparsification 减少截断误差，结合 LoRA 微调恢复精度；对单流 DiT 引入选择性层跳过。
• 硬件优化：实现定制 CUDA Kernel，利用稀疏 Tensor Core 加速线性层。
• 评估：以 FID、IS、sFID 等指标衡量生成质量，同时测量推理延迟与吞吐量。

关键发现

• 激活稀疏更鲁棒：DiT 激活天然具有高稀疏度，对 N:M 剪枝的容忍度远高于权重。
• 生成质量无损：在 50% 激活稀疏下，RT-Lynx 保持了与原始稠密模型几乎一致的 FID，质量损失可忽略。
• 实际加速：线性层平均加速 1.55×，得益于对稀疏 Tensor Core 的高效利用与定制 Kernel 实现。

与基线对比

• 相比传统权重稀疏化（如 ASP、SparseGPT），激活稀疏不剪除关键权重，避免了模型容量下降，因此在同等 2:4 稀疏度下，RT-Lynx 的 FID 显著更优。
• 该结果验证了“从权重稀疏转向激活稀疏”的范式价值，为扩散模型推理优化提供了新方向。

行业影响

落地场景

RT-Lynx 提出的激活稀疏化 (Activation Sparsification) 结合 N:M 半结构化稀疏 (Semi-structured Sparsity)，专为 Diffusion Transformer (DiT) 图像生成模型加速设计，可直接应用于以下产品与业务：

• 云端图像生成 API：如面向设计师、营销团队的 AI 生图服务，需处理大量并发请求。
• 视频生成与编辑工具：DiT 常用于视频扩散模型，稀疏化可降低单帧推理延迟，提升实时预览体验。
• 移动端/边缘设备图像生成：在手机、平板等设备上本地运行小型 DiT 模型，加速可改善交互响应。
• 游戏/虚拟世界资产创建：基于扩散模型的纹理、角色生成，节省算力成本。

商业价值

该方法在不损失生成质量的前提下，线性层平均加速 1.55×，且利用现有 GPU Tensor Core 对 2:4 结构化稀疏的硬件支持，能直接降低推理成本。

• 降本：同等吞吐量下所需 GPU 数量减少，或相同硬件可支撑更多用户，直接降低云服务运营成本。
• 体验提升：推理延迟降低，改善用户交互感受，尤其对实时或近实时生成场景（如对话式图像编辑）至关重要。
• 差异化优势：与主流的权重量化/蒸馏相比，激活稀疏化对生成式模型更鲁棒，提供了一种新的加速路径，可避开容量损失问题。

与现有产品/工作流的接口

RT-Lynx 的集成主要位于模型推理引擎层，可通过以下方式嵌入现有技术栈：

• 推理引擎：将优化后的 CUDA kernel 集成到 PyTorch/TensorRT 中，对已有的 DiT 模型实现一键式稀疏化与加速，无需模型重新训练（仅需少量 LoRA 微调）。
• 模型部署平台：在 ML Serving 平台（如 NVIDIA Triton）中提供稀疏化加速插件，下游应用无感知。
• 开发流程：现有模型训练完成后，增加一个后处理步骤：对激活层施加 N:M 稀疏化掩码 + 误差补偿，再用数百步 LoRA 微调恢复质量，即可生成加速版模型。

具体落地 Use Case

1. 创意内容平台：某平台提供 AI 商品图生成服务（根据产品图生成营销海报）。每日百万级请求，部署 RT-Lynx 后，单卡吞吐量提升约 50%，在不增加硬件投入的情况下可应对峰值流量，同时保持图像质量不降，用户满意度不变。
2. 自动驾驶仿真数据生成：使用 DiT 生成多样的街景图像用于感知模型训练。加速推理可缩短数据生产周期，提高仿真迭代效率，且激活稀疏化对生成图像的多样性影响极小，保证数据有效性。

局限

• **依赖额外训练开销：** 方法需要 LoRA 微调和选择性层跳过，虽然微调成本低于完整训练，但仍需额外计算资源和时间，可能不适合零成本部署场景；且微调策略的跨模型泛化性未充分验证。
• **端到端加速比有限：** 文中报告的 1.55 倍加速仅针对线性层，实际推理 pipeline 还包含注意力、归一化和激活函数等操作，完整模型的 wall-clock 加速可能明显低于理论计算加速比；另外 N:M 稀疏加速严重依赖 NVIDIA Tensor Core 的硬件特性，在其他平台上收益不明。
• **应用场景较窄：** 仅验证于 DiT 架构的图像生成扩散模型，未涉及文本、视频或多模态扩散模型；方法中的单流 DiT 层跳过、误差补偿策略可能无法直接迁移到其他 backbone，通用性有待进一步证明。