ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

CUTLASS Python 接口全解析:从 Python 中编译与运行高性能 GEMM 内核

CUTLASS Python 接口全解析:从 Python 中编译与运行高性能 GEMM 内核 CUTLASS Python 接口全解析从 Python 中编译与运行高性能 GEMM 内核【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass本文以仓库 python/README.md 为主线系统讲解 CUTLASS 提供的两个 Python 包——cutlass_cppgenCUTLASS Python 接口与cutlass_library内核枚举与 C 代码生成工具。你将掌握如何用几行 Python 代码声明、编译并运行 GEMM 等算子、如何导出 PyTorch CUDA 扩展以及这些高层 API 背后与仓库源码的对应关系从而在深度学习框架中快速落地 CUTLASS 内核。CUTLASS 完整 GEMM 分解层次图一、两个 Python 包的分工python/目录下包含两个与 CUTLASS 相关联的 Python 包cutlass_cppgenCUTLASS Python 接口允许直接在 Python 中编译并运行 CUTLASS 内核。注意它此前名为cutlass改名是为了与 CuTe Python DSL见 python/CuTeDSL区分开。cutlass_library用于枚举并发射emitCUTLASS 内核 C 代码的工具集被 CUTLASS CMake 系统用来构建一个可供 CUTLASS profiler 剖析的内核库。两者的安装入口分别是 python/setup_cutlass.py 与 python/setup_library.py。从源码结构看cutlass_cppgen是面向最终用户的高层接口而cutlass_library是其底层依赖——在 python/cutlass_cppgen/init.py 中直接import cutlass_library并从中重新导出DataType、LayoutType、MathOperation、TileDescription、EpilogueScheduleType等核心类型。二、快速上手三行代码跑通一个 GEMMcutlass_cppgen的入口是cutlass模块旧名称保留为包名最简用法如下import cutlass import numpy as np plan cutlass.op.Gemm(elementnp.float16, layoutcutlass.LayoutType.RowMajor) A, B, C, D [np.ones((1024, 1024), dtypenp.float16) for i in range(4)] plan.run(A, B, C, D)这段代码完成了四件事声明一个 FP16、RowMajor 布局的 GEMM 算子由接口自动挑选一套合理的内核默认配置编译出 CUTLASS 内核用 NumPy 张量也可以是 PyTorch / CuPy 张量执行一次 1024×1024 的矩阵乘。运行后会得到与直接手写 CUTLASS C 模板一致的D A * B C语义。在 python/cutlass_cppgen/op/gemm.py 的模块文档中官方对Gemm接口的定位有明确说明它让用户无需指定大量配置参数即可实例化、编译、运行 CUTLASS GEMM同时明确指出该接口不以最优性能为目标——若追求极致性能需要逐个指定并调优每个配置参数。2.1 数据类型的绑定与等价构造数据与布局类型在Gemm对象构造时被绑定生命周期内不可更改。以下两种构造方式是等价的# 显式指定各操作数类型可分别设置 element_A/element_B/element_C/element_D/element_accumulator plan cutlass_cppgen.op.Gemm(element_Atorch.float32, element_Btorch.float32, element_Ctorch.float32, element_Dtorch.float32, element_accumulatortorch.float32, layoutcutlass_cppgen.LayoutType.RowMajor) # 简写当所有操作数类型一致时只用 element 参数 plan cutlass_cppgen.op.Gemm(elementtorch.float32, layoutcutlass_cppgen.LayoutType.RowMajor)运行时传入的张量形状可以不同只要数据类型匹配例如A0 torch.rand((128, 256), devicecuda) B0 torch.rand((256, 64), devicecuda) C0 torch.zeros((128, 64), devicecuda) D0 torch.zeros((128, 64), devicecuda) plan.run(A0, B0, C0, D0) A1 torch.rand((32, 128), devicecuda) B1 torch.rand((128, 256), devicecuda) C1 torch.zeros((32, 256), devicecuda) D1 torch.zeros((32, 256), devicecuda) plan.run(A1, B1, C1, D1)2.2 编译与执行解耦Gemm支持把内核编译与执行解耦适合先编译、后多次运行的场景plan cutlass_cppgen.op.Gemm(elementnp.float32, layoutcutlass_cppgen.LayoutType.RowMajor) plan.compile() # 做其他工作... plan.run(A0, B0, C0, D0) # 做其他工作... plan.run(A1, B1, C1, D1)2.3 融合激活函数与异步执行在 pre-SM90 内核上可以轻松融合逐元素激活函数如 ReLUplan cutlass_cppgen.op.Gemm(elementnp.float32, layoutcutlass_cppgen.LayoutType.RowMajor) plan.activation cutlass_cppgen.epilogue.relu算子也可异步执行run()返回的参数对象支持显式同步plan cutlass_cppgen.op.Gemm(elementnp.float32, layoutcutlass_cppgen.LayoutType.RowMajor) args plan.run() # 做其他工作... args.sync()三、设计目标易用性优先CUTLASS Python 接口把易用放在首位具体体现为高层算子接口声明一个算子只需少量参数智能默认配置根据已指定参数自动选择合理的默认内核配置配置枚举为用户枚举在给定环境下已知可用的配置组合友好的错误信息尽可能抛出描述性强的 Python 运行时异常而不是 C 编译期错误框架导出简化将 CUTLASS 内核导出为框架扩展如 PyTorch CUDA 扩展的流程。3.1 明确的非目标Non-goals接口官方声明不打算做以下三件事理解这一点对合理使用至关重要不负责选择最优内核配置接口倾向于易用性而非最大可配置性其默认参数选择在部分场景下达不到最高性能。追求极致性能的用户应当通过剖析不同参数组合自行挑选或改用内置内核选择启发式的库如 cuBLAS。不作为 CUTLASS 内核的高性能容器接口不追求最小化 Python 函数围绕内核运行的开销。需要部署内核的用户应当直接使用接口发射出的 C 代码或使用 CUTLASS 发射器自动创建框架扩展如 PyTorch CUDA 扩展。不充当 Python 到 CUDA 内核的 JIT 编译引擎接口让 CUTLASS 能在 Python 代码中被使用可以被框架用作 JIT 编译的底层但自身不立志成为这样的框架。3.2 与 PyCUTLASS 的对比CUTLASS Python 接口构建在 CUTLASS 的 PyCUTLASS 库之上。PyCUTLASS 允许以与 CUTLASS C 接口几乎相同的配置空间来声明、编译和运行 GEMM、卷积与 grouped GEMM 算子——灵活性接近 C 接口但代价是必须像填写 C 模板参数那样指定大量配置参数。相比之下CUTLASS Python 接口提供更高级的声明、发射与编译 API无需穷举定义模板参数。这一设计在 python/cutlass_cppgen/library_defaults.py 中体现为KernelsForDataType类它按操作数 A/B/C 的 alignment 分组维护已知可用的内核集合并提供default_operation()等方法在未指定具体内核时返回默认选择见 python/cutlass_cppgen/library_defaults.py。此外该文件维护的生成器计算能力列表为[50, 60, 61, 70, 75, 80, 90, 100]python/cutlass_cppgen/library_defaults.py即接口覆盖 SM50 到 SM100 的内核生成能力。四、当前支持的功能截至本文所述版本cutlass_cppgen内部版本号为4.7.0见 python/cutlass_cppgen/init.py接口支持以下算子GEMM通用矩阵乘融合逐元素 epilogue 的 GEMM如 ReLU仅 pre-SM90 内核Stream K swizzling仅 pre-SM90 内核Grouped GEMM仅 pre-SM90 内核。对应的算子实现在 python/cutlass_cppgen/op 目录下gemm.pyGemm、gemm_grouped.pyGroupedGemm、conv.pyConv2d、Conv2dFprop、Conv2dDgrad、Conv2dWgrad并由 python/cutlass_cppgen/init.py 统一导出。五、环境准备与安装5.1 推荐环境NGC PyTorch Docker 容器官方推荐通过 NGC PyTorch Docker 容器使用该接口docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.08-py3 -p 8888:8888接口已在CUDA 11.8、12.0、12.1以及Python 3.8 与 3.9上经过测试。5.2 可选环境变量安装前可选择性设置以下环境变量CUTLASS_PATH克隆的 CUTLASS 仓库路径CUDA_INSTALL_PATHCUDA 安装路径。若未设置安装过程会按以下规则推断CUTLASS_PATH若本地安装则为当前目录上一级$(pwd)/..若通过 pip 安装则为cutlass_library所在位置的source目录。源码中该逻辑为CUTLASS_PATH os.getenv(CUTLASS_PATH, cutlass_library.source_path)python/cutlass_cppgen/init.pyCUDA_INSTALL_PATH$PATH中第一个nvcc所在目录即which nvcc | awk -F/bin/nvcc {print $1}。源码中_cuda_install_path_from_nvcc()通过which nvcc定位并校验该目录存在python/cutlass_cppgen/init.py。注意cuda-python的版本必须与CUDA_INSTALL_PATH中的 CUDA 版本匹配。接口在初始化 CUDA 上下文时会调用check_cuda_versions()校验 Python CUDA 版本不低于 nvcc 版本否则抛出异常python/cutlass_cppgen/init.py。此外可通过CUTLASS_CUDA_DEVICE_ID环境变量指定使用的 GPU 设备号python/cutlass_cppgen/init.py。5.3 安装方式方式一PyPI 稳定版。官方稳定版本通过nvidia-cutlassPyPI 包发布。任何其他名为cutlass的包均与 NVIDIA CUTLASS 无关安装时务必认准包名pip install nvidia-cutlass方式二源码安装。进入 CUTLASS 仓库根目录执行pip install .方式三可编辑安装开发模式。若希望修改接口源码后改动即时生效pip install -e .5.4 安装验证安装成功后运行以下代码验证使用小尺寸张量以加快验证速度import cutlass import numpy as np plan cutlass.op.Gemm(elementnp.float16, layoutcutlass.LayoutType.RowMajor) A, B, C, D [np.ones((128, 128), dtypenp.float16) for i in range(4)] plan.run(A, B, C, D)5.5 运行时内存池可选在 Python 3.9 且安装了rmm的环境下接口会自动启用 RMM 内存池use_rmm True见 python/cutlass_cppgen/init.py并通过get_memory_pool()惰性创建一个初始 1 GiB、上限 4 GiB 的 CUDA 内存池python/cutlass_cppgen/init.py用于管理内核运行时的设备内存分配。六、导出深度学习框架 CUDA 扩展接口提供了将 CUTLASS 内核导出为深度学习框架 CUDA 扩展的工具。目前支持导出PyTorch CUDA 扩展类似模式也可推广到其他框架。发射器实现位于 python/cutlass_cppgen/emit/pytorch.py由 python/cutlass_cppgen/init.py 导入。当前可导出为 PyTorch CUDA 扩展的算子包括GEMMGrouped GEMMConv2d官方示例 notebookgrouped GEMM 导出到 PyTorch 扩展位于 examples/python/deprecated/02_pytorch_extension_grouped_gemm.ipynb当前仓库中已移入deprecated目录示例思路仍具参考价值。七、Jupyter Notebook 示例使用 CUTLASS Python 接口的 Jupyter notebook 示例集中在 examples/python含CuTeDSL、advanced_compiler_control、deprecated三个子目录。从python/目录启动jupyter-lab ../examples/python八、构建接口文档接口使用 Sphinx 生成文档。构建前先安装附加依赖sudo apt-get install pandoc pip install --upgrade Sphinx furo pandoc myst-parser sphinx-copybutton nbsphinx nbsphinx-link sphinx-inline-tabs构建文档的前提是先按上文安装一节安装好 CUTLASS Python 接口然后执行sphinx-apidoc -o docs_src/source/ cutlass/ cutlass/backend* cd docs_src make html mv _build/* ../docs生成的文档会输出到仓库根目录的docs/对应源码目录 python/docs_src 与 python/docs。九、cutlass_library内核枚举与代码生成工具python/cutlass_library 包含枚举与发射 CUTLASS C 内核的实用工具被 CUTLASS CMake 系统用来构造可供 CUTLASS profiler 剖析的内核库。其模块结构覆盖了完整算子族gemm_operation.py、conv2d_operation.py、conv3d_operation.py、symm_operation.py、trmm_operation.py、rank_2k_operation.py、rank_k_operation.py以及面向不同架构的sm90_shapes.py、sm90_utils.py、sm100_shapes.py、sm100_utils.py等还有heuristics.py/heuristics_provider.py提供内核选择启发式逻辑。单独安装cutlass_library包python setup_library.py develop --user安装 CUTLASS Python 接口包时cutlass_library会被自动安装。此外也可以不安装模块而直接运行 python/cutlass_library/generator.py 脚本完成内核生成。十、小结CUTLASS Python 接口cutlass_cppgen与cutlass_library共同构成了 CUTLASS 的 Python 生态前者以少参数声明 智能默认配置 运行时友好错误的方式把 CUTLASS C 模板的复杂性封装在几行 Python 代码之下后者则承担内核枚举与 C 代码发射的底层职责。理解其易用优先、性能交给用户调优的定位再配合本文的安装、验证、扩展导出流程即可在 PyTorch 等框架中快速获得可编译、可运行的 CUTLASS 内核。【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表