nncase：基于e-graph的端到端LLM编译器，突破异构存储架构性能瓶颈

关键词：LLM 编译、 e-graph、异构存储架构、统一分布式编译、自动优化、端到端编译框架

本文转载自知乎账号：郑启航[1]
原文链接：https://zhuanlan.zhihu.com/p/1989088940733510928

nncase: An End-to-End Compiler for Efficient LLM Deployment on Heterogeneous Storage Architectures
论文地址：https://arxiv.org/pdf/2512.21571
项目地址：https://github.com/kendryte/nncase

很高兴为大家介绍我们的工作：nncase，一个支持 LLM 多存储架构端到端部署的 AI 编译器。

具体而言，nncase 是一个支持动态 shape、自动向量化布局、自动分布式切分、自动多内存层级 Fusion Tiling 的 AI 编译器，可以将 LLM 模型端到端编译为 Single Kernel 执行，达到超越推理框架的性能。

图1：nncase 架构概览，展示了 nncase 编译器的端到端编译流程，通过五个核心阶段连接高层模型定义与底层硬件原语，实现异构存储架构下 LLM 的高效部署：先通过基于 e-graph 的优化阶段（含自动向量化、自动分布式策略搜索）解决布局冲突与阶段排序问题；再经自动调度阶段（结合 MCTS 与 MINLP 优化循环结构和参数）提升片上缓存复用率；最后通过代码生成阶段，依托缓冲调度与 NTT 库实例化高效内核，最终实现 “一次编译，多端适配”，为异构硬件下 LLM 部署提供统一优化路径。

我们设计 nncase 的动机在于：

作为一个工具链团队，需要向后向前兼容各种硬件。
当前模型与硬件的规模化已成趋势，编译器自然也需要具备规模化（scaling）的能力，以支持不同场景、不同规格硬件的编译优化。我们相信不远的将来依旧会出现组合爆炸的问题，这是人力难以应对的。
考虑到多芯片/多晶粒、UMA/NUMA 等不同的硬件架构组合，需要以计算资源为粒度，从分布式视角将计算图进行拆分，并针对不同层级的存储与互联做相应的编译优化。

图9：单核心场景（1T）下不同框架的LLM token 生成吞吐量对比。其中 llama.cpp 代表手动优化基准，其余框架代表基于编译器的实现方案。基于 AMD Ryzen 9 5900X 单线程实验，清晰呈现性能层级：llama.cpp 因手动优化内核领先，nncase 紧随其后，远超 Intel IPEX 与 MLC LLM。具体来看，Qwen3-0.6B-F32 配置下，nncase 达 8.7 tokens/s，较 Intel IPEX 高 15%；F16 精度时提升至 13.87 tokens/s；Qwen3-1.7B-F16 仍保持 5.09 tokens/s，这验证了其单核心下自动向量化与调度模块的有效性，虽略逊手动优化，但在编译器方案中表现最优。

一、背景与动机

1.1 数据布局优化

目前通用的计算架构都会存在张量/向量级别的计算单元：

传统做法会采用在单个 kernel 内部进行向量化/张量化，但这对单个算子的调度要求较高，类似 BLIS 需要进行在线打包（online packing）。并且也错失了多个 kernel 之间存在数据布局复用的情况。
也有一些类似 TPP MLIR 的方案，尝试以张量为粒度进行数据布局重写，但它并没有充分考虑矩阵单元和向量单元之间数据布局转换的开销。

1.2 分布式策略搜索

分布式策略搜索是充分释放芯片并行算力的关键，但目前大部分工作都在高层框架中完成，没有在一个编译器中结合上下游的数据布局与算子调度来同步完成。

1.3 kernel 调度优化

调度优化直接决定了算力的有效利用率，但不同硬件平台在指令集、片上缓存结构、计算单元架构等方面的固有差异，使得“单一调度策略适配全硬件平台”的目标难以达成。

目前相关工作基本分为两类：

基于学习的调度方法（如 Looper、Ansor、Meta Scheduler），这类方法通常需要真实硬件进行学习，且学习时间过长，不适合快速部署场景。
基于分析建模的方法（如 Tuna、Chimera），这类方法目前难以支持复杂融合算子的求解，且部分工作以标量为最小优化粒度，难以实现寄存器级别的精细优化，因此其优化效果距离手工优化水平仍有较大差距。

1.4 核心观察与解决方案

数据布局优化需全局权衡转换开销，但传统方法易受“阶段排序”问题影响，错失更优方案。为此，nncase 实现了基于 e-graph 的模式重写系统，将所有可行的数据布局方案存入 e-graph，并基于 Roofline 模型评估开销，最终提取出全局最优的数据布局方案。
分布式策略本质是路径提取问题。nncase 通过自定义 e-graph 的构建过程，将分布式策略搜索纳入图优化范畴，复用计算/通信开销模型与提取器来寻找最优解。同时，结合内存约束避免内存溢出（OOM），并针对硬件架构进行内存分配优化。例如，在统一内存架构（UMA）下，可实现逻辑广播而物理上仅存储一份数据，或逻辑切分而物理紧凑排列。
调度优化的核心挑战是粒度过细与搜索耗时过长。nncase 提出了 nncase 张量模板库，通过手写的微内核向编译器隐藏寄存器级优化的复杂性。并以微内核形成的计算块作为基本调度单元，借鉴 Welder/TileFlow 的思想，将搜索空间划分为结构维度与参数维度。
- 结构维度：反映计算块的组织结构，对应循环融合、循环重排等优化。
- 参数维度：反映计算块自身的参数，如块大小、缓冲区位置。
  随后，系统采用蒙特卡洛树搜索求解结构层优化问题，并借助解析建模完成参数层优化。该方法在保障寄存器级优化性能的同时，有效规避了传统方法搜索时间过长的缺陷。

二、nncase 编译器架构

2.1 nncase 架构概览

nncase 采用基于 e-graph 的重写系统来解决各类图优化问题。

编译器输入计算图与硬件配置，首先在图优化阶段进行自动向量化与自动分布式策略搜索。随后，将支持仿射表示的运算转换为携带仿射映射信息的层次化计算块图。

在自动调度阶段，首先使用蒙特卡洛树搜索自由组合各种计算块图的结构及其仿射映射属性。接着，针对每个具体的计算块图，采用混合整数非线性规划建模，求解其最优的计算块大小、缓冲区放置方案与总耗时，从而动态地将融合、分块与缓冲区规划结合。

最后，将最优的计算块图及其参数生成 C++ 代码，通过 C++ 实例化手写的 NTT 库，并最终借助传统编译器生成可执行文件。

nncase 编译器大量使用求解器作为核心优化手段：
* 例如，e-graph 的提取器本质上是一个加权部分最大可满足性问题求解器。
* 自动调度阶段采用混合整数非线性规划求解器。
* 缓冲区调度则被转化为装箱问题，并采用约束规划可满足性求解器进行求解。

2.2 等式饱和

本节阐述为何选择 e-graph 作为重写系统的基础。

图2：传统项重写与基于 e-graph 的重写对比。图中：(a) 为待优化的原始计算图，目标是通过表1中的转置优化规则消除冗余操作。传统重写（b、c）受阶段排序影响显著：若贪心地先应用右侧二元转置合并规则，将得到次优结果（c，残留冗余转置）；而必须按特定顺序（先应用左侧二元转置合并规则，再应用一元转置合并规则）才能得到全消除冗余的最优结果（b）。基于 e-graph 的重写（d-f）则通过等式饱和过程保留所有等价节点（d），结合 Roofline 成本模型与加权部分最大可满足性问题求解，最终提取出全局最优子图（f），彻底规避了传统方法的局部最优陷阱。

假设我们有图 (a) 所示的计算图以及表1所列的优化规则。可以发现，规则的应用顺序会直接影响最终优化结果。如图 (b) 所示，若先向下移动左侧的转置操作，会在右侧和下方插入新的转置；幸运的是，左侧和下方新插入的转置可与原始转置抵消，最终得到一个转置的优化结果。

然而，如果先向下移动右侧的转置，如图 (c) 所示，那么下方和左侧新插入的转置只能与原有转置合并，最终仍残留两个转置，得到次优结果。

e-graph 通过 等式饱和 机制解决上述问题。它通过反复应用重写规则，将计算图中所有可能的等价节点都纳入 e-graph 中，最终结合成本模型与提取器，选出全局最优的计算子图。

表1：转置优化相关重写规则。该规则集是 nncase 解决转置优化问题的核心，为“等式饱和”机制提供关键支撑。例如，“消除空转置”规则可直接剔除计算图中无实际作用的多重转置操作，避免冗余数据移动；“融合连续转置”规则则能合并连续的转置变换，减少中间计算步骤。这些规则在 e-graph 中并非按固定顺序执行，而是同时生效，保留所有等价的转置优化路径，从而有效规避传统编译器因“阶段排序问题”导致的冗余转置残留，为后续基于 Roofline 成本模型提取全局最优子图奠定基础。

2.3 自动向量化

基于 e-graph 的类 Attention 子图自动向量化示意图

基于 e-graph 实现自动向量化布局转换的核心思路并不复杂，主要是在计算图中插入打包、打包后运算、解包等操作，并配合折叠规则进行优化。当然，具体实现涉及许多细节处理。其中最主要的挑战在于，需要确保微内核能够支持各类打包后的运算操作。

2.4 自动分布式策略搜索

我们基于 OneFlow 的 SBP（Split, Broadcast, Partial）抽象实现了自动分布式策略搜索。SBP 能够精确描述张量在不同设备间的分布状态。

图4：SBP 抽象与设备布局示意图。(a) 展示 1D 设备布局下 [2,2] 张量的三种 SBP（Split、Broadcast、Partial）策略实例；(b) 呈现 2D 嵌套拓扑（两组设备，每组含两个物理设备）及对应的 2D SBP 配置，直观体现 SBP 对不同拓扑下张量分布式状态的表征能力。

在集成 SBP 表示时，我们发现具有相同 SBP 属性的计算节点是等价的。基于此，我们可以在 e-graph 中构建 Boxing Sharding -> Computation -> Boxing UnSharding 的节点序列，并通过折叠（Fold）Sharding/UnSharding 操作来实现 SBP 属性的传播。这使我们得以复用 e-graph 的重写系统，并为 Boxing 操作引入基于 alpha-beta 的成本模型，与计算操作进行统一的成本评估。下图展示了一个简单的矩阵乘法分布式切分策略搜索逻辑：

分布式策略搜索的 e-graph 可视化。

2.5 自动调度

我们设计了以 Tile（分块）为粒度的调度设计空间，该空间由结构维度和参数维度构成。

图7：计算内核的设计空间，拆分为结构维度与参数维度。左侧展示结构维度（y 轴）的层级化 Tile 图，含算子融合逻辑；右侧及下方呈现参数维度（x 轴）的配置差异，对比不同分块大小与缓冲区位置对数据局部性的影响，绿色框为优化配置，红色框为低效配置。

结构维度的采样决定了具体的算子融合方式与循环顺序，而参数维度的采样则决定了每个 Tile 的大小以及缓冲区的位置。

具体的调度流程示意图。

调度流程首先将标准的计算图 Lower 到 Affine 表示形式。然后，结合硬件参数配置，将单个 Affine 操作转换为嵌套的 Tile Graph 表示。这本质上是一种符号化的后分块融合（Post-Tiling Fusion）过程。在 Tile Graph 的结构设计上，我们借鉴了部分 Welder/Tile Flow 的思路。

我们的核心区别在于，采用蒙特卡洛树搜索（MCTS）配合多面体分析来探索各种程序组织形式，算子融合是其中的一个子集。MCTS 搜索树中的每个节点对应一个 Tile Graph，随后对该 Tile Graph 进行树遍历以构建混合整数非线性规划（MINLP）模型。通过 MCTS 的迭代过程，有效解决了早期难以评估融合顺序（如 ABC 还是 BCD）的问题。

MCTS 搜索过程示意图。

对具体的 Tile Graph 进行 MINLP 建模是一个复杂的过程。我们的关键创新在于：使用离散变量支持缓冲区位置（Buffer Location）的搜索，并将存储层级与计算层级解耦（类似于 Halide 中的 compute_at 和 store_at 调度原语）。我们通过公式化模型评估内存足迹、内存复用和微内核计算开销，最终通过一次求解直接得到最优的 Tile 大小、各层级缓冲区大小及偏移量。

MINLP 建模与求解示意图。

2.6 代码生成

在代码生成前，我们会进行缓冲区别名分析，并结合 SBP 属性进行内存规划。

图8：利用 NNCase 张量模板库（NTT，Tensor Template Library）原语生成的 C++ 代码。代码通过ntt::fixed_shape_v等静态维度接口，将编译期可知的张量形状（如模型权重维度）提前解析为立即值，消除运行时形状计算开销；同时调用ntt::vector等显式向量类型，精准匹配Auto Vectorize模块确定的SIMD指令需求。此外，代码隐式集成NTTD（分布式扩展）的拓扑感知逻辑，可直接对接Auto Distribution模块输出的SBP策略，确保生成的代码在单核心或多核心场景下，均能高效复用片上缓存、减少数据移动，成为连接前端优化与硬件执行的关键纽带。

