PLAIN: Leveraging High Internal Bandwidth in PIM for Accelerating Large Language Model Inference via Mixed-Precision Quantization
TL;DR 精炼摘要
本文提出了PLAIN,这是一种新颖的软硬件协同设计框架,旨在通过混合精度量化加速大型语言模型推理。该框架通过优化参数的量化精度,结合PIM的特性,有效提升了推理性能,实现了5.03倍和1.69倍的性能提升,同时保持模型质量损失在可接受范围内。
摘要
DRAM-based processing-in-memory (DRAM-PIM) has gained commercial prominence in recent years. However, its integration for deep learning acceleration, particularly for large language models (LLMs), poses inherent challenges. Existing DRAM-PIM systems are limited in computational capabilities, primarily supporting element-wise and general matrix-vector multiplication (GEMV) operations, which contribute only a small portion of the execution time in LLM workloads. As a result, current systems still require powerful host processors to manage compute-heavy operations. To address these challenges and expand the applicability of commodity DRAM-PIMs in accelerating LLMs, we introduce PLAIN, a novel software/hardware co-design framework for PIM-enabled systems. PLAIN leverages the distribution locality of parameters and the unique characteristics of PIM to achieve optimal trade-offs between inference cost and model quality. Our framework includes three key innovations: 1) firstly, we propose a novel quantization algorithm that determines the optimal precision of parameters within each layer, considering both algorithmic and hardware characteristics to optimize hardware mapping; 2) PLAIN strategically utilizes both GPUs and PIMs, leveraging the high internal memory bandwidth within HBM for attention layers and the powerful compute capability of conventional systems for fully connected (FC) layers; 3) PLAIN integrates a workload-aware dataflow scheduler that efficiently arranges complex computations and memory access for mixed-precision tensors, optimizing execution across different hardware components. Experiments show PLAIN outperforms the conventional GPU with the same memory parameters and the state-of-the-art PIM accelerator, achieving a 5.03× and 1.69× performance boost, with negligible model quality loss.
思维导图
论文精读
中文精读
1. 论文基本信息
1.1. 标题
PLAIN: Leveraging High Internal Bandwidth in PIM for Accelerating Large Language Model Inference via Mixed-Precision Quantization (PLAIN:利用 PIM 中的高内部带宽通过混合精度量化加速大型语言模型推理)
1.2. 作者
Yiwei Hu, Fangxin Liu, Zongwu Wang, Yilong Zhao, Tao Yang, Li Jiang, Haibing Guan 等。 作者主要隶属于上海交通大学(SJTU),其中 Li Jiang 为通讯作者。
1.3. 发表期刊/会议
本论文在 Published at (UTC): 2025-10-26T00:00:00.000Z 处注明了发布时间,表明其计划在 2025 年发布。根据其内容和格式,推测其目前为预印本(pre-print)或已投稿至某个顶级计算机体系结构或机器学习系统会议,并已被接受但尚未正式出版。
1.4. 发表年份
2025年(根据 Published at (UTC) 信息)
1.5. 摘要
DRAM(动态随机存取存储器)内存内计算(Processing-in-Memory, PIM)近年来在商业上取得了显著进展。然而,将其集成到深度学习(Deep Learning)加速,特别是大型语言模型(Large Language Models, LLMs)中,面临固有的挑战。现有 DRAM-PIM 系统的计算能力有限,主要支持逐元素(element-wise)操作和通用矩阵向量乘法(General Matrix-Vector Multiplication, GEMV),而这些操作在 LLM 工作负载中只占执行时间的一小部分。因此,目前的系统仍然需要强大的主机处理器来管理计算密集型操作。
为了应对这些挑战并扩展商用 DRAM-PIM 在加速 LLMs 方面的适用性,本文引入了 PLAIN,一个新颖的针对 PIM 系统设计的软硬件协同设计(software/hardware co-design)框架。PLAIN 利用参数的分布局部性(distribution locality)和 PIM 的独特特性,在推理成本和模型质量之间实现最佳权衡。该框架包含三个关键创新点:1)提出了一种新颖的量化算法,在考虑算法和硬件特性的前提下,确定每层参数的最佳精度,以优化硬件映射;2)PLAIN 策略性地利用 GPU 和 PIM,发挥 HBM(高带宽内存)内部的高内存带宽用于注意力层(attention layers),并利用传统系统强大的计算能力处理全连接层(fully connected, FC layers);3)PLAIN 集成了一个工作负载感知的数据流调度器(workload-aware dataflow scheduler),用于高效地编排混合精度张量(mixed-precision tensors)的复杂计算和内存访问,优化跨不同硬件组件的执行。实验结果表明,PLAIN 在相同内存参数下,优于传统的 GPU 和最先进的 PIM 加速器,实现了 5.03 倍和 1.69 倍的性能提升,且模型质量损失可忽略不计。
1.6. 原文链接
/files/papers/69571ce38c5983e9f07b96e1/paper.pdf
2. 整体概括
2.1. 研究背景与动机
大型语言模型(LLMs)在各种语言任务中展现出卓越的能力,但其庞大的模型规模和对计算资源的巨大需求对部署构成了严峻挑战。例如,GPT-3 拥有 1750 亿参数,需要 326 GB 的 FP16 内存,而高端 GPU(如 A100 和 H100)仅提供最高 80 GB 内存,这形成了明显的内存不匹配。推理过程中的“内存墙”(memory wall)问题是一个关键瓶颈,尤其是在词元(token)生成阶段。GPU 计算单元与 DRAM(动态随机存取存储器)之间有限的带宽导致硬件利用率低下。例如,在 NVIDIA DGX A100 上运行 GPT-3,尽管拥有 640 GB HBM3 和 26.8 TB/s 的带宽,计算单元利用率仍可能低于 1%。
尽管加速器的计算性能显著提升,但内存容量和带宽并未同步增长。这凸显了在 LLM 部署中对更内存高效解决方案的迫切需求。
DRAM-based Processing-in-Memory (DRAM-PIM) 提出了一种有前景的解决方案,通过将计算单元放置在内存 Bank 附近,显著提高内存带宽,使其非常适合 LLM 等带宽受限(bandwidth-bound)的工作负载。然而,现有 DRAM-PIM 应用主要针对 GEMV(通用矩阵向量乘法)等内存密集型操作,其在 LLM 加速方面的潜力尚未被充分发掘。
模型量化(Model Quantization) 是另一种减少内存占用和数据移动的方法,通过使用更少的比特来表示权重(weights)和激活值(activations)。虽然 6-bit 量化在效率和准确性之间取得了良好平衡,但主流硬件缺乏对 6-bit 执行的原生支持,通常偏爱 4-bit 或 8-bit 格式,这限制了其实际部署。
为了扩展 DRAM-PIM 在深度神经网络(DNNs)中的应用,迫切需要一个 PIM 友好的框架。将 DRAM-PIM 架构与低精度( 6 比特)量化相结合,为加速 LLMs 提供了一个有前景的方法。然而,实现峰值性能仍需解决以下三个关键挑战:
- C1:平衡精度与压缩率。 现有量化算法难以在精度和压缩率之间取得最佳平衡,高效量化同时不显著降低精度至关重要。
- C2:与 DRAM-PIM 架构不兼容。 现有量化框架并非针对 DRAM-PIM 设计,导致其内存内计算能力未能得到充分利用。要充分发挥 DRAM-PIM 的优势,量化必须与其架构特性协同设计。
- C3:实现实际加速。 尽管量化在理论上具有优势,但将其转化为 DRAM-PIM 上的实际加速却很困难。有效协调量化策略与 DRAM-PIM 架构是实现这些收益的关键。
2.2. 核心贡献/主要发现
为应对上述挑战,本文提出了 PLAIN,一个软硬件协同设计框架,旨在通过 DRAM-PIM 架构实现混合精度量化(MPQ)。PLAIN 的核心贡献和主要发现包括:
- 提出新颖的量化算法: PLAIN 引入了一种局部性感知(Locality-Aware)的自适应量化(Adaptive Quantization)方法,根据组件的敏感性分配不同精度,从而在精度和压缩率之间取得平衡,解决了挑战 C1。该算法考虑了算法和硬件特性,以优化硬件映射。
- DRAM-PIM 友好的推理引擎: PLAIN 设计了一个 DRAM-PIM 友好的推理引擎(DRAM-PIM-Friendly Inference Engine),专为混合精度计算量身定制,最大限度地提高内存带宽和计算效率,解决了挑战 C2。它策略性地利用 GPU 和 PIM,将 HBM 内部的高内存带宽用于注意力层,而常规系统的强大计算能力用于全连接层。
- 自适应工作负载调度器: PLAIN 开发了一个自适应工作负载调度器(Adaptive Workload Scheduler),平衡跨核心的混合位操作,确保量化收益有效地转化为实际加速,解决了挑战 C3。该调度器高效地编排复杂计算和混合精度张量的内存访问,优化跨不同硬件组件的执行。
- 显著的性能和能耗提升: 实验结果表明,PLAIN 在相同内存参数下,相对于传统 GPU 和最先进的 PIM 加速器,实现了 5.03 倍和 1.69 倍的性能提升,同时模型质量损失可忽略不计。在能耗方面,PLAIN 也实现了显著的降低。
3. 预备知识与相关工作
3.1. 基础概念
- 大型语言模型 (Large Language Models, LLMs): 一类基于深度学习的神经网络模型,通常采用 Transformer 架构,拥有数十亿甚至数千亿参数,能够理解、生成和处理人类语言,并在各种自然语言处理(NLP)任务中表现出色。
- Transformer 架构 (Transformer Architecture): 一种基于自注意力机制(self-attention mechanism)的神经网络架构,是 LLMs 的基础。它主要由编码器(encoder)和解码器(decoder)组成,LLMs 通常是纯解码器架构。
- 预填充阶段 (Prefill Stage): LLM 推理的第一个阶段,处理输入的提示(prompt),生成所有输入词元对应的键(Key)和值(Value)矩阵,并将其缓存起来,用于后续的解码阶段。
- 解码阶段 (Decoding Stage): LLM 推理的第二个阶段,模型根据输入提示和缓存的键值对,逐个生成新的词元,直到达到停止条件。
- 注意力机制 (Attention Mechanism): Transformer 架构中的核心组件,允许模型在处理序列时,根据相关性对输入序列的不同部分赋予不同的权重,从而聚焦于最重要的信息。其核心计算包括查询(Query, Q)、键(Key, K)和值(Value, V)矩阵之间的交互。
- 前馈网络 (Feed-Forward Network, FFN): Transformer 层中的另一个关键组件,通常由两个线性层(linear layers)和一个激活函数(activation function)组成,用于对注意力机制的输出进行非线性变换。
- 内存墙 (Memory Wall): 指处理器(如 CPU 或 GPU)与主内存之间数据传输速率(带宽)的增长速度远低于处理器计算能力增长速度的现象。这导致在内存密集型任务中,系统性能受限于数据传输速度,而非处理器本身的计算能力。
- 内存内计算/处理 (Processing-in-Memory, PIM): 一种新型计算范式,旨在将计算单元直接集成到内存芯片或内存模块中,从而大幅减少数据在处理器和内存之间移动的距离和时间,有效缓解内存墙问题。
- DRAM-PIM (DRAM-based Processing-in-Memory): 基于 DRAM 技术实现的内存内计算,通过在标准的 DRAM 芯片或内存模块中添加少量计算逻辑,使其能够执行一些基本的计算操作,如逐元素操作、GEMV 等。
- 高带宽内存 (High Bandwidth Memory, HBM): 一种高性能的 3D 堆叠内存技术。它通过将多个 DRAM 芯片垂直堆叠并使用硅通孔(Through-Silicon Via, TSV)连接,实现极高的内存带宽和更低的功耗,常用于 GPU 和其他高性能计算系统。
- 模型量化 (Model Quantization): 一种模型压缩技术,通过将模型参数(如权重、激活值)从高精度浮点格式(如 FP32 或 FP16)转换为低精度整数格式(如 INT8 或 INT4),以减少模型大小、内存占用和计算量,从而加速推理过程。
- 混合精度量化 (Mixed-Precision Quantization, MPQ): 一种更高级的量化策略,它根据模型中不同层或不同参数的敏感性,为其分配不同的位宽。例如,对精度敏感的层使用较高位宽(如 INT8),对精度不敏感的层使用较低位宽(如 INT4),以在保持模型精度的同时实现更好的压缩和加速效果。
- SmoothQuant: 一种针对 LLM 的后训练量化(Post-Training Quantization)方法。它通过在推理前对激活值中的异常值(outliers)进行通道级(per-channel)的尺度缩放,并将这些尺度因子吸收(fold)到相邻的权重中,从而使激活值的分布更平滑、更易于低位宽量化,减少量化误差。
- 屋顶模型 (Roofline Model): 一种图形化的性能模型,用于分析和预测计算系统的性能瓶颈。它通过绘制峰值计算性能和峰值内存带宽的曲线,并结合工作负载的算术强度(arithmetic intensity,即每字节内存访问的浮点运算次数),来识别工作负载是受计算限制(compute-bound)还是受内存限制(memory-bound)。
- 通用矩阵向量乘法 (General Matrix-Vector Multiplication, GEMV): 一种基本的线性代数运算,形式为 ,其中 是矩阵, 是向量, 是结果向量。在深度学习中常用于全连接层或处理批次大小为 1 的推理场景。
- 通用矩阵乘法 (General Matrix Multiplication, GEMM): 一种基本的线性代数运算,形式为 ,其中
A, B, C都是矩阵。在深度学习中是计算密集型操作的核心,例如全连接层和注意力机制中的 QKV 计算。 - 困惑度 (Perplexity, PPL): 评估语言模型质量的指标,衡量模型预测一个样本序列(如一段文本)的能力。困惑度越低,表示模型对语言的建模能力越好,预测结果越接近真实分布。
3.2. 前人工作
论文在引言和背景部分提及了多项相关工作和现有挑战:
- DRAM-PIM 的现有应用及其局限性:
- 现有 DRAM-PIM 方案(如三星 HBM-PIM [22, 23, 24] 和 SK Hynix GDDR6-AiM [25])已在商业上取得进展,主要针对内存密集型操作,如 GEMV。
- 但这些 PIM 系统的计算能力通常有限,主要支持逐元素操作和 GEMV,这在 LLM 工作负载中只占一小部分执行时间。对于 LLM 中更复杂的计算密集型操作(如大型矩阵乘法),仍然需要强大的主机处理器。例如,AttAcc [8] 是一种异构 PIM 系统,它利用 GPU 进行预填充和 QKV 生成,PIM 用于解码阶段的多头注意力,但其设计仍有待优化。
- 主流量化方法的局限性:
- 模型量化是减少内存占用和加速推理的有效手段。
- 6-bit 量化被认为在效率和精度之间取得了良好平衡,但主流硬件缺乏对 6-bit 执行的原生支持,通常只支持 4-bit 或 8-bit 格式,这限制了 6-bit 量化的实际部署。
- 传统的量化方法(如 W4A4, W8A8)在不使用 SmoothQuant 等技术处理激活值异常值时,在大型 LLM 上可能导致显著的精度损失(如 Table II 所示)。
- 现有混合精度量化 (MPQ) 技术与 DRAM-PIM 架构的不兼容性:
- HAQ [27] 等混合精度量化技术使用强化学习来搜索每层的最佳位宽,并受限于硬件指标(如延迟、能耗)。
- LLM.int8() [28] 针对激活值中的大异常值,提出逐向量(vector-wise)量化,将异常值维度隔离为 FP16。
- 然而,这些现有细粒度混合精度技术并非专为 DRAM-PIM 设计,未能充分利用 DRAM-PIM 独特的架构约束和内存内计算能力,导致次优的性能。工作负载不平衡是 DRAM-PIM 中已映射混合精度模型面临的挑战,这阻碍了量化理论收益的实现。
3.3. 技术演进
LLMs 的快速发展带来了对计算和内存资源的巨大需求。随着模型规模的指数级增长,传统的 CPU/GPU 架构面临“内存墙”瓶颈,数据传输成为性能限制因素。PIM 技术通过将计算逻辑移至内存附近,显著提升了内存带宽,被视为解决这一瓶颈的有前景方案。同时,模型量化技术通过降低数据精度,进一步减少了内存占用和数据传输量。
PLAIN 的工作正处于 PIM 和量化技术交叉演进的脉络中。早期 PIM 方案主要侧重于通用计算或简单的内存密集型任务,计算能力有限。同时,量化技术不断演进,从简单的全模型定点量化发展到混合精度量化和针对 LLMs 异常值的特殊处理(如 SmoothQuant)。PLAIN 的创新点在于,它认识到简单地结合 PIM 和现有量化技术并不能达到最佳效果,必须通过软硬件协同设计,将 PIM 架构的特性(如高内部带宽、Bank 级并行性)与量化算法(如局部性感知、硬件感知)紧密结合,并设计高效的调度和重叠机制,才能真正发挥两者的协同优势,从而推动 LLM 在 PIM 平台上的高效部署。
3.4. 差异化分析
PLAIN 与现有方法的关键区别和创新点在于其软硬件协同设计理念,以及针对 LLM 推理场景的具体优化:
- 量化算法的硬件感知性:
- 现有工作: 许多混合精度量化方法(如 HAQ, LLM.int8())主要关注算法层面,或仅对通用硬件(如 GPU)进行有限的硬件感知。它们通常未考虑 DRAM-PIM 架构的独特约束(如有限的计算单元、高内部带宽、Bank 级并行性)。
- PLAIN 创新: PLAIN 提出了一个新颖的量化算法,它不仅考虑算法层面的精度敏感性(通过熵度量),还引入了硬件成本函数,将量化位宽与 DRAM-PIM 的计算和通信开销直接关联。这种“局部性感知自适应量化”方法,使其能够生成最适合 PIM 硬件映射的混合精度配置。
- PIM 架构的 LLM 优化:
- 现有工作: 现有 DRAM-PIM 系统(如 HBM-PIM, GDDR6-AiM)计算能力有限,主要支持 GEMV,对 LLM 中计算密集型的注意力层效率不高。AttAcc 虽尝试将注意力层卸载到 PIM,但其调度和量化集成可能不够精细。
- PLAIN 创新: PLAIN 设计了一个 DRAM-PIM 友好的推理引擎,通过引入 Bank-PIM 和 Stack-PIM 单元,专门为 LLM 的注意力层(QKV 投影、注意力分数、注意力上下文)进行计算卸载和加速。它充分利用了 HBM 的高内部带宽,同时保留 GPU 处理计算密集型全连接层的优势,实现了 GPU 与 PIM 的高效协同。
- 细粒度工作负载调度和重叠机制:
- 现有工作: 混合精度模型在 PIM 上可能导致工作负载不平衡,因为不同精度的操作计算量不同。现有调度机制可能无法有效处理这种异构性。
- PLAIN 创新: PLAIN 引入了位拆分词元 (Bit-wise splitting token) 策略,将 8-bit 激活值拆分为两个 4-bit 段,使得不同精度的计算可以在统一的 4-bit 乘法器上执行,从而解决了计算单元利用率不平衡的问题。此外,其自适应工作负载调度器结合了 Bank-PIM 调度表和贪婪调度算法,确保了 Bank 级并行性和负载平衡。更重要的是,PLAIN 设计了精密的重叠(overlapping)机制(如权重加载与计算重叠、量化/反量化与通信重叠、Softmax 与通信重叠),最大限度地隐藏了各种开销,将理论上的量化收益转化为实际的端到端加速。
4. 方法论
PLAIN 框架是一个软硬件协同设计,旨在优化 DRAM-PIM 系统上大型语言模型(LLMs)的推理性能。它通过结合新颖的混合精度量化算法、专门的 PIM 硬件架构以及高效的调度和重叠机制来实现这一目标。
4.1. 方法原理
PLAIN 的核心思想是利用 DRAM-PIM 架构的高内部带宽和 Bank 级并行性,结合针对 LLM 特性设计的混合精度量化方案,来克服 LLM 推理中的内存墙和计算效率问题。其背后的理论基础和直觉可以概括为:
- 局部性感知与硬件映射: LLM 的参数和激活值具有特定的分布特性,其中一些部分对精度更敏感,而另一些则可以承受较低的精度。PIM 硬件具有特定的计算能力和数据访问模式。通过局部性感知(Locality-Aware) 的量化算法,可以根据这些分布特性和硬件约束,智能地选择每层甚至每个张量的最佳位宽。这使得量化在保持模型质量的同时,能够最大程度地优化数据在 PIM 上的存储和计算映射。
- 异构计算的优势互补: LLM 包含多种类型的层,例如注意力层和全连接层。注意力层通常是带宽密集型(bandwidth-intensive),特别是键值缓存(KV cache)的访问。而全连接层则可能是计算密集型(compute-intensive)。PLAIN 认识到 GPU 在通用计算能力上的优势和 PIM 在高内部带宽上的优势,因此将注意力层卸载到 PIM 以利用其高带宽,而将全连接层保留在 GPU 上以利用其强大的计算能力,从而实现优势互补。
- 精细化调度与开销隐藏: 混合精度量化和异构计算会引入数据转换、通信和计算负载不平衡等新的开销。PLAIN 采用位拆分(Bit-wise splitting) 和动态调度(dynamic scheduling) 策略来平衡不同精度操作的工作负载,并利用重叠(overlapping)机制来隐藏数据传输、量化/反量化和 Softmax 等操作的延迟,确保 PIM 内部和 PIM 与 GPU 之间的高效数据流和并行执行,从而将理论上的性能优势转化为实际的端到端加速。
4.2. 核心方法详解 (逐层深入)
4.2.1. MIXQ-PIM 算法
MIXQ-PIM 算法是 PLAIN 框架中的核心量化算法,旨在实现 DRAM-PIM 友好的混合精度量化,以平衡推理成本和模型质量。
4.2.1.1. 量化粒度 (Quantization Granularity)
在大型语言模型中,激活值(activations)中的异常值(outliers)是低位宽量化(low-bit quantization)的主要障碍,因为它们会压缩大部分值到狭窄的范围内,导致精度损失。为了解决这个问题,PLAIN 采用了 SmoothQuant [9] 技术。
SmoothQuant 的核心思想是通过每通道尺度因子(per-channel scaling factors) 将激活值的异常值重新映射到权重中。具体来说,对于一个输入激活张量 和权重矩阵 的乘法运算,可以进行如下变换:
-
符号解释:
- : 原始的输出张量。
- : 原始的输入激活张量。
- : 原始的权重矩阵。
- : 一个对角矩阵,其对角线元素是每个通道的尺度因子(scaling factor)。这些尺度因子用于平滑激活值中的异常值。
- : 对角矩阵 的逆矩阵。
- : 经过平滑处理后的激活张量。
- : 经过尺度因子调整后的权重矩阵。
-
目的分析: 这个变换将激活值中的“难量化”部分转移到了权重中。由于权重通常是离线(offline)量化和存储的,这种转移可以使激活值的分布更平滑,从而更容易进行低位宽量化,同时保持最终计算结果的等价性。
经过平滑处理后,低误差的元素被重新组织到不同的通道中,使得混合精度量化更具硬件效率。这种转换允许 PLAIN 对激活值采用逐词元(token-wise)粒度进行量化,对权重采用逐多通道(multi-channel-wise)粒度进行量化。这两种粒度方案与 DRAM-PIM 的数据流良好对齐,可以在不引入复杂运行时开销的情况下提高精度控制。
4.2.1.2. 量化配置搜索 (Quantization Configuration Searching)
为了实现 DRAM-PIM 友好的量化,PLAIN 采用均匀对称量化(uniform and symmetric quantization),这种方法通过简单的尺度缩放和乘法将张量映射到整数,从而简化计算并降低功耗。目标是结合 INT8 和 INT4 精度以平衡准确性和效率。关键挑战在于如何在不进行重训练或昂贵的学习过程的情况下,以细粒度选择精度。PLAIN 提出了一种轻量级的、基于熵的启发式方法。
a) 权重精度熵 (Weight Precision Entropy): PLAIN 通过计算原始分布和量化分布之间的 KL 散度(KL-divergence)来量化权重误差:
-
符号解释:
- : 第 个权重块的量化熵。
- : KL 散度,用于衡量两个概率分布之间的差异。
- : 第 个权重块的原始浮点(FP)分布。
- : 第 个权重块量化为整数(INT)后的分布。
- INT: 表示量化后的位宽,此处从 INT4 或 INT8 中选择。
-
目的分析: 权重被划分为通道级块 以获得足够的统计支持。通过计算 KL 散度,PLAIN 能够量化不同位宽(INT4 或 INT8)下权重量化引入的误差。这种方法避免了重训练的需要,同时实现了细粒度的混合精度。
b) 简化激活熵 (Simplify Activation Entropy): 由于在运行时进行激活值分布拟合不切实际,PLAIN 引入了一个基于数据范围缩放的简化代理来近似激活值的量化误差: 其中,
-
符号解释:
- : 第 个词元(token)的简化激活熵。
- : 原始浮点激活张量 的绝对平均值经过
sigmoid函数处理。 - : 仿射因子(affine factor),表示原始浮点激活值范围与量化整数激活值范围的比值。
- 和 : 第 个词元的原始浮点激活值的最大值和最小值。
- 和 : 第 个词元的量化整数激活值的最大值和最小值。
-
目的分析: 这个操作可以由硬件在结果收集期间高效支持。激活值使用逐词元粒度,因此 表示词元数量。这个简化熵作为激活量化误差的代理,避免了复杂的运行时分布拟合。
c) 硬件感知配置搜索 (Hardware-aware configuration Searching): 较低的位宽会加剧量化效应,使熵更大。为了考虑硬件成本,PLAIN 定义了一个字节级(byte-level)成本函数:
-
符号解释:
- : 总的硬件成本。
- : 权重的硬件成本。
- : 激活值的硬件成本。
- : 权重张量的大小(例如,参数数量)。
- : 权重张量量化后的位宽。
- : 激活张量的大小。
- : 激活张量量化后的位宽。
-
目的分析: 权重可以离线加载到 PIM 中,因此只对计算成本进行惩罚。激活值需要进行 HBM 间通信,因此其通信成本也应被惩罚(在实际实现中,可能会以平方项的形式加入通信成本,以反映其更大的影响)。
最终的位宽选择通过最小化以下损失函数来确定:
-
符号解释:
- : 混合精度量化的总损失函数。
Entropy: 综合了权重精度熵和简化激活熵。- : 一个因子,用于调整熵和成本之间的相对权重。
-
目的分析: 该算法通过比较 INT4 和 INT8 的损失来选择最佳位宽,旨在在量化误差和硬件成本之间找到最佳平衡点。
4.2.1.3. PLAIN 推理过程 (Inference Process in PLAIN)
Figure 4 展示了 PLAIN 中两个词元(token)的推理数据流。

