Skip to content

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

Description

@YangXu1990uiuc

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:

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

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions