• 状态（Status）： Draft
• 作者（Authors）： @yucaopanmu
• 创建日期（Created）： 2026-09-16
• 更新日期（Updated）： 2026-09-24
• 相关 Issue/PR：
  • pytorch：#46265
  • op-plugin：#5879

---

1. 概述

1.1 简介

torch.nested.to_padded_tensor 用于将 NestedTensor（嵌套张量 / 不规则张量）转换为补齐（padding）后的稠密 Tensor，是 NestedTensor 生态中的基础算子。当前在 Ascend NPU 上，该算子对应的 NestedTensorPrivateUse1 dispatch key 下没有注册实现，调用时会报 Could not run 'aten::to_padded_tensor' with arguments from the 'NestedTensornpu' backend。

本方案采用 双仓联动 的方式打通该路径：

• pytorch（torch_npu）仓 在 codegen 层（torchnpugen/gen_backend_stubs.py）为 NestedTensorRegister 补充 to_padded_tensor 的注册，并引入 op_plugin/OpApiInterface.h 头文件，将算子注册到 NestedTensorPrivateUse1 dispatch key。
• op-plugin 仓 在 op_api 命名空间中提供 to_padded_tensor 的 kernel 实现，复用 PyTorch 原生设备泛化实现 at::native::NestedTensor_to_padded_tensor_generic，完成实际的补齐转换。

1.2 动机

• NestedTensor 尚未在 NPU 上提供完整的算子集合，to_padded_tensor 是高频的预处理 / 校验算子。
• 直接调用报错，导致上层基于 NestedTensor 的模型无法在 NPU 上运行。
• 复用 PyTorch 原生设备泛化实现，无需在 NPU 上重新开发 kernel，成本低、可维护性好。
• 相比早期「将适配逻辑放在 codegen 生成代码里」的方案，将 kernel 下沉到 op-plugin 仓更清晰：pytorch 仓只负责 dispatch key 的注册，op-plugin 仓负责算子实现，职责边界更清晰，也符合「pytorch 生成 dispatch key、op-plugin 消费」的分工。

1.3 目标

• 在 NestedTensorPrivateUse1 下注册 to_padded_tensor，使 torch.nested.to_padded_tensor 在 NPU 后端不再报错。
• 由 op-plugin 提供 op_api::to_padded_tensor 符号，供 pytorch codegen 生成的注册代码调用。
• 非目标：不实现 NPU 专用 kernel，不做性能优化，不覆盖其他 NestedTensor 算子；aclnn_extension（cpp_extension）路径隔离、不实现 to_padded_tensor。

---

2. 用例分析

单元测试用例（test/test_base_ops/test_nested_to_padded_tensor.py）覆盖以下场景，均以 CPU / NPU 双份逐位一致的 NestedTensor 对比验证：

用例	输入子张量形状	参数	验证点
fp32 32-byte 对齐	(32, 32) × 3	缺省	对齐形状的基本转换
fp32 非 32-byte 对齐	(5, 7) × 3	缺省	非对齐形状的转换
fp16	(32, 16) × 2	缺省	半精度 dtype
非零 padding	(4, 6) × 3	padding=1.5	补齐值 + 缺省 output_size
指定 output_size	(4, 6) × 2	output_size=[3, 8, 8]	显式指定输出尺寸
3D	(2, 3, 4) × 2	缺省	三维 NestedTensor
不规则形状	(2, 3), (3, 4), (1, 5)	缺省	真正的 padding 补齐
不规则形状 + padding	(2, 3), (3, 4), (1, 5)	padding=0.5	不规则 + 非零补齐值
1D	(3,), (5,), (7,)	缺省	一维 NestedTensor

DFX 要求：

