From 98df094cbc505bca211165bc9e4fdf7f0c72871d Mon Sep 17 00:00:00 2001 From: Will Lin Date: Mon, 6 Jul 2026 22:53:09 +0000 Subject: [PATCH] [bugfix]: coerce integer sdpa masks and keep torch 2D mask semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _normalize_attn_mask_for_sdpa passed tokenizer-produced int64 0/1 masks straight to F.scaled_dot_product_attention, which rejects long dtypes — the HYWorld pipeline crashed on exactly the SDPA-fallback machines the fastcheck stabilization targeted; coerce non-bool/non-float masks to bool. Lift the [batch, key_len] -> [batch, 1, 1, key_len] reinterpretation out of the impl into SDPAMetadataBuilder.build, where that semantic is known, so 2D masks reaching the impl keep torch's documented [query_len, key_len] broadcast meaning. Reuse the identical metadata for HYWorld's prope attention instead of rebuilding it, and note that the metadata does not choose the executing kernel. Co-Authored-By: Claude Fable 5 --- fastvideo/attention/backends/sdpa.py | 20 +++++++++++++++++--- fastvideo/models/dits/hyworld/hyworld.py | 12 +++++------- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/fastvideo/attention/backends/sdpa.py b/fastvideo/attention/backends/sdpa.py index a99c3e71e..6bf1bc101 100644 --- a/fastvideo/attention/backends/sdpa.py +++ b/fastvideo/attention/backends/sdpa.py @@ -50,6 +50,13 @@ 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, :] return SDPAMetadata(current_timestep=current_timestep, attn_mask=attn_mask) @@ -62,19 +69,26 @@ def _normalize_attn_mask_for_sdpa( return None attn_mask = attn_mask.to(device=query.device) + # F.scaled_dot_product_attention only accepts bool or float masks; + # tokenizers commonly produce int64 0/1 padding masks. + if attn_mask.dtype != torch.bool and not attn_mask.dtype.is_floating_point: + attn_mask = attn_mask != 0 + key_len = key.shape[-2] if attn_mask.shape[-1] > key_len: raise ValueError("Invalid attention mask length for SDPA: " f"expected at most {key_len}, got {attn_mask.shape[-1]}") if attn_mask.shape[-1] < key_len: + # Front-pad as "attend": double-stream layouts (HYWorld) prepend + # non-text tokens the tokenizer mask does not cover. 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: - return attn_mask[:, None, None, :] if attn_mask.dim() == 3: return attn_mask[:, None, :, :] - if attn_mask.dim() == 4: + 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. return attn_mask raise ValueError(f"Unsupported attention mask shape for SDPA: {attn_mask.shape}") diff --git a/fastvideo/models/dits/hyworld/hyworld.py b/fastvideo/models/dits/hyworld/hyworld.py index 402c49632..b16cd72e5 100644 --- a/fastvideo/models/dits/hyworld/hyworld.py +++ b/fastvideo/models/dits/hyworld/hyworld.py @@ -173,6 +173,9 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): img_v_prope = img_v_prope.permute(0, 2, 1, 3) # [batch, seqlen, num_heads, head_dim] # end hyworld + # The metadata only carries the text padding mask; the executing + # attention kernel is still whatever the layer's selector picked + # (flash-attn when installed), not necessarily torch SDPA. from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder attn_metadata = SDPAMetadataBuilder().build( current_timestep=0, @@ -192,14 +195,9 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock): ) # begin hyworld - # attention with prope - from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder - attn_metadata_prope = SDPAMetadataBuilder().build( - current_timestep=0, - attn_mask=encoder_attention_mask, - ) + # attention with prope (same text mask, so reuse the metadata) # NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope - with set_forward_context(current_timestep=0, attn_metadata=attn_metadata_prope): + with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): img_attn_prope, _ = self.attn( img_q_prope, img_k_prope,