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.

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