Merge remote-tracking branch 'origin/diffusion-fp16-accum' into fold-integration

studio/backend/core/inference/diffusion_speed.py

@@ -69,6 +69,8 @@ def snapshot_backend_flags() -> Optional[dict]:
     matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
     if matmul is not None and hasattr(matmul, "allow_tf32"):
         state["matmul_tf32"] = bool(matmul.allow_tf32)
+    if matmul is not None and hasattr(matmul, "allow_fp16_accumulation"):
+        state["matmul_fp16_accum"] = bool(matmul.allow_fp16_accumulation)
     cudnn = getattr(torch.backends, "cudnn", None)
     if cudnn is not None:
         if hasattr(cudnn, "allow_tf32"):
@@ -95,9 +97,9 @@ def restore_backend_flags(state: Optional[dict]) -> None:
             except Exception:  # noqa: BLE001 — best-effort per-flag restore
                 pass
 
-    _set(
-        getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32"
-    )
+    matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
+    _set(matmul, "allow_tf32", "matmul_tf32")
+    _set(matmul, "allow_fp16_accumulation", "matmul_fp16_accum")
     cudnn = getattr(torch.backends, "cudnn", None)
     _set(cudnn, "allow_tf32", "cudnn_tf32")
     _set(cudnn, "benchmark", "cudnn_benchmark")
@@ -179,6 +181,7 @@ def apply_speed_optims(
         "channels_last": False,
         "cudnn_benchmark": False,
         "tf32": False,
+        "fp16_accum": False,
         "fused_qkv": False,
         "compiled": False,
         "compiled_dequant": False,
@@ -203,6 +206,21 @@ def apply_speed_optims(
     if on_cuda:
         applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger)
 
+    # Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts, whose
+    # fp32-accumulate rate is halved; datacenter HBM parts gain nothing and keep the
+    # safer fp32 accumulate). Only affects fp16 matmuls -- bf16 loads were measured
+    # bit-identical with the flag on across every family (36/36 same-seed A/B cases),
+    # so on the quality-neutral tiers the flag engages only when the compute dtype is
+    # NOT fp16. On an fp16 pipeline (the pre-Ampere fallback dtype) the same harness
+    # measured real same-seed drift (mean 2-5% on SDXL / FLUX), so fp16 compute gets
+    # the 2x accumulate only under ``max``, the tier that already trades exactness for
+    # speed. Guarded by a per-family deny-list fed by the overflow validation harness
+    # and the UNSLOTH_DISABLE_FP16_ACCUM kill switch.
+    if on_cuda:
+        applied["fp16_accum"] = _enable_fp16_accumulation(
+            family, logger, dtype = getattr(target, "dtype", None), speed_mode = mode
+        )
+
     # --- the compile lever, remapped per tier ----------------------------------------
     # default = LIGHT compile: for a GGUF model, compile ONLY the dequant op chain
     #   (~70-80% of eager GGUF time) -- cheap, VRAM-free, resolution-invariant; the
@@ -330,6 +348,61 @@ def _enable_tf32(logger: Any) -> bool:
         return False
 
 
+# Families the overflow validation harness (scripts/fp16_accum_validate.py) found to
+# produce non-finite activations or NEW black frames under fp16 accumulation. Empty by
+# measurement: across all six families the harness found no overflow anywhere -- bf16
+# loads are bit-identical with the flag on, and fp16 loads stay finite (their same-seed
+# drift is why fp16 compute is additionally gated to the ``max`` tier below).
+_FP16_ACCUM_DENY: frozenset[str] = frozenset()
+
+
+def _enable_fp16_accumulation(
+    family: Any,
+    logger: Any,
+    *,
+    dtype: Any = None,
+    speed_mode: Optional[str] = None,
+) -> bool:
+    """Turn on fp16-accumulated fp16 GEMMs for consumer GPUs, where they run ~2x the
+    fp32-accumulate rate (datacenter HBM parts are not throughput-nerfed, so they keep
+    the safer default). Gated on: the torch build exposing the flag (2.10+), a
+    consumer-class device, the family not being deny-listed by the overflow harness,
+    the UNSLOTH_DISABLE_FP16_ACCUM kill switch being unset, and -- when the pipeline
+    compute dtype IS fp16, the only case where the accumulator width changes results --
+    the ``max`` tier (measured same-seed drift: mean 2-5%; bf16 loads are bit-identical
+    so they engage on any tier). The caller's snapshot/restore pair returns the
+    process-wide flag to its prior value on unload."""
+    import os
+
+    if os.environ.get("UNSLOTH_DISABLE_FP16_ACCUM", "").strip().lower() in (
+        "1",
+        "true",
+        "yes",
+        "on",
+    ):
+        return False
+    name = str(getattr(family, "name", family or "")).lower()
+    if name in _FP16_ACCUM_DENY:
+        return False
+    if str(dtype).replace("torch.", "") == "float16" and speed_mode != SPEED_MAX:
+        return False
+    try:
+        import torch
+
+        matmul = torch.backends.cuda.matmul
+        if not hasattr(matmul, "allow_fp16_accumulation"):
+            return False
+        from .diffusion_transformer_quant import _is_consumer_gpu
+
+        if not _is_consumer_gpu():
+            return False
+        matmul.allow_fp16_accumulation = True
+        return True
+    except Exception as exc:  # noqa: BLE001 — optimisation only
+        _warn(logger, "fp16_accum", exc)
+        return False
+
+
 def _fuse_qkv(pipe: Any, logger: Any) -> bool:
     for owner in (pipe, getattr(pipe, "transformer", None)):
         fn = getattr(owner, "fuse_qkv_projections", None)

studio/backend/tests/test_diffusion_speed.py

@@ -211,6 +211,7 @@ def test_speed_off_applies_nothing(monkeypatch):
         "fused_qkv": False,
         "compiled": False,
         "compiled_dequant": False,
+        "fp16_accum": False,
     }
     assert pipe.vae.mem_format is None and pipe.compiled is False
     # off must not touch any process-wide flag (bit-identical reference path).
@@ -351,3 +352,154 @@ def test_apply_tolerates_missing_optims(monkeypatch):
         bare, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX
     )
     assert applied["channels_last"] is False and applied["fused_qkv"] is False
+
+
+# ── fp16 accumulation (consumer fp16-GEMM fast path) ──────────────────────────
+
+
+def _stub_torch_fp16_accum(
+    monkeypatch,
+    *,
+    consumer = True,
+    with_flag = True,
+):
+    torch = types.ModuleType("torch")
+    torch.bfloat16 = "bfloat16"
+    torch.channels_last = "channels_last"
+    matmul_attrs = {"allow_tf32": False}
+    if with_flag:
+        matmul_attrs["allow_fp16_accumulation"] = False
+    torch.backends = types.SimpleNamespace(
+        cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(**matmul_attrs)),
+        cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False),
+    )
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    import core.inference.diffusion_transformer_quant as tq
+
+    monkeypatch.setattr(tq, "_is_consumer_gpu", lambda device = None: consumer)
+    return torch
+
+
+def test_snapshot_captures_fp16_accum_when_present(monkeypatch):
+    torch = _stub_torch_fp16_accum(monkeypatch)
+    torch.backends.cuda.matmul.allow_fp16_accumulation = True
+    snap = snapshot_backend_flags()
+    assert snap["matmul_fp16_accum"] is True
+    torch.backends.cuda.matmul.allow_fp16_accumulation = False
+    restore_backend_flags(snap)
+    assert torch.backends.cuda.matmul.allow_fp16_accumulation is True
+
+
+def test_snapshot_skips_fp16_accum_on_older_torch(monkeypatch):
+    _stub_torch_fp16_accum(monkeypatch, with_flag = False)
+    snap = snapshot_backend_flags()
+    assert "matmul_fp16_accum" not in snap
+    restore_backend_flags(snap)  # nothing to restore, no error
+
+
+def test_fp16_accum_engages_on_consumer_cuda(monkeypatch):
+    torch = _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    _stub_gguf_accel(monkeypatch)
+    applied = apply_speed_optims(
+        _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default"
+    )
+    assert applied["fp16_accum"] is True
+    assert torch.backends.cuda.matmul.allow_fp16_accumulation is True
+
+
+def test_fp16_accum_skipped_on_datacenter(monkeypatch):
+    torch = _stub_torch_fp16_accum(monkeypatch, consumer = False)
+    _stub_gguf_accel(monkeypatch)
+    applied = apply_speed_optims(
+        _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default"
+    )
+    assert applied["fp16_accum"] is False
+    assert torch.backends.cuda.matmul.allow_fp16_accumulation is False
+
+
+def test_fp16_accum_respects_kill_switch(monkeypatch):
+    _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    _stub_gguf_accel(monkeypatch)
+    monkeypatch.setenv("UNSLOTH_DISABLE_FP16_ACCUM", "1")
+    applied = apply_speed_optims(
+        _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default"
+    )
+    assert applied["fp16_accum"] is False
+
+
+@pytest.mark.parametrize("value", ["TRUE", "Yes", "On", " true "])
+def test_fp16_accum_kill_switch_is_case_insensitive(monkeypatch, value):
+    # The documented safety escape hatch must honor the common boolean spellings, not only
+    # lowercase "1"/"true"/"yes": an operator setting UNSLOTH_DISABLE_FP16_ACCUM=TRUE to stop
+    # fp16-accumulation drift would otherwise be silently ignored.
+    _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    _stub_gguf_accel(monkeypatch)
+    monkeypatch.setenv("UNSLOTH_DISABLE_FP16_ACCUM", value)
+    applied = apply_speed_optims(
+        _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default"
+    )
+    assert applied["fp16_accum"] is False
+
+
+def test_fp16_accum_respects_family_deny_list(monkeypatch):
+    _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    _stub_gguf_accel(monkeypatch)
+    monkeypatch.setattr(ds_mod, "_FP16_ACCUM_DENY", frozenset({"fragile-family"}))
+    fam = types.SimpleNamespace(supports_torch_compile = True, name = "fragile-family")
+    applied = apply_speed_optims(_Pipe(), _target(), is_gguf = True, family = fam, speed_mode = "default")
+    assert applied["fp16_accum"] is False
+
+
+def test_fp16_accum_skipped_when_flag_missing(monkeypatch):
+    _stub_torch_fp16_accum(monkeypatch, consumer = True, with_flag = False)
+    _stub_gguf_accel(monkeypatch)
+    applied = apply_speed_optims(
+        _Pipe(), _target(), is_gguf = True, family = _family(), speed_mode = "default"
+    )
+    assert applied["fp16_accum"] is False
+
+
+def test_fp16_accum_not_touched_off_cuda(monkeypatch):
+    torch = _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    applied = apply_speed_optims(
+        _Pipe(),
+        _target(device = "mps"),
+        is_gguf = False,
+        family = _family(),
+        speed_mode = "eager",
+    )
+    assert applied["fp16_accum"] is False
+    assert torch.backends.cuda.matmul.allow_fp16_accumulation is False
+
+
+def test_fp16_accum_denied_on_fp16_dtype_below_max(monkeypatch):
+    # fp16 compute is where the accumulator width actually changes results (measured
+    # same-seed drift, mean 2-5%): the quality-neutral tiers must refuse it.
+    torch = _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    _stub_gguf_accel(monkeypatch)
+    for mode in ("eager", "default"):
+        applied = apply_speed_optims(
+            _Pipe(),
+            _target(dtype = "float16"),
+            is_gguf = True,
+            family = _family(),
+            speed_mode = mode,
+        )
+        assert applied["fp16_accum"] is False
+    assert torch.backends.cuda.matmul.allow_fp16_accumulation is False
+
+
+def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch):
+    # max already trades exactness for speed (conv algos, max-autotune), so the 2x
+    # fp16 accumulate joins that tier for fp16 pipelines.
+    torch = _stub_torch_fp16_accum(monkeypatch, consumer = True)
+    _stub_gguf_accel(monkeypatch)
+    applied = apply_speed_optims(
+        _Pipe(with_compile = True, with_fuse = True),
+        _target(dtype = "float16"),
+        is_gguf = True,
+        family = _family(),
+        speed_mode = "MAX",
+    )
+    assert applied["fp16_accum"] is True
+    assert torch.backends.cuda.matmul.allow_fp16_accumulation is True