[TorchInductor] `_sfdp_pattern_25` and `_sfdp_pattern_26` miss FP16 Q/K/V with an FP32 attention mask
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
`_sfdp_pattern_25` and `_sfdp_pattern_26` do not cover the mixed-precision case where Q/K/V are FP16 but the attention mask is FP32. Both patterns match their FP16-mask baselines, but changing only the mask dtype to FP32 prevents the SDPA fusion.
## Minimal reproducer
The following tests use the same attention computation, with FP16 Q/K/V, input shape `(1, 4, 2, 8)`, mask shape `(1, 1, 4, 4)`, and inference dropout.
```python
"""_sfdp_pattern_25 candidate: FP16 Q/K/V + FP32 mask."""
from __future__ import annotations
import torch
import torch.nn.functional as F
class Model(torch.nn.Module):
def forward(self, q, k, v, mask):
q, k, v = (x.permute(0, 2, 1, 3) for x in (q, k, v))
score = q @ k.permute(0, 1, 3, 2) + mask
return (
F.dropout(
score.float().softmax(-1).to(q.dtype),
p=0.1,
training=False,
)
@ v
)
def make_case(device):
torch.manual_seed(1845)
qkv = tuple(
torch.randn(
1, 4, 2, 8,
device=device,
dtype=torch.float16,
)
for _ in range(3)
)
mask = torch.zeros(
1, 1, 4, 4,
device=device,
dtype=torch.float32,
)
return Model().to(device).eval(), (*qkv, mask)
def main():
device = torch.device("cpu")
model, inputs = make_case(device)
with torch.no_grad():
eager = model(*inputs)
torch._dynamo.reset()
compiled = torch.compile(model, fullgraph=True)
with torch.no_grad():
actual = compiled(*inputs)
torch.testing.assert_close(actual, eager, rtol=1e-2, atol=1e-2)
if __name__ == "__main__":
main()
```
```python
"""_sfdp_pattern_26 candidate: FP16 Q/K/V + FP32 mask."""
from __future__ import annotations
import torch
import torch.nn.functional as F
class Model(torch.nn.Module):
def forward(self, q, k, v, mask):
q, k, v = (x.permute(0, 2, 1, 3) for x in (q, k, v))
score = q @ k.permute(0, 1, 3, 2) + mask
output = (
F.dropout(
score.float().softmax(-1).to(q.dtype),
p=0.1,
training=False,
)
@ v
)
return output, k, v
def make_case(device):
torch.manual_seed(1846)
qkv = tuple(
torch.randn(
1, 4, 2, 8,
device=device,
dtype=torch.float16,
)
for _ in range(3)
)
mask = torch.zeros(
1, 1, 4, 4,
device=device,
dtype=torch.float32,
)
return Model().to(device).eval(), (*qkv, mask)
def main():
device = torch.device("cpu")
model, inputs = make_case(device)
with torch.no_grad():
eager = model(*inputs)
torch._dynamo.reset()
compiled_model = torch.compile(model, fullgraph=True)
with torch.no_grad():
actual = compiled_model(*inputs)
for actual_tensor, eager_tensor in zip(actual, eager):
torch.testing.assert_close(
actual_tensor,
eager_tensor,
rtol=1e-2,
atol=1e-2,
)
if __name__ == "__main__":
main()
```
## Observed behavior
**Instrumentation Results:**
```text
=== _sfdp_pattern_25 candidate.py ===
no HIT _sfdp_pattern_25
=== _sfdp_pattern_26 candidate.py ===
no HIT _sfdp_pattern_26
```
Each baseline differs from its candidate only in the attention-mask dtype.
```diff
- mask = torch.zeros(1, 1, 4, 4, dtype=torch.float16, device=device)
+ mask = torch.zeros(1, 1, 4, 4, dtype=torch.float32, device=device)
```
The FP16-mask baselines hit their corresponding patterns.
```text
=== _sfdp_pattern_25 baseline.py ===
HIT _sfdp_pattern_25
=== _sfdp_pattern_26 baseline.py ===
HIT _sfdp_pattern_26
```
**Graph Results:**
The post-joint graphs confirm the instrumentation results. The FP16-mask baselines contain `aten._scaled_dot_product_flash_attention_for_cpu.default`, whereas the FP32-mask candidates retain the unfused `bmm → add → softmax → bmm` computation. Pattern 26 shows the same behavior.
## Root Cause
In `torch/_inductor/fx_passes/fuse_attention.py`, patterns 25 and 26 are registered with `m()`/`m_bs1()`, whose dtype follows Q/K/V. They have no `m_float()`/`m_bs1_float()` candidates, so their serialized FP16 patterns require an FP16 mask even though `_sfdp_params_check` permits an FP32 mask.
The same mixed-dtype registration appears to be missing from patterns 5, 6, 14, and 29. In contrast, 16, 19, 21, 22, 24 already register explicit FP32-mask candidates and do not have this particular coverage gap. Only patterns 25 and 26 were executed here. Patterns 5, 6, 14, and 29 have a similar registration gap based on source inspection, but their mixed-dtype variants were not independently validated.
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 in torch/_inductor/fx_passes/fuse_attention.py by comparing the registration of patterns 25 and 26 with the existing m_float() and m_bs1_float() candidates. Run the supplied FP16 Q/K/V plus FP32-mask reproducers and inspect the post-joint graph. Done means both cases hit their corresponding patterns and produce the fused CPU scaled-dot-product attention operation; consider the similarly noted patterns only after validating them.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100