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.
This commit is contained in:
Daniel Han 2026-07-09 09:24:28 +00:00
commit be04ba00f4
9 changed files with 201 additions and 15 deletions

View file

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