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.
This commit is contained in:
Daniel Han
parent
8b385c44fb
commit
f24384b4e9
12 changed files with 2037 additions and 31 deletions
469
studio/backend/core/inference/diffusion_arch_patches.py
Normal file
469
studio/backend/core/inference/diffusion_arch_patches.py
Normal file
|
|
@ -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)
|
||||
338
studio/backend/core/inference/diffusion_compile_cache.py
Normal file
338
studio/backend/core/inference/diffusion_compile_cache.py
Normal file
|
|
@ -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)
|
||||
203
studio/backend/core/inference/diffusion_eager_patches.py
Normal file
203
studio/backend/core/inference/diffusion_eager_patches.py
Normal file
|
|
@ -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)
|
||||
132
studio/backend/core/inference/diffusion_gguf_compile.py
Normal file
132
studio/backend/core/inference/diffusion_gguf_compile.py
Normal file
|
|
@ -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)
|
||||
70
studio/backend/core/inference/diffusion_patch_backend.py
Normal file
70
studio/backend/core/inference/diffusion_patch_backend.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
252
studio/backend/tests/test_diffusion_arch_patches.py
Normal file
252
studio/backend/tests/test_diffusion_arch_patches.py
Normal file
|
|
@ -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()
|
||||
189
studio/backend/tests/test_diffusion_compile_cache.py
Normal file
189
studio/backend/tests/test_diffusion_compile_cache.py
Normal file
|
|
@ -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
|
||||
183
studio/backend/tests/test_diffusion_eager_patches.py
Normal file
183
studio/backend/tests/test_diffusion_eager_patches.py
Normal file
|
|
@ -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)
|
||||
73
studio/backend/tests/test_diffusion_gguf_compile.py
Normal file
73
studio/backend/tests/test_diffusion_gguf_compile.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue