intel / intel/torch-xpu-ops

[inductor][distributed] Inductor cudagraphs silently no-op on XPU: cudagraph_trees needs XPU allocator checkpoint bindings (`_xpu_getCheckpointState` et al.)

Open
#5,338 0 comments 0 reactions 1 assignee Claimed by @jemitche1 View on GitHub
bug module: distributed
Dominant language
Python
Stars
113
Forks
128
Avg merge
5d 9h
Merged PRs (30d)
112

Description

### 🐛 Describe the bug

`torch.compile` with `torch._inductor.config.triton.cudagraphs = True` silently does
nothing on XPU. The flag *is* read and inductor reaches its device check, which only
clears CUDA, so cudagraphs are skipped — and a test that asserts "graphs work" passes
without ever capturing one.

```bash
python - <<'EOF'
import torch
with torch._inductor.config.patch({"triton.cudagraphs": True}):
f = torch.compile(lambda x: (x + 1).sum(), backend="inductor", fullgraph=True)
inp = torch.randn(2, 8, device="xpu:1")
for _ in range(3):
out = f(inp)
print("out:", out.device)
EOF
```

```
W torch/_inductor/cudagraph_utils.py:401] [__cudagraphs] skipping cudagraphs due to skipping cudagraphs due to multiple devices: device(type='xpu', index=1)
out: xpu:1
```

The "multiple devices" wording is misleading — there is a single XPU device; the check
simply does not accept a non-CUDA device.

There are two blockers, in order:

**1. The gate.** `torch/_inductor/cudagraph_utils.py`,
`check_multiple_devices_or_any_cpu_nodes()` only clears a device whose `.type == "cuda"`:

```python
if (
len(device_node_mapping) == 1
and next(iter(device_node_mapping.keys())).type == "cuda"
):
return None
```

**2. The runtime.** With that check relaxed to accept `xpu`, capture then dies:

```
File "torch/_inductor/cudagraph_trees.py", line 2377, in __init__
with graph_capture_lock, torch.cuda.device(device_index):
File "torch/cuda/__init__.py", line 181, in _exchange_device
raise RuntimeError("PyTorch was compiled without CUDA support")
```

`cudagraph_trees.py` has 64 CUDA call sites across 22 distinct APIs. Most already map
1:1 onto `torch.xpu` — `synchronize`, `device`, `current_stream`, `Stream`, `stream`,
`graph`, `graph_pool_handle`, `current_device`, `memory`, `memory_snapshot`,
`is_initialized` are all present, plus `XPUGraph` where CUDA has `CUDAGraph` — and the
three core pool hooks exist as `torch._C._xpu_beginAllocateCurrentThreadToPool`,
`_xpu_endAllocateToPool`, `_xpu_releasePool`.

What is missing on the XPU side:

| CUDA binding | XPU equivalent | Used at | Needed for |
|---|---|---|---|
| `_cuda_getCheckpointState` | **missing** | `cudagraph_trees.py:1612`, in `_add_first_outputs` | **first recording of every graph** |
| `_cuda_setCheckpointPoolState` | **missing** | `:3257` | restoring allocator state when re-entering a recorded graph |
| `_cuda_cudaCachingAllocator_raw_delete` | **missing** | `:3267` | freeing outputs during that restore |
| `_cuda_checkPoolLiveAllocations` | **missing** | `:2215` | live-allocation asserts |
| `_cuda_isHistoryEnabled` | **missing** (XPU has `_xpu_recordMemoryHistory`) | `:203` | memory-history debug |

Because `_cuda_getCheckpointState` sits on the *first recording* path, this cannot be
worked around in Python: XPU cannot capture even a single graph through
`cudagraph_trees` today.

### Suggested fix

1. **XPU allocator (this repo):** add the missing bindings above — at minimum
`_xpu_getCheckpointState`, `_xpu_setCheckpointPoolState` and an XPU `raw_delete`,
since those three gate capture and replay.
2. **PyTorch inductor (separate PR):** make `cudagraph_utils` accept the current
accelerator and route `cudagraph_trees` through
`torch.get_device_module(device_type)`, aliasing `CUDAGraph`/`XPUGraph`. A
`torch._inductor.config.triton.xpugraphs` switch would let tests select the XPU path
explicitly while `triton.cudagraphs` remains the CUDA one. No such option exists
today — `hasattr(config.triton, "xpugraphs")` is `False`, and patching it raises
`AttributeError: torch._inductor.config.triton.xpugraphs does not exist`.

### Test impact

While porting `test/distributed/tensor/test_compile_on_one_rank.py` to be
device-agnostic (pytorch/pytorch#114850), 29 of its 30 tests were enabled on XPU. The
one exception is
`TestCompileOnOneRankDeviceAsParameter::test_cudagraphs_under_coor_runs_on_nonzero_device`,
which had to be marked `@skipIfXpu(msg="inductor cudagraphs is CUDA-only")`: with the
skip removed it *passes in 2.7s without recording a graph* — green but vacuous, which
is worse than an explicit skip.

Note that `torch.xpu.XPUGraph`, `torch.xpu.graph` and `torch.xpu.make_graphed_callables`
all work in eager mode. The gap is only that inductor never routes to them.

Related: #3570 skips the xccl collective graph-replay UTs (`test_allreduce_graph_replay`
et al.), which is c10d-level graph capture; this issue is about inductor's graph path.

### Versions

```
torch 2.15.0a0+git28789c4 (source build, commit 28789c4f989)
torch.version.xpu 20260100
triton-xpu 3.8.0+git1e2d42a0
device 4x Intel(R) Data Center GPU Max 1100
backend xccl
```

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.