video/image: honor explicit Speed=off for companions + trim, probe explicit TE kernels, bench fidelity

Address the Codex review round on the video/quant work:

- Companion auto-quant now honors an explicit Speed=off. Both loaders already pin the DiT dense
  under an explicit off (bit-exact reference), but the unset text-encoder / VAE quant still promoted
  to auto and silently fp8/int8'd the companions, breaking the bit-exact request. An UNSET speed
  still auto-quantises; an explicit companion scheme still forces it.
- The HunyuanVideo joint-attention trim is a speed lever (it swaps to the fused SDPA kernel), so gate
  it on a non-off speed tier exactly like the adjacent attention-backend selection -- the off path
  keeps the stock dense-mask attention.
- Explicit torchao text-encoder modes (int8 / fp8_dynamic / nvfp4) now run the same kernel smoke
  test the auto ladder uses. They could clear the capability gate yet fail the real GEMM on a build
  where quantize_ wraps the encoder but the kernel is broken; the caster's try/except only covers the
  cast, not the first forward, so the load would report engaged then crash at generation. Now it
  falls back to dense. Layerwise fp8 has no torchao GEMM, so the probe is a no-op for it.
- The trim pre-hook's fallback restores the caller's original kwargs (it may have emptied the image
  stream / trimmed a text stream before failing), so the stock dense-mask path runs on exactly what
  it expects, matching the empty-prompt guard.
- video_speedmem_bench mirrors the loader: installs the Hunyuan trim before the backend set (gated on
  an active tier) and skips the auto int8 quant when it is the fp8-denied memory fallback and dense
  fits resident, so the shipped/auto rows measure what the loader actually runs.

Tests: TE explicit-mode kernel probe (+ layerwise-fp8 bypass), trim mid-trim restore, and loader-level
speed=off companion suppression + trim skip for both backends. 262 backend tests pass; ruff clean.

studio/backend/tests/test_diffusion_attention_trim.py

@@ -154,6 +154,31 @@ def test_trim_pre_hook_never_raises_sets_flag_false():
     assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is False for b in dit.transformer_blocks)
 
 
+def test_trim_pre_hook_restores_inputs_on_midtrim_failure():
+    # A later stream trips the trim AFTER earlier inputs were already mutated (image emptied, mllm
+    # trimmed). The fallback must restore the caller's ORIGINAL kwargs so the stock dense-mask path
+    # (flag False) runs on exactly what it expects -- never a half-trimmed mix.
+    dit = _fake_dit()
+    img = torch.zeros(1, 5, 3)
+    mllm = torch.arange(4.0).reshape(1, 4, 1)
+    mllm_mask = torch.tensor([[1, 1, 0, 0]])
+    byt5 = torch.ones(1, 3, 1)
+    kwargs = {
+        "image_embeds": img,
+        "encoder_hidden_states": mllm,
+        "encoder_attention_mask": mllm_mask,
+        "encoder_hidden_states_2": byt5,
+        "encoder_attention_mask_2": "oops",  # malformed -> _trim_stream raises after mllm is trimmed
+    }
+    _, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
+    assert out["image_embeds"] is img  # emptied then restored
+    assert out["encoder_hidden_states"] is mllm  # trimmed then restored
+    assert out["encoder_attention_mask"] is mllm_mask
+    assert out["encoder_hidden_states_2"] is byt5
+    assert out["encoder_attention_mask_2"] == "oops"
+    assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is False for b in dit.transformer_blocks)
+
+
 def test_trim_pre_hook_absent_stream_not_written_back():
     # If encoder_hidden_states is absent from kwargs (a caller passing it positionally), the hook
     # must NOT write it back as None (that would collide: "got multiple values for argument") and

studio/backend/tests/test_diffusion_backend.py

@@ -2273,6 +2273,31 @@ def test_speed_off_load_suppresses_auto_dtype_quant(fake_runtime, tmp_path, monk
     assert _FakeTransformer.last["path"]  # GGUF from_single_file was used, not a dense build
 
 
+def test_speed_off_load_suppresses_auto_companion_quant(fake_runtime, tmp_path, monkeypatch):
+    # Mirror the DiT suppression for the companions: an explicit Speed="off" load with TE/VAE left at
+    # auto must keep them dense (mode "off"), not promote them to auto-quant and silently fp8/int8 the
+    # text encoder + VAE, which would break the bit-exact request. Unset speed still auto-quantises.
+    from core.inference import diffusion as dmod
+
+    te_modes: list = []
+    vae_modes: list = []
+    monkeypatch.setattr(
+        dmod, "quantize_text_encoders", lambda pipe, target, *, mode, **kw: te_modes.append(mode)
+    )
+    monkeypatch.setattr(
+        dmod, "quantize_vae", lambda pipe, target, *, mode, **kw: vae_modes.append(mode)
+    )
+    (tmp_path / "m.gguf").write_bytes(b"x")
+    backend = DiffusionBackend()
+    backend.load_pipeline(
+        str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "off"
+    )
+    assert te_modes == ["off"] and vae_modes == ["off"]  # dense, not auto
+    backend.unload()
+    backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image")
+    assert te_modes[-1] == "auto" and vae_modes[-1] == "auto"  # promoted when speed is not off
+
+
 def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatch):
     # transformer_quant + a CUDA resident plan -> load the DENSE transformer from the
     # base repo, place it on the device, quantise it, and report the engaged scheme.

studio/backend/tests/test_diffusion_precision.py

@@ -84,6 +84,9 @@ def _stub_casters(monkeypatch, recorder):
     dtq.make_filter_fn = lambda min_features, exclude = (), *, require_bf16 = False: (
         lambda module, fqn = "": True
     )
+    # The explicit-torchao path now runs the same kernel smoke test the auto ladder uses; pass it
+    # by default so these caster tests exercise the cast, not a broken-kernel fallback.
+    dtq._smoke_probe = lambda tq, device: True
     monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq)
 
 
@@ -215,6 +218,7 @@ def test_quantize_int8_uses_family_keep_bf16_schedule(monkeypatch):
     # int8 for a family with a measured schedule routes to the selective caster with
     # that family's (skip_first, skip_last); qwen-image keeps first+last 6 blocks bf16.
     _stub_torch(monkeypatch, cc = (10, 0))
+    monkeypatch.setattr(dp, "_te_scheme_probe", lambda scheme, device: True)
     calls: list = []
     monkeypatch.setattr(
         dp, "_cast_int8_selective", lambda enc, tgt, first, last: calls.append((enc, first, last))
@@ -245,6 +249,7 @@ def test_quantize_fp8_dynamic_uses_compute_caster(monkeypatch):
     # fp8_dynamic routes to the torchao per-row compute caster (not the layerwise one)
     # and needs no per-family schedule.
     _stub_torch(monkeypatch, cc = (9, 0))
+    monkeypatch.setattr(dp, "_te_scheme_probe", lambda scheme, device: True)
     calls: list = []
     monkeypatch.setattr(dp, "_cast_fp8_dynamic", lambda enc, tgt: calls.append(enc))
     te = object()
@@ -254,6 +259,31 @@ def test_quantize_fp8_dynamic_uses_compute_caster(monkeypatch):
     assert calls == [te]
 
 
+def test_quantize_explicit_torchao_probes_kernel(monkeypatch):
+    # An EXPLICIT torchao TE mode (int8 / fp8_dynamic / nvfp4) clears the capability gate but must
+    # still run the real GEMM smoke test the auto ladder uses: on a build where quantize_ wraps the
+    # encoder yet the kernel is broken, report dense (None) instead of crashing on the first forward.
+    _stub_torch(monkeypatch, cc = (10, 0))
+    monkeypatch.setattr(dp, "_te_scheme_probe", lambda scheme, device: False)
+    monkeypatch.setattr(dp, "_cast_fp8_dynamic", lambda *a: pytest.fail("must not cast on probe fail"))
+    monkeypatch.setattr(dp, "_cast_nvfp4", lambda *a: pytest.fail("must not cast on probe fail"))
+    monkeypatch.setattr(dp, "_cast_int8_selective", lambda *a: pytest.fail("must not cast on probe fail"))
+    pipe = types.SimpleNamespace(text_encoder = object())
+    assert quantize_text_encoders(pipe, _target(), mode = "fp8_dynamic") is None
+    assert quantize_text_encoders(pipe, _target(), mode = "nvfp4") is None
+    assert quantize_text_encoders(pipe, _target(), mode = "int8", family = "qwen-image") is None
+
+
+def test_te_scheme_probe_bypasses_layerwise_fp8():
+    # Layerwise fp8 has no torchao GEMM (not in _TE_SMOKE_SCHEME), so the probe is a no-op (True)
+    # for it and never vetoes it -- this is why the explicit-torchao veto above leaves plain fp8
+    # casting untouched. The torchao schemes DO carry a smoke scheme.
+    assert dp._te_scheme_probe(TE_QUANT_FP8, "cuda") is True
+    assert TE_QUANT_FP8 not in dp._TE_SMOKE_SCHEME
+    for scheme in (TE_QUANT_FP8_DYNAMIC, TE_QUANT_INT8, TE_QUANT_NVFP4):
+        assert scheme in dp._TE_SMOKE_SCHEME
+
+
 def test_quantize_int8_unsupported_hw_is_noop(monkeypatch):
     # int8 on pre-Ampere silicon (no int8 tensor cores) applies nothing.
     _stub_torch(monkeypatch, cc = (7, 5))

studio/backend/tests/test_video_backend.py

@@ -1129,6 +1129,56 @@ def test_video_speed_off_suppresses_auto_dtype_quant(fake_runtime, monkeypatch):
     assert calls == [True]
 
 
+def test_video_speed_off_suppresses_auto_companion_quant(fake_runtime, monkeypatch):
+    # Mirror the DiT suppression above for the companions: an explicit Speed="off" (bit-exact) load
+    # with TE/VAE left at auto must NOT promote them to auto-quant -- that would fp8/int8 the text
+    # encoder + VAE and silently break the bit-exact request. Unset speed still auto-quantises.
+    import core.inference.video as video_mod
+
+    te_modes: list = []
+    vae_modes: list = []
+    monkeypatch.setattr(
+        video_mod, "quantize_text_encoders", lambda pipe, target, *, mode, **kw: te_modes.append(mode)
+    )
+    monkeypatch.setattr(
+        video_mod, "quantize_vae", lambda pipe, target, *, mode, **kw: vae_modes.append(mode)
+    )
+    backend = VideoBackend()
+    backend.load_pipeline(
+        "Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline", speed_mode = "off"
+    )
+    assert te_modes == ["off"] and vae_modes == ["off"]  # dense, not auto
+    backend.unload()
+    backend.load_pipeline("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
+    assert te_modes[-1] == "auto" and vae_modes[-1] == "auto"  # promoted when speed is not off
+
+
+def test_video_speed_off_skips_hunyuan_trim(fake_runtime, monkeypatch):
+    # The HunyuanVideo joint-attention trim is a speed lever (swaps to the fused SDPA kernel), so an
+    # explicit Speed="off" (bit-exact reference) keeps the stock dense-mask attention -- like the
+    # attention backend below it, which also honors speed=off. Unset/active speed installs it.
+    import core.inference.video as video_mod
+
+    trim_calls: list = []
+    monkeypatch.setattr(
+        video_mod,
+        "install_hunyuan_attention_trim",
+        lambda view, family, **kw: trim_calls.append(True) or False,
+    )
+    backend = VideoBackend()
+    backend.load_pipeline(
+        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+        model_kind = "pipeline",
+        speed_mode = "off",
+    )
+    assert trim_calls == []  # not installed on the bit-exact path
+    backend.unload()
+    backend.load_pipeline(
+        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", model_kind = "pipeline"
+    )
+    assert trim_calls == [True]  # installed once (single DiT) when speed is active
+
+
 def test_video_step_cache_auto_from_default_schedule(fake_runtime, tmp_path):
     # Unset step cache is AUTO, decided from the model's default schedule: Wan's
     # 50-step default engages FBCache at load; the LTX distilled 8-step default