CANN TileLang转AscendC指南

CANN TileLang转AscendC指南 TileLang 设计转换到 AscendC Kernel 关键原则【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills本文档讨论如何将已经完成的 TileLang 设计系统地转换为等价的 AscendC 实现。本文中的 DSL 特指 TileLang。后文中的 tiling、绑定层、主 kernel 类、子模块拆分与同步关系都应以 TileLang 尤其是 tile-level 设计为直接来源。先看 Mapping再查 API 文档转换时应先阅读references/TileLang-AscendC-API-Mapping.md先确认 TileLang API 到 AscendC API 的映射关系再去阅读asc-devkit/docs/目录下对应的具体 API 文档。 知识库入口api_reference/INDEX.md禁止在 C 中直接调用 torch / ATen 计算接口pybind11.cpp中禁止使用torch::*、torch::nn::functional::*、ATen库或任何at::*计算接口来实现或替代核心计算包括但不限于at::einsum、at::matmul、at::softmax、at::bmm等。绑定层只负责参数检查、输出与 workspace 分配、tiling 填充和 kernel launch不允许把核心计算留在 C/ATen 侧。TileLang 到 AscendC 转换总览一个从 TileLang 设计转换得到的 AscendC 实现通常包含 4 部分Host 侧准备xxx_tiling.hpybind11.cpp这两部分共同构成 AscendC 的 host 侧准备逻辑通常需要结合 TileLang kernel 的 host 信息一起整理。xxx_tiling.h负责定义 shape、block size、tile size、workspace 深度等参数pybind11.cpp负责 Python 接口、输入校验、输出与 workspace 分配、tiling 构造和 kernel launch。公共工具如kernel_common.h、workspace_queue.h、matmul_tile.h,vector_tile.h提供调度、数据搬运、workspace 管理等通用能力。Kernel 入口xxx.cpp定义__global__ __aicore__kernel 和extern Claunch 函数。主 Kernel 类与计算子模块一个或多个*.h主Kernel类负责Init()/Process()主流程管理 GM tensor、调度和流水。若 TileLang 中存在多个T.prim_func将对应的主Kernel类拆到多个独立头文件中例如xxx_merge_n_kernel.h、xxx_single_row_kernel.h。若算子属于 C/V 融合算子或者 TileLang 设备侧存在多个有明确职责分工的Scope则可在这一部分下继续按计算阶段拆分子模块例如matmul.h、leakyrelu.h通常每个Scope对应一个子模块职责应与原 TileLang 设计中的计算阶段一一对应。对于纯 Vector 算子或者虽然有 host / queue / buffer 管理但设备侧只有单个 Vector 计算阶段 / 单个 VectorScope的简单算子主Kernel类本身通常就承载全部计算逻辑不再额外拆分子模块。T.prim_func到 AscendC Kernel 的映射规则TileLang 中有多少个T.prim_funcAscendC 侧就至少要有多少个独立的 kernel 实现单元。不要把多个T.prim_func折叠进同一个Kernel类里再靠运行时分支区分。具体要求每个T.prim_func都应对应一个独立的主Kernel类。每个T.prim_func都应对应至少一个独立的__global__ __aicore__kernel 入口和一个匹配的extern Claunch 函数。如果同一个T.prim_func需要按 dtype 分成多个实现例如 fp16 / fp32 / int8则可以在该prim_func之下派生出更多extern入口但主Kernel类的个数仍应首先与T.prim_func的个数对齐。Host 侧pybind11.cpp负责根据 shape、dtype 或其他 trace-time 条件选择调用哪个extern入口这种选择逻辑不应反向合并掉prim_func级别的结构差异。例如若 TileLang 提供merge_n和single_row两个T.prim_func则 AscendC 至少应有两个主Kernel类、两个__global__ __aicore__入口和两个extern Claunch 函数若两者还各自支持多种 dtype则extern数量可以更多但主Kernel类仍至少是两个。第一章Host 侧准备摘要完整实现细节与代码示例见references/dsl2Ascendc_host.md。要点Tiling 参数一致性所有 kernel 组件Cube/Vector/Host必须使用同一组 baseM/baseN/baseK 常量参数不匹配会导致错误的内存访问。Tiling Struct在 Host 侧预计算nTiles、nTilesPerH等派生量写入 tiling struct避免 kernel 里重复除法。绑定层职责pybind11.cpp负责参数检查、输出分配、workspace 分配、tiling 构造、kernel launch。绑定函数只接收 DSL 显式输入张量不接收输出和 workspace。模块名推荐_op_name_ext不要与任务目录同名。Workspace只要 DSL 声明了 workspace 参数或workspace_idx就必须分配 workspace。第二章公共工具摘要1. tile 层公共工具matmul_tile.h/vector_tile.h这一类文件主要承载 tile 级的数据搬运、分块计算封装和局部流水组织是把 DSL 里的 tile-level 设计落到 AscendC 时最常见的公共工具。纯 Vector 算子通常需要构建vector_tile.h一类工具可以参考rms_norm的 kernel 写法例如archive_tasks/rms_norm/kernel/下对vector_tile.h的使用纯 Cube 算子通常需要构建matmul_tile.h一类工具可以参考matmul_leakyrelu的 kernel 写法例如archive_tasks/matmul_leakyrelu/kernel/matmul_tile.hC/V 融合算子通常两类工具都要结合具体分工一起看2.workspace_queue.h与跨核同步这一部分主要对应 AIC / AIV 之间通过 workspace 传递中间结果、并通过 cross-core flag 建立 producer / consumer 协同的场景常见于 C/V 融合算子。纯 Vector 算子和纯 Cube 算子默认不要阅读references/dsl2Ascendc_cross_core_sync.mdC/V 融合算子或 DSL 中出现跨核协同 / workspace 生产者-消费者关系 / cross-core flag必须阅读references/dsl2Ascendc_cross_core_sync.md第三章Kernel 入口摘要KERNEL_TYPE 与 DSL vec_num 对应关系DSLvec_numKERNEL_TYPE每个 block 组成1KERNEL_TYPE_MIX_AIC_1_11 AIC 1 AIV2KERNEL_TYPE_MIX_AIC_1_21 AIC 2 AIV代码结构要求Process()封装工作负载循环调用CopyInX/ComputeX/CopyOutX每个阶段函数定义为__aicore__ inline第四章主 Kernel 类与计算子模块摘要完整实现细节与代码示例见纯 Vector 算子references/dsl2Ascendc_compute_vector.md纯 Cube 算子references/dsl2Ascendc_compute_cube.mdC/V 融合算子先看references/dsl2Ascendc_compute_cv.md再结合references/dsl2Ascendc_compute_cube.md和references/dsl2Ascendc_compute_vector.md常见陷阱速查表问题症状解决方法local UB buffer 该用TQue却写成TBuf结果系统性错误或生命周期错乱在T.serial中的输入/输出 buffer 用TQueTQue depth0 用了返回值形式 API编译报错VECIN/VECOUT 必须用引用形式TBuf DataCopy 后缺少 PipeBarrier结果随机错误插入PipeBarrierPIPE_MTE2()Fixpipe 未同步流同步超时CrossCoreSetFlag 使用PIPE_FIX内层循环中重复 DeQue结果错误每个外层迭代只 DeQue 一次Fixpipe dstStride baseNtile 间数据覆盖dstStride 应设为完整行宽 NWholeReduceMax 参数错误编译报错改用ReduceMax(dst, src, workBuf, count)使用不存在的 Divs编译报错改用Muls(dst, src, 1.0f/scaleVal, count)scan 算子cumsum 等fp16 2Ddim0验证失败NPUtorch.cumsum对 fp16 2Ddim0使用非确定性并行扫描参考输出本身不一致见下方「Scan 类算子转译注意事项」monkey-patchtorch.cumsum 混合 accumulation 精度Scan 类算子cumsum / cumprod 等转译注意事项1. NPUtorch.cumsumfp16 2Ddim0的已知 bug当算子属于 scan 类inclusive scan、prefix sum 等且参考实现调用torch.cumsum时必须注意NPU 上存在以下已知问题对float16、2D tensor、沿dim0strided scan的场景torch.cumsum内部使用非确定性并行扫描导致小 tensor多次运行结果随机波动~0.1-0.5% mismatch大 tensor系统性偏离正确值~10%同一 tensor 沿dim1contiguous scan则使用确定性串行扫描结果稳定。2.model_new_ascendc.py中的标准 Workaround转译 scan 类算子时应在model_new_ascendc.py中实施以下模式以 cumsum 为例步骤 AMonkey-patchtorch.cumsum在模块顶部拦截torch.cumsum将 2D fp16dim0自动转译为cumsum(x.T, dim1).T迫使参考模型走稳定的 contiguous scan 路径_original_cumsum torch.cumsum def _patched_cumsum(input, dim, *args, **kwargs): if input.dim() 2 and input.dtype torch.float16 and dim in (0, -2): return _original_cumsum(input.T, dim1).T return _original_cumsum(input, dim, *args, **kwargs) torch.cumsum _patched_cumsum步骤 BKernel 内部必须实现混合 accumulation 精度硬性要求仅在 kernel 中固定使用 fp32 accumulation 或固定使用 fp16 accumulation 都无法覆盖全部 case必须在 kernel 中通过 tiling 参数如useFp32Acc支持两种模式切换并在 Python wrapper 中根据 scan 长度L动态选择小 tensorscan 长度L 512NPU 参考走纯 fp16 串行扫描kernel 必须切换为fp16 accumulation。实现方式每步先将 fp32 acc cast 到 fp16再 cast 回 fp32 与输入相加确保逐元素舍入行为与参考一致。大 tensorL 512NPU 参考在 fp16 路径下仍表现出类似 fp32 的行为kernel 使用fp32 accumulation全程 fp32 累加最后统一 cast 到 fp16利用大数值下rtol容忍度较宽的特点通过验证。wrapper 中判断逻辑示例is_last_dim dim_pos x.ndim - 1 is_4d_non_last x.ndim 4 and dim_pos 1 is_2d_dim0 x.ndim 2 and dim_pos 0 use_fp32 (x.dtype torch.float16) and not ( is_last_dim or is_4d_non_last or (is_2d_dim0 and x.shape[dim] 512) )步骤 CFp16 输出 cast 模式硬性要求将 fp32 acc cast 到 fp16 时必须使用AscendC::RoundMode::CAST_NONE截断。经验证CAST_ROUND会导致 fp16 小 tensor case 产生额外 ~0.1-0.5% mismatch而CAST_NONE与 NPU PyTorch 的舍入行为最接近。3. 通用化要点bf16 与 fp16 的区别经实测NPUtorch.cumsum对bfloat16的 2Ddim0场景没有非确定性并行扫描 bug多次运行结果稳定且dim0与dim1-of-transpose结果完全一致。因此 monkey-patch 和混合 accumulation 精度策略仅针对 fp16bf16 可保持常规 fp32 accumulation 实现。但 fp16 输出 cast 为CAST_NONE的建议对 bf16 同样适用。以上模式不仅限于cumsum任何参考实现依赖torch.cumsum的 scan 类算子如cumprod在 fp16 2Ddim0场景下均应检查并应用相同 workaround。若 TileLang 设计中的 scan 算子后来被转译为 AscendC转译时应确保 kernel 支持useFp32Acc切换并在 wrapper 中根据原始 stride / shape 信息动态选择精度策略。若算子本身不是 scan 类但内部包含 prefix sum 作为子步骤如某些 sort / argsort 实现也应审视该子步骤是否触发相同的 NPU bug。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考