1. 概述
mcTriton在Triton的基础上增加了MXMACA® 后端,支持使用MXMACA硬件加速Triton中各类计算任务。
当前发布包含以下版本安装包:
Triton2.1.0 + Python 3.8/3.10
Triton3.0.0 + Python 3.8/3.10
支持x86_64/Arm架构下Ubuntu 20/22以及CentOS 8/9系统上运行。mcTriton的具体支持功能,参见 3 功能支持。
2. 快速安装
2.1. 基于pip安装
2.1.1. 环境准备
Python
匹配目标安装包的Python版本(例如Python 3.8或者Python 3.10)
MXMACA环境
安装Driver软件包以及MXMACA SDK软件包
2.1.2. 开始安装
可在PyTorch安装压缩包中获取Triton的wheel安装包。 解压PyTorch安装压缩包,可以获得Triton和PyTorch的wheel安装包(Triton 2.1版本wheel包位于PyTorch 2.1版本的安装压缩包中;Triton 3.0版本wheel包位于PyTorch 2.4版本的安装压缩包中)。
PyTorch安装压缩包以maca-pytorch${pytorch_version}-py${python_version}-${release_version}-${arch_info}.tar.xz格式命名。 Triton和PyTorch的wheel安装包分别以triton-*.whl和torch-*.whl格式命名。
操作步骤
安装方式同标准whl包。
安装mcPyTorch:
python -m pip install torch-\*.whl
安装mcTriton:
python -m pip install triton-\*.whl
2.1.3. 验证安装
操作步骤
运行前设置环境变量:
export MACA_PATH=/opt/maca/ export LD_LIBRARY_PATH=${MACA_PATH}/lib:${MACA_PATH}/mxgpu_llvm/lib: ${MACA_PATH}/ompi/lib:${LD_LIBRARY_PATH} export MACA_CLANG_PATH=${MACA_PATH}/mxgpu_llvm/bin
执行以下命令:
python -c "import triton"
没有报错信息,表明Triton包安装成功。
2.1.4. 如何卸载
卸载方式同标准whl包。
操作步骤
卸载mcTriton:
python -m pip uninstall triton
卸载mcPyTorch:
python -m pip uninstall torch
3. 功能支持
3.1. 基本功能
mcTriton以兼容Triton的原生使用方式为设计目标。大部分情况下,用户可以参考官方文档 获得mcTriton的使用方式。
3.2. 扩展功能
本章介绍mcTriton相较官方Triton提供的额外功能。
3.2.1. triton.Config
出于性能调优的目的,mcTriton扩展了triton.Config中的配置选项,提供给用户做不同场景下的性能调优。
3.2.1.1. pipeline
pipeline 选项用于指定在for循环中 triton.language.dot 场景下使用的不同优化方式。具体使用方法:
Triton 2.1
triton.Config({…},…,pipeline="basic",…)Triton 3.0
triton.Config({…,'pipeline': "basic",…},…)
支持参数:
basic:基本pipeline。basic-prefetch:基于basic模式,增加额外的数据预取环节。增大寄存器压力,减轻访存延迟压力。cpasync:使用cp.async功能在pipeline中将数据直接从global memory拷贝到shared memory,不经过寄存器中转。
3.2.2. triton.autotune
triton.autotune 是Triton中用于自动调优的装饰器。
用户提供可选的 triton.Config 列表,在运行时,Triton会执行“tune”过程,从这些配置中选取最合适的配置来运行Kernel,并将结果缓存下来。
为了进一步提升性能,mcTriton提供将自动调优(autotune)结果持久化至硬盘的功能。 通过此功能,程序可以在每次重新运行时跳过“tune”阶段,直接使用之前运行时保存的自动调优结果。操作步骤如下所示。
开启自动调优结果持久化功能的环境变量。
export TRITON_ENABLE_PERSISTENT_AUTOTUNE_CONFIGS=1
(可选)配置持久化路径。未配置时默认保存至$HOME/.triton/cache/configs。
export TRITON_AUTOTUNE_CONFIG_PATH=$USER_DEFINED_PERSISTENT_PATH
运行Kernel。
Kernel运行后会将自动调优结果保存到持久化路径下,后续程序再次运行时优先从该路径读取之前的缓存结果。
4. 常用环境变量
4.1. TRITON_DISABLE_SWIZZLE
默认为 OFF。控制关闭shared memory的swizzle功能。
设置:
export TRITON_DISABLE_SWIZZLE=1取消设置:
unset TRITON_DISABLE_SWIZZLE