Ascend for Pytorch基于graphImplInterface的图捕获实现方案（非cuda迁移变体）

sabijun

35人浏览 · 2026-09-16 10:52:00

sabijun · 2026-09-16 10:52:00 发布

作者：昇腾实战派
知识地图：https://blog.csdn.net/Lumos_Lovegood/article/details/161601003

背景概述

在 PyTorch 的 torch.compile(mode="reduce-overhead") 优化路径中，图捕获（Graph Capture）是实现性能提升的关键环节。当前 PyTorch 社区对图捕获的实现主要围绕 CUDA 设备展开，但随着 NPU 等其他加速器的接入需求日益增长，我们需要设计一套可扩展的图捕获架构，使得不同硬件后端能够以统一的方式接入这一优化路径。

本文介绍的是“使用 GraphImplInterface 契约、但不迁移 CUDA”的实现变体。该方案的核心思路是：NPU 及其他已注册后端通过 at::accelerator::Graph / NPUGraphImpl 走新的契约路径，而 CUDA 保留原有的 torch.cuda.CUDAGraph 老路径，通过 if device.type=="cuda" 分叉实现共存。这种设计在保证 CUDA 零回归的前提下，为 NPU 等后端提供了接入图捕获的能力。

架构设计

四层架构概览

整个图捕获系统分为四个层次，CUDA 不进入 L1 注册表：

L4  Inductor 消费层
    compile_fx.cudagraphify(device=...)
    cudagraph_trees.py: Manager/Node 经 DeviceInterface 取运行时
    图对象构造: if device.type=="cuda": torch.cuda.CUDAGraph()  ← 老路分叉
                else:                    torch.accelerator.Graph()
    cudagraph_utils.py 资格门: CUDA 特判 True; 其余查能力位

L3  Dynamo 设备契约层（Python 扩展点）
    DeviceInterface.is_graph_capture_supported()
      CUDA: 硬编码 True（不查注册表）
      NPU : torch._C._accelerator_hasGraphImpl("npu")
    DeviceInterface.GraphOps: pool 路由 / checkpoint / prepare ...
      CUDA 转发 torch._C._cuda_*（零变化）；NPU 转发 torch_npu._C._npu_*

L2  Python 门面层
    torch.accelerator.Graph（已有）
    torch._C._accelerator_hasGraphImpl（新增绑定）

L1  C++ 契约层
    GraphImplInterface / GraphImplRegistry（已有）
    XPUGraphImpl（已有）    NPUGraphImpl（torch_npu 侧新增）
    ❌ 无 CUDAGraphImpl —— CUDA 不在注册表

与其它方案的对比

维度	本变体（GraphImplInterface + 不迁 CUDA）	完整方案（GraphImplInterface + 迁 CUDA）	最小方案（鸭子类型 GraphOps）
图对象层（非 CUDA）	`torch.accelerator.Graph` → `NPUGraphImpl`（需 C++）	`torch.accelerator.Graph` → `NPUGraphImpl`（需 C++）	`di.GraphOps.create_graph()` → 复用 `torch_npu.npu.NPUGraph`（零 C++）
图对象层（CUDA）	`torch.cuda.CUDAGraph` 老路径（if 分叉）	`torch.accelerator.Graph` → `CUDAGraphImpl`	`di.GraphOps.create_graph()` → 转发 `torch.cuda.CUDAGraph`
CUDA 是否迁契约	否（不写 `CUDAGraphImpl`）	是	否
核心消费路径	图对象层 `if cuda` 分叉；其余统一 `di`/`GraphOps`	全统一 `accelerator.Graph`	全统一 `di.GraphOps`
NPU 是否需 `NPUGraphImpl`	需要	需要	不需要

整体调用链

torch.compile(mode="reduce-overhead")
│
└─ compile_fx_inner
   │
   ├─ check_multiple_devices_or_any_cpu_nodes
   │  单设备 d 合格 ⇔ (d.type=="cuda") or get_interface_for_device(d.type).is_graph_capture_supported()
   │
   ├─ check_caching_allocator_for_cudagraphs
   │  → di.GraphOps.caching_allocator_enabled()
   │
   └─ cudagraphify（签名加 device）
      │
      └─ cudagraph_trees.cudagraphify
         │
         └─ get_container(device).get_tree_manager()
            │
            └─ CUDAGraphTreeManager（= GraphTreeManager 别名）
               │  持有 di = get_interface_for_device(device.type)
               │
               ├─ 图对象构造 + 捕获（★唯一 CUDA 分叉点★）
               │  if device.type == "cuda":
               │      graph = torch.cuda.CUDAGraph()
               │      with torch.cuda.graph(graph, pool, stream, "thread_local"): ...
               │  else:
               │      graph = torch.accelerator.Graph(pool, "thread_local")
               │      with di.stream(stream), _graph_capture(graph): ...
               │
               ├─ 流/同步/设备上下文（所有后端统一）
               │  di.synchronize / di.Stream / di.current_stream / di.stream / di.device
               │
               └─ 内存池路由 + checkpoint + 捕获前清理（所有后端统一）
                  di.GraphOps.begin_allocate_current_thread_to_pool / end / release
                  di.GraphOps.get_checkpoint_state / set_checkpoint_pool_state /
                              check_pool_live_allocations / raw_delete
                  di.GraphOps.prepare_for_capture / graph_pool_handle / memory_snapshot

