安装配置环境系统ManjaroGPUNVIDIA GeForce MX450CUDA 最高支持版本13.1这里没有使用NVIDIA Triton Docker容器而是直接使用了12source .venv/bin/activate pip install torch triton -i https://pypi.tuna.tsinghua.edu.cn/simpleHello, world !来看一下 Triton 版的 Helloworld1234567891011121314151617181920212223import torch import triton import triton.language as tl triton.jit def vector_add_kernel(a_ptr, b_ptr, c_ptr): offsets tl.arange(0, 16) a tl.load(a_ptr offsets) b tl.load(b_ptr offsets) c a b tl.store(c_ptr offsets, c) def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int): grid (1,) vector_add_kernel[grid](a, b, c) if __name__ __main__: N 16 a torch.randn(N, devicecuda) b torch.randn(N, devicecuda) triton_output torch.empty_like(a) solve(a, b , triton_output, N) print(Answer:, triton_output)实现了一个简单的固定向量维度为16的向量加法算子。看不懂没关系。下面我们进行分析分析先来复习一下向量加法若dim(a⃗)dim(b⃗)ndim(a)dim(b)n那么有a⃗b⃗(a1,a2,a3,…,an)(b1,b2,b3,…,bn)(a1b1,a2b2,a3b3,…,anbn)ab(a1,a2,a3,…,an)(b1,b2,b3,…,bn)(a1b1,a2b2,a3b3,…,anbn)如果用我们常规的思维CPU思维求解肯定是这么计算这样的加法的伪代码12for i in range(0, n): c[i] a[i] b[i]这样计算向量如果有 nn 维度那么 CPU 就要计算 nn 次。为了加速这样的运算可以考虑考虑用多核一起计算并且汇总到c上。这里专门这样并行计算的 GPU 就派上用场了。假如 GPU有 kk 个线程那么我们可以这样命令前 nn 个线程12345678910111213第一个线程 计算 a[0] b[0] 第二个线程 计算 a[1] b[1] 第三个线程 计算 a[2] b[2] ... 第n个线程 计算 a[n - 1] b[n - 1]放在 CUDA 里面我们可以为单个线程写脚本像1234int i threadIdx.x; // 线程自己知道自己的编号 c[i] a[i] b[i]; // 只算 1 个标量但是很显然这会非常繁琐。成为这样的微操大师对于暂时没有想法产出凹到机制、榨干硬件的算子的我们没有必要。Triton所操纵的不是线程级别的代码而是向量级的代码。现在我们回到代码本身。显然代码由3部分构成123456789101112131415161718192021222324252627import torch import triton import triton.language as tl # 第一部分算子内核本体 triton.jit def vector_add_kernel(a_ptr, b_ptr, c_ptr): offsets tl.arange(0, 16) a tl.load(a_ptr offsets) b tl.load(b_ptr offsets) c a b tl.store(c_ptr offsets, c) # 第二部分调用算子的函数 def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int): grid (1,) vector_add_kernel[grid](a, b, c) # 实际开发场景中算子的应用 if __name__ __main__: N 16 a torch.randn(N, devicecuda) b torch.randn(N, devicecuda) triton_output torch.empty_like(a) solve(a, b , triton_output, N) print(Answer:, triton_output)算子内核章暂且先忽略这个函数是如何运行的先看看他的思路首先可知参数是3个指针指向加向量 a⃗a 、b⃗b 和结果输出向量 c⃗c 。然后我们使用arange获取了 a⃗a 、b⃗b 维度的索引一维到十六维同样储存在张量里面。易知a⃗a 、b⃗b 的指针地址加上第 nn 维元素的偏移等于指向该元素地址的指针显然offsets是一个向量真的么此处可为了理解类比为数组而一个整数加上一个向量等于一个新向量原向量的每一位加上该整数。所以我们使用了tl.load加载了 a⃗a 、b⃗b每个元素真的吗的数值。紧接着我们进行加法运算然后将结果储存在 cc 中的每一位上同样真的吗如果第一眼看上去事先不了解 Triton似乎可以得出一个结论工作流程main函数调用solve-solve调用vector_add_kernel-vector_add_kernel调用了向量的加法 那么也就是说程序最后还是执行了不知道用哪个数学库里面的向量加法还是我们这个“向量加法”算子进入了一个无限循环这是一种 CPU 编程思维如果你尝试在vector_add_kernel中加入一个调试信息就会发现实际上这个调试信息输出了近百次其实vector_add_kernel并不能被定义为一个普通的函数。它被triton.jit修饰是一个内核。这个内核被并不会直接被 Python 的解释器调用而是会在 Triton 的 JIT即时编译下被转为对每个GPU中线程的具体命令。详情见Triton学习 · 番外篇 · 编译流程也就是说你内核写的算子代码向量级的代码并不会被GPU中的某个线程所看到。单个特定的线程所看到的是一个针对数组中某个指定索引元素的操作。如下你写的给编译器看的1234offsets tl.arange(0, 16) # 向量 [0,1,2,...,15] a tl.load(a_ptr offsets) # 向量加载 c a b # 向量运算 tl.store(c_ptr offsets, c) # 向量存储编译器翻译后每个线程看到的举例、伪代码12345678910111213141516// Thread 0 看到的 a load(a_ptr 0) c a b store(c_ptr 0, c) // Thread 1 看到的 a load(a_ptr 1) c a b store(c_ptr 1, c) // Thread 5 看到的 a load(a_ptr 5) c a b store(c_ptr 5, c) // 并不代表真实的索引分配这样每个线程所进行的就是纯正的标量数值运算而非向量运算。offsets这个向量根本不存在于最终的线程代码中它被拆解成了每个线程各自的具体数值。
Triton学习 · Part 1 · Hello, world!
安装配置环境系统ManjaroGPUNVIDIA GeForce MX450CUDA 最高支持版本13.1这里没有使用NVIDIA Triton Docker容器而是直接使用了12source .venv/bin/activate pip install torch triton -i https://pypi.tuna.tsinghua.edu.cn/simpleHello, world !来看一下 Triton 版的 Helloworld1234567891011121314151617181920212223import torch import triton import triton.language as tl triton.jit def vector_add_kernel(a_ptr, b_ptr, c_ptr): offsets tl.arange(0, 16) a tl.load(a_ptr offsets) b tl.load(b_ptr offsets) c a b tl.store(c_ptr offsets, c) def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int): grid (1,) vector_add_kernel[grid](a, b, c) if __name__ __main__: N 16 a torch.randn(N, devicecuda) b torch.randn(N, devicecuda) triton_output torch.empty_like(a) solve(a, b , triton_output, N) print(Answer:, triton_output)实现了一个简单的固定向量维度为16的向量加法算子。看不懂没关系。下面我们进行分析分析先来复习一下向量加法若dim(a⃗)dim(b⃗)ndim(a)dim(b)n那么有a⃗b⃗(a1,a2,a3,…,an)(b1,b2,b3,…,bn)(a1b1,a2b2,a3b3,…,anbn)ab(a1,a2,a3,…,an)(b1,b2,b3,…,bn)(a1b1,a2b2,a3b3,…,anbn)如果用我们常规的思维CPU思维求解肯定是这么计算这样的加法的伪代码12for i in range(0, n): c[i] a[i] b[i]这样计算向量如果有 nn 维度那么 CPU 就要计算 nn 次。为了加速这样的运算可以考虑考虑用多核一起计算并且汇总到c上。这里专门这样并行计算的 GPU 就派上用场了。假如 GPU有 kk 个线程那么我们可以这样命令前 nn 个线程12345678910111213第一个线程 计算 a[0] b[0] 第二个线程 计算 a[1] b[1] 第三个线程 计算 a[2] b[2] ... 第n个线程 计算 a[n - 1] b[n - 1]放在 CUDA 里面我们可以为单个线程写脚本像1234int i threadIdx.x; // 线程自己知道自己的编号 c[i] a[i] b[i]; // 只算 1 个标量但是很显然这会非常繁琐。成为这样的微操大师对于暂时没有想法产出凹到机制、榨干硬件的算子的我们没有必要。Triton所操纵的不是线程级别的代码而是向量级的代码。现在我们回到代码本身。显然代码由3部分构成123456789101112131415161718192021222324252627import torch import triton import triton.language as tl # 第一部分算子内核本体 triton.jit def vector_add_kernel(a_ptr, b_ptr, c_ptr): offsets tl.arange(0, 16) a tl.load(a_ptr offsets) b tl.load(b_ptr offsets) c a b tl.store(c_ptr offsets, c) # 第二部分调用算子的函数 def solve(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, N: int): grid (1,) vector_add_kernel[grid](a, b, c) # 实际开发场景中算子的应用 if __name__ __main__: N 16 a torch.randn(N, devicecuda) b torch.randn(N, devicecuda) triton_output torch.empty_like(a) solve(a, b , triton_output, N) print(Answer:, triton_output)算子内核章暂且先忽略这个函数是如何运行的先看看他的思路首先可知参数是3个指针指向加向量 a⃗a 、b⃗b 和结果输出向量 c⃗c 。然后我们使用arange获取了 a⃗a 、b⃗b 维度的索引一维到十六维同样储存在张量里面。易知a⃗a 、b⃗b 的指针地址加上第 nn 维元素的偏移等于指向该元素地址的指针显然offsets是一个向量真的么此处可为了理解类比为数组而一个整数加上一个向量等于一个新向量原向量的每一位加上该整数。所以我们使用了tl.load加载了 a⃗a 、b⃗b每个元素真的吗的数值。紧接着我们进行加法运算然后将结果储存在 cc 中的每一位上同样真的吗如果第一眼看上去事先不了解 Triton似乎可以得出一个结论工作流程main函数调用solve-solve调用vector_add_kernel-vector_add_kernel调用了向量的加法 那么也就是说程序最后还是执行了不知道用哪个数学库里面的向量加法还是我们这个“向量加法”算子进入了一个无限循环这是一种 CPU 编程思维如果你尝试在vector_add_kernel中加入一个调试信息就会发现实际上这个调试信息输出了近百次其实vector_add_kernel并不能被定义为一个普通的函数。它被triton.jit修饰是一个内核。这个内核被并不会直接被 Python 的解释器调用而是会在 Triton 的 JIT即时编译下被转为对每个GPU中线程的具体命令。详情见Triton学习 · 番外篇 · 编译流程也就是说你内核写的算子代码向量级的代码并不会被GPU中的某个线程所看到。单个特定的线程所看到的是一个针对数组中某个指定索引元素的操作。如下你写的给编译器看的1234offsets tl.arange(0, 16) # 向量 [0,1,2,...,15] a tl.load(a_ptr offsets) # 向量加载 c a b # 向量运算 tl.store(c_ptr offsets, c) # 向量存储编译器翻译后每个线程看到的举例、伪代码12345678910111213141516// Thread 0 看到的 a load(a_ptr 0) c a b store(c_ptr 0, c) // Thread 1 看到的 a load(a_ptr 1) c a b store(c_ptr 1, c) // Thread 5 看到的 a load(a_ptr 5) c a b store(c_ptr 5, c) // 并不代表真实的索引分配这样每个线程所进行的就是纯正的标量数值运算而非向量运算。offsets这个向量根本不存在于最终的线程代码中它被拆解成了每个线程各自的具体数值。
