mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-10 00:08:58 +00:00
video: skip padded text tokens in HunyuanVideo-1.5 joint attention
HunyuanVideo-1.5's DiT runs a joint [video; text] self-attention and, on every block and step, builds a dense [B,1,N,N] boolean mask so the video never attends to the padded text. A dense bool attn_mask disables every fused SDPA kernel (flash rejects it; cuDNN and memory-efficient fall back), so the attention runs the slow math-style path: at the production shape (121 frames, 480p, N about 50k) one attention call is ~421ms with the mask vs ~19ms with attn_mask=None. The text is ~99.5% padding (a t2v prompt fills ~9 of ~1985 slots), so nearly all of that cost is spent masking padding. install_hunyuan_attention_trim installs an eager forward pre-hook that drops the all-zero image stream (t2v) and trims the mllm/byt5 text streams to their globally-valid columns, plus a null-mask attention processor that runs attn_mask=None once no partially-padded column remains (the batch-1 / per-guidance-branch case) and otherwise delegates to the stock dense-mask processor. The model already zeroes and masks the padded text and discards its attention output (only the video split feeds proj_out), so removing it is exact for the video; the only numeric change is the SDPA kernel (masked fallback to fused). Measured on a B200: 23.3s to 1.3s per DiT forward at 121 frames (~18x with regional compile, 0 graph breaks); per-forward cosine 0.99998 vs stock; equal distance to an fp32 reference (LPIPS fp32-vs-stock 0.292, fp32-vs-trim 0.307), so it is not less accurate than the current bf16 default. Wired auto-on for HunyuanVideo-1.5 in the video loader, before the attention backend set so the requested kernel pins onto the new processors; a no-op for every other family and reversible (stock dense-mask path on any anomaly). Adds hermetic tests and the diagnostic/validation scripts.
This commit is contained in:
Daniel Han
parent
e58f30be5f
commit
a5928064a0
8 changed files with 1117 additions and 2 deletions
194
scripts/hunyuan_attn_diag.py
Normal file
194
scripts/hunyuan_attn_diag.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
"""Diagnose HunyuanVideo-1.5 joint-attention cost: capture the REAL joint sequence length
|
||||
and per-text-stream padding, then time SDPA three ways at those exact shapes --
|
||||
(a) dense [B,1,N,N] bool mask (current default)
|
||||
(b) attn_mask=None (flash path; only valid if no padding remains)
|
||||
(c) dense mask at trimmed N (text padding removed, mask still built)
|
||||
so we know whether the win is the N-reduction (trim) or the mask-elimination (null).
|
||||
|
||||
Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_attn_diag.py [--repo ...] [--frames 121]
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
|
||||
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def _import_diffusers():
|
||||
# diffusers eagerly imports bitsandbytes through its quantizers; the bnb build here is
|
||||
# mismatched (cuda130), so disable the availability flag before importing (same trick as
|
||||
# scripts/video_speedmem_bench.py).
|
||||
import diffusers.utils.import_utils as iu
|
||||
|
||||
iu._bitsandbytes_available = False
|
||||
import diffusers
|
||||
|
||||
return diffusers
|
||||
|
||||
|
||||
CAP: dict = {}
|
||||
|
||||
|
||||
class _StopCapture(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _block_pre_hook(module, args, kwargs):
|
||||
# HunyuanVideo15TransformerBlock.forward(hidden_states, encoder_hidden_states, temb,
|
||||
# attention_mask, image_rotary_emb) -- positional per transformer:786-792.
|
||||
def _get(i, name):
|
||||
if name in kwargs:
|
||||
return kwargs[name]
|
||||
return args[i] if i < len(args) else None
|
||||
|
||||
hs = _get(0, "hidden_states")
|
||||
ehs = _get(1, "encoder_hidden_states")
|
||||
amask = _get(3, "attention_mask")
|
||||
if hs is None or ehs is None:
|
||||
return None
|
||||
CAP["n_video"] = int(hs.shape[1])
|
||||
CAP["n_text"] = int(ehs.shape[1])
|
||||
CAP["heads"] = int(getattr(module.attn, "heads", 0))
|
||||
CAP["dim_head"] = int(hs.shape[-1] // max(CAP["heads"], 1))
|
||||
CAP["batch"] = int(hs.shape[0])
|
||||
CAP["dtype"] = hs.dtype
|
||||
if amask is not None:
|
||||
m = amask.bool()
|
||||
CAP["text_valid_per_batch"] = m.sum(dim=1).tolist()
|
||||
CAP["text_cols_valid_any"] = int(m.any(dim=0).sum()) # what our global-trim would keep
|
||||
raise _StopCapture
|
||||
|
||||
|
||||
def _model_pre_hook(module, args, kwargs):
|
||||
# capture the raw per-stream padding breakdown before the reorder
|
||||
def g(name):
|
||||
return kwargs.get(name)
|
||||
|
||||
for key, mkey in (("encoder_hidden_states", "encoder_attention_mask"),
|
||||
("encoder_hidden_states_2", "encoder_attention_mask_2")):
|
||||
s = g(key)
|
||||
m = g(mkey)
|
||||
if s is not None:
|
||||
CAP.setdefault("streams", {})[key] = {
|
||||
"len": int(s.shape[1]),
|
||||
"valid": (m.bool().sum(dim=1).tolist() if m is not None else None),
|
||||
}
|
||||
ie = g("image_embeds")
|
||||
if ie is not None:
|
||||
CAP["image_embeds_len"] = int(ie.shape[1])
|
||||
CAP["image_is_t2v"] = bool(torch.all(ie == 0).item())
|
||||
return None
|
||||
|
||||
|
||||
def _time_sdpa(q, k, v, mask, iters=30):
|
||||
# q,k,v: [B, H, N, D]
|
||||
torch.cuda.synchronize()
|
||||
for _ in range(3): # warmup
|
||||
F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(iters):
|
||||
F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
|
||||
torch.cuda.synchronize()
|
||||
return (time.perf_counter() - t0) / iters * 1e3 # ms
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
|
||||
ap.add_argument("--frames", type=int, default=121)
|
||||
ap.add_argument("--width", type=int, default=832)
|
||||
ap.add_argument("--height", type=int, default=480)
|
||||
args = ap.parse_args()
|
||||
|
||||
diffusers = _import_diffusers()
|
||||
dev = "cuda:0"
|
||||
print(f"loading {args.repo} ...", flush=True)
|
||||
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16)
|
||||
pipe = pipe.to(dev)
|
||||
|
||||
pipe.transformer.register_forward_pre_hook(_model_pre_hook, with_kwargs=True)
|
||||
pipe.transformer.transformer_blocks[0].register_forward_pre_hook(_block_pre_hook, with_kwargs=True)
|
||||
|
||||
print("running 1 capture step ...", flush=True)
|
||||
try:
|
||||
pipe(
|
||||
prompt="a cat playing piano",
|
||||
num_frames=args.frames,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_inference_steps=1,
|
||||
)
|
||||
except _StopCapture:
|
||||
pass
|
||||
except Exception as exc: # noqa: BLE001
|
||||
# the StopCapture may surface wrapped; if we captured, continue
|
||||
if "n_video" not in CAP:
|
||||
raise
|
||||
print(f"(generation aborted after capture: {type(exc).__name__})", flush=True)
|
||||
|
||||
print("\n===== CAPTURED SHAPES =====", flush=True)
|
||||
for kk in ("batch", "n_video", "n_text", "heads", "dim_head", "dtype",
|
||||
"text_valid_per_batch", "text_cols_valid_any", "image_embeds_len",
|
||||
"image_is_t2v", "streams"):
|
||||
if kk in CAP:
|
||||
print(f" {kk}: {CAP[kk]}", flush=True)
|
||||
|
||||
B = CAP["batch"]
|
||||
H = CAP["heads"]
|
||||
D = CAP["dim_head"]
|
||||
n_video = CAP["n_video"]
|
||||
n_text = CAP["n_text"]
|
||||
N = n_video + n_text
|
||||
# trimmed joint length if we drop globally-invalid text columns
|
||||
keep_text = CAP.get("text_cols_valid_any", n_text)
|
||||
N_trim = n_video + keep_text
|
||||
dtype = CAP["dtype"]
|
||||
print(f"\n joint N = {N} (video {n_video} + text {n_text}); "
|
||||
f"trimmed N = {N_trim} (text kept {keep_text})", flush=True)
|
||||
|
||||
def mk(n):
|
||||
return torch.randn(B, H, n, D, device=dev, dtype=dtype)
|
||||
|
||||
# (a) dense mask over full N (current). Build [B,1,N,N] bool (mostly True).
|
||||
print("\n===== SDPA TIMING (ms/call, real shapes) =====", flush=True)
|
||||
q, k, v = mk(N), mk(N), mk(N)
|
||||
dense = torch.ones(B, 1, N, N, dtype=torch.bool, device=dev)
|
||||
# emulate text padding: last (n_text - keep_text) columns invalid
|
||||
if n_text - keep_text > 0:
|
||||
dense[:, :, :, n_video + keep_text:] = False
|
||||
dense[:, :, n_video + keep_text:, :] = False
|
||||
t_dense = _time_sdpa(q, k, v, dense)
|
||||
mask_gb = dense.numel() / 1e9
|
||||
print(f" (a) dense [B,1,N,N] mask N={N:>6} : {t_dense:7.3f} ms (mask {mask_gb:.2f} GB)", flush=True)
|
||||
|
||||
# (b) no mask over full N (upper bound of flash path if all valid)
|
||||
t_none = _time_sdpa(q, k, v, None)
|
||||
print(f" (b) attn_mask=None N={N:>6} : {t_none:7.3f} ms ({t_dense/t_none:.2f}x vs a)", flush=True)
|
||||
|
||||
# (c) trimmed N, dense all-True mask (text padding removed but mask still built)
|
||||
qt, kt, vt = mk(N_trim), mk(N_trim), mk(N_trim)
|
||||
dense_t = torch.ones(B, 1, N_trim, N_trim, dtype=torch.bool, device=dev)
|
||||
t_dense_trim = _time_sdpa(qt, kt, vt, dense_t)
|
||||
print(f" (c) dense mask @trimmed N={N_trim:>6} : {t_dense_trim:7.3f} ms ({t_dense/t_dense_trim:.2f}x vs a)", flush=True)
|
||||
|
||||
# (d) trimmed N, no mask (trim + null: the full proposed fast path)
|
||||
t_none_trim = _time_sdpa(qt, kt, vt, None)
|
||||
print(f" (d) no mask @trimmed N={N_trim:>6} : {t_none_trim:7.3f} ms ({t_dense/t_none_trim:.2f}x vs a)", flush=True)
|
||||
|
||||
print("\n Interpretation:", flush=True)
|
||||
print(f" trim-only ceiling (a->c): {(1-t_dense_trim/t_dense)*100:5.1f}% attn saving", flush=True)
|
||||
print(f" null-only ceiling (a->b): {(1-t_none/t_dense)*100:5.1f}% attn saving", flush=True)
|
||||
print(f" trim+null (a->d): {(1-t_none_trim/t_dense)*100:5.1f}% attn saving", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
124
scripts/hunyuan_trim_e2e.py
Normal file
124
scripts/hunyuan_trim_e2e.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
"""End-to-end pixel validation of the HunyuanVideo-1.5 attention trim: same seed, stock vs trim,
|
||||
per-frame LPIPS + mean luma (black-frame guard), plus a full-resolution trim gen for real
|
||||
wall-clock. Stock is only run at MODEST settings (a full 121-frame stock gen is ~19 min).
|
||||
|
||||
Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_e2e.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_REPO_ROOT / "studio" / "backend"))
|
||||
OUT = _REPO_ROOT / "outputs" / "video_speedmem"
|
||||
OUT.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def _import_diffusers():
|
||||
import diffusers.utils.import_utils as iu
|
||||
|
||||
iu._bitsandbytes_available = False
|
||||
import diffusers
|
||||
|
||||
return diffusers
|
||||
|
||||
|
||||
def _gen(pipe, seed, frames, steps, w, h):
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
out = pipe(prompt="a cat playing piano on a stage, cinematic",
|
||||
num_frames=frames, width=w, height=h,
|
||||
num_inference_steps=steps, generator=g, output_type="np")
|
||||
torch.cuda.synchronize()
|
||||
wall = (time.perf_counter() - t0) * 1e3
|
||||
frames_np = out.frames[0] # [F,H,W,C] in [0,1]
|
||||
return np.asarray(frames_np), wall
|
||||
|
||||
|
||||
def _luma(frames):
|
||||
# BT.601 luma over [F,H,W,C] in [0,1]
|
||||
r, gg, b = frames[..., 0], frames[..., 1], frames[..., 2]
|
||||
return float((0.299 * r + 0.587 * gg + 0.114 * b).mean())
|
||||
|
||||
|
||||
def _lpips_mean(loss_fn, a, b, stride=4):
|
||||
vals = []
|
||||
for i in range(0, len(a), stride):
|
||||
ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
|
||||
tb = torch.from_numpy(b[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
|
||||
with torch.no_grad():
|
||||
vals.append(loss_fn(ta.cuda(), tb.cuda()).item())
|
||||
return float(np.mean(vals)) if vals else None
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
|
||||
ap.add_argument("--cmp-frames", type=int, default=25)
|
||||
ap.add_argument("--cmp-steps", type=int, default=50)
|
||||
ap.add_argument("--cmp-w", type=int, default=832)
|
||||
ap.add_argument("--cmp-h", type=int, default=480)
|
||||
ap.add_argument("--full-frames", type=int, default=121)
|
||||
ap.add_argument("--full-steps", type=int, default=50)
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
args = ap.parse_args()
|
||||
|
||||
diffusers = _import_diffusers()
|
||||
from core.inference.diffusion_attention import install_hunyuan_attention_trim
|
||||
from core.inference.video_families import detect_video_family
|
||||
import lpips
|
||||
|
||||
print(f"loading {args.repo} ...", flush=True)
|
||||
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to("cuda")
|
||||
fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
|
||||
loss_fn = lpips.LPIPS(net="alex").cuda()
|
||||
|
||||
fr, st, w, hh = args.cmp_frames, args.cmp_steps, args.cmp_w, args.cmp_h
|
||||
# ---- STOCK x2 (nondeterminism floor) ----
|
||||
print(f"\n[stock#1] gen {fr}f/{st}steps {w}x{hh} ...", flush=True)
|
||||
stock1, w_stock = _gen(pipe, args.seed, fr, st, w, hh)
|
||||
print(f"[stock#1] wall={w_stock:.0f} ms luma={_luma(stock1):.4f} frames={stock1.shape}", flush=True)
|
||||
print(f"[stock#2] gen (same seed, measures run-to-run nondeterminism) ...", flush=True)
|
||||
stock2, _ = _gen(pipe, args.seed, fr, st, w, hh)
|
||||
|
||||
# ---- TRIM (same seed) ----
|
||||
engaged = install_hunyuan_attention_trim(pipe, fam, logger=None)
|
||||
print(f"\ninstall_hunyuan_attention_trim engaged={engaged}", flush=True)
|
||||
trim, w_trim = _gen(pipe, args.seed, fr, st, w, hh)
|
||||
print(f"[trim ] wall={w_trim:.0f} ms luma={_luma(trim):.4f} ({w_stock/w_trim:.2f}x vs stock)", flush=True)
|
||||
|
||||
floor = _lpips_mean(loss_fn, stock1, stock2)
|
||||
lp = _lpips_mean(loss_fn, stock1, trim)
|
||||
print(f"\nLPIPS(stock#1, stock#2) = {floor:.5f} <- nondeterminism floor", flush=True)
|
||||
print(f"LPIPS(stock#1, trim ) = {lp:.5f} <- trim vs stock", flush=True)
|
||||
print(f" => trim is {'WITHIN' if lp <= floor * 1.5 + 0.002 else 'ABOVE'} the nondeterminism floor", flush=True)
|
||||
|
||||
# ---- FULL-RES TRIM (wall-clock + black-frame guard; stock at full-res is ~19 min, skipped) ----
|
||||
print(f"\n[trim-full] gen {args.full_frames}f/{args.full_steps}steps {args.cmp_w}x{args.cmp_h} ...", flush=True)
|
||||
full, w_full = _gen(pipe, args.seed, args.full_frames, args.full_steps, args.cmp_w, args.cmp_h)
|
||||
print(f"[trim-full] wall={w_full/1000:.1f} s luma={_luma(full):.4f} frames={full.shape}", flush=True)
|
||||
|
||||
try:
|
||||
from PIL import Image
|
||||
Image.fromarray((trim[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_cmp.png")
|
||||
Image.fromarray((full[0] * 255).astype("uint8")).save(OUT / "vid_hunyuan_trim_full.png")
|
||||
print(f"\nsaved sample frames to {OUT}", flush=True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f"(png save skipped: {exc})", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
105
scripts/hunyuan_trim_fp32ref.py
Normal file
105
scripts/hunyuan_trim_fp32ref.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
"""Is the Hunyuan attention trim LESS accurate, or just DIFFERENT? Neither bf16-masked (stock)
|
||||
nor bf16-trim is ground truth. Compare BOTH against an fp32 reference (same seed): if the trim is
|
||||
as close to fp32 as stock is, the 0.14 LPIPS stock-vs-trim is a benign bf16-kernel resample, not a
|
||||
quality loss. If trim is clearly farther from fp32 than stock, it is a real regression.
|
||||
|
||||
Run: CUDA_VISIBLE_DEVICES=1 python scripts/hunyuan_trim_fp32ref.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_REPO_ROOT / "studio" / "backend"))
|
||||
|
||||
|
||||
def _import_diffusers():
|
||||
import diffusers.utils.import_utils as iu
|
||||
|
||||
iu._bitsandbytes_available = False
|
||||
import diffusers
|
||||
|
||||
return diffusers
|
||||
|
||||
|
||||
def _gen(pipe, seed, fr, st, w, h):
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
out = pipe(prompt="a cat playing piano on a stage, cinematic",
|
||||
num_frames=fr, width=w, height=h, num_inference_steps=st,
|
||||
generator=g, output_type="np")
|
||||
return np.asarray(out.frames[0])
|
||||
|
||||
|
||||
def _lpips_mean(loss_fn, a, b, stride=3):
|
||||
vals = []
|
||||
for i in range(0, len(a), stride):
|
||||
ta = torch.from_numpy(a[i]).permute(2, 0, 1).unsqueeze(0).float().cuda() * 2 - 1
|
||||
tb = torch.from_numpy(b[i]).permute(2, 0, 1).unsqueeze(0).float().cuda() * 2 - 1
|
||||
with torch.no_grad():
|
||||
vals.append(loss_fn(ta, tb).item())
|
||||
return float(np.mean(vals)) if vals else None
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
|
||||
ap.add_argument("--frames", type=int, default=25)
|
||||
ap.add_argument("--steps", type=int, default=30)
|
||||
ap.add_argument("--w", type=int, default=832)
|
||||
ap.add_argument("--h", type=int, default=480)
|
||||
ap.add_argument("--seed", type=int, default=1234)
|
||||
args = ap.parse_args()
|
||||
|
||||
diffusers = _import_diffusers()
|
||||
from core.inference.diffusion_attention import install_hunyuan_attention_trim
|
||||
from core.inference.video_families import detect_video_family
|
||||
import lpips
|
||||
|
||||
fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
|
||||
loss_fn = lpips.LPIPS(net="alex").cuda()
|
||||
fr, st, w, h, seed = args.frames, args.steps, args.w, args.h, args.seed
|
||||
|
||||
# ---- bf16 stock + bf16 trim on one pipe ----
|
||||
print(f"loading bf16 {args.repo} ...", flush=True)
|
||||
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to("cuda")
|
||||
print(f"[bf16 stock] gen {fr}f/{st}steps ...", flush=True)
|
||||
stock = _gen(pipe, seed, fr, st, w, h)
|
||||
install_hunyuan_attention_trim(pipe, fam, logger=None)
|
||||
print("[bf16 trim ] gen ...", flush=True)
|
||||
trim = _gen(pipe, seed, fr, st, w, h)
|
||||
del pipe
|
||||
gc.collect(); torch.cuda.empty_cache()
|
||||
|
||||
# ---- fp32 reference (stock masked attention, upcast) ----
|
||||
print("loading fp32 reference ...", flush=True)
|
||||
pipe32 = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.float32).to("cuda")
|
||||
print(f"[fp32 gold ] gen {fr}f/{st}steps ...", flush=True)
|
||||
gold = _gen(pipe32, seed, fr, st, w, h)
|
||||
|
||||
d_stock = _lpips_mean(loss_fn, gold, stock)
|
||||
d_trim = _lpips_mean(loss_fn, gold, trim)
|
||||
d_st = _lpips_mean(loss_fn, stock, trim)
|
||||
print("\n===== ACCURACY vs fp32 reference =====", flush=True)
|
||||
print(f" LPIPS(fp32, bf16-stock) = {d_stock:.5f}", flush=True)
|
||||
print(f" LPIPS(fp32, bf16-trim ) = {d_trim:.5f}", flush=True)
|
||||
print(f" LPIPS(bf16-stock, trim) = {d_st:.5f}", flush=True)
|
||||
if d_stock is not None and d_trim is not None:
|
||||
verdict = "NOT less accurate (trim ~= stock vs fp32)" if d_trim <= d_stock * 1.25 + 0.01 \
|
||||
else "LESS accurate (trim farther from fp32 than stock)"
|
||||
print(f"\n VERDICT: {verdict}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
155
scripts/hunyuan_trim_validate.py
Normal file
155
scripts/hunyuan_trim_validate.py
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
"""Validate the HunyuanVideo-1.5 padded-text attention trim: accuracy (stock vs trimmed forward
|
||||
output on the SAME real inputs), per-forward speed, and torch.compile compatibility -- all at the
|
||||
real production shape (default 121 frames / 480p).
|
||||
|
||||
Run: CUDA_VISIBLE_DEVICES=3 python scripts/hunyuan_trim_validate.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
|
||||
|
||||
import torch
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(_REPO_ROOT / "studio" / "backend"))
|
||||
|
||||
|
||||
def _import_diffusers():
|
||||
import diffusers.utils.import_utils as iu
|
||||
|
||||
iu._bitsandbytes_available = False
|
||||
import diffusers
|
||||
|
||||
return diffusers
|
||||
|
||||
|
||||
CAP: dict = {}
|
||||
|
||||
|
||||
class _Stop(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _capture_hook(module, args, kwargs):
|
||||
CAP["kwargs"] = {k: v for k, v in kwargs.items()}
|
||||
CAP["args"] = args
|
||||
raise _Stop
|
||||
|
||||
|
||||
def _forward(transformer, no_grad=True):
|
||||
ctx = torch.no_grad() if no_grad else torch.enable_grad()
|
||||
with ctx:
|
||||
out = transformer(*CAP["args"], **CAP["kwargs"])
|
||||
return out[0] if isinstance(out, tuple) else out.sample
|
||||
|
||||
|
||||
def _median_ms(fn, iters=8, warmup=2):
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
ts = []
|
||||
for _ in range(iters):
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
ts.append((time.perf_counter() - t0) * 1e3)
|
||||
ts.sort()
|
||||
return ts[len(ts) // 2]
|
||||
|
||||
|
||||
def _compare(a, b):
|
||||
a = a.float().flatten()
|
||||
b = b.float().flatten()
|
||||
cos = torch.nn.functional.cosine_similarity(a, b, dim=0).item()
|
||||
max_abs = (a - b).abs().max().item()
|
||||
denom = a.abs().max().item() or 1.0
|
||||
return cos, max_abs, max_abs / denom
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--repo", default="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v")
|
||||
ap.add_argument("--frames", type=int, default=121)
|
||||
ap.add_argument("--width", type=int, default=832)
|
||||
ap.add_argument("--height", type=int, default=480)
|
||||
ap.add_argument("--compile", action="store_true", help="also test regional compile of blocks")
|
||||
args = ap.parse_args()
|
||||
|
||||
diffusers = _import_diffusers()
|
||||
from core.inference.diffusion_attention import install_hunyuan_attention_trim
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
dev = "cuda:0"
|
||||
print(f"loading {args.repo} ...", flush=True)
|
||||
pipe = diffusers.DiffusionPipeline.from_pretrained(args.repo, torch_dtype=torch.bfloat16).to(dev)
|
||||
fam = detect_video_family(args.repo) or detect_video_family("hunyuanvideo-1.5")
|
||||
print(f"family: {getattr(fam, 'name', None)} / transformer_class={getattr(fam, 'transformer_class', None)}", flush=True)
|
||||
|
||||
h = pipe.transformer.register_forward_pre_hook(_capture_hook, with_kwargs=True)
|
||||
print("capturing one real forward input ...", flush=True)
|
||||
try:
|
||||
pipe(prompt="a cat playing piano", num_frames=args.frames,
|
||||
width=args.width, height=args.height, num_inference_steps=1)
|
||||
except _Stop:
|
||||
pass
|
||||
h.remove()
|
||||
|
||||
k = CAP["kwargs"]
|
||||
ehs = k.get("encoder_hidden_states")
|
||||
m = k.get("encoder_attention_mask")
|
||||
ie = k.get("image_embeds")
|
||||
print(f"\ncaptured: encoder_hidden_states={tuple(ehs.shape)} mask_valid={m.bool().sum(1).tolist()}"
|
||||
f" image_embeds={tuple(ie.shape) if ie is not None else None}"
|
||||
f" image_all_zero={bool(torch.all(ie==0).item()) if ie is not None else None}", flush=True)
|
||||
|
||||
# ---- STOCK forward (reference) + timing ----
|
||||
transformer = pipe.transformer
|
||||
out_stock = _forward(transformer).detach().clone()
|
||||
t_stock = _median_ms(lambda: _forward(transformer))
|
||||
print(f"\nSTOCK forward: {t_stock:8.2f} ms out={tuple(out_stock.shape)}", flush=True)
|
||||
|
||||
# ---- install trim, re-run same inputs ----
|
||||
engaged = install_hunyuan_attention_trim(pipe, fam, logger=None)
|
||||
print(f"install_hunyuan_attention_trim engaged = {engaged}", flush=True)
|
||||
out_trim = _forward(transformer).detach().clone()
|
||||
t_trim = _median_ms(lambda: _forward(transformer))
|
||||
|
||||
cos, max_abs, rel = _compare(out_stock, out_trim)
|
||||
print(f"TRIM forward: {t_trim:8.2f} ms ({t_stock/t_trim:.2f}x faster)", flush=True)
|
||||
print(f"\nACCURACY stock-vs-trim: cosine={cos:.8f} max_abs={max_abs:.4e} rel_max={rel:.4e}", flush=True)
|
||||
finite = bool(torch.isfinite(out_trim).all().item())
|
||||
print(f"trim output finite: {finite}", flush=True)
|
||||
|
||||
if args.compile:
|
||||
print("\ncompiling blocks (compile_repeated_blocks, mode=default, dynamic=True) ...", flush=True)
|
||||
try:
|
||||
for _a in ("recompile_limit", "cache_size_limit"):
|
||||
if hasattr(torch._dynamo.config, _a):
|
||||
setattr(torch._dynamo.config, _a, 64)
|
||||
transformer.compile_repeated_blocks(fullgraph=False, dynamic=True)
|
||||
out_c = _forward(transformer).detach().clone() # triggers compile
|
||||
t_c = _median_ms(lambda: _forward(transformer), iters=5, warmup=1)
|
||||
cos_c, ma_c, rel_c = _compare(out_stock, out_c)
|
||||
cnt = torch._dynamo.utils.counters
|
||||
print(f"TRIM+COMPILE forward: {t_c:8.2f} ms ({t_stock/t_c:.2f}x vs stock)", flush=True)
|
||||
print(f" accuracy vs stock: cosine={cos_c:.8f} max_abs={ma_c:.4e}", flush=True)
|
||||
print(f" dynamo recompiles={sum(cnt['recompiles'].values()) if 'recompiles' in cnt else '?'}"
|
||||
f" graph_breaks={sum(cnt['graph_break'].values()) if 'graph_break' in cnt else 0}", flush=True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
import traceback
|
||||
print(f"COMPILE FAILED: {type(exc).__name__}: {exc}", flush=True)
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
65
scripts/sdpa_mask_backend_probe.py
Normal file
65
scripts/sdpa_mask_backend_probe.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
"""Which SDPA backends tolerate a dense bool attn_mask, and at what cost, at Hunyuan's real
|
||||
joint shape (B=1, H=16, N=50345, D=128, bf16)? Decides whether nulling the all-True mask is
|
||||
the real win on the PRODUCTION cuDNN path (not just the native math fallback)."""
|
||||
import time
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
B, H, N, D = 1, 16, 50345, 128
|
||||
dev, dt = "cuda:0", torch.bfloat16
|
||||
|
||||
|
||||
def mk():
|
||||
return torch.randn(B, H, N, D, device=dev, dtype=dt)
|
||||
|
||||
|
||||
def timed(fn, iters=20):
|
||||
torch.cuda.synchronize()
|
||||
for _ in range(3):
|
||||
try:
|
||||
fn()
|
||||
except Exception as e: # noqa: BLE001
|
||||
return f"UNSUPPORTED ({type(e).__name__})"
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
for _ in range(iters):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
return (time.perf_counter() - t0) / iters * 1e3
|
||||
|
||||
|
||||
q, k, v = mk(), mk(), mk()
|
||||
dense = torch.ones(B, 1, N, N, dtype=torch.bool, device=dev)
|
||||
|
||||
backends = {
|
||||
"default(dispatch)": None,
|
||||
"MATH": [SDPBackend.MATH],
|
||||
"FLASH": [SDPBackend.FLASH_ATTENTION],
|
||||
"EFFICIENT": [SDPBackend.EFFICIENT_ATTENTION],
|
||||
"CUDNN": [SDPBackend.CUDNN_ATTENTION],
|
||||
}
|
||||
|
||||
print(f"shape B={B} H={H} N={N} D={D} {dt}\n")
|
||||
print(f"{'backend':<20}{'mask=dense(ms)':>18}{'mask=None(ms)':>18}")
|
||||
for name, bk in backends.items():
|
||||
def run_dense():
|
||||
if bk is None:
|
||||
return F.scaled_dot_product_attention(q, k, v, attn_mask=dense)
|
||||
with sdpa_kernel(bk):
|
||||
return F.scaled_dot_product_attention(q, k, v, attn_mask=dense)
|
||||
|
||||
def run_none():
|
||||
if bk is None:
|
||||
return F.scaled_dot_product_attention(q, k, v, attn_mask=None)
|
||||
with sdpa_kernel(bk):
|
||||
return F.scaled_dot_product_attention(q, k, v, attn_mask=None)
|
||||
|
||||
dms = timed(run_dense)
|
||||
nms = timed(run_none)
|
||||
d_s = f"{dms:.2f}" if isinstance(dms, float) else dms
|
||||
n_s = f"{nms:.2f}" if isinstance(nms, float) else nms
|
||||
print(f"{name:<20}{d_s:>18}{n_s:>18}")
|
||||
|
|
@ -387,3 +387,287 @@ def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None:
|
|||
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, EVERY step, materialises a dense [B,1,N,N] boolean mask (N = N_video + N_text) so
|
||||
# the video never attends to the padded text tokens. But a dense bool attn_mask DISABLES
|
||||
# every fused SDPA kernel (flash rejects it outright; cuDNN/efficient fall back), forcing the
|
||||
# slow math-style path: measured 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
|
||||
# paid purely to mask out padding. And the text is ~99.5% padding: a t2v prompt fills only ~9
|
||||
# of ~1985 text slots (image 729 + byt5 256 + mllm 1000 tokens, almost all zero-padded).
|
||||
#
|
||||
# The fix is exact, not approximate: the model already zero-fills + 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. We do it in an eager forward pre-hook
|
||||
# (outside the regionally-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-guidance-branch call) -- flag the DiT so the attention
|
||||
# processor skips the dense mask and runs the fused (cuDNN/flash) path. The only numeric change
|
||||
# is the SDPA kernel (masked fallback -> fused), a rounding-level difference on par with the
|
||||
# already-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 _null_mask_processor_cls():
|
||||
"""Build (once, lazily) a HunyuanVideo15AttnProcessor2_0 subclass whose ``__call__`` skips
|
||||
the dense-mask construction and runs attn_mask=None when the DiT is flagged (padding already
|
||||
removed by the pre-hook); otherwise it delegates to the stock processor unchanged, so a
|
||||
mixed-padding batch and any future diffusers change to the base processor 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, which
|
||||
# is fine for a secondary stream e.g. an unused byt5 in t2v -- it contributes no tokens)
|
||||
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 only valid because the padding
|
||||
was removed HERE, on the same call. It fires on ``module(...)`` (``__call__``) -- the diffusers
|
||||
pipeline, guider, cache_context and regional compile all go through ``__call__``. 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 (stock behaviour)."""
|
||||
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 always
|
||||
# -padded image tokens; is_t2v stays True inside forward (all() of empty is vacuously
|
||||
# True), so the model still runs its text-to-video path.
|
||||
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 diffusers pipeline always does). Never write
|
||||
# a key back that was absent -- a positionally-passed encoder_hidden_states would then
|
||||
# collide ("got multiple values for argument"). If the REQUIRED primary stream is absent
|
||||
# we cannot vouch for the mask, so drop the fast path; an absent optional byt5 just
|
||||
# contributes nothing and 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 text stream (mllm) flows through the TokenRefiner's own attention; never
|
||||
# hand it a 0-length sequence (pathological empty prompt). Revert everything and take the
|
||||
# stock dense-mask path in that rare case.
|
||||
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
|
||||
|
||||
for blk in getattr(module, "transformer_blocks", []):
|
||||
attn = getattr(blk, "attn", None)
|
||||
if attn is not None:
|
||||
setattr(attn, _NULL_ATTN_FLAG, null_ok)
|
||||
|
||||
return args, kwargs
|
||||
except Exception: # noqa: BLE001 — optimisation only; never break the forward
|
||||
for blk in getattr(module, "transformer_blocks", []):
|
||||
attn = getattr(blk, "attn", None)
|
||||
if attn is not None:
|
||||
setattr(attn, _NULL_ATTN_FLAG, False)
|
||||
return args, kwargs
|
||||
|
||||
|
||||
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 attention backend already pinned."""
|
||||
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 kernel swap is the only numeric change. Returns True when engaged. No-op (False)
|
||||
for any other family, an unexpected transformer/processor class, or on any failure -- the
|
||||
stock dense-mask path stays in place, so correctness never depends on this optimisation.
|
||||
|
||||
Call BEFORE apply_attention_backend so the requested kernel is pinned 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
|
||||
if getattr(dit, "_unsloth_trim_hook", None) is None:
|
||||
try:
|
||||
handle = dit.register_forward_pre_hook(_hunyuan_trim_pre_hook, with_kwargs=True)
|
||||
dit._unsloth_trim_hook = handle
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
_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
|
||||
|
|
|
|||
|
|
@ -42,7 +42,11 @@ from typing import Any, Optional
|
|||
|
||||
from loggers import get_logger
|
||||
|
||||
from .diffusion_attention import apply_attention_backend, select_attention_backend
|
||||
from .diffusion_attention import (
|
||||
apply_attention_backend,
|
||||
install_hunyuan_attention_trim,
|
||||
select_attention_backend,
|
||||
)
|
||||
from .diffusion_cache import (
|
||||
FBCACHE_MIN_STEPS,
|
||||
TC_AUTO,
|
||||
|
|
@ -1336,6 +1340,7 @@ class VideoBackend:
|
|||
else:
|
||||
cache_reason = "requested"
|
||||
attention_engaged = None
|
||||
attention_trim_engaged = False
|
||||
speed_optims: tuple = ()
|
||||
for view in views:
|
||||
# apply_attention_backend / apply_speed_optims both act on ``view.transformer``;
|
||||
|
|
@ -1344,6 +1349,11 @@ class VideoBackend:
|
|||
# first pass; a dense torchao transformer on the pipeline path is not a GGUF one,
|
||||
# so is_gguf keys off the load kind (gguf) AND no quant having engaged.
|
||||
gguf_transformer = kind == "gguf" and transformer_quant_engaged is None
|
||||
# HunyuanVideo-1.5 only: drop the ~99% zero-padded text tokens from the joint
|
||||
# attention so it runs the fused (cuDNN/flash) SDPA kernel instead of the dense-mask
|
||||
# fallback (~18x/DiT-forward at 121 frames, cosine ~1.0). Must precede the backend set
|
||||
# so the requested kernel pins onto the new processors. No-op for every other family.
|
||||
trim = install_hunyuan_attention_trim(view, fam, logger=logger)
|
||||
engaged = apply_attention_backend(
|
||||
view,
|
||||
select_attention_backend(
|
||||
|
|
@ -1365,7 +1375,10 @@ class VideoBackend:
|
|||
)
|
||||
if view is pipe:
|
||||
attention_engaged = engaged
|
||||
speed_optims = tuple(k for k, v in applied.items() if v)
|
||||
attention_trim_engaged = trim
|
||||
speed_optims = tuple(k for k, v in applied.items() if v) + (
|
||||
("hunyuan_attn_trim",) if trim else ()
|
||||
)
|
||||
with self._generate_lock:
|
||||
# A cancelled/superseded load must not place weights on the GPU the arbiter
|
||||
# may already have handed to another backend; recheck right before placement
|
||||
|
|
@ -1411,6 +1424,14 @@ class VideoBackend:
|
|||
attention_engaged or "native",
|
||||
"cuDNN fused attention on NVIDIA when a speed profile is active",
|
||||
),
|
||||
"attention_trim": (
|
||||
None,
|
||||
"on" if attention_trim_engaged else "off",
|
||||
"HunyuanVideo-1.5: padded text tokens dropped so joint attention runs "
|
||||
"the fused SDPA kernel (~18x per DiT forward, cosine ~1.0)"
|
||||
if attention_trim_engaged
|
||||
else "not applicable (non-Hunyuan family)",
|
||||
),
|
||||
"transformer_cache": (
|
||||
None if cache_auto else transformer_cache,
|
||||
cache_engaged or "off",
|
||||
|
|
|
|||
|
|
@ -434,3 +434,170 @@ def test_install_failure_falls_back_to_native(monkeypatch):
|
|||
monkeypatch.setattr(att, "_active_attention_backend", lambda: "native")
|
||||
t = _FakeTransformer(fail = True)
|
||||
assert apply_attention_backend(_pipe(t), "sage") is None
|
||||
|
||||
|
||||
# ── HunyuanVideo-1.5 padded-text attention trim ─────────────────────────────────────
|
||||
# _trim_stream / _hunyuan_trim_pre_hook use real torch tensor ops, so these run on CPU torch.
|
||||
import torch # noqa: E402
|
||||
|
||||
|
||||
def test_trim_stream_drops_trailing_padding():
|
||||
# right-padded (valid prefix): drop the globally-invalid tail, keep valid, flag all_valid.
|
||||
states = torch.arange(6.0).reshape(1, 6, 1)
|
||||
mask = torch.tensor([[1, 1, 1, 0, 0, 0]])
|
||||
out_s, out_m, all_valid = att._trim_stream(states, mask)
|
||||
assert out_s.shape == (1, 3, 1)
|
||||
assert torch.equal(out_s[0, :, 0], torch.tensor([0.0, 1.0, 2.0]))
|
||||
assert out_m.shape == (1, 3) and all_valid is True
|
||||
|
||||
|
||||
def test_trim_stream_layout_agnostic_drops_only_global_padding():
|
||||
# left-padded (valid suffix): any(dim=0) keeps positions valid for at least one element,
|
||||
# so the leading globally-invalid columns are dropped regardless of padding side.
|
||||
states = torch.arange(4.0).reshape(1, 4, 1)
|
||||
mask = torch.tensor([[0, 0, 1, 1]])
|
||||
out_s, out_m, all_valid = att._trim_stream(states, mask)
|
||||
assert torch.equal(out_s[0, :, 0], torch.tensor([2.0, 3.0])) and all_valid is True
|
||||
|
||||
|
||||
def test_trim_stream_full_mask_is_noop():
|
||||
states = torch.ones(1, 4, 2)
|
||||
mask = torch.ones(1, 4, dtype=torch.long)
|
||||
out_s, out_m, all_valid = att._trim_stream(states, mask)
|
||||
assert out_s.shape == (1, 4, 2) and all_valid is True
|
||||
|
||||
|
||||
def test_trim_stream_none_mask_passthrough():
|
||||
states = torch.ones(1, 4, 2)
|
||||
out_s, out_m, all_valid = att._trim_stream(states, None)
|
||||
assert out_s is states and out_m is None and all_valid is True
|
||||
|
||||
|
||||
def test_trim_stream_mixed_batch_not_all_valid():
|
||||
# batch>1 with different valid sets: the union is kept, but a column valid for only one
|
||||
# element remains partially padded -> all_valid False -> caller keeps the dense mask.
|
||||
states = torch.ones(2, 4, 1)
|
||||
mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]]) # elem1 has 2 valid, elem2 has 3
|
||||
out_s, out_m, all_valid = att._trim_stream(states, mask)
|
||||
assert out_s.shape == (2, 3, 1) # dropped the last col (invalid for both)
|
||||
assert all_valid is False
|
||||
|
||||
|
||||
def _fake_dit(n_blocks=2):
|
||||
blocks = [types.SimpleNamespace(attn=types.SimpleNamespace()) for _ in range(n_blocks)]
|
||||
return types.SimpleNamespace(transformer_blocks=blocks)
|
||||
|
||||
|
||||
def test_trim_pre_hook_empties_t2v_image_and_trims_and_flags():
|
||||
dit = _fake_dit()
|
||||
kwargs = {
|
||||
"image_embeds": torch.zeros(1, 5, 3), # all-zero -> t2v -> emptied
|
||||
"encoder_hidden_states": torch.arange(4.0).reshape(1, 4, 1),
|
||||
"encoder_attention_mask": torch.tensor([[1, 1, 0, 0]]),
|
||||
"encoder_hidden_states_2": torch.arange(3.0).reshape(1, 3, 1),
|
||||
"encoder_attention_mask_2": torch.tensor([[1, 0, 0]]),
|
||||
}
|
||||
args, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
|
||||
assert out["image_embeds"].shape == (1, 0, 3) # image tokens dropped
|
||||
assert out["encoder_hidden_states"].shape == (1, 2, 1) # mllm trimmed to 2 valid
|
||||
assert out["encoder_hidden_states_2"].shape == (1, 1, 1) # byt5 trimmed to 1 valid
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is True for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_trim_stream_all_invalid_yields_empty_but_valid():
|
||||
# A fully-padded secondary stream (e.g. unused byt5 in t2v) trims to 0 length and reports
|
||||
# all_valid True (vacuous) so it does NOT drop the fast path -- it just contributes no tokens.
|
||||
states = torch.ones(1, 5, 2)
|
||||
mask = torch.zeros(1, 5, dtype=torch.long)
|
||||
out_s, out_m, all_valid = att._trim_stream(states, mask)
|
||||
assert out_s.shape == (1, 0, 2) and all_valid is True
|
||||
|
||||
|
||||
def test_trim_pre_hook_byt5_all_invalid_keeps_fast_path():
|
||||
# The real t2v case: byt5 is entirely padding (valid=0). It must be emptied WITHOUT dropping
|
||||
# the null-mask fast path, since mllm still carries the prompt.
|
||||
dit = _fake_dit()
|
||||
kwargs = {
|
||||
"image_embeds": torch.zeros(1, 5, 3),
|
||||
"encoder_hidden_states": torch.arange(4.0).reshape(1, 4, 1),
|
||||
"encoder_attention_mask": torch.tensor([[1, 1, 1, 0]]),
|
||||
"encoder_hidden_states_2": torch.ones(1, 6, 1),
|
||||
"encoder_attention_mask_2": torch.zeros(1, 6, dtype=torch.long), # all padding
|
||||
}
|
||||
_, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
|
||||
assert out["encoder_hidden_states"].shape == (1, 3, 1)
|
||||
assert out["encoder_hidden_states_2"].shape == (1, 0, 1) # byt5 emptied
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is True for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_trim_pre_hook_empty_primary_reverts_and_disables():
|
||||
# Pathological empty prompt: mllm has 0 valid tokens. The TokenRefiner must not get a
|
||||
# 0-length sequence -> revert all inputs to original and take the stock dense-mask path.
|
||||
dit = _fake_dit()
|
||||
mllm = torch.ones(1, 4, 1)
|
||||
kwargs = {
|
||||
"image_embeds": torch.zeros(1, 5, 3),
|
||||
"encoder_hidden_states": mllm,
|
||||
"encoder_attention_mask": torch.zeros(1, 4, dtype=torch.long), # 0 valid
|
||||
}
|
||||
_, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
|
||||
assert out["encoder_hidden_states"] is mllm # reverted (not emptied)
|
||||
assert out["image_embeds"].shape == (1, 5, 3) # image revert too
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is False for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_trim_pre_hook_keeps_i2v_image():
|
||||
dit = _fake_dit()
|
||||
img = torch.ones(1, 5, 3) # nonzero -> i2v -> kept
|
||||
kwargs = {
|
||||
"image_embeds": img,
|
||||
"encoder_hidden_states": torch.arange(4.0).reshape(1, 4, 1),
|
||||
"encoder_attention_mask": torch.tensor([[1, 1, 1, 1]]),
|
||||
}
|
||||
_, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
|
||||
assert out["image_embeds"] is img # not emptied
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is True for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_trim_pre_hook_mixed_batch_flags_false():
|
||||
dit = _fake_dit()
|
||||
kwargs = {
|
||||
"image_embeds": torch.zeros(2, 2, 3),
|
||||
"encoder_hidden_states": torch.ones(2, 4, 1),
|
||||
"encoder_attention_mask": torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]]),
|
||||
}
|
||||
_, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is False for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_trim_pre_hook_never_raises_sets_flag_false():
|
||||
# A malformed mask (not a tensor) must not break the forward: flag False, no exception.
|
||||
dit = _fake_dit()
|
||||
kwargs = {"encoder_hidden_states": torch.ones(1, 2, 1), "encoder_attention_mask": "oops"}
|
||||
args, out = att._hunyuan_trim_pre_hook(dit, (), kwargs)
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is False for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_trim_pre_hook_absent_stream_not_written_back():
|
||||
# If encoder_hidden_states is absent from kwargs (a caller passing it positionally), the hook
|
||||
# must NOT write it back as None (that would collide: "got multiple values for argument") and
|
||||
# must drop the fast path (flag False) rather than null a mask it never verified.
|
||||
dit = _fake_dit()
|
||||
kwargs = {"image_embeds": torch.zeros(1, 4, 3)} # no encoder_hidden_states key
|
||||
_, out = att._hunyuan_trim_pre_hook(dit, (torch.ones(1, 5, 1),), kwargs)
|
||||
assert "encoder_hidden_states" not in out
|
||||
assert all(getattr(b.attn, att._NULL_ATTN_FLAG) is False for b in dit.transformer_blocks)
|
||||
|
||||
|
||||
def test_install_trim_noop_for_non_hunyuan_family():
|
||||
fam = types.SimpleNamespace(transformer_class="WanTransformer3DModel")
|
||||
pipe = types.SimpleNamespace(transformer=types.SimpleNamespace())
|
||||
assert att.install_hunyuan_attention_trim(pipe, fam) is False
|
||||
|
||||
|
||||
def test_install_trim_noop_when_transformer_class_mismatch():
|
||||
# Family claims Hunyuan but the loaded module isn't -> no processors touched, no diffusers
|
||||
# import; returns False rather than swapping an unknown attention processor.
|
||||
fam = types.SimpleNamespace(transformer_class="HunyuanVideo15Transformer3DModel")
|
||||
pipe = types.SimpleNamespace(transformer=types.SimpleNamespace()) # class name mismatch
|
||||
assert att.install_hunyuan_attention_trim(pipe, fam) is False
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue