企业官网建设流程全解析

PyTorch NestedTensor 源码解析：从数据结构到自定义算子实现指南

【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch

本文基于 PyTorch 仓库中 NestedTensor 模块的官方开发者文档 aten/src/ATen/native/nested/README.md，系统梳理 NestedTensor 的底层存储结构（nested_sizes_、nested_strides_、storage_offsets_ 与共享 Storage）、C++/Python 代码组织结构，以及在新算子上实现 NestedTensor 支持的两条技术路线（基于map_nt的稠密化高效实现与基于map_nested_tensor的组件级遍历实现），并结合仓库源码给出可验证的关键文件位置。读完后你可以理解 PyTorch 如何用单一扁平存储表达变长/不规则（ragged）张量序列，并能判断新算子应选用哪种实现策略。

NestedTensor 是什么：面向不规则形状的数据结构

NestedTensor 是 torch.Tensor 的泛化，旨在简化不同形状、不同长度数据的处理。其核心用法是表示一个由 N 个张量组成的列表：列表中的每个张量（称为tensor_component）具有相同的维数（ndim），所有组件被展平并合并进单一的 NestedTensor，同时保留足以重建各个原始组件的元信息。

文档明确列出了支撑这一表达的四个核心组成部分：

• nested_sizes_：形状为n_tensor_components x n_dims的二维张量，记录每个组件在各维度上的尺寸；
• nested_strides_：形状为n_tensor_components x n_dims的二维张量，记录每个组件的步长；
• storage_offsets_：一维张量，记录每个组件在扁平化存储中的起始偏移；
• storage_：实际承载全部展平后组件数据的存储对象，定义在c10::TensorImp之上。

从源码结构看，这些元信息直接体现在NestedTensorImpl的构造函数签名中。aten/src/ATen/NestedTensorImpl.h 中定义了struct TORCH_API NestedTensorImpl : public c10::TensorImpl，其完整构造入口为：

explicit NestedTensorImpl( Storage storage, c10::DispatchKeySet key_set, const caffe2::TypeMeta data_type, at::Tensor nested_sizes, at::Tensor nested_strides, at::Tensor storage_offsets);

这说明 NestedTensor 在类型系统上是一个TensorImpl的子类实例：它继承自 c10::TensorImpl，因此可以无缝接入 PyTorch 的 Dispatch、Autograd 与版本计数（version counter）机制（shallow_copy_and_detach等钩子在 NestedTensorImpl.h 中有对应的 override）。

为什么 C++ 侧通常用 wrap_buffer 而不是直接构造

文档特别指出：在 C++ 中构造 NestedTensor 时，你通常不会直接使用NestedTensorImpl的构造函数，而是使用定义在 aten/src/ATen/native/nested/NestedTensorUtils.h 的wrap_buffer函数。它是对构造函数的一个薄封装，职责是确保输入 buffer 为连续内存。原因如下：从稠密张量构造 NestedTensor 时，系统会对输入张量的 Storage 做浅拷贝（shallow copy）；如果输入张量不满足tensor.numel() == tensor.storage.numel()（例如它是某个大 Storage 的视图），浅拷贝 Storage 会覆盖到不属于该张量的内存区域，导致未定义行为。

源码中wrap_buffer的防护检查与文档描述一一对应：

inline at::Tensor wrap_buffer(const at::Tensor& buffer, const at::Tensor& nested_sizes) { TORCH_CHECK( buffer.dim() == 1, "Expected given buffer to be 1dim, but got ", buffer.dim(), " instead."); TORCH_CHECK( buffer.is_contiguous(), "Expected given buffer to be contiguous."); return at::detail::make_tensor<NestedTensorImpl>( buffer, nested_sizes); }

即 buffer 必须是一维且连续的，随后通过两参构造（buffer + nested_sizes，由 NestedTensorImpl.h 声明，strides 与 offsets 可由 sizes 推断）创建嵌套张量。同一头文件 NestedTensorUtils.h#L51-L59 还提供了显式传入nested_strides与storage_offsets的重载，供需要精确控制视图语义的场合使用。

与之配合的是读取方向的接口get_buffer()，定义在 NestedTensorImpl.h#L83-L88：它将 NestedTensor 以共享同一storage_impl的一维连续张量形式返回（即一个 view），并先断言嵌套张量本身是连续的。这正是后文map_nt高效实现的基石——操作稠密 buffer 时零拷贝共享存储。

代码结构总览：C++ 与 Python 的分层

文档将 NestedTensor 代码分为 C++ 与 Python 两部分：C++ 位于 aten/src/ATen/native/nested/，Python 侧位于 torch/nested/init.py。C++ 部分按算子类别拆分为如下文件（均已在仓库中逐一核实存在）：

此外，目录中还存在文档未逐一点名的 NestedTensorAliases.cpp 与 NestedTensorBackward.cpp，从命名与内容可推断前者负责别名/元信息类接口，后者负责反向算子注册——这也是map_nt的使用者之一（见下文）。

实现新算子路线一：稠密化的高效实现（map_nt）

文档将可实现的算子分为两大类。第一类是可以高效地"把 NestedTensor 视为稠密张量、将 raggedness（不规则性）折叠进稠密维度"来实现的函数。典型代表是逐元素的一元算子（abs、log、relu）与二元算子（add、mul、div）。

文档给出的map_nt定义如下：

template <typename Func> Tensor map_nt(const Tensor& nt, Func f) { auto* nt_impl = get_nested_tensor_impl(nt); const auto& sizes = nt_impl->get_nested_sizes(); return at::detail::make_tensor<NestedTensorImpl>(f(nt_impl->get_buffer()), sizes); }

其执行步骤可拆解为：

1. 从输入 NestedTensor 取得NestedTensorImpl；
2. 取得该 NestedTensor 的nested_sizes；
3. 调用get_buffer()，获得一个与输入 NestedTensor共享存储的扁平稠密张量；
4. 在稠密张量上调用函数f（即直接复用普通稠密算子的 kernel）；
5. 用f的输出与输入 sizes 构造一个新的 NestedTensor。

这套范式的精髓在于算子实现被归约到稠密 kernel 调用：逐元素算子天然与"元素属于哪个组件"无关，因此在折叠视图上执行是安全的，且没有任何逐组件循环。在 NestedTensorUnaryOps.cpp 中可以看到宏将其批量套用到各一元算子上（return map_nt(self, at::op_name);），NestedTensorBackward.cpp 与 NestedTensorMath.h 中也有同类使用，印证了文档"一元/二元算子经 map_nt 高效实现"的说法。

条件规整下的进一步特化

文档还指出：某些重要的函数在满足特定规整性条件（regularity，例如各组件在某维度上尺寸一致）时，可以通过访问底层 buffer 并以特殊方式查看（view）它来获得高效实现。文档给出的范例是linear，其实现位于 aten/src/ATen/native/nested/NestedTensorTransformerFunctions.cpp。从源码结构看，这类实现先检测规整条件是否成立，成立时把 buffer 按稠密形状解释（等价于把变长维度折叠进 batch/序列维度）直接喂给 GEMM 类 kernel，不成立时再退回逐组件处理——这也解释了为什么 Transformer 场景（序列长度不齐的 attention/linear）是 NestedTensor 的重点优化对象，且 cuda/NestedTensorTransformerFunctions.cu 提供了对应的 CUDA 特化。

实现新算子路线二：逐组件遍历（map_nested_tensor 与性能权衡）

第二类函数无法通过稠密视图高效实现，代表性例子是对 ragged 维度做softmax_nested。文档解释了原因：当对 ragged 维度计算 softmax 时，问题边界不可平凡分离——把 ragged 维度折叠进另一个维度后，无法确定哪些元素属于哪个 tensor_component；因此只能对每个组件单独应用 softmax。其实现位于 aten/src/ATen/native/nested/NestedTensorMath.cpp。

逐组件遍历的朴素做法是：迭代所有组件，对每个组件各启动一次 CUDA kernel。文档明确指出这对 GPU 而言非常低效，理想情况是启动一个对全部组件并行操作的单一 kernel。

当性能不是首要目标、只是想快速获得算子覆盖时，文档推荐使用 NestedTensorUtils.h 中的map_nested_tensor：

// This function is meant to ease rapid operator coverage for // NestedTensor kernels. It is not meant to be efficient. Use it judiciously. template <class F, class... A> inline at::Tensor map_nested_tensor(F&& fn, A... a) { return wrap_tensor_node( impl::map(std::forward<F>(fn), impl::get_nested_tensor_structure(a)...), std::nullopt, std::nullopt, std::nullopt, std::nullopt); }

源码注释与文档口径一致："用于快速扩展 NestedTensor kernel 的算子覆盖，不追求高效，谨慎使用"。其内部流程为：impl::get_nested_tensor_structure提取各输入张量的嵌套结构 →impl::map对组件列表逐一应用fn→ 用wrap_tensor_node将结果组件重新包装为 NestedTensor。文档给出的定位是：CUDA 上每个组件都会 launch 一个新 kernel，不宜作为高性能实现；但它是 CPU 实现的良好基线。

测试与验证：算子必须配套单测

文档的 Testing 章节要求：NestedTensor 的单测位于 test/test_nestedtensor.py，为 NestedTensor 新增算子时必须补充对应单测；这些测试运行在 CI 上，测试失败会导致 PR 无法合并。

从测试文件名可以推断覆盖范围相当完整：除构造与结构转换外（如test_nested_tensor_from_padded_cpu、test_nested_tensor_from_mask_and_to_padded_cpu、test_nested_tensor_to_padded_tensor_cpu），还包括数学与线性代数路径（test_nested_tensor_softmax_cpu、test_nested_tensor_linear_cpu、test_nested_tensor_bmm_gradcheck_cpu、test_nested_tensor_matmul_gradcheck_cpu）以及视图/索引/梯度行为（test_nested_tensor_reshape_gradcheck_cpu、test_indexing_backward_cpu）。这些测试类名与 test/test_nestedtensor.py 中的用例一一对应，可作为新算子补充测试时的参照模板。

小结：选对实现路线的关键判据

综合文档与源码证据，为 NestedTensor 实现新算子时可以按以下判据决策：

1. 逐元素/逐分量独立（如点态激活、逐元素二元运算）：走稠密化路线，用 buffer 视图复用稠密 kernel，零额外开销；
2. 规整性可检测（如各组件某维尺寸一致）：参考 NestedTensorTransformerFunctions.cpp 中linear的模式，规整时稠密化、否则降级；
3. 边界跨组件不可分离（如 ragged 维 softmax）：必须逐组件计算；性能敏感场景应编写融合 kernel（参考 cuda/ 目录下的特化实现），否则以map_nested_tensor作为正确性基线与 CPU 实现。

最后需要说明：原文档中 "Triton" 与 "Best Practices" 两个小节在当前仓库版本中是空占位章节（仅有标题、无正文），因此本文不对其做内容虚构，相关演进请以仓库后续提交为准。

【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch