TeaCache：让扩散模型少算几步，但尽量不掉画质

latent_t + timestep_t + prompt_condition
        ↓
DiT / Transformer
        ↓
预测噪声 / velocity / residual
        ↓
scheduler 更新 latent

step 50：完整计算 Transformer
step 49：完整计算 Transformer
step 48：完整计算 Transformer
...
step 1 ：完整计算 Transformer

step 50：完整计算 Transformer，缓存 residual
step 49：判断变化小，复用 step 50 的 residual
step 48：判断变化小，继续复用
step 47：变化累计变大，重新完整计算并更新缓存

timestep embedding 调制后的 noisy input

rel_l1 = mean(abs(current_modulated_input - previous_modulated_input)) \
         / mean(abs(previous_modulated_input))

部分 denoising step 里的 Transformer 主体计算

scheduler 更新
VAE decode
text encoder
prompt 编码
CPU/GPU 数据搬运

rel_l1_thresh

阈值越小：更保守，完整计算更多，质量更稳，速度提升较小
阈值越大：更激进，缓存复用更多，速度更快，质量风险更高

质量优先：0.10 ~ 0.20
均衡配置：0.20 ~ 0.40
速度优先：0.50 ~ 0.80

0.2 → 0.3 → 0.4 → 对比质量和耗时

1. DiT / Transformer-based diffusion 模型
2. 图像生成、视频生成、音频扩散生成
3. denoising steps 较多的推理任务
4. 对 1.5x~2x 加速有价值，同时能容忍极小质量波动的生产服务
5. 单卡加速场景

1. 4~8 step 的蒸馏模型
2. 强编辑、强控制、强文字生成场景
3. 对画质零损失要求极高的任务
4. 主要瓶颈不在 Transformer 的 pipeline

from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="tea_cache",
)

outputs = omni.generate(
    "A cat sitting on a windowsill",
    OmniDiffusionSamplingParams(num_inference_steps=50),
)

vllm serve Qwen/Qwen-Image --omni --port 8091 \
  --cache-backend tea_cache \
  --cache-config '{"rel_l1_thresh": 0.2}'

pip install -U diffusers transformers accelerate safetensors scipy pillow

import gc
import time
import types

import torch
from IPython.display import display
from diffusers import DiTPipeline, DPMSolverMultistepScheduler
from diffusers.models.modeling_outputs import Transformer2DModelOutput


MODEL_ID = "facebook/DiT-XL-2-256"
NUM_STEPS = 25
CLASS_ID = 207# ImageNet class id.


class SimpleTeaCacheState:
    """Minimal TeaCache-style state for DiT transformer calls."""

    def __init__(self, rel_l1_thresh=0.20, num_steps=25):
        self.rel_l1_thresh = rel_l1_thresh
        self.num_steps = num_steps
        self.reset()

    def reset(self):
        self.step_idx = 0
        self.accumulated_rel_l1 = 0.0
        self.previous_input = None
        self.previous_residual = None
        self.previous_sample = None
        self.full_compute_steps = 0
        self.cached_steps = 0

    @torch.no_grad()
    def should_compute(self, hidden_states: torch.Tensor) -> bool:
        is_first = self.step_idx == 0
        is_last = self.step_idx >= self.num_steps - 1

        if is_first or is_last:
            self.previous_input = hidden_states.detach().float()
            self.accumulated_rel_l1 = 0.0
            returnTrue

        if self.previous_input isNoneor self.previous_sample isNone:
            self.previous_input = hidden_states.detach().float()
            returnTrue

        current = hidden_states.detach().float()
        previous = self.previous_input

        denom = previous.abs().mean().clamp_min(1e-6)
        rel_l1 = (current - previous).abs().mean() / denom
        self.accumulated_rel_l1 += float(rel_l1.item())

        if self.accumulated_rel_l1 < self.rel_l1_thresh:
            returnFalse

        self.accumulated_rel_l1 = 0.0
        self.previous_input = current
        returnTrue

    def next_step(self):
        self.step_idx += 1


def clear_cuda_cache():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()


def load_dit_pipeline(device):
    dtype = torch.float16 if device == "cuda"else torch.float32
    pipe = DiTPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        token=False,
    )
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    return pipe.to(device)


def enable_simple_teacache_for_dit(pipe, rel_l1_thresh=0.20, num_steps=25):
    state = SimpleTeaCacheState(
        rel_l1_thresh=rel_l1_thresh,
        num_steps=num_steps,
    )

    transformer = pipe.transformer
    original_forward = transformer.forward

    def cached_forward(self, hidden_states, timestep, class_labels=None, **kwargs):
        compute_full = state.should_compute(hidden_states)

        ifnot compute_full and state.previous_sample isnotNone:
            if state.previous_residual isnotNoneand state.previous_residual.shape == hidden_states.shape:
                sample = hidden_states + state.previous_residual.to(
                    device=hidden_states.device,
                    dtype=hidden_states.dtype,
                )
            else:
                sample = state.previous_sample.to(
                    device=hidden_states.device,
                    dtype=hidden_states.dtype,
                )
            state.cached_steps += 1
            state.next_step()
            return Transformer2DModelOutput(sample=sample)

        out = original_forward(
            hidden_states=hidden_states,
            timestep=timestep,
            class_labels=class_labels,
            **kwargs,
        )

        state.previous_sample = out.sample.detach()
        if out.sample.shape == hidden_states.shape:
            state.previous_residual = out.sample.detach() - hidden_states.detach()
        else:
            state.previous_residual = None
            print("Warning: output sample shape differs from input hidden_states shape; cannot compute residual for caching.")
        state.full_compute_steps += 1
        state.next_step()
        return out

    transformer.forward = types.MethodType(cached_forward, transformer)
    return state


def run_generation(pipe, device):
    generator = torch.Generator(device=device).manual_seed(42)
    return pipe(
        class_labels=[CLASS_ID],
        num_inference_steps=NUM_STEPS,
        generator=generator,
    ).images[0]


def main():
    device = "cuda"if torch.cuda.is_available() else"cpu"
    print(f"device = {device}")

    pipe = load_dit_pipeline(device)

    clear_cuda_cache()
    t0 = time.time()
    image = run_generation(pipe, device)
    baseline_time = time.time() - t0
    baseline_image = image
    baseline_image.save("dit_baseline.png")
    print("Baseline image:")
    display(baseline_image)
    print(f"[Baseline] time = {baseline_time:.2f}s")

    del pipe
    clear_cuda_cache()

    pipe = load_dit_pipeline(device)
    state = enable_simple_teacache_for_dit(
        pipe,
        rel_l1_thresh=0.20,
        num_steps=NUM_STEPS,
    )

    clear_cuda_cache()
    t0 = time.time()
    image = run_generation(pipe, device)
    cached_time = time.time() - t0
    cached_image = image
    cached_image.save("dit_simple_teacache.png")
    print("SimpleTeaCache image:")
    display(cached_image)

    print(f"[SimpleTeaCache] time = {cached_time:.2f}s")
    print(f"full_compute_steps = {state.full_compute_steps}")
    print(f"cached_steps       = {state.cached_steps}")
    print(f"speedup            = {baseline_time / cached_time:.2f}x")


if __name__ == "__main__":
    main()

[Baseline] time = 1.42s
[SimpleTeaCache] time = 0.57s
full_compute_steps = 9
cached_steps       = 16
speedup            = 2.48x

rel_l1_thresh = 0.20

0.10：更稳，速度提升较小
0.20：均衡起点
0.30：更快，但可能有质量波动
0.50：偏激进，容易出画质问题

rel_l1_thresh = 0.10

图像细节糊
视频运动不稳定
人物脸部漂移
文字生成质量下降
编辑任务不稳定

cache_config = {"rel_l1_thresh": 0.1}

previous_residual
previous_input
accumulated_rel_l1
step_idx

每个 request 独立 TeaCacheState
请求结束后 reset
CFG cond/uncond 分支分别维护 cache
多线程/异步场景避免共享状态