Cache-DiT mental model

No cache is the reference. Cache-DiT skips safe DiT block work.

A DiT pipeline repeatedly denoises latent features. Without Cache-DiT, every denoising step runs every transformer block. Cache-DiT wraps the blocks, measures cross-step feature similarity, and reuses cached outputs when the current step is close enough.

normal compute Fn front blocks cache reuse / skip Bn back blocks cache miss → recompute
20–50typical denoising steps: the same DiT backbone is called again and again
>95%transformer backbone can dominate DiT latency / memory in the report’s model
1.6×–2.9×DBCache-only speedup examples on FLUX.1-dev L20, depending on cache aggressiveness

1. Reference path: normal DiT, no Cache-DiT

Every denoising step recomputes the full transformer stack. This is the baseline behavior and the easiest correctness reference.

Pipeline level

baseline
📝
Prompt
🧠
Text encoder
🌫️
Noisy latent
🧱
DiT blocks × T
Denoising step tall blocks run
Denoising step t+1all blocks run again
Cost model: T steps × L blocks. No step-level or block-level reuse is used.

Transformer block detail

full recompute
for step in denoising_steps:\n h = latent\n for block in transformer_blocks:\n h = block(h, t, cond) # attention + FFN\n latent = scheduler_update(h)
SimpleNo extra state or cache decision.
AccurateEvery block sees the exact current-step hidden state.
ExpensiveTemporal redundancy across neighboring steps is not used.

2. Cache-DiT setup: instrument first, then decide at runtime

Cache-DiT is inserted as a PyTorch-native middleware layer. It first discovers model-specific block layout, maps each block to a known forward pattern, then wraps block forward calls so runtime caching can happen without rewriting the whole model.

🧩
Diffusers pipe
🔎
BlockAdapter
🧾
Forward pattern
🎣
Wrap forward()
⚡
DBCache runtime
cache_dit.enable_cache(pipe)\n# instrumentation → cache hooks → optional parallelism / quantization / offload\n# during inference: compute, skip, or reuse per block / per step

3. Runtime comparison: no cache vs DBCache

DBCache keeps first Fn blocks and last Bn blocks computed. Middle blocks are cacheable: if the L1 residual difference is below threshold τ, reuse previous-step cache; otherwise recompute and refresh the cache.

No Cache-DiT

reference
Step t−1full compute
Step tfull compute
Step t+1full compute
Every row is identical in compute pattern: all 12 example blocks run.

Cache-DiT DBCache

skip safe middle work
Fn always compute
middle cache decision
Bn always compute
Warmup / full stepbuild reference cache
Cache-allowed stepL1 diff < τ → reuse
Cache-allowed stepL1 diff ≥ τ → some misses
Middle blocks are approximation candidates; front/back blocks protect quality and reduce drift.

4. Interactive toy view: change Fn, Bn, threshold, and SCM mode

This is a simplified visual model, not a kernel benchmark. It helps build intuition for why larger cache regions increase speed but also increase quality risk.

Lower ratio roughly means more skipped middle blocks.

28-step SCM timeline

fast
Green = warmup/full reference, blue = full compute, purple = cache-allowed step.

One cache-allowed step

12-block toy transformer
L1 diff < τmiddle block output is similar to previous step → reuse cached hidden/residual
⇄
L1 diff ≥ τmiddle block changed too much → run the real block and update cache

5. Report numbers: no-cache baseline as reference

Latency examples below are from the Cache-DiT technical report’s FLUX.1-dev, L20, 28-step DBCache configuration table. Lower latency means fewer expensive block forwards were executed.

Baseline no cache: 24.85s. F1B0 with aggressive threshold τ=0.20 is fastest in this table, but more aggressive cache settings can have more quality risk.

6. Key takeaway

No Cache-DiT

exact but expensive\nT denoising steps × L transformer blocks\nall attention + FFN recomputed

Cache-DiT

approximate but controlled\nwarmup/full steps build cache\nFn blocks anchor current step\nmiddle blocks skip/reuse when safe\nBn blocks repair accumulated error
Important: this is a diffusion feature cache, not an LLM KV cache. It reuses hidden states or residual paths across denoising steps, not attention KV tensors across generated tokens.
Created for learning purpose. Data and mechanism summarized from the Cache-DiT technical report; visualization uses a simplified 12-block toy model.