705 lines
32 KiB
Python
705 lines
32 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Select the diffusion transformer's attention backend.
|
|
|
|
diffusers' ``transformer.set_attention_backend(name)`` dispatcher swaps the SDPA kernel,
|
|
validating hardware/package at set time (default ``native`` = ``F.scaled_dot_product_attention``).
|
|
Attention is bandwidth-bound, so a better kernel is a real win orthogonal to weight quantisation
|
|
(it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile.
|
|
|
|
auto - the best *exact* backend for the device. On NVIDIA CUDA that is cuDNN fused attention
|
|
(``_native_cudnn``), ~1.18x end-to-end on B200, LPIPS ~0.004 (below the noise floor).
|
|
Elsewhere stays ``native``. Only upgrades when a speed profile is active, so
|
|
``speed_mode=off`` stays bit-identical.
|
|
native - force the default SDPA (bit-identical reference).
|
|
cudnn - cuDNN fused attention (exact; NVIDIA).
|
|
flash / flash3 / flash4 - FlashAttention 2 / 3 (Hopper) / 4 (SM100); exact, kernel-gated.
|
|
sage - SageAttention (INT8 QK); quantized, small quality cost, consumer-friendly.
|
|
xformers / aiter - memory-efficient (NVIDIA) / AITER (AMD ROCm).
|
|
|
|
Best-effort: an unavailable backend falls back to the diffusers default. torch/diffusers lazy.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
ATTN_AUTO = "auto"
|
|
ATTN_NATIVE = "native"
|
|
|
|
# User-facing alias -> the diffusers dispatcher backend name.
|
|
_ALIASES: dict[str, str] = {
|
|
"native": "native",
|
|
"sdpa": "native",
|
|
"cudnn": "_native_cudnn",
|
|
"flash": "flash",
|
|
"flash2": "flash",
|
|
"flash3": "_flash_3_hub",
|
|
"flash4": "flash_4_hub",
|
|
"sage": "sage",
|
|
"xformers": "xformers",
|
|
"aiter": "aiter",
|
|
}
|
|
ATTN_ALIASES = (ATTN_AUTO,) + tuple(dict.fromkeys(_ALIASES))
|
|
|
|
|
|
def normalize_attention_backend(value: Optional[str]) -> Optional[str]:
|
|
"""Lower/strip a requested backend; None / "" / "auto" -> "auto". Raises ValueError for an
|
|
unsupported alias so a bad request is rejected cheaply."""
|
|
if value is None:
|
|
return ATTN_AUTO
|
|
normalized = str(value).strip().lower()
|
|
if not normalized:
|
|
return ATTN_AUTO
|
|
if normalized not in ATTN_ALIASES:
|
|
raise ValueError(
|
|
f"Unsupported attention_backend '{value}'. Use one of: {', '.join(ATTN_ALIASES)}."
|
|
)
|
|
return normalized
|
|
|
|
|
|
# Backends diffusers validates only by package at set time but whose kernels need a specific
|
|
# CUDA arch at run time (so an explicit request on the wrong card sets fine then crashes
|
|
# mid-generation). Gate by a (min, max-exclusive) capability range: FA3 is Hopper-SM90 only
|
|
# (upper bound, so flash3 on a B200 drops to native), FA4 is Blackwell+ (no upper bound).
|
|
_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = {
|
|
"flash": ((8, 0), None), # FlashAttention 2 -> Ampere (SM80)+
|
|
"_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only
|
|
"flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+
|
|
}
|
|
|
|
|
|
def _cuda_capability() -> Optional[tuple[int, int]]:
|
|
"""(major, minor) compute capability of the active CUDA device, or None if unknown."""
|
|
try:
|
|
import torch
|
|
if not torch.cuda.is_available():
|
|
return None
|
|
return tuple(torch.cuda.get_device_capability()) # type: ignore[return-value]
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
|
|
|
|
def _backend_arch_supported(backend: str) -> bool:
|
|
"""False only when ``backend`` needs a CUDA arch outside this device's range. Unknown
|
|
capability returns True (never block on a guess; the run-time failure falls back to native)."""
|
|
bounds = _ARCH_CAPABILITY.get(backend)
|
|
if bounds is None:
|
|
return True
|
|
have = _cuda_capability()
|
|
if have is None:
|
|
return True
|
|
low, high = bounds
|
|
return have >= low and (high is None or have < high)
|
|
|
|
|
|
def _is_cuda_nvidia(target: Any) -> bool:
|
|
"""CUDA device on an NVIDIA (non-ROCm) build -- where cuDNN attention applies."""
|
|
if getattr(target, "device", None) != "cuda":
|
|
return False
|
|
try:
|
|
import torch
|
|
return getattr(torch.version, "hip", None) is None
|
|
except Exception: # noqa: BLE001
|
|
return False
|
|
|
|
|
|
def select_attention_backend(
|
|
target: Any, requested: Optional[str], *, speed_active: bool
|
|
) -> Optional[str]:
|
|
"""The dispatcher backend name to apply, or None to leave the diffusers default.
|
|
|
|
An explicit alias is honored (apply falls back if its kernel is unavailable). ``auto``
|
|
upgrades to cuDNN on NVIDIA CUDA only when a speed profile is active (so ``off`` stays
|
|
bit-identical); elsewhere returns None (native)."""
|
|
alias = normalize_attention_backend(requested)
|
|
if alias != ATTN_AUTO:
|
|
backend = _ALIASES[alias]
|
|
if backend == "native":
|
|
return None
|
|
# AITER is the AMD ROCm kernel: honor it on a ROCm CUDA target, drop it elsewhere (else
|
|
# the NVIDIA-only guard below would drop the one backend that only works on ROCm).
|
|
if backend == "aiter":
|
|
if getattr(target, "device", None) == "cuda" and not _is_cuda_nvidia(target):
|
|
return backend
|
|
return None
|
|
# cuDNN / flash* / sage are CUDA+NVIDIA-only; elsewhere the first generation crashes.
|
|
if not _is_cuda_nvidia(target):
|
|
return None
|
|
# An arch-gated kernel (flash3/flash4) on a card that can't run it sets fine then crashes.
|
|
if not _backend_arch_supported(backend):
|
|
return None
|
|
# cuDNN fused SDPA needs Ampere+ (SM80); gate an explicit request like the auto path.
|
|
if backend == "_native_cudnn" and not _cudnn_attention_supported():
|
|
return None
|
|
return backend
|
|
# auto
|
|
if speed_active and _is_cuda_nvidia(target) and _cudnn_attention_supported():
|
|
return "_native_cudnn"
|
|
return None
|
|
|
|
|
|
def _cudnn_attention_supported() -> bool:
|
|
"""cuDNN fused SDPA needs Ampere+ (SM80); on pre-SM80 cards (T4/V100) diffusers accepts it
|
|
then fails at generation, so gate the upgrade on capability. Unknown capability allows it."""
|
|
have = _cuda_capability()
|
|
return have is None or have >= (8, 0)
|
|
|
|
|
|
def attention_backend_supported_on_device(backend: Optional[str], device_index: int) -> bool:
|
|
"""Whether an already-resolved dispatcher backend can actually RUN on CUDA ``device_index``.
|
|
|
|
``select_attention_backend`` arch-gates a backend against the ACTIVE device, but a CFG-parallel
|
|
replica lives on a possibly HETEROGENEOUS second GPU (FA3 is Hopper-SM90 only, FA4 needs
|
|
Blackwell-SM100, cuDNN needs Ampere+). Installing the primary-resolved backend there without
|
|
re-checking would set fine then crash on the replica's first attention kernel. Re-applies the
|
|
same arch gate to a specific device index. None (native) is always fine; an unqueryable
|
|
capability returns True (best-effort, matching ``_backend_arch_supported``)."""
|
|
if backend is None:
|
|
return True
|
|
try:
|
|
import torch
|
|
have = tuple(torch.cuda.get_device_capability(device_index)) # type: ignore[assignment]
|
|
except Exception: # noqa: BLE001 -- unqueryable device: don't block on a guess
|
|
return True
|
|
bounds = _ARCH_CAPABILITY.get(backend)
|
|
if bounds is not None:
|
|
low, high = bounds
|
|
if not (have >= low and (high is None or have < high)):
|
|
return False
|
|
if backend == "_native_cudnn" and have < (8, 0):
|
|
return False
|
|
return True
|
|
|
|
|
|
# Optional-kernel backends installable on demand: dispatcher name -> (probe module, pip
|
|
# package). Wheels only (--only-binary=:all:): a source build needs a CUDA toolchain a Studio
|
|
# host may lack; no wheel means a native fallback. cuDNN/native ship with torch.
|
|
_INSTALLABLE_BACKENDS: dict[str, tuple[str, str]] = {
|
|
"sage": ("sageattention", "sageattention"),
|
|
"flash": ("flash_attn", "flash-attn"),
|
|
"_flash_3_hub": ("kernels", "kernels"), # FA3/FA4 from the HF kernels hub
|
|
"flash_4_hub": ("kernels", "kernels"),
|
|
"xformers": ("xformers", "xformers"),
|
|
}
|
|
|
|
# On-demand install gate (mirrors UNSLOTH_DIFFUSION_SD_CPP_INSTALL):
|
|
# auto (default) / 1 - install the missing package when a gated backend is requested
|
|
# 0 - never install; a missing kernel falls back to native
|
|
_ATTENTION_INSTALL_ENV = "UNSLOTH_DIFFUSION_ATTENTION_INSTALL"
|
|
|
|
# Packages a pip install was already attempted for in THIS process (success or failure). The
|
|
# loader pre-installs outside its locks, then re-resolves under _generate_lock where apply would
|
|
# otherwise call pip a SECOND time -- a no-wheel/offline host would re-run the full 600s install
|
|
# holding the load lock, blocking unload/cancel. A recorded attempt makes the retry a no-op.
|
|
_INSTALL_ATTEMPTED: set[str] = set()
|
|
|
|
|
|
def _kernels_hub_compatible(logger: Any = None) -> bool:
|
|
"""Whether installing the ``kernels`` package is SAFE next to the resident huggingface_hub.
|
|
Every ``kernels`` release (>= 0.13) needs hub >= 1.0's strict-dataclass API, and with an older
|
|
hub the breakage is NOT contained: ``import kernels`` raises at module scope and diffusers
|
|
imports ``kernels`` whenever installed, so EVERY later pipeline import crashes until it is
|
|
uninstalled (measured: hub 0.36 + kernels 0.13/0.16 both brick the HunyuanVideo-1.5 import).
|
|
So hub < 1.0 stacks must not auto-install it; the hub backend falls back to native. An
|
|
undeterminable hub version allows the install (previous behaviour)."""
|
|
try:
|
|
from importlib.metadata import version
|
|
return int(version("huggingface_hub").split(".", 1)[0]) >= 1
|
|
except Exception: # noqa: BLE001 -- unknown hub -> keep the previous behaviour
|
|
return True
|
|
|
|
|
|
def _ensure_attention_backend_installed(backend: str, logger: Any = None) -> None:
|
|
"""Best-effort wheel-only install of the package ``backend`` needs, when allowed.
|
|
|
|
Called after arch gating, so only for a backend that could work here. Failure is swallowed:
|
|
the subsequent set_attention_backend raises on the missing package and falls back to native."""
|
|
import importlib.util
|
|
import os
|
|
|
|
spec = _INSTALLABLE_BACKENDS.get(backend)
|
|
if spec is None:
|
|
return
|
|
module, package = spec
|
|
gate = os.environ.get(_ATTENTION_INSTALL_ENV, "auto").strip().lower()
|
|
if gate in ("0", "false", "no", "off"):
|
|
return
|
|
if package == "kernels" and not _kernels_hub_compatible(logger):
|
|
if logger is not None:
|
|
logger.warning(
|
|
"diffusion.attention: not installing 'kernels' for backend=%s -- the "
|
|
"resident huggingface_hub is < 1.0 and a kernels install would break "
|
|
"every later diffusers pipeline import; using the default backend",
|
|
backend,
|
|
)
|
|
return
|
|
try:
|
|
if importlib.util.find_spec(module) is not None:
|
|
return
|
|
except Exception: # noqa: BLE001 — a broken install probes as missing; try the install
|
|
pass
|
|
# Attempt each install once per process (see _INSTALL_ATTEMPTED): else the in-lock apply path
|
|
# re-runs the whole install under _generate_lock and blocks unload/cancel.
|
|
if package in _INSTALL_ATTEMPTED:
|
|
return
|
|
_INSTALL_ATTEMPTED.add(package)
|
|
import subprocess
|
|
import sys
|
|
|
|
if logger is not None:
|
|
logger.info(
|
|
"diffusion.attention: installing %s for backend=%s (wheel-only)", package, backend
|
|
)
|
|
try:
|
|
subprocess.run(
|
|
# --no-deps: install ONLY this kernel wheel. xformers/flash-attn pin an exact torch,
|
|
# so normal resolution would replace the running torch/triton. Without deps an
|
|
# ABI-incompatible kernel just fails to import -> native fallback.
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"pip",
|
|
"install",
|
|
"--only-binary",
|
|
":all:",
|
|
"--no-deps",
|
|
package,
|
|
],
|
|
capture_output = True,
|
|
timeout = 600,
|
|
check = True,
|
|
)
|
|
# The import system caches directory listings, so the next find_spec can miss the wheel
|
|
# just installed (mtime resolution). Invalidate the finder caches so it's picked up now.
|
|
importlib.invalidate_caches()
|
|
except Exception as exc: # noqa: BLE001 — no wheel / no network -> native fallback
|
|
if logger is not None:
|
|
# CalledProcessError.str() shows only the exit code; the real reason is in stderr.
|
|
# Surface it so the native fallback is diagnosable.
|
|
stderr = getattr(exc, "stderr", None)
|
|
if stderr:
|
|
if isinstance(stderr, bytes):
|
|
stderr = stderr.decode("utf-8", errors = "replace")
|
|
logger.warning(
|
|
"diffusion.attention: could not install %s; pip failed with: %s",
|
|
package,
|
|
stderr.strip() or str(exc),
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"diffusion.attention: could not install %s (%s); falling back to default",
|
|
package,
|
|
exc,
|
|
)
|
|
|
|
|
|
def _attention_dits(pipe: Any) -> list:
|
|
"""Every DiT the denoise loop runs: the primary ``transformer`` plus a second expert some
|
|
families carry (Ideogram's ``unconditional_transformer``, an MoE ``transformer_2``). The
|
|
backend must be set on ALL of them, else the second DiT keeps the native default."""
|
|
dits: list = []
|
|
for attr in ("transformer", "transformer_2", "unconditional_transformer"):
|
|
m = getattr(pipe, attr, None)
|
|
if m is not None and m not in dits:
|
|
dits.append(m)
|
|
return dits
|
|
|
|
|
|
def apply_attention_backend(
|
|
pipe: Any,
|
|
backend: Optional[str],
|
|
*,
|
|
logger: Any = None,
|
|
) -> Optional[str]:
|
|
"""Set ``backend`` on EVERY denoiser DiT via the diffusers dispatcher.
|
|
|
|
Returns the backend engaged, or None when left at native (``backend`` was None or the kernel
|
|
was unavailable -> graceful fallback, never a load failure).
|
|
|
|
diffusers keeps a process-wide active backend that ``set_attention_backend`` also updates, and
|
|
a fresh transformer's processors follow it (default None). So a load wanting native must
|
|
restore it explicitly, else it inherits a backend an earlier load pinned (e.g. cuDNN under a
|
|
speed profile), breaking the ``off`` guarantee. Best-effort."""
|
|
setters = [
|
|
s
|
|
for s in (getattr(t, "set_attention_backend", None) for t in _attention_dits(pipe))
|
|
if callable(s)
|
|
]
|
|
if not setters:
|
|
return None
|
|
if backend is not None:
|
|
_ensure_attention_backend_installed(backend, logger)
|
|
engaged = False
|
|
for fn in setters:
|
|
try:
|
|
fn(backend)
|
|
engaged = True
|
|
except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below
|
|
_warn(logger, backend, exc)
|
|
if engaged:
|
|
# set_attention_backend also pins the backend process-wide. Each DiT's processors now
|
|
# keep it locally, so reset the global to native ONCE, else a later unconfigured
|
|
# component inherits this kernel.
|
|
_reset_global_backend_to_native(logger)
|
|
if logger is not None:
|
|
logger.info("diffusion.attention: backend=%s", backend)
|
|
return backend
|
|
# No backend requested, or every set failed: pin native so a stale process-wide backend can't
|
|
# leak in. Fresh DiTs follow the global, so one reset via any setter covers them all.
|
|
_restore_native_backend(setters[0], logger)
|
|
return None
|
|
|
|
|
|
def _active_attention_backend() -> Optional[str]:
|
|
"""The diffusers process-wide active attention backend name, or None if undeterminable."""
|
|
try:
|
|
from diffusers.models.attention_dispatch import _AttentionBackendRegistry
|
|
|
|
# get_active_backend() returns (AttentionBackendName, fn) or None; take element 0 and
|
|
# read its .value ("native"), not off the tuple (which never compares equal to a name).
|
|
active = _AttentionBackendRegistry.get_active_backend()
|
|
if active is None:
|
|
return None
|
|
name = active[0] if isinstance(active, tuple) else active
|
|
return getattr(name, "value", str(name))
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
|
|
|
|
def _reset_global_backend_to_native(logger: Any) -> None:
|
|
"""Reset the process-wide active backend to native after a successful per-transformer set, so
|
|
a later unconfigured component doesn't inherit this kernel (the DiT's own processors keep it).
|
|
Best-effort: if the diffusers internals move, the prior (leaking) behavior is unchanged."""
|
|
if _active_attention_backend() == ATTN_NATIVE:
|
|
return
|
|
try:
|
|
from diffusers.models.attention_dispatch import (
|
|
AttentionBackendName,
|
|
_AttentionBackendRegistry,
|
|
)
|
|
_AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE)
|
|
except Exception: # noqa: BLE001 — best-effort; leave the global as-is on any change
|
|
pass
|
|
|
|
|
|
def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None:
|
|
"""Force the native default when the global active backend isn't already native."""
|
|
if _active_attention_backend() == ATTN_NATIVE:
|
|
return # already native -> avoid redundant work and an extra dispatcher warning
|
|
try:
|
|
set_backend_fn(ATTN_NATIVE)
|
|
except Exception as exc: # noqa: BLE001 — best-effort restore
|
|
_warn(logger, ATTN_NATIVE, exc)
|
|
|
|
|
|
def _warn(logger: Any, what: str, exc: Exception) -> None:
|
|
if logger is not None:
|
|
logger.warning("diffusion.attention: %s unavailable (%s); using default", what, exc)
|
|
|
|
|
|
# --------------------------------------------------------------------------------------
|
|
# HunyuanVideo-1.5 joint-attention padding trim (accuracy-exact speed win)
|
|
#
|
|
# HunyuanVideo15AttnProcessor2_0 runs a JOINT [video ; text] self-attention and, on EVERY block
|
|
# and step, materialises a dense [B,1,N,N] boolean mask so the video never attends to padded text.
|
|
# But a dense bool attn_mask DISABLES every fused SDPA kernel (flash rejects it; cuDNN/efficient
|
|
# fall back) and forces the slow math path: on a B200 at the production shape (N~=50k, 121 frames
|
|
# 480p) the SAME attention is 421 ms WITH the dense mask vs 19 ms with attn_mask=None -- a ~22x tax
|
|
# purely to mask padding. And the text is ~99.5% padding: a t2v prompt fills only ~9 of ~1985 slots
|
|
# (image 729 + byt5 256 + mllm 1000, almost all zero-padded).
|
|
#
|
|
# The fix is exact: the model already masks the padded text and DISCARDS its attention output (only
|
|
# the video split feeds proj_out), so removing the padded tokens before attention changes nothing
|
|
# for the video. Done in an eager forward pre-hook (outside the compiled blocks): drop the all-zero
|
|
# image stream (t2v), trim the mllm/byt5 streams to their globally-valid columns, and -- when
|
|
# nothing partially-padded remains (the common batch-1 / per-branch call) -- flag the DiT so the
|
|
# processor skips the dense mask and runs the fused path. The only numeric change is the SDPA kernel
|
|
# (masked -> fused), on par with the shipped cuDNN backend swap. Mixed-padding batches fall back to
|
|
# the stock dense mask.
|
|
_HUNYUAN15_TRANSFORMER_CLS = "HunyuanVideo15Transformer3DModel"
|
|
_HUNYUAN15_PROCESSOR_CLS = "HunyuanVideo15AttnProcessor2_0"
|
|
_NULL_ATTN_FLAG = "_unsloth_null_attn_mask"
|
|
|
|
_NULL_PROCESSOR_CACHE: dict = {}
|
|
|
|
|
|
def _set_hunyuan_null_mask(module: Any, enabled: bool) -> None:
|
|
"""Set the null-mask flag on every block's attention of ``module``. The flag is valid ONLY for
|
|
the forward whose pre-hook removed the padding, so a post-hook clears it back to False after
|
|
each call (see the module note and _hunyuan_trim_post_hook)."""
|
|
for blk in getattr(module, "transformer_blocks", []):
|
|
attn = getattr(blk, "attn", None)
|
|
if attn is not None:
|
|
setattr(attn, _NULL_ATTN_FLAG, enabled)
|
|
|
|
|
|
def _null_mask_processor_cls():
|
|
"""Build (once, lazily) a HunyuanVideo15AttnProcessor2_0 subclass whose ``__call__`` runs
|
|
attn_mask=None when the DiT is flagged (padding already removed by the pre-hook); otherwise it
|
|
delegates to the stock processor, so a mixed-padding batch and future diffusers changes stay
|
|
correct."""
|
|
cached = _NULL_PROCESSOR_CACHE.get("cls")
|
|
if cached is not None:
|
|
return cached
|
|
|
|
import torch
|
|
from diffusers.models.attention_dispatch import dispatch_attention_fn
|
|
from diffusers.models.transformers.transformer_hunyuan_video15 import (
|
|
HunyuanVideo15AttnProcessor2_0,
|
|
)
|
|
|
|
class _HunyuanNullMaskProcessor(HunyuanVideo15AttnProcessor2_0):
|
|
def __call__(
|
|
self,
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states = None,
|
|
attention_mask = None,
|
|
image_rotary_emb = None,
|
|
):
|
|
# Fast path only when the pre-hook removed all padding (attn_mask redundant); a
|
|
# constant python bool so torch.compile const-folds the branch (no graph break).
|
|
if not getattr(attn, _NULL_ATTN_FLAG, False):
|
|
return super().__call__(
|
|
attn,
|
|
hidden_states,
|
|
encoder_hidden_states = encoder_hidden_states,
|
|
attention_mask = attention_mask,
|
|
image_rotary_emb = image_rotary_emb,
|
|
)
|
|
|
|
# Null path = the stock body with the mask block removed and attn_mask=None.
|
|
query = attn.to_q(hidden_states)
|
|
key = attn.to_k(hidden_states)
|
|
value = attn.to_v(hidden_states)
|
|
|
|
query = query.unflatten(2, (attn.heads, -1))
|
|
key = key.unflatten(2, (attn.heads, -1))
|
|
value = value.unflatten(2, (attn.heads, -1))
|
|
|
|
query = attn.norm_q(query)
|
|
key = attn.norm_k(key)
|
|
|
|
if image_rotary_emb is not None:
|
|
from diffusers.models.embeddings import apply_rotary_emb
|
|
query = apply_rotary_emb(query, image_rotary_emb, sequence_dim = 1)
|
|
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim = 1)
|
|
|
|
if encoder_hidden_states is not None:
|
|
encoder_query = attn.add_q_proj(encoder_hidden_states)
|
|
encoder_key = attn.add_k_proj(encoder_hidden_states)
|
|
encoder_value = attn.add_v_proj(encoder_hidden_states)
|
|
|
|
encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
|
|
encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
|
|
encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
|
|
|
|
if attn.norm_added_q is not None:
|
|
encoder_query = attn.norm_added_q(encoder_query)
|
|
if attn.norm_added_k is not None:
|
|
encoder_key = attn.norm_added_k(encoder_key)
|
|
|
|
query = torch.cat([query, encoder_query], dim = 1)
|
|
key = torch.cat([key, encoder_key], dim = 1)
|
|
value = torch.cat([value, encoder_value], dim = 1)
|
|
|
|
hidden_states = dispatch_attention_fn(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_mask = None,
|
|
dropout_p = 0.0,
|
|
is_causal = False,
|
|
backend = self._attention_backend,
|
|
parallel_config = self._parallel_config,
|
|
)
|
|
|
|
hidden_states = hidden_states.flatten(2, 3)
|
|
hidden_states = hidden_states.to(query.dtype)
|
|
|
|
if encoder_hidden_states is not None:
|
|
enc_len = encoder_hidden_states.shape[1]
|
|
hidden_states, encoder_hidden_states = (
|
|
hidden_states[:, :-enc_len],
|
|
hidden_states[:, -enc_len:],
|
|
)
|
|
if getattr(attn, "to_out", None) is not None:
|
|
hidden_states = attn.to_out[0](hidden_states)
|
|
hidden_states = attn.to_out[1](hidden_states)
|
|
if getattr(attn, "to_add_out", None) is not None:
|
|
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
|
|
return hidden_states, encoder_hidden_states
|
|
|
|
return hidden_states
|
|
|
|
_NULL_PROCESSOR_CACHE["cls"] = _HunyuanNullMaskProcessor
|
|
return _HunyuanNullMaskProcessor
|
|
|
|
|
|
def _trim_stream(states, mask):
|
|
"""Drop the columns of a [B, S, D] text stream + its [B, S] mask that are padding for EVERY
|
|
batch element (globally invalid). Returns (states, mask, all_valid): all_valid is True when
|
|
the trimmed stream has NO partially-padded column left (so it needs no attention mask)."""
|
|
import torch
|
|
|
|
if states is None or mask is None or mask.dim() != 2:
|
|
return states, mask, True # nothing to mask -> treat as no-padding
|
|
mb = mask.bool()
|
|
keep = mb.any(dim = 0) # column valid for at least one batch element
|
|
if not bool(keep.all()):
|
|
states = states[:, keep]
|
|
mask = mask[:, keep]
|
|
mb = mb[:, keep]
|
|
# All remaining slots valid for every element (vacuously True for a 0-length stream, fine
|
|
# for an unused secondary stream e.g. byt5 in t2v).
|
|
all_valid = bool(mb.all().item())
|
|
return states, mask, all_valid
|
|
|
|
|
|
def _hunyuan_trim_pre_hook(module, args, kwargs):
|
|
"""Eager forward pre-hook: strip padded text tokens so the joint attention runs fused.
|
|
|
|
- Drop the image stream when it is entirely zero (t2v): those ~729 tokens are pure padding.
|
|
- Trim the mllm/byt5 text streams to their globally-valid columns.
|
|
- Flag every block's attention so the null-mask processor skips the dense mask when nothing
|
|
partially-padded remains (the batch-1 / per-guidance-branch case); otherwise leave the
|
|
flag False and the stock dense-mask path handles the residual padding correctly.
|
|
|
|
This hook is the correctness choke point: the null-mask flag is valid only because the padding
|
|
was removed HERE, on the same call. It fires on ``module(...)`` (``__call__``), which the
|
|
pipeline/guider/cache_context/compile all use. Do NOT invoke a hooked DiT via
|
|
``module.forward(...)`` directly: that skips pre-hooks, so a stale True flag would null the mask
|
|
over un-trimmed padding and corrupt the output.
|
|
|
|
Best-effort: any anomaly leaves the inputs untouched and the flag False."""
|
|
import torch
|
|
|
|
original = dict(kwargs)
|
|
try:
|
|
null_ok = True
|
|
|
|
image = kwargs.get("image_embeds")
|
|
if image is not None and image.numel() > 0 and bool(torch.all(image == 0).item()):
|
|
# All-zero image == "no image" (t2v). Emptying the token axis removes the 729 padded
|
|
# image tokens; is_t2v stays True in forward (all() of empty is vacuously True).
|
|
kwargs["image_embeds"] = image[:, :0]
|
|
|
|
for skey, mkey, required in (
|
|
("encoder_hidden_states", "encoder_attention_mask", True),
|
|
("encoder_hidden_states_2", "encoder_attention_mask_2", False),
|
|
):
|
|
# Only touch streams passed by keyword (the pipeline always does); never write back an
|
|
# absent key (a positional encoder_hidden_states would collide). An absent REQUIRED
|
|
# primary stream drops the fast path; an absent optional byt5 is fine.
|
|
if skey not in kwargs:
|
|
null_ok = null_ok and not required
|
|
continue
|
|
states, mask, all_valid = _trim_stream(kwargs.get(skey), kwargs.get(mkey))
|
|
kwargs[skey] = states
|
|
kwargs[mkey] = mask
|
|
null_ok = null_ok and all_valid
|
|
|
|
# The primary mllm stream flows through the TokenRefiner's own attention; never hand it a
|
|
# 0-length sequence (pathological empty prompt). Revert and take the stock dense-mask path.
|
|
primary = kwargs.get("encoder_hidden_states")
|
|
if primary is not None and primary.dim() == 3 and primary.shape[1] == 0:
|
|
kwargs.clear()
|
|
kwargs.update(original)
|
|
null_ok = False
|
|
|
|
_set_hunyuan_null_mask(module, null_ok)
|
|
return args, kwargs
|
|
except Exception: # noqa: BLE001 — optimisation only; never break the forward
|
|
# We may have trimmed some kwargs before failing. Restore the caller's untrimmed inputs so
|
|
# the stock dense-mask path (flag False) runs on exactly what it expects.
|
|
kwargs.clear()
|
|
kwargs.update(original)
|
|
_set_hunyuan_null_mask(module, False)
|
|
return args, kwargs
|
|
|
|
|
|
def _hunyuan_trim_post_hook(module, _args, output):
|
|
"""Clear the null-mask flag after each hooked forward, scoping the authorisation to exactly the
|
|
call whose pre-hook removed the padding. Registered with ``always_call=True`` so the flag is
|
|
also cleared when the forward raises -- otherwise a latched True would null the mask over
|
|
un-trimmed padding on any later direct ``module.forward(...)``. Returns the output unchanged."""
|
|
_set_hunyuan_null_mask(module, False)
|
|
return output
|
|
|
|
|
|
def _install_null_processors(dit: Any, logger: Any) -> bool:
|
|
"""Swap every stock block attention processor on ``dit`` for the null-mask subclass. Only
|
|
touches blocks whose processor is exactly the stock class (so a diffusers change or an
|
|
already-installed run is a no-op). Preserves any pinned attention backend."""
|
|
try:
|
|
cls = _null_mask_processor_cls()
|
|
except Exception as exc: # noqa: BLE001 — diffusers moved / unavailable -> skip
|
|
_warn(logger, "hunyuan_attn_trim", exc)
|
|
return False
|
|
installed = 0
|
|
for blk in getattr(dit, "transformer_blocks", []):
|
|
attn = getattr(blk, "attn", None)
|
|
proc = getattr(attn, "processor", None) if attn is not None else None
|
|
if proc is None:
|
|
continue
|
|
if isinstance(proc, cls):
|
|
installed += 1 # already ours (idempotent)
|
|
continue
|
|
if type(proc).__name__ != _HUNYUAN15_PROCESSOR_CLS:
|
|
continue # unknown processor -> leave it alone
|
|
new = cls()
|
|
# carry over any backend/parallel config the stock processor already held
|
|
new._attention_backend = getattr(proc, "_attention_backend", None)
|
|
new._parallel_config = getattr(proc, "_parallel_config", None)
|
|
try:
|
|
attn.set_processor(new)
|
|
except Exception: # noqa: BLE001 — fall back to direct assignment
|
|
attn.processor = new
|
|
installed += 1
|
|
return installed > 0
|
|
|
|
|
|
def install_hunyuan_attention_trim(
|
|
pipe: Any,
|
|
family: Any,
|
|
*,
|
|
logger: Any = None,
|
|
) -> bool:
|
|
"""HunyuanVideo-1.5 only: make the joint attention skip padded text tokens (see module note).
|
|
|
|
Installs a null-mask processor on every denoiser DiT block plus an eager pre-hook that trims the
|
|
padded text/image streams each forward. Bit-exact for the video output (the fused-vs-masked SDPA
|
|
swap is the only numeric change). Returns True when engaged; No-op (False) for any other family,
|
|
an unexpected class, or any failure -- the stock dense-mask path stays, so correctness never
|
|
depends on this. Call BEFORE apply_attention_backend so the kernel pins onto the new processor."""
|
|
if getattr(family, "transformer_class", None) != _HUNYUAN15_TRANSFORMER_CLS:
|
|
return False
|
|
engaged = False
|
|
for dit in _attention_dits(pipe):
|
|
if type(dit).__name__ != _HUNYUAN15_TRANSFORMER_CLS:
|
|
continue
|
|
if not _install_null_processors(dit, logger):
|
|
continue
|
|
# Installation (and every idle period between generations) starts in the conservative
|
|
# state: the flag is only ever True inside the exact forward its pre-hook trimmed.
|
|
_set_hunyuan_null_mask(dit, False)
|
|
if getattr(dit, "_unsloth_trim_hook", None) is None:
|
|
pre_handle = None
|
|
try:
|
|
pre_handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs = True)
|
|
# always_call: clear the flag even when the forward raises, so an exception can
|
|
# never leave the null-mask authorisation latched for a later direct forward.
|
|
post_handle = dit.register_forward_hook(_hunyuan_trim_post_hook, always_call = True)
|
|
dit._unsloth_trim_hook = (pre_handle, post_handle)
|
|
except Exception as exc: # noqa: BLE001 — optimisation only
|
|
if pre_handle is not None:
|
|
pre_handle.remove()
|
|
_set_hunyuan_null_mask(dit, False)
|
|
_warn(logger, "hunyuan_attn_trim", exc)
|
|
continue
|
|
engaged = True
|
|
if engaged and logger is not None:
|
|
logger.info("diffusion.attention: hunyuan padded-text trim engaged")
|
|
return engaged
|