[inductor] Attention-shaped pattern rewrite ignores the softmax dim, silently computing attention over the last axis
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
An attention-shaped chain whose softmax is taken over a **non-last** dimension is
rewritten by the post-grad pattern matcher as if the softmax were over the last
dimension. The compiled function silently returns the result of a *different*
computation — no error, no warning, no graph break.
```python
import torch, torch.nn as nn
class M(nn.Module):
def __init__(self, sm_dim):
super().__init__()
self.dropout = nn.Dropout(0.23)
self.sm_dim = sm_dim
def forward(self, q, k, v):
qk = torch.matmul(q, k.transpose(-2, -1))
s = qk.div(128)
s = torch.softmax(s, dim=self.sm_dim) # dim=0, NOT the last dim
s = self.dropout(s)
return s.matmul(v)
torch.manual_seed(0)
q, k, v = torch.randn(4,4,2,8), torch.randn(4,4,3,8), torch.randn(4,4,3,8)
def ev(m):
m.train(False)
with torch.no_grad():
return m(q, k, v)
def cp(m):
m.train(False); torch._dynamo.reset()
with torch.no_grad():
return torch.compile(m, backend="inductor")(q, k, v)
rel = lambda a, b: ((a-b).abs().max() / a.abs().max()).item()
print(rel(ev(M(0)), cp(M(0)))) # 0.3462 eager vs compiled, dim=0
print(rel(cp(M(0)), ev(M(-1)))) # 6.5e-08 compiled(dim=0) == eager(dim=-1)
```
### The compiled `dim=0` model computes the `dim=-1` model
The second number is the point: the compiled output of the `dim=0` model is
numerically identical to the **eager output of the `dim=-1` model**, so the
rewrite has silently replaced the user's softmax axis with the last axis.
### It is the pattern matcher
Disabling the pattern matcher removes the divergence entirely, and a model that
really does use `dim=-1` is unaffected:
| configuration | max relative diff (eager vs compiled) |
|---|---|
| `softmax(dim=0)`, `pattern_matcher=True` | **8.72** (shapes `(50,50,4,64)`) / **0.346** (shapes `(4,4,2,8)`) |
| `softmax(dim=0)`, `pattern_matcher=False` | 7.0e-08 |
| `softmax(dim=-1)`, `pattern_matcher=True` | 1.7e-07 |
```python
import torch._inductor.config as cfg
cfg.pattern_matcher = False # divergence disappears
```
`nn.Dropout` is in eval mode throughout (`m.train(False)`), so it is an identity
and the comparison is deterministic — the eager output is bit-stable over 200
repeated runs.
### Note on related reports
This is a different root cause from #195320 (which is about a permuted `K` being
bound as if it were `key.transpose(-2,-1)`). Here `k.transpose(-2, -1)` is the
canonical form and the operands are bound correctly; the only mismatched
attribute is the softmax dimension, which the rewrite does not check.
### Versions
Reproduced on `2.14.0+cpu`, CPU, inductor backend.
cc @ezyang @gchanan @kadeng @msaroufim @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @muchulee8 @amjames @aakhundov @coconutruben @jataylo @drisspg @liangel-02 @howardzhang-cv
Contributor guide
Research direction
Start by reproducing the eager-versus-compiled comparison with cfg.pattern_matcher enabled and disabled. Then inspect the post-grad pattern matcher for attention-shaped chains, focusing on whether it preserves the softmax dimension. Done means non-last dimensions match eager execution while dim=-1 remains unaffected, with regression coverage for both cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100