huggingface / huggingface/diffusers

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

オープン
#14,342 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

# 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.

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

src/diffusers/models/attention_dispatch.py の _cudnn_attention_forward_op から始め、次に同じ boolean マスクの処理について _native_flash_attention_forward_op を調べます。部分的にマスクされた行と完全にマスクされた行のケースを使って、両方のパスを torch.nn.functional.scaled_dot_product_attention と比較し、issue に記載されたパラメータ化された回帰テストを追加します。boolean マスクが公開関数と一致し、boolean でないマスクを変更しないことが完了の条件です。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning, performance
issue の種類
バグ
難易度
3/5
見積もり時間
1〜2日
活発さ
静か
明瞭さ
明確に書かれている
初心者へのやさしさ
74/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。