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 Profiling2.1 Benchmarking做一次 warm_uptorch.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 30702.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(fmatmul(dim{dim}), run_operation2(dimdim, operationlambda 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 MLPclass 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, deviceget_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.2Profilingbenchmark 过于粗粒度只能表示代码所用时长。而 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_stackwith_stack, # with_stack 如果为 True会记录每个操作的 Python 调用栈用于后面生成堆栈跟踪可视化 experimental_configtorch._C._profiler._ExperimentalConfig(verboseTrue) # 启用更详细的实验性配置让 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_bycuda_time_total, max_name_column_width80, row_limit10) # 按 CUDA 总耗时从高到低排序cuda_time_total 是该操作在所有调用中花费的 GPU 总时间 # max_name_column_width80限制操作名称列的最大宽度避免过长。 # row_limit10只显示耗时最长的前 10 个操作 # Write stack trace visualization # 用于生成火焰图 if with_stack: text_path fvar/stacks_{description}.txt svg_path fvar/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, devicedevice) y torch.randn(dim, dim, devicedevice) # 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(dim2048, operationadd_function)) print(add_profile)表头中Self指不包括其子调用Total指包括子调用①. aten::add 是 Pytorch 底层 C 核心库中的加法操作函数。atenA Tensor ENgine②.Unrecognized③. 第三项一大串被 C名称修饰Name Mangling后的符号名还原后为at::native::vectorized_elementwise_kernel4, at::native::CUDAFunctor_addfloat, float, ...。 这是一个 CUDA 向量化逐元素内核负责 float 型加法操作。 外函数 vertorized_elementwise_kernel 进行逐元素向量化内部第一个参数 4 代表向量化宽度即每个 CUDA 线程一次处理 4 个元素第二个参数 CUDAFunctor_add 指明执行的操作为加法。这是执行加法操作的核心部分。④. cudaLaunchKernelCPU 接受指令并发送给 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(dim2048, operationmatmul_function)) print(matmul_profile)分析方法差不多前六项的 Self CPU 之和为 aten::matmul 的 CPU total。①.aten:matmul矩阵乘法入口只负责调度实际运算位于子函数②.aten:mm二维矩阵乘法底层函数③.Unrecognized④.ampere_sgemm_128x64_nnNVDIA的线性代数库 cuBLAS 根据矩阵大小以及 GPU 硬件选出的高性能计算方案⑤. cudaOccupancyMaxActiveBlocksPerMultiprocessor计算给定内核函数在 GPU 的一个 SM 上最多能同时有多少个线程块⑥. cudaLaunchKernel同上⑦. cudaDeviceSynchronize同上将矩阵维度从 2048改为 128结果如下可以观察到第四项执行了不同的指令调用了不同的计算内核。在高抽象层矩阵乘法被视为一个整体操作。但在底层实现时根据矩阵维度以及硬件配置的差异系统实际调用的矩阵乘法运算内核可能完全不同这会导致相当大的性能差异。Torch Compile 工具内置了一个能够对硬件上的矩阵乘法性能进行微基准测试micro benchmark然后为模型选择性能最高的矩阵乘法子程序subroutines——得到 10% 左右的效率优化。具体见下 3.42.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(dim2048, operationcdist_function)) print(cdist_profile)2.2.5 gelu()gelu_function lambda a, b: torch.nn.functional.gelu(ab) gelu_profile profile(matmul, run_operation2(dim2048, operationgelu_function)) print(gelu_profile)2.2.6 softmax()softmax_function lambda a, b: torch.nn.functional.softmax(ab, dim-1) softmax_profile profile(matmul, run_operation2(dim2048, operationsoftmax_function)) print(softmax_profile)2.2.7 MLPif torch.cuda.is_available(): mlp_profile profile(mlp, run_mlp(dim2048, num_layers64, batch_size1024, num_steps2), with_stackTrue) else: mlp_profile profile(mlp, run_mlp(dim128, num_layers16, batch_size128, num_steps2), with_stackTrue) print(mlp_profile)2.3 NVIDIA —— Nsight System粗略看了一下这东西不简单工程实践上的东西先略过没自己用过听老哥讲根本听不懂在代码运行过程中 CPU 进度要远远快于 GPU 。例如在迭代过程中打印损失值 loss会影响 CPU 与 GPU 的运行状态。由于打印操作是发生在 CPU 上的所以 CPU 必须等待 GPU 计算成损失结果才能继续往下进行。故这种情况下两者进度同步 CPU 有大量的空转时间。3. CUDA Kernels3.1 pytorch manualgeluPytorch 内部的 GeLu 实现方式如下def pytorch_gelu(x:torch.Tensor): return torch.nn.functional.gelu(x,approximatetanh) 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, devicedevice) return lambda:operation(x) # benchmark pytorch_time benchmark(pytorch_gelu, run_operation1(dim16384, operationpytorch_gelu)) manual_time benchmark(manual_gelu,run_operation1(dim16384, operationmanual_gelu)) # profiling pytorch_table profile(pytorch_gelu,run_operation1(dim16384, operationpytorch_gelu)) manual_table profile(manual_gelu,run_operation1(dim16384, operationmanual_gelu))pytorch_tablemanual_tablemanual_gelu 执行了大量运算触发了多个 CUDA kernel没有实现 Computation Fesion数据被搬运太多次。 而 pytorch_gelu 只采用了一个 kernel 就完成了计算所以出现了 8 倍的 benchmark 时间差异。3.2 Write a Kernel geluCUDA 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_kernelnum_blocks, block_size(x.data_ptrfloat(), y.data_ptrfloat(), 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_okTrue) # 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_kernelnum_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_ptrfloat(), y.data_ptrfloat(), 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( nameinline_gelu, cuda_sources[cuda_gelu_src], cpp_sources[cpp_gelu_src], functions[gelu], # 暴露给 Python 的函数名 extra_cflags[-O2], extra_cuda_cflags[-O2], verboseTrue, build_directoryvar/cuda_gelu, ) print(✅ 编译成功) # 7. 测试 x torch.randn(10, devicecuda, dtypetorch.float32) y module.gelu(x) # 调用编译好的函数 print(f输入: {x}) print(f输出: {y}) except Exception as e: print(f❌ 编译失败: {e})1. pip install ninja2.解决报错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 geluTriton 的优点可以用纯 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_kernelnum_blocks, block_size(x.data_ptrfloat(), y.data_ptrfloat(), num_elements); triton_gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZEblock_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(axis0) # 可以理解成 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, maskmask) # Approx gelu is 0.5 * x * (1 tanh(sqrt(2/pi) * (x 0.044715 * x^3))) # Compute (tl.tanh doesnt 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, maskmask) # 类似于 SIMT许多 threads 同时拿到这个命令然后计算不同的数据Triton 编译后得到 PTX再经过一步变成 GPU 最终执行的机器码。代码还是跑不通没招了3.4 Torch Compilegelu能将未优化的 PyTorch 代码自动转换为更高效的版本它自动尝试进行内核融合等优化操作。compiled_gelu torch.compile(manual_gelu) compiled_time benchmark(compiled_gelu, run_operation1(dim16384, operationcompiled_gelu)) compiled_table profile(compiled_gelu, run_operation1(dim16384, operationcompiled_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_ptrx, y_ptry, x_row_stridex.stride(0), y_row_stridey.stride(0), num_colsN, BLOCK_SIZEblock_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, maskcol_offsets num_cols, otherfloat(-inf)) # otherfloat(-inf)被 mask 掉的元素用负无穷填充这样在后续 max 操作中不会影响结果 # Compute x_row x_row - tl.max(x_row, axis0) numerator tl.exp(x_row) denominator tl.sum(numerator, axis0) 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, maskcol_offsets num_cols)