diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index a339e0776b..4faf450cbe 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -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: diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 38d8a4832a..aa5c970114 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -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: diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index b20ba963ec..7fd7765ec5 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -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. diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 55e05257e6..4ab91eb577 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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: diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index c198d033de..64a4d792b9 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -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, diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 3b06b91c12..979d0a089c 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -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 == [] diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 8f9de41cfc..b13a35d153 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -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(