Skip to content

Convert bool attn_mask to additive bias in cuDNN forward op - #14349

Open
guptaishaan wants to merge 1 commit into
huggingface:mainfrom
guptaishaan:fix-14342
Open

Convert bool attn_mask to additive bias in cuDNN forward op#14349
guptaishaan wants to merge 1 commit into
huggingface:mainfrom
guptaishaan:fix-14342

Conversation

@guptaishaan

Copy link
Copy Markdown

Fixes #14342

_cudnn_attention_forward_op passes attn_mask into
torch.ops.aten._scaled_dot_product_cudnn_attention as attn_bias. That slot is an additive bias
in the query dtype, and unlike F.scaled_dot_product_attention the raw ATen op does not convert a
torch.bool mask first, so a boolean mask is consumed as non-float data and the masked positions
are not suppressed. This patch converts a boolean mask to a -inf/0 bias before the call, which
is what F.scaled_dot_product_attention does. It sits before ctx.attn_mask = attn_mask, so the
backward op gets the converted bias too.

-inf rather than torch.finfo(dtype).min: they are not interchangeable on a fully masked row
(zeros vs mean(V)), and -inf is the value that matches boolean-mask semantics on every backend,
which is what this wrapper is emulating.

New tests/models/test_attention_dispatch.py drives the forward op with a bool mask, once with a
partially masked row and once with a fully masked row, and compares against
F.scaled_dot_product_attention given the same mask. Both cases fail before the patch and pass
after.

Verified

1x A40 (sm_86), torch 2.9.0+cu126, cuDNN 9.10.2, bf16, B=1 H=2 S=16 D=64:

  • Standalone repro, max abs error against F.scaled_dot_product_attention goes from 2.12 (partial)
    and 1.15 (fully masked row) to exactly 0.0 on both, against a reference magnitude of ~1.8.
  • Pre-fix, the bool-mask output is bit identical to the no mask output for one head and garbage
    for the other, i.e. on this stack the mask is dropped rather than applied as an additive 0/1. Same
    numbers under torch 2.13.0+cu126 driving the ATen op directly.
  • pytest -m "attention or context_parallel" over test_attention_dispatch.py,
    test_attention_processor.py and test_models_transformer_flux.py: 7 passed, 41 skipped. The 8
    failures are *_hub backends failing to fetch a kernel build for this system and reproduce on a
    clean tree.
  • ruff check, ruff format --check, utils/check_copies.py all clean.

Not verified

  • Multi-GPU ulysses/ring end to end with a boolean mask. The fix is tested at the forward op, which
    is the layer the bug is at.
  • Gradients. The conversion lands in ctx.attn_mask, which is what the ATen backward wants, but the
    backward through this path has separate open problems (double K/V transpose, and
    compute_log_sumexp=return_lse leaving the backward with an empty lse), so a grad check here
    would not be meaningful yet.
  • bf16 on sm_86 only. No fp16, no sm_89+, no H100.
  • _native_flash_attention_forward_op does not have this bug, since
    aten::_scaled_dot_product_flash_attention has no bias argument. It accepts attn_mask and
    silently drops it, which is a separate problem and is left alone here.

Reported by @YangXu1990uiuc, who also identified the fix and the -inf vs finfo.min distinction.

`_cudnn_attention_forward_op` passes `attn_mask` into
`torch.ops.aten._scaled_dot_product_cudnn_attention` as `attn_bias`. That
slot is an additive bias in the query dtype, and the raw ATen op does not
convert a `torch.bool` mask the way `F.scaled_dot_product_attention` does,
so a boolean mask is consumed as non-float data and the masked positions
are not suppressed. On an A40 with torch 2.9 the bool-mask output is bit
identical to the no-mask output for one head and garbage for the other.

Convert a boolean mask to a `-inf`/`0` bias before the call. The
conversion happens before `ctx.attn_mask` is saved, so the backward op
gets the converted bias as well. `-inf` rather than `finfo(dtype).min`,
since the two differ on a fully masked row and `-inf` is what boolean
mask semantics mean on every backend.

Adds `tests/models/test_attention_dispatch.py`, which drives the forward
op with a partially masked row and a fully masked row and compares
against `F.scaled_dot_product_attention` with the same mask.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant