一、CUDA性能简易优化
这里是关于CUDA性能优化的基础教程，主要介绍了GPU的基础知识，帮助理解如何设计和优化深度学习中的层，使其能充分利用GPU性能。
CUDA性能优化简单教程，本篇介绍性能优化背景。
想知道实际中如何优化特定的层，或者某一层怎么设计才可能充分利用GPU，我们需要了解一些GPU的基础知识。
以下教程主要来源自NVIDIA官方：GPU Performance Background User's Guide[1]，主要是讲解深度学习中，网络中一些operation在GPU中是如何执行的，已经一些和性能相关的细节注意点。
具体内容如下：
- GPU的基本结构 (GPU Architecture Fundamentals[2])
- OP（operation）是如何被拆分设计为并行计算的 (GPU Execution Model[3])
- 如何通过算术强度（arithmetic intensity）估算性能限制 (Understanding Performance[4])
- 深度学习操作的大概分类及其各自的性能限制(DNN Operation Categories[5])
1. GPU的基础结构
GPU是一种高度并行的处理器架构，由processing elements and a memory hierarchy组成，和我们经常提到的SM和显存有直接关系。
在高层次上，NVIDIA® GPU由多个流多处理器（SM）组成，具有片上的L2缓存和高带宽DRAM。SM执行算术和其他指令；通过L2缓存从DRAM访问数据和代码。例如，NVIDIA A100 GPU包含108个SM，40 MB L2缓存以及来自80 GB HBM2内存的高达2039 GB/s的带宽。
GPU架构
每个SM都有自己的instruction schedulers和各种instruction execution pipelines。在现代神经网络中，Multiply-add是最常见的操作，作为全连接和卷积层的核心块操作，两者都可以视为向量点积（vector dot-products.）的集合。
下表显示了NVIDIA Ampere GPU架构上单个SM的每个时钟的乘加操作的各种数据类型。每个乘加操作包括两个操作，因此可以将表中的吞吐量乘以2，以获得FLOP counts per clock。
要获取GPU的FLOPS具体值，我们可以将这些乘以SM的数量和SM时钟速率。例如，带有108个SM和1.41 GHz时钟速率的A100 GPU具有156 TF32 TFLOPS和312 FP16 TFLOPS的峰值密集吞吐量（实际能达到的吞吐量取决于后续讨论的多个因素，这只是理论值）。
Multiply-add operations per clock per SM
另外，FP16操作可以在Tensor Cores或NVIDIA CUDA®核心中执行。NVIDIA Turing™架构可以在Tensor Cores或CUDA核心中执行INT8操作。Tensor Cores在NVIDIA Volta™ GPU架构中引入，用于加速机器学习和科学应用的矩阵乘法和累加操作。这些指令在小的矩阵块（例如4x4块）上操作。
Tensor Cores还可以在比输入更高的精度下进行乘积计算和累加。例如，在使用FP16（半精度浮点数）输入进行训练时，张量核心可以在不损失精度的情况下进行乘积计算，并以FP32（单精度浮点数）进行累加。
满足特定操作的可以在Tensor Core中执行，不满足的就在普通Cuda Core中执行（比如一个简单的FP16精度的element-wise的加法操作）。
2. GPU Execution Model
为了利用其并行资源，GPU通常会同时执行许多线程。
了解线程数与GPU性能之间的关系，有两个关键概念：
- GPU使用2级线程层次结构来执行函数。给定函数的线程被分组成相同大小的线程块，并启动一组线程块来执行函数。
- GPU通过切换到其他线程的执行来隐藏依赖指令的延迟。因此，为了有效利用GPU，所需的线程数量远远高于核心数或指令流水线数。
两级线程层次结构是因为GPU拥有许多流式多处理器（SM），每个SM又具有执行许多线程的流水线，并使其线程通过共享内存和同步进行通信。在运行时，一个线程块会被放置在一个SM上执行，使线程块中的所有线程能够高效地进行通信和同步。
使用单个线程块启动一个函数只会使一个SM工作，因此为了充分利用拥有多个SM的GPU，需要启动许多线程块。由于一个SM可以同时执行多个线程块，通常希望线程块的数量是SM数量的几倍。这样做的原因是为了尽量减少“尾部”效应，即在函数执行结束时只剩下少数活跃的线程块，从而在那段时间内GPU利用率不足，如下图所示。
GPU Execution Model
Utilization of an 8-SM GPU when 12 thread blocks with an occupancy of 1 block/SM at a time are launched for execution. Here, the blocks execute in 2 waves, the first wave utilizes 100% of the GPU, while the 2nd wave utilizes only 50%.
我们使用波（wave）这个术语来指代一组可以并行运行的线程块。以多个wave的线程块来启动函数执行是最高效的方式——在尾波（最后一个波次）中花费的时间占比较小，从而最小化尾效应，因此也减少了处理尾效应的需求（这种效应会导致总体执行效率降低，因为GPU的一部分核心在等待少数几个线程块完成时处于空闲状态）。对于高端GPU来说，通常只有当启动少于300个线程块的情况下，才需要考虑尾效应。
3. Understanding Performance
在给定的处理器上执行函数的性能受以下三个因素之一的限制：
- 内存带宽
- 数学带宽
- 延迟
以下称呼的数学，意思就是某个op的算术强度，就是这个操作需要多少flops才能实现
考虑一个简化模型，其中一个函数从内存中读取其输入，执行数学运算，然后将其输出写入内存。假设：
是显存访问和读取的时间, 是在数学算子上计算的时间。如果我们假设这两个操作可以在不同线程中执行, 可以重叠（overlap），那么总时间就是 )。两个当中谁大谁就是限制当前op性能的原因。如果是数学那么就是math limited, 如果是显存那就是memory limited。
在显存或数学运算中花费多少时间取决于算法及其实现方式，以及处理器的带宽。内存时间等于在内存中访问的字节数除以处理器的内存带宽。数学时间等于操作次数除以处理器的数学带宽。因此，在给定处理器上，如果一个给定算法是受到数学限制，则：
通过简单的代数，不等式可以重新排列为：
左侧是算法实现操作次数和访问字节数的比值，也就是俗称的arithmetic intensity。右侧是处理器计算能力和内存带宽之比，有时被称为_ops:byte比率_。因此，在给定处理器上，如果算法的算术强度高于处理器的ops:byte比率，则该算法受到数学限制。相反地，如果其算术强度低于处理器的ops:byte比率，则该算法受到内存限制。
让我们考虑一些深度神经网络的具体例子，列在下面的表格中。对于这些示例，我们将比较算法的算术强度与在NVIDIA Volta V100 GPU上的ops:byte比率。V100具有125 FP16 Tensor TFLOPS的峰值数学速率，约为900 GB/s 的芯片外存储器带宽和3.1 TB/s 的芯片内L2带宽，因此其ops:byte比率介于40到139之间，取决于操作数据（片内或片外存储器）的来源。
Operation | Arithmetic Intensity | Usually limited by... |
Linear layer (4096 outputs, 1024 inputs, batch size 512) | 315 FLOPS/B | arithmetic |
Linear layer (4096 outputs, 1024 inputs, batch size 1) | 1 FLOPS/B | memory |
Max pooling with 3x3 window and unit stride | 2.25 FLOPS/B | memory |
ReLU activation | 0.25 FLOPS/B | memory |
Layer normalization | < 10 FLOPS/B | memory |
如表所示，许多常见操作的算术强度很低——有时每从内存读取并写入两字节元素只执行一次操作。请注意，这种分析是一种简化，因为我们只计算了_算法_中使用的操作。实际上，函数还包含算法中未明确表达的操作指令，如内存访问指令、地址计算指令、控制流指令等。
算术强度和每字节操作数分析假设工作负载足够大，能够充分利用给定处理器的数学和内存管道。然而，如果工作负载不够大，或没有足够的并行性，处理器将被低效利用，性能将受到延迟的限制。例如，考虑启动一个单线程，它将访问16字节并执行16000次数学操作。尽管算术强度为1000 FLOPS/B，并且在V100 GPU上的执行应该是数学限制的，但只创建一个线程会严重低效利用GPU，使其几乎所有的数学管道和执行资源处于空闲状态。
此外，算术强度计算假设输入和输出仅从内存中访问一次。对于算法实现多次读取输入元素的情况并不少见，这会有效降低算术强度。因此，算术强度是一种初级近似；如果需要更准确的分析，需要看profiler。
4. DNN Operation Categories
我们平常的深度学习网络中有很多基础的op层，大概分为三类：
4.1. Elementwise Operations
Elementwise operations 可以是一元或二元操作；关键在于此类层对张量中的每个元素执行数学运算，独立于张量中所有其他元素。
例如，ReLU层对输入张量中的每个x返回max(0, x)。同样，两个张量的逐元素加法计算每个输出和值与其他和值无关。这类层包括大多数非线性函数（sigmoid、tanh等）、比例、偏置、加法等。
这些层往往受到内存限制，因为它们每访问一个字节只执行少量操作。有关激活函数的更多详细信息，请参见《优化内存限制层用户指南》中的激活[6]部分。
4.2. Reduction Operations
reduce operations也就是常见的规约操作。
例如，池化层在输入张量的某些邻域内计算值。Batch normalization在每个输出元素的操作中使用之前，计算张量的均值和标准差。除了池化和归一化层外，SoftMax也属于归约类别。典型的归约操作算术强度较低，因此受到内存限制。有关池化层的更多详细信息，可以在池化[7]中找到。
4.3. Dot-Product Operations
该类操作可以表示为两个tensor的点积，通常是一个权重（学习参数）张量和一个激活张量。也就是常说的矩阵乘法。
包括全连接层（fully-connected layers），可以很自然地表示为矩阵向量和矩阵矩阵乘积。
卷积也可以理解为就是做 matrix-vector and matrix-matrix multiplies，常见的卷积实现im2col就是典型的矩阵乘。不过矩阵乘处于算术瓶颈还是显存瓶颈也分情况，如果batch比较小或者输入比较小，那么还是memory bound，只有当batch大的时候才会是compute bound
5. 总结
如何大概预估我们某个op在某个GPU上的性能受限在哪，有以下几点方法：
- 弄清楚这个GPU上的 SM 数量，并确定 GPU 的 ops:bytes 比例
- 预估当前op算法的运算强度（arithmetic intensity）
- 通过估算线程块的数量和大小，判断是否有足够的并行度来充分利用 GPU。如果线程块的数量至少是 SM 数量的 4 倍，并且每个线程块由几百个线程组成，那么可能就有足够的并行度
- 最可能的性能瓶颈是：
- 如果没有足够的并行度，则是GPU常见的延迟问题，延迟无法掩盖
- 如果有足够的并行度且算法运算强度高于 GPU ops:bytes 比例，则是数学运算瓶颈
- 如果有足够的并行度且算法运算强度低于 GPU ops:bytes 比例，则是内存瓶颈
说白也就是看Memory-bound还是Compute-bound
参考资料
[1]
GPU Performance Background User's Guide: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html
[2]GPU Architecture Fundamentals: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#gpu-arch
[3]GPU Execution Model: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#gpu-execution
[4]Understanding Performance: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#understand-perf
[5]DNN Operation Categories: https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html#dnn-op-cat
[6]激活: https://docs.nvidia.com/deeplearning/performance/dl-performance-memory-bound/index.html#activations
[7]池化: https://docs.nvidia.com/deeplearning/performance/dl-performance-memory-bound/index.html#pooling
开发板商城 天皓智联 相关AI视觉设备
二、CUDA性能检查清单
文中详细讨论了如何合并全局内存访问、最大化占用率、理解内存与计算限制、最小化线程分化、通过Tiling重用数据、私有化内存使用、线程粗化、以及通过算法重写来提升性能。此外，文章还提到了使用ncu工具来分析和优化kernel性能，并通过实际的CUDA代码示例来阐释这些概念。最后，文章提到了Flash Attention算法作为通过数学角度重写算法以提升性能的一个例子。
大佬的课程笔记，欢迎关注：https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/cuda-mode
CUDA-MODE课程笔记 第8课: CUDA性能检查清单
课程笔记
这节课实际上算是CUDA-MODE 课程笔记 第一课: 如何在 PyTorch 中 profile CUDA kernels 这节课更细节的讲解。另外关于nsight compute相关指标细节解释可以参考 CUDA-MODE 第一课课后实战（上），CUDA-MODE 第一课课后实战（下） 这两篇笔记。
将GPU用于计算，我们最关心的肯定是性能。比较幸运的是，当我们掌握一些性能优化技巧之后它往往会经常被用到。这节课将会系统的介绍这些性能优化技巧。
本节课的课件和代码都在 https://github.com/cuda-mode/lectures 开源，我们可以用nvcc编译lecture8下面的cu文件，然后使用ncu进行profile。此外，这里的方法遵循了 https://arxiv.org/pdf/1804.06826.pdf 这篇论文的风格。up主也非常推荐大家阅读下这篇论文，用claude 3.5问了一下paper的主要内容，如下截图：
可以看到这篇论文主要是对Volta架构的GPU的架构细节进行分析，对性能优化是很重要的。
这张Slides从物理学的角度分析了下SRAM和DRAM的区别：
- DRAM由1个晶体管和1个电容器构成；SRAM: 由6个晶体管构成
- SRAM 比 DRAM 更快，但也更贵；SRAM 占用更多空间且发热更多；实际上SRAM就对应了GPU的Shared Memory，而DRAM对应的则是Shared Memory。
这里的youtube链接作者Bill是NVIDIA的首席科学家，他解释了很多为什么GPU设计成现在这个样子，并且由浅入深，基础细节讲的非常清楚。
这里的"性能检查清单"（Performance checklist），列出了一系列优化GPU程序性能的策略和技巧：
- 合并全局内存访问（Coalesced Global Memory Access）
- 最大化占用率（Maximize occupancy）
- 理解是内存受限还是计算受限（Understand if memory or compute bound）
- 最小化线程分化（Minimize control divergence）
- Tiling以更好的重用数据（Tiling of reused data）
- 私有化（Privatization）
- Thread Coarsening
- 使用更好的数学方法重写算法（Rewrite your algorithm using better math）
这里的Privatization指的应该就是Shared Memory/寄存器优化全局内存读取，而Coarsening大概指的就是一个线程应该完成多少任务，一般情况下我们让一个线程完成的任务尽量少，但是在Compute Bound情况下，让一个线程执行更多的工作可以让程序运行得更快。最后一点更好的数学方法重写算法的经典例子就是Flash Attention。
这张Slides讲述了GPU内存访问延迟的相关内容，下面的Figure3和表格都来自 https://arxiv.org/pdf/2208.11174 ，这个表格（Table IV），列出了不同类型内存的访问延迟（以时钟周期为单位）：
- 全局内存（Global memory）: 290 cycles
- L2 缓存: 200 cycles
- L1 缓存: 33 cycles
- 共享内存（Shared Memory）: 读取23 cycles，写入19 cycles
我后面也找到了这个paper里面做micro benchmark的代码：https://www.stuffedcow.net/research/cudabmk?q=research/cudabmk ，后面如果有空继续阅读下这篇 paper 以及测试代码。
这张Slides讲述了延迟（latency）在计算机系统中的重要性和一些相关概念。
- 标题 "It's the latency stupid" 强调了延迟的重要性。
- 吞吐量（Throughput）和延迟（Latency）的对比：
- 吞吐量容易提高，但延迟却很难降低。
- 举例说明：即使你可以并行使用80条电话线，每条线传输一个比特，但100毫秒的延迟仍然存在。
- 量化（Quantization）技术：
- 用于减少数据包大小的一种方法。
- 例如，Bolo（可能是某个系统或协议）尽可能使用字节（byte）而不是16位或32位字来减少数据包大小。
- 底部提供了一个网址链接，包含更多关于这个话题的详细讨论。
这张Slides开始介绍内存合并（Memory Coalescing）的概念。我们无法减少延迟，但可以通过读取连续的内存元素来隐藏延迟。Slides建议在进行案例研究时要关注以下三个方面：
- DRAM Throughput（DRAM吞吐量）
- Duration（持续时间）
- L1 cache throughput（L1缓存吞吐量）
这里说的内存合并的案例就是 https://github.com/cuda-mode/lectures/blob/main/lecture_008/coalesce.cu 这里所展示的。代码如下：
#include <iostream>
#include <cuda_runtime.h>
__global__ void copyDataNonCoalesced(float *in, float *out, int n) {
int index = blockIdx.x * blockDim.x + threadIdx.x;
if (index < n) {
out[index] = in[(index * 2) % n];
}
}
__global__ void copyDataCoalesced(float *in, float *out, int n) {
int index = blockIdx.x * blockDim.x + threadIdx.x;
if (index < n) {
out[index] = in[index];
}
}
void initializeArray(float *arr, int n) {
for(int i = 0; i < n; ++i) {
arr[i] = static_cast<float>(i);
}
}
int main() {
const int n = 1 << 24; // Increase n to have a larger workload
float *in, *out;
cudaMallocManaged(&in, n * sizeof(float));
cudaMallocManaged(&out, n * sizeof(float));
initializeArray(in, n);
int blockSize = 128; // Define block size
// int blockSize = 1024; // change this when talking about occupancy
int numBlocks = (n + blockSize - 1) / blockSize; // Ensure there are enough blocks to cover all elements
// Launch non-coalesced kernel
copyDataNonCoalesced<<<numBlocks, blockSize>>>(in, out, n);
cudaDeviceSynchronize();
initializeArray(out, n); // Reset output array
// Launch coalesced kernel
copyDataCoalesced<<<numBlocks, blockSize>>>(in, out, n);
cudaDeviceSynchronize();
cudaFree(in);
cudaFree(out);
return 0;
}这里段程序比较简单，用于演示内存合并（Memory Coalescing）的概念和其对性能的影响。它主要做了以下事情：
- 定义了两个CUDA kernel：
- copyDataNonCoalesced kernel：非合并内存访问模式，以非连续的方式读取输入数组（使用 (index * 2) % n 作为索引），这种访问模式会导致非合并的内存访问，可能降低性能。
- copyDataCoalesced kernel：合并内存访问模式，以连续的方式读取输入数组（直接使用 index 作为索引），这种访问模式允许合并内存访问，可以提高性能。
- 主函数：
- 分配统一内存（Unified Memory）用于输入和输出数组，初始化输入数组。
- 设置CUDA网格和块的大小，分别运行非合并和合并的kernel，在每次kernel执行后使用 cudaDeviceSynchronize() 确保GPU操作完成。
接着使用nvcc -o benchmark coalesce.cu来编译程序，然后执行ncu benchmark来Profile程序。
对于copyDataNonCoalesced kernel来说，DRAM内存吞吐量大约是89%，L1 Cache的吞吐量是30%，kernel的执行时间是764us。
对于copyDataCoalesced kernel来说，L1 Cache的吞吐量大约是37%，DRAM内存吞吐量是82%，执行时间是558us。
我们可以看到合并内存访问的kernel是有明显的性能提升的。可以预见，随着输入数据量的增大合并内存访问的优势会更明显。ncu的结果里面还提示计算的理论occupancy（100.0%）和实测的实际occupancy占用（77%）之间的差异可能是由于 kernel 执行期间的warp调度开销或工作负载不平衡导致的。在同一kernel 的不同块之间以及块内的不同 warps 之间都可能发生负载不平衡。把上面程序中的int blockSize = 128改成int blockSize = 1024再次用ncu profile，可以发现occupancy提升到了85.94%。
这张Slides讨论了GPU中的占用率（Occupancy）问题，主要内容如下：
- 两种quantization问题：
- a) Tile quantization：矩阵维度不能被线程块Tile大小整除。
- b) Wave quantization：Tile总数不能被GPU上的SM（流多处理器）数量整除。
- 性能图表比较和分析：
- 左图(a)：cuBLAS v10 上 NN GEMM 的性能
- 右图(b)：cuBLAS v11 上 NN GEMM 的性能
- 两图都是在 M = 1024, N = 1024 的矩阵维度下进行的测试
- 左图(a)显示性能呈现明显的阶梯状，有大幅波动。
- 右图(b)显示性能波动较小，整体更加平滑。我们可以看到cuBLAS v11 可能采用了更好的调度策略或优化技术，减少了由于Tile和Wave Quantization 导致的性能波动。
这张Slides讲解了在PyTorch中使用padding（填充）来解决Tensor Core矩阵乘法维度要求的问题。具体内容如下：
- 在PyTorch环境中，使用padding是解决某些问题的方法。
- 表格展示了不同cuBLAS和cuDNN版本下，使用Tensor Core的数据精度要求。这些要求适用于矩阵维度M、N和K。
- 版本区分：
- 左列：cuBLAS < 11.0 和 cuDNN < 7.6.3 的旧版本
- 右列：cuBLAS ≥ 11.0 和 cuDNN ≥ 7.6.3 的新版本
- 数据类型的要求：
- INT8：旧版本要求16的倍数；新版本总是可用，但16的倍数最高效，在A100上128的倍数最佳。
- FP16：旧版本要求8的倍数；新版本总是可用，但8的倍数最高效，在A100上64的倍数最佳。
- TF32：旧版本不适用；新版本总是可用，但4的倍数最高效，在A100上32的倍数最佳。
- FP64：旧版本不适用；新版本总是可用，但2的倍数最高效，在A100上16的倍数最佳。
新版本的cuBLAS和cuDNN提供了更灵活的Tensor Core使用条件。而A100 GPU可能需要更大的倍数来获得最佳性能。Padding可以用来将矩阵维度调整为这些推荐的倍数，以提高性能。
在CUDA中提升Occupancy的一个方法是修改kernel。
CUDA Occupancy calculator工具可以帮我们自动计算达到更好Occupancy的kernel启动参数，在上一节合并访存的.cu中调用这个Api结果显示，对于T4 GPU，最优的配置是网格大小为40，块大小为1024。代码见：https://github.com/cuda-mode/lectures/blob/main/lecture_008/occupancy.cu
在对这个程序进行ncu的时候有新的问题，那就是下面所展示的：
警告（WRN）：内存利用率高于计算利用率：请查看内存工作负载分析部分以识别DRAM瓶颈。检查内存重放（合并）指标，以确保您正在有效利用传输的字节。同时考虑是否可以通过每次内存访问执行更多工作（kernel融合）或是否有可以（重新）计算的值。
接下来开始讨论这个问题
讨论之前需要先了解一下这张Slides展示的Roofline模型，它决定了一个cuda kernel是compute bound还是memory bound。
这张Slides讲解了算术强度（Arithmetic intensity）的概念及其在处理器性能分析中的应用。这个slides来自gtc2019的一个讲解。
左侧指标是数学运算和内存操作的算法混合，称为算术强度。右侧指标是处理器的ops/byte比率。例如，V100 GPU可以执行125/0.9=139 FLOPS/B。比较算术强度和ops/byte比率可以指出算法受什么因素限制。
下面还给出了操作类型及其算术强度表格：
- Residual addition（残差加法）：0.166，受内存限制
- ReLU activation（ReLU激活）：0.25，受内存限制
- Batch normalization（批量归一化）：O(10)，受内存限制
- Convolution（卷积）：1-10000+（假设FP16数据），可能受内存或数学运算限制
链接：https://developer.download.nvidia.com/video/gputechconf/gtc/2019/presentation/s9926-tensor-core-performance-the-ultimate-guide.pdf
这张slides讲解了ReLU（Rectified Linear Unit）函数的算术强度分析：
- ReLU函数定义：f(x) = max(0, x)，应用于向量的每个元素。
- 操作描述：对每个元素进行1次读取、1次比较操作，可能还有1次写入。
- 数据类型：假设使用float32，即每个数占4字节（32位）。
- 计算分析：
- 操作数（Ops）：1（每个元素一次比较操作）
- 字节数（Byte）：2 * 4 = 8（读取和可能的写入，每次4字节）
- 算术强度计算：
- 最坏情况：1/8（当每个元素都需要写入时）
- 最好情况：1/4（当不需要写入时，只有读取操作） 结论：1/4 < 1，表明ReLU操作受内存带宽限制（Memory bound）
这张Slides对Float16的ReLU进行了算术强度分析，可以看打这种情况下最坏的算术强度是1/4，而不是Float32时的1/8，因此量化是可以提高计算强度的。
这张Slides讲解了矩阵乘法（Matmul）的算术强度分析。其中：
- FLOPS（浮点运算次数）计算：
- 对C中的每个输出元素，需要A的一行和B的一列做点积
- 需要N次乘法和N次加法
- 总FLOPS = M * K * 2N
- 字节数计算：
- 加载矩阵A和B：MN + NK
- 写入输出矩阵C：MK
- 总字节数 = MN + NK + MK
- 算术强度（AI）计算：
- AI = 2MNK / (MN + NK + MK)
- 结论：
- 对于大型矩阵，计算受限（Compute bound）
- 否则，带宽受限（Bandwidth bound）
这张Slides总结了如何优化不同类型的kernels：
- 带宽受限的kernel（Bandwidth Bound Kernels）优化策略：
- Fuse（融合）：合并多个操作以减少内存访问
- Quantize（量化）：使用更小的数据类型来减少内存传输
- Compile（编译）：可能指使用特定的编译技术来优化内存访问模式
- 计算受限的kernel（Compute Bound Kernels）优化策略：
- Write a better algorithm（编写更好的算法）：这意味着需要从算法层面进行优化
关于矩阵乘法Tiling减少全局内存访问请查看以前的CUDA-MODE 课程笔记 第四课: PMPP 书的第4-5章笔记 。
这张Slides对应这里的代码：https://github.com/cuda-mode/lectures/blob/main/lecture_008/divergence.cu ，主要是对下面2个kernel进行分析：
__global__ void processArrayWithDivergence(int *data, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < N) {
if (data[idx] % 2 == 0) {
data[idx] = data[idx] * 2; // 注意这个分支比下面的分支要慢，可能一个Warp里执行这个分支的线程会落后，Warp里的其它线程必须等待这些线程计算完成
} else {
data[idx] = data[idx] + 1;
}
}
}
__global__ void processArrayWithoutDivergence(int *data, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < N) {
int isEven = !(data[idx] % 2); // 这里做的事情和上面相同，但是规避了线程分化问题
data[idx] = isEven * (data[idx] * 2) + (!isEven) * (data[idx] + 1);
}
}- 控制分歧（control divergence）与占用率（occupancy）有关，但如果条件语句导致大量线程闲置，这是不好的。
- processArrayWithDivergence 耗时 0.074272 毫秒; processArrayWithoutDivergence 耗时 0.024704 毫秒;这表明去除control divergence可以显著提高性能（约3倍）。
- "ncu --set full divergence" 用这行命令来设置线程control divergence分析。
对于compute bound的kernel，让线程可以做更多工作，可能会更快。
- 性能比较：
- 运行命令：main ~/lecturex ./benchmark
- VecAdd 执行时间：0.245600 ms
- VecAddCoarsened 执行时间：0.015264 ms
- 关键观察：
- VecAddCoarsened启动了一半的线程数量
- 尽管线程数减少，但执行速度显著提高（约16倍）
这里的代码在 https://github.com/cuda-mode/lectures/blob/main/lecture_008/coarsening.cu 。
这也许可以解释Lecture 7中为什么对于Int4 Weight Only量化的高效kernel实现比普通的fp16的Kernel跑得更快。
这张Slides讨论了在GPU编程中的"私有化"（Privatization）技术。要点为：
- 将部分更新应用到数据的私有副本上，然后再写回全局或共享内存。
- 示例：
- 滑动窗口算法（Sliding window algorithm）
- 图示：1 2 [3] [4] [5] 6 7
- 这表明算法在一个局部窗口内进行操作。
- Privatization的优势：
- 更高的占用率（Higher occupancy）
- 更高的计算SM吞吐量（Higher compute SM throughput）
- 更低的DRAM吞吐量（Lower DRAM throughput）
这个滑动窗口算法对应的例子就是Mistral等大模型里面的滑动窗口注意力算法，
解释下这个图：
- 左侧矩阵：Vanilla Attention
- 展示了传统注意力机制，每个token都可以关注到所有其他token。
- 矩阵是下三角形，表示每个token只能关注到它自己和之前的token。
- 中间矩阵：Sliding Window Attention
- 展示了滑动窗口注意力机制，每个token只关注固定窗口大小内的相邻token。
- 这里窗口大小W=3，可以看到每个token只与前后3个token有连接。
- 右侧图：Effective Context Length
- 展示了多层滑动窗口注意力如何扩大有效上下文长度。
- 每一层都可以使信息向前传播W个token。
总结来说，传统注意力的操作数量与序列长度成二次方关系，内存使用随token数线性增长。在推理时，这会导致更高的延迟和更低的吞吐量，因为缓存可用性降低。滑动窗口注意力通过限制每个token最多只关注前一层的W个token来缓解这个问题。虽然窗口外的token不直接参与注意力计算，但它们仍然可以影响下一个词的预测。在每个注意力层，信息可以向前传播W个token。经过k个注意力层后，信息可以向前传播最多k × W个token。
继续讨论Privatization技术，https://github.com/cuda-mode/lectures/blob/main/lecture_008/privatization.cu 这里的核心代码是：
// CUDA kernel for vector addition without privatization
__global__ void vectorAdd(const float *a, const float *b, float *result, int n) {
int index = threadIdx.x + blockIdx.x * blockDim.x;
if (index < n) {
result[index] = a[index] + b[index];
}
}
// CUDA kernel for vector addition with privatization
__global__ void vectorAddPrivatized(const float *a, const float *b, float *result, int n) {
int index = threadIdx.x + blockIdx.x * blockDim.x;
if (index < n) {
float a_private = a[index]; // Load into private memory
float b_private = b[index]; // Load into private memory
result[index] = a_private + b_private;
}
}我们把a[index]，b[index]加载到private memory里面避免对全局内存的直接操作，但在这个VectorAdd的例子中没有加速。
但是在下面的滑动窗口求和的例子中通过把global memory加载到shared memory中，然后进行累加时求和操作就是在shared memory中进行操作。代码链接：https://github.com/cuda-mode/lectures/blob/main/lecture_008/privatization2.cu
// Kernel without privatization: Direct global memory access
__global__ void windowSumDirect(const float *input, float *output, int n, int windowSize) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int halfWindow = windowSize / 2;
if (idx < n) {
float sum = 0.0f;
for (int i = -halfWindow; i <= halfWindow; ++i) {
int accessIdx = idx + i;
if (accessIdx >= 0 && accessIdx < n) {
sum += input[accessIdx];
}
}
output[idx] = sum;
}
}
// Kernel with privatization: Preload window elements into registers
__global__ void windowSumPrivatized(const float *input, float *output, int n, int windowSize) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int halfWindow = windowSize / 2;
__shared__ float sharedData[1024]; // Assuming blockDim.x <= 1024
// Load input into shared memory (for demonstration, assuming window fits into shared memory)
if (idx < n) {
sharedData[threadIdx.x] = input[idx];
__syncthreads(); // Ensure all loads are complete
float sum = 0.0f;
for (int i = -halfWindow; i <= halfWindow; ++i) {
int accessIdx = threadIdx.x + i;
// Check bounds within shared memory
if (accessIdx >= 0 && accessIdx < blockDim.x && (idx + i) < n && (idx + i) >= 0) {
sum += sharedData[accessIdx];
}
}
output[idx] = sum;
}
}作者最后讲的一点是以Flash Attention为例，如果你可以从数学的角度重写算法，有可能可以让代码的性能大幅度提升。比如Flash Attention利用Safe Softmax的数学形式分块计算Attention。这部分讲解已经非常多了，大家可以参考 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/README.md 里面收集的Flash Attention相关的资料，最后这几张Slides就不再赘述了。
总结
这节课相当于对Lecture 1更系统的补充，重点介绍GPU kernel优化的一些实用技术和分析工具ncu。课程深入探讨了“性能检查清单”，概述了关键的优化策略：
- 合并全局内存访问：确保线程访问连续的内存位置，以最大限度地提高内存带宽利用率。
- 最大化占用率：优化kernel启动参数，以充分利用 GPU 的处理能力。
- 理解内存与计算限制：分析kernel特征以确定限制因素（内存带宽或计算能力），并应用适当的优化技术。
- 最小化线程的分化：避免导致warp内的线程采用不同执行路径的条件语句，从而导致性能下降。
- Tiling以重用数据：组织数据访问模式，以最大限度地提高缓存中的数据局部性和重用率。
- 私有化：利用私有内存（寄存器或共享内存）来减少全局内存访问并提高占用率。
- 线程粗化：调整线程粒度以平衡工作负载并最小化线程开销。
- 算法重写：探索替代的数学公式或算法以提高计算效率。
up主还通过cuda代码示例和使用 ncu 工具的分析结果来说明这些概念。它展示了内存合并、占用率优化和控制分歧最小化的好处。它还介绍了用于分析kernel性能瓶颈的 Roofline 模型，以及算术强度的概念，以确定kernel是受内存限制还是受计算限制。接着up主讨论了私有化技术，重点介绍了在滑动窗口算法中使用共享内存来改善数据局部性和减少全局内存访问。最后，简单描述了一下通过从数学角度重写算法来获得显著性能提升的潜力，并以 Flash Attention 为主要示例，这里就没有截最后几张slides了，毕竟Flash Attention的资料非常多了。
三、手写CUDA卷积算子
这里主要介绍如何利用CUDA实现一个2D卷积算子，实现过程较为简单，最终的实现效果可以在较小的尺寸下取得比cudnn快较大的性能。实测在以下参数配置下可以达到平均
1.2倍cudnn的性能。
CUDA介绍（from chatGPT）
现在深度学习大行其道，作为深度学习的基础软件设施，学习cuda也是很有意义的。本篇文章主要介绍如何利用CUDA实现一个2D卷积算子，实现过程较为简单，最终的实现效果可以在较小的尺寸下取得比cudnn快较大的性能。实测在以下参数配置下可以达到平均1.2倍cudnn的性能（娱乐结果，还与cudnn配置有关，更小更快）。
TIPS: 跳过cudnn初始化的时间，99轮平均时间
const int inC = 6;
const int inH = 768;
const int inW = 512;
const int kernelH = 6;
const int kernelW = 6;
const int outC = 6;
const int outH = inH - kernelH + 1;
const int outW = inW - kernelW + 1;1 卷积操作通俗介绍
1.1 数据布局(data layout)
卷积操作主要针对图像进行运算，我们常见的RGB即为三通道的二维图像，那么就可以通过一个一维数组存储所有的数据，再按照不同的布局去索引对应的数据，现在主要使用nchw和nhwc两种数据布局，其中
n - batch size 也可以理解为"图像"数量
c - channel num 即我们说的通道数量
h - height 图像高度，每个通道的高度宽度是一致的
w - width 图像宽度
那么显然nchw就是逐个通道的读取图像，nhwc即对所有通道的同样位置读取数据后，再切换到下一个为止
一个是优先通道读取，一个是优先位置读取
还有一种CHWN布局，感觉比较奇怪，并未过多了解
详细的可以参考英伟达官方文档Developer Guide : NVIDIA Deep Learning cuDNN Documentation (https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html)
nchw layout
nhwc layout
本文是按照nchw数据格式来进行算子的实现的。
1.2 直接卷积
相信大家都或多或少听过卷积，可以通过gpt的回答来直观地认识卷积操作
最基本的直接卷积操作是十分简单的，你可以想象一个滑动的矩阵窗口在原矩阵上移动，对应位置进行点积，得到结果后求和放到目标矩阵上，可以用以下图像直观地理解这一过程，向老师称为对对碰:)
图源：国科大模式识别课程
你会注意到上述过程中怎么没有什么channel的参与，只有一个矩阵
多输入通道的情况下，就是对每个通道的相同位置分别与卷积核进行对对碰，结果累加作为输出矩阵值；
多输入多输出通道，即对每个输出通道都进行上述操作
对于通道的理解建议参考[@双手插袋]的文章CNN卷积核与通道讲解 (https://zhuanlan.zhihu.com/p/251068800)
那么我们需要知道的是直接卷积操作其实就是原矩阵与卷积核间的对对碰，产生所谓的特征图feature map，十分的简单，这也方便我们对其进行并行任务划分
注意到上述文章中并没有提到padding和stride，本篇文章并没有针对padding和stride的实现
padding
padding是作为对图像的填充，可以发现上面的特征图尺寸缩小了一圈，是因为直接卷积势必会造成这一结果
通过padding可以加强图像边缘特征，避免边缘特征被忽略
stride
stride可以简单的理解为跨步，即上面的小窗口在矩阵上滑动的步长，默认为1
即上述图像中下一次卷积的中心应该是4为中心的3*3子矩阵
如果你设置为2，那么下一次是3为中心的3*3子矩阵了
1.3 其他卷积计算方法
除去直接卷积，也有一些其他方法进行卷积，感兴趣的读者可以自行了解，仅举以下几例参考
Img2col
即把图像展开为一个行向量组，卷积核/滤波器（kernel/filter）展开为一列或多列向量，转化为矩阵乘去计算卷积结果
FFT method
利用傅里叶变换的频域变换去做卷积，这样做的优势是计算量会小很多
Winograd Algorithm
也是一种将图像变换到另外一个空间再去做运算再做变换得到结果，会减少很多乘法运算
2 整体实现思路
2.1 block与thread划分
首先我们需要考虑如何对代表图像的多通道矩阵来进行block与thread的划分，这一部分是有说法的
不同的切分方式会让block在SM上的流转效率有很大的差别
本文仅提供一个十分草率的切分，我们都清楚目前在英伟达的GPU上，任务的调度最小单元是warp
一个warp以32个线程为一组，故通过8*4的block来进行矩阵的切分，每个block里共32个位置
这样可以保证每个block上到SM时不用去与其他的block拼接线程，产生额外开销
注意我这里用的是位置，并不是元素，32个线程，每个线程去负责一个位置的计算
以16*20的矩阵为例，对其进行划分的结果如下图所示，(x,y)是笛卡尔坐标系，与行主序不同
2.2 数据转移
- 关于位置和规模(size)
那么为什么说一个block有32个位置，而不是32个元素呢，首先注意到卷积操作虽然遍历到了原矩阵的所有元素
但是你按中心点的位序去数的话（以卷积核3*3为例），结果应该是这个样子
注意这里仅示意卷积中心点范围，请与后文作区分
按3*3矩阵的中心来看，中心正好是去掉外面一圈的位置，按照左上角元素来看，恰好应该是（左上角，右下角）
这样一个区间，参数解释如下
row_num 原矩阵中一行元素的数目
inH inW 原矩阵的H W
kernelH kernelW 卷积核的H W
outH outW 输出矩阵的H W
当然你也可以用中心点而不是左上角的元素作为窗口的标识来设计算法
恰巧你上面算出来的这个范围也正是你得到的feature map的下标范围
我们也可以得到输出矩阵的规模为
请注意大小和位置下标的区别，一个从1开始数一个从0开始数
- 一个block的数据转移
确定了整体的尺寸，那么我们来看一个block需要的数据尺寸是多少呢
显然你可以发现，对于输出矩阵进行block划分是更合理的，这样可以保证一个block
32个位置恰好对应输出矩阵的32个位置，而不用过多的去考虑输出矩阵的排布
那么对于上述提到的划分，可以通过下图来直观感受block划分在原矩阵的效果
22*18的in产生20*16的out
那么一个block用到的元素范围应该是哪些呢，我们要做的是卷积操作，每个中心点应该有对应卷积核大小的矩阵参与运算，那么以(0,0)和(4,1)的block为例，给出他们的涉及原矩阵范围如下图所示
蓝色为一个block需要用到的原矩阵元素
那么我们可以确定一个block，8×4的情况下需要读取10×6的原矩阵的元素，也是+kernelH-1来确定的
那么对应输出矩阵就是一个萝卜一个坑了，不需要额外考虑
这样就确定了一个block需要从GMEM到SMEM的元素范围
至于怎么转移，我们在代码实现中讲述，当然你可以单独指定某几个进程去完成所有的转移任务
2.3 计算逻辑
- 不考虑channel
不考虑channel的情况下，即单输入通道单输出通道单卷积核这样最简单的情况
我们只需要做三件事
① 将block对应的数据转移到SMEM中
② 利用线程的tid去计算对应输出矩阵位置的结果
③ 将结果写回输出矩阵
- 只考虑inC
这种情况下我们要做的额外的事儿就多一点
加一层循环，让每个线程计算多个in channel的数据，并累加起来作为结果
需要用到一个寄存器来存储这个中间结果
- 考虑inC与outC
其实要做的事情也就比上面多一点，就是开大点空间
让线程去存储多个outC的中间结果，分别累加
最后写回的时候也分别写回即可
3 详细实现过程
3.1 整体实现思路
主要从自己的角度出发去还原怎样一步步构造出这样一个初级的算法
首先实现一个最简单的版本，CPU串行版本，并保证CPU串行版本可以获取正确的结果
此后再在其基础上进行并行化的改造，而直接卷积运算的过程其实相对是比较简单的
我们在不考虑padding与stride的情况下，是可以不借助任何参考资料来直接完成第一版代码的
3.1.1 CPU串行版本的卷积算子
#define element_type float
#define OFFSET(row, col, ld) ((row) * (ld) + (col))
/*
@brief: 串行卷积实现 CPU代码 NCHW
@param in inC inH inW: 输入矩阵(数组) channel height width
@param out outC outH outW: 输出矩阵 channel height width
@param kernel kernelH kernelW: 卷积核 height width
*/
void serial_convolution(element_type *in, element_type *out, element_type *kernel, int batch_size,
int inC, int inH, int inW,
int outC, int outH, int outW,
int kernelH, int kernelW)
{
float val;
int out_pos, in_pos, kernel_pos;
for (int oc = 0; oc < outC; oc++) // 每个输出通道
{
// 对一个位置的操作 用当前输入channel卷积去对相应的输出channel
// 保证每个outChannel都是多inChannel累积的结果
for (int i = 0; i < outH; i++)
{
for (int j = 0; j < outW; j++)
{
val = 0; // 避免累积和需要多次读取写入
out_pos = oc * outH * outW + OFFSET(i, j, outW);
for (int ic = 0; ic < inC; ic++) // 对每个输入通道
{
for (int ii = 0; ii < kernelH; ii++)
{
for (int jj = 0; jj < kernelW; jj++)
{
in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);
kernel_pos = oc * kernelH * kernelW + OFFSET(ii, jj, kernelW);
val += in[in_pos] * kernel[kernel_pos];
}
}
}
out[out_pos] = val; // 与cudnn计算结果为相反数
}
}
}
}这是我最终完成的CPU串行版本代码，可以发现套了足足有5层循环
在我们传统观念中，这可是 O(n5)O(n^5)O(n^5) 的最笨算法了
不过没有关系，我们关注的并不是他的性能，cuda上也不会去跑这一版代码
我们需要关注的是怎么样能得到正确的结果，且如何设计循环的嵌套关系来使用尽量少的访存次数
使用尽量多的本地中间结果，这样可以尽可能地减少我们的算法在访存方面的消耗
要明白GPU上的线程如果去读GMEM上的数据需要几百个时钟周期，读SMEM需要几十个时钟周期
读取SM上的寄存器需要的时钟周期会更少！
因此我们需要竭力优化的一部分是如何减少访存，多用本地存储来代替
另一方面这也是因为计算本身是十分简单的点积，不太可能去做出更大的优化
3.1.2 循环顺序设计
逐层去观察循环的嵌套顺序，发现是
outC-->H-->W--->inC-->kernelH-->kernelW
这样的计算顺序不一定是最优化的，笔者也没有进行详细的计算论证，但是这样的计算顺序是出于以下角度考虑
① 多通道卷积结果的维度/通道数/feature map数就是我们的outC，是我们最终要写回的out矩阵的维度，将其放在最外层循环，作用是：
- 一次循环内完成这个out channel中的所有计算，再接着进行下一个outC的计算
- 由于out数据是在一维数组中存储，且为nchw格式，那么不同outC中的数据跨度其实是很大的，连续的完成一个outC的内容可以更好的利用局部性原理
- 个人理解逐个outC的计算是很是一种比较直观和自然(方便想象与理解)的角度
- 串行过程中我们可以使用尽量少的中间变量去维护中间结果，如果你先遍历inC后遍历outC的话，其实你是需要维护outC个中间变量的
这样的顺序也是在后面做并行化改造过程中逐渐发现的一个较为合理的顺序，我们可以在后文中更加直观的感受到这样设计的优势
② 出于nchw布局的涉及，H W的顺序是基本固定的，当然你也可以先W后H，不过一般是行主序存储.. 还是先H比较快一些
③ inC为何出现在H W之后？请回顾多通道卷积的过程，一个feature map的值是由多个inC与kernel分别点击累加形成的，如果你将inC放置在H W之前的话，在下方的代码中，你是不是就需要设置height×width个中间变量来存储这里的val值呢？
in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);
kernel_pos = oc * kernelH * kernelW + OFFSET(ii, jj, kernelW);
val += in[in_pos] * kernel[kernel_pos];将inC放置在H W之后，是相当于在一个outC上进行计算，对不同inC同样的位置分别计算得到了val的准确值，最终写回，这样在串行的版本中，我们只需要一个float即可存储好中间结果来避免空间的浪费！
TIPS：注意上方对于下标的计算，我们以两个位序举例说明
in_pos = ic * inH * inW + OFFSET(i + ii, j + jj, inW);nchw的数据布局格式下，这里是默认n为1的，注意本文所有的实现都是建立在n假设为1的情况，其实n为更大值也不是很有意义，这样的布局下，下一张图像在计算意义上是没有任何差别的，无非是你将数据的起始地址跳过一大部分，切到下一张图像
说回这个式子，其中ic为in channel，inH inW分别是输入矩阵的高度与宽度，后面宏定义的OFFSET其实就是简略写法，你也可以写成(i+ii)*inW + j + jj
in_pos的含义是在当前循环变量下输入矩阵的位置
同理，out_pos的计算是一样的
out_pos = oc * outH * outW + OFFSET(i, j, outW);ii和jj是相对于卷积核的相对位置循环变量，输出位置是用不到他们的
进行并行化改造
其实当你把串行版本设计明白后，你对于并行化改造的想法也差不多有个七七八八了
主要是出于以下三个角度去设计并优化的
① 尽量减少访存次数（当然不是不访问），尤其是减少访问GMEM的次数，善用SMEM与register
（对于GMEM SMEM和register等访存层次相关知识不熟的读者可以去了解一下CUDA的存储层次）
② 此外要划分明确各个线程要负责的任务区域和他的行为应达到的效果，做好下标计算
③ 计算行为是很快的，我们要尽可能去掩盖访存延迟，让线程去火力全开计算（预取prefetch）
下面的章节都是在并行化改造过程中的一些细节，代码其实是一版版写出来的，这里是对最终版本进行说明
（所谓的一版版就是划分出不同块，分别测试是否与预期一致，再去完成下面的块）
3.2 线程任务均分
这部分其实是源于 @有了琦琦的棍子 在GMEM讲解中的数据转移部分，基本算是照抄了
十分感谢前辈，不过还不知道这种方法的确切名字，目前暂时称为均分，其实思想是很朴素的
我们的block设计的是8*4的大小，对应32个线程，但是涉及到in矩阵的数据可不只是32个元素，那么
我们需要尽可能地平均分配任务给线程，保证每个线程承担差不多的任务量来达到更好的平均性能
差不多是因为，不太可能都是整除的情况
这部分主要通过图示讲解，自己设计的过程中大多是通过纸笔演算确定下标的
首先确定一些变量，注意CUDA的笛卡尔坐标系和笔者的行号row和列号col的区别
int block_row = blockIdx.y;
int block_col = blockIdx.x;
int thread_row = threadIdx.y, thread_col = threadIdx.x;
int tid = thread_row * threadW + thread_col;由于要重复使用inC内的数据，我们肯定是要开一个SMEM去存储这部分数据的，那么就有一个GMEM->SMEM的数据转移过程，以8×4的block和3×3的kernel为例，我们可以得到如下的景象
其中橙色部分是我们的block，一个tid(thread id)是一个线程，也是block中的一个位置，也是outC中的一个位置
那么白色部分就是我们在block范围之外但会用到的数据，这部分数据可以看到像两条网格
那么我们怎么把这些数据从GMEM转移到SMEM呢，首先我们考虑（以下部分为自己笨拙的思考过程）
方案① 边缘线程负责白色区域
橙色为仅负责自己的位置，紫色负责3个位置，红色负责9个
看起来是不是好像也还行，只要我们通过thread_row和thread_col判断一下当前进程是否在边缘
对这些进程进行单独的编码就可以了，不过在写代码前可以先算一笔账
这个网格共有10×6=60个元素，我们有32个线程，那么最好的情况下，是每个线程负责
60/32=1.875个元素，也就是花费1.875个单位时间（这里的单位时间是抽象概念，假定为每个线程处理每个元素的时间）
那么可以看一下这种划分方式下，每个线程平均负责的元素为
后面的项是权重，前面的项如 说明这个线程处理9个线程，那么花费的时间应当是9倍，所以性能应当是九分之一（相当于只处理一个元素的线程），且线程是warp调度的，32个线程里面有这么一个拖后腿分子，想必并行情况下整体花费时间是取决于这个31号线程的
这个方案的效率是理想情况的一半都不到，说明这种方案是不太可行的，写出来效果也不一定好呢，换！
方案② 平均划分
其实笔者也想过一些其他奇怪的方法，但是感觉平均思想似乎是最佳的，那么何不一步到胃呢？
我们先来定义一些变量，后面再来逐步解释
// 分块边界 boundary是限制正常范围 edge是需要补的范围
int row_boundary = outH / BLOCK_HEIGHT - 1,
col_boundary = outW / BLOCK_WIDTH - 1;
int row_edge = outH % BLOCK_HEIGHT, col_edge = outW % BLOCK_WIDTH;
···
int single_trans_ele_num = 4; // 线程一次转移的数据数
int cur_in_block_height = BLOCK_HEIGHT + KERNEL_HEIGHT - 1, // 读入in的block height
cur_in_block_width = BLOCK_WIDTH + KERNEL_WIDTH - 1, // 读入in的block width
in_tile_thread_per_row, // 以tile为单位转移数据，一行需要的thread数
in_tile_row_start, // tile的行起始位置
in_tile_col, // tile的列
in_tile_row_stride; // tile行跨度
// 修正边缘block尺寸
if (block_row == row_boundary)
{
cur_in_block_height = BLOCK_HEIGHT + row_edge + kernelH - 1;
}
if (block_col == col_boundary)
{
cur_in_block_width = BLOCK_WIDTH + col_edge + kernelW - 1;
}
in_tile_thread_per_row = cur_in_block_width / single_trans_ele_num;
in_tile_row_start = tid / in_tile_thread_per_row;
in_tile_col = tid % in_tile_thread_per_row * single_trans_ele_num;
in_tile_row_stride = thread_num_per_block / in_tile_thread_per_row;3.2.1 “block”设计与修正
不要急着头大，我们逐个说明，首先看顶头部分的变量，是关于限制范围的
因为我们要首先确定一个block内的线程要负责多少元素呢，因此需要界定这样的范围
我们前面只提到了block涉及到的in范围是扩大了一圈的，其实你的in矩阵相对于out矩阵也是多了一圈的
当多的这么一圈不能构成新的block时，那么注定我们的block网格是不能覆盖到out矩阵的！
我们还是上图比较直观
咱们的block网格只有16×20这么大，out矩阵有18×22这么大，明显可以看到蓝色的两条
是不足以构成新的block的，那么还有红色的部分，就是in矩阵的大小了，可以看到有20×24这么大
而我们的block是建立在out矩阵上的，所以我们起码也要覆盖到蓝色矩阵的所有范围吧
那么在不修改block尺寸的情况下，最简单的方法就是人为地去修正这些特定block的大小啦
修正后的block应该是这个样子的
修正后的block把out全覆盖了~
怎么修正呢？无非就是利用block位序去判断并修改尺寸啦，即这两行代码
// 修正边缘block尺寸
if (block_row == row_boundary)
{
cur_in_block_height = BLOCK_HEIGHT + row_edge + kernelH - 1;
}
if (block_col == col_boundary)
{
cur_in_block_width = BLOCK_WIDTH + col_edge + kernelW - 1;
}结合图片，是不是这些变量的概念就清晰了起来
注意我们所有变量都是有一个in的标识，这是标注in矩阵的范围
out矩阵的划分自然是有out的标识，且步骤都是一样的，只不过需要补的范围不太一样罢了
3.2.2 线程行为指定
还有一段代码我们没有解释，是这一段（thread_num_per_block本文默认为32，没有修改）
in_tile_thread_per_row = cur_in_block_width / single_trans_ele_num;
in_tile_row_start = tid / in_tile_thread_per_row;
in_tile_col = tid % in_tile_thread_per_row * single_trans_ele_num;
in_tile_row_stride = thread_num_per_block / in_tile_thread_per_row;这段我觉得是最抽象的部分也恰恰是最为精华的设计，首先要明确，是通过行里面的小片/tile作为线程处理的最小单元来进行设计的
其实变量名已经做了一部分的解释，可以大概解释为如下的含义
in_tile_thread_per_row 一行里面会有多少个tile
in_tile_row_start 当前线程负责的tile的起始行号
in_tile_col 当前线程负责的列号
in_tile_row_stride 如果还有元素要处理，那么需要跳过的行数/stride
好像不是那么的直观，我们再上一张图
左面是我们的block与in矩阵的关系，我们要把他都转移过来，且利用了fetch_float4的向量指令（也是single_trans_ele_num设置为4的原因）
以7号线程为例，当前的in_block为10×6大小，那么上面四个变量的值分别为1,7,0,32
这个例子比较简单，可以发现一行其实是有一个半的tile的，那么需要一点点小小的修正来让每个线程
读取4+2个元素，这点小小的修正我们可以看代码
那么再来一个复杂的例子，假设我们在考虑out矩阵的事情，那么一个线程负责一个元素的话
请问这种方式对嘛？
是不是直观上你感觉应该是这样的，他可以丝滑的衔接好每个元素，完成我们的分配~
那么给出我们利用这个均分思想让每个线程负责任务的代码如下，大家再想一想分配后的图像
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
/*do something*/
}浅浅一个for循环，只不过所有条件都是我们仔细设计的，循环内部就是每个线程根据这些位序
去对应的显存位置上对数据一通操作罢了
那么注意部分，线程在跨过一个stride时，这个单位是不是row？那么意味着0号线程在下次任务会踩到30号的位置！如下图所示
实际上的线程分配
这样才是正确的线程操作顺序，当然由于我们是通过CUDA并行计算的，实际上上半部分是并行的，下半部分是在0-29号线程完成了上面的任务后才进行计算的（注意他们是32个一组/warp调度上来执行的）
这样其实有个小隐患，30号和31号以及0,1号会对这两个位置上重复进行操作，如果他们的行为不一致的话
会导致我们的结果出错，本例中他们的行为是一致的，故无所谓先后
通过这样的机制，我们可以指定每个线程负责的元素位置以及个数（tile大小），灵活地应用于不同的任务！
3.3 预取机制
这部分就是很基本的数据预取，计算的效率远远大于访存，计算时读取数据进来，完成基本的运算
（复杂运算也不是一行代码可以解决的）
再把结果存到对应位置，我们发现是不是即使是计算你也需要访存，节省访存开销是十分重要的
整体的数据传输逻辑是GMEM->SMEM->register->GMEM->MEM
并没有使用到Constant Memory和Texture Memory，那么结合数据预取的机制下
整体的框架如下方伪代码所示
初始化我们所需要的所有变量并修正block规模;
分配好shared memory用于加速访存;
// 预读取第一个channel的数据
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
把in中的数据从GMEM转到SMEM;
}
// 预读取第一个kernel的数据 可以使用很简单的读取策略 因为数据很少
if (thread_row >= 0 && thread_row < KERNEL_HEIGHT && thread_col == 0)
{
把kernel的数据从GMEM转到SMEM;
}
__syncthreads();
// 这里oc在外ic在内的设计是为了方便写回
for (int oc = 0; oc < outC; oc++)
{
for (int ic = 0; ic < inC; ic++)
{
// i,j 是相当于当前block起始位置而言
// 用ic的每个block去对oc的kernel进行计算
for (int i = 0; i < cur_out_block_height && (out_tile_row_start + i) < cur_out_block_height;
i += out_tile_row_stride)
{
计算当前ic与oc的结果，存到register;
}
// 读取下一个in channel数据 3,932,160
if (ic + 1 < inC)
{
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
读取下一个channel的数据;
}
}
__syncthreads();
}
if (oc + 1 < outC)
{
读取下一个kernel数据;
}
__syncthreads();
// 注意这样的循环顺序下已经完成了一个outC的计算
for (int i = 0; i < cur_out_block_height && (out_tile_row_start + i) < cur_out_block_height; i += out_tile_row_stride;)
{
写回当前outC的数据;
}
// 预读取下一个in channel数据 需要注意这时候要从头读了
for (int i = 0; i < cur_in_block_height && in_tile_row_start < cur_in_block_height;
i += in_tile_row_stride)
{
读取第一个channel的数据;
}
}到这里其实我们就完成了大部分内容了，整体骨架就是这样，其余就是一些细节上的下标计算问题了
3.4 一些杂项却又需要细节
3.4.1 中间结果存储设计
可以看到我们的伪代码中循环顺序是先oc再ic
可以想象一下，如果你先ic再oc的话，这样确实是我们只需要遍历一遍ic，oc多次遍历
但是我们也要考虑写回部分，写回你还需要单独再去写，理论上先ic的话会快一些
这里就不给大家放图了，读者可以自己想象一下两种计算顺序的区别
需要注意的是
线程能利用的硬件资源是有限的，一个warp共用一个SM上的寄存器，具体到每个线程大概32-255个寄存器（来源于chatGPT,不严谨，需要核实，后面gpt又说v100一个线程可以用800个..）
总之我们还是能少用就少用几个
当register存不下我们这些中间变量，就会放到local memory中
所谓的local memory是位于GMEM上的，如果发生这种情况，每次读取中间结果
你还得跑到GMEM上去访存，是非常之浪费时间的
两种循环其实需要的register数目都是oc×2（2是因为你一个线程要负责好几个位置的）
出于修正考虑，哥们儿直接开4倍，保证不会越界
3.4.2 下标计算
这部分其实，你串行算的明白，你并行就算的明白，我们举几个例子来说明一下
FETCH_FLOAT4(load_reg[0]) =
FETCH_FLOAT4(in[begin_pos + OFFSET(in_tile_row_start + i, in_tile_col, inW)]);
s_in[in_tile_row_start + i][in_tile_col] = load_reg[0];
s_in[in_tile_row_start + i][in_tile_col + 1] = load_reg[1];
s_in[in_tile_row_start + i][in_tile_col + 2] = load_reg[2];
s_in[in_tile_row_start + i][in_tile_col + 3] = load_reg[3];这里是利用向量指令去一次读取4个32位数据，s_in是开在SMEM上的，in是GMEM上的一位数据
那么可以看这个后面的下标
begin_pos 代表当前block的起始位序
OFFSET 是一个宏定义，代表行×一行元素数目
in[xxx] 下标其实就是当前block位置+block内的位置
再看一个写入中间结果的位置
temp_pos = i / out_tile_row_stride + j +
oc * (cur_out_block_height / out_tile_row_stride + 1);这里要考虑到线程是在计算它负责的第几个元素，那么就要用i / out_tile_row_stride来判断
如果处理多个元素，那你还得用j来控制一下当前是第几个元素
还要考虑到不同的oc，一个oc内负责的元素有cur_out_block_height / out_tile_row_stride +1这么多个
我们再看一个
out_pos = oc * outH * outW +
block_row * BLOCK_HEIGHT * outW + block_col * BLOCK_WIDTH +
OFFSET(out_tile_row_start + i, out_tile_col + j, outW);首先略过几个oc的范围，再计算当前block的起始位置，再计算上block内的相对位置
每个下标都要明白其计算的含义，本例中有很多公共表达式没有提取出来提前计算，会影响一定性能
3.5 完整代码
完整代码已上传：
https://github.com/Pegessi/conv2d_direct
3.6 性能测试
虽然是娱乐测试，但是也严谨一点，可以发现这个代码会受channel数目影响很大
代码还有一点小bug，不过不影响你执行，大家可能会发现（亟待修复）
不同数据规模下性能在cudnn的1/10到10倍上下横跳，有空给大家测一下放个完整的图。
四、大模型手撕CUDA
在CUDA编程中实现不同矩阵运算和操作的优化技巧。
前段时间参加了一些大模型的面试，大部分都要手撕CUDA，因此也整体复习了一遍CUDA优化相关的内容，整理了一些高频题的基本写法，保存在这里也便于日后自己复习。当然，有些代码不一定是最优化解，比如GEMM，想要在面试短短的30分钟内写一个好的GEMM Kernel，是有些难度的。印象比较深刻的是，其中有一场面试2个多小时，一个小时问项目，剩下一个小时在写GEMM，说实话，如果不是事先有准备过一些，直接上手写优化版还是会有点慌。就自己的经验而言，命中率还挺高（目前还没有遇到这些题目之外的），考虑深度优化的，一般也不会让你在面试短短几十分钟手撸出来。要是遇到有面试官让你手撸一个FlashAttention，那说明，你们实在是没有缘分，还是尽早好聚好散的好，或者提前结束面试，把时间省下来，出去吃顿烧烤也不错...，附GitHub链接：
https://github.com/DefTruth/cuda-learn-note
TIPS: 文章整理为方便自己复习，不喜欢的请自动跳过哈。
01 高频面试题汇总简介
相关kernel如下。也就是不到1000行代码，建议背下来，我个人是喜欢背记，背的过程中基本就慢慢理解所有细节。当然，每个人的学习方法都不一样哈，自己觉得舒服就行。
- sgemm naive, sgemm + block-tile + k-tile + vec4
- sgemv k32/k128/k16 kernel
- warp/block reduce sum/max, block all reduce + vec4
- dot product, dot product + vec4
- elementwise, elementwise + vec4
- histogram, histogram + vec4
- softmax, softmax + vec4 (grid level memory fence)
- sigmoid, sigmoid + vec4
- relu, relu + vec4
- layer_norm, layer_norm + vec4
- rms_norm, rms_norm + vec4
- ....
题内话，大模型相关的岗位，手撕CUDA的概率非常大，leetcode反而写的少，就前段时间个人的经验，基本是4:1的比例，还是建议好好复习下CUDA。当然，这些只是最简单的kernel实现，比如flash_attn，FMHA这些优化手段，就不在这篇文章里写了，面试中基本都会问到。FlashAttention系列原理详解，可以看我写的另一篇文章：
DefTruth：[FlashAttention系列][2w字] 原理详解: 从Online-Softmax到FlashAttention-1/2/FlashDecoding629(https://zhuanlan.zhihu.com/p/668888063)
02 sgemm naive, sgemm + block-tile + k-tile + vec4
#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#include <vector>
#include <algorithm>
#include <cuda_runtime.h>
#define WARP_SIZE 32
#define INT4(value) (reinterpret_cast<int4*>(&(value))[0])
#define FLOAT4(value) (reinterpret_cast<float4*>(&(value))[0])
// SGEMM: Block Tile + K Tile, with smem
// Block Tile (BM, BN) + K Tile (BK=32)
// grid((N + BN - 1) / BN, (M + BM - 1) / BM), block(BN, BM)
// a: MxK, b: KxN, c: MxN, compute: c = a * b, all row major
__global__ void sgemm(float* a, float* b, float* c, int M, int N, int K) {
// [1] Block Tile: 32x32的block处理c上一块32x32的元素计算
// [2] K Tile: 使用共享内存，并将K分块为BK大小的块
constexpr int BM = 32;
constexpr int BN = 32;
constexpr int BK = 32;
__shared__ float s_a[BM][BK], s_b[BK][BN];
int bx = blockIdx.x;
int by = blockIdx.y;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tid = threadIdx.y * blockDim.x + tx; // tid within the block
// load values to shared memory, 32x32 threads working together
// to fetch data along the row direction of a and b both for s_a
// and s_b 32x32x4x2=8KB, we use 32x32 threads within block to
// load 32x32 elements from global memory to shared memory, namely,
// each thread will load 1 element.
int load_smem_a_m = tid / 32; // 0~31, tid / 32, tid / BM, threadIdx.y
int load_smem_a_k = tid % 32; // 0~31, tid % 32, tid % BK, threadIdx.x
int load_smem_b_k = tid / 32; // 0~31, tid / 32, tid / BK, threadIdx.y
int load_smem_b_n = tid % 32; // 0~31, tid % 32, tid % BN, threadIdx.x
int load_gmem_a_m = by * BM + load_smem_a_m; // global row of a and c
int load_gmem_b_n = bx * BN + load_smem_b_n; // global col of b and c
// if (load_gmem_a_m >= M || load_gmem_b_n >= N) return;
float sum = 0.f;
for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {
int load_gmem_a_k = bk * BK + load_smem_a_k;
int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
s_a[load_smem_a_m][load_smem_a_k] = a[load_gmem_a_addr];
int load_gmem_b_k = bk * BK + load_smem_b_k;
int load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n;
s_b[load_smem_b_k][load_smem_b_n] = b[load_gmem_b_addr];
__syncthreads();
#pragma unroll
for (int k = 0; k < BK; ++k) {
int comp_smem_a_m = load_smem_a_m;
int comp_smem_b_n = load_smem_b_n;
sum += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];
}
__syncthreads();
}
int store_gmem_c_m = load_gmem_a_m;
int store_gmem_c_n = load_gmem_b_n;
int store_gmem_c_addr = store_gmem_c_m * N + store_gmem_c_n;
c[store_gmem_c_addr] = sum;
}
// SGEMM: Block Tile + Thread Tile + K Tile + Vec4, with smem
// BK:TILE_K=8 BM=BN=128
// TM=TN=8 增加计算密度 BM/TM=16 BN/TN=16
// dim3 blockDim(BN/TN, BM/TM);
// dim3 gridDim((N + BN - 1) / BN, (M + BM - 1) / BM)
__global__ void sgemm_thread_tile_vec4(
float* a, float* b, float* c, int M, int N, int K) {
// [1] Block Tile: 一个16x16的block处理C上大小为128X128的一个目标块
// [2] Thread Tile: 每个thread负责计算TM*TN(8*8)个元素，增加计算密度
// [3] K Tile: 将K分块，每块BK大小，迭代(K+BK-1/BK)次，
// 每次计算TM*TN个元素各自的部分乘累加
// [4] Vectorize: 减少load和store指令，使用float4
constexpr int BM = 128;
constexpr int BN = 128;
constexpr int BK = 8;
constexpr int TM = 8;
constexpr int TN = 8;
int bx = blockIdx.x;
int by = blockIdx.y;
int tx = threadIdx.x;
int ty = threadIdx.y;
int tid = threadIdx.y * blockDim.x + tx; // tid within the block
__shared__ float s_a[BM][BK], s_b[BK][BN]; // 2*128*8*4=8KB
// 0. 先计算shared memory中的索引
// tid和需要加载的smem s_a[BM][BK] 之间的索引关系 BM=128 BK=8 按行读取 A行主序
// 对于s_a每行8个数据，每个线程读取4个，需要2个线程；总共128行，需要128x2刚好256线程
int load_smem_a_m = tid / 2; // tid/2 (128/8)*(128/8)=256 threads per block, tid/2->[0,128), BM=128 0~127
int load_smem_a_k = (tid % 2 == 0) ? 0 : 4; // (tid%2 == 0) ? 0 : 4, col of s_a 0,4
// tid和需要加载的smem s_b[BK][BN] 之间的索引关系 BK=8 BN=128 按行读取 B行主序
// 对于s_b每行128个数据，每个线程读4个数据，需要32个线程；总共8行，需要32x8=256个线程
int load_smem_b_k = tid / 32; // tid/32, row of s_b 256/32=8 行 0~7
int load_smem_b_n = (tid % 32) * 4; // (tid % 32) * 4, col of s_b 0,4,...,124
// 1. 再计算全局内存中的索引
// 要加载到s_a中的元素对应到A全局内存中的行数 每个block负责出C中大小为BM*BN的块
int load_gmem_a_m = by * BM + load_smem_a_m; // global row of a and c
int load_gmem_b_n = bx * BN + load_smem_b_n; // global col of b and c
float r_c[TM][TN] = {0.0}; // 8x8
// 2. 先对K进行分块，每块BK大小
for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {
// 加载数据到共享内存smem s_a BM*BK 128*8 vectorize float4
int load_gmem_a_k = bk * BK + load_smem_a_k; // global col of a
int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
FLOAT4(s_a[load_smem_a_m][load_smem_a_k]) = FLOAT4(a[load_gmem_a_addr]);
// 加载数据到共享内存smem s_b BK*BN 8*128 vectorize float4
int load_gmem_b_k = bk * BK + load_smem_b_k; // global row of b
int load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n;
FLOAT4(s_b[load_smem_b_k][load_smem_b_n]) = FLOAT4(b[load_gmem_b_addr]);
__syncthreads();
#pragma unroll
for (int k = 0; k < BK; k++) {
// 3. 每个线程负责计算BM*BN(12x128)中的TM*TN(8x8)个元素
#pragma unroll
for (int m = 0; m < TM; m++) {
#pragma unroll
for (int n = 0; n < TN; n++) {
// k from 0~7，0 ~ BK, ty and tx range from 0 to 15, 16x8=128
int comp_smem_a_m = ty * TM + m; // 128*8 128/TM(8)=16 M方向 16线程
int comp_smem_b_n = tx * TN + n; // 8*128 128/TN(8)=16 N方向 16线程
r_c[m][n] += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];
}
}
}
__syncthreads();
}
#pragma unroll
for (int m = 0; m < TM; ++m) {
int store_gmem_c_m = by * BM + ty * TM + m;
#pragma unroll
for (int n = 0; n < TN; n += 4) {
int store_gmem_c_n = bx * BN + tx * TN + n;
int store_gmem_c_addr = store_gmem_c_m * N + store_gmem_c_n;
FLOAT4(c[store_gmem_c_addr]) = FLOAT4(r_c[m][n]);
}
}
}这里gemm的实现比较简单，只使用了CUDA Cores，并且只实现Block Tile + K Tile以及Block Tile + K Tile+Thread Tile+向量化的版本。主要在于如何加载gmem中的数据到smem，也就是把全局内存中的数据索引mapping到共享内存中的。核心思维：把一个block中的线程id按照线性来理解，然后把这个线性的id和全局内存索引以及共享内存索引进行匹配。比如Block Tile + K Tile的实现，block内一共32x32个Threads，需要加载到smem的数据也是32x32，那么，最简单的做法，只需要每个线程加载一个互不重复数据即可。NOTE，本文的gemm kernel修改自：紫气东来：CUDA（三）：通用矩阵乘法：从入门到熟练(https://zhuanlan.zhihu.com/p/657632577)
03 warp/block reduce sum/max
// Warp Reduce Sum
template<const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum(float val) {
#pragma unroll
for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
val += __shfl_xor_sync(0xffffffff, val, mask);
}
return val;
}
// Warp Reduce Max
template<const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_max(float val) {
#pragma unroll
for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, mask));
}
return val;
}
// Block reduce sum/max/min device helper for Layer/RMS Norm/Softmax etc.
// grid 1D block 1D, grid(N/128), block(128)
template<const int NUM_THREADS=128>
__device__ __forceinline__ float block_reduce_sum(float val) {
// always <= 32 warps per block (limited by 1024 threads per block)
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
int warp = threadIdx.x / WARP_SIZE;
int lane = threadIdx.x % WARP_SIZE;
static __shared__ float shared[NUM_WARPS];
val = warp_reduce_sum<WARP_SIZE>(val);
if (lane == 0) shared[warp] = val;
__syncthreads();
val = (lane < NUM_WARPS) ? shared[lane] : 0.0f;
val = warp_reduce_sum<NUM_WARPS>(val);
return val;
}
template<const int NUM_THREADS=128>
__device__ __forceinline__ float block_reduce_max(float val) {
// always <= 32 warps per block (limited by 1024 threads per block)
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
int warp = threadIdx.x / WARP_SIZE;
int lane = threadIdx.x % WARP_SIZE;
static __shared__ float shared[NUM_WARPS];
val = warp_reduce_max<WARP_SIZE>(val);
if (lane == 0) shared[warp] = val;
__syncthreads();
val = (lane < NUM_WARPS) ? shared[lane] : -FLT_MAX;
val = warp_reduce_max<NUM_WARPS>(val);
return val;
}warp reduce几乎已经成为大部分reduce kernel的标准写法了，比如vLLM中，就是这种经典的写法。所以，先搞懂warp reduce（也就是搞懂各种warp functions的用法），再去写其他kernel，思路就会容易很多。需要注意的是，warp函数处理的是寄存器上的数据，也就是说，此时，没必要先加载数据到smem，再进行reduce，直接加载到寄存器即可（以前犯过这个小错误...）。Warp Functions建议参考：jhang：CUDA编程入门之Warp-Level Primitives(https://zhuanlan.zhihu.com/p/572820783)
04 block all reduce + vec4
// Block All Reduce Sum
// grid(N/128), block(128)
// a: Nx1, y=sum(a)
template<const int NUM_THREADS = 128>
__global__ void block_all_reduce_sum(float* a, float* y, int N) {
int tid = threadIdx.x;
int idx = blockIdx.x * NUM_THREADS + tid;
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
__shared__ float reduce_smem[NUM_WARPS];
// keep the data in register is enougth for warp operaion.
float sum = (idx < N) ? a[idx] : 0.0f;
int warp = tid / WARP_SIZE;
int lane = tid % WARP_SIZE;
// perform warp sync reduce.
sum = warp_reduce_sum<WARP_SIZE>(sum);
// warp leaders store the data to shared memory.
if (lane == 0) reduce_smem[warp] = sum;
__syncthreads(); // make sure the data is in shared memory.
// the first warp compute the final sum.
sum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;
if (warp == 0) sum = warp_reduce_sum<NUM_WARPS>(sum);
if (tid == 0) atomicAdd(y, sum);
}
// Block All Reduce Sum + float4
// grid(N/128), block(128/4)
// a: Nx1, y=sum(a)
template<const int NUM_THREADS = 128/4>
__global__ void block_all_reduce_sum_vec4(float* a, float* y, int N) {
int tid = threadIdx.x;
int idx = (blockIdx.x * NUM_THREADS + tid) * 4;
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
__shared__ float reduce_smem[NUM_WARPS];
float4 reg_a = FLOAT4(a[idx]);
// keep the data in register is enougth for warp operaion.
float sum = (idx < N) ? (reg_a.x + reg_a.y + reg_a.z + reg_a.w) : 0.0f;
int warp = tid / WARP_SIZE;
int lane = tid % WARP_SIZE;
// perform warp sync reduce.
sum = warp_reduce_sum<WARP_SIZE>(sum);
// warp leaders store the data to shared memory.
if (lane == 0) reduce_smem[warp] = sum;
__syncthreads(); // make sure the data is in shared memory.
// the first warp compute the final sum.
sum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;
if (warp == 0) sum = warp_reduce_sum<NUM_WARPS>(sum);
if (tid == 0) atomicAdd(y, sum);
}block all reduce是在warp reduce的基础上进行的，reduce_smem这部分的共享内存申请无法避免，这是用来同步每个warp之间得到局部结果。注意，最后，还需要atomicAdd做一个block级别的原子操作，以得到全局的和。float4向量化优化访存，可以减缓WarpScheduler发送指令的压力。
05 sgemv k32/k128/k16 kernel
// SGEMV: Warp SGEMV K32
// 假设K为32的倍数，每个warp负责一行
// grid(M/4), block(32,4) blockDim.x=32=K, blockDim.y=4
// a: MxK, x: Kx1, y: Mx1, compute: y = a * x
__global__ void sgemv_k32(float* a, float* x, float* y, int M, int K) {
int tx = threadIdx.x; // 0~31
int ty = threadIdx.y; // 0~4
int bx = blockIdx.x; // 0~M/4
int lane = tx % WARP_SIZE; // 0~31
int m = bx * blockDim.y + ty; // (0~M/4) * 4 + (0~3)
if (m < M) {
float sum = 0.0f;
int NUM_WARPS = (K + WARP_SIZE - 1) / WARP_SIZE;
#pragma unroll
for (int w = 0; w < NUM_WARPS; ++w) {
// 若NUM_WARPS>=2，先将当前行的数据累加到第一个warp中
int k = w * WARP_SIZE + lane;
sum += a[m * K + k] * x[k];
}
sum = warp_reduce_sum<WARP_SIZE>(sum);
if (lane == 0) y[m] = sum;
}
}
// SGEMV: Warp SGEMV K128 + Vec4
// 假设K为128的倍数 float4
// grid(M/4), block(32,4) blockDim.x=32=K, blockDim.y=4
// a: MxK, x: Kx1, y: Mx1, compute: y = a * x
__global__ void sgemv_k128(float* a, float* x, float* y, int M, int K) {
// 每个线程负责4个元素，一个warp覆盖128个元素
int tx = threadIdx.x; // 0~31
int ty = threadIdx.y; // 0~3
int bx = blockIdx.x; // 0~M/4
int lane = tx % WARP_SIZE; // 0~31
int m = blockDim.y * bx + ty; // (0~M/4) * 4 + (0~3)
if (m < M) {
float sum = 0.0f;
// process 4*WARP_SIZE elements per warp.
int NUM_WARPS = (((K + WARP_SIZE - 1) / WARP_SIZE) + 4 - 1) / 4;
#pragma unroll
for (int w = 0; w < NUM_WARPS; ++w) {
int k = (w * WARP_SIZE + lane) * 4;
float4 reg_x = FLOAT4(x[k]);
float4 reg_a = FLOAT4(a[m * K + k]);
sum += (reg_a.x * reg_x.x + reg_a.y * reg_x.y
+ reg_a.z * reg_x.z + reg_a.w * reg_x.w);
}
sum = warp_reduce_sum<WARP_SIZE>(sum);
if(lane == 0) y[m] = sum;
}
}
// SGEMV: Warp SGEMV K16
// 假设K为16 < 32,每个warp负责2行，每行有16个元素
// NUM_THREADS=128, NUM_WARPS=NUM_THREADS/WARP_SIZE;
// NUM_ROWS=NUM_WARPS * ROW_PER_WARP, grid(M/NUM_ROWS), block(32,NUM_WARPS)
// a: MxK, x: Kx1, y: Mx1, compute: y = a * x
template<const int ROW_PER_WARP = 2>
__global__ void sgemv_k16(float* A, float* x, float* y, int M, int K) {
constexpr int K_WARP_SIZE = (WARP_SIZE + ROW_PER_WARP - 1) / ROW_PER_WARP;
int tx = threadIdx.x; // 0~31
int ty = threadIdx.y; // 0~NUM_WARPS
int bx = blockIdx.x; // 0~M/NUM_ROWS (NUM_ROWS=NUM_WARPS * ROW_PER_WARP)
int lane = tx % WARP_SIZE; // 0~31
int k = lane % K_WARP_SIZE; // 0~15
// gloabl row of a: MxK and y:Mx1, blockDim.y=NUM_WARPS
int m = (blockDim.y * bx + ty) * ROW_PER_WARP + lane / K_WARP_SIZE;
if (m < M) {
float sum = A[m * K + k] * x[k];
sum = warp_reduce_sum<K_WARP_SIZE>(sum);
// 注意是k == 0，而不是lane == 0
if(k == 0) y[m] = sum;
}
}估计有些大佬倒立都能写sgemv的各种优化版了，核心思路其实也是基于warp reduce，考虑K的不同情况进行优化。本文的sgemv kernel修改自：有了琦琦的棍子：深入浅出GPU优化系列：gemv优化(https://zhuanlan.zhihu.com/p/494144694)
06 dot product, dot product + vec4
// Dot Product
// grid(N/128), block(128)
// a: Nx1, b: Nx1, y=sum(elementwise_mul(a,b))
template<const int NUM_THREADS = 128>
__global__ void dot(float* a, float* b, float* y, int N) {
int tid = threadIdx.x;
int idx = blockIdx.x * NUM_THREADS + tid;
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
__shared__ float reduce_smem[NUM_WARPS];
// keep the data in register is enougth for warp operaion.
float prod = (idx < N) ? a[idx] * b[idx] : 0.0f;
int warp = tid / WARP_SIZE;
int lane = tid % WARP_SIZE;
// perform warp sync reduce.
prod = warp_reduce_sum<WARP_SIZE>(prod);
// warp leaders store the data to shared memory.
if (lane == 0) reduce_smem[warp] = prod;
__syncthreads(); // make sure the data is in shared memory.
// the first warp compute the final sum.
prod = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;
if (warp == 0) prod = warp_reduce_sum<NUM_WARPS>(prod);
if (tid == 0) atomicAdd(y, prod);
}
// Dot Product + Vec4
// grid(N/128), block(128/4)
// a: Nx1, b: Nx1, y=sum(elementwise_mul(a,b))
template<const int NUM_THREADS = 128/4>
__global__ void dot_vec4(float* a, float* b, float* y, int N) {
int tid = threadIdx.x;
int idx = (blockIdx.x * NUM_THREADS + tid) * 4;
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
__shared__ float reduce_smem[NUM_WARPS];
float4 reg_a = FLOAT4(a[idx]);
float4 reg_b = FLOAT4(b[idx]);
float prod = (idx < N) ? (reg_a.x * reg_b.x + reg_a.y * reg_b.y
+ reg_a.z * reg_b.z + reg_a.w * reg_b.w) : 0.0f;
int warp = tid / WARP_SIZE;
int lane = tid % WARP_SIZE;
// perform warp sync reduce.
prod = warp_reduce_sum<WARP_SIZE>(prod);
// warp leaders store the data to shared memory.
if (lane == 0) reduce_smem[warp] = prod;
__syncthreads(); // make sure the data is in shared memory.
// the first warp compute the final sum.
prod = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;
if (warp == 0) prod = warp_reduce_sum<NUM_WARPS>(prod);
if (tid == 0) atomicAdd(y, prod);
}dot product kernel的核心就是block reduce，不多说了。
07 elementwise, elementwise + vec4
// ElementWise Add
// grid(N/128), block(128)
// a: Nx1, b: Nx1, c: Nx1, c = elementwise_add(a, b)
__global__ void elementwise_add(float* a, float* b, float* c, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < N) c[idx] = a[idx] + b[idx];
}
// ElementWise Add + Vec4
// grid(N/128), block(128/4)
// a: Nx1, b: Nx1, c: Nx1, c = elementwise_add(a, b)
__global__ void elementwise_add_vec4(float* a, float* b, float* c, int N) {
int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);
if (idx < N) {
float4 reg_a = FLOAT4(a[idx]);
float4 reg_b = FLOAT4(b[idx]);
float4 reg_c;
reg_c.x = reg_a.x + reg_b.x;
reg_c.y = reg_a.y + reg_b.y;
reg_c.z = reg_a.z + reg_b.z;
reg_c.w = reg_a.w + reg_b.w;
FLOAT4(c[idx]) = reg_c;
}
}elementwise可以考虑加点向量化进行访存优化。
08 histogram, histogram + vec4
// Histogram
// grid(N/128), block(128)
// a: Nx1, y: count histogram
__global__ void histogram(int* a, int* y, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < N) atomicAdd(&(y[a[idx]]), 1);
}
// Histogram + Vec4
// grid(N/128), block(128/4)
// a: Nx1, y: count histogram
__global__ void histogram_vec4(int* a, int* y, int N) {
int idx = 4 * (blockIdx.x * blockDim.x + threadIdx.x);
if (idx < N) {
int4 reg_a = INT4(a[idx]);
atomicAdd(&(y[reg_a.x]), 1);
atomicAdd(&(y[reg_a.y]), 1);
atomicAdd(&(y[reg_a.z]), 1);
atomicAdd(&(y[reg_a.w]), 1);
}
}统计频数直方图，很简单，两行代码搞定。
09 softmax, softmax + vec4 (grid level memory fence)
// Softmax x: N, y: N
// grid(N/128), block(K=128)
template<const int NUM_THREADS = 128>
__global__ void softmax(float* x, float* y, float* total, int N) {
const int tid = threadIdx.x;
const int idx = blockIdx.x * blockDim.x + tid;
constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
__shared__ float reduce_smem[NUM_WARPS];
float sum = (idx < N) ? expf(x[idx]) : 0.0f;
int warp = tid / WARP_SIZE;
int lane = tid % WARP_SIZE;
sum = warp_reduce_sum<WARP_SIZE>(sum);
if (lane == 0) reduce_smem[warp] = sum;
__syncthreads();
// compute the final sum in each warp
sum = (lane < NUM_WARPS) ? reduce_smem[lane] : 0.0f;
sum = warp_reduce_sum<NUM_WARPS>(sum); // sum(e^x_0,...,e^x_n-1)
// get the total sum of all blocks.
if (tid == 0) atomicAdd(total, sum);
__threadfence(); // grid level memory fence 注意这里需要网格级别的内存同步
// e^x_i/sum(e^x_0,...,e^x_n-1)
if (idx < N) y[idx] = block_smem[tid] / (*total);
}
// Softmax x: N, y: N
// grid(N/128), block(K=128)
template<const int NUM_THREADS = 128>
__global__ void softmax_v2(float* x, float* y, float* total, int N) {
const int tid = threadIdx.x;
const int idx = blockIdx.x * blockDim.x + tid;
float exp_val = (idx < N) ? expf(x[idx]) : 0.0f;
float sum = block_reduce_sum<NUM_THREADS>(exp_val);
// get the total sum of all blocks.
if (tid == 0) atomicAdd(total, sum);
__threadfence(); // grid level memory fence 注意这里需要网格级别的内存同步
// e^x_i/sum(e^x_0,...,e^x_n-1)
if (idx < N) y[idx] = exp_val / (*total);
}
// Softmax Vec4 x: N, y: N
// grid(N/128), block(128/4)
template<const int NUM_THREADS = 128/4>
__global__ void softmax_v2_vec4(float* x, float* y, float* total, int N) {
const int tid = threadIdx.x;
const int idx = (blockIdx.x * blockDim.x + tid) * 4;
float4 reg_x = FLOAT4(x[idx]);
float4 reg_exp;
reg_exp.x = (idx < N) ? expf(reg_x.x) : 0.0f;
reg_exp.y = (idx < N) ? expf(reg_x.y) : 0.0f;
reg_exp.z = (idx < N) ? expf(reg_x.z) : 0.0f;
reg_exp.w = (idx < N) ? expf(reg_x.w) : 0.0f;
float exp_val = (reg_exp.x + reg_exp.y + reg_exp.z + reg_exp.w);
float sum = block_reduce_sum<NUM_THREADS>(exp_val);
// get the total sum of all blocks.
if (tid == 0) atomicAdd(total, sum);
__threadfence(); // grid level memory fence 注意这里需要网格级别的内存同步
// e^x_i/sum(e^x_0,...,e^x_n-1)
if (idx < N) {
float4 reg_y;
reg_y.x = reg_exp.x / (*total);
reg_y.y = reg_exp.y / (*total);
reg_y.z = reg_exp.z / (*total);
reg_y.w = reg_exp.w / (*total);
FLOAT4(y[idx]) = reg_y;
}
}softmax稍微要注意的就是内存同步的问题，这里，你需要做一个网格级别的同步，而不能仅仅是block级别，否则拿不到全局的exp sum作为分母项。因此使用 __threadfence 这个网格及内存同步操作。不过效率我还没测过，实在要高效的话，可能得整成FA2那样的 1-pass + online softmax的实现。不过，如果是面试的话，就不要太为难自己了... ，但是FA1/FA2的论文很经典，强烈建议多读几遍。
10 sigmoid, sigmoid + vec4
// Sigmoid x: N, y: N y=1/(1+exp(-x))
// grid(N/128), block(K=128)
__global__ void sigmoid(float* x, float* y, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < N) y[idx] = 1.0f / (1.0f + expf(-x[idx]));
}
// Sigmoid x: N, y: N y=1/(1+exp(-x)) Vec4
// grid(N/128), block(128/4)
__global__ void sigmoid_vec4(float* x, float* y, int N) {
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
if (idx < N) {
float4 reg_x = FLOAT4(x[idx]);
float4 reg_y;
reg_y.x = 1.0f / (1.0f + expf(-reg_x.x));
reg_y.y = 1.0f / (1.0f + expf(-reg_x.y));
reg_y.z = 1.0f / (1.0f + expf(-reg_x.z));
reg_y.w = 1.0f / (1.0f + expf(-reg_x.w));
FLOAT4(y[idx]) = reg_y;
}
}// Relu x: N, y: N y=max(0,x)
// grid(N/128), block(K=128)
__global__ void relu(float* x, float* y, int N) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < N) y[idx] = fmaxf(0.0f, x[idx]);
}
// Relu x: N, y: N y=max(0,x) Vec4
// grid(N/128/4), block(128/4)
__global__ void relu_vec4(float* x, float* y, int N) {
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
if (idx < N) {
float4 reg_x = FLOAT4(x[idx]);
float4 reg_y;
reg_y.x = fmaxf(0.0f, reg_x.x);
reg_y.y = fmaxf(0.0f, reg_x.y);
reg_y.z = fmaxf(0.0f, reg_x.z);
reg_y.w = fmaxf(0.0f, reg_x.w);
FLOAT4(y[idx]) = reg_y;
}
}// Layer Norm: x: NxK(K=128<1024), y': NxK, y'=x-mean(x)/std(x) each row
// mean(x) = sum(x)/K, 1/std(x) = rsqrtf( sum( (x-mean(x))^2 )/K ) each row
// grid(N*K/K), block(K<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g + b (g: scale, b: bias)
template<const int NUM_THREADS=128>
__global__ void layer_norm(float* x, float* y, float g, float b, int N, int K) {
int tid = threadIdx.x; // 0..K-1
int bid = blockIdx.x; // 0..N-1
int idx = bid * blockDim.x + threadIdx.x;
const float epsilon = 1e-5f;
__shared__ float s_mean; // shared within block
__shared__ float s_variance; // shared within block
float value = (idx < N * K) ? x[idx] : 0.0f; // load once only
float sum = block_reduce_sum<NUM_THREADS>(value);
if (tid == 0) s_mean = sum / (float) K;
// wait for s_mean in shared memory to be ready for all threads
__syncthreads();
float variance = (value - s_mean) * (value - s_mean);
variance = block_reduce_sum<NUM_THREADS>(variance);
if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);
// wait for s_variance in shared memory to be ready for all threads
__syncthreads();
if (idx < N * K) y[idx] = ((value - s_mean) * s_variance) * g + b;
}
// Layer Norm Vec4: x: NxK(K=128<1024), y': NxK, y'=x-mean(x)/std(x) each row
// mean(x) = sum(x)/K, 1/std(x) = rsqrtf( sum( (x-mean(x))^2 )/K ) each row
// grid(N*K/K), block(K/4<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g + b (g: scale, b: bias)
template<const int NUM_THREADS=128/4>
__global__ void layer_norm_vec4(float* x, float* y, float g, float b, int N, int K) {
int tid = threadIdx.x; // 0..K-1
int bid = blockIdx.x; // 0..N-1
int idx = (bid * blockDim.x + threadIdx.x) * 4;
const float epsilon = 1e-5f;
__shared__ float s_mean; // shared within block
__shared__ float s_variance; // shared within block
float4 reg_x = FLOAT4(x[idx])
float value = (idx < N * K) ? (reg_x.x + reg_x.y
+ reg_x.z + reg_x.w) : 0.0f;
float sum = block_reduce_sum<NUM_THREADS>(value);
if (tid == 0) s_mean = sum / (float) K;
// wait for s_mean in shared memory to be ready for all threads
__syncthreads();
float4 reg_x_hat;
reg_x_hat.x = reg_x.x - s_mean;
reg_x_hat.y = reg_x.y - s_mean;
reg_x_hat.z = reg_x.z - s_mean;
reg_x_hat.w = reg_x.w - s_mean;
float variance = reg_x_hat.x * reg_x_hat.x + reg_x_hat.y * reg_x_hat.y
+ reg_x_hat.z * reg_x_hat.z + reg_x_hat.w * reg_x_hat.w;
variance = block_reduce_sum<NUM_THREADS>(variance);
if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);
// wait for s_variance in shared memory to be ready for all threads
__syncthreads();
float4 reg_y;
reg_y.x = reg_x_hat.x * s_variance * g + b;
reg_y.y = reg_x_hat.y * s_variance * g + b;
reg_y.z = reg_x_hat.z * s_variance * g + b;
reg_y.w = reg_x_hat.w * s_variance * g + b;
if (idx < N * K) FLOAT4(y[idx]) = reg_y;
}layer norm实现的核心同样也是block reduce和warp reduce，然后再整点向量化...
// RMS Norm: x: NxK(K=128<1024), y': NxK, y'=x/rms(x) each row
// 1/rms(x) = rsqrtf( sum(x^2)/K ) each row
// grid(N*K/K), block(K<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g (g: scale)
template<const int NUM_THREADS=128>
__global__ void rms_norm(float* x, float* y, float g, int N, int K) {
int tid = threadIdx.x; // 0..K-1
int bid = blockIdx.x; // 0..N-1
int idx = bid * blockDim.x + threadIdx.x;
const float epsilon = 1e-5f;
__shared__ float s_variance; // shared within block
float value = (idx < N * K) ? x[idx] : 0.0f; // load once only
float variance = value * value;
variance = block_reduce_sum<NUM_THREADS>(variance);
if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);
// wait for s_variance in shared memory to be ready for all threads
__syncthreads();
if (idx < N * K) y[idx] = (value * s_variance) * g;
}
// RMS Norm Vec4: x: NxK(K=128<1024), y': NxK, y'=x/rms(x) each row
// 1/rms(x) = rsqrtf( sum(x^2)/K ) each row
// grid(N*K/K), block(K/4<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g (g: scale)
template<const int NUM_THREADS=128/4>
__global__ void rms_norm_vec4(float* x, float* y, float g, int N, int K) {
int tid = threadIdx.x; // 0..K-1
int bid = blockIdx.x; // 0..N-1
int idx = (bid * blockDim.x + threadIdx.x) * 4;
const float epsilon = 1e-5f;
__shared__ float s_variance; // shared within block
float4 reg_x = FLOAT4(x[idx]);
float variance = (idx < N * K) ? (reg_x.x * reg_x.x + reg_x.y * reg_x.y
+ reg_x.z * reg_x.z + reg_x.w * reg_x.w) : 0.0f;
variance = block_reduce_sum<NUM_THREADS>(variance);
if (tid == 0) s_variance = rsqrtf(variance / (float) K + epsilon);
// wait for s_variance in shared memory to be ready for all threads
__syncthreads();
float4 reg_y;
reg_y.x = reg_x.x * s_variance * g;
reg_y.y = reg_x.y * s_variance * g;
reg_y.z = reg_x.z * s_variance * g;
reg_y.w = reg_x.w * s_variance * g;
if (idx < N * K) FLOAT4(y[idx]) = reg_y;
}rms norm实现的核心同样也是block reduce和warp reduce...，然后再加点float4向量化什么的。
struct Box {
float x1, y1, x2, y2, score;
float area() const {return (std::abs(x2 - x1 + 1)) * std::abs(y2 - y1 + 1); }
float iou_of(const Box& other) const{
float inner_x1 = x1 > other.x1 ? x1 : other.x1;
float inner_y1 = y1 > other.y1 ? y1 : other.y1;
float inner_x2 = x2 < other.x2 ? x2 : other.x2;
float inner_y2 = y2 < other.y2 ? y2 : other.y2;
float inner_h = inner_y2 - inner_y1 + 1.0f;
float inner_w = inner_x2 - inner_x1 + 1.0f;
float inner_area = inner_h * inner_w;
return (inner_area / (area() + tbox.area() - inner_area));
}
}
void hard_nms(std::vector<Box> &input, std::vector<Box> &output, float iou_threshold){
if (input.empty()) return;
std::sort(input.begin(), input.end(),[](Box& a, Box& b) { return a.score > b.score; });
int box_num = input.size();
std::vector<int> merged(box_num, 0);
for (int i = 0; i < box_num; ++i) {
if (merged[i]) continue;
merged[i] = 1;
for (int j = i + 1; j < box_num; ++j) {
if (merged[j]) continue;
float iou = input[i].iou_of(input[j]);
if (iou > iou_threshold) merged[j] = 1;
}
output.push_back(input[i]);
}
}CV相关的经常会要手撕NMS，也记录下。
15 总结
可以发现，大部分kernel的基本写法都是依赖warp reduce和block reduce的，基本上只要熟练应用warp functions各种场景的写法，应该问题不大；softmax需要考虑网格级同步的问题，或者online softmax以及FlashAttention；sgemm的优化是个很大的课题，不是案例中写的这么简单，但是入门的话，基本就是tiling的思想以及如何做索引之间的mapping；sgemv的优化则主要考虑K不同的值（因为M为1了），比如K=16,64,128等情况下，如何按照warp来处理；relu、sigmoid等都是elementwise的操作，很好实现，可以再考虑加点向量化优化访存；layer norm和rms norm在数学上其实也是挺清晰简单的，落实到cuda kernel时，只要按照逐个token来处理，headdim没有超过1024的情况下（一个block最多可以放1024个threads），可以放到一个block处理，这样并行化就很好写。当然，核心还是warp reduce和block reduce；NMS是乱入的，没有CUDA版本，别问了...
最后，附GitHub repo:
https://github.com/DefTruth/cuda-learn-note/
