尧图网站建设 
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_tensorNestedTensorImpl( 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 部分按算子类别拆分为如下文件均已在仓库中逐一核实存在文件职责NestedTensorImpl.h / NestedTensorImpl.cppNestedTensor 数据结构本体及其方法sizes/strides/offsets 访问、buffer 视图、view 支持等NestedTensorUtils.h / NestedTensorUtils.cpp操作 NestedTensor 的工具函数包括下文详述的map_nested_tensor与wrap_bufferNestedTensorUnaryOps.cpp一元算子可经map_nt高效实现的函数NestedTensorBinaryOps.h / NestedTensorBinaryOps.cpp二元算子可经NestedTensor_elementwise_Tensor高效实现的函数该辅助函数位于 cpp 文件中NestedTensorFactories.cppNestedTensor 的工厂函数如empty_likeNestedTensorMath.h / NestedTensorMath.cpp数学函数如 softmax、embeddingNestedTensorMatmul.cpp矩阵乘类函数matmul、linear、bmmNestedTensorTransformerFunctions.h / .cpp支撑 BetterTransformer 工作流的功能函数cuda/CUDA 实现如 NestedTensorMatmul.cu、NestedTensorTransformerFunctions.cu 等此外目录中还存在文档未逐一点名的 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_tensorNestedTensorImpl(f(nt_impl-get_buffer()), sizes); }其执行步骤可拆解为从输入 NestedTensor 取得NestedTensorImpl取得该 NestedTensor 的nested_sizes调用get_buffer()获得一个与输入 NestedTensor共享存储的扁平稠密张量在稠密张量上调用函数f即直接复用普通稠密算子的 kernel用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::forwardF(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 实现新算子时可以按以下判据决策逐元素/逐分量独立如点态激活、逐元素二元运算走稠密化路线用 buffer 视图复用稠密 kernel零额外开销规整性可检测如各组件某维尺寸一致参考 NestedTensorTransformerFunctions.cpp 中linear的模式规整时稠密化、否则降级边界跨组件不可分离如 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch NestedTensor 源码解析：从数据结构到自定义算子实现指南
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_tensorNestedTensorImpl( 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 部分按算子类别拆分为如下文件均已在仓库中逐一核实存在文件职责NestedTensorImpl.h / NestedTensorImpl.cppNestedTensor 数据结构本体及其方法sizes/strides/offsets 访问、buffer 视图、view 支持等NestedTensorUtils.h / NestedTensorUtils.cpp操作 NestedTensor 的工具函数包括下文详述的map_nested_tensor与wrap_bufferNestedTensorUnaryOps.cpp一元算子可经map_nt高效实现的函数NestedTensorBinaryOps.h / NestedTensorBinaryOps.cpp二元算子可经NestedTensor_elementwise_Tensor高效实现的函数该辅助函数位于 cpp 文件中NestedTensorFactories.cppNestedTensor 的工厂函数如empty_likeNestedTensorMath.h / NestedTensorMath.cpp数学函数如 softmax、embeddingNestedTensorMatmul.cpp矩阵乘类函数matmul、linear、bmmNestedTensorTransformerFunctions.h / .cpp支撑 BetterTransformer 工作流的功能函数cuda/CUDA 实现如 NestedTensorMatmul.cu、NestedTensorTransformerFunctions.cu 等此外目录中还存在文档未逐一点名的 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_tensorNestedTensorImpl(f(nt_impl-get_buffer()), sizes); }其执行步骤可拆解为从输入 NestedTensor 取得NestedTensorImpl取得该 NestedTensor 的nested_sizes调用get_buffer()获得一个与输入 NestedTensor共享存储的扁平稠密张量在稠密张量上调用函数f即直接复用普通稠密算子的 kernel用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::forwardF(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 实现新算子时可以按以下判据决策逐元素/逐分量独立如点态激活、逐元素二元运算走稠密化路线用 buffer 视图复用稠密 kernel零额外开销规整性可检测如各组件某维尺寸一致参考 NestedTensorTransformerFunctions.cpp 中linear的模式规整时稠密化、否则降级边界跨组件不可分离如 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考