From c998183cc25da90e16dde2fe62ba4768ef62c4ac Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Jul 2026 10:47:03 +0000 Subject: [PATCH] feat(video): step caching for HunyuanVideo-1.5 (MagCache auto, FBCache registry) + int8 trim fix 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. --- scripts/video_speedmem_bench.py | 79 ++++- .../backend/core/inference/diffusion_cache.py | 269 ++++++++++++++--- .../inference/diffusion_transformer_quant.py | 35 ++- studio/backend/core/inference/video.py | 18 +- studio/backend/tests/test_diffusion_cache.py | 270 ++++++++++++++++++ .../tests/test_diffusion_transformer_quant.py | 16 +- 6 files changed, 637 insertions(+), 50 deletions(-) diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index 9f774d1655..96b39bff7e 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -268,6 +268,29 @@ _CONFIGS: dict[str, dict[str, Any]] = { "ditint8_nocache": dict( te = "none", vae = "none", dit = "int8", speed = "default", attn = "native", cache = "off" ), + # Hunyuan padded-text trim isolation: "cudnn" above is the same stack WITH the trim + # (it auto-engages under any active speed tier), so trim_off isolates its win. The + # trim key defaults True everywhere else; only this row forces it off. + "trim_off": dict( + te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off", + trim = False, + ), + # Compile isolation at matched attention/trim: eager tier (channels_last + cudnn + # benchmark, NO compile) vs the "cudnn" row (default tier = regional compile). + "eager_trim": dict( + te = "none", vae = "none", dit = "none", speed = "eager", attn = "auto", cache = "off" + ), + # int8 DiT baseline at the production attention stack (cudnn + trim), cache off, + # directly comparable to the "cudnn" dense row. + "int8_cudnn": dict( + te = "none", vae = "none", dit = "int8", speed = "default", attn = "auto", cache = "off" + ), + # The full companion-quant stack WITHOUT step caching: te/vae auto + dit auto (dense + # -fit skip on Hunyuan) + compile + cudnn + trim. The stacked-best candidate default + # when the family's auto cache policy stays off. + "shipped_nocache": dict( + te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "auto", cache = "off" + ), } @@ -335,6 +358,7 @@ def _apply_levers( fam_obj, force_fp32_vae: bool, default_steps: int, + cache_threshold: Optional[float] = None, logger = None, ) -> dict: """Apply the configured levers with the loader's own argument values, in the loader's order: @@ -355,7 +379,7 @@ def _apply_levers( ) from core.inference.diffusion_cache import ( apply_step_cache, - TC_FBCACHE, + auto_cache_mode, FBCACHE_MIN_STEPS, ) @@ -436,14 +460,18 @@ def _apply_levers( # Step cache FIRST (compile keys fullgraph off an active cache); per expert. cache_active = False if cfg["cache"] == "auto": - cache_request = TC_FBCACHE if default_steps >= FBCACHE_MIN_STEPS else None + # Per-family auto mode, exactly like the loader (video.py): MagCache for the + # HunyuanVideo-1.5 families, FBCache elsewhere. + cache_request = auto_cache_mode(fam_name) if default_steps >= FBCACHE_MIN_STEPS else None if cache_request is not None: for v in views: engaged["cache"] = apply_step_cache( v, mode = cache_request, - threshold = None, + threshold = cache_threshold, quant_active = dit_quant_active, + family = fam_name, + steps = default_steps, logger = logger, ) cache_active = engaged["cache"] not in (None, "off") @@ -453,7 +481,7 @@ def _apply_levers( # text tokens so the fused SDPA kernel runs (~18x/DiT-forward, cosine ~1.0). A speed lever, so # gated on an active tier like the loader; no-op for every non-Hunyuan family. trim_engaged = False - if speed_active: + if speed_active and cfg.get("trim", True): for v in views: trim_engaged = install_hunyuan_attention_trim(v, fam_obj, logger = logger) or trim_engaged engaged["attn_trim"] = trim_engaged @@ -493,13 +521,15 @@ def _timed_video( dit_quant_active, default_steps, guidance_via_guider = False, + cache_threshold = None, + family = None, logger = None, ): - """One clip generation. Re-checks FBCache per generation (maybe_toggle_step_cache) exactly - like the loader, then times total + per-step. Returns (output, total_s, [per_step_ms]).""" + """One clip generation. Re-checks the step cache per generation (maybe_toggle_step_cache) + exactly like the loader, then times total + per-step. Returns (output, total_s, [per_step_ms]).""" import torch - from core.inference.diffusion_cache import maybe_toggle_step_cache, FBCACHE_MIN_STEPS + from core.inference.diffusion_cache import auto_cache_mode, maybe_toggle_step_cache if cache_mode == "auto": # Toggle on EVERY expert view, exactly like the loader's per-view recheck @@ -512,7 +542,13 @@ def _timed_video( for v in views: try: maybe_toggle_step_cache( - v, steps = steps, quant_active = dit_quant_active, threshold = None, logger = logger + v, + steps = steps, + quant_active = dit_quant_active, + threshold = cache_threshold, + mode = auto_cache_mode(family), + family = family, + logger = logger, ) except Exception: pass @@ -592,6 +628,7 @@ def _run_config( seed: int, iters: int, out: Path, + cache_threshold: Optional[float] = None, logger = None, ): import numpy as np @@ -627,6 +664,7 @@ def _run_config( fam_obj = fam_obj, force_fp32_vae = force_fp32, default_steps = default_steps, + cache_threshold = cache_threshold, logger = logger, ) pipe = pipe.to("cuda") @@ -638,6 +676,7 @@ def _run_config( cache_mode = cfg["cache"] # warmup (pays the one-time compile / autotune) + warmup_t0 = time.perf_counter() _timed_video( pipe, steps = steps, @@ -650,8 +689,11 @@ def _run_config( dit_quant_active = dit_active, default_steps = default_steps, guidance_via_guider = gvg, + cache_threshold = cache_threshold, + family = family, logger = logger, ) + warmup_s = time.perf_counter() - warmup_t0 _reset_peak() dts, steps_ms = [], [] last_out = None @@ -668,6 +710,8 @@ def _run_config( dit_quant_active = dit_active, default_steps = default_steps, guidance_via_guider = gvg, + cache_threshold = cache_threshold, + family = family, logger = logger, ) dts.append(dt) @@ -700,6 +744,14 @@ def _run_config( ) except Exception: pass + # Persist EVERY config's frames too, so a row generated before the reference exists + # (parallel per-GPU processes) can be LPIPS-rescored offline instead of publishing null. + if arrs: + try: + import numpy as _np + _np.savez_compressed(out / f"frames_{family}_{name}.npz", *arrs) + except Exception: + pass row = { "config": name, @@ -713,9 +765,13 @@ def _run_config( "cache": engaged["cache"] or "off", "effective_speed": engaged.get("_effective_speed"), "speed_optims": engaged["speed_optims"], + "attn_trim": engaged.get("attn_trim", False), + "cache_threshold": cache_threshold, + "cache_marker": getattr(getattr(pipe, "transformer", None), "_unsloth_step_cache", None), "load_peak_gb": round(load_peak, 2), "weights_gb": round(weights_gb, 2), "gen_peak_gb": round(gen_peak, 2), + "warmup_s": round(warmup_s, 3), "gen_latency_s": round(_median(dts), 3), "per_step_ms": round(_median(steps_ms), 1), "n_frames": len(arrs), @@ -739,6 +795,12 @@ def main(argv = None) -> int: ap.add_argument("--seed", type = int, default = 42) ap.add_argument("--iters", type = int, default = 3) ap.add_argument("--out", default = "outputs/video_speedmem") + ap.add_argument( + "--cache-threshold", + type = float, + default = None, + help = "FBCache residual-diff threshold override (None -> the production default)", + ) args = ap.parse_args(argv) import logging @@ -793,6 +855,7 @@ def main(argv = None) -> int: seed = args.seed, iters = args.iters, out = out, + cache_threshold = args.cache_threshold, logger = logger, ) if n == "reference": diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index ad5eb50127..b14000d56c 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -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: diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 6ca1d39618..974b88fe3a 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -82,6 +82,33 @@ _INT8_EXCLUDE_NAME_TOKENS = ( "time_embed", ) +# int8 PER-FAMILY name exclusions, on top of _INT8_EXCLUDE_NAME_TOKENS. The attention trim +# (diffusion_attention's pre-hook) shrinks HunyuanVideo-1.5's text streams from their padded +# lengths to the VALID token counts, so every text-stream Linear runs at a tiny M the int8 +# dynamic path cannot handle (both failures measured on B200): +# - M = 0 (the t2v byt5 / image streams trim to zero tokens): torchao returns the input +# UNPROJECTED (a quantized 1472->2048 Linear maps [1, 0, 1472] to [1, 0, 1472]), so the +# 2048-wide cond-type add crashes -> context_embedder_2 / image_embedder; +# - M <= 16 (a short prompt, or the empty negative prompt's ~6 tokens): torch._int_mm +# requires M > 16 and raises -> the TokenRefiner (context_embedder.*) and every block's +# context-stream projections (add_q/k/v_proj, to_add_out, ff_context). +# All of these run at M = text tokens (tens) vs the video stream's M ~ 32k+, so keeping +# them bf16 costs nothing measurable; the attention/FFN video-stream linears keep the full +# int8 coverage. "context_embedder" also matches "context_embedder_2" (substring check). +_HUNYUAN15_INT8_EXCLUDES = ( + "context_embedder", + "image_embedder", + "add_q_proj", + "add_k_proj", + "add_v_proj", + "to_add_out", + "ff_context", +) +_INT8_FAMILY_EXCLUDE_NAME_TOKENS: dict[str, tuple[str, ...]] = { + "hunyuanvideo-1.5": _HUNYUAN15_INT8_EXCLUDES, + "hunyuanvideo-1.5-720p": _HUNYUAN15_INT8_EXCLUDES, +} + # fp8 (and the other per-row scaled_mm schemes) PER-FAMILY name exclusions. Per-row fp8 scales # each activation ROW by row_amax / 448; a row whose amax is 0 -- a PADDING token in a @@ -111,7 +138,9 @@ def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tupl """Name tokens to exclude from quantisation for ``scheme`` (optionally family-specific). int8 (torch._int_mm, M>16) skips the M=1 modulation / conditioning-embedder projections (see - _INT8_EXCLUDE_NAME_TOKENS) on every family. The per-row scaled_mm schemes (fp8 / mxfp8 / nvfp4) + _INT8_EXCLUDE_NAME_TOKENS) on every family, plus the per-family zero-token embedders whose + M=0 input torchao passes through unprojected (_INT8_FAMILY_EXCLUDE_NAME_TOKENS, HunyuanVideo + -1.5 under the attention trim). The per-row scaled_mm schemes (fp8 / mxfp8 / nvfp4) exclude nothing by default, but on families whose zero-padded conditioning sequence would divide by a zero row scale they skip the offending input embedder (see _FP8_FAMILY_EXCLUDE_NAME_TOKENS). ``family=None`` preserves the historical behaviour (int8 tokens, or () otherwise). Shared by the @@ -120,7 +149,9 @@ def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tupl that then crashes (int8 M=1 -> _int_mm) or infs (fp8 padding row -> scaled_mm) at the first denoise step.""" if scheme == TQ_INT8: - return _INT8_EXCLUDE_NAME_TOKENS + return _INT8_EXCLUDE_NAME_TOKENS + _INT8_FAMILY_EXCLUDE_NAME_TOKENS.get( + str(family or "").strip().lower(), () + ) return _FP8_FAMILY_EXCLUDE_NAME_TOKENS.get(str(family or "").strip().lower(), ()) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index f0863bd4d4..a9c8bf959b 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -50,8 +50,8 @@ from .diffusion_attention import ( from .diffusion_cache import ( FBCACHE_MIN_STEPS, TC_AUTO, - TC_FBCACHE, apply_step_cache, + auto_cache_mode, maybe_toggle_step_cache, normalize_transformer_cache, ) @@ -1317,10 +1317,16 @@ class VideoBackend: # GGUF checkpoints and torchao-quantised DiTs both need the higher quantised # threshold for the cache to still trigger over the quant noise. cache_quant_active = kind == "gguf" or transformer_quant_engaged is not None - default_cache_steps: Optional[int] = None + # Computed for every request: the auto step-count policy keys on it, and an + # explicit magcache request needs it too (the ratio curve is interpolated over + # the configured step count). + default_cache_steps, _ = default_video_generation_params(gguf_filename, repo_id, base) if cache_auto: - default_cache_steps, _ = default_video_generation_params(gguf_filename, repo_id, base) - cache_request = TC_FBCACHE if default_cache_steps >= FBCACHE_MIN_STEPS else None + # The engaged CACHE MODE is per-family: MagCache where FBCache's uncapped + # skipping derails the trajectory (HunyuanVideo-1.5), FBCache elsewhere. + cache_request = ( + auto_cache_mode(fam.name) if default_cache_steps >= FBCACHE_MIN_STEPS else None + ) cache_engaged = None for view in views: engaged = apply_step_cache( @@ -1332,6 +1338,8 @@ class VideoBackend: # (diffusion.py): both an engaged transformer_quant AND a GGUF checkpoint # (quantized weights) count as quant-active here (cache_quant_active, L1172). quant_active = cache_quant_active, + family = fam.name, + steps = default_cache_steps, logger = logger, ) if view is pipe: @@ -1715,6 +1723,8 @@ class VideoBackend: steps = steps, quant_active = state.cache_quant_active, threshold = state.cache_threshold, + mode = auto_cache_mode(fam.name), + family = fam.name, logger = logger, ) if toggled != state.transformer_cache: diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 1794482239..560a075a11 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -356,3 +356,273 @@ def test_toggle_noop_without_cache_support(monkeypatch): def test_toggle_noop_without_transformer(): assert maybe_toggle_step_cache(types.SimpleNamespace(), steps = 28) is None + + +# ── FBCache block-metadata registration (HunyuanVideo-1.5) ───────────────────────── +from core.inference.diffusion_cache import ( # noqa: E402 + _ensure_block_metadata_registered, + _invalidate_child_registry_cache, +) + + +def _stub_hunyuan15_registry(monkeypatch): + """Stub the two diffusers modules the registration helper imports: the FBCache + metadata registry (diffusers.hooks._helpers) and the HunyuanVideo-1.5 transformer + module carrying the block class. Returns (registry_cls, block_cls).""" + + class _Metadata: + def __init__( + self, + return_hidden_states_index = None, + return_encoder_hidden_states_index = None, + ): + self.return_hidden_states_index = return_hidden_states_index + self.return_encoder_hidden_states_index = return_encoder_hidden_states_index + + class _Registry: + registry: dict = {} + + @classmethod + def get(cls, model_class): + if model_class not in cls.registry: + raise ValueError(f"Model class {model_class} not registered.") + return cls.registry[model_class] + + @classmethod + def register(cls, model_class, metadata): + cls.registry[model_class] = metadata + + class HunyuanVideo15TransformerBlock: # the sentinel block class + pass + + helpers = types.ModuleType("diffusers.hooks._helpers") + helpers.TransformerBlockMetadata = _Metadata + helpers.TransformerBlockRegistry = _Registry + monkeypatch.setitem(sys.modules, "diffusers.hooks._helpers", helpers) + + blocks = types.ModuleType( + "diffusers.models.transformers.transformer_hunyuan_video15" + ) + blocks.HunyuanVideo15TransformerBlock = HunyuanVideo15TransformerBlock + monkeypatch.setitem( + sys.modules, + "diffusers.models.transformers.transformer_hunyuan_video15", + blocks, + ) + return _Registry, HunyuanVideo15TransformerBlock + + +class HunyuanVideo15Transformer3DModel(_MixinTransformer): + """A CacheMixin-style fake whose CLASS NAME keys the extra-metadata table.""" + + +def test_hunyuan15_block_metadata_is_registered(monkeypatch): + registry, block_cls = _stub_hunyuan15_registry(monkeypatch) + _ensure_block_metadata_registered(HunyuanVideo15Transformer3DModel()) + meta = registry.registry[block_cls] + # The 1.5 dual-stream block returns (hidden_states, encoder_hidden_states) -- the + # same layout as the natively registered HunyuanVideo 1.0 block. + assert meta.return_hidden_states_index == 0 + assert meta.return_encoder_hidden_states_index == 1 + + +def test_hunyuan15_registration_defers_to_a_native_one(monkeypatch): + # A diffusers release that ships the registration natively must win: the helper + # probes TransformerBlockRegistry.get first and never overwrites. + registry, block_cls = _stub_hunyuan15_registry(monkeypatch) + native = object() + registry.registry[block_cls] = native + _ensure_block_metadata_registered(HunyuanVideo15Transformer3DModel()) + assert registry.registry[block_cls] is native + + +def test_registration_noop_for_other_families(monkeypatch): + registry, _ = _stub_hunyuan15_registry(monkeypatch) + _ensure_block_metadata_registered(_MixinTransformer()) + assert registry.registry == {} + + +def test_registration_failure_is_swallowed(monkeypatch): + # diffusers internals moved / import fails -> best-effort no-op; enable_cache then + # surfaces its own error and the load runs uncached, exactly as before the patch. + monkeypatch.setitem(sys.modules, "diffusers.hooks._helpers", None) + _ensure_block_metadata_registered(HunyuanVideo15Transformer3DModel()) # no raise + + +# ── stale child-registry invalidation after enable_cache ─────────────────────────── +def test_invalidate_child_registry_cache_clears_stale_list(): + # An UNCACHED generation's cache_context call froze an EMPTY child list on the + # transformer's HookRegistry; enable_cache installs block hooks _set_context would + # then never reach ("No context is set" on the first cached forward). The helper + # drops the stale cache so the next cache_context rebuilds it over the new hooks. + t = _MixinTransformer() + t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = []) + _invalidate_child_registry_cache(t) + assert t._diffusers_hook._child_registries_cache is None + + +def test_invalidate_child_registry_cache_noops(): + _invalidate_child_registry_cache(_MixinTransformer()) # no _diffusers_hook + t = _MixinTransformer() + t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = None) + _invalidate_child_registry_cache(t) # nothing cached yet + assert t._diffusers_hook._child_registries_cache is None + + +def test_apply_step_cache_registers_and_invalidates_for_hunyuan15(monkeypatch): + _stub_diffusers(monkeypatch) + registry, block_cls = _stub_hunyuan15_registry(monkeypatch) + t = HunyuanVideo15Transformer3DModel() + t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = []) + engaged = apply_step_cache(_pipe(t), mode = "fbcache") + assert engaged == TC_FBCACHE + assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD + assert block_cls in registry.registry # metadata registered before enable_cache + assert t._diffusers_hook._child_registries_cache is None # stale cache dropped + + +# ── magcache mode (per-family auto cache) ────────────────────────────────────────── +from core.inference.diffusion_cache import ( # noqa: E402 + DEFAULT_MAGCACHE_THRESHOLD, + MAGCACHE_MAX_SKIP_STEPS, + MAGCACHE_RETENTION_RATIO, + TC_MAGCACHE, + _MAGCACHE_FAMILY_RATIOS, + auto_cache_mode, +) + + +class _MagConfig: + def __init__( + self, + threshold, + max_skip_steps, + retention_ratio, + num_inference_steps, + mag_ratios, + ): + self.threshold = threshold + self.max_skip_steps = max_skip_steps + self.retention_ratio = retention_ratio + self.num_inference_steps = num_inference_steps + self.mag_ratios = mag_ratios + + +def _stub_diffusers_with_magcache(monkeypatch): + _stub_diffusers(monkeypatch) + hooks = sys.modules["diffusers.hooks"] + hooks.MagCacheConfig = _MagConfig + + +def test_normalize_accepts_magcache(): + assert normalize_transformer_cache("magcache") == TC_MAGCACHE + assert normalize_transformer_cache("MagCache") == TC_MAGCACHE + + +def test_auto_cache_mode_per_family(): + # HunyuanVideo-1.5: FBCache free-runs (no cap / no error budget) and derails the + # trajectory (measured LPIPS 0.54 at its default threshold), so auto engages the + # bounded MagCache there; every other family keeps the measured FBCache default. + assert auto_cache_mode("hunyuanvideo-1.5") == TC_MAGCACHE + assert auto_cache_mode("hunyuanvideo-1.5-720p") == TC_MAGCACHE + assert auto_cache_mode("HunyuanVideo-1.5-720p") == TC_MAGCACHE + for other in (None, "", "flux", "wan2.2-ti2v-5b", "ltx-2", "z-image"): + assert auto_cache_mode(other) == TC_FBCACHE + + +def test_magcache_families_have_calibrated_ratios(): + # Every family the auto policy routes to magcache must ship a calibrated curve, or + # the auto default silently runs uncached (apply_step_cache checks the table). + from core.inference.diffusion_cache import _FAMILY_AUTO_CACHE_MODE + + for fam, mode in _FAMILY_AUTO_CACHE_MODE.items(): + if mode == TC_MAGCACHE: + ratios = _MAGCACHE_FAMILY_RATIOS[fam] + assert len(ratios) == 50 # the default 50-step schedule they were calibrated on + assert all(0.5 < r < 1.5 for r in ratios) + + +def test_magcache_engages_with_family_curve(monkeypatch): + _stub_diffusers_with_magcache(monkeypatch) + t = _MixinTransformer() + engaged = apply_step_cache( + _pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50 + ) + assert engaged == TC_MAGCACHE + cfg = t.enabled_with + assert cfg.threshold == DEFAULT_MAGCACHE_THRESHOLD + assert cfg.max_skip_steps == MAGCACHE_MAX_SKIP_STEPS + assert cfg.retention_ratio == MAGCACHE_RETENTION_RATIO + assert cfg.num_inference_steps == 50 + assert cfg.mag_ratios == list(_MAGCACHE_FAMILY_RATIOS["hunyuanvideo-1.5-720p"]) + # The marker carries the step count so the auto toggle re-engages on a change. + assert t._unsloth_step_cache == f"magcache@{DEFAULT_MAGCACHE_THRESHOLD}#s50" + + +def test_magcache_without_calibration_runs_uncached(monkeypatch): + # No silent FBCache fallback: the family was routed to magcache exactly because + # FBCache derails it, so an uncalibrated family must run uncached instead. + _stub_diffusers_with_magcache(monkeypatch) + t = _MixinTransformer() + assert apply_step_cache(_pipe(t), mode = "magcache", family = "flux", steps = 50) is None + assert t.enabled_with is None + + +def test_magcache_without_steps_runs_uncached(monkeypatch): + _stub_diffusers_with_magcache(monkeypatch) + t = _MixinTransformer() + assert ( + apply_step_cache(_pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p") is None + ) + assert t.enabled_with is None + + +def test_magcache_explicit_threshold_wins(monkeypatch): + _stub_diffusers_with_magcache(monkeypatch) + t = _MixinTransformer() + apply_step_cache( + _pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 30, + threshold = 0.24, + ) + assert t.enabled_with.threshold == 0.24 + + +def test_toggle_engages_family_magcache(monkeypatch): + _stub_diffusers_with_magcache(monkeypatch) + t = _ToggleTransformer() + mode = maybe_toggle_step_cache( + _pipe(t), steps = 30, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" + ) + assert mode == TC_MAGCACHE and t.enables == 1 + assert t.enabled_with.num_inference_steps == 30 + + +def test_toggle_magcache_reengages_on_step_change(monkeypatch): + # MagCache interpolates its calibrated curve over the CONFIGURED step count, so a + # step-count change must disable + re-enable; the same count stays idempotent. + _stub_diffusers_with_magcache(monkeypatch) + t = _ToggleTransformer() + maybe_toggle_step_cache( + _pipe(t), steps = 30, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" + ) + maybe_toggle_step_cache( + _pipe(t), steps = 30, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" + ) + assert t.enables == 1 and t.disables == 0 # idempotent at the same count + mode = maybe_toggle_step_cache( + _pipe(t), steps = 50, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" + ) + assert mode == TC_MAGCACHE and t.disables == 1 and t.enables == 2 + assert t.enabled_with.num_inference_steps == 50 + + +def test_toggle_magcache_disengages_below_bar(monkeypatch): + _stub_diffusers_with_magcache(monkeypatch) + t = _ToggleTransformer() + maybe_toggle_step_cache( + _pipe(t), steps = 30, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" + ) + mode = maybe_toggle_step_cache( + _pipe(t), steps = 8, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" + ) + assert mode is None and t.disables == 1 diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index ff29c7f4e6..a6617d73ce 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -461,8 +461,22 @@ def test_exclude_tokens_for_scheme_family(): # Hunyuan is not localisable (fp8 stays denied), and an unknown family gets nothing. assert exclude_tokens_for_scheme(TQ_FP8, "hunyuanvideo-1.5") == () assert exclude_tokens_for_scheme(TQ_FP8, "z-image") == () - # int8 is family-independent (its zero-row handling needs no per-family skip). + # int8 tolerates zero ROWS (per-token; no divide), so most families need no extra skip... assert exclude_tokens_for_scheme(TQ_INT8, "wan2.2-ti2v-5b") == _INT8_EXCLUDE_NAME_TOKENS + # ...but the attention trim shrinks HunyuanVideo-1.5's text streams to their VALID token + # counts, where the int8 dynamic path fails two ways (both measured): a zero-token (M=0) + # input passes through UNPROJECTED (byt5/image embedders on t2v -> cond-type add crash), + # and torch._int_mm requires M > 16 (the ~6-token empty negative prompt crashes the + # TokenRefiner and the blocks' context-stream projections). All the text-stream linears + # must stay bf16; the M ~ 32k video-stream linears keep the int8 coverage. + from core.inference.diffusion_transformer_quant import _HUNYUAN15_INT8_EXCLUDES + + for fam in ("hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"): + assert ( + exclude_tokens_for_scheme(TQ_INT8, fam) + == _INT8_EXCLUDE_NAME_TOKENS + _HUNYUAN15_INT8_EXCLUDES + ) + assert "context_embedder" in _HUNYUAN15_INT8_EXCLUDES # covers context_embedder_2 too # ── apply ───────────────────────────────────────────────────────────────────────