diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index df1ebece71..288ac59944 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -342,9 +342,16 @@ def _apply_levers( via _SecondExpertView, exactly like the loader, so A14B latency + accuracy are real.""" from core.inference.diffusion_precision import quantize_text_encoders from core.inference.diffusion_vae_quant import quantize_vae - from core.inference.diffusion_transformer_quant import quantize_transformer + from core.inference.diffusion_transformer_quant import ( + quantize_transformer, + is_int8_memory_fallback, + ) from core.inference.diffusion_speed import apply_speed_optims, snapshot_backend_flags - from core.inference.diffusion_attention import select_attention_backend, apply_attention_backend + from core.inference.diffusion_attention import ( + select_attention_backend, + apply_attention_backend, + install_hunyuan_attention_trim, + ) from core.inference.diffusion_cache import ( apply_step_cache, TC_FBCACHE, @@ -367,8 +374,15 @@ def _apply_levers( if getattr(pipe, "transformer_2", None) is not None: views.append(_SecondExpertView(pipe)) - # DiT quant (pipeline kind, resident): mutates each expert's transformer in place. - if cfg["dit"] not in ("none", "off"): + # DiT quant (pipeline kind, resident): mutates each expert's transformer in place. Mirror the + # loader's dense-fit skip: for an AUTO request on an int8-fallback family (HunyuanVideo-1.5, where + # fp8 is black-framed so auto lands on int8, a memory-only lever ~7% slower AND less accurate than + # dense+compile), run the dense DiT instead when it fits resident -- and the benchmark always + # loads resident (no offload). Explicit int8/fp8 configs are honored (the whole point of the + # sweep). Without this the "shipped"/"ditquant" auto rows would measure int8 where the loader runs + # dense, overstating the shipped cost on Hunyuan. + dense_fit_skip = cfg["dit"] == "auto" and is_int8_memory_fallback(tgt, fam_name) + if cfg["dit"] not in ("none", "off") and not dense_fit_skip: schemes = [ quantize_transformer(v, tgt, mode = cfg["dit"], family = fam_name, logger = logger) for v in views @@ -423,6 +437,16 @@ def _apply_levers( ) cache_active = engaged["cache"] not in (None, "off") + # HunyuanVideo-1.5 joint-attention trim (per expert), BEFORE the backend set so the requested + # kernel pins onto the new processors -- exactly the loader's order. Drops the ~99% zero-padded + # text tokens so the fused SDPA kernel runs (~18x/DiT-forward, cosine ~1.0). A speed lever, so + # gated on an active tier like the loader; no-op for every non-Hunyuan family. + trim_engaged = False + if speed_active: + for v in views: + trim_engaged = install_hunyuan_attention_trim(v, fam_obj, logger = logger) or trim_engaged + engaged["attn_trim"] = trim_engaged + # Attention (per expert). backend = select_attention_backend(tgt, cfg["attn"], speed_active = speed_active) for v in views: diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index b698dd71f4..b922a729db 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1217,17 +1217,23 @@ class DiffusionBackend: normalize_transformer_cache(transformer_cache) normalize_te_quant(text_encoder_quant) normalize_vae_quant(vae_quant) + # An explicit Speed="off" (bit-exact reference) load pins the companions dense too, mirroring + # the transformer_quant default below (load_pipeline): promoting an UNSET TE/VAE to auto-quant + # here would silently fp8/int8 the text encoder + VAE and break the bit-exact request -- an + # auto DEFAULT overriding the EXPLICIT off control. Only an EXPLICIT off suppresses; an unset + # speed still auto-quantises, and an explicit companion scheme still forces it. + speed_off = speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF # text_encoder_quant tri-state, mirroring transformer_quant: UNSET (None / "") -> auto, # which picks the best accurate TE scheme for this GPU + family (fp8_dynamic / int8 / # layerwise fp8) or stays dense when none qualifies. An explicit "none"/"off" pins the - # encoder dense; an explicit scheme forces it. So the shipped default is auto. + # encoder dense; an explicit scheme forces it. So the shipped default is auto (dense under off). if text_encoder_quant is None or str(text_encoder_quant).strip() == "": - text_encoder_quant = TE_QUANT_AUTO + text_encoder_quant = "off" if speed_off else TE_QUANT_AUTO # vae_quant tri-state, same contract: UNSET -> auto (fp8_dynamic conv compute on resident # fp8-GEMM silicon that passes the conv probe, else layerwise fp8, else dense); none/off -> - # dense; an explicit scheme forces it. + # dense; an explicit scheme forces it. Also pinned dense under an explicit Speed="off". if vae_quant is None or str(vae_quant).strip() == "": - vae_quant = VAE_QUANT_AUTO + vae_quant = "off" if speed_off else VAE_QUANT_AUTO # For a full pipeline the repo itself supplies every component, so it is its # own base; the single-file kinds resolve the companion base diffusers repo. base = ( diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index e70a542e5d..613850b9d9 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -602,6 +602,11 @@ def _hunyuan_trim_pre_hook(module, args, kwargs): return args, kwargs except Exception: # noqa: BLE001 — optimisation only; never break the forward + # We may have trimmed some kwargs (image_embeds / a text stream) before failing. Restore the + # caller's untrimmed inputs so the stock dense-mask path (flag False below) runs on exactly + # what it expects -- the same restore the empty-prompt guard above does. + kwargs.clear() + kwargs.update(original) for blk in getattr(module, "transformer_blocks", []): attn = getattr(blk, "attn", None) if attn is not None: diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 460824ed4b..a44de06fb5 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -253,6 +253,16 @@ def quantize_text_encoders( return None if not te_quant_supported(target, mode): return None + # Mirror select_te_quant_scheme's auto path: an EXPLICIT torchao TE mode (int8 / fp8_dynamic / + # nvfp4) can clear the capability gate above yet fail the real GEMM on a torchao/torch build + # where quantize_ wraps the encoder but the kernel is broken -- and the caster's try/except only + # catches the cast, not the first prompt-encoder forward. Run the same kernel smoke test the auto + # ladder uses so a failing kernel falls back to dense here instead of crashing at generation. + # Layerwise fp8 (no torchao GEMM) has no smoke scheme, so the probe is a no-op (True) for it. + device = str(getattr(target, "device", "cuda")) + if not _te_scheme_probe(mode, device): + _note(logger, f"text-encoder '{mode}' failed the kernel smoke test; staying dense") + return None if mode == TE_QUANT_INT8: first, last = skip # type: ignore[misc] diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 9febd1b62e..9952d633c7 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -921,15 +921,21 @@ class VideoBackend: vae_quant = vae_quant, ) kind = resolve_video_model_kind(gguf_filename, model_kind) + # An explicit Speed="off" (bit-exact reference) load pins the companions dense too, mirroring + # the transformer_quant default below: promoting an UNSET TE/VAE to auto-quant would silently + # fp8/int8 the text encoder + VAE and break the bit-exact request (an auto DEFAULT overriding + # the EXPLICIT off control). Only an EXPLICIT off suppresses; an unset speed still auto + # -quantises, and an explicit companion scheme still forces it. + speed_off = speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF # text_encoder_quant tri-state (mirrors the image backend + transformer_quant): UNSET # (None / "") -> auto (pick the best accurate TE scheme for this GPU + family); an explicit # "none"/"off" pins the encoder dense; a scheme forces it. So the shipped default is auto. if text_encoder_quant is None or str(text_encoder_quant).strip() == "": - text_encoder_quant = TE_QUANT_AUTO + text_encoder_quant = "off" if speed_off else TE_QUANT_AUTO # vae_quant tri-state, same contract. The vae_force_fp32 families (Wan) keep the VAE dense # regardless (quantize_vae's force_fp32 gate), so auto is safe as the shipped default. if vae_quant is None or str(vae_quant).strip() == "": - vae_quant = VAE_QUANT_AUTO + vae_quant = "off" if speed_off else VAE_QUANT_AUTO base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo) with self._lock: @@ -985,10 +991,9 @@ class VideoBackend: # precision to auto-quant here would engage int8/fp8 + regional compile and silently # break the user's bit-exact request (an auto DEFAULT overriding an EXPLICIT control), # and the quant path below would then also force effective_speed back to default. - # Suppress the auto default when speed was explicitly pinned off, mirroring the image - # backend (diffusion.py); otherwise auto (the dense-capable default) applies. "off" - # normalizes to None (no dense quant), keeping the dense bf16 path. - speed_off = speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF + # Suppress the auto default when speed was explicitly pinned off (speed_off, computed + # above with the companions), mirroring the image backend (diffusion.py); otherwise auto + # (the dense-capable default) applies. "off" normalizes to None (no dense quant). transformer_quant = "off" if speed_off else TQ_AUTO # ── memory plan: family-table resident estimate + frames-aware headroom. @@ -1353,7 +1358,13 @@ class VideoBackend: # 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. - trim = install_hunyuan_attention_trim(view, fam, logger = logger) + # 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( diff --git a/studio/backend/tests/test_diffusion_attention_trim.py b/studio/backend/tests/test_diffusion_attention_trim.py index 6aae6bdedb..03142d8889 100644 --- a/studio/backend/tests/test_diffusion_attention_trim.py +++ b/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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 2d18ba0772..897a0b3c86 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/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. diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 2bba13fea8..548894fa6f 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/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)) diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index c6928ba936..7008834dc6 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/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