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：

// python_variable.cpp:2204
if (check_has_torch_dispatch(obj)) {
    var.unsafeGetTensorImpl()->set_python_dispatch(true);
}

// c10/core/TensorImpl.h:1991
void set_python_dispatch(bool k) {
    if (k) key_set_ = key_set_.add(c10::python_ks);
}

由于携带了 Python key，会被 dispatch 的 PythonFallbackKernel 拦截，通过 PyInterpreter.cpp:305 的 ConcretePyInterpreterVTable::dispatch 执行。

主要作用是提取 Python 对象，通过 handle_torch_function_no_python_arg_parser 分发：

  • Step 1（:558）：dispatch_on_mode —— 如果有 TorchDispatchMode 活跃，先调 mode 的 __torch_dispatch__。（这个用法是 with xxx，去 call __torch_dispatch__，比如 xray）
  • Step 2（:577）：dispatch_on_subclass —— 调 subclass（DTensor）的 __torch_dispatch__。（这个用法是 Tensor 子类，比如 DTensor）
static py::object dispatch_on_subclass(...) {
    for (auto& arg : overloaded_args) {     // 遍历所有 DTensor 参数
        // 第 322 行：从 arg 上取 "__torch_dispatch__" 属性
        py::object torch_function =
            PyObject_FastGetAttrString(arg, torch_function_name_str);
            //                              ↑ "__torch_dispatch__"

        // ... 跳过 disabled 的情况 ...

        // 第 349-369 行：先检查有没有注册的 per-op dispatch rule
        // 对于普通 DTensor 没有，所以跳过

        ret = py::reinterpret_steal<py::object>(PyObject_CallFunctionObjArgs(
            torch_function.ptr(),       // DTensor.__torch_dispatch__ (bound method)
            torch_api_function,         // torch.ops.aten.add.Tensor
            py_types.ptr(),             // (DTensor,)
            args,                       // (dtensor_a, dtensor_b)
            kwargs,                     // {}
            NULL));

从这里，C++ 通过 CPython 的 C API 去调用 Python 解释器执行这个 Python 函数，这也是为什么旧版本慢——每次算子调用都要从 C++ “绕回” Python 解释器去执行 __torch_dispatch__，而新版的 fast path 就是为了尽量不绕回 Python。

2.1.2 新版本改动

    // 检查完 per-op dispatch rule 后，加了一个 if 判断：

+   if (!is_torch_function && is_dtensor(arg)) {
+       ret = dispatchDTensorOp(
+           *opt_op, torch_api_function, args, kwargs, opt_stack);
+   } else {
        ret = PyObject_CallFunctionObjArgs(
            torch_function.ptr(),
            torch_api_function,
            py_types.ptr(),
            args, kwargs, NULL);
+   }

dispatchDTensorOp 是一个与旧版本 __torch_dispatch__ 等价的计算流程，只不过用 C++ 改写了。概览：

dispatchDTensorOp(op, args, kwargs, stack)
│
├─ A: custom_op_handlers 命中？ → 直接调 handler，return
│
├─ B: C++ cache 命中？
│   ├─ YES → cached_sharding 有值，跳到 D
│   └─ NO  → 进入 C
│
├─ C: 调 Python propagation，缓存结果
│   └─ decompose 了？ → 直接 return 结果
│
├─ D: 执行本地算子
│   ├─ 不需要 redistribute 且非随机算子？
│   │   → C++ callBoxed()                    ← 全 C++，零 Python
│   └─ 否则 → 调 Python slow path
│
└─ E: 包装结果
    ├─ 普通/view 算子？ → wrap + C++ alias 处理
    └─ inplace/out 算子？ → 调 Python tail 处理

阶段 A：custom_op_handlers（优先级最高）

这里的逻辑不变，依然是全局优先级最高的，调用的也是 Python 方法。

// C++ 侧
const auto custom_op_handlers =
    op_dispatcher.attr(dtensor_interned_strings._custom_op_handlers);
PyObject* custom_op_handler =
    PyDict_GetItemWithError(custom_op_handlers.ptr(), py_op.ptr());
if (custom_op_handler) {
    auto result = checked_vectorcall(
        custom_op_handler, py_op.ptr(), args.ptr(), kwargs.ptr());
    stack->clear();
    return result;        // 命中就直接返回，后面全跳过
}

# 等价的 Python 逻辑
if op_call in self._custom_op_handlers:
    return self._custom_op_handlers[op_call](op_call, args, kwargs)

阶段 B：C++ Cache 查找

这一部分会尝试读缓存，如果存在 native_sharding_propagator_cache，就无需重新执行阶段 C。

create_native_op_schema 遍历 stack 上的参数，为每个 DTensor 提取 DTensorSpec（placement + shape），组成 NativeOpSchema 作为 hash key。然后在线程局部的 C++ HashMap 中查找。

  • 命中：cached_sharding 有值 → 跳过整个阶段 C（不调任何 Python）
  • 未命中：cached_sharding 为空 → 进入阶段 C
  • 构建 key 失败（如有 SymInt）：opt_native_op_schema 为 nullopt → 也进入阶段 C
// 尝试在 C++ 中构建一个精简的缓存 key
auto opt_native_op_schema = create_native_op_schema(op, py_op, stack);
if (opt_native_op_schema.has_value()) {
    native_sharding_propagator_cache =
        &get_thread_local_native_sharding_propagator_cache();
    cached_sharding =
        native_sharding_propagator_cache->find(opt_native_op_schema->first);
}

阶段 C：回退 Python 做 sharding propagation

if (!cached_sharding) {
    // 1. 调 Python 构建 OpInfo（unwrap DTensor 参数）
    py_op_info = checked_vectorcall(
        op_dispatcher.attr("unwrap_to_op_info").ptr(),
        py_op.ptr(), args.ptr(), kwargs.ptr());

    // 2. 调 Python 做 sharding propagation
    py::object sharding = checked_vectorcall(
        op_dispatcher.attr("_propagate_op_sharding_dispatch_slow_path").ptr(),
        py_op.ptr(), args.ptr(), kwargs.ptr(), py_op_info.ptr(),
        /*try_cache*/ !opt_native_op_schema.has_value() ? Py_True : Py_False);

# 等价的 Python 逻辑
op_info = self.unwrap_to_op_info(op_call, args, kwargs)
self.sharding_propagator.propagate(op_info)

如果算子没有注册策略但能 decompose，propagation 内部会 decompose 并直接返回结果。C++ 这里通过类型检查区分"返回的是 OutputSharding 还是直接结果"：

    // 3. propagation 有时直接返回 decompose 的结果（不是 OutputSharding）
    if (!py::isinstance(sharding, get_output_sharding_class())) {
        stack->clear();
        return sharding;  // 直接返回 tensor 结果
    }

# 等价的 Python 逻辑
except NotImplementedError:
    if has_kernel_for(CompositeImplicitAutograd):
        out = op_call.decompose(*args, **kwargs)
        return out

这是新版的核心收益：首次 miss 时计算结果，缓存起来，后续相同 op + placement 组合直接命中阶段 B：

    // 4. 缓存结果，下次就走阶段 B 的快路径
    cached_sharding = sharding;
    if (opt_native_op_schema.has_value()) {
        native_sharding_propagator_cache->insert(
            std::move(opt_native_op_schema->first), std::move(sharding));
    }

阶段 D：获取 compute_mesh + 执行本地算子

// 获取 compute_mesh
py::object compute_mesh;
if (opt_native_op_schema.has_value()) {
    compute_mesh = std::move(opt_native_op_schema->second);  // C++ 已经拿到了
} else {
    get_py_op_info_if_needed();                               // 回退 Python 取
    compute_mesh = py_op_info.attr("compute_mesh");
}

// 执行本地算子
const bool local_results_success =
    get_local_results(op, cached_sharding, compute_mesh, participating, stack);
if (local_results_success) {
    // C++ 快路径成功：stack 上已经是本地计算结果
    py_local_results = torch::jit::createPyObjectForStack(std::move(*stack));
} else {
    // 回退 Python（需要 redistribute 或是随机算子）
    get_py_op_info_if_needed();
    py_local_results = checked_vectorcall(
        get_dtensor_get_local_results_slow_path().ptr(), ...);
}

阶段 E：包装结果

C++ 侧有一系列限定条件，如果不满足，还是会回退 Python 实现：

if (!is_inplace_op && !is_out_variant_op &&
    !(output_spec.is_none() &&
      (op.operator_name().name == "aten::equal" &&
       is_default_overload(op.operator_name().overload_name)))) {
    // C++ fast path: wrap + alias 处理
} else {
    // 回退 Python tail 处理
}

2.2 Cache 详解

2.2.1 Cache 的数据结构

NativeShardingPropagatorCache
  └── c10::FastMap<NativeOpSchema, py::object>
                    ↑ key              ↑ value
组成说明
KeyNativeOpSchema — 描述"这个算子，配上这些 placement 和参数，是什么样的调用"
Valuepy::object — 一个 Python OutputSharding 对象，包含 sharding propagation 的完整结果（output_spec、redistribute_schema 等）
存储位置thread_local，每个线程一份独立的 cache（无锁，无跨线程竞争）

2.2.2 Cache Key — NativeOpSchema 由什么组成

NativeOpSchema 的相等和哈希由三部分决定（python_variable.cpp:1117-1123）：

bool operator==(const NativeOpSchema& rhs) const {
    return op_ == rhs.op_                             // (1) 算子
        && args_schema_len_ == rhs.args_schema_len_   // (2) 参数数量
        && comparison_key_ == rhs.comparison_key_;     // (3) 参数特征
}

具体来说：

  1. op_ — c10::OperatorHandle，唯一标识一个算子，比如 aten::add.Tensor、aten::mm.default。
  2. args_schema_len_ — 参数列表长度。
  3. comparison_key_ — SmallVector<IValueOrDTensorSpec, 8>，这是核心。遍历所有参数，对每个参数按类型生成 key 元素：
参数类型放入 comparison_key 的内容
DTensordtensor._spec（即 DTensorSpec Python 对象，包含 placements + mesh + tensor_meta）
普通 Tensor为其构造一个 replicate 的 DTensorSpec
int/float/bool 等标量仅当 idx >= static_argnum 时放入（即非静态参数才参与比较）
List（如 torch.cat 的输入列表）先放列表长度，再逐元素递归处理

关键点：DTensorSpec 的比较用的是 Python 的 __eq__，所以只要 placements 和 mesh 相同，就视为等价。tensor 的具体 shape/stride/数据不参与 key 比较（前提是 RuntimeSchemaInfo 将对应参数标为 static）。

2.2.3 命中条件

Cache HIT 条件（所有条件同时满足）：

  1. 同一个算子（OperatorHandle 相同，比如都是 aten::add.Tensor）
  2. 所有 DTensor 参数的 DTensorSpec 相同（placements 和 mesh 一致）
  3. 非 static 标量参数值相同（比如 alpha=1 vs alpha=2 是不同 key）
  4. 参数个数相同
  5. 没有 SymInt（有 SymInt 直接跳过 cache）

Cache MISS 场景：

  • 第一次调用某个算子 + placement 组合 → 必然 miss
  • 同一算子但 placement 不同：torch.add(Shard(0), Shard(0)) vs torch.add(Replicate(), Replicate()) → miss
  • 同一算子但非 static 参数变化：某些算子的 dim 参数等
  • DTensor 含 SymInt → 直接 bail out 到 Python path，不查 cache
  • RuntimeSchemaInfo 标记 needs_pytree=True 且 fast path 未启用 → bail out

2.2.4 Cache Value — OutputSharding 包含什么

Cache 存储的 OutputSharding Python 对象包含：

字段说明
output_spec输出的 DTensorSpec（placements + mesh）
redistribute_schema如果输入需要重新分布，描述怎么 redistribute（None 说明不需要）
needs_redistribute是否需要在执行前 redistribute 输入

当 cache hit 时，后续流程直接看 output_sharding 决定：

  • 无需 redistribute → C++ 直接执行 local op（get_local_results :1369），全程不碰 Python
  • 需要 redistribute → 回退到 Python slow path 做 redistribute + 执行