资讯动态

内存受限算子优化实战:从带宽瓶颈到算子融合的加速方法

发布时间:2026/10/17 2:06:53 来源:尧图企业网站定制
1. 什么是内存受限mem bound算子它到底限制了什么1.1 核心定义当数据搬运成为瓶颈前几天和一位做推理引擎的朋友线上交流他说自己花了一整个月优化各种卷积和矩阵乘算子端到端延迟却几乎没怎么降。我看了一眼他贴出来的性能剖析截图DRAM 吞吐量接近硬件上限而 SM 里的计算单元占用率低得可怜立刻意识到问题出在哪儿——他一直在优化计算受限的算子但真正拖后腿的是那批内存受限memory-bound业内常简写为 mem bound算子。所谓内存受限算子是指执行时间的瓶颈不在于“算得慢”而在于“数据搬运太慢”。一个算子从输入到输出如果计算强度arithmetic intensity即 FLOPs/Byte 的比值很低意味着每做一次浮点运算都要从内存里搬好几个字节的数据那么 GPU 或 CPU 的算力再强也没用DDR 或 HBM 的带宽才是真正的水龙头。这类算子在我们日常的模型里到处都是逐元素的 ReLU、Sigmoid、Cast 类型转换、残差相加、Channel 维度上的 Normalization、Softmax 的中间阶段、Reduce 求和、数据拷贝与 Layout 转换。它们往往单个看起来都很“轻”但串联起来却能在端到端推理里贡献超过一半的延迟尤其是在 batch size 比较小的时候。1.2 计算受限与内存受限的本质差异要彻底理解 mem bound先要建立一个简单的模型。任何算子在硬件上的执行粗略来说就是在“搬运数据”和“执行计算”两件事之间循环。我们可以用下面的对比表来感受两类算子的差异维度计算受限compute-bound内存受限memory-bound典型代表大尺寸卷积、矩阵乘、FFT激活函数、归一化、Copy、Reduce计算强度高通常 10 FLOP/Byte低通常 1~2 FLOP/Byte瓶颈资源乘加单元、ALU、Tensor CoreDRAM 带宽、缓存带宽优化方向减少冗余计算、提高计算密度减少访存量、提升有效带宽性能特征利用率高硬件忙利用率低硬件在等待这个表看起来很理论但在实际调优中非常有用。我记得有一次优化某个检测模型的前处理部分把输入图像从 uint8 转成 float32 再除以 255然后再做归一化减均值乘方差。这个流程三个算子单独看都不复杂但写成三个 kernel 执行的话每张图至少要把整块图像数据从全局内存读三遍、写三遍。剖析结果一出来计算单元几乎全程在看戏HBM 带宽却已经跑满了。这种情况下再怎么增加 CUDA core 的数量也救不回来唯一的出路就是减少内存访问次数。一个更直观的生活化类比是假如你要把一百箱货物从仓库搬进车间加工加工台的速度极快每箱只需要一秒钟但从仓库到车间的通道很窄搬一箱要花十秒钟。这时你的生产效率不是由加工台决定的而是由那条通道决定的。内存受限算子就是这么个状态计算单元就是加工台内存带宽就是那条通道。1.3 为什么 mem bound 算子在真实场景里比想象中更多很多人有一个误区觉得模型里大部分时间都应该花在卷积和 attention 这种“大算子”上。但如果你用 profiler 去统计一个真实部署模型的算子耗时分布往往会发现一个惊人的事实卷积和 GEMM 都已经被各种加速库优化得非常好了剩下的时间大头反而集中在那些不起眼的小算子上面。比如 NLP 模型里的 LayerNorm、Softmax、GELU、残差连接这些算子频繁出现而且每次处理的数据量都不大。之前在某语音识别模型上做过一次统计Transformer 结构里非 GEMM 算子占总延迟的比例可以到 40%~60%。这些算子的计算量很小但每个都要完整地读一遍激活值、写一遍激活值在 decode 阶段 batch 又小数据根本喂不饱计算单元。这时候你就是一个典型的 mem bound 场景。2. 先搞清楚带宽的账再谈优化2.1 硬件视角GPU 的存储层级和带宽要优化内存受限算子首先要明白数据在这些硬件上是如何流动的。以一块常见的 NVIDIA GPU 为例从离计算单元最近的地方开始寄存器Register每个 SM 里都有海量的寄存器文件访问延迟只有几个周期带宽极大但容量有限总共不过几百 KB 级别。共享内存Shared Memory片上 SRAM延迟几十个周期带宽远高于全局内存但容量通常只有几十 KB 到两百多 KB。L1 / L2 缓存L1 紧贴 SML2 是芯片上最后一道缓存防线。L2 缓存命中时能大幅降低对 DRAM 的访问压力。全局内存Global Memory / DRAMHBM、GDDR 等显存容量最大但带宽是最紧张的资源延迟通常要几百个周期。一个算子在执行时数据从 DRAM 一级一级搬到寄存器这个过程如果被频繁触发DRAM 带宽就会被耗尽。所以内存受限优化的第一原则很简单能在片上解决的数据访问就别往片外跑能少访问一遍的数据就坚决不要访问两遍。我在实际开发中经常用“带宽预算”这个概念来指导优化。假设你的 GPU 实测可持续带宽是 B GB/s某个算子的输入输出总数据量是 D GB那么理论上最低耗时就是 D / B 秒。你不可能突破这个下限除非减少 D。所以优化的本质就是压缩 D而不是去调什么编译器优化等级。2.2 用 profiler 数据判断到底是不是 mem bound口说无凭用性能分析工具量化一下才靠谱。在 CUDA 生态里最常用的是 Nsight Computencu和 Nsight Systemsnsys。如果你做的是 AI 框架层面可以用 PyTorch Profiler、TensorRT 的 perf 工具或者简单点用 CUDA Events 手工计时。对单个 kernelncu 的Memory Throughput指标非常关键。它给出的是 kernel 执行期间内存子系统L1、L2、DRAM的利用率。如果 DRAM Throughput 长期在 90% 以上而 SM 的 Compute Utilization 在 30% 以下这基本就是板上钉钉的 mem bound。另外还有Memory Chart和Warp Stall Sampling可以看到 warp 到底是在等数据还是在等计算。这里我提供一个非常简单的经验判断法看 kernel 的“计算强度”。把 kernel 执行的总浮点运算次数FLOPs除以实际从 DRAM 读写的数据量Bytes如果结果小于硬件的“拐点值”很多 GPU 在 5~10 FLOP/Byte 之间那这个 kernel 大概率是内存受限。反之如果远大于这个值那就是计算受限。搞清楚这一点方向上就不会出大错。2.3 一个容易被人忽略的点实测带宽不等于理论带宽写代码之前还要有个心理预期你实际能拿到的内存带宽远低于硬件规格书上的理论峰值。理论峰值是工程上几乎达不到的数字实际可持续带宽因为访存模式、页表开销、DRAM 刷新等因素能跑出理论值的 70%~85% 就已经非常好了。所以在优化 mem bound 算子时不要好高骛远地盯着理论带宽去设计目标而是先写一个最简单的 copy kernel比如cudaMemcpy或者设备端大数组复制当作 “bandwidth test”测出你这个硬件在当前情况下能稳定达到的带宽上限。把这个实测值当作天花板再去评估算子优化空间。这个习惯帮我避过不少坑有时前面优化半天最后发现瓶颈其实在 PCIe 传输或者 CPU 侧的数据搬运GPU 内核本身早就到顶了。3. 加速 mem bound 算子的核心方法论3.1 算子融合把多次搬运压缩成一次内存受限优化的第一选择永远不是让单个 kernel 变得更快而是从整体结构上减少 kernel 的个数。算子融合就是把多个内存受限算子合并到一个 kernel 里执行让中间结果留在寄存器或共享内存里不落回全局内存。举个例子常见的BiasAdd ReLU Cast三段操作。如果拆开写每个算子都要完整读写一遍数据如果融合成一个 kernel数据从全局内存读一次在寄存器里完成加偏置、激活、类型转换最后写回一次。内存访问从 3 读 3 写降为 1 读 1 写理论上就能缩短近三倍的访存时间。实际效果还会受到各种开销影响但思路是完全正确的。在 PyTorch 里很多人没意识到torch.add和torch.relu如果分两次调用每次都会启动一个 kernel。而写成torch.relu(x bias)这种复合表达式PyTorch 有时会借助fusion机制合并执行。但更可靠的做法是用专门的算子融合编译器比如 TVM、Triton、TensorRT 或者自定义 CUDA kernel。实际工作中我倾向于对高频的模式手写融合 kernel因为可控性最强。3.2 向量化访问让每次访存指令发挥最大价值在 GPU 上一个线程一次能处理的数据宽度是有限的。默认情况下如果你写的是 C 标量代码编译器可能会生成 32 位或 64 位的访存指令。但当你有连续内存访问、且数据对齐时使用float4、double2、int4这类向量类型一次指令就能搬运 16 字节显著减少指令数提高访存效率。这个操作的原理在于现代内存系统对连续的大块访问非常友好向量化可以让请求的事务数量减少减少地址计算和指令调度的开销同时提升 DRAM 的有效带宽利用率。我给大家看一个非常经典的例子。假设你要对一个大数组做逐元素y[i] x[i] * 2.0f普通写法每个线程处理一个元素__global__ void scale_kernel(const float* x, float* y, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) { y[i] x[i] * 2.0f; } }改成每线程处理四个元素__global__ void scale_kernel_vec4(const float4* x, float4* y, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n / 4) { float4 v x[i]; v.x * 2.0f; v.y * 2.0f; v.z * 2.0f; v.w * 2.0f; y[i] v; } }别小看这个改动。对纯 Copy 类和逐元素类算子向量化通常能带来 20%~50% 的带宽提升因为事务数少、地址对齐好。前提是内存地址按 16 字节对齐CUDA 分配器通常满足这一点并且数组长度是向量宽度的整数倍。如果不能整除可以对剩余部分单独处理或者让尾部用掩码方式处理不要让边界条件拖垮主体性能。3.3 合并访存模式让相邻线程访问相邻地址GPU 全局内存是按“事务”来读取的硬件会把同一个 warp32 个线程发出的访存请求合并成尽可能少的几个内存事务。理想情况是线程 i 访问地址 i、线程 i1 访问地址 i1这样 32 个线程的访问落在连续区间内内存控制器一次就能搞定。反之如果相邻线程访问相隔很远的内存地址硬件就需要拆成很多次事务吞吐量会断崖式下跌。这就是为什么在写 mem bound kernel 时你要时刻问自己同一个 warp 里 32 个线程在访问哪些地址它们连续吗一个常见的反例是在二维矩阵处理里用错了索引方式导致每个线程负责一整行而相邻线程访问的行之间相隔一个 stride。这样访存完全乱序带宽能损失好几倍。调整方法通常就是让线程的threadIdx.x对应最内层连续变化的索引。如果你要处理的是矩阵常见策略是把矩阵按行展平然后让 block 中的连续线程按线性 ID 去取连续元素而不是让一个线程去处理一行。这种“网格跨度循环”grid-stride loop配合展平索引既能保证合并访存又能处理任意规模的数据。3.4 选择合适的 block 大小和 grid 配置很多初学者以为 block 越大越好或者 256、512 随便选一个就行。但内存受限算子的性能跟线程配置的关系虽然不像计算受限那么大却也不可忽视。关键原因是 GPU 需要足够的并行度去掩盖内存延迟。一个粗略的经验法则是让总的线程数至少能覆盖内存带宽需求。如果每个线程访问的数据太少内存系统可能没有被充分“压满”如果每个线程访问的数据太多又有可能引入过多的寄存器压力或者导致负载不均衡。在实践中我会从每个线程处理 4 个元素、每个 block 256 个线程开始测用ncu看 DRAM 利用率再尝试 128、512、1024 这些配置。如果 DRAM 利用率稳定在 90% 以上说明访存压力已经足够再调配置收益不大。顺带提一句__launch_bounds__可以用来限制编译器对寄存器数量的分配避免因为寄存器占用过高导致占用率不足。3.5 减少数据精度和数据量从源头压缩访问量从带宽公式就能推导出一个最简单的结论数据量越小搬运时间越短。所以如果能从上游直接拿到低精度输入为什么还要先转成 float32 再做计算呢比如量化模型里权重是 int8激活值是 int8中间结果如果也能保持在 int8/int16 进行计算访存量就比 float32 少一倍到四倍。很多推理框架之所以在 CPU 或 GPU 上做 INT8 推理能获得接近两倍的加速很大一部分功劳就来自内存带宽压力的降低而不是单纯的计算变快。另一个思路是尽量避免在算子链的中间阶段产生大的中间 tensor。比如在训练框架中有些操作会为了反向传播保留中间激活值导致内存占用和访问量飙升。如果能在不影响数值结果的前提下用重计算recompute替代部分存储也可以变相减少内存流量。3.6 用共享内存和 L2 缓存做数据复用对于需要跨线程共享数据的算子比如矩阵的行操作、通道维度的 Reduce可以先把数据从全局内存批量加载到共享内存再由所有线程从共享内存中读取。这样可以避免每个线程都直接访问全局内存把 DRAM 的流量分摊到片上。但要清醒认识到一点共享内存不是灵丹妙药。如果数据本身没有被复用只是每个线程各取各的那搬进共享内存反而多了一次写一次读亏得更多。只有存在明显的复用模式比如多个线程需要读取同一个邻域、同一个 tile 或者同一条通道的数据时共享内存才值回票价。L2 缓存的作用也常常被忽略。某些算子如果访问的是几十 MB 内的数据而 GPU 的 L2 有 40MB 甚至更大那么第二遍访问就有可能在 L2 命中。此时不需要逐字节地压 DRAM。举个实际例子在 Transformer 推理中K/V cache 通常不大如果我们能把同一 token 序列的多次计算控制在 L2 能覆盖的范围内能显著减少 DRAM 流量。这点在做 decode 阶段优化时尤其有用适当调整 batch size 和序列长度让参与运算的数据集“塞进”L2效果立竿见影。3.7 访存与计算重叠流水线思维在端到端层面mem bound 算子的加速不只靠内核本身。现代 GPU 支持多个 kernel 并发执行也支持内存拷贝与计算的重叠。如果两个 kernel 之间没有依赖关系理论上可以让一个 kernel 在等待内存数据时另一个 kernel 的计算已经在进行。在 CUDA 里用多个 stream 就可以实现这种重叠。具体的做法是把不依赖的数据切分成多个 chunk分配到不同的 stream 里每个 stream 内部按“拷入 - 计算 - 拷出”三个阶段流水化。这里的核心是保证 stream 之间有足够的独立性否则 GPU 的硬件调度器也无法把它们拉开并行。很多人用多 stream 没有效果就是因为 kernel 之间实际上存在隐式依赖比如都去写同一个输出缓冲区或者都依赖同一个全局内存计数器。同步开销也得注意。stream 之间如果频繁做cudaStreamSynchronize同步带来的等待会直接吞掉流水线带来的收益。更精细的做法是用事件event管理 stream 间的依赖让硬件自己调度。4. 实操案例拆解一个 LayerNorm 的完整加速过程4.1 基线实现与剖析数据理论讲多了容易飘还是用一个真实度很高的案例来落地。假设我们要实现一个针对[batch, seq_len, hidden]形状的 LayerNormhidden 4096batch × seq_len 1024。朴素的 PyTorch 实现长这样def layer_norm(x, gamma, beta, eps1e-5): mean x.mean(dim-1, keepdimTrue) var x.var(dim-1, unbiasedFalse, keepdimTrue) return (x - mean) / torch.sqrt(var eps) * gamma beta这里的mean和var各启动一个 reduction kernel随后的归一化又启动一个大 kernel。实际上PyTorch 在这个场景下会调用一个高度优化的 CUDA 实现不会真的有四个 kernel但我们可以借此理解优化思路。用 ncu 对自定义的朴素实现进行剖析看到的核心数据大致是这样的DRAM 吞吐量约 62% 的理论峰值计算利用率约 12%Kernel 执行时间约 180 微秒访存指令占比很高大量时间花在加载 x 上判定明显的 mem bound。而且 DRAM 利用率只有 62%说明数据还没有把内存系统压满异常点在于 kernel 内部有多次读取中间结果、多次写回的情况。4.2 第一轮优化算子融合减少全局内存往返第一件事就是把“算均值”、“算方差”、“归一化”这三件事合并到一个 kernel 里。 LayerNorm 的核心逻辑是每个样本一行 hidden 个元素需要求出均值和方差然后再逐元素归一化。如果不做融合第一个 kernel 要完整读一遍 x 来算均值第二个 kernel 又要读一遍 x 算方差第三个 kernel 再读一遍 x 做归一化。三次读取 4096 × 1024 个 float这就是十几 MB 的读流量全部压向 DRAM。一个典型的融合策略是两遍法加一遍法。如果 hidden 维度足够大一个 block 处理一行数据先做一次 block 内的 reduce 算出均值再做一次 reduce 算方差最后统一做归一化。这样 x 最多读两遍。如果想做到只读一遍那就需要引入 Welford 在线算法一次遍历同时累积均值和方差。实际代码里通常使用 Welford 来避免数值稳定性问题。融合之后ncu 显示的时间和吞吐量有显著变化DRAM 吞吐量拉到 88%Kernel 执行时间约 95 微秒节省差不多一半的时间第一轮优化收益最大的原因很简单我们砍掉了多余的全局内存访问这是 mem bound 优化里最划算的动作。4.3 第二轮优化向量化与负载均衡Baseline 融合后的 kernel 可能是一个线程处理一个元素或者一个 block 串行处理一大批元素。下一步是引入float4向量化访问同时仔细设计 block 的切分逻辑。假设 hidden 4096我们可以把一个 block 设为 256 个线程每个线程用float4处理 4 个元素这样一个 block 正好覆盖 1024 个元素。但这样需要 4 个 block 才能处理一行 4096 个元素跨 block 做 reduce 就很麻烦。另一个方法是让一个 block 处理一行中的 4096 个元素但 block 内部先做分段的局部 reduce最后通过__shfl指令归并。这里有一个很重要的权衡block 太大可能导致 reduce 阶段的同步开销变大但太小又会导致每个线程处理的元素变少、访存请求不够密集。实际测试后我用 grid 覆盖所有行每个 block 使用 256 个线程处理一行的一个 tile比如 1024 个元素通过两次循环遍历处理完整行既照顾了向量化也让 reduce 的复杂度可控。这一轮优化的收益来自于访存指令数的减少和 load/store 效率的提升DRAM 吞吐量从 88% 提升到 93%Kernel 执行时间约 78 微秒4.4 第三轮优化针对小 batch 场景的 L2 优化如果上面的场景改成 decode 阶段batch × seq_len 可能很小比如等于 8那整个 LayerNorm 的输入就只有 32KB 左右。此时 DRAM 带宽远没有跑满延迟可能在 kernel 启动和同步上。优化思路要完全换掉减少 kernel 启动次数把多个小算子进一步融合成一个大算子。使用 CUDA Graph 捕获一组操作减少 launch 开销。如果数据小到能放进 L2确保访问模式对 L2 友好避免无谓地冲刷缓存。这种小算子的场景花时间调内核内部的意义已经不大了更大的瓶颈在 CPU 的调度延迟和 GPU 的 kernel 启动间隙上。我们可以把十几个小算子打包成一个 CUDA Graph实测可以省掉大量单次 launch 的微秒级开销。对于大批次的训练场景启动开销被数据量摊薄这个优化就没那么明显。4.5 优化后的经验总结经过三轮优化LayerNorm 从 180 微秒降到 78 微秒收益主要来自访问量减少和访存指令效率提升。放在真实模型里这种优化直接影响到端到端吞吐。更大规模的工程还有一招如果你能把它和前面的 GEMM 或者后面的注意力融合成一个超大的自定义 kernel那中间结果的读写就能彻底省掉。不过这种方案工程复杂度高、调式难度大而且对不同模型结构适配成本不一样。我的建议是先用低成本的融合和向量化吃到大头收益再评估是否值得上更大的融合方案。5. 内存受限算子优化的常见误区与实战排查5.1 误区一以为所有 mem bound 都是 DRAM 带宽不够不少人一看 mem bound 就去查 DRAM 带宽但这只是最表层的解读。实际上内存受限有时是 L2 带宽不够、共享内存 bank conflict、或者 TLB miss 过多导致的。DRAM 利用率不高但 kernel 依然很慢的情况我遇到过很多次。排查方法很简单在 ncu 里区分看L1/TEX Cache Throughput、L2 Cache Throughput、DRAM Throughput三项指标。如果 DRAM 没跑满而 L1 或 L2 已经接近 100%那就不是 DRAM 的事而是你的访问模式在缓存层出了问题。比如共享内存 bank conflict 会让 L1 的带宽被无效的多次访问拖垮此时调整数据在共享内存中的布局或者填充 padding往往能立竿见影。5.2 误区二盲目增大 block 数量或线程数量有一种错误直觉是 mem bound 算子反正不费算力线程越多越好。事实上线程太多、每个线程处理的数据太少有时会增加调度和访存请求的数量导致性能反而下降。尤其是当每个线程处理 1 个元素时warp 里的访存事务可能非常密集但指令调度的开销也被放大了。实践方法是固定一个合理的每个线程处理元素数常见 4~16 之间慢慢调整 block 大小观察 DRAM 吞吐量和时延的变化。当吞吐量不再明显上升时就已经到了硬件的“饱和点”再堆线程只会增加开销。5.3 误区三只优化计算内核忽略内存拷贝和设备端数据传输如果你是在做工程部署特别需要注意 H2DHost to Device和 D2HDevice to Host的拷贝。很多时候你在 GPU kernel 里省下来的时间会被每一次从 CPU 同步数据到 GPU 的 PCIe 传输时间吃掉。我见过一个在线推理服务前处理在 CPU 上做把处理好的数据拷贝到 GPU然后 GPU 跑模型再把结果拷回来。数据分析下来PCIe 传输时间占了总延迟的一半以上。这种场景问题的关键不在 mem bound 算子而在数据流设计。解决思路是让前处理也搬上 GPU或者采用异步传输、双缓冲流水线让拷贝与计算重叠。优化的顺序永远是先放大处再抠细节。5.4 排查流程一个可复用的脚本化思路我自己在工程中遇到性能问题时基本都会走一套固定流程用 nsys 做全局视角 profile找出耗时 Top10 的 kernel 和 API 调用。对 Top 候选用 ncu 跑一次细节分析记录 DRAM/L2/L1 吞吐、计算利用率、指令分布。判断每个 kernel 属于计算受限还是内存受限贴上标签。对内存受限的先检查有没有多余的全局内存读写、访存模式是否合并、能否向量化。对计算受限的再去考虑算法改进、Tensor Core 使用、减少冗余计算。每次优化只改一个变量重新剖析对比确保每一轮改进都有数据支撑。这套流程看上去不炫酷但极其有效。它保证你不会在一个方向上投入太多精力而忽略了更明显的瓶颈。5.5 实测技巧巧用 bandwidth test 校准目标最后分享一个很实用的小技巧。在开始优化前我都会先写一个最简单的数据复制 kernel用它来测出当前硬件实际能达到的最高带宽。__global__ void copy_kernel(const float4* src, float4* dst, int n4) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n4) { dst[i] src[i]; } }测量这个 kernel 的耗时后用总字节数输入加输出除以时间就得到该算子在当前环境下的“实用带宽上限”。后面所有 mem bound 优化都以这个数字为参考如果优化后的 kernel 带宽已经贴近它就不要再挤了应该把精力转向减少访问量或者优化别的算子。我在几次调优过程中发现所谓“优化了半天提不上去”的 case很多其实是已经撞到了这个实际带宽天花板。这时候你想的不是继续压榨内核而是换一种思路比如算子融合、精度降低、或者把部分计算挪到数据量更小的地方去执行。6. 从算子到系统内存受限优化在真实项目中的坐标6.1 算子加速的上限取决于数据流设计单独优化一个算子收益终究有限。真正体现系统能力的地方在于把整个模型的数据流梳理清楚。一张输入图片从内存进入 GPU经过预处理、网络推理、后处理中间每一步的数据格式、精度、存储位置都决定了内存受限算子能否进一步提速。比如在 NLP 模型的 decode 阶段K/V cache 的读取模式极其频繁且重复。如果能把这个数据的布局改成更利于顺序访问的形式或者让请求的 batch 凑成更大的连续块相当于把一连串的内存受限小算子全部照顾到了。相反如果你只盯着单个算子去抠全局的数据流设计不合理很快你就会发现改来改去都只在边缘摩擦。6.2 不同硬件平台上的 mem bound 差异很多人以为内存受限算子在 GPU 上面的优化经验和 CPU 上的可以完全照搬其实不然。CPU 的内存受限问题通常会和 cache line、NUMA 亲和性、SIMD 指令宽度、内存通道数都搅在一起。GPU 上合不合适CPU 上不一定合适GPU 上强调 vectorized load 和 warp 合并访存CPU 上则更关注缓存分块和预取。我做过一个 ARM 端的算子移植原本在 GPU 上表现极佳的 float4 向量化在 ARM NEON 平台上要改成 128 位寄存器指令而 GPU 上的 block 归约逻辑到了 CPU 上要换成 OpenMP 并行 局部归约。底层原则一致都是减少内存访问、提升单位指令的数据吞吐量但实现细节完全不同。所以如果你在移植别人的优化代码之前先不要问“这段 kernel 写得对不对”而是先问“这个硬件平台的数据访问特性和优化 target 是什么”。否则很容易出现 GPU 上 3 倍加速的算子到 CPU 上反而比朴素实现更慢的诡异情况。6.3 和 AI 框架的相互配合做工程落地时你没法完全绕开 AI 框架去手写所有算子。PyTorch、TensorFlow 这些框架自带了很多优化过的算子但它们的融合能力、调度方式各有不同。有些框架对torch.jit.script或者其他编译器的支持并不总是稳定的你花大力气写的融合 kernel 可能因为一个 dispatch 判断没走对路径就根本不会被调用。我的建议是在实际框架里集成自定义算子时先用框架自带的 profiler 确认代码确实执行了你写的那个 CUDA kernel而不是 fallback 到默认实现。见过太多人改了半天的算子最后发现因为 shape 动态变化或者 device 不一致框架压根没调用他们的实现。这类问题浪费时间而且不容易察觉。6.4 团队协作中的性能沉淀最后多聊一句工程管理上的体会。内存受限算子的优化很多经验和硬件强相关换一块 GPU 或者换一个架构最优参数可能就得重调。所以我强烈建议团队里把每个优化过的算子连同它的剖析数据、硬件环境、最优参数一起沉淀下来。我现在自己写代码时都会在算子旁边留一个注释块记录当时的 ncu 关键指标和优化思路。半年后回头再看这些注释比任何文档都管用。从单个 kernel 到整条推理链路mem bound 算子加速的本质始终是把数据搬运的量减到最少把每次搬运的效率提到最高。理解了这一条你就不会被各种花哨的优化技巧带偏方向。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