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.
cuDNN attention CP path consumes a bool
attn_maskas an additive 0/1 biasSplit 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_oppassesattn_maskstraight intotorch.ops.aten._scaled_dot_product_cudnn_attention:The public
F.scaled_dot_product_attentionconverts a boolean mask to an additive-inf/0bias before dispatching. The raw ATen op does not. So atorch.boolmaskis consumed as an additive 0/1: "masked" positions get
+0.0and "keep" positionsget
+1.0, which is neither a mask nor a no-op — it perturbs the scores of exactly thepositions that were supposed to be kept, and does not suppress the ones that were
supposed to be dropped.
_native_flash_attention_forward_optakes the sameattn_maskargument through thesame wrapper, so it is worth checking whether it has the same exposure.
Repro
With the conversion removed (i.e. current
main), comparing the op againstF.scaled_dot_product_attentiongiven the same boolean mask, bf16,B=1 H=2 S=16 D=64,one partially masked row and one fully masked row:
With a
-inf/0conversion in front of the call, both pass.Fix
Convert before the call, matching what
F.scaled_dot_product_attentiondoes:One note on the fill value, since it is easy to get wrong:
-infandtorch.finfo(dtype).minare not interchangeable in general — on a fully masked rowthe CPU math path gives zeros for
-infandmean(V)forfinfo.min. Through thecuDNN op specifically I measured no difference between the two, so this is not a
cuDNN-visible bug;
-infis simply the value that matches boolean semantics on everybackend, 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.