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.
This commit is contained in:
Daniel Han
parent
d54e2bc5f0
commit
c998183cc2
6 changed files with 637 additions and 50 deletions
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(), ())
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ───────────────────────────────────────────────────────────────────────
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue