From 7dbdd28161eb236830f9e9abe15e4e0e7fff88fc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Jul 2026 14:29:14 +0000 Subject: [PATCH] perf(video): accuracy-first round 2 for HunyuanVideo-1.5: compile parity, cache quality presets, dual-GPU CFG Cuts the shipped default's LPIPS vs the bit-exact reference from 0.224 to 0.139 while going faster (24.9 s to 21.2 s at 720p/33f/30 steps, 22.7x vs reference), and makes the remaining speed/accuracy trade a user knob. - inductor precision parity: set emulate_precision_casts=True for the regional compile (fused pointwise kernels kept fp32 intermediates where eager rounds to bf16 between ops); full-clip LPIPS vs bit-exact 0.221 to 0.052 at zero speed cost. Snapshot/restored with the other process-wide backend flags. - cache x compile composition fix: diffusers cache hooks are torch.compiler.disable'd, so every COMPUTED step ran eager (1.69 vs 1.09 s/step) under MagCache/FBCache in both enable orders. Re-point each hook's fn_ref.original_forward at a torch.compile'd wrapper of the same bound method (armed only where the speed layer compiled the block; restored before every disable_cache so the uncached path stays pristine). Balanced MagCache at 50 steps: 1.48x to 2.17x, identical skip counts, bit-identical uncached rerun after enable/disable cycles. - transformer_cache_quality knob (quality|balanced|fast; API + UI + bench) mapping to (threshold, max_skip_steps, retention_ratio). Auto resolves to the near-lossless quality preset (0.06, 2, 0.3; 1.63-1.64x at pairwise LPIPS 0.05-0.09) for the HunyuanVideo-1.5 families and to balanced (the pre-knob values, byte-identical behaviour) everywhere else. - TE auto-quant resolves dense for HunyuanVideo-1.5: TE fp8_dynamic alone moves the clip to LPIPS 0.236 vs bit-exact for zero speed win (the quantised encoder perturbs the conditioning and the trajectory amplifies it chaotically); VAE fp8 stays in auto (0.053, at the compile floor). Explicit schemes honored. - dual-GPU CFG branch parallelism (new diffusion_cfg_parallel.py): transformer proxy + DiT replica on the most-free second CUDA device + worker thread, branch-routed off the pipeline's own cache_context names. Auto engages only where measured bit-identical (eager tier: max abs diff 0.0, 1.66x); the compiled stack is explicit cfg_parallel=on (1.52x over the sequential default; per-device compiled artifacts differ by 1 bf16 ulp/step, documented in the resolved record). Fail-soft gates: family allowlist, guider CFG, pipeline kind, dense DiT, no offload, free-VRAM check; single-GPU loads are untouched and the memory plan stays single-device. - video API: the transformer_cache literal now accepts auto/magcache (an explicit magcache request was rejected at the pydantic layer); the mxfp8 family deny records the round-2 measurement (block-32 MX scaling fixes the zero-row collapse, no black frames, but is latency-neutral at LPIPS 0.37: fails both ship bars). Measured on B200 via the production lever path (video_speedmem_bench.py, which gained a --cache-quality lever and companion-quant isolation configs). Tests: 441 passing across the video inference suite (32 new for cfg-parallel, 20 for presets/arming, 3 for the inductor flag, 2 for TE auto-dense); ruff clean. --- scripts/video_speedmem_bench.py | 39 +- .../backend/core/inference/diffusion_cache.py | 199 +++++- .../core/inference/diffusion_cfg_parallel.py | 610 ++++++++++++++++++ .../core/inference/diffusion_precision.py | 18 + .../backend/core/inference/diffusion_speed.py | 42 ++ .../inference/diffusion_transformer_quant.py | 6 + studio/backend/core/inference/video.py | 126 ++++ studio/backend/models/inference.py | 48 +- studio/backend/routes/video.py | 2 + studio/backend/tests/test_diffusion_cache.py | 266 ++++++++ .../tests/test_diffusion_cfg_parallel.py | 446 +++++++++++++ .../backend/tests/test_diffusion_precision.py | 24 + studio/backend/tests/test_diffusion_speed.py | 78 +++ studio/frontend/src/features/video/api.ts | 7 +- .../src/features/video/video-page.tsx | 40 +- 15 files changed, 1929 insertions(+), 22 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_cfg_parallel.py create mode 100644 studio/backend/tests/test_diffusion_cfg_parallel.py diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index 7bce5cab02..c1ea7df163 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -250,6 +250,16 @@ _CONFIGS: dict[str, dict[str, Any]] = { "te_fbcache": dict( te = "auto", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "auto" ), + # Companion-quant accuracy isolation vs the bit-exact reference (uncached, so the cache + # cannot mask it): TE-only and VAE-only on top of the trim+cudnn+compile stack. With the + # compile rounding fixed (emulate_precision_casts), the companions are the next-largest + # divergence source, and only one of them should pay for it. + "diag_te_nocache": dict( + te = "auto", vae = "none", dit = "none", speed = "default", attn = "auto", cache = "off" + ), + "diag_vae_nocache": dict( + te = "none", vae = "auto", dit = "none", speed = "default", attn = "auto", cache = "off" + ), "ditfp8_fbcache": dict( te = "none", vae = "none", dit = "auto", speed = "default", attn = "auto", cache = "auto" ), @@ -362,6 +372,7 @@ def _apply_levers( force_fp32_vae: bool, default_steps: int, cache_threshold: Optional[float] = None, + cache_quality: Optional[str] = None, logger = None, ) -> dict: """Apply the configured levers with the loader's own argument values, in the loader's order: @@ -383,6 +394,8 @@ def _apply_levers( from core.inference.diffusion_cache import ( apply_step_cache, auto_cache_mode, + auto_cache_quality, + normalize_cache_quality, FBCACHE_MIN_STEPS, ) @@ -467,6 +480,9 @@ def _apply_levers( # HunyuanVideo-1.5 families, FBCache elsewhere. cache_request = auto_cache_mode(fam_name) if default_steps >= FBCACHE_MIN_STEPS else None if cache_request is not None: + # Quality preset resolution, exactly like the loader (video.py): an unset + # request takes the family's measured auto default. + quality = normalize_cache_quality(cache_quality) or auto_cache_quality(fam_name) for v in views: engaged["cache"] = apply_step_cache( v, @@ -475,6 +491,7 @@ def _apply_levers( quant_active = dit_quant_active, family = fam_name, steps = default_steps, + quality = quality, logger = logger, ) cache_active = engaged["cache"] not in (None, "off") @@ -525,6 +542,7 @@ def _timed_video( default_steps, guidance_via_guider = False, cache_threshold = None, + cache_quality = None, family = None, logger = None, ): @@ -532,7 +550,12 @@ def _timed_video( exactly like the loader, then times total + per-step. Returns (output, total_s, [per_step_ms]).""" import torch - from core.inference.diffusion_cache import auto_cache_mode, maybe_toggle_step_cache + from core.inference.diffusion_cache import ( + auto_cache_mode, + auto_cache_quality, + maybe_toggle_step_cache, + normalize_cache_quality, + ) if cache_mode == "auto": # Toggle on EVERY expert view, exactly like the loader's per-view recheck @@ -551,6 +574,8 @@ def _timed_video( threshold = cache_threshold, mode = auto_cache_mode(family), family = family, + quality = normalize_cache_quality(cache_quality) + or auto_cache_quality(family), logger = logger, ) except Exception: @@ -632,6 +657,7 @@ def _run_config( iters: int, out: Path, cache_threshold: Optional[float] = None, + cache_quality: Optional[str] = None, logger = None, ): import numpy as np @@ -668,6 +694,7 @@ def _run_config( force_fp32_vae = force_fp32, default_steps = default_steps, cache_threshold = cache_threshold, + cache_quality = cache_quality, logger = logger, ) pipe = pipe.to("cuda") @@ -693,6 +720,7 @@ def _run_config( default_steps = default_steps, guidance_via_guider = gvg, cache_threshold = cache_threshold, + cache_quality = cache_quality, family = family, logger = logger, ) @@ -714,6 +742,7 @@ def _run_config( default_steps = default_steps, guidance_via_guider = gvg, cache_threshold = cache_threshold, + cache_quality = cache_quality, family = family, logger = logger, ) @@ -770,6 +799,7 @@ def _run_config( "speed_optims": engaged["speed_optims"], "attn_trim": engaged.get("attn_trim", False), "cache_threshold": cache_threshold, + "cache_quality": cache_quality, "cache_marker": getattr(getattr(pipe, "transformer", None), "_unsloth_step_cache", None), "load_peak_gb": round(load_peak, 2), "weights_gb": round(weights_gb, 2), @@ -804,6 +834,12 @@ def main(argv = None) -> int: default = None, help = "FBCache residual-diff threshold override (None -> the production default)", ) + ap.add_argument( + "--cache-quality", + default = None, + choices = ("quality", "balanced", "fast"), + help = "Step-cache quality preset (None -> the family's production auto default)", + ) args = ap.parse_args(argv) import logging @@ -859,6 +895,7 @@ def main(argv = None) -> int: iters = args.iters, out = out, cache_threshold = args.cache_threshold, + cache_quality = args.cache_quality, logger = logger, ) if n == "reference": diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 0b40272a72..9c3d56cc2b 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -50,6 +50,75 @@ 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. +_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), +} + + +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. @@ -285,6 +354,101 @@ def _invalidate_child_registry_cache(transformer: Any) -> None: 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``), and every cache + ``new_forward`` is ``@torch.compiler.disable``d because its skip decision is + data-dependent Python. The disable is recursive, so the compute branch's call into + ``original_forward`` runs EAGER and the block's regional compile artifact + (``_compiled_call_impl``) is never reached: measured 1.69 vs 1.09 s/step on + HunyuanVideo-1.5-720p, i.e. the cache forfeited the whole 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: identical skip counts, balanced MagCache 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 @@ -314,12 +478,15 @@ def apply_step_cache( quant_active: bool = False, family: Optional[str] = None, steps: Optional[int] = None, + quality: 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. The magcache mode additionally needs ``family`` (to look up the + 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). Best-effort: never raises for an incompatible model.""" @@ -331,14 +498,13 @@ def apply_step_cache( transformer = getattr(pipe, "transformer", None) if transformer is None: return None + quality = normalize_cache_quality(quality) or CQ_BALANCED if mode == TC_MAGCACHE: - thr = threshold if threshold is not None else DEFAULT_MAGCACHE_THRESHOLD + preset_thr, mag_skip, mag_retention = _MAGCACHE_QUALITY_PRESETS[quality] + thr = threshold if threshold is not None else preset_thr else: - thr = ( - threshold - if threshold is not None - else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD) - ) + dense_thr, quant_thr = _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 @@ -381,8 +547,8 @@ def apply_step_cache( config: Any = MagCacheConfig( threshold = thr, - max_skip_steps = MAGCACHE_MAX_SKIP_STEPS, - retention_ratio = MAGCACHE_RETENTION_RATIO, + max_skip_steps = mag_skip, + retention_ratio = mag_retention, num_inference_steps = int(steps), mag_ratios = list(ratios), ) @@ -402,6 +568,11 @@ def apply_step_cache( # 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 @@ -411,7 +582,10 @@ def apply_step_cache( 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. + # 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 @@ -471,6 +645,9 @@ def _disengage_step_cache( 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: @@ -489,6 +666,7 @@ def maybe_toggle_step_cache( threshold: Optional[float] = None, mode: str = TC_FBCACHE, family: Optional[str] = None, + quality: Optional[str] = None, logger: Any = None, ) -> Optional[str]: """Generation-time enable/disable for an AUTO cache decision, keyed on the actual @@ -521,6 +699,7 @@ def maybe_toggle_step_cache( quant_active = quant_active, family = family, steps = steps, + quality = quality, logger = logger, ) if not want and engaged: diff --git a/studio/backend/core/inference/diffusion_cfg_parallel.py b/studio/backend/core/inference/diffusion_cfg_parallel.py new file mode 100644 index 0000000000..856a05e3b5 --- /dev/null +++ b/studio/backend/core/inference/diffusion_cfg_parallel.py @@ -0,0 +1,610 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Dual-GPU CFG branch parallelism for guider-driven video pipelines. + +A CFG denoise step runs the SAME DiT twice -- once per guidance branch -- and the +diffusers modular pipelines (HunyuanVideo-1.5) run those branches sequentially: +``for guider_state_batch in guider_state: ... self.transformer(...)``, with both noise +predictions only consumed AFTER the loop in ``self.guider(guider_state)``. On a +multi-GPU host the second GPU is idle the whole time, so routing one branch to a full +DiT replica there almost halves the denoise wall time: measured on 2x B200 +(HunyuanVideo-1.5-720p, 1280x720/33f/30 steps, trim + cuDNN + regional compile) +1.77x end-to-end uncached and 1.72x with the family's auto MagCache -- where the +output is BIT-IDENTICAL to the single-GPU run (final-latent max abs diff 0.0, +byte-identical frames), because the replica runs the identical checkpoint with +identical kernels on an identical device and only the branch's placement changes. + +Mechanism (no pipeline fork): a proxy object replaces ``pipe.transformer``. The +pipeline opens ``transformer.cache_context(name)`` before each branch, which tells the +proxy WHICH branch is being dispatched; the first branch (``pred_cond``) runs on the +replica via a persistent worker thread (the call returns a placeholder immediately), +the second runs on the primary from the pipeline's own thread, and a patched +``guider.forward`` resolves the placeholder -- joining the worker and issuing the +replica -> primary copy only after BOTH branches' kernels are queued, so the copy is +the only cross-device sync. Thread dispatch (not plain async CUDA) is required +because the MagCache hooks' ``.item()`` skip decisions block the CPU mid-branch. + +Accuracy policy: parallel output is BIT-IDENTICAL only when both branches run the +same kernels -- i.e. the EAGER (uncompiled) stack, where it was verified to the byte +(with and without the step cache; cudnn.benchmark on or off). Under regional compile +the two devices' separately compiled inductor artifacts differ by ~1 bf16 ulp per +step (a reduction kernel with a different summation order; invariant to trace order, +cudnn.benchmark and the deterministic flags), and a 30-step trajectory amplifies +that chaotically (measured: mean ~12/255 per-pixel delta, composition and luma +preserved) -- the same class as the compile-vs-eager divergence this branch's +emulate_precision_casts work just eliminated. So the AUTO policy engages only on an +eager-tier load, where notably eager+parallel (~0.85 s/step both branches) even +beats the compiled sequential default (~1.09 s/step) at the eager tier's better +accuracy; the compiled stack requires an explicit "on" that accepts the fp-noise +divergence for the 1.72x. Generations that may (re)compile either module (first run +after load, a cache toggle, a shape change) are dispatched inline: dynamo +compilation patches ``Module.__call__`` process-wide, so a replica compile +concurrent with primary execution corrupts numerics or crashes. + +Two process-global thread-safety landmines this module also handles: +- diffusers' ``_native_cudnn`` attention backend enters the process-global + ``torch.nn.attention.sdpa_kernel(...)`` context per call; two threads racing its + save/restore can leave the process stuck cudnn-only (the VAE's mask-carrying SDPA + then dies with "No available kernel"). While engaged, the backend is swapped for a + direct ``aten._scaled_dot_product_cudnn_attention`` call -- the exact kernel the + composite dispatches to, so numerics are unchanged -- and restored on teardown. +- see the inline-dispatch compile serialization above. + +Best-effort throughout: any gate or build failure leaves the pipe exactly as loaded +(single device), reported through the resolved record. torch / diffusers imported +lazily. +""" + +from __future__ import annotations + +import contextlib +import queue +import threading +from typing import Any, Optional + +CFG_PARALLEL_OFF = "off" +CFG_PARALLEL_AUTO = "auto" +CFG_PARALLEL_ON = "on" +CFG_PARALLEL_MODES = (CFG_PARALLEL_OFF, CFG_PARALLEL_AUTO, CFG_PARALLEL_ON) + +# Families the AUTO policy may engage on (keyed like diffusion_cache's family tables): +# the guider-driven HunyuanVideo-1.5 pipelines, where bit-identity and the 1.7x win +# were measured. An explicit "on" skips this list (the mechanical gates still apply). +_CFG_PARALLEL_FAMILY_ALLOW = frozenset({"hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"}) + +# Secondary-device VRAM bar: the replica's weight bytes plus activation/workspace +# headroom (measured replica-device peak 17.4 GB for the 15.6 GB bf16 HV15 DiT at +# 720p/33f -- ~1.8 GB activations; the margin also absorbs the CUDA context and +# cudnn workspaces). A co-tenant process can grab the GPU between the check and the +# load, so the build stays best-effort regardless. +_REPLICA_HEADROOM_BYTES = int(4.5 * (1 << 30)) + + +def normalize_cfg_parallel(value: Optional[str]) -> str: + """Lower/strip a requested cfg_parallel mode; None / "" -> auto (the gate decides). + Raises ValueError for an unsupported value so a bad request is rejected cheaply.""" + if value is None: + return CFG_PARALLEL_AUTO + normalized = str(value).strip().lower() + if not normalized or normalized == CFG_PARALLEL_AUTO: + return CFG_PARALLEL_AUTO + if normalized in ("none",): + return CFG_PARALLEL_OFF + if normalized not in CFG_PARALLEL_MODES: + raise ValueError( + f"Unsupported cfg_parallel '{value}'. Use one of: {', '.join(CFG_PARALLEL_MODES)}." + ) + return normalized + + +class _PendingPred: + """Placeholder for a branch prediction still being produced by the worker thread. + The pipeline stores ``self.transformer(...)[0]`` per branch and only reads it in + the guider combine, so returning ``(pending,)`` satisfies the unwrap; the patched + guider forward resolves it -- joining the worker, then issuing the cross-device + copy from the MAIN thread so it is stream-ordered after the other branch's + kernels.""" + + def __init__(self) -> None: + self.event = threading.Event() + self.value: Any = None + self.error: Optional[BaseException] = None + + def resolve(self, device: Any) -> Any: + self.event.wait() + if self.error is not None: + raise self.error + v = self.value + return v.to(device, non_blocking = True) if v.device != device else v + + +class _ReplicaView: + """Present the replica as ``pipe.transformer`` to the production lever helpers + (trim installer / attention backend / regional compile), exactly like video.py's + ``_SecondDiTView`` presents an MoE second expert. Everything else reads through + to the real pipe.""" + + def __init__(self, pipe: Any, replica: Any) -> None: + object.__setattr__(self, "_pipe", pipe) + object.__setattr__(self, "transformer", replica) + + def __getattr__(self, name: str) -> Any: + return getattr(object.__getattribute__(self, "_pipe"), name) + + +class CFGParallelProxy: + """Stands in for ``pipe.transformer``: routes the ``pred_cond`` branch to the + replica DiT on the secondary device while the other branch runs on the primary; + every other attribute delegates to the primary. Cache mutations fan out to BOTH + modules so the per-branch cache state matches the single-GPU run exactly. + + ``enabled`` / ``dispatch`` are resolved per generation by ``plan_generation`` + (see the module note's accuracy policy); with ``enabled`` False the proxy is a + pure passthrough, i.e. exactly the sequential path.""" + + def __init__( + self, + primary: Any, + replica: Any, + guider: Any, + *, + compiled: bool, + explicit_on: bool, + logger: Any = None, + ) -> None: + import torch + + self._primary = primary + self._replica = replica + self._guider = guider + self._logger = logger + self._compiled = bool(compiled) + self._explicit_on = bool(explicit_on) + self._ctx: Optional[str] = None + self.enabled = False + self.dispatch = "inline" + # Replica cache state fell out of sync (an enable/disable half-failed): stop + # routing to it -- the sequential passthrough is always correct. + self._broken = False + # A generation only "settles" a (shape, steps, cache) key once it COMPLETES; + # until then every dispatch stays inline so a (re)compile of either module is + # serialized (dynamo patches Module.__call__ process-wide during tracing). + self._settled_key: Optional[tuple] = None + self._pending_key: Optional[tuple] = None + # id-keyed cache for constant-per-generation inputs (text embeds are ~200 MB; + # re-copying them every step would waste PCIe/NVLink time). + self._const_cache: dict = {} + self._p_dev = next(primary.parameters()).device + self._r_dev = next(replica.parameters()).device + self._jobs: queue.Queue = queue.Queue() + self._worker = threading.Thread( + target = self._worker_loop, daemon = True, name = "cfg-parallel-replica" + ) + self._worker.start() + # The guider combines the branch predictions on the primary device; resolving + # the replica's pending branch THERE puts the replica -> primary copy after + # both branches' queued kernels (the only cross-device sync per step). + self._orig_guider_forward = guider.forward + p_dev = self._p_dev + + def _resolve(pred: Any) -> Any: + if isinstance(pred, _PendingPred): + return pred.resolve(p_dev) + if isinstance(pred, torch.Tensor) and pred.device != p_dev: + return pred.to(p_dev, non_blocking = True) + return pred + + orig_forward = self._orig_guider_forward + + def _device_homogenising_forward(pred_cond, pred_uncond = None, **kw): + return orig_forward(_resolve(pred_cond), _resolve(pred_uncond), **kw) + + guider.forward = _device_homogenising_forward + + # ── delegation ──────────────────────────────────────────────────────────── + def __getattr__(self, name: str) -> Any: + primary = self.__dict__.get("_primary") + if primary is None: + raise AttributeError(name) + return getattr(primary, name) + + def modules(self) -> list: + """Both modules' submodules, so helpers that walk ``transformer.modules()`` + (the cache-hook inner arming in diffusion_cache) reach the replica's blocks + too -- otherwise its computed steps would run eager and the slower branch + would erase the parallel win.""" + mods = list(self._primary.modules()) + try: + mods += list(self._replica.modules()) + except Exception: # noqa: BLE001 -- a torch-less fake in tests + pass + return mods + + # ── cache fan-out (keeps per-branch cache state identical to single-GPU) ── + def enable_cache(self, config: Any) -> None: + self._primary.enable_cache(config) + try: + self._replica.enable_cache(config) + _invalidate_registry(self._replica) + self._broken = False + except Exception: + # A half-cached pair would skip differently per branch; re-raise so the + # caller's best-effort path disables BOTH (disable_cache below fans out). + self._broken = True + raise + + def disable_cache(self) -> None: + self._primary.disable_cache() + try: + self._replica.disable_cache() + self._broken = False + except Exception as exc: # noqa: BLE001 -- replica out of sync: stop routing + self._broken = True + _warn(self._logger, "cfg-parallel replica disable_cache", exc) + + def _reset_stateful_cache(self, *args: Any, **kwargs: Any) -> None: + for module in (self._primary, self._replica): + reset = getattr(module, "_reset_stateful_cache", None) + if callable(reset): + try: + reset(*args, **kwargs) + except Exception: # noqa: BLE001 -- reset is best-effort + pass + + @contextlib.contextmanager + def cache_context(self, name: str): + self._ctx = name + try: + if self.enabled and self.dispatch == "inline": + # Thread dispatch enters the replica's context inside the worker + # instead, so it stays open across the whole worker-side forward. + with self._primary.cache_context(name), self._replica.cache_context(name): + yield + else: + with self._primary.cache_context(name): + yield + finally: + self._ctx = None + + # ── per-generation policy ───────────────────────────────────────────────── + def plan_generation( + self, *, cache_engaged: bool, steps: int, width: int, height: int, frames: int + ) -> dict: + """Resolve routing + dispatch for the next generation (call AFTER the cache + toggle so the engaged state is current). Returns the plan for logging.""" + # The engaged-cache marker may live on the proxy (a post-install toggle) or on + # the primary (pre-install engage); the delegating getattr covers both. + marker = getattr(self, "_unsloth_step_cache", None) + key = (int(steps), int(width), int(height), int(frames), str(marker)) + # Bit-identity matrix (measured on 2x B200): the eager stack is byte-identical + # (cache on or off; both branches run the same eager kernels), while ANY + # compiled stack drifts ~1 ulp/step across the devices' separately compiled + # artifacts, amplified over the trajectory -- the cache state does not change + # that (its computed steps run the per-device compiled inners), so identity + # keys on the KERNELS alone and only an explicit "on" accepts compiled drift. + lossless = not self._compiled + cfg_active = getattr(self._guider, "num_conditions", 2) > 1 + self.enabled = cfg_active and not self._broken and (lossless or self._explicit_on) + self.dispatch = "thread" if key == self._settled_key else "inline" + self._pending_key = key + plan = { + "enabled": self.enabled, + "dispatch": self.dispatch, + "lossless": lossless, + "cache_engaged": bool(cache_engaged), + } + if self._logger is not None: + self._logger.info("diffusion.cfg_parallel: plan %s", plan) + return plan + + def note_generation_done(self) -> None: + """Commit the settled key after a COMPLETED generation; a cancelled/failed one + keeps the next dispatch inline (its compiles may not have finished).""" + self._settled_key = self._pending_key + + # ── the branch router ───────────────────────────────────────────────────── + def _move(self, v: Any) -> Any: + import torch + + if not isinstance(v, torch.Tensor): + return v + hit = self._const_cache.get(id(v)) + if hit is not None and hit[0] is v: + return hit[1] + moved = v.to(self._r_dev, non_blocking = True) + if v.numel() * v.element_size() >= (1 << 20): + if len(self._const_cache) > 16: + self._const_cache.clear() # new-latents ids churn; constants re-promote + self._const_cache[id(v)] = (v, moved) + return moved + + def _worker_loop(self) -> None: + import torch + + while True: + job = self._jobs.get() + if job is None: # shutdown sentinel + return + ctx_name, args, kwargs, pending = job + try: + with torch.inference_mode(), self._replica.cache_context(ctx_name): + pending.value = self._replica(*args, **kwargs)[0] + except BaseException as exc: # noqa: BLE001 -- surfaced at resolve() + pending.error = exc + finally: + pending.event.set() + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + if self.enabled and self._ctx == "pred_cond": + # Input copies are issued from the MAIN thread: their source-stream event + # is recorded BEFORE the other branch's kernels are queued on the primary. + args = tuple(self._move(a) for a in args) + kwargs = {k: self._move(v) for k, v in kwargs.items()} + if self.dispatch == "thread": + pending = _PendingPred() + self._jobs.put((self._ctx, args, kwargs, pending)) + return (pending,) + return self._replica(*args, **kwargs) + return self._primary(*args, **kwargs) + + def shutdown(self) -> None: + try: + self._jobs.put(None) + self._worker.join(timeout = 5.0) + except Exception: # noqa: BLE001 -- daemon thread; process exit reaps it + pass + + +# ── thread-safe cuDNN attention (scoped to an engaged proxy) ────────────────────── +_CUDNN_PATCH_STATE: dict = {} + + +def _install_threadsafe_cudnn_attention(logger: Any = None) -> bool: + """Swap diffusers' ``_native_cudnn`` backend for a direct aten call (see module + note). Idempotent; returns True when the patch is (already) installed.""" + if _CUDNN_PATCH_STATE: + return True + try: + import torch + from diffusers.models import attention_dispatch as ad + + orig = ad._native_cudnn_attention + + def _threadsafe_cudnn_attention( + query, key, value, attn_mask = None, dropout_p = 0.0, is_causal = False, + scale = None, enable_gqa = False, return_lse = False, _parallel_config = None, + ): + # Fall back to the stock (context-managed) path for the shapes/options the + # direct op does not cover; the video DiT hot path never takes them. + if _parallel_config is not None or return_lse or enable_gqa or dropout_p: + return orig( + query, key, value, attn_mask = attn_mask, dropout_p = dropout_p, + is_causal = is_causal, scale = scale, enable_gqa = enable_gqa, + return_lse = return_lse, _parallel_config = _parallel_config, + ) + q, k, v = (x.permute(0, 2, 1, 3).contiguous() for x in (query, key, value)) + out = torch.ops.aten._scaled_dot_product_cudnn_attention( + q, k, v, attn_mask, False, 0.0, is_causal, False, scale = scale + )[0] + return out.permute(0, 2, 1, 3) + + ad._native_cudnn_attention = _threadsafe_cudnn_attention + patched_keys = [] + for key, fn in list(getattr(ad._AttentionBackendRegistry, "_backends", {}).items()): + if fn is orig: + ad._AttentionBackendRegistry._backends[key] = _threadsafe_cudnn_attention + patched_keys.append(key) + _CUDNN_PATCH_STATE.update(module = ad, orig = orig, keys = patched_keys) + if logger is not None: + logger.info("diffusion.cfg_parallel: thread-safe cudnn attention installed") + return True + except Exception as exc: # noqa: BLE001 -- without it, parallel dispatch is unsafe + _warn(logger, "thread-safe cudnn attention", exc) + return False + + +def _restore_threadsafe_cudnn_attention() -> None: + if not _CUDNN_PATCH_STATE: + return + try: + ad = _CUDNN_PATCH_STATE["module"] + orig = _CUDNN_PATCH_STATE["orig"] + ad._native_cudnn_attention = orig + for key in _CUDNN_PATCH_STATE["keys"]: + ad._AttentionBackendRegistry._backends[key] = orig + except Exception: # noqa: BLE001 -- best-effort restore + pass + _CUDNN_PATCH_STATE.clear() + + +# ── gate + build ────────────────────────────────────────────────────────────────── +def _pick_secondary_device(primary_index: int) -> tuple[Optional[int], int]: + """(most-free visible CUDA device != primary, its free bytes).""" + import torch + + best, best_free = None, -1 + for idx in range(torch.cuda.device_count()): + if idx == primary_index: + continue + try: + free, _total = torch.cuda.mem_get_info(idx) + except Exception: # noqa: BLE001 -- device unqueryable: skip it + continue + if free > best_free: + best, best_free = idx, free + return best, best_free + + +def maybe_enable_cfg_parallel( + pipe: Any, + fam: Any, + *, + requested: Optional[str], + kind: str, + transformer_source: Optional[str], + hf_token: Optional[str], + dtype: Any, + quant_engaged: Optional[str], + offload_active: bool, + compiled: bool, + attention_backend: Optional[str], + speed_active: bool, + logger: Any = None, +) -> tuple[Optional[CFGParallelProxy], str]: + """Gate, build and install the CFG-parallel proxy on ``pipe``. Returns + ``(proxy, reason)`` when engaged or ``(None, reason)`` explaining which gate + failed. Best-effort: never raises; a miss leaves the pipe untouched.""" + mode = normalize_cfg_parallel(requested) + if mode == CFG_PARALLEL_OFF: + return None, "disabled by request" + explicit_on = mode == CFG_PARALLEL_ON + fam_name = str(getattr(fam, "name", "") or "").strip().lower() + if not explicit_on and fam_name not in _CFG_PARALLEL_FAMILY_ALLOW: + return None, "family not in the measured allowlist" + if not explicit_on and compiled: + # Cross-device compiled artifacts differ by ~1 ulp/step and the trajectory + # amplifies it (see module note): only the eager tier is bit-identical, so + # auto refuses to spend a replica's VRAM on a stack it would never route. + return None, ( + "compiled stack diverges across devices (~1 ulp/step, amplified over the " + "trajectory); auto parallelises only the eager tier -- request " + "cfg_parallel=on to accept fp-noise divergence for the ~1.7x" + ) + if not bool(getattr(fam, "guidance_via_guider", False)): + return None, "pipeline is not guider-driven (no per-branch cache_context)" + if kind != "pipeline": + return None, f"'{kind}' load has no clean second transformer source" + if quant_engaged: + return None, f"quantized DiT ({quant_engaged}) replica is unvalidated" + if offload_active: + return None, "offload plan moves the DiT; a pinned replica would defeat it" + try: + import torch + except Exception: # noqa: BLE001 + return None, "torch unavailable" + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + return None, "needs 2+ CUDA devices" + primary = getattr(pipe, "transformer", None) + if primary is None: + return None, "pipe has no transformer" + # The branch routing keys off the pipeline's per-branch cache_context calls. + from .diffusion_cache import _ensure_block_metadata_registered, _pipeline_opens_cache_context + + if not _pipeline_opens_cache_context(pipe): + return None, "pipeline opens no cache_context (branches are not identifiable)" + try: + p_dev = next(primary.parameters()).device + if p_dev.type != "cuda": + return None, f"primary DiT is on {p_dev.type}, not cuda" + weight_bytes = sum(p.numel() * p.element_size() for p in primary.parameters()) + secondary, free = _pick_secondary_device(p_dev.index or 0) + need = weight_bytes + _REPLICA_HEADROOM_BYTES + if secondary is None: + return None, "no queryable secondary CUDA device" + if free < need: + return None, ( + f"secondary cuda:{secondary} has {free / 2**30:.1f} GiB free, " + f"needs {need / 2**30:.1f} GiB for the DiT replica" + ) + except Exception as exc: # noqa: BLE001 -- any probe failure: stay single-device + _warn(logger, "cfg-parallel gating", exc) + return None, "device probe failed" + + # ── build the replica and mirror the primary's levers ── + try: + replica = type(primary).from_pretrained( + transformer_source, + subfolder = "transformer", + torch_dtype = dtype, + token = hf_token or None, + ).to(f"cuda:{secondary}") + replica.eval() + except Exception as exc: # noqa: BLE001 -- download/VRAM race: stay single-device + _warn(logger, "cfg-parallel replica load", exc) + return None, "replica load failed" + try: + from .diffusion_attention import apply_attention_backend, install_hunyuan_attention_trim + + view = _ReplicaView(pipe, replica) + if speed_active: + install_hunyuan_attention_trim(view, fam, logger = logger) + if attention_backend is not None: + apply_attention_backend(view, attention_backend, logger = logger) + if compiled: + from .diffusion_speed import _compile_repeated_blocks + + # Same tier the primary got (the video default tier); a cache may engage + # or toggle on this DiT, so fullgraph stays off exactly like the loader. + _compile_repeated_blocks(view, logger, cache_active = True) + if not _install_threadsafe_cudnn_attention(logger): + raise RuntimeError("thread-safe attention patch failed") + # The proxy's class name hides the transformer's from the cache metadata + # registration probe, so register while the real class is still visible. + _ensure_block_metadata_registered(primary, logger) + guider = getattr(pipe, "guider", None) + if guider is None or not callable(getattr(guider, "forward", None)): + raise RuntimeError("pipe has no patchable guider") + proxy = CFGParallelProxy( + primary, + replica, + guider, + compiled = compiled, + explicit_on = explicit_on, + logger = logger, + ) + pipe.transformer = proxy + except Exception as exc: # noqa: BLE001 -- roll the replica back; stay single-device + _warn(logger, "cfg-parallel install", exc) + try: + del replica + torch.cuda.empty_cache() + except Exception: # noqa: BLE001 + pass + return None, "replica install failed" + if logger is not None: + logger.info( + "diffusion.cfg_parallel: engaged (replica on cuda:%d, %.1f GiB weights)", + secondary, + weight_bytes / 2**30, + ) + return proxy, f"engaged: DiT replica on cuda:{secondary}" + + +def teardown_cfg_parallel(pipe: Any, proxy: Any, logger: Any = None) -> None: + """Restore the pipe to its single-device shape and free the replica's VRAM. + Safe to call with a half-built or foreign object; never raises.""" + try: + primary = getattr(proxy, "_primary", None) + if primary is not None and getattr(pipe, "transformer", None) is proxy: + pipe.transformer = primary + guider = getattr(proxy, "_guider", None) + orig_fwd = getattr(proxy, "_orig_guider_forward", None) + if guider is not None and orig_fwd is not None: + guider.forward = orig_fwd + shutdown = getattr(proxy, "shutdown", None) + if callable(shutdown): + shutdown() + if getattr(proxy, "_replica", None) is not None: + proxy._replica = None + proxy._const_cache = {} + _restore_threadsafe_cudnn_attention() + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if logger is not None: + logger.info("diffusion.cfg_parallel: torn down (replica freed)") + except Exception as exc: # noqa: BLE001 -- teardown is best-effort + _warn(logger, "cfg-parallel teardown", exc) + + +def _invalidate_registry(module: Any) -> None: + from .diffusion_cache import _invalidate_child_registry_cache + + _invalidate_child_registry_cache(module) + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.cfg_parallel: %s unavailable (%s); running single-device", what, exc) diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 0531ab3868..d931668723 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -131,6 +131,14 @@ _TE_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( # rarer case where even keep-bf16 int8 (or fp8) misses the bar for a specific encoder. _TE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {} +# 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"} +) + # 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. _TE_SMOKE_SCHEME = {TE_QUANT_FP8_DYNAMIC: "fp8", TE_QUANT_INT8: "int8", TE_QUANT_NVFP4: "nvfp4"} @@ -209,6 +217,16 @@ def select_te_quant_scheme( requested = normalize_te_quant(requested) if requested is None or requested != TE_QUANT_AUTO: return requested + # AUTO resolves dense for these families regardless of hardware. TE quant perturbs + # the CONDITIONING and a multi-step video trajectory amplifies that chaotically: on + # HunyuanVideo-1.5-720p (B200, 720p/33f/30 steps) TE fp8_dynamic ALONE moves the + # clip to LPIPS 0.236 vs the bit-exact reference while the rest of the shipped + # stack sits at 0.052-0.053, for ZERO speed win (35.48 vs 35.36 s e2e -- the TE + # runs once per generation). The ~6.7 GB of weight savings is not worth being the + # single dominant accuracy cost of the default stack. (VAE fp8 stays in auto: + # measured 0.053, at the compile floor -- decode-only, no trajectory to amplify.) + if (family or "").strip().lower() in _TE_AUTO_DENSE_FAMILIES: + return None from .diffusion_transformer_quant import _capability, _is_consumer_gpu cap = _capability() diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 9a1ec7ecc5..f29b3952d7 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -77,6 +77,9 @@ def snapshot_backend_flags() -> Optional[dict]: state["cudnn_tf32"] = bool(cudnn.allow_tf32) if hasattr(cudnn, "benchmark"): state["cudnn_benchmark"] = bool(cudnn.benchmark) + inductor_cfg = _inductor_config() + if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"): + state["inductor_emulate_precision_casts"] = bool(inductor_cfg.emulate_precision_casts) return state @@ -103,6 +106,20 @@ def restore_backend_flags(state: Optional[dict]) -> None: cudnn = getattr(torch.backends, "cudnn", None) _set(cudnn, "allow_tf32", "cudnn_tf32") _set(cudnn, "benchmark", "cudnn_benchmark") + _set(_inductor_config(), "emulate_precision_casts", "inductor_emulate_precision_casts") + + +def _inductor_config() -> Any: + """``torch._inductor.config`` or None. Resolved as attributes off the imported torch + module (real torch exposes ``_inductor`` directly after ``import torch``) rather + than a submodule import, so a stubbed/partial torch (tests, exotic builds) cleanly + reports None instead of picking a stale real module out of ``sys.modules``.""" + try: + import torch + + return getattr(getattr(torch, "_inductor", None), "config", None) + except Exception: # noqa: BLE001 — no inductor -> nothing to snapshot/set + return None def normalize_speed_mode(value: Optional[str]) -> str: @@ -342,6 +359,16 @@ def _compile_repeated_blocks( for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver if hasattr(dynamo_cfg, _limit_attr): setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64)) + # Match eager's intermediate rounding inside inductor's fused pointwise kernels: + # by default they keep chains in fp32 where eager materialises bf16 between ops, + # a per-forward rounding delta (max abs ~0.09 on the HunyuanVideo-1.5 DiT) that a + # multi-step denoise amplifies chaotically -- measured full-clip LPIPS vs the + # bit-exact reference drops 0.221 -> 0.052 at ZERO speed cost (1.093 vs 1.089 + # s/step on a B200, 720p/33f). Process-global, so snapshot_backend_flags carries + # it and unload restores the prior value. + inductor_cfg = _inductor_config() + if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"): + inductor_cfg.emulate_precision_casts = True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) return False @@ -354,6 +381,21 @@ def _compile_repeated_blocks( engaged = True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) + continue + # A step cache engaged BEFORE this compile (the production load order) has + # already wrapped each block's forward in a @torch.compiler.disable'd hook, so + # the compute branch would run eager on every non-skipped step and forfeit the + # regional compile entirely (measured 1.69 vs 1.09 s/step on HunyuanVideo-1.5). + # Re-point the hooks' inner forward at compiled wrappers; no-op when no cache + # hooks are installed. The toggle path (cache engaged after load) is armed by + # apply_step_cache instead. Lazy import: diffusion_cache imports nothing from + # this module, but keep the dependency one-directional at import time. + try: + from .diffusion_cache import _compile_hooked_block_inners + + _compile_hooked_block_inners(transformer, logger) + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "cache-hook inner compile", exc) return engaged diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 974b88fe3a..bb7b8f99e8 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -230,6 +230,12 @@ _FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = { # on only the main blocks is still 100% non-finite). No small exclude set exists, so fp8 stays # denied and auto lands on int8 (clean, per-token). The 480p and 720p repacks share the DiT + # activations, so both deny. (ltx-2 fp8 measures clean on the same stack, so it is absent.) + # mxfp8 was ALSO separately measured here (B200, 720p/33f, Blackwell-blog selective recipe: + # block-32 MX, min(K,N) >= 1024, text-stream linears excluded, 434/784 Linears quantized): + # block scaling does fix the zero-row collapse (no black frames, all finite), but it is exactly + # latency-neutral (1.00x e2e at 30 AND 50 steps; only the FFN up-proj GEMM wins, 1.66x, and it + # is too small a share of the trimmed forward) at LPIPS 0.37-0.38 vs the same dense stack -- + # fails both ship bars (>= 1.1x, <= 0.05), so the deny stays on measurement, not association. "hunyuanvideo-1.5": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), "hunyuanvideo-1.5-720p": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), } diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index a9c8bf959b..55e05257e6 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -52,9 +52,17 @@ from .diffusion_cache import ( TC_AUTO, apply_step_cache, auto_cache_mode, + auto_cache_quality, maybe_toggle_step_cache, + normalize_cache_quality, normalize_transformer_cache, ) +from .diffusion_cache import _disengage_step_cache +from .diffusion_cfg_parallel import ( + maybe_enable_cfg_parallel, + normalize_cfg_parallel, + teardown_cfg_parallel, +) from .diffusion_device import resolve_diffusion_device_target from .diffusion_memory import ( apply_memory_plan, @@ -286,6 +294,9 @@ class _VideoLoadState: # Inputs the generation-time toggle re-applies (quantised threshold + override). cache_quant_active: bool = False cache_threshold: Optional[float] = None + # The resolved cache quality preset ("quality" | "balanced" | "fast"); the toggle + # re-applies it so a mid-session re-engage keeps the requested preset. + cache_quality: Optional[str] = None # Dense transformer quant actually engaged ("int8" | "fp8" | "nvfp4" | "mxfp8") or # None. Mirrors the image backend's _LoadState.transformer_quant: on a pipeline-kind # load the dense DiT(s) can be torchao-quantised in place onto the low-precision @@ -299,6 +310,12 @@ class _VideoLoadState: # convolutional decoder (Conv2d/Conv3d) shrinks in place; the vae_force_fp32 families # (Wan) never quantise (force_fp32 -> dense). Mirrors the image backend's _LoadState. vae_quant: Optional[str] = None + # Dual-GPU CFG branch parallelism: "on" when a DiT replica on a second CUDA device + # runs the pred_cond branch (bit-identical under the family's auto step cache; + # ~1.7x e2e). The handle is the installed proxy -- generate() plans each run's + # dispatch on it and _teardown_state frees the replica through it. + cfg_parallel: Optional[str] = None + cfg_parallel_handle: Any = None resolved: Optional[dict] = None @@ -548,9 +565,11 @@ class VideoBackend: attention_backend: Optional[str] = None, transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, + transformer_cache_quality: Optional[str] = None, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None, vae_quant: Optional[str] = None, + cfg_parallel: Optional[str] = None, model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" @@ -586,9 +605,11 @@ class VideoBackend: attention_backend = attention_backend, transformer_cache = transformer_cache, transformer_cache_threshold = transformer_cache_threshold, + transformer_cache_quality = transformer_cache_quality, transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant, vae_quant = vae_quant, + cfg_parallel = cfg_parallel, model_kind = model_kind, _load_token = token, ), @@ -909,9 +930,11 @@ class VideoBackend: attention_backend: Optional[str] = None, transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, + transformer_cache_quality: Optional[str] = None, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None, vae_quant: Optional[str] = None, + cfg_parallel: Optional[str] = None, model_kind: Optional[str] = None, _load_token: Optional[int] = None, _base_local_dir: Optional[str] = None, @@ -1314,6 +1337,15 @@ class VideoBackend: # so both denoisers cache; the engaged mode is identical across experts. cache_request = normalize_transformer_cache(transformer_cache) cache_auto = transformer_cache is None or cache_request == TC_AUTO + # Cache quality preset tri-state: unset / "auto" -> the family's measured + # default (the near-lossless "quality" preset for HunyuanVideo-1.5, "balanced" + # -- the pre-knob behaviour -- elsewhere); an explicit preset is honored + # verbatim. Resolved here so the generation-time toggle re-applies it. + cache_quality_requested = normalize_cache_quality(transformer_cache_quality) + cache_quality = cache_quality_requested or auto_cache_quality(fam.name) + # Validate the dual-GPU CFG request cheaply here; the gate itself must run + # after placement (it keys on the post-plan device layout + free VRAM). + normalize_cfg_parallel(cfg_parallel) # GGUF checkpoints and torchao-quantised DiTs both need the higher quantised # threshold for the cache to still trigger over the quant noise. cache_quant_active = kind == "gguf" or transformer_quant_engaged is not None @@ -1340,6 +1372,7 @@ class VideoBackend: quant_active = cache_quant_active, family = fam.name, steps = default_cache_steps, + quality = cache_quality, logger = logger, ) if view is pipe: @@ -1434,6 +1467,48 @@ class VideoBackend: except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only logger.warning("video.vae_tiling_failed: %s", exc) + # ── dual-GPU CFG branch parallelism (auto on the measured families, else + # opt-in). AFTER placement so the memory plan stays single-device: the DiT + # replica lives entirely on a SECOND CUDA device and is gated on that + # device's free VRAM; a single-GPU host, an offload plan, or a quantised + # DiT all fall through to today's single-device path with the reason + # surfaced in the resolved record. + cfg_parallel_proxy, cfg_parallel_reason = maybe_enable_cfg_parallel( + pipe, + fam, + requested = cfg_parallel, + kind = kind, + transformer_source = _base_local_dir or base, + hf_token = hf_token, + dtype = dtype, + quant_engaged = transformer_quant_engaged, + offload_active = offload_policy != "none", + compiled = "compiled" in speed_optims, + attention_backend = attention_engaged, + speed_active = effective_speed != SPEED_OFF, + logger = logger, + ) + if cfg_parallel_proxy is not None and cache_engaged: + # The load engaged the step cache on the primary BEFORE the proxy + # existed; re-engage THROUGH the proxy so the replica carries the same + # hooks and each branch's cache state matches the single-GPU run + # exactly (the bit-identity precondition). Cheap: hook install only. + _disengage_step_cache( + cfg_parallel_proxy._primary, + reason = "re-engaging through the cfg-parallel proxy", + logger = logger, + ) + cache_engaged = apply_step_cache( + pipe, + mode = cache_request, + threshold = transformer_cache_threshold, + quant_active = cache_quant_active, + family = fam.name, + steps = default_cache_steps, + quality = cache_quality, + logger = logger, + ) + resolved = build_resolved_record( { "memory_mode": ( @@ -1468,6 +1543,19 @@ class VideoBackend: cache_engaged or "off", cache_reason, ), + "transformer_cache_quality": ( + transformer_cache_quality, + cache_quality, + "requested speed/accuracy preset (threshold + skip budget)" + if cache_quality_requested is not None + else "auto: the family's measured preset (near-lossless " + "'quality' for HunyuanVideo-1.5, 'balanced' elsewhere)", + ), + "cfg_parallel": ( + cfg_parallel, + "on" if cfg_parallel_proxy is not None else "off", + cfg_parallel_reason, + ), "transformer_quant": ( transformer_quant, transformer_quant_engaged or "off", @@ -1530,9 +1618,12 @@ class VideoBackend: cache_auto = cache_may_toggle, cache_quant_active = cache_quant_active, cache_threshold = transformer_cache_threshold, + cache_quality = cache_quality, transformer_quant = transformer_quant_engaged, text_encoder_quant = text_encoder_quant_engaged, vae_quant = vae_quant_engaged, + cfg_parallel = "on" if cfg_parallel_proxy is not None else None, + cfg_parallel_handle = cfg_parallel_proxy, resolved = resolved, ) # Ownership of the globals transferred to _state / _teardown_state. @@ -1725,6 +1816,7 @@ class VideoBackend: threshold = state.cache_threshold, mode = auto_cache_mode(fam.name), family = fam.name, + quality = state.cache_quality, logger = logger, ) if toggled != state.transformer_cache: @@ -1742,6 +1834,26 @@ class VideoBackend: ) if state.transformer_cache: self._reset_step_cache(pipe) + # Dual-GPU CFG parallelism: resolve this generation's routing AFTER the + # cache toggle (the plan keys on the engaged cache state -- parallel is + # bit-identical only with the cache on or an uncompiled stack) and + # serialize any run that may (re)compile. Planning must never fail a + # generation: a planner error just pins the sequential passthrough. + cfg_proxy = state.cfg_parallel_handle + if cfg_proxy is not None: + try: + cfg_proxy.plan_generation( + cache_engaged = bool(state.transformer_cache), + steps = steps, + width = width, + height = height, + frames = frames, + ) + except Exception: # noqa: BLE001 -- fall back to single-device + try: + cfg_proxy.enabled = False + except Exception: # noqa: BLE001 + pass try: with torch.inference_mode(), progress_ctx: output = pipe(**kwargs) @@ -1760,6 +1872,13 @@ class VideoBackend: raise RuntimeError(VIDEO_CANCELLED_MSG) from None if cancel.is_set(): raise RuntimeError(VIDEO_CANCELLED_MSG) + if cfg_proxy is not None: + # The (shape, steps, cache) key settles only on a COMPLETED run, so + # a cancelled/failed one keeps the next dispatch compile-safe. + try: + cfg_proxy.note_generation_done() + except Exception: # noqa: BLE001 + pass self._gen.update(phase = "export", eta_seconds = None) video_frames = output.frames[0] @@ -1852,6 +1971,11 @@ class VideoBackend: from . import diffusion_gguf_compile diffusion_gguf_compile.uninstall_all() + # Free the CFG-parallel replica on ITS device and restore the pipe's + # single-device shape (proxy out, guider forward + attention backend + # restored) before the pipe itself is dropped. + if state.cfg_parallel_handle is not None: + teardown_cfg_parallel(state.pipe, state.cfg_parallel_handle, logger = logger) del state clear_gpu_cache() @@ -1895,6 +2019,7 @@ class VideoBackend: "transformer_quant": None, "text_encoder_quant": None, "vae_quant": None, + "cfg_parallel": None, "has_audio": False, "defaults": None, "resolved": None, @@ -1924,6 +2049,7 @@ class VideoBackend: "transformer_quant": state.transformer_quant, "text_encoder_quant": state.text_encoder_quant, "vae_quant": state.vae_quant, + "cfg_parallel": state.cfg_parallel, "has_audio": fam.has_audio, "defaults": { "steps": default_steps, diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 587c9cd530..f6406f73a8 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2336,19 +2336,43 @@ class VideoLoadRequest(BaseModel): "attention; xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An unavailable " "kernel falls back to the default.", ) - transformer_cache: Optional[Literal["off", "fbcache"]] = Field( + transformer_cache: Optional[Literal["off", "auto", "fbcache", "magcache"]] = Field( None, - description = "Opt-in step caching (off by default). fbcache = First-Block-Cache: " - "reuse the transformer tail across denoise steps when the first block's residual " - "barely changes. Engages on many-step schedules only; incompatible models run " - "uncached.", + description = "Step caching (null/auto: the family's measured mode engages on " + "many-step schedules, re-checked per generation). fbcache = First-Block-Cache: reuse " + "the transformer tail across denoise steps when the first block's residual barely " + "changes. magcache = MagCache: skip whole steps from a per-family calibrated " + "magnitude curve with a bounded error budget (the auto mode for HunyuanVideo-1.5, " + "where FBCache derails the trajectory; needs a calibrated family curve, else runs " + "uncached). Incompatible models run uncached.", ) transformer_cache_threshold: Optional[float] = Field( None, ge = 0.0, le = 1.0, - description = "FBCache residual threshold (higher = skips more steps = faster, lower " - "quality). null auto-picks the family default.", + description = "Step-cache residual threshold (higher = skips more steps = faster, " + "lower quality). null auto-picks the engaged mode's family default.", + ) + transformer_cache_quality: Optional[Literal["auto", "quality", "balanced", "fast"]] = Field( + None, + description = "Step-cache speed/accuracy preset. quality = near-lossless (lower " + "threshold + tighter skip budget, smaller speedup); balanced = the measured family " + "defaults; fast = more skipping for more speed at a visible quality cost. null/auto " + "picks the family's measured default: quality for HunyuanVideo-1.5 (1.6x at half the " + "drift of balanced), balanced elsewhere. An explicit transformer_cache_threshold " + "overrides the preset's threshold; the preset still sets the MagCache skip cap / " + "retention window.", + ) + cfg_parallel: Optional[Literal["off", "auto", "on"]] = Field( + None, + description = "Dual-GPU CFG branch parallelism: run the two guidance branches " + "concurrently, one on a DiT replica on a second CUDA device (~1.7x end-to-end on " + "HunyuanVideo-1.5, replica ~20 GB VRAM). null/auto engages only where the output is " + "bit-identical to single-GPU: the measured families on an EAGER speed tier (each " + "compiled stack's per-device inductor artifacts drift ~1 ulp/step, which a clip " + "trajectory amplifies). on = engage wherever mechanically possible, including the " + "compiled stack, accepting that fp-noise divergence (composition/brightness " + "preserved). off = never.", ) transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = ( Field( @@ -2545,7 +2569,15 @@ class VideoStatusResponse(BaseModel): description = "Attention backend engaged via the diffusers dispatcher (e.g. " "_native_cudnn), or null for the default SDPA", ) - transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") + transformer_cache: Optional[str] = Field( + None, description = "Step cache engaged: fbcache | magcache | null" + ) + cfg_parallel: Optional[str] = Field( + None, + description = "Dual-GPU CFG branch parallelism engaged: 'on' (DiT replica on a second " + "CUDA device runs one guidance branch) | null (single-device). The resolved record " + "carries the gate reason.", + ) transformer_quant: Optional[str] = Field( None, description = "Dense transformer quant engaged on a pipeline load: int8 | fp8 | nvfp4 | " diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index cf421a6923..7883791901 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -124,9 +124,11 @@ async def load_video_model( attention_backend = request.attention_backend, transformer_cache = request.transformer_cache, transformer_cache_threshold = request.transformer_cache_threshold, + transformer_cache_quality = request.transformer_cache_quality, transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, vae_quant = request.vae_quant, + cfg_parallel = request.cfg_parallel, model_kind = request.model_kind, ) return VideoStatusResponse(**status_dict) diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 9952953bab..e35c703425 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -609,3 +609,269 @@ def test_toggle_magcache_disengages_below_bar(monkeypatch): _pipe(t), steps = 8, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p" ) assert mode is None and t.disables == 1 + + +# ── cache quality presets (speed/accuracy knob) ──────────────────────────────────── +from core.inference.diffusion_cache import ( # noqa: E402 + CACHE_QUALITY_LEVELS, + CQ_BALANCED, + CQ_FAST, + CQ_QUALITY, + _FBCACHE_QUALITY_THRESHOLDS, + _MAGCACHE_QUALITY_PRESETS, + normalize_cache_quality, +) + + +def test_normalize_cache_quality_unset_and_auto_are_none(): + for value in (None, "", " ", "auto", "AUTO"): + assert normalize_cache_quality(value) is None + + +def test_normalize_cache_quality_levels_and_casing(): + assert normalize_cache_quality("quality") == CQ_QUALITY + assert normalize_cache_quality(" Balanced ") == CQ_BALANCED + assert normalize_cache_quality("FAST") == CQ_FAST + + +def test_normalize_cache_quality_rejects_unknown(): + with pytest.raises(ValueError): + normalize_cache_quality("ultra") + + +def test_quality_preset_tables_cover_every_level(): + # A missing preset row would KeyError at engage time; the tables and the public + # levels tuple must stay in lockstep. + assert set(_MAGCACHE_QUALITY_PRESETS) == set(CACHE_QUALITY_LEVELS) + assert set(_FBCACHE_QUALITY_THRESHOLDS) == set(CACHE_QUALITY_LEVELS) + + +def test_balanced_presets_match_the_preknob_defaults(): + # "balanced" IS the pre-knob shipped behaviour: a load without the knob must be + # byte-identical to the round-1 defaults. + assert _MAGCACHE_QUALITY_PRESETS[CQ_BALANCED] == ( + DEFAULT_MAGCACHE_THRESHOLD, + MAGCACHE_MAX_SKIP_STEPS, + MAGCACHE_RETENTION_RATIO, + ) + assert _FBCACHE_QUALITY_THRESHOLDS[CQ_BALANCED] == ( + DEFAULT_FBCACHE_THRESHOLD, + QUANT_FBCACHE_THRESHOLD, + ) + + +def test_magcache_quality_preset_engages_conservative_params(monkeypatch): + # Calibrated on HunyuanVideo-1.5-720p (50 steps): thr 0.06 / cap 2 / retention 0.3 = + # 1.11x at pairwise LPIPS 0.057 vs balanced's 1.49x at 0.126. + _stub_diffusers_with_magcache(monkeypatch) + t = _MixinTransformer() + engaged = apply_step_cache( + _pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50, + quality = "quality", + ) + assert engaged == TC_MAGCACHE + thr, cap, retention = _MAGCACHE_QUALITY_PRESETS[CQ_QUALITY] + assert t.enabled_with.threshold == thr + assert t.enabled_with.max_skip_steps == cap + assert t.enabled_with.retention_ratio == retention + + +def test_magcache_explicit_threshold_beats_the_preset(monkeypatch): + # The preset still supplies the skip cap / retention window, but a pinned threshold + # wins (the documented contract of transformer_cache_threshold). + _stub_diffusers_with_magcache(monkeypatch) + t = _MixinTransformer() + apply_step_cache( + _pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50, + quality = "fast", threshold = 0.05, + ) + assert t.enabled_with.threshold == 0.05 + assert t.enabled_with.max_skip_steps == _MAGCACHE_QUALITY_PRESETS[CQ_FAST][1] + + +def test_fbcache_quality_preset_thresholds(monkeypatch): + _stub_diffusers(monkeypatch) + dense_thr, quant_thr = _FBCACHE_QUALITY_THRESHOLDS[CQ_QUALITY] + t = _MixinTransformer() + apply_step_cache(_pipe(t), mode = "fbcache", quality = "quality") + assert t.enabled_with.threshold == dense_thr + t2 = _MixinTransformer() + apply_step_cache(_pipe(t2), mode = "fbcache", quality = "quality", quant_active = True) + assert t2.enabled_with.threshold == quant_thr + + +def test_apply_step_cache_rejects_bad_quality(monkeypatch): + _stub_diffusers(monkeypatch) + with pytest.raises(ValueError): + apply_step_cache(_pipe(_MixinTransformer()), mode = "fbcache", quality = "bogus") + + +def test_toggle_threads_quality_through(monkeypatch): + _stub_diffusers_with_magcache(monkeypatch) + t = _ToggleTransformer() + maybe_toggle_step_cache( + _pipe(t), steps = 30, mode = TC_MAGCACHE, family = "hunyuanvideo-1.5-720p", + quality = "quality", + ) + assert t.enabled_with.threshold == _MAGCACHE_QUALITY_PRESETS[CQ_QUALITY][0] + assert t.enabled_with.max_skip_steps == _MAGCACHE_QUALITY_PRESETS[CQ_QUALITY][1] + + +# ── compiled cache-hook inners (regional compile x step cache composition) ────────── +import functools # noqa: E402 + +from core.inference.diffusion_cache import ( # noqa: E402 + _compile_hooked_block_inners, + _restore_hooked_block_inners, + auto_cache_quality, +) + + +def test_auto_cache_quality_per_family(): + assert auto_cache_quality("hunyuanvideo-1.5") == CQ_QUALITY + assert auto_cache_quality("HunyuanVideo-1.5-720p") == CQ_QUALITY + for other in (None, "", "flux", "wan2.2-ti2v-5b", "ltx-2"): + assert auto_cache_quality(other) == CQ_BALANCED + + +class _BoundInner: + """Provides a plain bound method for fn_ref.original_forward (__self__ present).""" + + def forward(self, *args, **kwargs): + return "eager" + + +def _hooked_block(*, compiled = True, hook_name = "mag_cache_block_hook", bound = True): + inner = _BoundInner() + orig = inner.forward if bound else functools.partial(_BoundInner.forward, inner) + hook = types.SimpleNamespace(fn_ref = types.SimpleNamespace(original_forward = orig)) + block = types.SimpleNamespace( + _diffusers_hook = types.SimpleNamespace(hooks = {hook_name: hook}), + _compiled_call_impl = object() if compiled else None, + ) + return block, hook, orig + + +def _fake_dit(blocks): + return types.SimpleNamespace(modules = lambda: [types.SimpleNamespace()] + blocks) + + +def _stub_torch_compile(monkeypatch): + compiled_calls = [] + + def _compile(fn, **kwargs): + compiled_calls.append((fn, kwargs)) + wrapper = lambda *a, **k: fn(*a, **k) # noqa: E731 + wrapper._unsloth_test_compiled_of = fn + return wrapper + + torch = types.ModuleType("torch") + torch.compile = _compile + monkeypatch.setitem(sys.modules, "torch", torch) + return compiled_calls + + +def test_arming_swaps_inner_for_compiled_wrapper(monkeypatch): + calls = _stub_torch_compile(monkeypatch) + block, hook, orig = _hooked_block() + assert _compile_hooked_block_inners(_fake_dit([block])) == 1 + assert hook.fn_ref.original_forward is not orig + assert hook.fn_ref.original_forward._unsloth_test_compiled_of is orig + assert hook._unsloth_orig_inner is orig + # The inner compile must match the cache-active tier: graph-breakable + dynamic. + assert calls[0][1] == {"fullgraph": False, "dynamic": True} + + +def test_arming_is_idempotent(monkeypatch): + _stub_torch_compile(monkeypatch) + block, hook, _ = _hooked_block() + dit = _fake_dit([block]) + assert _compile_hooked_block_inners(dit) == 1 + once = hook.fn_ref.original_forward + assert _compile_hooked_block_inners(dit) == 0 # marker short-circuits + assert hook.fn_ref.original_forward is once + + +def test_arming_skips_uncompiled_blocks(monkeypatch): + # An eager-tier load has no _compiled_call_impl: the hook must stay untouched + # (compiling the inner would ADD compile where the user chose eager). + _stub_torch_compile(monkeypatch) + block, hook, orig = _hooked_block(compiled = False) + assert _compile_hooked_block_inners(_fake_dit([block])) == 0 + assert hook.fn_ref.original_forward is orig + + +def test_arming_skips_partial_captured_inner(monkeypatch): + # A stacked hook chain (e.g. group offload) captures a functools.partial, not the + # plain bound method; arming would compile the wrong layer of the chain. + _stub_torch_compile(monkeypatch) + block, hook, orig = _hooked_block(bound = False) + assert _compile_hooked_block_inners(_fake_dit([block])) == 0 + assert hook.fn_ref.original_forward is orig + + +def test_arming_covers_every_cache_hook_family(monkeypatch): + _stub_torch_compile(monkeypatch) + names = ( + "mag_cache_leader_block_hook", + "mag_cache_block_hook", + "fbc_leader_block_hook", + "fbc_block_hook", + ) + blocks = [_hooked_block(hook_name = n)[0] for n in names] + assert _compile_hooked_block_inners(_fake_dit(blocks)) == len(names) + + +def test_restore_puts_the_exact_original_back(monkeypatch): + _stub_torch_compile(monkeypatch) + block, hook, orig = _hooked_block() + dit = _fake_dit([block]) + _compile_hooked_block_inners(dit) + _restore_hooked_block_inners(dit) + assert hook.fn_ref.original_forward is orig + assert hook._unsloth_orig_inner is None + + +def test_restore_tolerates_fakes_without_modules(): + _restore_hooked_block_inners(_MixinTransformer()) # no .modules(): no-op + + +def test_disengage_restores_inners_before_disable(monkeypatch): + # remove_hook splices fn_ref.original_forward back into module.forward, so the + # compiled wrapper must be swapped out BEFORE disable_cache runs. + from core.inference import diffusion_cache as dc_mod + + order = [] + + class _T(_MixinTransformer): + def disable_cache(self): + order.append("disable") + + def modules(self): + order.append("restore-walk") + return [] + + t = _T() + t._unsloth_step_cache = "magcache@0.12#s50" + assert dc_mod._disengage_step_cache(t, reason = "test") is True + assert order == ["restore-walk", "disable"] + + +def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch): + # The generation-time toggle engages the cache AFTER the load already compiled the + # blocks; apply_step_cache must arm the fresh hooks itself. + _stub_diffusers_with_magcache(monkeypatch) + _stub_torch_compile(monkeypatch) + block, hook, orig = _hooked_block() + + class _T(_MixinTransformer): + def modules(self): + return [block] + + t = _T() + engaged = apply_step_cache( + _pipe(t), mode = "magcache", family = "hunyuanvideo-1.5-720p", steps = 50 + ) + assert engaged == TC_MAGCACHE + assert hook.fn_ref.original_forward is not orig + assert hook._unsloth_orig_inner is orig diff --git a/studio/backend/tests/test_diffusion_cfg_parallel.py b/studio/backend/tests/test_diffusion_cfg_parallel.py new file mode 100644 index 0000000000..80f2746d73 --- /dev/null +++ b/studio/backend/tests/test_diffusion_cfg_parallel.py @@ -0,0 +1,446 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for dual-GPU CFG branch parallelism (``diffusion_cfg_parallel``). + +torch is stubbed via ``sys.modules`` (the module imports it lazily), the DiT modules are +fakes that record calls, and the guider is a plain namespace -- so the gating matrix, the +proxy's routing/fan-out semantics, the per-generation dispatch policy, and teardown are +all exercised without a GPU, a replica download, or diffusers.""" + +from __future__ import annotations + +import contextlib +import sys +import types + +import pytest + +from core.inference.diffusion_cfg_parallel import ( + CFG_PARALLEL_AUTO, + CFG_PARALLEL_OFF, + CFG_PARALLEL_ON, + CFGParallelProxy, + _pick_secondary_device, + maybe_enable_cfg_parallel, + normalize_cfg_parallel, + teardown_cfg_parallel, +) + + +# ── normalisation ───────────────────────────────────────────────────────────────── +def test_normalize_unset_and_auto(): + for value in (None, "", " ", "auto", "AUTO"): + assert normalize_cfg_parallel(value) == CFG_PARALLEL_AUTO + + +def test_normalize_modes_and_casing(): + assert normalize_cfg_parallel("off") == CFG_PARALLEL_OFF + assert normalize_cfg_parallel("none") == CFG_PARALLEL_OFF + assert normalize_cfg_parallel(" ON ") == CFG_PARALLEL_ON + + +def test_normalize_rejects_unknown(): + with pytest.raises(ValueError): + normalize_cfg_parallel("both") + + +# ── fakes ───────────────────────────────────────────────────────────────────────── +class _FakeDevice: + def __init__(self, type_ = "cuda", index = 0): + self.type = type_ + self.index = index + + +class _FakeTensor: + """Just enough tensor for the proxy's _move / guider resolve paths.""" + + def __init__(self, device, tag = "t", nbytes = 8): + self.device = device + self.tag = tag + self._nbytes = nbytes + + def numel(self): + return self._nbytes + + def element_size(self): + return 1 + + def to(self, device, non_blocking = False): + return _FakeTensor(device, tag = self.tag, nbytes = self._nbytes) + + +class _FakeDiT: + def __init__(self, device_index = 0, fail_enable = False): + self._device = _FakeDevice(index = device_index) + self.fail_enable = fail_enable + self.enabled_with = None + self.disables = 0 + self.resets = 0 + self.contexts: list = [] + self.calls: list = [] + self._mods = [self, types.SimpleNamespace(name = f"block{device_index}")] + + def parameters(self): + return iter([types.SimpleNamespace( + numel = lambda: 100, element_size = lambda: 2, device = self._device + )]) + + def modules(self): + return list(self._mods) + + def enable_cache(self, config): + if self.fail_enable: + raise RuntimeError("replica enable boom") + self.enabled_with = config + + def disable_cache(self): + self.disables += 1 + + def _reset_stateful_cache(self): + self.resets += 1 + + @contextlib.contextmanager + def cache_context(self, name): + self.contexts.append(name) + yield + + def __call__(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return (_FakeTensor(self._device, tag = "pred"),) + + +def _stub_torch(monkeypatch, *, device_count = 2, free = None): + torch = types.ModuleType("torch") + torch.Tensor = _FakeTensor + free = free if free is not None else {} + + def _mem_get_info(idx): + return free.get(idx, (64 << 30, 80 << 30)) + + torch.cuda = types.SimpleNamespace( + is_available = lambda: device_count > 0, + device_count = lambda: device_count, + mem_get_info = _mem_get_info, + empty_cache = lambda: None, + ) + torch.inference_mode = contextlib.nullcontext + monkeypatch.setitem(sys.modules, "torch", torch) + return torch + + +def _make_proxy(monkeypatch, *, compiled = False, explicit_on = False, fail_enable = False): + _stub_torch(monkeypatch) + primary = _FakeDiT(device_index = 0) + replica = _FakeDiT(device_index = 1, fail_enable = fail_enable) + guider = types.SimpleNamespace(forward = lambda *a, **k: ("combined", a, k), num_conditions = 2) + proxy = CFGParallelProxy( + primary, replica, guider, compiled = compiled, explicit_on = explicit_on + ) + return proxy, primary, replica, guider + + +# ── gating matrix ───────────────────────────────────────────────────────────────── +class _CtxPipe: + """A pipeline whose __call__ opens transformer.cache_context (the branch signal).""" + + def __init__(self, transformer): + self.transformer = transformer + self.guider = types.SimpleNamespace(forward = lambda *a, **k: None) + + def __call__(self): + with self.transformer.cache_context("pred_cond"): + pass + + +def _fam(name = "hunyuanvideo-1.5-720p", guider = True): + return types.SimpleNamespace(name = name, guidance_via_guider = guider) + + +def _gate(monkeypatch, pipe, fam, **overrides): + # compiled=False = the eager tier, the only stack auto parallelises (bit-identity). + kwargs = dict( + requested = None, + kind = "pipeline", + transformer_source = "repo", + hf_token = None, + dtype = "bf16", + quant_engaged = None, + offload_active = False, + compiled = False, + attention_backend = "_native_cudnn", + speed_active = True, + ) + kwargs.update(overrides) + return maybe_enable_cfg_parallel(pipe, fam, **kwargs) + + +def test_gate_disabled_by_request(monkeypatch): + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(), requested = "off") + assert proxy is None and reason == "disabled by request" + + +def test_gate_family_allowlist(monkeypatch): + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(name = "wan2.2-ti2v-5b")) + assert proxy is None and "allowlist" in reason + + +def test_gate_auto_refuses_compiled_stack(monkeypatch): + # The per-device inductor artifacts drift ~1 ulp/step; auto is bit-identical-only, + # so a compiled load never engages (and never spends the replica VRAM). + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(), compiled = True) + assert proxy is None and "cfg_parallel=on" in reason + + +def test_gate_requires_guider_pipeline(monkeypatch): + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(guider = False)) + assert proxy is None and "guider" in reason + + +def test_gate_requires_pipeline_kind(monkeypatch): + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(), kind = "gguf") + assert proxy is None and "second transformer source" in reason + + +def test_gate_skips_quantized_dit(monkeypatch): + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(), quant_engaged = "int8") + assert proxy is None and "int8" in reason + + +def test_gate_skips_offload(monkeypatch): + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(), offload_active = True) + assert proxy is None and "offload" in reason + + +def test_gate_needs_two_gpus(monkeypatch): + _stub_torch(monkeypatch, device_count = 1) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam()) + assert proxy is None and "2+ CUDA devices" in reason + + +def test_gate_needs_secondary_vram(monkeypatch): + # 1 GiB free on the only other device < weights + headroom -> stay single-device. + _stub_torch(monkeypatch, free = {1: (1 << 30, 80 << 30)}) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam()) + assert proxy is None and "free" in reason and "needs" in reason + + +def test_gate_replica_load_failure_is_soft(monkeypatch): + # Every gate passes; the replica from_pretrained blows up (download / VRAM race): + # the load must proceed single-device, never raise. + _stub_torch(monkeypatch) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam()) + assert proxy is None and reason == "replica load failed" + + +def test_explicit_on_skips_family_allowlist(monkeypatch): + # "on" bypasses the measured-family list; it still fails soft at the replica load + # (the fake DiT class has no from_pretrained), proving the gate ORDER. + _stub_torch(monkeypatch) + proxy, reason = _gate( + monkeypatch, _CtxPipe(_FakeDiT()), _fam(name = "some-future-family"), requested = "on" + ) + assert proxy is None and reason == "replica load failed" + + +def test_pick_secondary_prefers_most_free(monkeypatch): + _stub_torch( + monkeypatch, + device_count = 3, + free = {1: (10 << 30, 80 << 30), 2: (40 << 30, 80 << 30)}, + ) + idx, free = _pick_secondary_device(0) + assert idx == 2 and free == 40 << 30 + + +# ── proxy semantics ─────────────────────────────────────────────────────────────── +def test_proxy_delegates_reads_to_primary(monkeypatch): + proxy, primary, _, _ = _make_proxy(monkeypatch) + primary.some_flag = "x" + assert proxy.some_flag == "x" + proxy.shutdown() + + +def test_proxy_modules_covers_both(monkeypatch): + # The cache-hook inner arming walks transformer.modules(); missing the replica's + # blocks would leave its computed steps eager and erase the parallel win. + proxy, primary, replica, _ = _make_proxy(monkeypatch) + mods = proxy.modules() + for m in primary.modules() + replica.modules(): + assert any(m is x for x in mods) + proxy.shutdown() + + +def test_enable_cache_fans_out(monkeypatch): + proxy, primary, replica, _ = _make_proxy(monkeypatch) + proxy.enable_cache({"threshold": 0.12}) + assert primary.enabled_with == {"threshold": 0.12} + assert replica.enabled_with == {"threshold": 0.12} + proxy.disable_cache() + assert primary.disables == 1 and replica.disables == 1 + proxy.shutdown() + + +def test_replica_enable_failure_reraises_and_breaks(monkeypatch): + # A half-cached pair would skip differently per branch; the raise lets the caller's + # best-effort path disable both, and _broken pins the sequential passthrough. + proxy, primary, _, _ = _make_proxy(monkeypatch, fail_enable = True) + with pytest.raises(RuntimeError): + proxy.enable_cache({}) + assert primary.enabled_with == {} # primary was hooked before the replica failed + plan = proxy.plan_generation( + cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33 + ) + assert plan["enabled"] is False + proxy.shutdown() + + +def test_reset_stateful_cache_fans_out(monkeypatch): + proxy, primary, replica, _ = _make_proxy(monkeypatch) + proxy._reset_stateful_cache() + assert primary.resets == 1 and replica.resets == 1 + proxy.shutdown() + + +def test_cache_context_enters_both_only_when_parallel_inline(monkeypatch): + proxy, primary, replica, _ = _make_proxy(monkeypatch) + proxy.enabled, proxy.dispatch = True, "inline" + with proxy.cache_context("pred_cond"): + pass + assert primary.contexts == ["pred_cond"] and replica.contexts == ["pred_cond"] + proxy.enabled = False + with proxy.cache_context("pred_uncond"): + pass + assert replica.contexts == ["pred_cond"] # sequential: primary only + proxy.shutdown() + + +def test_routing_pred_cond_to_replica_inline(monkeypatch): + proxy, primary, replica, _ = _make_proxy(monkeypatch) + proxy.enabled, proxy.dispatch = True, "inline" + with proxy.cache_context("pred_cond"): + proxy("latents") + with proxy.cache_context("pred_uncond"): + proxy("latents") + assert len(replica.calls) == 1 and len(primary.calls) == 1 + proxy.shutdown() + + +def test_routing_passthrough_when_disabled(monkeypatch): + proxy, primary, replica, _ = _make_proxy(monkeypatch) + proxy.enabled = False + with proxy.cache_context("pred_cond"): + proxy("latents") + assert len(primary.calls) == 1 and len(replica.calls) == 0 + proxy.shutdown() + + +def test_thread_dispatch_resolves_through_guider(monkeypatch): + proxy, primary, replica, guider = _make_proxy(monkeypatch) + proxy.enabled, proxy.dispatch = True, "thread" + with proxy.cache_context("pred_cond"): + out = proxy("latents") + # The worker resolves the pending prediction; the patched guider forward joins it + # and hands a primary-device tensor to the original combine. + combined, args, _ = guider.forward(out[0], _FakeTensor(_FakeDevice(index = 0))) + assert combined == "combined" + assert args[0].device.index == 0 # replica output copied to the primary device + assert len(replica.calls) == 1 + proxy.shutdown() + + +# ── per-generation dispatch policy ────────────────────────────────────────────────── +def test_plan_parallel_on_eager_settles_to_thread(monkeypatch): + proxy, _, _, _ = _make_proxy(monkeypatch, compiled = False) + plan = proxy.plan_generation( + cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33 + ) + assert plan["enabled"] is True and plan["dispatch"] == "inline" # first run: compile-safe + proxy.note_generation_done() + plan = proxy.plan_generation( + cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33 + ) + assert plan["dispatch"] == "thread" # settled key: full overlap + proxy.shutdown() + + +def test_plan_sequential_for_compiled_stack_even_with_cache(monkeypatch): + # Compiled per-device artifacts drift regardless of the cache state (its computed + # steps run the per-device compiled inners): auto stays sequential, an explicit + # "on" accepts the fp-noise divergence. + proxy, _, _, _ = _make_proxy(monkeypatch, compiled = True) + for cache_engaged in (True, False): + plan = proxy.plan_generation( + cache_engaged = cache_engaged, steps = 30, width = 1280, height = 720, frames = 33 + ) + assert plan["enabled"] is False and plan["lossless"] is False + proxy.shutdown() + proxy_on, _, _, _ = _make_proxy(monkeypatch, compiled = True, explicit_on = True) + plan = proxy_on.plan_generation( + cache_engaged = False, steps = 10, width = 1280, height = 720, frames = 33 + ) + assert plan["enabled"] is True and plan["lossless"] is False + proxy_on.shutdown() + + +def test_plan_parallel_for_eager_stack(monkeypatch): + proxy, _, _, _ = _make_proxy(monkeypatch, compiled = False) + plan = proxy.plan_generation( + cache_engaged = False, steps = 10, width = 1280, height = 720, frames = 33 + ) + assert plan["enabled"] is True and plan["lossless"] is True + proxy.shutdown() + + +def test_plan_requires_cfg_conditions(monkeypatch): + # guidance ~1 collapses the guider to one condition: nothing to overlap. + proxy, _, _, guider = _make_proxy(monkeypatch) + guider.num_conditions = 1 + plan = proxy.plan_generation( + cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33 + ) + assert plan["enabled"] is False + proxy.shutdown() + + +def test_shape_change_forces_inline_once(monkeypatch): + proxy, _, _, _ = _make_proxy(monkeypatch) + proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33) + proxy.note_generation_done() + plan = proxy.plan_generation( + cache_engaged = True, steps = 30, width = 960, height = 544, frames = 33 + ) + assert plan["dispatch"] == "inline" # new shape may recompile: serialize + proxy.shutdown() + + +def test_cancelled_generation_stays_inline(monkeypatch): + proxy, _, _, _ = _make_proxy(monkeypatch) + proxy.plan_generation(cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33) + # No note_generation_done (cancel/failure): the same key must stay inline. + plan = proxy.plan_generation( + cache_engaged = True, steps = 30, width = 1280, height = 720, frames = 33 + ) + assert plan["dispatch"] == "inline" + proxy.shutdown() + + +# ── teardown ────────────────────────────────────────────────────────────────────── +def test_teardown_restores_pipe_and_guider(monkeypatch): + proxy, primary, _, guider = _make_proxy(monkeypatch) + orig_forward = proxy._orig_guider_forward + pipe = types.SimpleNamespace(transformer = proxy) + teardown_cfg_parallel(pipe, proxy) + assert pipe.transformer is primary + assert guider.forward is orig_forward + assert proxy._replica is None + + +def test_teardown_tolerates_foreign_object(): + teardown_cfg_parallel(types.SimpleNamespace(transformer = None), object()) diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 0ac365102b..3b06b91c12 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -577,3 +577,27 @@ def test_quantize_text_encoders_auto_resolves_and_applies(monkeypatch): mode = quantize_text_encoders(pipe, _target(), mode = "auto", family = "qwen-image") assert mode == TE_QUANT_FP8_DYNAMIC assert calls == [te] + + +def test_select_te_auto_resolves_dense_for_hunyuanvideo15(monkeypatch): + # HunyuanVideo-1.5 (both repacks): TE quant perturbs the conditioning and the video + # trajectory amplifies it chaotically (measured LPIPS 0.236 vs bit-exact from TE + # fp8_dynamic ALONE, vs 0.052 for the rest of the stack) at zero speed win, so the + # AUTO default keeps the encoder dense on ANY hardware. + _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 = "hunyuanvideo-1.5") is None + assert select_te_quant_scheme(_target(), "auto", family = "HunyuanVideo-1.5-720p") is None + # Other families keep the normal ladder on the same stubbed hardware. + assert select_te_quant_scheme(_target(), "auto", family = "qwen-image") == 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). + _stub_tq_select(monkeypatch, cc = (10, 0), consumer = False) + _allow_te(monkeypatch, {TE_QUANT_FP8_DYNAMIC}) + assert ( + select_te_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5-720p") + == TE_QUANT_FP8_DYNAMIC + ) diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 594de55983..33843acf70 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -531,3 +531,81 @@ def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch): ) assert applied["fp16_accum"] is True assert torch.backends.cuda.matmul.allow_fp16_accumulation is True + + +# ── inductor precision-cast emulation (compile-vs-eager numeric parity) ───────── + + +def _stub_inductor_config(monkeypatch, torch, *, emulate = False): + """Attach a fake ``_inductor.config`` to the stubbed torch module (diffusion_speed + resolves it as attributes off the imported torch, never via sys.modules -- so the + real torch._inductor lingering in sys.modules cannot leak into stubbed tests).""" + cfg = types.SimpleNamespace(emulate_precision_casts = emulate) + torch._inductor = types.SimpleNamespace(config = cfg) + return cfg + + +def test_regional_compile_enables_emulate_precision_casts(monkeypatch): + # Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds + # to bf16 between ops; over a multi-step denoise that compounds to a visible drift + # (LPIPS 0.221 vs bit-exact on HunyuanVideo-1.5-720p). emulate_precision_casts + # restores eager's rounding at zero measured speed cost (LPIPS 0.052), so the + # regional compile path must switch it on. + torch = _stub_torch(monkeypatch) + _stub_gguf_accel(monkeypatch) + cfg = _stub_inductor_config(monkeypatch, torch, emulate = False) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT + ) + assert applied["compiled"] is True + assert cfg.emulate_precision_casts is True + + +def test_snapshot_restores_emulate_precision_casts(monkeypatch): + # The flag is process-global, so the unload path must restore the pre-load value + # exactly like the TF32 / cudnn.benchmark globals. + torch = _stub_torch(monkeypatch) + cfg = _stub_inductor_config(monkeypatch, torch, emulate = False) + snap = snapshot_backend_flags() + assert snap["inductor_emulate_precision_casts"] is False + cfg.emulate_precision_casts = True + restore_backend_flags(snap) + assert cfg.emulate_precision_casts is False + + +def test_missing_inductor_config_is_tolerated(monkeypatch): + # A build without torch._inductor (or with the flag renamed) must neither break the + # snapshot nor the compile path. + _stub_torch(monkeypatch) # the stub torch has no _inductor attribute + _stub_gguf_accel(monkeypatch) + snap = snapshot_backend_flags() + assert "inductor_emulate_precision_casts" not in snap + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT + ) + assert applied["compiled"] is True + + +def test_regional_compile_arms_cache_hook_inners(monkeypatch): + # The production load order engages the step cache BEFORE compile, so the regional + # compile pass must re-arm the already-installed cache hooks with compiled inner + # forwards (otherwise every computed step runs eager under the hook's + # torch.compiler.disable; measured 1.69 vs 1.09 s/step on HunyuanVideo-1.5-720p). + _stub_torch(monkeypatch) + _stub_gguf_accel(monkeypatch) + from core.inference import diffusion_cache as dc_mod + + armed = [] + monkeypatch.setattr( + dc_mod, + "_compile_hooked_block_inners", + lambda transformer, logger = None: armed.append(transformer) or 1, + ) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT + ) + assert applied["compiled"] is True + assert armed == [pipe.transformer] diff --git a/studio/frontend/src/features/video/api.ts b/studio/frontend/src/features/video/api.ts index 28231197cf..08dd03ef25 100644 --- a/studio/frontend/src/features/video/api.ts +++ b/studio/frontend/src/features/video/api.ts @@ -102,8 +102,13 @@ export interface VideoLoadRequest { | "sage" | "xformers" | "aiter"; - transformer_cache?: "off" | "fbcache"; + transformer_cache?: "off" | "fbcache" | "magcache"; transformer_cache_threshold?: number; + // Step-cache speed/accuracy preset (omit for the backend default, "balanced"). + transformer_cache_quality?: "quality" | "balanced" | "fast"; + // Dual-GPU CFG branch parallelism (omit for auto: engages on measured families when a + // second GPU with enough free VRAM is available; bit-identical with the step cache on). + cfg_parallel?: "off" | "auto" | "on"; // Dense DiT precision on full-pipeline loads (omit for the hardware-ladder auto; // "none" pins plain bf16). GGUF / single-file checkpoints carry their own precision. transformer_quant?: "none" | "fp8" | "int8" | "nvfp4" | "mxfp8"; diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index 2a38842ec6..c8bf04e52c 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -517,7 +517,13 @@ export function VideoPage({ active = true }: { active?: boolean }) { const [attentionBackend, setAttentionBackend] = useState< "auto" | "native" | "cudnn" | "flash3" | "sage" >("auto"); - const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache">("auto"); + const [transformerCache, setTransformerCache] = useState<"auto" | "off" | "fbcache" | "magcache">( + "auto", + ); + const [cacheQuality, setCacheQuality] = useState<"auto" | "quality" | "balanced" | "fast">( + "auto", + ); + const [cfgParallel, setCfgParallel] = useState<"auto" | "off" | "on">("auto"); const [transformerQuant, setTransformerQuant] = useState< "auto" | "none" | "fp8" | "int8" | "nvfp4" | "mxfp8" >("auto"); @@ -939,7 +945,9 @@ export function VideoPage({ active = true }: { active?: boolean }) { speed_mode: speedMode === "auto" ? undefined : speedMode, attention_backend: attentionBackend === "auto" ? undefined : attentionBackend, transformer_cache: transformerCache === "auto" ? undefined : transformerCache, + transformer_cache_quality: cacheQuality === "auto" ? undefined : cacheQuality, transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant, + cfg_parallel: cfgParallel === "auto" ? undefined : cfgParallel, }); } catch (err) { dismissLoadToast(); @@ -959,7 +967,9 @@ export function VideoPage({ active = true }: { active?: boolean }) { speedMode, attentionBackend, transformerCache, + cacheQuality, transformerQuant, + cfgParallel, ], ); @@ -1233,7 +1243,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { /> } value={transformerCache} onValueChange={(v) => setTransformerCache(v as typeof transformerCache)} @@ -1241,6 +1251,32 @@ export function VideoPage({ active = true }: { active?: boolean }) { ["auto", "Auto"], ["off", "Off"], ["fbcache", "First-Block-Cache"], + ["magcache", "MagCache"], + ]} + /> + } + value={cacheQuality} + onValueChange={(v) => setCacheQuality(v as typeof cacheQuality)} + options={[ + ["auto", "Auto"], + ["quality", "Quality"], + ["balanced", "Balanced"], + ["fast", "Fast"], + ]} + /> + } + value={cfgParallel} + onValueChange={(v) => setCfgParallel(v as typeof cfgParallel)} + options={[ + ["auto", "Auto"], + ["off", "Off"], + ["on", "On"], ]} /> {status?.loaded && canReapply && (