背景是前前段时间在看 sglang diffusion 相关部分,大概研究了一下并做了些笔记,同时嫌自己的 blog 确实太空了,故而分享。

但我的研究不晚于 2026-02-10 ,可能存在因版本更新等原因造成的不准确描述,同时我这方面的学习也刚刚起步还算是才疏学浅,还请注意甄别。

sglang 目前支持了 TeaCachecache-dit , 但关于 TeaCache 的文档描述有些问题把 TeaCache 含义与论文地址打错了。我提的 pr 也一直无人理会…

TeaCache (Timestep Embedding Aware Cache)

论文地址:https://arxiv.org/pdf/2411.19108

简单来说,是通过实验发现 diffusion 模型 timestep embedingtimestep embedding modulated noisy inputs 均与输出呈现出了强相关性。因此可以通过这两个值判断是否使用缓存,其中 timestep embedding modulated noisy inputs 表现出了更强的相关性,性能上也更加优异,但我猜测可能是性能差异并不明显,同时实现的复杂度与计算的复杂度等因素,在实际的实现中,更多是采用 timestep embeding 作为判断。包括 TeaCache 库, ComfyUI插件sglang 都仅使用了 timestep embedding

具体一些,比较当前时间步与前一时间步 timestep embeding 之间的 l1 距离进行多项式缩放之后累加与阈值判断。所跳过的是整个 transforms layer,跳过所有的 transform blocks。缓存的是 transforms layer 的残差,也就是变化趋势。使用缓存是用当前的输入加上缓存的残差。

相邻时间步的变化不是一致的会存在一个标量误差,因此引入了多项式缩放进行拟合。

累计 l1 距离超过阈值则重新进行计算缓存。

diffusion 相关性 性能比较

cache-dit

cache-dit 从缓存粒度上来看可以说是 TeaCache 的增强,但从判断方式来看更像是 AdaCahe 的简化,是作用在某个或某连续几个的 transform 块,计算比较的是不同时间步下特征残差的 l1 距离。缓存的是跳过块的残差。

当前 sglang 中支持了 DBCache DBPrune 两种缓存方式,与 TaylorSeer SCM 两种增强机制。

DBCache

将 transformers layer 按顺序分成了三大块: FN MN BN 。而不再粗暴地整个跳过。

完整地计算 FN, 计算 FN 残差,计算与缓存的 FN 块残差的 l1 距离,没有缩放也没有累加,直接与阈值比较。

若使用缓存,则跳过 MN 块的计算;若不使用缓存,则更新 FN 块残差缓存,重新计算 MN 更新 MN 残差缓存。

最后计算 BN 块。

DBPrune

DBPrune 又是 DBCache 的进一步增强。
缓存的粒度从多个连续的 transformer 块也就是 MN 块,变成了每个 transform 块
类似的,计算当前时间步当前块与上一时间步当前块的 hidden_states l1 距离,与阈值比较判断是否使用缓存

同时引入了动态阈值,取当前时间步前 5 块 diff 的平均值,若 ori_threshold < avg_diff < max_threshold 则取平均值做阈值。

TaylorSeer&SCM

只是两种增强机制。

TaylorSeer 使用泰勒展开预测变化趋势,替代所缓存的残差。但基于 特征值的变化在相邻步数上是平滑的 这一假设,因此仅适用于推理步数较多的模型,不适用于蒸馏后的模型。

SCM (Step Computation Masking),对缓存粒度更加精细地控制,可将推理过程按步骤进行分段并控制每段使用缓存的步数。同样仅使用于推理步数较多的模型。


我的一些参考文章: