pytorch - 💡(How to fix) Fix [RFC] XPUGraph Trees [1 participants]

🚀 The feature, motivation and pitch

XPUGraph was introduced in PyTorch 2.11. To accelerate model performance with torch.compile on Intel GPUs, this RFC proposes XPUGraph trees feature, which reduces host overhead in torch.compile by leveraging XPUGraph.

Approaches

We considered two approaches:

1. Generalize XPUGraph Trees and CUDAGraph Trees on the Inductor side.
2. Implement XPUGraph Tree independently, isolating detailed implementation for different hardware backends

Approach 1 — Generalize Graph Trees for XPU, CUDA, and other accelerators

This approach aims to minimize implementation divergence by introducing a unified abstraction layer for Graph Trees. Specifically, we propose a new abstract interface that encapsulates device-specific graph and memory APIs. Each backend (e.g., XPU, CUDA, and other accelerators) implements this interface by providing concrete bindings to its underlying runtime (such as stream management, graph instantiation, replay, and memory handling).

This chart shows GraphInterface design details

<img width="679" height="444" alt="Image" src="https://github.com/user-attachments/assets/f53dad72-66ca-48d3-a407-74fdb41eccf9" />

Under this design, the Graph Trees logic in torch.compile remains backend-agnostic, while device-specific behaviors are isolated behind a well-defined abstraction boundary.

Pros:

• Single code path reduces overall maintenance overhead and simplifies long-term code maintenance.
• Provides a consistent programming and execution model across different accelerators
• Scales well to future backends with minimal changes to core logic

Cons:

• Requires cross-backend alignment for any new feature, increasing coordination overhead
• Potential abstraction overhead that may limit backend-specific optimizations in some cases

Approach 2 —— Implementation for XPU: Almost 1:1 mapping between XPU and CUDA

This approach implements Graph Trees for XPU by closely mirroring the existing CUDA Graph Trees design, with an almost 1:1 mapping at the API and call stack levels. For example, components such as cudagraph_trees and cudagraphify are replicated as xpugraph_trees and xpugraphify, preserving the same control flow, graph capture semantics, and execution patterns.

Under this design, the XPU backend directly reuses the proven CUDA Graph Trees architecture, adapting only the device-specific runtime calls (e.g., stream, graph capture, and replay APIs) to the XPU execution model. This enables rapid development and reduces the need for significant refactoring in the existing torch.compile integration.

Pros:

• Low risk to the existing Graph Trees implementation, as it would not touch CUDAGraph trees code.
• Minimal architectural overhead, avoiding the introduction of new abstraction layers.

Cons:

• Code duplication: maintains two parallel implementations with largely identical logic, increasing maintenance burden.
• Limited scalability: does not generalize well if additional backends require Graph Trees support.

CC @EikanWang @gujinghui