Merge diffusion-fp16-accum: torchao probe stub in precision-mode tests

studio/backend/core/inference/diffusion_speed.py

@@ -357,7 +357,11 @@ _FP16_ACCUM_DENY: frozenset[str] = frozenset()
 
 
 def _enable_fp16_accumulation(
-    family: Any, logger: Any, *, dtype: Any = None, speed_mode: Optional[str] = None
+    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

studio/backend/tests/test_diffusion_dit_trainer.py

@@ -230,10 +230,15 @@ def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
 
 def _patch_capability(monkeypatch, capability):
     # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor
-    # core capability (fp8 needs sm89+, mxfp8 needs sm100+).
+    # core capability (fp8 needs sm89+, mxfp8 needs sm100+). The torchao probe is stubbed
+    # functional so these tests exercise the CAPABILITY gate on hosts without torchao
+    # (the CPU-only CI runner does not install it).
     import torch
+
+    import core.training.diffusion_train_common as dtc
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
     monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
+    monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True)
 
 
 def test_train_precision_modes_blackwell_lists_mxfp8(monkeypatch):

studio/backend/tests/test_diffusion_speed.py

@@ -357,7 +357,12 @@ def test_apply_tolerates_missing_optims(monkeypatch):
 # ── fp16 accumulation (consumer fp16-GEMM fast path) ──────────────────────────
 
 
-def _stub_torch_fp16_accum(monkeypatch, *, consumer = True, with_flag = True):
+def _stub_torch_fp16_accum(
+    monkeypatch,
+    *,
+    consumer = True,
+    with_flag = True,
+):
     torch = types.ModuleType("torch")
     torch.bfloat16 = "bfloat16"
     torch.channels_last = "channels_last"
@@ -427,9 +432,7 @@ def test_fp16_accum_respects_family_deny_list(monkeypatch):
     _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"
-    )
+    applied = apply_speed_optims(_Pipe(), _target(), is_gguf = True, family = fam, speed_mode = "default")
     assert applied["fp16_accum"] is False