Ember: A Compiler for Efficient Embedding Operations on Decoupled Access-Execute Architectures

📄 arXiv: 2504.09870v1 📥 PDF

作者: Marco Siracusa, Olivia Hsu, Victor Soria-Pardos, Joshua Randall, Arnaud Grasset, Eric Biscondi, Doug Joseph, Randy Allen, Fredrik Kjolstad, Miquel Moretó Planas, Adrià Armejach

分类: cs.AR, cs.LG, cs.PL

发布日期: 2025-04-14

备注: 14 pages, 19 figures, under review

💡 一句话要点

Ember编译器：为解耦访问-执行架构优化嵌入操作

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 编译器 嵌入查找 解耦访问-执行架构 DAE 性能优化 推荐系统 图学习

📋 核心要点

  1. 现有方法在处理推荐模型等应用中的不规则嵌入查找时效率较低，成为性能瓶颈。
  2. Ember编译器通过为解耦访问-执行架构生成优化代码，从而加速嵌入查找过程。
  3. 实验表明，使用Ember编译的代码在DAE架构上比GPU有显著的性能和能效提升。

📝 摘要（中文）

不规则嵌入查找是推荐模型、稀疏大语言模型和图学习模型中的关键瓶颈。本文首先证明，通过将这些查找卸载到专用访问单元，解耦访问-执行（DAE）处理器在端到端模型上实现了比GPU高2.6倍的性能和6.4倍的性能/瓦特。然后，我们提出了Ember编译器，用于从PyTorch和TensorFlow自动生成优化的DAE代码。与其他DAE编译器相反，Ember具有多个专门为不同优化级别设计的中间表示。通过这种方式，Ember可以实现所有优化，以匹配手写代码的性能，从而充分释放DAE架构在大规模上的潜力。

🔬 方法详解

问题定义：论文旨在解决推荐系统、稀疏大语言模型和图学习模型中普遍存在的、由不规则嵌入查找导致的性能瓶颈问题。现有方法，如在GPU上执行这些查找，效率较低，限制了整体模型的性能和能效。

核心思路：论文的核心思路是利用解耦访问-执行（DAE）架构的优势，将嵌入查找操作卸载到专门的访问单元。DAE架构允许访问单元独立于执行单元进行数据访问，从而减少了数据访问的延迟，提高了整体性能。Ember编译器的目标是自动生成针对DAE架构优化的代码，充分发挥其潜力。

技术框架：Ember编译器的整体框架包括前端、中间表示（IR）和后端。前端负责解析PyTorch和TensorFlow等框架的代码，并将其转换为Ember的内部表示。中间表示是Ember的核心，它包含多个专门为不同优化级别设计的IR。后端负责将优化后的IR转换为DAE架构上的可执行代码。

关键创新：Ember的关键创新在于其多层次的中间表示。不同的IR允许Ember在不同的优化阶段应用不同的优化策略，从而实现更精细的控制和更高的性能。与其他DAE编译器不同，Ember能够实现与手写代码相当的性能，充分释放DAE架构的潜力。

关键设计：Ember的关键设计包括：1) 针对DAE架构的特定优化，例如数据预取和缓存管理；2) 多层次的中间表示，允许在不同的优化阶段应用不同的优化策略；3) 自动代码生成，减少了手动编写优化代码的工作量。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，通过Ember编译器生成的代码在DAE架构上实现了显著的性能提升。具体来说，与GPU相比，Ember在端到端模型上实现了2.6倍的性能提升和6.4倍的性能/瓦特提升。此外，Ember能够生成与手写代码性能相当的代码，充分释放了DAE架构的潜力。

🎯 应用场景

Ember编译器及其优化的DAE架构在推荐系统、自然语言处理（特别是稀疏大语言模型）和图学习等领域具有广泛的应用前景。通过加速嵌入查找操作，可以显著提高这些模型的训练和推理速度，从而支持更大规模的模型和更复杂的应用。这有助于提升用户体验，并推动相关领域的发展。

📄 摘要（原文）

Irregular embedding lookups are a critical bottleneck in recommender models, sparse large language models, and graph learning models. In this paper, we first demonstrate that, by offloading these lookups to specialized access units, Decoupled Access-Execute (DAE) processors achieve 2.6$\times$ higher performance and 6.4$\times$ higher performance/watt than GPUs on end-to-end models. Then, we propose the Ember compiler for automatically generating optimized DAE code from PyTorch and TensorFlow. Conversely from other DAE compilers, Ember features multiple intermediate representations specifically designed for different optimization levels. In this way, Ember can implement all optimizations to match the performance of hand-written code, unlocking the full potential of DAE architectures at scale.