金融软件网站建设公司免费推广的网站有哪些
【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 tl@triton.jit
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 tl@triton.jit
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