• 兼容性： 对外 API 签名与 PyTorch 保持一致，不改变用户调用方式。
• 可维护性： 改动分别落在 pytorch 的 codegen 生成逻辑与 op-plugin 的 opapi kernel，遵循各自仓库已有风格。
• 可测试性： 两仓改动一起编包后，用新包运行单元测试 test/test_base_ops/test_nested_to_padded_tensor.py（run_tests 输出 OK 表示 9 个用例全部通过、CPU / NPU 数值一致），并用 torchapi_id4942_to_padded_tensor.py 脚本在 NPU 上做端到端验证。

约束：

• 仅在 NestedTensorPrivateUse1 key 下生效，不影响其他后端。
• pytorch 仓依赖 op-plugin 提供的 op_api::to_padded_tensor 符号（跨仓接口依赖）。

---

3. 方案设计

3.1 总体方案

将算子注册与算子实现拆分到两个仓库：

1. op-plugin 仓：
   • 在 op_plugin/config/op_plugin_functions.yaml 中注册 to_padded_tensor 的 op_api 接口（版本 [v2.1, newest]）。
   • 新增 op_plugin/ops/opapi/NestedTensorToPaddedTensorKernelNpuOpApi.cpp，在 op_api 命名空间提供 to_padded_tensor 实现，委托给 at::native::NestedTensor_to_padded_tensor_generic。
   • 在 test/core_tests/torch_npu_OpApi_schema_all.json 中登记 to_padded_tensor 基线（版本 ["v2.1", "newest"]）。
2. pytorch（torch_npu）仓：
   • 在 torchnpugen/gen_backend_stubs.py 新增 _nestedtensor_register_header()，为生成的 NestedTensorRegister.cpp 无条件引入 op_plugin/OpApiInterface.h（提供 op_api::to_padded_tensor 声明），并在非 aclnn_extension 场景另引入 op_plugin/OpInterface.h（提供主仓 op_plugin:: 命名空间核函数声明，aclnn_extension 场景不涉及、无需引入）。
   • 新增 _nestedtensor_extra_impls()（对齐已有 _quantized_extra_impls()），返回 unbind.int / values / _nested_tensor_size 三个原生注册，以及条件注册的 to_padded_tensor。
   • 将 SPECIAL_REGISTERS["nestedtensor"] 的 header 改为 _nestedtensor_register_header()、extra_impls 改为 _nestedtensor_extra_impls()。
   • to_padded_tensor 仅在 not _is_aclnn_extension_codegen() 时注册，aclnn_extension 路径隔离、不实现——该算子归属主仓、复用 PyTorch 原生实现，无需进入面向用户自定义算子的 extension 通道。

完整调用链路：

torch.nested.to_padded_tensor(nested, padding, output_size)
        │  schema: aten::to_padded_tensor(Tensor, float, SymInt[]?)
        ▼
  dispatcher (按 dispatch key 查找)
        │  NestedTensorPrivateUse1
        ▼
  NestedTensorRegister.cpp
        │  m.impl("to_padded_tensor", TORCH_FN(op_api::to_padded_tensor))
        ▼
  op_api::to_padded_tensor(self, padding, OptionalIntArrayRef output_size)
        │
        ▼
  at::native::NestedTensor_to_padded_tensor_generic(self, padding, output_size)
        │  (PyTorch 原生设备泛化实现)
        ▼
  padded Tensor

3.2 技术选型

方案	说明	结论
A. 在 codegen 层内联匿名适配函数	在 gen_backend_stubs.py 生成 SymInt[]? → OptionalIntArrayRef 的适配函数并注册	弃用：把算子实现逻辑混入 codegen，职责不清
B. op-plugin 提供 opapi kernel + pytorch codegen 注册	算子实现下沉到 op-plugin，pytorch 只做 dispatch key 注册	采用
C. Python 层 patch	绕过 C++ dispatch	弃用：难以覆盖所有调用路径

选择 B 的理由：

• 与 op-plugin 现有 opapi 目录下的算子实现风格一致（如 NestedTensorToPaddedTensorKernelNpuOpApi.cpp）。
• 与 pytorch 现有 SPECIAL_REGISTERS["nestedtensor"] 中 unbind.int、values、_nested_tensor_size 的注册方式一致。
• 职责清晰：pytorch 生成 dispatch key 注册，op-plugin 消费并提供算子符号。

3.3 功能与性能设计

• 功能流程：
  1. NestedTensorRegister.cpp 在 NestedTensorPrivateUse1 下注册 to_padded_tensor。
  2. dispatcher 命中 op_api::to_padded_tensor。
  3. 复用 at::native::NestedTensor_to_padded_tensor_generic 完成转换。
• 性能： 复用 PyTorch 原生设备泛化实现，算子自动在 NestedTensor buffer 所在设备（NPU）上执行；本方案不以性能为目标。

3.4 安全隐私与 DFX 设计

• 无新增安全 / 隐私面。
• 兼容性：保持原有 schema 不变，用户无感知。
• 可维护性：算子实现集中在 op-plugin 的 opapi 目录，后续如需迁移到 NPU kernel，只需替换 op_api::to_padded_tensor 的实现，无需改动 pytorch 仓。

3.5 编程与调用设计

3.5.1 编程模型基本设计

• 开发环境： pytorch（torch_npu）源仓库 + op-plugin 源仓库；pytorch 侧通过 generate_code.sh 生成代码。
• 开发约束： op_api::to_padded_tensor 与生成代码中 TORCH_FN(op_api::to_padded_tensor) 的签名必须一致。
• 可验收设计：
  • 单元测试： 两仓改动一起编包后，用新包运行 python test/test_base_ops/test_nested_to_padded_tensor.py，run_tests 输出 OK 表示 9 个用例全部通过、CPU / NPU 输出数值一致。
  • 端到端验证： 运行脚本 torchapi_id4942_to_padded_tensor.py，分别验证「32-byte 对齐」与「非 32-byte 对齐」两类 NestedTensor，确认输出中无 error 字段且 target_api 为 torch.nested.to_padded_tensor。

3.5.2 接口定义与设计

接口描述： aten::to_padded_tensor(Tensor self, float padding, SymInt[]? output_size=None) -> Tensor

改动内容：

1. op-plugin —— op_plugin/config/op_plugin_functions.yaml

在 official 区域新增 to_padded_tensor（该算子带 SymInt 参数，登记在 official 区域，不放入 symint 区域）：

- func: to_padded_tensor(Tensor self, float padding, SymInt[]? output_size=None) -> Tensor
  op_api: [v2.1, newest]

版本说明： to_padded_tensor 的 schema 含 SymInt[]?，版本登记为 [v2.1, newest]，与 topk 等带符号整型参数的 op_api 算子保持一致（区别于不带 SymInt 的普通算子常用的 all_version）。条目位置保持在 threshold_backward 之后、topk 之前，与基线 JSON 顺序一致。

SymInt 登记说明： 只需放 official、不登记 symint，也无需新增同名重载。symint 段不是生成列表，而是标记集合（命中才给算子名追加 _symint 后缀）；判定规则为——kernel 用 int64_t / IntArrayRef / OptionalIntArrayRef 实现则只放 official，用 c10::SymInt / SymIntArrayRef 实现（需保留符号）才在 symint 段补一条同名标记。本 kernel 是非 symint 签名 at::OptionalIntArrayRef，放 official（has_symint=False）生成的 op_api::to_padded_tensor(..., OptionalIntArrayRef) 与实现一致；若误加 symint 标记，会生成引用不存在的 op_api::to_padded_tensor_symint 的 wrapper 而编译/链接报错。又因该 schema 是 aten 唯一签名、无 int[]? 重载，非 symint kernel 经 dispatcher boxed/unboxed 物化即可服务 symint schema，故无需额外重载。

2. op-plugin —— test/core_tests/torch_npu_OpApi_schema_all.json

在 threshold_backward 与 topk 之间新增基线条目，供兼容性测试 test_op_func_compatibility 校验：

"func: to_padded_tensor(Tensor self, float padding, SymInt[]? output_size=None) -> Tensor": {
  "version": ["v2.1", "newest"]
},

