研究摘要：
深度学习模型的规模不断扩大，对 GPU 的计算能力和内存带宽提出了更高的要求。算子融合（Operator Fusion）作为一项关键优化技术，能够通过减少内存访问和提升数据局部性来缓解 GPU 的性能瓶颈。然而，现有的方法在应对计算与内存混合受限算子融合时，存在搜索空间不完整、内存访问冗余及调优耗时过长等问题。为此，论文提出了一种名为 MCFuser 的创新框架，通过高级平铺表达、DAG 分析和启发式搜索方法，为 MBCI 算子生成高性能融合内核。在多项基准测试中，MCFuser 展现了显著的性能提升与调优时间减少，成为深度学习编译器领域的创新解决方案。
研究背景
随着深度学习模型规模的不断扩大，计算和内存需求呈指数级增长。在这一背景下，如何提高深度学习模型的执行效率成为研究的重点问题。算子融合（Operator Fusion）是一种关键优化策略，通过减少内存访问、提升数据局部性，有效缓解 GPU 内存带宽瓶颈。然而，传统算子融合主要针对计算密集型或内存密集型的单一算子，对于计算和内存混合受限（MBCI，Memory-Bound Compute-Intensive）的算子融合缺乏有效的解决方案。而这些场景在自注意力模块等深度学习模型中却尤为重要。
表 1 现有工作对比
研究现状
目前，多数研究在算子融合领域已取得显著进展。例如，AStitch 专注于内存密集型算子的融合，而 DNNFusion 结合了单个计算密集型算子与周边内存密集型算子的优化。然而，这些方法普遍忽略了计算密集型算子之间的融合潜力。对于 MBCI 算子，虽然一些方法（如 Ansor 和 Chimera）尝试优化，但受限于搜索空间和冗余内存访问问题，未能充分发挥硬件潜能。
此外，由于张量维度动态变化的特性，许多传统计算密集型算子在特定条件下会变为内存受限型。这种动态特性进一步凸显了高效融合策略的重要性，但当前方法在这方面的探索仍不充分。
图 1 Tiling 表达式与搜索空间
研究方法
论文提出了一种名为 MCFuser 的创新框架，专注于高效融合 MBCI 运算符链，并有效解决了传统方法在搜索空间不完整、内存访问冗余和调优效率低下等方面的挑战。MCFuser 通过引入高级平铺（tiling）表达式，对可能的融合结构进行了全面定义，既涵盖了深度平铺也包括扁平平铺模式，从而确保了搜索空间的完整性，避免了遗漏潜在的优化机会。此外，该框架利用有向无环图（DAG）分析技术，优化了内存访问路径，显著减少了冗余的内存传输，将不必要的数据移动从 GPU 的全局内存转移到共享内存中，有效缓解了内存带宽的瓶颈。
为了进一步提升搜索效率，MCFuser 设计了一套完善的搜索空间剪枝策略，提出了如去除冗余候选、避免多余内存填充、限制共享内存使用等规则。这些规则不仅大幅减少了无效候选项，还显著缩小了搜索范围，使得优化过程更加高效。此外，框架还结合了一个性能分析模型，该模型通过深入分析内存访问和计算延迟来快速评估每个候选方案的性能表现，避免了传统方法中依赖硬件运行数据进行调优的高昂成本。与此同时，MCFuser 将这一分析模型与启发式搜索算法相结合，从庞大的搜索空间中自动定位最优融合配置，在保证性能提升的同时显著缩短了调优时间。
综上所述，MCFuser 通过全面的搜索空间构建、内存访问优化、搜索空间剪枝和性能模型的高效结合，不仅实现了 MBCI 运算符的高效融合，还在性能和资源利用效率方面展现出显著的优势，为深度学习编译器领域提供了一个极具参考价值的解决方案。
图 2 端到端测试
研究成果
MCFuser 框架在 NVIDIA A100 和 RTX 3080 GPU 上进行了全面验证，展现了卓在实际测试中，MCFuser 在多种场景下表现优异。对于 GEMM 链和自注意力模块的子图融合，MCFuser 在 A100 和 RTX3080 上的性能均显著优于 PyTorch 和其他主流框架。同时，如图5所示，在端到端网络（如 BERT-Small 和 BERT-Base）的测试中，MCFuser 能够轻松集成到现有的 TVM 和 Triton 编译器栈中，实现对 MBCI 子图的高效优化。MCFuser 的创新在于突破了传统框架将计算密集型算子视为融合边界的限制，通过针对 MBCI 算子的深入优化，实现了数据局部性和计算性能的双重提升。这种方法不仅显著缩短了调优时间，还为深度学习模型的高效训练提供了一种通用且灵活的解决方案。
相关研究成果已收录于 CCF-A 类顶级会议 SC 24（The International Conference for High Performance Computing, Networking, Storage, and Analysis），论文题为 MCFuser: High-Performance and Rapid Fusion of Memory-Bound Compute-Intensive Operators。