Dtensor新特性FastPath
DTensor FastPath 技术分析 1. 概述 DTensor 的核心工作流是:拆包参数 → 推断输出分片策略 → 可能做 redistribute → 执行本地算子 → 包装结果。 在 PyTorch 2.9 及之前版本中,这条路径完全在 Python 中执行,入口是 __torch_dispatch__。dispatch 是存在开销的,往往被 __torch_dispatch__ 拦截后在 Python 侧的调库开销更大,甚至可能超过一些小算子的执行时间。 关键 PR:https://github.com/pytorch/pytorch/pull/167051 核心思想:将操作搬到 C++,在 python_arg_parser.cpp 的 dispatch_on_subclass 层面拦截 DTensor,绕过 __torch_dispatch__ 协议,直接进入 C++ 实现的 dispatchDTensorOp。 2. 核心机制 2.1 __torch_dispatch__ 变更至 C++ 2.1.1 旧版本 __torch_dispatch__ 拦截逻辑 在旧版中,DTensor 是一个标准的 Tensor subclass,算子分发走的是 PyTorch 的标准 __torch_dispatch__ 协议。先梳理一下 2.9 之前是怎么把 DTensor 全部通过 __torch_dispatch__ 处理的。 当 DTensor 被构造时(通过 _make_dtensor),C++ 会检测它定义了 __torch_dispatch__,于是在 tensor 的 DispatchKeySet 上设置 Python + PythonTLSSnapshot 两个 key: ...