From 10db6a477764141f9ae6c68ff202a2db38c532cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 08:01:04 +0000 Subject: [PATCH] Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). --- scripts/diffusion_bench.py | 11 +++ studio/backend/core/inference/diffusion.py | 9 ++- .../inference/diffusion_transformer_quant.py | 73 ++++++++++++++++--- studio/backend/models/inference.py | 8 ++ studio/backend/routes/inference.py | 1 + studio/backend/tests/test_diffusion_routes.py | 16 ++++ .../tests/test_diffusion_transformer_quant.py | 65 ++++++++++++++++- 7 files changed, 170 insertions(+), 13 deletions(-) diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 96bbd86d02..4f53567651 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -217,6 +217,9 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: speed_mode = args.speed_mode, text_encoder_quant = args.text_encoder_quant, transformer_quant = args.transformer_quant, + transformer_quant_fast_accum = {"auto": None, "on": True, "off": False}[ + args.fp8_fast_accum + ], ) _wait_for_load(backend) _cuda_sync() @@ -299,6 +302,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "cpu_offload": args.cpu_offload, "text_encoder_quant": args.text_encoder_quant, "transformer_quant": args.transformer_quant, + "fp8_fast_accum": args.fp8_fast_accum, }, } @@ -466,6 +470,13 @@ def _build_parser() -> argparse.ArgumentParser: "quantise it onto the low-precision tensor cores (faster than GGUF, higher " "VRAM). auto picks per GPU; falls back to GGUF if unsupported / no VRAM", ) + p.add_argument( + "--fp8-fast-accum", + default = "auto", + choices = ["auto", "on", "off"], + help = "fp8 accumulate: auto picks by GPU class (fast on consumer, precise on " + "data-center); on/off force it", + ) p.add_argument( "--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload" ) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 38e03e7ced..076a296124 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -276,6 +276,7 @@ class DiffusionBackend: speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -308,6 +309,7 @@ class DiffusionBackend: speed_mode = speed_mode, text_encoder_quant = text_encoder_quant, transformer_quant = transformer_quant, + transformer_quant_fast_accum = transformer_quant_fast_accum, _load_token = token, ), daemon = True, @@ -430,6 +432,7 @@ class DiffusionBackend: speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -496,6 +499,7 @@ class DiffusionBackend: hf_token, target, transformer_quant, + transformer_quant_fast_accum, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build logger.warning( @@ -601,6 +605,7 @@ class DiffusionBackend: hf_token: Optional[str], target: DiffusionDeviceTarget, mode: Optional[str], + fast_accum: Optional[bool] = None, ) -> tuple[Any, str]: """Build the opt-in fast pipeline: load the DENSE bf16 transformer from the base repo (``subfolder="transformer"``), assemble the pipeline, place it on the device, @@ -618,7 +623,9 @@ class DiffusionBackend: pipe_kwargs["token"] = hf_token pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) pipe.to(device) - scheme = quantize_transformer(pipe, target, mode = mode, logger = logger) + scheme = quantize_transformer( + pipe, target, mode = mode, fast_accum = fast_accum, logger = logger + ) if scheme is None: raise RuntimeError("transformer quant unsupported for this device/scheme") return pipe, scheme diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index a7d612ff7f..c77af06f7b 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -59,6 +59,41 @@ _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( # Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once. _SMOKE_CACHE: dict[tuple[str, str], bool] = {} +# Data-center GPU model tokens (un-nerfed FP32 accumulate). Matched as whole tokens of +# torch.cuda.get_device_name(), so the workstation "A4000" is not mistaken for the +# data-center "A40". Anything not here -- GeForce, workstation RTX, or an unknown name -- +# is treated as consumer-class (FP32-accumulate halved). See developer.nvidia.com/cuda/gpus. +_DATACENTER_GPU_TOKENS = frozenset({ + "B200", "B100", "GB200", "GB300", "GB10", # Blackwell data center + "H200", "H100", "H800", "H20", # Hopper data center + "A100", "A800", "A30", "A40", "A16", "A10", "A2", # Ampere data center + "L40", "L40S", "L4", "L20", "L2", # Ada data center + "V100", "P100", "P40", "T4", # legacy data center +}) + + +def _is_consumer_gpu(device: Any = None) -> bool: + """Whether the active GPU is consumer / workstation class (GDDR), where fp8 FP32 + accumulate is throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center + HBM parts (recognised by name token) are not nerfed and return False, so they keep + the higher-precision default accumulate for free. Heuristic on the device name: a + GeForce / TITAN name is always consumer; a recognised data-center token is not; + anything else (workstation RTX, unknown) defaults to consumer -- the safe choice, + since fast accumulate is free on data-center and a win on consumer. Best-effort: + True on any probe failure.""" + try: + import re + + import torch + + name = torch.cuda.get_device_name(device).upper() + except Exception: # noqa: BLE001 — no torch / no device -> assume consumer + return True + if "GEFORCE" in name or "TITAN" in name: + return True + tokens = set(re.split(r"[^A-Z0-9]+", name)) + return not (tokens & _DATACENTER_GPU_TOKENS) + def normalize_transformer_quant(value: Optional[str]) -> Optional[str]: """Lower/strip a requested transformer quant; None / "" / "none" / "off" -> None. @@ -162,9 +197,17 @@ def _smoke_probe(scheme: str, device: str) -> bool: return ok -def _make_quant_config(scheme: str) -> Any: +def _resolve_fast_accum(fast_accum: Optional[bool]) -> bool: + """The fp8 ``use_fast_accum`` to apply. ``None`` auto-detects by GPU class + (consumer / workstation -> fast; data-center -> precise); an explicit bool forces it.""" + return _is_consumer_gpu() if fast_accum is None else bool(fast_accum) + + +def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: """The torchao dynamic-activation config for ``scheme`` (lazy import; prototype - import for the Blackwell fp4 / mx schemes is inside the branch that needs it).""" + import for the Blackwell fp4 / mx schemes is inside the branch that needs it). + + ``fast_accum`` applies to fp8 only: None auto-detects by GPU class, True/False force it.""" from torchao.quantization import ( Float8DynamicActivationFloat8WeightConfig, Int8DynamicActivationInt8WeightConfig, @@ -173,15 +216,17 @@ def _make_quant_config(scheme: str) -> Any: if scheme == TQ_INT8: return Int8DynamicActivationInt8WeightConfig() if scheme == TQ_FP8: - # Lock fast (FP16) accumulate. On consumer Blackwell (e.g. RTX 50xx) the fp8 - # tensor cores run at ~838 TFLOPS with FP16 accumulate but only ~419 with FP32, - # so the fast-accum path is a free ~2x there. torchao already defaults it on; - # set it explicitly so a future default change can't silently halve consumer - # throughput. (Negligible numeric effect for diffusion's short reductions.) + # Choose fp8 accumulate by GPU class (unless forced). On consumer / workstation + # cards (GDDR) the fp8 tensor cores run ~2x faster with FP16 (fast) accumulate + # than FP32 (e.g. ~838 vs ~419 TFLOPS on RTX 50xx), so fast accumulate is a real + # win there. Data-center HBM parts default to the higher-precision accumulate. + # fast accumulate is a precision (not overflow) tradeoff and stays below the fp8 + # quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations). try: from torchao.float8 import Float8MMConfig + return Float8DynamicActivationFloat8WeightConfig( - mm_config = Float8MMConfig(use_fast_accum = True) + mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum)) ) except Exception: # noqa: BLE001 — older torchao without the explicit knob return Float8DynamicActivationFloat8WeightConfig() @@ -226,12 +271,16 @@ def quantize_transformer( *, mode: Optional[str], min_features: int = DEFAULT_MIN_LINEAR_FEATURES, + fast_accum: Optional[bool] = None, logger: Any = None, ) -> Optional[str]: """Quantise ``pipe.transformer``'s FLOP-heavy linears in place with the arch-chosen dynamic scheme. Returns the scheme actually engaged, or None when disabled / unsupported / failed -- the caller then loads GGUF instead. Best-effort: it never - raises for an ordinary unsupported environment (a failure leaves the module dense).""" + raises for an ordinary unsupported environment (a failure leaves the module dense). + + ``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None + auto-detects (fast on consumer, precise on data-center), True/False force it.""" scheme = select_transformer_quant_scheme(target, mode) if scheme is None: return None @@ -241,7 +290,11 @@ def quantize_transformer( try: from torchao.quantization import quantize_ - quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(min_features)) + quantize_( + transformer, + _make_quant_config(scheme, fast_accum = fast_accum), + filter_fn = make_filter_fn(min_features), + ) # Runtime-only marker (torchao tensors are not safetensors-serializable; this # backend is inference-only, so this is purely diagnostic). try: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b4593e96e8..d0d62393dc 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1730,6 +1730,14 @@ class DiffusionLoadRequest(BaseModel): "Ampere int8); an explicit scheme forces it. Needs CUDA + bf16 + room " "for the dense load; falls back to GGUF otherwise.", ) + transformer_quant_fast_accum: Optional[bool] = Field( + None, + description = "fp8 only: FP8 matmul accumulate. null auto-detects by GPU class " + "(fast FP16 accumulate on consumer/workstation cards, where FP32 " + "accumulate is ~2x slower; precise FP32 accumulate on data-center " + "HBM cards, which are not nerfed). true/false force it. Negligible " + "quality effect (below the fp8 quant noise floor); no overflow risk.", + ) class DiffusionGenerateRequest(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4451d6f040..ec84b4ee7f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10086,6 +10086,7 @@ async def load_diffusion_model( speed_mode = request.speed_mode, text_encoder_quant = request.text_encoder_quant, transformer_quant = request.transformer_quant, + transformer_quant_fast_accum = request.transformer_quant_fast_accum, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index f6d55a50c6..9d66edd750 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -327,6 +327,22 @@ def test_transformer_quant_threads_through_to_backend(client, monkeypatch): assert backend.last_load_kwargs.get("transformer_quant") == "auto" +def test_transformer_quant_fast_accum_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_quant": "fp8", + "transformer_quant_fast_accum": False, + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False + + def test_invalid_transformer_quant_returns_422_without_eviction(client): # An unsupported transformer_quant is rejected by the request schema (Literal), so # the GPU is never acquired and no chat model is evicted. diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index c89f562c62..3ee17bc816 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -207,6 +207,50 @@ def test_smoke_probe_caches_and_tolerates_failure(monkeypatch): assert tq._smoke_probe(TQ_FP8, "cuda") is False +# ── consumer-vs-datacenter detection (fp8 fast-accumulate gate) ────────────────── + + +def _stub_device_name(monkeypatch, name): + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(get_device_name=lambda device=None: name) + monkeypatch.setitem(sys.modules, "torch", torch) + + +@pytest.mark.parametrize("name", [ + "NVIDIA GeForce RTX 5090", + "NVIDIA GeForce RTX 4090", + "NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40 + "NVIDIA RTX 6000 Ada Generation", + "NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC) +]) +def test_is_consumer_gpu_true(monkeypatch, name): + _stub_device_name(monkeypatch, name) + assert tq._is_consumer_gpu() is True + + +@pytest.mark.parametrize("name", [ + "NVIDIA B200", + "NVIDIA H100 80GB HBM3", + "NVIDIA A100-SXM4-80GB", + "NVIDIA A40", # data-center Ampere (distinct token from RTX A4000) + "NVIDIA L40S", + "NVIDIA L4", + "Tesla V100-SXM2-16GB", +]) +def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name): + _stub_device_name(monkeypatch, name) + assert tq._is_consumer_gpu() is False + + +def test_is_consumer_gpu_defaults_true_on_probe_failure(monkeypatch): + # No torch / no device name available -> assume consumer (safe: fast accum is free + # on data center and a win on consumer). + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace() # no get_device_name + monkeypatch.setitem(sys.modules, "torch", torch) + assert tq._is_consumer_gpu() is True + + # ── filter ────────────────────────────────────────────────────────────────────── @@ -230,9 +274,25 @@ def test_make_filter_fn(monkeypatch): # ── apply ─────────────────────────────────────────────────────────────────────── +def test_resolve_fast_accum(monkeypatch): + # None auto-detects by GPU class; an explicit bool forces it. + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: True) + assert tq._resolve_fast_accum(None) is True + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: False) + assert tq._resolve_fast_accum(None) is False + assert tq._resolve_fast_accum(True) is True # forced on (e.g. on a data-center card) + assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card) + + def test_quantize_transformer_applies_and_marks(monkeypatch): monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8) - monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: f"{scheme}cfg") + seen: dict = {} + + def _mk(scheme, fast_accum = None): + seen["scheme"], seen["fast_accum"] = scheme, fast_accum + return f"{scheme}cfg" + + monkeypatch.setattr(tq, "_make_quant_config", _mk) recorder: list = [] tqz = types.ModuleType("torchao.quantization") tqz.quantize_ = lambda module, config, filter_fn = None: recorder.append( @@ -242,10 +302,11 @@ def test_quantize_transformer_applies_and_marks(monkeypatch): transformer = types.SimpleNamespace() pipe = types.SimpleNamespace(transformer = transformer) - assert quantize_transformer(pipe, _target(), mode = "fp8") == TQ_FP8 + assert quantize_transformer(pipe, _target(), mode = "fp8", fast_accum = False) == TQ_FP8 assert len(recorder) == 1 and recorder[0][0] is transformer and recorder[0][1] == "fp8cfg" assert callable(recorder[0][2]) # a filter_fn was passed assert transformer._unsloth_runtime_quant == TQ_FP8 # diagnostic marker set + assert seen["fast_accum"] is False # the override is forwarded into the config def test_quantize_transformer_none_when_unsupported(monkeypatch):