pytorch / pytorch/pytorch

AOTInductor: CUDA illegal memory access compiling multi-level module with dynamic-shape interpolates + torch.cond (torch 2.11 & 2.13)

Open
#195,604 0 comments 0 reactions 1 assignee Claimed by @desertfire View on GitHub
bot-triaged module: aotinductor module: crash module: dynamic shapes module: higher order operators module: inductor module: pt2-dispatcher oncall: export oncall: pt2 release triage triage review
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

`torch._inductor.aoti_compile_and_package` crashes with a device-side CUDA fault (`cudaErrorIllegalAddress`, sometimes `cudaErrorInvalidAddressSpace`) while compiling (benchmarking generated kernels for) a multi-level iterative module exported with dynamic shapes. `torch.export.export` succeeds; eager execution of the module is correct.

The module is a coarse-to-fine Jacobi solver (pymatting-style foreground/background estimation used in production image matting): 13 unrolled pyramid levels; each level does bilinear `F.interpolate` to symbolically-computed sizes (`2 + (dim0 - 2 + d - 1) // d`), elementwise sweep math, and a `torch.cond` (tensor predicate) choosing the per-level sweep count at runtime.

Deterministic across two releases:

| torch | environment | result |
|---|---|---|
| 2.7 (NGC 25.02) | H100 | inductor codegen fails earlier: `NameError('zuf0 is not defined')` in a generated Triton kernel — the symfloat scale argument of a symbolic-size bilinear interpolate is not plumbed into the kernel |
| 2.11.0a0+a6c236b9fd (NGC 26.03) | H100 | export OK, AOTI compile → `cudaErrorIllegalAddress` during `codegen_and_compile` |
| 2.13.0a0+9186a08b2c (NGC 26.07) | H100 | same (also seen as `cudaErrorInvalidAddressSpace`) |

Additional observations:

- `CUDA_LAUNCH_BLOCKING=1` does not localize the fault: the context is poisoned and the error surfaces at `preserve_rng_state` exit inside `codegen_and_compile`.
- `aot_inductor.dump_aoti_minifier=True` cannot minify: the first fault kills the CUDA context, so the minifier's subsequent re-runs all fail (no repro dumped). Subprocess isolation for the minifier would have helped here.
- Possibly related (same failure family, different trigger): #174608.

Repro (self-contained):

