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}")
|
||||
Loading…
Add table
Add a link
Reference in a new issue