efc7a44747
Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline).
103 lines
4.4 KiB
Python
103 lines
4.4 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Opt-in step caching for the diffusion transformer (First-Block-Cache).
|
|
|
|
Across denoising steps a DiT's output changes little once the trajectory settles, so most of
|
|
the transformer can be reused. First-Block-Cache (FBCache) computes the first block, and if
|
|
its residual barely changed from the previous step (within ``threshold``) it skips the
|
|
remaining blocks and reuses their cached output. diffusers ships it natively
|
|
(``transformer.enable_cache(FirstBlockCacheConfig(...))`` for CacheMixin models, or the
|
|
standalone ``apply_first_block_cache`` hook).
|
|
|
|
Measured on Flux.1-dev (28 steps, 1024px, B200): ~1.4x on top of torch.compile (2.83 ->
|
|
2.03 s) at LPIPS ~0.08 vs the no-cache output -- deep inside the speed-for-quality bar.
|
|
|
|
OFF by default and a deliberate per-load opt-in, because the win scales with step count: a
|
|
few-step distilled model (e.g. Z-Image-Turbo at ~8 steps) has almost no headroom and a
|
|
single skipped step is a large fraction of the trajectory, so caching is for many-step
|
|
models (Flux / Qwen-Image). It composes with torch.compile only with ``fullgraph=False``
|
|
(the cache's compiler-disabled decision is a graph break), which the speed layer switches to
|
|
automatically when a cache is engaged. Best-effort: an incompatible model (e.g. a transformer
|
|
whose block signature the hook does not recognise) is caught and the load proceeds uncached.
|
|
torch / diffusers imported lazily.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
TC_OFF = "off"
|
|
TC_FBCACHE = "fbcache"
|
|
TC_MODES = (TC_FBCACHE,)
|
|
|
|
# FBCache residual thresholds: higher skips more steps (faster, lower quality). The dense
|
|
# bf16 default; a quantised transformer shifts the residual distribution, so it needs a
|
|
# higher threshold for the cache to trigger at all (per ParaAttention's fp8 guidance).
|
|
DEFAULT_FBCACHE_THRESHOLD = 0.08
|
|
QUANT_FBCACHE_THRESHOLD = 0.12
|
|
|
|
|
|
def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
|
|
"""Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled).
|
|
|
|
Raises ValueError for an unsupported value so a bad request is rejected cheaply."""
|
|
if value is None:
|
|
return None
|
|
normalized = str(value).strip().lower().replace("-", "_")
|
|
if not normalized or normalized in ("none", "off"):
|
|
return None
|
|
if normalized not in TC_MODES:
|
|
raise ValueError(
|
|
f"Unsupported transformer_cache '{value}'. Use one of: off, {', '.join(TC_MODES)}."
|
|
)
|
|
return normalized
|
|
|
|
|
|
def apply_step_cache(
|
|
pipe: Any,
|
|
*,
|
|
mode: Optional[str],
|
|
threshold: Optional[float] = None,
|
|
quant_active: bool = False,
|
|
logger: Any = None,
|
|
) -> Optional[str]:
|
|
"""Engage step caching on ``pipe.transformer``. Returns the mode actually engaged, or
|
|
None when disabled / unsupported (the load then runs uncached). ``threshold`` overrides
|
|
the default; ``quant_active`` raises the default so the cache still triggers on a
|
|
quantised transformer. Best-effort: never raises for an incompatible model."""
|
|
mode = normalize_transformer_cache(mode)
|
|
if mode is None:
|
|
return None
|
|
transformer = getattr(pipe, "transformer", None)
|
|
if transformer is None:
|
|
return None
|
|
thr = threshold if threshold is not None else (
|
|
QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD
|
|
)
|
|
try:
|
|
from diffusers import FirstBlockCacheConfig
|
|
|
|
config = FirstBlockCacheConfig(threshold=thr)
|
|
enable_cache = getattr(transformer, "enable_cache", None)
|
|
if callable(enable_cache):
|
|
enable_cache(config)
|
|
else:
|
|
from diffusers.hooks import apply_first_block_cache
|
|
|
|
apply_first_block_cache(transformer, config)
|
|
try:
|
|
transformer._unsloth_step_cache = f"{mode}@{thr}"
|
|
except Exception: # noqa: BLE001 — marker is best-effort
|
|
pass
|
|
if logger is not None:
|
|
logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
|
|
return mode
|
|
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
|
|
_warn(logger, mode, exc)
|
|
return None
|
|
|
|
|
|
def _warn(logger: Any, what: str, exc: Exception) -> None:
|
|
if logger is not None:
|
|
logger.warning("diffusion.cache: %s unavailable (%s); running uncached", what, exc)
|