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

feat(video): step caching for HunyuanVideo-1.5 (MagCache auto, FBCache registry) + int8 trim fix

Browse source 
HunyuanVideo-1.5 loads previously logged 'fbcache unavailable (Model class
HunyuanVideo15TransformerBlock not registered)': diffusers 0.39 ships FBCache
block metadata for HunyuanVideo 1.0 but not 1.5, although the 1.5 DiT is fully
cache-shaped (CacheMixin, homogeneous residual-additive dual-stream blocks,
cache_context per guidance branch). Register the missing metadata at engage
time (deferring to a native registration when a future diffusers ships one).

Measured on a B200 (720p t2v, 1280x720, 33 frames, seed 42), FBCache is fast
but not shippable for this family: 1.44x at 30 steps / 2.41x at 50 steps, at
LPIPS 0.43-0.54 vs the same uncached stack with a +5..8 luma drift (no skip
cap or error budget, so the trajectory derails into a different clip). MagCache
(same registry metadata, also dispatched via enable_cache) is bounded by
design and lands the win: 1.49x end-to-end at 50 steps at LPIPS 0.147 with the
same composition, 1.21x at LPIPS 0.071 on the 480p model. Ship magcache as the
per-family AUTO cache mode for hunyuanvideo-1.5 / -720p with 50-step
calibrated per-family mag_ratios (cond/uncond curves agree within 0.014,
30 vs 50-step calibration within 0.027 after interpolation); every other
family keeps fbcache, and explicit fbcache/magcache requests are honored.
The auto toggle re-engages magcache on a step-count change so the ratio curve
is re-interpolated over the actual schedule.

Two production bugs fixed along the way:
- diffusers' HookRegistry caches its child-registry list, so enabling a cache
  AFTER any uncached generation (the auto off-to-on toggle) left the new block
  hooks without a context ('No context is set' on the first cached forward).
  Invalidate the stale cache after every enable_cache.
- An explicit int8 DiT request crashed under the padded-text trim: torchao's
  int8 dynamic path returns a zero-token (M=0) input unprojected (t2v byt5 /
  image streams -> cond-type add shape crash) and torch._int_mm requires
  M > 16 (an empty negative prompt trims to ~6 tokens -> TokenRefiner crash).
  Add per-family int8 excludes for the text-stream linears (context_embedder*,
  image_embedder, add_q/k/v_proj, to_add_out, ff_context); they run at tens of
  tokens vs the ~32k video stream, so the exclusion costs nothing measurable.

Bench: trim lever key (trim_off / eager_trim isolation configs), int8_cudnn +
shipped_nocache rows, --cache-threshold, warmup timing, per-config frame
persistence for offline LPIPS rescoring, and the loader's per-family auto
cache mode mirrored. Full 720p matrix recorded: reference 481.6s ->
trim+cudnn+compile 35.4s -> shipped default with TE/VAE quant + magcache
24.9s (19.4x, peak VRAM 89.4 -> 81.8 GB), int8 latency-neutral (dense auto
policy confirmed), compile 1.56x per step, trim 13.5x per step at production
shapes.

