huggingface / huggingface/diffusers

Non-hub varlen attention backends mispack keys for non-contiguous masks (missed in #14114/#14115)

Open
#14,139 1 comment 0 reactions 0 assignees View on GitHub
bug models
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

The three non-hub varlen backends pack keys/values by slicing `key[b, :valid_len]`, where `valid_len = attn_mask.sum(dim=1)` is only the *count* of valid keys. This assumes the valid keys are contiguous from index 0. For a left-padded or interior-gap mask, this packs the wrong keys: attention runs over padding positions and drops real keys, producing silently wrong output with no error.

The hub variants were fixed for exactly this in #14114 / #14115 (and #13999 / #14102), but the three non-hub backends were not updated.

**Affected** (`src/diffusers/models/attention_dispatch.py`):
- `_flash_varlen_attention` (`flash_varlen`) at [L2872](https://github.com/huggingface/diffusers/blob/408ce8db29b9430786937dc9dec98d212116846f/src/diffusers/models/attention_dispatch.py#L2872)
- `_flash_varlen_attention_3` (`_flash_varlen_3`) at [L3124](https://github.com/huggingface/diffusers/blob/408ce8db29b9430786937dc9dec98d212116846f/src/diffusers/models/attention_dispatch.py#L3124)
- `_sage_varlen_attention` (`sage_varlen`) at [L3757](https://github.com/huggingface/diffusers/blob/408ce8db29b9430786937dc9dec98d212116846f/src/diffusers/models/attention_dispatch.py#L3757)

**Correct behavior:** gather by the mask's true indices, as the hub variants already do ([L2789](https://github.com/huggingface/diffusers/blob/408ce8db29b9430786937dc9dec98d212116846f/src/diffusers/models/attention_dispatch.py#L2789)): `indices_k = attn_mask.flatten().nonzero().flatten()`, then index the flattened key/value.

**Packing divergence** (CPU, no kernels needed):
```python
import torch
# joint mask [valid_text, PAD, image...], an interior gap (as QwenImage builds):
attn_mask = torch.tensor([[True, True, False, True, True]]) # 2 text + 1 pad + 2 image
valid_len = attn_mask.sum(1) # tensor([4]), count only

wrong = torch.arange(5)[: valid_len.item()] # [0,1,2,3]: keeps PAD@2, drops key@4
correct = attn_mask.flatten().nonzero().flatten() # [0,1,3,4]
```

**Real trigger:**
QwenImage builds a joint mask `cat([encoder_hidden_states_mask, image_mask])` ([transformer_qwenimage.py:508](https://github.com/huggingface/diffusers/blob/408ce8db29b9430786937dc9dec98d212116846f/src/diffusers/models/transformers/transformer_qwenimage.py#L508)); with right-padded text and an all-ones image block, the mask has an interior gap. Users who select a varlen backend (e.g. via `set_attention_backend`) with batched, differing-length prompts get corrupted output. The default backend isn't varlen, so it is opt-in but silent.

**Related:**
#13809 fixes only the hub FA3 varlen path. #12870 adds separate `*_split` backends and does not change this packing, so this fix is independent.


I can send PR for this. The three backends diverged because the varlen packing is duplicated inline with no single source of truth, so the hub fix never reached them. I'll fix the non-hub backends, lift the shared packing into one helper (scoped to non-hub, to stay clear of #12870) so it cannot recur, and add a CPU regression matrix across backends and mask shapes.

Let me know ~

Contributor guide

Open the contributing guide

Research direction

Start in src/diffusers/models/attention_dispatch.py at _flash_varlen_attention, _flash_varlen_attention_3, and _sage_varlen_attention, then compare their packing with the hub variant around L2789. Reproduce the issue with the CPU mask example and verify completion with regression coverage for the three backends and the described mask shapes, including the QwenImage mask construction in transformer_qwenimage.py.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.