ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

perf(video): accuracy-first round 2 for HunyuanVideo-1.5: compile parity, cache quality presets, dual-GPU CFG

Browse source 
Cuts the shipped default's LPIPS vs the bit-exact reference from 0.224 to 0.139
while going faster (24.9 s to 21.2 s at 720p/33f/30 steps, 22.7x vs reference),
and makes the remaining speed/accuracy trade a user knob.

- inductor precision parity: set emulate_precision_casts=True for the regional
  compile (fused pointwise kernels kept fp32 intermediates where eager rounds to
  bf16 between ops); full-clip LPIPS vs bit-exact 0.221 to 0.052 at zero speed
  cost. Snapshot/restored with the other process-wide backend flags.
- cache x compile composition fix: diffusers cache hooks are
  torch.compiler.disable'd, so every COMPUTED step ran eager (1.69 vs 1.09
  s/step) under MagCache/FBCache in both enable orders. Re-point each hook's
  fn_ref.original_forward at a torch.compile'd wrapper of the same bound method
  (armed only where the speed layer compiled the block; restored before every
  disable_cache so the uncached path stays pristine). Balanced MagCache at 50
  steps: 1.48x to 2.17x, identical skip counts, bit-identical uncached rerun
  after enable/disable cycles.
- transformer_cache_quality knob (quality|balanced|fast; API + UI + bench)
  mapping to (threshold, max_skip_steps, retention_ratio). Auto resolves to the
  near-lossless quality preset (0.06, 2, 0.3; 1.63-1.64x at pairwise LPIPS
  0.05-0.09) for the HunyuanVideo-1.5 families and to balanced (the pre-knob
  values, byte-identical behaviour) everywhere else.
- TE auto-quant resolves dense for HunyuanVideo-1.5: TE fp8_dynamic alone moves
  the clip to LPIPS 0.236 vs bit-exact for zero speed win (the quantised encoder
  perturbs the conditioning and the trajectory amplifies it chaotically); VAE
  fp8 stays in auto (0.053, at the compile floor). Explicit schemes honored.
- dual-GPU CFG branch parallelism (new diffusion_cfg_parallel.py): transformer
  proxy + DiT replica on the most-free second CUDA device + worker thread,
  branch-routed off the pipeline's own cache_context names. Auto engages only
  where measured bit-identical (eager tier: max abs diff 0.0, 1.66x); the
  compiled stack is explicit cfg_parallel=on (1.52x over the sequential
  default; per-device compiled artifacts differ by 1 bf16 ulp/step, documented
  in the resolved record). Fail-soft gates: family allowlist, guider CFG,
  pipeline kind, dense DiT, no offload, free-VRAM check; single-GPU loads are
  untouched and the memory plan stays single-device.
- video API: the transformer_cache literal now accepts auto/magcache (an
  explicit magcache request was rejected at the pydantic layer); the mxfp8
  family deny records the round-2 measurement (block-32 MX scaling fixes the
  zero-row collapse, no black frames, but is latency-neutral at LPIPS 0.37:
  fails both ship bars).

Measured on B200 via the production lever path (video_speedmem_bench.py, which
gained a --cache-quality lever and companion-quant isolation configs). Tests:
441 passing across the video inference suite (32 new for cfg-parallel, 20 for
presets/arming, 3 for the inductor flag, 2 for TE auto-dense); ruff clean.
This commit is contained in:
Daniel Han 2026-07-10 14:29:14 +00:00
commit 7dbdd28161
15 changed files with 1929 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

39
scripts/video_speedmem_bench.py
View file

@ -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":

199
studio/backend/core/inference/diffusion_cache.py
View file

@ -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:

610
studio/backend/core/inference/diffusion_cfg_parallel.py Normal file
View file

@ -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)

18
studio/backend/core/inference/diffusion_precision.py
View file

@ -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()

42
studio/backend/core/inference/diffusion_speed.py
View file

@ -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

6
studio/backend/core/inference/diffusion_transformer_quant.py
View file

@ -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}),
}

126
studio/backend/core/inference/video.py
View file

@ -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,

48
studio/backend/models/inference.py
View file

@ -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 | "

2
studio/backend/routes/video.py
View file

@ -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)

266
studio/backend/tests/test_diffusion_cache.py
View file

@ -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

446
studio/backend/tests/test_diffusion_cfg_parallel.py Normal file
View file

@ -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())

24
studio/backend/tests/test_diffusion_precision.py
View file

@ -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
)

78
studio/backend/tests/test_diffusion_speed.py
View file

@ -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]

7
studio/frontend/src/features/video/api.ts
View file

@ -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";

40
studio/frontend/src/features/video/video-page.tsx
View file

@ -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 }) {
/>
<AdvancedSelect
label="Step cache"
hint="First-Block-Cache reuses the transformer tail across steps for many-step models. Auto turns it on at 20+ steps and off for few-step distilled models, re-checked per clip."
hint="Reuses transformer work across denoise steps for many-step models. Auto picks the family's measured mode (MagCache on HunyuanVideo-1.5, First-Block-Cache elsewhere) at 20+ steps and turns it off for few-step distilled models, re-checked per clip."
badge={<ResolvedBadge status={status} controlKey="transformer_cache" />}
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"],
]}
/>
<AdvancedSelect
label="Cache quality"
hint="Speed/accuracy preset for the step cache. Quality skips conservatively (near-lossless, smaller speedup); Balanced skips more for more speed; Fast is the most aggressive. Auto picks the family's measured default (Quality on HunyuanVideo-1.5, Balanced elsewhere)."
badge={<ResolvedBadge status={status} controlKey="transformer_cache_quality" />}
value={cacheQuality}
onValueChange={(v) => setCacheQuality(v as typeof cacheQuality)}
options={[
["auto", "Auto"],
["quality", "Quality"],
["balanced", "Balanced"],
["fast", "Fast"],
]}
/>
<AdvancedSelect
label="Dual-GPU CFG"
hint="Runs the two guidance branches concurrently, one on a DiT replica on a second GPU (~1.7x on HunyuanVideo-1.5; the replica takes ~20 GB VRAM there). Auto engages only where the output is bit-identical to single-GPU (the Eager speed tier); On also parallelizes the compiled stack, accepting fp-noise-level divergence."
badge={<ResolvedBadge status={status} controlKey="cfg_parallel" />}
value={cfgParallel}
onValueChange={(v) => setCfgParallel(v as typeof cfgParallel)}
options={[
["auto", "Auto"],
["off", "Off"],
["on", "On"],
]}
/>
{status?.loaded && canReapply && (