RT-Lynx:把 GEMM 稀疏化放在正确位置以加速扩散模型推理 RT-Lynx: Putting the GEMM Sparsity In a Right Way for Diffusion Models

完整流程分四步。第一步是 Norm-Compensated Activation Sparsification:对每个 token 的激活向量 X(维度 D_in,通常是 4 的倍数以便 2:4 分组),先用 Top-K(每 4 选 2)得到稀疏向量 X̃,然后计算 rescale 因子 $s = \|X\|_2 / (\|\tilde{X}\|_2 + \epsilon)$,其中 ϵ=1e-8 保证数值稳定,得到 S(X) = s·X̃。这一步的输入是稠密激活 X,输出是稀疏但等范数的 S(X),且因为是 closed-form,计算开销可忽略。第二步是 LoRA 补偿:把 S(X) 送入稀疏主分支得到 $Y_s = S(X) \cdot W^T$,同时原始稠密 X 送入 R=64 的 LoRA 分支得到 $Y_r = X \cdot (L_A L_B)^T$,最终输出 $Y = Y_s + Y_r$。LoRA 的训练目标是最小化 $\|X \cdot W^T - S(X) \cdot W^T - X \cdot (L_A L_B)^T\|_2^2$,即让 LoRA 学到'dense 输出减去 sparse 输出的残差',论文报告 2k 步即可收敛。第三步是 Selective Layer Skipping:对 double-stream 的 Qwen-Image,QKV/W_up/W_down 都做激活稀疏化;但对 single-stream 的 FLUX 与 Z-Image,稀疏化残差难以完全被 R=64 LoRA 吸收,因此跳过部分层——FLUX 跳过单流路径的 attn.o_proj 和 mlp.down,Z-Image 跳过 attn.o_proj 和 mlp.up——其余层正常应用 RT-Lynx。第四步是 CUDA Kernel 优化:自研 kernel 把(1)mask 判定、(2)Top-K 选择、(3)数据压缩到 SpTC 兼容布局、(4)streamK 风格稀疏矩阵乘这四步融合,全部在 register/SMEM 层级完成,稀疏化 overhead <10%;同时把 Ys 的稀疏矩阵乘和 Yr 的稠密矩阵乘在 on-chip 做累加,避免 LoRA 中间张量写回显存,降低访存与同步开销。整个算法在论文 Algorithm 1 中以伪代码形式给出,6 行即可描述清楚。

技术新颖性可以拆成四点。第一,'激活比权重更稀疏'这个观察虽然不新,但论文第一次系统地用 DiT 上的 RFE 度量给出了量化证据(图 2c),并把它从观察升级为'应该把稀疏化对象换成激活'的方法论主张,这是研究范式(paradigm shift)层面的贡献。第二,Norm-Compensated rescale 看似简单(s = ‖X‖/‖X̃‖),但放在 N:M 稀疏化语境里它承担了'保持输出尺度不变'的关键作用,论文实验显示单独 SA-NC 就能把 Qwen-Image sDCI FID 从 48.59 降到 37.56(-10.93),IR 从 0.4724 拉到 0.8399,说明这一行公式换来了将近一半的精度恢复,工程性价比极高。第三,R=64 的极小秩 LoRA 选择体现了一种克制美学:大多数方法把 LoRA 做得越大越准(Slim 用 R=0.1d),RT-Lynx 反向论证'只需补偿稀疏化引入的残差,这个残差是低秩的',因此用极小秩就能搞定,这让 LoRA 分支的额外延迟近乎可忽略。第四,CUDA kernel 的四步融合与 on-chip 累加虽然属于系统优化范畴,但论文报告 sparse overhead 从 40%+ 压到 <10%,这是把'理论稀疏'变成'实际加速'的临门一脚,工程价值不亚于算法本身的创新。

