在大语言模型推理优化领域,内核融合粒度的选择始终是一个核心难题。过于保守的融合会导致频繁的内核启动开销,而激进的融合又可能因寄存器压力过大反而降低性能。传统编译器通常依赖成本模型在编译时做出这些决策,但预测模型与实际硬件行为之间的偏差往往导致次优的结果。Luminal 编译器采取了一条截然不同的技术路线:将融合决策推迟到运行时,通过搜索式优化在真实硬件上经验性地找出最优配置,从而在融合策略上实现了对传统静态分析的超越。
静态融合分析的困境与局限
现代深度学习编译器的融合策略通常遵循一个标准范式:首先对算子依赖图进行静态分析,识别出可以安全融合的操作序列;然后通过成本模型估算不同融合方案的预期性能;最后选择成本最低的方案进行代码生成。这种方法的优势在于编译时即可完成所有决策,不引入运行时开销,但它的致命弱点在于成本模型的准确性高度依赖于对硬件特性的精确建模。当面对复杂的注意力机制、动态形状变化或多 GPU 通信模式时,静态预测往往与实际执行效果产生显著偏差。
以 Transformer 中的注意力计算为例,一次完整的推理涉及多次矩阵乘法、softmax 操作、缩放运算和注意力权重的应用。传统编译器需要决定是将这些操作全部融合为单个内核,还是分成多个阶段执行。融合决策依赖于对内存带宽利用率、寄存器占用量、线程块 occupancy 等因素的准确估算,而这些因素之间存在复杂的相互作用。即便使用经验丰富的工程团队手工调优的成本模型,也难以覆盖所有可能的模型配置和硬件组合。更关键的是,模型架构的快速演进不断带来新的算子组合模式,使得成本模型的维护成本居高不下。
Luminal 的设计者意识到,试图在编译时预测所有可能的执行路径本质上是一项不可能完成的任务。不同 GPU 架构的内存层次结构差异显著,同一融合策略在不同批大小和序列长度下的表现可能截然不同。更重要的是,现代 GPU 的执行特性受到运行时资源调度状态的深刻影响,而这些状态在编译时根本无法获知。与其在编译时依赖可能出错的预测,不如将决策权交给运行时,让实际执行结果来指导选择。
搜索式优化的核心机制
Luminal 采用了基于搜索的编译器架构,其核心思想是构建一个包含多种可能融合策略的搜索空间,然后在实际硬件上执行这些变体,通过测量真实性能来选择最优配置。这种方法将编译器的角色从「做出最佳决策」转变为「探索决策空间并验证结果」,从根本上解决了静态分析的不确定性问题。编译过程不再是单一路径的优化,而变成了对多个候选配置的并行探索与实证比较。
具体而言,Luminal 的编译流程分为两个主要阶段。首先是搜索空间的构建阶段,编译器将计算图转换为多种可能的内核变体,每种变体代表一种不同的融合策略、 tiling 布局或内存访问模式。这个搜索空间的构造基于一组预定义的重写规则,这些规则编码了已知的优化模式,但具体选择哪些规则组合则留给搜索过程决定。Luminal 的设计哲学是让搜索来发现复杂的优化,而不是依赖手工编写的启发式规则。
第二个阶段是搜索执行阶段,系统在目标硬件上依次运行搜索空间中的候选变体,记录每次执行的延迟或吞吐量指标。由于搜索空间可能包含大量候选,直接穷举搜索是不现实的。Luminal 采用迭代式的搜索策略,每轮选择当前最优配置的邻域进行进一步探索,同时引入早停机制来跳过明显较慢的变体。搜索的迭代次数可以通过配置参数控制,牺牲一定的编译时间来换取更优的运行时性能。这种可调控的搜索深度使得用户可以根据实际需求在编译时间和运行性能之间做出权衡。
多级内核搜索与自动优化发现
Luminal 2.0 引入了多级内核搜索机制,进一步扩展了搜索式优化的能力边界。在之前的版本中,搜索主要作用于单层融合策略的选择,而多级搜索则将搜索范围扩展到了内核组织的多个层次。这种多层搜索能够发现跨层次的协同优化机会,而这些机会往往难以通过单层分析来识别。例如,某个 tiling 策略单独看可能不是最优的,但配合特定的融合决策后反而能产生最佳效果。
多级搜索的一个重要成果是能够自动发现类似 Flash Attention 的复杂优化。Flash Attention 的核心思想是通过分块计算和重计算来减少内存访问,但其具体实现涉及复杂的 tiling 参数选择和调度策略。传统编译器需要由专家将 Flash Attention 实现为专用算子,而 Luminal 的多级搜索能够从原始的计算图出发,自动搜索出具有类似效果的内核配置。这意味着 Luminal 不需要为每种新型注意力机制编写专用代码,而是通过搜索来「发明」等效的优化实现。
搜索式方法的另一个显著优势是其自适应能力。当模型在新的硬件平台上运行时,Luminal 只需要重新执行搜索过程即可获得针对该平台的优化配置,而不需要重新调优成本模型或编写平台特定的优化规则。这种硬件无关性大大简化了跨平台部署的复杂度。对于需要支持多种 GPU 型号的推理服务而言,这种特性尤为有价值。同一个模型定义可以通过搜索在不同硬件上获得各自的最优融合策略,无需维护多套平台特定的优化代码。
运行时决策的工程实践考量
将融合决策推迟到运行时并非没有代价,搜索过程本身需要消耗时间和计算资源。Luminal 通过多种工程手段来控制这一开销。搜索空间的结构化设计确保了搜索过程能够快速收敛到较优区域,而不是在全局空间中漫无目的地探索。配置缓存机制允许保存和复用之前的搜索结果,使得相同模型的重复编译可以跳过搜索阶段。此外,搜索过程本身也可以并行化,利用多个 CPU 核心同时评估不同的候选配置。
搜索式优化的另一个实践考量是搜索结果的可重现性。由于搜索过程涉及实际执行,硬件状态的微小变化可能导致结果的波动。Luminal 通过多次执行取平均值的策略来减少这种噪声的影响,同时提供种子参数来支持确定性的搜索行为。对于需要严格可重现的生产环境,这些特性确保了搜索结果在不同运行之间的一致性。
在实际部署场景中,搜索通常在模型初始化阶段完成一次,此后生成的优化配置会被缓存用于后续的推理请求。对于推理服务而言,较长的模型生命周期使得一次性的搜索开销可以分摊到大量的推理请求上,从而获得显著的平均性能提升。Luminal Cloud 和本地部署两种模式都支持这种缓存机制,用户可以根据自己的场景选择合适的部署策略。
搜索式内核融合代表了深度学习编译器设计的一种新范式。它放弃了在编译时做出确定性最优决策的企图,转而接受运行时探索的必要性,以实证结果取代预测模型。通过将优化的负担从分析能力转移到搜索能力,Luminal 在复杂多变的硬件环境中实现了更加稳健和自适应的性能优化。这种方法论的转变对于应对未来硬件的快速演进和模型架构的持续创新具有重要的参考价值。
资料来源:
- Luminal 官方仓库:https://github.com/luminal-ai/luminal
- Luminal 编译器文档:https://docs.luminalai.com/docs/compilers