3. op-plugin —— 新增 op_plugin/ops/opapi/NestedTensorToPaddedTensorKernelNpuOpApi.cpp

#include "op_plugin/OpApiInterface.h"
#include <ATen/native/nested/NestedTensorMath.h>

namespace op_api {
at::Tensor to_padded_tensor(
    const at::Tensor& self,
    double padding,
    at::OptionalIntArrayRef output_size) {
  // 调用 PyTorch 原生实现
  return at::native::NestedTensor_to_padded_tensor_generic(
      self, padding, output_size);
}
} // namespace op_api

全流程与设备泛化说明： op_api::to_padded_tensor 只是薄封装，转换由 at::native::NestedTensor_to_padded_tensor_generic 完成：先在 CPU 读 nested_sizes / max_size 推导各子张量的 split_sizes（纯元数据，不碰数据），再对 get_buffer() 依次执行 dispatcher 算子——at::split_with_sizes 切分并还原子张量原始形状、at::constant_pad_nd 用 padding 补齐到 output_size（缺省取 max_size）、at::stack 沿第 0 维堆叠成 dense 张量返回。三个算子在 NPU 上分别落到 aclnnSplitWithSize、aclnnConstantPadNd、aclnnStack，并按 dispatch key 自动跟随 buffer 所在设备（NPU）执行，数据全程不离开 NPU，CPU 只做形状元数据推导，所以NPU的kernel 实现无需 .cpu() / .to(npu)。

4. pytorch —— torchnpugen/gen_backend_stubs.py

新增头文件生成函数：

def _nestedtensor_register_header() -> str:
    """生成嵌套张量注册所需的头文件,提供 op_api::to_padded_tensor 等声明"""
    headers = [
        '#include "op_plugin/OpApiInterface.h"',
    ]
    if not _is_aclnn_extension_codegen():
        headers.append('#include "op_plugin/OpInterface.h"')
    return "\n".join(headers) + "\n"

新增 extra_impls 生成函数（格式对齐已有 _quantized_extra_impls()）：

def _nestedtensor_extra_impls() -> List[str]:
    extra_impls = []
    extra_impls.extend([
        'm.impl("unbind.int", TORCH_FN(at::native::NestedTensor_unbind));',
        'm.impl("values", TORCH_FN(at::native::values_nested));',
        'm.impl("_nested_tensor_size", TORCH_FN(at::native::_nested_tensor_size));',
    ])
    if not _is_aclnn_extension_codegen():
        extra_impls.append('m.impl("to_padded_tensor", TORCH_FN(op_api::to_padded_tensor));')
    return extra_impls

SPECIAL_REGISTERS["nestedtensor"] 配置：

'nestedtensor': SpecialRegisterConfig(
    dispatch_key="NestedTensorPrivateUse1",
    filename="NestedTensorRegister",
    header=_nestedtensor_register_header(),
    extra_impls=_nestedtensor_extra_impls(),
),

aclnn_extension 路径隔离说明：to_padded_tensor 的注册（m.impl("to_padded_tensor", TORCH_FN(op_api::to_padded_tensor))）依赖 op-plugin 经 codegen 生成的 op_plugin/OpApiInterface.h 中的 op_api::to_padded_tensor 声明。而 aclnn_extension（cpp_extension）场景下该头由 npu_custom.yaml 生成，不含此声明，若无条件注册会产生 'to_padded_tensor' is not a member of 'op_api' 的编译错误。因此通过 _is_aclnn_extension_codegen() 条件生成：仅在非 aclnn_extension 场景注册 to_padded_tensor，aclnn_extension 路径隔离、不做实现；unbind.int / values / _nested_tensor_size 三个原生 NestedTensor 注册保持不变。

---

4. 测试设计

测试采用「两仓改动一起编包」的方式验证：pytorch 侧生成 NestedTensorRegister.cpp 并引入 op-plugin 头文件，op-plugin 侧提供 op_api::to_padded_tensor，两者一起编译后，用验证脚本在 NPU 上做端到端验证。

• 单元测试： 新增 op-plugin 仓测试 test/test_base_ops/test_nested_to_padded_tensor.py，参考 test_stack.py 等标准算子的 TestCase / run_tests 结构：用确定性构造 torch.arange(numel).reshape(shape) 生成 CPU / NPU 两份逐位一致的 NestedTensor，经 cpu_op_exec / npu_op_exec 分别调用 to_padded_tensor 后，用 assertRtolEqual 对比 CPU 与 NPU 输出。覆盖 9 个用例：

  用例	输入子张量形状	参数
  fp32 32-byte 对齐	(32, 32) × 3	缺省
  fp32 非 32-byte 对齐	(5, 7) × 3	缺省
  fp16	(32, 16) × 2	缺省
  非零 padding	(4, 6) × 3	padding=1.5
  指定 output_size	(4, 6) × 2	output_size=[3, 8, 8]
  3D	(2, 3, 4) × 2	缺省
  不规则形状	(2, 3), (3, 4), (1, 5)	缺省
  不规则形状 + padding	(2, 3), (3, 4), (1, 5)	padding=0.5
  1D	(3,), (5,), (7,)	缺省

• 端到端测试： 使用脚本 torchapi_id4942_to_padded_tensor.py，在 NPU 上分别构造「32-byte 对齐」与「非 32-byte 对齐」两类 NestedTensor，调用 nested.to_padded_tensor(padding=0.0)，校验结果无 error 字段、输出 dense 张量正确；脚本同时统计 driver / pta 显存占用。
• 兼容性门禁： op-plugin 仓运行 test_op_func_compatibility，确认新增 to_padded_tensor 已登记到基线 torch_npu_OpApi_schema_all.json。
• 编译验证： 两仓全量 codegen 后编译，确认 NestedTensorRegister.cpp 正常生成、op_api::to_padded_tensor 符号可正确链接。

---

5. 缺点和风险

• 风险： kernel 采用非 symint 签名注册到 symint schema，依赖 dispatcher 自动物化 SymInt；trace / export 等带符号尺寸场景需实际验证。
• 负面影响： 仅新增一个 operator 注册与一个 kernel，不影响既有功能。
• 实现成本： 代码量小，维护成本低。
• 兼容性： 保持 PyTorch 对外 schema 不变，无 breaking change。

应对措施：如出现 NPU 基础算子缺失，可在 torch_npu 侧补齐对应算子；如符号尺寸场景异常，再评估改为 OptionalSymIntArrayRef 签名并显式转换。

---

6. 现有技术

• 参考 PyTorch 上游 NestedTensor_to_padded_tensor_generic（ATen/native/nested/NestedTensorMath.h）。
• 参考 torch_npu 现有 SPECIAL_REGISTERS["nestedtensor"] 中 unbind.int / values / _nested_tensor_size 的注册范式。
• 参考 op-plugin 现有 op_plugin/ops/opapi 目录下算子的实现风格。

---

7. 未解决问题

• 非 symint kernel 注册到 symint schema 在带符号尺寸（trace / export）场景下的行为，待用例确认。单元测试使用具体形状，未覆盖符号尺寸；如异常，按第 5 节应对措施评估改为 OptionalSymIntArrayRef 签名并显式转换。

---

附录

• 参考资料链接：
  • PyTorch 上游 NestedTensorMath.h
• 术语表：
  • NestedTensor：嵌套张量。
  • to_padded_tensor：转换为 padding 后的稠密张量算子。
  • NestedTensorPrivateUse1：NestedTensor 在 NPU 上的私有 dispatch key。
  • SymInt：符号化整数，PyTorch 用以表示 trace 期间的符号尺寸。
  • op_api：op-plugin 中面向高层 API 的算子命名空间。
• 文档更新计划： 代码合入后补充 COMPATIBILITY.md 中该算子支持说明。