fix(bench): mirror production contracts in the quant/video benchmarks

Reject partial dual-DiT quantization in video_speedmem_bench (the loader fails
that load all-or-none; a mixed quantized/dense row is unloadable), toggle the
generation-time FBCache recheck on every expert view like the loader's per-view
iteration, and rescore lpips_vs_reference in a post-pass so a --configs order
that lists reference late no longer publishes null.

In quant_speedmem_bench, track per-encoder engagement via a weight-storage
fingerprint so a partial multi-encoder cast cannot certify a still-dense
encoder with a ~1.0 cosine, and load vae_force_fp32 families (Wan) at fp32
with a matching latent dtype so the dense VAE row measures what production
runs.

Gate the attention-trim tests with pytest.importorskip so a no-torch
environment keeps the backend test suite collectable.

scripts/quant_speedmem_bench.py

@@ -286,6 +286,15 @@ def _te_hidden_refs(bag, toks, device):
     return refs
 
 
+def _te_param_sig(te) -> tuple:
+    """Weight-storage fingerprint (per-param impl class + dtype). It changes iff a caster
+    actually rewrote this encoder's weights: layerwise fp8 restores float8 storage, the
+    torchao modes swap in tensor subclasses. quantize_text_encoders returns ONE scheme for
+    the whole bag even when a per-encoder cast failed, so this is how a bench tells which
+    encoders really engaged."""
+    return tuple((type(p).__name__, str(p.dtype)) for p in te.parameters())
+
+
 def measure_te_accuracy(
     family: str,
     *,
@@ -312,6 +321,11 @@ def measure_te_accuracy(
     rows: list[dict] = []
     for scheme in schemes:
         bag = _load_text_encoders(repo, device)
+        dense_sigs = {
+            attr: _te_param_sig(getattr(bag, attr))
+            for attr in _TE_ATTRS
+            if getattr(bag, attr, None) is not None
+        }
         engaged = quantize_text_encoders(bag, _target(), mode = scheme, family = family, logger = logger)
         if engaged is None:
             # The scheme was skipped (unsupported GPU / build, family deny, or a failed kernel
@@ -335,6 +349,24 @@ def measure_te_accuracy(
             continue
         cur = _te_hidden_refs(bag, toks, device)
         for attr, ref_vecs in refs.items():
+            # The caster is best-effort PER encoder (a failure leaves that one dense), yet
+            # the returned scheme covers the whole bag. Scoring a still-dense encoder against
+            # the dense reference would read ~1.0 and falsely certify the scheme for it, so
+            # record any encoder whose weight storage did not change as NOT engaged instead.
+            if _te_param_sig(getattr(bag, attr)) == dense_sigs.get(attr):
+                rows.append(
+                    {
+                        "family": family,
+                        "encoder": attr,
+                        "scheme": engaged,
+                        "cosine": None,
+                        "min_cosine": None,
+                        "relL2": None,
+                        "pass": None,
+                        "engaged": False,
+                    }
+                )
+                continue
             q_vecs = cur.get(attr, [])
             if not q_vecs:
                 continue
@@ -382,11 +414,15 @@ def _encode_once(bag) -> None:
 
 
 # ── VAE loading + latent shape (reused from quant_accuracy_sweep) ──────────────
-def _load_vae(repo: str, device: str):
+def _load_vae(repo: str, device: str, *, force_fp32: bool = False):
     import torch
 
     diffusers = _import_diffusers()
-    vae = diffusers.AutoModel.from_pretrained(repo, subfolder = "vae", torch_dtype = torch.bfloat16)
+    # vae_force_fp32 families (Wan) store AND run the VAE in fp32 in production (the loader
+    # pins a per-component dtype), and quantize_vae stays dense for them regardless. Loading
+    # bf16 here would measure the memory/latency of a truncated VAE production never runs.
+    dtype = torch.float32 if force_fp32 else torch.bfloat16
+    vae = diffusers.AutoModel.from_pretrained(repo, subfolder = "vae", torch_dtype = dtype)
     return vae.to(device).eval()
 
 
@@ -434,7 +470,10 @@ def _make_latent(vae, device: str):
     g = torch.Generator().manual_seed(1234)
     shape = (1, channels, 3, 32, 32) if is_3d else (1, channels, 64, 64)
     z = torch.randn(shape, generator = g, dtype = torch.float32)
-    return z.to(device = device, dtype = torch.bfloat16)
+    # Match the VAE's parameter dtype (fp32 for the vae_force_fp32 families, bf16 otherwise),
+    # like the pipelines do before decode; an fp32 conv rejects a bf16 latent outright.
+    dtype = next(vae.parameters()).dtype
+    return z.to(device = device, dtype = dtype)
 
 
 # ── measurement primitives ────────────────────────────────────────────────────
@@ -524,9 +563,10 @@ def _measure_vae_scheme(
     from core.inference.diffusion_vae_quant import quantize_vae
 
     device = "cuda"
+    force_fp32 = bool(_FAMILIES.get(family, {}).get("vae_force_fp32", False))
     _empty()
     _reset_peak()
-    vae = _load_vae(repo, device)
+    vae = _load_vae(repo, device, force_fp32 = force_fp32)
     z = _make_latent(vae, device)
     _sync()
     mem_dense = _alloc_gb()
@@ -536,7 +576,6 @@ def _measure_vae_scheme(
 
     # quantize_vae reads pipe.vae, so hand it a bag exposing .vae (it mutates that module in place).
     # Pass the family's force_fp32 (Wan) so the real dense-only behaviour is reflected.
-    force_fp32 = bool(_FAMILIES.get(family, {}).get("vae_force_fp32", False))
     engaged = quantize_vae(
         types.SimpleNamespace(vae = vae),
         _target(),

scripts/video_speedmem_bench.py

@@ -385,6 +385,16 @@ def _apply_levers(
             quantize_transformer(v, tgt, mode = cfg["dit"], family = fam_name, logger = logger)
             for v in views
         ]
+        # All-or-none across experts, mirroring the production loader (video.py): the first
+        # expert is mutated in place, so a second-expert miss cannot fall back to dense and
+        # the loader fails that load. Rejecting here keeps the benchmark from publishing a
+        # dit_scheme + timings for a mixed quantized/dense pipeline users cannot actually load.
+        n_engaged = sum(1 for s in schemes if s is not None)
+        if 0 < n_engaged < len(views):
+            raise RuntimeError(
+                f"dit quant '{cfg['dit']}' engaged on only {n_engaged}/{len(views)} experts; "
+                "production rejects this partial state, so the row would be unloadable"
+            )
         engaged["dit"] = schemes[0]
         engaged["dit_experts"] = schemes
         _empty()
@@ -489,12 +499,20 @@ def _timed_video(
     from core.inference.diffusion_cache import maybe_toggle_step_cache, FBCACHE_MIN_STEPS
 
     if cache_mode == "auto":
-        try:
-            maybe_toggle_step_cache(
-                pipe, steps = steps, quant_active = dit_quant_active, threshold = None, logger = logger
-            )
-        except Exception:
-            pass
+        # Toggle on EVERY expert view, exactly like the loader's per-view recheck
+        # (video.py iterates _views_for): toggling only the primary pipe would disable
+        # FBCache on pipe.transformer while transformer_2 stays cached, measuring a
+        # mixed cache state production never runs on a dual-expert MoE.
+        views = [pipe]
+        if getattr(pipe, "transformer_2", None) is not None:
+            views.append(_SecondExpertView(pipe))
+        for v in views:
+            try:
+                maybe_toggle_step_cache(
+                    v, steps = steps, quant_active = dit_quant_active, threshold = None, logger = logger
+                )
+            except Exception:
+                pass
 
     g = torch.Generator(device = "cuda").manual_seed(seed)
     step_ts: list[float] = []
@@ -754,6 +772,7 @@ def main(argv = None) -> int:
                     ref_arrs = [z[k] for k in z.files]
             except Exception:
                 ref_arrs = None
+    row_arrs: list = []
     for n in names:
         row, arrs = _run_config(
             n,
@@ -772,11 +791,20 @@ def main(argv = None) -> int:
             ref_arrs = arrs
         row["lpips_vs_reference"] = _mean_lpips(ref_arrs, arrs) if ref_arrs is not None else None
         rows.append(row)
+        row_arrs.append(arrs)
         print(
             f"  [{n}] {json.dumps({k: row[k] for k in ('dit_scheme','te_scheme','attn','cache','effective_speed','weights_gb','gen_peak_gb','gen_latency_s','per_step_ms','lpips_vs_reference')})}",
             flush = True,
         )
 
+    # --configs accepts an arbitrary order, so rows finalized BEFORE the reference clip was
+    # generated (e.g. --configs shipped,reference) scored lpips_vs_reference as None; rescore
+    # them now that the reference frames exist instead of silently publishing null.
+    if ref_arrs is not None:
+        for r, arrs in zip(rows, row_arrs):
+            if r["lpips_vs_reference"] is None:
+                r["lpips_vs_reference"] = _mean_lpips(ref_arrs, arrs)
+
     # speedups relative to reference (if present)
     ref_lat = next((r["gen_latency_s"] for r in rows if r["config"] == "reference"), None)
     for r in rows:

studio/backend/tests/test_diffusion_attention_trim.py

@@ -12,9 +12,13 @@ from __future__ import annotations
 
 import types
 
-import torch
+import pytest
 
-import core.inference.diffusion_attention as att
+# Skip at collection (not abort) when torch is absent so the rest of the backend suite
+# stays collectable, matching how the policy tests next door gate their heavy imports.
+torch = pytest.importorskip("torch")
+
+import core.inference.diffusion_attention as att  # noqa: E402
 
 
 def test_trim_stream_drops_trailing_padding():