From 9027195db0a8abde6af730f60aab5041307dd21d Mon Sep 17 00:00:00 2001 From: Functionhx <2994114386@qq.com> Date: Fri, 10 Jul 2026 03:35:18 +0800 Subject: [PATCH 1/2] Fix non-hub varlen attention backends mispacking keys for non-contiguous masks The three non-hub varlen backends (_flash_varlen_attention, _flash_varlen_attention_3, _sage_varlen_attention) used a per-batch slice pattern (key[b, :valid_len]) that fails when the attention mask is non-contiguous. The hub backends (FLASH_VARLEN_HUB, _FLASH_3_VARLEN_HUB) already received this fix via the nonzero gather pattern. Root cause: The per-batch slice assumes a contiguous mask layout, but non-contiguous masks have different cu_seqlens structures. The slice picks up wrong positions from each batch element. Fix: When a mask is present, use attn_mask.flatten().nonzero().flatten() to gather valid positions from reshaped key/value tensors, which correctly handles any mask contiguity. When no mask, use simple flatten(), which is always contiguous. This matches the exact pattern used in the hub backends. Closes #14139. Signed-off-by: Yuchen Fan --- src/diffusers/models/attention_dispatch.py | 72 ++++++++++------------ 1 file changed, 33 insertions(+), 39 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 9414c151fd67..f49e408e91d7 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -3049,22 +3049,20 @@ def _flash_varlen_attention( if attn_mask is not None: attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device) ) - ) - - key_valid, value_valid = [], [] - for b in range(batch_size): - valid_len = seqlens_k[b] - key_valid.append(key[b, :valid_len]) - value_valid.append(value[b, :valid_len]) + indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten() + key_packed = key.reshape(-1, *key.shape[2:])[indices_k] + value_packed = value.reshape(-1, *value.shape[2:])[indices_k] + else: + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device) + ) + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) query_packed = query.flatten(0, 1) - key_packed = torch.cat(key_valid, dim=0) - value_packed = torch.cat(value_valid, dim=0) out = flash_attn_varlen_func( q=query_packed, @@ -3331,22 +3329,20 @@ def _flash_varlen_attention_3( if attn_mask is not None: attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device) ) - ) - - key_valid, value_valid = [], [] - for b in range(batch_size): - valid_len = seqlens_k[b] - key_valid.append(key[b, :valid_len]) - value_valid.append(value[b, :valid_len]) + indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten() + key_packed = key.reshape(-1, *key.shape[2:])[indices_k] + value_packed = value.reshape(-1, *value.shape[2:])[indices_k] + else: + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device) + ) + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) query_packed = query.flatten(0, 1) - key_packed = torch.cat(key_valid, dim=0) - value_packed = torch.cat(value_valid, dim=0) result = flash_attn_3_varlen_func( q=query_packed, @@ -3964,22 +3960,20 @@ def _sage_varlen_attention( if attn_mask is not None: attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device) ) - ) - - key_valid, value_valid = [], [] - for b in range(batch_size): - valid_len = seqlens_k[b] - key_valid.append(key[b, :valid_len]) - value_valid.append(value[b, :valid_len]) + indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten() + key_packed = key.reshape(-1, *key.shape[2:])[indices_k] + value_packed = value.reshape(-1, *value.shape[2:])[indices_k] + else: + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device) + ) + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) query_packed = query.flatten(0, 1) - key_packed = torch.cat(key_valid, dim=0) - value_packed = torch.cat(value_valid, dim=0) out = sageattn_varlen( q=query_packed, From 5d3c551ea9c06abeec33aa3443ec3238406732c8 Mon Sep 17 00:00:00 2001 From: Functionhx <2994114386@qq.com> Date: Fri, 10 Jul 2026 03:39:51 +0800 Subject: [PATCH 2/2] Fix FlaxDiffusionPipeline.download() on Windows: same os.path.join backslash bug Same root cause as the DiffusionPipeline fix: os.path.join(k, "*") produces backslash paths on Windows (e.g. text_encoder\*), which fnmatch does not normalize, causing pattern mismatch and silently skipped subfolder downloads. Fix: use f"{k}/*" for forward-slash paths, consistent with pipeline_utils.py. Fixes huggingface/diffusers#14142 Signed-off-by: Yuchen Fan --- src/diffusers/pipelines/pipeline_flax_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pipeline_flax_utils.py b/src/diffusers/pipelines/pipeline_flax_utils.py index 8c29db8aa45e..b81d81265bf5 100644 --- a/src/diffusers/pipelines/pipeline_flax_utils.py +++ b/src/diffusers/pipelines/pipeline_flax_utils.py @@ -336,7 +336,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None ) # make sure we only download sub-folders and `diffusers` filenames folder_names = [k for k in config_dict.keys() if not k.startswith("_")] - allow_patterns = [os.path.join(k, "*") for k in folder_names] + allow_patterns = [f"{k}/*" for k in folder_names] allow_patterns += [FLAX_WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, cls.config_name] ignore_patterns = ["*.bin", "*.safetensors"] if not from_pt else []