该图像是PLAIN执行数据流的示意图,展示了Weight-Activation计算(QKV投影)和Activation-Activation计算(OK矩阵乘法)的流程。图中包含多个Bank-PIMs和数据流的连接,显示了计量化、Tokens、缓冲区及复杂计算的结构。
Fig. 4: The execution dataflow of PLAIN.
推理过程步骤如下:
- 权重加载: 每个 PIM 单元首先加载当前层的权重。
- 词元量化与调度: 词元被混合精度量化。例如,
token0使用 INT4 精度,而token1使用 INT8 精度,需要更多的计算资源。为了平衡工作负载,token1被拆分并分派到三个 PIM 组进行并行计算。 - 计算: 计算使用 4-bit 乘法器和加法树执行,生成 INT32 格式的输出,以安全支持累加和部分和融合(partial sum fusion)而不会溢出。
- QKV 矩阵量化: 随后,查询(Query, Q)、键(Key, K)和值(Value, V)矩阵被单独量化。由于 GEMM(通用矩阵乘法)的布局特性,V 矩阵沿不同维度进行量化。
- 激活计算: Q 和 K 矩阵被加载到 PIM 中进行激活计算。
- 结果处理与反量化: Q 和 K 矩阵计算的结果与 V 矩阵相乘,然后进行最终的反量化(dequantization)步骤。
- Softmax 操作: 在 Stack-PIM 上完成 Softmax 操作后,对 S(注意力分数)和 v(加权值)的计算以相同的方式进行。
- 多阶段量化-反量化: 每个注意力层经历三个量化-反量化阶段。PLAIN 的设计确保这些操作以最小的开销得到支持。
4.2.2. PLAIN 硬件架构 (PLAIN Hardware Architecture)
Figure 5 展示了 PLAIN 的硬件架构,它被集成到 HBM(高带宽内存)堆栈中,与 XPU(例如 GPU 或 NPU)形成一个异构系统。

