前一阵子帮客户优化一个BERT模型的推理链路问题现象很典型GPU利用率看着不低内存也够但端到端延迟就是压不下去。我拿了profiler一看发现一张子图里连续排了八九个逐元素算子每个算子都单独把中间结果写回HBM、下一个算子再读回来一来一回纯属“搬砖”。这种情况下你再怎么调卷积、调GEMM都白搭因为瓶颈根本不在算力而在数据搬运。MetaInfer这套AI推理优化框架里的mHC算子融合优化就是专门解决这类问题的工具我基于海光K100AI的开发机把它完整调通了一遍今天把思路、代码、踩坑记录一起整理出来给你一份能直接参考的实战笔记。1. 从“算得快”到“跑得快”算子融合到底解决了什么1.1 一次推理链路里的性能账单先别看太多理论咱们直接算一笔账。一个典型的Transformer层抛开中间复杂的reshape、transpose光算子就有GEMM、BiasAdd、GELU、LayerNorm、Softmax、Residual Add这些。每个算子单独跑意味着GPU要把上一轮的输出整个写到显存里下一轮再整个读出来。中间张量越大这笔开销越肉疼。举个例子一个batch size 512、hidden size 4096的中间张量float32下就是8MB。一次“写8MB 读8MB”的访存流量是16MB。假设K100AI上实测有效带宽在1.2TB/s左右一个单纯逐元素算子跑完就要十几微秒。听起来不多那连续七八个算子就是一百多微秒而且这些算子大部分都是“读一个数、算一下、写一个数”计算量小到可以忽略时间全耗在来回搬运上。算力再强数据在显存里来回倒腾性能也上不去。推理优化的第一课就是先让数据少搬家再让计算跑得更快。1.2 融合的底层逻辑减少访存与启动开销算子融合的核心理念很朴素把多个算子的计算逻辑合并到一个kernel里中间结果不落地留在寄存器里直接喂给下一个计算步骤。如果按naive方式实现比如GELU的bias加法和tanh激活单独两个kernel中间结果必须先写回HBM再被第二个kernel读回来融合后一个线程处理一个元素bias加完紧接着算激活中间结果始终在寄存器里HBM流量直接砍半。这里必须科普一个概念Roofline模型通俗说就是“屋顶模型”。把算子的性能天花板看成两者取小——峰值算力决定的计算上限和峰值带宽决定的访存上限。对于逐元素算子比如AB每个元素只做一次浮点加法但读写要搬2个float。算下来算术强度不到0.5 FLOP/Byte而现代加速卡的带宽天花板显然比算力天花板低得多所以这类算子注定是访存密集型的。融合的价值就是从公式分子分母上同时做文章减少分母访存量同时减少kernel数量启动开销。kernel启动开销也别小看。CPU提交一个kernel到GPU执行中间包含命令下发、同步、队列调度实测在K100AI上单个kernel的启动收尾开销通常在5到15微秒。一个推理图几十个算子光启动开销就可能占掉几百微秒。融合把kernel数量降下来这部分时间也实打实省掉了。1.3 什么时候不要融合先分清访存密集与计算密集融合不是万能的我见过不少团队吃了“无脑融合”的亏。核心要分清算子到底是访存密集还是计算密集。卷积、矩阵乘这类算子算术强度高时间主要花在计算上。对它们来说融合的收益主要来自“尾巴”——也就是GEMM后面的BiasAdd、激活函数这类逐元素操作。通常做法是把激活和bias加进去做进GEMM的epilogue阶段而不是把两个大计算算子硬揉在一起。硬融两个GEMM反而会破坏原有tiling策略让寄存器压力飙升最后得不偿失。而LayerNorm、Softmax这类带规约的算子融合的收益就非常大。它们既有访存密集的部分归一化又有通信密集的规约过程多一个kernel就多一轮全量读加全量写。这部分我会在第三章节细讲。一句话总结判断标准如果算子序列里大部分是逐元素操作融合是必修课如果全是计算密集大算子先老实调底层kernel别指望融合救场。2. MetaInfer与mHC这套优化模块的设计思路2.1 mHC是什么混合计算场景的算子融合模块MetaInfer是一款面向推理场景的引擎工作方式和常见的高性能推理框架类似加载训练好的模型对计算图做静态优化再生成面向特定硬件调优过的执行计划。在整个优化栈里图优化是承上启下的关键层而mHCmeta Hybrid Compute的缩写你可以理解为MetaInfer内部负责混合计算融合的模块就是图优化里最核心的一块拼图。mHC处理的“混合计算”指的是一个计算图里既有访存密集的逐元素算子、又有规约算子、还有计算密集的GEMM/卷积算子形态混合、访存特征各异。mHC的设计目标就是根据这些算子的数据流特征把适合合并的算子分组融合成新的“超级算子”让执行计划更贴近硬件的运行方式。这套模块的定位类似推理引擎里的“编译器优化器”。训练框架导出的原始模型就像一个人手写、充满冗余的代码mHC做的工作就是把它编译成高效、内联、少掉中间变量的机器码。它不改变模型的数学意义只改变计算的执行方式。2.2 为什么按“数据流形态”分类而不是全图一把梭第一次接触mHC的人通常会问既然融合这么好为什么不把整张图都合成一个kernel答案很简单算力硬件不允许。一个kernel在GPU上执行时所有线程执行的是同一段代码。如果一张图里既有GEMM又有Softmax还有LayerNorm它们的并行模式和访存模式完全不同硬塞进一个kernel要么让不同线程做不同分支导致warp divergence线程发散要么寄存器被各种中间数据塞爆要么根本无法设计出合理的同步策略。强行全图融合的结果大概率比不融合还慢。mHC的处理思路是“分类融合”先把图里的算子按数据流形态分组再对每组采用专门的融合模板。逐元素链BiasAdd、激活、残差Add、乘常数这类算子数据流是一对一天然适合直接内联。规约类LayerNorm、Softmax、ReduceMean这类算子内部有跨线程通信需要专门的block归约策略。计算算子epilogueGEMM/卷积归一化激活残差重点是把尾部的逐元素计算融进主计算kernel的收尾阶段。布局转换类reshape、transpose、pad、slice这类算子有时可以直接消除让上下游共享内存视图省掉一整轮拷贝。这种“先分类、再模板化”的思路非常实用。因为每一类融合都有成熟的实现套路可复用、可测试、可单独调优比一次性写一个巨型融合器要稳妥得多。我在K100AI上调通mHC之后最大的体会就是设计上少一点浪漫多一点分类学工程反而顺得多。2.3 K100AI硬件特性如何反向影响融合设计融合策略不是凭空设计的硬件特性决定了很多细节。海光K100AI是一类面向AI计算的加速卡生态上走的是ROCm/HIP路线这意味着编写kernel用的语言是HIP-C性能分析和基础库用的是ROCm工具链。和常见GPU的习惯类似又有很多细节差异下面几点直接影响融合实现。第一是wavefront概念。HIP编程里一个wavefront相当于硬件一次调度的线程束通常是64个线程部分指令支持wave32。设计block大小时要尽量让线程数是wavefront大小的整数倍比如256就是4个wavefront避免出现调度浪费。第二是共享内存LDS容量。做规约融合时比如LayerNorm需要把部分数据暂存在LDS里做跨线程求和。K100AI每个CU的LDS容量是有限的所以block size不能随意加大否则会顶爆LDS。我一般的做法是block取256或512规约分两段做先在每个线程内做局部累加再用shuffle或LDS做跨线程归约最后只剩几十个线程参与最终求和。第三是访存对齐和向量化。K100AI的HBM带宽很高但带宽只有在连续访存、宽度足够时才能发挥出来。融合kernel里尽量用float2/float4做向量化读写一次搬128位数据访问效率提升明显。我在实测里发现同样一个融合算子float4版本比float1版本带宽利用率能高20%到30%。第四是内核并行度。推理场景下一个小算子的数据量可能只有几十万元素如果融合后kernel的grid规模太小没法铺满所有计算单元性能反而会掉。这时候就要考虑把多个batch维度拆开或者主动压低block size来抬高grid规模保证每个CU都忙起来。3. mHC三类核心融合模式与原理拆解3.1 逐元素融合最常见也最容易吃满带宽逐元素融合是mHC最基础的模式处理的是那些“读一个、算一下、写一个”的算子。典型例子是bias加法和激活函数比如Transformer里的LinearBiasGELU的尾段ResNet里的ConvBNReLU。实现思路非常直接每个线程处理一段连续数据把n个算子的计算依次在寄存器里做完最后写一次结果。代码上要注意几点用__restrict__修饰指针让编译器放心做乱序优化用向量化访存尽量让grid和block的配置覆盖到所有元素。一个不那么显眼但很关键的细节是数值函数的选择。HIP里有很多数学函数tanh是精确版本__tanhf是快速版本tanhf是两者之间。逐元素融合里如果用快速版本速度确实快但精度会漂移。做推理部署时如果模型是fp32且对精度敏感建议先用精确函数做对照再决定要不要换快速版本。我在GPT类模型上试过部分求值函数换成快速版后输出差异在1e-4量级多数任务没问题但一旦下游有敏感判断逻辑就得谨慎。3.2 规约类融合LayerNorm和Softmax的经典优化规约类算子和逐元素算子最大的区别是它需要“全组数据都算完才知道结果”。LayerNorm要先算整行均值、方差再做归一化Softmax要先找最大值、算指数和再归一化。如果用三个独立kernel实现中间结果至少要写回显存两轮。mHC的做法是把这个过程压缩到一个kernel里。拿LayerNorm举例一次推理里的张量形状通常是[B, H]对最后一维做规约。第一遍遍历该行数据用Welford方法或朴素两遍法算出mean和rstd第二遍遍历同一行直接做归一化。因为数据分块在block里中间两个block之间还能共享LDS实现方式就是让一个block处理一行或几行数据把行数据加载到LDS后做并行归约。这里有一个常见的性能坑如果直接让每个block处理一行而hidden size又比较大比如4096一个block算一行时每个线程要遍历多个元素这时就要注意LDS用量和访存合并。我的处理是block size取256每行数据分成多段并行归约最后在LDS里做跨线程合并。优化后效果非常明显原来三个kernel变成零个——不是变成一个原来三次全量读写变成一次读加一次写访存流量直接降到三分之一。Softmax是一样的套路不同的只是先做max再用exp和为分母。需要注意的是为了数值稳定Safe Softmax必须保留“先减去max”这一步融合时千万不要为了省一次遍历把这个保护去掉否则大数据量下很容易出现NaN。3.3 计算算子epilogue融合把尾巴融进主kernelGEMM和卷积这类计算密集算子的融合用的是“不同姿势”。不是在GEMM之后再跑一个激活kernel而是在GEMM的计算内核收尾阶段把bias加法和激活函数直接塞进去。原理不复杂GEMM本来就是分块计算的每个输出元素在寄存器里累积出结果之后需要写回显存。在写回之前寄存器里的数值顺手做一次bias加、做一次激活、加一次残差一分钱额外访存都不用花。相比单独跑一个激活kernel等于免费白赚了这一段性能。具体实现上如果直接用rocBLAS/hipBLAS的GEMM接口有些版本支持epilogue参数可以传一个回调函数或模板参数把激活函数塞进GEMM的收尾阶段。如果底层库不支持也可以在自定义HIP kernel里实现一个简单的tiled GEMM把epilogue逻辑写在累加循环结束之后。实际操作里融合卷积BatchNorm的收益更可观。BatchNorm在推理阶段可以折叠成“缩放平移”的线性操作它和前面的卷积融合后等于把卷积的结果先缩放再平移再激活。这一步不只能省掉一个kernel还省掉了BN层在归一化时的均值方差计算节省效果在CNN类模型上尤其明显。这就是mHC里常见的“先常量折叠再算子融合”组合拳。3.4 收益估算用Roofline模型算清楚省了多少流量理论讲再多不如自己算一遍。假设有一个中间张量大小是512×4096float32总数据量8MB。现在有七个逐元素算子串在一起每个都要完整读一遍、写一遍。不融合时每个算子访存是“读8MB写8MB16MB”七个算子就是112MB。按K100AI实测有效带宽1.2TB/s算光访存时间就要大约93微秒再加上七个kernel的启动开销按每个10微秒算又多了70微秒。合计160微秒以上。融合成一个kernel后中间数据不出寄存器整个过程中只发生一次完整读和一次完整写访存量16MB约13微秒加上1个kernel启动10微秒总共不超过25微秒。这还没算kernel之间同步等待的时间损失。同样的算子序列融合后快了将近6倍。方案访存量访存耗时kernel启动总计约不融合7个kernel112MB93μs70μs163μs融合1个kernel16MB13μs10μs23μs加速比7倍7倍7个→1个约7倍这种估算在动手前特别有用。你先拿Roofline模型把瓶颈定性再估算融合能砍掉多少访存流量就知道值不值得花精力去写融合kernel。很多场景下还没写代码收益就已经算明白了。4. 基于海光K100AI的完整实操记录4.1 开发环境与性能基线先交代一下我手上这套环境K100AI加速卡ROCm驱动为K100AI配套版本开发包包含HIP运行时、rocBLAS/hipBLAS以及MetaInfer的release包。语言方面用HIP-C写kernelPython侧只是做模型分析和数据比对。拿到机器第一件事先确认硬件和驱动状态rocm-smi # 查看DCU设备与温度、利用率 rocminfo # 查看设备架构信息 hipcc --version # 确认HIP编译器版本然后是建立性能基线。不要一上来就做融合先跑一遍原始模型把每个算子的耗时、访存流量、kernel数量记下来。我用的是MetaInfer自带的profiler也可以配合ROCm的profiling工具抓kernel级时间线。注意推理测试一定要“预热”先跑几十轮让显存分配、缓存状态稳定下来再取中位数或均值否则第一轮和第二轮的数据差异会误导你。4.2 写一个融合的BiasGELU算子先从一个最简单的融合开始把一个bias加法和GELU激活合并成一个kernel。这个算子小逻辑简单但能完整展示融合的思路。GELU有很多近似形式这里用tanh近似版__global__ void FusedBiasGeluKernel(const float* __restrict__ x, const float* __restrict__ bias, float* __restrict__ y, int rows, int cols) { int idx blockIdx.x * blockDim.x threadIdx.x; int total rows * cols; if (idx total) return; int c idx % cols; float v x[idx] bias[c]; v 0.5f * v * (1.0f tanhf(0.79788456f * v * (1.0f 0.044715f * v * v))); y[idx] v; }这是标量版本逻辑清楚适合理解。实际部署时我强烈建议改成向量化版本把连续四个元素打包成float4一次读、一次算、一次写利用128位访存指令。唯一的前提是cols是4的倍数且内存对齐Transformer类模型的hidden size基本都是4的倍数这个条件很容易满足。向量化还有一种写法是用grid-stride loop让每个线程处理多段数据。好处是社会层数grid可以固定不用跟着数据量变化频繁调整对推理场景的启动稳定性有帮助。kernel的launch配置我通常这样算block取256线程需要的wavefront数就是4个grid取ceil(total / (256 * 4))这里的4是vector宽度。算出来之后先跑一遍再微调vector宽度观察带宽利用率。4.3 在MetaInfer里注册并使用融合算子单独有kernel还不够要在MetaInfer里让引擎自动替换原来的算子序列。这一步mHC已经封装好了不需要手写匹配逻辑你只需要打开对应的融合选项。下面是一个风格化的示例接口不同版本的MetaInfer API命名会有差异以你自己的SDK文档为准meta_infer::Graph graph meta_infer::Load(/data/models/bert_base.onnx); meta_infer::FusionOptions opts; opts.enable true; opts.fusion_pattern meta_infer::FusionPattern::kElementwiseChain; opts.fusion_pattern meta_infer::FusionPattern::kReduceFusion; opts.fusion_pattern meta_infer::FusionPattern::kConvEpilogueFusion; auto engine meta_infer::Engine::Build(graph, opts, meta_infer::Device::K100AI); engine-Compile(); engine-SaveEngine(bert_base_mhc.mif);关键点在于融合动作是发生在编译期的引擎会先对计算图做模式匹配找到满足条件的算子子图再替换为对应的融合kernel。所以如果模型是动态shape某些融合模板可能因为无法确定静态边界而被跳过这时候要么固定shape要么改用支持动态shape的融合模板。如果项目里走的是PyTorch也可以考虑先用torch.compile做一级图捕获导出为静态图后再交给MetaInfer做融合和部署。两套工具的侧重点不同torch负责拿到干净的静态计算图MetaInfer负责把图变成高效执行计划配合起来效果最好。4.4 用Profiler验证结果与调参要点融合做完了不profiler一切都是玄学。我的验证流程固定四步先抓kernel时间线、再对比访存流量、随后做数值比对、最后看整体延迟。kernel时间线这一步很关键。融合前时间线上能看到一串短kernel每个都有明显的启动和结束融合后一串短kernel变成一个长kernelkernel数量大幅下降。如果发现融合后的kernel运行时间长于原来所有kernel之和那说明kernel自身有问题多半是寄存器溢出或者block配置不合理。数值比对也不能省。拿原始模型输出当reference融合模型的输出做逐元素对比统计最大绝对误差和相对误差。fp32模型一般要求最大误差在1e-5量级fp16/BF16模型要放宽到1e-2量级具体以业务需求为准。出现异常先别急着改kernel先把快速数学函数替换成精确版再比一次。最后看端到端延迟。注意一次推理里可能同时存在多个融合算子profiler输出的single kernel时间不能简单相加要看真实执行时间线因为GPU上kernel之间本来就有重叠。端到端延迟才是用户真正关心的指标。4.5 实测数据对比融合前与融合后的差距我在一个类BERT结构上做的实测数据如下给你一个直观印象。模型单层包含GeMM、BiasAdd、GELU、LayerNorm、Residual Add序列长度128、hidden 768、batch 32共对比了那一段子图不含中间attention的矩阵乘。指标未开启mHC开启mHC变化kernel数量187-61%中间张量访存总量约210MB约90MB-57%子图平均耗时0.86ms0.51ms-40%端到端单次推理4.2ms3.5ms-17%能看到子图收益比端到端收益大因为整张图里还有attention内部的矩阵乘、各种非融合算子它们不会因为这部分融合而加速。这也说明一个现实融合优化做得再好也只是推理优化的其中一块要和算子库调优、显存复用、并发调度配合才能把整张图吃干榨净。5. 常见问题与排查技巧实录5.1 融合后精度不对先按这个顺序查融合后精度漂移大概率是下面几个原因之一按顺序排查效率最高。第一快速数学函数。tanhf、__expf这类函数精度低于对应精确版。把kernel里所有快速版本换回精确版本再比对一次。如果误差消失问题就在这里。要不要保持快速版本取决于你的精度预算。第二数值运算顺序。融合后浮点运算顺序变了比如原来的(ab)c变成a(bc)对fp32来说结果是会有微小差异的。多数场景在误差允许范围内不用管如果业务对误差极其敏感就只能保留原来的拆分方式。第三量化尺度因子。在量化模型里反量化dequant常出现在卷积或Linear之后。mHC做融合时如果没处理好scale和zero_point的先后顺序比如先融合后反量化数值就会错。检查时重点看量化算子在融合后的位置是否合理。第四规约顺序。LayerNorm或Softmax里如果改变了规约顺序比如分块归约后合并在大数据量下可能产生较大误差。Welford算法本身稳定但如果你手动分块后有精度问题试试把分块粒度调小或改用Kahan补偿。5.2 融合后反而变慢别急先怀疑这三个位置融合理论上应该变快但实测变慢的情况我见多了主要集中在三个位置。第一寄存器溢出。融合的算子越多寄存器里要同时存的数据就越多。编译器一旦发现寄存器放不下就会把数据spill到本地内存本地内存本质是显存速度比寄存器慢几个数量级。这时性能不升反降非常正常。解决办法是少融几个算子或者用__launch_bounds__限制最大线程数和寄存器数强制编译器做取舍。第二block配置不对。有些逐元素算子数据量小如果你仍按大vector宽度去融grid规模小到没法铺满所有CU性能自然差。对这种算子减少vector宽度、缩小block size让grid规模大一点往往立竿见影。第三融合把一个访存密集型算子和一个计算密集型算子捆在一起。两者的最优并行度和寄存器需求完全不同硬放一个kernel会互相拖累。正确的姿势是用异步方式或者干脆分开。5.3 LDS溢出和内核调度问题怎么处理规约类融合最容易踩LDS的坑。LayerNorm每个block要缓存一整行或半行数据到LDS如果block size是256每个float占4字节光数据缓冲就要1KB再加归约用的临时数组几十个block同时运行就会把LDS顶爆表现为kernel编译失败或者运行时错误。我的处理套路是先做两段归约每个线程先把自己负责的元素在寄存器里累加再用LDS做block内归约最后只需要一个很小的跨线程合并段。LDS只存部分和不存原始数据用量能降一个数量级。调度问题则更多体现在“融合kernel太少、别的大算子太多”的场景。K100AI上执行kernel是流水式的多个kernel在队列里排队发射。如果某个融合kernel占用的资源过大后续kernel可能在等待资源释放造成流水气泡。这时调整融合粒度拆成两个中等kernel整体流水反而更顺。记住一个原则融合的目标是吞吐最大化不是kernel数量最小化。5.4 避坑速查表直接复制到团队文档里问题现象常见原因排查/解决方法融合后误差超预算快速数学函数换精确版重新比对融合后误差超预算量化因子顺序错误检查dequant在融合后的位置融合后变慢寄存器溢出减融合算子数或加__launch_bounds__融合后变慢grid规模太小降vector宽度、降block sizekernel编译失败LDS溢出两段归约LDS只存部分和动态shape下融合失效静态图边界不确定固定shape或换动态模板profiling里kernel时间变长访存未向量化改用float4等宽位访存端到端收益不明显其他非融合算子拖后腿结合算子库调优、显存复用一起做写在后面几个我自己摸索出来的习惯最后分享三个实战习惯。第一个习惯是“先算后写”动手做融合前先按第三章的估算方式算出理论收益如果预期加速比低于1.5倍果断放弃把时间留给更值得的算子。第二个习惯是“保留开关”不要一股脑把全图融合打开我倾向于在代码里保留一套逐类开关Ai推理模型的业务形态经常变有些任务适合逐元素融合有些适合规约融合灵活比激进更重要。第三个习惯是在K100AI上调融合kernel时多对比几轮再下结论。GPU性能波动比想象中大温度、频率、其他任务干扰都会影响测量结果。我通常跑满100轮预热再取中间50轮的中位数作为最终数据一次测量就拍板很容易被波动带偏。其实算子融合这个技术本身不神秘核心就是把访存开销和启动开销压到最低。但真正把它做扎实需要你对硬件、对数据流、对数值稳定都有感觉。希望这篇笔记能帮你少走几段我走过的弯路下次在K100AI上调推理模型时记得先看看你的算子里有多少“搬砖”的把它们先融掉。