关键改动详解

PR-2：能力查询与 CUDA 特判

文件：torch/csrc/DeviceAccelerator.cpp

新增 _accelerator_hasGraphImpl 绑定：

m.def("_accelerator_hasGraphImpl", [](std::optional<std::string> device_type) {
  const auto dt = device_type.has_value()
      ? c10::Device(device_type.value()).type()
      : at::accelerator::getAccelerator(/*check=*/true);
  return at::has_graph_impl(dt);
}, py::arg("device_type") = std::nullopt);

文件：torch/_dynamo/device_interface.py

CUDA 不能查注册表（未注册），必须硬编码 True：

class DeviceInterface:
    @classmethod
    def is_graph_capture_supported(cls) -> bool:
        return False

class CudaInterface(DeviceInterface):
    @classmethod
    def is_graph_capture_supported(cls) -> bool:
        return True     # ★特判★ CUDA 走老路支持图捕获，但未入注册表

class XpuInterface(DeviceInterface):
    @classmethod
    def is_graph_capture_supported(cls) -> bool:
        return torch._C._accelerator_hasGraphImpl(cls.device.type)

PR-3：资格门设备无关化

文件：torch/_inductor/cudagraph_utils.py

# check_multiple_devices_or_any_cpu_nodes
- if (len(device_node_mapping) == 1
-         and next(iter(device_node_mapping.keys())).type == "cuda"):
-     return None
+ if len(device_node_mapping) == 1:
+     device = next(iter(device_node_mapping.keys()))
+     from torch._dynamo.device_interface import get_interface_for_device
+     if get_interface_for_device(device.type).is_graph_capture_supported():
+         return None

# check_caching_allocator_for_cudagraphs
- if torch.cuda.is_available() and not torch._C._cuda_cudaCachingAllocator_is_enabled():
+ di = get_interface_for_device(torch.accelerator.current_accelerator().type)
+ if di.is_available() and not di.GraphOps.caching_allocator_enabled():

PR-4：DeviceInterface.GraphOps 扩展点

文件：torch/_dynamo/device_interface.py

class CudaInterface(DeviceInterface):
    class GraphOps(DeviceInterface.GraphOps):
        graph_pool_handle = staticmethod(torch.cuda.graph_pool_handle)
        begin_allocate_current_thread_to_pool = staticmethod(torch._C._cuda_beginAllocateCurrentThreadToPool)
        end_allocate_to_pool = staticmethod(torch._C._cuda_endAllocateToPool)
        release_pool = staticmethod(torch._C._cuda_releasePool)
        get_checkpoint_state = staticmethod(torch._C._cuda_getCheckpointState)
        set_checkpoint_pool_state = staticmethod(torch._C._cuda_setCheckpointPoolState)
        check_pool_live_allocations = staticmethod(torch._C._cuda_checkPoolLiveAllocations)
        raw_delete = staticmethod(torch._C._cuda_cudaCachingAllocator_raw_delete)
        caching_allocator_enabled = staticmethod(torch._C._cuda_cudaCachingAllocator_is_enabled)
        @staticmethod
        def prepare_for_capture():
            ...  # 搬移 cudagraph_trees.py 的 clear_cublas_manager + disable_conv_cache_emptying

PR-5a：机械替换（流/同步/设备/allocator/checkpoint）

文件：torch/_inductor/cudagraph_trees.py

现调用	替换为
`torch.cuda.synchronize()`	`di.synchronize()`
`torch.cuda.Stream()` / `current_stream()`	`di.Stream()` / `di.current_stream()`
`torch.cuda.stream(...)` / `device(idx)`	`di.stream(...)` / `di.device(idx)`
`_cuda_beginAllocateCurrentThreadToPool` 等	`di.GraphOps.*`
`_cuda_getCheckpointState` 等	`di.GraphOps.*`
`clear_cublas_manager()+disable_conv_cache_emptying()`	`di.GraphOps.prepare_for_capture()`
`torch.cuda.memory_snapshot()`	`di.GraphOps.memory_snapshot()`

PR-5b：图对象构造——CUDA 分叉

文件：torch/_inductor/cudagraph_trees.py

if device.type == "cuda":
    # CUDA 老路径，一行不改
    self.graph = torch.cuda.CUDAGraph()
    self.cuda_graphs_thread_pool = torch.cuda.graph_pool_handle()
    capture_cm = torch.cuda.graph(self.graph, pool=..., stream=self.stream,
                                  capture_error_mode="thread_local")
else:
    # 注册后端走 accelerator.Graph
    self.graph = torch.accelerator.Graph(pool=..., capture_error_mode="thread_local")
    self.cuda_graphs_thread_pool = di.GraphOps.graph_pool_handle()
    capture_cm = _graph_capture(self.graph)

with di.stream(self.stream), capture_cm:
    static_outputs = model(inputs)

NPU 侧接入方案

NPU 侧需要完成以下三步：

1. C++ NPUGraphImpl 实现

struct NPUGraphImpl final : public at::GraphImplInterface {
  explicit NPUGraphImpl(const at::GraphImplArgs& args = {}) {}
  void capture_begin(MempoolId_t pool, at::GraphCaptureMode mode) override {
    graph_.capture_begin(pool, toAclCaptureMode(mode));
  }
  void capture_end() override { graph_.capture_end(); }
  void instantiate()  override { /* CANN 在 capture_end 内已实例化 → no-op */ }
  void replay() override { graph_.replay(); }
  void reset()  override { graph_.reset(); }
  MempoolId_t pool() const override { return graph_.pool(); }
  void enable_debug_mode() override { graph_.enable_debug_mode(); }
  void debug_dump(const std::string& p) override { graph_.debug_dump(p); }
 private:
  c10_npu::NPUGraph graph_;
};
REGISTER_GRAPH_IMPL(NPU, NPUGraphImpl)

2. NpuInterface 实现

class NpuInterface(DeviceInterface):
    @classmethod
    def is_graph_capture_supported(cls):
        return torch._C._accelerator_hasGraphImpl("npu")
    class GraphOps(DeviceInterface.GraphOps):
        graph_pool_handle = staticmethod(torch_npu.npu.graphs.graph_pool_handle)
        begin_allocate_current_thread_to_pool = staticmethod(torch_npu._C._npu_beginAllocateCurrentThreadToPool)
        end_allocate_to_pool = staticmethod(torch_npu._C._npu_endAllocateToPool)
        release_pool = staticmethod(torch_npu._C._npu_releasePool)
        get_checkpoint_state = staticmethod(torch_npu._C._npu_getCheckpointState)
        set_checkpoint_pool_state = staticmethod(torch_npu._C._npu_setCheckpointPoolState)

3. 适配层注意事项

NPUGraphImpl 需要弥合三处签名差异：

capture_begin 的 mode 参数映射
capture_end 后 CANN 自动实例化，instantiate 为 no-op
keep_graph 参数无语义，设为 true 时 WARN

CUDA 零回归验证

逻辑	现（直绑 CUDA）	本变体
图对象构造/捕获/回放	`torch.cuda.CUDAGraph` + `torch.cuda.graph(...)`	完全不变（`if cuda` 分支保原路）
内存池路由 + checkpoint	`torch._C._cuda_*`	`di.GraphOps.*`（一行转发原符号）
流/同步/设备	`torch.cuda.*`	`di.`（CUDA 实现即 `torch.cuda.`）
资格判断	`== "cuda"` 字面量	`is_graph_capture_supported()`，CUDA 仍硬编码 True

不迁移 CUDA 的代价

重新引入 == "cuda" 特判：核心路径出现 if device.type=="cuda" 分叉，与社区"拆掉 ==cuda"的方向相悖
能力位契约不对称：CUDA 硬编码 True、NPU 查注册表
else 分支失去 in-tree 共验后端：NPU 成为新消费路径的唯一验证者
instantiate/keep_graph 等语义边界只由 NPU no-op 路径探

PR 拆分与合入判据

PR	内容	合入门槛
PR-2	`_accelerator_hasGraphImpl` + `is_graph_capture_supported`（CUDA 硬编码 True）	能力查询贯通；CUDA 返回 True
PR-3	资格门 + allocator 检查 + ir.py（CUDA 特判保留）	假后端能力位 True 时放行
PR-4	`DeviceInterface.GraphOps` + CUDA/XPU 转发	逐符号 `assertIs`
PR-5a	机械替换（di.* / GraphOps.*，含 CUDA）	`test_cudagraph_trees.py` 零改动通过
PR-5b	图对象 `if cuda: 老路 else: accelerator.Graph`	非 CUDA 分支端到端（靠 NPU/XPU）
PR-6	设备泛型测试	NPU/XPU 单独验 else 分支

总结

本变体通过 GraphImplInterface 契约为 NPU 等后端提供了接入图捕获的标准化路径，同时通过 if cuda 分叉保留了 CUDA 的原有实现，实现了"零回归"的目标。虽然引入了 ==cuda 特判等代价，但相比完整方案省去了 CUDA 迁移的前置工作，适合作为渐进式改造的中间步骤。NPU 侧需要完成 NPUGraphImpl 的 C++ 实现和 NpuInterface 的 Python 扩展，即可接入 torch.compile(mode="reduce-overhead") 优化路径。# 背景概述

在 PyTorch 的 Inductor 图捕获（cudagraph）机制中，原有的实现高度耦合于 CUDA 设备，大量代码直接使用 torch.cuda.* 和 torch._C._cuda_* 符号。随着 XPU、NPU 等其他加速器后端逐步接入 PyTorch 生态，这种硬编码的 CUDA 依赖成为设备无关化的主要障碍。

本文旨在描述如何通过引入设备接口抽象层（DeviceInterface）和 GraphOps 契约，将图捕获相关的资格判断、内存池管理、checkpoint 操作等核心逻辑从 CUDA 硬编码改造为设备无关的通用实现。改造的核心思路是：CUDA 保留原有行为不变（零回归），其他后端通过注册机制获得图捕获能力。

1. `torch/_dynamo/device_interface.py`：设备接口抽象层

1.1 `DeviceInterface` 基类：能力位与 `GraphOps` 契约

目的：定义“设备是否支持 Inductor 图捕获”的共享能力位（供资格门与图对象分发消费），以及 cudagraph_trees 所需的 allocator/checkpoint 运行时扩展点 GraphOps（契约层）。默认实现使未适配的后端安全地“不支持图捕获”。default_generators 供 ir.py 的 GeneratorState 设备无关化。

1.2 `CudaInterface`：能力位硬编码 True + `GraphOps` 转发 + `default_generators`

修改前（raise_if_triton_unavailable 之后直接是 get_mtia_stream）

        elif "nvidia" not in triton.backends.backends:
            raise RuntimeError("triton not built with the 'nvidia' backend")

get_mtia_stream: Callable[[int], int] | None

修改后

        elif "nvidia" not in triton.backends.backends:
            raise RuntimeError("triton not built with the 'nvidia' backend")

@classmethod
    def is_graph_capture_supported(cls) -> bool:
        # 本变体不迁移 CUDA：CUDA 保留 torch.cuda.CUDAGraph 老路、未入 C++ 注册表，
        # 不能由 _accelerator_hasGraphImpl 推导 → 硬编码 True 保持 CUDA 既有资格。
        return True

@staticmethod
    def default_generators() -> Any:
        return torch.cuda.default_generators

class GraphOps(DeviceInterface.GraphOps):
        graph_pool_handle = staticmethod(torch.cuda.graph_pool_handle)
        # torch._C._cuda_* 仅 CUDA 构建存在 → 运行时包装懒绑定（对齐 get_cuda_stream 条件导入模式）
        @staticmethod
        def begin_allocate_current_thread_to_pool(device, pool) -> None:
            torch._C._cuda_beginAllocateCurrentThreadToPool(device, pool)
        @staticmethod
        def end_allocate_to_pool(device, pool) -> None:
            torch._C._cuda_endAllocateToPool(device, pool)
        @staticmethod
        def release_pool(device, pool) -> None:
            torch._C._cuda_releasePool(device, pool)
        @staticmethod
        def get_checkpoint_state(device, pool) -> Any:
            return torch._C._cuda_getCheckpointState(device, pool)
        @staticmethod
        def set_checkpoint_pool_state(device, state, stale_storages, storages_to_add_deleters_to) -> None:
            torch._C._cuda_setCheckpointPoolState(device, state, stale_storages, storages_to_add_deleters_to)
        @staticmethod
        def check_pool_live_allocations(device, pool, expected_live_allocations) -> bool:
            return torch._C._cuda_checkPoolLiveAllocations(device, pool, expected_live_allocations)
        @staticmethod
        def raw_delete(ptr) -> None:
            torch._C._cuda_cudaCachingAllocator_raw_delete(ptr)
        @staticmethod
        def caching_allocator_enabled() -> bool:
            return torch._C._cuda_cudaCachingAllocator_is_enabled()
        @staticmethod
        @contextlib.contextmanager
        def prepare_for_capture() -> Any:
            # 复刻 cudagraph_trees clear_cublas_manager + disable_conv_cache_emptying（同符号，零行为变化）
            torch._C._cuda_clearCublasWorkspaces()
            prev_conv = torch._C._cuda_get_conv_benchmark_empty_cache()
            torch._C._cudnn_set_conv_benchmark_empty_cache(False)
            try:
                yield
            finally:
                torch._C._cudnn_set_conv_benchmark_empty_cache(prev_conv)
                torch._C._cuda_clearCublasWorkspaces()
        memory_snapshot = staticmethod(torch.cuda.memory_snapshot)

get_mtia_stream: Callable[[int], int] | None

目的：给 CUDA 提供图捕获能力位与 GraphOps 内置实现。

