跳转至

Nova：面向深度学习的端到端 MLIR 编译器

文章背景与核心概要

现代深度学习模型在规模化训练和推理时，其性能很大程度上取决于高层数学运算映射到底层硬件上的效率。传统的深度学习框架通常依赖不透明、预编译的硬编码算子库（如 Attention 算子），这导致执行模型缺乏全图视角，难以实现硬件利用率的最大化。

为了弥补这一鸿沟，本文介绍了 Nova 的最新进展——一个基于 MLIR（多级中间表示）构建的自动化、端到端 JIT（即时编译）编译器。Nova 通过直接从计算图结构中综合出细粒度内核，实现了对硬件映射的绝对控制。它能够捕获急切执行（eager execution），将前向和反向传播统一到一个单一的值语义方言中，并执行激进的全图优化。

通过将复杂的因果注意力子图、逐元素（element-wise）操作以及受内存带宽限制的归一化操作融合成单个内核，Nova 显著减少了全局内存的往返读写，同时严格保持了数值等价性。在 Ada 6000 GPU 上训练完整的 GPT-2 架构的评估中，Nova 展现出了卓越的端到端吞吐量，平均达到每秒 421K 个 Token，超越了原生急切执行（406K）以及 PyTorch 的 torch.compile（405K）。

概述与总结

Nova 是一个专为深度学习设计的自动化端到端 JIT（即时编译）编译器，基于 MLIR（多层中间表示）构建。通过直接从计算图的结构中综合出细粒度内核，Nova 实现了对硬件映射的绝对控制。

Nova is an automated, end-to-end JIT (Just-In-Time) compiler designed for deep learning, built using MLIR (Multi-Level Intermediate Representation). By synthesizing fine-grained kernels directly from the computation graph's structure, Nova achieves absolute control over hardware mapping.

与依赖刚性、不透明、预编译内核库（如 Attention 相关库）的传统张量框架不同，Nova 可以捕获急切执行，将前向和反向传递统一到单个值语义方言中，并执行积极的全图优化。这使得复杂的因果注意力子图、逐元素操作和受内存限制的归一化能够融合成单个内核，在严格保持数值平价的同时，大幅减少全局内存往返。

Unlike traditional tensor frameworks that rely on rigid, opaque, pre-compiled kernel libraries (such as those for Attention), Nova captures eager executions, unifies forward and backward passes into a single value-semantic dialect, and executes aggressive whole-graph optimizations. This enables complex causal attention sub-graphs, element-wise operations, and memory-bound normalizations to be fused into single kernels, drastically reducing global memory roundtrips while strictly maintaining numerical parity.

论文元数据

  • arXiv 标识符： arXiv:2608.00029 (v2)
  • 主学科： 人工智能 (cs.AI)
  • 次学科： 硬件架构 (cs.AR)、机器学习 (cs.LG)、程序设计语言 (cs.PL)
  • 提交日期： 2026年7月15日 (最后修订于：2026年8月22日)
  • DOI： 10.48550/arXiv.2608.00029
  • arXiv Identifier: arXiv:2608.00029 (v2)
  • Primary Subject: Artificial Intelligence (cs.AI)
  • Secondary Subjects: Hardware Architecture (cs.AR), Machine Learning (cs.LG), Programming Languages (cs.PL)
  • Submission Date: July 15, 2026 (Last revised: August 22, 2026)
  • DOI: 10.48550/arXiv.2608.00029

作者

  • Adwaid Suresh
  • Aparna A
  • Harshini V M
  • Jona Delcy C A
  • Killi Uma Maheswara Rao
  • Ram Charan Golla
  • Surendra Vendra

Authors

  • Adwaid Suresh
  • Aparna A
  • Harshini V M
  • Jona Delcy C A
  • Killi Uma Maheswara Rao
  • Ram Charan Golla
  • Surendra Vendra

摘要

大规模深度学习模型的性能在很大程度上取决于高层数学运算如何有效地映射到底层物理硬件上。虽然高层张量框架提供了灵活的抽象，但它们的执行模型本质上缺乏最大化硬件利用率所需的全局图可见性，这往往迫使开发者在处理 Attention 等复杂操作时依赖不透明的手写内核库。

The performance of deep learning models at scale relies heavily on how effectively high-level mathematical operations are mapped to underlying physical hardware. While high-level tensor frameworks provide flexible abstractions, their execution models inherently lack the whole-graph visibility required to maximize hardware utilization, often forcing a reliance on opaque, hand-written kernel libraries for complex operations like Attention.

为了弥补这一鸿沟，我们推出了 Nova 的下一个迭代版本——一个自动化的端到端 JIT 编译器，它通过直接从计算结构中综合细粒度内核，实现了对硬件映射的绝对控制。在这项工作中，我们扩展了 Nova 的编译流水线，以原生支持完整的 Transformer 架构。通过捕获急切执行并将前向和反向传播统一到一个值语义方言中，Nova 释放了激进的全图优化能力。

To bridge this gap, we present the next iteration of Nova, an automated end-to-end JIT compiler that achieves absolute control over hardware mapping by synthesizing fine-grained kernels directly from the computation's structure. In this work, we extend Nova's compilation pipeline to natively support full Transformer architectures. By capturing eager executions and unifying forward and backward passes into a single value-semantic dialect, Nova unlocks aggressive whole-graph optimizations.

与依赖刚性的预编译库调用不同，Nova 专注于广泛的跨算子融合，将复杂的因果注意力子图、逐元素操作和内存密集型归一化直接坍缩为单一的融合内核，从而大幅减少全局内存往返。

Rather than relying on rigid, pre-compiled library calls, Nova focuses on extensive cross-operator fusions, collapsing complex causal attention sub-graphs, element-wise operations, and memory-bound normalizations directly into single fused kernels to drastically reduce global memory roundtrips.

在 Ada 6000 GPU 上训练完整的 GPT-2 架构的评估中，Nova 表现出了优越的端到端吞吐量，平均达到 421K 词元/秒，而我们自己的急切执行为 406K，PyTorch 的 torch.compile 为 405K。通过编译器原生的融合大幅减少内存受限的开销，Nova 在严格保持数值等价性的同时，在现代硬件上实现了高效的完整大模型（LLM）编译。

In our evaluations training a full GPT-2 architecture on Ada 6000 GPUs, Nova demonstrates superior end-to-end throughput, averaging 421K tokens/second compared to 406K for our own eager execution and 405K for PyTorch's torch.compile. By drastically reducing memory-bound overheads through compiler-native fusion, Nova enables efficient full LLM compilation on modern hardware while strictly maintaining numerical parity.

全文与访问链接

外部参考与引用