背景是前前段时间在看 sglang diffusion 相关部分,大概研究了一下并做了些笔记,同时嫌自己的 blog 确实太空了,故而分享。
但我的研究不晚于 2026-02-10 ,可能存在因版本更新等原因造成的不准确描述,同时我这方面的学习也刚刚起步还算是才疏学浅,还请注意甄别。
sglang 目前支持了 TeaCache 与 cache-dit , 但关于 TeaCache 的文档描述有些问题把 TeaCache 含义与论文地址打错了。我提的 pr 也一直无人理会…
TeaCache (Timestep Embedding Aware Cache)
论文地址:https://arxiv.org/pdf/2411.19108
简单来说,是通过实验发现 diffusion 模型 timestep embeding 与 timestep 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 距离超过阈值则重新进行计算缓存。
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),对缓存粒度更加精细地控制,可将推理过程按步骤进行分段并控制每段使用缓存的步数。同样仅使用于推理步数较多的模型。
我的一些参考文章: