huggingface / huggingface/diffusers
Non-hub varlen attention backends mispack keys for non-contiguous masks (missed in #14114/#14115)
- 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
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