expandable_segments: `empty_cache()` with another device current unmaps segments under queued work (`unmapHandles` synchronizes the null stream of the wrong device) → illegal memory access
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.6k
- PR merge metrics
- PR metrics pending
Description
### 🐛 Describe the bug
With `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, `torch.cuda.empty_cache()` called
while the current device differs from the device whose segment is being released unmaps
expandable-segment memory that still has queued kernels writing to it; the process gets a
sticky `CUDA error: an illegal memory access was encountered`. The only variable in the
reproduction is the current device at the moment of the call. The freed tensor is released
on its own allocation stream, which the allocator documents as needing no `record_stream`
("it already correctly manages the life cycle of tensors on only one stream",
[_tensor_docs.py#L3977-L3982](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/torch/_tensor_docs.py#L3977-L3982)),
so the release path, not the caller, has to protect it.
### Minimal reproduction (two GPUs, one thread, a few seconds)
```python
"""expandable_segments: empty_cache() with the WRONG current device unmaps memory under queued work.
CUDA_VISIBLE_DEVICES=0,1 python repro.py other -> illegal memory access
CUDA_VISIBLE_DEVICES=0,1 python repro.py victim -> control, clean
"""
import os
import sys
if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ: # must be set before CUDA initializes
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch
mode = sys.argv[1] if len(sys.argv) > 1 else "other"
victim = 1
sweep_device = 0 if mode == "other" else victim
torch.cuda.synchronize(0) # create cuda:0's context up front, not inside the sweep
torch.cuda.set_device(victim)
src = torch.ones(48 * 1024 * 1024, device=f"cuda:{victim}", dtype=torch.uint8) # a large-pool segment
for i in range(3):
torch.cuda._sleep(200_000_000) # ~100 ms queued on the victim's DEFAULT stream
x = torch.empty_like(src)
x.copy_(src) # writes into x, queued behind the sleep
del x # freed while its writer is still queued (a same-stream free records no event)
torch.cuda.set_device(sweep_device) # the ONLY variable: the sweeper's current device
torch.cuda.empty_cache() # unmapHandles: cudaStreamSynchronize(nullptr) resolves on THIS device
torch.cuda.set_device(victim)
torch.cuda.synchronize(victim) # the fault surfaces here (sticky error)
print(f"{mode}: iteration {i} clean", flush=True)
```
Observed, `other` (two runs, identical; it faults on the first iteration — no "iteration 0
clean" line is printed; full traceback):
```
Traceback (most recent call last):
File "repro.py", line 29, in
torch.cuda.synchronize(victim) # the fault surfaces here (sticky error)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".../site-packages/torch/cuda/__init__.py", line 1281, in synchronize
return torch._C._cuda_synchronize()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
torch.AcceleratorError: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
The CUDA driver logged these messages, which may provide useful details:
Returning 700 (CUDA_ERROR_ILLEGAL_ADDRESS) from cuCtxSynchronize_v2
```
Observed, `victim` (control):
```
victim: iteration 0 clean
victim: iteration 1 clean
victim: iteration 2 clean
```
The same script with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False` (the classic
`cudaFree` path) in mode `other` is clean for all three iterations: the driver synchronizes
`cudaFree` against the pointer's own device regardless of the caller's current device. Real
programs hit the expandable case through `torch.cuda.graph.__enter__`, which synchronizes
only the current device and then calls the global `torch.cuda.empty_cache()`
([graphs.py#L1276](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/torch/cuda/graphs.py#L1276));
on a device-pinned thread that sweeps every other device's allocator with the wrong current
device. Two details of the repro matter: the tensor is 48 MiB so it occupies a large-pool
segment of its own (a small-pool tensor shares a page with `src` and nothing is unmapped), and
`cuda:0`'s context is created before the loop (creating it inside the first sweep is slow
enough to hide the window).
### Analysis
> **AI-assisted analysis** — contained here per `AI_POLICY.md`; see the disclosure below.
> Line numbers are pinned to `main` @ `bfa6fd479747b1223840f2db5c00d98517b50e9f`.
>
> The safety synchronization in `ExpandableSegment::unmapHandles`
> ([CUDACachingAllocator.cpp#L1008-L1021](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/c10/cuda/CUDACachingAllocator.cpp#L1008-L1021))
> targets the wrong device in this case (excerpt; the trailing comment is an annotation):
>
> ```cpp
> if (stream_) {
> C10_CUDA_CHECK(cudaStreamSynchronize(*stream_)); // annotation: no device guard on this branch
> } else {
> cuda::CUDAGuard device_guard(device_);
> C10_CUDA_CHECK(cudaDeviceSynchronize());
> }
> // ... cuMemUnmap_ / cuMemRelease_ ...
> ```
>
> Three facts compose into the fault:
>
> 1. **`stream_` is an engaged optional holding `nullptr` for default-stream segments.**
> `stream_` is `std::optional`
> ([#L1085](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/c10/cuda/CUDACachingAllocator.cpp#L1085));
> the creation site passes the block's raw `cudaStream_t`
> ([#L3499-L3500](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/c10/cuda/CUDACachingAllocator.cpp#L3499-L3500)),
> so the optional is engaged even when the handle is the null default stream. `if (stream_)`
> is then true and the unguarded branch runs `cudaStreamSynchronize(nullptr)`.
> 2. **The null stream handle resolves against the calling thread's current device.** Unlike a
> real `cudaStream_t`, which carries its device, `nullptr` means the legacy default stream of
> whatever device is current at call time. The only `CUDAGuard` in the function is in the
> `else` branch, which default-stream segments never reach.
> 3. **`empty_cache()` sweeps every device from the calling thread with no device guard on the
> path.** `NativeCachingAllocator::emptyCache` loops over all device allocators
> ([#L4949-L4952](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/c10/cuda/CUDACachingAllocator.cpp#L4949-L4952))
> → `release_cached_blocks` → `release_blocks` → `unmap_block`
> ([#L4303-L4306](https://github.com/pytorch/pytorch/blob/bfa6fd479747b1223840f2db5c00d98517b50e9f/c10/cuda/CUDACachingAllocator.cpp#L4303-L4306))
> → `ExpandableSegment::unmap` → `unmapHandles`; none of them sets the device.
>
> Net effect: a thread whose current device is X, sweeping device Y's allocator, "protects" the
> unmap of Y's default-stream segments by synchronizing X's default stream — a no-op for Y —
> and then `cuMemUnmap`s Y's pages while Y's queued kernels still reference them. The
> per-device allocator mutex is held across the walk, so this is not a free-list race; the
> synchronization itself is aimed at the wrong device.
>
> **Expected behavior / fix.** Hoist the device guard above the branch so the null handle
> resolves against the segment's own device:
>
> ```cpp
> cuda::CUDAGuard device_guard(device_); // covers BOTH branches
> if (stream_) {
> C10_CUDA_CHECK(cudaStreamSynchronize(*stream_));
> } else {
> C10_CUDA_CHECK(cudaDeviceSynchronize());
> }
> ```
>
> Possibly related: #144025 (`empty_cache()` touches `cuda:0` when another device is current).
**Disclosure.** An AI assistant helped me draft the analysis quoted above and the script. I ran
the script myself on the machine in Versions — `other` twice, `victim` once, and the
`expandable_segments:False` control — and checked each cited line against the pinned
permalinks. My own conclusion matches the analysis: the synchronize before `cuMemUnmap` is
aimed at the caller's current device instead of the segment's device, and the one-line guard
hoist is the fix I would expect.
### Versions
`collect_env` could not read the driver version; `nvidia-smi` reports 610.43.02.
```
Collecting environment information...
PyTorch version: 2.14.0+cu130
Is debug build: False
CUDA used to build PyTorch: 13.0
ROCM used to build PyTorch: N/A
OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04.2) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35
Python version: 3.12.13 (main, Apr 7 2026, 20:45:25) [Clang 22.1.1 ] (64-bit runtime)
Python platform: Linux-5.15.0-126-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to:
GPU models and configuration:
GPU 0: NVIDIA H200
GPU 1: NVIDIA H200
GPU 2: NVIDIA H200
GPU 3: NVIDIA H200
GPU 4: NVIDIA H200
GPU 5: NVIDIA H200
GPU 6: NVIDIA H200
GPU 7: NVIDIA H200
Nvidia driver version: Could not collect
cuDNN version: Could not collect
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False
Caching allocator config: N/A
CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 57 bits virtual
Byte Order: Little Endian
CPU(s): 160
On-line CPU(s) list: 0-159
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) Platinum 8460Y+
CPU family: 6
Model: 143
Thread(s) per core: 2
Core(s) per socket: 40
Socket(s): 2
Stepping: 8
Frequency boost: enabled
CPU max MHz: 2001.0000
CPU min MHz: 800.0000
BogoMIPS: 4000.00
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single intel_ppin cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities
Virtualization: VT-x
L1d cache: 3.8 MiB (80 instances)
L1i cache: 2.5 MiB (80 instances)
L2 cache: 160 MiB (80 instances)
L3 cache: 210 MiB (2 instances)
NUMA node(s): 2
NUMA node0 CPU(s): 0-39,80-119
NUMA node1 CPU(s): 40-79,120-159
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected
Versions of relevant libraries:
[pip3] nvidia-cublas==13.1.1.3
[pip3] nvidia-cuda-cupti==13.0.85
[pip3] nvidia-cuda-nvrtc==13.0.88
[pip3] nvidia-cuda-runtime==13.0.96
[pip3] nvidia-cudnn-cu13==9.24.0.43
[pip3] nvidia-cufft==12.0.0.61
[pip3] nvidia-curand==10.4.0.35
[pip3] nvidia-cusolver==12.0.4.66
[pip3] nvidia-cusparse==12.6.3.3
[pip3] nvidia-cusparselt-cu13==0.8.1
[pip3] nvidia-nccl-cu13==2.30.7
[pip3] nvidia-nvjitlink==13.3.33
[pip3] nvidia-nvtx==13.0.85
[pip3] torch==2.14.0+cu130
[pip3] triton==3.8.0
[conda] Could not collect
```
cc @ptrblck @msaroufim @eqy @tinglvv @nWEIdia
Contributor guide
Assessment
This issue has not been assessed yet.