From 8bdd307889cbe56390a55fef87ffe5c894caf45c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 05:38:41 +0000 Subject: [PATCH] Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. --- scripts/nvfp4_probe.py | 6 ++++- studio/backend/core/inference/diffusion.py | 23 +++++++++++++++++++ .../inference/diffusion_transformer_quant.py | 9 +++++++- .../backend/tests/test_diffusion_backend.py | 3 +++ 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/scripts/nvfp4_probe.py b/scripts/nvfp4_probe.py index 546e32881a..902cf8a2d2 100644 --- a/scripts/nvfp4_probe.py +++ b/scripts/nvfp4_probe.py @@ -16,7 +16,7 @@ import numpy as np BASE = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" -OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_images") +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_images" def _psnr(a, b): @@ -84,11 +84,15 @@ def main(argv = None) -> int: p.add_argument("--seed", type = int, default = 42) p.add_argument("--iters", type = int, default = 3) p.add_argument("--min-feat", type = int, default = 512) + p.add_argument("--out-dir", default = None, help = "image output dir (default: repo outputs/)") args = p.parse_args(argv) steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat import torch import torch.nn as nn + global OUT + if args.out_dir: + OUT = Path(args.out_dir).expanduser() OUT.mkdir(parents = True, exist_ok = True) def filt(mod, fqn = ""): diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 65e7ce8faa..e5cf4e20e8 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -45,6 +45,7 @@ from .diffusion_memory import ( snapshot_device_memory, ) from .diffusion_speed import ( + SPEED_DEFAULT, SPEED_OFF, apply_speed_optims, resolve_speed_mode, @@ -534,6 +535,18 @@ class DiffusionBackend: # the quant noise floor), dense models stay bit-identical `off`. An # explicit speed_mode (incl. "off") is honored verbatim. effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename)) + # A torchao-quantized dense transformer runs its matmuls through the + # regional torch.compile; UNcompiled (eager) it is ~30x slower and would + # lose to the GGUF fallback. A dense model otherwise resolves to `off`, so + # force at least `default` (regional compile) whenever the quant engaged, + # or the opt-in "fast" path silently commits an eager, pathologically slow + # pipeline. + if transformer_quant_engaged is not None and effective_speed == SPEED_OFF: + logger.info( + "diffusion.transformer_quant: forcing speed_mode=default " + "(quantized transformer must be compiled; eager is ~30x slower)" + ) + effective_speed = SPEED_DEFAULT # Opt-in speed optims run BEFORE placement (channels_last / compile # must precede CPU offload). Snapshot the process-wide backend flags # first so unload can restore them: TF32 / cudnn.benchmark are global, @@ -547,6 +560,16 @@ class DiffusionBackend: speed_mode = effective_speed, logger = logger, ) + if transformer_quant_engaged is not None and not speed_applied.get("compiled"): + # Promotion above could not engage compile (e.g. the family is not + # compile-friendly, or compile_repeated_blocks failed): the quantized + # transformer is now running eager, which is far slower than the GGUF + # path it replaced. Surface it loudly rather than hiding the regression. + logger.warning( + "diffusion.transformer_quant: %s engaged but the transformer is NOT " + "compiled; eager torchao quant is ~30x slower than GGUF here", + transformer_quant_engaged, + ) # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), # also before placement so the offload hooks move the smaller weights. te_quant = quantize_text_encoders( diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 296c30576e..6e1341983f 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -257,7 +257,14 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: return Float8DynamicActivationFloat8WeightConfig() if scheme == TQ_NVFP4: from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig - return NVFP4DynamicActivationNVFP4WeightConfig() + # Select the CUTLASS FP4 path, not the default Triton kernel: torchao defaults + # use_triton_kernel=True, which needs MSLK installed. On a Blackwell box with the + # CUTLASS FP4 extension but no MSLK, the default would make the smoke probe fail + # and silently fall back to GGUF instead of using the FP4 tensor cores. + try: + return NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel = False) + except TypeError: # older torchao without the knob + return NVFP4DynamicActivationNVFP4WeightConfig() if scheme == TQ_MXFP8: import torch from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 634f322639..d5760362d1 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -947,6 +947,9 @@ def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatc transformer_quant = "fp8", ) assert status["transformer_quant"] == "fp8" + # No speed_mode was given, but a quantized transformer is ~30x slower eager, so the + # backend promotes it to `default` (regional compile) instead of the dense `off`. + assert status["speed_mode"] == "default" assert calls["from_pretrained"] == 1 and calls["quantize"] == 1 assert calls["quant_mode"] == "fp8" assert calls["fp_kwargs"]["subfolder"] == "transformer" # dense transformer subfolder