该图像是PLAIN硬件架构示意图,展示了DRAM-PIM系统中的各个组件,包括主机CPU、PIM模块、存储单元和调度器等,旨在优化深度学习推理过程中的计算效率与内存带宽。
Fig. 5: PLAIN hardware architecture.
-
异构系统集成:
- 系统包含主机 CPU、XPU、多个 HBM 堆栈以及支持 PLAIN 的内存。
- 当 XPU 处理非注意力层时,PLAIN 作为传统的 HBM 内存运行,避免了通信开销。中间结果被写回内存,使 GPU 能够重叠计算与内存操作,提高请求级并行度。
-
HBM 扩展与 PIM 单元:
- 为了实现 LLM 中快速的注意力推理,PLAIN 扩展了标准的 HBM,使其具备内存内计算能力。
- 每个 HBM 堆栈由八个 3D 堆叠的 DRAM 裸片(die)和一个缓冲裸片(buffer die)通过 TSV(硅通孔)连接。
- PLAIN 的设计在最小化 DRAM 更改的同时,增加了计算逻辑,并支持 DRAM 和 PIM 两种模式。它引入了两种类型的 PIM 单元:
Bank-PIM和Stack-PIM。
-
Bank-PIM 单元:
- 位置: Bank-PIM 单元放置在 DRAM Bank 之间,并通过 I/O 边界连接(参见 Figure 5 右侧)。
- 组成: 每个 Bank-PIM 包含 -bit 寄存器文件(register files)、4-bit 乘法器、一个加法树(adder tree)和一个 32-bit 输出寄存器。
- 数据流: 数据通过 512-bit 总线从奇/偶 Bank 行缓冲区(row buffers)读取,结果通过专用的结果总线发送到缓冲裸片。
- 功能: Bank-PIM 执行量化后的 GEMM/GEMV 操作,用于注意力层的三个阶段:(1) QKV 投影(
tokenweight),(2) 注意力分数(query),以及 (3) 注意力上下文(scorevalue)。 - 并行性: 每个 PIM 从不同的 Bank 读取操作数,以避免访问冲突。计算调度基于矩阵乘法中的第一个操作数。
-
Stack-PIM 单元:
- 位置: Stack-PIM 位于缓冲裸片上(每个 HBM 堆栈一个)。
- 功能: 管理量化、反量化、Softmax、数据累加和调度。
- 组成: 包含用于缩放因子(scaling factors)和调度表(scheduling tables)的缓冲区,以及 Softmax 单元、量化单元和反量化单元。
- 数据处理: 反量化单元对累加结果进行尺度缩放,并将其转发给 Softmax 或量化单元。量化单元应用 PLAIN 的量化算法,并将输出分派给 Bank-PIMs 或将缩放因子存储在缓冲区内存中。Softmax 使用量化输入的中间最大值进行。
- 控制与协同: 与 Bank-PIM 不同,Stack-PIM 集成了控制和计算功能。主机 CPU 通过 DRAM 控制器使用调度表协调执行。Stack-PIM 在 Bank-PIM 计算之前处理量化,并在结果累加之后处理反量化。Softmax 步骤在注意力分数阶段之后执行,以最小的开销完成一个注意力传递。
4.2.3. 调度设计 (Schedule Design)
PLAIN 单元的调度算法旨在最大限度地提高计算资源利用率和系统吞吐量。通过利用 DRAM 固有的 Bank 级并行性,目标是饱和内部带宽并最小化执行停顿。
4.2.3.1. 位拆分词元 (Bit-wise splitting token)
理想情况下,所有 Bank-PIM 单元应同时完成任务。然而,混合精度量化(例如 INT8 与 INT4 词元)引入了不均匀的计算负载,可能导致资源利用率低下。PLAIN 通过采用位拆分(bit-wise splitting) 策略来解决此问题,以统一跨精度级别的计算模型。
如图 Figure 6 所示:

该图像是示意图,展示了PLAIN方法中将8位权重和激活向量进行位分割的过程。通过将8位向量拆分为两个4位向量并分布到多个PIMs中,经过乘法和移位操作后,将结果累加至32位加法器输出。
Fig. 6: Assuming weight and activation vectors are INT8, PLAIN performs a bit-wise split to convert the INT8 vector into two INT4 vectors, distributing them across multiple PIMs. After calculations, Stack-PIM shifts and accumulates the results in the buffer die.
- INT8 拆分: 8-bit 激活值被拆分为高 4-bit 和低 4-bit 段。
- 并行处理: 每个 4-bit 段由共享的 4-bit 乘法器独立处理,并分布到多个 PIMs。
- 统一计算模型: 尽管权重以通道粒度使用混合精度进行量化,但只需要少量的硬件支持(即移位器)即可支持 4-bit 和 8-bit 计算。
- 中间结果处理: 所有中间结果存储在 32-bit 寄存器中,允许安全地移位部分结果而不会溢出。最终输出在缓冲裸片上累加。
- 优势: 这种方法实现了 Bank-PIM 架构在不同精度级别上的统一,并保持了负载平衡。
4.2.3.2. Bank-PIM 调度表 (Bank-PIM schedule table)
高效的交错(interleaving)是 Bank-PIMs 和 Stack-PIM 之间在长读/写周期中保持高吞吐量的关键。同一组内 Bank 之间的冲突会显著降低带宽。为缓解此问题,PLAIN 引入了一个存储在缓冲裸片上的调度表,并集成到 DRAM 控制器中。
Figure 7 展示了 Bank-PIM 调度表的结构:
| Bank ID | PIM ID | Group ID | Token ID | Token Type | Scaling Factor |
| 0-1 | 0 | 1 | 1 | INT8-HIGH | #FLOAT 16 |
| 2-3 | 8 | 4 | 4 | INT4 | #FLOAT 16 |
| * | ** | ** | * | ** | *** |
| 128-129 | 64 | 1 | 1 | INT8-LOW | #FLOAT 16 |
| *** | * | *** | *** | * | ** |
| 1022-1023 | 511 | NONE | NONE | NONE | NONE |
Fig. 7: Bank-PIM schedule table stored on the buffer die. Continuous arrangement of bank addresses reduces scaling factor lookup time complexity to O(1) for quantization units.
- 表结构: 调度表按 Bank ID 排序,Bank ID 顺序地从上到下、从左到右索引 Bank,不进行交错。
- 共享顺序: 这种顺序也被累加和缩放因子缓冲区共享,使得反量化单元无需索引查找即可访问相应的参数。
- 动态分配:
PIM ID代表 Bank 的交错顺序,而Group ID在运行时动态分配。 - 记录信息: 表中还记录了词元位宽(
Token Type),用于准确地累加结果。
4.2.3.3. Bank-PIM 调度算法 (Bank-PIM schedule algorithm)
PLAIN 的调度算法旨在平衡 Bank-PIMs 之间工作负载,同时最小化与 Stack-PIM 的通信。
- 输出行分配: 每个 Bank-PIM 被分配一个完整或部分输出行,以避免 Stack-PIM 中的额外累加。根据 GEMM/GEMV 原理,每个 Bank-PIM 至少处理一个词元。
- 并行度提升: 为了增加并行度,每个 INT8 词元被拆分为两个 INT4 词元,并分发到不同的 Bank-PIMs。然而,仅词元级分区不足,尤其是在输入长度较短的生成阶段。
- 贪婪调度策略: 为此,PLAIN 提出了一种贪婪调度策略,进一步分区计算矩阵。词元被拆分,Bank-PIMs 根据 4-bit 词元的数量进行分组。每个组通过通道级分区(channel-wise partitioning)共同计算一个词元。这种方法确保了 PIM 之间的均衡利用和高效映射。
- QKV 生成阶段: 权重预加载到所有 Bank-PIMs。虽然分发工作会增加激活通信,但它显著降低了计算延迟。
- 分数和上下文阶段: 在分数和上下文阶段,通信(特别是 K 和 V)占主导地位,但 PLAIN 的策略仍然受益于减少每 Bank 的计算时间以及重叠技术。
- LLM 生成优化: 对于 LLM 生成,序列长度通常为 12。由于 KV 缓存的存在,计算开销得以降低。K 矩阵的重新分发是一次性的,而 V 矩阵需要重新量化和动态分发。PLAIN 的重叠方案(详见 Section 4.2.4)缓解了这一成本,从而提高了性能。
4.2.4. 重叠与并行 (Overlapping And Parallelism)
虽然调度算法充分利用了 Bank 级并行性,但如果 Bank-PIMs 和 Stack-PIM 之间顺序执行,协调仍将是瓶颈。幸运的是,不同组件的工作负载通常是独立的,并且可以在没有数据依赖或硬件冲突的情况下重叠执行。PLAIN 通过以下机制利用这种并行性:
-
数据独立性: 不同组件处理不相交的数据子集。
-
命令隔离: DRAM 和 PIM 命令避免在 C/A 总线上发生冲突。
-
独占执行: 每个 Bank 在任何时刻都仅在 DRAM 模式或 PIM 模式下独占运行。
Figure 8 展示了 PLAIN 系统中不同组件的重叠时间线。
该图像是图表,展示了 PLAIN 系统中不同组件的重叠时间线。包括了加权加载、量化与通信重叠,以及 Softmax 与通信的实现,旨在优化大语言模型的推理效率。
Fig. 8: Overlapping timeline of different components in PLAIN system. B-PIM: Bank-PIM; S-PIM: Stack-PIM; Com: communication; Cal: calculation; Soft: Softmax.
4.2.4.1. 权重加载重叠 (Weight Loading Overlap)
- 场景: 在 QKV 生成阶段,权重被预加载到每个 Bank-PIM 中的两个 Bank 之一。
- 问题: 随着模型规模的增长,权重容量可能超出单个 Bank,可能导致执行停滞。
- 解决方案: PLAIN 采用重叠方案。如图 Figure 8(a) 所示,当一个 Bank 进行计算时(例如 Bank 0),另一个 Bank(例如 Bank 1)加载权重,反之亦然。PLIAN 将 GEMM 分解为 GEMV,以重用本地寄存器文件中的缓存词元,从而为权重传输创造空闲时间。
- Ping-Pong 缓冲区: 采用 Ping-Pong 缓冲区方案,在层与层之间交替 Bank 的角色:奇数 Bank 进入 DRAM 模式,而偶数 Bank 进入 PIM 模式。
- 大规模模型: 对于大型模型,权重分布到多个 Bank,并将 Bank-PIMs 分组为虚拟 PIM,以扩展此方法同时保持 Ping-Pong 效率。
4.2.4.2. 量化和通信重叠 (Quantization and Communication Overlap)
- 场景: 在注意力分数和上下文阶段之前,激活值在缓冲裸片上聚合,Stack-PIM 在此执行反/量化操作。
- 问题: 这些操作是计算密集型的,其开销可与全精度 GEMM 相媲美。
- 解决方案: 为减少其延迟,PLAIN 将这些操作与通信和计算重叠。如图 Figure 8(b) 所示,Bank-PIMs 在计算过程中通过结果总线发送中间结果。Stack-PIM 以固定间隔收集这些结果,通过访问片上缓冲区中相应的缩放因子,立即应用反量化,无需额外的查找开销。同时,它并行更新量化单元(例如 min/max 值)。
- 效果: 除了最终结果传输,所有反/量化步骤都与正在进行的计算重叠,从而最大限度地减少了它们对关键路径的影响。
4.2.4.3. Softmax 和通信重叠 (Softmax and Communication Overlap)
- 场景: Softmax 操作在 Stack-PIM 上以全精度执行,位于注意力分数和上下文阶段之间。
- 问题: 这是一个关键的计算步骤。
- 解决方案: PLAIN 将此操作与量化和 KV 缓存通信重叠。在生成阶段,K 和 V 存储在 PIMs 中。由于 K 使用逐词元量化,它可以直接重用。然而,V 是逐通道量化,必须重新量化和重新分发。如图 Figure 8(c) 所示,Stack-PIM 在执行 Softmax 的同时开始量化 V。一旦 Softmax 完成,生成的注意力分数被发送到 Bank-PIMs。
- 效果: 这种重叠隐藏了 Softmax 和 V 量化的延迟,进一步提高了端到端吞吐量。
5. 实验设置
5.1. 数据集
实验使用了 WikiText-103 数据集。
- 来源与特点: 该数据集包含超过 1 亿个词元,这些词元是从经过验证的维基百科优秀和特色文章中提取的。
- 用途:
- 计算语言困惑度(PPL): 用于评估量化算法对解码器模型推理性能的影响。
- 文本生成任务: 为了弥补 PPL 计算(无生成阶段)与实际推理(有生成阶段)之间数据流的差异,数据集被分割成对话长度的片段,用作文本生成任务的提示。
5.2. 评估指标
对论文中出现的每一个评估指标,提供以下完整说明:
-
困惑度 (Perplexity, PPL)
- 概念定义: 困惑度是衡量语言模型好坏的一个重要指标。它量化了模型预测文本序列中下一个词元的难度。通俗来说,困惑度越低,表示模型对文本的预测能力越强,模型对语言的理解和生成能力越好。在信息论中,困惑度与交叉熵损失(Cross-Entropy Loss)密切相关。
- 数学公式:
困惑度可以由给定文本序列的平均负对数似然(即交叉熵损失)计算得出。对于一个包含 个词元的序列 ,其困惑度定义为:
等价地,通过序列的平均交叉熵损失
L(W)计算: 其中,平均交叉熵损失L(W)定义为: - 符号解释:
- : 一个包含 个词元的文本序列,表示为 。
- : 文本序列中词元的总数量。
- : 在给定前
i-1个词元 的条件下,模型预测第 个词元 的概率。 - : 自然对数的底数,约等于
2.71828。 - : 自然对数。
L(W): 整个序列的平均交叉熵损失。
-
加速比 (Speedup)
- 概念定义: 加速比是衡量一个优化方法相对于原始(或基线)方法性能提升的倍数。它通常被定义为基线方法完成任务所需时间与优化方法完成相同任务所需时间的比值。加速比越大,表示优化效果越显著。
- 数学公式:
- 符号解释:
- : 基线方法完成特定任务所需的总执行时间。
- : 提议的优化方法完成相同任务所需的总执行时间。
-
能耗 (Energy Consumption)
- 概念定义: 能耗是指系统在执行特定计算任务期间所消耗的电能总量。它通常通过对系统的瞬时功率进行时间积分来获得,或者在功率恒定或平均功率已知的情况下,通过功率乘以执行时间来估算。在论文中,通常会提供归一化的能耗,即相对于基线系统的能耗比例,以直观展示节能效果。能耗越低,表示系统在完成任务时越节能。
- 数学公式: 能量 的基本计算公式是功率 乘以时间 : 更一般地,如果功率随时间变化,则为功率对时间的积分: 归一化能耗通常表示为:
- 符号解释:
- : 总能耗,通常单位为焦耳 (J) 或瓦时 (Wh)。
- : 系统在任务执行期间的平均功耗(瓦特, W)。
P(t): 系统在时间 的瞬时功耗。- : 任务的执行时间(秒, s)。
- : 提议方法完成任务所需的能耗。
- : 基线方法完成任务所需的能耗。
-
面积开销 (Area Overhead)
- 概念定义: 面积开销衡量了在集成现有硬件系统时,为了实现新的功能(如 PLAIN 中的 PIM 计算单元和缓冲区)所额外占用的芯片物理面积。它通常以绝对面积值(如 )或相对于原有芯片面积的百分比来表示。面积开销是衡量硬件设计复杂性和制造成本的重要指标。更小的面积开销意味着更高的集成度、更低的制造成本以及更少的封装挑战。
- 数学公式: 面积开销通常直接以数值给出,或者以百分比形式计算:
- 符号解释:
- : 为实现新功能而新增的硬件组件的总面积。
- : 原有系统(或芯片)的总面积。
5.3. 对比基线
为了全面评估 PLAIN 的性能,论文将其与以下三种基线系统进行了比较:
- GPU:
- 描述: 代表传统的 GPU 推理方案。
- 配置: 使用 Huggingface 和 PyTorch 编译实际的 LLM 推理工作负载,并在 Nvidia A100 GPU 上测量端到端延迟,采用 FP16 (半精度浮点数) 精度。
- SmoothQuant [9]:
- 描述: 一种先进的后训练量化方法,旨在解决 LLM 激活值中的异常值问题。
- 配置: 将模型权重和激活值量化为 INT8 (W8A8),并利用 A100 GPU 的 Tensor Cores 进行低精度推理。自注意力层使用 Cutlass 库进行编译。
- AttAcc [8]:
- 描述: 一种异构 PIM 系统,旨在加速基于 Transformer 的生成模型推理。
- 配置: AttAcc 将 GPU 用于预填充(prefill)和 QKV(查询、键、值)生成阶段,而 PIM 用于解码(decoding)阶段的多头注意力(multi-head attention)计算。论文中将其硬件配置调整为与 PLAIN 相同的量级,以确保公平比较。
5.4. 实验配置
以下是 PLAIN 实验的详细配置,涵盖了 HBM 组织和 LLM 模型配置:
以下是原文 Table I 的结果:
| HBM Organization | |
| HBM Organization | 4 Banks per Bank-group, 4 BGs per Pseudo channel, 8 pCHs per die |
| HBM Timing Parameter | Frequency = 1GHz, tRP = 19, tRCD = 19, tRAS = 45, tRRDL = 4, tWR = 8, tCCD_S = 2, tCCD_L = 4, tREFI = 5070, tFAW = 39 |
| LLM configuration | |
| Model | Layers Hidden_size |
| GPT2-large | 36 1280 |
| GPT2-xl OPT-6.7b | 48 1600 32 4096 |
| OPT-13b | 40 5120 |
| LLaMA-2-7b | 32 4096 |
| LLaMA-2-13b | 40 5120 |
- HBM 组织:
- 每个 Bank 组包含 4 个 Bank。
- 每个伪通道(Pseudo channel)包含 4 个 Bank 组。
- 每个裸片(die)包含 8 个伪通道。
- HBM 时序参数:
- 频率(Frequency)= 1GHz
- tRP = 19 (行预充电时间)
- tRCD = 19 (行到列延迟)
- tRAS = 45 (行激活时间)
- tRRDL = 4 (不同 Bank 组之间行激活到行激活的延迟)
- tWR = 8 (写恢复时间)
- tCCD_S = 2 (同一 Bank 组内列到列延迟)
- tCCD_L = 4 (不同 Bank 组之间列到列延迟)
- tREFI = 5070 (刷新周期时间)
- tFAW = 39 (激活命令窗口)
- LLM 模型配置:
-
GPT2-large: 36 层,隐藏层大小 1280。
-
GPT2-xl OPT-6.7b: 48 层,隐藏层大小 1600;32 层,隐藏层大小 4096。
-
OPT-13b: 40 层,隐藏层大小 5120。
-
LLaMA-2-7b: 32 层,隐藏层大小 4096。
-
LLaMA-2-13b: 40 层,隐藏层大小 5120。
仿真环境: 论文通过修改 Ramulator2 [30] 创建了一个内部模拟器。实验分为两个阶段:首先,在 GPU 上编译 LLM 推理工作负载,生成经过 PLAIN 算法量化后的激活值精度轨迹。然后,该轨迹用于模拟内存访问,生成系统的周期精确推理时间。PLAIN 在标准 HBM3 中添加了算术单元和缓冲区,以计算注意力层并处理量化/反量化。这些组件的能耗和面积消耗通过 Verilog(用于算术单元)、Synopsys Design Compiler(用于综合)和 Cacti 7.0 [31](用于 DRAM 和缓冲裸片中的缓冲区)进行计算。此外,还考虑了 HBM3 操作(包括激活和读取)的功耗 [32, 33, 34]。
-
6. 实验结果与分析
6.1. 核心结果分析
6.1.1. 量化结果
论文通过条带滑动窗口(striped sliding window)方法评估了 PLAIN 算法在多个 LLM 上的表现,序列长度和步长均设置为 1024。
以下是原文 Table II 的结果:
| Model | GPT2-large GPT2-xl OPT-6.7b OPT-13b LLAMA2-7b LLAMA2-13b | |||||
| FP16 | 19.42 | 17.38 | 12.42 | 10.12 | 6.45 | 5.43 |
| W4A4 | 1164.16 | 983.13 | inf | inf | inf | 6388.79 |
| W8A8 | 19.71 | 18.29 | 30.27 | 4272.0 | 6.27 | 5.52 |
| SmoothQuant (W8A8) | - | - | 12.42 | 11.90 | 6.14 | 5.50 |
| SmoothQuant (W6A6) | - | 15.25 | 14.73 | 9.12 | 7.80 | |
| PLAIN (W6A6) | 20.56 | 18.20 | 13.48 | 9.95 | 6.91 | 5.92 |
- 小型模型表现: 对于像 GPT2-large 和 GPT2-xl 这样的小型模型,激活值中的异常值(outliers)很少见,因此
SmoothQuant并非必需。在这种情况下,朴素的 W8A8 量化已经达到了 FP16 级别的性能(例如,GPT2-large 的 FP16 为 19.42,W8A8 为 19.71)。PLAIN 在此基础上进一步降低了权重和激活值的位宽,但只带来了可忽略的性能下降(GPT2-large 的 PLAIN (W6A6) 为 20.56)。 - 大型模型挑战与 SmoothQuant 的必要性: 随着模型尺寸的增加(如 OPT、LLaMA 模型),激活值异常值变得突出。此时,单独使用 W8A8 量化会导致显著的精度损失(例如,OPT-13b 的 FP16 为 10.12,W8A8 猛增到 4272.0,LLaMA2-7b 的 W8A8 也略高于 FP16)。对于这些大型模型,
SmoothQuant变得至关重要,它能够将 W8A8 的性能恢复到接近 FP16 水平(例如,OPT-6.7b 的 SmoothQuant (W8A8) 为 12.42,与 FP16 的 12.42 完全一致)。 - PLAIN 的优势: PLAIN 算法在 W6A6 配置下,对于大部分模型都能保持与 FP16 相当甚至更好的困惑度(例如 OPT-13b,FP16 是 10.12,PLAIN (W6A6) 是 9.95),同时比
SmoothQuant (W6A6)表现更优(例如 OPT-6.7b,SmoothQuant (W6A6) 是 15.25,PLAIN (W6A6) 是 13.48)。这表明 PLAIN 能够有效地在更低的位宽下保持模型质量,超越了仅使用SmoothQuant的传统量化方法。W4A4 的困惑度普遍非常高,表明其精度损失过大,不适用于这些模型。
6.1.2. 延迟/加速比
Figure 9 展示了 PLAIN 相较于 GPU-FP16 的加速比,以及与 GPUSmoothQuant 和 AttAcc-PIM 的比较。实验使用批处理大小 1,输入词元长度 64,输出词元长度 64。

该图像是图表,展示了不同模型(如GPU-FP16、GPUSmoothQuant、AttAcc和PLAIN)在推理速度上的比较。数据显示,PLAIN在GPT2-xl模型中的加速比可达到5.03倍,表明其在大语言模型推理中的优势。
Fig. 9: Speedup comparison results of GPU-FP16, GPUSmoothQuant, AttAcc-PIM, PLAIN. We use batch size 1, input token length 64 and output token length 64.
- 相对于 GPU-FP16 的显著加速: PLAIN 在各种模型尺寸和类型上都实现了 4-5 倍的加速,平均加速比为 4.41 倍。这显著优于
GPUSmoothQuant(通常在 1.5-2.5 倍加速)和AttAcc-PIM(通常在 2.5-3.5 倍加速)。 - 加速原因分析:
- 数据位宽减少: 通过混合精度量化降低了数据位宽,从而减少了计算负载和数据传输量。
- 架构优化: PLAIN 的软硬件协同设计,特别是专门的 PIM 架构,最大限度地减少了权重移动和量化/反量化带来的开销。
- 高效调度和重叠: PLAIN 精心设计的调度器和重叠机制有效地隐藏了数据传输和前处理/后处理的延迟,确保了 PIM 内部和与 GPU 之间的高效并行执行。
- 与其他 PIM 方案的对比: PLAIN 甚至优于最先进的
AttAcc-PIM,例如在 GPT2-xl 模型上实现了 5.03 倍的加速,而AttAcc-PIM约为 3.2 倍。这表明 PLAIN 的协同设计在处理 LLM 推理方面更有效。
6.1.3. 能耗
Figure 10 展示了 PLAIN 和 AttAcc 相较于 GPU 的归一化能耗。实验配置与加速比实验相同。

该图像是一个柱状图,展示了GPU、AttAcc-PIM和PLAIN在不同模型(如GPT2-large、OPT-6.7B等)中归一化能耗的比较。数据表明PLAIN在特定模型上能显著减少能耗,达到0.22倍。
Fig. 10: Normalized energy consumption results of GPU, AttAcc and PLAIN. Experiment configuration is the same as speedup results.
- AttAcc 的能耗降低:
AttAcc在相同输入和输出长度下显示出持续的能耗降低(大约在 0.4-0.6 倍之间),这得益于 PIM 减少了数据移动。 - PLAIN 的显著优势: PLAIN 实现了显著更大的能耗降低,对于大多数模型能耗降至 GPU 的 0.2-0.3 倍(例如 LLaMA2-7b 达到 0.22 倍)。
- 能耗降低原因分析:
- 低位宽计算: PLAIN 通过对权重和激活值使用更低的位宽,最大限度地减少了计算量,同时抵消了量化带来的开销。这直接减少了执行每个操作所需的能量。
- 卸载 QKV 生成: PLAIN 将 QKV 生成卸载到 PIM,显著减少了权重数据在 GPU 和内存之间的移动,从而降低了通信能耗和总体的系统能耗。 PLAIN 在能耗方面的优异表现进一步证实了其软硬件协同设计在实现高效 LLM 推理方面的有效性。
6.2. 数据呈现 (表格)
本节已在“5. 实验设置”和“6.1. 核心结果分析”中完整转录了原文的 Table I 和 Table II。
6.3. 消融实验/参数分析
6.3.1. 消融研究 (Ablation Study)
Figure 11 展示了在 OPT-6.7B 模型上(输入/输出词元长度均为 64),移除 PLAIN 各个优化组件后的加速比,以 GPU-FP16 为基线。

该图像是图表,展示了在OPT-6.7B模型下不同配置的加速效果比较。OPT-MixQ、OPT-kcache、OPT-kcache-balance以及OPT-kcache-balance-overlap的加速性能分别为不同的数值,其中OPT-kcache-balance-overlap实现了0.85倍和1.56倍的加速。
Fig. 11: PLAIN's ablation study in OPT-6.7B model. kcache: key vectors cache; balance: split a INT8 token to 2 INT4 token; overlap: overlapping quantization and softmax time with communication time
-
基础 PIM 架构加速: 如果仅使用 PIM 架构本身的 Bank 级并行性进行混合精度计算,其加速比约为 1.62 倍。这表明 PIM 架构本身提供了基础的性能提升。
-
重叠优化(overlap): 移除重叠优化后(包括量化与通信重叠、权重加载与计算重叠、Softmax 与值通信重叠),加速比从约 4.5 倍大幅下降到 1.62 倍。这表明重叠优化对性能提升贡献巨大,若无这些优化,PIM 的收益将大打折扣,因为数据传输和前后处理开销会成为主要瓶制。
-
工作负载平衡(balance): 在重叠优化基础上,如果移除工作负载平衡(即不将 INT8 激活值拆分为两个 INT4 词元),加速比从 4.5 倍下降到 2.95 倍(约 1.55 倍的下降)。这突出强调了工作负载平衡的重要性。如果没有平衡机制,INT4 PIM 单元可能在 INT8 任务执行时处于空闲状态,导致资源利用率低下。
-
Kcache 优化: Kcache 优化(在预填充阶段预加载键向量到 Bank-PIMs)在所有优化都存在的情况下,带来了额外的 0.2 倍加速。这表明缓存重用对性能也有积极影响,尤其是在解码阶段。
结论: PLAIN 的各项优化组件都是至关重要的。特别是重叠机制和工作负载平衡策略,对将 PIM 的理论优势转化为实际的端到端加速贡献最大。
6.3.2. 敏感性分析 (Sensitivity Analysis)
Figure 12 展示了 PLAIN 在不同输出词元长度下的加速比。

该图像是图表,展示了不同新令牌数量下,PLAIN 相比于其他方法(如 GPU-FP16、GPU-smoothquant 和 AttAcc)的加速效果。其中,PLAIN 在 64 和 128 个最大新令牌时,分别取得了 4.11 倍和 4.91 倍的提升。
Fig. 12: PLAIN's speedup over different tokens.
-
短输出长度(16 词元): 对于短输出(例如 16 词元),预填充阶段(prefill phase)占据主导地位。预填充通常是计算密集型(compute-bound)的,因为需要处理整个输入序列。在这种情况下,PIM 的内存带宽优势不如 GPU 的峰值 FLOPS 优势明显,因此加速效果有限。
-
长输出长度(64 或 128 词元): 随着输出长度的增加(例如 64 或 128 词元),解码阶段(decoding phase)变得更为重要。解码阶段涉及许多注意力层中的 GEMV 操作。这些操作通常是内存密集型(memory-bound)的(算术强度约为 1),因为每次生成一个新词元都需要访问大量缓存的键值对。PIM 处理内存密集型任务的优势在此场景下得到了充分发挥,从而显著提升了 PLAIN 的性能。例如,在 128 词元时,PLAIN 实现了约 4.91 倍的加速。
结论: PLAIN 的性能提升在内存密集型的解码阶段更为显著,这与其利用 PIM 高内部带宽的设计目标相符。
6.3.3. 扩展性分析 (Scalability Analysis)
Figure 13 展示了随着 HBM 堆栈数量增加,PLAIN 的可扩展性。

