论文题目：AttentionLego: An Open-Source Building Block For Spatially-Scalable Large Language Model AcceleratorWith Processing-In-Memory Technology
论文链接：https://arxiv.org/abs/2401.11459
Code已开源：https://bonany.cc/attentionleg
自注意力模块是Transformer架构中的关键组成部分，但是，使用GPU进行计算时，自注意力模块对I/O带宽的需求非常高，这导致了在内存和处理单元之间传输中间计算结果时的瓶颈。
AttentionLego 的设计思想是创建一个可扩展的、基于自注意力机制的大型语言模型（LLM）加速器。这个加速器的核心是一个名为AttentionLego的基本构建块，它利用了处理内存（Processing-In-Memory, PIM）技术来提高计算效率。PIM技术将处理单元和内存集成在同一物理芯片上，从而消除了处理器和内存之间数据传输的需求，显著减少了延迟并提高了性能。
AttentionLego设计思想
AttentionLego 的名称中的“Lego”暗示了其模块化和可扩展性。就像乐高积木（Lego）可以组合成各种结构一样，AttentionLego旨在通过重复使用其基本构建块来构建更大规模的LLM处理器。这种设计允许开发者根据需要轻松地扩展硬件架构，以适应不同规模的LLM应用。
具体来说，AttentionLego由以下几个主要部分组成：
- 输入处理模块 （ Input Process module ）：负责计算输入权重矩阵（WQ, WK, WV）与输入向量的乘积。
- 分数模块 （ Score module ）：计算查询（Q）和键（K）向量的点积，生成QKT矩阵。
- Softmax模块 ：计算softmax非线性激活函数，用于确定注意力权重。
- DMA模块 （ Direct Memory Access ）：控制模块间以及与外部存储器之间的数据传输。
- 顶层控制器 （ Top controller ）：协调不同模块之间的通信、数据流和功能操作，确保整个系统按照设计要求正常运行。
AttentionLego架构
推荐阅读
- • TinyGPT-V：小而强大，轻量级骨架打造高效多模态大语言模型
- • 2024：ToB、Agent、多模态
- • TA们的RAG真正投产了吗？（上）
- • 2023年最新LLM发展时间线一览（ChatGPT、LLaMA等）
欢迎关注我的公众号“ PaperAgent ”， 每天一篇大模型（LLM）文章来锻炼我们的思维，简单的例子，不简单的方法，提升自己。