From a5928064a07f00a8c57a4aaa9bdf91bd710cdb86 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 05:09:49 +0000 Subject: [PATCH] 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. --- scripts/hunyuan_attn_diag.py | 194 ++++++++++++ scripts/hunyuan_trim_e2e.py | 124 ++++++++ scripts/hunyuan_trim_fp32ref.py | 105 +++++++ scripts/hunyuan_trim_validate.py | 155 ++++++++++ scripts/sdpa_mask_backend_probe.py | 65 ++++ .../core/inference/diffusion_attention.py | 284 ++++++++++++++++++ studio/backend/core/inference/video.py | 25 +- .../backend/tests/test_diffusion_attention.py | 167 ++++++++++ 8 files changed, 1117 insertions(+), 2 deletions(-) create mode 100644 scripts/hunyuan_attn_diag.py create mode 100644 scripts/hunyuan_trim_e2e.py create mode 100644 scripts/hunyuan_trim_fp32ref.py create mode 100644 scripts/hunyuan_trim_validate.py create mode 100644 scripts/sdpa_mask_backend_probe.py diff --git a/scripts/hunyuan_attn_diag.py b/scripts/hunyuan_attn_diag.py new file mode 100644 index 000000000..4a7210cfb --- /dev/null +++ b/scripts/hunyuan_attn_diag.py @@ -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() diff --git a/scripts/hunyuan_trim_e2e.py b/scripts/hunyuan_trim_e2e.py new file mode 100644 index 000000000..2992e584c --- /dev/null +++ b/scripts/hunyuan_trim_e2e.py @@ -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() diff --git a/scripts/hunyuan_trim_fp32ref.py b/scripts/hunyuan_trim_fp32ref.py new file mode 100644 index 000000000..87372e656 --- /dev/null +++ b/scripts/hunyuan_trim_fp32ref.py @@ -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() diff --git a/scripts/hunyuan_trim_validate.py b/scripts/hunyuan_trim_validate.py new file mode 100644 index 000000000..97936ce07 --- /dev/null +++ b/scripts/hunyuan_trim_validate.py @@ -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() diff --git a/scripts/sdpa_mask_backend_probe.py b/scripts/sdpa_mask_backend_probe.py new file mode 100644 index 000000000..3dd83edea --- /dev/null +++ b/scripts/sdpa_mask_backend_probe.py @@ -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}") diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py index 8e60e1a5f..30849ada6 100644 --- a/studio/backend/core/inference/diffusion_attention.py +++ b/studio/backend/core/inference/diffusion_attention.py @@ -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 diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 3a181eaea..5b36ab677 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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", diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py index da200f2c4..13a183d25 100644 --- a/studio/backend/tests/test_diffusion_attention.py +++ b/studio/backend/tests/test_diffusion_attention.py @@ -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