AMD publishes open-source TeaCache implementation for Stable Diffusion NPU acceleration on Ryzen AI

AMD has published an open-source implementation of TeaCache optimized for Stable Diffusion 3 and 3.5 running on Ryzen AI neural processing units through ONNX Runtime. The technique addresses a fundamental inefficiency in diffusion models: the denoising loop performs full forward passes at every timestep even when adjacent steps produce nearly identical intermediate activations. TeaCache measures relative L1 distance between timestep embeddings to dynamically skip redundant computation, reusing cached attention outputs and feature maps when changes fall below a configurable threshold.

The implementation supports both naive and rescaled caching strategies, with the latter applying polynomial rescaling to reduce bias in change detection. AMD's adaptation differs from the original video diffusion paper by embedding caching decisions inside individual CachedTransformerBlocks rather than controlling them at the loop level, making the design more portable across pipelines. The code runs within a static ONNX-based inference pipeline on Ryzen AI NPUs and requires no retraining or architectural modifications to existing models.

How TeaCache works in AMD's ONNX pipeline

• Caching granularity: AMD implements caching at the transformer block level via CachedTransformerBlocks that replace JointTransformerBlocks, allowing individual blocks to decide whether to reuse cached latents rather than applying a global loop-level policy.
• Change detection metric: TeaCache tracks relative L1 distance between intermediate representations across timesteps, comparing the current timestep embedding against the last one that triggered a full forward pass.
• Two caching modes: Naive caching uses raw L1 distances against threshold δ, while rescaled caching applies a polynomial function f(L1_rel(F, t)) to reduce bias, with coefficients optimizable per-model.
• ONNX Runtime integration: The implementation operates within a static ONNX-based pipeline on Ryzen AI NPUs, with the CPU handling input preparation and the NPU executing the backbone network where caching decisions occur.
• No retraining required: The wrapper-based approach requires no architectural modifications or fine-tuning of Stable Diffusion 3/3.5 weights, making it applicable to existing deployed models.

Why it matters

This release gives Linux and Windows developers a practical, vendor-documented path to reduce diffusion model latency on consumer AMD hardware without proprietary dependencies. By open-sourcing the block-level caching strategy and providing ONNX Runtime integration steps, AMD enables independent optimization work on Ryzen AI NPUs and potentially other backends, reducing the gap between local NPU inference and GPU-based deployment for generative AI workloads.