pytorch / pytorch/pytorch

[TorchInductor] `_sfdp_pattern_25` and `_sfdp_pattern_26` miss FP16 Q/K/V with an FP32 attention mask

Open
#195,784 2 comments 0 reactions 0 assignees View on GitHub
bot-triaged enhancement 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

`_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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.