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为例，不同卡有细微差异，但量级一致）：

这里最重要的数字是：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管理。当同一Warp（32个线程）访问全局内存时，硬件会把它们的地址收集起来，尽可能合并成尽量少的事务（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个Bank（bank宽度通常为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][TILE+1]，这样一行末尾多出来的一个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 用“共享内存重排”拯救非合并全局访问

前面提到矩阵转置的典型非合并访问，实际工程里最常用的就是分块转置法。流程大致是：

1. 每个Block从全局内存中加载一块连续的数据进共享内存，注意这一步按行连续读（合并）
2. 设置__syncthreads()同步
3. 从共享内存按列读出，写入全局内存

关键在第三步，因为共享内存的反列访问本身有Bank Conflict，所以要加上Padding。最终的效果是全局内存两边都是合并访问，共享内存上的冲突也通过Padding消掉了。实际操作下来，带宽能从原来的非合并版本的20%提升到70%以上，具体数字取决于转置规模。

4. AI训练推理里的访存优化：从PyTorch到Tensor Core的映射

在纯CUDA层面优化过之后，回到AI框架层面，你会发现很多深度学习框架的“魔改”其实就是在内存访问模式上做文章。PyTorch本身会自动优化一部分，但如果你对它的底层布局没有概念，很容易写出“看起来正确但极慢”的模型代码。

4.1 张量布局：NCHW与NHWC各有各的访存逻辑

在CV模型里，张量通常有两种布局：

• NCHW：PyTorch默认布局，通道维在第二维，空间维在最内层。对于每一个像素点，它的所有通道是连续并列的（在H和W维度上）。
• NHWC：TensorFlow/CuDNN一些加速库更偏好的布局，空间位置在内存上连续，通道维在最后。

哪个更好？没有绝对。如果你做卷积，CuDNN和cuDNN风格kernel通常对NHWC更友好，因为它能让Warp内相邻线程访问相邻的像素（空间连续），而不是跨通道跳跃。在做BatchNorm或LayerNorm时，如果统计维度是Channel，NCHW意味着一个Warp访问的是同一个batch下连续空间点的同一种通道数据，比较符合合并访问；但如果做PixelShuffle这种跨维度重排，NHWC会好写很多。

工程建议是：别管模型里怎么写，用Profiler看kernel的访存行为，确认瓶颈在哪。如果模型大部分时间花在通道维逐元素操作（如激活、BN）上，NCHW通常没有大问题；但如果你用TensorRT或ONNX Runtime部署，它们内部极大概率会转成NHWC或更自定义的布局，输入数据如果不用to(memory_format=torch.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 Compute（ncu）里，打开Kernel Detail 页面，核心看三组数据：

• Memory Throughput：包括DRAM吞吐、L1/L2 Cache吞吐。如果DRAM Throughput接近100%，基本可以断定是访存带宽瓶颈。
• L1/L2 Hit Rate：如果全局内存访问大量能落到L1/L2，说明空间局部性不错；但如果命中率低且DRAM吞吐爆炸，说明你的访存模式没有利用好缓存。
• Memory/Bound：Ncu直接给出该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 到 float4

GPU的单个线程可以按向量类型加载数据。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 两阶段Kernel：Transpose + Computation 的融合思路

最后是“不用共享内存也能做转置”的融合思路：将全流程拆成两个阶段，第一阶段把非连续访问的数据按warp tile连续读入寄存器并写到一个scratchpad（临时缓冲）中；第二阶段再从scratchpad按新布局读取并计算。这个思路在FlashAttention里也出现过：QK^T结果不写回全局内存，而是传给Softmax后直接与V相乘。

举这个例子的核心是想说：访存优化并不总是纯粹的“减少全局访问”，有时候通过增加一次额外的内存读写，来换取计算阶段完全合并的访问模式，整体收益反而更高。这需要你在“多读一次数据”和“串行化非合并访问”之间做权衡，而profiler就是你的天平。

7. 访存优化不是玄学：把“减少搬运”变成默认思维

访存模式优化并不是独立的招数，它更像一种系统层面的思维习惯。在写AI算子、调推理引擎、甚至是看PyTorch模型结构时，常问自己三个问题，能解决90%的访存问题：

1. 每个数据从HBM读进来之后，被完整地利用了多少次？如果只用了一次，那考虑能否复用或融合
2. Warp里的32个线程访问的地址是连续的吗？如果不是，考虑重排数据或分块处理
3. 有没有中间结果在HBM里写了又读？如果有，考虑能否在片上完成融合计算

我之前优化过一个视频推理管线，印象特别深。模型结构不复杂，但整体延迟很高。Profiler拉出来一看，真正吃时间的不是卷积、不是注意力，而是每一帧输入在做标准化时都把整个视频帧的数据从HBM读一遍、写一遍，然后下一个算子又读一遍。后来把标准化、缩放、减均值融合成一个kernel，几十个算子合一，整体吞吐直接翻倍。这就是“打通访存任督二脉”最直观的收益。

GPU的内存带宽是稀缺资源，算力可以堆晶体管，但物理距离和数据线路的带宽是有限制的。所以性能工程做到最后，很多时候做的不是“增加计算指令”，而是“减少数据的无效旅行”。希望这篇笔记能帮你建立起访存优化的体系框架，下次遇到性能问题，至少能分清是算计不过来，还是数据搬不过来。