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
|
|
@ -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