pytorch / pytorch/pytorch

`torch.compile` crashes on `stft`/`fftn`/`rfftn`/`ifftn` backward — meta kernel stride mismatch

Open
#182,200 0 comments 0 reactions 1 assignee Claimed by @aorenste View on GitHub
bot-triaged module: aotdispatch module: crash module: fft module: meta tensors module: pt2-dispatcher oncall: pt2 triaged
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.