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 && (