Skip to content

trunk/610343e5acab74781f4ac043633d32465f8a86e5

trunk/610343e5acab74781f4ac043633d32465f8a86e5: [inductor] Skip SDPA fusion when the matched scale is a Tensor (#191496)
@HussainNizamani HussainNizamani tagged this 26 Aug 14:28 
Fixes #191203

## What happens

Writing attention as `q @ k.transpose(-2, -1) / scale` with a 0-d tensor scale works in eager but crashes under torch.compile:

```
RuntimeError: aten::scaled_dot_product_attention() Expected a value of type
'Optional[float]' for argument 'scale' but instead found type 'FakeTensor'.
Cast error details: cannot extract float from tensor with meta storage
```

fuse_attention.py has two families of sfdp patterns. Patterns 1-4, 11-17 and 28-30 go through `_sfdp_extra_check`, which inspects the div/mul node that produces the scale and rejects the match when it is not a float or int. Patterns 5-10 and 18-20 instead capture the scale through the scalar_workaround mechanism as `match.kwargs["inv_scale"]`, and nothing ever type checks it. So the pattern matches, the replacement builds `1.0 / inv_scale` with a tensor in it, and forwards that into scaled_dot_product_attention's `float? scale` argument, which fails during fake tensor tracing.

## The fix

`_sfdp_params_check` is the one check every sfdp pattern runs (the second family uses it directly as its extra_check, and `_sfdp_extra_check` calls it as its last step), so the guard lives there:

```python
if "inv_scale" in match.kwargs and not isinstance(
    match.kwargs["inv_scale"], (float, int)
):
    return False
```

With a real scalar scale nothing changes and the fusion fires as before. With a tensor scale the match is rejected, no rewrite happens, and the model compiles as the unfused graph, which is the same fallback the `_sfdp_extra_check` family already uses for non scalar scales.

## Testing

- New regression test (cpu and gpu registrations): the issue's repro shape with a tensor scale compiles, matches eager, and asserts no fusion fired. Without the fix it fails with the exact error above.
- All 58 CPU tests in test/inductor/test_fused_attention.py pass (1 pre-existing skip). lintrunner clean.
- Patterns 6-10 only have GPU-side positive tests, so a CUDA run of test_sdpa_rewriter_6 to 10 would be a welcome extra check that scalar scale fusion there is unaffected.

P.S. Claude Assisted, all changes reviewed and tested by me.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/191496
Approved by: https://github.com/jansel, https://github.com/eellison
Assets 2
Loading