独立分析可补充以下四点弱点及改进方向。第一,LoRA 的训练数据规模与泛化性未充分披露:论文说 LoRA 用 20k 随机采样的用户 prompt 训练 2k 步收敛,但没说明这 20k prompt 是否覆盖了 MJHQ 和 sDCI 的领域分布(后者是基准测试集),如果训练数据分布与评测分布有显著漂移,RT-Lynx 的'反超 dense 基线'结果可能存在数据泄露风险,改进方向是在评测时使用'未见过的 prompt 子集'并报告 train/test domain gap。第二,激活稀疏化层数与稀疏度选择是基于经验而非自动搜索:论文固定 2:4 + 选定 layer,但 DiT 各层对稀疏化的敏感度差异巨大(图 2c 的 RFE 曲线),更系统的做法是引入 layer-wise sparsity ratio(每层根据 RFE 自适应选择 2:4 或 1:4 或不稀疏),用 NAS 或可微搜索自动决定;这能把 1.55× linear speedup 进一步推向理论上限。第三,RT-Lynx Kernel 仅评估 H20,缺乏跨硬件对比:虽然 H20 是当下 DiT 推理的主要硬件,但 RTX 4090/5090 的消费级用户群巨大,稀疏化加速是否能复现、稀疏 overhead 是否仍在 10% 以内,需要补充实验;另外稀疏 GEMM 在 H100 (FP8 加速) 上的相对优势也值得验证。第四,论文 Table 6 显示 RT-Lynx 单用只有 1.20×~1.24× 端到端加速,这个数字对很多用户'不够性感',因为哪怕不上任何加速,仅靠 TeaCache (l=0.6) 就能拿 2.64×;论文应该更明确地定位 RT-Lynx 为'算子级加速',主打'叠加在 cache/蒸馏/量化之上'的复合收益,而非单独使用的替代方案,这一点在 Abstract 和 Introduction 中未充分强调,容易让读者高估单用价值。

未来研究方向可分作者提出与基于成果可延伸两条线。作者明确指出的方向:(1) 把 RT-Lynx 推广到视频 DiT(Wan、Sora、CogVideoX),论文实验只覆盖文生图,但视频 DiT 的 token 数更多、稀疏化 overhead 比例可能更高,值得专门 kernel 优化;(2) 与更多 attention 加速(如 Sliding Tile Attention、Linear Attention)正交组合,论文只在 SpargeAttn 上验证过;(3) 与 FP4/FP8 等更激进数值精度叠加。基于成果可延伸的方向:(4) 把 Selective Layer Skipping 升级为'可学习的 layer mask',用 NAS 或可微松弛搜索每层是否稀疏化,有望逼近理论加速极限;(5) 把激活稀疏化推广到 diffusion LLM(如 LLaDA、Mercury 等),这类模型的 token 空间与 DiT 类似但更接近纯文本,值得验证是否能复用相同的'天然稀疏'论点;(6) 用 RT-Lynx 思路改造 MoE 路由,既然 DiT 激活在每 token 上只有 5%~10% 通道显著激活,这与 MoE 的'专家激活'稀疏性高度同构,稀疏 GEMM 可能反过来降低 MoE 路由的通信与计算成本;(7) 探索 1:4 甚至 1:8 等更激进稀疏度,论文仅在 2:4 上验证,但部分层(如 attention V 投影)可能容许更高稀疏度,带来更大加速。

复现评估方面,论文给出的可复现信息相对充分但仍有改进空间。代码与模型层面,作者在 Abstract 提到'first work to achieve lossless N:M sparsification for DiT',但论文未在正文或附录明确声明是否开源 CUDA kernel、LoRA 权重、训练脚本,从 arXiv 论文惯例看这种工程性工作通常会开源,但需要等到正式发表或作者 GitHub 公开才能确认;若有开源,kernel 代码会包含四步融合的具体 register allocation、warp 分工策略、streamK 配置,这些是论文 Figure 5(c) 之外的重要实现细节。数据层面,作者使用自收集的 20k prompt 训练 LoRA,这个数据集未公开,复现者需要自行准备近似分布的 prompt;评测使用 MJHQ-30K 和 sDCI 两个公开 benchmark,各采样 5000 prompt,这个规模与采样方式是确定的,可复现。算力层面,所有训练在 NVIDIA H20 GPU 上完成(单卡或多卡未明确),NVIDIA Driver 580.82.07 + CUDA 13.0 是相对新的环境,可能对老硬件不友好;CUDA kernel 调优高度依赖 GPU 架构,在 H20 上跑出的 1.88× speedup 在 A100/H100/RTX 4090 上未必能复现到同一水平。复现难度评估:中等偏高——核心创新点(Norm-Compensated + R=64 LoRA)的算法部分是简单的,半天就能跑通,但要复现 CUDA kernel 的具体性能数字需要扎实的 GPU programming 经验,预计一个熟练工程师复现算法+主要实验需要 2~3 周,完整 kernel 复现 1~2 个月。