From 514f4c60fe196af33b08c93237bf6a061e766350 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 10:06:52 +0000 Subject: [PATCH] Harden the video speed stack: cache quality pin, device identity, transactional caches, quant safety - Wan2.2-A14B step cache: pin the balanced FBCache threshold to 0.08 even when quant is active (per-family override in diffusion_cache.py). Auto-fp8 made the generic quant promotion (0.12) the family's effective default at pairwise LPIPS 0.128, over the 0.08 quality gate the balanced preset is held to. Measured operating point with fp8 actually engaged (1280x720/81f/50 steps, B200): fb@0.08 = 1.08x at 0.129 vs the old fb@0.12 = 2.58x at 0.181; documented in the preset table. Explicit thresholds and the fast preset are unaffected. - MagCache curves: validated the shipped 33-frame calibrations at the production 121-frame default for hunyuanvideo-1.5-720p, hunyuanvideo-1.5 (480p) and wan2.2-ti2v-5b. Fresh 121-frame calibrations differ by <= 0.024 max abs entry and produce byte-identical frames at the auto presets (hv720 quality 1.69x at LPIPS 0.042, hv480 quality 1.66x at 0.018, wan5b balanced 1.74x at 0.026, all pairwise vs the same-load uncached stack), so the curves ship unchanged with the frame-count transfer documented next to them. - Dual-GPU CFG parallelism: the secondary-device pick now prefers a device whose name and compute capability match the primary, and the gate declines a mismatched pair in auto mode (eager kernel selection is arch-dependent, so the advertised bit-identity cannot hold across different GPU models); an explicit cfg_parallel=on proceeds but is downgraded to lossless=False with a warning. - A14B expert step cache is now all-or-none, mirroring the transactional quant loop: a mixed outcome (cache engaged on one expert but not the other) is rolled back and reported uncached with the failure reason, on both the load path and the generation-time auto toggle. - Partial torchao quantization is no longer reported as dense: after an in-place quantize_/caster failure, the DiT / text encoder / VAE is scanned for leftover torchao tensor-subclass parameters and the load fails with a clear error when any are found (a half-quantized module cannot run as dense, and offload's Module.to() crashes on torchao tensors). Failures that swapped nothing keep the best-effort dense fallback. - Cleanup: apply_attention_backend / apply_speed_optims / the attention trim are called once on the pipe (they already fan out over every DiT internally), so the second A14B expert no longer passes through them twice; the stale dual-DiT helper comment is rewritten to match the two helper shapes. Tests: device-identity picker/gate/lossy-plan coverage, per-family threshold pin scoping, all-or-none rollback in both failure directions, and partial-quant detection for all three quant modules. --- .../backend/core/inference/diffusion_cache.py | 39 +++- .../core/inference/diffusion_cfg_parallel.py | 74 ++++++- .../core/inference/diffusion_precision.py | 8 + .../inference/diffusion_transformer_quant.py | 50 +++++ .../core/inference/diffusion_vae_quant.py | 8 + studio/backend/core/inference/video.py | 180 +++++++++++------- studio/backend/tests/test_diffusion_cache.py | 31 +++ .../tests/test_diffusion_cfg_parallel.py | 87 ++++++++- .../backend/tests/test_diffusion_precision.py | 39 ++++ .../tests/test_diffusion_transformer_quant.py | 69 +++++++ .../backend/tests/test_diffusion_vae_quant.py | 35 ++++ studio/backend/tests/test_video_backend.py | 91 +++++++++ 12 files changed, 630 insertions(+), 81 deletions(-) diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 2c42a40385..ff04b4a64c 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -67,7 +67,9 @@ CACHE_QUALITY_LEVELS = (CQ_QUALITY, CQ_BALANCED, CQ_FAST) # (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. +# three rather than just the threshold. Re-measured at the production 121-frame default +# (same protocol): quality = 1.69x at 0.042 (720p) / 1.66x at 0.018 (480p) -- the +# 33-frame operating points transfer (see the ratio-curve note below). _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), @@ -84,6 +86,24 @@ _FBCACHE_QUALITY_THRESHOLDS: dict[str, tuple[float, float]] = { CQ_FAST: (QUANT_FBCACHE_THRESHOLD, 0.15), } +# Per-family FBCache threshold overrides: (family, preset) -> (dense, quant-active). +# Wan2.2-A14B: on reference hardware the UNSET precision auto-promotes to fp8, which +# made the generic quant-active promotion (0.08 -> 0.12) the family's EFFECTIVE +# default -- but fb@0.12 measures pairwise LPIPS 0.128 (2.88x, dense probe, B200, +# 1280x720/33f/50 steps), far over the <= 0.08 quality gate the balanced preset is +# held to, and the round-3 adjudication that kept FBCache as the auto mode assumed +# the 0.08 point (1.28x/0.098). No compliant faster point exists (fb@0.04 is SLOWER +# than uncached at 0.016; MagCache was measured worse), so balanced pins 0.08 with +# quant active too: quality first, the fast points stay one explicit preset away. +# Operating point with the shipped quant actually engaged (fp8 DiTs, production +# shape 1280x720/81f/50 steps, pairwise vs the same fp8 uncached load, B200): +# fb@0.08 = 1.08x at LPIPS 0.129, vs the old effective default fb@0.12 = 2.58x at +# 0.181 -- fp8's residual noise makes the quantised cache trip the quality gate at +# ANY speedup, so 0.08 is the least-drift cache-on point, not a compliant one. +_FAMILY_FBCACHE_THRESHOLDS: dict[tuple[str, str], tuple[float, float]] = { + ("wan2.2-t2v-a14b", CQ_BALANCED): (DEFAULT_FBCACHE_THRESHOLD, DEFAULT_FBCACHE_THRESHOLD), +} + def normalize_cache_quality(value: Optional[str]) -> Optional[str]: """Lower/strip a requested cache quality; None / "" / "auto" -> None (the loader @@ -133,6 +153,16 @@ FBCACHE_MIN_STEPS = 20 # 50-step curve within 0.027 after nearest-interpolation, so ONE curve per family is # enough -- diffusers interpolates it to the actual step count. Conditional-branch curve # per the MagCache calibration guidance. +# +# Frame-count transfer VALIDATED at the production default (121 frames): these curves +# were calibrated on 33-frame clips, and recalibrating each family at its production +# shape (121f / 50 steps / default resolution, B200) moves the curve by <= 0.024 max +# abs entry diff (hv720 0.019, hv480 0.021, wan5b 0.024) -- small enough that the auto +# preset's skip schedule is UNCHANGED: the 33f-curve and fresh-121f-curve runs produced +# byte-identical frames on every family, so the 33f curves ship as-is. Measured at 121f +# with the shipped curves (pairwise LPIPS vs the same-load uncached compiled stack): +# hv720 quality 1.69x at 0.042, hv480 quality 1.66x at 0.018, wan5b balanced 1.74x at +# 0.026 -- each inside the gate its 33f adjudication was held to. _MAGCACHE_720P_RATIOS = ( 1.0, 1.0226, @@ -617,7 +647,12 @@ def apply_step_cache( preset_thr, mag_skip, mag_retention = _MAGCACHE_QUALITY_PRESETS[quality] thr = threshold if threshold is not None else preset_thr else: - dense_thr, quant_thr = _FBCACHE_QUALITY_THRESHOLDS[quality] + # A family override wins over the generic preset table (Wan2.2-A14B pins its + # balanced threshold to the quality-gated 0.08 even when quant is active). + dense_thr, quant_thr = _FAMILY_FBCACHE_THRESHOLDS.get( + (str(family or "").strip().lower(), quality), + _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 diff --git a/studio/backend/core/inference/diffusion_cfg_parallel.py b/studio/backend/core/inference/diffusion_cfg_parallel.py index 7db963f6b2..7873643794 100644 --- a/studio/backend/core/inference/diffusion_cfg_parallel.py +++ b/studio/backend/core/inference/diffusion_cfg_parallel.py @@ -151,6 +151,7 @@ class CFGParallelProxy: *, compiled: bool, explicit_on: bool, + device_match: bool = True, logger: Any = None, ) -> None: import torch @@ -161,6 +162,10 @@ class CFGParallelProxy: self._logger = logger self._compiled = bool(compiled) self._explicit_on = bool(explicit_on) + # False when the replica sits on a DIFFERENT GPU model/arch than the primary + # (explicit-on only; auto declines the mismatch at the gate): eager kernels + # differ across archs, so lossless must never be reported for the pair. + self._device_match = bool(device_match) self._ctx: Optional[str] = None self.enabled = False self.dispatch = "inline" @@ -292,7 +297,9 @@ class CFGParallelProxy: # 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 + # A replica on a DIFFERENT device model/arch (explicit-on only) runs different + # eager kernels too, so it is never lossless either. + lossless = not self._compiled and self._device_match 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" @@ -447,11 +454,34 @@ def _restore_threadsafe_cudnn_attention() -> None: # ── gate + build ────────────────────────────────────────────────────────────────── -def _pick_secondary_device(primary_index: int) -> tuple[Optional[int], int]: - """(most-free visible CUDA device != primary, its free bytes).""" +def _device_identity(idx: int) -> Optional[tuple]: + """(device name, compute capability) for CUDA device ``idx``, or None when the + props cannot be queried (a stubbed/old torch): identity is then treated as + unknown and the check stays best-effort rather than blocking the engage.""" import torch - best, best_free = None, -1 + try: + return ( + str(torch.cuda.get_device_name(idx)), + tuple(torch.cuda.get_device_capability(idx)), + ) + except Exception: # noqa: BLE001 -- unqueryable props: identity unknown + return None + + +def _pick_secondary_device(primary_index: int) -> tuple[Optional[int], int, bool]: + """(secondary CUDA device != primary, its free bytes, identity-match flag). + + The bit-identity contract needs the SAME kernels on both branches, and eager + kernel selection is arch-dependent (cuDNN heuristics, SM-count-dependent tiling), + so the picker prefers the most-free device whose (name, capability) MATCH the + primary's; only when no matching device exists does it fall back to the most-free + mismatched one (so an explicit ``on`` can still engage, lossy). An unqueryable + identity counts as a match (best-effort, the pre-check behaviour).""" + import torch + + primary_id = _device_identity(primary_index) + best, best_free, best_match = None, -1, False for idx in range(torch.cuda.device_count()): if idx == primary_index: continue @@ -459,9 +489,11 @@ def _pick_secondary_device(primary_index: int) -> tuple[Optional[int], int]: 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 + candidate_id = _device_identity(idx) + match = primary_id is None or candidate_id is None or candidate_id == primary_id + if (match, free) > (best_match, best_free): + best, best_free, best_match = idx, free, match + return best, best_free, best_match def maybe_enable_cfg_parallel( @@ -532,10 +564,35 @@ def maybe_enable_cfg_parallel( 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) + primary_index = p_dev.index or 0 + secondary, free, device_match = _pick_secondary_device(primary_index) need = weight_bytes + _REPLICA_HEADROOM_BYTES if secondary is None: return None, "no queryable secondary CUDA device" + if not device_match: + # A different GPU model/arch runs different eager kernels (cuDNN + # heuristics, SM-count-dependent tiling), so the byte-identity the AUTO + # policy promises cannot hold across the pair. Auto declines; an explicit + # "on" proceeds but is downgraded to lossy (plan_generation reports + # lossless=False) with a warning. + primary_name = (_device_identity(primary_index) or ("unknown",))[0] + secondary_name = (_device_identity(secondary) or ("unknown",))[0] + if not explicit_on: + return None, ( + f"secondary cuda:{secondary} ({secondary_name}) is a different device " + f"than the primary ({primary_name}); eager kernels are arch-dependent, " + "so bit-identity cannot hold -- request cfg_parallel=on to accept the " + "divergence" + ) + if logger is not None: + logger.warning( + "diffusion.cfg_parallel: replica device cuda:%d (%s) differs from the " + "primary (%s); explicit on proceeds but the output is NOT bit-identical " + "to the single-GPU run (lossless=False)", + secondary, + secondary_name, + primary_name, + ) if free < need: return None, ( f"secondary cuda:{secondary} has {free / 2**30:.1f} GiB free, " @@ -589,6 +646,7 @@ def maybe_enable_cfg_parallel( guider, compiled = compiled, explicit_on = explicit_on, + device_match = device_match, logger = logger, ) pipe.transformer = proxy diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index f3b2b25705..ebbba22e1a 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -362,6 +362,14 @@ def quantize_text_encoders( caster(encoder, target) cast.append(attr) except Exception as exc: # noqa: BLE001 — leave this encoder dense + # The torchao casters mutate the encoder in place module-by-module, so a + # mid-pass failure may have left it PARTIALLY quantized -- a state that + # cannot run as the dense encoder this fallback would report (and that + # offload's Module.to() hard-crashes on). Fail the load for that; a clean + # miss (nothing swapped, e.g. layerwise fp8) stays best-effort dense. + from .diffusion_transformer_quant import raise_if_partially_quantized + + raise_if_partially_quantized(encoder, what = f"text_encoder_quant {mode}:{attr}", exc = exc) _warn(logger, f"{mode}:{attr}", exc) return mode if cast else None diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index bb7b8f99e8..82cff2fcec 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -574,6 +574,51 @@ def make_filter_fn( return filter_fn +def torchao_quantized_param_fqns(module: Any) -> list[str]: + """FQNs of ``module`` parameters whose tensors are torchao subclasses. + + torchao's ``quantize_`` mutates the module in place, swapping Linear/Conv weights + to tensor subclasses one submodule at a time -- so an exception mid-pass (OOM on a + 14B DiT, a layer a kernel rejects) leaves the EARLIER layers quantized. A module in + that state must fail the load: it cannot run as the dense module the caller would + otherwise report (offload's ``Module.to()`` hard-crashes on torchao tensors, and + mixed dense/quant weights are an unvalidated numeric state). Detection keys on the + tensor class's module path ("torchao" in ``type(t).__module__``), which covers every + torchao subclass family (affine / float8 / nvfp4 / mx) without importing torchao. + Best-effort: an unscannable module reports no leftovers (the pre-check behaviour).""" + names: list[str] = [] + named = getattr(module, "named_parameters", None) + if not callable(named): + return names + try: + for name, param in named(): + for tensor in (param, getattr(param, "data", None)): + if tensor is None: + continue + if "torchao" in (getattr(type(tensor), "__module__", "") or ""): + names.append(name) + break + except Exception: # noqa: BLE001 -- scan is best-effort; report what was found + pass + return names + + +def raise_if_partially_quantized(module: Any, *, what: str, exc: Exception) -> None: + """After an in-place ``quantize_``/caster failure: if ``module`` was left PARTIALLY + quantized (some torchao tensor-subclass params present), raise so the load fails + with a clear error instead of the caller reporting a dense fallback and running a + half-quantized module. A clean failure (no torchao params) returns, keeping the + best-effort dense-fallback contract.""" + leftover = torchao_quantized_param_fqns(module) + if not leftover: + return + raise RuntimeError( + f"{what} failed midway and left {len(leftover)} parameter(s) quantized " + f"(e.g. '{leftover[0]}'); the module is partially quantized and cannot fall " + f"back to dense -- reload the model (original error: {exc})" + ) from exc + + def quantize_transformer( pipe: Any, target: Any, @@ -625,6 +670,11 @@ def quantize_transformer( pass return scheme except Exception as exc: # noqa: BLE001 — leave the transformer dense -> GGUF fallback + # quantize_ swaps weights module-by-module, so a mid-pass failure may have left + # some layers quantized: that state cannot run as dense (offload's Module.to() + # crashes on torchao tensors; mixed precision is unvalidated), so fail the load + # instead of reporting a dense fallback. A clean miss stays best-effort dense. + raise_if_partially_quantized(transformer, what = f"transformer_quant {scheme}", exc = exc) _warn(logger, scheme, exc) return None diff --git a/studio/backend/core/inference/diffusion_vae_quant.py b/studio/backend/core/inference/diffusion_vae_quant.py index b9acefd863..24fddbbd4c 100644 --- a/studio/backend/core/inference/diffusion_vae_quant.py +++ b/studio/backend/core/inference/diffusion_vae_quant.py @@ -336,6 +336,14 @@ def quantize_vae( _cast_vae_fp8(vae, target) return mode except Exception as exc: # noqa: BLE001 — leave the VAE dense + # fp8_dynamic's quantize_ swaps conv/linear weights module-by-module, so a + # mid-pass failure may have left the VAE PARTIALLY quantized -- fail the load + # for that instead of reporting a dense fallback (mixed dense/quant weights + # are unvalidated and offload's Module.to() crashes on torchao tensors). A + # clean miss (nothing swapped, e.g. layerwise fp8) stays best-effort dense. + from .diffusion_transformer_quant import raise_if_partially_quantized + + raise_if_partially_quantized(vae, what = f"vae_quant {mode}", exc = exc) _warn(logger, mode, exc) return None diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index b939b4455c..87fc720647 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -408,16 +408,18 @@ def _progress(phase: Optional[str], **extra: Any) -> dict[str, Any]: # ── dual-DiT (Wan2.2-A14B MoE) helpers ──────────────────────────────────────── # -# The imported optimisation helpers (apply_speed_optims / apply_attention_backend / -# apply_step_cache) and the dense quantiser all read ``pipe.transformer`` and act on -# that ONE denoiser -- correct for every single-DiT family (LTX-2, Wan2.2-TI2V-5B). # Wan2.2-A14B is a dual-expert MoE: ``transformer`` handles the high-noise steps and # ``transformer_2`` the low-noise steps (pipeline_wan.py routes by boundary_ratio), so # an optimisation applied only to ``transformer`` would leave the second expert eager / -# unquantised / on the wrong attention kernel for half the schedule. Rather than fork -# each helper, present the second DiT to them AS ``pipe.transformer`` via a thin proxy -# and call the helper a second time, so the helpers stay untouched and single-DiT loads -# are bit-identical (the proxy is only built for is_moe families). +# unquantised / uncached for half the schedule. Two helper shapes exist: +# - apply_speed_optims / apply_attention_backend / install_hunyuan_attention_trim fan +# out over EVERY denoiser DiT internally (_denoiser_dits / _attention_dits cover a +# present ``transformer_2``), so the loader calls them ONCE on the pipe. +# - apply_step_cache and the dense quantiser read ``pipe.transformer`` and act on +# that ONE denoiser (the cache additionally needs a per-expert calibrated curve). +# Rather than fork them, present the second DiT AS ``pipe.transformer`` via a thin +# proxy view and call the helper once per expert, so the helpers stay untouched and +# single-DiT loads are bit-identical (the proxy is only built for is_moe families). def _transformer_names(pipe: Any, fam: VideoFamily) -> tuple[str, ...]: @@ -470,6 +472,42 @@ def _views_for(pipe: Any, fam: VideoFamily) -> tuple[Any, ...]: return (pipe,) +def _step_cache_all_or_none( + pipe: Any, + fam: VideoFamily, + engage_fn: Any, + *, + logger: Any, +) -> tuple[Optional[str], Optional[str]]: + """Run ``engage_fn(view, expert_name)`` (apply_step_cache or the auto toggle) over + every expert and enforce ALL-OR-NONE, mirroring the transactional quant loop: on a + mixed outcome (the cache engaged on some experts but not all -- apply_step_cache + never raises, it returns None and rolls back only that one transformer) the engaged + expert(s) are disengaged and the cache is reported off with the failure surfaced. + Without this, a second-expert failure would leave the FIRST expert cached while + status keys on it and reports the whole MoE as cached (half the schedule then runs + uncached at a mismatched quality/perf point). Returns (mode-or-None, failure + reason-or-None); a single-DiT family can never see a mixed outcome, so its + behaviour is unchanged.""" + results: list[tuple[Any, str, Optional[str]]] = [] + for view, expert_name in zip(_views_for(pipe, fam), _transformer_names(pipe, fam)): + results.append((view, expert_name, engage_fn(view, expert_name))) + engaged = [(view, name, mode) for view, name, mode in results if mode is not None] + if engaged and len(engaged) < len(results): + missing = ", ".join(name for _, name, mode in results if mode is None) + for view, name, _mode in engaged: + _disengage_step_cache( + getattr(view, "transformer", None), + reason = f"all-or-none rollback: cache did not engage on {missing}", + logger = logger, + ) + return None, ( + f"step cache engaged on only {len(engaged)}/{len(results)} experts " + f"({missing} failed); disengaged all experts and running uncached" + ) + return (engaged[0][2] if engaged else None), None + + class VideoBackend: """One loaded video pipeline; loads swap it atomically (same model as images).""" @@ -1476,20 +1514,22 @@ class VideoBackend: cache_request = ( auto_cache_mode(fam.name) if default_cache_steps >= FBCACHE_MIN_STEPS else None ) - cache_engaged = None - # Each view is zipped with the pipe attribute it exposes as ``transformer`` (the + # Each expert view passes the pipe attribute it exposes as ``transformer`` (the # expert-view iteration contract): a dual-expert MoE's second view passes # expert="transformer_2" so MagCache resolves THAT expert's calibrated curve -- - # the experts split the schedule at the boundary timestep, so their curves differ. - for view, expert_name in zip(views, _transformer_names(pipe, fam)): - engaged = apply_step_cache( + # the experts split the schedule at the boundary timestep, so their curves + # differ. All-or-none across experts (mirroring the quant loop above): a mixed + # outcome is rolled back and reported uncached instead of leaving one expert + # cached while status reports the whole MoE as cached. + def _engage_load_cache(view: Any, expert_name: str) -> Optional[str]: + return apply_step_cache( view, mode = cache_request, threshold = transformer_cache_threshold, # A quantized transformer's block residuals are larger, so it needs the # higher FBCache trigger threshold to cache at all. Mirror the image path # (diffusion.py): both an engaged transformer_quant AND a GGUF checkpoint - # (quantized weights) count as quant-active here (cache_quant_active, L1172). + # (quantized weights) count as quant-active here (cache_quant_active). quant_active = cache_quant_active, family = fam.name, steps = default_cache_steps, @@ -1497,14 +1537,18 @@ class VideoBackend: expert = expert_name, logger = logger, ) - if view is pipe: - cache_engaged = engaged + + cache_engaged, cache_partial_reason = _step_cache_all_or_none( + pipe, fam, _engage_load_cache, logger = logger + ) # The auto decision can flip at generation time, but only on a DiT that # supports caching at all (a non-CacheMixin transformer can never engage). cache_may_toggle = cache_auto and callable( getattr(getattr(pipe, "transformer", None), "enable_cache", None) ) - if cache_auto: + if cache_partial_reason: + cache_reason = cache_partial_reason + elif cache_auto: if cache_engaged: cache_reason = ( f"auto: {default_cache_steps}-step default schedule reaches " @@ -1519,37 +1563,32 @@ class VideoBackend: ) else: cache_reason = "requested" - attention_engaged = None - attention_trim_engaged = False - speed_optims: tuple = () # A dense torchao transformer on the pipeline path is not a GGUF one, so is_gguf # keys off the load kind (gguf) AND no quant having engaged. gguf_transformer = kind == "gguf" and transformer_quant_engaged is None - for view in views: - # apply_attention_backend acts on ``view.transformer``; calling it once per - # view sets the kernel on each expert. The engaged values match across - # experts (same device/family/mode), so record the first pass. - # HunyuanVideo-1.5 only: drop the ~99% zero-padded text tokens from the joint - # attention so it runs the fused (cuDNN/flash) SDPA kernel instead of the dense-mask - # fallback (~18x/DiT-forward at 121 frames, cosine ~1.0). Must precede the backend set - # so the requested kernel pins onto the new processors. No-op for every other family. - # A speed lever like the attention backend below, so honor an explicit Speed="off" (the - # bit-exact reference path keeps the stock dense-mask attention). - trim = ( - install_hunyuan_attention_trim(view, fam, logger = logger) - if effective_speed != SPEED_OFF - else False - ) - engaged = apply_attention_backend( - view, - select_attention_backend( - target, attention_backend, speed_active = effective_speed != SPEED_OFF - ), - logger = logger, - ) - if view is pipe: - attention_engaged = engaged - attention_trim_engaged = trim + # install_hunyuan_attention_trim and apply_attention_backend fan out over every + # denoiser DiT internally (diffusion_attention._attention_dits covers + # ``transformer`` AND a present ``transformer_2``), so ONE pipe-level call + # covers a dual-expert MoE -- a per-view loop would pass the second expert + # through them twice (idempotent, but wasted work). + # Trim is HunyuanVideo-1.5 only: drop the ~99% zero-padded text tokens from the joint + # attention so it runs the fused (cuDNN/flash) SDPA kernel instead of the dense-mask + # fallback (~18x/DiT-forward at 121 frames, cosine ~1.0). Must precede the backend set + # so the requested kernel pins onto the new processors. No-op for every other family. + # A speed lever like the attention backend below, so honor an explicit Speed="off" (the + # bit-exact reference path keeps the stock dense-mask attention). + attention_trim_engaged = ( + install_hunyuan_attention_trim(pipe, fam, logger = logger) + if effective_speed != SPEED_OFF + else False + ) + attention_engaged = apply_attention_backend( + pipe, + select_attention_backend( + target, attention_backend, speed_active = effective_speed != SPEED_OFF + ), + logger = logger, + ) # Pre-warmed torch.compile cache (Mega-cache), mirroring the image backend: when a # compiled tier will run, point inductor at a per-fingerprint dir and load a matching # bundle BEFORE the first compiled forward. Measured on HunyuanVideo-1.5-480p (B200): @@ -1587,23 +1626,25 @@ class VideoBackend: # failed or cancelled load must restore TORCHINDUCTOR_CACHE_DIR itself # (_run_load's error handler, token-scoped like the globals). self._precommit_compile_cache = (_load_token, compile_ctx) - for view in views: - applied = apply_speed_optims( - view, - target, - is_gguf = gguf_transformer, - family = fam, - speed_mode = effective_speed, - # An auto cache that could still engage mid-session also drops - # fullgraph: enabling FBCache under a fullgraph-compiled DiT would - # crash the first cached generation. - cache_active = cache_engaged is not None or cache_may_toggle, - offload_active = plan.offload_policy != "none", - ) - if view is pipe: - speed_optims = tuple(k for k, v in applied.items() if v) + ( - ("hunyuan_attn_trim",) if attention_trim_engaged else () - ) + # apply_speed_optims fans out over every denoiser DiT internally + # (diffusion_speed._denoiser_dits: the regional compile and the qkv fusion + # cover ``transformer_2``; the VAE/global levers are pipe-level), so one + # pipe-level call covers a dual-expert MoE. + applied = apply_speed_optims( + pipe, + target, + is_gguf = gguf_transformer, + family = fam, + speed_mode = effective_speed, + # An auto cache that could still engage mid-session also drops + # fullgraph: enabling FBCache under a fullgraph-compiled DiT would + # crash the first cached generation. + cache_active = cache_engaged is not None or cache_may_toggle, + offload_active = plan.offload_policy != "none", + ) + speed_optims = tuple(k for k, v in applied.items() if v) + ( + ("hunyuan_attn_trim",) if attention_trim_engaged else () + ) with self._generate_lock: # A cancelled/superseded load must not place weights on the GPU the arbiter # may already have handed to another backend; recheck right before placement @@ -2281,11 +2322,12 @@ class VideoBackend: # drops it. Explicit choices never toggle. Runs per view so a dual-DiT # MoE toggles both experts. if state.cache_auto: - toggled = state.transformer_cache - for view, expert_name in zip( - _views_for(pipe, fam), _transformer_names(pipe, fam) - ): - toggled = maybe_toggle_step_cache( + # All-or-none across MoE experts, exactly like the load path: a + # mixed toggle outcome (one expert engaged, one not) is rolled + # back and reported uncached instead of keying status on + # whichever expert happened to be toggled last. + def _toggle_cache(view: Any, expert_name: str) -> Optional[str]: + return maybe_toggle_step_cache( view, steps = steps, quant_active = state.cache_quant_active, @@ -2296,6 +2338,10 @@ class VideoBackend: expert = expert_name, logger = logger, ) + + toggled, toggle_partial_reason = _step_cache_all_or_none( + pipe, fam, _toggle_cache, logger = logger + ) if toggled != state.transformer_cache: # _VideoLoadState is frozen (loads swap it as one unit); this # tracks the pipe-level toggle that already happened so @@ -2304,7 +2350,7 @@ class VideoBackend: entry = (state.resolved or {}).get("transformer_cache") if isinstance(entry, dict): entry["value"] = toggled or "off" - entry["reason"] = ( + entry["reason"] = toggle_partial_reason or ( f"auto: {steps}-step generation " + ("reaches" if toggled else "is below") + f" {FBCACHE_MIN_STEPS}" diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py index 0597a4461b..7714c2a11a 100644 --- a/studio/backend/tests/test_diffusion_cache.py +++ b/studio/backend/tests/test_diffusion_cache.py @@ -143,6 +143,37 @@ def test_explicit_threshold_overrides_quant(monkeypatch): assert t.enabled_with.threshold == 0.2 +def test_a14b_balanced_pins_quant_threshold(monkeypatch): + # Wan2.2-A14B's effective default is quant-active (unset precision auto-promotes to + # fp8 on reference hardware), and the generic 0.08 -> 0.12 promotion violates the + # family's <= 0.08 quality gate (fb@0.12 measured pairwise LPIPS 0.128) -- so the + # balanced preset pins 0.08 with quant active too. Other families keep the table. + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + apply_step_cache(_pipe(t), mode = "fbcache", quant_active = True, family = "wan2.2-t2v-a14b") + assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD + other = _MixinTransformer() + apply_step_cache(_pipe(other), mode = "fbcache", quant_active = True, family = "ltx-2") + assert other.enabled_with.threshold == QUANT_FBCACHE_THRESHOLD + + +def test_a14b_pin_scope_is_balanced_only(monkeypatch): + # The pin is (family, balanced)-scoped: an explicit threshold still wins, and the + # explicit "fast" preset keeps the generic quant table (the 2.9x point stays one + # request away). + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + apply_step_cache( + _pipe(t), mode = "fbcache", threshold = 0.12, quant_active = True, family = "wan2.2-t2v-a14b" + ) + assert t.enabled_with.threshold == 0.12 + fast = _MixinTransformer() + apply_step_cache( + _pipe(fast), mode = "fbcache", quality = "fast", quant_active = True, family = "wan2.2-t2v-a14b" + ) + assert fast.enabled_with.threshold == 0.15 + + def test_non_cachemixin_runs_uncached(monkeypatch): # A transformer without enable_cache (e.g. Z-Image) must NOT install the standalone hook # -- its pipeline opens no cache_context, so it runs uncached instead of crashing at gen. diff --git a/studio/backend/tests/test_diffusion_cfg_parallel.py b/studio/backend/tests/test_diffusion_cfg_parallel.py index 9c08501076..0dd5b70d1f 100644 --- a/studio/backend/tests/test_diffusion_cfg_parallel.py +++ b/studio/backend/tests/test_diffusion_cfg_parallel.py @@ -132,20 +132,30 @@ def _stub_torch( *, device_count = 2, free = None, + names = None, + caps = None, + with_identity = True, ): torch = types.ModuleType("torch") torch.Tensor = _FakeTensor free = free if free is not None else {} + names = names if names is not None else {} + caps = caps if caps is not None else {} def _mem_get_info(idx): return free.get(idx, (64 << 30, 80 << 30)) - torch.cuda = types.SimpleNamespace( + cuda_kwargs = dict( is_available = lambda: device_count > 0, device_count = lambda: device_count, mem_get_info = _mem_get_info, empty_cache = lambda: None, ) + if with_identity: + # Homogeneous by default so the pre-identity-gate tests keep engaging. + cuda_kwargs["get_device_name"] = lambda idx: names.get(idx, "Fake GPU") + cuda_kwargs["get_device_capability"] = lambda idx: caps.get(idx, (9, 0)) + torch.cuda = types.SimpleNamespace(**cuda_kwargs) torch.inference_mode = contextlib.nullcontext monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -157,12 +167,20 @@ def _make_proxy( compiled = False, explicit_on = False, fail_enable = False, + device_match = True, ): _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) + proxy = CFGParallelProxy( + primary, + replica, + guider, + compiled = compiled, + explicit_on = explicit_on, + device_match = device_match, + ) return proxy, primary, replica, guider @@ -289,8 +307,69 @@ def test_pick_secondary_prefers_most_free(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 + idx, free, match = _pick_secondary_device(0) + assert idx == 2 and free == 40 << 30 and match is True + + +# ── device identity (bit-identity needs the SAME kernels -> the same arch) ───────── +def test_pick_secondary_prefers_identity_match_over_free(monkeypatch): + # cuda:2 has the most free VRAM but is a different model; cuda:1 matches the + # primary, so the picker takes it (bit-identity beats headroom). + _stub_torch( + monkeypatch, + device_count = 3, + free = {1: (30 << 30, 80 << 30), 2: (60 << 30, 80 << 30)}, + names = {0: "NVIDIA B200", 1: "NVIDIA B200", 2: "NVIDIA H100"}, + ) + idx, free, match = _pick_secondary_device(0) + assert idx == 1 and free == 30 << 30 and match is True + + +def test_pick_secondary_falls_back_to_mismatch(monkeypatch): + # No matching device exists: the most-free mismatched one is still returned (an + # explicit "on" can engage it, lossy), flagged match=False. + _stub_torch(monkeypatch, names = {0: "NVIDIA B200", 1: "NVIDIA H100"}) + idx, _free, match = _pick_secondary_device(0) + assert idx == 1 and match is False + + +def test_pick_secondary_unknown_identity_counts_as_match(monkeypatch): + # A torch without queryable device props (identity unknown) must stay best-effort: + # the check never blocks what the pre-identity-gate behaviour allowed. + _stub_torch(monkeypatch, with_identity = False) + idx, _free, match = _pick_secondary_device(0) + assert idx == 1 and match is True + + +def test_gate_auto_declines_device_mismatch(monkeypatch): + _stub_torch(monkeypatch, names = {0: "NVIDIA B200", 1: "NVIDIA H100"}) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam()) + assert proxy is None + assert "different device" in reason and "cfg_parallel=on" in reason + + +def test_gate_auto_declines_capability_mismatch(monkeypatch): + # Same marketing name, different compute capability: still a different arch. + _stub_torch(monkeypatch, caps = {0: (10, 0), 1: (9, 0)}) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam()) + assert proxy is None and "different device" in reason + + +def test_gate_explicit_on_allows_device_mismatch(monkeypatch): + # Explicit "on" passes the identity gate (downgraded to lossy); it then fails soft + # at the replica load (the fake DiT has no from_pretrained), proving the gate order. + _stub_torch(monkeypatch, names = {0: "NVIDIA B200", 1: "NVIDIA H100"}) + proxy, reason = _gate(monkeypatch, _CtxPipe(_FakeDiT()), _fam(), requested = "on") + assert proxy is None and reason == "replica load failed" + + +def test_plan_lossy_on_device_mismatch(monkeypatch): + # An explicit-on engage across mismatched devices must never report lossless, even + # on the (otherwise byte-identical) eager tier. + proxy, _, _, _ = _make_proxy(monkeypatch, explicit_on = True, device_match = False) + plan = proxy.plan_generation(cache_engaged = False, steps = 30, width = 1280, height = 720, frames = 33) + assert plan["enabled"] is True and plan["lossless"] is False + proxy.shutdown() # ── proxy semantics ─────────────────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 301eff4d27..2b90162d88 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -721,3 +721,42 @@ def test_fp8_dynamic_filter_skips_zero_row_linear(monkeypatch): live = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.5, 0.5]])) assert ff(dead, "text_model.encoder.layers.2.self_attn.out_proj") is False assert ff(live, "text_model.encoder.layers.2.mlp.fc1") is True + + +# ── partial in-place cast detection (fails the load, not a silent dense report) ──── +class _TorchaoLikeTensor: + """Detection keys on the tensor class's module path ("torchao" in __module__).""" + + +_TorchaoLikeTensor.__module__ = "torchao.quantization.linear_activation_quantized_tensor" + + +class _PartiallyCastEncoder: + def __init__(self): + self._swapped = False + + def named_parameters(self): + if self._swapped: + yield ("model.layers.0.mlp.up_proj.weight", _TorchaoLikeTensor()) + yield ("model.layers.1.mlp.up_proj.weight", types.SimpleNamespace()) + + +def test_quantize_partial_cast_failure_fails_load(monkeypatch): + # The caster mutates the encoder in place module-by-module: a mid-pass failure + # that left torchao params behind must raise (the encoder cannot run as the dense + # module a best-effort fallback would report), unlike the clean failure above. + _stub_torch(monkeypatch) + hooks = types.ModuleType("diffusers.hooks") + casting = types.ModuleType("diffusers.hooks.layerwise_casting") + casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",) + + def _swap_one_then_boom(module, **kwargs): + module._swapped = True + raise RuntimeError("encoder cast failed mid-pass") + + hooks.apply_layerwise_casting = _swap_one_then_boom + monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks) + monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting) + pipe = types.SimpleNamespace(text_encoder = _PartiallyCastEncoder()) + with pytest.raises(RuntimeError, match = "partially quantized"): + quantize_text_encoders(pipe, _target(), mode = "fp8") diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index a6617d73ce..10e4bba73e 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -736,3 +736,72 @@ def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch): assert ( filt(big, "blocks.0.attn2.to_k") is True ) # cross-attn K/V stay fp8 (embedder bias rescues rows) + + +# ── partial in-place quant detection (torchao_quantized_param_fqns) ──────────────── +class _TorchaoLikeTensor: + """Stands in for a torchao tensor subclass: detection keys on the class's module + path, so the fake just claims a torchao __module__.""" + + +_TorchaoLikeTensor.__module__ = "torchao.dtypes.affine_quantized_tensor" + + +class _MutableDiT: + """A fake transformer whose quantize_ pass 'swapped' one weight before failing.""" + + def __init__(self): + self._swapped = False + + def named_parameters(self): + if self._swapped: + yield ("blocks.0.attn.to_q.weight", _TorchaoLikeTensor()) + yield ("blocks.1.attn.to_q.weight", types.SimpleNamespace()) + + +def test_torchao_param_scan_detects_swapped_weights(): + dit = _MutableDiT() + assert tq.torchao_quantized_param_fqns(dit) == [] + dit._swapped = True + assert tq.torchao_quantized_param_fqns(dit) == ["blocks.0.attn.to_q.weight"] + # Unscannable object -> no leftovers reported (best-effort, the pre-check path). + assert tq.torchao_quantized_param_fqns(object()) == [] + + +def test_quantize_transformer_partial_failure_raises(monkeypatch): + # quantize_ swaps weights module-by-module, so a mid-pass exception (OOM on a 14B + # DiT, a layer a kernel rejects) can leave earlier layers quantized. That module + # cannot run as dense (offload's Module.to() crashes on torchao tensors), so the + # load must FAIL with a clear error instead of reporting a dense fallback. + monkeypatch.setattr( + tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8 + ) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg") + tqz = types.ModuleType("torchao.quantization") + + def _convert_one_then_boom(module, config, filter_fn = None): + module._swapped = True # the in-place swap of the first submodule + raise RuntimeError("OOM mid-conversion") + + tqz.quantize_ = _convert_one_then_boom + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + pipe = types.SimpleNamespace(transformer = _MutableDiT()) + with pytest.raises(RuntimeError, match = "partially quantized"): + quantize_transformer(pipe, _target(), mode = "fp8") + + +def test_quantize_transformer_clean_failure_still_falls_back_dense(monkeypatch): + # A failure that swapped NOTHING keeps the best-effort contract: dense fallback. + monkeypatch.setattr( + tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8 + ) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg") + tqz = types.ModuleType("torchao.quantization") + + def _boom(module, config, filter_fn = None): + raise RuntimeError("failed before any swap") + + tqz.quantize_ = _boom + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + pipe = types.SimpleNamespace(transformer = _MutableDiT()) + assert quantize_transformer(pipe, _target(), mode = "fp8") is None diff --git a/studio/backend/tests/test_diffusion_vae_quant.py b/studio/backend/tests/test_diffusion_vae_quant.py index b0a45fa49c..a872926b18 100644 --- a/studio/backend/tests/test_diffusion_vae_quant.py +++ b/studio/backend/tests/test_diffusion_vae_quant.py @@ -540,3 +540,38 @@ def test_quantize_vae_fp8_dynamic_probes_conv3d_for_video_vae(monkeypatch): == VAE_QUANT_FP8_DYNAMIC ) assert len(calls) == 1 + + +# ── partial in-place quant detection (fails the load, not a silent dense report) ─── +class _TorchaoLikeTensor: + """Detection keys on the tensor class's module path ("torchao" in __module__).""" + + +_TorchaoLikeTensor.__module__ = "torchao.dtypes.affine_quantized_tensor" + + +class _PartiallyQuantizedVae: + def __init__(self): + self._swapped = False + + def named_parameters(self): + if self._swapped: + yield ("decoder.up_blocks.0.conv.weight", _TorchaoLikeTensor()) + yield ("decoder.up_blocks.1.conv.weight", types.SimpleNamespace()) + + +def test_quantize_vae_partial_cast_failure_fails_load(monkeypatch): + # fp8_dynamic's quantize_ swaps conv/linear weights module-by-module: a mid-pass + # failure that left torchao params behind must raise instead of reporting a dense + # fallback (mixed weights are unvalidated; offload's Module.to() crashes on them). + _stub_torch(monkeypatch, cc = (10, 0)) + _allow_vae(monkeypatch, {VAE_QUANT_FP8}) + + def _swap_one_then_boom(v, t): + v._swapped = True + raise RuntimeError("mid-pass conv failure") + + monkeypatch.setattr(vq, "_cast_vae_fp8", _swap_one_then_boom) + pipe = types.SimpleNamespace(vae = _PartiallyQuantizedVae()) + with pytest.raises(RuntimeError, match = "partially quantized"): + quantize_vae(pipe, _target(), mode = "fp8") diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 68d711c4e9..084c47fa23 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -2079,3 +2079,94 @@ def test_begin_generate_preempts_running_prewarm(fake_runtime, monkeypatch): time.sleep(0.02) assert backend._generate_job_active is False backend.unload() + + +# ── all-or-none step cache across MoE experts (_step_cache_all_or_none) ──────────── +def _moe_pipe_and_fam(): + fam = types.SimpleNamespace(is_moe = True, name = "wan2.2-t2v-a14b") + t1 = types.SimpleNamespace(tag = "expert1") + t2 = types.SimpleNamespace(tag = "expert2") + pipe = types.SimpleNamespace(transformer = t1, transformer_2 = t2) + return pipe, fam, t1, t2 + + +def test_step_cache_all_or_none_rolls_back_second_expert_failure(monkeypatch): + # apply_step_cache never raises -- a second-expert failure returns None for that + # expert while the FIRST stays cached. The helper must disengage the engaged + # expert and report the cache off (all-or-none, mirroring the quant loop). + import core.inference.video as video + + pipe, fam, t1, _t2 = _moe_pipe_and_fam() + disengaged: list = [] + monkeypatch.setattr( + video, + "_disengage_step_cache", + lambda transformer, *, reason, logger = None: disengaged.append((transformer, reason)) or True, + ) + calls: list = [] + + def engage(view, expert_name): + calls.append(expert_name) + return "fbcache" if expert_name == "transformer" else None + + mode, reason = video._step_cache_all_or_none(pipe, fam, engage, logger = None) + assert calls == ["transformer", "transformer_2"] + assert mode is None + assert reason is not None and "1/2" in reason and "transformer_2" in reason + assert len(disengaged) == 1 and disengaged[0][0] is t1 + + +def test_step_cache_all_or_none_rolls_back_first_expert_failure(monkeypatch): + # Mirror image: only the SECOND expert engaged -> it is the one disengaged. + import core.inference.video as video + + pipe, fam, _t1, t2 = _moe_pipe_and_fam() + disengaged: list = [] + monkeypatch.setattr( + video, + "_disengage_step_cache", + lambda transformer, *, reason, logger = None: disengaged.append(transformer) or True, + ) + mode, reason = video._step_cache_all_or_none( + pipe, fam, + lambda view, expert_name: "magcache" if expert_name == "transformer_2" else None, + logger = None, + ) + assert mode is None and reason is not None + assert disengaged == [t2] + + +def test_step_cache_all_or_none_uniform_outcomes(monkeypatch): + # Both experts engaged -> the mode is reported with no rollback; neither engaged + # -> plain uncached with no failure reason (the pre-existing best-effort path). + import core.inference.video as video + + pipe, fam, _t1, _t2 = _moe_pipe_and_fam() + monkeypatch.setattr( + video, + "_disengage_step_cache", + lambda transformer, *, reason, logger = None: pytest.fail("no rollback on uniform outcome"), + ) + assert video._step_cache_all_or_none( + pipe, fam, lambda view, expert_name: "fbcache", logger = None + ) == ("fbcache", None) + assert video._step_cache_all_or_none( + pipe, fam, lambda view, expert_name: None, logger = None + ) == (None, None) + + +def test_step_cache_all_or_none_single_dit(monkeypatch): + # A single-DiT family runs the engage exactly once and can never see a mixed + # outcome -- behaviour identical to the pre-helper loop. + import core.inference.video as video + + fam = types.SimpleNamespace(is_moe = False, name = "wan2.2-ti2v-5b") + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace(), transformer_2 = None) + calls: list = [] + + def engage(view, expert_name): + calls.append((view, expert_name)) + return "magcache" + + assert video._step_cache_all_or_none(pipe, fam, engage, logger = None) == ("magcache", None) + assert calls == [(pipe, "transformer")]