diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 7cb1074425..cb6873a661 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -55,6 +55,11 @@ DEFAULT_MIN_LINEAR_FEATURES = 512 # (0.81x end-to-end on Z-Image 1024px) AND notably less accurate (LPIPS 0.166 vs fp8's # 0.044), because FP4's per-forward quant overhead is not amortised and the format is # coarser. So nvfp4 is kept as an explicit opt-in, never the auto pick for diffusion. +# +# This order is the DATA-CENTER preference. On a consumer / workstation GPU it is reordered +# to put int8 first (see ``_prefer_consumer_scheme``): consumer cards halve fp8/fp16 FP32- +# accumulate throughput, while int8 runs full-rate (int32 accumulate is not nerfed), so int8 +# is as fast or faster than fp8 on every consumer NVIDIA / AMD / Intel part. _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( ((10, 0), (TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8)), # Blackwell sm_100+ ((8, 9), (TQ_FP8, TQ_INT8)), # Ada sm_89 / Hopper sm_90 @@ -72,13 +77,15 @@ _DATACENTER_GPU_TOKENS = frozenset( { "B200", "B100", + "B300", # Blackwell Ultra data center "GB200", "GB300", "GB10", # Blackwell data center "H200", "H100", "H800", - "H20", # Hopper data center + "H20", + "GH200", # Grace-Hopper superchip (data center) "A100", "A800", "A30", @@ -99,15 +106,22 @@ _DATACENTER_GPU_TOKENS = frozenset( ) +# Professional parts the rest of the backend treats as datacenter-class (see llama_cpp.py +# _DATACENTER_GPU_RE, which applies the same FP32-accum tuning to them). Matched as phrases +# because the marker spans tokens ("RTX PRO 6000", "RTX 6000 ADA"), so they must not be +# misread as consumer (which would put int8 ahead of fp8 and pick fast accumulate). +_PROFESSIONAL_GPU_MARKERS = ("RTX PRO 6000", "RTX 6000 ADA") + + 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.""" + """Whether the active GPU is consumer-class (GDDR), where fp8 FP32 accumulate is + throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center HBM parts and + professional parts (recognised by name) are not nerfed and return False, so they keep + the higher-precision default accumulate and fp8 first. Heuristic on the device name: a + GeForce / TITAN name is always consumer; a recognised data-center token or professional + marker is not; anything else (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 @@ -117,6 +131,8 @@ def _is_consumer_gpu(device: Any = None) -> bool: return True if "GEFORCE" in name or "TITAN" in name: return True + if any(marker in name for marker in _PROFESSIONAL_GPU_MARKERS): + return False tokens = set(re.split(r"[^A-Z0-9]+", name)) return not (tokens & _DATACENTER_GPU_TOKENS) @@ -168,13 +184,24 @@ def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Op return None for floor, schemes in _AUTO_LADDER: if cap >= floor: - for scheme in schemes: + for scheme in _prefer_consumer_scheme(schemes, device): if _scheme_supported(scheme, device): return scheme return None return None +def _prefer_consumer_scheme(schemes: tuple[str, ...], device: Any) -> tuple[str, ...]: + """Reorder an arch tier's schemes for the GPU class. On a consumer / workstation card + move int8 to the front: consumer parts halve fp8/fp16 FP32-accumulate throughput, while + int8 runs at full rate (int32 accumulate is not nerfed), so int8 is as fast or faster + than fp8 on every consumer NVIDIA / AMD / Intel GPU (and the only path on pre-Ada + consumer without fp8 tensor cores). Data-center HBM parts keep fp8 first.""" + if TQ_INT8 in schemes and schemes[0] != TQ_INT8 and _is_consumer_gpu(device): + return (TQ_INT8,) + tuple(s for s in schemes if s != TQ_INT8) + return schemes + + def _capability() -> Optional[tuple[int, int]]: try: import torch diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index d2adaa9f7c..d527723b2e 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -39,6 +39,7 @@ def _stub_torch( cc = (10, 0), with_fp8 = True, cuda_available = True, + device_name = "NVIDIA B200", ): torch = types.ModuleType("torch") torch.bfloat16 = "bfloat16" @@ -48,6 +49,9 @@ def _stub_torch( torch.cuda = types.SimpleNamespace( is_available = lambda: cuda_available, get_device_capability = lambda *a: cc, + # data-center name by default so the ladder tests get the data-center order; + # consumer tests pass a GeForce name (or monkeypatch _is_consumer_gpu). + get_device_name = lambda *a: device_name, ) monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -104,11 +108,49 @@ def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch): assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 +def test_auto_consumer_blackwell_prefers_int8(monkeypatch): + # Consumer Blackwell (RTX 50xx): fp8 FP32-accumulate is throughput-halved while int8 is + # full-rate, so auto prefers int8 even though fp8 is available (the data-center default). + _stub_torch(monkeypatch, cc = (10, 0), device_name = "NVIDIA GeForce RTX 5090") + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + # int8 unavailable -> falls back to the rest of the tier (fp8 next). + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + + +def test_auto_consumer_ada_prefers_int8(monkeypatch): + # Consumer Ada (RTX 4090): int8 runs ~2x fp8's nerfed FP32-accumulate rate. + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA GeForce RTX 4090") + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_workstation_unknown_prefers_int8(monkeypatch): + # Unknown / workstation name -> treated as consumer (the safe default) -> int8 first. + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA RTX A5000") + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_professional_rtx_prefers_fp8(monkeypatch): + # Professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) are classified datacenter + # by the rest of the backend, so auto keeps fp8 first (not int8) -- matching llama_cpp. + for device_name, cc in ( + ("NVIDIA RTX PRO 6000 Blackwell Server Edition", (10, 0)), + ("NVIDIA RTX 6000 Ada Generation", (8, 9)), + ): + _stub_torch(monkeypatch, cc = cc, device_name = device_name) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + + def test_auto_ada_hopper_prefers_fp8(monkeypatch): - _stub_torch(monkeypatch, cc = (8, 9)) + # Data-center Ada (L40S) / Hopper (H100): not nerfed -> fp8 first. + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA L40S") _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 - _stub_torch(monkeypatch, cc = (9, 0)) # Hopper + _stub_torch(monkeypatch, cc = (9, 0), device_name = "NVIDIA H100 80GB HBM3") # Hopper assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 @@ -224,7 +266,7 @@ def _stub_device_name(monkeypatch, 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 RTX A5000", # workstation: A5000 token, not professional/datacenter "NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC) ], ) @@ -237,12 +279,16 @@ def test_is_consumer_gpu_true(monkeypatch, name): "name", [ "NVIDIA B200", + "NVIDIA B300", # Blackwell Ultra (matches llama_cpp datacenter regex) + "NVIDIA GH200 480GB", # Grace-Hopper superchip (was misread as consumer) "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", + "NVIDIA RTX PRO 6000 Blackwell Server Edition", # professional -> datacenter-class + "NVIDIA RTX 6000 Ada Generation", # professional -> datacenter-class ], ) def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name):