a5928064a0
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.
65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
#!/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}")
|