cuda-12.4.0 devel docker 中源码安装 OpenAI triton

 0，宿主机安装cuda sdk 和 kmd dirver 等

可参考 nv official docs；

1，准备 docker 容器

下载docker image:
 

$ sudo docker pull nvidia/cuda:12.6.2-devel-ubuntu20.04

创建容器：
 

sudo docker run --gpus all -it --name cuda_LHL_01  -v /home/hongleili/ex_triton/tmp1:/root/ex_triton/tmp1   nvidia/cuda:12.6.2-devel-ubuntu20.04

2，安装 git vim python3 wget pip

 

apt-get update
apt install git vim python3 wget
apt install python3-pip

3, 源码安装 cmake

# apt install libssl-dev
# git clone https://github.com/Kitware/CMake.git
# cd CMake/
# git checkout v3.31.6./bootstrap
make -j
make install

4, 安装 pybind11

 

pip install ninja cmake wheel pybind11  -i https://pypi.tuna.tsinghua.edu.cn/simple

5, 下载 triton

git clone https://github.com/triton-lang/triton.gitcd triton/git checkout v2.1.0pip install -e python

安装：

pip install -e python  -i https://pypi.tuna.tsinghua.edu.cn/simple

成功提示如下，

 下图这个输出会停留较长时间，

6，安装 pytorch

查找 torch 的 pip 网址：

https://pypi.org/search

安装 torch：

不指定 cuda 版本，

pip install torch==2.4.1 -i https://pypi.tuna.tsinghua.edu.cn/simple

指定 cuda 版本为12.6，

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

也可以从源码安装 torch

7，测试 triton 环境

基本原则：

1）Triton 目前只支持 PyTorch 算子融合, 所以在使用triton时，需要事先import torch 包；
2）在 核函数 中, 我们只能使用 tl 中提供的关键字和函数, 或者被 triton.jit 装饰过的函数（类似 __device__ 修饰）,不能调用其他 triton 或 pytorch 等的函数；
3）而在 核函数 之外, 可以使用 triton 提供的函数，但是，不能使用 tl 中提供的函数；