From 8ecfd30923b1e8d0a46cc6eb56c213cc795d92e1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 25 Sep 2026 16:22:17 -0700 Subject: [PATCH] Studio: whole-model NVFP4 for the video families with hosted pre-quantized denoisers (#10729) * studio: whole-model NVFP4 for the video families with hosted pre-quantized denoisers Wan2.2-TI2V-5B, Wan2.2-T2V-A14B (both experts), HunyuanVideo-1.5 480p and 720p run NVFP4 on the torchao backend from hosted pre-quantized denoisers, so a load never quantises 5 to 28 GB of dense weights on the fly. Every admitted linear (min 512 features, in and out divisible by 16) is NVFP4; nothing is per-layer here. - one shared auto walk with a per-family head in front of it, and divisible_for_scheme - ZeroRowSafeLinear, a live bug: torchao's nvfp4 activation path raises on HunyuanVideo-1.5's empty image stream on every text-to-video render - loader gates on the packed-weight fingerprint, the component and require_divisible - conventional video pre-quant seeding, all components or none, with the memory plan taking the artifact's resident size rather than the dense build peak - registry rows for the four hosted NVFP4 repos nvfp4 stays an explicit opt-in here; no family gets an auto row on this branch. The tooling that builds and fingerprints these artifacts is not in this repo. It lives in the private scripts repo with the rest of the quant generation tooling. * Tighten the comments and docstrings on the video prequant path Collapse the long single-line comments this branch added to the repo's wrap width, shorten the prose docstrings to the parts that say why, and drop the test docstrings that only restate the test name. Comments, docstrings and whitespace only: no code changes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Credit the resident video pipeline's memory in the seed plan * Pin the residency check in the two-expert download plan test * Credit only the target card's resident bytes and fingerprint built checkpoints The residency plan counted every accelerator storage of the old pipeline as reclaimable on the card the new load targets, so switching GPUs could skip the dense denoiser shards on a card that never held them. Count only the target's. The builder never wrote metadata["fingerprint"], so the loader's packed-weight check had nothing to compare against for artifacts it produced. * Fingerprint torchao 0.18 Int8Tensor prequant weights torchao >= 0.18 int8 quantization produces Int8Tensor, which the packed fingerprint payload map did not list, so every int8 weight was skipped and the recorded fingerprint verified nothing. Hash its qdata, scale and optional zero point / activation slots, with a stub test and a real torchao round trip that catches a flipped byte. * Apply the GEMM alignment floor in the offline prequant builder and record it The runtime quantize_ skips Linears whose in/out features miss the scheme's tiling floor, but build_prequant_checkpoint.py did not, so a ragged layer could be baked as a quantized GEMM the runtime leaves dense. The builder now passes divisible_for_scheme(scheme) to make_filter_fn and stamps require_divisible so the loader can verify it. * Trim comments * Studio: NVFP4 kill switch for image and video, default off The hosted unsloth/*-NVFP4 checkpoints are still private, so NVFP4 is disabled in general until UNSLOTH_NVFP4_DIFFUSION is set to a truthy value (1, true, yes, on). The default lives in one constant, NVFP4_DIFFUSION_DEFAULT in core/inference/diffusion_nvfp4_flag.py, so turning it on for everyone later is a one-line change. With the switch on, behaviour is unchanged. While it is off: - auto never picks nvfp4 for the dense transformer (ladder and family prefer rows), and no smoke probe is spent on it; - an explicit transformer_quant or text_encoder_quant of nvfp4 gets a 400 "NVFP4 is disabled in this build" from every image and video load and download-plan route, ahead of the precision fallback; - no hosted NVFP4 denoiser or encoder is resolved, seeded, planned or sized, and the Hub helpers never ask about a *-NVFP4 repo; - /images/info estimates and the /api/system dense ladder omit nvfp4, and /api/system reports nvfp4_diffusion so the Precision and Text encoder selects hide NVFP4. The offline builder does not consult the switch. Existing NVFP4 test modules run with it on through a conftest fixture; the default-off behaviour, and that the enabled path is still reachable end to end, are covered in test_nvfp4_diffusion_flag.py. * Studio: keep the diffusion bench NVFP4-capable under the kill switch An explicit --transformer-quant or --text-encoder-quant nvfp4 turns the switch on for the bench itself, like the other developer tools. The switch tests ask auto the way a load does, with a hosted checkpoint where the selector requires one, so the enabled-path check stays meaningful on the branches that add that requirement. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: reset a hidden NVFP4 pick to Auto when the backend switch is off A Precision or Text encoder select reseeded to nvfp4 by an earlier load kept that value after a restart turned UNSLOTH_NVFP4_DIFFUSION off: the option was filtered out, so the select rendered blank and the next load or Reapply sent nvfp4 and got a 400. Once /api/system reports the switch off, a held nvfp4 now falls back to auto on the image and video pages. Nothing is reset before system info arrives, since the switch reads off by default until then. * Studio: keep the hosted video denoiser seed for an explicit scheme under Speed off Only an auto precision is rewritten to off under Speed off; an explicit scheme is still quantized (and upgrades the speed), so suppressing the seed staged the dense denoiser for a runtime quantize. Matches diffusion.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: update the Speed off video plan test for the kept explicit seed * Studio: refuse an NVFP4 checkpoint loaded directly as the model while NVFP4 is off With UNSLOTH_NVFP4_DIFFUSION off, a hosted NVFP4 repo (or any repo or local path whose pre-quant metadata records scheme nvfp4) picked as the model itself named no precision, so none of the scheme gates saw it. The image and video load and download-plan routes now 400 such a pick with the same disabled message style. Detection reads the recorded scheme of a local file, a local directory or a cached Hub snapshot, then the family NVFP4 repo tables, then the -NVFP4 name. Switched on, the pick proceeds as before. * Studio: format the NVFP4 checkpoint refusal tests * Studio: skip pipeline roots and the thread hop in the NVFP4 checkpoint refusal A diffusers pipeline loads from its component folders, so a stray root artifact no longer makes its directory read as an NVFP4 checkpoint. With the switch on the route helper returns before handing anything to a thread. * Studio: stamp the built denoiser component and drop an unused metadata reader * Studio: keep the NVFP4 flag and speed-off video plan tests off the runner's own GPU * Studio: trim comments in the NVFP4 video layer * Studio: check the hosted video seed rows before spending a quant smoke probe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: let the prequant builder target the second Wan A14B expert Add --component (default transformer) for the subfolder, metadata and upload name, and resolve video families so the Wan and HunyuanVideo denoisers can be built. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- scripts/build_prequant_checkpoint.py | 78 +- scripts/diffusion_bench.py | 2 + scripts/verify_prequant_backend.py | 1 + studio/backend/core/inference/diffusion.py | 12 + .../core/inference/diffusion_auto_policy.py | 4 + .../core/inference/diffusion_families.py | 4 + .../inference/diffusion_inference_info.py | 3 + .../core/inference/diffusion_nvfp4_flag.py | 98 +++ .../core/inference/diffusion_precision.py | 8 +- .../core/inference/diffusion_prequant.py | 407 +++++++++- .../core/inference/diffusion_quant_pad.py | 74 ++ .../core/inference/diffusion_te_prequant.py | 4 + .../inference/diffusion_transformer_quant.py | 143 +++- studio/backend/core/inference/video.py | 713 ++++++++++++++---- .../core/inference/video_denoiser_prequant.py | 140 ++++ .../backend/core/inference/video_families.py | 44 +- studio/backend/main.py | 11 + studio/backend/routes/inference.py | 46 ++ studio/backend/routes/video.py | 14 + studio/backend/tests/conftest.py | 40 + .../tests/test_build_prequant_checkpoint.py | 115 +++ .../backend/tests/test_diffusion_prequant.py | 320 +++++++- .../backend/tests/test_diffusion_quant_pad.py | 113 +++ .../tests/test_diffusion_transformer_quant.py | 253 +++++++ .../tests/test_nvfp4_diffusion_flag.py | 621 +++++++++++++++ studio/backend/tests/test_video_backend.py | 566 ++++++++++++++ studio/backend/tests/test_video_families.py | 103 +++ studio/backend/tests/test_video_prequant.py | 494 +++++++++++- .../src/features/images/images-page.tsx | 51 +- .../src/features/video/video-page.tsx | 28 +- studio/frontend/src/hooks/use-gpu-info.ts | 4 + studio/frontend/src/hooks/use-host-class.ts | 9 + studio/frontend/src/hooks/use-system.ts | 2 + studio/frontend/src/lib/nvfp4-options.ts | 20 + .../dense-quant-schemes-identity.test.ts | 1 + studio/frontend/tests/nvfp4-options.test.ts | 58 ++ 36 files changed, 4366 insertions(+), 238 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_nvfp4_flag.py create mode 100644 studio/backend/core/inference/video_denoiser_prequant.py create mode 100644 studio/backend/tests/test_nvfp4_diffusion_flag.py create mode 100644 studio/frontend/src/lib/nvfp4-options.ts create mode 100644 studio/frontend/tests/nvfp4-options.test.ts diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index b0e203b91a..bbdb1cc712 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -28,6 +28,17 @@ from typing import Any, Optional, Sequence BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend" +# Mirrors core.inference.diffusion_prequant.DEFAULT_PREQUANT_COMPONENT; the backend joins sys.path only in main(). +DEFAULT_COMPONENT = "transformer" + + +def resolve_build_family(base: str, override: Optional[str] = None) -> Optional[Any]: + """The image family for ``base`` / ``override``, else the video family (Wan, HunyuanVideo), else None.""" + from core.inference.diffusion_families import detect_family + from core.inference.video_families import detect_video_family + + return detect_family(base, override = override) or detect_video_family(base, override = override) + def convrot_refusal( group: int, rotatable: Sequence[str], not_divisible: Sequence[str] @@ -55,6 +66,7 @@ def upload_destination( safetensors: bool = False, override: Optional[str] = None, upload_repo: Optional[str] = None, + component: str = DEFAULT_COMPONENT, ) -> str: """The repo-root filename this build should publish under. @@ -87,6 +99,26 @@ def upload_destination( "would publish an artifact nothing can open. Rename the upload, or change --out." ) return override + if component and component != DEFAULT_COMPONENT: + # A second denoiser (Wan A14B's transformer_2) resolves ONLY its task-specific row, with no + # fallback, so any other name would publish an artifact the loader never asks for. + from core.inference.diffusion_families import family_prequant_filename + + specific = family_prequant_filename(fam, scheme, task = component) + if specific is None or specific == family_prequant_filename(fam, scheme): + raise ValueError( + f"family {getattr(fam, 'name', fam)!r} declares no prequant_filenames entry for " + f"({scheme!r}, {component!r}), so the loader would never ask for this component. " + "Add the entry to the family table, or pass --upload-filename." + ) + wanted = ".safetensors" if safetensors else ".pt" + if not specific.lower().endswith(wanted): + raise ValueError( + f"family {getattr(fam, 'name', fam)!r} declares {specific!r} for " + f"({scheme!r}, {component!r}), which does not end in {wanted!r} like --out. " + "Rename --out, or pass --upload-filename." + ) + return specific from core.inference.diffusion_prequant import prequant_filename if not rotated and not safetensors: @@ -129,6 +161,20 @@ def upload_destination( return preferred +def quant_filter_settings(scheme: str, family: Optional[str]) -> dict: + """Runtime ``quantize_`` filter inputs for ``scheme``; one dict feeds filter and metadata.""" + from core.inference.diffusion_transformer_quant import ( + _REQUIRE_BF16_SCHEMES, + divisible_for_scheme, + exclude_tokens_for_scheme, + ) + return { + "exclude_name_tokens": list(exclude_tokens_for_scheme(scheme, family)), + "require_bf16": scheme in _REQUIRE_BF16_SCHEMES, + "require_divisible": divisible_for_scheme(scheme), + } + + def main(argv = None) -> int: p = argparse.ArgumentParser() p.add_argument( @@ -142,6 +188,12 @@ def main(argv = None) -> int: "directory name differs from the Hub repo. Must still name this family's base model.", ) p.add_argument("--scheme", required = True, help = "quant scheme: int8 | fp8 | nvfp4 | mxfp8") + p.add_argument( + "--component", + default = DEFAULT_COMPONENT, + help = "denoiser subfolder to quantise, recorded in the checkpoint (e.g. transformer_2 for " + "Wan2.2 A14B's second expert)", + ) p.add_argument( "--out", required = True, @@ -179,7 +231,6 @@ def main(argv = None) -> int: import torchao import diffusers - from core.inference.diffusion_families import detect_family from core.inference.diffusion_prequant import prequant_format_for # Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant). @@ -187,10 +238,8 @@ def main(argv = None) -> int: FP8_GRANULARITY, TQ_FP8, TQ_SCHEMES, - _REQUIRE_BF16_SCHEMES, _make_quant_config, _resolve_fast_accum, - exclude_tokens_for_scheme, make_filter_fn, ) from torchao.quantization import quantize_ @@ -199,7 +248,8 @@ def main(argv = None) -> int: if scheme not in TQ_SCHEMES: print(f"error: --scheme must be one of {TQ_SCHEMES} (not 'auto')", flush = True) return 2 - fam = detect_family(args.base, override = args.family) + component = (args.component or "").strip() or DEFAULT_COMPONENT + fam = resolve_build_family(args.base, override = args.family) if fam is None: print(f"error: unknown family '{args.family}'", flush = True) return 2 @@ -268,29 +318,29 @@ def main(argv = None) -> int: safetensors = is_safetensors_out, override = args.upload_filename, upload_repo = args.upload_repo, + component = component, ) except ValueError as exc: print(f"error: {exc}", flush = True) return 2 print(f"== build prequant ({fam.name}/{scheme}, min_feat={args.min_features}) ==", flush = True) - print(f" loading dense transformer from {args.base} (subfolder=transformer) ...", flush = True) + print(f" loading dense transformer from {args.base} (subfolder={component}) ...", flush = True) t0 = time.time() transformer = transformer_cls.from_pretrained( - args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token + args.base, subfolder = component, torch_dtype = torch.bfloat16, token = args.hf_token ).to("cuda") print(f" quantising in place ({scheme}) ...", flush = True) - # Mirror the runtime exclusions: int8 skips the M=1 modulation projections (torch._int_mm needs M>16) plus - # per-family ones; family=None bakes linears the runtime rejects. - exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name) - # fp8 / mxfp8 need bf16 weights, so skip non-bf16 Linears; nvfp4 handles fp32. Mirrors the runtime gate. - require_bf16 = scheme in _REQUIRE_BF16_SCHEMES + filter_settings = quant_filter_settings(scheme, fam.name) + exclude_name_tokens = tuple(filter_settings["exclude_name_tokens"]) + require_bf16 = filter_settings["require_bf16"] # fp8 bakes the accumulate mode in; record it so the loader can reject a contradicting request. fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None filter_fn = make_filter_fn( args.min_features, exclude_name_tokens = exclude_name_tokens, require_bf16 = require_bf16, + require_divisible = filter_settings["require_divisible"], ) # ConvRot, BEFORE quantize_: rotating the weights is only worth anything if the quantizer then sees the rotated @@ -332,14 +382,20 @@ def main(argv = None) -> int: # Let the loader reject a checkpoint that would not match the runtime path. "exclude_name_tokens": list(exclude_name_tokens), "require_bf16": require_bf16, + "require_divisible": filter_settings["require_divisible"], "fast_accum": fast_accum, "torch_dtype": args.dtype, "quant_backend": "torchao", "transformer_class": fam.transformer_class, + # The subfolder built above: the loader refuses it as another denoiser (e.g. transformer_2). + "component": component, "torch_version": torch.__version__, "torchao_version": getattr(torchao, "__version__", "?"), "diffusers_version": diffusers.__version__, } + from core.inference.diffusion_prequant import packed_weight_fingerprint + + metadata["fingerprint"] = packed_weight_fingerprint(state_dict) # fp8 granularity: lets the loader reject a stale per-tensor checkpoint (runtime needs per-row). if scheme == TQ_FP8: metadata["fp8_granularity"] = FP8_GRANULARITY diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 73daa1cc35..6d8a3f7498 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -598,6 +598,8 @@ def _build_parser() -> argparse.ArgumentParser: def main(argv: Optional[list[str]] = None) -> int: args = _build_parser().parse_args(argv) + if "nvfp4" in (args.transformer_quant, args.text_encoder_quant): + os.environ["UNSLOTH_NVFP4_DIFFUSION"] = "1" if isinstance(args.gguf, str): args.gguf = args.gguf.strip() or None if bool(args.write_baseline) == bool(args.compare): diff --git a/scripts/verify_prequant_backend.py b/scripts/verify_prequant_backend.py index 9ea2f7a66c..17d9b60526 100644 --- a/scripts/verify_prequant_backend.py +++ b/scripts/verify_prequant_backend.py @@ -84,6 +84,7 @@ def _lpips(ref, arr): def run(mode, steps, seed, res): sys.path.insert(0, str(BACKEND)) + os.environ["UNSLOTH_NVFP4_DIFFUSION"] = "1" import torch import diffusers from core.inference.diffusion_prequant import ( diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 57b33faa4f..3a78f9e05b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1877,7 +1877,15 @@ class DiffusionBackend: ``begin_load`` runs inside ``acquire_for``, which evicts chat under the arbiter lock before the register callback, and after ``select_and_activate_engine``, which unloads the resident model on an engine switch. + + NVFP4 while the NVFP4 switch is off raises ``ValueError`` (the route's 400) first, ahead of + the opt-in fallback: a disabled scheme is refused, never swapped. """ + from .diffusion_nvfp4_flag import refuse_disabled_nvfp4 + + refuse_disabled_nvfp4( + transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant + ) if precision_fallback_allowed(): return pinned = normalize_transformer_quant(transformer_quant) @@ -3405,6 +3413,10 @@ class DiffusionBackend: candidate answer, and skipping them all reports the miss, which keeps the shards.""" if source is None or getattr(source, "kind", None) != "repo": return None + from .diffusion_nvfp4_flag import nvfp4_blocked, nvfp4_repo_blocked + + if nvfp4_blocked(scheme) or nvfp4_repo_blocked(source.location): + return None from huggingface_hub import HfApi info = HfApi(token = hf_token or None).model_info(source.location, files_metadata = True) diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index afc232cb2a..63e5c99456 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -231,6 +231,10 @@ def estimate_dense_quant( ) -> Optional[DenseQuantEstimate]: """Estimate the candidate's footprint from the family table, or None when the family (or scheme factor) is unknown.""" + from .diffusion_nvfp4_flag import nvfp4_blocked + + if nvfp4_blocked(scheme): + return None components = family_bf16_components_gb(fam, base_repo) factor = _QUANT_STEADY_FACTOR.get(scheme) if components is None or factor is None: diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 9b68946f52..5713df6ee7 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -20,6 +20,8 @@ from pathlib import Path, PurePosixPath from typing import NamedTuple, Optional, Sequence from utils.paths.path_utils import is_appledouble_metadata +from .diffusion_nvfp4_flag import nvfp4_blocked + # Runtime->route contract: the /images/generate route matches these messages EXACTLY for a 409 (vs a 500), so both # engines raise them verbatim. @@ -1183,6 +1185,8 @@ def family_prequant_repo( close enough that planning around it costs nothing, since the base_model_id validation refuses the artifact well after the plan was made. A base whose weights really differ belongs in ``prequant_excluded_bases``, which returns None here instead.""" + if nvfp4_blocked(scheme): + return None # Both tables are keyed on lowercased upstream ids. base = canonical_base(base_repo).lower() if base: diff --git a/studio/backend/core/inference/diffusion_inference_info.py b/studio/backend/core/inference/diffusion_inference_info.py index ade971c1b0..56b92c5707 100644 --- a/studio/backend/core/inference/diffusion_inference_info.py +++ b/studio/backend/core/inference/diffusion_inference_info.py @@ -15,6 +15,7 @@ from __future__ import annotations from typing import Any from .diffusion_auto_policy import _FAMILY_BF16_GB, _QUANT_STEADY_FACTOR +from .diffusion_nvfp4_flag import nvfp4_blocked def _round1(value: float) -> float: @@ -33,6 +34,8 @@ def family_inference_infos() -> list[dict[str, Any]]: companions_gb = text_encoders_gb + vae_gb estimated = {"bf16": _round1(transformer_gb + companions_gb)} for scheme, factor in _QUANT_STEADY_FACTOR.items(): + if nvfp4_blocked(scheme): + continue estimated[scheme] = _round1(transformer_gb * factor + companions_gb) infos.append( { diff --git a/studio/backend/core/inference/diffusion_nvfp4_flag.py b/studio/backend/core/inference/diffusion_nvfp4_flag.py new file mode 100644 index 0000000000..519111e4cd --- /dev/null +++ b/studio/backend/core/inference/diffusion_nvfp4_flag.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The one NVFP4 switch for Studio image/video; read at call time, never cached, so tests can monkeypatch it.""" + +from __future__ import annotations + +import os +from typing import Iterable, Optional + +NVFP4_DIFFUSION_ENV = "UNSLOTH_NVFP4_DIFFUSION" + +# Flip to True once the hosted *-NVFP4 repos are public. +NVFP4_DIFFUSION_DEFAULT = False + +NVFP4_SCHEME = "nvfp4" + +_TRUTHY = frozenset({"1", "true", "yes", "on"}) +_FALSY = frozenset({"0", "false", "no", "off"}) + + +def nvfp4_diffusion_enabled() -> bool: + """Whether NVFP4 may be used; unrecognised values take ``NVFP4_DIFFUSION_DEFAULT``.""" + raw = os.environ.get(NVFP4_DIFFUSION_ENV) + if raw is None: + return NVFP4_DIFFUSION_DEFAULT + value = raw.strip().lower() + if value in _TRUTHY: + return True + if value in _FALSY: + return False + return NVFP4_DIFFUSION_DEFAULT + + +def is_nvfp4(value: Optional[str]) -> bool: + """Whether ``value`` names the NVFP4 scheme, in any case or spacing.""" + if value is None: + return False + return str(value).strip().lower().replace("-", "_") == NVFP4_SCHEME + + +def nvfp4_blocked(scheme: Optional[str]) -> bool: + """True when ``scheme`` is NVFP4 and the switch is off: the one test every gate asks.""" + return is_nvfp4(scheme) and not nvfp4_diffusion_enabled() + + +def nvfp4_repo_blocked(repo_id: Optional[str]) -> bool: + """True for a hosted ``*-NVFP4`` repo while the switch is off (Hub-helper backstop).""" + if nvfp4_diffusion_enabled() or not repo_id: + return False + return str(repo_id).strip().rstrip("/").lower().endswith("-nvfp4") + + +def without_nvfp4(schemes: Iterable[str]) -> tuple[str, ...]: + """``schemes`` in order, minus NVFP4 while the switch is off.""" + if nvfp4_diffusion_enabled(): + return tuple(schemes) + return tuple(scheme for scheme in schemes if not is_nvfp4(scheme)) + + +def nvfp4_disabled_message(control: str = "transformer_quant") -> str: + """The refusal a request naming NVFP4 gets while the switch is off.""" + return ( + f"{control}='nvfp4' could not be used: NVFP4 is disabled in this build. Pick another " + f"precision, or set {NVFP4_DIFFUSION_ENV}=1 on the Studio backend to enable it." + ) + + +def refuse_disabled_nvfp4( + *, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None +) -> None: + """Raise ``ValueError`` (400) for NVFP4 while off, before any silent fallback could swap it.""" + if nvfp4_diffusion_enabled(): + return + if is_nvfp4(transformer_quant): + raise ValueError(nvfp4_disabled_message("transformer_quant")) + if is_nvfp4(text_encoder_quant): + raise ValueError(nvfp4_disabled_message("text_encoder_quant")) + + +def nvfp4_checkpoint_disabled_message() -> str: + """The refusal a load or plan of an NVFP4 checkpoint gets while the switch is off.""" + return ( + "model_path could not be used: it is an NVFP4 checkpoint and NVFP4 is disabled in this " + f"build. Pick another model, or set {NVFP4_DIFFUSION_ENV}=1 on the Studio backend to " + "enable it." + ) + + +def refuse_disabled_nvfp4_checkpoint(*model_paths: Optional[str]) -> None: + """Raise ``ValueError`` (400) for an NVFP4 checkpoint model path while off; it bypasses the scheme gates.""" + if nvfp4_diffusion_enabled(): + return + from .diffusion_prequant import declares_nvfp4_checkpoint + + for path in model_paths: + if path and declares_nvfp4_checkpoint(path): + raise ValueError(nvfp4_checkpoint_disabled_message()) diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index d017be23f5..2c5f909441 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -38,6 +38,8 @@ from functools import lru_cache from core._torchao_stub import is_stubbed, torch_is_rocm +from .diffusion_nvfp4_flag import nvfp4_blocked, nvfp4_disabled_message + TE_QUANT_FP8 = "fp8" TE_QUANT_NVFP4 = "nvfp4" TE_QUANT_INT8 = "int8" @@ -79,6 +81,8 @@ def normalize_te_quant(value: Optional[str]) -> Optional[str]: raise ValueError( f"Unsupported text_encoder_quant '{value}'. Use one of: {', '.join(TE_QUANT_MODES)}." ) + if nvfp4_blocked(normalized): + raise ValueError(nvfp4_disabled_message("text_encoder_quant")) return normalized @@ -110,7 +114,7 @@ def resolve_te_quant_request( """ if not te_quant_is_auto(value): return normalize_te_quant(value), False - if auto_scheme is None: + if auto_scheme is None or nvfp4_blocked(auto_scheme): return None, False # Validate the family's own field rather than trusting it: a typo here would otherwise reach # quantize_text_encoders as an unknown mode on every default load of that family. @@ -166,6 +170,8 @@ def te_quant_supported(target: Any, mode: str) -> bool: Blackwell sm_100+ (nvfp4).""" if getattr(target, "device", None) != "cuda": return False + if nvfp4_blocked(mode): + return False # Torchao modes cannot use the Windows stub or ROCm's non-SM capability values. Plain fp8 is only a dtype cast and # remains supported. if mode in _TE_TORCHAO_MODES and (is_stubbed("torchao") or torch_is_rocm()): diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index ce44c4dcf0..2c61ce2f7a 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -34,6 +34,8 @@ from dataclasses import dataclass from collections.abc import Sequence from typing import Any, Optional +from .diffusion_nvfp4_flag import nvfp4_blocked + # torch.save dict layout tag; bump on an on-disk change so old/foreign artifacts are rejected PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1" @@ -47,6 +49,8 @@ PREQUANT_FORMAT_ROTATED = "unsloth_prequant_transformer_state_dict_v2" PREQUANT_FORMATS = (PREQUANT_FORMAT, PREQUANT_FORMAT_ROTATED) +DEFAULT_PREQUANT_COMPONENT = "transformer" + def prequant_format_for(metadata: Any) -> str: """The on-disk format tag an offline builder should stamp for ``metadata``.""" @@ -432,23 +436,31 @@ def prequant_repo_filename( repo_id: str, scheme: str, suffix: str = ".pt", + *, + component: Optional[str] = None, ) -> str: - """The model-name checkpoint filename for ``scheme`` in ``repo_id``: the hosted repos are named - -FP8 (or -INT8 / -quantized) and carry -.pt files, e.g. - unsloth/Z-Image-Turbo-FP8 -> Z-Image-Turbo-INT8.pt / Z-Image-Turbo-FP8.pt. + """Filename for ``scheme`` in ``repo_id``; a non-default ``component`` inserts --. ``suffix`` picks the container. It defaults to ``.pt`` so every existing caller keeps naming the artifact it already names; ``derived_prequant_filenames`` is what puts the safetensors spelling of the same name ahead of it.""" model = repo_id.rsplit("/", 1)[-1] - for drop in ("-fp8", "-int8", "-quantized"): + for drop in ("-fp8", "-int8", "-nvfp4", "-mxfp8", "-quantized"): if model.lower().endswith(drop): model = model[: -len(drop)] break + part = (component or "").strip() + if part and part != DEFAULT_PREQUANT_COMPONENT: + return f"{model}-{part}-{scheme.upper()}{suffix}" return f"{model}-{scheme.upper()}{suffix}" -def derived_prequant_filenames(repo_id: str, scheme: str) -> tuple[str, ...]: +def derived_prequant_filenames( + repo_id: str, + scheme: str, + *, + component: Optional[str] = None, +) -> tuple[str, ...]: """The names to try for ``(repo_id, scheme)``, best first, safetensors AHEAD of the pickle. Preferring safetensors is a policy decision rather than a detail: it needs no constructor @@ -458,11 +470,14 @@ def derived_prequant_filenames(repo_id: str, scheme: str) -> tuple[str, ...]: picked up with no code change, and a repo that never does keeps resolving exactly what it resolves today, because the ``.pt`` names stay in the chain behind it. """ - return ( - prequant_repo_filename(repo_id, scheme, ".safetensors"), - prequant_repo_filename(repo_id, scheme, ".pt"), - prequant_filename(scheme), + names = ( + prequant_repo_filename(repo_id, scheme, ".safetensors", component = component), + prequant_repo_filename(repo_id, scheme, ".pt", component = component), ) + part = (component or "").strip() + if part and part != DEFAULT_PREQUANT_COMPONENT: + return names + return names + (prequant_filename(scheme),) def resolve_prequant_source( @@ -488,6 +503,8 @@ def resolve_prequant_source( checkpoints at the root, so there is no directory to prepend; a repo that nested them would 404 on the primary AND on the fallback and the load would silently fall back to dense. """ + if nvfp4_blocked(scheme): + return None override = (path_override or "").strip() if override: return PrequantSource(kind = "path", location = override, filename = None) @@ -591,6 +608,294 @@ def local_prequant_scheme(path: str) -> Optional[str]: return scheme +def hosted_nvfp4_repo_ids() -> frozenset: + """Lowercased repo ids any image/video family registers for nvfp4, read off the family tables.""" + from dataclasses import fields, is_dataclass + + from .diffusion_nvfp4_flag import is_nvfp4 + + families: list = [] + for module in ("diffusion_families", "video_families"): + try: + mod = __import__(f"{__package__}.{module}", fromlist = ["_FAMILIES"]) + families.extend(getattr(mod, "_FAMILIES", ()) or ()) + except Exception: # noqa: BLE001 - a table that fails to import names no repo + continue + repos = set() + for fam in families: + if not is_dataclass(fam): + continue + for field in fields(fam): + value = getattr(fam, field.name, None) + if not isinstance(value, tuple): + continue + for row in value: + if not isinstance(row, tuple) or not row or not isinstance(row[-1], str): + continue + if "/" in row[-1] and any(isinstance(x, str) and is_nvfp4(x) for x in row[:-1]): + repos.add(row[-1].strip().lower()) + return frozenset(repos) + + +# Root only: a diffusers pipeline keeps its weights in component folders. +_PREQUANT_PROBE_SUFFIXES = (".safetensors", ".pt", ".pth") +_PREQUANT_PROBE_LIMIT = 16 + + +def _cached_snapshot_dirs(repo_id: str) -> list: + """The local snapshot folders of Hub repo ``repo_id``, newest first. No network.""" + import os + + parts = repo_id.strip().strip("/").split("/") + if len(parts) != 2 or not all(parts): + return [] + roots = [] + try: + from utils.hf_cache_settings import active_hf_hub_cache + roots.append(active_hf_hub_cache()) + except Exception: # noqa: BLE001 + pass + try: + from huggingface_hub import constants + roots.append(constants.HF_HUB_CACHE) + except Exception: # noqa: BLE001 + pass + found = [] + for root in dict.fromkeys(r for r in roots if r): + snapshots = os.path.join(root, "models--" + "--".join(parts), "snapshots") + try: + entries = [os.path.join(snapshots, name) for name in os.listdir(snapshots)] + except OSError: + continue + entries = [e for e in entries if os.path.isdir(e)] + entries.sort(key = lambda e: os.path.getmtime(e), reverse = True) + found.extend(entries) + return found + + +def _artifacts_declare_nvfp4(folder: str) -> bool: + """Whether a pre-quant artifact at the root of ``folder`` records ``scheme = nvfp4``.""" + import os + + from .diffusion_nvfp4_flag import is_nvfp4 + + try: + names = sorted(os.listdir(folder)) + except OSError: + return False + if "model_index.json" in names: + return False + probed = 0 + for name in names: + if not name.lower().endswith(_PREQUANT_PROBE_SUFFIXES): + continue + path = os.path.join(folder, name) + if not os.path.isfile(path): + continue + if is_nvfp4(local_prequant_scheme(path)): + return True + probed += 1 + if probed >= _PREQUANT_PROBE_LIMIT: + break + return False + + +def declares_nvfp4_checkpoint(model_path: Optional[str]) -> bool: + """Whether ``model_path`` is an NVFP4 prequant checkpoint: recorded metadata first, else registry or name. Never raises.""" + import os + + raw = str(model_path or "").strip() + if not raw: + return False + try: + local = os.path.expanduser(raw) + if os.path.isfile(local): + if not local.lower().endswith(_PREQUANT_PROBE_SUFFIXES): + return False + from .diffusion_nvfp4_flag import is_nvfp4 + return is_nvfp4(local_prequant_scheme(local)) + if os.path.isdir(local): + if _artifacts_declare_nvfp4(local): + return True + else: + key = raw.strip("/").lower() + if key in hosted_nvfp4_repo_ids(): + return True + if any(_artifacts_declare_nvfp4(snap) for snap in _cached_snapshot_dirs(raw)): + return True + except Exception: # noqa: BLE001 - a probe that fails is "not known to be NVFP4" + pass + return raw.rstrip("/\\").lower().endswith("-nvfp4") + + +# Bump on any payload-order or hash change, so two recipes are never compared under one name. +FINGERPRINT_ALGO = "md5-packed-v1" + +# Payload attrs per torchao class NAME, hash order; unlisted reads "not covered", never "equal". +_FINGERPRINT_PAYLOAD: dict = { + "NVFP4Tensor": ("qdata", "scale", "per_tensor_scale"), + "Float8Tensor": ("qdata", "scale"), + "Int8Tensor": ( + "qdata", + "scale", + "zero_point", + "act_quant_scale", + "act_quant_zero_point", + "act_pre_scale", + ), + "MXTensor": ("qdata", "scale"), + "LinearActivationQuantizedTensor": ("original_weight_tensor",), + "AffineQuantizedTensor": ("tensor_impl",), + "PlainAQTTensorImpl": ("int_data", "scale", "zero_point"), +} + + +def _packed_bytes(tensor: Any, torch: Any) -> bytes: + """Raw bytes of ``tensor``; a uint8 view reinterprets, so fp8 / fp4 / bf16 hash exactly.""" + t = tensor.detach().contiguous() + if t.dim() == 0: + t = t.reshape(1) + if t.dtype is not torch.uint8: + t = t.view(torch.uint8) + return t.cpu().numpy().tobytes() + + +def _hash_packed_payload(tensor: Any, digest: Any, torch: Any) -> bool: + """Hash slot names + bytes into ``digest``; False if uncovered or no slot present.""" + names = _FINGERPRINT_PAYLOAD.get(type(tensor).__name__) + if names is None: + return False + hashed = False + for name in names: + value = getattr(tensor, name, None) + if value is None: + continue + digest.update(name.encode("utf-8")) + if type(value).__name__ in _FINGERPRINT_PAYLOAD: + if not _hash_packed_payload(value, digest, torch): + return False + hashed = True + continue + digest.update(_packed_bytes(value, torch)) + hashed = True + return hashed + + +def packed_weight_fingerprint(state_dict: Any, *, select: Any = None) -> dict: + """md5 of every ``.weight``'s packed bytes by fqn; not the pickle, so re-saves match.""" + import hashlib + + import torch + + modules: dict = {} + skipped: list = [] + items = state_dict.items() if hasattr(state_dict, "items") else () + for key, tensor in items: + if key != "weight" and not str(key).endswith(".weight"): + continue + if select is not None and not select(str(key)): + continue + digest = hashlib.md5() + try: + covered = _hash_packed_payload(tensor, digest, torch) + except Exception: # noqa: BLE001 -- an unreadable payload is uncovered, never a raise + covered = False + if covered: + modules[key] = digest.hexdigest() + else: + skipped.append(key) + return { + "algo": FINGERPRINT_ALGO, + "count": len(modules), + "modules": modules, + "skipped": skipped, + } + + +FINGERPRINT_MODE_ENV = "UNSLOTH_PREQUANT_FINGERPRINT" +FINGERPRINT_MODES = ("full", "sample", "off") +FINGERPRINT_SAMPLE_RATE = 8 + + +def _fingerprint_mode() -> str: + """``full`` (default) / ``sample`` / ``off``. An unrecognised value reads as the default.""" + import os + + mode = (os.environ.get(FINGERPRINT_MODE_ENV) or "").strip().lower() + return mode if mode in FINGERPRINT_MODES else "full" + + +def _fingerprint_sampled(fqn: str) -> bool: + """Stable 1-in-8 by md5 of the fqn; ``hash()`` is randomised per process by PYTHONHASHSEED.""" + import hashlib + return hashlib.md5(fqn.encode("utf-8")).digest()[0] % FINGERPRINT_SAMPLE_RATE == 0 + + +def _verify_packed_fingerprint( + state_dict: Any, + metadata: Any, + *, + logger: Any = None, +) -> bool: + """Refuse on a fingerprint mismatch; no block, or one this build cannot compute, passes.""" + block = (metadata or {}).get("fingerprint") + expected = (block or {}).get("modules") if isinstance(block, dict) else None + if not expected: + return True + mode = _fingerprint_mode() + if mode == "off": + if logger is not None: + logger.debug( + "diffusion.prequant: fingerprint check disabled (%s=off)", FINGERPRINT_MODE_ENV + ) + return True + try: + actual = ( + packed_weight_fingerprint( + state_dict, + select = _fingerprint_sampled if mode == "sample" else None, + ).get("modules") + or {} + ) + except Exception as exc: # noqa: BLE001 -- an uncomputable fingerprint checks nothing + _warn(logger, "fingerprint", exc) + return True + if not actual: + _warn( + logger, + "fingerprint", + RuntimeError( + f"this build recognised none of the {len(expected)} quantized weights the " + "checkpoint fingerprinted (a torchao payload rename?); loading it unverified" + ), + ) + return True + checked = [k for k in expected if mode != "sample" or _fingerprint_sampled(k)] + differing = sorted(k for k in checked if actual.get(k) != expected[k]) + counted = mode != "sample" and len(actual) != len(expected) + if not differing and not counted: + if logger is not None: + logger.info( + "diffusion.prequant: fingerprint verified (%d of %d quantized weights, %s)", + len(checked), + len(expected), + mode, + ) + return True + if logger is not None: + logger.error( + "diffusion.prequant: fingerprint MISMATCH (%d of %d checked weights differ, %d " + "weights present vs %d recorded): %s. The checkpoint does not hold the bytes it was " + "built with; falling back to the dense path", + len(differing), + len(checked), + len(actual), + len(expected), + ", ".join(differing[:5]) or "counts only", + ) + return False + + def usable_prequant_source( fam: Any, scheme: str, @@ -782,6 +1087,7 @@ def load_prequantized_transformer( cache_dir: Optional[str] = None, prepare_model: Optional[Any] = None, config_subfolder: str = "transformer", + component: Optional[str] = None, local_files_only: bool = False, logger: Any = None, ) -> Optional[Any]: @@ -791,24 +1097,8 @@ def load_prequantized_transformer( lands under huggingface_hub's import-time constant, so a mid-session cache change re-downloads into a root Unsloth no longer reads. - ``config_subfolder`` is where the DENOISER CONFIG lives inside ``base``, defaulting to the - universal ``transformer``. A family hosting several denoiser partitions in one repo overrides it - with the one this checkpoint belongs to (MiniMax-H3's ``transformer_ref``): the scoped download - stages only that partition, so reading the config from the other one would send an otherwise - fully staged load back to the Hub. - - ``prepare_model`` (optional) is called as ``prepare_model(transformer, metadata)`` on the - freshly built skeleton, AFTER ``from_config`` and BEFORE ``load_state_dict``. That window is the - only one where a family can reshape the module to match how the checkpoint was baked (a swapped - submodule, a patched attention class): earlier there is no module, and later ``strict=True`` has - already rejected the mismatch. It gets the checkpoint's own metadata so it can key on what was - baked rather than on today's defaults. A raising callback falls out to the outer handler below, - i.e. a warning and a dense fallback, never a failed load. - - A checkpoint that declares an ACTIVATION ROTATION (``diffusion_convrot``) has the matching - online half installed here, on exactly the fqns it records. That is unconditional and central - rather than a family opt-in, because the one failure mode worth designing against is the silent - one: rotated weights met by unrotated activations render wrong pixels and raise nothing. + ``component`` is checked (experts share every other key); ``prepare_model`` may reshape before + ``load_state_dict``; a recorded ``diffusion_convrot`` always gets its online half installed. Returns the placed transformer, or None on any problem (missing / mismatched / unreadable checkpoint, unsupported meta-init, or a rotation this build cannot apply exactly) so the caller @@ -843,10 +1133,19 @@ def load_prequantized_transformer( # hand back the same dict, so every check below applies to them equally. ckpt = _load_prequant_checkpoint(path, map_location = "cpu") if not _validate_checkpoint( - ckpt, scheme, base, logger, min_features = min_features, fast_accum = fast_accum + ckpt, + scheme, + base, + logger, + min_features = min_features, + fast_accum = fast_accum, + component = component, ): return None state_dict = ckpt["state_dict"] + # The only check reading what the artifact HOLDS: corruption after build passes the rest. + if not _verify_packed_fingerprint(state_dict, ckpt.get("metadata") or {}, logger = logger): + return None _pin_kernel_preference(state_dict, logger) # Read from the root that actually supplied the checkpoint: after a mid-session cache change the pinned root @@ -892,14 +1191,19 @@ def load_prequantized_transformer( apply_activation_rotation(transformer, metadata, logger = logger) + # assign=True shares the checkpoint tensors; a live ref keeps the CPU copy past to(device). + del state_dict + del ckpt + transformer = transformer.to(device) # Same small-M row padding the runtime quantise path applies, and for the same reason: a checkpoint built # under the current exclusion set QUANTISES the family's small-M linears, so without the wrappers they would # raise inside _int_mm the moment the compiled scope reaches them. After load_state_dict, since wrapping # reparents the Linears; after .to() so the granularity probe reads the device tensors the GEMM will see. - from .diffusion_transformer_quant import apply_small_m_padding + from .diffusion_transformer_quant import apply_small_m_padding, apply_zero_row_guard apply_small_m_padding(transformer, scheme, metadata.get("family"), logger = logger) + apply_zero_row_guard(transformer, scheme, metadata.get("family"), logger = logger) # from_config starts in TRAIN mode while the dense/GGUF paths use from_pretrained (eval()'d). Match it so # train/eval-sensitive layers cannot make prequant inference diverge. try: @@ -1206,14 +1510,9 @@ def _validate_checkpoint( logger: Any, min_features: Optional[int] = None, fast_accum: Optional[bool] = None, + component: Optional[str] = None, ) -> bool: - """Reject a checkpoint that is the wrong format / scheme / base model / filter. ``min_features`` - (when given) is the runtime Linear-feature threshold: a different ``--min-features`` - quantises a different set of Linears, so assign=True would silently install a mismatched - model while status still reports the scheme. Reject it. ``fast_accum`` (fp8 only): when the - caller forces it and the checkpoint baked a different value, the loaded kernels would ignore - the request, so reject and let the dense path honor it. A checkpoint predating a metadata - field (absent) is accepted for back-compat.""" + """Reject a wrong format / scheme / base / filter / denoiser; absent fields predate them.""" if not isinstance(ckpt, dict) or ckpt.get("format") not in PREQUANT_FORMATS: _warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format")) return False @@ -1307,6 +1606,42 @@ def _validate_checkpoint( ), ) return False + ckpt_divisible = meta.get("require_divisible") + if ckpt_divisible is not None: + from .diffusion_transformer_quant import divisible_for_scheme + expected_divisible = divisible_for_scheme(scheme) + if int(ckpt_divisible) != expected_divisible: + _warn( + logger, + scheme, + ValueError( + f"checkpoint require_divisible {ckpt_divisible!r} != {expected_divisible!r}" + ), + ) + return False + elif logger is not None: + logger.debug( + "diffusion.prequant: checkpoint records no require_divisible (built before the " + "field); accepting it for %s", + scheme, + ) + if component: + ckpt_component = meta.get("component") + if ckpt_component is not None and str(ckpt_component) != str(component): + _warn( + logger, + scheme, + ValueError( + f"checkpoint component {ckpt_component!r} != {component!r}; this is another " + "denoiser of the same family and would load clean and render wrong" + ), + ) + return False + if ckpt_component is None and logger is not None: + logger.debug( + "diffusion.prequant: checkpoint records no component; accepting it for %r", + component, + ) # fp8 fast-accum is baked into the saved kernels; only enforce when the caller forces it. if fast_accum is not None: ckpt_fa = meta.get("fast_accum") diff --git a/studio/backend/core/inference/diffusion_quant_pad.py b/studio/backend/core/inference/diffusion_quant_pad.py index 206f28cddb..2a6ae4aa31 100644 --- a/studio/backend/core/inference/diffusion_quant_pad.py +++ b/studio/backend/core/inference/diffusion_quant_pad.py @@ -219,6 +219,80 @@ class PadToMinM(nn.Module): return f"min_m = {self.min_m}, pad_to = {self.pad_to}" +class ZeroRowSafeLinear(nn.Module): + """Answer an EMPTY activation here: torchao NVFP4's whole-input ``max()`` raises on it.""" + + def __init__(self, inner: nn.Linear) -> None: + super().__init__() + self.inner = inner + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + # The bias add broadcasts over zero elements, kept so the result matches F.linear's. + out = x.new_zeros((*x.shape[:-1], self.inner.out_features)) + bias = getattr(self.inner, "bias", None) + return out if bias is None else out + bias + return self.inner(x) + + def __getattr__(self, name: str) -> Any: + try: + return super().__getattr__(name) + except AttributeError: + if name == "inner": + raise + inner = self._modules.get("inner") + if inner is None: + raise + return getattr(inner, name) + + def state_dict(self, *args: Any, **kwargs: Any) -> Any: # type: ignore[override] + """Emit the inner Linear's tensors under the wrapper's prefix, loadable unwrapped.""" + destination = kwargs.pop("destination", args[0] if args else None) + prefix = kwargs.pop("prefix", args[1] if len(args) > 1 else "") + keep_vars = kwargs.pop("keep_vars", args[2] if len(args) > 2 else False) + if destination is None: + return self.inner.state_dict(prefix = prefix, keep_vars = keep_vars) + self.inner.state_dict(destination = destination, prefix = prefix, keep_vars = keep_vars) + return destination + + def _load_from_state_dict( + self, + state_dict: Any, + prefix: str, + local_metadata: Any, + strict: bool, + missing_keys: list, + unexpected_keys: list, + error_msgs: list, + ) -> None: + """Accept the unwrapped key names ``state_dict`` above writes, and hand them to ``inner``.""" + for key in [k for k in state_dict if k.startswith(prefix)]: + leaf = key[len(prefix) :] + if not leaf or leaf.startswith("inner."): + continue + state_dict[prefix + "inner." + leaf] = state_dict.pop(key) + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + +def wrap_zero_row_linears(model: nn.Module, fqns: Iterable[str]) -> tuple[str, ...]: + """Wrap each quantized Linear in ``fqns`` with ``ZeroRowSafeLinear`` (idempotent).""" + done: list[str] = [] + for fqn in sorted(set(fqns)): + parent_name, _, leaf = fqn.rpartition(".") + try: + parent = model.get_submodule(parent_name) if parent_name else model + module = getattr(parent, leaf) + except AttributeError: + continue + if not is_quantized_linear(module): + continue + setattr(parent, leaf, ZeroRowSafeLinear(module)) + done.append(fqn) + return tuple(done) + + def padding_is_bitwise_exact( module: Any, m: int, diff --git a/studio/backend/core/inference/diffusion_te_prequant.py b/studio/backend/core/inference/diffusion_te_prequant.py index dfa2d02277..e9943bc719 100644 --- a/studio/backend/core/inference/diffusion_te_prequant.py +++ b/studio/backend/core/inference/diffusion_te_prequant.py @@ -229,6 +229,10 @@ def family_te_prequant_repo(fam: Any, scheme: str, component: str) -> Optional[s Reads the family's ``te_prequant_repos`` (scheme, component, repo_id) triples; the field is optional on both DiffusionFamily and VideoFamily, so one resolver serves both loaders. """ + from .diffusion_nvfp4_flag import nvfp4_blocked + + if nvfp4_blocked(scheme): + return None for entry in getattr(fam, "te_prequant_repos", ()) or (): try: entry_scheme, entry_component, repo_id = entry diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index ee0a65f551..a4b1a13954 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -26,11 +26,13 @@ import os as _os import re as _re import sys as _sys import threading as _threading +from dataclasses import dataclass from typing import Any, Optional # stdlib-only module (no torch), so this stays inside the "imported lazily" promise above. from core._torchao_stub import is_stubbed, torch_is_rocm from .diffusion_native_quant import int8_act_disabled, int8_act_requested, is_native_linear +from .diffusion_nvfp4_flag import nvfp4_blocked, nvfp4_disabled_message, without_nvfp4 from .diffusion_torchao_patches import install_torchao_int_mm_patch # Also runs in the spawned smoke-probe child, which imports this module and nothing else of the backend. @@ -197,6 +199,45 @@ def apply_small_m_padding( return wrapped +# nvfp4 raises on an empty activation; HunyuanVideo-1.5's attention trim hits it every t2v render. +_HUNYUAN15_NVFP4_ZERO_ROW_TOKENS = ("image_embedder", "context_embedder_2") +_NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS: dict[str, tuple[str, ...]] = { + "hunyuanvideo-1.5": _HUNYUAN15_NVFP4_ZERO_ROW_TOKENS, + "hunyuanvideo-1.5-720p": _HUNYUAN15_NVFP4_ZERO_ROW_TOKENS, +} + + +def zero_row_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tuple[str, ...]: + """Tokens needing the empty-activation guard; nvfp4 only (fp8 / mxfp8 reduce per row).""" + if scheme != TQ_NVFP4: + return () + return _NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS.get(str(family or "").strip().lower(), ()) + + +def apply_zero_row_guard( + transformer: Any, + scheme: str, + family: Optional[str] = None, + *, + logger: Any = None, +) -> tuple[str, ...]: + """Wrap this family's zero-row-reachable quantized Linears after quantize; not best-effort.""" + tokens = zero_row_tokens_for_scheme(scheme, family) + if not tokens: + return () + from .diffusion_quant_pad import matching_linear_fqns, wrap_zero_row_linears + + wrapped = wrap_zero_row_linears(transformer, matching_linear_fqns(transformer, tokens)) + if wrapped and logger is not None: + logger.info( + "diffusion.transformer_quant: zero-row guarded %d %s linears on %s", + len(wrapped), + scheme, + family, + ) + return wrapped + + def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tuple[str, ...]: """Name tokens to exclude from quantisation for ``scheme`` (optionally family-specific). int8 (M>16) skips the M=1 modulation / conditioning-embedder projections @@ -212,6 +253,14 @@ def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tupl return () +_SCHEME_DIVISIBLE: dict[str, int] = {TQ_FP8: 16, TQ_NVFP4: 16, TQ_MXFP8: 32} + + +def divisible_for_scheme(scheme: str) -> int: + """The in/out feature alignment ``scheme``'s GEMM requires, 0 when it has none.""" + return _SCHEME_DIVISIBLE.get(scheme, 0) + + # Per-arch preference for ``auto``, best first. int8 leads every tier: it runs full-rate on consumer and # workstation cards (which halve fp8 FP32 accumulate) and is within a few percent of fp8 elsewhere. # nvfp4 is explicit opt-in and kept OUT of the ladder, being slower and less accurate at DiT shapes. @@ -239,6 +288,18 @@ _FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = { } +@dataclass(frozen = True) +class _AutoPrefer: + """Family head of the ``auto`` order; off below ``floor`` and on consumer GPUs by default.""" + + floor: tuple[int, int] + schemes: tuple[str, ...] + consumer_ok: bool = False + + +_FAMILY_AUTO_PREFER: dict[str, _AutoPrefer] = {} + + # Schemes denied for TRAINING on top of the inference table. Training holds a stricter bar because the evidence above # is rendering evidence: it says a frozen fp8 forward reconstructs the bf16 image, not that a LoRA converges when its # frozen linears are fp8. Nobody has run that, so qwen fp8 stays out of the Train UI until someone does. Delete the @@ -279,6 +340,8 @@ def explain_unusable_scheme(family: Optional[str], scheme: str) -> str: has to separate them: a family the scheme is measured to break, a torchao that cannot run at all in this install, and a GPU without the kernels. Only the last is a hardware limit, and only the last is what the message used to say.""" + if nvfp4_blocked(scheme): + return nvfp4_disabled_message() if family_denies_scheme(family, scheme): return ( f"'{scheme}' is ruled out for family '{family}' by the measured accuracy gate (it " @@ -428,6 +491,8 @@ def normalize_transformer_quant(value: Optional[str]) -> Optional[str]: raise ValueError( f"Unsupported transformer_quant '{value}'. Use one of: {', '.join(TQ_MODES)}." ) + if nvfp4_blocked(normalized): + raise ValueError(nvfp4_disabled_message("transformer_quant")) return normalized @@ -759,14 +824,11 @@ def select_transformer_quant_scheme( cap = _capability() if cap is None: return None - for floor, schemes in _AUTO_LADDER: - if cap >= floor: - for scheme in schemes: - if _family_denied(family, scheme): - continue - if _scheme_supported(scheme, device): - return scheme - return None + for scheme in _auto_scheme_order(family, device, cap): + if _family_denied(family, scheme): + continue + if _scheme_supported(scheme, device): + return scheme return None @@ -803,7 +865,7 @@ def dense_quant_host_capable(target: Any) -> bool: card = _smoke_cache_device_key(str(getattr(target, "device", "cuda"))) for floor, schemes in _AUTO_LADDER: if cap >= floor: - return any(_SMOKE_CACHE.get((scheme, card), True) for scheme in schemes) + return any(_SMOKE_CACHE.get((scheme, card), True) for scheme in without_nvfp4(schemes)) return False @@ -820,14 +882,33 @@ def auto_scheme_candidates(target: Any, family: Optional[str] = None) -> tuple[s cap = _capability() if cap is None: return () + return tuple( + scheme + for scheme in _auto_scheme_order(family, device, cap) + if not _family_denied(family, scheme) and _scheme_supported(scheme, device) + ) + + +def _auto_scheme_order(family: Optional[str], device: Any, cap: tuple[int, int]) -> tuple[str, ...]: + """``auto``'s scheme order for this family and GPU, best first; empty below every tier.""" + tier: tuple[str, ...] = () for floor, schemes in _AUTO_LADDER: if cap >= floor: - return tuple( - scheme - for scheme in schemes - if not _family_denied(family, scheme) and _scheme_supported(scheme, device) - ) - return () + tier = schemes + break + if not tier: + return () + prefer = _FAMILY_AUTO_PREFER.get(str(family or "").strip().lower()) + head: tuple[str, ...] = () + if prefer is not None and cap >= prefer.floor: + if prefer.consumer_ok or not _is_consumer_gpu(device): + head = prefer.schemes + order: list[str] = [] + # Filtered here so the NVFP4 switch reaches every reader of the order. + for scheme in without_nvfp4(head + tier): + if scheme not in order: + order.append(scheme) + return tuple(order) def auto_scheme_candidates_cached(target: Any, family: Optional[str] = None) -> tuple[str, ...]: @@ -842,15 +923,13 @@ def auto_scheme_candidates_cached(target: Any, family: Optional[str] = None) -> cap = _capability() if cap is None: return () - card = _smoke_cache_device_key(str(getattr(target, "device", "cuda"))) - for floor, schemes in _AUTO_LADDER: - if cap >= floor: - return tuple( - scheme - for scheme in schemes - if not _family_denied(family, scheme) and _SMOKE_CACHE.get((scheme, card), True) - ) - return () + device = str(getattr(target, "device", "cuda")) + card = _smoke_cache_device_key(device) + return tuple( + scheme + for scheme in _auto_scheme_order(family, device, cap) + if not _family_denied(family, scheme) and _SMOKE_CACHE.get((scheme, card), True) + ) def _capability() -> Optional[tuple[int, int]]: @@ -869,6 +948,8 @@ def _scheme_supported( unproven_ok: bool = False, ) -> bool: """CUDA + (for fp8) the fp8 dtype + a cached quantise+matmul smoke test for ``scheme``.""" + if nvfp4_blocked(scheme): + return False try: import torch if not torch.cuda.is_available(): @@ -966,7 +1047,14 @@ def _child_probe_table(device: str) -> Optional[dict[str, Optional[bool]]]: # and binds the child to this process's lifetime, so a wedged probe cannot outlive the backend. proc = ctx.Process( target = run_without_native_path_secret, - args = (__name__, "_child_probe_entry", {}, device, TQ_SCHEMES, queue), + args = ( + __name__, + "_child_probe_entry", + {}, + device, + without_nvfp4(TQ_SCHEMES), + queue, + ), daemon = True, ) proc.start() @@ -1463,9 +1551,7 @@ def quantize_transformer( # non-bf16 ones. "lora_" keeps a baked adapter's side path high precision. Runtime only: NOT part of # exclude_tokens_for_scheme, whose list is baked into prequant metadata. exclude = exclude_tokens_for_scheme(scheme, family) + ("lora_",) - # GEMM tiling floors per scheme: scaled_mm needs 16-aligned dims, MX block scaling 32. int8's _int_mm has no - # such floor and keeps the historical filter. - divisible = {TQ_FP8: 16, TQ_NVFP4: 16, TQ_MXFP8: 32}.get(scheme, 0) + divisible = divisible_for_scheme(scheme) quantize_( transformer, _make_quant_config(scheme, fast_accum = fast_accum), @@ -1480,6 +1566,7 @@ def quantize_transformer( # here means the transformer is quantized but not safely compilable, so it falls into the except below and the # caller loads GGUF. apply_small_m_padding(transformer, scheme, family, logger = logger) + apply_zero_row_guard(transformer, scheme, family, logger = logger) try: transformer._unsloth_runtime_quant = scheme except Exception: # noqa: BLE001 - marker is best-effort diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index e94de608a2..ec7a42fda1 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -31,6 +31,7 @@ family's official base repos, or a local path the user explicitly picked. from __future__ import annotations import contextlib +import dataclasses import functools import inspect import os @@ -166,6 +167,7 @@ from .video_families import ( validate_video_request_shape, video_family_prequant_available, video_family_prequant_repo, + video_family_prequant_resident_gb, video_family_prequant_schemes, ) from .video_frames import uint8_video_frames @@ -258,7 +260,15 @@ def assert_video_precision_available( Public because the ROUTE has to make this call itself, before it takes the GPU: the copy in ``begin_load`` runs inside ``acquire_for``, which evicts chat under the arbiter lock before the - register callback.""" + register callback. + + NVFP4 while the NVFP4 switch is off raises ``ValueError`` (the route's 400) first, ahead of the + opt-in fallback: a disabled scheme is refused, never swapped.""" + from .diffusion_nvfp4_flag import refuse_disabled_nvfp4 + + refuse_disabled_nvfp4( + transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant + ) if precision_fallback_allowed(): return pinned = normalize_transformer_quant(transformer_quant) @@ -880,7 +890,7 @@ class _VideoLoadState: # MiniMax-H3 only: the pre-quantized denoiser was PINNED to the device and taken out of the offload rotation, so it # is resident alongside whatever component is running. The memory preflight has to know, because that turns its # floor from a max into a sum. - h3_denoiser_pinned: bool = False + denoiser_pinned: bool = False resolved: Optional[dict] = None @@ -1221,6 +1231,178 @@ def _h3_auto_denoiser_scheme( return H3_AUTO_FALLBACK_SCHEME +def _planned_denoiser_components(fam: Any, kind: str) -> tuple[str, ...]: + """The denoiser component(s) a CONVENTIONAL seeded load covers, or ``()``.""" + try: + if kind != "pipeline" or getattr(fam, "modular_workflow", None): + return () + from .video_denoiser_prequant import denoiser_components + return denoiser_components(fam) + except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense shards + return () + + +def _video_auto_denoiser_scheme( + fam: Any, + *, + target: Any, + requested: Optional[str], + base_repo: Optional[str], + speed_mode: Optional[str] = None, +) -> Optional[str]: + """The scheme a CONVENTIONAL video load would seed from a hosted checkpoint, or None.""" + # Only an AUTO precision becomes "off" under Speed="off"; an EXPLICIT scheme is honored, so it keeps its seed. + auto = requested is None or str(requested).strip().lower() in ("", "auto") + if auto and speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF: + return None + try: + if getattr(fam, "modular_workflow", None): + return None + # Registry before the smoke probe: a scheme with no hosted row can never seed, so it must not spawn probes. + hosted = video_family_prequant_schemes(fam) + if not hosted or (not auto and normalize_transformer_quant(requested) not in hosted): + return None + scheme = select_transformer_quant_scheme( + target, requested, family = getattr(fam, "name", None) + ) + if scheme is None or scheme == TQ_AUTO: + return None + from .video_denoiser_prequant import denoiser_prequant_sources + + # EVERY component or none: partial coverage drops shards the dense fallback then has to open. + if denoiser_prequant_sources(fam, scheme, base_repo) is None: + return None + from .diffusion_prequant import restricted_prequant_load_supported + + if not restricted_prequant_load_supported(scheme): + return None + return scheme + except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense denoiser + return None + + +# The planner DECIDED against the seed (not "never ran"); a load must not re-take it after teardown. +DENOISER_SEED_DECLINED = "__declined__" + + +def _target_ordinal(target: Any) -> Optional[int]: + """Card index of a device target (unindexed CUDA = current); None on a single-device backend.""" + ordinal = getattr(target, "ordinal", None) + if ordinal is not None or getattr(target, "device", None) != "cuda": + return ordinal + try: + import torch + return int(torch.cuda.current_device()) + except Exception: # noqa: BLE001 -- unanswerable: count every card, the old reading + return None + + +def _pipeline_device_mib(pipe: Any, ordinal: Optional[int] = None) -> int: + """Accelerator MiB the pipe holds on ``ordinal`` (or every card), counting each storage once.""" + if pipe is None: + return 0 + seen: set[int] = set() + total = 0 + + def _add(t: Any) -> None: + nonlocal total + flatten = getattr(t, "__tensor_flatten__", None) + if callable(flatten) and type(t).__name__ not in ("Tensor", "Parameter"): + try: + for name in flatten()[0]: + _add(getattr(t, name)) + return + except Exception: # noqa: BLE001 -- fall through to the plain reading + pass + try: + device = getattr(t, "device", None) + if getattr(device, "type", "cpu") in ("cpu", "meta"): + return + if ordinal is not None and (getattr(device, "index", None) or 0) != ordinal: + return + storage = t.untyped_storage() + if storage.data_ptr() in seen: + return + seen.add(storage.data_ptr()) + total += int(storage.nbytes()) + except Exception: # noqa: BLE001 -- an unreadable tensor adds nothing + pass + + components = getattr(pipe, "components", None) + modules = components.values() if isinstance(components, dict) else () + for module in modules: + for fn in ("parameters", "buffers"): + try: + for t in getattr(module, fn)(): + _add(getattr(t, "data", t)) + except Exception: # noqa: BLE001 -- not a torch module + break + return total // (1024 * 1024) + + +def _video_seed_stays_resident( + fam: Any, + *, + target: Any, + scheme: str, + memory_mode: Optional[str], + text_encoder_quant: Optional[str], + base_repo: Optional[str], + reclaimable_mib: int = 0, +) -> bool: + """True when an artifact-sized plan keeps the denoiser resident; torchao rejects offload.""" + components = getattr(fam, "bf16_components_gb", None) + if not components: + return True + measured = video_family_prequant_resident_gb(fam, scheme) + factor = _QUANT_STEADY_FACTOR.get(scheme) + if measured: + denoiser_gb = float(measured) + elif factor is not None: + denoiser_gb = components[0] * factor + else: + return True + import torch + + from .diffusion_te_prequant import te_prequant_budget_scale + + dtype = getattr(target, "dtype", None) + dtype_scale = ( + 2.0 if getattr(target, "device", None) != "cpu" and dtype is torch.float32 else 1.0 + ) + te_scale = te_prequant_budget_scale( + fam, + te_quant_mode = text_encoder_quant, + target = target, + base = base_repo, + ) + vae_scale = 1.0 if getattr(fam, "vae_force_fp32", False) else dtype_scale + mib_per_gb = 1000.0**3 / (1024.0 * 1024.0) + model_dense_mib = int( + (denoiser_gb + components[1] * te_scale * dtype_scale + components[2] * vae_scale) + * mib_per_gb + ) + device_memory = settled_snapshot_device_memory(target) + if reclaimable_mib > 0 and device_memory.free_mib is not None: + free = device_memory.free_mib + int(reclaimable_mib) + if device_memory.total_mib is not None: + free = min(free, device_memory.total_mib) + device_memory = dataclasses.replace(device_memory, free_mib = free) + planned = plan_diffusion_memory( + target = target, + device_memory = device_memory, + model_dense_mib = model_dense_mib, + runtime_headroom_mib = estimate_video_runtime_mib( + width = fam.resolution_presets[0][0], + height = fam.resolution_presets[0][1], + num_frames = fam.default_num_frames, + ), + companion_dense_mib = None, + requested_mode = normalize_memory_mode(memory_mode), + ) + return planned.offload_policy == "none" + + def _progress(phase: Optional[str], **extra: Any) -> dict[str, Any]: return {"phase": phase, **extra} @@ -1847,28 +2029,50 @@ class VideoBackend: text_encoder_quant = kwargs.get("text_encoder_quant"), speed_mode = kwargs.get("speed_mode"), h3_task = kwargs.get("h3_task"), - # This decides the file set AND the memory policy, so it has to read the card the pipeline will land on: - # the default one can be larger or smaller than the pick. gpu_ordinal = kwargs.get("gpu_ordinal"), ) + video_auto_denoiser = self._video_planned_auto_denoiser_scheme( + fam, + base = base, + kind = kind, + transformer_quant = kwargs.get("transformer_quant"), + speed_mode = kwargs.get("speed_mode"), + memory_mode = kwargs.get("memory_mode"), + text_encoder_quant = kwargs.get("text_encoder_quant"), + gpu_ordinal = kwargs.get("gpu_ordinal"), + ) + video_seed_declined = video_auto_denoiser == DENOISER_SEED_DECLINED + if video_seed_declined: + video_auto_denoiser = None + conventional_denoiser = kind == "pipeline" and not getattr( + fam, "modular_workflow", None + ) + requested_denoiser = None if conventional_denoiser else kwargs.get("transformer_quant") skip_transformer_weights = self._denoiser_prequant_verified( fam, - h3_auto_denoiser or kwargs.get("transformer_quant"), + h3_auto_denoiser or video_auto_denoiser or requested_denoiser, base, kwargs.get("h3_task"), kwargs.get("hf_token"), local_files_only = local_files_only, + kind = kind, + ) + skip_transformer_components = ( + _planned_denoiser_components(fam, kind) if skip_transformer_weights else () ) - # Handed to the loader only when the dense shards really are gone from the pull. Then the choice is already - # committed and the loader must not re-take it against a reading that has moved since -- there would be no - # dense denoiser left to fall back to. When the plan kept them, this stays None and the loader decides for - # itself as before. kwargs["_h3_auto_denoiser_planned"] = ( h3_auto_denoiser if h3_auto_denoiser and skip_transformer_weights else None ) - # And for MiniMax-H3's conditioner, whose hosted quantized artifact replaces the base repo's 62 GB dense - # text_encoder/ shards outright. Verified, not just name-matched: this scheme decides whether the base pull - # DROPS the dense encoder, and a derivative that the seed will decline must keep its own. + kwargs["_video_auto_denoiser_planned"] = ( + DENOISER_SEED_DECLINED + if video_seed_declined + else ( + video_auto_denoiser + if video_auto_denoiser and skip_transformer_weights + else None + ) + ) + kwargs["_denoiser_prequant_skipped"] = skip_transformer_components h3_te_scheme = self._h3_te_quant_scheme_verified( fam, kwargs.get("text_encoder_quant"), @@ -1884,6 +2088,7 @@ class VideoBackend: kind, te_sources = te_sources, skip_transformer_weights = skip_transformer_weights, + skip_transformer_components = skip_transformer_components, h3_task = kwargs.get("h3_task"), h3_te_scheme = h3_te_scheme, local_files_only = local_files_only, @@ -1903,6 +2108,17 @@ class VideoBackend: H3_TE_QUANT_REPO, H3_LEGACY_TE_QUANT_REPO, ) + # Dense shards dropped: a delete of the seeded repo mid-fetch strands the load. + if skip_transformer_weights: + claimed = self._denoiser_prequant_repo_ids( + fam, + h3_auto_denoiser or video_auto_denoiser or requested_denoiser, + base, + kwargs.get("h3_task"), + ) + self._loading.asset_repos = tuple( + dict.fromkeys(self._loading.asset_repos + claimed) + ) self._loading.expected_bytes = expected # Checkpoint downloads outside the lock so an unload can preempt the multi-GB pull; companions pre-download # the same way. @@ -1957,6 +2173,7 @@ class VideoBackend: ltx23 = True, te_sources = te_sources, skip_transformer_weights = skip_transformer_weights, + skip_transformer_components = skip_transformer_components, h3_task = kwargs.get("h3_task"), h3_te_scheme = h3_te_scheme, local_files_only = local_files_only, @@ -1980,6 +2197,19 @@ class VideoBackend: cancel_event = cancel_event, local_files_only = local_files_only, ) + # The denoiser artifact too: the injection that would fetch it has no cancel event. + if skip_transformer_weights: + self._fetch_denoiser_prequant( + self._denoiser_prequant_source_list( + fam, + h3_auto_denoiser or video_auto_denoiser or requested_denoiser, + base, + kwargs.get("h3_task"), + ), + kwargs.get("hf_token"), + cancel_event = cancel_event, + local_files_only = local_files_only, + ) base_local = self._predownload_base( base, kwargs.get("hf_token"), @@ -1987,6 +2217,7 @@ class VideoBackend: ltx23 = ltx23, skip_te_components = te_skipped + h3_te_skipped, skip_transformer_weights = skip_transformer_weights, + skip_transformer_components = skip_transformer_components, # Picks the H3 denoiser partition: fl2va stages transformer/, ref2va stages the separate 66.28 GB # transformer_ref/ that load_components(workflow="ref2va") opens h3_task = kwargs.get("h3_task"), @@ -2638,29 +2869,14 @@ class VideoBackend: h3_task: Optional[str], hf_token: Optional[str], local_files_only: bool = False, + kind: str = "pipeline", ) -> bool: """``_denoiser_prequant_covered`` plus the Hub check, for the decisions that COMMIT. The pure probe answers from the registry, which is right where it must not raise or touch - the network. It is NOT enough to drop the dense shards from the pull: a task-specific row - gets no filename fallback, so a Ref2VA artifact that is renamed, unpublished or gated - resolves to nothing while the registry still says the scheme is covered. Skipping on that - word alone stages neither denoiser, and the bf16 fallback the loader documents then has - nothing to open offline (or pulls 66 GB inline, outside this load's progress, cancel and - disk budget). - - Same rule as ``_h3_te_quant_scheme_verified`` next door, and the same fail-closed - direction: unanswerable keeps the dense shards. - - ``local_files_only`` swaps the Hub probe for a cache probe rather than refusing the load. - Refusing would break the one case the flag exists for -- a fully cached model coming up - offline -- and the question the probe answers is not actually about the Hub: it is - "is there a replacement denoiser to open instead of the dense shards?". Offline nothing is - downloaded either way, so the honest reading is whether the artifact is already on disk, - and that is exactly as strict as the online one in the direction that matters (an absent - checkpoint keeps the dense shards, so the loader's bf16 fallback still has something to - open).""" - if not self._denoiser_prequant_covered(fam, transformer_quant, base, h3_task): + the network, but a renamed / gated task artifact still reads covered there. Unanswerable + keeps the dense shards; ``local_files_only`` answers from the cache.""" + if not self._denoiser_prequant_covered(fam, transformer_quant, base, h3_task, kind = kind): return False from .diffusion_prequant import restricted_prequant_load_supported @@ -2750,6 +2966,67 @@ class VideoBackend: except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense shards return None + def _video_planned_auto_denoiser_scheme( + self, + fam: Any, + *, + base: Optional[str], + kind: str, + transformer_quant: Optional[str], + speed_mode: Optional[str], + memory_mode: Optional[str] = None, + text_encoder_quant: Optional[str] = None, + gpu_ordinal: Optional[int] = None, + ) -> Optional[str]: + """The plan's seed scheme; ``DENOISER_SEED_DECLINED`` if the load would offload.""" + try: + if kind != "pipeline" or getattr(fam, "modular_workflow", None): + return None + import torch + + # SCOPED, not pinned: this pooled to_thread worker must not keep this request's card. + with diffusion_device_scope(gpu_ordinal): + target = ( + resolve_diffusion_device_target() + if gpu_ordinal is None + else resolve_diffusion_device_target(ordinal = gpu_ordinal) + ) + if getattr(fam, "fp16_incompatible", False) and target.dtype is torch.float16: + # The loader promotes fp16 to float32 here, so read the dtype the load will. + return None + scheme = _video_auto_denoiser_scheme( + fam, + target = target, + requested = transformer_quant, + base_repo = base, + speed_mode = speed_mode, + ) + if scheme is None: + return None + resident = self._state + if not _video_seed_stays_resident( + fam, + target = target, + scheme = scheme, + memory_mode = memory_mode, + text_encoder_quant = text_encoder_quant, + base_repo = base, + reclaimable_mib = _pipeline_device_mib( + getattr(resident, "pipe", None), + ordinal = _target_ordinal(target), + ), + ): + logger.info( + "video.denoiser_prequant: an artifact-sized plan for %s still offloads on " + "this card, and offload moves the DiT, so the dense denoiser shards are " + "kept", + scheme, + ) + return DENOISER_SEED_DECLINED + return scheme + except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense shards + return None + def _h3_te_quant_scheme_verified( self, fam: Any, @@ -2886,39 +3163,113 @@ class VideoBackend: return () return ("text_encoder",) + def _fetch_denoiser_prequant( + self, + sources: list[Any], + hf_token: Optional[str], + *, + cancel_event: Optional[threading.Event] = None, + local_files_only: bool = False, + ) -> None: + """Pre-fetch hosted denoiser checkpoint(s) under the load's cancel event; best effort.""" + cancel = cancel_event if cancel_event is not None else self._cancel_event + from core.inference.diffusion_prequant import candidate_filenames_of + from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback + + for source in sources: + if getattr(source, "kind", None) != "repo": + continue + names = list(dict.fromkeys(candidate_filenames_of(source))) + for index, name in enumerate(names): + try: + hf_hub_download_with_xet_fallback( + source.location, + name, + hf_token, + cancel_event = cancel, + cache_dir = hub_cache_dir(), + reuse_other_cache_root = True, + local_files_only = local_files_only, + ) + break + except Exception as exc: # noqa: BLE001 -- the injection re-resolves and falls back + if cancel.is_set(): + raise + if index == len(names) - 1: + logger.warning( + "video.denoiser_prequant_fetch_failed: %s/%s: %s", + source.location, + name, + exc, + ) + @staticmethod def _denoiser_prequant_covered( fam: Any, transformer_quant: Optional[str], base: Optional[str], h3_task: Optional[str] = None, + *, + kind: str = "pipeline", ) -> bool: - """True when a hosted PRE-QUANTIZED denoiser checkpoint replaces the dense DiT shards. - - The one place the plan, the byte estimate and the scoped pull all ask, so none of them can - stage weights another one skipped. Only a modular-workflow family qualifies: every other - family quantises its own dense transformer on device and still needs those shards. - - ``h3_task`` names the partition, because coverage is per (scheme, task): a scheme whose - hosted checkpoint is the OTHER partition covers nothing here, and answering True for it - would drop the very shards that load then has to open. - - Never raises and never touches the network -- a pure registry lookup -- because it runs on - the download-planning path, where a failure would cost the whole plan.""" + """True when hosted pre-quantized denoisers replace ALL dense DiT shards; never raises.""" try: - if not getattr(fam, "modular_workflow", None): - return False scheme = normalize_transformer_quant(transformer_quant) - # "auto" leaves the choice to the backend, which keeps the released bfloat16 DENOISER for a modular workflow - # (the conditioner default is a separate decision, made in _h3_te_quant_scheme), so it must NOT drop the - # dense shards that load then wants. The hosted checkpoints were measured and left opt-in: see - # _load_h3_modular_pipeline. + # An unresolved "auto" must NOT drop dense shards; the backend settles it first. if scheme is None or scheme == TQ_AUTO: return False - return video_family_prequant_available(fam, scheme, task = h3_task, base_repo = base) + if getattr(fam, "modular_workflow", None): + return video_family_prequant_available(fam, scheme, task = h3_task, base_repo = base) + if kind != "pipeline": + return False + from .video_denoiser_prequant import denoiser_prequant_sources + + return denoiser_prequant_sources(fam, scheme, base) is not None except Exception: # noqa: BLE001 -- an unanswerable probe keeps the dense shards return False + @staticmethod + def _denoiser_prequant_source_list( + fam: Any, + transformer_quant: Optional[str], + base: Optional[str], + h3_task: Optional[str] = None, + ) -> list[Any]: + """Every ``PrequantSource`` a seeded load would open, or ``[]``; never raises.""" + scheme = (transformer_quant or "").strip().lower() + if scheme in ("", "auto", "off", "none"): + return [] + try: + from .diffusion_prequant import resolve_prequant_source + from .video_denoiser_prequant import denoiser_prequant_sources + + if getattr(fam, "modular_workflow", None): + source = resolve_prequant_source(fam, scheme, base_repo = base, task = h3_task) + return [source] if source is not None else [] + resolved = denoiser_prequant_sources(fam, scheme, base) + return list(resolved.values()) if resolved else [] + except Exception as exc: # noqa: BLE001 -- a bad registry entry must not sink the plan + logger.warning("video.denoiser_prequant_unresolved: %s", exc) + return [] + + @staticmethod + def _denoiser_prequant_repo_ids( + fam: Any, + transformer_quant: Optional[str], + base: Optional[str], + h3_task: Optional[str] = None, + ) -> tuple[str, ...]: + """Repo id(s) a seeded denoiser is fetched from, for the in-flight delete guard.""" + return tuple( + dict.fromkeys( + src.location + for src in VideoBackend._denoiser_prequant_source_list( + fam, transformer_quant, base, h3_task + ) + if getattr(src, "kind", None) == "repo" and getattr(src, "location", None) + ) + ) + @staticmethod def _denoiser_prequant_hub_files( fam: Any, @@ -2927,46 +3278,35 @@ class VideoBackend: api: Any, h3_task: Optional[str] = None, ) -> tuple[Optional[str], list[tuple[str, int]]]: - """``(repo_id, [(rfilename, size)])`` for the hosted pre-quantized denoiser, or - ``(None, [])``. + """``(repo_id, [(rfilename, size)])`` for all hosted denoisers or ``(None, [])``.""" + sources = VideoBackend._denoiser_prequant_source_list(fam, transformer_quant, base, h3_task) + if not sources or any(getattr(src, "kind", None) != "repo" for src in sources): + return None, [] + locations = {src.location for src in sources} + if len(locations) != 1: + logger.warning("video.denoiser_prequant_split_repos: %s", sorted(locations)) + return None, [] + location = sources[0].location + from core.inference.diffusion_nvfp4_flag import nvfp4_repo_blocked - The denoiser mirror of ``_te_prequant_hub_files``. ``_denoiser_prequant_covered`` already - drops the dense DiT shards from the plan, so without this the checkpoint that replaces - them is in no entry at all: the byte total under-reports by the size of the artifact - (20.25 GB for H3 int8), the disk preflight passes on a volume that cannot hold it, and an - offline stage completes without the one file the load needs. - - Same failure rule as the encoder helper: an unreachable repo yields no files and is only - logged, because the plan is a staging hint and the load still falls back to dense.""" - scheme = (transformer_quant or "").strip().lower() - if scheme in ("", "auto", "off", "none"): + if nvfp4_repo_blocked(location): return None, [] try: - from .diffusion_prequant import resolve_prequant_source - - # Task-keyed, like the coverage probe: staging the other partition's checkpoint would advertise the right - # byte count for the wrong file - source = resolve_prequant_source(fam, scheme, base_repo = base, task = h3_task) - except Exception as exc: # noqa: BLE001 -- a bad registry entry must not sink the plan - logger.warning("video.denoiser_prequant_unresolved: %s", exc) + info = api.model_info(location, files_metadata = True) + except Exception as exc: # noqa: BLE001 -- unavailable prequant means the dense DiT + logger.warning("video.denoiser_prequant_unavailable: %s: %s", location, exc) return None, [] - if source is None or getattr(source, "kind", None) != "repo": - return None, [] - # The root name first, then the legacy scheme name: resolve_prequant_source hands back both and the load tries - # them in that order, so the plan must stage whichever one exists. from core.inference.diffusion_prequant import candidate_filenames_of - wanted = list(candidate_filenames_of(source)) - try: - info = api.model_info(source.location, files_metadata = True) - except Exception as exc: # noqa: BLE001 -- unavailable prequant means the dense DiT - logger.warning("video.denoiser_prequant_unavailable: %s: %s", source.location, exc) - return None, [] by_name = {s.rfilename: int(s.size or 0) for s in (info.siblings or [])} - for name in wanted: - if name in by_name: - return source.location, [(name, by_name[name])] - return None, [] + files: list[tuple[str, int]] = [] + for src in sources: + wanted = list(candidate_filenames_of(src)) + found = next((n for n in wanted if n in by_name), None) + if found is None: + return None, [] + files.append((found, by_name[found])) + return location, files @staticmethod def _denoiser_prequant_cached_repo( @@ -2977,37 +3317,23 @@ class VideoBackend: ) -> Optional[str]: """The hosted pre-quantized denoiser repo when its checkpoint is ALREADY cached, else None. - The offline twin of ``_denoiser_prequant_hub_files``, for a load that may not reach the - Hub. ``resolve_prequant_source`` is pure registry work and ``try_to_load_from_cache`` - (behind ``_hub_file_is_cached``) is network-free, so this answers the same question -- is - there a replacement denoiser to open? -- from disk alone. - - Both candidate names are tried, in the order the load tries them, and both cache roots are - searched, because that is what the fetch would have resolved through. No size to - corroborate against without the Hub, so a cached name is taken at face value: the loader - opening a damaged checkpoint falls back to dense, which is the same failure an online - size mismatch would have produced one step later.""" - try: - from .diffusion_prequant import resolve_prequant_source - source = resolve_prequant_source( - fam, - (transformer_quant or "").strip().lower(), - base_repo = base, - task = h3_task, - ) - except Exception as exc: # noqa: BLE001 -- a bad registry entry keeps the dense shards - logger.warning("video.denoiser_prequant_unresolved: %s", exc) - return None + Offline twin of ``_denoiser_prequant_hub_files``; a cached name is taken at face value.""" + sources = VideoBackend._denoiser_prequant_source_list(fam, transformer_quant, base, h3_task) # A local override is on disk by definition; only a hosted checkpoint has a cache to probe - if source is None or getattr(source, "kind", None) != "repo": + if not sources or any(getattr(src, "kind", None) != "repo" for src in sources): return None from core.inference.diffusion import DiffusionBackend from core.inference.diffusion_prequant import candidate_filenames_of - for name in candidate_filenames_of(source): - if DiffusionBackend._hub_file_is_cached(source.location, name): - return source.location + cached: list[str] = [] + for src in sources: + for name in candidate_filenames_of(src): + if DiffusionBackend._hub_file_is_cached(src.location, name): + cached.append(src.location) + break + if len(cached) == len(sources): + return cached[0] # No log here: the caller reports the same "keeping its dense denoiser shards" outcome for a miss, and logging # it twice would read as two separate decisions. return None @@ -3088,6 +3414,7 @@ class VideoBackend: ltx23: bool = False, skip_te_components: tuple[str, ...] = (), skip_transformer_weights: bool = False, + skip_transformer_components: tuple[str, ...] = (), h3_task: Optional[str] = None, ) -> list[tuple[str, int]]: """The (rfilename, size) list a load actually needs from the base repo. @@ -3110,9 +3437,7 @@ class VideoBackend: - the dense weight shards of the denoiser partition this load opens, under ``skip_transformer_weights``, supplied instead by a hosted PRE-QUANTIZED denoiser checkpoint (H3's transformer is 66.3 GB of its base repo). - ``transformer/config.json`` is kept for the same reason the pre-cast - encoders keep theirs: the pre-quant loader meta-inits the DiT from it, - so dropping it would break the very load that made the skip safe. + ``transformer/config.json`` stays; ``skip_transformer_components`` names the covered ones. ``h3_task`` picks the H3 denoiser partition. The base repo ships two, and a load only ever brings up one: ``transformer/`` for fl2va (which also covers text-only) and @@ -3132,6 +3457,11 @@ class VideoBackend: # will actually open: a reference load stages transformer_ref/ and nothing else, so dropping "transformer/" # drops nothing and the plan carries the full 66.28 GB the pre-quantized checkpoint exists to replace. h3_skip_component = h3_denoiser_prefix.rstrip("/") + skip_components = ( + (tuple(skip_transformer_components) or (h3_skip_component,)) + if skip_transformer_weights + else () + ) h3_prefixes = ( "audio_scheduler/", "audio_vae/", @@ -3167,7 +3497,7 @@ class VideoBackend: continue # is_prequant_covered_weight is component-agnostic (it matches "/" against a weight suffix), so # the encoder helper serves the denoiser verbatim. - if skip_transformer_weights and is_prequant_covered_weight(name, (h3_skip_component,)): + if skip_components and is_prequant_covered_weight(name, skip_components): continue if name.startswith("text_encoder/diffusion_pytorch_model"): continue @@ -3188,6 +3518,7 @@ class VideoBackend: ltx23: bool = False, te_sources: Optional[dict[str, Any]] = None, skip_transformer_weights: bool = False, + skip_transformer_components: tuple[str, ...] = (), h3_task: Optional[str] = None, h3_te_scheme: Optional[str] = None, local_files_only: bool = False, @@ -3233,6 +3564,7 @@ class VideoBackend: ltx23 = ltx23, skip_te_components = skip_te_components, skip_transformer_weights = skip_transformer_weights, + skip_transformer_components = skip_transformer_components, h3_task = h3_task, ) ) @@ -3300,6 +3632,22 @@ class VideoBackend: ) or transformer_quant ) + video_planned = self._video_planned_auto_denoiser_scheme( + fam, + base = base, + kind = kind, + transformer_quant = transformer_quant, + speed_mode = load_kwargs.get("speed_mode"), + memory_mode = load_kwargs.get("memory_mode"), + text_encoder_quant = text_encoder_quant, + gpu_ordinal = load_kwargs.get("gpu_ordinal"), + ) + if video_planned == DENOISER_SEED_DECLINED: + video_planned = None + if kind == "pipeline" and not getattr(fam, "modular_workflow", None): + transformer_quant = video_planned + else: + transformer_quant = video_planned or transformer_quant # Only the header tells an LTX-2.3 checkpoint from 2.0 and it is not on disk yet, so narrow the base pull by # NAME: a wrong guess costs an inline pull, the wide base list costs gigabytes. ltx23 = self._pick_looks_like_ltx23(fam, repo_id, gguf_filename, kind) @@ -3447,9 +3795,10 @@ class VideoBackend: skip_transformer_weights = ( dq_repo is not None and self._denoiser_prequant_covered( - fam, transformer_quant, base, h3_task + fam, transformer_quant, base, h3_task, kind = kind ) ), + skip_transformer_components = _planned_denoiser_components(fam, kind), h3_task = h3_task, ), revision = getattr(info, "sha", None), @@ -3752,6 +4101,7 @@ class VideoBackend: ltx23: bool = False, skip_te_components: tuple[str, ...] = (), skip_transformer_weights: bool = False, + skip_transformer_components: tuple[str, ...] = (), h3_task: Optional[str] = None, cancel_event: Optional[threading.Event] = None, local_files_only: bool = False, @@ -3788,6 +4138,7 @@ class VideoBackend: ltx23 = ltx23, skip_te_components = skip_te_components, skip_transformer_weights = skip_transformer_weights, + skip_transformer_components = skip_transformer_components, h3_task = h3_task, ) if not any(name == "model_index.json" for name, _ in files): @@ -3949,6 +4300,8 @@ class VideoBackend: _base_local_dir: Optional[str] = None, _te_prequant_skipped: tuple[str, ...] = (), _h3_auto_denoiser_planned: Optional[str] = None, + _video_auto_denoiser_planned: Optional[str] = None, + _denoiser_prequant_skipped: tuple[str, ...] = (), ) -> dict[str, Any]: fam = self.validate_load_request( repo_id, @@ -4110,10 +4463,16 @@ class VideoBackend: num_frames = fam.default_num_frames, ) - def _plan_for_te_scale(scale: float, *, log: bool) -> tuple[Any, Any, bool]: + def _plan_for_te_scale( + scale: float, + *, + log: bool, + denoiser_gb: Optional[float] = None, + ) -> tuple[Any, Any, bool]: """``(plan, bf16_plan, quant_replanned)`` for a text-encoder budget of ``scale`` x its bf16 size. Pure and cheap (``plan_diffusion_memory`` is arithmetic), so the - dense-encoder plan can be rebuilt below if the pre-cast injection does not land.""" + dense-encoder plan can be rebuilt below if the pre-cast injection does not land. + ``denoiser_gb`` replaces the bf16 DiT term for a SEEDED load (steady and peak).""" text_encoder_gb = components[1] * scale if components is not None else 0.0 # dtype_scale doubles bf16 terms when the fp16 promotion lands fp32 on an accelerator. A vae_force_fp32 # family is the one component it must NOT touch: its table term is already recorded at fp32 and assembly @@ -4135,8 +4494,11 @@ class VideoBackend: components[1], ) if kind == "pipeline": + transformer_gb = ( + components[0] * dtype_scale if denoiser_gb is None else float(denoiser_gb) + ) model_dense_mib = ( - int((components[0] * dtype_scale + scaled_companions_gb) * mib_per_gb) + int((transformer_gb + scaled_companions_gb) * mib_per_gb) if components is not None else None ) @@ -4162,15 +4524,10 @@ class VideoBackend: # need, so re-plan with the scheme factor and keep resident if it fits. dense_plan = planned replanned_for_quant = False - # NOT re-triggered by a unified-memory shortfall, deliberately. This re-plan prices the quantised DiT's - # STEADY size, but the video path has no pre-quantised artifact: the DiT is always materialised dense - # (from_pretrained / from_single_file) and quantize_transformer rewrites it in place afterwards, so the - # build PEAK is the bf16 figure whatever scheme is picked. On discrete VRAM the steady state is the right - # thing to plan placement against -- an oversized peak raises a catchable torch OOM and offload is still - # there -- but on unified memory the peak is what the OS kills for, with no exception to catch, so the - # refusal below has to keep reading the dense plan. + # Not for a unified-memory shortfall: runtime quant still peaks at the dense bf16 DiT. if ( kind == "pipeline" + and denoiser_gb is None and planned.offload_policy != "none" and normalize_transformer_quant(transformer_quant) is not None and ( @@ -4206,12 +4563,42 @@ class VideoBackend: replanned_for_quant = True return planned, dense_plan, replanned_for_quant - plan, bf16_plan, quant_replanned = _plan_for_te_scale(te_scale, log = True) + denoiser_seed_scheme: Optional[str] = None + if kind == "pipeline" and _video_auto_denoiser_planned != DENOISER_SEED_DECLINED: + denoiser_seed_scheme = _video_auto_denoiser_planned or _video_auto_denoiser_scheme( + fam, + target = target, + requested = transformer_quant, + base_repo = base, + speed_mode = speed_mode, + ) + denoiser_seed_sources: dict[str, Any] = {} + if denoiser_seed_scheme is not None: + from .video_denoiser_prequant import denoiser_prequant_sources + denoiser_seed_sources = denoiser_prequant_sources(fam, denoiser_seed_scheme, base) or {} + if not denoiser_seed_sources: + denoiser_seed_scheme = None + denoiser_seed_gb: Optional[float] = None + if denoiser_seed_scheme is not None and components is not None: + measured = video_family_prequant_resident_gb(fam, denoiser_seed_scheme) + factor = _QUANT_STEADY_FACTOR.get(denoiser_seed_scheme) + if measured: + denoiser_seed_gb = float(measured) + elif factor is not None: + denoiser_seed_gb = components[0] * factor + plan, bf16_plan, quant_replanned = _plan_for_te_scale( + te_scale, log = True, denoiser_gb = denoiser_seed_gb + ) + if denoiser_seed_scheme is not None and plan.offload_policy != "none": + # Offload hooks use Module.to(), which torchao tensors reject: re-plan at bf16. + denoiser_seed_scheme = None + denoiser_seed_gb = None + plan, bf16_plan, quant_replanned = _plan_for_te_scale(te_scale, log = False) + settled_te_scale = te_scale # On unified memory the plan's 'none' policy is a placement, not a fit: there is no offload tier left to fall # back to, so an oversized load is killed by the OS with no torch OOM to catch. Refuse now, after the eviction - # above and before any weight is materialised -- on the dense plan, which is what the DiT build peaks at (see - # above). + # above and before any weight is materialised, on the plan the build actually peaks at. raise_on_unified_memory_shortfall(plan, family = getattr(fam, "name", None), logger = logger) # ── build the pipeline. @@ -4255,11 +4642,50 @@ class VideoBackend: "video.te_prequant_budget: no pre-cast encoder was injected; re-planning " "memory at the dense bf16 text-encoder size" ) - plan, bf16_plan, quant_replanned = _plan_for_te_scale(1.0, log = False) + settled_te_scale = 1.0 + plan, bf16_plan, quant_replanned = _plan_for_te_scale( + settled_te_scale, log = False, denoiser_gb = denoiser_seed_gb + ) + if denoiser_seed_scheme is not None and plan.offload_policy != "none": + denoiser_seed_scheme = None + denoiser_seed_gb = None + plan, bf16_plan, quant_replanned = _plan_for_te_scale(settled_te_scale, log = False) + denoiser_injected: dict[str, Any] = {} + if denoiser_seed_scheme is not None: + from .video_denoiser_prequant import denoiser_prequant_pipe_kwargs + + denoiser_injected = denoiser_prequant_pipe_kwargs( + fam, + base, + scheme = denoiser_seed_scheme, + dtype = dtype, + device = device, + hf_token = hf_token, + target = target, + local_files_only = local_files_only, + cache_dir = hub_cache_dir(), + logger = logger, + ) + pipe_kwargs.update(denoiser_injected) + if not denoiser_injected: + logger.warning( + "video.denoiser_prequant: no pre-quantized denoiser was seeded; re-planning " + "memory at the dense bf16 DiT size" + ) + denoiser_seed_scheme = None + denoiser_seed_gb = None + plan, bf16_plan, quant_replanned = _plan_for_te_scale(settled_te_scale, log = False) + raise_on_unified_memory_shortfall( + plan, family = getattr(fam, "name", None), logger = logger + ) # Injection is best-effort (a missing checkpoint falls back to the dense encoder), but the scoped pre-download # already dropped those shards and from_pretrained cannot re-fetch from a local snapshot dir, so top them up # here. ltx23 never reaches this: it gets the hub id. - missing_dense = [c for c in (_te_prequant_skipped or ()) if c not in pipe_kwargs] + missing_dense = [ + c + for c in tuple(_te_prequant_skipped or ()) + tuple(_denoiser_prequant_skipped or ()) + if c not in pipe_kwargs + ] if missing_dense and _base_local_dir: logger.warning( "video.te_prequant: %s not injected; restoring the dense weights the " @@ -4358,12 +4784,14 @@ class VideoBackend: # Why the quant did not engage, in the caller's terms; threaded into `resolved`. transformer_quant_decline: Optional[str] = None transformer_quant_decline_status = RESOLVED_FELL_BACK + transformer_quant_source: Optional[str] = None video_offload = plan.offload_policy != "none" native_scheme = ( native_quant_scheme( target, transformer_quant_pinned, family = fam.name, offload = video_offload ) - if kind == "pipeline" + # A seeded denoiser is a torchao checkpoint; the torchao-free path must not claim it. + if kind == "pipeline" and not denoiser_injected else None ) native_kwargs = ( @@ -4379,6 +4807,7 @@ class VideoBackend: kind == "pipeline" and normalize_transformer_quant(transformer_quant) == TQ_AUTO and not quant_replanned + and not denoiser_injected and plan.offload_policy == "none" # An unmeasured budget also plans "none" without proving a fit. and plan.estimates.get("safe_device_budget_mib") is not None @@ -4406,7 +4835,19 @@ class VideoBackend: # holding a working GPU learns nothing about why it declined. transformer_quant_decline = dense_transformer_unsupported_reason(target) transformer_quant_decline_status = RESOLVED_UNSUPPORTED - if ( + if denoiser_injected: + transformer_quant_engaged = denoiser_seed_scheme + transformer_quant_source = ( + f"hosted pre-quantized {denoiser_seed_scheme} denoiser checkpoint " + f"({', '.join(sorted({src.location for src in denoiser_seed_sources.values()}))})" + ) + logger.info( + "video.transformer_quant: %s pre-quantized denoiser(s) seeded for %s; skipping " + "the runtime quantise", + denoiser_seed_scheme, + fam.name, + ) + elif ( kind == "pipeline" and normalize_transformer_quant(transformer_quant) is not None and dense_transformer_supported(target) @@ -4742,11 +5183,13 @@ class VideoBackend: transformer_quant_requested, transformer_quant_engaged or "off", # Honest framing: the shipped torchao schemes cut load time and resident memory ~2x, but - # per-step GEMMs are at best bf16 parity. - native_quant_reason( + # per-step GEMMs are at best bf16 parity. A seeded load names its checkpoint instead. + transformer_quant_source + or native_quant_reason( getattr(pipe, "transformer", None), transformer_quant_engaged ) - if transformer_quant_engaged is not None and native_scheme is not None + if transformer_quant_engaged is not None + and (transformer_quant_source or native_scheme is not None) else "DiT(s) quantised (halves resident weights; hosted checkpoints cut " "load time; per-step speed is roughly bf16 parity)" if transformer_quant_engaged is not None @@ -5594,7 +6037,7 @@ class VideoBackend: # a dense fallback recorded as int8 would under-state the floor by 39 GB and let a doomed generation # start. text_encoder_quant = text_encoder_quant_engaged, - h3_denoiser_pinned = denoiser_pinned, + denoiser_pinned = denoiser_pinned, resolved = resolved, ) self._precommit_globals = None @@ -6244,7 +6687,7 @@ class VideoBackend: state.text_encoder_quant, bf16_gb = H3_TEXT_ENCODER_BF16_GB ), transformer_gb = h3_transformer_resident_gb(state.transformer_quant), - transformer_pinned = bool(getattr(state, "h3_denoiser_pinned", False)), + transformer_pinned = bool(getattr(state, "denoiser_pinned", False)), ) if available_vram_gb + 0.25 < required_vram_gb: raise RuntimeError( diff --git a/studio/backend/core/inference/video_denoiser_prequant.py b/studio/backend/core/inference/video_denoiser_prequant.py new file mode 100644 index 0000000000..3084b2aad7 --- /dev/null +++ b/studio/backend/core/inference/video_denoiser_prequant.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Seed a conventional video pipeline with hosted pre-quantized denoisers, all-or-nothing.""" + +from __future__ import annotations + +from typing import Any, Optional + +MOE_SECOND_DENOISER = "transformer_2" + + +def denoiser_components(fam: Any) -> tuple[str, ...]: + """Pipeline attribute name(s) of this family's denoiser(s), from the registry.""" + if getattr(fam, "is_moe", False): + return ("transformer", MOE_SECOND_DENOISER) + return ("transformer",) + + +def denoiser_prequant_sources( + fam: Any, scheme: Optional[str], base_repo: Optional[str] +) -> Optional[dict[str, Any]]: + """``{component: PrequantSource}`` for EVERY denoiser, or None (never partial); never raises.""" + wanted = (scheme or "").strip().lower() + if wanted in ("", "auto", "off", "none"): + return None + try: + from .diffusion_nvfp4_flag import nvfp4_blocked + + if nvfp4_blocked(wanted): + return None + from .diffusion_prequant import resolve_prequant_source + from .video_families import video_family_prequant_task_specific + + sources: dict[str, Any] = {} + for component in denoiser_components(fam): + task = None if component == "transformer" else component + if task is not None and not video_family_prequant_task_specific(fam, wanted, task): + # Without its own row, ``task`` resolves expert 1's file: loads fine, wrong weights. + return None + sources[component] = resolve_prequant_source( + fam, + wanted, + base_repo = base_repo, + task = task, + ) + if sources[component] is None: + return None + return sources or None + except Exception: # noqa: BLE001 -- an unanswerable registry keeps the dense denoiser + return None + + +def denoiser_prequant_pipe_kwargs( + fam: Any, + base_repo: str, + *, + scheme: Optional[str], + dtype: Any, + device: str, + hf_token: Optional[str] = None, + target: Any = None, + local_files_only: bool = False, + cache_dir: Optional[str] = None, + logger: Any = None, +) -> dict[str, Any]: + """Overrides for assembly, or ``{}`` (dense bf16 + runtime quant) unless all denoisers seed.""" + try: + import gc + + sources = denoiser_prequant_sources(fam, scheme, base_repo) + if not sources: + return {} + if target is not None: + from .diffusion_transformer_quant import dense_transformer_supported + if not dense_transformer_supported(target): + return {} + import diffusers + + from .diffusion_prequant import load_prequantized_transformer + from .diffusion_transformer_quant import DEFAULT_MIN_LINEAR_FEATURES + + transformer_cls = getattr(diffusers, str(getattr(fam, "transformer_class", "")), None) + if transformer_cls is None: + _warn( + logger, + scheme or "", + ValueError( + f"transformer class {getattr(fam, 'transformer_class', None)!r} " + "not found in diffusers" + ), + ) + return {} + seeded: dict[str, Any] = {} + for component, source in sources.items(): + module = load_prequantized_transformer( + transformer_cls, + base_repo, + source, + device = device, + dtype = dtype, + hf_token = hf_token, + scheme = scheme, + min_features = DEFAULT_MIN_LINEAR_FEATURES, + config_subfolder = component, + component = component, + cache_dir = cache_dir, + local_files_only = local_files_only, + logger = logger, + ) + if module is None: + seeded.clear() + gc.collect() + if logger is not None: + logger.info( + "video.denoiser_prequant: no usable %s checkpoint for %s %s, so the " + "dense denoiser(s) are loaded and quantised at runtime instead", + scheme, + getattr(fam, "name", None), + component, + ) + return {} + seeded[component] = module + gc.collect() + if logger is not None: + logger.info( + "video.denoiser_prequant: seeded %s pre-quantized denoiser(s) for %s (%s)", + len(seeded), + getattr(fam, "name", None), + scheme, + ) + return seeded + except Exception as exc: # noqa: BLE001 -- seeding is an optimisation, never a blocker + _warn(logger, "pipe_kwargs", exc) + return {} + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("video.denoiser_prequant: %s failed: %s", what, exc) diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py index 82387d7775..fdb29fa68d 100644 --- a/studio/backend/core/inference/video_families.py +++ b/studio/backend/core/inference/video_families.py @@ -21,6 +21,8 @@ import re from dataclasses import dataclass, field from typing import Optional +from .diffusion_nvfp4_flag import nvfp4_blocked, without_nvfp4 + # The request model's ceiling on num_frames, declared HERE so the shape gate and the bound cannot drift: the gate's # refusal names the lattice point above the request, and suggesting one the request model would itself reject is a # dead end. VideoGenerateRequest imports this for its `le`. @@ -100,6 +102,7 @@ class VideoFamily: # parameters), so 0.55 x 66.3 GB over-states it by 16 GB and a hard refusal turns away a load that fits. Measured # from Hub file metadata (2026-08-09): MiniMax-H3-FP8.pt 20,260,192,855 bytes, MiniMax-H3-INT8.pt 20,253,894,865. prequant_resident_gb: Optional[float] = None + prequant_resident_gb_by_scheme: tuple[tuple[str, float], ...] = field(default_factory = tuple) # Per-variant overrides as (base_repo, scheme, repo_id), keyed on the LOWERCASED upstream base id. A pre-quantized # checkpoint is baked from ONE base's weights and the loader refuses it for any other base, so a variant that ships # its own denoiser needs its own entry; a variant without one falls through to prequant_repos and, if that @@ -249,6 +252,9 @@ _FAMILIES: tuple[VideoFamily, ...] = ( pipeline_class = "WanPipeline", transformer_class = "WanTransformer3DModel", base_repo = "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + prequant_repos = (("nvfp4", "unsloth/Wan2.2-TI2V-5B-NVFP4"),), + prequant_filenames = (("nvfp4", "Wan2.2-TI2V-5B-NVFP4.pt"),), + prequant_resident_gb_by_scheme = (("nvfp4", 2.9),), # "wan2.2-5b"/"wan-ti2v" are the picker/GGUF short ids; "wan2.2-ti2v" catches the repo stem aliases = ("wan2.2-5b", "wan-ti2v", "wan2.2-ti2v", "wan-ti2v-5b"), has_audio = False, @@ -276,6 +282,13 @@ _FAMILIES: tuple[VideoFamily, ...] = ( pipeline_class = "WanPipeline", transformer_class = "WanTransformer3DModel", base_repo = "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + prequant_repos = (("nvfp4", "unsloth/Wan2.2-T2V-A14B-NVFP4"),), + prequant_filenames = ( + ("nvfp4", "Wan2.2-T2V-A14B-NVFP4.pt"), + ("nvfp4", "transformer_2", "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"), + ), + # BOTH experts: the plan subtracts one denoiser term and this family builds two. + prequant_resident_gb_by_scheme = (("nvfp4", 16.2),), aliases = ("wan2.2-14b", "wan-t2v", "wan2.2-t2v", "wan-t2v-a14b", "wan-a14b"), has_audio = False, # is_moe drives the dual-DiT optimisation layers; cfg2_kwarg names the pipeline kwarg for transformer_2's @@ -306,6 +319,9 @@ _FAMILIES: tuple[VideoFamily, ...] = ( pipeline_class = "HunyuanVideo15Pipeline", transformer_class = "HunyuanVideo15Transformer3DModel", base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", + prequant_repos = (("nvfp4", "unsloth/HunyuanVideo-1.5-NVFP4"),), + prequant_filenames = (("nvfp4", "HunyuanVideo-1.5-Diffusers-480p_t2v-NVFP4.pt"),), + prequant_resident_gb_by_scheme = (("nvfp4", 4.8),), # No bare "hunyuanvideo" alias: it would also claim the incompatible 1.0 repos. aliases = ("hunyuanvideo-1-5", "hunyuanvideo1.5", "hunyuanvideo1-5", "hv15"), has_audio = False, @@ -331,6 +347,10 @@ _FAMILIES: tuple[VideoFamily, ...] = ( pipeline_class = "HunyuanVideo15Pipeline", transformer_class = "HunyuanVideo15Transformer3DModel", base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v", + # RTN, not GPTQ: the 720p transformer is separately trained; 480p corrections do not apply. + prequant_repos = (("nvfp4", "unsloth/HunyuanVideo-1.5-NVFP4"),), + prequant_filenames = (("nvfp4", "HunyuanVideo-1.5-Diffusers-720p_t2v-NVFP4.pt"),), + prequant_resident_gb_by_scheme = (("nvfp4", 4.8),), aliases = ("hunyuanvideo-1.5-diffusers-720p_t2v", "hv15-720p"), has_audio = False, guidance_via_guider = True, @@ -422,6 +442,8 @@ def video_family_prequant_repo( the refusal path of a load request, and a table typo must not turn a legitimate pick into a 500. A family object that predates these fields simply has no hosted checkpoint. """ + if nvfp4_blocked(scheme): + return None base = _prequant_base_key(base_repo) if base: for entry in getattr(fam, "prequant_variant_repos", ()) or (): @@ -439,6 +461,26 @@ def video_family_prequant_repo( return None +def video_family_prequant_resident_gb(fam: VideoFamily, scheme: str) -> Optional[float]: + """The MEASURED resident size in decimal GB of this family's hosted ``scheme`` denoiser.""" + if nvfp4_blocked(scheme): + return None + for entry in getattr(fam, "prequant_resident_gb_by_scheme", ()) or (): + if not isinstance(entry, (tuple, list)) or len(entry) != 2: + continue + entry_scheme, resident_gb = entry + if entry_scheme == scheme and resident_gb: + try: + return float(resident_gb) + except (TypeError, ValueError): # a malformed row is "not measured", never a 500 + continue + measured = getattr(fam, "prequant_resident_gb", None) + try: + return float(measured) if measured else None + except (TypeError, ValueError): + return None + + def video_family_prequant_task_specific(fam: VideoFamily, scheme: str, task: str) -> bool: """True when the family names an artifact for exactly this ``(scheme, task)`` pair. @@ -499,7 +541,7 @@ def video_family_prequant_schemes(fam: VideoFamily, task: Optional[str] = None) schemes.append(entry[1]) if task: schemes = [s for s in schemes if video_family_prequant_available(fam, s, task = task)] - return tuple(schemes) + return without_nvfp4(schemes) def snap_num_frames(fam: VideoFamily, num_frames: int) -> int: diff --git a/studio/backend/main.py b/studio/backend/main.py index 9de3aa5a49..6c087b5284 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -2318,6 +2318,15 @@ def _dense_quant_supported() -> bool: return bool(_dense_quant_capability) +def _nvfp4_diffusion_enabled() -> bool: + """Whether image and video generation may offer NVFP4 (``UNSLOTH_NVFP4_DIFFUSION``).""" + try: + from core.inference.diffusion_nvfp4_flag import nvfp4_diffusion_enabled + return nvfp4_diffusion_enabled() + except Exception: # noqa: BLE001 -- a capability read must never fail a status request + return False + + def _dense_quant_schemes() -> list[str]: """The scheme ladder for ``/api/system``, a pure read of already-resolved state: the polled route must never import torch, and the entry beside it refreshed both in one pass.""" @@ -2424,6 +2433,8 @@ def get_system_info( # pure read of that same pass. "dense_quant_supported": _dense_quant_supported(), "dense_quant_schemes": _dense_quant_schemes(), + # Torch-free env read, safe on this polled route. + "nvfp4_diffusion": _nvfp4_diffusion_enabled(), } diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8fbceab552..b832570eb5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -39951,6 +39951,44 @@ async def _selected_gpu_ordinal(gpu_ids, *, allow_ranking: bool = True) -> Optio ) +def _refuse_disabled_nvfp4_request(request: Any) -> None: + """400 a load or plan naming NVFP4 while the NVFP4 switch (``UNSLOTH_NVFP4_DIFFUSION``) is off. + + First thing in every image and video load / plan route, ahead of the precision gates and their + opt-in silent fallback, so the request is refused outright, never swapped for another scheme, + and nothing is resolved, planned or fetched for it.""" + from core.inference.diffusion_nvfp4_flag import refuse_disabled_nvfp4 + try: + refuse_disabled_nvfp4( + transformer_quant = getattr(request, "transformer_quant", None), + text_encoder_quant = getattr(request, "text_encoder_quant", None), + ) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) + + +async def _refuse_disabled_nvfp4_checkpoint(request: Any) -> None: + """400 a load or plan whose model is itself an NVFP4 checkpoint while the NVFP4 switch is off. + + The scheme gates above never see such a pick: it names no precision, it IS one. Runs after the + account checks, since it reads the local path or cached snapshot the request names.""" + from core.inference.diffusion_nvfp4_flag import ( + nvfp4_diffusion_enabled, + refuse_disabled_nvfp4_checkpoint, + ) + + if nvfp4_diffusion_enabled(): + return + try: + await asyncio.to_thread( + refuse_disabled_nvfp4_checkpoint, + getattr(request, "model_path", None), + getattr(request, "base_repo", None), + ) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) + + @studio_router.post("/images/download-plan", response_model = DiffusionDownloadPlanResponse) async def diffusion_download_plan( request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) @@ -39960,6 +39998,7 @@ async def diffusion_download_plan( Validates the same way /images/load does, so an unloadable pick fails here rather than after a multi-GB download.""" + _refuse_disabled_nvfp4_request(request) if account_access.managed_account(): await asyncio.to_thread(account_access.require_media_references, request) if account_access.managed_account(): @@ -39971,6 +40010,7 @@ async def diffusion_download_plan( request = request.model_copy( update = {"hf_token": account_access.account_hf_token(request.hf_token)} ) + await _refuse_disabled_nvfp4_checkpoint(request) from core.inference.diffusion import ( get_diffusion_backend, resolve_local_single_file, @@ -40099,9 +40139,13 @@ def _assert_native_precision_unset( Raises RuntimeError, which the route maps to 409 alongside the diffusers refusals.""" from core.inference.diffusion_auto_policy import precision_fallback_allowed + from core.inference.diffusion_nvfp4_flag import refuse_disabled_nvfp4 from core.inference.diffusion_precision import normalize_te_quant from core.inference.diffusion_transformer_quant import TQ_AUTO, normalize_transformer_quant + refuse_disabled_nvfp4( + transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant + ) if precision_fallback_allowed(): return pinned = normalize_transformer_quant(transformer_quant) @@ -40162,6 +40206,7 @@ async def load_diffusion_model_gated( Media auto-switch awaits this rather than the route so the idle unload can tell an API-loaded pipeline from one the user picked on the Images page. """ + _refuse_disabled_nvfp4_request(request) if account_access.managed_account(): await asyncio.to_thread(account_access.require_media_references, request) account_access.require_idle_other_accounts() @@ -40174,6 +40219,7 @@ async def load_diffusion_model_gated( request = request.model_copy( update = {"hf_token": account_access.account_hf_token(request.hf_token)} ) + await _refuse_disabled_nvfp4_checkpoint(request) # Tested at ENTRY because `begin_load` returns before the worker moves a byte, so the only # fact available is that a repo ALREADY cached is not one this load will fetch. The WRITE is # deferred to the launch below, since the validation in between 400s without moving a byte. diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index b6acd8b2e6..3963170ef7 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -144,6 +144,12 @@ async def video_download_plan( ): """The repos + files this pick needs, so the frontend stages them through the Hub download manager instead of the load downloading inline. Mirrors /images/download-plan.""" + from routes.inference import ( + _refuse_disabled_nvfp4_checkpoint, + _refuse_disabled_nvfp4_request, + ) + + _refuse_disabled_nvfp4_request(request) if account_access.managed_account(): await asyncio.to_thread(account_access.require_media_references, request) if account_access.managed_account(): @@ -155,6 +161,7 @@ async def video_download_plan( request = request.model_copy( update = {"hf_token": account_access.account_hf_token(request.hf_token)} ) + await _refuse_disabled_nvfp4_checkpoint(request) from core.inference.diffusion import resolve_local_single_file from core.inference.video import ( assert_video_precision_available, @@ -269,6 +276,12 @@ async def load_video_model_gated( """Everything ``POST /video/load`` does, plus who asked for it. Media auto-switch awaits this rather than the route so the idle unload can tell an API-loaded pipeline from one the user picked on the Video page. """ + from routes.inference import ( + _refuse_disabled_nvfp4_checkpoint, + _refuse_disabled_nvfp4_request, + ) + + _refuse_disabled_nvfp4_request(request) if account_access.managed_account(): await asyncio.to_thread(account_access.require_media_references, request) account_access.require_idle_other_accounts() @@ -281,6 +294,7 @@ async def load_video_model_gated( request = request.model_copy( update = {"hf_token": account_access.account_hf_token(request.hf_token)} ) + await _refuse_disabled_nvfp4_checkpoint(request) # Same as the image load: tested at entry, because `begin_load` returns before the worker # moves a byte, and written at the launch below, because the validation in between 400s # without starting one and the record would be permanent. diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index 866b6df82f..e794196f62 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -1239,3 +1239,43 @@ def _drop_the_idle_reload_stash_between_tests(): yield finally: _keepwarm._set_last_unloaded(None) + + +# Run with the NVFP4 switch on; test_nvfp4_diffusion_flag and test_build_prequant_checkpoint must stay off. +_NVFP4_ENABLED_TEST_MODULES = frozenset( + { + "test_dense_quant_rocm_gate_9396", + "test_diffusion_auto_policy", + "test_diffusion_backend", + "test_diffusion_inference_info", + "test_diffusion_lora", + "test_diffusion_native_quant", + "test_diffusion_pipeline_prequant", + "test_diffusion_precision", + "test_diffusion_prequant", + "test_diffusion_quant_pad", + "test_diffusion_routes", + "test_diffusion_te_prequant", + "test_diffusion_transformer_quant", + "test_train_precision_scheme_contract", + "test_video_backend", + "test_video_families", + "test_video_h3_te_quant", + "test_video_prequant", + "test_video_routes", + "test_xformers_stub_diffusion_parity", + } +) +_NVFP4_ENABLED_TEST_PREFIX = "test_diffusion_nvfp4_" + + +@pytest.fixture(autouse = True) +def _nvfp4_diffusion_enabled_for_nvfp4_tests(request, monkeypatch): + """Switch NVFP4 on for the modules above; every other module sees the default (off).""" + module = getattr(request, "module", None) + name = getattr(module, "__name__", "").rsplit(".", 1)[-1] + if name in _NVFP4_ENABLED_TEST_MODULES or name.startswith(_NVFP4_ENABLED_TEST_PREFIX): + monkeypatch.setenv("UNSLOTH_NVFP4_DIFFUSION", "1") + else: + monkeypatch.delenv("UNSLOTH_NVFP4_DIFFUSION", raising = False) + yield diff --git a/studio/backend/tests/test_build_prequant_checkpoint.py b/studio/backend/tests/test_build_prequant_checkpoint.py index 22789d04e4..5d51b04819 100644 --- a/studio/backend/tests/test_build_prequant_checkpoint.py +++ b/studio/backend/tests/test_build_prequant_checkpoint.py @@ -253,3 +253,118 @@ def test_the_recorded_base_must_be_the_canonical_id_not_just_the_same_tail(capsy assert fam.transformer_class in str(exc), exc continue assert rc != 2 or "is not" not in capsys.readouterr().out, accepted + + +def _build_nvfp4( + monkeypatch, + tmp_path, + fam = None, + extra_argv = (), + loaded = None, +): + """Run the build on a tiny aligned + ragged dense model; returns (quantized names, metadata).""" + torch = pytest.importorskip("torch") + pytest.importorskip("torchao") + import torchao.quantization as tq + + import core.inference.diffusion_transformer_quant as dtq + + fam = fam or detect_family("Tongyi-MAI/Z-Image-Turbo") + assert fam is not None + + class _Dense(torch.nn.Module): + def __init__(self): + super().__init__() + self.aligned = torch.nn.Linear(1024, 1024, dtype = torch.bfloat16) + self.ragged = torch.nn.Linear(1024, 1000, dtype = torch.bfloat16) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + if loaded is not None: + loaded.append(kwargs.get("subfolder")) + return cls() + + def to(self, *args, **kwargs): + return self + + fake_diffusers = types.ModuleType("diffusers") + fake_diffusers.__version__ = "0" + setattr(fake_diffusers, fam.transformer_class, _Dense) + monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers) + + quantized = [] + + def _fake_quantize(model, config, filter_fn): + quantized.extend(n for n, m in model.named_modules() if n and filter_fn(m, n)) + + monkeypatch.setattr(tq, "quantize_", _fake_quantize) + monkeypatch.setattr(dtq, "_make_quant_config", lambda *a, **k: object()) + + build = _script() + out = tmp_path / "out.pt" + argv = ["--base", "b", "--family", fam.name, "--scheme", "nvfp4", "--out", str(out)] + assert build.main(argv + list(extra_argv)) == 0 + return quantized, torch.load(out, weights_only = False)["metadata"] + + +def test_the_build_skips_ragged_linears_and_records_the_alignment_floor(monkeypatch, tmp_path): + """A build skips Linears off the GEMM tiling floor and stamps ``require_divisible``.""" + import core.inference.diffusion_transformer_quant as dtq + + quantized, meta = _build_nvfp4(monkeypatch, tmp_path) + assert quantized == ["aligned"] + assert meta["require_divisible"] == dtq.divisible_for_scheme("nvfp4") == 16 + + +def test_the_build_stamps_its_denoiser_so_it_cannot_load_as_the_second_expert( + monkeypatch, tmp_path +): + """The build reads ``subfolder="transformer"``; published under a transformer_2 name, the loader must refuse it.""" + from core.inference.diffusion_prequant import PREQUANT_FORMAT, _validate_checkpoint + + _, meta = _build_nvfp4(monkeypatch, tmp_path) + assert meta["component"] == "transformer" + ckpt = {"format": PREQUANT_FORMAT, "metadata": {**meta, "base_model_id": ""}, "state_dict": {}} + assert _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer") + assert not _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer_2") + + +def test_the_second_expert_builds_from_its_own_subfolder_and_loads_as_transformer_2( + monkeypatch, tmp_path +): + """``--component transformer_2`` reads that subfolder, stamps it, and publishes under the name the loader asks for.""" + from core.inference.diffusion_prequant import PREQUANT_FORMAT, _validate_checkpoint + + wan = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers", override = "wan2.2-t2v-a14b") + assert wan is not None + loaded = [] + _, meta = _build_nvfp4( + monkeypatch, tmp_path, fam = wan, extra_argv = ["--component", "transformer_2"], loaded = loaded + ) + assert loaded == ["transformer_2"] + assert meta["component"] == "transformer_2" and meta["family"] == wan.name + ckpt = {"format": PREQUANT_FORMAT, "metadata": {**meta, "base_model_id": ""}, "state_dict": {}} + assert _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer_2") + assert not _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer") + + build = _script() + assert ( + build.upload_destination(wan, "nvfp4", rotated = False, component = "transformer_2") + == "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt" + ) + # The default component keeps the name it has always published under. + assert build.upload_destination(wan, "nvfp4", rotated = False) == "transformer_nvfp4.pt" + # A component the family declares no row for would land where nothing ever looks. + with pytest.raises(ValueError, match = "transformer_2"): + build.upload_destination( + detect_family("Tongyi-MAI/Z-Image-Turbo"), + "nvfp4", + rotated = False, + component = "transformer_2", + ) + + +def test_a_default_build_still_reads_the_transformer_subfolder(monkeypatch, tmp_path): + loaded = [] + _, meta = _build_nvfp4(monkeypatch, tmp_path, loaded = loaded) + assert loaded == ["transformer"] and meta["component"] == "transformer" diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 39db4d73ec..14ef77d165 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -88,6 +88,18 @@ def test_prequant_repo_filename_convention(): ) assert prequant_repo_filename("org/Some-Model-quantized", "fp8") == "Some-Model-FP8.pt" assert prequant_repo_filename("org/PlainRepo", "int8") == "PlainRepo-INT8.pt" + assert ( + prequant_repo_filename("unsloth/Wan2.2-TI2V-5B-NVFP4", "nvfp4") == "Wan2.2-TI2V-5B-NVFP4.pt" + ) + assert prequant_repo_filename("unsloth/Model-MXFP8", "mxfp8") == "Model-MXFP8.pt" + assert ( + prequant_repo_filename("unsloth/Wan2.2-T2V-A14B-NVFP4", "nvfp4", component = "transformer_2") + == "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt" + ) + assert ( + prequant_repo_filename("unsloth/Wan2.2-T2V-A14B-NVFP4", "nvfp4", component = "transformer") + == "Wan2.2-T2V-A14B-NVFP4.pt" + ) def test_resolve_variant_base_picks_variant_repo(): @@ -302,6 +314,225 @@ def test_usable_source_repo_unaffected_by_allowlist(monkeypatch, restricted_load assert src is not None and src.kind == "repo" and src.location == "org/hosted-fp8" +_UINT8 = object() # the stub torch's uint8, so the fingerprint's dtype view is a no-op here + + +class _Bytes: + """Enough of a tensor for the fingerprint: it views the payload as uint8 and hashes it.""" + + dtype = _UINT8 + + def __init__(self, payload): + self._payload = payload + + def detach(self): + return self + + def contiguous(self): + return self + + def dim(self): + return 1 + + def view(self, dtype): + return self + + def cpu(self): + return self + + def numpy(self): + import numpy as np + return np.frombuffer(self._payload, dtype = np.uint8) + + +class Float8Tensor: + """A quantized weight as far as the fingerprint is concerned: the class NAME is the key.""" + + def __init__(self, qdata): + self.qdata = _Bytes(qdata) + self.scale = _Bytes(b"scale") + + +class _Recorder: + """A logger that keeps what it was told, so a refusal can be asserted to NAME the weight.""" + + def __init__(self): + self.lines: list = [] + + def _record(self, msg, *args): + self.lines.append(msg % args if args else msg) + + debug = info = warning = error = _record + + @property + def text(self): + return "\n".join(self.lines) + + +def _fingerprinted_ckpt(payloads, *, recorded = None): + """A checkpoint whose state dict holds ``payloads`` and whose metadata fingerprints them.""" + from core.inference.diffusion_prequant import packed_weight_fingerprint + + state_dict = {fqn: Float8Tensor(data) for fqn, data in payloads.items()} + ckpt = _good_ckpt() + ckpt["state_dict"] = state_dict + ckpt["metadata"]["fingerprint"] = recorded or packed_weight_fingerprint(state_dict) + return ckpt + + +def test_a_flipped_byte_in_a_packed_weight_is_refused_and_named(monkeypatch, tmp_path): + good = _fingerprinted_ckpt( + {"blocks.0.attn1.to_q.weight": b"q0", "blocks.1.attn1.to_q.weight": b"q1"} + ) + logger = _Recorder() + assert _load(monkeypatch, tmp_path, good, logger = logger) is not None + corrupt = _fingerprinted_ckpt( + {"blocks.0.attn1.to_q.weight": b"q0", "blocks.1.attn1.to_q.weight": b"q1"} + ) + corrupt["state_dict"]["blocks.1.attn1.to_q.weight"].qdata = _Bytes(b"Q1") + logger = _Recorder() + assert _load(monkeypatch, tmp_path, corrupt, logger = logger) is None + assert "blocks.1.attn1.to_q.weight" in logger.text + assert "blocks.0.attn1.to_q.weight" not in logger.text + + +def test_a_checkpoint_without_a_fingerprint_block_still_loads(monkeypatch, tmp_path): + ckpt = _good_ckpt() + ckpt["state_dict"] = {"blocks.0.attn1.to_q.weight": Float8Tensor(b"q0")} + assert _load(monkeypatch, tmp_path, ckpt) is not None + ckpt["metadata"]["fingerprint"] = {"algo": "md5-packed-v1", "count": 0, "modules": {}} + assert _load(monkeypatch, tmp_path, ckpt) is not None + + +def test_the_fingerprint_mode_env_picks_how_much_is_checked(monkeypatch, tmp_path): + from core.inference.diffusion_prequant import ( + FINGERPRINT_MODE_ENV, + _fingerprint_sampled, + ) + + sampled = [f"blocks.{i}.attn1.to_q.weight" for i in range(64)] + checked = [fqn for fqn in sampled if _fingerprint_sampled(fqn)] + assert 0 < len(checked) < len(sampled) + assert checked == [fqn for fqn in sampled if _fingerprint_sampled(fqn)] + + payloads = {fqn: fqn.encode("utf-8") for fqn in sampled} + corrupt = _fingerprinted_ckpt(payloads) + missed = next(fqn for fqn in sampled if not _fingerprint_sampled(fqn)) + corrupt["state_dict"][checked[0]].qdata = _Bytes(b"flipped") + monkeypatch.setenv(FINGERPRINT_MODE_ENV, "sample") + assert _load(monkeypatch, tmp_path, corrupt) is None + only_missed = _fingerprinted_ckpt(payloads) + only_missed["state_dict"][missed].qdata = _Bytes(b"flipped") + assert _load(monkeypatch, tmp_path, only_missed) is not None + monkeypatch.setenv(FINGERPRINT_MODE_ENV, "full") + assert _load(monkeypatch, tmp_path, only_missed) is None + monkeypatch.setenv(FINGERPRINT_MODE_ENV, "off") + assert _load(monkeypatch, tmp_path, only_missed) is not None + monkeypatch.setenv(FINGERPRINT_MODE_ENV, "sometimes") + assert _load(monkeypatch, tmp_path, only_missed) is None + + +class _CountedBytes(_Bytes): + """A payload that records the fqn whose bytes were actually read.""" + + def __init__(self, fqn, payload, hashed): + super().__init__(payload) + self._fqn = fqn + self._hashed = hashed + + def numpy(self): + self._hashed.append(self._fqn) + return super().numpy() + + +def test_a_class_whose_payload_slots_all_read_none_is_not_fingerprinted(): + from core.inference.diffusion_prequant import packed_weight_fingerprint + + renamed = Float8Tensor(b"q0") + renamed.qdata = None + renamed.scale = None + fingerprint = packed_weight_fingerprint({"blocks.0.attn1.to_q.weight": renamed}) + assert fingerprint["modules"] == {} + assert fingerprint["skipped"] == ["blocks.0.attn1.to_q.weight"] + assert fingerprint["count"] == 0 + + +class Int8Tensor: + """torchao 0.18's int8 weight, as far as the fingerprint is concerned.""" + + def __init__(self, qdata): + self.qdata = _Bytes(qdata) + self.scale = _Bytes(b"scale") + self.zero_point = None + self.act_quant_scale = None + self.act_quant_zero_point = None + self.act_pre_scale = None + + +def test_torchao_018_int8_weights_are_fingerprinted(): + from core.inference.diffusion_prequant import packed_weight_fingerprint + + fqn = "blocks.0.attn1.to_q.weight" + good = packed_weight_fingerprint({fqn: Int8Tensor(b"q0")}) + assert good["count"] == 1 and good["skipped"] == [] + flipped = packed_weight_fingerprint({fqn: Int8Tensor(b"q1")}) + assert flipped["modules"][fqn] != good["modules"][fqn] + + +def test_real_torchao_int8_round_trip_fingerprint_catches_corruption(): + torch = pytest.importorskip("torch") + pytest.importorskip("torchao") + import io + + from torchao.quantization import quantize_ + + from core.inference.diffusion_prequant import ( + _verify_packed_fingerprint, + packed_weight_fingerprint, + ) + from core.inference.diffusion_transformer_quant import _make_quant_config + + torch.manual_seed(0) + model = torch.nn.Sequential(torch.nn.Linear(64, 64, bias = False)).to(torch.bfloat16) + quantize_(model, _make_quant_config("int8")) + fingerprint = packed_weight_fingerprint(model.state_dict()) + assert fingerprint["count"] == 1, fingerprint + assert fingerprint["skipped"] == [] + + buf = io.BytesIO() + torch.save(model.state_dict(), buf) + buf.seek(0) + reloaded = torch.load(buf, weights_only = False) + assert _verify_packed_fingerprint(reloaded, {"fingerprint": fingerprint}) + + weight = reloaded["0.weight"] + inner = getattr(weight, "qdata", None) + if inner is None: # torchao <= 0.17: the AffineQuantizedTensor chain + inner = weight.original_weight_tensor.tensor_impl.int_data + inner.view(-1)[0] ^= 1 + assert not _verify_packed_fingerprint(reloaded, {"fingerprint": fingerprint}) + + +def test_sample_mode_hashes_only_the_weights_it_compares(monkeypatch, tmp_path): + from core.inference.diffusion_prequant import FINGERPRINT_MODE_ENV, _fingerprint_sampled + + fqns = [f"blocks.{i}.attn1.to_q.weight" for i in range(64)] + ckpt = _fingerprinted_ckpt({fqn: fqn.encode("utf-8") for fqn in fqns}) + hashed: list = [] + for fqn, tensor in ckpt["state_dict"].items(): + tensor.qdata = _CountedBytes(fqn, fqn.encode("utf-8"), hashed) + tensor.scale = _CountedBytes(fqn, b"scale", hashed) + monkeypatch.setenv(FINGERPRINT_MODE_ENV, "sample") + assert _load(monkeypatch, tmp_path, ckpt) is not None + assert sorted(set(hashed)) == sorted(fqn for fqn in fqns if _fingerprint_sampled(fqn)) + + +def test_a_build_that_recognises_no_quantized_weight_loads_unverified(monkeypatch, tmp_path): + ckpt = _fingerprinted_ckpt({"blocks.0.attn1.to_q.weight": b"q0"}) + ckpt["state_dict"] = {"blocks.0.attn1.to_q.weight": object()} + assert _load(monkeypatch, tmp_path, ckpt) is not None + + # ── load_prequantized_transformer ──────────────────────────────────────────────── class _FakeTransformer: calls: dict = {} @@ -380,6 +611,7 @@ def _stub_torch_accelerate( seen["safe_globals"] = list(entries) torch.load = _load + torch.uint8 = _UINT8 # A stub without this namespace would let a regression to an unrestricted load pass silently. torch.serialization = types.SimpleNamespace(add_safe_globals = _add_safe_globals) monkeypatch.setitem(sys.modules, "torch", torch) @@ -413,6 +645,8 @@ def _load( exists = True, allow_local = True, fast_accum = None, + component = None, + logger = None, ): _FakeTransformer.calls = {} _stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises) @@ -434,7 +668,8 @@ def _load( hf_token = None, scheme = scheme, fast_accum = fast_accum, - logger = None, + component = component, + logger = logger, ) @@ -590,6 +825,38 @@ def test_load_require_bf16_nvfp4_true_is_none(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, ckpt, scheme = "nvfp4") is None +def test_load_require_divisible_mismatch_is_none(monkeypatch, tmp_path): + ckpt = _good_ckpt(scheme = "fp8") + ckpt["metadata"]["require_divisible"] = 0 + assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is None + + +def test_load_require_divisible_match_ok(monkeypatch, tmp_path): + ckpt = _good_ckpt(scheme = "fp8") + ckpt["metadata"]["require_divisible"] = 16 + assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is not None + + +def test_load_require_divisible_absent_is_accepted(monkeypatch, tmp_path): + ckpt = _good_ckpt(scheme = "fp8") + assert "require_divisible" not in ckpt["metadata"] + assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is not None + + +def test_load_component_mismatch_is_none(monkeypatch, tmp_path): + ckpt = _good_ckpt() + ckpt["metadata"]["component"] = "transformer_2" + assert _load(monkeypatch, tmp_path, ckpt, component = "transformer") is None + assert _load(monkeypatch, tmp_path, ckpt, component = "transformer_2") is not None + + +def test_load_component_absent_is_accepted(monkeypatch, tmp_path): + ckpt = _good_ckpt() + assert _load(monkeypatch, tmp_path, ckpt, component = "transformer") is not None + ckpt["metadata"]["component"] = "transformer_2" + assert _load(monkeypatch, tmp_path, ckpt) is not None + + def test_resolve_checkpoint_path_expands_user(monkeypatch, tmp_path): # The allowlist gate expands ~, so the existence check must too, or a "~/..." checkpoint that passed the gate is silently skipped. import os @@ -1890,6 +2157,57 @@ def test_the_floor_check_ignores_dense_and_unreadable_state_dicts(): assert pq._fp8_activation_floor_present({}, None) is True +def test_the_checkpoint_is_released_before_the_device_copy(monkeypatch, tmp_path): + """assign=True shares the checkpoint's tensors: unreference them before ``.to(device)``.""" + import weakref + + seen: dict = {} + + class _StateDict(dict): + """A weak-referenceable state dict: plain dicts cannot be weakly referenced.""" + + class _ReleaseProbe(_FakeTransformer): + def load_state_dict( + self, + sd, + strict = True, + assign = False, + ): + _FakeTransformer.calls["load_state_dict"] = {"strict": strict, "assign": assign} + seen["state_dict"] = weakref.ref(sd) + + def to(self, device): + seen["alive_at_move"] = seen["state_dict"]() is not None + return super().to(device) + + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, None) + monkeypatch.setattr( + pq, + "_torch_load_prequant", + lambda path, **kwargs: { + "format": PREQUANT_FORMAT, + "metadata": {"scheme": "int8", "base_model_id": "Tongyi-MAI/Z-Image-Turbo"}, + "state_dict": _StateDict(weight = object()), + }, + ) + monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path)) + path = tmp_path / "ckpt.pt" + path.write_bytes(b"x") + + out = load_prequantized_transformer( + _ReleaseProbe, + "Tongyi-MAI/Z-Image-Turbo", + PrequantSource(kind = "path", location = str(path), filename = None), + device = "cuda", + dtype = "bfloat16", + scheme = "int8", + ) + + assert out is not None + assert seen["alive_at_move"] is False + + def test_the_plan_only_commits_to_a_hosted_name_this_install_can_open(monkeypatch): """The call that drops the released dense shards must ask both questions, not one. diff --git a/studio/backend/tests/test_diffusion_quant_pad.py b/studio/backend/tests/test_diffusion_quant_pad.py index 72d3cd20b5..93c352af46 100644 --- a/studio/backend/tests/test_diffusion_quant_pad.py +++ b/studio/backend/tests/test_diffusion_quant_pad.py @@ -23,11 +23,13 @@ from core.inference.diffusion_quant_pad import ( # noqa: E402 DEFAULT_PAD_TO, INT_MM_MIN_M, PadToMinM, + ZeroRowSafeLinear, activation_granularity_is_per_row, is_quantized_linear, matching_linear_fqns, padding_is_bitwise_exact, wrap_small_m_linears, + wrap_zero_row_linears, ) @@ -306,6 +308,86 @@ def test_is_quantized_linear_only_accepts_a_torchao_weight(): assert is_quantized_linear(nn.LayerNorm(8)) is False +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize("shape", [(1, 0, 8), (0, 8), (2, 0, 8)]) +def test_an_empty_activation_comes_back_at_the_projected_width(bias, shape): + inner = _RecordingLinear(8, 6, bias = bias).to(torch.bfloat16) + wrapped = ZeroRowSafeLinear(inner) + x = torch.zeros(*shape, dtype = torch.bfloat16) + with torch.no_grad(): + got, want = wrapped(x), torch.nn.functional.linear(x, inner.weight, inner.bias) + assert got.shape == (*shape[:-1], 6) == want.shape + assert got.dtype == want.dtype == torch.bfloat16 + assert torch.equal(got, want) + assert not hasattr(inner, "seen"), "the inner GEMM must not be called with an empty input" + + +def test_the_guard_is_inert_on_a_non_empty_activation(): + torch.manual_seed(0) + inner = _RecordingLinear(8, 6) + wrapped = ZeroRowSafeLinear(inner) + for shape in ((1, 8), (5, 8), (2, 7, 8)): + x = torch.randn(*shape) + with torch.no_grad(): + assert torch.equal(wrapped(x), inner(x)) + assert inner.seen.shape == x.shape + + +def test_the_guard_passes_attributes_and_keys_through(): + plain, guarded = _Tiny(), _Tiny() + guarded.context_embedder = ZeroRowSafeLinear(guarded.context_embedder) + assert guarded.context_embedder.in_features == 8 + assert guarded.context_embedder.out_features == 6 + assert guarded.context_embedder.weight is guarded.context_embedder.inner.weight + with pytest.raises(AttributeError): + guarded.context_embedder.definitely_not_a_linear_attribute + + assert sorted(guarded.state_dict()) == sorted(plain.state_dict()) + assert "context_embedder.inner.weight" not in guarded.state_dict() + assert plain.load_state_dict(dict(guarded.state_dict()), strict = True) + assert guarded.load_state_dict(dict(plain.state_dict()), strict = True) + assert torch.equal(guarded.context_embedder.weight, plain.context_embedder.weight) + + +def test_wrapping_for_zero_rows_is_idempotent_and_skips_dense_linears(): + model = _Tiny() + assert wrap_zero_row_linears(model, ["context_embedder"]) == () + assert isinstance(model.context_embedder, nn.Linear) + + model.context_embedder = _fake_quantized_linear() + assert wrap_zero_row_linears(model, ["context_embedder"]) == ("context_embedder",) + assert isinstance(model.context_embedder, ZeroRowSafeLinear) + assert wrap_zero_row_linears(model, ["context_embedder"]) == () + assert not isinstance(model.context_embedder.inner, ZeroRowSafeLinear) + + +def test_the_two_wrappers_never_stack(): + model = _Tiny() + model.context_embedder = _fake_quantized_linear() + assert wrap_small_m_linears(model, ["context_embedder"]) == ("context_embedder",) + assert wrap_zero_row_linears(model, ["context_embedder"]) == () + assert isinstance(model.context_embedder, PadToMinM) + assert model.context_embedder(torch.randn(0, 8)).shape == (0, 6) + + +def test_the_guard_list_is_read_with_the_same_substring_rule(): + class _Projection(nn.Module): + def __init__(self): + super().__init__() + self.linear_1 = _fake_quantized_linear(8, 8) + self.linear_2 = _fake_quantized_linear(8, 6) + + def forward(self, x): + return self.linear_2(self.linear_1(x)) + + model = _Tiny() + model.image_embedder = _Projection() + fqns = matching_linear_fqns(model, ("image_embedder",)) + assert fqns == ("image_embedder.linear_1", "image_embedder.linear_2") + assert wrap_zero_row_linears(model, fqns) == fqns + assert model.image_embedder(torch.zeros(1, 0, 8)).shape == (1, 0, 6) + + # ── the real thing ──────────────────────────────────────────────────────────── @@ -324,3 +406,34 @@ def test_int8_padding_is_bitwise_exact_on_a_real_quantized_linear(): assert activation_granularity_is_per_row(lin) is True for m in (1, 10, 13, 16, 17, 19, 64): assert padding_is_bitwise_exact(lin, m), f"padding changed the kept rows at M = {m}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason = "nvfp4 dynamic quant needs Blackwell") +def test_the_zero_row_guard_on_a_real_nvfp4_linear(): + """A real nvfp4 Linear survives a zero-row call once wrapped.""" + pytest.importorskip("torchao") + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("nvfp4 needs sm_100+") + from torchao.quantization import quantize_ + + from core.inference.diffusion_transformer_quant import _make_quant_config, make_filter_fn + + lin = nn.Linear(1152, 3072, dtype = torch.bfloat16).cuda().eval() + quantize_(lin, _make_quant_config("nvfp4"), filter_fn = make_filter_fn(0)) + assert is_quantized_linear(lin) + + empty = torch.zeros(1, 0, 1152, dtype = torch.bfloat16, device = "cuda") + with pytest.raises(RuntimeError, match = "numel"): + lin(empty) + + guarded = ZeroRowSafeLinear(lin) + out = guarded(empty) + assert out.shape == (1, 0, 3072) and out.dtype == torch.bfloat16 + + torch.manual_seed(0) + for m in (1, 5, 7): + x = torch.randn(1, m, 1152, dtype = torch.bfloat16, device = "cuda") + with torch.no_grad(): + got = guarded(x) + assert torch.equal(got, lin(x)) + assert torch.isfinite(got).all() diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 6308a8aae9..d24ebc7114 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -1375,6 +1375,107 @@ def test_apply_small_m_padding_is_inert_without_a_pad_list(monkeypatch): apply_small_m_padding(object(), TQ_INT8, "minimax-h3") +def test_the_zero_row_guard_is_nvfp4_only(): + from core.inference.diffusion_transformer_quant import zero_row_tokens_for_scheme + for scheme in (TQ_FP8, TQ_INT8, TQ_MXFP8, "auto"): + for family in ("hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"): + assert zero_row_tokens_for_scheme(scheme, family) == (), (scheme, family) + + +def test_both_hunyuan_tiers_guard_their_trimmable_streams(): + from core.inference.diffusion_transformer_quant import zero_row_tokens_for_scheme + for family in ("hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"): + tokens = zero_row_tokens_for_scheme(TQ_NVFP4, family) + assert "image_embedder" in tokens + assert "context_embedder_2" in tokens + + +def test_an_unlisted_family_has_no_zero_row_guard(): + """Scoped to the families whose attention path is measured to produce an empty activation.""" + from core.inference.diffusion_transformer_quant import zero_row_tokens_for_scheme + for family in ("z-image", "qwen-image", "wan2.2-ti2v-5b", "minimax-h3", None): + assert zero_row_tokens_for_scheme(TQ_NVFP4, family) == () + + +def test_the_guard_list_does_not_double_up_with_the_pad_or_exclude_lists(): + from core.inference.diffusion_transformer_quant import ( + _NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS, + exclude_tokens_for_scheme, + pad_tokens_for_scheme, + zero_row_tokens_for_scheme, + ) + for family in _NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS: + guarded = zero_row_tokens_for_scheme(TQ_NVFP4, family) + assert guarded + assert exclude_tokens_for_scheme(TQ_NVFP4, family) == () + assert pad_tokens_for_scheme(TQ_NVFP4, family) == () + + +def test_quantize_transformer_guards_after_padding(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_NVFP4}) + order = [] + tqz = types.ModuleType("torchao.quantization") + tqz.quantize_ = lambda module, config, filter_fn = None: order.append("quantize_") + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg") + monkeypatch.setattr( + tq, + "apply_small_m_padding", + lambda transformer, scheme, family = None, logger = None: (order.append("pad") or ()), + ) + monkeypatch.setattr( + tq, + "apply_zero_row_guard", + lambda transformer, scheme, family = None, logger = None: ( + order.append(("guard", scheme, family)) or () + ), + ) + transformer = types.SimpleNamespace() + pipe = types.SimpleNamespace(transformer = transformer) + assert ( + quantize_transformer(pipe, _target(), mode = "nvfp4", family = "hunyuanvideo-1.5") == TQ_NVFP4 + ) + assert order == ["quantize_", "pad", ("guard", TQ_NVFP4, "hunyuanvideo-1.5")] + assert transformer._unsloth_runtime_quant == TQ_NVFP4 + + +def test_a_guard_failure_fails_the_whole_quantise(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_NVFP4}) + tqz = types.ModuleType("torchao.quantization") + tqz.quantize_ = lambda module, config, filter_fn = None: None + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg") + + def _boom( + transformer, + scheme, + family = None, + logger = None, + ): + raise RuntimeError("cannot resolve the guard list") + + monkeypatch.setattr(tq, "apply_zero_row_guard", _boom) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert quantize_transformer(pipe, _target(), mode = "nvfp4", family = "hunyuanvideo-1.5") is None + assert not hasattr(pipe.transformer, "_unsloth_runtime_quant") + + +def test_apply_zero_row_guard_is_inert_without_a_guard_list(monkeypatch): + """apply_zero_row_guard is inert, and imports nothing, without a guard list.""" + from core.inference.diffusion_transformer_quant import apply_zero_row_guard + + stub = types.ModuleType("core.inference.diffusion_quant_pad") # no names to import + monkeypatch.setitem(sys.modules, "core.inference.diffusion_quant_pad", stub) + + assert apply_zero_row_guard(object(), TQ_NVFP4, "z-image") == () + assert apply_zero_row_guard(object(), TQ_FP8, "hunyuanvideo-1.5") == () + assert apply_zero_row_guard(object(), TQ_NVFP4, None) == () + with pytest.raises(ImportError): + apply_zero_row_guard(object(), TQ_NVFP4, "hunyuanvideo-1.5") + + def test_the_training_deny_is_a_superset_of_the_inference_deny(): # The two tables are separate because rendering evidence is not training evidence, but the # relationship must only ever go one way: anything inference refuses, training refuses too. @@ -1437,6 +1538,158 @@ def test_the_candidate_list_agrees_with_the_selector_on_the_winner(monkeypatch): assert (candidates[0] if candidates else None) == chosen, (cc, family) +def _prefer( + monkeypatch, + row, + family = "fake-video", +): + """Install one ``_FAMILY_AUTO_PREFER`` row.""" + monkeypatch.setattr(tq, "_FAMILY_AUTO_PREFER", {family: row}) + + +def _nvfp4_head(**kw): + return tq._AutoPrefer(floor = (10, 0), schemes = (TQ_NVFP4,), **kw) + + +def test_a_prefer_row_leads_the_tier_on_datacenter_blackwell(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8}) + _prefer(monkeypatch, _nvfp4_head()) + assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_NVFP4 + assert tq.auto_scheme_candidates(_target(), "fake-video") == ( + TQ_NVFP4, + TQ_INT8, + TQ_FP8, + TQ_MXFP8, + ) + assert select_transformer_quant_scheme(_target(), "auto", family = "z-image") == TQ_INT8 + assert tq.auto_scheme_candidates(_target()) == (TQ_INT8, TQ_FP8, TQ_MXFP8) + + +def test_a_prefer_row_is_dropped_below_its_capability_floor(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA L40S") + _allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_INT8}) + _prefer(monkeypatch, _nvfp4_head()) + assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_INT8 + assert tq.auto_scheme_candidates(_target(), "fake-video") == (TQ_INT8, TQ_FP8) + + +def test_a_head_the_probe_rejects_falls_through_to_the_tier(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_MXFP8, TQ_INT8}) + _prefer(monkeypatch, _nvfp4_head()) + assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_INT8 + assert tq.auto_scheme_candidates(_target(), "fake-video") == (TQ_INT8, TQ_FP8, TQ_MXFP8) + + +def test_the_deny_list_outranks_a_prefer_row(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8}) + _prefer(monkeypatch, _nvfp4_head(), family = "qwen-image") + assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8 + assert TQ_NVFP4 not in tq.auto_scheme_candidates(_target(), "qwen-image") + + +def test_a_prefer_row_needs_consumer_ok_to_apply_to_a_consumer_gpu(monkeypatch): + _stub_torch(monkeypatch, cc = (12, 0), device_name = "NVIDIA GeForce RTX 5090") + _allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8}) + _prefer(monkeypatch, _nvfp4_head()) + assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_INT8 + _prefer(monkeypatch, _nvfp4_head(consumer_ok = True)) + assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_NVFP4 + assert tq.auto_scheme_candidates(_target(), "fake-video") == ( + TQ_NVFP4, + TQ_INT8, + TQ_FP8, + TQ_MXFP8, + ) + + +def test_a_head_already_in_the_tier_is_not_listed_twice(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_MXFP8, TQ_INT8}) + _prefer(monkeypatch, tq._AutoPrefer(floor = (10, 0), schemes = (TQ_INT8,))) + assert tq.auto_scheme_candidates(_target(), "fake-video") == (TQ_INT8, TQ_FP8, TQ_MXFP8) + + +def test_a_capability_below_every_tier_has_no_order_even_with_a_head(monkeypatch): + _stub_torch(monkeypatch, cc = (7, 5)) + _allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_INT8}) + _prefer(monkeypatch, tq._AutoPrefer(floor = (7, 0), schemes = (TQ_NVFP4,))) + assert tq._auto_scheme_order("fake-video", "cuda", (7, 5)) == () + assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") is None + assert tq.auto_scheme_candidates(_target(), "fake-video") == () + + +@pytest.mark.parametrize("family", [None, "fake-video", "z-image", "qwen-image"]) +@pytest.mark.parametrize( + "allowed", + [ + {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8}, + {TQ_FP8, TQ_MXFP8, TQ_INT8}, + {TQ_NVFP4, TQ_INT8}, + {TQ_INT8}, + set(), + ], +) +def test_the_candidate_head_stays_the_selector_winner_with_a_prefer_row( + monkeypatch, family, allowed +): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, allowed) + _prefer(monkeypatch, _nvfp4_head()) + candidates = tq.auto_scheme_candidates(_target(), family) + chosen = select_transformer_quant_scheme(_target(), "auto", family = family) + assert (candidates[0] if candidates else None) == chosen, (family, allowed) + + +def test_the_shipped_prefer_table_keeps_the_ladder_as_it_was(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + _allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8}) + assert tq._FAMILY_AUTO_PREFER == {} + for family in (None, "hunyuanvideo-1.5", "hunyuanvideo-1.5-720p", "wan2.2-ti2v-5b"): + assert select_transformer_quant_scheme(_target(), "auto", family = family) == TQ_INT8 + + +def test_divisible_for_scheme_matches_each_gemm(): + """scaled_mm (fp8, and torchao's nvfp4 path) needs 16-aligned dims and MX block scaling 32.""" + from core.inference.diffusion_transformer_quant import divisible_for_scheme + + assert divisible_for_scheme(TQ_FP8) == 16 + assert divisible_for_scheme(TQ_NVFP4) == 16 + assert divisible_for_scheme(TQ_MXFP8) == 32 + assert divisible_for_scheme(TQ_INT8) == 0 + assert divisible_for_scheme("auto") == 0 + + +def test_quantize_transformer_filters_on_the_scheme_alignment(monkeypatch): + torch_stub = _stub_torch(monkeypatch, cc = (10, 0)) + + class _Linear: + def __init__(self, in_features, out_features): + self.in_features, self.out_features = in_features, out_features + self.weight = types.SimpleNamespace(dtype = torch_stub.bfloat16) + + torch_stub.nn = types.SimpleNamespace(Linear = _Linear) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg") + tqz = types.ModuleType("torchao.quantization") + seen: list = [] + tqz.quantize_ = lambda module, config, filter_fn = None: seen.append(filter_fn) + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + + filters = {} + for scheme in (TQ_NVFP4, TQ_MXFP8): + _allow(monkeypatch, {scheme}) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert quantize_transformer(pipe, _target(), mode = scheme) == scheme + filters[scheme] = seen[-1] + + aligned_16 = _Linear(528, 2048) # 528 % 16 == 0, 528 % 32 == 16 + assert filters[TQ_NVFP4](aligned_16, "image_embedder.linear_2") is True + assert filters[TQ_MXFP8](aligned_16, "image_embedder.linear_2") is False + assert filters[TQ_NVFP4](_Linear(520, 2048), "blocks.0.ff.net.0") is False + + def test_the_pre_eviction_gate_does_not_refuse_on_an_indeterminate_probe(monkeypatch): # The route-level precision gate asks the selector, not the probe, so unproven_ok has to reach # through select_transformer_quant_scheme for the leniency to exist where it matters. diff --git a/studio/backend/tests/test_nvfp4_diffusion_flag.py b/studio/backend/tests/test_nvfp4_diffusion_flag.py new file mode 100644 index 0000000000..62bc8e34f7 --- /dev/null +++ b/studio/backend/tests/test_nvfp4_diffusion_flag.py @@ -0,0 +1,621 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The NVFP4 switch (``UNSLOTH_NVFP4_DIFFUSION``), default OFF.""" + +from __future__ import annotations + +import asyncio +import inspect +import types + +import pytest + +from core.inference import diffusion_nvfp4_flag as flag +from core.inference import diffusion_transformer_quant as tq + +ENV = "UNSLOTH_NVFP4_DIFFUSION" +DISABLED = "NVFP4 is disabled in this build" + + +@pytest.fixture(autouse = True) +def _default_off(monkeypatch): + monkeypatch.delenv(ENV, raising = False) + yield + + +def _enable(monkeypatch): + monkeypatch.setenv(ENV, "1") + + +def test_the_default_is_off_and_lives_in_one_constant(): + assert flag.NVFP4_DIFFUSION_DEFAULT is False + assert flag.NVFP4_DIFFUSION_ENV == ENV + assert flag.nvfp4_diffusion_enabled() is False + + +@pytest.mark.parametrize("requested", ["auto", "fp8", "int8"]) +def test_a_video_family_without_a_usable_hosted_seed_spends_no_smoke_probe(monkeypatch, requested): + import torch + + import core.inference.video as vid + from core.inference.video_families import _FAMILIES + + calls = [] + monkeypatch.setattr( + tq, + "_child_probe_table", + lambda card: calls.append(card) or {s: None for s in tq.TQ_SCHEMES}, + ) + monkeypatch.setattr( + tq, "_smoke_probe", lambda s, d, unproven_ok = False: calls.append(s) or unproven_ok + ) + monkeypatch.setattr(tq, "dense_transformer_supported", lambda t: True) + monkeypatch.setattr(tq, "_capability", lambda: (10, 0)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + target = types.SimpleNamespace(device = "cuda", dtype = torch.bfloat16, ordinal = 0) + for fam in _FAMILIES: + if getattr(fam, "modular_workflow", None): + continue + assert ( + vid._video_auto_denoiser_scheme( + fam, target = target, requested = requested, base_repo = fam.base_repo + ) + is None + ) + assert calls == [] + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", " yes ", "on", "On"]) +def test_truthy_values_enable(monkeypatch, value): + monkeypatch.setenv(ENV, value) + assert flag.nvfp4_diffusion_enabled() is True + + +@pytest.mark.parametrize("value", ["0", "false", "no", "off", "", " ", "garbage", "2", "enable"]) +def test_everything_else_disables(monkeypatch, value): + monkeypatch.setenv(ENV, value) + assert flag.nvfp4_diffusion_enabled() is False + + +def test_flipping_the_default_is_the_one_line_change(monkeypatch): + """Unset (and unrecognised) follow the constant; an explicit 0 still turns it off.""" + monkeypatch.setattr(flag, "NVFP4_DIFFUSION_DEFAULT", True) + assert flag.nvfp4_diffusion_enabled() is True + monkeypatch.setenv(ENV, "0") + assert flag.nvfp4_diffusion_enabled() is False + monkeypatch.setenv(ENV, "garbage") + assert flag.nvfp4_diffusion_enabled() is True + + +def test_the_switch_is_read_at_call_time(monkeypatch): + assert flag.nvfp4_blocked("nvfp4") + _enable(monkeypatch) + assert not flag.nvfp4_blocked("nvfp4") + monkeypatch.setenv(ENV, "0") + assert flag.nvfp4_blocked("NVFP4") + + +def test_only_nvfp4_is_blocked(): + for scheme in ("int8", "fp8", "mxfp8", "fp8_dynamic", "auto", None, ""): + assert not flag.nvfp4_blocked(scheme) + assert flag.without_nvfp4(("int8", "nvfp4", "fp8")) == ("int8", "fp8") + assert flag.nvfp4_repo_blocked("unsloth/Wan2.2-TI2V-5B-NVFP4") + assert not flag.nvfp4_repo_blocked("unsloth/Wan2.2-TI2V-5B-FP8") + assert not flag.nvfp4_repo_blocked(None) + + +def _blackwell(monkeypatch, *, supported = ("int8", "fp8", "nvfp4", "mxfp8")): + """A datacenter sm_100 host whose smoke probe passes ``supported``, nothing allocated.""" + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(tq, "_capability", lambda: (10, 0)) + monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda device = None: False) + monkeypatch.setattr(tq, "_child_probe_table", lambda device: None) + monkeypatch.setattr( + tq, "_smoke_probe", lambda scheme, device, unproven_ok = False: scheme in supported + ) + monkeypatch.setattr(tq, "_SMOKE_CACHE", {}) + monkeypatch.setattr( + tq, + "_AUTO_LADDER", + (((10, 0), (tq.TQ_INT8, tq.TQ_FP8, tq.TQ_NVFP4, tq.TQ_MXFP8)),), + ) + monkeypatch.setattr( + tq, + "_FAMILY_AUTO_PREFER", + {"wan2.2-t2v-a14b": tq._AutoPrefer(floor = (10, 0), schemes = (tq.TQ_NVFP4,))}, + ) + return types.SimpleNamespace(device = "cuda", dtype = None) + + +def _auto(target, family): + """``auto`` on ``family`` with a hosted checkpoint for every scheme, so only the switch blocks nvfp4.""" + kwargs = {} + if "has_prequant" in inspect.signature(tq.select_transformer_quant_scheme).parameters: + kwargs["has_prequant"] = lambda scheme: True + return tq.select_transformer_quant_scheme(target, "auto", family = family, **kwargs) + + +def test_auto_never_offers_nvfp4(monkeypatch): + target = _blackwell(monkeypatch) + fam = "wan2.2-t2v-a14b" + assert tq.TQ_NVFP4 not in tq._auto_scheme_order(fam, "cuda", (10, 0)) + assert tq.TQ_NVFP4 not in tq.auto_scheme_candidates(target, fam) + assert tq.TQ_NVFP4 not in tq.auto_scheme_candidates_cached(target, fam) + assert _auto(target, fam) == tq.TQ_INT8 + + +def test_a_host_that_only_runs_nvfp4_has_nothing_automatic(monkeypatch): + target = _blackwell(monkeypatch, supported = ("nvfp4",)) + assert _auto(target, "wan2.2-t2v-a14b") is None + probed = [] + monkeypatch.setattr(tq, "_smoke_probe", lambda scheme, *a, **k: probed.append(scheme) or True) + assert tq._scheme_supported(tq.TQ_NVFP4, "cuda") is False + assert probed == [] + + +def test_the_child_probe_is_not_asked_about_nvfp4(monkeypatch): + asked = [] + + class _Stop(Exception): + pass + + def _spawn(*args, **kwargs): + asked.append(kwargs.get("args") or args) + raise _Stop + + import multiprocessing as mp + + class _Ctx: + def Queue(self): + return types.SimpleNamespace(close = lambda: None, join_thread = lambda: None) + + def Process(self, *args, **kwargs): + return _spawn(*args, **kwargs) + + monkeypatch.setattr(mp, "get_context", lambda kind: _Ctx()) + monkeypatch.setattr(tq, "_CHILD_PROBE_UNAVAILABLE", False) + monkeypatch.setattr(tq, "_close_probe_child", lambda proc, queue: True) + assert tq._child_probe_table("cuda") is None + assert asked, "the spawn was never attempted" + schemes = asked[0][4] + assert tq.TQ_NVFP4 not in schemes and tq.TQ_INT8 in schemes + + +def test_an_explicit_nvfp4_transformer_quant_is_refused_not_swapped(monkeypatch): + target = _blackwell(monkeypatch) + with pytest.raises(ValueError, match = DISABLED): + tq.normalize_transformer_quant("nvfp4") + with pytest.raises(ValueError, match = DISABLED): + tq.normalize_transformer_quant(" NVFP4 ") + with pytest.raises(ValueError, match = DISABLED): + tq.select_transformer_quant_scheme(target, "nvfp4", family = "z-image") + assert DISABLED in tq.explain_unusable_scheme("z-image", "nvfp4") + assert tq.normalize_transformer_quant("fp8") == tq.TQ_FP8 + assert tq.select_transformer_quant_scheme(target, "mxfp8", family = "z-image") == tq.TQ_MXFP8 + + +def test_the_precision_gates_refuse_nvfp4_even_under_the_silent_fallback(monkeypatch): + """The opt-in fallback swaps a declined scheme for another one; a disabled scheme is refused.""" + from core.inference.diffusion import DiffusionBackend + from core.inference.video import assert_video_precision_available + + monkeypatch.setenv("UNSLOTH_DIFFUSION_ALLOW_PRECISION_FALLBACK", "1") + with pytest.raises(ValueError, match = DISABLED): + DiffusionBackend.assert_precision_available( + object(), None, model_kind = "pipeline", transformer_quant = "nvfp4" + ) + with pytest.raises(ValueError, match = DISABLED): + DiffusionBackend.assert_precision_available( + object(), None, model_kind = "pipeline", text_encoder_quant = "nvfp4" + ) + with pytest.raises(ValueError, match = DISABLED): + assert_video_precision_available(None, model_kind = "pipeline", transformer_quant = "nvfp4") + from routes.inference import _assert_native_precision_unset + + with pytest.raises(ValueError, match = DISABLED): + _assert_native_precision_unset(transformer_quant = "nvfp4") + + +def _run(coro): + return asyncio.run(coro) + + +@pytest.mark.parametrize( + "field,control", + [("transformer_quant", "transformer_quant"), ("text_encoder_quant", "text_encoder_quant")], +) +def test_every_load_and_plan_route_400s_an_nvfp4_request(monkeypatch, field, control): + from fastapi import HTTPException + + from models.inference import DiffusionLoadRequest, VideoLoadRequest + from routes import inference as image_routes + from routes import video as video_routes + + touched = [] + monkeypatch.setattr( + image_routes.account_access, "managed_account", lambda: touched.append("acct") or False + ) + monkeypatch.setattr( + video_routes.account_access, "managed_account", lambda: touched.append("acct") or False + ) + image = DiffusionLoadRequest(model_path = "Tongyi-MAI/Z-Image-Turbo", **{field: "nvfp4"}) + video = VideoLoadRequest(model_path = "Wan-AI/Wan2.2-TI2V-5B-Diffusers", **{field: "nvfp4"}) + calls = [ + image_routes.diffusion_download_plan(image, current_subject = "u"), + image_routes.load_diffusion_model_gated(image, "u"), + video_routes.video_download_plan(video, current_subject = "u"), + video_routes.load_video_model_gated(video, "u"), + ] + for call in calls: + with pytest.raises(HTTPException) as exc: + _run(call) + assert exc.value.status_code == 400 + assert DISABLED in exc.value.detail + assert f"{control}='nvfp4'" in exc.value.detail + assert touched == [] + + +def test_text_encoder_nvfp4_is_refused(): + from core.inference import diffusion_precision as dp + + with pytest.raises(ValueError, match = DISABLED): + dp.normalize_te_quant("nvfp4") + with pytest.raises(ValueError, match = DISABLED): + dp.resolve_te_quant_request("nvfp4", None) + assert dp.resolve_te_quant_request(None, "nvfp4") == (None, False) + assert dp.resolve_te_quant_request("auto", "fp8") == ("fp8", True) + target = types.SimpleNamespace(device = "cuda", dtype = None) + assert dp.te_quant_supported(target, dp.TE_QUANT_NVFP4) is False + assert dp.normalize_te_quant("fp8_dynamic") == dp.TE_QUANT_FP8_DYNAMIC + + +def test_no_hosted_nvfp4_text_encoder_is_resolved(): + from core.inference.diffusion_te_prequant import family_te_prequant_repo + fam = types.SimpleNamespace( + te_prequant_repos = (("nvfp4", "text_encoder", "unsloth/Some-Model-NVFP4"),) + ) + assert family_te_prequant_repo(fam, "nvfp4", "text_encoder") is None + + +_WAN_T2V = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" + + +def _wan_a14b(): + from core.inference.video_families import detect_video_family + return detect_video_family(_WAN_T2V) + + +def test_no_hosted_nvfp4_repo_is_resolved_for_any_family(): + from core.inference.diffusion_families import family_prequant_repo + from core.inference.diffusion_prequant import resolve_prequant_source + from core.inference.video_families import ( + _FAMILIES as VIDEO_FAMILIES, + video_family_prequant_available, + video_family_prequant_repo, + video_family_prequant_resident_gb, + video_family_prequant_schemes, + ) + from core.inference.video_denoiser_prequant import denoiser_prequant_sources + + hosted = [f for f in VIDEO_FAMILIES if any(s == "nvfp4" for s, _r in f.prequant_repos)] + assert hosted, "the registry no longer hosts an NVFP4 denoiser; retarget this test" + for fam in hosted: + assert video_family_prequant_repo(fam, "nvfp4") is None + assert family_prequant_repo(fam, "nvfp4") is None + assert resolve_prequant_source(fam, "nvfp4") is None + assert resolve_prequant_source(fam, "nvfp4", path_override = "/x/local-NVFP4.pt") is None + assert denoiser_prequant_sources(fam, "nvfp4", fam.base_repo) is None + assert not video_family_prequant_available(fam, "nvfp4") + assert video_family_prequant_resident_gb(fam, "nvfp4") is None + assert "nvfp4" not in video_family_prequant_schemes(fam) + + +class _RecordingApi: + def __init__(self, repos): + self.repos = repos + self.asked: list[str] = [] + + def model_info( + self, + repo_id, + files_metadata = False, + token = None, + ): + self.asked.append(repo_id) + return types.SimpleNamespace(siblings = self.repos[repo_id]) + + +def _sib(name, size): + return types.SimpleNamespace(rfilename = name, size = size) + + +_WAN_SIBLINGS = { + _WAN_T2V: [ + _sib("model_index.json", 1000), + _sib("transformer/config.json", 1000), + _sib("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000), + _sib("transformer_2/config.json", 1000), + _sib("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000), + _sib("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000), + _sib("vae/diffusion_pytorch_model.safetensors", 500_000_000), + ], + "unsloth/Wan2.2-T2V-A14B-NVFP4": [ + _sib("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000), + _sib("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000), + ], +} + + +def _wan_plan(monkeypatch): + """A Wan2.2-A14B pipeline plan on a host whose auto selector lands on nvfp4, Hub mocked.""" + import core.inference.video as video_mod + from core.inference import diffusion_prequant + from core.inference.diffusion import DiffusionBackend + + api = _RecordingApi(_WAN_SIBLINGS) + monkeypatch.setattr("huggingface_hub.HfApi", lambda *a, **k: api) + monkeypatch.setattr( + DiffusionBackend, + "_hub_file_is_cached", + staticmethod(lambda repo_id, filename, revision = None, expected_size = None, **kw: False), + ) + monkeypatch.setattr(video_mod, "select_transformer_quant_scheme", lambda *a, **k: "nvfp4") + monkeypatch.setattr( + diffusion_prequant, "restricted_prequant_load_supported", lambda *a, **k: True + ) + monkeypatch.setattr(video_mod, "_video_seed_stays_resident", lambda fam, **kw: True) + monkeypatch.setattr( + video_mod, + "resolve_diffusion_device_target", + lambda *a, **k: types.SimpleNamespace(backend = "cuda", device = "cuda", dtype = None), + ) + from core.inference.video import VideoBackend + + plan = VideoBackend().download_plan(_WAN_T2V, model_kind = "pipeline") + staged = {f for e in plan["entries"] for f in e["files"]} + return plan, staged, api + + +def test_the_video_download_plan_asks_no_nvfp4_repo_and_keeps_the_dense_shards(monkeypatch): + plan, staged, api = _wan_plan(monkeypatch) + assert not any(repo.lower().endswith("-nvfp4") for repo in api.asked), api.asked + assert not any("NVFP4" in f for f in staged) + assert "transformer/diffusion_pytorch_model.safetensors" in staged + assert "transformer_2/diffusion_pytorch_model.safetensors" in staged + + +def test_the_video_seed_hub_probe_never_reaches_an_nvfp4_repo(): + from core.inference.video import VideoBackend + + api = _RecordingApi(_WAN_SIBLINGS) + repo, files = VideoBackend._denoiser_prequant_hub_files(_wan_a14b(), "nvfp4", _WAN_T2V, api) + assert (repo, files) == (None, []) + assert api.asked == [] + + +def test_the_image_hub_entry_never_asks_about_an_nvfp4_repo(monkeypatch): + from core.inference.diffusion import DiffusionBackend + + asked = [] + monkeypatch.setattr( + "huggingface_hub.HfApi", lambda *a, **k: types.SimpleNamespace(model_info = asked.append) + ) + source = types.SimpleNamespace( + kind = "repo", location = "unsloth/Z-Image-Turbo-NVFP4", filename = "x-NVFP4.pt" + ) + assert DiffusionBackend._prequant_source_hub_entry(source, None, scheme = "nvfp4") is None + assert DiffusionBackend._prequant_source_hub_entry(source, None) is None + assert asked == [] + + +def test_the_info_payload_and_estimates_omit_nvfp4(monkeypatch): + from core.inference.diffusion_auto_policy import estimate_dense_quant + from core.inference.diffusion_families import detect_family + from core.inference.diffusion_inference_info import family_inference_infos + + infos = family_inference_infos() + assert infos + for info in infos: + assert "nvfp4" not in info["estimated_resident_gb"] + assert "int8" in info["estimated_resident_gb"] + assert estimate_dense_quant(detect_family("Tongyi-MAI/Z-Image-Turbo"), "nvfp4") is None + assert estimate_dense_quant(detect_family("Tongyi-MAI/Z-Image-Turbo"), "int8") is not None + + +def test_api_system_reports_the_switch(monkeypatch): + import main + + assert main._nvfp4_diffusion_enabled() is False + _enable(monkeypatch) + assert main._nvfp4_diffusion_enabled() is True + + +def test_enabled_the_nvfp4_path_is_reachable_end_to_end(monkeypatch): + _enable(monkeypatch) + from core.inference import diffusion_precision as dp + from core.inference.diffusion_inference_info import family_inference_infos + from core.inference.video_denoiser_prequant import denoiser_prequant_sources + from core.inference.video_families import video_family_prequant_repo + + target = _blackwell(monkeypatch) + assert _auto(target, "wan2.2-t2v-a14b") == "nvfp4" + assert tq.normalize_transformer_quant("nvfp4") == "nvfp4" + assert tq.select_transformer_quant_scheme(target, "nvfp4", family = "z-image") == "nvfp4" + assert dp.normalize_te_quant("nvfp4") == "nvfp4" + fam = _wan_a14b() + assert video_family_prequant_repo(fam, "nvfp4") == "unsloth/Wan2.2-T2V-A14B-NVFP4" + sources = denoiser_prequant_sources(fam, "nvfp4", _WAN_T2V) + assert sources and set(sources) == {"transformer", "transformer_2"} + assert all(s.location == "unsloth/Wan2.2-T2V-A14B-NVFP4" for s in sources.values()) + assert all("nvfp4" in info["estimated_resident_gb"] for info in family_inference_infos()) + + +def test_enabled_the_video_plan_stages_the_hosted_nvfp4_denoisers(monkeypatch): + _enable(monkeypatch) + plan, staged, api = _wan_plan(monkeypatch) + assert "unsloth/Wan2.2-T2V-A14B-NVFP4" in api.asked + assert {"Wan2.2-T2V-A14B-NVFP4.pt", "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"} <= staged + assert "transformer/diffusion_pytorch_model.safetensors" not in staged + + +def test_enabled_the_routes_let_nvfp4_through_the_switch(monkeypatch): + _enable(monkeypatch) + from routes.inference import _refuse_disabled_nvfp4_request + _refuse_disabled_nvfp4_request( + types.SimpleNamespace(transformer_quant = "nvfp4", text_encoder_quant = "nvfp4") + ) + + +def _prequant_dir(root, name, scheme): + """A directory holding one tiny safetensors pre-quant artifact that records ``scheme``.""" + import json + + import torch + from safetensors.torch import save_file + + from core.inference.diffusion_prequant import PREQUANT_FORMAT + from core.inference.prequant_safetensors import UNSLOTH_FORMAT_KEY, UNSLOTH_METADATA_KEY + + folder = root / name + folder.mkdir() + save_file( + {"w": torch.zeros(2)}, + str(folder / "model.safetensors"), + metadata = { + UNSLOTH_FORMAT_KEY: PREQUANT_FORMAT, + UNSLOTH_METADATA_KEY: json.dumps({"scheme": scheme}), + }, + ) + return folder + + +def _checkpoint_route_calls(image_path, video_path, monkeypatch): + from models.inference import DiffusionLoadRequest, VideoLoadRequest + from routes import inference as image_routes + from routes import video as video_routes + + monkeypatch.setattr(image_routes.account_access, "managed_account", lambda: False) + monkeypatch.setattr(video_routes.account_access, "managed_account", lambda: False) + monkeypatch.setattr(image_routes.account_access, "require_idle_other_accounts", lambda: None) + monkeypatch.setattr(video_routes.account_access, "require_idle_other_accounts", lambda: None) + image = DiffusionLoadRequest(model_path = image_path) + video = VideoLoadRequest(model_path = video_path) + return [ + lambda: image_routes.diffusion_download_plan(image, current_subject = "u"), + lambda: image_routes.load_diffusion_model_gated(image, "u"), + lambda: video_routes.video_download_plan(video, current_subject = "u"), + lambda: video_routes.load_video_model_gated(video, "u"), + ] + + +HOSTED_NVFP4 = [ + "unsloth/Z-Image-Turbo-NVFP4", + "unsloth/FLUX.1-schnell-NVFP4", + "unsloth/Qwen-Image-2512-NVFP4", + "unsloth/Wan2.2-TI2V-5B-NVFP4", + "unsloth/Wan2.2-T2V-A14B-NVFP4", + "unsloth/HunyuanVideo-1.5-NVFP4", +] + + +@pytest.mark.parametrize("repo", HOSTED_NVFP4) +def test_a_hosted_nvfp4_repo_as_the_model_is_refused(repo): + with pytest.raises(ValueError, match = DISABLED): + flag.refuse_disabled_nvfp4_checkpoint(repo) + + +def test_every_family_registered_nvfp4_repo_is_recognised_by_the_table(): + from core.inference.diffusion_prequant import hosted_nvfp4_repo_ids + from core.inference.video_families import _FAMILIES + + registered = { + repo.lower() + for fam in _FAMILIES + for scheme, repo in fam.prequant_repos + if scheme == "nvfp4" + } + assert registered and registered <= hosted_nvfp4_repo_ids() + + +def test_a_cached_repo_is_judged_by_its_metadata_not_its_name(tmp_path, monkeypatch): + from core.inference import diffusion_prequant as dpq + + snap = _prequant_dir(tmp_path, "snap", "nvfp4") + monkeypatch.setattr(dpq, "_cached_snapshot_dirs", lambda repo: [str(snap)]) + assert dpq.declares_nvfp4_checkpoint("someone/renamed-checkpoint") is True + with pytest.raises(ValueError, match = DISABLED): + flag.refuse_disabled_nvfp4_checkpoint("someone/renamed-checkpoint") + + +def test_a_local_dir_whose_metadata_declares_nvfp4_is_refused(tmp_path): + folder = _prequant_dir(tmp_path, "my-checkpoint", "nvfp4") + with pytest.raises(ValueError, match = DISABLED): + flag.refuse_disabled_nvfp4_checkpoint(str(folder)) + with pytest.raises(ValueError, match = DISABLED): + flag.refuse_disabled_nvfp4_checkpoint(str(folder / "model.safetensors")) + + +def test_the_routes_400_an_nvfp4_checkpoint_as_the_model(tmp_path, monkeypatch): + from fastapi import HTTPException + folder = _prequant_dir(tmp_path, "local-z-image", "nvfp4") + for image_path, video_path in ( + ("unsloth/Z-Image-Turbo-NVFP4", "unsloth/Wan2.2-TI2V-5B-NVFP4"), + (str(folder), str(folder)), + ): + for call in _checkpoint_route_calls(image_path, video_path, monkeypatch): + with pytest.raises(HTTPException) as exc: + _run(call()) + assert exc.value.status_code == 400 + assert DISABLED in exc.value.detail + assert "NVFP4 checkpoint" in exc.value.detail + assert str(tmp_path) not in exc.value.detail + + +@pytest.mark.parametrize( + "path", + [ + "Tongyi-MAI/Z-Image-Turbo", + "unsloth/Z-Image-Turbo-FP8", + "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + "unsloth/Qwen-Image-2512-GGUF", + "", + None, + ], +) +def test_non_nvfp4_models_are_not_refused(path): + flag.refuse_disabled_nvfp4_checkpoint(path) + + +def test_a_local_dir_with_another_scheme_is_not_refused(tmp_path): + folder = _prequant_dir(tmp_path, "my-fp8-checkpoint", "fp8") + flag.refuse_disabled_nvfp4_checkpoint(str(folder)) + plain = tmp_path / "plain-pipeline" + plain.mkdir() + (plain / "model_index.json").write_text("{}") + flag.refuse_disabled_nvfp4_checkpoint(str(plain)) + pipeline = _prequant_dir(tmp_path, "pipeline-with-extra", "nvfp4") + (pipeline / "model_index.json").write_text("{}") + flag.refuse_disabled_nvfp4_checkpoint(str(pipeline)) + + +def test_enabled_an_nvfp4_checkpoint_as_the_model_proceeds(tmp_path, monkeypatch): + _enable(monkeypatch) + folder = _prequant_dir(tmp_path, "my-checkpoint", "nvfp4") + for path in (*HOSTED_NVFP4, str(folder)): + flag.refuse_disabled_nvfp4_checkpoint(path) + + reached = [] + + def _stop(*args, **kwargs): + reached.append(True) + raise RuntimeError("reached the backend") + + monkeypatch.setattr("core.inference.diffusion.get_diffusion_backend", _stop) + calls = _checkpoint_route_calls("unsloth/Z-Image-Turbo-NVFP4", str(folder), monkeypatch) + with pytest.raises(RuntimeError, match = "reached the backend"): + _run(calls[0]()) + assert reached == [True] diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 7642e0db82..5c0df4690d 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -9245,6 +9245,572 @@ def test_h3_generate_non_oom_error_leaves_the_graphs_alone(fake_runtime): assert handle.resets == 0 +def _stub_denoiser_seed( + monkeypatch, + *, + scheme = "nvfp4", + components = ("transformer", "transformer_2"), + seeded = True, + repo = "unsloth/Wan2.2-T2V-A14B-NVFP4", + plan_scheme = "nvfp4", +): + """Stub the whole seeding route: plan-time scheme, registry lookup and checkpoint load.""" + import core.inference.video as video_mod + import core.inference.video_denoiser_prequant as dq + + if plan_scheme is not None: + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: plan_scheme) + source = types.SimpleNamespace( + kind = "repo", location = repo, filename = "x.pt", fallback_filename = None + ) + monkeypatch.setattr( + dq, "denoiser_prequant_sources", lambda fam, s, base: {c: source for c in components} + ) + modules = {c: object() for c in components} + calls = [] + + def _fake_pipe_kwargs(fam, base, **kwargs): + calls.append({"base": base, **kwargs}) + return dict(modules) if seeded else {} + + monkeypatch.setattr(dq, "denoiser_prequant_pipe_kwargs", _fake_pipe_kwargs) + return calls, modules + + +def test_wan_a14b_prequant_seeds_both_dits_instead_of_quantising_them(fake_runtime, monkeypatch): + """The prequant twin of ``test_wan_a14b_dense_quant_applies_to_both_dits``.""" + import core.inference.video as video_mod + + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + quantised = [] + + def _fake_quant( + view, + target, + *, + mode, + family, + logger = None, + ): + quantised.append(view.transformer) + return "int8" + + monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant) + calls, modules = _stub_denoiser_seed(monkeypatch) + + backend = VideoBackend() + status = backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + + assembled = _FakeWanPipelineSingle.last + assert assembled["transformer"] is modules["transformer"] + assert assembled["transformer_2"] is modules["transformer_2"] + assert quantised == [], "a seeded DiT must not be quantised again" + assert status["transformer_quant"] == "nvfp4" + resolved = status["resolved"]["transformer_quant"] + assert "unsloth/Wan2.2-T2V-A14B-NVFP4" in resolved["reason"] + assert resolved["requested"] == "nvfp4" and resolved["value"] == "nvfp4" + assert resolved["status"] == "applied" + assert calls and calls[0]["scheme"] == "nvfp4" + + +def test_a_checkpoint_that_will_not_load_falls_back_to_the_dense_quant(fake_runtime, monkeypatch): + import core.inference.video as video_mod + + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + quantised = [] + + def _fake_quant( + view, + target, + *, + mode, + family, + logger = None, + ): + quantised.append(view.transformer) + return "nvfp4" + + monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant) + _stub_denoiser_seed(monkeypatch, seeded = False) + + backend = VideoBackend() + status = backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + pipe = backend._state.pipe + assert quantised == [pipe.transformer, pipe.transformer_2] + assert status["transformer_quant"] == "nvfp4" + assert "unsloth/" not in status["resolved"]["transformer_quant"]["reason"] + + +def test_seeding_is_skipped_entirely_under_offload(fake_runtime, monkeypatch): + import core.inference.video as video_mod + + monkeypatch.setenv("UNSLOTH_DIFFUSION_ALLOW_PRECISION_FALLBACK", "1") + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None) + calls, _modules = _stub_denoiser_seed(monkeypatch) + real_plan = video_mod.plan_diffusion_memory + monkeypatch.setattr( + video_mod, + "plan_diffusion_memory", + lambda **kwargs: dataclasses.replace(real_plan(**kwargs), offload_policy = "model"), + ) + _stub_apply_memory_plan(monkeypatch, video_mod) + + backend = VideoBackend() + status = backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + assert calls == [], "no checkpoint may be fetched for a load that cannot keep it resident" + assert status["transformer_quant"] is None + assert "moves the DiT" in status["resolved"]["transformer_quant"]["reason"] + + +def _plans_for(monkeypatch, video_mod): + """Every ``model_dense_mib`` the memory planner was asked to judge, in order.""" + real_plan = video_mod.plan_diffusion_memory + seen: list = [] + + def _record(**kwargs): + seen.append(kwargs.get("model_dense_mib")) + return real_plan(**kwargs) + + monkeypatch.setattr(video_mod, "plan_diffusion_memory", _record) + return seen + + +def test_the_memory_plan_prices_a_seeded_denoiser_at_the_measured_row(fake_runtime, monkeypatch): + import core.inference.video as video_mod + + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None) + monkeypatch.setattr(video_mod, "video_family_prequant_resident_gb", lambda fam, scheme: 8.0) + _stub_denoiser_seed(monkeypatch) + seen = _plans_for(monkeypatch, video_mod) + + backend = VideoBackend() + backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + fam = _detect_load_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers", None, None) + mib_per_gb = 1000.0**3 / (1024.0 * 1024.0) + companions = fam.bf16_components_gb[1] + fam.bf16_components_gb[2] + assert seen == [int((8.0 + companions) * mib_per_gb)] + assert seen[0] < int((fam.bf16_components_gb[0] + companions) * mib_per_gb) + + +def test_a_failed_seed_replans_at_bf16_and_refuses_again(fake_runtime, monkeypatch): + """The artifact-sized budget is valid only once the checkpoint is in hand.""" + import core.inference.video as video_mod + + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: "nvfp4") + monkeypatch.setattr(video_mod, "video_family_prequant_resident_gb", lambda fam, scheme: 8.0) + _stub_denoiser_seed(monkeypatch, seeded = False) + seen = _plans_for(monkeypatch, video_mod) + refusals: list = [] + monkeypatch.setattr( + video_mod, + "raise_on_unified_memory_shortfall", + lambda plan, family = None, logger = None: refusals.append( + plan.model_dense_mib if hasattr(plan, "model_dense_mib") else None + ), + ) + + backend = VideoBackend() + backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + fam = _detect_load_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers", None, None) + mib_per_gb = 1000.0**3 / (1024.0 * 1024.0) + companions = fam.bf16_components_gb[1] + fam.bf16_components_gb[2] + assert seen == [ + int((8.0 + companions) * mib_per_gb), + int((fam.bf16_components_gb[0] + companions) * mib_per_gb), + ] + assert len(refusals) == 2 + + +def test_the_planned_scheme_is_what_the_load_seeds(fake_runtime, monkeypatch): + import core.inference.video as video_mod + + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None) + calls, _modules = _stub_denoiser_seed(monkeypatch, plan_scheme = None) + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: None) + + backend = VideoBackend() + status = backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + _video_auto_denoiser_planned = "nvfp4", + ) + assert [c["scheme"] for c in calls] == ["nvfp4"] + assert status["transformer_quant"] == "nvfp4" + + +def test_a_seed_the_plan_declined_is_not_re_taken_by_the_load(fake_runtime, monkeypatch): + """The load honours the plan's decline: re-deciding here would fetch the artifact inline.""" + import core.inference.video as video_mod + + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: "nvfp4") + calls, _modules = _stub_denoiser_seed(monkeypatch, plan_scheme = None) + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4") + + backend = VideoBackend() + status = backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + _video_auto_denoiser_planned = video_mod.DENOISER_SEED_DECLINED, + ) + assert calls == [], "a seed the plan declined may not be fetched inline by the load" + assert status["transformer_quant"] == "nvfp4" + assert "unsloth/" not in status["resolved"]["transformer_quant"]["reason"] + + +_A14B_SIBLINGS = [ + _sibling("model_index.json", 10), + _sibling("transformer/config.json", 1), + _sibling("transformer/diffusion_pytorch_model-00001-of-00002.safetensors", 14), + _sibling("transformer/diffusion_pytorch_model-00002-of-00002.safetensors", 14), + _sibling("transformer_2/config.json", 1), + _sibling("transformer_2/diffusion_pytorch_model-00001-of-00002.safetensors", 14), + _sibling("transformer_2/diffusion_pytorch_model-00002-of-00002.safetensors", 14), + _sibling("text_encoder/model-00001-of-00001.safetensors", 11), + _sibling("vae/diffusion_pytorch_model.safetensors", 3), + _sibling("scheduler/scheduler_config.json", 1), +] + + +def test_base_download_files_drops_both_experts_and_keeps_both_configs(): + info = types.SimpleNamespace(siblings = _A14B_SIBLINGS) + + dense = dict(VideoBackend._base_download_files(info, "pipeline")) + seeded = dict( + VideoBackend._base_download_files( + info, + "pipeline", + skip_transformer_weights = True, + skip_transformer_components = ("transformer", "transformer_2"), + ) + ) + assert not any(n.endswith(".safetensors") and n.startswith("transformer") for n in seeded) + assert seeded["transformer/config.json"] == 1 + assert seeded["transformer_2/config.json"] == 1 + assert seeded["text_encoder/model-00001-of-00001.safetensors"] == 11 + assert sum(dense.values()) - sum(seeded.values()) == 56 + + +def test_base_download_files_keeps_the_h3_partition_default(): + info = types.SimpleNamespace(siblings = _H3_SIBLINGS) + references = dict( + VideoBackend._base_download_files( + info, "pipeline", skip_transformer_weights = True, h3_task = "ref2va" + ) + ) + assert not any(n.startswith("transformer_ref/diffusion_pytorch_model") for n in references) + assert "transformer_ref/config.json" in references + + +def test_the_download_plan_stages_both_experts_artifacts(monkeypatch): + import core.inference.video as video_mod + import core.inference.video_denoiser_prequant as dq + + _plan_api( + monkeypatch, + { + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": [ + _PlanSibling("model_index.json", 1000), + _PlanSibling("transformer/config.json", 1000), + _PlanSibling("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000), + _PlanSibling("transformer_2/config.json", 1000), + _PlanSibling("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000), + _PlanSibling("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000), + _PlanSibling("vae/diffusion_pytorch_model.safetensors", 500_000_000), + ], + "unsloth/Wan2.2-T2V-A14B-NVFP4": [ + _PlanSibling("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000), + _PlanSibling("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000), + ], + }, + ) + sources = { + "transformer": types.SimpleNamespace( + kind = "repo", + location = "unsloth/Wan2.2-T2V-A14B-NVFP4", + filename = "Wan2.2-T2V-A14B-NVFP4.pt", + fallback_filename = None, + ), + "transformer_2": types.SimpleNamespace( + kind = "repo", + location = "unsloth/Wan2.2-T2V-A14B-NVFP4", + filename = "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", + fallback_filename = None, + ), + } + monkeypatch.setattr(dq, "denoiser_prequant_sources", lambda fam, scheme, base: sources) + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4") + monkeypatch.setattr(video_mod, "_video_seed_stays_resident", lambda fam, **kw: True) + + plan = VideoBackend().download_plan( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + + staged = {f for e in plan["entries"] for f in e["files"]} + assert "Wan2.2-T2V-A14B-NVFP4.pt" in staged + assert "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt" in staged + assert not any( + f.endswith("diffusion_pytorch_model.safetensors") and f.startswith("transformer") + for f in staged + ) + assert "transformer/config.json" in staged and "transformer_2/config.json" in staged + assert ( + plan["required_bytes"] + == 8_000_000_000 + 8_100_000_000 + 1000 * 3 + 11_000_000_000 + 500_000_000 + ) + + +def test_an_explicit_scheme_under_speed_off_stages_the_hosted_experts(monkeypatch): + """An EXPLICIT scheme under speed_mode="off" is still quantized, so the plan keeps its seed.""" + import core.inference.video as video_mod + import core.inference.video_denoiser_prequant as dq + + _plan_api( + monkeypatch, + { + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": [ + _PlanSibling("model_index.json", 1000), + _PlanSibling("transformer/config.json", 1000), + _PlanSibling("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000), + _PlanSibling("transformer_2/config.json", 1000), + _PlanSibling("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000), + _PlanSibling("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000), + _PlanSibling("vae/diffusion_pytorch_model.safetensors", 500_000_000), + ], + "unsloth/Wan2.2-T2V-A14B-NVFP4": [ + _PlanSibling("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000), + _PlanSibling("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000), + ], + }, + ) + sources = { + "transformer": types.SimpleNamespace( + kind = "repo", + location = "unsloth/Wan2.2-T2V-A14B-NVFP4", + filename = "Wan2.2-T2V-A14B-NVFP4.pt", + fallback_filename = None, + ), + "transformer_2": types.SimpleNamespace( + kind = "repo", + location = "unsloth/Wan2.2-T2V-A14B-NVFP4", + filename = "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", + fallback_filename = None, + ), + } + monkeypatch.setattr(dq, "denoiser_prequant_sources", lambda fam, scheme, base: sources) + monkeypatch.setattr( + video_mod, "select_transformer_quant_scheme", lambda target, requested, **kw: requested + ) + monkeypatch.setattr(video_mod, "_video_seed_stays_resident", lambda fam, **kw: True) + + plan = VideoBackend().download_plan( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + speed_mode = "off", + ) + + staged = {f for e in plan["entries"] for f in e["files"]} + assert "transformer/diffusion_pytorch_model.safetensors" not in staged + assert "transformer_2/diffusion_pytorch_model.safetensors" not in staged + assert {"Wan2.2-T2V-A14B-NVFP4.pt", "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"} <= staged + + +def _cuda_plan_target(monkeypatch, video_mod, *, free_gib): + """Point the planning path at a cuda card of ``free_gib``, off the test host's own hardware.""" + import torch + + from core.inference.diffusion_device import DiffusionDeviceTarget + from core.inference.diffusion_memory import DeviceMemory + + target = DiffusionDeviceTarget( + device = "cuda", + dtype = torch.bfloat16, + backend = "cuda", + vendor = "nvidia", + supports_model_cpu_offload = True, + supports_default_torch_compile = True, + supports_pinned_transfer = True, + ) + monkeypatch.setattr(video_mod, "resolve_diffusion_device_target", lambda **kw: target) + monkeypatch.setattr( + video_mod, + "settled_snapshot_device_memory", + lambda t, *a, **k: DeviceMemory( + backend = "cuda", + device = "cuda", + memory_kind = "discrete_vram", + free_mib = int(free_gib * 1024), + total_mib = int(free_gib * 1024), + ), + ) + + +def _a14b_plan(monkeypatch): + """The A14B repo pair the seeded-plan tests resolve against.""" + import core.inference.video_denoiser_prequant as dq + + _plan_api( + monkeypatch, + { + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": [ + _PlanSibling("model_index.json", 1000), + _PlanSibling("transformer/config.json", 1000), + _PlanSibling("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000), + _PlanSibling("transformer_2/config.json", 1000), + _PlanSibling("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000), + _PlanSibling("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000), + _PlanSibling("vae/diffusion_pytorch_model.safetensors", 500_000_000), + ], + "unsloth/Wan2.2-T2V-A14B-NVFP4": [ + _PlanSibling("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000), + _PlanSibling("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000), + ], + }, + ) + sources = { + "transformer": types.SimpleNamespace( + kind = "repo", + location = "unsloth/Wan2.2-T2V-A14B-NVFP4", + filename = "Wan2.2-T2V-A14B-NVFP4.pt", + fallback_filename = None, + ), + "transformer_2": types.SimpleNamespace( + kind = "repo", + location = "unsloth/Wan2.2-T2V-A14B-NVFP4", + filename = "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", + fallback_filename = None, + ), + } + monkeypatch.setattr(dq, "denoiser_prequant_sources", lambda fam, scheme, base: sources) + + +def test_a_plan_that_still_offloads_at_artifact_size_stages_the_dense_experts(monkeypatch): + """A card the artifact-sized plan still offloads on cannot seed: stage the dense experts.""" + import core.inference.video as video_mod + + _a14b_plan(monkeypatch) + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4") + _cuda_plan_target(monkeypatch, video_mod, free_gib = 24) + + plan = VideoBackend().download_plan( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + + staged = {f for e in plan["entries"] for f in e["files"]} + assert "transformer/diffusion_pytorch_model.safetensors" in staged + assert "transformer_2/diffusion_pytorch_model.safetensors" in staged + assert not any(f.endswith(".pt") for f in staged) + + +def test_a_card_the_artifact_fits_on_still_stages_the_artifacts(monkeypatch): + """Where the artifact-sized plan stays resident the load seeds, so dense shards stay out.""" + import core.inference.video as video_mod + + _a14b_plan(monkeypatch) + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4") + _cuda_plan_target(monkeypatch, video_mod, free_gib = 180) + + plan = VideoBackend().download_plan( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + + staged = {f for e in plan["entries"] for f in e["files"]} + assert "Wan2.2-T2V-A14B-NVFP4.pt" in staged + assert "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt" in staged + assert not any( + f.endswith("diffusion_pytorch_model.safetensors") and f.startswith("transformer") + for f in staged + ) + + +def test_an_offloading_memory_mode_stages_the_dense_experts_on_any_card(monkeypatch): + """An explicit offload memory_mode stages the dense experts even on a roomy card.""" + import core.inference.video as video_mod + + _a14b_plan(monkeypatch) + monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4") + _cuda_plan_target(monkeypatch, video_mod, free_gib = 180) + + plan = VideoBackend().download_plan( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + memory_mode = "low_vram", + ) + + staged = {f for e in plan["entries"] for f in e["files"]} + assert "transformer/diffusion_pytorch_model.safetensors" in staged + assert not any(f.endswith(".pt") for f in staged) + + +def test_a_dense_encoder_fallback_that_forces_offload_also_drops_the_seed( + fake_runtime, monkeypatch +): + """A failed pre-cast encoder re-plans at bf16, which can offload: re-decide the seed there.""" + import core.inference.diffusion_te_prequant as te + import core.inference.video as video_mod + + monkeypatch.setenv("UNSLOTH_DIFFUSION_ALLOW_PRECISION_FALLBACK", "1") + monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None) + calls, _modules = _stub_denoiser_seed(monkeypatch) + monkeypatch.setattr(te, "te_prequant_budget_scale", lambda fam, **kwargs: 0.5) + monkeypatch.setattr(te, "te_prequant_pipe_kwargs", lambda fam, base, **kwargs: {}) + real_plan = video_mod.plan_diffusion_memory + seen: list = [] + + def _plan(**kwargs): + seen.append(kwargs.get("model_dense_mib")) + planned = real_plan(**kwargs) + return planned if len(seen) == 1 else dataclasses.replace(planned, offload_policy = "model") + + monkeypatch.setattr(video_mod, "plan_diffusion_memory", _plan) + _stub_apply_memory_plan(monkeypatch, video_mod) + + backend = VideoBackend() + status = backend.load_pipeline( + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + model_kind = "pipeline", + transformer_quant = "nvfp4", + ) + + assert calls == [], "no checkpoint may be seeded into a load that will offload the DiT" + assert status["transformer_quant"] is None + + # ── Progress truth: the bar must track the GPU, not the host queue ───────────── # Measured on MiniMax-H3 (960x544x124, 30 steps, int8 DiT + int8 TE, speed=default with a captured # CUDA graph, B200): the 29 host-side scheduler.step calls land in 0.9 s and the GPU is still diff --git a/studio/backend/tests/test_video_families.py b/studio/backend/tests/test_video_families.py index a4d5e83922..ee236dd95a 100644 --- a/studio/backend/tests/test_video_families.py +++ b/studio/backend/tests/test_video_families.py @@ -412,3 +412,106 @@ def test_minimax_h3_offers_every_advertised_aspect_ratio(): assert width * height <= H3_CANVAS_MAX_PIXELS, (width, height) rule_width, rule_height = h3_canvas_for_aspect(width, height) assert width <= rule_width and height <= rule_height, (width, height) + + +def _resident_fam(**kwargs): + from core.inference.video_families import VideoFamily + + base = dict( + name = "test-video", + pipeline_class = "TestPipeline", + transformer_class = "TestTransformer3DModel", + base_repo = "org/test-video", + ) + base.update(kwargs) + return VideoFamily(**base) + + +def test_the_per_scheme_row_wins_and_the_float_is_the_fallback(): + from core.inference.video_families import video_family_prequant_resident_gb + + fam = _resident_fam( + prequant_resident_gb = 20.3, + prequant_resident_gb_by_scheme = (("nvfp4", 8.1), ("fp8", 13.6)), + ) + assert video_family_prequant_resident_gb(fam, "nvfp4") == pytest.approx(8.1) + assert video_family_prequant_resident_gb(fam, "fp8") == pytest.approx(13.6) + assert video_family_prequant_resident_gb(fam, "int8") == pytest.approx(20.3) + + +def test_an_unmeasured_family_reports_nothing_rather_than_guessing(): + import types + + from core.inference.video_families import video_family_prequant_resident_gb + + assert video_family_prequant_resident_gb(_resident_fam(), "nvfp4") is None + assert video_family_prequant_resident_gb(types.SimpleNamespace(), "nvfp4") is None + fam = _resident_fam(prequant_resident_gb_by_scheme = (("nvfp4",), ("nvfp4", "big"))) + assert video_family_prequant_resident_gb(fam, "nvfp4") is None + + +def test_the_h3_measurement_still_answers_through_the_helper(): + from core.inference.video_families import video_family_prequant_resident_gb + + fam = detect_video_family("MiniMaxAI/MiniMax-H3") + assert video_family_prequant_resident_gb(fam, "int8") == pytest.approx(fam.prequant_resident_gb) + assert video_family_prequant_resident_gb(fam, "fp8") == pytest.approx(fam.prequant_resident_gb) + + +def test_every_hosted_nvfp4_denoiser_carries_its_measured_resident_size(): + from core.inference.video_families import ( + _FAMILIES, + video_family_prequant_resident_gb, + ) + + expected = { + "wan2.2-ti2v-5b": 2.9, + "wan2.2-t2v-a14b": 16.2, + "hunyuanvideo-1.5": 4.8, + "hunyuanvideo-1.5-720p": 4.8, + } + hosting = { + fam.name + for fam in _FAMILIES + if any(scheme == "nvfp4" for scheme, _repo in fam.prequant_repos) + } + assert hosting == set(expected) + for fam in _FAMILIES: + if fam.name not in expected: + continue + assert video_family_prequant_resident_gb(fam, "nvfp4") == pytest.approx(expected[fam.name]) + + +def test_the_measured_nvfp4_size_is_a_4_bit_fraction_of_the_term_it_replaces(): + from core.inference.video_families import ( + _FAMILIES, + video_family_prequant_resident_gb, + ) + + seen = 0 + for fam in _FAMILIES: + if not any(scheme == "nvfp4" for scheme, _repo in fam.prequant_repos): + continue + measured = video_family_prequant_resident_gb(fam, "nvfp4") + if measured is None or not fam.bf16_components_gb: + continue + seen += 1 + assert ( + 0.2 * fam.bf16_components_gb[0] < measured < 0.4 * fam.bf16_components_gb[0] + ), fam.name + assert seen == 4 + + +def test_the_a14b_row_prices_both_experts_not_one(): + from core.inference.video_families import ( + detect_video_family, + video_family_prequant_resident_gb, + ) + + fam = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers") + assert fam.is_moe + names = {entry[-1] for entry in fam.prequant_filenames} + assert len(names) == 2 + measured = video_family_prequant_resident_gb(fam, "nvfp4") + assert measured == pytest.approx(16.2) + assert measured > 0.25 * fam.bf16_components_gb[0] diff --git a/studio/backend/tests/test_video_prequant.py b/studio/backend/tests/test_video_prequant.py index c1c4033693..d1d2405272 100644 --- a/studio/backend/tests/test_video_prequant.py +++ b/studio/backend/tests/test_video_prequant.py @@ -11,6 +11,8 @@ still run (and still get tested) in an environment where diffusers cannot be imp from __future__ import annotations +import types + import pytest from core.inference.diffusion_prequant import ( @@ -1118,7 +1120,7 @@ def test_the_fallback_is_resolved_before_the_download_is_planned(monkeypatch): verified = src.index("_denoiser_prequant_verified") predownload = src.index("_predownload_base") assert planned < verified < predownload - assert 'h3_auto_denoiser or kwargs.get("transformer_quant")' in src + assert "h3_auto_denoiser or video_auto_denoiser or requested_denoiser" in src def _h3_placement_probe( @@ -1176,6 +1178,496 @@ def test_an_unreadable_card_keeps_the_rotation(monkeypatch): assert vid._h3_dense_denoiser_fits(vid._h3_dense_denoiser_resident_bytes(None), None) is False +def _moe_fam(**kwargs): + """A dual-expert MoE family (Wan2.2-A14B's shape) hosting both experts in one repo.""" + return _fam( + is_moe = True, + prequant_repos = (("nvfp4", "org/test-quantized"),), + prequant_filenames = (("nvfp4", "transformer_2", "test-transformer_2-NVFP4.pt"),), + **kwargs, + ) + + +def test_the_components_a_family_seeds_are_its_denoisers(): + from core.inference.video_denoiser_prequant import denoiser_components + assert denoiser_components(_fam()) == ("transformer",) + assert denoiser_components(_fam(is_moe = True)) == ("transformer", "transformer_2") + + +def test_the_second_expert_is_addressed_through_the_task_slot(): + """The second expert is addressed through the task slot, with no filename fallback.""" + from core.inference.video_denoiser_prequant import denoiser_prequant_sources + + sources = denoiser_prequant_sources(_moe_fam(), "nvfp4", "org/test-video") + assert set(sources) == {"transformer", "transformer_2"} + from core.inference.diffusion_prequant import candidate_filenames_of + + assert candidate_filenames_of(sources["transformer"]) == ( + "test-NVFP4.safetensors", + "test-NVFP4.pt", + "transformer_nvfp4.pt", + ) + assert candidate_filenames_of(sources["transformer_2"]) == ("test-transformer_2-NVFP4.pt",) + + +def test_every_component_or_none(): + """Partial coverage is not coverage: every component or none.""" + from core.inference.video_denoiser_prequant import denoiser_prequant_sources + + single = _fam(prequant_repos = (("nvfp4", "org/test-NVFP4"),)) + assert set(denoiser_prequant_sources(single, "nvfp4", "org/test-video")) == {"transformer"} + assert ( + denoiser_prequant_sources( + _fam(is_moe = True, prequant_repos = (("nvfp4", "org/x"),)), "nvfp4", "org/test-video" + ) + is None + ) + assert denoiser_prequant_sources(_moe_fam(), "int8", "org/test-video") is None + for scheme in (None, "", "auto", "off", "none"): + assert denoiser_prequant_sources(_moe_fam(), scheme, "org/test-video") is None + + +def _stub_prequant_loader(monkeypatch, outcomes): + """Record every ``load_prequantized_transformer`` call and return ``outcomes`` in order.""" + import gc + import sys + import types + + import core.inference.diffusion_prequant as pq + + diffusers = types.ModuleType("diffusers") + diffusers.TestTransformer3DModel = object + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + + events: list = [] + calls: list = [] + + def _fake_load(transformer_cls, base, source, **kwargs): + calls.append({"base": base, "source": source, **kwargs}) + events.append(("load", kwargs.get("component"))) + return outcomes[len(calls) - 1] + + monkeypatch.setattr(pq, "load_prequantized_transformer", _fake_load) + monkeypatch.setattr(gc, "collect", lambda *a, **k: events.append(("collect", None))) + return calls, events + + +def test_seeding_loads_every_expert_into_its_own_component(monkeypatch): + """Seeding loads every expert into its own component, from its own config.""" + from core.inference.video_denoiser_prequant import denoiser_prequant_pipe_kwargs + + first, second = object(), object() + calls, events = _stub_prequant_loader(monkeypatch, [first, second]) + + seeded = denoiser_prequant_pipe_kwargs( + _moe_fam(), + "org/test-video", + scheme = "nvfp4", + dtype = "bfloat16", + device = "cuda:0", + hf_token = "tok", + cache_dir = "/cache", + ) + assert seeded == {"transformer": first, "transformer_2": second} + assert [c["config_subfolder"] for c in calls] == ["transformer", "transformer_2"] + assert [c["component"] for c in calls] == ["transformer", "transformer_2"] + assert [c["source"].filename for c in calls] == [ + "test-NVFP4.safetensors", + "test-transformer_2-NVFP4.pt", + ] + from core.inference.diffusion_transformer_quant import DEFAULT_MIN_LINEAR_FEATURES + + assert {c["min_features"] for c in calls} == {DEFAULT_MIN_LINEAR_FEATURES} + assert events == [ + ("load", "transformer"), + ("collect", None), + ("load", "transformer_2"), + ("collect", None), + ] + + +def test_a_second_expert_that_will_not_load_seeds_nothing(monkeypatch): + """A second expert that will not load seeds nothing and releases what was loaded.""" + from core.inference.video_denoiser_prequant import denoiser_prequant_pipe_kwargs + + calls, events = _stub_prequant_loader(monkeypatch, [object(), None]) + + assert ( + denoiser_prequant_pipe_kwargs( + _moe_fam(), + "org/test-video", + scheme = "nvfp4", + dtype = "bfloat16", + device = "cuda:0", + ) + == {} + ) + assert len(calls) == 2 + assert events[-1] == ("collect", None) + + +def test_seeding_declines_a_host_that_cannot_run_the_quant(monkeypatch): + import types + + import core.inference.diffusion_transformer_quant as tq + from core.inference.video_denoiser_prequant import denoiser_prequant_pipe_kwargs + + calls, _events = _stub_prequant_loader(monkeypatch, [object(), object()]) + monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: False) + assert ( + denoiser_prequant_pipe_kwargs( + _moe_fam(), + "org/test-video", + scheme = "nvfp4", + dtype = "bfloat16", + device = "cuda:0", + target = types.SimpleNamespace(device = "cpu"), + ) + == {} + ) + assert calls == [] + + +def _video_auto( + monkeypatch, + *, + scheme, + fam = None, + **over, +): + """``_video_auto_denoiser_scheme`` with the device selector stubbed to ``scheme``.""" + from core.inference import video as vid + + monkeypatch.setattr( + vid, "select_transformer_quant_scheme", lambda target, requested, family = None: scheme + ) + kw = dict( + target = None, + requested = "auto", + base_repo = "org/test-video", + speed_mode = None, + ) + kw.update(over) + return vid._video_auto_denoiser_scheme(fam if fam is not None else _moe_fam(), **kw) + + +def test_the_conventional_auto_scheme_needs_an_artifact_for_every_expert(monkeypatch): + assert _video_auto(monkeypatch, scheme = "nvfp4") == "nvfp4" + assert _video_auto(monkeypatch, scheme = "int8") is None + assert ( + _video_auto( + monkeypatch, scheme = "nvfp4", fam = _fam(is_moe = True, prequant_repos = (("nvfp4", "org/x"),)) + ) + is None + ) + assert _video_auto(monkeypatch, scheme = None) is None + + +def test_speed_off_and_the_modular_workflow_are_never_seeded_here(monkeypatch): + assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = "off") is None + assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = " OFF ") is None + assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = "off", requested = None) is None + assert _video_auto(monkeypatch, scheme = "nvfp4", fam = _moe_fam(modular_workflow = "fl2va")) is None + + +def test_an_explicit_scheme_keeps_its_seed_under_speed_off(monkeypatch): + """An explicit scheme under Speed="off" is honored, so the plan keeps the hosted checkpoint.""" + assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = "off", requested = "nvfp4") == "nvfp4" + from core.inference.diffusion_transformer_quant import select_transformer_quant_scheme + + assert select_transformer_quant_scheme(object(), "off") is None + + +def test_an_install_that_cannot_open_a_checkpoint_keeps_the_dense_denoiser(monkeypatch): + import core.inference.diffusion_prequant as pq + monkeypatch.setattr(pq, "restricted_prequant_load_supported", lambda scheme = None: False) + assert _video_auto(monkeypatch, scheme = "nvfp4") is None + + +def test_the_conventional_coverage_probe_reads_every_component(): + from core.inference.video import VideoBackend + + fam = _moe_fam() + assert VideoBackend._denoiser_prequant_covered(fam, "nvfp4", "org/test-video") + assert not VideoBackend._denoiser_prequant_covered(fam, "int8", "org/test-video") + assert not VideoBackend._denoiser_prequant_covered(fam, "auto", "org/test-video") + assert not VideoBackend._denoiser_prequant_covered(fam, None, "org/test-video") + assert not VideoBackend._denoiser_prequant_covered( + fam, "nvfp4", "org/test-video", None, kind = "gguf" + ) + half = _fam(is_moe = True, prequant_repos = (("nvfp4", "org/x"),)) + assert not VideoBackend._denoiser_prequant_covered(half, "nvfp4", "org/test-video") + + +def _planned_denoiser_request(monkeypatch, fam, **load_kwargs): + """The scheme the planner hands ``_denoiser_prequant_verified`` (stubbed, as is the build).""" + from core.inference import video as vid + + backend = vid.VideoBackend() + backend._load_token = 1 + backend._loading = vid._VideoLoadingState(repo_id = fam.base_repo, base_repo = fam.base_repo) + monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: fam) + monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None) + monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None) + monkeypatch.setattr(backend, "_video_planned_auto_denoiser_scheme", lambda *a, **k: None) + monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: None) + monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None) + seen: list = [] + + def _verified(_fam, transformer_quant, *a, **k): + seen.append(transformer_quant) + return False + + monkeypatch.setattr(backend, "_denoiser_prequant_verified", _verified) + backend._run_load( + repo_id = fam.base_repo, + local_files_only = True, + _load_token = 1, + **load_kwargs, + ) + assert seen, "the plan never reached the denoiser probe" + return seen[0] + + +def test_a_conventional_plan_drops_no_shard_the_load_will_not_seed(monkeypatch): + """Only the planned seed drops shards: a conventional load never forwards the raw request.""" + from core.inference.video_families import detect_video_family + + wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + assert wan is not None and not wan.modular_workflow + assert ( + _planned_denoiser_request(monkeypatch, wan, transformer_quant = "nvfp4", speed_mode = "off") + is None + ) + h3 = detect_video_family("MiniMaxAI/MiniMax-H3") + assert h3 is not None and h3.modular_workflow + assert ( + _planned_denoiser_request( + monkeypatch, h3, transformer_quant = "int8", speed_mode = "off", h3_task = "fl2va" + ) + == "int8" + ) + + +def test_a_seed_the_plan_declined_is_pinned_into_the_load(monkeypatch): + """The plan's decline travels as its own value: the load cannot re-decide on a roomier card.""" + from core.inference import video as vid + from core.inference.video_families import detect_video_family + + wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + assert wan is not None and not wan.modular_workflow + backend = vid.VideoBackend() + backend._load_token = 1 + backend._loading = vid._VideoLoadingState(repo_id = wan.base_repo, base_repo = wan.base_repo) + monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: wan) + monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None) + monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None) + monkeypatch.setattr( + backend, + "_video_planned_auto_denoiser_scheme", + lambda *a, **k: vid.DENOISER_SEED_DECLINED, + ) + monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None) + probed: list = [] + + def _verified(_fam, transformer_quant, *a, **k): + probed.append(transformer_quant) + return False + + monkeypatch.setattr(backend, "_denoiser_prequant_verified", _verified) + loaded: list = [] + monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: loaded.append(kwargs)) + backend._run_load( + repo_id = wan.base_repo, + local_files_only = True, + _load_token = 1, + transformer_quant = "nvfp4", + ) + assert probed == [None], "a declined seed drops no dense shard from the pull" + assert loaded and loaded[0]["_video_auto_denoiser_planned"] == vid.DENOISER_SEED_DECLINED + + +def test_the_seeded_denoiser_repo_is_claimed_against_a_concurrent_delete(monkeypatch): + """The seeded denoiser's repo (neither repo_id nor base_repo) joins the in-flight claim.""" + from core.inference import video as vid + from core.inference.video_families import detect_video_family + + wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + assert wan is not None and not wan.modular_workflow + backend = vid.VideoBackend() + backend._load_token = 1 + backend._loading = vid._VideoLoadingState(repo_id = wan.base_repo, base_repo = wan.base_repo) + monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: wan) + monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None) + monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None) + monkeypatch.setattr(backend, "_video_planned_auto_denoiser_scheme", lambda *a, **k: "nvfp4") + monkeypatch.setattr(backend, "_denoiser_prequant_verified", lambda *a, **k: True) + monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None) + claimed: list = [] + monkeypatch.setattr( + backend, "load_pipeline", lambda **kwargs: claimed.extend(backend.loading_repo_ids()) + ) + backend._run_load( + repo_id = wan.base_repo, + local_files_only = True, + _load_token = 1, + transformer_quant = "nvfp4", + ) + assert "unsloth/Wan2.2-TI2V-5B-NVFP4" in claimed + + +def test_both_experts_of_an_moe_resolve_to_one_claimed_repo(): + """The two A14B experts are two files in ONE repo, so the claim names it once.""" + from core.inference.video_families import detect_video_family + + a14b = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers") + assert a14b is not None and a14b.is_moe + assert VideoBackend._denoiser_prequant_repo_ids(a14b, "nvfp4", a14b.base_repo) == ( + "unsloth/Wan2.2-T2V-A14B-NVFP4", + ) + assert VideoBackend._denoiser_prequant_repo_ids(a14b, "int8", a14b.base_repo) == () + assert VideoBackend._denoiser_prequant_repo_ids(a14b, "auto", a14b.base_repo) == () + + +def test_the_seeded_denoiser_artifact_is_fetched_under_the_load_cancel_event(monkeypatch): + """The seeded artifact is prefetched under the load's cancel event, not by the injection.""" + import threading + + from core.inference import video as vid + from core.inference.video_families import detect_video_family + + wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + assert wan is not None and not wan.modular_workflow + backend = vid.VideoBackend() + backend._load_token = 1 + backend._loading = vid._VideoLoadingState(repo_id = wan.base_repo, base_repo = wan.base_repo) + monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: wan) + monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None) + monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None) + monkeypatch.setattr(backend, "_video_planned_auto_denoiser_scheme", lambda *a, **k: "nvfp4") + monkeypatch.setattr(backend, "_denoiser_prequant_verified", lambda *a, **k: True) + monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None) + monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: None) + fetched: list = [] + import utils.hf_xet_fallback as xet + + monkeypatch.setattr( + xet, + "hf_hub_download_with_xet_fallback", + lambda repo, filename, token = None, **kwargs: fetched.append( + (repo, filename, kwargs.get("cancel_event")) + ), + ) + cancel = threading.Event() + + backend._run_load( + repo_id = wan.base_repo, + local_files_only = True, + _load_token = 1, + _cancel_event = cancel, + transformer_quant = "nvfp4", + ) + + assert [(r, f) for r, f, _ in fetched] == [ + ("unsloth/Wan2.2-TI2V-5B-NVFP4", "Wan2.2-TI2V-5B-NVFP4.pt") + ] + assert fetched[0][2] is cancel + + +def test_the_seed_plan_credits_the_resident_pipelines_memory(monkeypatch): + """The plan credits the resident pipeline's bytes, not other tenants', capped by the total.""" + import core.inference.video as vid + from core.inference.diffusion_memory import DeviceMemory + from core.inference.video_families import detect_video_family + + fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + seen: list = [] + monkeypatch.setattr( + vid, + "settled_snapshot_device_memory", + lambda _t: DeviceMemory("cuda", "cuda", "discrete_vram", 10_000, 32_000), + ) + monkeypatch.setattr( + vid, + "plan_diffusion_memory", + lambda **kw: seen.append(kw["device_memory"].free_mib) + or types.SimpleNamespace(offload_policy = "none"), + ) + target = types.SimpleNamespace(device = "cuda", dtype = None) + for credit in (0, 15_000, 40_000): + vid._video_seed_stays_resident( + fam, + target = target, + scheme = "fp8", + memory_mode = None, + text_encoder_quant = None, + base_repo = None, + reclaimable_mib = credit, + ) + assert seen == [10_000, 25_000, 32_000] + + +def test_the_planner_hands_the_resident_pipelines_bytes_to_the_seed_plan(monkeypatch): + import core.inference.video as vid + from core.inference.video_families import detect_video_family + + fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") + backend = vid.VideoBackend() + resident_pipe = object() + backend._state = types.SimpleNamespace(pipe = resident_pipe) + monkeypatch.setattr(vid, "_video_auto_denoiser_scheme", lambda fam, **kw: "fp8") + monkeypatch.setattr( + vid, + "_pipeline_device_mib", + lambda pipe, ordinal = None: ( + captured.setdefault("ordinal", ordinal), + 12_345 if pipe is resident_pipe else 0, + )[1], + ) + captured: dict = {} + monkeypatch.setattr( + vid, "_video_seed_stays_resident", lambda fam, **kw: captured.update(kw) or True + ) + monkeypatch.setattr( + vid, + "resolve_diffusion_device_target", + lambda **_k: types.SimpleNamespace(device = "cuda", dtype = None, ordinal = None), + ) + assert ( + backend._video_planned_auto_denoiser_scheme( + fam, base = None, kind = "pipeline", transformer_quant = "auto", speed_mode = None + ) + == "fp8" + ) + assert captured["reclaimable_mib"] == 12_345 + + +def test_only_the_resident_bytes_on_the_target_card_are_credited(): + """Only bytes on the target card are credited; a storage shared by two tensors counts once.""" + import core.inference.video as vid + + def tensor(index, ptr, nbytes): + storage = types.SimpleNamespace(data_ptr = lambda: ptr, nbytes = lambda: nbytes) + return types.SimpleNamespace( + device = types.SimpleNamespace(type = "cuda", index = index), + untyped_storage = lambda: storage, + ) + + tensors = [tensor(0, 1, 3 << 20), tensor(0, 1, 3 << 20), tensor(1, 2, 5 << 20)] + module = types.SimpleNamespace(parameters = lambda: iter(tensors), buffers = lambda: iter(())) + pipe = types.SimpleNamespace(components = {"transformer": module}) + assert vid._pipeline_device_mib(pipe) == 8 + assert vid._pipeline_device_mib(pipe, ordinal = 0) == 3 + assert vid._pipeline_device_mib(pipe, ordinal = 1) == 5 + assert vid._pipeline_device_mib(pipe, ordinal = 2) == 0 + + +def test_the_credit_is_read_on_the_card_the_load_targets(): + import core.inference.video as vid + assert vid._target_ordinal(types.SimpleNamespace(device = "cuda", ordinal = 3)) == 3 + assert vid._target_ordinal(types.SimpleNamespace(device = "mps", ordinal = None)) is None + + def test_unreadability_is_reported_only_when_it_is_what_kept_bfloat16(monkeypatch): try: import importlib diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index ea8532889e..cdb7fb7c8c 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -78,7 +78,13 @@ import { curatedArtifactTakesDenseQuant, loadSpecFor, } from "@/features/model-picker/components/model-selector/model-catalog"; -import { useDenseQuantSchemes, useHostClass } from "@/hooks/use-host-class"; +import { + useDenseQuantSchemes, + useHostClass, + useNvfp4Diffusion, + useNvfp4DiffusionKnown, +} from "@/hooks/use-host-class"; +import { nvfp4SelectionFallback, withNvfp4Option } from "@/lib/nvfp4-options"; import type { ModelOption, ModelSelectorChangeMeta, @@ -1275,6 +1281,8 @@ export function ImagesPage({ const { isMobile, pinned } = useSidebar(); const hostClass = useHostClass(); const denseQuantSchemes = useDenseQuantSchemes(); + const nvfp4Diffusion = useNvfp4Diffusion(); + const nvfp4DiffusionKnown = useNvfp4DiffusionKnown(); const imageModels = useImageModels(hostClass, denseQuantSchemes); const { rootStyle: railRootStyle } = useMediaRailWidth("images"); const [quant, setQuant] = useState(galleryCache.quant); @@ -1435,6 +1443,10 @@ export function ImagesPage({ const [attentionBackend, setAttentionBackend] = useState<"auto" | "native" | "cudnn" | "flash3" | "sage">( "auto", ); + useEffect(() => { + setTransformerQuant((v) => nvfp4SelectionFallback(v, nvfp4DiffusionKnown, nvfp4Diffusion)); + setTextEncoderQuant((v) => nvfp4SelectionFallback(v, nvfp4DiffusionKnown, nvfp4Diffusion)); + }, [nvfp4Diffusion, nvfp4DiffusionKnown, transformerQuant, textEncoderQuant]); const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto"); // "auto", or the physical index to pin this load to; offered only on a multi-card CUDA/ROCm // host. Persisted, unlike the selects around it: status carries the device a pipeline is on @@ -4409,12 +4421,15 @@ export function ImagesPage({ // The explicit low-precision schemes need the dense tensor-core path, which a Mac or // CPU-only host cannot run, so the picker does not list what the loader would refuse. ...(hostOffersDensePrecision(hostClass) - ? ([ - ["fp8", "FP8"], - ["int8", "INT8"], - ["nvfp4", "NVFP4 (Blackwell)"], - ["mxfp8", "MXFP8 (Blackwell)"], - ] as [string, string][]) + ? withNvfp4Option( + [ + ["fp8", "FP8"], + ["int8", "INT8"], + ["nvfp4", "NVFP4 (Blackwell)"], + ["mxfp8", "MXFP8 (Blackwell)"], + ] as [string, string][], + nvfp4Diffusion, + ) : []), ]} /> @@ -4434,16 +4449,18 @@ export function ImagesPage({ badge={} value={textEncoderQuant} onValueChange={(v) => setTextEncoderQuant(v as typeof textEncoderQuant)} - options={[ - ["auto", "Default"], - // The opt-out. Reachable only since a family default can pick a scheme on its own: with - // "Default" meaning bf16 everywhere, omitting the field WAS the dense request. - ["none", "Dense (bf16)"], - ["fp8", "FP8 (storage)"], - ["fp8_dynamic", "FP8 (compute)"], - ["int8", "INT8"], - ["nvfp4", "NVFP4 (Blackwell)"], - ]} + options={withNvfp4Option( + [ + ["auto", "Default"], + // Opt-out: now that a family default can pick a scheme, omitting the field is no longer the dense request. + ["none", "Dense (bf16)"], + ["fp8", "FP8 (storage)"], + ["fp8_dynamic", "FP8 (compute)"], + ["int8", "INT8"], + ["nvfp4", "NVFP4 (Blackwell)"], + ] as [string, string][], + nvfp4Diffusion, + )} /> (galleryCache.quant); // Starts from the last prompt generated with; the example is only a placeholder. @@ -1023,6 +1031,9 @@ function VideoGenerator({ const [transformerQuant, setTransformerQuant] = useState< "auto" | "none" | "fp8" | "int8" | "nvfp4" | "mxfp8" >("auto"); + useEffect(() => { + setTransformerQuant((v) => nvfp4SelectionFallback(v, nvfp4DiffusionKnown, nvfp4Diffusion)); + }, [nvfp4Diffusion, nvfp4DiffusionKnown, transformerQuant]); // The last load descriptor, so "Reapply" can reload the same model with new advanced options. const lastLoad = useRef<({ repoId: string } & VideoLoadOptions) | null>(null); // Render-safe mirror of whether a page-initiated load supplied a complete Reapply target. @@ -3499,12 +3510,15 @@ function VideoGenerator({ // The explicit low-precision schemes need the dense tensor-core path, which a Mac or // CPU-only host cannot run, so the picker does not list what the loader would refuse. ...(hostOffersDensePrecision(hostClass) - ? ([ - ["fp8", "FP8"], - ["int8", "INT8"], - ["nvfp4", "NVFP4 (Blackwell)"], - ["mxfp8", "MXFP8 (Blackwell)"], - ] as [string, string][]) + ? withNvfp4Option( + [ + ["fp8", "FP8"], + ["int8", "INT8"], + ["nvfp4", "NVFP4 (Blackwell)"], + ["mxfp8", "MXFP8 (Blackwell)"], + ] as [string, string][], + nvfp4Diffusion, + ) : []), ]} /> diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index 77202fc3d0..8c72680ef6 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -52,6 +52,8 @@ export interface GpuInfo { /** The dense quant schemes the backend says this host can run, best first ("fp8", "int8"). Empty * until system info arrives, and on a backend too old to report the field. */ denseQuantSchemes: readonly string[]; + /** False until system info arrives, and on backends that do not report it. */ + nvfp4Diffusion: boolean; name: string; memoryTotalGb: number; memorySharedGb: number; @@ -89,6 +91,7 @@ const DEFAULT_GPU: GpuInfo = { backend: "", denseQuantSupported: false, denseQuantSchemes: [], + nvfp4Diffusion: false, name: "Unknown", memoryTotalGb: 0, memorySharedGb: 0, @@ -116,6 +119,7 @@ function toGpuInfo( backend: data?.device_backend ?? "", denseQuantSupported: data?.dense_quant_supported === true, denseQuantSchemes: normalizeDenseQuantSchemes(data?.dense_quant_schemes), + nvfp4Diffusion: data?.nvfp4_diffusion === true, cpuCore: data?.cpu?.physical_count ?? 0, cpuThread: data?.cpu?.logical_count ?? 0, systemRamAvailableGb: data?.memory?.available_gb ?? 0, diff --git a/studio/frontend/src/hooks/use-host-class.ts b/studio/frontend/src/hooks/use-host-class.ts index cdd05d9129..9fb0731abe 100644 --- a/studio/frontend/src/hooks/use-host-class.ts +++ b/studio/frontend/src/hooks/use-host-class.ts @@ -25,6 +25,15 @@ export function useHostClass(): HostClass { ); } +export function useNvfp4Diffusion(): boolean { + return useGpuInfo().nvfp4Diffusion; +} + +/** Whether `/api/system` answered, so a false `useNvfp4Diffusion()` is not just the unloaded default. */ +export function useNvfp4DiffusionKnown(): boolean { + return useGpuInfo().budgetKnown; +} + /** The dense quant schemes this host can run, best first. */ export function useDenseQuantSchemes(): readonly string[] { return useGpuInfo().denseQuantSchemes; diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts index a3e0c5e6db..2083b6031a 100644 --- a/studio/frontend/src/hooks/use-system.ts +++ b/studio/frontend/src/hooks/use-system.ts @@ -59,6 +59,8 @@ export interface SystemInfoResponse { /** The dense quant schemes this host can run, best first. Absent on older backends, where readers * default to [] and name no precision. */ dense_quant_schemes?: string[]; + /** Absent on older backends, where readers treat it as off. */ + nvfp4_diffusion?: boolean; uptime_seconds: number | null; cpu: { logical_count: number; diff --git a/studio/frontend/src/lib/nvfp4-options.ts b/studio/frontend/src/lib/nvfp4-options.ts new file mode 100644 index 0000000000..251ac73206 --- /dev/null +++ b/studio/frontend/src/lib/nvfp4-options.ts @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export function withNvfp4Option( + options: readonly T[], + nvfp4Enabled: boolean, +): T[] { + if (nvfp4Enabled) return [...options]; + return options.filter(([value]) => value.trim().toLowerCase() !== "nvfp4"); +} + +/** A held `nvfp4` falls back to `auto` once the backend says it is off; before `known`, nothing resets. */ +export function nvfp4SelectionFallback( + value: T, + nvfp4Known: boolean, + nvfp4Enabled: boolean, +): T | "auto" { + if (!nvfp4Known || nvfp4Enabled) return value; + return value.trim().toLowerCase() === "nvfp4" ? "auto" : value; +} diff --git a/studio/frontend/tests/dense-quant-schemes-identity.test.ts b/studio/frontend/tests/dense-quant-schemes-identity.test.ts index 0678de71a7..e91d109ec6 100644 --- a/studio/frontend/tests/dense-quant-schemes-identity.test.ts +++ b/studio/frontend/tests/dense-quant-schemes-identity.test.ts @@ -42,6 +42,7 @@ function snapshot( backend: "cuda", denseQuantSupported: schemes.length > 0, denseQuantSchemes: schemes, + nvfp4Diffusion: false, name: "NVIDIA GeForce RTX 4090", memoryTotalGb: 24, memorySharedGb: 0, diff --git a/studio/frontend/tests/nvfp4-options.test.ts b/studio/frontend/tests/nvfp4-options.test.ts new file mode 100644 index 0000000000..02571c1afb --- /dev/null +++ b/studio/frontend/tests/nvfp4-options.test.ts @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import assert from "node:assert/strict"; +import { test } from "node:test"; +import { nvfp4SelectionFallback, withNvfp4Option } from "../src/lib/nvfp4-options.ts"; + +const OPTIONS: [string, string][] = [ + ["fp8", "FP8"], + ["int8", "INT8"], + ["nvfp4", "NVFP4 (Blackwell)"], + ["mxfp8", "MXFP8 (Blackwell)"], +]; + +test("NVFP4 is hidden while the backend switch is off", () => { + assert.deepEqual( + withNvfp4Option(OPTIONS, false).map(([value]) => value), + ["fp8", "int8", "mxfp8"], + ); +}); + +test("NVFP4 is offered, in place, when the backend enables it", () => { + assert.deepEqual(withNvfp4Option(OPTIONS, true), OPTIONS); +}); + +test("the input list is never mutated", () => { + const copy = OPTIONS.map((o) => [...o]); + withNvfp4Option(OPTIONS, false); + assert.deepEqual(OPTIONS, copy); +}); + +test("a held NVFP4 pick falls back to auto once the backend reports the switch off", () => { + assert.equal(nvfp4SelectionFallback("nvfp4", true, false), "auto"); + assert.equal(nvfp4SelectionFallback("NVFP4", true, false), "auto"); +}); + +test("other picks, an enabled switch, and a not-yet-known switch keep the selection", () => { + assert.equal(nvfp4SelectionFallback("fp8", true, false), "fp8"); + assert.equal(nvfp4SelectionFallback("none", true, false), "none"); + assert.equal(nvfp4SelectionFallback("nvfp4", true, true), "nvfp4"); + assert.equal(nvfp4SelectionFallback("nvfp4", false, false), "nvfp4"); +}); + +test("both pages reset the hidden NVFP4 selects through the fallback", async () => { + const { readFile } = await import("node:fs/promises"); + const images = await readFile(new URL("../src/features/images/images-page.tsx", import.meta.url), "utf8"); + const video = await readFile(new URL("../src/features/video/video-page.tsx", import.meta.url), "utf8"); + for (const setter of ["setTransformerQuant", "setTextEncoderQuant"]) { + assert.match( + images, + new RegExp(`${setter}\\(\\(v\\) => nvfp4SelectionFallback\\(v, nvfp4DiffusionKnown, nvfp4Diffusion\\)\\)`), + ); + } + assert.match( + video, + /setTransformerQuant\(\(v\) => nvfp4SelectionFallback\(v, nvfp4DiffusionKnown, nvfp4Diffusion\)\)/, + ); +});