1.Review of GPUs
- 每个 Thread 独享一组 Register
- 每个 Block 占有一块 shared_memory ( shared_memory 位于 SM 内部,但即使一个 SM 上被分配了多个 Block,这些 Block 之间也不能共享 shared_memory,所以一个 SM 内的多个 Block 之间的通信也比较昂贵)
- Block 内部共享 shared_memory,通信代价低;跨 Block 代价高。
所有需要交互的数据都应尽量保持在同一个 Block 或计算单元内部,确保运算速度达到极致。
2. Benchmarking and Profiling
2.1 Benchmarking
做一次 warm_up;torch.cuda.synchronize() 保持 GPU 和 CPU 状态同步( Assignment 1 里写过)
2.1.1 sleep()
mean_time = benchmark("sleep", lambda: time.sleep(50/1000)) print("mean_time = ", mean_time)结果为(RTX 3070):
2.1.2 矩阵乘法
def run_operation2(dim, operation): a = torch.randn(dim, dim) b = torch.randn(dim, dim) def run(): return operation(a, b) return run if torch.cuda.is_available(): dims = (1024, 2048, 4096, 8192, 16384) # @inspect dims else: dims = (1024, 2048) # @inspect dims matmul_results = [] for dim in dims: # @ inspect dim result = benchmark(f"matmul(dim={dim})", run_operation2(dim=dim, operation=lambda a, b: a @ b)) matmul_results.append((dim, result)) # @inspect matmul_results for dim, time_ms in matmul_results: print(f"{dim:<10} {time_ms:<15.4f}") # :<10 ———— 左对齐,占用10个字符宽度 # :<15.4f ———— 左对齐,总宽度15格,保留4位小数结果为:
发现了一个比较有意思的点:对于自己的 3070 小破卡,基本上时间差都是 8 倍:因为对于维度为 n 的方阵来说,矩阵乘法的复杂度为O(n^3)。所以维度为 2 倍,所需时间为 8 倍。但对于 H100 来说,前期不遵循 8 倍关系,只有最后 8192 → 16384 大致遵循。猜测是因为 H100 计算力太强,直到 8192 维,Roofline 还停留在左侧的内存带宽瓶颈,没有完全释放计算性能。
2.1.3 MLP
class MLP(nn.Module): """Simple MLP: linear -> GeLU -> linear -> GeLU -> ... -> linear -> GeLU""" def __init__(self, dim: int, num_layers: int): super().__init__() self.layers = nn.ModuleList([nn.Linear(dim, dim) for _ in range(num_layers)]) def forward(self, x: torch.Tensor): for layer in self.layers: x = layer(x) x = torch.nn.functional.gelu(x) return x def run_mlp(dim: int, num_layers: int, batch_size: int, num_steps: int) -> Callable: # Callable 表示返回的是一个可调用的函数 # Define a model (with random weights) model = MLP(dim, num_layers).to(get_device()) # Define an input (random) x = torch.randn(batch_size, dim, device=get_device()) def run(): # Run the model `num_steps` times (note: no optimizer updates) for step in range(num_steps): # Forward y = model(x).mean() # Backward y.backward() return run一个简单的 MLP 模型,甚至没计算损失函数,只是用了 model(x) 的平均值 .mean() 来计算梯度。
Test 1:
Test 2:在 Test 1 的基础上,引入新变量 scale,分别与 run_mlp 的四个参数相乘。发现运行时间与 num_layers 以及 num_steps 呈现线性相关。
2.2Profiling
benchmark 过于粗粒度,只能表示代码所用时长。而 profiling 不仅能精确显示时间消耗在哪些函数,还能追溯到代码与 PyTorch 接口的交互,追踪从高层到底层的调用,看到底层实际执行的命令,可以更直观地理解程序如何在硬件上实际执行。
def profile(description: str, run: Callable, num_warmups: int = 1, with_stack: bool = False): # Warmup for _ in range(num_warmups): run() if torch.cuda.is_available(): torch.cuda.synchronize() # Wait for CUDA threads to finish (important!) # Run the code with the profiler with torch.profiler.profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], # 同时监测 CPU 端和 GPU 端的所有操作 with_stack=with_stack, # with_stack 如果为 True,会记录每个操作的 Python 调用栈(用于后面生成堆栈跟踪可视化) experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True) # 启用更详细的实验性配置,让 profiler 输出更多底层信息 ) as prof: run() if torch.cuda.is_available(): torch.cuda.synchronize() # Wait for CUDA threads to finish (important!) # Print out table table = prof.key_averages().table(sort_by="cuda_time_total", max_name_column_width=80, row_limit=10) # 按 CUDA 总耗时从高到低排序(cuda_time_total 是该操作在所有调用中花费的 GPU 总时间) # max_name_column_width=80:限制操作名称列的最大宽度,避免过长。 # row_limit=10:只显示耗时最长的前 10 个操作 # Write stack trace visualization # 用于生成火焰图 if with_stack: text_path = f"var/stacks_{description}.txt" svg_path = f"var/stacks_{description}.svg" prof.export_stacks(text_path, "self_cuda_time_total") return table2.2.1 sleep()
sleep_function = lambda: time.sleep(50 / 1000) sleep_profile = profile("sleep", sleep_function) print(sleep_profile)2.2.2 矩阵加法
def run_operation2(dim: int, operation: Callable) -> Callable: # Setup: create two random dim x dim matrices x = torch.randn(dim, dim, device=device) y = torch.randn(dim, dim, device=device) # Return a function to perform the operation return lambda: operation(x, y) add_function = lambda a, b: a + b add_profile = profile("add", run_operation2(dim=2048, operation=add_function)) print(add_profile)表头中Self指不包括其子调用,Total指包括子调用
①. aten::add 是 Pytorch 底层 C++ 核心库中的加法操作函数。(aten:A Tensor ENgine)
②.Unrecognized
③. 第三项一大串:被 C++名称修饰(Name Mangling)后的符号名,还原后为:at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctor_add<float, float>, ...>。 这是一个 CUDA 向量化逐元素内核,负责 float 型加法操作。 外函数 vertorized_elementwise_kernel 进行逐元素向量化,内部第一个参数 4 代表向量化宽度,即每个 CUDA 线程一次处理 4 个元素;第二个参数 CUDAFunctor_add 指明执行的操作为加法。这是执行加法操作的核心部分。
④. cudaLaunchKernel:CPU 接受指令并发送给 GPU 的过程。
注意到这四项的 Self CPU 时间和恰好等于 aten::add 的 CPU total,说明它们为 aten::add 的子函数
⑤. cudaDeviceSynchronize:等待 GPU 完成计算并传回数据
2.2.3矩阵乘法
matmul_function = lambda a, b: a @ b matmul_profile = profile("matmul", run_operation2(dim=2048, operation=matmul_function)) print(matmul_profile)分析方法差不多,前六项的 Self CPU 之和为 aten::matmul 的 CPU total。
①.aten:matmul:矩阵乘法入口,只负责调度,实际运算位于子函数;
②.aten:mm:二维矩阵乘法底层函数
③.Unrecognized
④.ampere_sgemm_128x64_nn:NVDIA的线性代数库 cuBLAS 根据矩阵大小以及 GPU 硬件选出的高性能计算方案
⑤. cudaOccupancyMaxActiveBlocksPerMultiprocessor:计算给定内核函数在 GPU 的一个 SM 上最多能同时有多少个线程块
⑥. cudaLaunchKernel:同上
⑦. cudaDeviceSynchronize:同上
将矩阵维度从 2048改为 128,结果如下:
可以观察到第四项执行了不同的指令,调用了不同的计算内核。在高抽象层,矩阵乘法被视为一个整体操作。但在底层实现时,根据矩阵维度以及硬件配置的差异,系统实际调用的矩阵乘法运算内核可能完全不同,这会导致相当大的性能差异。
Torch Compile 工具:内置了一个能够对硬件上的矩阵乘法性能进行微基准测试(micro benchmark),然后为模型选择性能最高的矩阵乘法子程序(subroutines)——得到 10% 左右的效率优化。(具体见下 3.4)
2.2.4 torch.cdist()
无论是加法还是乘法, CPU 和 GPU 之间为一对一的关系:一个 CPU 操作对应一个 GPU 操作。
torch.cdist 计算的是两组矩阵之间的欧式距离,即两组词向量的逐对距离度量。
cdist_function = lambda a, b: torch.cdist(a,b) cdist_profile = profile("matmul", run_operation2(dim=2048, operation=cdist_function)) print(cdist_profile)2.2.5 gelu()
gelu_function = lambda a, b: torch.nn.functional.gelu(a+b) gelu_profile = profile("matmul", run_operation2(dim=2048, operation=gelu_function)) print(gelu_profile)2.2.6 softmax()
softmax_function = lambda a, b: torch.nn.functional.softmax(a+b, dim=-1) softmax_profile = profile("matmul", run_operation2(dim=2048, operation=softmax_function)) print(softmax_profile)2.2.7 MLP
if torch.cuda.is_available(): mlp_profile = profile("mlp", run_mlp(dim=2048, num_layers=64, batch_size=1024, num_steps=2), with_stack=True) else: mlp_profile = profile("mlp", run_mlp(dim=128, num_layers=16, batch_size=128, num_steps=2), with_stack=True) print(mlp_profile)2.3 NVIDIA —— Nsight System
粗略看了一下这东西不简单,工程实践上的东西,先略过
没自己用过听老哥讲根本听不懂
在代码运行过程中, CPU 进度要远远快于 GPU 。例如在迭代过程中打印损失值 loss,会影响 CPU 与 GPU 的运行状态。由于打印操作是发生在 CPU 上的,所以 CPU 必须等待 GPU 计算成损失结果,才能继续往下进行。故这种情况下,两者进度同步, CPU 有大量的空转时间。
3. CUDA Kernels
3.1 pytorch & manual(gelu)
Pytorch 内部的 GeLu 实现方式如下:
def pytorch_gelu(x:torch.Tensor): return torch.nn.functional.gelu(x,approximate="tanh") x = torch.tensor([1.]) y1 = pytorch_gelu(x)利用 tanh 来近似计算 GeLu 以加快计算速度,没有使用 GeLu 的精确定义用标准高斯分布的累计分布函数 CDF 来计算。
原始方法计算:
def manual_gelu(x: torch.Tensor): return 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x * x * x))) y2 = manual_gelu(x)计算两者结果,并分别进行 benchmark 以及 profiling 。
def run_operation1(dim:int,operation:Callable)->Callable: x = torch.randn(dim, dim, device=device) return lambda:operation(x) # benchmark pytorch_time = benchmark("pytorch_gelu", run_operation1(dim=16384, operation=pytorch_gelu)) manual_time = benchmark("manual_gelu",run_operation1(dim=16384, operation=manual_gelu)) # profiling pytorch_table = profile("pytorch_gelu",run_operation1(dim=16384, operation=pytorch_gelu)) manual_table = profile("manual_gelu",run_operation1(dim=16384, operation=manual_gelu))pytorch_table:
manual_table:
manual_gelu 执行了大量运算,触发了多个 CUDA kernel(没有实现 Computation Fesion,数据被搬运太多次)。 而 pytorch_gelu 只采用了一个 kernel 就完成了计算,所以出现了 8 倍的 benchmark 时间差异。
3.2 Write a Kernel (gelu)
CUDA is an extension of C/C++ with APIs for managing GPUs and a programming model for expressing parallelism.
语言层面: C/C++ 扩展
语言层面:提供操作 GPU 的 API
编程模型:实现并行
编写 CUDA kernel 函数并启动它时,它会自动在 GPU 的数千个线程上并行执行,对向量或矩阵的所有元素同时进行计算。在 CUDA 的编程模型中,Grid 是顶层容器,包含若干个 Block;每个 Block 包含若干个 Thread。例如在二维 Grid 中,每个 Block 通过 (blockIdx.x, blockIdx.y) 来定位;二维 Block 中,每个 Thread 通过 (threadIdx.x, threadIdx.y) 定位。(这些参数不是通过参数传给 Kernel 的,Kernel 可以直接访问 CUDA 提供的内置变量 threadIdx、blockIdx、blockDim、gridDim,通过这些变量,每个线程可以计算出自己在全局数据中的唯一索引,从而处理对应的数据元素)
调试 CUDA 时,需要将环境变量设置为 os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
只有这样才能正确调试 CUDA,不过这牺牲了运行时性能,系统能够返回详细的错误信息。
#include <math.h> #include <torch/extension.h> #include <c10/cuda/CUDAException.h> __global__ void gelu_kernel(float* in, float* out, int num_elements) { // Get the index into the tensor int i = blockIdx.x * blockDim.x + threadIdx.x; // 算出当前线程要处理张量中的哪个位置 if (i < num_elements) { // 防止超出张量范围(因为总线程数 ≥ 元素数) // Do the actual computation out[i] = 0.5 * in[i] * (1.0 + tanh(0.79788456 * (in[i] + 0.044715 * in[i] * in[i] * in[i]))); } } // __global__ 不能有返回值,用 out 指针写回结果 inline unsigned int cdiv(unsigned int a, unsigned int b) { // inline: 建议编译器消除函数调用开销,直接展开代码 // Compute ceil(a / b) return (a + b - 1) / b; // -1 防止原本就能整除 (a = 9, b = 3) } // Host函数:发生在CPU侧 torch::Tensor gelu(torch::Tensor x) { // 检查张量在 GPU 上且连续 TORCH_CHECK(x.device().is_cuda()); TORCH_CHECK(x.is_contiguous()); // 在 GPU 上分配一块和 x 同样大小(形状)的内存,但不初始化 torch::Tensor y = torch::empty_like(x); // Determine grid (elements divided into blocks) int num_elements = x.numel(); // 获取元素总数 例如 [3, 4] 的张量 → num_elements = 12 int block_size = 1024; // Number of threads in a block int num_blocks = cdiv(num_elements, block_size); //用向上取整除法算出需要启动多少个 Block 才能覆盖所有元素 // Launch the kernel gelu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), num_elements); // 在 GPU 上启动 num_blocks 个 Block,每个 Block 有 1024 个线程,总共 num_blocks × 1024 个线程并行执行 // 由于 gelu_kernel 由核心修饰符 __global__ 修饰,说明 gelu_kernel 由 CPU 调用, GPU 个线程并行执行 // gelu_kernel 需要 <<<>>> 传参 // 所有被 __global__ 修饰的核函数都需要传这两个参数 C10_CUDA_KERNEL_LAUNCH_CHECK(); //检查内核启动是否成功。如果内核里有 bug ,这里会立即捕获并报错 return y; }(搞了一整天没编译成功 .cu 和 .cpp 文件,先搁置吧。心累,这种实践上的东西 AI 也是张口就来,解决不了)贴一个完整代码在下面,benchmark 和 profiling 也只有略过。
import torch import os from torch.utils.cpp_extension import load_inline # 1. 确保目录存在 def ensure_directory_exists(path): os.makedirs(path, exist_ok=True) # 2. CUDA 源代码(核函数实现) cuda_gelu_src = """ #include <math.h> __global__ void gelu_kernel(float* in, float* out, int num_elements) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < num_elements) { float x = in[i]; out[i] = 0.5f * x * (1.0f + tanhf(0.79788456f * (x + 0.044715f * x * x * x))); } } // 这个函数被 C++ 封装代码调用 void launch_gelu_kernel(float* in, float* out, int num_elements) { int block_size = 1024; int num_blocks = (num_elements + block_size - 1) / block_size; gelu_kernel<<<num_blocks, block_size>>>(in, out, num_elements); } """ # 3. C++ 封装代码(PyTorch 绑定) cpp_gelu_src = """ #include <torch/extension.h> #include <c10/cuda/CUDAException.h> // 声明 CUDA 函数(在 cuda_sources 中实现) void launch_gelu_kernel(float* in, float* out, int num_elements); // 被 Python 调用的函数 torch::Tensor gelu(torch::Tensor x) { TORCH_CHECK(x.is_cuda(), "x must be on GPU"); TORCH_CHECK(x.dtype() == torch::kFloat32, "x must be float32"); auto y = torch::empty_like(x); int num_elements = x.numel(); launch_gelu_kernel(x.data_ptr<float>(), y.data_ptr<float>(), num_elements); C10_CUDA_KERNEL_LAUNCH_CHECK(); return y; } """ # 4. 安全检查 if not torch.cuda.is_available(): print("CUDA 不可用,跳过编译") exit(1) # 5. 确保编译目录存在 ensure_directory_exists("var/cuda_gelu") # 6. 使用 load_inline 编译 try: module = load_inline( name="inline_gelu", cuda_sources=[cuda_gelu_src], cpp_sources=[cpp_gelu_src], functions=["gelu"], # 暴露给 Python 的函数名 extra_cflags=["-O2"], extra_cuda_cflags=["-O2"], verbose=True, build_directory="var/cuda_gelu", ) print("✅ 编译成功!") # 7. 测试 x = torch.randn(10, device='cuda', dtype=torch.float32) y = module.gelu(x) # 调用编译好的函数 print(f"输入: {x}") print(f"输出: {y}") except Exception as e: print(f"❌ 编译失败: {e}")1. pip install ninja;
2.解决报错:subprocess.CalledProcessError: Command '['where', 'cl']' returned non-zero exit status 1 Pytorch 找不到 Visual Studio 的 C++ 编译器 cl.exe 解决方法链接
3.解决报错:subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 2. 将 torch.utils.pp_extension 中的
command = ['ninja', '-v'] -----> command = ['ninja', '--version']
4.:新问题:搞不定了,摆烂了
文件也成功生成了
3.3 Triton Kernels (gelu)
Triton 的优点:可以用纯 Python 写 GPU Kernel,无需关心线程管理,只需专注于线程块的设计;能自动处理许多底层细节,比如自动调整内存访问模式。代码是以 Tile 为中心写的,编译器负责将这些 Tile 调度到 SM 上。跨 SM 的数据共享需要手动处理,跨 SM 的并行 Triton 能够自动实现。
首先需要安装 triton(由于自己是 Windows 环境,需要进行一些设置)
import torch import triton import triton.language as tl # 原本用 C++ 实现的 CPU 侧的 Host 函数 # torch::Tensor gelu(torch::Tensor x) def triton_gelu(x: torch.Tensor): assert x.is_cuda assert x.is_contiguous() # Allocate output tensor y = torch.empty_like(x) # Determine grid (elements divided into blocks) num_elements = x.numel() block_size = 1024 # Number of threads num_blocks = triton.cdiv(num_elements, block_size) # 这一步的实现与 CUDA 稍有区别 # gelu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), num_elements); triton_gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=block_size) # Triton 要求 grid 参数是一个元组,用来支持多维网格 # (M,) 一维:启动 M 个 block,索引为 0, 1, 2, ..., M-1 # (M,N) 二维:启动 M*N 个 block,索引为 (0,0), (0,1), ..., (M-1, N-1) # (M,N,K) 三维 return y @triton.jit def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr): # BLOCK_SIZE: tl.constexpr:编译时常量(在编译时确定,不能运行时改变) pid = tl.program_id(axis=0) # 可以理解成 blockId block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) # 计算索引 # CUDA 需要计算出每个 thread 的索引 # Triton kernel 计算的offsets 是一个向量,长度为 block_size # Handle boundary mask = offsets < num_elements # Read x = tl.load(x_ptr + offsets, mask=mask) # Approx gelu is 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) # Compute (tl.tanh doesn't exist, use tanh(a) = (exp(2a) - 1) / (exp(2a) + 1) a = 0.79788456 * (x + 0.044715 * x * x * x) exp = tl.exp(2 * a) tanh = (exp - 1) / (exp + 1) y = 0.5 * x * (1 + tanh) # Store tl.store(y_ptr + offsets, y, mask=mask) # 类似于 SIMT,许多 threads 同时拿到这个命令,然后计算不同的数据Triton 编译后得到 PTX,再经过一步变成 GPU 最终执行的机器码。
代码还是跑不通,没招了
3.4 Torch Compile(gelu)
能将未优化的 PyTorch 代码自动转换为更高效的版本,它自动尝试进行内核融合等优化操作。
compiled_gelu = torch.compile(manual_gelu) compiled_time = benchmark("compiled_gelu", run_operation1(dim=16384, operation=compiled_gelu)) compiled_table = profile("compiled_gelu", run_operation1(dim=16384, operation=compiled_gelu))(还是跑不了,这章听得我难受,配环境变量搞了一整天,啥也没成,还不知道原因)
3.5 Triton Kernel (softmax)
def triton_softmax(x: torch.Tensor): # Allocate output tensor y = torch.empty_like(x) # Determine grid M, N = x.shape # Number of rows x number of columns block_size = triton.next_power_of_2(N) # 将 N 向上取整到 2 的幂 num_blocks = M # Each block is a row 每行独立计算 Softmax,一行对应一个 block # Launch kernel triton_softmax_kernel[(M,)]( x_ptr=x, y_ptr=y, x_row_stride=x.stride(0), y_row_stride=y.stride(0), num_cols=N, BLOCK_SIZE=block_size ) # 传递了行跨度 (stride) 以便在内存中正确跳转到不同行 return y @triton.jit def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride, num_cols, BLOCK_SIZE: tl.constexpr): assert num_cols <= BLOCK_SIZE # Process each row independently row_idx = tl.program_id(0) col_offsets = tl.arange(0, BLOCK_SIZE) # Read from global memory x_start_ptr = x_ptr + row_idx * x_row_stride x_ptrs = x_start_ptr + col_offsets x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf")) # other=float("-inf"):被 mask 掉的元素用负无穷填充,这样在后续 max 操作中不会影响结果 # Compute x_row = x_row - tl.max(x_row, axis=0) numerator = tl.exp(x_row) denominator = tl.sum(numerator, axis=0) y_row = numerator / denominator # Block 内部会有一个比较 max 以及求和的操作 # 这两步 triton 会自行处理 # 虽然单个线程看起来需要等待其他线程的结果,但 Triton 在 Block 层面实现了同步 # Write back to global memory y_start_ptr = y_ptr + row_idx * y_row_stride y_ptrs = y_start_ptr + col_offsets tl.store(y_ptrs, y_row, mask=col_offsets < num_cols)