CANN/catlass Torch测试框架

Torch Catlass【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlasstests/optest是 CATLASS 示例算子接入 PyTorch 的测试框架。框架将 AscendC/CATLASS kernel 通过 C extension 注册为torch.ops.catlass.*并由 Python 包torch_catlass提供调用入口。详细设计见 docs/design.md。架构总览框架按职责分为五层Python API (torch_catlass.ops.*) - Python loader (torch_catlass/__init__.py) - PyTorch C extension (src/) - Kernel adapter (src/include/template/) - Kernel 实现 (kernels/)Python 层负责用户接口与动态库加载。C extension 层负责torch.ops.catlass.*注册和 NPU dispatch。adapter 层负责 Tensor 参数到 kernel ABI 参数的转换。kernel 层负责 prebuilt kernel 与 JIT kernel 的实现和执行。目录结构tests/optest/ ├── CMakeLists.txt ├── build.sh ├── pyproject.toml ├── docs/ │ └── design.md ├── include/ │ ├── catlass_kernel.h │ ├── catlass_kernel_jit.h │ └── catlass_kernel_prebuilt.h ├── src/ │ ├── catlass_torch.cpp │ └── include/common/ ├── kernels/ │ ├── 00_basic_matmul/ │ ├── include/ │ └── jit/ ├── torch_catlass/ │ ├── __init__.py │ └── ops/ ├── tests/ │ └── test_00_basic_matmul.py └── utils/昇腾与架构识别说明本项目面向昇腾 NPU 的 CANN/Ascend C 开发栈相关术语和 SoC 信息以昇腾公开开发文档为准。Python loader 侧通过torch_npu.npu.get_device_name()映射 arch id用于加载torch_catlass/lib/arch/下的预构建库。JIT 编译侧通过GetCurrentNPUArch()AscendC 平台 API获取当前 SoC 对应架构用于设置编译参数。当前代码中的映射Ascend910B.*/Ascend910_93-2201Ascend950-3510。公共 ABI 与接口分层include/catlass_kernel.h仅作为聚合头接口按模块拆分为catlass_kernel_jit.hmatmul/gemm/gemv 等 JIT 接口参数语义为TParams Params。catlass_kernel_prebuilt.hflash-attention/conv/mla 等 prebuilt 接口。其中TParams表示编译期模板参数dtype/layout/transpose 等。Params表示运行期参数shape、buffer address 等。构建与测试环境要求Python 3.11PyTorch torch-npu与本机 CANN/驱动版本匹配CMake 3.16可用昇腾 NPU 设备与环境编译bash build.sh测试pytest tests/ -v使用示例import torch import torch_catlass a torch.randn(1024, 1024, dtypetorch.float16, devicenpu) b torch.randn(1024, 1024, dtypetorch.float16, devicenpu) c torch_catlass.basic_matmul(a, b, outDTypefloat16)开发指南新增算子流程在include/catlass_kernel_jit.h或include/catlass_kernel_prebuilt.h预留 ABI带 example 编号说明。在kernels/实现 kernelJIT template 或 prebuilt。在src/增加 C op 适配和注册。在torch_catlass/ops/增加 Python wrapper。在tests/增加 pytest 用例。【免费下载链接】catlass本项目是CANN的算子模板库提供NPU上高性能矩阵乘及其相关融合类算子模板样例。项目地址: https://gitcode.com/cann/catlass创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考