多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

catlass_cppgen 使用指南:用 Python 描述算子,快速生成高性能 C++ 核函数

catlass_cppgen 使用指南:用 Python 描述算子,快速生成高性能 C++ 核函数 catlass_cppgen 使用指南用 Python 描述算子快速生成高性能 C 核函数【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi在昇腾平台上手写一个高性能 GEMM 算子往往要在 Tile 形状、调度策略、多架构适配之间反复试错后处理逻辑还得手工往核函数里拼。catlass_cppgen 是面向 Ascend 的算子代码生成框架你用 Python 描述张量形状与算子参数它负责产出优化后的 C 核函数代码调参只改 Python 侧配置。它替你把哪些脏活干了一句话定位在 Python 里声明输入什么张量、算什么算子框架输出目标架构下可用的 C 核函数模板与参数绑定代码。核心能力点只描述不绑定OpTensor 用 shape、stride 和数据类型刻画输入无需先备好真实数据多架构覆盖构造算子时声明目标架构AtlasA2/A3、Ascend950生成对应实现调优参数外置两级 TileShape分块大小与 DispatchPolicy线程级调度策略都能在 Python 侧调后处理成图EVGEpilogue Visitor Graph把核函数尾部的偏置、激活等逻辑交给一张访问者图承载工作流四步从算子规划到核函数整条链路是固定的用 OpTensor 描述 A、B 等输入张量构造Gemm/GroupGemm算子对象调用get_kernels()拿到候选 Kernelkernel.tune(...)指定 Tile 形状和调度策略kernel.gen_kernel_template()输出 C 核函数模板gen_params_device()生成参数绑定代码相当于把算子当图纸交给一条流水线你给规格它出货。最小可运行示例from catlass_cppgen.op.gemm import Gemm from catlass_cppgen.common.op_tensor import OpTensor from catlass_cppgen.common.data_type import DataType from catlass_cppgen.catlass.layout.layout import RowMajor from catlass_cppgen.catlass.gemm_coord import GemmShape from catlass_cppgen.catlass.arch.arch import Arch # 1. 只声明张量形状、步长和类型不绑定真实数据 a OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT) b OpTensor.from_shape_stride((256, 384), (384, 1), DataType.FLOAT) # 2. 声明算子拉出候选 Kernel 列表 gemm Gemm(atlas_archArch.Ascend950, elementDataType.FLOAT, layoutRowMajor, Aa, Bb) kernel gemm.get_kernels()[0] # 3. 调优两级 Tile 形状后生成 C 核函数模板 kernel.tune(GemmShape(128, 256, 64), GemmShape(128, 256, 64)) print(kernel.gen_kernel_template())能力地图按场景挑 Kernel基础 GEMM覆盖二维场景。BasicMatmulKernel处理最普通的 A(M,K) × B(K,N)固定 alpha1、beta0可挂可选 Bias输入升到三维、各 batch 共享维度时用BatchedMatmulKernel。K 维很长时计算量会堆在少数尾块上切 K 类 Kernel 就是把 K 维分摊到多核并行MultiCoreSplitkMatmulKernel是标准多核 Split-KTailMultiCoreSplitkMatmulKernel针对尾块做了优化变体StreamkMatmulKernel用 Stream-K 调度动态分块避免最后一个 tile 空等。Group GEMM面向 M 轴切分的分组矩阵乘GroupedMatmulSliceMKernel接收三维的 A/B 和一个 INT64 的groupListVectorLayout 向量各组 M 维度可以不同。EVG 后处理解决核函数收尾那几行激活、偏置不好维护的问题。BasicMatmulTlaVisitorKernel支持 EVG你在 Python 里写一个 epilogue 函数可用 add/sub/mul/div 四则、relu/silu/sigmoid/leakyRelu/Prelu 激活、max/min 选择、cast 类型转换、constant 常量多节点可自由串联也支持行广播比如 (1,256) 的 bias 广播到 (128,256) 的 accum。安装与目录速览从源码构建后安装 wheelpip install build python -m build pip install dist/catlass_cppgen-*.whl开发场景直接装编辑模式pip install -e .仓库分几块catlass_cppgen/是主体op算子基类、kernel特化类、common张量与类型、catlass架构与布局组件docs/是 API 文档tests/按特性分目录放单测。进阶定向选 Kernel、调优与 EVG 拼法不想默认取第一个 Kernel 时按类型定向拿from catlass_cppgen.kernel.gemm import BasicMatmulKernel kernel find_kernel_by_type(kernels, BasicMatmulKernel)tune()接收两级 TileShape可追加dispatch_policy指定调度如MmadPingpong(arch_tagArch.Ascend950)。EVG 则通过算子构造时的evg_config传入fn_src放 epilogue 函数源码字符串example_inputs给出每个变量的张量规格evg_config { fn_src: def epilogue(accum, bias):\n return relu(accum bias), example_inputs: { accum: OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT), bias: OpTensor.from_shape_stride((1, 256), (256, 1), DataType.FLOAT), }, } gemm Gemm(atlas_archArch.Ascend950, evg_configevg_config, Aa, Bb) assert gemm.get_kernels()[0].is_support_evgAPI 细节看仓库内文档kernel_api.md、evg_api.md、optensor_api.md。谁适合用如果你要在昇腾上批量产出或调优 GEMM 类高性能算子又不想把时间耗在核函数手拼上catlass_cppgen 值得装进工具链。入门建议从 kernel_api.md 读起再对照tests/op/下的用例改参数验证行为。【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表