能力位硬编码 True：这是“不迁移 CUDA”变体的关键——CUDA 不在 GraphImplRegistry，_accelerator_hasGraphImpl("cuda") 会返回 False，故不能查注册表，必须硬编码以保持 CUDA 既有资格。
GraphOps 一行转发原 torch._C._cuda_* 符号：保证消费层统一走 di.GraphOps 后，CUDA 调用的底层 C 符号不变（行为零变化）。用运行时包装而非 staticmethod(torch._C._cuda_...)，避免 CPU-only 构建在类定义期因符号缺失而 AttributeError。
prepare_for_capture 内联 cublas/conv 清理：避免 _dynamo 反向依赖 _inductor，且使用与原 cudagraph_trees 相同的符号。

1.3 `XpuInterface`：能力位查注册表

修改前（raise_if_triton_unavailable 之后直接是 CpuDeviceProperties）

        if "intel" not in triton.backends.backends:
            raise RuntimeError("triton not built with the 'intel' backend")

@dataclass
class CpuDeviceProperties:

修改后

        if "intel" not in triton.backends.backends:
            raise RuntimeError("triton not built with the 'intel' backend")

@classmethod
    def is_graph_capture_supported(cls) -> bool:
        # XPU 已迁入 GraphImplInterface 契约（XPUGraphImpl 已注册），注册表是唯一事实源。
        return torch._C._accelerator_hasGraphImpl(cls.device.type)

@dataclass
class CpuDeviceProperties:

目的：XPU 已按契约注册 XPUGraphImpl，其能力位以 C++ 注册表为准。与 CUDA 的硬编码形成对照，体现“注册后端查注册表、未迁 CUDA 特判”的不对称契约。

注：本变体未给 XpuInterface 补 GraphOps（XPU 的 Inductor 图树启用属延后项，其 torch._C._xpu_* checkpoint 符号未必齐全）；且该方法依赖 C++ 绑定，需重编译后才可调用。

2. `torch/_inductor/cudagraph_utils.py`

2.1 `check_multiple_devices_or_any_cpu_nodes`（资格门）

修改前（:361）

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

修改后

    if len(device_node_mapping) == 1:
        device = next(iter(device_node_mapping.keys()))
        from torch._dynamo.device_interface import get_interface_for_device

        # 设备无关资格门：单设备图合格 iff 该设备支持图捕获。
        # CUDA 经 CudaInterface 硬编码 is_graph_capture_supported 仍合格；
        # 注册后端（XPU/NPU）查注册表。
        if get_interface_for_device(device.type).is_graph_capture_supported():
            return None

目的：把资格判断从 == "cuda" 字面量改为查能力位，让 NPU（及任何支持图捕获的后端）能通过单设备资格门。CUDA 因硬编码 True 而行为不变。

2.2 `check_caching_allocator_for_cudagraphs`

修改前（:371）

    if (
        torch.cuda.is_available()
        # pyrefly: ignore [missing-attribute]
        and not torch._C._cuda_cudaCachingAllocator_is_enabled()
    ):
        return format_default_skip_message(
            "cudagraph capture requires the caching allocator; "
            "current allocator is uncached"
        )
    return None

修改后

    current_accelerator = torch.accelerator.current_accelerator()
    if current_accelerator is None:
        return None
    from torch._dynamo.device_interface import get_interface_for_device

    di = get_interface_for_device(current_accelerator.type)
    if di.is_available() and not di.GraphOps.caching_allocator_enabled():
        return format_default_skip_message(
            "cudagraph capture requires the caching allocator; "
            "current allocator is uncached"
        )
    return None

目的：把“caching allocator 是否开启”的检查从硬编码 CUDA 改为按当前 accelerator 经 GraphOps.caching_allocator_enabled() 分发。NPU 用自己的 allocator 状态判定，CUDA 经转发保持原判定。

3. `torch/_inductor/ir.py`（GeneratorState 断言设备无关化）

两处结构相同（args_flat 分支 :7172、example_inputs 分支 :7242）。

3.1 `:7172` 附近

修改前

                    device_index = arg.device.index
                    if not (arg.device.type == "cuda" and device_index is not None):
                        raise AssertionError(
                            'Expected arg.device.type == "cuda" and device_index is not None'
                        )
                    real_non_tensor_args.append(
                        torch.cuda.default_generators[device_index].clone_state()
                    )

修改后

                    device_index = arg.device.index
                    if device_index is None:
                        raise AssertionError(
                            "Expected GeneratorState device_index is not None"
                        )
                    from torch._dynamo.device_interface import (
                        get_interface_for_device,
                    )

                    di = get_interface_for_device(arg.device.type)
                    real_non_tensor_args.append(
                        di.default_generators()[device_index].clone_state()
                    )

3.2 `:7242` 附近

修改前

                device_index = x.device.index
                if not (x.device.type == "cuda" and device_index is not None):
                    raise AssertionError(
                        'Expected x.device.type == "cuda" and device_index is not None'
                    )
                example_args.append(
                    torch.cuda.default_generators[device_index].clone_state()
                )

