`torch.compile` crashes on `stft`/`fftn`/`rfftn`/`ifftn` backward — meta kernel stride mismatch
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
## 🐛 Describe the bug
`torch.compile` crashes with an `AssertionError` (stride mismatch) in the backward pass for several FFT operations. The forward pass compiles and runs correctly; only backward crashes.
Affected ops and their crash boundaries:
| Op | Crashes when | Minimal crashing input |
|----|-------------|----------------------|
| `torch.stft` | Always (any n_fft) | `(128,)` signal, n_fft=32 |
| `torch.fft.fftn` | Transforming >= 3 dims | `(2, 4, 4, 4)` input |
| `torch.fft.rfftn` | Transforming >= 2 dims | `(4, 8, 8)` input |
| `torch.fft.ifftn` | Transforming >= 3 dims | `(2, 8, 8, 8)` complex input |
Lower-dimensional variants (`fft`, `rfft`, `fft2`, `rfft2`, `irfftn`, `hfft`, `ihfft`) work correctly.
This is an **aot_autograd-level** issue — both `inductor` and `aot_eager` backends crash.
**Related**: #106623, #145977. The ops listed in #145977 (`rfft`, `rfft2`, etc.) appear to be fixed on current nightly, but `stft`/`fftn`/`rfftn`/`ifftn` still crash.
## Minimal reproducer
```python
import torch
# 1. stft backward
x = torch.randn(128, device="cuda", requires_grad=True)
@torch.compile(backend="inductor")
def f_stft(a):
return torch.stft(a, n_fft=32, hop_length=16, return_complex=True).abs().sum()
f_stft(x).backward() # AssertionError: stride mismatch
```
```python
# 2. fftn backward (4D input)
x = torch.randn(2, 4, 4, 4, device="cuda", requires_grad=True)
@torch.compile(backend="inductor")
def f_fftn(a):
return torch.fft.fftn(a).abs().sum()
f_fftn(x).backward() # AssertionError: stride mismatch
```
```python
# 3. rfftn backward (3D input)
x = torch.randn(4, 8, 8, device="cuda", requires_grad=True)
@torch.compile(backend="inductor")
def f_rfftn(a):
return torch.fft.rfftn(a).abs().sum()
f_rfftn(x).backward() # AssertionError: stride mismatch
```
```python
# 4. ifftn backward (4D complex input)
x = torch.randn(2, 8, 8, 8, device="cuda", dtype=torch.cfloat, requires_grad=True)
@torch.compile(backend="inductor")
def f_ifftn(a):
return torch.fft.ifftn(a).abs().sum()
f_ifftn(x).backward() # AssertionError: stride mismatch
```
All of these work correctly in eager mode.
## Root cause analysis
The meta kernels for `_fft_r2c` and `_fft_c2c` (used by these ops) return strides inconsistent with what eager execution produces. During backward, the transposed/mismatched strides cause `assert_size_stride` failures.
For example, `_fft_c2c` on a `(2, 4, 4, 4)` input may produce meta strides like `(64, 1, 16, 4)` instead of the correct `(64, 16, 4, 1)`.
This is the same family as #106623 (FFT meta strides). The fix for `rfft`/`rfft2`/etc. in #145977 appears to have addressed some variants but not these.
## Crash boundary details
Tested systematically across 17 FFT variants:
| Op | Status |
|----|--------|
| `fft`, `rfft`, `ifft`, `irfft` | OK |
| `fft2`, `rfft2`, `ifft2`, `irfft2` | OK |
| `hfft`, `ihfft` | OK |
| `irfftn` | OK |
| `fftn` (2D transform) | OK |
| `ifftn` (2D transform) | OK |
| **`stft`** | **CRASH** |
| **`fftn`** (3D+ transform) | **CRASH** |
| **`rfftn`** (2D+ transform) | **CRASH** |
| **`ifftn`** (3D+ transform) | **CRASH** |
### Versions
## Versions
- **PyTorch**: 2.13.0.dev20260429+cu126
- **Triton**: 3.2.0+git4b3bb1e8
- **CUDA**: 12.6
- **GPU**: Tesla T4 (sm_75)
- **Python**: 3.11
- **Also tested**: CPU backend (same crash)
cc @mruberry @ezyang @eellison @bdhirsh @bobrenjc93 @aorenste @chauhang @penguinwu
Contributor guide
Assessment
This issue has not been assessed yet.