huggingface / huggingface/diffusers

cuDNN attention CP path consumes a bool attn_mask as an additive 0/1 bias

Ouverte
#14,342 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
Langage dominant
Python
Étoiles
34.5k
Forks
7.3k
Merge moyen
3 j 3 h
PR mergées (30 j)
91

Description

# cuDNN attention CP path consumes a bool `attn_mask` as an additive 0/1 bias

Split out of #14338 so that issue can be closed cleanly by #14341, which fixes the
K/V transpose but not this.

## What's wrong

`_cudnn_attention_forward_op` passes `attn_mask` straight into
`torch.ops.aten._scaled_dot_product_cudnn_attention`:

```python
# src/diffusers/models/attention_dispatch.py
out, lse, ... = torch.ops.aten._scaled_dot_product_cudnn_attention(
query=query, key=key, value=value,
attn_bias=attn_mask, # <-- additive bias slot
...
)
```

The public `F.scaled_dot_product_attention` converts a boolean mask to an additive
`-inf`/`0` bias before dispatching. The raw ATen op does not. So a `torch.bool` mask
is consumed as an additive **0/1**: "masked" positions get `+0.0` and "keep" positions
get `+1.0`, which is neither a mask nor a no-op — it perturbs the scores of exactly the
positions that were supposed to be kept, and does not suppress the ones that were
supposed to be dropped.

`_native_flash_attention_forward_op` takes the same `attn_mask` argument through the
same wrapper, so it is worth checking whether it has the same exposure.

## Repro

With the conversion removed (i.e. current `main`), comparing the op against
`F.scaled_dot_product_attention` given the same boolean mask, bf16, `B=1 H=2 S=16 D=64`,
one partially masked row and one fully masked row:

```
partial FAIL
fully_masked_row FAIL
```

With a `-inf`/`0` conversion in front of the call, both pass.

## Fix

Convert before the call, matching what `F.scaled_dot_product_attention` does:

```python
if attn_mask is not None and attn_mask.dtype == torch.bool:
attn_mask = torch.zeros_like(attn_mask, dtype=query.dtype).masked_fill_(
attn_mask.logical_not(), float("-inf")
)
```

One note on the fill value, since it is easy to get wrong: `-inf` and
`torch.finfo(dtype).min` are **not** interchangeable in general — on a fully masked row
the CPU math path gives zeros for `-inf` and `mean(V)` for `finfo.min`. Through the
cuDNN op specifically I measured no difference between the two, so this is not a
cuDNN-visible bug; `-inf` is simply the value that matches boolean semantics on every
backend, which is what this wrapper is emulating.

I have this fix plus a parametrized regression test (partial / fully-masked row) rebased
on top of #14341, and can open it as a PR whenever a maintainer would like — happy to
wait until #14341 lands to avoid conflicting with it.

Found during an audit of cuDNN SDPA integrations across the ecosystem, run by the NVIDIA
cuDNN team.

Guide de contribution

Ouvrir le guide de contribution

Piste de recherche

Commencez dans src/diffusers/models/attention_dispatch.py à _cudnn_attention_forward_op, puis examinez _native_flash_attention_forward_op pour le même traitement des masques booléens. Comparez les deux chemins avec torch.nn.functional.scaled_dot_product_attention en utilisant les cas de lignes partiellement et entièrement masquées, et ajoutez le test de régression paramétré décrit dans l’issue. Le travail est terminé lorsque les masques booléens correspondent à la fonction publique sans modifier les masques non booléens.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Évaluation

Stack technique
python, pytorch
Domaine
machine-learning, performance
Type d'issue
Bug
Difficulté
3/5
Temps estimé
1-2 jours
Activité
Calme
Clarté
Clairement spécifiée
Accessibilité débutants
74/100

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.