修改后

                device_index = x.device.index
                if device_index is None:
                    raise AssertionError(
                        "Expected GeneratorState device_index is not None"
                    )
                from torch._dynamo.device_interface import get_interface_for_device

                di = get_interface_for_device(x.device.type)
                example_args.append(
                    di.default_generators()[device_index].clone_state()
                )

目的：解除 GeneratorState（含 RNG 生成器占位符的 AOTI 场景）对 CUDA 的硬断言，改为经 di.default_generators() 取设备默认生成器。保留“device_index 不为 None”的真实前置断言。CUDA 经 CudaInterface.default_generators() 返回 torch.cuda.default_generators，行为不变；NPU 提供自己的实现即可走通。

4. `torch/csrc/DeviceAccelerator.cpp`（新增能力查询绑定，需重编译）

4.1 新增头文件

修改前

#include <c10/core/AllocatorConfig.h>
#include <torch/csrc/DeviceAccelerator.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/utils/device_lazy_init.h>

修改后

#include <ATen/core/GraphImplInterface.h>
#include <c10/core/AllocatorConfig.h>
#include <c10/core/Device.h>
#include <torch/csrc/DeviceAccelerator.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/utils/device_lazy_init.h>

目的：引入 at::has_graph_impl（GraphImplInterface.h）与 c10::Device(std::string) 构造（Device.h）。

4.2 新增 `_accelerator_hasGraphImpl` 绑定

修改前（_accelerator_getDefaultGenerator 之后直接是 Accelerator Graph class binding）

  m.def("_accelerator_getDefaultGenerator", [](c10::DeviceIndex device_index) {
    const auto device_type = at::accelerator::getAccelerator(true).value();
    torch::utils::maybe_initialize_device(device_type);
    return at::accelerator::getDefaultGenerator(device_index);
  });

// Accelerator Graph class binding

修改后

  m.def("_accelerator_getDefaultGenerator", [](c10::DeviceIndex device_index) {
    const auto device_type = at::accelerator::getAccelerator(true).value();
    torch::utils::maybe_initialize_device(device_type);
    return at::accelerator::getDefaultGenerator(device_index);
  });

// Whether a backend-agnostic graph capture implementation (GraphImplInterface)
  // is registered for the given device type (default: current accelerator).
  // Consumed by DeviceInterface.is_graph_capture_supported for registered backends
  // (XPU/NPU). CUDA is intentionally not registered in the no-CUDA-migration
  // variant, so CudaInterface hardcodes its capability bit instead of using this.
  m.def(
      "_accelerator_hasGraphImpl",
      [](std::optional<std::string> device_type) {
        const auto dt = device_type.has_value()
            ? c10::Device(device_type.value()).type()
            : at::accelerator::getAccelerator(/*check=*/true).value();
        return at::has_graph_impl(dt);
      },
      py::arg("device_type") = std::nullopt);

// Accelerator Graph class binding

目的：暴露 Python 可查询的“某设备类型是否注册了 GraphImplInterface 实现”接口，供 XpuInterface/NpuInterface 的 is_graph_capture_supported 消费。CUDA 因不入注册表而不走此查询（硬编码 True）。

⚠️ 此为 C++ 改动，需 pip install -e . -v --no-build-isolation 重新编译后才生效；未重编译前，XpuInterface/NpuInterface.is_graph_capture_supported 调用会 AttributeError。

5. `torch/_inductor/cudagraph_trees.py`（PR-5a/5b，消费层主体）

5.0 总体方案：从当前 accelerator 解析 `di`

cudagraph_trees 本就是“进程单 accelerator”假设。为避免把 device_type 穿透进所有类/自由函数的签名，新增两个模块级 helper 从当前 accelerator 解析设备接口：

def _graph_device_type() -> str:
    acc = torch.accelerator.current_accelerator()
    return acc.type if acc is not None else "cuda"

def _graph_device_interface(device_type: str | None = None) -> Any:
    from torch._dynamo.device_interface import get_interface_for_device
    return get_interface_for_device(device_type or _graph_device_type())

并新增图捕获上下文分叉 helper（CUDA 走 torch.cuda.graph；注册后端走 torch.accelerator.Graph，其本身是“当前流上开始/结束捕获”的上下文管理器）：

@contextlib.contextmanager
def _graph_capture_context(di, device_type, graph, stream, pool, capture_error_mode="thread_local"):
    if device_type == "cuda":
        with torch.cuda.graph(graph, stream=stream, pool=pool, capture_error_mode=capture_error_mode):
            yield
    else:
        with di.stream(stream), graph:  # accelerator.Graph 携带构造期 pool/mode
            yield

目的：一处解析、全局复用；CUDA 的 di.*/di.GraphOps.* 均转发原符号（零行为变化），仅图对象层 if 分叉。

5.1 三个类：`init` 存 `self.device_interface` / `self.device_type`

CUDAWarmupNode / CUDAGraphNode / CUDAGraphTreeManager 的 __init__ 各加：

self.device_type = _graph_device_type()
self.device_interface = _graph_device_interface(self.device_type)

