514f4c60fe
- Wan2.2-A14B step cache: pin the balanced FBCache threshold to 0.08 even when quant is active (per-family override in diffusion_cache.py). Auto-fp8 made the generic quant promotion (0.12) the family's effective default at pairwise LPIPS 0.128, over the 0.08 quality gate the balanced preset is held to. Measured operating point with fp8 actually engaged (1280x720/81f/50 steps, B200): fb@0.08 = 1.08x at 0.129 vs the old fb@0.12 = 2.58x at 0.181; documented in the preset table. Explicit thresholds and the fast preset are unaffected. - MagCache curves: validated the shipped 33-frame calibrations at the production 121-frame default for hunyuanvideo-1.5-720p, hunyuanvideo-1.5 (480p) and wan2.2-ti2v-5b. Fresh 121-frame calibrations differ by <= 0.024 max abs entry and produce byte-identical frames at the auto presets (hv720 quality 1.69x at LPIPS 0.042, hv480 quality 1.66x at 0.018, wan5b balanced 1.74x at 0.026, all pairwise vs the same-load uncached stack), so the curves ship unchanged with the frame-count transfer documented next to them. - Dual-GPU CFG parallelism: the secondary-device pick now prefers a device whose name and compute capability match the primary, and the gate declines a mismatched pair in auto mode (eager kernel selection is arch-dependent, so the advertised bit-identity cannot hold across different GPU models); an explicit cfg_parallel=on proceeds but is downgraded to lossless=False with a warning. - A14B expert step cache is now all-or-none, mirroring the transactional quant loop: a mixed outcome (cache engaged on one expert but not the other) is rolled back and reported uncached with the failure reason, on both the load path and the generation-time auto toggle. - Partial torchao quantization is no longer reported as dense: after an in-place quantize_/caster failure, the DiT / text encoder / VAE is scanned for leftover torchao tensor-subclass parameters and the load fails with a clear error when any are found (a half-quantized module cannot run as dense, and offload's Module.to() crashes on torchao tensors). Failures that swapped nothing keep the best-effort dense fallback. - Cleanup: apply_attention_backend / apply_speed_optims / the attention trim are called once on the pipe (they already fan out over every DiT internally), so the second A14B expert no longer passes through them twice; the stale dual-DiT helper comment is rewritten to match the two helper shapes. Tests: device-identity picker/gate/lossy-plan coverage, per-family threshold pin scoping, all-or-none rollback in both failure directions, and partial-quant detection for all three quant modules.
880 lines
39 KiB
Python
880 lines
39 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Opt-in step caching for the diffusion transformer (First-Block-Cache).
|
|
|
|
Across denoising steps a DiT's output changes little once the trajectory settles, so most of
|
|
the transformer can be reused. First-Block-Cache (FBCache) computes the first block, and if
|
|
its residual barely changed from the previous step (within ``threshold``) it skips the
|
|
remaining blocks and reuses their cached output. diffusers ships it natively
|
|
(``transformer.enable_cache(FirstBlockCacheConfig(...))`` for CacheMixin models, or the
|
|
standalone ``apply_first_block_cache`` hook).
|
|
|
|
Measured on Flux.1-dev (28 steps, 1024px, B200): ~1.4x on top of torch.compile (2.83 ->
|
|
2.03 s) at LPIPS ~0.08 vs the no-cache output -- deep inside the speed-for-quality bar.
|
|
|
|
OFF by default and a deliberate per-load opt-in, because the win scales with step count: a
|
|
few-step distilled model (e.g. Z-Image-Turbo at ~8 steps) has almost no headroom and a
|
|
single skipped step is a large fraction of the trajectory, so caching is for many-step
|
|
models (Flux / Qwen-Image). It composes with torch.compile only with ``fullgraph=False``
|
|
(the cache's compiler-disabled decision is a graph break), which the speed layer switches to
|
|
automatically when a cache is engaged. Best-effort: an incompatible model (e.g. a transformer
|
|
whose block signature the hook does not recognise) is caught and the load proceeds uncached.
|
|
torch / diffusers imported lazily.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
TC_OFF = "off"
|
|
TC_AUTO = "auto"
|
|
TC_FBCACHE = "fbcache"
|
|
TC_MAGCACHE = "magcache"
|
|
TC_MODES = (TC_FBCACHE, TC_MAGCACHE)
|
|
|
|
# FBCache residual thresholds: higher skips more steps (faster, lower quality). The dense
|
|
# bf16 default; a quantised transformer shifts the residual distribution, so it needs a
|
|
# higher threshold for the cache to trigger at all (per ParaAttention's fp8 guidance).
|
|
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
|
|
|
|
# ── cache quality presets ──────────────────────────────────────────────────────────
|
|
# A user-facing speed/accuracy knob over the step cache's internals (threshold, skip cap,
|
|
# retention window). "balanced" is exactly the pre-knob shipped behaviour; "quality"
|
|
# trades most of the cache speedup for a near-lossless clip; "fast" skips more
|
|
# aggressively. An explicit transformer_cache_threshold always overrides the preset's
|
|
# threshold (the preset still supplies the magcache skip cap / retention window).
|
|
CQ_QUALITY = "quality"
|
|
CQ_BALANCED = "balanced"
|
|
CQ_FAST = "fast"
|
|
CACHE_QUALITY_LEVELS = (CQ_QUALITY, CQ_BALANCED, CQ_FAST)
|
|
|
|
# MagCache preset -> (threshold, max_skip_steps, retention_ratio). Calibrated on
|
|
# HunyuanVideo-1.5-720p (B200, 1280x720, 33 frames, 50 steps, pairwise LPIPS vs the same
|
|
# uncached trim+cudnn+compile stack, WITH the compiled hook inners below): quality
|
|
# (0.06, 2, 0.3) = 1.64x at LPIPS 0.050 (30 steps: 1.63x at 0.093) vs balanced
|
|
# (0.12, 3, 0.2) = 2.17x at LPIPS 0.129 (30 steps: 2.02x at 0.201). Skip counts bind on
|
|
# the cap + retention window below threshold ~0.12, which is why quality tightens all
|
|
# three rather than just the threshold. Re-measured at the production 121-frame default
|
|
# (same protocol): quality = 1.69x at 0.042 (720p) / 1.66x at 0.018 (480p) -- the
|
|
# 33-frame operating points transfer (see the ratio-curve note below).
|
|
_MAGCACHE_QUALITY_PRESETS: dict[str, tuple[float, int, float]] = {
|
|
CQ_QUALITY: (0.06, 2, 0.3),
|
|
CQ_BALANCED: (DEFAULT_MAGCACHE_THRESHOLD, MAGCACHE_MAX_SKIP_STEPS, MAGCACHE_RETENTION_RATIO),
|
|
CQ_FAST: (0.24, MAGCACHE_MAX_SKIP_STEPS, MAGCACHE_RETENTION_RATIO),
|
|
}
|
|
|
|
# FBCache preset -> threshold (dense, quant-active). "balanced" keeps the measured
|
|
# defaults (0.08 dense / 0.12 quantised); "quality" halves the trigger so the cache only
|
|
# reuses when the first-block residual is nearly static; "fast" uses the quantised
|
|
# threshold everywhere.
|
|
_FBCACHE_QUALITY_THRESHOLDS: dict[str, tuple[float, float]] = {
|
|
CQ_QUALITY: (0.04, 0.06),
|
|
CQ_BALANCED: (DEFAULT_FBCACHE_THRESHOLD, QUANT_FBCACHE_THRESHOLD),
|
|
CQ_FAST: (QUANT_FBCACHE_THRESHOLD, 0.15),
|
|
}
|
|
|
|
# Per-family FBCache threshold overrides: (family, preset) -> (dense, quant-active).
|
|
# Wan2.2-A14B: on reference hardware the UNSET precision auto-promotes to fp8, which
|
|
# made the generic quant-active promotion (0.08 -> 0.12) the family's EFFECTIVE
|
|
# default -- but fb@0.12 measures pairwise LPIPS 0.128 (2.88x, dense probe, B200,
|
|
# 1280x720/33f/50 steps), far over the <= 0.08 quality gate the balanced preset is
|
|
# held to, and the round-3 adjudication that kept FBCache as the auto mode assumed
|
|
# the 0.08 point (1.28x/0.098). No compliant faster point exists (fb@0.04 is SLOWER
|
|
# than uncached at 0.016; MagCache was measured worse), so balanced pins 0.08 with
|
|
# quant active too: quality first, the fast points stay one explicit preset away.
|
|
# Operating point with the shipped quant actually engaged (fp8 DiTs, production
|
|
# shape 1280x720/81f/50 steps, pairwise vs the same fp8 uncached load, B200):
|
|
# fb@0.08 = 1.08x at LPIPS 0.129, vs the old effective default fb@0.12 = 2.58x at
|
|
# 0.181 -- fp8's residual noise makes the quantised cache trip the quality gate at
|
|
# ANY speedup, so 0.08 is the least-drift cache-on point, not a compliant one.
|
|
_FAMILY_FBCACHE_THRESHOLDS: dict[tuple[str, str], tuple[float, float]] = {
|
|
("wan2.2-t2v-a14b", CQ_BALANCED): (DEFAULT_FBCACHE_THRESHOLD, DEFAULT_FBCACHE_THRESHOLD),
|
|
}
|
|
|
|
|
|
def normalize_cache_quality(value: Optional[str]) -> Optional[str]:
|
|
"""Lower/strip a requested cache quality; None / "" / "auto" -> None (the loader
|
|
resolves it per family via ``auto_cache_quality``). Raises ValueError for an
|
|
unsupported value."""
|
|
if value is None:
|
|
return None
|
|
normalized = str(value).strip().lower()
|
|
if not normalized or normalized == "auto":
|
|
return None
|
|
if normalized not in CACHE_QUALITY_LEVELS:
|
|
raise ValueError(
|
|
f"Unsupported transformer_cache_quality '{value}'. Use one of: auto, "
|
|
f"{', '.join(CACHE_QUALITY_LEVELS)}."
|
|
)
|
|
return normalized
|
|
|
|
|
|
# Families whose UNSET cache quality resolves to the near-lossless "quality" preset
|
|
# instead of "balanced". HunyuanVideo-1.5 (both repacks) measured with the compiled
|
|
# cache inners (see _compile_hooked_block_inners): quality = 1.63-1.64x at pairwise
|
|
# LPIPS 0.05 (50 steps) / 0.09 (30 steps) vs balanced's 2.02-2.17x at 0.13-0.20 --
|
|
# the accuracy-first default keeps most of the speedup at under half the drift, and
|
|
# balanced / fast stay one explicit request away. Families without a measured quality
|
|
# point keep balanced (their pre-knob behaviour).
|
|
_FAMILY_AUTO_CACHE_QUALITY: dict[str, str] = {
|
|
"hunyuanvideo-1.5": CQ_QUALITY,
|
|
"hunyuanvideo-1.5-720p": CQ_QUALITY,
|
|
}
|
|
|
|
|
|
def auto_cache_quality(family: Optional[str]) -> str:
|
|
"""The cache quality preset an UNSET request resolves to for ``family``."""
|
|
return _FAMILY_AUTO_CACHE_QUALITY.get(str(family or "").strip().lower(), CQ_BALANCED)
|
|
|
|
|
|
# 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.
|
|
#
|
|
# Frame-count transfer VALIDATED at the production default (121 frames): these curves
|
|
# were calibrated on 33-frame clips, and recalibrating each family at its production
|
|
# shape (121f / 50 steps / default resolution, B200) moves the curve by <= 0.024 max
|
|
# abs entry diff (hv720 0.019, hv480 0.021, wan5b 0.024) -- small enough that the auto
|
|
# preset's skip schedule is UNCHANGED: the 33f-curve and fresh-121f-curve runs produced
|
|
# byte-identical frames on every family, so the 33f curves ship as-is. Measured at 121f
|
|
# with the shipped curves (pairwise LPIPS vs the same-load uncached compiled stack):
|
|
# hv720 quality 1.69x at 0.042, hv480 quality 1.66x at 0.018, wan5b balanced 1.74x at
|
|
# 0.026 -- each inside the gate its 33f adjudication was held to.
|
|
_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,
|
|
)
|
|
# Wan2.2-TI2V-5B, calibrated at 1280x704 / 33 frames / 50 steps on the family base
|
|
# checkpoint (B200, trim-less compiled stack). Cond/uncond branches agree within
|
|
# 0.0008, so one (conditional) curve serves both CFG contexts.
|
|
_MAGCACHE_WAN5B_RATIOS = (
|
|
1.0,
|
|
0.9906,
|
|
0.9996,
|
|
0.9936,
|
|
0.9968,
|
|
0.9958,
|
|
0.9956,
|
|
0.9953,
|
|
0.9957,
|
|
0.9954,
|
|
0.9941,
|
|
0.9958,
|
|
0.9933,
|
|
0.9938,
|
|
0.9948,
|
|
0.9936,
|
|
0.9948,
|
|
0.9925,
|
|
0.994,
|
|
0.9927,
|
|
0.9913,
|
|
0.9919,
|
|
0.9918,
|
|
0.9907,
|
|
0.989,
|
|
0.9901,
|
|
0.9892,
|
|
0.9903,
|
|
0.9884,
|
|
0.9868,
|
|
0.9851,
|
|
0.9848,
|
|
0.9849,
|
|
0.9831,
|
|
0.9818,
|
|
0.9804,
|
|
0.9781,
|
|
0.9756,
|
|
0.9733,
|
|
0.9717,
|
|
0.9688,
|
|
0.9646,
|
|
0.9611,
|
|
0.9559,
|
|
0.9503,
|
|
0.9443,
|
|
0.938,
|
|
0.9315,
|
|
0.9227,
|
|
0.9208,
|
|
)
|
|
|
|
# All curves are calibrated at the family's default 50-step schedule. A single-DiT
|
|
# curve therefore has 50 entries and MagCacheConfig interpolates it to the actual step
|
|
# count. A dual-expert MoE (Wan2.2-A14B) runs each expert on a SLICE of the schedule
|
|
# (the boundary_ratio split) and the MagCache hook counts each expert's OWN forwards
|
|
# from 0, so each expert carries its own curve, keyed "family::transformer_2" for the
|
|
# second expert, whose length is the number of steps that expert ran during the 50-step
|
|
# calibration; engage-time scales it proportionally to the requested step count (the
|
|
# boundary split is a fixed fraction of the schedule for a given checkpoint).
|
|
_MAGCACHE_CALIBRATION_STEPS = 50
|
|
|
|
_MAGCACHE_FAMILY_RATIOS: dict[str, tuple[float, ...]] = {
|
|
"hunyuanvideo-1.5": _MAGCACHE_480P_RATIOS,
|
|
"hunyuanvideo-1.5-720p": _MAGCACHE_720P_RATIOS,
|
|
"wan2.2-ti2v-5b": _MAGCACHE_WAN5B_RATIOS,
|
|
}
|
|
|
|
|
|
def _magcache_ratio_key(family: Optional[str], expert: Optional[str]) -> str:
|
|
"""The `_MAGCACHE_FAMILY_RATIOS` key for a (family, expert) pair: the bare family
|
|
name for the primary ``transformer``, ``family::expert`` for a second expert."""
|
|
fam = str(family or "").strip().lower()
|
|
exp = str(expert or "").strip().lower()
|
|
if exp in ("", "transformer"):
|
|
return fam
|
|
return f"{fam}::{exp}"
|
|
|
|
|
|
# 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. On Wan2.2-TI2V-5B both modes
|
|
# stay composition-true, but MagCache dominates the accuracy/speed frontier (B200,
|
|
# 1280x704/33f/50 steps, pairwise LPIPS vs the same uncached compiled stack): balanced
|
|
# MagCache 1.65x at 0.034 vs FBCache 0.08 at 1.49x/0.031, and at the fast points 1.73x
|
|
# at 0.044 vs 1.71x at 0.083 -- FBCache's error grows unboundedly past its threshold
|
|
# while MagCache's budget caps it. On Wan2.2-A14B (dual-expert MoE) the OPPOSITE holds
|
|
# (B200, 1280x720/33f/50 steps, per-expert calibrated curves, same pairwise protocol):
|
|
# FBCache 0.12 at 2.88x/0.128 dominates balanced MagCache (1.80x/0.145) and FBCache
|
|
# 0.08 sits at 1.28x/0.098 vs MagCache quality's 1.14x/0.074 -- the 16-step high-noise
|
|
# expert leaves MagCache too few forwards to skip within its error budget -- so the
|
|
# family keeps the FBCache default and no calibrated curve ships (an explicit magcache
|
|
# request runs uncached with a warning rather than engaging a measured-worse mode).
|
|
# 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,
|
|
"wan2.2-ti2v-5b": 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).
|
|
|
|
Raises ValueError for an unsupported value so a bad request is rejected cheaply."""
|
|
if value is None:
|
|
return None
|
|
normalized = str(value).strip().lower().replace("-", "_")
|
|
if not normalized or normalized in ("none", "off"):
|
|
return None
|
|
if normalized == TC_AUTO:
|
|
return TC_AUTO
|
|
if normalized not in TC_MODES:
|
|
raise ValueError(
|
|
f"Unsupported transformer_cache '{value}'. Use one of: off, auto, "
|
|
f"{', '.join(TC_MODES)}."
|
|
)
|
|
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), ...)
|
|
#
|
|
# LTX-2 is DELIBERATELY absent: its LTX2VideoTransformerBlock is also unregistered in
|
|
# diffusers 0.39, but it returns (hidden_states, audio_hidden_states) -- a JOINT
|
|
# video+audio stream -- while both cache hook families cache/skip only the single
|
|
# ``hidden_states`` stream and, on a skipped step, fetch the parameter literally named
|
|
# ``encoder_hidden_states`` (the TEXT embeddings) for the second return slot. A naive
|
|
# registration would therefore feed text embeddings into the next block's audio input
|
|
# on every skipped step. Step caching for LTX-2 needs a dual-stream cache
|
|
# implementation, not a metadata entry; until then the family runs uncached (verified:
|
|
# enable_cache raises "not registered" and the load proceeds uncached, and the
|
|
# distilled LTX-2.3 checkpoints run 8-step schedules below FBCACHE_MIN_STEPS anyway).
|
|
_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
|
|
|
|
|
|
# diffusers' cache hook registry names whose compute branch we re-point at a compiled
|
|
# inner forward (leader = the measuring first block, block = the remaining ones); both
|
|
# hook families share the fn_ref layout.
|
|
_CACHE_HOOK_NAMES = (
|
|
"mag_cache_leader_block_hook",
|
|
"mag_cache_block_hook",
|
|
"fbc_leader_block_hook",
|
|
"fbc_block_hook",
|
|
)
|
|
|
|
|
|
def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
|
|
"""Restore the regional compile on cache-hooked blocks' COMPUTED steps.
|
|
|
|
``enable_cache`` replaces each block's ``forward`` with the hook's ``new_forward``
|
|
(stashing the pre-hook bound method in ``fn_ref.original_forward``), whose skip
|
|
decision is data-dependent Python: MagCache ``@torch.compiler.disable``s the whole
|
|
``new_forward`` (recursive -- the compute branch runs EAGER), and even FBCache's
|
|
traceable ``new_forward`` graph-breaks around its disabled threshold decision,
|
|
which on some archs (measured: Qwen-Image) drops the compute branch's call into
|
|
``original_forward`` out of the compiled region -- the block's regional compile
|
|
artifact (``_compiled_call_impl``) is never reached and the cache forfeits the
|
|
compile win on every non-skipped step. An explicitly ``torch.compile``d callable
|
|
re-enables
|
|
dynamo for its own extent even inside a disabled frame, so re-pointing
|
|
``fn_ref.original_forward`` at a compiled wrapper of the same bound method restores
|
|
compiled compute steps while the skip decision stays eager exactly as designed.
|
|
Measured (B200, scripts/image_speedmem_bench.py): Qwen-Image FBCache computed steps
|
|
91.8 -> 71.2 ms (= the uncached compiled rate), 1.21x end to end; FLUX.1-dev is
|
|
neutral (its FBCache ``new_forward`` happens to trace, so computed steps were
|
|
already compiled -- same-process armed vs unarmed latents bit-identical); on the
|
|
video DiT balanced MagCache went 39.4 -> 26.9 s at 50 steps.
|
|
|
|
Only blocks the speed layer actually compiled are armed (``_compiled_call_impl``
|
|
guard -- eager tiers stay untouched), and only when ``original_forward`` is a plain
|
|
bound method (a stacked hook chain, e.g. offload, captures a partial and is
|
|
skipped). Idempotent via the ``_unsloth_orig_inner`` marker; best-effort. Returns
|
|
the number of hooks armed."""
|
|
try:
|
|
import torch
|
|
except Exception: # noqa: BLE001 -- no torch, nothing to arm
|
|
return 0
|
|
armed = 0
|
|
try:
|
|
for module in transformer.modules():
|
|
registry = getattr(module, "_diffusers_hook", None)
|
|
if registry is None or getattr(module, "_compiled_call_impl", None) is None:
|
|
continue
|
|
hooks = getattr(registry, "hooks", None) or {}
|
|
for name in _CACHE_HOOK_NAMES:
|
|
hook = hooks.get(name)
|
|
fn_ref = getattr(hook, "fn_ref", None) if hook is not None else None
|
|
orig = getattr(fn_ref, "original_forward", None)
|
|
if orig is None or getattr(hook, "_unsloth_orig_inner", None) is not None:
|
|
continue
|
|
if getattr(orig, "__self__", None) is None:
|
|
continue # not the plain bound method; arming would miss the block
|
|
# fullgraph=False / dynamic=True: a cache is active by definition (its
|
|
# decision points graph-break) and this matches the default tier the
|
|
# regional compile used. Dynamo caches per code object, so re-arming
|
|
# after a toggle is effectively free (~0.03 s).
|
|
fn_ref.original_forward = torch.compile(orig, fullgraph = False, dynamic = True)
|
|
hook._unsloth_orig_inner = orig
|
|
armed += 1
|
|
except Exception as exc: # noqa: BLE001 -- best-effort: the cache still works eager
|
|
_warn(logger, "cache-hook inner compile", exc)
|
|
return armed
|
|
if armed and logger is not None:
|
|
logger.info(
|
|
"diffusion.cache: %d cache-hooked block(s) armed with compiled inner forwards",
|
|
armed,
|
|
)
|
|
return armed
|
|
|
|
|
|
def _restore_hooked_block_inners(transformer: Any) -> None:
|
|
"""Undo ``_compile_hooked_block_inners``: put the plain bound methods back and clear
|
|
the markers. MUST run before ``disable_cache`` -- ``remove_hook`` splices
|
|
``fn_ref.original_forward`` back into ``module.forward``, and leaving the compiled
|
|
wrapper there would pin a stale compiled callable onto the uncached path."""
|
|
try:
|
|
modules = list(transformer.modules())
|
|
except Exception: # noqa: BLE001 -- not a torch module (tests/fakes): nothing armed
|
|
return
|
|
for module in modules:
|
|
registry = getattr(module, "_diffusers_hook", None)
|
|
if registry is None:
|
|
continue
|
|
hooks = getattr(registry, "hooks", None) or {}
|
|
for name in _CACHE_HOOK_NAMES:
|
|
hook = hooks.get(name)
|
|
orig = getattr(hook, "_unsloth_orig_inner", None) if hook is not None else None
|
|
if orig is None:
|
|
continue
|
|
try:
|
|
hook.fn_ref.original_forward = orig
|
|
hook._unsloth_orig_inner = None
|
|
except Exception: # noqa: BLE001 -- per-hook best-effort
|
|
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
|
|
does NOT guarantee it: Flux Kontext / img2img / inpaint / controlnet reuse the CacheMixin
|
|
FluxTransformer2DModel but never open a cache_context. Read from the pipeline ``__call__``
|
|
source, resolved off the instance so a per-expert proxy view (``_SecondDiTView``)
|
|
delegates to the real pipe; if it cannot be read, report False so the cache stays off."""
|
|
import inspect
|
|
|
|
call = getattr(pipe, "__call__", None)
|
|
if call is None:
|
|
return False
|
|
try:
|
|
src = inspect.getsource(call)
|
|
except (OSError, TypeError):
|
|
return False
|
|
# Match the actual call `cache_context(` -- a bare mention in a comment/docstring lacks
|
|
# the paren, so this does not false-positive on prose.
|
|
return "cache_context(" in src
|
|
|
|
|
|
def apply_step_cache(
|
|
pipe: Any,
|
|
*,
|
|
mode: Optional[str],
|
|
threshold: Optional[float] = None,
|
|
quant_active: bool = False,
|
|
family: Optional[str] = None,
|
|
steps: Optional[int] = None,
|
|
quality: Optional[str] = None,
|
|
expert: Optional[str] = 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 FBCache default so the cache still triggers on
|
|
a quantised transformer. ``quality`` picks the preset parameter set (threshold + the
|
|
magcache skip cap / retention window); an explicit ``threshold`` still wins over the
|
|
preset's threshold. 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); a dual-expert
|
|
MoE caller passes ``expert`` (the pipe attribute the view exposes, e.g.
|
|
"transformer_2") so each expert gets ITS OWN calibrated curve -- the experts split
|
|
the schedule at the boundary timestep, and the hook counts each expert's own
|
|
forwards from 0, so one shared full-schedule curve would be misaligned for both.
|
|
Best-effort: never raises for an incompatible model."""
|
|
mode = normalize_transformer_cache(mode)
|
|
if mode is None or mode == TC_AUTO:
|
|
# AUTO must be resolved by the loader (step-count policy) before reaching the
|
|
# engage call; treat a stray auto as off rather than crashing the load.
|
|
return None
|
|
transformer = getattr(pipe, "transformer", None)
|
|
if transformer is None:
|
|
return None
|
|
quality = normalize_cache_quality(quality) or CQ_BALANCED
|
|
if mode == TC_MAGCACHE:
|
|
preset_thr, mag_skip, mag_retention = _MAGCACHE_QUALITY_PRESETS[quality]
|
|
thr = threshold if threshold is not None else preset_thr
|
|
else:
|
|
# A family override wins over the generic preset table (Wan2.2-A14B pins its
|
|
# balanced threshold to the quality-gated 0.08 even when quant is active).
|
|
dense_thr, quant_thr = _FAMILY_FBCACHE_THRESHOLDS.get(
|
|
(str(family or "").strip().lower(), quality),
|
|
_FBCACHE_QUALITY_THRESHOLDS[quality],
|
|
)
|
|
thr = threshold if threshold is not None else (quant_thr if quant_active else dense_thr)
|
|
# 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
|
|
# the first generation -- so a model without enable_cache runs uncached per the
|
|
# best-effort contract instead of being reported as cached and then failing.
|
|
enable_cache = getattr(transformer, "enable_cache", None)
|
|
if not callable(enable_cache):
|
|
_warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
|
|
return None
|
|
# A CacheMixin transformer is necessary but NOT sufficient: the First-Block-Cache hook
|
|
# raises "No context is set" on the first forward unless the PIPELINE wraps its denoise
|
|
# loop in transformer.cache_context(...). Flux Kontext / img2img / inpaint / controlnet
|
|
# reuse the CacheMixin FluxTransformer2DModel yet their __call__ opens no cache_context,
|
|
# so engaging FBCache there would crash every default generation -- run uncached instead.
|
|
if not _pipeline_opens_cache_context(pipe):
|
|
_warn(
|
|
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:
|
|
if mode == TC_MAGCACHE:
|
|
ratio_key = _magcache_ratio_key(family, expert)
|
|
ratios = _MAGCACHE_FAMILY_RATIOS.get(ratio_key)
|
|
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 '{ratio_key}'"),
|
|
)
|
|
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
|
|
|
|
# A full-schedule curve (one entry per calibration step) interpolates to the
|
|
# requested step count directly. An expert SUB-curve (dual-expert MoE) covers
|
|
# only that expert's slice of the calibration schedule, and the hook indexes
|
|
# it by the expert's own forward count, so scale its configured step count by
|
|
# the same steps/calibration ratio: the boundary split is a fixed fraction of
|
|
# the schedule, so the expert runs ~len(ratios) * steps / 50 forwards.
|
|
num_steps = int(steps)
|
|
if len(ratios) != _MAGCACHE_CALIBRATION_STEPS:
|
|
num_steps = max(1, round(len(ratios) * int(steps) / _MAGCACHE_CALIBRATION_STEPS))
|
|
config: Any = MagCacheConfig(
|
|
threshold = thr,
|
|
max_skip_steps = mag_skip,
|
|
retention_ratio = mag_retention,
|
|
num_inference_steps = num_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)
|
|
# If the blocks are already regionally compiled (the generation-time toggle
|
|
# path: compile ran at load), re-point the fresh hooks' compute branch at
|
|
# compiled inners; the load path (cache before compile) is armed by
|
|
# _compile_repeated_blocks instead. No-op when nothing is compiled.
|
|
_compile_hooked_block_inners(transformer, logger)
|
|
try:
|
|
transformer._unsloth_step_cache = marker
|
|
except Exception: # noqa: BLE001 — marker is best-effort
|
|
pass
|
|
if logger is not None:
|
|
logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr)
|
|
return mode
|
|
except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached
|
|
# enable_cache can fail after hooking some blocks; drop any partial hooks so
|
|
# the reported-uncached model doesn't actually run half-cached. Any armed
|
|
# compiled inners must be restored FIRST (remove_hook splices original_forward
|
|
# back into module.forward).
|
|
_restore_hooked_block_inners(transformer)
|
|
try:
|
|
transformer.disable_cache()
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
_warn(logger, mode, exc)
|
|
return None
|
|
|
|
|
|
def effective_denoise_steps(steps: int, strength: Optional[float]) -> int:
|
|
"""The number of steps diffusers ACTUALLY denoises for a request.
|
|
|
|
An image-conditioned workflow with ``strength`` < 1 (img2img / upscale / inpaint) runs
|
|
only a fraction of ``num_inference_steps``: diffusers' ``get_timesteps`` computes
|
|
``init_timestep = min(int(num_inference_steps * strength), num_inference_steps)`` and
|
|
denoises exactly ``init_timestep`` steps -- the product is FLOORED, not rounded. The auto
|
|
step-cache policy must key on THIS count -- e.g. a 28-step upscale at strength 0.35 runs
|
|
``int(9.8) = 9`` real steps, exactly the short trajectory FBCache should stay off (each
|
|
skipped step is a large quality hit). ``strength`` None (txt2img / reference) or >= 1 -> the
|
|
full count.
|
|
"""
|
|
s = int(steps)
|
|
if strength is None or float(strength) >= 1.0:
|
|
return s
|
|
return max(1, min(int(s * float(strength)), s))
|
|
|
|
|
|
def effective_request_strength(
|
|
request_strength: Optional[float],
|
|
has_init_image: bool,
|
|
pipe_accepts_strength: bool,
|
|
pipe_default_strength: Any,
|
|
) -> Optional[float]:
|
|
"""The strength the pipe will ACTUALLY apply, for keying the auto step-cache policy.
|
|
|
|
Only image-conditioned pipelines that take ``strength`` apply it (txt2img / a pipe without
|
|
the kwarg run the full trajectory -> None). When the request omits ``strength`` the loader
|
|
does NOT pass the kwarg, so the pipe runs its OWN signature default (< 1 for every img2img /
|
|
inpaint pipeline here, e.g. 0.6); the policy must key on that default, not the full step
|
|
count, or FBCache engages on a fraction of the advertised steps. A non-numeric default
|
|
(``inspect.Parameter.empty``) falls back to the full count (None).
|
|
"""
|
|
if not (has_init_image and pipe_accepts_strength):
|
|
return None
|
|
if request_strength is not None:
|
|
return 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:
|
|
# Before remove_hook splices fn_ref.original_forward back into module.forward:
|
|
# the compiled inner wrappers must not leak onto the uncached path.
|
|
_restore_hooked_block_inners(transformer)
|
|
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,
|
|
quality: Optional[str] = None,
|
|
expert: Optional[str] = None,
|
|
logger: Any = None,
|
|
) -> Optional[str]:
|
|
"""Generation-time enable/disable for an AUTO cache decision, keyed on the actual
|
|
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
|
|
# endswith, not substring: "#s5" would match inside "#s50".
|
|
and not str(engaged).endswith(f"#s{int(steps)}")
|
|
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 = mode,
|
|
threshold = threshold,
|
|
quant_active = quant_active,
|
|
family = family,
|
|
steps = steps,
|
|
quality = quality,
|
|
expert = expert,
|
|
logger = logger,
|
|
)
|
|
if not want and engaged:
|
|
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:
|
|
if logger is not None:
|
|
logger.warning("diffusion.cache: %s unavailable (%s); running uncached", what, exc)
|