很多做AI性能优化的人一开始都会陷入一个误区觉得GPU是“算力怪兽”优化重点就该放在FLOPs上。等真正用ncu一测却发现很多kernel的SM利用率低得可怜Tensor Core在那闲着瓶颈根本不在计算而在数据搬运。做GPU性能工程最重要的一课就是GPU是拿延迟换吞吐的设备但它的计算单元再快也得等数据从显存里搬进来。访存模式Memory Access Pattern就是这套数据流系统的“任督二脉”。入口堵了后面算得再快也是空转。这篇文章是《AI系统性能工程学习笔记》系列的第七篇我结合自己在模型推理、训练kernel优化中的实际踩坑经验把GPU访存的底层机制、常见陷阱和优化手段串一遍重点说清楚“为什么这么改就能快”以及“如何用工具证明它确实变快了”。适合正在做PyTorch算子优化、自定义CUDA kernel或者纯粹被显存带宽困扰的同行参考。1. 先搞懂GPU到底“慢”在哪一份各层级带宽/延迟的对照认知访存优化这件事说白了就是在不同的存储层级之间做数据搬运时间的博弈。想要看懂后续的优化手法第一步要把GPU的内存层级结构刻在脑子里。NVIDIA这边的典型架构从内到外大概是寄存器Register → 共享内存Shared Memory → L1 Cache → L2 Cache → 显存HBM/HBM2e/HBM3。每一层之间的带宽和延迟差异是数量级的差距知道了这个你就明白为什么“数据放对地方”比“多跑几个线程”更重要。1.1 带宽差距的真实数字我整理了一份典型数据以A100 80GB为例不同卡有细微差异但量级一致存储层级典型容量典型带宽典型延迟寄存器每个SM 256KB极高TB/s级别几乎0共享内存/L1每个SM 192KB可配置约几十TB/s级别20-30 cycleL2 Cache40MB约5-6TB/s200-300 cycleHBM显存80GB约2TB/s400-800 cycle这里最重要的数字是HBM显存带宽通常只有2TB/s左右而SM内部共享内存的带宽比它高一个数量级。你可以把共享内存理解成GPU的“CPU L1 Cache”但它可以手动管理所以懂的人和你说的“手工缓存优化”本质就是把数据提前搬到这一步来做。1.2 算力与带宽的“剪刀差”还有一个必须掌握的概念是计算强度Arithmetic Intensity也就是“每字节数据要执行多少次运算”。现代GPU峰值算力FP16/Fp32高得吓人但显存带宽的增长幅度远不如算力。举个例子A100在FP16 Tensor Core下算力高达312 TFLOPS显存带宽2TB/s这意味着你要维持算力跑满每读一个字节的数据至少要执行156次FP16运算。计算强度低于这个拐点的操作叫做访存受限Memory-Bound你再怎么优化kernel的指令流水线也白搭数据搬不过来。深度学习里的很多算子其实都属于访存受限型。比如逐元素的激活函数ReLU、Sigmoid归一化层LayerNorm、BatchNorm的统计部分残差连接相加权重衰减、梯度Clipping这些操作数据量巨大、计算量极小。遇到这类kernel优化的核心就一句话最大化从显存读出来的每个字节的利用率尽量减少不必要的数据搬动。这也是为什么像PyTorch会做“算子融合”比如把激活融合进前一个卷积中的原因本质就是少在HBM和SM之间跑几趟。2. 合并访问决定带宽利用率的“第一性原理”搞清楚了层级差距第二个绕不开的概念是合并访问Coalesced Access。很多人写CUDA kernel算法上看着没问题矩阵乘法算得比CPU快但带宽用不满跑个带宽测试发现只有理论值的20%。这时候第一怀疑对象就是访存模式没有做合并。2.1 合并访问是什么一次缓存行搬运的“团购”GPU的显存控制器以缓存行/Sector为粒度接收访问请求。以A100为例它的L2缓存行和全局内存传输粒度通常按32字节的Sector管理。当同一Warp32个线程访问全局内存时硬件会把它们的地址收集起来尽可能合并成尽量少的事务Memory Transaction去访问。最理想的情况是一个Warp连续访问一段对齐的128字节一个Cache Line硬件只用一次事务就搞定。就像团购一样32个人买同一栋楼的房子只要一趟车就能把全部人送到。而如果这32个线程去访问天南海北的地址每人都要单独发一趟“货车”那运输资源的浪费就是灾难级的。给个代码示例对比。最“教科书级”的错误写法是线程索引和数组访问维度错位// 错误示范线程x对应矩阵的列索引但矩阵按行存储 // 假设矩阵是 row-major并要转置访问 __global__ void badAccessKernel(float* matrix, float* out, int width) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; out[row * width col] matrix[col * width row]; // 非合并按列跳着读 }当线程按threadIdx.x连续变化时读取的地址是在列方向上跳跃的间隔是width * sizeof(float)字节。每一个Warp内32个线程访问的根本不是连续的128字节而是散布在几十个不同的内存页里硬件被迫产生大量内存事务带宽利用率暴跌。正确写法通常有两种一是保持线程遍历连续地址让输出变为非合并二是用共享内存做转置缓冲区先把不规则访问变成块状连续访问再统一读写。训练和推理中很多“transpose”算子慢就是因为生成的访存模式违背了合并原则。2.2 对齐的重要性从256字节到128字节的边界坑除了连续性**对齐Alignment**是另一个容易忽略的细节。GPU通常对32字节、64字节、128字节的访问粒度做合并优化。如果你分配的内存起始地址没有按128字节对齐那么同一个缓存行可能跨越两个甚至多个内存事务边界导致一次合并访问被拆成两次。实际遇到的情况是自己写了一个自定义的Allocator给每个tensor分配的偏移是随意的结果所有kernel带宽直接掉了一半。NVIDIA官方推荐用cudaMalloc通常会自动对齐但如果你实现内存池复用、或者自己写一些奇怪的数据打包就要用到posix_memalign或cudaMallocAsync确保16字节以上对齐。涉及到自定义结构体时确保结构体大小是“最大成员对齐值”的整数倍避免数组元素之间出现未对齐的“坑”。结论永远假设你的Warp需要连续访问连续地址且起始地址对齐到128字节的整数倍。这比任何花哨的优化技巧都重要。3. 别让共享内存“卡脖子”Bank Conflict 与 Padding 的攻防战当你已经把全局内存访问优化到合并了下一步常常是往kernel里加共享内存用来做数据复用、tile缓存。共享内存的带宽虽然比HBM高得多但它并不是无限并发读写。它的底层被分成了32个Bankbank宽度通常为4字节如果一个Warp内多个线程同时访问不同的Bank硬件可以并行处理。但如果两个线程同时访问的是同一个Bank的相同地址会触发Broadcast机制不惩罚如果访问的是同一Bank内的不同地址就会发生Bank Conflict每个冲突会串行化处理。3.1 一个经典陷阱共享内存二维数组的列访问我最早在优化Cuda矩阵乘法tile时遇到了这个坑。定义一个共享内存数组__shared__ float tile[TILE][TILE];读取时如果用的是tile[threadIdx.y][threadIdx.x]这种按列读取的模式那么同一Warp内不同线程的地址在同一列上落到共享内存地址上恰好映射到同一个Bank上瞬间引发32路冲突。一句话总结共享内存的行方向访问最顺畅列方向访问要小心。应对方案就是Padding加填充把维度从[TILE][TILE]改成[TILE][TILE1]这样一行末尾多出来的一个float会把后续行的Bank对齐打散列访问时同一Warp内的地址会错开冲突自然消除。3.2 更隐蔽的冲突数据布局与索引计算Bank Conflict不一定出现在经典的二维矩阵里。比如你在做一个粒子模拟每个线程处理一个粒子粒子属性存成float positions[N][3]你让float3 p positions[tid]实际上结构体内部连续存放3个float。这三个float落在连续的Bank上对单个线程读取一个float3来说没毛病但如果所有线程在不同粒子的同一个分量之间做广播运算Bank分布就会很微妙。这类问题建议用一个小工具思路去自查在kernel里多加两个时钟计数分别测“不加共享内存”和“加共享内存”两个版本耗时如果加完之后时间没有明显下降就要怀疑是共享内存上的冲突把收益抵消了。用cuda-gdb或Nsight Compute可以查shared_st_bank_conflict指标这个指标很直观。3.3 用“共享内存重排”拯救非合并全局访问前面提到矩阵转置的典型非合并访问实际工程里最常用的就是分块转置法。流程大致是每个Block从全局内存中加载一块连续的数据进共享内存注意这一步按行连续读合并设置__syncthreads()同步从共享内存按列读出写入全局内存关键在第三步因为共享内存的反列访问本身有Bank Conflict所以要加上Padding。最终的效果是全局内存两边都是合并访问共享内存上的冲突也通过Padding消掉了。实际操作下来带宽能从原来的非合并版本的20%提升到70%以上具体数字取决于转置规模。4. AI训练推理里的访存优化从PyTorch到Tensor Core的映射在纯CUDA层面优化过之后回到AI框架层面你会发现很多深度学习框架的“魔改”其实就是在内存访问模式上做文章。PyTorch本身会自动优化一部分但如果你对它的底层布局没有概念很容易写出“看起来正确但极慢”的模型代码。4.1 张量布局NCHW与NHWC各有各的访存逻辑在CV模型里张量通常有两种布局NCHWPyTorch默认布局通道维在第二维空间维在最内层。对于每一个像素点它的所有通道是连续并列的在H和W维度上。NHWCTensorFlow/CuDNN一些加速库更偏好的布局空间位置在内存上连续通道维在最后。哪个更好没有绝对。如果你做卷积CuDNN和cuDNN风格kernel通常对NHWC更友好因为它能让Warp内相邻线程访问相邻的像素空间连续而不是跨通道跳跃。在做BatchNorm或LayerNorm时如果统计维度是ChannelNCHW意味着一个Warp访问的是同一个batch下连续空间点的同一种通道数据比较符合合并访问但如果做PixelShuffle这种跨维度重排NHWC会好写很多。工程建议是别管模型里怎么写用Profiler看kernel的访存行为确认瓶颈在哪。如果模型大部分时间花在通道维逐元素操作如激活、BN上NCHW通常没有大问题但如果你用TensorRT或ONNX Runtime部署它们内部极大概率会转成NHWC或更自定义的布局输入数据如果不用to(memory_formattorch.channels_last)对齐就会在边界处产生不必要的拷贝。4.2 非连续张量与contiguous()的代价PyTorch里最常见的访存坑之一是非连续张量non-contiguous tensor。切片操作、转置、某些reshape都会产生一个“视图view”底层数据没有拷贝但形状和stride都变了。这时如果你调用某些算子框架可能会隐式执行contiguous()也就是做一次完整的数据拷贝把旧存储里的数据按新的逻辑顺序重新排列。那次做数据预处理优化发现GPU上有个算子在跑DALI迁移后的代码时特别慢Profiler显示一个copy_占据了将近40%的耗时最后发现就是代码里一行tensor.transpose(1,2)没加.contiguous()导致的。之后推理引擎在准备batch时我直接规定“所有输入必须已经channels_last连续”避免在关键路径上做隐式拷贝。所以看到代码里有大量连续调用的transpose()view()组合时务必检查一下是不是会触发copy。对于优化有两个思路提前一次性布局好不要在每个Batch里反复转置拷贝如果之前的kernel可以实现跨stride访问就直接传入非连续张量避免拷贝不过通常只有专门的kernel才支持4.3 FlashAttention与算子融合为什么不“读三遍”注意力机制中QK^T矩阵计算出来之后要跟V矩阵相乘中间还有Softmax。标准实现里softmax的结果要写回全局内存然后下一遍kernel再读回来。FlashAttention的核心思想就是不把中间结果写回HBM而是在片上利用SRAM共享内存完成分块计算并在kernel内直接做融合。这背后的本质就是访存优化减少HBM读写次数。对AI框架开发者而言我们不需要手写FlashAttention但需要理解为什么某些“算子融合”能获得巨大加速。PyTorch的torch.compile做的是同一件事把逐元素操作融合成一个kernel避免中间张量落回HBM。这篇笔记的实用建议是尽量用torch.compile如果不行就手写一个简单fusion kernel把“读一次数据做完所有逐元素操作”变成地基。5. 定位访存瓶颈的工具与实测别再凭感觉优化优化之前先做Profiling这是性能工程的第一纪律。用直觉猜瓶颈常常会翻车工具能直接告诉你kernel是计算受限还是访存受限以及具体败在哪个环节。5.1 三个核心指标DRAM Throughput, L1/L2 Hit Rate, Memory Bound在Nsight Computencu里打开Kernel Detail 页面核心看三组数据Memory Throughput包括DRAM吞吐、L1/L2 Cache吞吐。如果DRAM Throughput接近100%基本可以断定是访存带宽瓶颈。L1/L2 Hit Rate如果全局内存访问大量能落到L1/L2说明空间局部性不错但如果命中率低且DRAM吞吐爆炸说明你的访存模式没有利用好缓存。Memory/BoundNcu直接给出该kernel是Compute-Bound还是Memory-Bound的判定以及对应的理论瓶颈SM Occupancy。实测中我见过一个Stable Diffusion的UNet自定义算子ncu显示DRAM Throughput 89%SM利用率却只有35%说明纯粹在等数据。该算子做的是x x a * attn_w之类的残差计算强度低到不可救药。解决方案不是写更快的加法而是把多个逐元素算子用融合方式减少访问次数。5.2 一个完整排查例子为什么conv模型的第一个卷积特别慢有一次做优化发现ResNet推理时第一个卷积kernel的耗时是后面的两倍。ncu数据显示该kernel的L2命中率极低DRAM吞吐也不算高但Memory Throughput整体只有60%多。后来发现是因为输入图像的Layout 是NCHW而模型把权重转成了OIHW第一个卷积的Warp读取输入时每个通道的数据在全局内存里间隔很远导致访问无法合并并且造成大量的TLB miss页表缓存缺失。解决办法很简单把输入重新整理成连续的四维张量按NHWC的格式排布通过channels_last结合torch.compile后第一个卷积的速度几乎提升了50%。这个例子典型说明了即使PyTorch看起来在跑同一个卷积实际的数据排布也会让底层kernel走完全不同的访存路径。5.3 微调带宽测试写一个最简单的带宽探针不要只依赖框架自己写一个简单的CUDA带宽测试kernel有助于建立直觉。比如一个纯读取kernel__global__ void readOnlyKernel(const float* in, float* out, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) { out[i] in[i] * 2.0f; } }连续访问数据理论上应该达到接近显存峰值带宽。如果这个kernel的实测带宽达不到理论值的80%就要检查内存分配对齐、访问粒度、以及是否开启了ECC等。这个探针是很多性能工程师的“基准线”后续所有访存优化都和这个基准对比。5.4 PyTorch Profiler中的内存视图框架侧用torch.profiler也有帮助。torch.profiler.profile(activities[ProfilerActivity.CUDA])在结果里能看到每个算子的memory_bandwidth和memory_usage能快速定位到偷走时间的大算子。训练脚本里还可以打开torch.profiler.tensorboard_trace_handler把trace丢到TensorBoard里看GPU utilization曲线和kernel的耗时分布。多数时候模型里最耗时的算子往往不是参数量最大的层而是大量小的、访存受限的逐元素操作。6. 把“带宽换算力”的思维落实到代码里三大通用优化模板最后分享几个我在实际项目里反复用到的通用优化模板。它们不是某个特定算法的优化而是“访存优化”这个通用思想的代码化表达。6.1 Grid-Stride Loop让计算强度提高的经典结构Grid-Stride Loop的核心思想是让固定数量的线程循环处理多个元素而不是为每个数据元素启动一个线程。这样不仅增加了每个线程的数据复用还减少了Block调度开销。对于机器学习数据预处理、纯逐元素kernel这种方式经常能提高缓存命中率。__global__ void gridStrideKernel(const float* in, float* out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; int stride gridDim.x * blockDim.x; for (int i idx; i n; i stride) { out[i] in[i] * 2.0f 1.0f; } }这段代码的访存模式仍然是合并的但循环让它比一次性起大量block的版本更稳定。对于小数据量场景它避免了几十万个小小block的启动开销。对于大数据量场景它让L2 Cache能持续服务相邻数据而不是一个线程只读一次就退出。6.2 向量化访问从 float 到 float4GPU的单个线程可以按向量类型加载数据。float4一次读16字节对齐的访存比一个个读4字节要高效。对于大数据量的纯拷贝/变换任务把数据按float4处理几乎总能带来20%到40%的带宽提升。__global__ void vectorizedKernel(const float4* in, float4* out, int n4) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n4) { float4 val in[i]; val.x * 2.0f; val.y * 2.0f; val.z * 2.0f; val.w * 2.0f; out[i] val; } }需要注意数据长度必须是4的倍数起始地址按照16字节对齐。实际工程里可以考虑把尾数剩余的部分单独用一个scalar kernel处理。在PyTorch扩展里编译期如果知道张量是连续的TensorAccessor的default_accessor会帮你做向量化编译但自定义kernel就得靠自己了。6.3 两阶段KernelTranspose Computation 的融合思路最后是“不用共享内存也能做转置”的融合思路将全流程拆成两个阶段第一阶段把非连续访问的数据按warp tile连续读入寄存器并写到一个scratchpad临时缓冲中第二阶段再从scratchpad按新布局读取并计算。这个思路在FlashAttention里也出现过QK^T结果不写回全局内存而是传给Softmax后直接与V相乘。举这个例子的核心是想说访存优化并不总是纯粹的“减少全局访问”有时候通过增加一次额外的内存读写来换取计算阶段完全合并的访问模式整体收益反而更高。这需要你在“多读一次数据”和“串行化非合并访问”之间做权衡而profiler就是你的天平。7. 访存优化不是玄学把“减少搬运”变成默认思维访存模式优化并不是独立的招数它更像一种系统层面的思维习惯。在写AI算子、调推理引擎、甚至是看PyTorch模型结构时常问自己三个问题能解决90%的访存问题每个数据从HBM读进来之后被完整地利用了多少次如果只用了一次那考虑能否复用或融合Warp里的32个线程访问的地址是连续的吗如果不是考虑重排数据或分块处理有没有中间结果在HBM里写了又读如果有考虑能否在片上完成融合计算我之前优化过一个视频推理管线印象特别深。模型结构不复杂但整体延迟很高。Profiler拉出来一看真正吃时间的不是卷积、不是注意力而是每一帧输入在做标准化时都把整个视频帧的数据从HBM读一遍、写一遍然后下一个算子又读一遍。后来把标准化、缩放、减均值融合成一个kernel几十个算子合一整体吞吐直接翻倍。这就是“打通访存任督二脉”最直观的收益。GPU的内存带宽是稀缺资源算力可以堆晶体管但物理距离和数据线路的带宽是有限制的。所以性能工程做到最后很多时候做的不是“增加计算指令”而是“减少数据的无效旅行”。希望这篇笔记能帮你建立起访存优化的体系框架下次遇到性能问题至少能分清是算计不过来还是数据搬不过来。
很多做AI性能优化的人一开始都会陷入一个误区觉得GPU是“算力怪兽”优化重点就该放在FLOPs上。等真正用ncu一测却发现很多kernel的SM利用率低得可怜Tensor Core在那闲着瓶颈根本不在计算而在数据搬运。做GPU性能工程最重要的一课就是GPU是拿延迟换吞吐的设备但它的计算单元再快也得等数据从显存里搬进来。访存模式Memory Access Pattern就是这套数据流系统的“任督二脉”。入口堵了后面算得再快也是空转。这篇文章是《AI系统性能工程学习笔记》系列的第七篇我结合自己在模型推理、训练kernel优化中的实际踩坑经验把GPU访存的底层机制、常见陷阱和优化手段串一遍重点说清楚“为什么这么改就能快”以及“如何用工具证明它确实变快了”。适合正在做PyTorch算子优化、自定义CUDA kernel或者纯粹被显存带宽困扰的同行参考。1. 先搞懂GPU到底“慢”在哪一份各层级带宽/延迟的对照认知访存优化这件事说白了就是在不同的存储层级之间做数据搬运时间的博弈。想要看懂后续的优化手法第一步要把GPU的内存层级结构刻在脑子里。NVIDIA这边的典型架构从内到外大概是寄存器Register → 共享内存Shared Memory → L1 Cache → L2 Cache → 显存HBM/HBM2e/HBM3。每一层之间的带宽和延迟差异是数量级的差距知道了这个你就明白为什么“数据放对地方”比“多跑几个线程”更重要。1.1 带宽差距的真实数字我整理了一份典型数据以A100 80GB为例不同卡有细微差异但量级一致存储层级典型容量典型带宽典型延迟寄存器每个SM 256KB极高TB/s级别几乎0共享内存/L1每个SM 192KB可配置约几十TB/s级别20-30 cycleL2 Cache40MB约5-6TB/s200-300 cycleHBM显存80GB约2TB/s400-800 cycle这里最重要的数字是HBM显存带宽通常只有2TB/s左右而SM内部共享内存的带宽比它高一个数量级。你可以把共享内存理解成GPU的“CPU L1 Cache”但它可以手动管理所以懂的人和你说的“手工缓存优化”本质就是把数据提前搬到这一步来做。1.2 算力与带宽的“剪刀差”还有一个必须掌握的概念是计算强度Arithmetic Intensity也就是“每字节数据要执行多少次运算”。现代GPU峰值算力FP16/Fp32高得吓人但显存带宽的增长幅度远不如算力。举个例子A100在FP16 Tensor Core下算力高达312 TFLOPS显存带宽2TB/s这意味着你要维持算力跑满每读一个字节的数据至少要执行156次FP16运算。计算强度低于这个拐点的操作叫做访存受限Memory-Bound你再怎么优化kernel的指令流水线也白搭数据搬不过来。深度学习里的很多算子其实都属于访存受限型。比如逐元素的激活函数ReLU、Sigmoid归一化层LayerNorm、BatchNorm的统计部分残差连接相加权重衰减、梯度Clipping这些操作数据量巨大、计算量极小。遇到这类kernel优化的核心就一句话最大化从显存读出来的每个字节的利用率尽量减少不必要的数据搬动。这也是为什么像PyTorch会做“算子融合”比如把激活融合进前一个卷积中的原因本质就是少在HBM和SM之间跑几趟。2. 合并访问决定带宽利用率的“第一性原理”搞清楚了层级差距第二个绕不开的概念是合并访问Coalesced Access。很多人写CUDA kernel算法上看着没问题矩阵乘法算得比CPU快但带宽用不满跑个带宽测试发现只有理论值的20%。这时候第一怀疑对象就是访存模式没有做合并。2.1 合并访问是什么一次缓存行搬运的“团购”GPU的显存控制器以缓存行/Sector为粒度接收访问请求。以A100为例它的L2缓存行和全局内存传输粒度通常按32字节的Sector管理。当同一Warp32个线程访问全局内存时硬件会把它们的地址收集起来尽可能合并成尽量少的事务Memory Transaction去访问。最理想的情况是一个Warp连续访问一段对齐的128字节一个Cache Line硬件只用一次事务就搞定。就像团购一样32个人买同一栋楼的房子只要一趟车就能把全部人送到。而如果这32个线程去访问天南海北的地址每人都要单独发一趟“货车”那运输资源的浪费就是灾难级的。给个代码示例对比。最“教科书级”的错误写法是线程索引和数组访问维度错位// 错误示范线程x对应矩阵的列索引但矩阵按行存储 // 假设矩阵是 row-major并要转置访问 __global__ void badAccessKernel(float* matrix, float* out, int width) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; out[row * width col] matrix[col * width row]; // 非合并按列跳着读 }当线程按threadIdx.x连续变化时读取的地址是在列方向上跳跃的间隔是width * sizeof(float)字节。每一个Warp内32个线程访问的根本不是连续的128字节而是散布在几十个不同的内存页里硬件被迫产生大量内存事务带宽利用率暴跌。正确写法通常有两种一是保持线程遍历连续地址让输出变为非合并二是用共享内存做转置缓冲区先把不规则访问变成块状连续访问再统一读写。训练和推理中很多“transpose”算子慢就是因为生成的访存模式违背了合并原则。2.2 对齐的重要性从256字节到128字节的边界坑除了连续性**对齐Alignment**是另一个容易忽略的细节。GPU通常对32字节、64字节、128字节的访问粒度做合并优化。如果你分配的内存起始地址没有按128字节对齐那么同一个缓存行可能跨越两个甚至多个内存事务边界导致一次合并访问被拆成两次。实际遇到的情况是自己写了一个自定义的Allocator给每个tensor分配的偏移是随意的结果所有kernel带宽直接掉了一半。NVIDIA官方推荐用cudaMalloc通常会自动对齐但如果你实现内存池复用、或者自己写一些奇怪的数据打包就要用到posix_memalign或cudaMallocAsync确保16字节以上对齐。涉及到自定义结构体时确保结构体大小是“最大成员对齐值”的整数倍避免数组元素之间出现未对齐的“坑”。结论永远假设你的Warp需要连续访问连续地址且起始地址对齐到128字节的整数倍。这比任何花哨的优化技巧都重要。3. 别让共享内存“卡脖子”Bank Conflict 与 Padding 的攻防战当你已经把全局内存访问优化到合并了下一步常常是往kernel里加共享内存用来做数据复用、tile缓存。共享内存的带宽虽然比HBM高得多但它并不是无限并发读写。它的底层被分成了32个Bankbank宽度通常为4字节如果一个Warp内多个线程同时访问不同的Bank硬件可以并行处理。但如果两个线程同时访问的是同一个Bank的相同地址会触发Broadcast机制不惩罚如果访问的是同一Bank内的不同地址就会发生Bank Conflict每个冲突会串行化处理。3.1 一个经典陷阱共享内存二维数组的列访问我最早在优化Cuda矩阵乘法tile时遇到了这个坑。定义一个共享内存数组__shared__ float tile[TILE][TILE];读取时如果用的是tile[threadIdx.y][threadIdx.x]这种按列读取的模式那么同一Warp内不同线程的地址在同一列上落到共享内存地址上恰好映射到同一个Bank上瞬间引发32路冲突。一句话总结共享内存的行方向访问最顺畅列方向访问要小心。应对方案就是Padding加填充把维度从[TILE][TILE]改成[TILE][TILE1]这样一行末尾多出来的一个float会把后续行的Bank对齐打散列访问时同一Warp内的地址会错开冲突自然消除。3.2 更隐蔽的冲突数据布局与索引计算Bank Conflict不一定出现在经典的二维矩阵里。比如你在做一个粒子模拟每个线程处理一个粒子粒子属性存成float positions[N][3]你让float3 p positions[tid]实际上结构体内部连续存放3个float。这三个float落在连续的Bank上对单个线程读取一个float3来说没毛病但如果所有线程在不同粒子的同一个分量之间做广播运算Bank分布就会很微妙。这类问题建议用一个小工具思路去自查在kernel里多加两个时钟计数分别测“不加共享内存”和“加共享内存”两个版本耗时如果加完之后时间没有明显下降就要怀疑是共享内存上的冲突把收益抵消了。用cuda-gdb或Nsight Compute可以查shared_st_bank_conflict指标这个指标很直观。3.3 用“共享内存重排”拯救非合并全局访问前面提到矩阵转置的典型非合并访问实际工程里最常用的就是分块转置法。流程大致是每个Block从全局内存中加载一块连续的数据进共享内存注意这一步按行连续读合并设置__syncthreads()同步从共享内存按列读出写入全局内存关键在第三步因为共享内存的反列访问本身有Bank Conflict所以要加上Padding。最终的效果是全局内存两边都是合并访问共享内存上的冲突也通过Padding消掉了。实际操作下来带宽能从原来的非合并版本的20%提升到70%以上具体数字取决于转置规模。4. AI训练推理里的访存优化从PyTorch到Tensor Core的映射在纯CUDA层面优化过之后回到AI框架层面你会发现很多深度学习框架的“魔改”其实就是在内存访问模式上做文章。PyTorch本身会自动优化一部分但如果你对它的底层布局没有概念很容易写出“看起来正确但极慢”的模型代码。4.1 张量布局NCHW与NHWC各有各的访存逻辑在CV模型里张量通常有两种布局NCHWPyTorch默认布局通道维在第二维空间维在最内层。对于每一个像素点它的所有通道是连续并列的在H和W维度上。NHWCTensorFlow/CuDNN一些加速库更偏好的布局空间位置在内存上连续通道维在最后。哪个更好没有绝对。如果你做卷积CuDNN和cuDNN风格kernel通常对NHWC更友好因为它能让Warp内相邻线程访问相邻的像素空间连续而不是跨通道跳跃。在做BatchNorm或LayerNorm时如果统计维度是ChannelNCHW意味着一个Warp访问的是同一个batch下连续空间点的同一种通道数据比较符合合并访问但如果做PixelShuffle这种跨维度重排NHWC会好写很多。工程建议是别管模型里怎么写用Profiler看kernel的访存行为确认瓶颈在哪。如果模型大部分时间花在通道维逐元素操作如激活、BN上NCHW通常没有大问题但如果你用TensorRT或ONNX Runtime部署它们内部极大概率会转成NHWC或更自定义的布局输入数据如果不用to(memory_formattorch.channels_last)对齐就会在边界处产生不必要的拷贝。4.2 非连续张量与contiguous()的代价PyTorch里最常见的访存坑之一是非连续张量non-contiguous tensor。切片操作、转置、某些reshape都会产生一个“视图view”底层数据没有拷贝但形状和stride都变了。这时如果你调用某些算子框架可能会隐式执行contiguous()也就是做一次完整的数据拷贝把旧存储里的数据按新的逻辑顺序重新排列。那次做数据预处理优化发现GPU上有个算子在跑DALI迁移后的代码时特别慢Profiler显示一个copy_占据了将近40%的耗时最后发现就是代码里一行tensor.transpose(1,2)没加.contiguous()导致的。之后推理引擎在准备batch时我直接规定“所有输入必须已经channels_last连续”避免在关键路径上做隐式拷贝。所以看到代码里有大量连续调用的transpose()view()组合时务必检查一下是不是会触发copy。对于优化有两个思路提前一次性布局好不要在每个Batch里反复转置拷贝如果之前的kernel可以实现跨stride访问就直接传入非连续张量避免拷贝不过通常只有专门的kernel才支持4.3 FlashAttention与算子融合为什么不“读三遍”注意力机制中QK^T矩阵计算出来之后要跟V矩阵相乘中间还有Softmax。标准实现里softmax的结果要写回全局内存然后下一遍kernel再读回来。FlashAttention的核心思想就是不把中间结果写回HBM而是在片上利用SRAM共享内存完成分块计算并在kernel内直接做融合。这背后的本质就是访存优化减少HBM读写次数。对AI框架开发者而言我们不需要手写FlashAttention但需要理解为什么某些“算子融合”能获得巨大加速。PyTorch的torch.compile做的是同一件事把逐元素操作融合成一个kernel避免中间张量落回HBM。这篇笔记的实用建议是尽量用torch.compile如果不行就手写一个简单fusion kernel把“读一次数据做完所有逐元素操作”变成地基。5. 定位访存瓶颈的工具与实测别再凭感觉优化优化之前先做Profiling这是性能工程的第一纪律。用直觉猜瓶颈常常会翻车工具能直接告诉你kernel是计算受限还是访存受限以及具体败在哪个环节。5.1 三个核心指标DRAM Throughput, L1/L2 Hit Rate, Memory Bound在Nsight Computencu里打开Kernel Detail 页面核心看三组数据Memory Throughput包括DRAM吞吐、L1/L2 Cache吞吐。如果DRAM Throughput接近100%基本可以断定是访存带宽瓶颈。L1/L2 Hit Rate如果全局内存访问大量能落到L1/L2说明空间局部性不错但如果命中率低且DRAM吞吐爆炸说明你的访存模式没有利用好缓存。Memory/BoundNcu直接给出该kernel是Compute-Bound还是Memory-Bound的判定以及对应的理论瓶颈SM Occupancy。实测中我见过一个Stable Diffusion的UNet自定义算子ncu显示DRAM Throughput 89%SM利用率却只有35%说明纯粹在等数据。该算子做的是x x a * attn_w之类的残差计算强度低到不可救药。解决方案不是写更快的加法而是把多个逐元素算子用融合方式减少访问次数。5.2 一个完整排查例子为什么conv模型的第一个卷积特别慢有一次做优化发现ResNet推理时第一个卷积kernel的耗时是后面的两倍。ncu数据显示该kernel的L2命中率极低DRAM吞吐也不算高但Memory Throughput整体只有60%多。后来发现是因为输入图像的Layout 是NCHW而模型把权重转成了OIHW第一个卷积的Warp读取输入时每个通道的数据在全局内存里间隔很远导致访问无法合并并且造成大量的TLB miss页表缓存缺失。解决办法很简单把输入重新整理成连续的四维张量按NHWC的格式排布通过channels_last结合torch.compile后第一个卷积的速度几乎提升了50%。这个例子典型说明了即使PyTorch看起来在跑同一个卷积实际的数据排布也会让底层kernel走完全不同的访存路径。5.3 微调带宽测试写一个最简单的带宽探针不要只依赖框架自己写一个简单的CUDA带宽测试kernel有助于建立直觉。比如一个纯读取kernel__global__ void readOnlyKernel(const float* in, float* out, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) { out[i] in[i] * 2.0f; } }连续访问数据理论上应该达到接近显存峰值带宽。如果这个kernel的实测带宽达不到理论值的80%就要检查内存分配对齐、访问粒度、以及是否开启了ECC等。这个探针是很多性能工程师的“基准线”后续所有访存优化都和这个基准对比。5.4 PyTorch Profiler中的内存视图框架侧用torch.profiler也有帮助。torch.profiler.profile(activities[ProfilerActivity.CUDA])在结果里能看到每个算子的memory_bandwidth和memory_usage能快速定位到偷走时间的大算子。训练脚本里还可以打开torch.profiler.tensorboard_trace_handler把trace丢到TensorBoard里看GPU utilization曲线和kernel的耗时分布。多数时候模型里最耗时的算子往往不是参数量最大的层而是大量小的、访存受限的逐元素操作。6. 把“带宽换算力”的思维落实到代码里三大通用优化模板最后分享几个我在实际项目里反复用到的通用优化模板。它们不是某个特定算法的优化而是“访存优化”这个通用思想的代码化表达。6.1 Grid-Stride Loop让计算强度提高的经典结构Grid-Stride Loop的核心思想是让固定数量的线程循环处理多个元素而不是为每个数据元素启动一个线程。这样不仅增加了每个线程的数据复用还减少了Block调度开销。对于机器学习数据预处理、纯逐元素kernel这种方式经常能提高缓存命中率。__global__ void gridStrideKernel(const float* in, float* out, int n) { int idx blockIdx.x * blockDim.x threadIdx.x; int stride gridDim.x * blockDim.x; for (int i idx; i n; i stride) { out[i] in[i] * 2.0f 1.0f; } }这段代码的访存模式仍然是合并的但循环让它比一次性起大量block的版本更稳定。对于小数据量场景它避免了几十万个小小block的启动开销。对于大数据量场景它让L2 Cache能持续服务相邻数据而不是一个线程只读一次就退出。6.2 向量化访问从 float 到 float4GPU的单个线程可以按向量类型加载数据。float4一次读16字节对齐的访存比一个个读4字节要高效。对于大数据量的纯拷贝/变换任务把数据按float4处理几乎总能带来20%到40%的带宽提升。__global__ void vectorizedKernel(const float4* in, float4* out, int n4) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n4) { float4 val in[i]; val.x * 2.0f; val.y * 2.0f; val.z * 2.0f; val.w * 2.0f; out[i] val; } }需要注意数据长度必须是4的倍数起始地址按照16字节对齐。实际工程里可以考虑把尾数剩余的部分单独用一个scalar kernel处理。在PyTorch扩展里编译期如果知道张量是连续的TensorAccessor的default_accessor会帮你做向量化编译但自定义kernel就得靠自己了。6.3 两阶段KernelTranspose Computation 的融合思路最后是“不用共享内存也能做转置”的融合思路将全流程拆成两个阶段第一阶段把非连续访问的数据按warp tile连续读入寄存器并写到一个scratchpad临时缓冲中第二阶段再从scratchpad按新布局读取并计算。这个思路在FlashAttention里也出现过QK^T结果不写回全局内存而是传给Softmax后直接与V相乘。举这个例子的核心是想说访存优化并不总是纯粹的“减少全局访问”有时候通过增加一次额外的内存读写来换取计算阶段完全合并的访问模式整体收益反而更高。这需要你在“多读一次数据”和“串行化非合并访问”之间做权衡而profiler就是你的天平。7. 访存优化不是玄学把“减少搬运”变成默认思维访存模式优化并不是独立的招数它更像一种系统层面的思维习惯。在写AI算子、调推理引擎、甚至是看PyTorch模型结构时常问自己三个问题能解决90%的访存问题每个数据从HBM读进来之后被完整地利用了多少次如果只用了一次那考虑能否复用或融合Warp里的32个线程访问的地址是连续的吗如果不是考虑重排数据或分块处理有没有中间结果在HBM里写了又读如果有考虑能否在片上完成融合计算我之前优化过一个视频推理管线印象特别深。模型结构不复杂但整体延迟很高。Profiler拉出来一看真正吃时间的不是卷积、不是注意力而是每一帧输入在做标准化时都把整个视频帧的数据从HBM读一遍、写一遍然后下一个算子又读一遍。后来把标准化、缩放、减均值融合成一个kernel几十个算子合一整体吞吐直接翻倍。这就是“打通访存任督二脉”最直观的收益。GPU的内存带宽是稀缺资源算力可以堆晶体管但物理距离和数据线路的带宽是有限制的。所以性能工程做到最后很多时候做的不是“增加计算指令”而是“减少数据的无效旅行”。希望这篇笔记能帮你建立起访存优化的体系框架下次遇到性能问题至少能分清是算计不过来还是数据搬不过来。