该图像是一个条形图,展示了随着堆叠数量(#Stack)增加,PLAIN的速度提升(Speedup)。从图中可以看出,堆叠数量从1增加到4时,速度提升逐渐增加,显现更多存储带宽带来的性能提升。
Fig. 13: The speedup of PLAIN with increasing stacks.
-
高效工作负载分布: PLAIN 能够将注意力层的工作负载高效地分布到多个 HBM 堆栈中,而不会增加堆栈间的通信开销。这得益于其优化的激活数据映射策略。
-
Bank-PIMs 与通信开销: 增加 Bank-PIMs 的数量会增加 Bank 到裸片(die)的通信开销。
-
近线性性能提升: 尽管存在通信开销,但由于 FLOPS(每秒浮点运算次数)的显著提升,PLAIN 实现了近乎线性的性能改进。这表明 PLAIN 能够有效地利用更多的 PIM 资源,扩展性良好。
-
硬件配置灵活性: PLAIN 的硬件配置具有灵活性,可以通过调整 HBM 堆栈内 DRAM 和 PIM 裸片的数量来调整容量、带宽和 FLOPS。
结论: PLAIN 具有良好的可扩展性,能够有效利用增加的 HBM 堆栈资源来进一步提升 LLM 推理性能。
6.4. 面积开销 (Area overhead)
PLAIN 在 HBM 堆栈内的总面积开销是:
- 每 DRAM 裸片(die)。
- 每缓冲裸片(buffer die)。 这导致每 DRAM 裸片额外增加了 6.93% 的面积。
具体组件的面积成本:
-
每个 DRAM 裸片包含 64 个 GEMV 单元,每个 GEMV 单元占用 (基于 1z-nm DRAM 工艺 [21])。
-
缓冲裸片包含量化单元、Softmax 单元、累加器和缓冲区:
- 量化单元:
- Softmax 单元:
- 累加器:
- 缓冲区:
-
缓冲裸片上的 Softmax 单元和累加器面积参考了先前的工作,所有处理单元都按 7nm 工艺 [36] 进行缩放。
结论: PLAIN 的面积开销相对较小,显示了其在硬件实现上的高效性和可行性。
7. 总结与思考
7.1. 结论总结
本文介绍了 PLAIN,一个创新的软硬件协同设计框架,旨在通过混合精度量化和 PIM(内存内计算)技术优化大型语言模型(LLMs)的推理过程。PLAIN 的核心贡献在于其对算法和硬件的深度集成:
-
硬件高效量化算法: 提出了一个新颖的量化算法,该算法能够根据参数的局部性分布和硬件特性,智能地确定每层的最佳精度,从而优化硬件映射,并有效地在不显著牺牲模型质量的情况下实现低位宽(如 W6A6)量化。
-
PIM 优化的异构架构: 设计了 PLAIN 架构,通过引入 Bank-PIM 和 Stack-PIM 单元,将 LLM 中内存带宽密集型的注意力层卸载到 HBM 内部,利用其高内部带宽进行加速,同时保留 GPU 处理计算密集型全连接层,实现了异构系统的高效协同。
-
精细化调度与重叠机制: 集成了工作负载感知的数据流调度器,通过位拆分词元(bit-wise splitting token)策略平衡混合精度操作的负载,并设计了多级重叠机制(包括权重加载、量化/反量化、Softmax 与通信和计算的重叠),最大限度地隐藏了各种开销,确保了 PIM 内部和 PIM 与 GPU 之间的高效并行执行。
实验评估表明,PLAIN 在性能、能耗和模型质量方面均优于现有方案。相对于传统的 GPU(FP16)和最先进的 PIM 加速器(AttAcc),PLAIN 实现了 5.03 倍和 1.69 倍的性能提升,同时能耗显著降低,且模型质量损失可忽略不计。消融研究和敏感性分析也验证了 PLAIN 各个优化组件的有效性以及其在不同场景下的适应性。
7.2. 局限性与未来工作
论文中未明确指出 PLAIN 的具体局限性,但根据其方法和实验,可以推断以下几点潜在局限性:
-
通用性限制: PLAIN 的设计是针对 DRAM-PIM,特别是 HBM-PIM 架构的。其特定的 Bank-PIM 和 Stack-PIM 单元以及调度策略可能不直接适用于其他类型的 PIM 架构(例如基于电阻式内存 RRAM-PIM 或其他非 DRAM PIM),这限制了其通用性。
-
编程模型复杂性: 软硬件协同设计虽然带来性能优势,但也意味着更复杂的编程模型和开发栈。为 PLAIN 编写高效的应用程序或将其集成到现有 ML 框架中可能需要专门的工具链和专业知识。
-
仅关注注意力层: 论文主要强调了注意力层在 PIM 上的加速。虽然全连接层仍在 GPU 上执行,但 LLM 中可能还存在其他类型的层或操作(例如激活函数),PIM 对这些操作的加速潜力或挑战未被详细探讨。
-
量化精度上限: 论文主要在 INT4/INT8 和 W6A6 范围内探索混合精度。随着更低位宽(如 INT2 或 INT1)量化技术的发展,PLAIN 是否能继续保持其优势,以及如何在更低的位宽下平衡精度和效率,可能需要进一步研究。
-
模型大小限制: 尽管 PLAIN 旨在加速大型 LLM,但其在内存内的计算能力仍然受限于 PIM 内部的有限资源。对于未来可能出现的万亿级参数模型,单栈 PIM 或多栈 PIM 的扩展性是否能持续满足需求仍需验证。
基于这些潜在局限性,未来工作可以包括:
-
扩展到更广范围的 PIM 架构: 探索 PLAIN 的设计原则如何应用于其他新兴的 PIM 硬件,提高其通用性。
-
更细粒度的量化和动态调整: 开发更高级的运行时量化策略,根据实际输入数据动态调整精度,以进一步提高效率。
-
支持更多 LLM 操作类型: 扩展 PIM 单元的计算能力,使其能够更高效地处理 LLM 中的其他复杂操作,甚至部分全连接层。
-
软件生态系统建设: 开发更友好的编程接口和编译器支持,降低 PLAIN 的使用门槛,促进其在实际应用中的部署。
7.3. 个人启发与批判
个人启发:
- 软硬件协同设计的力量: PLAIN 再次证明了在特定领域(如 LLM 推理)中,仅有软件优化或硬件加速是不足的,真正的突破往往来自于软硬件的深度协同设计。针对硬件特性设计算法,并为算法需求定制硬件功能,是实现极致性能的关键。
- PIM 在 LLM 中的巨大潜力: LLM 面临的“内存墙”问题日益突出,PIM 作为一种内存中心计算范式,其高内部带宽和减少数据移动的特性使其成为 LLM 推理加速的天然候选。PLAIN 的成功进一步验证了 PIM 在解决 LLM 瓶颈方面的巨大潜力。
- 精细化优化的价值: PLAIN 的调度设计和多级重叠机制是其成功的关键。这些精细到比特级和Bank级的优化,以及对各种开销的巧妙隐藏,使得理论上的收益能够转化为实际的端到端加速。这提醒我们在系统设计中,细节的把控至关重要。
- 混合精度量化的发展方向: 传统的混合精度量化更多关注层级或通道级。PLAIN 的硬件感知量化算法和位拆分策略,为混合精度量化提供了更精细、更实用的思路,尤其是在异构计算环境中。
批判:
- 架构特定性带来的部署挑战: 尽管 PLAIN 在性能上表现出色,但其高度定制化的 Bank-PIM 和 Stack-PIM 架构,以及对 HBM 内部结构的修改,可能使其在商业化部署和推广方面面临挑战。除非 PIM 技术成为主流并形成统一标准,否则这种高度定制化的解决方案难以被广泛采用。
- 量化精度的鲁棒性: 论文虽然展示了 PLAIN 在 W6A6 精度下的优异表现,但对于更低位宽(如 INT4)或在更广泛、多样化的数据集上,其量化策略的鲁棒性如何,仍需进一步验证。LLM 的量化通常非常敏感,微小的误差累积可能导致灾难性的结果。
- 软件栈和编程模型的复杂性: 软硬件协同设计通常伴随着软件开发复杂度的增加。论文中并未详细描述 PLAIN 的编程模型、编译器支持或运行时环境。如何让开发者能够方便、高效地利用 PLAIN 的硬件能力,将是其从研究走向实际应用的关键挑战。手动调度和优化可能难以扩展到更多模型和场景。
- 实际功耗和热管理问题: PIM 将计算单元放置在内存中,虽然减少了数据移动能耗,但也可能在内存芯片内部产生局部热点(hotspot)。论文虽然提供了能耗数据,但并未深入探讨实际的热管理挑战及其对性能和可靠性的潜在影响。
- 缺乏与更广泛 PIM 硬件的比较: 论文主要与 AttAcc 进行了 PIM 方面的比较。如果能与更多不同类型(如逻辑层 PIM 或电阻式 PIM)或更通用的 PIM 硬件进行对比,将更能凸显 PLAIN 在特定场景下的优势和不足。
相似论文推荐
基于向量语义检索推荐的相关论文。