```python
"""Standalone repro: AOTInductor emits a CUDA illegal-memory-access while compiling
(benchmarking generated kernels for) a multi-level iterative module with dynamic shapes.

The module is a coarse-to-fine Jacobi solver (pymatting-style foreground estimation):
13 unrolled pyramid levels; per level, bilinear resizes to symbolically-computed sizes,
elementwise sweep math, and a torch.cond choosing the sweep count at runtime.

Observed:
- torch 2.11 (NGC 26.03) and torch 2.13 (NGC 26.07), H100: torch.export succeeds,
aoti_compile_and_package dies with cudaErrorIllegalAddress during inductor's
codegen_and_compile (kernel benchmarking). CUDA_LAUNCH_BLOCKING=1 does not localize
it (context is poisoned; error surfaces at preserve_rng_state exit).
- aot_inductor.dump_aoti_minifier cannot minify: the IMA kills the CUDA context.
- torch 2.7 (NGC 25.02): fails earlier in codegen with NameError('zuf0 is not defined')
(symfloat kernel argument from the symbolic-size bilinear interpolate).

Run: python aoti_ima_repro.py
"""
import torch

class MultiLevelSolver(torch.nn.Module):
N_LEVELS = 13

def __init__(self, n_small_iteration=10, n_big_iteration=4, small_level_size=32):
super().__init__()
self.regularisation = torch.nn.Parameter(torch.tensor(1e-5))
self.n_small_iteration = n_small_iteration
self.n_big_iteration = n_big_iteration
self.small_level_size = small_level_size

def neighbors(self, x):
padded = torch.nn.functional.pad(x, [1, 1, 1, 1], mode="reflect")
return padded[:, :, 1:-1, 0:-2], padded[:, :, 1:-1, 2:], padded[:, :, 0:-2, 1:-1], padded[:, :, 2:, 1:-1]

def _run_sweeps(self, n_iteration, FB, aximage, w_l, w_r, w_t, w_b, b00, b01, b11):
for _ in range(n_iteration):
fb_l, fb_r, fb_t, fb_b = self.neighbors(FB)
unk = w_l * fb_l + w_r * fb_r + w_t * fb_t + w_b * fb_b + aximage
unk_f, unk_b = unk[0:1], unk[1:2]
FB = torch.clip(torch.cat([b00 * unk_f + b01 * unk_b, b01 * unk_f + b11 * unk_b]), 0, 1)
return FB

def forward(self, image, mask):
_, _, h0, w0 = image.shape
f = torch.nn.functional.interpolate(image, size=(1, 1), mode="bilinear")
FB = torch.cat([f, f], dim=0)

def small_branch(FB, aximage, w_l, w_r, w_t, w_b, b00, b01, b11):
return self._run_sweeps(self.n_small_iteration, FB, aximage, w_l, w_r, w_t, w_b, b00, b01, b11)

def big_branch(FB, aximage, w_l, w_r, w_t, w_b, b00, b01, b11):
return self._run_sweeps(self.n_big_iteration, FB, aximage, w_l, w_r, w_t, w_b, b00, b01, b11)

for k in range(1, self.N_LEVELS + 1):
divisor = 2 ** (self.N_LEVELS - k)
# branch-free >=2 clamp, symint-friendly
h = 2 + (h0 - 2 + divisor - 1) // divisor
w = 2 + (w0 - 2 + divisor - 1) // divisor

image_level = torch.nn.functional.interpolate(image, (h, w), mode="bilinear")
mask_level = torch.nn.functional.interpolate(mask, (h, w), mode="bilinear")
FB = torch.nn.functional.interpolate(FB, (h, w), mode="bilinear")

a0 = mask_level
a1 = 1.0 - a0
aximage = torch.cat((a0, a1), 0) * image_level
m_l, m_r, m_t, m_b = self.neighbors(mask_level)
w_l = self.regularisation + torch.abs(a0 - m_l)
w_r = self.regularisation + torch.abs(a0 - m_r)
w_t = self.regularisation + torch.abs(a0 - m_t)
w_b = self.regularisation + torch.abs(a0 - m_b)

gradient_sum = w_l + w_r + w_t + w_b
a00 = a0 * a0 + gradient_sum
a11 = a1 * a1 + gradient_sum
a01 = a0 * a1
inv_det = 1.0 / (a00 * a11 - a01 * a01)
b00, b01, b11 = inv_det * a11, inv_det * -a01, inv_det * a00

is_small = (torch.full((), h, dtype=torch.int64, device=image.device) <= self.small_level_size) & (
torch.full((), w, dtype=torch.int64, device=image.device) <= self.small_level_size
)
FB = torch.cond(is_small, small_branch, big_branch, (FB, aximage, w_l, w_r, w_t, w_b, b00, b01, b11))

return FB.chunk(2)

def main():
print("torch", torch.__version__)
model = MultiLevelSolver().cuda().eval()
sample = (torch.rand(1, 3, 3000, 3000).cuda(), torch.rand(1, 1, 3000, 3000).cuda())

with torch.no_grad():
out = model(*sample) # eager works
print("eager OK:", tuple(out[0].shape))

h = torch.export.Dim("h", min=2, max=6000)
w = torch.export.Dim("w", min=2, max=6000)
spatial = {2: h, 3: w}
with torch.no_grad():
ep = torch.export.export(model, args=sample, dynamic_shapes=(spatial, spatial), strict=False)
print("export OK")

torch._inductor.aoti_compile_and_package(ep, package_path="/tmp/multilevel_solver.pt2")
print("AOTI compile OK") # never reached on 2.11 / 2.13: cudaErrorIllegalAddress

if __name__ == "__main__":
main()
```

### Versions

Observed inside NVIDIA NGC containers on H100 80GB (driver 535 family):
- NGC 26.03 -> torch 2.11.0a0+a6c236b9fd.nv26.03, CUDA 13
- NGC 26.07 -> torch 2.13.0a0+9186a08b2c.nv26.07
- NGC 25.02 -> torch 2.7 (the earlier `zuf0` codegen NameError)

(We have not yet reproduced on a stock PyPI wheel — happy to run any variant that helps triage.)

cc @chauhang @penguinwu @ezyang @bobrenjc93 @aditvenk @laithsakka @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @ydwu4 @avikchaudhuri @zhxchen17 @tugsbayasgalan @angelayi @bdhirsh @aorenste @desertfire @yushangdi @iupaikov-amd

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.