Validated end to end through the real VideoBackend: load resolves
transformer_cache=magcache with trim + compile + cudnn, generation
re-interpolates 50 -> 30 steps, auto-disengages below 20 steps, re-engages
after an uncached generation, unload restores globals. Hermetic tests cover
the registration, the child-cache invalidation, magcache engage/threshold/
no-curve/no-steps paths, auto-mode routing, toggle re-interpolation, and the
family int8 excludes.
This commit is contained in:
Daniel Han 2026-07-10 10:47:03 +00:00
commit c998183cc2
6 changed files with 637 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -30,7 +30,8 @@ from typing import Any, Optional
TC_OFF = "off"
TC_AUTO = "auto"
TC_FBCACHE = "fbcache"
TC_MODES = (TC_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
@ -38,12 +39,68 @@ TC_MODES = (TC_FBCACHE,)
DEFAULT_FBCACHE_THRESHOLD = 0.08
QUANT_FBCACHE_THRESHOLD = 0.12
# MagCache (diffusers >= 0.39): skips whole steps from a PRE-CALIBRATED residual-magnitude
# curve with an accumulated-error budget, a consecutive-skip cap, and a no-skip retention
# window over the early steps -- so unlike FBCache the divergence from the uncached
# trajectory is bounded. Measured on HunyuanVideo-1.5-720p (B200, 50 steps, 720p clip):
# threshold 0.12 = 1.5x end-to-end at LPIPS 0.147 vs the same uncached stack with the SAME
# composition (FBCache at its 0.08 default reached 2.4x but LPIPS 0.54: a brighter,
# visibly different clip -- why the fbcache auto policy excludes this family).
DEFAULT_MAGCACHE_THRESHOLD = 0.12
MAGCACHE_MAX_SKIP_STEPS = 3
MAGCACHE_RETENTION_RATIO = 0.2
# The auto policy's step-count bar: FBCache's win scales with step count (each skipped
# step is a larger quality hit on a short trajectory), so auto engages it only at 20+
# steps -- full "dev"-style schedules (28+) qualify, distilled turbo models (4-9) never do.
FBCACHE_MIN_STEPS = 20
# Per-family MagCache magnitude-ratio curves (MagCacheConfig.mag_ratios), calibrated with
# diffusers' calibrate mode on the family base checkpoints at the default 50-step schedule
# (720p clip, B200). The curve is checkpoint-dependent but highly stable where it matters:
# the CFG cond/uncond branches differ by <= 0.014 and a 30-step calibration matches the
# 50-step curve within 0.027 after nearest-interpolation, so ONE curve per family is
# enough -- diffusers interpolates it to the actual step count. Conditional-branch curve
# per the MagCache calibration guidance.
_MAGCACHE_720P_RATIOS = (
1.0, 1.0226, 1.0093, 1.001, 1.0008, 1.0001, 0.9995, 1.0003, 0.9998, 0.9993, 0.9994, 0.9993,
0.9997, 1.0002, 0.9994, 0.9985, 0.9987, 0.9997, 0.9979, 0.9987, 0.9985, 0.9982, 0.9977, 0.998,
0.9979, 0.9971, 0.9968, 0.9967, 0.9964, 0.9965, 0.9959, 0.9954, 0.995, 0.9938, 0.9942, 0.9924,
0.9924, 0.9907, 0.9905, 0.9878, 0.9867, 0.9845, 0.9808, 0.9773, 0.9715, 0.9652, 0.9529,
0.9347, 0.9011, 0.83,
)
_MAGCACHE_480P_RATIOS = (
1.0, 1.0077, 1.0138, 1.0043, 1.0029, 0.9986, 0.9966, 1.0, 1.0006, 0.9996, 0.9993, 0.9986, 1.0,
0.9993, 0.9966, 0.9986, 0.9988, 0.9991, 0.998, 0.9977, 0.9976, 0.9971, 0.9973, 0.9969, 0.996,
0.9961, 0.9949, 0.9958, 0.9933, 0.9942, 0.9941, 0.9926, 0.9929, 0.9916, 0.9923, 0.9887, 0.99,
0.9882, 0.9865, 0.9833, 0.9827, 0.9791, 0.9763, 0.9718, 0.9657, 0.9563, 0.9454, 0.9264,
0.8967, 0.8382,
)
_MAGCACHE_FAMILY_RATIOS: dict[str, tuple[float, ...]] = {
"hunyuanvideo-1.5": _MAGCACHE_480P_RATIOS,
"hunyuanvideo-1.5-720p": _MAGCACHE_720P_RATIOS,
}
# Families whose AUTO step-cache decision engages MagCache instead of FBCache. On
# HunyuanVideo-1.5 FBCache free-runs (no skip cap, no error budget) and derails the
# trajectory (LPIPS 0.54 + a luma shift at its default threshold), while MagCache holds
# the same composition at 1.5x -- see the constants above. Every other family keeps the
# measured FBCache default. An EXPLICIT "fbcache"/"magcache" request always wins.
_FAMILY_AUTO_CACHE_MODE: dict[str, str] = {
"hunyuanvideo-1.5": TC_MAGCACHE,
"hunyuanvideo-1.5-720p": TC_MAGCACHE,
}
def auto_cache_mode(family: Optional[str]) -> str:
"""The cache mode the AUTO policy engages for ``family`` (mode only; the step-count
bar and the engage call are the caller's job). MagCache additionally needs a
calibrated ratio curve: a family routed here without one runs uncached (the
apply_step_cache magcache branch checks), never silently falls back to FBCache."""
return _FAMILY_AUTO_CACHE_MODE.get(str(family or "").strip().lower(), TC_FBCACHE)
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).
@ -64,6 +121,81 @@ def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
return normalized
# Transformer block classes whose FBCache metadata is missing from the installed
# diffusers. The First-Block-Cache hook reads each block's (hidden_states,
# encoder_hidden_states) return layout from TransformerBlockRegistry; diffusers 0.39
# registers the HunyuanVideo 1.0 blocks but not the 1.5 ones, so enable_cache raises
# "Model class HunyuanVideo15TransformerBlock not registered" on a DiT that is
# otherwise fully cache-compatible: CacheMixin, one homogeneous ``transformer_blocks``
# list of residual-additive dual-stream blocks returning (hidden_states,
# encoder_hidden_states) -- the exact layout of the registered 1.0 block. Keyed by the
# TRANSFORMER class name so only a family that needs the patch pays for it, and probed
# via TransformerBlockRegistry.get first so a diffusers release that ships the
# registration natively makes this a no-op.
# transformer class -> ((block module, block class, hs index, ehs index), ...)
_EXTRA_BLOCK_METADATA: dict[str, tuple[tuple[str, str, int, Optional[int]], ...]] = {
"HunyuanVideo15Transformer3DModel": (
(
"diffusers.models.transformers.transformer_hunyuan_video15",
"HunyuanVideo15TransformerBlock",
0,
1,
),
),
}
def _ensure_block_metadata_registered(transformer: Any, logger: Any = None) -> None:
"""Register the missing FBCache block metadata for ``transformer``'s family (see
``_EXTRA_BLOCK_METADATA``). Best-effort: a failure just leaves enable_cache to raise
its own error and the load runs uncached, exactly as before this patch."""
specs = _EXTRA_BLOCK_METADATA.get(type(transformer).__name__)
if not specs:
return
try:
import importlib
from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
for module_name, cls_name, hs_index, ehs_index in specs:
block_cls = getattr(importlib.import_module(module_name), cls_name)
try:
TransformerBlockRegistry.get(block_cls)
continue # a newer diffusers registers it natively
except ValueError:
pass
TransformerBlockRegistry.register(
block_cls,
TransformerBlockMetadata(
return_hidden_states_index = hs_index,
return_encoder_hidden_states_index = ehs_index,
),
)
if logger is not None:
logger.info("diffusion.cache: registered %s block metadata for fbcache", cls_name)
except Exception as exc: # noqa: BLE001 -- best-effort; enable_cache surfaces the real error
_warn(logger, "block metadata registration", exc)
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)."""
registry = getattr(transformer, "_diffusers_hook", None)
if registry is not None and getattr(registry, "_child_registries_cache", None) is not None:
try:
registry._child_registries_cache = None
except Exception: # noqa: BLE001 -- diffusers internals moved; leave as-is
pass
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
@ -91,12 +223,17 @@ def apply_step_cache(
mode: Optional[str],
threshold: Optional[float] = None,
quant_active: bool = False,
family: Optional[str] = None,
steps: Optional[int] = None,
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."""
the default; ``quant_active`` raises the FBCache default so the cache still triggers on
a quantised transformer. The magcache mode additionally needs ``family`` (to look up the
calibrated ratio curve) and ``steps`` (MagCache interpolates that curve over the
configured step count and sizes its no-skip retention window from it). 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
@ -105,11 +242,14 @@ def apply_step_cache(
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)
)
if mode == TC_MAGCACHE:
thr = threshold if threshold is not None else DEFAULT_MAGCACHE_THRESHOLD
else:
thr = (
threshold
if threshold is not None
else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD)
)
# Engage only via the transformer's native enable_cache (the diffusers CacheMixin path):
# the lower-level apply_first_block_cache hook would install on a non-CacheMixin
# transformer too (e.g. Z-Image), whose pipeline opens no cache_context and would crash
@ -129,16 +269,51 @@ def apply_step_cache(
logger, mode, RuntimeError("pipeline __call__ opens no cache_context; running uncached")
)
return None
# Some cache-compatible block classes are missing from the installed diffusers'
# FBCache metadata registry (HunyuanVideo-1.5); register them before enable_cache.
# Both hook families (FBCache / MagCache) read the same block metadata.
_ensure_block_metadata_registered(transformer, logger)
try:
try:
from diffusers import FirstBlockCacheConfig
except ImportError: # older diffusers exports it only from diffusers.hooks
from diffusers.hooks import FirstBlockCacheConfig
if mode == TC_MAGCACHE:
ratios = _MAGCACHE_FAMILY_RATIOS.get(str(family or "").strip().lower())
if ratios is None:
# No silent FBCache fallback: the family was routed to magcache exactly
# because FBCache derails it, so an uncalibrated checkpoint runs uncached.
_warn(
logger, mode,
RuntimeError(f"no calibrated mag_ratios for family '{family}'"),
)
return None
if not steps or int(steps) <= 0:
_warn(logger, mode, RuntimeError("magcache needs the step count to engage"))
return None
from diffusers.hooks import MagCacheConfig
config = FirstBlockCacheConfig(threshold = thr)
config: Any = MagCacheConfig(
threshold = thr,
max_skip_steps = MAGCACHE_MAX_SKIP_STEPS,
retention_ratio = MAGCACHE_RETENTION_RATIO,
num_inference_steps = int(steps),
mag_ratios = list(ratios),
)
# The curve is interpolated over the CONFIGURED step count, so the marker
# carries it: the auto toggle re-engages on a step-count change.
marker = f"{mode}@{thr}#s{int(steps)}"
else:
try:
from diffusers import FirstBlockCacheConfig
except ImportError: # older diffusers exports it only from diffusers.hooks
from diffusers.hooks import FirstBlockCacheConfig
config = FirstBlockCacheConfig(threshold = thr)
marker = f"{mode}@{thr}"
enable_cache(config)
# A prior uncached generation may have frozen an empty child-registry list on
# the transformer's HookRegistry; the block hooks just installed would then
# never receive the cache context. Must follow every enable_cache.
_invalidate_child_registry_cache(transformer)
try:
transformer._unsloth_step_cache = f"{mode}@{thr}"
transformer._unsloth_step_cache = marker
except Exception: # noqa: BLE001 — marker is best-effort
pass
if logger is not None:
@ -195,49 +370,73 @@ def effective_request_strength(
return pipe_default_strength if isinstance(pipe_default_strength, (int, float)) else None
def _disengage_step_cache(transformer: Any, *, reason: str, logger: Any = None) -> bool:
"""disable_cache + clear the marker; True when the transformer is now uncached."""
disable_cache = getattr(transformer, "disable_cache", None)
if not callable(disable_cache):
return False
try:
disable_cache()
transformer._unsloth_step_cache = None
if logger is not None:
logger.info("diffusion.cache: step cache disengaged (%s)", reason)
return True
except Exception as exc: # noqa: BLE001 -- keep the cache rather than crash
_warn(logger, "step cache disable", exc)
return False
def maybe_toggle_step_cache(
pipe: Any,
*,
steps: int,
quant_active: bool = False,
threshold: Optional[float] = None,
mode: str = TC_FBCACHE,
family: Optional[str] = None,
logger: Any = None,
) -> Optional[str]:
"""Generation-time enable/disable for an AUTO cache decision, keyed on the actual
step count: engage FBCache at ``FBCACHE_MIN_STEPS`` or more, run uncached below it.
Idempotent (the ``_unsloth_step_cache`` marker tracks the engaged state), so calling
it on every generation is cheap. Only the loader's auto path calls this; an explicit
user choice is never toggled. Returns the mode now active (or None when uncached)."""
step count: engage ``mode`` (the family's auto cache mode) at ``FBCACHE_MIN_STEPS``
or more, run uncached below it. Idempotent (the ``_unsloth_step_cache`` marker tracks
the engaged state), so calling it on every generation is cheap -- except a magcache
step-count change, which re-engages so the ratio curve is re-interpolated over the
actual schedule. Only the loader's auto path calls this; an explicit user choice is
never toggled. Returns the mode now active (or None when uncached)."""
transformer = getattr(pipe, "transformer", None)
if transformer is None:
return None
engaged = getattr(transformer, "_unsloth_step_cache", None)
want = int(steps) >= FBCACHE_MIN_STEPS
if (
want
and engaged
and mode == TC_MAGCACHE
and f"#s{int(steps)}" not in str(engaged)
and _disengage_step_cache(
transformer, reason = f"magcache re-interpolating for {steps} steps", logger = logger
)
):
engaged = None
if want and not engaged:
return apply_step_cache(
pipe,
mode = TC_FBCACHE,
mode = mode,
threshold = threshold,
quant_active = quant_active,
family = family,
steps = steps,
logger = logger,
)
if not want and engaged:
disable_cache = getattr(transformer, "disable_cache", None)
if callable(disable_cache):
try:
disable_cache()
transformer._unsloth_step_cache = None
if logger is not None:
logger.info(
"diffusion.cache: fbcache disengaged (auto: %s steps < %s)",
steps,
FBCACHE_MIN_STEPS,
)
return None
except Exception as exc: # noqa: BLE001 — keep the cache rather than crash
_warn(logger, "fbcache disable", exc)
return TC_FBCACHE
return TC_FBCACHE if engaged else None
if _disengage_step_cache(
transformer,
reason = f"auto: {steps} steps < {FBCACHE_MIN_STEPS}",
logger = logger,
):
return None
return mode
return mode if engaged else None
def _warn(logger: Any, what: str, exc: Exception) -> None: