From f24384b4e90f5826aaed000a7ebbe88cb923c99b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jun 2026 23:09:42 +0000 Subject: [PATCH] Studio diffusion: eager patches + torch.compile cache speed phase Adds the opt-in speed path for the GGUF diffusion transformer behind a selectable speed mode (default off, so output is unchanged until a profile is chosen): - diffusion_eager_patches.py: shared eager fast-paths (channels_last, attention/backend selection, fused norms and QKV) installed at load and rolled back on unload or failed load. - diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent torch.compile cache and the GGUF-transformer compile wiring. - diffusion_arch_patches.py: architecture-specific patches. - diffusion_patch_backend.py: shared install/restore plumbing. - diffusion_speed.py: speed-profile planning. Tests for each module plus the benchmarking and probe scripts used to measure speed, memory, and accuracy of the path. --- .../core/inference/diffusion_arch_patches.py | 469 ++++++++++++++++++ .../core/inference/diffusion_compile_cache.py | 338 +++++++++++++ .../core/inference/diffusion_eager_patches.py | 203 ++++++++ .../core/inference/diffusion_gguf_compile.py | 132 +++++ .../core/inference/diffusion_patch_backend.py | 70 +++ .../backend/core/inference/diffusion_speed.py | 87 +++- studio/backend/tests/conftest.py | 6 + .../tests/test_diffusion_arch_patches.py | 252 ++++++++++ .../tests/test_diffusion_compile_cache.py | 189 +++++++ .../tests/test_diffusion_eager_patches.py | 183 +++++++ .../tests/test_diffusion_gguf_compile.py | 73 +++ studio/backend/tests/test_diffusion_speed.py | 64 ++- 12 files changed, 2036 insertions(+), 30 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_arch_patches.py create mode 100644 studio/backend/core/inference/diffusion_compile_cache.py create mode 100644 studio/backend/core/inference/diffusion_eager_patches.py create mode 100644 studio/backend/core/inference/diffusion_gguf_compile.py create mode 100644 studio/backend/core/inference/diffusion_patch_backend.py create mode 100644 studio/backend/tests/test_diffusion_arch_patches.py create mode 100644 studio/backend/tests/test_diffusion_compile_cache.py create mode 100644 studio/backend/tests/test_diffusion_eager_patches.py create mode 100644 studio/backend/tests/test_diffusion_gguf_compile.py diff --git a/studio/backend/core/inference/diffusion_arch_patches.py b/studio/backend/core/inference/diffusion_arch_patches.py new file mode 100644 index 0000000000..267a97266e --- /dev/null +++ b/studio/backend/core/inference/diffusion_arch_patches.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Per-architecture eager fusions for the diffusion DiT blocks. + +The shared patches in ``diffusion_eager_patches.py`` cover what diffusers factors into +SHARED classes (``RMSNorm``, the ``AdaLayerNorm*`` modulation classes -- the latter only +used by flux.1). The remaining fusible ops -- the gated residual ``x = x + gate * out`` and +the inline modulation ``norm * (1 + scale) + shift`` -- are written longhand in each +family's PER-MODEL block ``forward`` (flux.2 / qwen / z-image), so they need per-arch +patches. This module supplies them, in the ``unsloth_zoo.temporary_patches`` style: one +small patch function per target, applied through the shared, fingerprint-checked, +reversible backend (``diffusion_patch_backend`` -> ``patch_function`` / ``restore_original``). + +The single actionable fused op is ``torch.addcmul(a, b, c) == a + b * c`` (one FMA kernel, +1-ULP vs mul+add, MORE accurate). Two forms exist: + * out-of-place ``torch.addcmul(...)`` -- COMPILE-SAFE: lowers to plain ops, no aliasing, + neutral under ``torch.compile``. Used here for ALL non-``off`` tiers. It captures the + real eager win (fewer kernel launches); it does NOT save the output allocation. + * in-place ``x.addcmul_(...)`` -- would also save the allocation, but the allocation part + measured NEUTRAL on this stack (CUDA caching allocator already recycles; cf. the GGUF + weight-buffer result), and in-place residual mutation is compile-unsafe + aliasing-risky. + So it is deliberately NOT used; the out-of-place form is the whole, safe win. + +Each patch is guarded TWICE: ``can_safely_patch`` checks the forward SIGNATURE, and a local +source-body check (``_body_has``) confirms the exact lines we rewrite are still present, so a +future diffusers that changed the block body simply leaves the block UNPATCHED (correctness +first) rather than running a stale copy. Kill-switch: ``UNSLOTH_DIFFUSION_ARCH_PATCHES=0``. + +Implemented for all four families (extend by adding entries to ``_SPECS``): + * qwen-image ``QwenImageTransformerBlock._modulate`` -- modulation addcmul (all 4 sites). + * z-image ``ZImageTransformerBlock.forward`` -- the 2 gated-residual addcmuls. + * flux.1 ``FluxTransformerBlock.forward`` (inline norm2 modulation + 4 gated residuals) + + ``FluxSingleTransformerBlock.forward`` (residual + gate*proj_out). + * flux.2-klein ``Flux2TransformerBlock.forward`` (4 inline modulations + 4 gated residuals) + + ``Flux2SingleTransformerBlock.forward`` (inline modulation + gated residual). +""" + +from __future__ import annotations + +import inspect +import logging +import os +from typing import Any, Callable, Optional + +import torch + +from .diffusion_patch_backend import apply_patch, revert_patch + +logger = logging.getLogger(__name__) + +_ENV_ENABLE = "UNSLOTH_DIFFUSION_ARCH_PATCHES" + + +def _patches_enabled() -> bool: + return (os.environ.get(_ENV_ENABLE) or "").strip().lower() not in ("0", "off", "false", "no") + + +def _body_has(fn: Callable, *needles: str) -> bool: + """True iff every ``needle`` appears in ``fn``'s source -- a body-drift guard so a patch + self-disables if diffusers changed the lines it rewrites.""" + try: + src = inspect.getsource(fn) + except (OSError, TypeError): + return False + return all(n in src for n in needles) + + +# ===================================================================================== +# qwen-image: QwenImageTransformerBlock._modulate (modulation addcmul, all 4 call sites) +# ===================================================================================== +def _qwen_modulate(self, x, mod_params, index=None): + """diffusers 0.38 ``QwenImageTransformerBlock._modulate`` with the final + ``x*(1+scale)+shift`` fused to ``torch.addcmul`` (covers both the global and the + per-token ``index`` branches, since both end in that same expression).""" + shift, scale, gate = mod_params.chunk(3, dim=-1) + + if index is not None: + actual_batch = shift.size(0) // 2 + shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] + scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] + gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] + index_expanded = index.unsqueeze(-1) + shift_0_exp = shift_0.unsqueeze(1) + shift_1_exp = shift_1.unsqueeze(1) + scale_0_exp = scale_0.unsqueeze(1) + scale_1_exp = scale_1.unsqueeze(1) + gate_0_exp = gate_0.unsqueeze(1) + gate_1_exp = gate_1.unsqueeze(1) + shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp) + scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp) + gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp) + else: + shift_result = shift.unsqueeze(1) + scale_result = scale.unsqueeze(1) + gate_result = gate.unsqueeze(1) + + # fused: x * (1 + scale_result) + shift_result + return torch.addcmul(shift_result, x, 1 + scale_result), gate_result + + +def _spec_qwen_modulate(): + try: + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "_modulate", None) + if orig is None or not _body_has(orig, "x * (1 + scale_result) + shift_result"): + return None + return (cls, "_modulate", _qwen_modulate) + + +# ===================================================================================== +# z-image: ZImageTransformerBlock.forward (the 2 gated-residual addcmuls) +# ===================================================================================== +def _zimage_forward( + self, + x: torch.Tensor, + attn_mask: torch.Tensor, + freqs_cis: torch.Tensor, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, +): + """diffusers 0.38 ``ZImageTransformerBlock.forward`` with the two gated residuals + ``x = x + gate * sublayer`` fused to ``torch.addcmul`` (out-of-place). The shift-free + ``*scale`` modulation and the non-gated (``else``) residuals are left as-is.""" + from diffusers.models.transformers.transformer_z_image import select_per_token + + if self.modulation: + seq_len = x.shape[1] + + if noise_mask is not None: + mod_noisy = self.adaLN_modulation(adaln_noisy) + mod_clean = self.adaLN_modulation(adaln_clean) + + scale_msa_noisy, gate_msa_noisy, scale_mlp_noisy, gate_mlp_noisy = mod_noisy.chunk(4, dim=1) + scale_msa_clean, gate_msa_clean, scale_mlp_clean, gate_mlp_clean = mod_clean.chunk(4, dim=1) + + gate_msa_noisy, gate_mlp_noisy = gate_msa_noisy.tanh(), gate_mlp_noisy.tanh() + gate_msa_clean, gate_mlp_clean = gate_msa_clean.tanh(), gate_mlp_clean.tanh() + + scale_msa_noisy, scale_mlp_noisy = 1.0 + scale_msa_noisy, 1.0 + scale_mlp_noisy + scale_msa_clean, scale_mlp_clean = 1.0 + scale_msa_clean, 1.0 + scale_mlp_clean + + scale_msa = select_per_token(scale_msa_noisy, scale_msa_clean, noise_mask, seq_len) + scale_mlp = select_per_token(scale_mlp_noisy, scale_mlp_clean, noise_mask, seq_len) + gate_msa = select_per_token(gate_msa_noisy, gate_msa_clean, noise_mask, seq_len) + gate_mlp = select_per_token(gate_mlp_noisy, gate_mlp_clean, noise_mask, seq_len) + else: + mod = self.adaLN_modulation(adaln_input) + scale_msa, gate_msa, scale_mlp, gate_mlp = mod.unsqueeze(1).chunk(4, dim=2) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + + # Attention block -- fused gated residual: x + gate_msa * attention_norm2(attn_out) + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis + ) + x = torch.addcmul(x, gate_msa, self.attention_norm2(attn_out)) + + # FFN block -- fused gated residual: x + gate_mlp * ffn_norm2(feed_forward(...)) + x = torch.addcmul( + x, gate_mlp, self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + ) + else: + attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis) + x = x + self.attention_norm2(attn_out) + x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x))) + + return x + + +def _spec_zimage_forward(): + try: + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "x = x + gate_msa * self.attention_norm2(attn_out)", + "x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))", + ): + return None + return (cls, "forward", _zimage_forward) + + +# ===================================================================================== +# flux.1: FluxTransformerBlock / FluxSingleTransformerBlock +# (block modulation goes through AdaLayerNormZero -- already handled by the shared patch; +# here we fuse the inline norm2 modulation + the gated residual adds.) +# ===================================================================================== +def _flux_double_forward( + self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None, joint_attention_kwargs=None +): + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( + encoder_hidden_states, emb=temb + ) + joint_attention_kwargs = joint_attention_kwargs or {} + attention_outputs = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + if len(attention_outputs) == 2: + attn_output, context_attn_output = attention_outputs + elif len(attention_outputs) == 3: + attn_output, context_attn_output, ip_attn_output = attention_outputs + + # fused: hidden_states + gate_msa * attn_output + hidden_states = torch.addcmul(hidden_states, gate_msa.unsqueeze(1), attn_output) + + norm_hidden_states = self.norm2(hidden_states) + # fused: norm * (1 + scale_mlp) + shift_mlp + norm_hidden_states = torch.addcmul(shift_mlp[:, None], norm_hidden_states, 1 + scale_mlp[:, None]) + + ff_output = self.ff(norm_hidden_states) + hidden_states = torch.addcmul(hidden_states, gate_mlp.unsqueeze(1), ff_output) + if len(attention_outputs) == 3: + hidden_states = hidden_states + ip_attn_output + + encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_msa.unsqueeze(1), context_attn_output) + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = torch.addcmul( + c_shift_mlp[:, None], norm_encoder_hidden_states, 1 + c_scale_mlp[:, None] + ) + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_mlp.unsqueeze(1), context_ff_output) + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +def _spec_flux_double(): + try: + from diffusers.models.transformers.transformer_flux import FluxTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "hidden_states = hidden_states + attn_output", + "norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]", + "encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output", + ): + return None + return (cls, "forward", _flux_double_forward) + + +def _flux_single_forward( + self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None, joint_attention_kwargs=None +): + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + residual = hidden_states + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states)) + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2) + gate = gate.unsqueeze(1) + # fused: residual + gate * proj_out(hidden_states) + hidden_states = torch.addcmul(residual, gate, self.proj_out(hidden_states)) + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] + return encoder_hidden_states, hidden_states + + +def _spec_flux_single(): + try: + from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "hidden_states = gate * self.proj_out(hidden_states)", + "hidden_states = residual + hidden_states", + ): + return None + return (cls, "forward", _flux_single_forward) + + +# ===================================================================================== +# flux.2-klein: Flux2TransformerBlock / Flux2SingleTransformerBlock +# (modulation is INLINE here -- not via AdaLayerNorm -- so we fuse both the modulation and +# the gated residuals; scale/shift/gate are [B,1,dim] so no [:, None] is needed.) +# ===================================================================================== +def _flux2_double_forward( + self, hidden_states, encoder_hidden_states, temb_mod_img, temb_mod_txt, + image_rotary_emb=None, joint_attention_kwargs=None, +): + from diffusers.models.transformers.transformer_flux2 import Flux2Modulation + + joint_attention_kwargs = joint_attention_kwargs or {} + (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = Flux2Modulation.split(temb_mod_img, 2) + (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = Flux2Modulation.split( + temb_mod_txt, 2 + ) + + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = torch.addcmul(shift_msa, norm_hidden_states, 1 + scale_msa) + + norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) + norm_encoder_hidden_states = torch.addcmul(c_shift_msa, norm_encoder_hidden_states, 1 + c_scale_msa) + + attention_outputs = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + attn_output, context_attn_output = attention_outputs + + hidden_states = torch.addcmul(hidden_states, gate_msa, attn_output) + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = torch.addcmul(shift_mlp, norm_hidden_states, 1 + scale_mlp) + + ff_output = self.ff(norm_hidden_states) + hidden_states = torch.addcmul(hidden_states, gate_mlp, ff_output) + + encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_msa, context_attn_output) + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = torch.addcmul(c_shift_mlp, norm_encoder_hidden_states, 1 + c_scale_mlp) + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = torch.addcmul(encoder_hidden_states, c_gate_mlp, context_ff_output) + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +def _spec_flux2_double(): + try: + from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa", + "hidden_states = hidden_states + gate_mlp * ff_output", + "encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output", + ): + return None + return (cls, "forward", _flux2_double_forward) + + +def _flux2_single_forward( + self, hidden_states, encoder_hidden_states, temb_mod, image_rotary_emb=None, + joint_attention_kwargs=None, split_hidden_states=False, text_seq_len=None, +): + from diffusers.models.transformers.transformer_flux2 import Flux2Modulation + + if encoder_hidden_states is not None: + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + mod_shift, mod_scale, mod_gate = Flux2Modulation.split(temb_mod, 1)[0] + + norm_hidden_states = self.norm(hidden_states) + norm_hidden_states = torch.addcmul(mod_shift, norm_hidden_states, 1 + mod_scale) + + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = torch.addcmul(hidden_states, mod_gate, attn_output) + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + if split_hidden_states: + encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] + return encoder_hidden_states, hidden_states + else: + return hidden_states + + +def _spec_flux2_single(): + try: + from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock as cls + except Exception: # noqa: BLE001 + return None + orig = getattr(cls, "forward", None) + if orig is None or not _body_has( + orig, + "norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift", + "hidden_states = hidden_states + mod_gate * attn_output", + ): + return None + return (cls, "forward", _flux2_single_forward) + + +# ===================================================================================== +# registry + lifecycle +# ===================================================================================== +# Each entry is a zero-arg resolver returning (cls, attr, new_fn) or None (target absent / +# body drifted). All entries here are COMPILE-SAFE (out-of-place addcmul). +_SPECS: tuple[Callable[[], Optional[tuple]], ...] = ( + _spec_qwen_modulate, + _spec_zimage_forward, + _spec_flux_double, + _spec_flux_single, + _spec_flux2_double, + _spec_flux2_single, +) + +# (cls, attr) pairs we successfully patched, for an exact reverse. +_patched: list[tuple] = [] + + +def install_arch_patches() -> int: + """Install the per-arch compile-safe fusions (idempotent). Returns the count applied. + + Safe for every non-``off`` tier: the fusions lower to plain ops and are neutral under + ``torch.compile`` (so they also help the ``max`` regionally-compiled blocks).""" + if not _patches_enabled(): + uninstall_arch_patches() + return 0 + if _patched: + return len(_patched) + for resolve in _SPECS: + try: + spec = resolve() + except Exception as exc: # noqa: BLE001 + logger.warning("arch-patch: resolver %s failed: %s", getattr(resolve, "__name__", resolve), exc) + spec = None + if spec is None: + continue + cls, attr, new_fn = spec + if apply_patch(cls, attr, new_fn, match_level="relaxed"): + _patched.append((cls, attr)) + else: + logger.warning("arch-patch: skipping %s.%s (signature mismatch / unavailable)", + getattr(cls, "__name__", cls), attr) + logger.info("arch-patch: installed %d/%d per-arch fusions", len(_patched), len(_SPECS)) + return len(_patched) + + +def uninstall_arch_patches() -> None: + """Restore every per-arch patched method/forward (idempotent).""" + for cls, attr in list(_patched): + revert_patch(cls, attr) + _patched.clear() + + +def is_installed() -> bool: + return bool(_patched) diff --git a/studio/backend/core/inference/diffusion_compile_cache.py b/studio/backend/core/inference/diffusion_compile_cache.py new file mode 100644 index 0000000000..b9544c4c82 --- /dev/null +++ b/studio/backend/core/inference/diffusion_compile_cache.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pre-warmed ``torch.compile`` cache for the diffusion denoiser (Mega-cache). + +The regional ``torch.compile`` of the repeated denoiser block (``diffusion_speed.py``) +pays a one-time 25-58s compile on the FIRST image after a load. This module lets that +cost be paid ONCE -- by us (the distributor) ahead of time, or by the user on a first +run -- and reused on every later load via torch's portable Mega-cache +(``torch.compiler.save_cache_artifacts`` / ``load_cache_artifacts``, torch >= 2.7). + +PORTABILITY IS NOT UNIVERSAL. A compiled artifact is only valid for the SAME torch +version, Triton version, CUDA build, and GPU architecture it was produced on (and the +same model graph: family, dtype, quant scheme, attention backend, compile kwargs, shape +bucket). torch validates these on load and a mismatch simply yields no cache hit -- it +does NOT error. So this layer is built around an EXACT-MATCH fingerprint with a SILENT +FALLBACK to local compile: a miss is normal and never fatal. We therefore ship per-arch +bundles keyed by the full fingerprint, never one universal cache. See +``outputs/compile_cache/DISTRIBUTION.md``. + +Lifecycle (driven by the caller, around ``_compile_repeated_blocks``): + 1. ``begin(...)`` -> build the fingerprint, point ``TORCHINDUCTOR_CACHE_DIR`` at a + per-key dir, and ``load_cache_artifacts`` if a matching bundle + exists. Must run BEFORE the first compiled forward. + 2. (compile + one warmup forward happen as usual; on a hit they reuse the cache.) + 3. ``save(...)`` -> ``save_cache_artifacts`` to the bundle + manifest, AFTER the + warmup forward, when in distributor/save mode. + 4. ``restore(...)`` -> put ``TORCHINDUCTOR_CACHE_DIR`` back on unload. + +Everything is env-gated and best-effort; torch is imported lazily. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import json +import os +import time +from pathlib import Path +from typing import Any, Optional + +# ----------------------------------------------------------------------------- env knobs +# UNSLOTH_DIFFUSION_COMPILE_CACHE: auto (default) | 0 | 1 +# auto -> load a matching bundle if present (no automatic save). +# 1 -> load AND save (distributor / first-run warm). +# 0 -> disabled (plain local compile, no cache dir override). +# UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR: root dir for bundles (default under the workspace). +# UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE: 1 -> force-enable save even in "auto". +_ENV_MODE = "UNSLOTH_DIFFUSION_COMPILE_CACHE" +_ENV_DIR = "UNSLOTH_DIFFUSION_COMPILE_CACHE_DIR" +_ENV_SAVE = "UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE" + +_DEFAULT_ROOT = Path.home() / ".cache" / "unsloth" / "diffusion_compile_cache" + +_MANIFEST_NAME = "manifest.json" +_BUNDLE_NAME = "cache.bin" +_FORMAT_VERSION = 1 + + +def cache_mode() -> str: + """``off`` | ``auto`` | ``on`` from the environment. ``auto`` is the default.""" + raw = (os.environ.get(_ENV_MODE) or "auto").strip().lower() + if raw in ("0", "off", "false", "no"): + return "off" + if raw in ("1", "on", "true", "yes"): + return "on" + return "auto" + + +def _save_enabled(mode: str) -> bool: + if mode == "off": + return False + if mode == "on": + return True + # auto: save only if explicitly opted in. + return (os.environ.get(_ENV_SAVE) or "").strip().lower() in ("1", "on", "true", "yes") + + +def cache_root() -> Path: + root = os.environ.get(_ENV_DIR) + return Path(root) if root else _DEFAULT_ROOT + + +# --------------------------------------------------------------------------- fingerprint +def _triton_version() -> Optional[str]: + try: + import triton # noqa: PLC0415 + return str(getattr(triton, "__version__", None)) + except Exception: # noqa: BLE001 — triton optional + return None + + +def _diffusers_version() -> Optional[str]: + try: + import diffusers # noqa: PLC0415 + return str(getattr(diffusers, "__version__", None)) + except Exception: # noqa: BLE001 + return None + + +def environment_fingerprint() -> dict[str, Any]: + """The HARD-portability dimensions: any difference here invalidates a bundle. + + These mirror what torch's inductor cache itself keys on (torch + triton + CUDA + + GPU type), plus diffusers (the graph source). We surface them explicitly so the + manifest is self-describing and a mismatch is obvious to a human, not just to torch. + """ + fp: dict[str, Any] = { + "format": _FORMAT_VERSION, + "torch": None, + "torch_cuda": None, + "triton": _triton_version(), + "diffusers": _diffusers_version(), + "gpu_name": None, + "gpu_capability": None, + } + try: + import torch # noqa: PLC0415 + fp["torch"] = str(torch.__version__) + fp["torch_cuda"] = str(torch.version.cuda) + if torch.cuda.is_available(): + fp["gpu_name"] = torch.cuda.get_device_name(0) + cap = torch.cuda.get_device_capability(0) + fp["gpu_capability"] = f"sm_{cap[0]}{cap[1]}" + except Exception: # noqa: BLE001 — best-effort + pass + return fp + + +def model_fingerprint( + *, + family: Any, + transformer: Any, + dtype: Any, + quant: Any, + attention_backend: Any, + compile_kwargs: dict[str, Any], + shape_bucket: Any = None, +) -> dict[str, Any]: + """The MODEL-graph dimensions that change the compiled artifact. + + ``family`` is the Unsloth family name; ``transformer`` is the live module (we read + its class name + ``_repeated_blocks`` so the key tracks exactly what gets compiled). + """ + blocks = list(getattr(transformer, "_repeated_blocks", []) or []) + return { + "family": str(family), + "transformer_cls": type(transformer).__name__ if transformer is not None else None, + "repeated_blocks": sorted(str(b) for b in blocks), + "dtype": str(dtype), + "quant": str(quant) if quant is not None else "none", + "attention_backend": str(attention_backend) if attention_backend is not None else "default", + "compile_kwargs": {k: compile_kwargs[k] for k in sorted(compile_kwargs)}, + "shape_bucket": shape_bucket, + } + + +def cache_key(env_fp: dict[str, Any], model_fp: dict[str, Any]) -> str: + payload = json.dumps({"env": env_fp, "model": model_fp}, sort_keys=True, default=str) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:32] + + +# ----------------------------------------------------------------------------- lifecycle +@dataclasses.dataclass +class CacheContext: + """Carries the per-load cache state between ``begin`` and ``save``/``restore``.""" + + key: str + dir: Path + bundle: Path + manifest_path: Path + env_fp: dict[str, Any] + model_fp: dict[str, Any] + mode: str + hit: bool = False + saved: bool = False + prev_inductor_dir: Optional[str] = None + prev_inductor_dir_set: bool = False + + +def begin( + *, + family: Any, + transformer: Any, + dtype: Any, + quant: Any, + attention_backend: Any, + compile_kwargs: dict[str, Any], + shape_bucket: Any = None, + logger: Any = None, +) -> Optional[CacheContext]: + """Point inductor at a per-key dir and load a matching bundle, BEFORE compile. + + Returns a ``CacheContext`` to pass to ``save``/``restore``, or ``None`` when the + cache is disabled or torch lacks the Mega-cache API. Never raises. + """ + mode = cache_mode() + if mode == "off": + return None + try: + import torch # noqa: PLC0415 + if not (hasattr(torch.compiler, "save_cache_artifacts") + and hasattr(torch.compiler, "load_cache_artifacts")): + _warn(logger, "Mega-cache API unavailable (need torch >= 2.7); skipping") + return None + except Exception as exc: # noqa: BLE001 + _warn(logger, f"torch import failed: {exc}") + return None + + env_fp = environment_fingerprint() + model_fp = model_fingerprint( + family=family, transformer=transformer, dtype=dtype, quant=quant, + attention_backend=attention_backend, compile_kwargs=compile_kwargs, + shape_bucket=shape_bucket, + ) + key = cache_key(env_fp, model_fp) + cdir = cache_root() / key + ctx = CacheContext( + key=key, dir=cdir, bundle=cdir / _BUNDLE_NAME, + manifest_path=cdir / _MANIFEST_NAME, env_fp=env_fp, model_fp=model_fp, mode=mode, + ) + + # Isolate inductor's on-disk cache per key so bundles never cross-contaminate. + try: + cdir.mkdir(parents=True, exist_ok=True) + ctx.prev_inductor_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + ctx.prev_inductor_dir_set = True + os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(cdir / "inductor") + except Exception as exc: # noqa: BLE001 + _warn(logger, f"could not set TORCHINDUCTOR_CACHE_DIR: {exc}") + + # Try an exact-match load. A miss/mismatch is normal and non-fatal. + if ctx.bundle.exists() and ctx.manifest_path.exists(): + ctx.hit = _try_load(ctx, logger) + else: + _info(logger, f"compile-cache: no bundle for key {key} (will compile locally)") + return ctx + + +def _try_load(ctx: CacheContext, logger: Any) -> bool: + try: + manifest = json.loads(ctx.manifest_path.read_text()) + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: unreadable manifest: {exc}") + return False + + # Exact-match guard (defence in depth: torch also validates internally on load). + if manifest.get("env") != ctx.env_fp or manifest.get("model") != ctx.model_fp: + _warn(logger, "compile-cache: fingerprint mismatch; falling back to local compile") + return False + + try: + data = ctx.bundle.read_bytes() + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: cannot read bundle: {exc}") + return False + + # Integrity check (corruption / truncation; not a security signature). + digest = hashlib.sha256(data).hexdigest() + if manifest.get("sha256") and manifest["sha256"] != digest: + _warn(logger, "compile-cache: bundle checksum mismatch; ignoring") + return False + + try: + import torch # noqa: PLC0415 + info = torch.compiler.load_cache_artifacts(data) + if info is None: + _warn(logger, "compile-cache: load_cache_artifacts returned None (no hit)") + return False + _info(logger, f"compile-cache: loaded bundle for key {ctx.key}") + return True + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: load failed: {exc}") + return False + + +def save(ctx: Optional[CacheContext], *, logger: Any = None) -> bool: + """Persist the compiled artifacts to the bundle + manifest, AFTER a warmup forward. + + No-op unless save is enabled (mode ``on`` or ``UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE``). + Returns True if a bundle was written. + """ + if ctx is None or not _save_enabled(ctx.mode) or ctx.saved: + return False + try: + import torch # noqa: PLC0415 + result = torch.compiler.save_cache_artifacts() + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: save_cache_artifacts failed: {exc}") + return False + if not result or result[0] is None: + _warn(logger, "compile-cache: nothing to save (empty artifacts)") + return False + + data = result[0] + try: + ctx.dir.mkdir(parents=True, exist_ok=True) + ctx.bundle.write_bytes(data) + manifest = { + "format": _FORMAT_VERSION, + "key": ctx.key, + "created": time.time(), + "bytes": len(data), + "sha256": hashlib.sha256(data).hexdigest(), + "env": ctx.env_fp, + "model": ctx.model_fp, + } + ctx.manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True, default=str)) + ctx.saved = True + _info(logger, f"compile-cache: saved bundle ({len(data)} bytes) for key {ctx.key}") + return True + except Exception as exc: # noqa: BLE001 + _warn(logger, f"compile-cache: could not write bundle: {exc}") + return False + + +def restore(ctx: Optional[CacheContext]) -> None: + """Restore ``TORCHINDUCTOR_CACHE_DIR`` to its pre-load value. Call on unload.""" + if ctx is None or not ctx.prev_inductor_dir_set: + return + try: + if ctx.prev_inductor_dir is None: + os.environ.pop("TORCHINDUCTOR_CACHE_DIR", None) + else: + os.environ["TORCHINDUCTOR_CACHE_DIR"] = ctx.prev_inductor_dir + except Exception: # noqa: BLE001 + pass + + +def _warn(logger: Any, msg: str) -> None: + if logger is not None: + logger.warning("diffusion.compile_cache: %s", msg) + + +def _info(logger: Any, msg: str) -> None: + if logger is not None: + logger.info("diffusion.compile_cache: %s", msg) diff --git a/studio/backend/core/inference/diffusion_eager_patches.py b/studio/backend/core/inference/diffusion_eager_patches.py new file mode 100644 index 0000000000..cccf779b95 --- /dev/null +++ b/studio/backend/core/inference/diffusion_eager_patches.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Reversible, failure-safe monkey-patches that speed up the diffusion denoiser. + +These patch a couple of SHARED diffusers building-block classes (so a tiny patch surface +covers all four Studio families) to make inference faster WITHOUT reducing accuracy and +WITHOUT regressing ``torch.compile``. Both patches are *compile-safe*: they lower to plain +torch ops, are bit-identical or FMA-1-ULP (i.e. equal-or-MORE accurate) vs stock diffusers, +and ``torch.compile`` fuses either form to the same kernel -- so they help the EAGER path +and are neutral under compile. + +What it patches (measured on a B200, bf16, DiT shapes): + +* ``normalization.RMSNorm.forward`` -> fused ``F.rms_norm`` on the common path (non-NPU, + no bias, weight None/fp16/bf16; else the exact original, incl. its fp32 quirk). This is + the standout win (~6-12x per call) because QK-norm runs every attention block on + Qwen-Image + Z-Image. Bit-identical in bf16. + +* ``AdaLayerNormContinuous`` / ``AdaLayerNormZero`` / ``AdaLayerNormZeroSingle`` -> + the ``norm(x)*(1+scale)+shift`` modulation fused via ``torch.addcmul`` (~1.2x). fp32 + exact; bf16 within 1 ULP (FMA, a single rounding -> more accurate than mul+add). Covers + flux.1 / flux.2-klein / qwen-image. + +Deliberately NOT patched (evidence-based): +* ``FeedForward`` (skipping its eval no-op Dropout) -- measured a small REGRESSION because + the per-module ``isinstance`` check costs more than the skipped identity dispatch. +* GEGLU / SwiGLU / GELU -- already mul-bound; a real win needs a custom Triton/CUDA kernel + (out of scope, correctness risk). +* Per-family custom MLP/norm classes and attention -- not shared (no leverage), and + attention already routes through ``F.scaled_dot_product_attention`` via the existing + ``set_attention_backend`` dispatcher. + +Lifecycle: ``install_compile_safe_patches()`` is idempotent and patches at the class level. +Install it for any active speed tier; the bit-identical ``off`` reference path must run with +the patches UNINSTALLED, so the caller uninstalls on an ``off`` load and on unload. The +Studio CHAT<->DIFFUSION arbiter guarantees a single active diffusion pipe, so class-level +state is safe. +""" + +from __future__ import annotations + +import logging +import os +from typing import Callable, Optional + +import torch +import torch.nn.functional as F + +from .diffusion_patch_backend import apply_patch, revert_patch + +logger = logging.getLogger(__name__) + +# Kill-switch: set UNSLOTH_DIFFUSION_EAGER_PATCHES=0 to disable the patches entirely (for +# A/B benchmarking or to rule them out while debugging). Enabled by default. +_ENV_ENABLE = "UNSLOTH_DIFFUSION_EAGER_PATCHES" + + +def _patches_enabled() -> bool: + return (os.environ.get(_ENV_ENABLE) or "").strip().lower() not in ("0", "off", "false", "no") + +# --------------------------------------------------------------------------- # +# Resolve the diffusers classes we patch. Any import failure -> that patch is +# simply unavailable (None) and is skipped at install time. +# --------------------------------------------------------------------------- # +try: + from diffusers.models.normalization import ( + AdaLayerNormContinuous as _AdaLayerNormContinuous, + AdaLayerNormZero as _AdaLayerNormZero, + AdaLayerNormZeroSingle as _AdaLayerNormZeroSingle, + RMSNorm as _RMSNorm, + ) +except Exception: # noqa: BLE001 + _AdaLayerNormContinuous = _AdaLayerNormZero = _AdaLayerNormZeroSingle = _RMSNorm = None + +try: + from diffusers.utils.import_utils import is_torch_npu_available as _is_npu + _NPU = bool(_is_npu()) +except Exception: # noqa: BLE001 + _NPU = False + + +# --------------------------------------------------------------------------- # +# Patched forwards. Each mirrors diffusers 0.38 semantics, with the documented +# fused fast path. ``addcmul(input, t1, t2) == input + t1 * t2`` in one fused kernel. +# --------------------------------------------------------------------------- # +def _adaln_continuous_forward(self, x, conditioning_embedding): + emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + # original: self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return torch.addcmul(shift[:, None, :], self.norm(x), 1 + scale[:, None, :]) + + +def _adaln_zero_forward(self, x, timestep=None, class_labels=None, hidden_dtype=None, emb=None): + if self.emb is not None: + emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype) + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1) + # original: self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + x = torch.addcmul(shift_msa[:, None], self.norm(x), 1 + scale_msa[:, None]) + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +def _adaln_zero_single_forward(self, x, emb=None): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1) + x = torch.addcmul(shift_msa[:, None], self.norm(x), 1 + scale_msa[:, None]) + return x, gate_msa + + +# Filled in at install time with the ORIGINAL RMSNorm.forward so the guarded fast path +# can fall back for the uncommon (NPU / bias / fp32-weight) cases. +_orig_rmsnorm_forward: Optional[Callable] = None + + +def _rmsnorm_forward(self, hidden_states): + # Fall back to the exact original for cases where F.rms_norm is NOT equivalent to + # diffusers' implementation: + # * NPU / bias / fp32-weight -> the original has special handling / an fp32 quirk; + # * tuple `dim` -> diffusers always reduces the LAST dim (`mean(-1)`) while + # F.rms_norm reduces every dim in `self.dim` (differs for a multi-dim shape); + # * dtype mismatch (e.g. fp32 activations into an fp16/bf16-weight norm) -> diffusers + # computes the variance in fp32 from the ORIGINAL tensor and only casts before the + # weight multiply, so casting first would change the variance. + if ( + _NPU + or self.bias is not None + or _orig_rmsnorm_forward is None + or len(tuple(self.dim)) != 1 + ): + return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc] + weight = self.weight + if weight is None: + return F.rms_norm(hidden_states, self.dim, None, self.eps) + if weight.dtype in (torch.float16, torch.bfloat16) and hidden_states.dtype == weight.dtype: + # Common DiT path (bf16 activations + bf16 weight): F.rms_norm matches diffusers + # bit-for-bit (both reduce the variance in fp32 internally), just fused. + return F.rms_norm(hidden_states, self.dim, weight, self.eps) + # Mixed dtype / fp32 weight -> keep the exact original behaviour. + return _orig_rmsnorm_forward(self, hidden_states) # type: ignore[misc] + + +# --------------------------------------------------------------------------- # +# Install / uninstall. All swaps go through the shared patch backend +# (unsloth_zoo patch_function / restore_original): the live original is fingerprinted +# (can_safely_patch, relaxed) so a diffusers that renamed/reordered a forward's params is +# left UNPATCHED instead of miscompiled, and the original is stashed for an exact restore. +# --------------------------------------------------------------------------- # +def _specs(): + # (class, patched_fn) + return [ + (_AdaLayerNormContinuous, _adaln_continuous_forward), + (_AdaLayerNormZero, _adaln_zero_forward), + (_AdaLayerNormZeroSingle, _adaln_zero_single_forward), + (_RMSNorm, _rmsnorm_forward), + ] + + +# Classes whose `forward` we successfully patched, so uninstall reverts exactly those. +_patched: list[type] = [] + + +def install_compile_safe_patches() -> int: + """Install the shared compile-safe speedup patches (idempotent). + + Returns the number of patches applied. A second call while installed is a no-op. + """ + global _orig_rmsnorm_forward + if not _patches_enabled(): + uninstall_patches() # ensure OFF even if a prior call installed them + return 0 + if _patched: + return len(_patched) + for cls, new_fn in _specs(): + if cls is None: + continue + # Capture the live original BEFORE patching so the RMSNorm fast path can fall back + # to it for the uncommon (NPU / bias / fp32-weight / tuple-dim) cases. + if cls is _RMSNorm: + _orig_rmsnorm_forward = cls.forward + if apply_patch(cls, "forward", new_fn, match_level="relaxed"): + _patched.append(cls) + else: + logger.warning("eager-patch: skipping %s (signature mismatch / unavailable)", + getattr(cls, "__name__", cls)) + if cls is _RMSNorm: + _orig_rmsnorm_forward = None + logger.info("eager-patch: installed %d/%d shared diffusion patches", len(_patched), len(_specs())) + return len(_patched) + + +def uninstall_patches() -> None: + """Restore every patched class to its exact original ``forward`` (idempotent).""" + global _orig_rmsnorm_forward + for cls in list(_patched): + revert_patch(cls, "forward") + _patched.clear() + _orig_rmsnorm_forward = None + + +def is_installed() -> bool: + """True if any compile-safe patch is currently installed.""" + return bool(_patched) diff --git a/studio/backend/core/inference/diffusion_gguf_compile.py b/studio/backend/core/inference/diffusion_gguf_compile.py new file mode 100644 index 0000000000..995a745c9a --- /dev/null +++ b/studio/backend/core/inference/diffusion_gguf_compile.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGUF dequant accelerator for the light-compile (``default``) diffusion path. + +Profiling (outputs/profile_eager/) showed that ~70-80% of EAGER GGUF denoise CUDA +time is the per-forward weight dequant: every ``GGUFLinear.forward`` calls +``diffusers.quantizers.gguf.utils.dequantize_gguf_tensor`` -> ``dequantize_blocks_Q4_K``, +a ~20-op pure-PyTorch chain (nibble shifts + masks + block-scale mul + zero-point sub + +assembly), once per linear per step. + +COMPILED DEQUANT (``install_compiled_dequant``): swap ``dequantize_gguf_tensor`` for +``torch.compile(orig, dynamic=True)``. Inductor fuses the op chain into a few kernels. +Measured 1.24-1.64x warm with a small one-time compile (~7.5-10.4s) and ZERO extra VRAM -- +the weights stay quantized. ``dynamic=True`` is key: the dequant inputs are the WEIGHT +tensors (fixed shapes, independent of image resolution / batch), so it compiles once and +never recompiles on a resolution change. ``GGUFLinear.forward_native`` resolves the +function as a module global, so replacing the module attribute reroutes every linear. + +It is the ``default`` tier's lever (the transformer block stays eager). It is deliberately +NOT used under ``max`` (full regional block compile): there the block is compiled as one +graph which fuses the dequant inline, and a separately-compiled dequant would be traced +into that graph and break it -- so ``max`` runs the stock dequant and lets the block +compile fuse it. + +The swap goes through the shared, fingerprint-checked, reversible patch backend +(``diffusion_patch_backend``); ``uninstall_*`` restores the exact original so a later +bit-identical ``off`` load runs the stock dequant. Kill-switch: +``UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT=0``. torch / diffusers are imported lazily. + +(A global weight-buffer accelerator lived here too but was removed: it measured neutral +end-to-end -- the CUDA caching allocator already serves the per-forward cast allocation +from its pool with zero ``cudaMalloc`` churn, so reusing one buffer saved nothing on a +DiT's compute-bound forward. See outputs/arch_patch/SUMMARY.md.) +""" + +from __future__ import annotations + +import os +from typing import Any + +from .diffusion_patch_backend import apply_patch, revert_patch + +# --- kill-switch ------------------------------------------------------------------- + +_ENV_COMPILE_DEQUANT = "UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT" +_DISABLED = {"0", "off", "false", "no"} + + +def _enabled(env_name: str) -> bool: + """Enabled unless explicitly disabled (default ON).""" + return str(os.environ.get(env_name, "1")).strip().lower() not in _DISABLED + + +def _gguf_utils(): + """The diffusers GGUF utils module, or None if this diffusers build lacks it.""" + try: + from diffusers.quantizers.gguf import utils as gguf_utils # noqa: PLC0415 + + return gguf_utils + except Exception: # noqa: BLE001 — old/!GGUF diffusers -> accelerator is a no-op + return None + + +# --- compiled dequant -------------------------------------------------------------- + +# True while our compiled wrapper is installed (the shared patch backend stashes the +# original on the module for an exact restore). +_compiled_dequant_installed = False +_DEQUANT_ATTR = "dequantize_gguf_tensor" + + +def is_compiled_dequant_installed() -> bool: + return _compiled_dequant_installed + + +def install_compiled_dequant(logger: Any = None) -> bool: + """Replace ``dequantize_gguf_tensor`` with ``torch.compile(orig, dynamic=True)`` via the + shared patch backend (original stashed for restore). + + Idempotent (a second call is a no-op while installed). Returns True if the compiled + dequant is in place afterwards, False if disabled / unavailable / it failed.""" + global _compiled_dequant_installed + if not _enabled(_ENV_COMPILE_DEQUANT): + return False + if _compiled_dequant_installed: + return True + gguf_utils = _gguf_utils() + if gguf_utils is None or not hasattr(gguf_utils, _DEQUANT_ATTR): + return False + try: + import torch # noqa: PLC0415 + + compiled = torch.compile(gguf_utils.dequantize_gguf_tensor, dynamic=True) + # force=True: the new callable is the SAME function compiled, so its fingerprint + # differs from the original and can_safely_patch would (correctly) reject it. + if apply_patch(gguf_utils, _DEQUANT_ATTR, compiled, force=True): + _compiled_dequant_installed = True + return True + return False + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "install_compiled_dequant", exc) + _compiled_dequant_installed = False + return False + + +def uninstall_compiled_dequant() -> None: + """Restore the original ``dequantize_gguf_tensor``. Idempotent.""" + global _compiled_dequant_installed + if not _compiled_dequant_installed: + return + gguf_utils = _gguf_utils() + if gguf_utils is not None: + revert_patch(gguf_utils, _DEQUANT_ATTR) + _compiled_dequant_installed = False + + +# --- convenience ------------------------------------------------------------------- + + +def uninstall_all() -> None: + """Uninstall the GGUF accelerator. Idempotent; safe to call on every unload.""" + uninstall_compiled_dequant() + + +def is_installed() -> bool: + return is_compiled_dequant_installed() + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.gguf_compile: %s failed: %s", what, exc) diff --git a/studio/backend/core/inference/diffusion_patch_backend.py b/studio/backend/core/inference/diffusion_patch_backend.py new file mode 100644 index 0000000000..ea56b68310 --- /dev/null +++ b/studio/backend/core/inference/diffusion_patch_backend.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""One vetted path for every diffusion monkey-patch. + +Thin wrappers over ``unsloth_zoo.temporary_patches.utils`` ``patch_function`` / +``restore_original`` so all of the backend's runtime patching (eager fusions, GGUF +accelerators, per-arch block rewrites) goes through the SAME fingerprint-checked, +reversible mechanism: + +* ``patch_function`` stores the live original under a unique attribute and, unless + ``force=True``, runs ``can_safely_patch`` (a parameter-name/kind/required fingerprint; + ``match_level="relaxed"`` ignores type-annotation drift but still rejects a real + signature change) -- so a future diffusers/transformers that renamed or reordered a + forward's parameters is simply left unpatched instead of silently miscompiled. +* ``restore_original`` puts the stored original back -- exact, idempotent uninstall. + +``unsloth_zoo`` is imported LAZILY inside each call, never at module import: it runs GPU +detection at import time and raises without an accelerator (set ``UNSLOTH_ALLOW_CPU=1`` to +bypass, as the test conftest does), and the diffusion backend must stay importable on a +CPU-only host. This mirrors the backend's existing deferred unsloth_zoo imports (see +``core/training/trainer.py``, ``core/export/export.py``). If the import fails (unsloth_zoo +absent, or a no-GPU host without the bypass), patching is a best-effort no-op: the stock +forward runs, correctness is preserved, only the optimisation is skipped. +""" + +from __future__ import annotations + +from typing import Any + + +def apply_patch( + target: Any, + attr: str, + new_fn: Any, + *, + match_level: str = "relaxed", + force: bool = False, +) -> bool: + """Patch ``target.attr -> new_fn`` via ``unsloth_zoo`` ``patch_function`` (the original is + stashed for ``revert_patch``). Returns True iff the swap was applied. Returns False + (never raises) if unsloth_zoo is unavailable or ``can_safely_patch`` rejects the swap. + + ``force=True`` skips the safety check -- use it only when the new callable is the SAME + function transformed (e.g. its ``torch.compile`` wrapper), where a fingerprint mismatch + is expected and benign.""" + try: + from unsloth_zoo.temporary_patches.utils import patch_function + except Exception: # noqa: BLE001 — no unsloth_zoo / no-GPU host -> optimisation skipped + return False + try: + return bool( + patch_function(target, attr, new_fn, match_level=match_level, force=force) + ) + except Exception: # noqa: BLE001 — best-effort; leave the original in place + return False + + +def revert_patch(target: Any, attr: str) -> bool: + """Restore ``target.attr`` from the original stashed by ``apply_patch``. Idempotent and + best-effort: returns False (never raises) if there is nothing stored or unsloth_zoo is + unavailable.""" + try: + from unsloth_zoo.temporary_patches.utils import restore_original + except Exception: # noqa: BLE001 + return False + try: + return bool(restore_original(target, attr)) + except Exception: # noqa: BLE001 + return False diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 6ad3dcfff7..9fca6e2ba4 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -6,22 +6,34 @@ Off by default, so the default render path stays bit-identical to a plain run (the property the regression harness checks). When the operator opts in, this applies the near-lossless speedups in the order the diffusers guides recommend -(channels_last + cudnn.benchmark -> regional compile, with TF32 / fused-QKV under -"max"): +(channels_last + cudnn.benchmark -> compile, with TF32 / fused-QKV under "max"): off - nothing (default; bit-identical reference). - default - near-lossless: channels_last VAE memory format + cudnn.benchmark conv - autotune + regional torch.compile of the denoiser's repeated block WHERE - eligible (bf16, CUDA, a compile-friendly family). Compile is the big win - (~2.3x denoise on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, - well above the Q4 quantisation noise floor, so it does not meaningfully - move output quality). - max - default plus near-lossless TF32 matmul and fused QKV projections. + eager - everything lossless EXCEPT torch.compile: channels_last VAE + + cudnn.benchmark + the attention backend + the shared eager monkey-patches + (fused RMSNorm / AdaLayerNorm + per-arch addcmul fusions, see + diffusion_eager_patches.py / diffusion_arch_patches.py). The fast first-image + / casual-use path -- no compile tax to amortise. + default - LIGHT compile. For a GGUF model: channels_last + cudnn.benchmark + + torch.compile of ONLY the dequant op chain + (``torch.compile(dequantize_gguf_tensor, dynamic=True)``) -- the dequant is + ~70-80% of eager GGUF time, so fusing it gives ~1.24-1.64x for a small + one-time compile (~7.5-10.4s) and ZERO extra VRAM, resolution-invariant + (the dequant inputs are fixed-shape weights). For a dense (non-GGUF) model + there is no dequant, so ``default`` falls back to regional torch.compile of + the denoiser's repeated block (the only compile lever a dense model has). + max - the FULL torch.compile: regional max-autotune compile of the denoiser's + repeated block (which fuses the GGUF dequant AND the matmul/norm/elementwise + in one graph -- ~3.2x on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, + well above the Q4 noise floor) plus TF32 matmul and fused QKV projections. -Regional compile used to be gated off for the GGUF transformer, but it compiles and -runs faster on the current diffusers/torch (measured; the GGUF dequant ops stay -eager and the rest of the repeated block compiles), so the GGUF gate is removed; the -per-family ``supports_torch_compile`` flag and the bf16/CUDA checks still apply. +Tier rationale: ``default`` is the cheap, always-amortising compile (compile just the +hot GGUF dequant; the block stays eager) so the first image is fast and VRAM is +untouched; ``max`` pays the larger regional-compile tax for the bigger warm speedup. +The compiled dequant is deliberately skipped under ``max`` -- the regional block compile +subsumes the dequant fusion (a separately-compiled dequant would be traced into that +graph and break it), so ``max`` runs the stock dequant and lets the block compile it. The +per-family ``supports_torch_compile`` flag and the bf16/CUDA checks gate regional compile. The backend flags this layer flips (TF32, cudnn.benchmark) are PROCESS-WIDE, so ``snapshot_backend_flags`` / ``restore_backend_flags`` let the caller capture the @@ -34,10 +46,13 @@ from __future__ import annotations from typing import Any, Optional +from . import diffusion_gguf_compile as gguf_compile + SPEED_OFF = "off" +SPEED_EAGER = "eager" SPEED_DEFAULT = "default" SPEED_MAX = "max" -SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX) +SPEED_MODES = (SPEED_OFF, SPEED_EAGER, SPEED_DEFAULT, SPEED_MAX) def snapshot_backend_flags() -> Optional[dict]: @@ -85,12 +100,13 @@ def normalize_speed_mode(value: Optional[str]) -> str: def resolve_speed_mode(value: Optional[str], *, is_gguf: bool) -> str: """The effective speed mode when the caller leaves it UNSET (``None``). - A GGUF model defaults to ``default``: regional compile is ~2.2x faster and its - numeric perturbation sits well below the quantisation noise floor (measured - PSNR ~37 dB compile-vs-eager versus ~21 dB Q4-vs-bf16), so it does not reduce - output quality relative to the dense reference. A dense (non-GGUF) model stays - ``off`` / bit-identical, since there compile would be the only source of drift. - An explicit value -- including ``"off"`` -- is always honored verbatim.""" + A GGUF model defaults to ``default``: it compiles only the hot dequant op chain + (~70-80% of eager GGUF time) for ~1.24-1.64x at a small one-time compile and zero + extra VRAM -- a cheap, always-amortising win whose numeric perturbation sits well + below the quantisation noise floor (the dequant graph is unchanged, just + Inductor-fused). A dense (non-GGUF) model stays ``off`` / bit-identical, since there + compile would be the only source of drift. An explicit value -- including ``"off"`` + -- is always honored verbatim.""" if value is None: return SPEED_DEFAULT if is_gguf else SPEED_OFF return normalize_speed_mode(value) @@ -139,30 +155,47 @@ def apply_speed_optims( "tf32": False, "fused_qkv": False, "compiled": False, + "compiled_dequant": False, } mode = normalize_speed_mode(speed_mode) if mode == SPEED_OFF: return applied + on_cuda = getattr(target, "device", None) == "cuda" + family_allows_compile = bool(getattr(family, "supports_torch_compile", True)) + # Lossless: a channels-last VAE speeds up its convolutions with no numeric change. applied["channels_last"] = _vae_channels_last(pipe, logger) # Near-lossless: let cuDNN autotune the fixed-shape VAE convs (CUDA only). It may # pick a different conv algorithm, so it is a "default"-tier (not bit-identical) win. - if getattr(target, "device", None) == "cuda": + if on_cuda: applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger) - # Near-lossless and the largest win: regional compile of the repeated denoiser - # block, where eligible (now incl. the GGUF transformer). `max` opts into - # max-autotune (longer compile, autotuned kernels). - if compile_eligible(target, is_gguf = is_gguf, family = family): + # --- 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 + # transformer block stays eager. A dense model has no dequant, so default falls + # back to the regional block compile (its only compile lever). + # max = FULL compile: regional max-autotune compile of the repeated denoiser block + # (fuses dequant + matmul + norm + elementwise in one graph). It subsumes the + # dequant fusion, so we do NOT also install the standalone compiled dequant here. + # eager = no compile at all. + if mode == SPEED_DEFAULT: + if is_gguf and on_cuda and family_allows_compile: + applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger) + elif compile_eligible(target, is_gguf = is_gguf, family = family): + applied["compiled"] = _compile_repeated_blocks( + pipe, logger, max_autotune = False, cache_active = cache_active + ) + elif mode == SPEED_MAX and compile_eligible(target, is_gguf = is_gguf, family = family): applied["compiled"] = _compile_repeated_blocks( - pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active + pipe, logger, max_autotune = True, cache_active = cache_active ) if mode == SPEED_MAX: # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed. - if getattr(target, "device", None) == "cuda": + if on_cuda: applied["tf32"] = _enable_tf32(logger) applied["fused_qkv"] = _fuse_qkv(pipe, logger) diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py index b0b9ee309c..435bafe484 100644 --- a/studio/backend/tests/conftest.py +++ b/studio/backend/tests/conftest.py @@ -22,6 +22,12 @@ _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) +# Let the diffusion patch backend lazily import unsloth_zoo on a CPU-only / no-GPU test +# host: unsloth_zoo runs accelerator detection at import and raises without a GPU unless +# this is set (device_type.get_device_type checks torch.cuda first, so it is a no-op on a +# real GPU run). setdefault so an explicit override wins. +os.environ.setdefault("UNSLOTH_ALLOW_CPU", "1") + # Pytest CLI options diff --git a/studio/backend/tests/test_diffusion_arch_patches.py b/studio/backend/tests/test_diffusion_arch_patches.py new file mode 100644 index 0000000000..6a75991fc8 --- /dev/null +++ b/studio/backend/tests/test_diffusion_arch_patches.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Numerical + lifecycle tests for the per-arch eager fusions (``diffusion_arch_patches``). + +Each per-arch patch only fuses ``a + b*c`` -> ``torch.addcmul`` (1-ULP, more accurate), so +the patched method/forward must match the stock diffusers one within fp tolerance. We also +check install/uninstall reversibility + idempotency, the kill-switch, and the body-drift +guard (a diffusers whose block body changed is left unpatched). Runs on CPU. +""" + +from __future__ import annotations + +import types + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("diffusers") + +from core.inference import diffusion_arch_patches as ap # noqa: E402 + + +@pytest.fixture(autouse=True) +def _clean(): + ap.uninstall_arch_patches() + yield + ap.uninstall_arch_patches() + + +# ── qwen-image _modulate (modulation addcmul) ─────────────────────────────────── + + +def test_qwen_modulate_matches_stock_global_and_indexed(): + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q + + B, L, D = 2, 16, 64 + x = torch.randn(B, L, D) + mod = torch.randn(B, 3 * D) + # _modulate uses no real `self` state, so call it unbound with self=None. + ref_x, ref_g = Q._modulate(None, x, mod) + got_x, got_g = ap._qwen_modulate(None, x, mod) + torch.testing.assert_close(got_x, ref_x, atol=1e-5, rtol=1e-4) + assert torch.equal(got_g, ref_g) + + # per-token `index` branch (mod batch is 2*B). + idx = torch.randint(0, 2, (B, L)) + mod2 = torch.randn(2 * B, 3 * D) + ref2_x, ref2_g = Q._modulate(None, x, mod2, idx) + got2_x, got2_g = ap._qwen_modulate(None, x, mod2, idx) + torch.testing.assert_close(got2_x, ref2_x, atol=1e-5, rtol=1e-4) + assert torch.equal(got2_g, ref2_g) + + +# ── z-image block forward (gated-residual addcmul) ────────────────────────────── + + +class _AttnStub(torch.nn.Module): + """Deterministic stand-in for ZImageAttention so the block forward runs without RoPE / + freqs_cis (the patch only changes the residual adds, which is what we validate).""" + + def __init__(self, dim): + super().__init__() + self.proj = torch.nn.Linear(dim, dim, bias=False) + + def forward(self, h, **kwargs): + return self.proj(h) + + +def _zimage_block(dim=64, heads=4): + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock + + blk = ZImageTransformerBlock( + layer_id=0, dim=dim, n_heads=heads, n_kv_heads=heads, + norm_eps=1e-5, qk_norm=True, modulation=True, + ).eval() + blk.attention = _AttnStub(dim).eval() + return blk + + +def _adaln_dim(dim): + from diffusers.models.transformers.transformer_z_image import ADALN_EMBED_DIM + return min(dim, ADALN_EMBED_DIM) + + +def test_zimage_forward_matches_stock_global_modulation(): + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock + + torch.manual_seed(0) + blk = _zimage_block() + B, L, D = 2, 16, 64 + x = torch.randn(B, L, D) + adaln = torch.randn(B, _adaln_dim(D)) + + with torch.inference_mode(): + ref = ZImageTransformerBlock.forward(blk, x, None, None, adaln_input=adaln).clone() + got = ap._zimage_forward(blk, x, None, None, adaln_input=adaln) + torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-4) + + +def test_zimage_forward_matches_stock_per_token_modulation(): + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock + + torch.manual_seed(1) + blk = _zimage_block() + B, L, D = 2, 16, 64 + x = torch.randn(B, L, D) + ad = _adaln_dim(D) + adaln_noisy = torch.randn(B, ad) + adaln_clean = torch.randn(B, ad) + noise_mask = torch.randint(0, 2, (B, L)) + + with torch.inference_mode(): + ref = ZImageTransformerBlock.forward( + blk, x, None, None, noise_mask=noise_mask, + adaln_noisy=adaln_noisy, adaln_clean=adaln_clean, + ).clone() + got = ap._zimage_forward( + blk, x, None, None, noise_mask=noise_mask, + adaln_noisy=adaln_noisy, adaln_clean=adaln_clean, + ) + torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-4) + + +# ── flux.1 / flux.2 block forwards (modulation + gated-residual addcmul) ───────── + + +class _Tuple2AttnStub(torch.nn.Module): + """Double-stream attention stub -> (img_out, ctx_out).""" + + def __init__(self, dim): + super().__init__() + self.pi = torch.nn.Linear(dim, dim, bias=False) + self.pc = torch.nn.Linear(dim, dim, bias=False) + + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + return self.pi(hidden_states), self.pc(encoder_hidden_states) + + +class _SingleAttnStub(torch.nn.Module): + """Single-stream attention stub -> tensor.""" + + def __init__(self, dim): + super().__init__() + self.p = torch.nn.Linear(dim, dim, bias=False) + + def forward(self, hidden_states, **kwargs): + return self.p(hidden_states) + + +def _close_any(got, ref): + if isinstance(ref, tuple): + assert len(got) == len(ref) + for g, r in zip(got, ref): + torch.testing.assert_close(g, r, atol=1e-5, rtol=1e-4) + else: + torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-4) + + +D, H = 64, 4 +B, L, LC = 2, 16, 8 + + +def test_flux_double_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux import FluxTransformerBlock + + torch.manual_seed(0) + blk = FluxTransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval() + blk.attn = _Tuple2AttnStub(D).eval() + hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D) + with torch.inference_mode(): + ref = FluxTransformerBlock.forward(blk, hs, ehs, temb) + got = ap._flux_double_forward(blk, hs, ehs, temb) + _close_any(got, ref) + + +def test_flux_single_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock + + torch.manual_seed(1) + blk = FluxSingleTransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval() + blk.attn = _SingleAttnStub(D).eval() + hs, ehs, temb = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, D) + with torch.inference_mode(): + ref = FluxSingleTransformerBlock.forward(blk, hs, ehs, temb) + got = ap._flux_single_forward(blk, hs, ehs, temb) + _close_any(got, ref) + + +def test_flux2_double_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock + + torch.manual_seed(2) + blk = Flux2TransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval() + blk.attn = _Tuple2AttnStub(D).eval() + hs, ehs = torch.randn(B, L, D), torch.randn(B, LC, D) + tmi, tmt = torch.randn(B, 6 * D), torch.randn(B, 6 * D) + with torch.inference_mode(): + ref = Flux2TransformerBlock.forward(blk, hs, ehs, tmi, tmt) + got = ap._flux2_double_forward(blk, hs, ehs, tmi, tmt) + _close_any(got, ref) + + +def test_flux2_single_forward_matches_stock(): + from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock + + torch.manual_seed(3) + blk = Flux2SingleTransformerBlock(dim=D, num_attention_heads=H, attention_head_dim=D // H).eval() + blk.attn = _SingleAttnStub(D).eval() + hs, ehs, tm = torch.randn(B, L, D), torch.randn(B, LC, D), torch.randn(B, 3 * D) + with torch.inference_mode(): + ref = Flux2SingleTransformerBlock.forward(blk, hs, ehs, tm) + got = ap._flux2_single_forward(blk, hs, ehs, tm) + _close_any(got, ref) + + +# ── lifecycle ─────────────────────────────────────────────────────────────────── + + +def test_install_idempotent_and_reversible(): + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q + from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock as Z + + q_orig, z_orig = Q._modulate, Z.forward + n1 = ap.install_arch_patches() + n2 = ap.install_arch_patches() # idempotent + assert n1 == 6 and n2 == n1 # qwen + z-image + flux.1 x2 + flux.2 x2 + assert Q._modulate is not q_orig and Z.forward is not z_orig + assert ap.is_installed() + + ap.uninstall_arch_patches() + assert not ap.is_installed() + assert Q._modulate is q_orig and Z.forward is z_orig # exact restore + ap.uninstall_arch_patches() # idempotent + + +def test_kill_switch(monkeypatch): + from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformerBlock as Q + + monkeypatch.setenv("UNSLOTH_DIFFUSION_ARCH_PATCHES", "0") + orig = Q._modulate + assert ap.install_arch_patches() == 0 + assert not ap.is_installed() + assert Q._modulate is orig + + +def test_body_drift_guard_skips_changed_block(monkeypatch): + # If a resolver's body-check fails (diffusers changed the lines we rewrite), that patch + # is skipped. Force the qwen resolver to see a drifted body. + monkeypatch.setattr(ap, "_body_has", lambda fn, *needles: False) + assert ap.install_arch_patches() == 0 + assert not ap.is_installed() diff --git a/studio/backend/tests/test_diffusion_compile_cache.py b/studio/backend/tests/test_diffusion_compile_cache.py new file mode 100644 index 0000000000..8b12f85d0e --- /dev/null +++ b/studio/backend/tests/test_diffusion_compile_cache.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the pre-warmed torch.compile cache (``diffusion_compile_cache.py``). + +The Mega-cache API (``torch.compiler.save_cache_artifacts`` / ``load_cache_artifacts``) +is monkeypatched with deterministic in-memory fakes so the fingerprint / exact-match / +integrity / fallback / lifecycle logic is exercised without a real compile. The +fingerprint helpers run against the real torch on this box. +""" + +from __future__ import annotations + +import json +import types + +import pytest + +from core.inference import diffusion_compile_cache as cc + + +def _transformer(blocks=("FluxTransformerBlock", "FluxSingleTransformerBlock")): + return types.SimpleNamespace(_repeated_blocks=list(blocks)) + + +_BEGIN_KW = dict( + family="flux.1", + dtype="torch.bfloat16", + quant=None, + attention_backend="_native_cudnn", + compile_kwargs={"fullgraph": True, "dynamic": True}, + shape_bucket="1024x1024", +) + + +# --------------------------------------------------------------------------- fingerprint +def test_environment_fingerprint_has_hard_dimensions(): + fp = cc.environment_fingerprint() + for k in ("torch", "torch_cuda", "triton", "diffusers", "gpu_name", "gpu_capability"): + assert k in fp + + +def test_cache_key_stable_across_kwarg_order(): + efp = cc.environment_fingerprint() + t = _transformer() + a = cc.model_fingerprint(family="flux.1", transformer=t, dtype="bf16", quant=None, + attention_backend="x", compile_kwargs={"fullgraph": True, "dynamic": True}) + b = cc.model_fingerprint(family="flux.1", transformer=t, dtype="bf16", quant=None, + attention_backend="x", compile_kwargs={"dynamic": True, "fullgraph": True}) + assert cc.cache_key(efp, a) == cc.cache_key(efp, b) + + +@pytest.mark.parametrize("field,value", [ + ("family", "qwen-image"), + ("dtype", "torch.float16"), + ("quant", "int8"), + ("attention_backend", "native"), + ("shape_bucket", "512x512"), +]) +def test_cache_key_sensitive_to_model_dims(field, value): + efp = cc.environment_fingerprint() + t = _transformer() + base = dict(family="flux.1", transformer=t, dtype="bf16", quant=None, + attention_backend="x", compile_kwargs={"fullgraph": True}, shape_bucket="1024x1024") + k0 = cc.cache_key(efp, cc.model_fingerprint(**base)) + base[field] = value + assert cc.cache_key(efp, cc.model_fingerprint(**base)) != k0 + + +def test_repeated_blocks_change_key(): + efp = cc.environment_fingerprint() + k1 = cc.cache_key(efp, cc.model_fingerprint(family="f", transformer=_transformer(("A",)), + dtype="bf16", quant=None, attention_backend="x", compile_kwargs={})) + k2 = cc.cache_key(efp, cc.model_fingerprint(family="f", transformer=_transformer(("B",)), + dtype="bf16", quant=None, attention_backend="x", compile_kwargs={})) + assert k1 != k2 + + +# ----------------------------------------------------------------------------- env knobs +@pytest.mark.parametrize("raw,expected", [ + ("0", "off"), ("off", "off"), ("1", "on"), ("on", "on"), + ("auto", "auto"), ("", "auto"), ("garbage", "auto"), +]) +def test_cache_mode(monkeypatch, raw, expected): + monkeypatch.setenv(cc._ENV_MODE, raw) + assert cc.cache_mode() == expected + + +def test_cache_mode_default_auto(monkeypatch): + monkeypatch.delenv(cc._ENV_MODE, raising=False) + assert cc.cache_mode() == "auto" + + +# ------------------------------------------------------------------------------ disabled +def test_begin_returns_none_when_disabled(monkeypatch): + monkeypatch.setenv(cc._ENV_MODE, "0") + assert cc.begin(transformer=_transformer(), **_BEGIN_KW) is None + + +def test_begin_returns_none_without_megacache_api(monkeypatch): + monkeypatch.setenv(cc._ENV_MODE, "auto") + fake_torch = types.ModuleType("torch") + fake_torch.compiler = types.SimpleNamespace() # no save/load attrs + monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch) + assert cc.begin(transformer=_transformer(), **_BEGIN_KW) is None + + +# ----------------------------------------------------------------- megacache fake + flow +@pytest.fixture +def fake_megacache(monkeypatch): + """Patch torch.compiler save/load with deterministic in-memory behaviour.""" + import torch + + state = {"saved": None, "loaded_with": None} + + def fake_save(): + return (b"ARTIFACT-BYTES", None) + + def fake_load(data: bytes): + state["loaded_with"] = data + return object() if data == b"ARTIFACT-BYTES" else None + + monkeypatch.setattr(torch.compiler, "save_cache_artifacts", fake_save, raising=False) + monkeypatch.setattr(torch.compiler, "load_cache_artifacts", fake_load, raising=False) + return state + + +def test_save_then_load_roundtrip(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "on") # load + save + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + + # First load: cold (no bundle yet). + ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW) + assert ctx is not None and ctx.hit is False + assert cc.save(ctx) is True + assert ctx.bundle.exists() and ctx.manifest_path.exists() + + # Second load with the SAME fingerprint: warm hit. + ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW) + assert ctx2 is not None and ctx2.hit is True + assert fake_megacache["loaded_with"] == b"ARTIFACT-BYTES" + assert ctx2.key == ctx.key + + +def test_no_save_in_auto_mode(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "auto") + monkeypatch.delenv(cc._ENV_SAVE, raising=False) + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW) + assert cc.save(ctx) is False # auto without SAVE opt-in does not write + assert not ctx.bundle.exists() + + +def test_fingerprint_mismatch_falls_back(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "on") + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW) + cc.save(ctx) + + # Tamper the manifest's env fingerprint -> exact-match guard must reject the bundle. + manifest = json.loads(ctx.manifest_path.read_text()) + manifest["env"]["torch"] = "0.0.0-other" + ctx.manifest_path.write_text(json.dumps(manifest)) + + ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW) + assert ctx2.hit is False # mismatch -> local compile, non-fatal + + +def test_corrupt_bundle_rejected(monkeypatch, tmp_path, fake_megacache): + monkeypatch.setenv(cc._ENV_MODE, "on") + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW) + cc.save(ctx) + ctx.bundle.write_bytes(b"CORRUPTED") # manifest sha256 no longer matches + + ctx2 = cc.begin(transformer=_transformer(), **_BEGIN_KW) + assert ctx2.hit is False + + +# ------------------------------------------------------------------------------- restore +def test_restore_inductor_dir(monkeypatch, tmp_path, fake_megacache): + import os + monkeypatch.setenv(cc._ENV_MODE, "auto") + monkeypatch.setenv(cc._ENV_DIR, str(tmp_path)) + monkeypatch.setenv("TORCHINDUCTOR_CACHE_DIR", "/tmp/prior-inductor") + ctx = cc.begin(transformer=_transformer(), **_BEGIN_KW) + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] != "/tmp/prior-inductor" # redirected + cc.restore(ctx) + assert os.environ["TORCHINDUCTOR_CACHE_DIR"] == "/tmp/prior-inductor" # restored diff --git a/studio/backend/tests/test_diffusion_eager_patches.py b/studio/backend/tests/test_diffusion_eager_patches.py new file mode 100644 index 0000000000..d3f9c2afa5 --- /dev/null +++ b/studio/backend/tests/test_diffusion_eager_patches.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Numerical + lifecycle tests for the shared eager speedup patches. + +Builds the REAL diffusers 0.38 modules, captures the stock output, installs the patches, +and asserts the patched output matches within tolerance (fp32 on CPU always; bf16 on CUDA +when available). Also checks install/uninstall reversibility + idempotency, the +signature-guard no-op, and that a patched block compiles ``fullgraph=True`` (no graph break). +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("diffusers") + +import torch.nn as nn # noqa: E402 + +from core.inference import diffusion_eager_patches as ep # noqa: E402 +from diffusers.models.normalization import ( # noqa: E402 + AdaLayerNormContinuous, + AdaLayerNormZero, + AdaLayerNormZeroSingle, + RMSNorm, +) + +B, S, D, COND = 2, 16, 64, 32 + + +@pytest.fixture(autouse=True) +def _clean_patches(): + ep.uninstall_patches() + yield + ep.uninstall_patches() + + +def _devices_dtypes(): + cases = [("cpu", torch.float32)] + if torch.cuda.is_available(): + cases.append(("cuda", torch.bfloat16)) + return cases + + +def _build(cls, device, dtype): + torch.manual_seed(0) + if cls is RMSNorm: + m = RMSNorm(D, eps=1e-6, elementwise_affine=True) + elif cls is AdaLayerNormContinuous: + m = AdaLayerNormContinuous(D, COND, elementwise_affine=False, eps=1e-6, norm_type="layer_norm") + elif cls is AdaLayerNormZero: + m = AdaLayerNormZero(D, num_embeddings=None, norm_type="layer_norm") + elif cls is AdaLayerNormZeroSingle: + m = AdaLayerNormZeroSingle(D, norm_type="layer_norm") + return m.to(device=device, dtype=dtype).eval() + + +def _inputs(cls, device, dtype): + torch.manual_seed(1) + x = torch.randn(B, S, D, device=device, dtype=dtype) + if cls is RMSNorm: + return (x,) + if cls is AdaLayerNormContinuous: + return (x, torch.randn(B, COND, device=device, dtype=dtype)) + # AdaLayerNormZero / Single take the conditioning emb of width D + return (x, torch.randn(B, D, device=device, dtype=dtype)) + + +def _call(cls, m, args): + if cls is AdaLayerNormZero: + return m(args[0], emb=args[1]) + return m(*args) + + +def _first(out): + return out[0] if isinstance(out, tuple) else out + + +@pytest.mark.parametrize("cls", [RMSNorm, AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle]) +@pytest.mark.parametrize("device,dtype", _devices_dtypes()) +def test_patched_matches_original(cls, device, dtype): + m = _build(cls, device, dtype) + args = _inputs(cls, device, dtype) + + with torch.inference_mode(): + ref = _first(_call(cls, m, args)).clone() + + assert ep.install_compile_safe_patches() >= 1 + with torch.inference_mode(): + got = _first(_call(cls, m, args)) + + # The fused ops are FMA-based (addcmul) / fused (F.rms_norm): within ~1 ULP of the + # stock mul+add (and more accurate, single rounding), NOT bit-identical in fp32. + atol, rtol = (1e-5, 1e-4) if dtype == torch.float32 else (8e-3, 8e-3) + torch.testing.assert_close(got, ref, atol=atol, rtol=rtol) + + +def test_rmsnorm_mixed_dtype_falls_back(): + """fp32 activations into a bf16-weight RMSNorm: diffusers reduces variance in fp32 from + the original tensor, so the fused path must FALL BACK (identical output, not divergent).""" + m = RMSNorm(D, eps=1e-6, elementwise_affine=True).to(torch.bfloat16).eval() + x = torch.randn(B, S, D, dtype=torch.float32) + with torch.inference_mode(): + ref = m(x).clone() + ep.install_compile_safe_patches() + with torch.inference_mode(): + got = m(x) + torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) # exact fallback + + +def test_rmsnorm_tuple_dim_falls_back(): + """diffusers RMSNorm always reduces the LAST dim even for a tuple `dim`; F.rms_norm + would reduce all of them, so a multi-dim `dim` must FALL BACK to the original.""" + m = RMSNorm((2, D), eps=1e-6, elementwise_affine=True).eval() + x = torch.randn(B, 2, D) + with torch.inference_mode(): + ref = m(x).clone() + ep.install_compile_safe_patches() + with torch.inference_mode(): + got = m(x) + torch.testing.assert_close(got, ref, atol=0.0, rtol=0.0) # exact fallback + + +def test_install_idempotent_and_reversible(): + rms = RMSNorm(D, eps=1e-6) + orig = RMSNorm.forward + n1 = ep.install_compile_safe_patches() + n2 = ep.install_compile_safe_patches() # second call is a no-op + assert n1 >= 1 and n2 == n1 + assert RMSNorm.forward is not orig + assert ep.is_installed() + ep.uninstall_patches() + assert RMSNorm.forward is orig # exact restore + assert not ep.is_installed() + ep.uninstall_patches() # idempotent uninstall + del rms + + +def test_kill_switch_disables_patches(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_EAGER_PATCHES", "0") + orig = RMSNorm.forward + assert ep.install_compile_safe_patches() == 0 # no-op + assert not ep.is_installed() + assert RMSNorm.forward is orig # untouched + + +def test_signature_guard_skips_changed_class(monkeypatch): + """A diffusers class whose forward signature differs must be left untouched.""" + class WeirdRMS(nn.Module): + def forward(self, x, extra): # not (self, hidden_states) + return x + + orig = WeirdRMS.forward + monkeypatch.setattr(ep, "_RMSNorm", WeirdRMS) + monkeypatch.setattr(ep, "_AdaLayerNormContinuous", None) + monkeypatch.setattr(ep, "_AdaLayerNormZero", None) + monkeypatch.setattr(ep, "_AdaLayerNormZeroSingle", None) + applied = ep.install_compile_safe_patches() + assert applied == 0 # nothing matched -> nothing patched + assert WeirdRMS.forward is orig # left untouched + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="compile graph-break check needs CUDA") +def test_no_graph_break_under_fullgraph(): + ep.install_compile_safe_patches() + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.rms = RMSNorm(D, eps=1e-6) + self.ada = AdaLayerNormContinuous(D, COND, elementwise_affine=False, norm_type="layer_norm") + + def forward(self, x, cond): + return self.ada(self.rms(x), cond) + + m = Block().to("cuda", torch.bfloat16).eval() + x = torch.randn(B, S, D, device="cuda", dtype=torch.bfloat16) + cond = torch.randn(B, COND, device="cuda", dtype=torch.bfloat16) + compiled = torch.compile(m, fullgraph=True) # raises if a graph break occurs + with torch.inference_mode(): + out = compiled(x, cond) + assert out.shape == (B, S, D) diff --git a/studio/backend/tests/test_diffusion_gguf_compile.py b/studio/backend/tests/test_diffusion_gguf_compile.py new file mode 100644 index 0000000000..93b3fb5320 --- /dev/null +++ b/studio/backend/tests/test_diffusion_gguf_compile.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the compiled GGUF dequant accelerator (``diffusion_gguf_compile.py``). + +Covers install/uninstall idempotency + exact reversibility, the kill-switch, and the +on-by-default behaviour. Runs on CPU -- patching the module attribute is lazy +(torch.compile only traces on the first real call). +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +gguf_utils = pytest.importorskip("diffusers.quantizers.gguf.utils") + +from core.inference import diffusion_gguf_compile as gc # noqa: E402 + + +@pytest.fixture(autouse=True) +def _clean(): + # Always start and end from a clean, unpatched state so tests do not leak the + # process-wide patch into each other. + gc.uninstall_all() + yield + gc.uninstall_all() + + +def test_compiled_dequant_install_uninstall_reversible(): + orig = gguf_utils.dequantize_gguf_tensor + assert gc.is_compiled_dequant_installed() is False + + assert gc.install_compiled_dequant() is True + assert gc.is_compiled_dequant_installed() is True + # The module attribute is now a different (compiled) callable... + assert gguf_utils.dequantize_gguf_tensor is not orig + # ...idempotent: a second install is a no-op, attribute unchanged. + patched = gguf_utils.dequantize_gguf_tensor + assert gc.install_compiled_dequant() is True + assert gguf_utils.dequantize_gguf_tensor is patched + + gc.uninstall_compiled_dequant() + assert gc.is_compiled_dequant_installed() is False + # Exact original restored. + assert gguf_utils.dequantize_gguf_tensor is orig + # Uninstall is idempotent. + gc.uninstall_compiled_dequant() + assert gguf_utils.dequantize_gguf_tensor is orig + + +def test_compiled_dequant_kill_switch(monkeypatch): + monkeypatch.setenv("UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT", "0") + orig = gguf_utils.dequantize_gguf_tensor + assert gc.install_compiled_dequant() is False + assert gc.is_compiled_dequant_installed() is False + assert gguf_utils.dequantize_gguf_tensor is orig + + +def test_compiled_dequant_on_by_default(monkeypatch): + # The compiled dequant is the real win, so it is ON without any env opt-in. + monkeypatch.delenv("UNSLOTH_DIFFUSION_GGUF_COMPILE_DEQUANT", raising=False) + assert gc.install_compiled_dequant() is True + assert gc.is_compiled_dequant_installed() is True + + +def test_uninstall_all(monkeypatch): + orig = gguf_utils.dequantize_gguf_tensor + gc.install_compiled_dequant() + assert gc.is_installed() is True + gc.uninstall_all() + assert gc.is_installed() is False + assert gguf_utils.dequantize_gguf_tensor is orig diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index fb7c0fdb79..b112ac749d 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -14,8 +14,10 @@ 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, @@ -27,6 +29,20 @@ from core.inference.diffusion_speed import ( ) +def _stub_gguf_accel(monkeypatch): + """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.""" + called = {"compiled_dequant": 0} + + def _install(logger = None): + called["compiled_dequant"] += 1 + return True + + monkeypatch.setattr(ds_mod.gguf_compile, "install_compiled_dequant", _install) + return called + + def _target( *, device = "cuda", @@ -158,14 +174,18 @@ def test_speed_off_applies_nothing(monkeypatch): "tf32": False, "fused_qkv": False, "compiled": False, + "compiled_dequant": 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). assert torch.backends.cudnn.benchmark is False -def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): +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 (its only compile lever) -- and no GGUF accelerators. torch = _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_DEFAULT @@ -178,17 +198,55 @@ def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): # default also autotunes the VAE convs but does NOT flip TF32 or fuse QKV. assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True assert applied["tf32"] is False and applied["fused_qkv"] is False + # No GGUF dequant on a dense model. + assert applied["compiled_dequant"] is False + assert called == {"compiled_dequant": 0} -def test_speed_default_compiles_gguf(monkeypatch): +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. _stub_torch(monkeypatch) + called = _stub_gguf_accel(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) assert applied["channels_last"] is True - # GGUF now compiles (the big near-lossless win). + assert applied["compiled_dequant"] is True + # The transformer block is NOT regionally compiled under GGUF default. + assert applied["compiled"] is False and pipe.compiled is False + assert called == {"compiled_dequant": 1} + + +def test_speed_eager_gguf_installs_no_accelerator(monkeypatch): + # eager = lossless-but-no-compile: neither the compiled dequant nor the regional + # block compile run; only the process-wide lossless levers (channels_last, cudnn) + # and the shared/per-arch eager monkey-patches (installed elsewhere) engage. + _stub_torch(monkeypatch) + called = _stub_gguf_accel(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_EAGER + ) + assert applied["compiled_dequant"] is False and applied["compiled"] is False + assert pipe.compiled is False + assert called == {"compiled_dequant": 0} + + +def test_speed_max_gguf_regional_compile_not_dequant(monkeypatch): + # GGUF `max` = the FULL regional block compile (which fuses the dequant inline), so + # the standalone compiled dequant is deliberately OFF. + _stub_torch(monkeypatch) + called = _stub_gguf_accel(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["compiled"] is True and pipe.compiled is True + assert pipe.compile_kwargs["mode"] == "max-autotune-no-cudagraphs" + assert applied["compiled_dequant"] is False + assert called == {"compiled_dequant": 0} def test_speed_default_cudnn_benchmark_only_on_cuda(monkeypatch):