特别值得一提的是我们的 NTT 库，它支持固定形状与动态形状的混合编程，以及基于类型的分布式切分。除了高效之外，它还能显著降低编写分布式内核的心智负担。

使用 NTTD（NTT 分布式扩展）编写的 SUMMA 算法示例。SUMMA 是一种分布式矩阵乘法算法，其基于切分类型信息，通过 remote 函数配合运行时能力，可以方便地访问对等设备上的张量。

三、性能实验评估

我们在 x86 架构上测试了非分布式场景。通过自动向量化与自动调度，nncase 达到了 llama.cpp 约 80% 的性能，超过了其他编译器方案。

MLC LLM 性能较弱，主要是因为其 Meta Schedule 在 LLM 场景下难以有效应用，实际使用的是预设的调度模板 DLight。而 nncase 则利用多面体编译技术处理动态形状问题，将自动调度的结果与微内核组合，尽可能逼近自动调度所搜索到的理论性能。

在 4核/8核的测试中，nncase 的性能略微超越了 llama.cpp。主要原因是 llama.cpp 采用基于共享内存的并行加速方案，该方案在每个算子计算结束后都需要进行一次全局同步，引入了额外开销。

而 nncase 则是以原生分布式的视角，尽量减少同步/通信开销为准则进行分布式搜索，其实际同步次数明显少于 llama.cpp。因此，即便在自动生成算子性能不及手写实现的情况下，nncase 仍能实现与之媲美的性能。

图10：多核心场景（4T/8T）下LLM token 生成吞吐量对比。基于 AMD Ryzen 9 5900X 平台实验，nncase 展现出卓越的扩展性，性能超过手动优化的库。Qwen3-0.6B-F16 配置下，4T 时 nncase（23.5 tokens/s）略超 llama.cpp（23.2 tokens/s），8T 仍保持领先；Qwen3-1.7B-F16 的 4T 性能较 1T 提升 74%，远超 llama.cpp 的 32%。这得益于其静态任务划分（规避 OpenMP 动态调度开销）与核心级分布式粒度优化，突破传统共享内存并行的效率瓶颈。

总结与展望

nncase 解决了 dynamic shape、auto vectorize、auto distribution、auto schedule 等关键问题，以 24 位贡献者开发的编译器，达到了 1396 位贡献者手写优化所取得的性能，是一个显著的成果。

然而，当前仍面临诸多挑战，例如 computation-communication overlap、flash/paged attention 等棘手问题尚未解决。

通算融合方面虽有过一些探索，但尚不完善 [Hands-On Polyhedral] 自制通算融合 DSL Part1[10]。

该项目因不可抗力原因已停止，但也为我们探索新方向提供了契机。

目前图编译器（graph compiler）未能成功的主要原因仍是过度依赖自动调度。现有优化中，一部分无法被有效表示，另一部分虽可通过调度实现，但过程漫长。因此，内核编译器（kernel compiler）目前更受关注。

图编译器需要开放其中间表示，为用户提供操作接口。然而，tvm.te 等接口的表现力仍显不足。基于分块（tile-based）的中间表示是一个有前景的切入点，它允许用户手动编写调度，由编译器管理资源，最终嵌入到计算图中进行整体优化，实现编译与计算的协同。

参考资料

[1] 郑启航: https://www.zhihu.com/people/28-18-94-5

[2] Term Rewriting, Equality Saturation and Deep Learning Compilers: https://zhuanlan.zhihu.com/p/437972991

[3] Equality Saturation 优化在 AI 编译器中遇到的挑战: https://zhuanlan.zhihu.com/p/605459519

[4] Glenside : 如何自动发现im2col布局转换?: https://zhuanlan.zhihu.com/p/456616977

[5] OneFlow源码解析：Eager模式下的SBP Signature推导: https://zhuanlan.zhihu.com/p/605108910

[6] AKG: 使用post-tiling fusion策略完成无副作用的内存优化 : https://zhuanlan.zhihu.com/p/535606722

[7] [Hands-On Polyhedral] 动手多面体编译: https://zhuanlan.zhihu.com/p/21062125487

[8] SunnyCase: https://www.zhihu.com/people/01688e8d818824d2d21b48434f74c7cc

[9] 分布式存储架构下的矩阵乘与编译器: https://zhuanlan.zhihu.com/p/6906641426

[10] [Hands-On Polyhedral] 自制通算融合DSL Part1: https://zhuanlan.zhihu.com/p/1954332062753493221

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：http://www.itsolotime.com/archives/16551