SuperAPI AI 工具集

tilelang-vector-ascend-ops-migration

编程与开发

查看项目

更新于 10-16 05:00

技能介绍

该技能用于把 TileLang 算子从 GPU(CUDA)实现迁移到华为昇腾 NPU,通常在 gpu/ 目录分析已有实现,生成对应的 npu/ 目录代码(若用户未指明存放位置则新建 npu 目录并提示位置),并自动适配 NPU 硬件约束与 API 差异。

主要能力与用法:分析 GPU 算子的输入输出形状、数据类型、核心计算逻辑、内存分配与并行策略;按迁移规则改写环境变量、装饰器与 Kernel 声明;进行数据类型和 API 映射;检查 UB 内存、对齐、分核等硬件约束;生成含 PyTorch 封装与测试的 NPU 代码。关键约定包括:设置 os.environ['TILELANG_ASCEND_MODE'] = 'Developer',使用 @tilelang.jit(target="npuir"),T.Kernel(..., is_npu=True),动态 shape 用 T.symbolic,T.fill→T.clear,clear_accum→initC,transpose→a_transpose/b_transpose,T.vconv→T.vcast,T.vcmp 签名 (src0, src1, dst, cmp_op),标量运算需经 T.vbrc 广播,T.copy 用起始索引加 size 处理边界。

适用场景:将适配 GPU 的 TileLang vector 算子迁移为适配 NPU 的算子。限制与注意:必须分块处理,UB 总量需小于 96KB(建议 85KB 以下),维度需对齐且输入 contiguous,测试须符合实际数值范围并以误差小于 1e-3 判定;调试前需先确认 tilelang-ascend 环境是否配置。迁移时需查阅 docs/TileLang-NPU-API-Reference.md、docs/GPU-To-NPU-Migration-Methods.md 及 references 下调试与硬件优化文档;涉及已知精度问题 API 的算子应绕过或报告无法迁移。具体环境配置、依赖版本与更完整的调试方法暂未提供。