
三步生成 CATLASS 高性能内核catlass_cppgen 代码生成框架上手指南【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi手写一个 GEMM 内核要面对的是一长串模板参数Tile 尺寸、调度策略、架构宏、布局描述配错任何一处都是难以定位的编译报错参数组合还会随矩阵规模不断膨胀。catlass_cppgen 是一个基于 Python 生成 CATLASS 高性能算子的代码生成框架——你在 Python 里声明张量形状、数据类型与目标架构它替你产出可编译的 C 内核代码。项目定位Python 侧声明C 侧落地一句话概括catlass_cppgen 把“描述算子”和“实现算子”拆开你只管前者。它有三个核心能力值得记住声明式参数定义用 OpTensor 描述输入的形状、步幅与数据类型即可无需绑定真实数据也不用关心底层内存一键内核生成从算子对象拿到 Kernel 后get_kernels()、tune()、gen_kernel_template() 三次调用就完成“规划 → 调优 → 出码”全流程后处理可扩展借助 EVGEpilogue Visitor Graph后处理访问者图把激活、Bias 相加、类型转换等计算挂到矩阵乘的尾声阶段无需手写融合代码。架构层面框架通过 Arch 枚举声明目标代际覆盖 AtlasA2/A3、Ascend950 等多种硬件架构。能力全景能生成哪些算子GEMM 家族是主力按优化策略分档每档对应一个 Kernel 特化类基础矩阵乘BasicMatmulKernelA、B 为二维输入固定 alpha1.0、beta0.0支持可选 Bias对应最常规的稠密 matmul批处理矩阵乘BatchedMatmulKernelA、B 为三维batch, M, K/batch, K, N所有批次共享同一组维度适合 batch 推理多核 Split-KMultiCoreSplitkMatmulKernel沿 K 方向多核切分K 很大而 M/N 偏小时用它提升效率同样支持可选 Bias尾块场景可用优化变体 TailMultiCoreSplitkMatmulKernelStream-KStreamkMatmulKernel采用 Stream-K 调度策略摊平负载适合负载不均的矩阵形状EVG Visitor 矩阵乘BasicMatmulTlaVisitorKernel面向 EVG 后处理框架的 matmul 变体。分组矩阵乘方面目前提供沿 M 轴切分的 GroupedMatmulSliceMKernel用于在一次调用内完成多组 M 维度各异的矩阵乘。后处理能力由 EVG 承载写法接近一段 Python 小函数可用构件包括二元运算 add / sub / mul / div如accum bias、激活函数 relu / silu / sigmoid / leakyRelu / Prelu、比较选择 max / min、类型转换 cast、常量 constant。多个节点可以串联成组合计算并支持行广播——例如偏置向量沿行方向展开到整块累加器。安装与快速上手安装步骤 拿到源码后任选一种方式装进环境开发模式改动即时生效pip install -e .构建分发包先pip install build再执行python -m build在dist/下得到.whl与.tar.gz然后pip install dist/catlass_cppgen-*.whl或以普通方式安装pip install .五行代码跑通基础 GEMM最小流程是描述输入张量 → 建立 Gemm 算子 → 取出 Kernel → 生成内核代码。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.arch.arch import Arch a OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT) b OpTensor.from_shape_stride((256, 384), (384, 1), DataType.FLOAT) gemm Gemm(atlas_archArch.Ascend950, Aa, Bb) kernel gemm.get_kernels()[0] print(kernel.gen_kernel_template())get_kernels() 返回若干候选 Kernel你可以按类型挑选例如指定 BasicMatmulKernel也可以直接取第一个。拿到 kernel 后gen_kernel_template() 输出核函数模板gen_params_device() 负责参数绑定的代码生成两者配合即是一份可用的内核。进阶玩法Group GEMM 与 EVG 后处理 两组进阶用法各由一个关键对象表达。Group GEMM先构造一个 INT64 类型的 groupList 张量VectorLayout(4)、shape 为 (4,)声明分组规模再传给 GroupGemm(atlas_arch..., Aa, Bb_3d, groupListgroupList) 建立算子取回 kernels 后照常用 tune(GemmShape(256, 256, 256), GemmShape(256, 256, 64)) 做 tiling 调优。EVG 后处理用一个函数头fn_src加一组 name:tensor 示例输入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), result: OpTensor.from_shape_stride((128, 256), (256, 1), DataType.FLOAT), }, } kernel Gemm(atlas_archArch.Ascend950, evg_configevg_config, Aa, Bb).get_kernels()[0] assert kernel.is_support_evgis_support_evg 为 True 即表示该 Kernel 支持 EVGKernel 侧还可以用 to_evg() 将后处理配置绑定上去再按需 tune 调整 Tile 形状。调优要点TileShape、DispatchPolicy 与架构标签⚙️ tune() 是调优的统一入口接收三类信息两级 GemmShapeTileShape如 GemmShape(128, 256, 64)分别描述宏块与原子粒度的 Tile 形状可选的 dispatch_policy调度策略对象如 MmadPingpong(arch_tagArch.Ascend950)arch_tag 架构标签经 Arch 枚举声明目标代际AtlasA2/A3、Ascend950 均可选。矩阵维度不同合理的 Tile 与策略组合差异很大这正是显式调优存在的意义。资源导航文档、测试与源码入口 想深入某个环节按路径查docs/kernel_api.mdKernel API 基础文档覆盖调优与特性查询docs/optensor_api.mdOpTensor API张量声明的完整输入方式docs/evg_api.mdEVG 后处理 API 参考tests/单元测试按 catlass、common、op 三个维度组织读测试用例是理解 API 最快的一条路源码侧catlass_cppgen/op/ 是算子入口gemm.py、group_gemm.pycatlass_cppgen/kernel/ 存放各 Kernel 特化类EVG 相关实现位于 catlass_cppgen/catlass/evg。回到完整链路Gemm / GroupGemm 负责算子规划get_kernels() 取出调优对象tune() / to_evg() 完成配置最后由生成方法输出 C——从声明到落地就差这一次调用链。【免费下载链接】YiA series of large language models trained from scratch by developers 01-ai项目地址: https://gitcode.com/GitHub_Trending/yi/Yi创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考