目的：类方法内统一用 self.device_interface.* 取运行时，用 self.device_type 做 if 分叉。

5.2 机械替换映射（前 → 后）

现调用	替换为	所在
`torch.cuda.synchronize()`	`self.device_interface.synchronize()` / `di.synchronize()`	各类方法 / `_use_cuda_memory_pool_manager` / `check_memory_pool`
`torch.cuda.Stream()` / `current_stream()` / `stream(...)` / `device(...)`	`di.Stream()` / `di.current_stream()` / `di.stream(...)` / `di.device(...)`	同上
`torch.cuda.graph_pool_handle()`	`di.GraphOps.graph_pool_handle()`	manager `__init__`
`torch._C._cuda_beginAllocateCurrentThreadToPool` / `_endAllocateToPool` / `_releasePool`	`di.GraphOps.begin_allocate_current_thread_to_pool` 等	`_use_cuda_memory_pool_manager`
`torch._C._cuda_getCheckpointState`	`self.device_interface.GraphOps.get_checkpoint_state`	`CUDAGraphNode.record`
`torch._C._cuda_setCheckpointPoolState`	`self.device_interface.GraphOps.set_checkpoint_pool_state`	manager
`torch._C._cuda_checkPoolLiveAllocations`	`di.GraphOps.check_pool_live_allocations`	`check_memory_pool`
`torch._C._cuda_cudaCachingAllocator_raw_delete`	`self.device_interface.GraphOps.raw_delete`	manager
`torch.cuda.memory_snapshot()`	`di.GraphOps.memory_snapshot()`（None → 返回空段，跳过慢路径断言）	`get_cudagraph_segments`

5.3 图对象 if 分叉（PR-5b）——两处构造 + 两处捕获

构造（CUDAGraphNode._record、CUDAGraphTreeManager.__init__）：

if self.device_type == "cuda":
    self.graph = torch.cuda.CUDAGraph()
else:
    self.graph = torch.accelerator.Graph(pool=<pool>, capture_error_mode="thread_local")

manager __init__ 中先取 pool 再建图（accelerator.Graph 构造期需要 pool），CUDA 无影响。

捕获：torch.cuda.graph(self.graph, stream=..., pool=..., capture_error_mode=...) → _graph_capture_context(self.device_interface, self.device_type, self.graph, self.stream, <pool>, "thread_local")。

5.4 CUDA 专属清理的处理：守卫而非合并

关键决策：clear_cublas_manager / disable_conv_cache_emptying / enable_history_recording（get_history_recording）未改为 GraphOps.prepare_for_capture，而是在函数顶部加 if _graph_device_type() != "cuda": yield/nullcontext; return 守卫。

目的/原因：warmup 路径用了两个清理（conv + cublas），而 record 路径只用了 cublas 一个；若统一替换为“合并二者”的 prepare_for_capture，会给 record 路径新增 conv 清理 → 改变 CUDA 行为。守卫方案让 CUDA 调用点字节不变、NPU 安全 no-op（NPU 无 cublas/cudnn 清理需求），且这些函数内的 torch._C._cuda_* 永不在非 CUDA 路径执行。

5.5 类型注解放宽

self.graph: torch.cuda.CUDAGraph | None / existing_cuda_graph: torch.cuda.CUDAGraph | None / stream: torch.cuda.Stream → Any | None / Any（TreeManagerContainer、CUDAWarmupNode、CUDAGraphNode）。

6. `torch/_inductor/compile_fx.py`（PR-5a/5b，旧非 trees 路径）

cudagraphify_impl（config.triton.cudagraph_trees=False 时的回退路径）的 warmup + record 段：

修改前

    torch.cuda.synchronize()
    stream = torch.cuda.Stream()
    stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        model(list(static_inputs))
    stream.synchronize()
    torch.cuda.current_stream().wait_stream(stream)
    torch.cuda.synchronize()
    # record
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"):
        static_outputs = model(list(static_inputs))

修改后

    from torch._dynamo.device_interface import get_interface_for_device
    _acc = torch.accelerator.current_accelerator()
    device_type = _acc.type if _acc is not None else "cuda"
    di = get_interface_for_device(device_type)
    # warmup
    di.synchronize()
    stream = di.Stream()
    stream.wait_stream(di.current_stream())
    with di.stream(stream):
        model(list(static_inputs))
    stream.synchronize()
    di.current_stream().wait_stream(stream)
    di.synchronize()
    # record
    if device_type == "cuda":
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"):
            static_outputs = model(list(static_inputs))
    else:
        graph = torch.accelerator.Graph(capture_error_mode="thread_local")
        with di.stream(stream), graph:
            static_outputs = model(list(static_inputs))

目的：让旧回退路径也设备无关。graph.replay() 对两种图对象通用，无需改动。

get_cuda_device_context（compile 期设备上下文）属设备守卫范畴，非 cudagraph 捕获点，未改。

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者，聚合华为+生态”的社区，内容涵盖鲲鹏、昇腾资源，帮助开发者快速获取所需的知识、经验、软件、工具、算力，支撑开发者易学、好用、成功，成为核心开发者。

更多推荐

国产化排队叫号系统能不能兼容飞腾、鲲鹏国产 CPU？

鲲鹏昇腾开发者社区

算子开发高级课程总结（个人学习总结，如有侵犯立即删除）

本文介绍昇腾AI处理器算子调试与优化技术，重点阐述CPU与NPU孪生调试机制：通过AscendC算子源码分别编译为CPU.so（GDB/printf调试）和NPUbin（Profiling/上板分析），实现功能精度与性能的分域调试。核心方法包括：CPU域使用printf、GDB单步调试定位逻辑错误；NPU域借助DumpTensor、msSanitizer（内存/竞争检测）和msProf（指令流水图

鲲鹏昇腾开发者社区

GPU训练迁到昇腾怎么改：torch.cuda 替换、transfer_to_npu 与迁移工具全景

只需导入torch_npu并使用transfer_to_npu，即可实现GPU训练脚本到昇腾NPU的自动迁移。官方提供三种方式——自动、工具、手工迁移，分别对应最小改动、离线转换与手动替换。迁移后务必进行精度对拍，排查时优先确认设备可见性与插件版本。通过昇腾知识图谱可快速检索迁移方案、算子替换清单及常见问题，高效完成代码适配。

鲲鹏昇腾开发者社区

所有评论(0)

查看更多评论

sabijun

@qq_52008485

已为社区贡献2条内容

Ascend for Pytorch基于graphImplInterface的图捕获实现方案（非cuda迁移变体）

sabijun

背景概述

架构设计

四层架构概览

与其它方案的对比

整体调用链

关键改动详解

PR-2：能力查询与 CUDA 特判

PR-3：资格门设备无关化

PR-4：DeviceInterface.GraphOps 扩展点

PR-5a：机械替换（流/同步/设备/allocator/checkpoint）

PR-5b：图对象构造——CUDA 分叉

NPU 侧接入方案

1. C++ NPUGraphImpl 实现

2. NpuInterface 实现

3. 适配层注意事项

CUDA 零回归验证

不迁移 CUDA 的代价

PR 拆分与合入判据

总结

1. torch/_dynamo/device_interface.py：设备接口抽象层

1.1 DeviceInterface 基类：能力位与 GraphOps 契约

1.2 CudaInterface：能力位硬编码 True + GraphOps 转发 + default_generators

1.3 XpuInterface：能力位查注册表

2. torch/_inductor/cudagraph_utils.py

2.1 check_multiple_devices_or_any_cpu_nodes（资格门）

2.2 check_caching_allocator_for_cudagraphs

3. torch/_inductor/ir.py（GeneratorState 断言设备无关化）

3.1 :7172 附近

3.2 :7242 附近

4. torch/csrc/DeviceAccelerator.cpp（新增能力查询绑定，需重编译）

4.1 新增头文件

4.2 新增 _accelerator_hasGraphImpl 绑定

5. torch/_inductor/cudagraph_trees.py（PR-5a/5b，消费层主体）

5.0 总体方案：从当前 accelerator 解析 di

5.1 三个类：__init__ 存 self.device_interface / self.device_type

5.2 机械替换映射（前 → 后）

5.3 图对象 if 分叉（PR-5b）——两处构造 + 两处捕获

5.4 CUDA 专属清理的处理：守卫而非合并

5.5 类型注解放宽

6. torch/_inductor/compile_fx.py（PR-5a/5b，旧非 trees 路径）

所有评论(0)

温馨提示：您尚未绑定手机号

sabijun

1. `torch/_dynamo/device_interface.py`：设备接口抽象层

1.1 `DeviceInterface` 基类：能力位与 `GraphOps` 契约

1.2 `CudaInterface`：能力位硬编码 True + `GraphOps` 转发 + `default_generators`

1.3 `XpuInterface`：能力位查注册表

2. `torch/_inductor/cudagraph_utils.py`

2.1 `check_multiple_devices_or_any_cpu_nodes`（资格门）

2.2 `check_caching_allocator_for_cudagraphs`

3. `torch/_inductor/ir.py`（GeneratorState 断言设备无关化）

3.1 `:7172` 附近

3.2 `:7242` 附近

4. `torch/csrc/DeviceAccelerator.cpp`（新增能力查询绑定，需重编译）

4.2 新增 `_accelerator_hasGraphImpl` 绑定

5. `torch/_inductor/cudagraph_trees.py`（PR-5a/5b，消费层主体）

5.0 总体方案：从当前 accelerator 解析 `di`

5.1 三个类：`init` 存 `self.device_interface` / `self.device_type`

6. `torch/_inductor/compile_fx.py`（PR-5a/5b，旧非 trees 路径）