ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override

Browse source 
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).
This commit is contained in:
Daniel Han 2026-06-26 08:01:04 +00:00
commit 10db6a4777
7 changed files with 170 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

11
scripts/diffusion_bench.py
View file

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

9
studio/backend/core/inference/diffusion.py
View file

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

73
studio/backend/core/inference/diffusion_transformer_quant.py
View file

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

8
studio/backend/models/inference.py
View file

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

1
studio/backend/routes/inference.py
View file

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

16
studio/backend/tests/test_diffusion_routes.py
View file

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

65
studio/backend/tests/test_diffusion_transformer_quant.py
View file

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