[TorchInductor] Masked SDPA fusion misses commuted mask addition
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
Masked SDPA fusion patterns encode the attention-mask addition as `scores + attn_mask` and miss the equivalent `attn_mask + scores` form. I reproduced this with `_sfdp_pattern_25` and `_sfdp_pattern_14`. The baseline is fused, while exchanging only the two operands leaves the attention computation unfused.
## Minimal reproducer
The reproducer is based on the upstream `_test_sdpa_rewriter_25` case. It uses FP16 inputs with shape `(4, 2, 16, 32)`, an attention mask with shape `(1, 1, 1, 2)`, and the same softmax and inference-dropout structure.
```python
"""Positive case: score + attn_mask."""
from __future__ import annotations
import torch
import torch.nn.functional as F
from torch._dynamo.utils import counters
class Model(torch.nn.Module):
def forward(self, q, k, v, mask, training):
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
scores = q @ k.permute(0, 1, 3, 2)
masked_score = scores + mask
return F.dropout(
masked_score.float().softmax(-1).type(v.dtype), p=0.1, training=training
) @ v
def main():
torch.manual_seed(1845)
device = torch.device("cpu")
model = Model().to(device).eval()
qkv = tuple(
torch.randn(4, 2, 16, 32, dtype=torch.half, device=device)
for _ in range(3)
)
inputs = (*qkv, torch.randn(1, 1, 1, 2, dtype=torch.half, device=device))
compiled_inputs = [x.detach().clone() for x in inputs]
training = False
torch.manual_seed(1846)
eager = model(*inputs, training)
torch._dynamo.reset()
counters.clear()
torch.manual_seed(1846)
compiled = torch.compile(model, backend="inductor", fullgraph=True)(
*compiled_inputs, training
)
torch.testing.assert_close(compiled, eager, atol=2e-3, rtol=2e-3)
print("fuse_attention:", counters["inductor"]["fuse_attention"])
if __name__ == "__main__":
main()
```
The commuted-operand case is identical except for the following single-line change:
```diff
- masked_score = scores + mask
+ masked_score = mask + scores
```
## Observed behavior
**Instrumentation**
The baseline matches `_sfdp_pattern_25`, while the commuted form does not.
```text
scores + mask: HIT _sfdp_pattern_25
mask + scores: no HIT _sfdp_pattern_25
```
**Graph**
The post-joint graphs confirm the instrumentation result:
```text
scores + mask:
aten._scaled_dot_product_flash_attention_for_cpu.default
mask + scores:
bmm → add → softmax decomposition → bmm
```
## Independent verification with pattern 14
I also repeated the same phenomenon for `_sfdp_pattern_14`.
```python
"""Positive control for _sfdp_pattern_14: scores + attn_mask."""
from __future__ import annotations
import torch
from torch._dynamo.utils import counters
class Model(torch.nn.Module):
def forward(self, query, key, value):
attn_mask = torch.ones(
query.size(1), key.size(1), dtype=torch.bool, device=query.device
).tril(diagonal=0)
attn_mask = attn_mask.masked_fill(
torch.logical_not(attn_mask), -float("inf")
)
q = query.permute(0, 2, 1, 3)
k = key.permute(0, 2, 1, 3)
v = value.permute(0, 2, 1, 3)
scores = torch.matmul(q, k.transpose(-2, -1)).div(3.0)
masked_score = scores + attn_mask
return masked_score.softmax(dim=-1).matmul(v)
def main():
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for this test")
torch.manual_seed(1012)
device = torch.device("cuda:0")
model = Model().to(device).eval()
inputs = tuple(
torch.randn(4, 2, 16, 32, device=device, requires_grad=True)
for _ in range(3)
)
with torch.no_grad():
eager = model(*inputs)
torch._dynamo.reset()
counters.clear()
with torch.no_grad():
compiled = torch.compile(model, backend="inductor", fullgraph=True)(*inputs)
torch.cuda.synchronize()
torch.testing.assert_close(compiled, eager, atol=2e-3, rtol=2e-3)
print("fuse_attention:", counters["inductor"]["fuse_attention"])
if __name__ == "__main__":
main()
```
The commuted case again changes only the operand order.
```diff
- masked_score = scores + mask
+ masked_score = mask + scores
```
The same behavior occurs with `_sfdp_pattern_14`. Exchanging only the operands of the mask addition changes the result from a pattern hit to a miss.
## Root Cause
The SDPA pattern encodes `aten.add.Tensor(score, attn_mask)` with a fixed operand order. The commutative `aten.add.Tensor(attn_mask, score)` form is therefore not matched. The shared SDPA parameter check also assumes that `add.args[1]` is the attention mask.
Several other SDPA patterns—including 5, 6, 16, 19, 21, 22, 24, 26, and 29—appear to encode the same fixed operand order and may therefore be affected. These patterns were not independently tested here.
Reproduced on PyTorch main commit `f07882e`.
### Versions
PyTorch main commit f07882e
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @drisspg @liangel-02 @howardzhang-cv
Contributor guide
Research direction
Start by locating the TorchInductor SDPA fusion pattern definitions and the shared SDPA parameter check referenced in the issue. Compare _sfdp_pattern_25 and _sfdp_pattern_14 with the upstream _test_sdpa_rewriter_25 case, then add coverage for both operand orders and verify that the commuted forms fuse without regressing existing patterns.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100