[bugfix]: keep sdpa metadata masks 2D; lift to 4D only inside the torch-sdpa impl

SDPAMetadataBuilder.build() lifted 2D [batch, key_len] padding masks to 4D
[batch, 1, 1, key_len], but the metadata is cross-backend: HYWorld and
HunyuanVideo15 build SDPAMetadata while the layer selector picks FLASH_ATTN
when installed. flash_attn.py's equal-length branch pads
attn_mask.shape[1] assuming 2D, so a 4D mask produced a garbage-length pad
and an out-of-bounds gather (CUDA device-side assert, SIGABRT in CI
transformer tests).

Store the mask 2D in the builder (the convention every in-tree
producer/consumer assumes) and do the 2D -> [batch, 1, 1, key_len] lift
inside _normalize_attn_mask_for_sdpa, the torch-sdpa-only consumer, so
torch.sdpa still does not misread 2D as its [query_len, key_len]
broadcast. The int -> bool coercion is unchanged. Adds a CPU-only
regression test for the contract.

fastvideo/attention/backends/sdpa.py

@@ -50,13 +50,11 @@ class SDPAMetadataBuilder(AttentionMetadataBuilder):
             current_timestep: int,
             attn_mask: torch.Tensor,
     ) -> SDPAMetadata:
-        # A 2D mask handed to the builder is a tokenizer-style
-        # [batch, key_len] padding mask (the HunyuanVideo/HYWorld call
-        # sites); lift it to a broadcastable [batch, 1, 1, key_len] HERE,
-        # where that semantic is known. Masks that reach the impl as 2D keep
-        # torch's documented [query_len, key_len] broadcast meaning.
-        if attn_mask is not None and attn_mask.dim() == 2:
-            attn_mask = attn_mask[:, None, None, :]
+        # Store the mask exactly as passed. The metadata is cross-backend:
+        # call sites (HYWorld, HunyuanVideo15) build SDPAMetadata while the
+        # layer's selector may pick FLASH_ATTN, whose impl requires the
+        # tokenizer-style 2D [batch, key_len] padding mask. Any reshaping for
+        # torch.sdpa happens inside the SDPA impl (_normalize_attn_mask_for_sdpa).
         return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask)
 
 
@@ -84,11 +82,14 @@ def _normalize_attn_mask_for_sdpa(
         valid_value = True if attn_mask.dtype == torch.bool else 0.0
         attn_mask = F.pad(attn_mask, (key_len - attn_mask.shape[-1], 0), value=valid_value)
 
+    if attn_mask.dim() == 2:
+        # In-tree producers pass 2D [batch, key_len] padding masks; lift to a
+        # broadcastable [batch, 1, 1, key_len] here so torch.sdpa does not
+        # reinterpret 2D as its documented [query_len, key_len] broadcast.
+        return attn_mask[:, None, None, :]
     if attn_mask.dim() == 3:
         return attn_mask[:, None, :, :]
-    if attn_mask.dim() in (2, 4):
-        # 2D keeps torch's documented [query_len, key_len] broadcast; batch
-        # padding masks are lifted to 4D by SDPAMetadataBuilder.build.
+    if attn_mask.dim() == 4:
         return attn_mask
     raise ValueError(f"Unsupported attention mask shape for SDPA: {attn_mask.shape}")
 

fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only regression test for the SDPA metadata mask contract.
+
+SDPAMetadata is cross-backend: HYWorld and HunyuanVideo15 build it while the
+layer's attention selector may pick FLASH_ATTN, whose impl requires the
+tokenizer-style 2D [batch, key_len] padding mask (flash_attn.py pads
+``attn_mask.shape[1]`` assuming 2D). The builder must therefore store 2D
+masks unchanged; only the torch-sdpa impl lifts them to 4D internally.
+
+Regression guard for the CI break where the builder lifted 2D -> 4D and the
+FLASH_ATTN consumer produced a garbage-length pad (CUDA device-side assert).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from fastvideo.attention.backends.sdpa import (SDPAMetadataBuilder, _normalize_attn_mask_for_sdpa)
+
+
+def test_builder_keeps_2d_padding_mask_2d() -> None:
+    """A tokenizer-style 2D int mask must be stored as-is (cross-backend contract)."""
+    mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.int64)
+    md = SDPAMetadataBuilder().build(current_timestep=0, attn_mask=mask)
+    assert md.attn_mask is not None
+    assert md.attn_mask.dim() == 2
+    assert md.attn_mask.shape == (2, 3)
+    assert torch.equal(md.attn_mask, mask)
+
+
+def test_normalize_lifts_2d_int_mask_to_4d_bool_for_torch_sdpa() -> None:
+    """The torch-sdpa consumer coerces int -> bool and lifts 2D -> [B,1,1,K]."""
+    batch, q_len, k_len, heads, head_dim = 2, 4, 3, 1, 8
+    query = torch.zeros(batch, heads, q_len, head_dim)
+    key = torch.zeros(batch, heads, k_len, head_dim)
+    mask = torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.int64)
+
+    out = _normalize_attn_mask_for_sdpa(mask, query, key)
+    assert out is not None
+    assert out.dtype == torch.bool
+    assert out.shape == (batch, 1, 1, k_len)
+    assert torch.equal(out[:, 0, 0, :], mask != 0)