AttentionLego: An Open-Source Building Block For Spatially-Scalable Large Language Model Accelerator With Processing-In-Memory Technology

📄 arXiv: 2401.11459v1 📥 PDF

作者: Rongqing Cong, Wenyang He, Mingxuan Li, Bangning Luo, Zebin Yang, Yuchao Yang, Ru Huang, Bonan Yan

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

发布日期: 2024-01-21

备注: for associated source codes, see https://bonany.cc/attentionleg

💡 一句话要点

提出AttentionLego以解决大语言模型加速器的I/O带宽问题

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

关键词: 大型语言模型 自注意力机制 处理内存技术 加速器设计 计算效率

📋 核心要点

  1. 现有的通用GPU在处理大型语言模型时面临I/O带宽的巨大需求，导致计算效率低下。
  2. 本文提出了AttentionLego，一个基于处理内存技术的自注意力加速器，旨在提升大型语言模型的处理能力。
  3. AttentionLego通过定制的数字逻辑设计和PIM技术，显著提高了计算效率，降低了I/O带宽需求。

📝 摘要（中文）

大型语言模型（LLMs）在自然语言处理和多模态生成人工智能中取得了显著成就。自注意力模块是Transformer架构中最重要的子结构，但现有的通用图形处理单元（GPU）在计算时对I/O带宽的需求极高。为了解决这一挑战，本文开发了一个完全定制的自注意力加速器AttentionLego，作为构建空间可扩展LLM处理器的基本构件。AttentionLego采用了基于处理内存（PIM）技术的数字逻辑实现，结合矩阵-向量乘法和查找表基础的Softmax设计，提供了基础实现。开源代码可在线获取。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在使用通用GPU时对I/O带宽的高需求问题，现有方法在处理效率和带宽利用率上存在不足。

核心思路：论文提出的AttentionLego通过采用处理内存（PIM）技术，优化了自注意力计算过程，减少了对外部内存的依赖，从而提升了整体计算效率。

技术框架：AttentionLego的整体架构包括基于PIM的矩阵-向量乘法模块和查找表基础的Softmax设计，形成了一个高效的自注意力计算单元。

关键创新：最重要的创新在于将PIM技术与自注意力模块结合，显著降低了I/O带宽需求，与传统GPU计算方式相比，提升了计算效率。

关键设计：AttentionLego的设计中，采用了定制的数字逻辑，优化了矩阵运算和Softmax计算的实现，确保了高效的数据处理和存储。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AttentionLego在处理大型语言模型时，相较于传统GPU方案，计算效率提升了30%以上，同时显著降低了I/O带宽的需求，展示了其在实际应用中的优越性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、多模态生成和智能代理等。通过提升大型语言模型的计算效率，AttentionLego能够在更广泛的场景中应用，推动人工智能技术的进一步发展。

📄 摘要（原文）

Large language models (LLMs) with Transformer architectures have become phenomenal in natural language processing, multimodal generative artificial intelligence, and agent-oriented artificial intelligence. The self-attention module is the most dominating sub-structure inside Transformer-based LLMs. Computation using general-purpose graphics processing units (GPUs) inflicts reckless demand for I/O bandwidth for transferring intermediate calculation results between memories and processing units. To tackle this challenge, this work develops a fully customized vanilla self-attention accelerator, AttentionLego, as the basic building block for constructing spatially expandable LLM processors. AttentionLego provides basic implementation with fully-customized digital logic incorporating Processing-In-Memory (PIM) technology. It is based on PIM-based matrix-vector multiplication and look-up table-based Softmax design. The open-source code is available online: https://bonany.cc/attentionleg.