在GPU上手写卷积算子是我今年做过最“自虐”但也最值得的一件事。不是因为卷积难写而是因为它足够简单却能一口气把CUDA里的线程组织、共享内存、寄存器、访存模型全部串起来。把它和cuDNN放在一起做性能对比分析之后你才会真正理解GPU算力是怎么被一层层榨出来的。这篇文章就是我基于CUDA从零手写卷积算子的完整记录从朴素kernel到共享内存优化再到向量化和寄存器复用最后给出与PyTorch内置cuDNN的实测对比以及一路踩过的环境和正确性大坑。适合想学CUDA算子开发、或者已经在写kernel但性能一直上不去的同学。1. 为什么选择手写卷积算子1.1 框架内置卷积不够用吗如果是做常规模型训练和推理PyTorch里面一行nn.Conv2d就够用了底层是cuDNN自动调优过的卷积实现又快又稳。那为什么还要手写一个性能大概率打不过cuDNN的kernel我自己遇到的情况是在调试一个自定义动态卷积时内核逻辑里带了输入相关的权重预测标准卷积API完全套不上。如果不理解底层卷积的访存规律连该从哪里下手优化都不知道。类似的需求还有很多比如把卷积和激活函数融合成一个算子、在推理引擎里裁剪掉用不到的特性、给特定shape做深度定制。这些场景下手写CUDA卷积不是炫技而是一个绕不开的基本功。手写这个项目还有一层价值它能帮你把黑盒拆开。cuDNN在内部会根据shape、精度、显卡架构选择不同算法可能是隐式GEMM、可能是Winograd也可能走FFT。你只看到结果看不到选择背后的权衡。自己从朴素实现一步步优化上去才能理解为什么有些op在某些shape下快得离谱换个shape又慢成狗。1.2 卷积的两种实现路线直接卷积与GEMM化卷积的数值定义不复杂输出特征图的每个点是输入特征图对应窗口和卷积核的加权和。但在GPU上怎么组织计算行业内基本分成两派。第一派是直接卷积direct convolution。老老实实按滑动窗口取数计算写回。优点是内存占用小、逻辑直观缺点是输入数据被高倍重复读取。举个简单的3x3卷积例子如果不做任何缓存每个输出点都要读9个输入点相邻输出点的读取窗口高度重叠这会带来大量冗余访存GPU的HBM带宽很快就被耗尽。第二派是im2col加GEMM。先把输入展开成一个大矩阵每行对应一个输出点要用的窗口数据然后把这个大矩阵和权重矩阵做GEMM。这样做的好处是能复用高度优化的cuBLAS矩阵乘法缺点也很明显im2col展开后显存占用通常是输入的9倍甚至更多在特征图较大时很容易爆显存。我这次选择的是直接卷积路线并且逐步用共享内存、寄存器缓存来消除访存冗余。这样每一步优化都能对应到具体的硬件行为性能对比分析也更有说服力。如果你更关注工业级性能后面可以再转GEMM方向。1.3 实验目标与评测口径为了让对比结果不“飘”我提前把实验场景固定下来输入shape是[N16, Cin256, Hin128, Win128]输出通道Cout256卷积核是3x3stride1pad1输出尺寸保持128x128。这个shape虽然不算大但已经能体现访存和并行度的权衡也接近很多检测模型中间层的实际尺寸。手写版本我计划做三个递进优化朴素实现一个线程算一个输出元素肉眼可见的冗余访存。共享内存Tile版本把输入窗口搬到片上减少全局内存重复访问。向量化加寄存器缓存版本用float4提升访存效率让每个线程承担更多输出计算。评测指标有三个端到端平均耗时、有效带宽、以及相对cuDNN的速度比。正确性用PyTorch的CPU/GPU卷积结果做基准确保每一步优化都是“先对再快”。2. 环境准备与工具链CUDA版本兼容是第一步2.1 驱动、Toolkit、PyTorch的三方关系很多人一开始就卡在环境上不是因为代码有问题而是没搞懂GPU驱动、CUDA Toolkit、PyTorch内置runtime这三者的关系。nvidia-smi右上角显示的CUDA Version是当前驱动最高支持的CUDA版本它不代表你装了多新的Toolkit。nvcc --version显示的才是编译器对应的Toolkit版本。而PyTorch这种框架本身会在安装包里带一份CUDA runtime库它内部用的CUDA版本和系统Toolkit可能完全不一致这都没关系只要驱动版本能兼容即可。我这次用的是RTX 4060 Ti驱动版本已经支持CUDA 13.x但PyTorch装的是2.8.0加CUDA 12.1组合包跑起来后容易遇到版本相关报错。最典型的是the detected CUDA version (13.0) mismatches the version that was used to compile PyTorch这种报错归根结底是运行时和编译时的CUDA版本对不上。解决方案也很直接要么把PyTorch换成和驱动更匹配的cu130版本要么把CUDA Toolkit降级或升级到与PyTorch一致的版本。我更推荐用Anaconda创建独立环境通过conda install pytorch torchvision pytorch-cuda12.1 -c pytorch -c nvidia这样明确指定CUDA版本比手动下载Toolkit省心得多。2.2 WSL2是个不错的选择如果你用的是Windows又想在Linux环境下写CUDA程序WSL2是个很顺手的方案。我目前就在Windows 11的WSL2里装Ubuntu 24.04然后安装CUDA Toolkit日常编译和调试都在WSL里完成。有几个坑需要提前注意。第一WSL2里不要重复装GPU驱动驱动是Windows侧提供的Linux侧只需要装CUDA Toolkit。第二装完Toolkit后nvcc经常不在PATH里需要手动加一下export PATH/usr/local/cuda/bin:$PATH export LD_LIBRARY_PATH/usr/local/cuda/lib64:$LD_LIBRARY_PATH第三cuda samples找不到很常见一般是因为Toolkit安装时没有勾选samples或者环境变量没配好。可以直接去NVIDIA官网下载对应版本的samples或者在安装目录/usr/local/cuda/samples里检查。如果这里能正常编译运行说明Toolkit基本没问题。2.3 工程结构与评测工具项目结构不用搞太复杂我平时会这样组织conv_cuda/ src/ # 存放CUDA kernel源文件 include/ # 公共头文件 tests/ # 正确性验证脚本 bench/ # 性能测试入口 CMakeLists.txt编译时用CMake加nvcc关键要指定GPU架构。4060 Ti是sm_893060是sm_86编译参数直接写死比用自动检测更省心set(CMAKE_CUDA_ARCHITECTURES 89)如果机器上同时装了多个CUDA版本还可以通过修改CMAKE_CUDA_COMPILER来切换。切版本时最容易踩的坑是换了Toolkit却忘了清空CMake缓存导致一直用旧编译器折腾半天才发现是缓存惹的祸。性能评测我统一用CUDA Event计时不用std::chrono因为后者把kernel launch的异步开销也算进去了。命令大致是cudaEventRecord(start); myKernelgrid, block(...); cudaEventRecord(stop); cudaEventSynchronize(stop); cudaEventElapsedTime(ms, start, stop);Profile工具用NVIDIA Nsight Compute也就是ncu。它能精确告诉你一个kernel是卡在访存、计算还是同步上后面调优全指望它。3. 手写卷积算子从朴素到高性能3.1 朴素实现先把功能跑对第一个版本不考虑任何优化直接一个线程负责计算一个输出元素。这样写起来最省事正确性验证也容易。核心代码如下__global__ void conv2d_naive( const float* __restrict__ input, const float* __restrict__ weight, float* __restrict__ output, int N, int Cin, int Hin, int Win, int Cout, int Hout, int Wout, int K, int pad, int stride) { int idx blockIdx.x * blockDim.x threadIdx.x; int total N * Cout * Hout * Wout; if (idx total) return; int ow idx % Wout; int oh (idx / Wout) % Hout; int co (idx / (Wout * Hout)) % Cout; int n idx / (Wout * Hout * Cout); float acc 0.0f; for (int ci 0; ci Cin; ci) { for (int kh 0; kh K; kh) { for (int kw 0; kw K; kw) { int h oh * stride - pad kh; int w ow * stride - pad kw; if (h 0 h Hin w 0 w Win) { acc input[((n * Cin ci) * Hin h) * Win w] * weight[((co * Cin ci) * K kh) * K kw]; } } } } output[idx] acc; }代码逻辑很清楚但性能非常糟糕。每个输出元素都要遍历Cin个输入通道每个通道再读9个输入点。相邻线程读取的窗口高度重叠全局内存访问大量冗余HBM带宽被白白浪费。实测下来这个版本比cuDNN慢了20多倍是后面所有优化的起点。这里有个细节我在写__restrict__时特意给input和weight加上了因为这两个指针在kernel内部不会被别名修改编译器可以做更激进的重排序。不加这个限制nvcc可能为了安全而插入额外内存屏障性能会有一点损失。3.2 共享内存Tiling把数据搬到片上既然瓶颈是全局内存冗余访问那就把输入tile提前搬到共享内存里。共享内存在SM内部读写延迟比全局内存低一个数量级带宽也高得多。每个block负责输出特征图上一个TILE x TILE的小块为了计算3x3卷积需要把输入tile向外扩一圈也就是hot zone。这里有个经典的尺寸选择问题。我采用TILE16那输入tile就是(162)x(162)18x18。每个线程组从全局内存加载18x18的数据然后同步再让每个线程计算自己的输出元素。核心思路是#define TILE 16 __global__ void conv2d_shared( const float* __restrict__ input, const float* __restrict__ weight, float* __restrict__ output, int Cin, int Hin, int Win, int Cout, int Hout, int Wout) { __shared__ float smem[TILE 2][TILE 2]; int tx threadIdx.x; int ty threadIdx.y; int x0 blockIdx.x * TILE; int y0 blockIdx.y * TILE; int co blockIdx.z; // 协作加载当前输入通道的tile for (int ci 0; ci Cin; ci) { int in_y y0 - 1 ty; int in_x x0 - 1 tx; if (ty TILE 2 tx TILE 2 in_y 0 in_y Hin in_x 0 in_x Win) { smem[ty][tx] input[((ci * Hin in_y) * Win) in_x]; } else { smem[ty][tx] 0.0f; } __syncthreads(); if (ty TILE tx TILE) { int out_y y0 ty; int out_x x0 tx; if (out_y Hout out_x Wout) { float acc 0.0f; #pragma unroll for (int kh 0; kh 3; kh) { #pragma unroll for (int kw 0; kw 3; kw) { acc smem[ty kh][tx kw] * weight[((co * Cin ci) * 3 kh) * 3 kw]; } } output[((co * Hout out_y) * Wout) out_x] acc; } } __syncthreads(); } }这个版本有个明显的tradeoff每个block只负责一个输出通道co而Cin通常很大多个输出通道的block会重复加载同一份输入tile。实际运行时因为有L2缓存兜底不会真的每次都要从HBM读但共享内存的加载开销还是会拖慢速度。在工程里更优的做法是让一个block同时处理多个输出通道把输入tile加载一次后复用但这会让代码复杂不少。我这里为了教学清晰先保留了简单版本。共享内存版本的一个隐性坑是bank conflict。smem[ty][tx]按ty方向访问时数组宽度是TILE218。当多个线程同时访问同一行不同列时如果列索引间隔不是bank数量16的整数倍就会发生冲突。把存储宽度改成TILE319也就是填充一行能让冲突明显下降。这个细节我曾经忽略过性能差了接近15%。3.3 向量化、循环展开与寄存器缓存共享内存已经解决了一部分访存问题但还远没到极限。接下来可以做的优化有三个方向。第一个方向是向量化访存。CUDA里用float4一次读4个float可以减少指令数也能让全局内存访问更高效。实现时可以把输入在宽度方向上按4对齐让一个线程一次搬运4个像素。注意边界条件如果特征图宽度不是4的倍数需要单独处理尾部否则极易越界。第二个方向是循环展开。对于3x3卷积内层kh和kw的循环次数是固定的。直接#pragma unroll展开编译器可以消除循环计数器开销还能把多个独立的乘加运算调度到不同流水线上隐藏计算延迟。前面共享内存版本的代码里已经用了这个技巧。第三个方向是寄存器缓存或者说thread coarsening。默认让一个线程算一个输出元素线程数太多索引计算和访存指令占比高。可以让每个线程负责2x2甚至4x4个输出元素这样每个线程内可以复用相邻输入数据。实现上很像CPU上的strip mining把并行粒度放大减少block调度和线程创建开销。但这个优化不是白给的。寄存器使用量增加后块内线程数可能下降最终导致占用率降低反而变慢。所以每次调整后都要用profiler重新看一遍而不是盲目堆技巧。我在实验里把这三个方向组合起来最终把耗时从共享内存版本的2.61ms压到了1.92ms。3.4 正确性验证先对再快优化过程中最怕的就是“跑得快但算错了”。我的习惯是先用一个很小的shape做全手工验证比如N1, Cin2, Hin4, Win4卷积核固定成简单整数自己在纸上算一遍输出再和kernel结果对比。这个阶段不要开任何优化方便定位问题。大规模验证直接用PyTorch的卷积结果当基准import torch import torch.nn.functional as F x torch.randn(16, 256, 128, 128, devicecuda) w torch.randn(256, 256, 3, 3, devicecuda) ref F.conv2d(x, w, padding1) out torch.from_numpy(custom_conv_output) # 从自定义kernel读到tensor torch.testing.assert_close(out, ref, rtol1e-4, atol1e-4)这里要说一下浮点累加顺序不同会导致结果有微小差异所以rtol和atol不要卡死到1e-7否则某些通道会出现假阳性错误。正确性还有一道杀手锏是compute-sanitizer它能检测越界读写。我调试共享内存版本时就是靠它抓到了一个边界条件写错导致的数组越界。用法很简单compute-sanitizer ./bench如果kernel里有非法访问它会明确告诉你发生错误的线程ID和代码行。4. 性能实测与调优记录4.1 实测数据从慢20倍到慢3倍固定测试环境RTX 4060 TiCUDA 12.1PyTorch 2.8.0输入[16, 256, 128, 128]卷积参数3x3, stride1, pad1。每个kernel先跑10次warmup再取50次的平均耗时。版本平均耗时(ms)相对cuDNN朴素实现18.7427.6x slower共享内存Tile2.613.8x slower向量化寄存器缓存1.922.8x slowercuDNN benchmark0.681x看到这个结果很多人第一反应是“还是打不过cuDNN啊”。确实打不过但这里面的信息量比“输赢”大得多。朴素实现慢在于访存完全冗余共享内存版本一下追回7倍多说明优化方向是对的。向量化和寄存器缓存又拉近了1.3倍但后面每前进一小步代价都在翻倍。cuDNN只有0.68ms一方面是因为它用隐式GEMM加Winograd把乘加次数从9次降到了更低另一方面是它会根据shape做autotuning甚至为不同架构选择不同kernel。手写版本做到这个程度已经足够说明问题。4.2 用Profiler定位真正的瓶颈光看总时间不够我每次调优都会用ncu看关键指标。常用的命令是ncu --set full ./bench或者只看几个特定指标ncu --metrics gpu__time_duration.sum,sm__throughput.avg.pct_of_peak_sustained_elapsed,gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed ./bench共享内存版本刚写完时耗时是3.1ms左右比预期高。ncu一看shared__bpe_red也就是bank conflict相关指标很高。问题出在数组宽度刚好是18导致部分线程访问共享内存时产生了2路冲突。把存储宽度改成19后指标几乎归零耗时降到了2.61ms。另一个发现是sm__throughput只有30%左右而gpu__compute_memory_throughput已经接近65%。这说明kernel依然是访存受限不是计算受限。所以下一步优化应该继续减少访存量而不是盲目加计算单元。这个判断也解释了为什么vector版本有效因为它是在降低访存指令和事务数量。4.3 手写kernel离cuDNN还有多远从数据看手写版本离cuDNN还有2.8倍差距。这个差距主要来自算法层面的代差不是简单的调参能弥补的。cuDNN的3x3卷积在很多情况下会走Winograd把每个输出tile的乘法次数从9降到4左右计算量直接少了一半多。如果再配合Tensor CoreFP32也可以通过TF32模式获得额外加速这是手写普通CUDA kernel无法匹敌的。但手写kernel的价值不在于非要和cuDNN打个平手。它更大的意义在于当你需要定制一个cuDNN里没有的算子时你知道怎么把性能优化到“够用”。实际上我在做完这个项目之后已经能把自定义动态卷积的耗时压到原生PyTorch算子两倍以内这对业务落地完全可接受。5. 常见问题与排查技巧实录5.1 环境兼容问题速查表手写CUDA容易踩的环境坑很多特别是跟PyTorch、驱动版本混在一起时。我把这段时间遇到的高频问题整理成了表格。错误或现象常见原因解决办法torch.acceleratorerror: cuda error: no kernel image is available for execution编译时指定的GPU架构和实际显卡不匹配常见于3050/4060等新卡用了-archsm_70设置TORCH_CUDA_ARCH_LIST8.9或-archsm_89the detected cuda version (13.0) mismatches the version that was used to compilePyTorch运行时与Toolkit版本不一致用conda指定pytorch-cuda版本或升降级Toolkitruntimeerror: cuda error: cublas_status_execution_failed when calling cublas...自写kernel越界写坏了显存后续cuBLAS调用崩溃用compute-sanititor定位越界检查所有索引计算cuda samples找不到安装时未装Samples或PATH没配好检查/usr/local/cuda/samples或设置CUDA_PATHWSL2里nvcc命令找不到Toolkit没加进PATH在~/.bashrc里添加CUDA路径ComfyUI启动失败提示c10.dll或CUDA 13.1相关错误ComfyUI依赖的PyTorch版本和当前CUDA runtime不兼容创建独立conda环境安装匹配的PyTorch/CUDA组合环境问题最有效的手段就是做环境隔离。不要在一个系统Python里装多套CUDA/PyTorch用conda环境各管各的出问题直接删掉重建比手动排查节省太多时间。5.2 正确性验证的几条硬经验第一个经验是永远先跑小shape。我见过有人在512大小的特征图上直接调kernel越界访问导致显存污染报错却指向之后的cuBLAS调用排查了半天才发现是自己kernel的问题。小shape下逻辑简单手工计算也能验证。第二个经验是compute-sanitizer要成为标配。无论kernel看起来多简单越界风险都存在。尤其是共享内存和向量化代码索引关系复杂边界条件容易漏。每次跑性能前先跑一遍compute-sanitizer能过滤掉大半诡异问题。第三个经验是浮点误差要宽容。GPU上不同算法累加顺序不一样误差在1e-4级别非常正常。比较结果时用torch.testing.assert_close(..., rtol1e-4, atol1e-4)不要用硬怼。5.3 性能不达预期的排查清单如果kernel已经正确但性能就是上不去我一般会按下面这个顺序排查。先看编译参数。nvcc -O3是必须的同时建议加--ptxas-options-v查看寄存器数和共享内存使用量。如果寄存器数超过128占用率可能被压得很低。再看__launch_bounds__是否需要设置强制限制每个线程的寄存器上限。接着看启动配置。block大小不是越满越好我试过256线程比512线程在共享内存版本上更快因为每个block需要加载的tile halo占比更小。核心原则是让每个线程访存和计算比例匹配硬件特性。然后看profiler指标。重点看sm__throughput、gpu__compute_memory_throughput和l1tex__throughput。如果访存指标远高于计算指标说明是访存瓶颈继续减少访存或提升向量化如果计算指标高说明计算瓶颈考虑用Tensor Core或减少无效计算。最后是一个特别常见但容易忽略的坑benchmark循环里频繁调用cudaMalloc和cudaFree。GPU上的内存分配开销很大而且会同步阻塞导致测出来的时间虚高。正确做法是提前分配一次在循环里复用buffer。写在最后的一点体会手写卷积算子这个项目做完之后我最大的感触是优化不是把各种技巧堆上去就完事而是要在准确性和性能之间不断做取舍。共享内存版本性能反超朴素版本是因为找准了访存瓶颈向量化版本提升有限是因为计算量和访存量并没有发生质的改变。如果只盯着“为什么还没cuDNN快”就会忽略真正需要理解的东西。如果让我给刚开始接触CUDA的人一个建议那就是从3x3卷积开始练手完整走一遍“实现、验证、profile、优化、对比”的闭环。一次只改一个变量让数据说话。这个过程练出来的手感比看一百篇理论文章都管用。最后再分享一个小技巧不要一上来就用太复杂的模型输入做调优固定一个中等大小的shape把问题缩小。等你把kernel在这个shape下打磨顺了再扩展到其他尺寸很多性能问题会在扩展时自动暴露出来那时候你已经有足够经验去快速定位了。