ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'image-generation' into video-diffusion-improvements

Browse source
This commit is contained in:
Daniel Han 2026-07-12 13:59:40 +00:00
commit e727b9e82f
95 changed files with 8911 additions and 5348 deletions
Download patch file Download diff file Expand all files Collapse all files

172
studio/backend/core/inference/diffusion_cache.py
View file

@ -3,24 +3,19 @@
"""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).
Once a DiT's trajectory settles its output changes little across steps, so most of the
transformer can be reused. FBCache computes the first block and, if its residual barely changed
from the previous step (within ``threshold``), skips the rest and reuses their cached output.
diffusers ships it natively (``transformer.enable_cache(FirstBlockCacheConfig(...))``).
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.
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 -- 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.
OFF by default: the win scales with step count, so a few-step distilled model (Z-Image-Turbo
~8 steps) has almost no headroom and caching is for many-step models (Flux / Qwen-Image). It
composes with torch.compile only at ``fullgraph=False`` (the cache's compiler-disabled decision
is a graph break), which the speed layer switches to automatically. Best-effort: an incompatible
model is caught and the load proceeds uncached. torch / diffusers imported lazily.
"""
from __future__ import annotations
@ -33,9 +28,9 @@ TC_FBCACHE = "fbcache"
TC_MAGCACHE = "magcache"
TC_MODES = (TC_FBCACHE, TC_MAGCACHE)
# 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).
# FBCache residual thresholds: higher skips more steps (faster, lower quality). Quantised
# transformers shift the residual distribution, so they need a higher threshold to trigger at
# all (per ParaAttention's fp8 guidance).
DEFAULT_FBCACHE_THRESHOLD = 0.08
QUANT_FBCACHE_THRESHOLD = 0.12
@ -360,10 +355,8 @@ def auto_cache_mode(family: Optional[str]) -> str:
def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
"""Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled),
"auto" -> TC_AUTO (the loader decides from the step count).
Raises ValueError for an unsupported value so a bad request is rejected cheaply."""
"""Lower/strip a cache mode; None / "" / "none" / "off" -> None, "auto" -> TC_AUTO (loader
decides from step count). Raises ValueError for an unsupported value."""
if value is None:
return None
normalized = str(value).strip().lower().replace("-", "_")
@ -440,14 +433,11 @@ def _ensure_block_metadata_registered(transformer: Any, logger: Any = None) -> N
def _invalidate_child_registry_cache(transformer: Any) -> None:
"""Drop the HookRegistry's cached child-registry list after (un)installing hooks.
``cache_context`` propagates the state context through ``_get_child_registries``,
which diffusers 0.39 caches on first use. An UNCACHED generation already calls
``cache_context`` (the pipeline wraps every denoise call), creating the
transformer-level registry with an EMPTY cached child list -- so a later
``enable_cache`` (the auto step-count toggle engaging FBCache mid-session) installs
block hooks that ``_set_context`` never reaches, and the first cached forward dies
with "No context is set". Invalidate the stale cache so the next ``cache_context``
rebuilds it over the freshly hooked blocks. Best-effort and cheap (one attribute)."""
``cache_context`` propagates state through ``_get_child_registries``, which diffusers 0.39
caches on first use. An uncached generation already calls it, creating an EMPTY cached child
list -- so a later ``enable_cache`` installs block hooks ``_set_context`` never reaches and the
first cached forward dies with "No context is set". Invalidate so the next ``cache_context``
rebuilds it over the freshly hooked blocks. Best-effort."""
registry = getattr(transformer, "_diffusers_hook", None)
if registry is not None and getattr(registry, "_child_registries_cache", None) is not None:
try:
@ -456,9 +446,8 @@ def _invalidate_child_registry_cache(transformer: Any) -> None:
pass
# diffusers' cache hook registry names whose compute branch we re-point at a compiled
# inner forward (leader = the measuring first block, block = the remaining ones); both
# hook families share the fn_ref layout.
# diffusers cache hook names whose compute branch we re-point at a compiled inner forward
# (leader = measuring first block, block = the rest); both share the fn_ref layout.
_CACHE_HOOK_NAMES = (
"mag_cache_leader_block_hook",
"mag_cache_block_hook",
@ -470,30 +459,21 @@ _CACHE_HOOK_NAMES = (
def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
"""Restore the regional compile on cache-hooked blocks' COMPUTED steps.
``enable_cache`` replaces each block's ``forward`` with the hook's ``new_forward``
(stashing the pre-hook bound method in ``fn_ref.original_forward``), whose skip
decision is data-dependent Python: MagCache ``@torch.compiler.disable``s the whole
``new_forward`` (recursive -- the compute branch runs EAGER), and even FBCache's
traceable ``new_forward`` graph-breaks around its disabled threshold decision,
which on some archs (measured: Qwen-Image) drops the compute branch's call into
``original_forward`` out of the compiled region -- the block's regional compile
artifact (``_compiled_call_impl``) is never reached and the cache forfeits the
compile win on every non-skipped step. An explicitly ``torch.compile``d callable
re-enables
dynamo for its own extent even inside a disabled frame, so re-pointing
``fn_ref.original_forward`` at a compiled wrapper of the same bound method restores
compiled compute steps while the skip decision stays eager exactly as designed.
Measured (B200, scripts/image_speedmem_bench.py): Qwen-Image FBCache computed steps
91.8 -> 71.2 ms (= the uncached compiled rate), 1.21x end to end; FLUX.1-dev is
neutral (its FBCache ``new_forward`` happens to trace, so computed steps were
already compiled -- same-process armed vs unarmed latents bit-identical); on the
video DiT balanced MagCache went 39.4 -> 26.9 s at 50 steps.
``enable_cache`` replaces each block's ``forward`` with the hook's ``new_forward`` (stashing
the bound method in ``fn_ref.original_forward``), whose skip decision is data-dependent Python:
MagCache ``@torch.compiler.disable``s the whole thing (compute runs EAGER), and even FBCache's
traceable ``new_forward`` graph-breaks around its disabled decision, which on some archs
(Qwen-Image) drops the compute call out of the compiled region -- so ``_compiled_call_impl`` is
never reached and the cache forfeits the compile win on every computed step. A ``torch.compile``d
callable re-enables dynamo for its own extent even inside a disabled frame, so re-pointing
``original_forward`` at a compiled wrapper restores compiled compute steps while the skip
decision stays eager. Measured: Qwen-Image FBCache computed steps 91.8 -> 71.2 ms (uncached
compiled rate), 1.21x end-to-end; FLUX.1-dev neutral (its new_forward traces); video DiT
MagCache 39.4 -> 26.9 s at 50 steps.
Only blocks the speed layer actually compiled are armed (``_compiled_call_impl``
guard -- eager tiers stay untouched), and only when ``original_forward`` is a plain
bound method (a stacked hook chain, e.g. offload, captures a partial and is
skipped). Idempotent via the ``_unsloth_orig_inner`` marker; best-effort. Returns
the number of hooks armed."""
Only speed-layer-compiled blocks are armed (``_compiled_call_impl`` guard) and only when
``original_forward`` is a plain bound method (a stacked hook chain is skipped). Idempotent via
``_unsloth_orig_inner``; best-effort. Returns the number armed."""
try:
import torch
except Exception: # noqa: BLE001 -- no torch, nothing to arm
@ -512,11 +492,10 @@ def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
if orig is None or getattr(hook, "_unsloth_orig_inner", None) is not None:
continue
if getattr(orig, "__self__", None) is None:
continue # not the plain bound method; arming would miss the block
# fullgraph=False / dynamic=True: a cache is active by definition (its
# decision points graph-break) and this matches the default tier the
# regional compile used. Dynamo caches per code object, so re-arming
# after a toggle is effectively free (~0.03 s).
continue # not a plain bound method; arming would miss the block
# fullgraph=False / dynamic=True: a cache is active (its decision graph-breaks) and
# this matches the default tier. Dynamo caches per code object, so re-arming after
# a toggle is ~free (~0.03 s).
fn_ref.original_forward = torch.compile(orig, fullgraph = False, dynamic = True)
hook._unsloth_orig_inner = orig
armed += 1
@ -532,10 +511,10 @@ def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
def _restore_hooked_block_inners(transformer: Any) -> None:
"""Undo ``_compile_hooked_block_inners``: put the plain bound methods back and clear
the markers. MUST run before ``disable_cache`` -- ``remove_hook`` splices
``fn_ref.original_forward`` back into ``module.forward``, and leaving the compiled
wrapper there would pin a stale compiled callable onto the uncached path."""
"""Undo ``_compile_hooked_block_inners``: restore the bound methods and clear the markers.
MUST run before ``disable_cache`` -- ``remove_hook`` splices ``original_forward`` back into
``module.forward``, so a leftover compiled wrapper would pin a stale callable on the uncached
path."""
try:
modules = list(transformer.modules())
except Exception: # noqa: BLE001 -- not a torch module (tests/fakes): nothing armed
@ -558,12 +537,10 @@ def _restore_hooked_block_inners(transformer: Any) -> None:
def _pipeline_opens_cache_context(pipe: Any) -> bool:
"""Whether the pipeline enters ``transformer.cache_context(...)`` in its denoise loop.
The First-Block-Cache hook requires it at run time, and a CacheMixin transformer alone
does NOT guarantee it: Flux Kontext / img2img / inpaint / controlnet reuse the CacheMixin
FluxTransformer2DModel but never open a cache_context. Read from the pipeline ``__call__``
source, resolved off the instance so a per-expert proxy view (``_SecondDiTView``)
delegates to the real pipe; if it cannot be read, report False so the cache stays off."""
"""Whether the pipeline enters ``transformer.cache_context(...)`` in its denoise loop. The
FBCache hook requires it at run time, and a CacheMixin transformer alone doesn't guarantee it
(Flux Kontext / img2img / inpaint / controlnet reuse FluxTransformer2DModel but open none).
Read from ``__call__`` source; False when unreadable so the cache stays off."""
import inspect
call = getattr(pipe, "__call__", None)
@ -573,8 +550,7 @@ def _pipeline_opens_cache_context(pipe: Any) -> bool:
src = inspect.getsource(call)
except (OSError, TypeError):
return False
# Match the actual call `cache_context(` -- a bare mention in a comment/docstring lacks
# the paren, so this does not false-positive on prose.
# Match the call `cache_context(` (the paren avoids a false positive on prose).
return "cache_context(" in src
@ -601,8 +577,7 @@ def apply_step_cache(
Best-effort: never raises for an incompatible model."""
mode = normalize_transformer_cache(mode)
if mode is None or mode == TC_AUTO:
# AUTO must be resolved by the loader (step-count policy) before reaching the
# engage call; treat a stray auto as off rather than crashing the load.
# AUTO is resolved by the loader before this; treat a stray auto as off.
return None
transformer = getattr(pipe, "transformer", None)
if transformer is None:
@ -628,11 +603,9 @@ def apply_step_cache(
if not callable(enable_cache):
_warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
return None
# A CacheMixin transformer is necessary but NOT sufficient: the First-Block-Cache hook
# raises "No context is set" on the first forward unless the PIPELINE wraps its denoise
# loop in transformer.cache_context(...). Flux Kontext / img2img / inpaint / controlnet
# reuse the CacheMixin FluxTransformer2DModel yet their __call__ opens no cache_context,
# so engaging FBCache there would crash every default generation -- run uncached instead.
# A CacheMixin transformer is necessary but not sufficient: the hook raises "No context is set"
# unless the PIPELINE wraps its denoise loop in cache_context(...). Flux Kontext / img2img /
# inpaint / controlnet reuse FluxTransformer2DModel yet open none, so run uncached instead.
if not _pipeline_opens_cache_context(pipe):
_warn(
logger, mode, RuntimeError("pipeline __call__ opens no cache_context; running uncached")
@ -689,10 +662,9 @@ def apply_step_cache(
# block hooks just installed would never receive the cache context. Must follow
# every enable_cache.
_invalidate_child_registry_cache(transformer)
# If the blocks are already regionally compiled (the generation-time toggle
# path: compile ran at load), re-point the fresh hooks' compute branch at
# compiled inners; the load path (cache before compile) is armed by
# _compile_repeated_blocks instead. No-op when nothing is compiled.
# If blocks are already regionally compiled (toggle path: compile ran at load), re-point
# the fresh hooks' compute branch at compiled inners; the load path is armed by
# _compile_repeated_blocks. No-op when nothing is compiled.
_compile_hooked_block_inners(transformer, logger)
try:
transformer._unsloth_step_cache = marker
@ -702,10 +674,9 @@ def apply_step_cache(
logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
return mode
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
# enable_cache can fail after hooking some blocks; drop any partial hooks so
# the reported-uncached model doesn't actually run half-cached. Any armed
# compiled inners must be restored FIRST (remove_hook splices original_forward
# back into module.forward).
# enable_cache can fail after hooking some blocks; drop partial hooks so the
# reported-uncached model isn't half-cached. Restore armed compiled inners FIRST
# (remove_hook splices original_forward back into module.forward).
_restore_hooked_block_inners(transformer)
try:
transformer.disable_cache()
@ -718,14 +689,11 @@ def apply_step_cache(
def effective_denoise_steps(steps: int, strength: Optional[float]) -> int:
"""The number of steps diffusers ACTUALLY denoises for a request.
An image-conditioned workflow with ``strength`` < 1 (img2img / upscale / inpaint) runs
only a fraction of ``num_inference_steps``: diffusers' ``get_timesteps`` computes
``init_timestep = min(int(num_inference_steps * strength), num_inference_steps)`` and
denoises exactly ``init_timestep`` steps -- the product is FLOORED, not rounded. The auto
step-cache policy must key on THIS count -- e.g. a 28-step upscale at strength 0.35 runs
``int(9.8) = 9`` real steps, exactly the short trajectory FBCache should stay off (each
skipped step is a large quality hit). ``strength`` None (txt2img / reference) or >= 1 -> the
full count.
An image-conditioned workflow with ``strength`` < 1 (img2img / upscale / inpaint) denoises
only ``init_timestep = min(int(num_inference_steps * strength), num_inference_steps)`` steps
-- FLOORED, not rounded. The auto step-cache policy keys on THIS count (e.g. a 28-step upscale
at strength 0.35 runs int(9.8) = 9 steps, the short trajectory FBCache should stay off).
``strength`` None or >= 1 -> the full count.
"""
s = int(steps)
if strength is None or float(strength) >= 1.0:
@ -741,12 +709,10 @@ def effective_request_strength(
) -> Optional[float]:
"""The strength the pipe will ACTUALLY apply, for keying the auto step-cache policy.
Only image-conditioned pipelines that take ``strength`` apply it (txt2img / a pipe without
the kwarg run the full trajectory -> None). When the request omits ``strength`` the loader
does NOT pass the kwarg, so the pipe runs its OWN signature default (< 1 for every img2img /
inpaint pipeline here, e.g. 0.6); the policy must key on that default, not the full step
count, or FBCache engages on a fraction of the advertised steps. A non-numeric default
(``inspect.Parameter.empty``) falls back to the full count (None).
Only image-conditioned pipelines taking ``strength`` apply it (else full trajectory -> None).
When the request omits it the loader doesn't pass the kwarg, so the pipe uses its OWN signature
default (< 1 for every img2img / inpaint pipeline, e.g. 0.6); the policy keys on that default,
else FBCache engages on a fraction of the advertised steps. A non-numeric default -> None.
"""
if not (has_init_image and pipe_accepts_strength):
return None