【CUDA】Triton

【CUDA】Triton

1. CUDA 与 Triton 的基本区别

CUDA 编程模型：

在传统的 CUDA 编程中，CUDA 是标量程序，带有阻塞线程（blocked threads）。

• 标量程序（Scalar Program）：表示我们直接为每个线程编写操作代码，每个线程处理一个数据元素。
• 阻塞线程（Blocked Threads）：为了高效地处理大数据集，线程被组织成线程块（thread blocks）。线程块之间的计算是隔离的，线程块内的线程通过共享内存（shared memory）进行协作。

在 CUDA 中，程序员直接处理线程和线程块。你需要自己编写线程级别的细节，如如何加载数据、如何存储结果、如何管理共享内存等。CUDA 程序员需要掌握这些低级细节，以便充分利用 GPU 的硬件特性。

Triton 编程模型：

Triton 是基于 CUDA 的高层抽象。Triton 是带有标量线程（scalar threads）的块化程序（blocked program）。

• 块化程序（Blocked Program）：表示你不再直接操作单个线程，而是关注线程块层次的优化。你编写的程序将自动被编译器转换为低级代码。
• 标量线程（Scalar Threads）：在 Triton 中，程序员不再需要管理线程级别的细节（如内存访问模式），编译器会自动为你处理这些复杂的操作，如数据加载、存储、共享内存使用等。

2. Triton vs CUDA：高层抽象 vs 低层控制

在 CUDA 中，程序员需要掌握线程级操作，包括如何组织线程、如何管理共享内存和同步等。程序员通常需要关注线程块之间如何交互，以及如何利用硬件特性（如共享内存、寄存器）来提高效率。简而言之，CUDA 提供的是较为底层的控制，允许程序员手动优化每个操作的细节。

而在 Triton 中，程序员更多关注高层操作，例如卷积、矩阵乘法等深度学习中的标准操作。Triton 会通过编译器自动处理低级细节，如数据加载、存储、内存调度、线程间同步等。Triton 让深度学习的 GPU 编程更像是高级语言编程，减少了繁琐的底层优化工作。

直观理解：

• CUDA：程序员编写的是标量程序，并通过线程块来管理计算工作。你需要手动优化每个线程的行为以及线程块之间的协调。这意味着你要深入理解 GPU 的硬件架构。
• Triton：程序员编写的是块化程序（关注大块操作），并使用标量线程。Triton 的编译器将自动处理线程级别的操作，使得程序员不再需要处理每个线程的细节，能更专注于高级操作。

3. Triton 的优势：简化深度学习编程

Triton 的一个重要特点是为深度学习程序员提供了一个更高层次的抽象，使得复杂的 GPU 编程变得更简洁。程序员不再需要掌握底层的线程管理或内存优化，可以像写 Python 代码一样编写高效的 GPU 核心操作（如卷积、矩阵乘法等）。这对于深度学习研究者来说是一个很大的优势，因为他们的工作更关注模型设计和算法优化，而不是 GPU 编程的低级细节。

举个例子，传统上使用 CUDA 实现 cuBLAS 或 cuDNN 等高效的深度学习库需要深入理解 GPU 架构和高效的内存管理策略，而 Triton 让 Python 程序员通过一个更直观的 API 编写出与这些库同样高效的代码。

4. 为什么不能直接跳过 CUDA 使用 Triton？

Triton 是建立在 CUDA 之上的，因此，理解 CUDA 的一些基本概念仍然是非常重要的。虽然 Triton 抽象了很多复杂的低级细节，但它的性能仍然依赖于底层 CUDA 的硬件特性。

• CUDA 是 Triton 的基础：Triton 的编译器最终会将代码转换成 CUDA 代码并在 GPU 上执行。因此，了解 CUDA 中的线程组织、内存模型等基础概念仍然是优化性能的重要基础。
• 自定义优化：对于一些特定场景，程序员可能需要手动优化代码，或针对特定硬件架构进行调优，这时可能需要直接编写 CUDA 内核，或者深入理解 CUDA 的底层特性来进一步优化 Triton 生成的代码。

