pytorch / pytorch/pytorch

[TorchInductor] Masked SDPA fusion misses commuted mask addition

Open
#195,782 1 comment 0 reactions 0 assignees View on GitHub
bot-triaged module: inductor module: performance module: sdpa oncall: pt2 triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.