
1. 从一行调用说起算子与内核的“相亲”现场你写下一行y relu(x)或者更复杂点y matmul(a, b) bias然后程序跑起来了结果也对。但中间到底发生了什么为什么同一个relu在 CPU 上跑得好好的换到 GPU 上也能跑为什么有些算子在小张量上飞快一到大张量就拉胯这些问题的答案都藏在“算子如何找到它的内核”这件事里。在 TensorPlay 这类深度学习框架中算子Operator是用户视角的计算单元比如加法、卷积、ReLU。而内核Kernel是真正在某个硬件后端上执行计算的代码比如一段 CUDA C、一段 C SIMD 指令、甚至是一段通过代码生成产出的机器码。一个算子可以有多个内核实现分别对应不同的设备、数据类型、内存布局和形状范围。分发器Dispatcher就是那个“红娘”它负责在运行时根据当前上下文为算子挑选出最合适的那个内核。这篇文章要聊的就是这场“相亲”的全流程从算子被调用到内核被选中再到最终执行中间经历了哪些关卡每一关的决策依据是什么以及为什么这么设计。如果你正在做算子开发、框架适配或者单纯好奇“我写的这行代码到底跑在哪”那这篇内容应该能给你一些实在的参考。2. 整体设计为什么需要分发器而不是硬编码2.1 硬编码的诱惑与陷阱最朴素的做法是在算子实现里直接写if device cuda: call_cuda_kernel() else: call_cpu_kernel()。小规模场景下这确实能跑。但很快你就会遇到几个绕不过去的问题。第一组合爆炸。设备有 CPU、GPU、NPU数据类型有 fp32、fp16、int8内存布局有 NCHW、NHWC再加上不同的形状范围小张量、大张量、特殊形状如果每个组合都硬编码一个分支代码会变成一团乱麻。第二扩展困难。每加一个新后端就要去改所有算子的实现这显然不可持续。第三无法动态优化。硬编码的分支是静态的没法根据运行时信息比如当前 GPU 的占用率、张量的实际形状做动态选择。所以成熟框架都会把“选内核”这件事从算子实现里抽出来交给一个独立的、可配置的、可扩展的分发器。2.2 分发器的核心职责分发器要解决的核心问题可以概括为一句话给定一个算子调用请求找到并执行最合适的内核。拆开来看它需要做几件事注册内核实现者把自己的内核“登记”到分发器附带描述信息比如支持哪些设备、哪些数据类型、哪些布局。匹配当算子被调用时分发器根据当前上下文设备、dtype、shape 等筛选出所有候选内核。排序与选择如果多个内核都匹配需要一个优先级规则来决定用哪个。通常越特化的内核优先级越高。回退如果没有完全匹配的内核要有回退策略比如自动把 fp16 转成 fp32 再算或者把不支持的布局转成支持的布局。缓存每次调用都重新匹配一遍太慢需要把匹配结果缓存起来下次直接命中。这套机制的设计目标很明确让算子实现者只关心计算逻辑让内核实现者只关心性能优化让分发器去处理所有“脏活累活”。2.3 一个关键取舍运行时分发 vs 编译时分发这里有一个重要的设计选择分发是在运行时做还是在编译时做运行时分发的优点是灵活可以根据实际输入动态选择内核支持 JIT 编译和自动调优。缺点是每次调用都有分发开销虽然可以通过缓存缓解但缓存查找本身也有成本。编译时分发的优点是零运行时开销内核在编译期就确定好了。缺点是灵活性差没法根据运行时信息做选择而且需要为每种可能的组合都生成代码编译时间会很长。TensorPlay 这类框架通常采用混合策略对于常见组合走编译时确定的快速路径对于特殊组合或需要 JIT 的场景走运行时分发。分发器本身也会做多级缓存把最常用的匹配结果放在最快能查到的位置。3. 核心细节内核注册、匹配与优先级3.1 内核注册把“简历”交到分发器手里内核注册是分发流程的起点。一个内核在注册时需要提供哪些信息这取决于分发器的匹配维度。常见的维度包括算子名称这个内核属于哪个算子比如relu、matmul。设备类型CPU、GPU、NPU 等。数据类型fp32、fp16、int8、bool 等。内存布局NCHW、NHWC、阻塞布局等。形状约束有些内核只支持特定形状范围比如“最后一维是 4 的倍数”。优先级当多个内核都匹配时谁优先。注册的代码通常长这样以伪代码示意REGISTER_KERNEL(relu) .device(DeviceType::CUDA) .dtype(DataType::FLOAT32) .layout(Layout::NCHW) .priority(100) .implement([](const Tensor x) { return cuda_relu_fp32_nchw(x); });这段注册代码的意思是为relu算子注册一个 CUDA 上的 fp32 NCHW 内核优先级 100。分发器在匹配时会把这些信息作为筛选条件。注意注册信息的粒度要适中。太粗会导致匹配不精确太细会导致注册表膨胀。实践中通常把“设备 dtype”作为主要维度布局和形状约束作为次要维度。3.2 匹配过程从调用请求到候选集当用户调用relu(x)时分发器拿到的是一个调用请求包含算子名称和输入张量的元信息。匹配过程大致分三步第一步按算子名称筛选。从注册表中取出所有属于relu的内核。这一步通常用哈希表实现O(1) 复杂度。第二步按硬性条件筛选。硬性条件包括设备类型、数据类型、布局。这些条件不满足内核直接出局。比如输入是 fp16那所有只支持 fp32 的内核就被过滤掉。第三步按软性条件排序。剩下的内核可能还有多个比如一个通用内核和一个特化内核都匹配。这时按优先级排序优先级高的先选。优先级通常由内核实现者指定特化程度越高的内核优先级越高。匹配的结果是一个有序的候选列表。分发器取第一个候选如果执行成功就返回如果失败比如形状不满足内核的内部约束就尝试下一个候选。3.3 优先级设计为什么特化内核优先优先级的设计有一个基本原则越特化的内核优先级越高。原因很简单。通用内核为了兼容各种情况通常做了很多妥协比如用最保守的循环顺序、不做向量化、不假设对齐。而特化内核针对特定形状、特定布局做了优化性能可能高出几倍甚至几十倍。举个例子。一个matmul算子可能有一个通用内核支持任意形状还有一个特化内核只支持 MNK64 的小矩阵。对于 64x64 的输入特化内核可能比通用内核快 5 倍。如果优先级设计反了先选通用内核那特化内核就白写了。但特化内核也有风险如果形状判断逻辑有 bug或者边界情况没处理好可能导致结果错误。所以分发器通常会在特化内核执行失败时自动回退到通用内核。这个回退机制是安全网不能省。3.4 回退策略当没有内核匹配时怎么办理想情况下每个算子调用都能找到完全匹配的内核。但现实中总有一些组合没有对应的内核实现。比如用户用了一个冷门 dtype或者一个非常规布局。这时分发器需要回退。常见的回退策略有几种类型提升把 fp16 转成 fp32用 fp32 内核算完再转回来。精度可能略有损失但至少能跑。布局转换把 NHWC 转成 NCHW用 NCHW 内核算完再转回来。多了一次内存拷贝但功能正确。设备回退GPU 上没有内核就回退到 CPU 执行。性能差很多但至少不报错。报错如果连回退都做不到就明确报错告诉用户哪个组合不支持。回退策略的选择需要权衡。类型提升和布局转换会引入额外开销设备回退可能慢几十倍。实践中通常优先做类型提升和布局转换设备回退作为最后手段并且会打警告日志。实操心得回退路径一定要有测试覆盖。很多框架的 bug 都出在回退路径上因为主路径天天跑回退路径很少跑一跑就出问题。4. 实操过程一个算子从调用到执行的完整链路4.1 调用入口用户代码如何触发分发用户代码里写的是y relu(x)。这行代码在框架内部会经历什么首先relu是一个算子对象它重载了__call__或者提供了一个forward方法。当用户调用时框架会收集输入张量的元信息设备、dtype、shape、layout、是否连续等。这些信息被打包成一个“调用上下文”。然后框架把这个上下文交给分发器。分发器根据算子名称和上下文走前面说的匹配流程找到内核并执行。这里有一个细节元信息的收集本身也有开销。如果每次调用都重新收集一遍对于小算子来说分发开销可能比计算本身还大。所以框架通常会把元信息缓存在张量对象里只在必要时更新。4.2 缓存机制如何避免重复匹配分发器的缓存通常分两级。第一级是“内核缓存”。键是算子名称设备dtype布局形状签名值是匹配到的内核指针。形状签名通常不是完整的 shape而是 shape 的某种摘要比如“所有维度都是 4 的倍数”或者“最后一维等于 64”。这样可以把形状相近的调用映射到同一个缓存项。第二级是“执行计划缓存”。对于复杂的算子匹配到内核后还需要生成执行计划比如分配临时内存、确定循环顺序、选择分块大小。这些计划也可以缓存避免每次重新生成。缓存的失效策略也很重要。当设备状态变化比如 GPU 内存不足、或者注册表更新比如动态加载了新内核时缓存需要部分或全部失效。实践中通常用版本号来标记缓存的有效性注册表更新时版本号加一缓存项版本号不匹配就重新匹配。4.3 代码生成当没有现成内核时有些算子组合太冷门或者形状太特殊没有现成的内核可用。这时框架可能会走代码生成路径在运行时生成一段代码编译成内核然后执行。代码生成的核心是模板 参数化。比如一个元素级算子可以生成一个循环循环体是算子的计算逻辑循环边界和向量化宽度作为参数。生成的代码可以是 C 源码编译成动态库也可以是 LLVM IRJIT 编译成机器码。代码生成的优点是灵活性极高可以覆盖任意形状和 dtype 组合。缺点是编译开销大第一次调用可能慢几百毫秒。所以代码生成的结果通常也会缓存第二次调用就直接用编译好的内核。注意代码生成路径的调试难度远高于普通内核。生成的代码往往没有符号信息出错时堆栈很难看。建议在代码生成时保留源码方便排查。4.4 执行与回写内核跑完之后内核执行完之后分发器还需要做一些收尾工作。如果走了回退路径比如类型提升需要把结果转换回用户期望的 dtype 和布局。如果内核是异步执行的比如 CUDA 内核需要处理同步和错误检查。如果内核执行失败需要触发回退尝试下一个候选内核。这些收尾工作看起来琐碎但缺一不可。很多“结果不对”的 bug都出在收尾环节比如忘了做类型转换或者异步错误没检查到。5. 常见问题与排查技巧实录5.1 内核匹配失败为什么我的算子找不到内核这是最常见的问题。表现是调用算子时直接报错提示“no kernel found for ...”。排查思路分几步第一步确认注册信息。检查内核是否真的注册了注册的设备、dtype、布局是否和调用上下文一致。常见错误是注册了 fp32 但调用的是 fp16或者注册了 NCHW 但输入是 NHWC。第二步确认注册时机。有些框架的注册是静态初始化如果内核所在的动态库没被加载注册就不会发生。检查动态库是否在调用前加载了。第三步确认形状约束。有些内核有形状约束比如“最后一维必须是 4 的倍数”。如果输入不满足内核会被过滤掉。检查输入形状是否满足所有约束。第四步看回退是否生效。如果框架配置了回退策略但回退也没找到内核就会报错。检查回退路径是否覆盖了当前组合。5.2 性能不达预期为什么选中的内核很慢有时候内核能找到但性能远低于预期。可能的原因有几个原因一选错了内核。优先级设计不合理导致通用内核被选中特化内核没被用上。检查优先级配置确认特化内核的优先级高于通用内核。原因二回退路径被触发。比如 fp16 内核没找到回退到 fp32多了一次类型转换。检查日志看是否走了回退。原因三缓存未命中。每次调用都重新匹配分发开销大。检查缓存是否生效形状签名是否过于精细导致缓存命中率低。原因四内核本身性能差。排除以上原因后可能是内核实现本身的问题。用 profiler 看内核执行时间确认瓶颈在内核内部还是分发环节。5.3 结果不正确分发对了但算错了结果不正确的原因通常更隐蔽。常见的有回退路径的转换逻辑有 bug。比如 fp16 转 fp32 时舍入方式不对或者布局转换时索引算错了。异步执行未同步。CUDA 内核是异步的如果没同步就读取结果可能读到未完成的数据。缓存了错误的执行计划。比如形状签名碰撞两个不同形状的调用映射到了同一个缓存项用了错误的执行计划。内核内部的边界处理有 bug。比如形状不是向量化宽度的整数倍时尾部处理错了。排查这类问题通常需要对比回退路径和主路径的结果逐步缩小范围。5.4 常见问题速查表问题现象可能原因排查方法解决思路报错“no kernel found”注册信息不匹配检查注册的设备/dtype/布局补充注册或调整回退策略性能远低于预期选错内核或走回退看日志确认选中内核调整优先级或补充特化内核结果不正确回退转换有 bug对比主路径和回退路径结果修复转换逻辑首次调用特别慢代码生成编译开销看是否走了 JIT 路径预热或缓存编译结果缓存命中率低形状签名过细统计缓存命中率粗化形状签名实操心得分发器的问题十有八九出在“以为匹配了但其实没匹配”或者“以为没匹配但其实匹配了”。加日志把每次匹配的候选列表和最终选择打出来能省很多排查时间。6. 内核开发的几个实战建议6.1 注册信息要精确但不要过度注册信息太粗匹配不精确可能选到不合适的内核。注册信息太细注册表膨胀匹配开销大。实践中建议把“设备 dtype”作为主要维度布局和形状约束作为次要维度。形状约束尽量用“范围”而不是“精确值”比如“最后一维是 4 的倍数”比“最后一维等于 64”更通用。6.2 优先级要拉开差距如果两个内核的优先级很接近分发器的选择可能不稳定今天选 A 明天选 B。建议优先级拉开差距特化内核给高优先级通用内核给低优先级中间留出足够的间隔。6.3 回退路径要测试回退路径是安全网但也是最容易出问题的地方。建议在 CI 里专门跑回退路径的测试确保回退逻辑正确。6.4 缓存要能失效缓存能提升性能但缓存失效逻辑如果不对会导致结果错误。建议用版本号标记缓存有效性注册表更新时版本号加一缓存项版本号不匹配就重新匹配。6.5 日志要能定位问题分发器的日志应该包含调用上下文、候选内核列表、最终选中的内核、是否走了回退。这些信息在排查问题时非常有用。7. 从分发器看框架设计的取舍分发器这个组件看起来只是“选个内核”但背后涉及很多设计取舍。灵活性与性能的取舍。运行时分发灵活但有开销编译时分发快但不灵活。混合策略是折中但实现复杂度高。通用性与特化的取舍。通用内核覆盖广但性能一般特化内核性能好但覆盖窄。分发器的优先级机制让两者共存但需要精心设计优先级。缓存与一致性的取舍。缓存提升性能但缓存失效逻辑复杂容易出 bug。版本号机制是常见解法但需要所有注册路径都正确更新版本号。回退与报错的取舍。回退让更多组合能跑但可能掩盖问题报错让问题暴露但用户体验差。实践中通常优先回退但打警告日志。这些取舍没有标准答案取决于框架的目标场景和用户群体。但理解这些取舍能帮助你在做算子开发和框架适配时做出更合理的决策。8. 一个具体的例子ReLU 算子的分发全流程为了把前面的内容串起来我们走一遍 ReLU 算子的完整分发流程。假设用户调用relu(x)x是一个 CUDA 上的 fp32 张量形状是[32, 64, 56, 56]布局是 NCHW内存连续。第一步收集上下文。框架收集到算子名称relu设备CUDAdtypefp32布局NCHW形状[32, 64, 56, 56]连续性true。第二步查缓存。用relu,CUDA,fp32,NCHW, 形状签名作为键查缓存。形状签名可能是“4D所有维度是 4 的倍数”。如果缓存命中直接拿到内核指针跳到第五步。第三步匹配候选。缓存未命中走匹配流程。从注册表中取出所有relu内核按设备、dtype、布局筛选剩下 CUDA fp32 NCHW 的内核。假设有两个一个通用内核优先级 50一个特化内核支持“最后一维是 4 的倍数”优先级 100。第四步排序选择。按优先级排序特化内核排前面。检查形状约束最后一维 56 是 4 的倍数满足。选中特化内核。第五步执行内核。调用特化内核传入输入张量和输出张量。内核内部做向量化计算每个线程处理 4 个元素。第六步收尾。内核执行成功检查异步错误返回结果。把匹配结果写入缓存下次直接命中。如果特化内核执行失败比如形状约束判断有 bug分发器会回退到通用内核重新执行。如果通用内核也失败报错。这个流程看起来简单但每一步都有细节。比如形状签名的设计、缓存的失效、回退的触发条件都需要仔细考虑。9. 代码生成在分发中的角色代码生成是分发器的一个重要补充。当没有现成内核时代码生成可以动态产出内核覆盖冷门组合。代码生成的典型流程是分发器发现没有匹配的内核触发代码生成代码生成器根据算子名称和上下文生成源码源码被编译成内核内核被注册到分发器分发器重新匹配这次能找到了。代码生成的关键是模板设计。模板要足够通用能覆盖各种形状和 dtype又要足够高效生成的代码性能不能太差。实践中通常用“模板 参数”的方式模板定义计算逻辑参数定义形状、dtype、向量化宽度等。代码生成的另一个问题是编译开销。第一次生成内核可能慢几百毫秒对于小算子来说不可接受。所以代码生成的结果通常缓存起来第二次调用直接用。注意代码生成路径的调试难度高。生成的代码往往没有符号信息出错时堆栈很难看。建议在代码生成时保留源码方便排查。10. 写在最后分发器的未来演进分发器这个组件从最早的硬编码分支到现在的多级缓存 代码生成已经演进了很多代。未来的方向可能是更智能的自动调优根据运行时性能数据动态调整内核优先级甚至自动生成针对特定形状的特化内核。但无论怎么演进核心问题不变在正确性、性能和灵活性之间找平衡。分发器就是那个平衡点它决定了算子调用最终跑在哪段代码上也决定了框架的整体性能上限。如果你在做算子开发建议花点时间理解分发器的匹配逻辑和缓存机制。很多时候性能问题不是内核写得不好而是分发器选错了内核。把分发器调对比优化内核本身见效更快。