perf(video): generalize round-2 levers to Wan2.2 and LTX-2: per-family step cache, per-expert MagCache, TE quant audit
Extends the HunyuanVideo-1.5 round-1/2 optimization levers to wan2.2-ti2v-5b, wan2.2-t2v-a14b (dual-expert MoE) and ltx-2, shipping only what beats the incumbent on the measured accuracy-speed frontier (B200, LPIPS(AlexNet) pairwise vs the same uncached compiled stack at identical seed/settings). - Wan2.2-TI2V-5B auto step cache switches FBCache to calibrated MagCache: balanced (0.12, 3, 0.2) measures 1.65x at pairwise LPIPS 0.034 vs the incumbent FBCache 0.08 at 1.49x/0.031, and 1.73x/0.044 vs 1.71x/0.083 at the fast points (FBCache error grows unboundedly past its threshold while MagCache's budget caps it). A 50-step calibrated curve ships; cond/uncond branches agree within 0.0008 so one curve serves both CFG contexts. - Per-expert MagCache plumbing for dual-expert MoEs: the experts split the schedule at the boundary timestep (Wan2.2-A14B: 16 + 34 of 50) and the hook counts each expert's own forwards from 0, so a shared full-schedule curve would be misaligned for both. apply_step_cache / maybe_toggle_step_cache / the loader now thread an expert name; a second expert resolves family::transformer_2 curves and sub-curves scale their configured step count by steps/50. Single-DiT behaviour unchanged. - Wan2.2-A14B keeps FBCache: with per-expert curves, 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; the 16-step high-noise expert starves MagCache's skip budget. No calibrated curve ships, so an explicit magcache request runs uncached with a warning instead of engaging a measured-worse mode. - Wan2.2-A14B TE auto quant resolves dense: TE fp8_dynamic alone costs pairwise LPIPS 0.1195 for a 1.03x once-per-generation encode (146.7 to 142.7 s e2e). Wan2.2-TI2V-5B shares the UMT5 encoder but stays quantized (0.0396 pairwise at a real 1.09x on its much faster DiT). - LTX-2 TE fp8_dynamic family-denied: torchao per-row compute fp8 on the Gemma3-27B encoder black-frames the whole clip (mean luma 137.9 to 0.0, LPIPS 0.78; reproduced compiled and eager), while layerwise fp8 is near-lossless (pairwise 0.0043) at the same shrink, so auto falls through to it and explicit fp8_dynamic requests are refused. - LTX-2 step caching deliberately stays unregistered, now documented on _EXTRA_BLOCK_METADATA: the block returns a joint (video, audio) stream pair and both cache hook families would substitute text embeddings into the audio slot on every skipped step; a dual-stream cache is required, and the distilled checkpoints run below FBCACHE_MIN_STEPS anyway. - Compile parity (emulate_precision_casts) verified family-neutral and kept global: wan5b 1.75x/0.0029 on vs 1.54x/0.0082 off; ltx2 1.308x/0.0013 vs 1.307x/0.0025; a14b 2711 vs 2717 ms/step. Cache-hook compile arming verified to generalize (wan5b fb@0.04 armed 1.216x vs raw 1.048x). Dual-GPU CFG stays HunyuanVideo-1.5-only: LTX-2 runs batch-CFG in one forward and the Wan pipelines consume each branch inline with no guider combine hook. - video_speedmem_bench gains epc_off (compile-parity isolation) and fbcache_explicit / magcache_explicit configs plus expert-aware cache application mirroring the loader. Measured via scripts/video_speedmem_bench.py and the round-3 single-load probes; full data and per-family decision table in outputs/video_families_optim_round3.md (workspace). Tests: 235 passing across the five video inference suite files (9 new: per-expert curve resolution and step scaling, uncalibrated-expert refusal, toggle expert threading, wan5b magcache auto load/toggle, ltx2 deny auto+explicit, a14b TE auto-dense); ruff clean.
This commit is contained in:
Daniel Han
parent
d879a90bfa
commit
d58141b611
7 changed files with 390 additions and 29 deletions
|
|
@ -304,6 +304,26 @@ _CONFIGS: dict[str, dict[str, Any]] = {
|
|||
"shipped_nocache": dict(
|
||||
te = "auto", vae = "auto", dit = "auto", speed = "default", attn = "auto", cache = "off"
|
||||
),
|
||||
# Compile-parity isolation: the same compiled stack as "cudnn" but with inductor's
|
||||
# emulate_precision_casts turned back OFF after the speed layer set it, so the row
|
||||
# measures the numeric + speed effect of the round-2 parity flag per family.
|
||||
"epc_off": dict(
|
||||
te = "none",
|
||||
vae = "none",
|
||||
dit = "none",
|
||||
speed = "default",
|
||||
attn = "auto",
|
||||
cache = "off",
|
||||
epc = False,
|
||||
),
|
||||
# Explicit cache modes at the dense compiled stack (bypass the family AUTO policy),
|
||||
# for FBCache-vs-MagCache head-to-head rows at identical settings.
|
||||
"fbcache_explicit": dict(
|
||||
te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "fbcache"
|
||||
),
|
||||
"magcache_explicit": dict(
|
||||
te = "none", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "magcache"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -475,15 +495,24 @@ def _apply_levers(
|
|||
|
||||
# Step cache FIRST (compile keys fullgraph off an active cache); per expert.
|
||||
cache_active = False
|
||||
if cfg["cache"] == "auto":
|
||||
if cfg["cache"] in ("auto", "fbcache", "magcache"):
|
||||
# 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
|
||||
# HunyuanVideo-1.5 families, FBCache elsewhere. An explicit "fbcache"/"magcache"
|
||||
# config value bypasses the auto policy (the head-to-head rows).
|
||||
if cfg["cache"] == "auto":
|
||||
cache_request = (
|
||||
auto_cache_mode(fam_name) if default_steps >= FBCACHE_MIN_STEPS else None
|
||||
)
|
||||
else:
|
||||
cache_request = cfg["cache"]
|
||||
if cache_request is not None:
|
||||
# Quality preset resolution, exactly like the loader (video.py): an unset
|
||||
# request takes the family's measured auto default.
|
||||
# request takes the family's measured auto default. Expert names zip with
|
||||
# the views (the loader's expert-view iteration contract) so a dual-expert
|
||||
# MoE resolves per-expert MagCache curves.
|
||||
quality = normalize_cache_quality(cache_quality) or auto_cache_quality(fam_name)
|
||||
for v in views:
|
||||
experts = ("transformer", "transformer_2")
|
||||
for v, expert in zip(views, experts):
|
||||
engaged["cache"] = apply_step_cache(
|
||||
v,
|
||||
mode = cache_request,
|
||||
|
|
@ -492,6 +521,7 @@ def _apply_levers(
|
|||
family = fam_name,
|
||||
steps = default_steps,
|
||||
quality = quality,
|
||||
expert = expert,
|
||||
logger = logger,
|
||||
)
|
||||
cache_active = engaged["cache"] not in (None, "off")
|
||||
|
|
@ -525,6 +555,16 @@ def _apply_levers(
|
|||
logger = logger,
|
||||
)
|
||||
engaged["_effective_speed"] = speed
|
||||
# emulate_precision_casts A/B: the speed layer sets the flag True inside
|
||||
# _compile_repeated_blocks; compile is lazy (first forward), so flipping it back off
|
||||
# here gives the pre-round-2 inductor numerics for the whole run ("epc_off" config).
|
||||
if not cfg.get("epc", True):
|
||||
try:
|
||||
import torch
|
||||
torch._inductor.config.emulate_precision_casts = False
|
||||
engaged["epc"] = False
|
||||
except Exception:
|
||||
pass
|
||||
return engaged
|
||||
|
||||
|
||||
|
|
@ -565,7 +605,7 @@ def _timed_video(
|
|||
views = [pipe]
|
||||
if getattr(pipe, "transformer_2", None) is not None:
|
||||
views.append(_SecondExpertView(pipe))
|
||||
for v in views:
|
||||
for v, expert in zip(views, ("transformer", "transformer_2")):
|
||||
try:
|
||||
maybe_toggle_step_cache(
|
||||
v,
|
||||
|
|
@ -574,7 +614,9 @@ def _timed_video(
|
|||
threshold = cache_threshold,
|
||||
mode = auto_cache_mode(family),
|
||||
family = family,
|
||||
quality = normalize_cache_quality(cache_quality) or auto_cache_quality(family),
|
||||
quality = normalize_cache_quality(cache_quality)
|
||||
or auto_cache_quality(family),
|
||||
expert = expert,
|
||||
logger = logger,
|
||||
)
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -237,19 +237,109 @@ _MAGCACHE_480P_RATIOS = (
|
|||
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. Every other family keeps the
|
||||
# measured FBCache default. An EXPLICIT "fbcache"/"magcache" request always wins.
|
||||
# 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,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -293,6 +383,17 @@ def normalize_transformer_cache(value: Optional[str]) -> Optional[str]:
|
|||
# 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": (
|
||||
(
|
||||
|
|
@ -480,6 +581,7 @@ def apply_step_cache(
|
|||
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
|
||||
|
|
@ -489,8 +591,12 @@ def apply_step_cache(
|
|||
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). Best-effort:
|
||||
never raises for an incompatible model."""
|
||||
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
|
||||
|
|
@ -531,14 +637,15 @@ def apply_step_cache(
|
|||
_ensure_block_metadata_registered(transformer, logger)
|
||||
try:
|
||||
if mode == TC_MAGCACHE:
|
||||
ratios = _MAGCACHE_FAMILY_RATIOS.get(str(family or "").strip().lower())
|
||||
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 family '{family}'"),
|
||||
RuntimeError(f"no calibrated mag_ratios for '{ratio_key}'"),
|
||||
)
|
||||
return None
|
||||
if not steps or int(steps) <= 0:
|
||||
|
|
@ -546,11 +653,22 @@ def apply_step_cache(
|
|||
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 = int(steps),
|
||||
num_inference_steps = num_steps,
|
||||
mag_ratios = list(ratios),
|
||||
)
|
||||
# The curve is interpolated over the CONFIGURED step count, so the marker
|
||||
|
|
@ -668,6 +786,7 @@ def maybe_toggle_step_cache(
|
|||
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
|
||||
|
|
@ -701,6 +820,7 @@ def maybe_toggle_step_cache(
|
|||
family = family,
|
||||
steps = steps,
|
||||
quality = quality,
|
||||
expert = expert,
|
||||
logger = logger,
|
||||
)
|
||||
if not want and engaged:
|
||||
|
|
|
|||
|
|
@ -126,16 +126,31 @@ _TE_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
|
|||
|
||||
# Text encoders whose activation ranges break a scheme at the MODEL level (measured hidden-state
|
||||
# cosine vs bf16, via scripts/diffusion_quant_builder.py). Populated from the accuracy sweep; a
|
||||
# denied scheme is skipped by ``auto`` and refused when requested explicitly. Empty by default:
|
||||
# int8 already gates on a per-family keep-bf16 schedule (``_TE_INT8_SKIP``), so this is for the
|
||||
# rarer case where even keep-bf16 int8 (or fp8) misses the bar for a specific encoder.
|
||||
_TE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {}
|
||||
# denied scheme is skipped by ``auto`` and refused when requested explicitly. int8 already gates
|
||||
# on a per-family keep-bf16 schedule (``_TE_INT8_SKIP``), so this is for the rarer case where a
|
||||
# scheme breaks the encoder outright.
|
||||
# ltx-2 / fp8_dynamic: torchao per-row compute-fp8 on the Gemma3-27B encoder BLACK-FRAMES the
|
||||
# whole clip (B200, measured pairwise vs the dense encoder at identical seed/settings: mean
|
||||
# luma 137.9 -> 0.0, LPIPS 0.78; reproduced at 1216x704/33f/40 steps compiled and 384x256/9f/10
|
||||
# steps eager). Layerwise fp8 on the same encoder is near-lossless (pairwise LPIPS 0.0043) at
|
||||
# the same ~2x shrink, so auto falls through to it -- the deny costs nothing.
|
||||
_TE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
|
||||
"ltx-2": frozenset({TE_QUANT_FP8_DYNAMIC}),
|
||||
}
|
||||
|
||||
# Families whose AUTO text-encoder quant resolves dense (see select_te_quant_scheme):
|
||||
# measured out-of-bar trajectory drift for zero speed win on the video families below.
|
||||
# Unlike the deny table this only steers the AUTO default; an explicit scheme request
|
||||
# (text_encoder_quant="fp8_dynamic") is still honored verbatim.
|
||||
_TE_AUTO_DENSE_FAMILIES: frozenset[str] = frozenset({"hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"})
|
||||
# wan2.2-t2v-a14b: TE fp8_dynamic ALONE moves the clip to pairwise LPIPS 0.1195 vs
|
||||
# the dense-TE stack (B200, 1280x720/33f/50 steps, identical seed) for 146.7 ->
|
||||
# 142.7 s e2e -- the UMT5 encoder runs once per generation, so the 1.03x is noise
|
||||
# next to being the dominant accuracy cost. The dual-expert MoE trajectory amplifies
|
||||
# the conditioning perturbation ~3x harder than the same encoder on wan2.2-ti2v-5b
|
||||
# (0.0396 pairwise, kept quantized there for a real 1.09x on its much faster DiT).
|
||||
_TE_AUTO_DENSE_FAMILIES: frozenset[str] = frozenset(
|
||||
{"hunyuanvideo-1.5", "hunyuanvideo-1.5-720p", "wan2.2-t2v-a14b"}
|
||||
)
|
||||
|
||||
# Map a TE torchao scheme to the transformer smoke-probe scheme (same torchao GEMM), so ``auto``
|
||||
# degrades gracefully when a build lacks a kernel. Layerwise fp8 has no torchao GEMM to probe.
|
||||
|
|
|
|||
|
|
@ -1360,7 +1360,11 @@ class VideoBackend:
|
|||
auto_cache_mode(fam.name) if default_cache_steps >= FBCACHE_MIN_STEPS else None
|
||||
)
|
||||
cache_engaged = None
|
||||
for view in views:
|
||||
# Each view is zipped with the pipe attribute it exposes as ``transformer`` (the
|
||||
# expert-view iteration contract): a dual-expert MoE's second view passes
|
||||
# expert="transformer_2" so MagCache resolves THAT expert's calibrated curve --
|
||||
# the experts split the schedule at the boundary timestep, so their curves differ.
|
||||
for view, expert_name in zip(views, _transformer_names(pipe, fam)):
|
||||
engaged = apply_step_cache(
|
||||
view,
|
||||
mode = cache_request,
|
||||
|
|
@ -1373,6 +1377,7 @@ class VideoBackend:
|
|||
family = fam.name,
|
||||
steps = default_cache_steps,
|
||||
quality = cache_quality,
|
||||
expert = expert_name,
|
||||
logger = logger,
|
||||
)
|
||||
if view is pipe:
|
||||
|
|
@ -1808,7 +1813,9 @@ class VideoBackend:
|
|||
# MoE toggles both experts.
|
||||
if state.cache_auto:
|
||||
toggled = state.transformer_cache
|
||||
for view in _views_for(pipe, fam):
|
||||
for view, expert_name in zip(
|
||||
_views_for(pipe, fam), _transformer_names(pipe, fam)
|
||||
):
|
||||
toggled = maybe_toggle_step_cache(
|
||||
view,
|
||||
steps = steps,
|
||||
|
|
@ -1817,6 +1824,7 @@ class VideoBackend:
|
|||
mode = auto_cache_mode(fam.name),
|
||||
family = fam.name,
|
||||
quality = state.cache_quality,
|
||||
expert = expert_name,
|
||||
logger = logger,
|
||||
)
|
||||
if toggled != state.transformer_cache:
|
||||
|
|
|
|||
|
|
@ -513,11 +513,18 @@ def test_normalize_accepts_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.
|
||||
# bounded MagCache there. Wan2.2-TI2V-5B: both modes hold composition, but MagCache
|
||||
# dominates the accuracy/speed frontier (1.65x at pairwise LPIPS 0.034 vs FBCache's
|
||||
# 1.49x at 0.031; 1.73x/0.044 vs 1.71x/0.083 at the fast points), so auto engages
|
||||
# MagCache with its calibrated curve. Wan2.2-A14B measured the OTHER way (FBCache
|
||||
# 0.12 at 2.88x/0.128 dominates balanced MagCache's 1.80x/0.145; the 16-step
|
||||
# high-noise expert starves MagCache's budget), so the MoE stays on FBCache. 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("wan2.2-ti2v-5b") == TC_MAGCACHE
|
||||
for other in (None, "", "flux", "wan2.2-t2v-a14b", "ltx-2", "z-image"):
|
||||
assert auto_cache_mode(other) == TC_FBCACHE
|
||||
|
||||
|
||||
|
|
@ -611,6 +618,118 @@ def test_toggle_magcache_disengages_below_bar(monkeypatch):
|
|||
assert mode is None and t.disables == 1
|
||||
|
||||
|
||||
# ── per-expert magcache curves (dual-expert MoE, Wan2.2-A14B) ───────────────────────
|
||||
from core.inference.diffusion_cache import ( # noqa: E402
|
||||
_MAGCACHE_CALIBRATION_STEPS,
|
||||
_magcache_ratio_key,
|
||||
)
|
||||
|
||||
|
||||
def test_magcache_ratio_key_primary_and_expert():
|
||||
# The primary transformer resolves the bare family key (back-compat with every
|
||||
# single-DiT family); a second expert resolves "family::expert".
|
||||
assert _magcache_ratio_key("wan2.2-t2v-a14b", None) == "wan2.2-t2v-a14b"
|
||||
assert _magcache_ratio_key("wan2.2-t2v-a14b", "transformer") == "wan2.2-t2v-a14b"
|
||||
assert (
|
||||
_magcache_ratio_key("Wan2.2-T2V-A14B", "transformer_2")
|
||||
== "wan2.2-t2v-a14b::transformer_2"
|
||||
)
|
||||
|
||||
|
||||
def test_magcache_expert_resolves_its_own_curve(monkeypatch):
|
||||
_stub_diffusers_with_magcache(monkeypatch)
|
||||
from core.inference import diffusion_cache as dc_mod
|
||||
|
||||
primary_curve = tuple([1.0] * 15)
|
||||
expert_curve = tuple([0.99] * 35)
|
||||
monkeypatch.setitem(dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe", primary_curve)
|
||||
monkeypatch.setitem(
|
||||
dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve
|
||||
)
|
||||
t = _MixinTransformer()
|
||||
engaged = apply_step_cache(
|
||||
_pipe(t), mode = "magcache", family = "fam-moe", steps = 50,
|
||||
expert = "transformer_2",
|
||||
)
|
||||
assert engaged == TC_MAGCACHE
|
||||
assert t.enabled_with.mag_ratios == list(expert_curve)
|
||||
|
||||
|
||||
def test_magcache_expert_subcurve_scales_step_count(monkeypatch):
|
||||
# An expert sub-curve covers only that expert's slice of the calibration schedule
|
||||
# (the hook counts the expert's OWN forwards from 0), so the configured step count
|
||||
# scales by steps / calibration-steps: a 35-of-50 sub-curve at a 30-step request
|
||||
# configures round(35 * 30 / 50) = 21 steps -- NOT the full 30.
|
||||
_stub_diffusers_with_magcache(monkeypatch)
|
||||
from core.inference import diffusion_cache as dc_mod
|
||||
|
||||
expert_curve = tuple([0.99] * 35)
|
||||
monkeypatch.setitem(
|
||||
dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve
|
||||
)
|
||||
t = _MixinTransformer()
|
||||
apply_step_cache(
|
||||
_pipe(t), mode = "magcache", family = "fam-moe", steps = 30,
|
||||
expert = "transformer_2",
|
||||
)
|
||||
assert t.enabled_with.num_inference_steps == round(35 * 30 / _MAGCACHE_CALIBRATION_STEPS)
|
||||
# At the calibration step count itself the sub-curve maps 1:1.
|
||||
t2 = _MixinTransformer()
|
||||
apply_step_cache(
|
||||
_pipe(t2), mode = "magcache", family = "fam-moe",
|
||||
steps = _MAGCACHE_CALIBRATION_STEPS, expert = "transformer_2",
|
||||
)
|
||||
assert t2.enabled_with.num_inference_steps == 35
|
||||
|
||||
|
||||
def test_magcache_full_curve_keeps_requested_steps(monkeypatch):
|
||||
# A full 50-entry curve interpolates to the requested count directly (the
|
||||
# single-DiT behaviour is unchanged by the expert plumbing).
|
||||
_stub_diffusers_with_magcache(monkeypatch)
|
||||
t = _MixinTransformer()
|
||||
apply_step_cache(
|
||||
_pipe(t), mode = "magcache", family = "wan2.2-ti2v-5b", steps = 30,
|
||||
expert = "transformer",
|
||||
)
|
||||
assert t.enabled_with.num_inference_steps == 30
|
||||
assert len(t.enabled_with.mag_ratios) == _MAGCACHE_CALIBRATION_STEPS
|
||||
|
||||
|
||||
def test_magcache_expert_without_curve_runs_uncached(monkeypatch):
|
||||
# A second expert with no calibrated sub-curve must run uncached, NOT silently
|
||||
# reuse the primary's curve (the experts split the schedule; the curves differ).
|
||||
_stub_diffusers_with_magcache(monkeypatch)
|
||||
from core.inference import diffusion_cache as dc_mod
|
||||
|
||||
monkeypatch.setitem(dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe", tuple([1.0] * 15))
|
||||
t = _MixinTransformer()
|
||||
assert (
|
||||
apply_step_cache(
|
||||
_pipe(t), mode = "magcache", family = "fam-moe", steps = 50,
|
||||
expert = "transformer_2",
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert t.enabled_with is None
|
||||
|
||||
|
||||
def test_toggle_threads_expert_through(monkeypatch):
|
||||
_stub_diffusers_with_magcache(monkeypatch)
|
||||
from core.inference import diffusion_cache as dc_mod
|
||||
|
||||
expert_curve = tuple([0.98] * 35)
|
||||
monkeypatch.setitem(
|
||||
dc_mod._MAGCACHE_FAMILY_RATIOS, "fam-moe::transformer_2", expert_curve
|
||||
)
|
||||
t = _ToggleTransformer()
|
||||
mode = maybe_toggle_step_cache(
|
||||
_pipe(t), steps = 50, mode = TC_MAGCACHE, family = "fam-moe",
|
||||
expert = "transformer_2",
|
||||
)
|
||||
assert mode == TC_MAGCACHE
|
||||
assert t.enabled_with.mag_ratios == list(expert_curve)
|
||||
|
||||
|
||||
# ── cache quality presets (speed/accuracy knob) ────────────────────────────────────
|
||||
from core.inference.diffusion_cache import ( # noqa: E402
|
||||
CACHE_QUALITY_LEVELS,
|
||||
|
|
|
|||
|
|
@ -592,6 +592,27 @@ def test_select_te_auto_resolves_dense_for_hunyuanvideo15(monkeypatch):
|
|||
assert select_te_quant_scheme(_target(), "auto", family = "qwen-image") == TE_QUANT_FP8_DYNAMIC
|
||||
|
||||
|
||||
def test_select_te_auto_resolves_dense_for_wan_a14b_but_not_wan_5b(monkeypatch):
|
||||
# Wan2.2-A14B: TE fp8_dynamic alone costs pairwise LPIPS 0.1195 vs the dense-TE
|
||||
# stack for a 1.03x once-per-generation encode (146.7 -> 142.7 s e2e), so AUTO
|
||||
# keeps the encoder dense. Wan2.2-TI2V-5B shares the UMT5 encoder but measured
|
||||
# in-bar (0.0396 pairwise) at a real 1.09x on its far faster DiT, so it keeps the
|
||||
# normal ladder.
|
||||
_stub_tq_select(monkeypatch, cc = (10, 0), consumer = False)
|
||||
_allow_te(monkeypatch, {TE_QUANT_FP8_DYNAMIC, TE_QUANT_INT8, TE_QUANT_FP8})
|
||||
assert select_te_quant_scheme(_target(), "auto", family = "wan2.2-t2v-a14b") is None
|
||||
assert select_te_quant_scheme(_target(), "auto", family = "Wan2.2-T2V-A14B") is None
|
||||
assert (
|
||||
select_te_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b")
|
||||
== TE_QUANT_FP8_DYNAMIC
|
||||
)
|
||||
# The auto-dense table steers only the DEFAULT; an explicit request stays verbatim.
|
||||
assert (
|
||||
select_te_quant_scheme(_target(), "fp8_dynamic", family = "wan2.2-t2v-a14b")
|
||||
== TE_QUANT_FP8_DYNAMIC
|
||||
)
|
||||
|
||||
|
||||
def test_select_te_explicit_scheme_still_honored_for_hunyuanvideo15(monkeypatch):
|
||||
# The auto-dense table steers only the DEFAULT; an explicit request stays verbatim
|
||||
# (select returns it as-is; quantize_text_encoders re-gates hardware support).
|
||||
|
|
@ -601,3 +622,28 @@ def test_select_te_explicit_scheme_still_honored_for_hunyuanvideo15(monkeypatch)
|
|||
select_te_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5-720p")
|
||||
== TE_QUANT_FP8_DYNAMIC
|
||||
)
|
||||
|
||||
|
||||
def test_select_te_auto_ltx2_denies_fp8_dynamic_falls_to_layerwise_fp8(monkeypatch):
|
||||
# LTX-2's Gemma3-27B encoder BLACK-FRAMES the whole clip under torchao per-row
|
||||
# compute fp8 (measured pairwise vs the dense encoder: mean luma 137.9 -> 0.0,
|
||||
# LPIPS 0.78), while layerwise fp8 is near-lossless (0.0043) at the same shrink --
|
||||
# so the family deny drops fp8_dynamic and auto falls through (int8 has no ltx-2
|
||||
# keep-bf16 schedule) to layerwise fp8.
|
||||
_stub_tq_select(monkeypatch, cc = (10, 0), consumer = False)
|
||||
_allow_te(monkeypatch, {TE_QUANT_FP8_DYNAMIC, TE_QUANT_INT8, TE_QUANT_FP8})
|
||||
assert select_te_quant_scheme(_target(), "auto", family = "ltx-2") == TE_QUANT_FP8
|
||||
|
||||
|
||||
def test_quantize_explicit_fp8_dynamic_refused_for_ltx2(monkeypatch):
|
||||
# The deny contract covers EXPLICIT requests too: black frames are a model-level
|
||||
# breakage, not a preference, so the encoder stays dense instead.
|
||||
_allow_te(monkeypatch, {TE_QUANT_FP8_DYNAMIC})
|
||||
calls: list = []
|
||||
monkeypatch.setattr(dp, "_cast_fp8_dynamic", lambda enc, tgt: calls.append(enc))
|
||||
pipe = types.SimpleNamespace(text_encoder = object())
|
||||
assert (
|
||||
quantize_text_encoders(pipe, _target(), mode = "fp8_dynamic", family = "ltx-2")
|
||||
is None
|
||||
)
|
||||
assert calls == []
|
||||
|
|
|
|||
|
|
@ -394,9 +394,16 @@ def fake_runtime(monkeypatch):
|
|||
diffusers.HunyuanVideo15Pipeline = _FakeHV15Pipeline
|
||||
diffusers.HunyuanVideo15Transformer3DModel = _FakeTransformer
|
||||
diffusers.FirstBlockCacheConfig = lambda threshold = None: ("fbcache", threshold)
|
||||
# diffusers.hooks.MagCacheConfig: the auto cache mode for the HunyuanVideo-1.5 and
|
||||
# Wan2.2-TI2V-5B families (calibrated curves); the fake records its kwargs so the
|
||||
# cache tests can assert the engaged mode + step count.
|
||||
diffusers_hooks = types.ModuleType("diffusers.hooks")
|
||||
diffusers_hooks.MagCacheConfig = lambda **kwargs: ("magcache", kwargs)
|
||||
diffusers.hooks = diffusers_hooks
|
||||
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
|
||||
monkeypatch.setitem(sys.modules, "diffusers.hooks", diffusers_hooks)
|
||||
monkeypatch.setattr("core.inference.video.clear_gpu_cache", lambda: None)
|
||||
# MP4 encode needs real frames + PyAV; the backend contract under test is the
|
||||
# byte handoff, so stub the encoder.
|
||||
|
|
@ -1249,12 +1256,13 @@ def test_video_speed_off_skips_hunyuan_trim(fake_runtime, monkeypatch):
|
|||
|
||||
|
||||
def test_video_step_cache_auto_from_default_schedule(fake_runtime, tmp_path):
|
||||
# Unset step cache is AUTO, decided from the model's default schedule: Wan's
|
||||
# 50-step default engages FBCache at load; the LTX distilled 8-step default
|
||||
# keeps it off. Both are re-checked per generation (toggle test below).
|
||||
# Unset step cache is AUTO, decided from the model's default schedule: Wan
|
||||
# TI2V-5B's 50-step default engages its auto mode (MagCache, calibrated curve) at
|
||||
# load; the LTX distilled 8-step default keeps it off. Both are re-checked per
|
||||
# generation (toggle test below).
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
||||
assert status["transformer_cache"] == "fbcache"
|
||||
assert status["transformer_cache"] == "magcache"
|
||||
assert status["resolved"]["transformer_cache"]["source"] == "auto"
|
||||
backend.unload()
|
||||
|
||||
|
|
@ -1276,11 +1284,14 @@ def test_video_step_cache_auto_toggles_on_actual_steps(fake_runtime):
|
|||
# it. An explicit "off" never toggles.
|
||||
backend = VideoBackend()
|
||||
backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
|
||||
assert backend.status()["transformer_cache"] == "fbcache"
|
||||
assert backend.status()["transformer_cache"] == "magcache"
|
||||
backend.generate(prompt = "a sloth", steps = 8)
|
||||
assert backend.status()["transformer_cache"] is None
|
||||
backend.generate(prompt = "a sloth", steps = 30)
|
||||
assert backend.status()["transformer_cache"] == "fbcache"
|
||||
assert backend.status()["transformer_cache"] == "magcache"
|
||||
# The re-engage interpolated the calibrated curve over the ACTUAL step count.
|
||||
cfg = backend._state.pipe.transformer.cache_config
|
||||
assert cfg[0] == "magcache" and cfg[1]["num_inference_steps"] == 30
|
||||
backend.unload()
|
||||
|
||||
backend.load_pipeline(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue