Convert bool attn_mask to additive bias in cuDNN forward op - #14349
Open
guptaishaan wants to merge 1 commit into
Open
Convert bool attn_mask to additive bias in cuDNN forward op#14349guptaishaan wants to merge 1 commit into
guptaishaan wants to merge 1 commit into
Conversation
`_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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #14342
_cudnn_attention_forward_oppassesattn_maskintotorch.ops.aten._scaled_dot_product_cudnn_attentionasattn_bias. That slot is an additive biasin the query dtype, and unlike
F.scaled_dot_product_attentionthe raw ATen op does not convert atorch.boolmask first, so a boolean mask is consumed as non-float data and the masked positionsare not suppressed. This patch converts a boolean mask to a
-inf/0bias before the call, whichis what
F.scaled_dot_product_attentiondoes. It sits beforectx.attn_mask = attn_mask, so thebackward op gets the converted bias too.
-infrather thantorch.finfo(dtype).min: they are not interchangeable on a fully masked row(zeros vs
mean(V)), and-infis the value that matches boolean-mask semantics on every backend,which is what this wrapper is emulating.
New
tests/models/test_attention_dispatch.pydrives the forward op with a bool mask, once with apartially masked row and once with a fully masked row, and compares against
F.scaled_dot_product_attentiongiven the same mask. Both cases fail before the patch and passafter.
Verified
1x A40 (sm_86), torch 2.9.0+cu126, cuDNN 9.10.2, bf16,
B=1 H=2 S=16 D=64:F.scaled_dot_product_attentiongoes from 2.12 (partial)and 1.15 (fully masked row) to exactly 0.0 on both, against a reference magnitude of ~1.8.
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"overtest_attention_dispatch.py,test_attention_processor.pyandtest_models_transformer_flux.py: 7 passed, 41 skipped. The 8failures are
*_hubbackends failing to fetch a kernel build for this system and reproduce on aclean tree.
ruff check,ruff format --check,utils/check_copies.pyall clean.Not verified
is the layer the bug is at.
ctx.attn_mask, which is what the ATen backward wants, but thebackward through this path has separate open problems (double K/V transpose, and
compute_log_sumexp=return_lseleaving the backward with an empty lse), so a grad check herewould not be meaningful yet.
_native_flash_attention_forward_opdoes not have this bug, sinceaten::_scaled_dot_product_flash_attentionhas no bias argument. It acceptsattn_maskandsilently drops it, which is a separate problem and is left alone here.
Reported by @YangXu1990uiuc, who also identified the fix and the
-infvsfinfo.mindistinction.