5. 资源学习：

• Triton 文档：可以通过官方文档详细了解 Triton 的编程模型、API 和最佳实践。
• OpenAI 博客：可以深入了解 Triton 的设计思想、底层实现以及与 CUDA 的关系。
• GitHub：可以查看 Triton 的源代码、示例程序和开源项目，帮助你理解如何使用 Triton 编写高效的深度学习程序。

总结：

• CUDA 是一种低级 GPU 编程框架，程序员需要自己处理线程调度、内存访问等底层优化细节。
• Triton 提供了一个更高层次的抽象，简化了深度学习 GPU 编程，让程序员能够专注于算法层次的开发，而不需要担心低级硬件细节。
• Triton 是建立在 CUDA 基础之上的，因此了解 CUDA 的基本概念对深入理解 Triton 及其性能优化非常重要。

---

Code

通过下面简易的代码，学习triton的使用方法，以及了解triton相较于cuda的高层次抽象。

vec_add.py

这个程序使用 Triton 和 PyTorch 实现了向量加法（x + y），并对比了两者的性能。Triton 是一个用于编写高效 GPU 内核的工具，类似于 CUDA 但更简单。程序的核心是一个 Triton 内核 add_kernel，它分块处理数据并支持边界检查。通过性能测试，程序比较了 Triton 和 PyTorch 原生加法操作的吞吐量（GB/s），并生成图表展示结果。最终目的是展示 Triton 在高性能计算中的优势。

import torch
import triton
import triton.language as [e-mail]
def add_kernel(x_ptr,  # *Pointer* to first input vector. 指向第一个输入向量的指针。y_ptr,  # *Pointer* to second input vector. 指向第二个输入向量的指针。output_ptr,  # *Pointer* to output vector. 指向输出向量的指针。n_elements,  # Size of the vector. 向量的大小。BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process. 每个程序应处理的元素数量。# NOTE: `constexpr` so it can be used as a shape value. 注意：`constexpr` 因此它可以用作形状值。):# There are multiple 'programs' processing different data. We identify which program# 有多个“程序”处理不同的数据。需要确定是哪一个程序：pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0. 使用 1D 启动网格，因此轴为 0。# This program will process inputs that are offset from the initial data.# 该程序将处理相对初始数据偏移的输入。# For instance, if you had a vector of length 256 and block_size of 64, the programs would each access the elements [0:64, 64:128, 128:192, 192:256].# 例如，如果有一个长度为 256, 块大小为 64 的向量，程序将各自访问 [0:64, 64:128, 128:192, 192:256] 的元素。# Note that offsets is a list of pointers:# 注意 offsets 是指针列表：block_start = pid * BLOCK_SIZEoffsets = block_start + tl.arange(0, BLOCK_SIZE)# Create a mask to guard memory operations against out-of-bounds accesses.# 创建掩码以防止内存操作超出边界访问。mask = offsets < n_elements# Load x and y from DRAM, masking out any extra elements in case the input is not a multiple of the block size.# 从 DRAM 加载 x 和 y，如果输入不是块大小的整数倍，则屏蔽掉任何多余的元素。x = tl.load(x_ptr + offsets, mask=mask)y = tl.load(y_ptr + offsets, mask=mask)output = x + y# Write x + y back to DRAM.# 将 x + y 写回 DRAM。tl.store(output_ptr + offsets, output, mask=mask)def add(x: torch.Tensor, y: torch.Tensor):# We need to preallocate the output.# 需要预分配输出。output = torch.empty_like(x)assert x.is_cuda and y.is_cuda and output.is_cudan_elements = output.numel()# The SPMD launch grid denotes the number of kernel instances that run in parallel.# SPMD（单程序多数据） 启动网格表示并行运行的内核实例的数量。# It is analogous to CUDA launch grids. It can be either Tuple[int], or Callable(metaparameters) -> Tuple[int].# 它类似于 CUDA 启动网格。它可以是 Tuple[int]，也可以是 Callable(metaparameters) -> Tuple[int]。# In this case, we use a 1D grid where the size is the number of blocks:# 在这种情况下，使用 1D 网格，其中大小是块的数量：grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )# NOTE:# 注意：#  - Each torch.tensor object is implicitly converted into a pointer to its first element.#  - 每个 torch.tensor 对象都会隐式转换为其第一个元素的指针。#  - `triton.jit`'ed functions can be indexed with a launch grid to obtain a callable GPU kernel.#  - `triton.jit` 函数可以通过启动网格索引来获得可调用的 GPU 内核。#  - Don't forget to pass meta-parameters as keywords arguments.#  - 不要忘记以关键字参数传递元参数。add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)# We return a handle to z but, since `torch.cuda.synchronize()` hasn't been called, the kernel is still running asynchronously at this point.# 返回 z 的句柄，但由于 `torch.cuda.synchronize()` 尚未被调用，此时内核仍在异步运行。return outputtorch.manual_seed(0)
size = 2**25
x = torch.rand(size, device='cuda')
y = torch.rand(size, device='cuda')@triton.testing.perf_report(triton.testing.Benchmark(x_names=['size'],  # Argument names to use as an x-axis for the plot. 用作绘图 x 轴的参数名称。x_vals=[2**i for i in range(12, 28, 1)],  # Different possible values for `x_name`. `x_name` 的不同可能值。x_log=True,  # x axis is logarithmic. x 轴为对数。line_arg='provider',  # Argument name whose value corresponds to a different line in the plot. 参数名称，其值对应于绘图中的不同线条。line_vals=['triton', 'torch'],  # Possible values for `line_arg`. `line_arg` 的可能值。line_names=['Triton', 'Torch'],  # Label name for the lines. 线条的标签名称。styles=[('blue', '-'), ('green', '-')],  # Line styles. 线条样式。ylabel='GB/s',  # Label name for the y-axis. y 轴标签名称。plot_name='vector-add-performance',  # Name for the plot. Used also as a file name for saving the plot. 绘图名称。也用作保存绘图的文件名。args={},  # Values for function arguments not in `x_names` and `y_name`. 不在 `x_names` 和 `y_name` 中的函数参数值。))
def benchmark(size, provider):x = torch.rand(size, device='cuda', dtype=torch.float32)y = torch.rand(size, device='cuda', dtype=torch.float32)quantiles = [0.5, 0.2, 0.8]if provider == 'torch':ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles)if provider == 'triton':ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y), quantiles=quantiles)# 3 * x.numel() * x.element_size() 因为在执行 x + y 操作时，实际上是执行了三个步骤：加载 x 张量到 GPU，加载 y 张量到 GPU，将 x + y 的结果写回到 GPU 内存gbps = lambda ms: 3 * x.numel() * x.element_size() / ms * 1e-6return gbps(ms), gbps(max_ms), gbps(min_ms)benchmark.run(print_data=True, show_plots=True,save_path='./')

可以看到，二者性能十分接近

softmax.cu

#include <cuda_runtime.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>__global__ void softmax_cuda(float* input, float* output, int B, int N) {int tid = blockIdx.x * blockDim.x + threadIdx.x;int bid = blockIdx.y;if (tid < N && bid < B) {// 简易实现，每个线程都在重复计算整个 batch 的最大值、指数和、softmax！int offset = bid * N;float max_val = input[offset];for (int i = 1; i < N; ++i) {max_val = max(max_val, input[offset + i]);}float sum = 0.0f;for (int i = 0; i < N; ++i) {sum += expf(input[offset + i] - max_val);}for (int i = 0; i < N; ++i) {output[offset + i] = expf(input[offset + i] - max_val) / sum;}}
}void softmax(float* x, int N) {float max = x[0];for (int i = 1; i < N; ++i) {if (x[i] > max) max = x[i];}float sum = 0.0f;for (int i = 0; i < N; ++i) {x[i] = exp(x[i] - max);sum += x[i];}for (int i = 0; i < N; ++i) {x[i] /= sum;}
}int main() {const int B = 32;    // 批量大小const int N = 1024;  // 行长float* x_cpu = (float*)malloc(B * N * sizeof(float));float* x_gpu = (float*)malloc(B * N * sizeof(float));float *d_input, *d_output;for (int i = 0; i < B * N; ++i) {x_cpu[i] = (float)rand() / RAND_MAX;x_gpu[i] = x_cpu[i];}cudaMalloc((void**)&d_input, B * N * sizeof(float));cudaMalloc((void**)&d_output, B * N * sizeof(float));cudaMemcpy(d_input, x_gpu, B * N * sizeof(float), cudaMemcpyHostToDevice);int theads_per_block = 256;int blocks_per_grid_x = (N + theads_per_block - 1) / theads_per_block;dim3 grid_dim(blocks_per_grid_x, B);// 二维网格，一维blocksoftmax_cuda<<<grid_dim, theads_per_block>>>(d_input, d_output, B, N);cudaMemcpy(x_gpu, d_output, B * N * sizeof(float), cudaMemcpyDeviceToHost);softmax(x_cpu, N);float max_diff = 0.0f;for (int i = 0; i < N; ++i) {float diff = fabsf(x_cpu[i] - x_gpu[i]);if (diff > max_diff) {max_diff = diff;}}printf("Maximum difference between CPU and GPU results (first batch): %e\n", max_diff);// Clean upfree(x_cpu);free(x_gpu);cudaFree(d_input);cudaFree(d_output);return 0;
}

softmax.py

import torch
import triton
import triton.language as [e-mail]
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols,BLOCK_SIZE: tl.constexpr,
):# Get the program IDrow_idx = tl.program_id(axis=0)# Compute the memory offsets for this rowrow_start_ptr = input_ptr + row_idx * input_row_strideout_row_start_ptr = output_ptr + row_idx * output_row_stride# Load the row into SRAMrow = tl.load(row_start_ptr + tl.arange(0, BLOCK_SIZE), mask=tl.arange(0, BLOCK_SIZE) < n_cols, other=-float('inf'))# other=-float('inf') 为了保证数值稳定性，在 softmax 计算中，指数运算前减去了最大值。如果某些列的值为负无穷，在指数运算后会变成零，不会影响最终的归一化结果。# Compute max for numerical stabilityrow_max = tl.max(row, axis=0)# Subtract max from row and exponentiatenumerator = tl.exp(row - row_max)# Compute sum for normalizationdenominator = tl.sum(numerator, axis=0)# Normalizesoftmax_output = numerator / denominator# Store the outputtl.store(out_row_start_ptr + tl.arange(0, BLOCK_SIZE), softmax_output, mask=tl.arange(0, BLOCK_SIZE) < n_cols)def triton_softmax(x):n_rows, n_cols = x.shapeoutput = torch.empty_like(x)# Determine the block sizeBLOCK_SIZE = triton.next_power_of_2(n_cols)BLOCK_SIZE = min(BLOCK_SIZE, 1024)  # Launch the Triton kernelgrid = (n_rows,)softmax_kernel[grid](output, x,x.stride(0), output.stride(0),n_cols, BLOCK_SIZE=BLOCK_SIZE)return output# Set up the input tensor
torch.manual_seed(0)
x = torch.randn(256, 1000, device='cuda')
# x = torch.tensor([[1.0, 2.0, 3.0]], device='cuda')
# Compute softmax using PyTorch
torch_result = torch.softmax(x, dim=1)# Compute softmax using Triton
triton_result = triton_softmax(x)# Compare results
max_diff = torch.max(torch.abs(torch_result - triton_result))
print(f"Maximum difference between PyTorch and Triton results: {max_diff:.2e}")# Check if results are close
is_close = torch.allclose(torch_result, triton_result, rtol=1e-5, atol=1e-5)
print(f"Results are close: {is_close}")

参考资料：https://github.com/Infatoshi/cuda-course/tree/master/08_Triton