Kelle: Co-design KV Caching and eDRAM for Efficient LLM Serving in Edge Computing

作者/机构: Tianhua Xia (纽约大学 Tandon 工程学院), Sai Qian Zhang (纽约大学 Tandon 工程学院)

A1 主要贡献

大型语言模型 (LLM) 在边缘设备上的运行对于降低延迟、改善实时处理和增强隐私至关重要。然而，在边缘设备上实现 LLM 面临挑战，尤其是在管理键值 (KV) 缓存方面。KV 缓存机制通过存储和重用先前计算的键 (Key) 和值 (Value) 向量，避免了对早期词元 (token) 的重复计算，从而在注意力计算中发挥着关键作用。但随着输入文本变长，KV 缓存的大小与序列长度成线性关系增长，导致显著的内存占用和数据访问成本，这在内存和计算能力有限的边缘设备上尤为棘手。

为解决此问题，本文探索使用嵌入式 DRAM (eDRAM) 作为边缘设备上 LLM 服务的主要片上存储介质。相较于 SRAM，eDRAM 具有更高的存储密度和更低的漏电功耗，但其主要缺点是需要周期性刷新以防止数据丢失，而刷新操作会消耗大量能量。

为了应对集成 eDRAM 的挑战，并充分利用其优势，本文提出了 Kelle，一个软硬件协同设计解决方案，专为在基于 eDRAM 的边缘系统上部署 LLM 而优化。核心贡献如下：

  • 提出 Kelle 算法-系统协同设计方案：专为基于 eDRAM 的边缘系统上的设备内 LLM 服务而设计。为了优化 eDRAM 的集成成本并提高 LLM 执行效率，引入了基于注意力的驱逐和重计算策略 (AERP) 以及二维自适应刷新策略 (2DRP)，以实现高效的 KV 缓存管理（第 4.1 和 4.2 节）。
  • 设计 Kelle 加速器：该加速器利用 eDRAM 作为主要的片上存储，并采用定制的内存布局。为了最大化效率，该加速器集成了专用的 eDRAM 控制器和脉动驱逐器 (systolic evictor)，以高效实现 AERP 和 2DRP 策略（第 5 节）。
  • 引入 Kelle 调度器：采用一种高效的计算模式来优化 eDRAM 数据的生命周期和 LLM 服务延迟，从而显著减少 eDRAM 的刷新能耗和内存流量（第 6 节）。
  • 评估结果：实验证明，与现有的基线解决方案相比，Kelle 加速器实现了 3.9 倍的加速和 4.5 倍的能效提升，同时对 LLM 的准确性影响可以忽略不计（第 7、8 节）。


图 1: (a) LLM 词元生成。(b) 用于存储中间数据的 KV 缓存，其中 N 表示词元索引。

A3 背景知识与关键洞察

2.1 LLM 工作流程

现代 LLM 的结构与服务流程。现代 LLM（如 Llama 系列【索引 74, Llama: Open and efficient foundation language models+2023+arXiv; 索引 75, Llama 2: Open foundation and fine-tuned chat models+2023+arXiv】、GPT 系列【索引 11, Language models are few-shot learners+2020+arXiv; 索引 65, Language models are unsupervised multitask learners+2019+OpenAI blog】）由一堆 transformer 解码器构成，每个解码器包含一个自注意力 (Self-Attention, SA) 模块和一个前馈网络 (Feed Forward Network, FFN)。在 LLM 服务过程中，SA 模块的输入首先与三个权重矩阵 $W_Q$、$W_K$ 和 $W_V$ 相乘，分别产生查询 (Q)、键 (K) 和值 (V) 向量。然后，生成的 Q 和 K 与 V 结合，经过乘法、softmax 和残差加法，生成 SA 模块的输出。该输出随后被传递给 FFN 进行进一步处理，FFN 通常包含一个标准的多层感知机 (MLP)【索引 63, Improving language understanding by generative pretraining+2018】或门控 MLP【索引 45, Pay attention to mlps+2021+NeurIPS; 索引 74, Llama: Open and efficient foundation language models+2023+arXiv; 索引 75, Llama 2: Open foundation and fine-tuned chat models+2023+arXiv】。FFN 由多个全连接 (FC) 层和一个中间激活函数（如 GeLU【索引 31, Gaussian error linear units (gelus)+2016+arXiv】）组成。LLM 服务涉及两个主要阶段：预填充 (pre-filling) 和解码 (decoding)。在预填充阶段，模型并行处理上下文词元。在解码阶段，模型基于当前和之前的词元预测下一个词元。这是通过将当前输入与之前词元的信息（以其键和值向量表示）相结合来完成的。这个过程以自回归的方式重复进行（如图 1 (a) 所示）。


图 2: 一个 KV 向量计算的例子。

2.2 KV 缓存

解码阶段的 KV 缓存机制。在解码阶段，每个新生成的词元的 KV 向量被存储在 KV 缓存中，以提高生成速度，如图 1 (b) 所示。通过这样做，可以避免在每次生成新词元时为之前的词元重新计算向量。具体来说，为了产生一个 LLM 模块的输出，第 N 个词元的输入向量（长度为 C）与 $W_Q$、$W_K$ 和 $W_V$ 相乘，生成三个向量：查询 $q_N$、键 $k_N$ 和值 $v_N$，每个向量的维度为 $1 \times D$，其中 D 表示通道大小，如图 1 (b) 所示。$q_N$ 以及其他 KV 向量随后沿通道维度被分割成多个部分，每个部分的维度为 $1 \times D_h$，其中 H 表示头的数量。对于头 h，得到的向量分别表示为 $q_N^h$、$k_N^h$、$v_N^h$。然后，在头 h 的 KV 向量中，$q_N^h$ 与每个键向量 $k_n^h$（$1 \le n \le N$）进行点积运算，结果通过一个 softmax 函数，产生一个维度为 $1 \times N$ 的注意力分数向量，表示为 $A_N^h$。接下来，注意力分数向量用于与每个值向量 $v_n^h$（$1 \le n \le N$）计算点积，产生一个长度为 $D_h$ 的结果向量 $y_N^h$。$y_N^h$ 随后在多个头之间拼接，得到的向量 $y_N$ 再与 $W_O$ 相乘。这个过程如图 2 所示，并可以通过以下方程表示：

$$y_N^h = \sum_{1 \le n \le N} (A_{N,n}^h \cdot v_n^h)$$

KV 缓存的计算特性与压缩技术。其中 $A_{N,n}^h$ 代表 $A_N^h$ 的第 n 个元素。根据方程 1 和方程 2，我们注意到 KV 向量对 $[k_n^h, v_n^h]$ 的相对顺序不影响解码计算。换句话说，如果我们交换两对 KV 向量的值（例如，交换 $[k_1^h, v_1^h]$ 和 $[k_2^h, v_2^h]$），使用方程 1 和 2 产生的结果 $y_N^h$ 保持不变。KV 缓存压缩技术大致可分为两类：词元丢弃【索引 27, Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs+2023+arXiv; 索引 44, MiniCache: KV Cache Compression in Depth Dimension for Large Language Models+2024+arXiv; 索引 47, Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time+2023+arXiv; 索引 83, Efficient Streaming Language Models with Attention Sinks+2024+arXiv; 索引 88, PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference+2024+arXiv; 索引 98, H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models+2023+arXiv】和 KV 缓存量化【索引 32, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization+2024+arXiv; 索引 48, KIVI: A Tuning-Free Asymmetric 2-bit Quantization for KV Cache+2024+arXiv; 索引 82, DFRot: Achieving Outlier-Free and Massive Activation-Free for Rotated LLMs with Refined Rotation+2024+arXiv】。词元丢弃策略识别并永久丢弃不重要的词元，使其之后无法访问。例如，StreamLLM【索引 83】识别出对 LLM 性能至关重要的初始词元（sink tokens）并保留最近的词元。H2O【索引 98】识别具有高累积注意力分数的重击词元 (heavy hitter tokens)。KIVI【索引 49】通过通道分组实现 2 位非对称量化。QuaRot【索引 6, Quarot: Outlier-free 4-bit inference in rotated llms+2024+arXiv】利用零样本 Hadamard 变换来减少模型中的异常值，从而实现 4 位量化。推测解码【索引 33, Speculative decoding and beyond: An in-depth survey of techniques+2025+arXiv; 索引 34, DREAM: Drafting with Refined Target Features and Entropy-Adaptive Cross-Attention Fusion for Multimodal Speculative Decoding+2025+arXiv; 索引 42, Fast inference from transformers via speculative decoding+2023+ICML; 索引 51, PipeSpec: Breaking Stage Dependencies in Hierarchical LLM Decoding+2025+arXiv】是另一种加速 LLM 的推理技术，它使用一个轻量级草稿模型提出多个词元，然后由完整模型选择性地验证。Kelle 可以与推测解码技术正交地结合使用。

2.3 嵌入式 DRAM

eDRAM 的特性与挑战。各种 eDRAM 电路设计【索引 28, A 1-Mbit Fully Logic-Compatible 3T Gain-Cell Embedded DRAM in 16-nm FinFET+2020+IEEE Solid-State Circuits Letters; 索引 89, A logic-compatible eDRAM computein-memory with embedded ADCs for processing neural networks+2020+IEEE Transactions on Circuits and Systems I: Regular Papers】已成为 SRAM 的替代方案，其中一些仅需两个晶体管。其中，3T-eDRAM 表现突出，其密度是 SRAM 的两倍以上，静态功耗降低了 3.5 倍【索引 13, Technology comparison for large last-level caches (L3Cs): Low-leakage SRAM, low write-energy STT-RAM, and refresh-optimized eDRAM+2013+HPCA; 索引 17, A 3T Gain Cell Embedded DRAM Utilizing Preferential Boosting for High Density and Low Power On-Die Caches+2011+IEEE Journal of Solid-State Circuits】。表 1 展示了 3T-eDRAM 和 SRAM 的比较，结果是使用 Destiny【索引 60, DESTINY: A tool for modeling emerging 3D NVM and eDRAM caches+2015+DATE】在 65nm 技术节点上模拟的。eDRAM 的优势，包括更高的存储密度、更低的访问延迟和能耗，使其成为 LLM 实现的有吸引力的选择。然而，尽管 eDRAM 提供了多个优势，一个严重的缺点是需要周期性刷新以防止由电荷泄漏引起的数据损坏。因此，eDRAM 最适合存储大量瞬时数据，从而可以避免频繁刷新。

eDRAM 在加速器中的应用。已有研究探索在加速器系统中使用 eDRAM 来促进 CNN 计算【索引 15, DaDianNao: A Machine-Learning Supercomputer+2014+MICRO; 索引 54, St-DRC: Stretchable DRAM refresh controller with no parity-overhead error correction scheme for energy-efficient DNNs+2019+DAC; 索引 76, RANA: Towards Efficient Neural Acceleration with Refresh-Optimized Embedded DRAM+2018+ISCA; 索引 97, CAMEL: Co-Designing AI Models and Embedded DRAMs for Efficient On-Device Learning+2023+arXiv】，并提出了减轻刷新功耗开销的方法。DaDianNao【索引 15】将 eDRAM 分区以减轻刷新失败，但未解决刷新能耗或数据保留的挑战。RANA【索引 76】在 CNN 训练过程中注入位保留错误，以减轻低刷新频率导致的精度下降。CAMEL【索引 97】优化 CNN 模型架构以缩短训练期间的数据生命周期。虽然先前研究已证明 eDRAM 在 CNN 推理和训练中的效率，但其在 LLM 中的潜力尚未被探索。相比之下，Kelle 专注于利用 eDRAM 最小化 LLM 中 KV 缓存的片外内存访问，这是一个先前未被探索的领域。

Table 1: SRAM 和 eDRAM 的比较。

2.4 边缘 LLM 加速器

现有边缘 LLM 加速技术。为了在边缘设备上部署 LLM，一些研究提出了提高量化 transformer 准确性的方法【索引 23, An efficient hardware accelerator for sparse transformer neural networks+2022+ISCAS; 索引 30, OliVe: Accelerating Large Language Models via Hardware-friendly Outlier-Victim Pair Quantization+2023+ISCA; 索引 40, Tender: Accelerating Large Language Models via Tensor Decomposition and Runtime Requantization+2024+ISCA; 索引 46, COMET: Towards Partical W4A4KV4 LLMs Serving+2024+arXiv; 索引 50, Hardware accelerator for multi-head attention and position-wise feed-forward in the transformer+2020+SOCC; 索引 57, Energy-efficient neural network accelerator based on outlier-aware low-precision computation+2018+ISCA; 索引 85, DynaX: Sparse Attention Acceleration with Dynamic X:M Fine-Grained Structured Pruning+2025+ASPLOS; 索引 92, Gobo: Quantizing attention-based nlp models for low latency and energy efficient inference+2020+MICRO; 索引 100, Transpim: A memory-based acceleration via software-hardware co-design for transformer+2022+HPCA】。Tender【索引 40】通过将缩放因子设为 2 的幂次方，提出了一种硬件友好的 LLM 量化方法。COMET【索引 46】为 4 位 LLM 量化设计了高效的混合精度 GPU 内核。其他工作，如 FlexGen【索引 68, Flexgen: High-throughput generative inference of large language models with a single gpu+2023+ICML】、InfiniGen【索引 41, {InfiniGen}: Efficient generative inference of large language models with dynamic {KV} cache management+2024+OSDI】、InstInfer【索引 56, Instinfer: In-storage attention offloading for cost-effective long-context llm inference+2024+arXiv】、LLM.npu【索引 86, Fast On-device LLM Inference with NPUs+2024+arXiv】，探索了在片上单元和主存之间进行模型卸载的策略，以在资源受限的设备中高效部署 LLM。Cambricon-LLM【索引 90, Cambricon-llm: A chiplet-based hybrid architecture for on-device inference of 70b llm+2024+MICRO】提出了一种基于小芯片的混合架构，带有 NPU 和专用的 NAND 闪存芯片，以实现高效的设备上推理。

3.1 扩展片上内存的优缺点

片上内存与 LLM 性能。正如早期研究【索引 90, Cambricon-llm: A chiplet-based hybrid architecture for on-device inference of 70b llm+2024+MICRO; 索引 93, Llmcompass: Enabling efficient hardware design for large language model inference+2024+ISCA; 索引 99, ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching+2024+arXiv】所示，LLM 的服务速度受到片外内存带宽的显著限制。特别是在 LLM 的解码阶段，访问 KV 缓存构成了最关键的瓶颈【索引 41, {InfiniGen}: Efficient generative inference of large language models with dynamic {KV} cache management+2024+OSDI; 索引 98, H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models+2023+arXiv; 索引 99, ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching+2024+arXiv】。一个直接的方法是扩大片上 SRAM 的大小，这可以减少昂贵的片外内存访问并提升系统性能【索引 15, DaDianNao: A Machine-Learning Supercomputer+2014+MICRO; 索引 76, RANA: Towards Efficient Neural Acceleration with Refresh-Optimized Embedded DRAM+2018+ISCA】。为了说明这一点，我们评估了两个边缘计算系统的延迟，它们分别配置了 4MB 和 8MB 的 SRAM，运行 LLaMA2-7B 模型并处理不同序列长度。测试在一个模拟平台上进行，该平台具有一个 32 × 32 的脉动阵列用于 8 位 MAC 操作，以及 16GB DRAM 和 64GB/s 带宽，这反映了类似于 Google Coral 边缘设备【索引 72, Google Coral Edge TPU Board Vs NVIDIA Jetson Nano Dev board Hardware Comparison+2020】的边缘张量处理单元 (TPU)。如图 3 (a) 所示，将 SRAM 大小加倍平均带来了 1.27 倍的加速。然而，在评估平台中将 SRAM 容量从 4MB 扩大到 8MB，功耗和芯片面积分别增加了 29% 和 26%。考虑到边缘环境有限的面积和功耗预算，增加 SRAM 大小会减少可用于其他关键组件的资源，导致系统性能次优【索引 14, Deep learning with edge computing: A review+2019+Proc. IEEE; 索引 39, Memory requirements for balanced computer architectures+1986+ACM SIGARCH Computer Architecture News; 索引 79, Exploration of balanced design in resource-constrained edge device for efficient CNNs+2022+IEEE Transactions on Circuits and Systems II: Express Briefs】。因此，我们得出以下观察：

观察 1：更大的片上内存缓解了 LLM 中的 KV 缓存瓶颈，但在使用 SRAM 作为片上存储的边缘设备中带来了面积和功耗的代价。


图 3: (a) 具有 4MB 与 8MB SRAM 的边缘系统在不同模型和序列长度下的归一化延迟。(b) 具有 8MB eDRAM 和 8MB SRAM 的边缘系统的面积分解。红线是面积预算。(c) 集成 eDRAM 的边缘系统的能耗分解。显示的是预填充长度为 512 的解码长度。在解码期间，一个 8MB 的 eDRAM 用于存储部分层的 KV 缓存。报告的 DRAM 能耗包括模型权重访问和从 eDRAM 卸载 KV 缓存的能耗。

3.2 集成 eDRAM 的优缺点

eDRAM 作为 SRAM 的替代方案。为了在不增加面积的情况下增加片上内存大小，一种方法是用 eDRAM 替代 SRAM。根据表 1，eDRAM 不仅在相同面积下提供超过两倍的容量，而且访问和漏电能耗也更低。如图 3 (b) 所示，具有 8MB eDRAM 的评估系统比具有 8MB SRAM 的系统占用更少的面积，从而在更小的芯片尺寸内实现了更低的 LLM 服务延迟。大量研究【索引 5, Mosaic: Exploiting the spatial locality of process variation to reduce refresh energy in on-chip eDRAM modules+2014+HPCA; 索引 16, eDRAM-based tiered-reliability memory with applications to low-power frame buffers+2014+ISLPED; 索引 84, High performance Si-MoS2 heterogeneous embedded DRAM+2024+Nature Communications】和商业产品【索引 22, Intel eDRAM Attacks Graphics in Pre 3-D IC Days+2013+EE Times; 索引 24, The 12-core power8™ processor with 7.6 tb/s io bandwidth, integrated voltage regulation, and resonant clocking+2014+IEEE Journal of Solid-State Circuits; 索引 28, A 1-Mbit Fully Logic-Compatible 3T Gain-Cell Embedded DRAM in 16-nm FinFET+2020+IEEE Solid-State Circuits Letters; 索引 37, Embedded Flash and DRAM Components+2024+Kingston Technology; 索引 80, POWER7™, a highly parallel, scalable multi-core high end server processor+2010+IEEE Journal of Solid-State Circuits】证明了将 eDRAM 集成为主要片上存储介质的可行性。然而，其在边缘设备中惠及 LLM 服务的潜力尚未被探索。尽管 eDRAM 提供了几个优势，但先前的研究【索引 76, RANA: Towards Efficient Neural Acceleration with Refresh-Optimized Embedded DRAM+2018+ISCA; 索引 97, CAMEL: Co-Designing AI Models and Embedded DRAMs for Efficient On-Device Learning+2023+arXiv】表明，刷新操作可能成为整个系统能耗的一个重要瓶颈。此外，当 eDRAM 用于存储生命周期较长的数据时，不频繁的刷新会增加读出错误的风险，如图 4 所示。保留失败率表示为具有保留错误的比特百分比，随刷新间隔的变化而变化。为了说明这个问题，我们使用一个 8MB 的 eDRAM 替换了第 3.1 节中描述的系统中的 4MB SRAM。eDRAM 的刷新间隔设置为 45µs，以确保没有数据损坏。我们评估了 eDRAM 系统在不同模型和序列长度下的能耗。如图 3 (c) 所示，未经优化，eDRAM 刷新操作占总能耗高达 46%，导致平均能耗增加 1.7 倍。

观察 2：在相同的芯片面积下，eDRAM 在边缘设备的 LLM 服务中可以比 SRAM 带来延迟优势。然而，为了充分利用其功耗优势，必须大幅减少 eDRAM 的刷新操作。


图 4: 105°C 下 65nm eDRAM 保留失败分布【索引 38, Analysis of Retention Time Distribution of Embedded DRAM - A New Method to Characterize Across-Chip Threshold Voltage Variation+2008+IEEE International Test Conference】。
图 5: Kelle 加速器概览。

3.3 Kelle: KV 缓存与 eDRAM 的协同设计

Kelle 的核心设计策略。为了最小化 eDRAM 的能耗，三个有效的策略是：降低数据刷新频率、减少存储的数据量和缩短数据生命周期。为了使 eDRAM 能够增强边缘设备的 LLM 服务性能，我们提出了 Kelle，一个硬件和算法协同设计解决方案，旨在最小化 eDRAM 刷新能耗和高效管理 KV 缓存。

eDRAM 刷新控制。降低数据刷新频率可能会增加保留失败的风险，导致数据损坏。这引出了一个关键问题：LLM 在多大程度上可以容忍 KV 缓存中的数据损坏而不影响准确性？带着这个问题，我们协同设计了 eDRAM 内存布局和控制器，并采用了二维自适应刷新策略 (2DRP)，该策略设置了细粒度的动态刷新间隔，详见第 4.2 节。

KV 缓存驱逐。一个更小的 KV 缓存可以显著减少 eDRAM 的数据存储需求，从而降低刷新能耗并提高系统性能。先前的研究已经观察到，驱逐不重要的词元不会损害生成质量。然而，为了识别不重要的词元，先前的工作要么需要对序列进行性能分析【索引 27, Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs+2023+arXiv; 索引 47, Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time+2023+arXiv】，要么需要额外的计算【索引 83, Efficient Streaming Language Models with Attention Sinks+2024+arXiv; 索引 98, H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models+2023+arXiv】。为了高效地管理 KV 缓存，我们提出了一种新颖的脉动驱逐器架构，以加速基于注意力的驱逐和重计算策略 (AERP) 的操作，如第 4.1 节所述。

KV 向量重计算。随着序列长度的增长，KV 缓存的优势在某个阈值会减弱，因为访问片外内存的时间可能超过重新计算部分 KV 张量的时间。重计算与 eDRAM 存储瞬时数据的优势非常契合，如第 4.1 节所示。然而，重计算和存储之间的平衡需要考虑硬件特性进行仔细的调度。我们提出了 Kelle 调度器，通过设计计算模式来减少 KV 向量的数据生命周期，具体描述在第 6 节。

A2 方法细节

本节介绍 Kelle 框架中使用的高效算法，其概览如图 5 所示。在执行期间，Kelle 利用基于注意力的驱逐和重计算策略 (AERP) 以及二维自适应刷新策略 (2DRP) 来管理 eDRAM 操作，具体将在第 4.1 节和第 4.2 节中描述。

4.1 基于注意力的驱逐和重计算策略 (AERP)

我们首先讨论在解码阶段 eDRAM 容量达到上限时的驱逐策略。

驱逐策略。对于容量有限、最多能容纳 $N'$ 个词元的 KV 缓存，在解码阶段，当第 (N'+1) 个词元到达时，需要从已有的词元 $n$（其中 $1 \le n \le N'$）中驱逐其 KV 向量 $[k_n^h, v_n^h]$。要驱逐的第 h 个头和第 n 个词元的 KV 向量是根据其重要性 $s_n^h$ 来选择的，该重要性通过将该词元与 KV 缓存中所有其他词元的注意力分数（方程 1）相加计算得出，如下所示：

$$s_{n}^{h}=\sum_{1\le i\le n}A_{n,i}^{h}$$

驱逐过程示例。驱逐过程的一个例子如图 6 所示。假设 KV 缓存的预算是存储总共 $N' = 4$ 个向量。我们考虑一个有三个注意力头的案例。为清晰起见，我们只描绘了第一个头的计算，并省略了头的标记。当 $[k_5, v_5]$ 到达时，首先使用方程 3 计算重要性分数，如图 6 (a) 所示。然后，重要性分数最小的词元（第三个词元）的 KV 向量被驱逐，如图 6 (b) 所示。利用 $y_N^h$ 的计算不受 KV 向量相对顺序影响这一事实，可以通过按顺序读取 KV 缓存中的 KV 向量来计算方程 1 和方程 2，而无需关心它们原始的词元索引。值得注意的是，同一词元 n 的重要性分数 $s_n^h$ 在不同的注意力头之间可能会有所不同。因此，KV 向量的驱逐模式在这些头 h 之间也会不同。

预填充阶段的驱逐策略。对于上下文词元长度为 $N_{pref}$ 的预填充阶段，所有上下文词元都是并行处理的。对于每个层内的每个头，第 N 个词元的重要性分数计算为 $s_N^h = \sum_{1 \le i \le N_{pref}} A_{N,i}^h$。具有最高 $N'$ 个 $s_N^h$ 的词元将被保留在 KV 缓存中用于解码操作。

保留关键 token。除了具有最高 $s_N^h$ 分数的词元外，初始词元和最近的词元也会被保留，因为它们对模型性能有已证的影响，这一点由先前的工作【索引 83, Efficient Streaming Language Models with Attention Sinks+2024+arXiv; 索引 98, H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models+2023+arXiv】证明，并得到我们实验的支持。


图 6: (a) 计算三个头中每个头的重要性分数。(b) 得分最低的词元的 KV 向量被新的 KV 向量替换。输入向量 x4 被存储，因为第四个词元在三个头中的两个头中是重要的。存储 x4 释放了一个 eDRAM 条目，从而降低了 eDRAM 的刷新成本。(c) 为节省 eDRAM 存储空间而重新计算第四个词元的 KV 向量。

重计算策略。如第 2.3 节所述，eDRAM 非常适合存储瞬时数据。虽然驱逐策略减少了模型执行期间需要保留的 KV 向量数量，但在 eDRAM 中存储生命周期较长的 KV 向量仍然会因所需的刷新操作而产生严重成本。为了减轻刷新成本，我们可以进一步应用重计算技术。具体来说，对于 KV 缓存中的一个词元子集 $S_{recomp}$，它们的 KV 向量将使用相应的输入向量 $x_n$ 重新计算，这些输入向量作为 $W_Q, W_K$ 和 $W_V$ 的输入，如图 1 (b) 所示。通过利用重计算，存储需求可以从维护两个向量（K 和 V）减少到仅持有一个向量（输入 x）。这种方法允许根据需要重新计算 K 和 V，有效减轻了 KV 向量的长数据生命周期问题。

重计算的条件。在解码阶段执行期间，首先通过将输入向量 $x_n$ 与 $W_K$ 和 $W_V$ 相乘来重新计算所有头 $h \in H$ 的 KV 向量 $k_n^h$ 和 $v_n^h$（图 6 (c)），然后用于解码过程。为了通过重计算节省 KV 缓存存储，输入向量 $x_n$ 的存储成本（即 $1 \times C$）必须小于重新计算的 KV 向量的存储成本。为了满足这一点，如果一个词元 n 的 KV 向量在不进行重计算的情况下会在至少 $\pi > 50\%$ 的头中被保留，那么它的 KV 向量将使用 $x_n$ 进行重计算，其中 $\pi$ 代表该词元的流行度。这种方法是合理的，因为与 KV 向量相关的存储成本（计算为 $2 \times N_h \times D_h$）将超过 $x_n$ 的大小（即 C）。如图 6 (b) 所示，第四个词元在三个头中的两个头中被认为是流行的，所以保留了输入向量 $x_4$，避免了存储 KV 向量。

AERP 整体流程。除了节省存储空间，重新计算的 KV 向量将是瞬时的，因为它们仅在短时间内用于计算方程 1，这进一步利用了 eDRAM 的优势。此外，由于计算引擎采用脉动阵列架构，重计算的额外成本将是最小的，这在第 5.2 节中讨论。在预填充阶段，首先计算每个头 h 中每个词元 n 的重要性分数 $s_n^h$。接下来，对于每个头 h，根据相应词元的重要性分数驱逐 KV 向量。在重要性分数高的词元中，那些 KV 向量被保留在至少 50% 的头中的（即流行词元），它们的输入向量 $x_n$ 被存储；否则，存储 KV 向量。在解码期间，通过计算流行度 $\pi$ 来动态确定每个新词元的存储格式。图 7 (a) 总结了整个 AERP 方案。尽管词元流行度在解码过程中可能会变化，但经验证据表明波动有限，即对超过 50% 的头重要的词元很少会降低其重要性。因此，在 Kelle 中，一旦一个词元以其输入向量的形式存储，其存储格式在整个解码过程中保持固定，除非它被驱逐。

4.2 二维自适应刷新策略 (2DRP)

LLM 对 KV 缓存数据损坏的容忍度。为了探究 LLM 在不影响准确性的情况下对 KV 缓存中数据损坏的容忍度，我们通过在 eDRAM 存储单元中引入位翻转错误来模拟保留失败。具体来说，我们使用 Wikitext-2 【索引 52, Pointer Sentinel Mixture Models+2016+arXiv】数据集评估了 LLaMA2-7B 模型在困惑度 (Perplexity, PPL) 上的影响。较低的 PPL 表示更好的性能。在执行期间，以均匀概率在 KV 缓存中引入位翻转错误。结果呈现在图 8 (a) 中，显示当错误率低于 $10^{-3}$ 时，PPL 的增加保持在最小，低于 0.1。然而，随着位翻转错误的持续增加，PPL显著上升。这表明 LLM 可以容忍一定程度的 KV 缓存错误。一个自然的后续问题是：是否可能开发一种更细粒度的刷新策略，能够在保持准确性的同时支持更低的刷新频率？


图 7: (a) AERP 总结，为简化仅显示一个头 h。(b) 二维自适应刷新策略。(c) 2DRP 的一个例子。k5[15:8] 表示第五个词元的键向量的第 8 到 15 位。颜色越深表示比特刷新越频繁，从而导致更低的保留错误率。

基于重要性的自适应刷新策略。在第 4.1 节中，词元根据其重要性分数被驱逐，如方程 3 所定义。我们假设类似的方法可以应用于 eDRAM 刷新策略，即为重要性较低的词元的 KV 向量或输入向量分配较低的刷新频率，为重要性较高的词元分配较高的频率。为了验证这一假设，我们实现了一个自适应刷新策略并重复了实验。词元根据其重要性分数分为两组，为简单起见，称为高分词元 (HST) 组和低分词元 (LST) 组。将 KV 向量中位保留失败（位翻转错误）的概率 P 分别应用于 HST 和 LST 组中相应词元的 KV 向量。图 8 (b) 的结果显示，当位保留失败影响 HST 组时，LLM 性能下降得更多，而不是 LST 组，这表明 HST 组中的词元需要更高的刷新频率，从而支持了我们的假设。


图 8: (a) 不同位翻转错误率 P 下的 PPL。(b) 当仅对 (a) HST vs. LST 和 (b) MSB vs. LSB 应用位翻转时，不同位翻转错误率下的 LLM 准确性，其中 P 表示错误率。较低的 PPL 反映了更好的性能，红色数字表示 PPL 值之间的差距。

基于比特位的自适应刷新策略。此外，可以合理地假设，较低有效位 (LSBs) 对保留失败错误的脆弱性低于较高有效位 (MSBs)，因为 LSBs 上的位翻转错误导致的值变化较小。对于 KV 向量中的每个值，我们向 MSBs（位 15-8）或 LSBs（位 7-0）引入位保留错误。结果如图 8 (c) 所示，揭示了在相同的位翻转错误率下，MSBs 对保留错误比 LSBs 更敏感，进一步支持了我们的假设。

2DRP 策略。基于上述观察，我们提出了一种称为二维自适应刷新策略 (2DRP) 的自适应刷新控制策略，如图 7 (b) 所示。该策略根据 KV 向量或输入向量中每个值的位位置以及每个词元的重要性分数来调整每个 eDRAM 单元的刷新频率。2DRP 的一个例子如图 7 (c) 所示，其中 KV 缓存最多可容纳 $N' = 3$ 个词元。刷新频率随着词元重要性和位位置重要性的增加而增加。在执行期间，动态计算 KV 和输入向量的重要性分数，并根据这些分数和位位置相应地分配刷新频率。

5 Kelle 边缘加速器

图 9 概述了 Kelle 加速器。它包含一个混合的 eDRAM-SRAM 内存子系统、一个可重构脉动阵列 (RSA) 和专用功能单元 (SFU)。权重被量化为 8 位，激活值和 KV 向量保持在 16 位，权重存储在 SRAM 中，而激活值和 KV 向量则保存在 eDRAM 中。在操作期间，脉动驱逐器累积注意力分数，eDRAM 控制器处理 KV 向量的驱逐和重计算，同时动态调整刷新频率，如第 4.1 和 4.2 节所述。RSA 中的每个处理元件 (PE) 执行 8 位乘法累加 (MAC) 操作。


图 9: Kelle 硬件加速器概览。

SFU 处理非线性操作，包括激活函数、softmax、归一化和位置嵌入。正如先前研究【索引 20, ViTALiTy: Unifying Low-rank and Sparse Approximation for Vision Transformer Acceleration with a Linear Taylor Attention+2022+arXiv; 索引 62, PICACHU: Plug-In CGRA Handling Upcoming Nonlinear Operations in LLMs+2025+ASPLOS; 索引 78, SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning+2020+HPCA; 索引 81, Hyft: A Reconfigurable Softmax Accelerator with Hybrid Numeric Format for both Training and Inference+2024+ISLPED】所示，非线性操作的能耗随着输入序列长度的增加而增加。在这些操作中，softmax 消耗大量资源。我们采用 Softermax【索引 71, Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformers+2021+arXiv】中的在线最大值计算来最小化内存访问。对于其他非线性操作，我们遵循计算流程并使用查找表 (LUT) 来执行计算。

5.1 内存子系统

内存子系统设计。图 10 展示了 Kelle 加速器的内存子系统。在此设计中，一个 2MB 的 SRAM 存储权重，而激活值和 KV 向量分别保存在一个 256KB 的激活 eDRAM 和一个 4MB 的 KV 缓存 eDRAM 中。Kelle 加速器通过将 KV 向量根据重要性分数和位位置分为四组来实现 2DRP，并相应地应用刷新频率。具体来说，HST 组中词元的 KV 向量的 MSBs（位 15-8）以最高频率刷新，而 LST 组中词元的 KV 向量的 LSBs（位 7-0）以最低频率刷新。为了支持 AERP，对于某些词元，输入向量被存储在 KV 缓存 eDRAM 中，而不是 KV 向量。这些输入向量然后被分为四组，并以与 KV 向量相同的方式进行控制，按重要性分数和位宽组织。为简单起见，我们使用 KV 向量来描述内存子系统设计，而不引用输入向量。

2DRP 的硬件实现。为了在 LLM 推理期间执行 2DRP，KV 向量的每个元素都按位分割并存储在不同的 eDRAM 库 (bank) 中。具体来说，对于 KV 向量，MSBs 和 LSBs 存储在不同的 KV 缓存 eDRAM 库中，分别称为 MSB 库和 LSB 库，在图 10 中用较深和较浅的颜色突出显示。每个词元的重要性分数使用方程 3 动态计算，并以 4 位精度存储在一个寄存器文件中，每个条目对应一个跨越四个库的 KV 向量。对应于同一词元的 KV 向量在不同的 eDRAM 库中共享相同的地址。该系统具有一个单一的驱逐控制器，用于管理所有四个库的 AERP，以及两个刷新控制器，分别负责在 MSB 和 LSB 库上执行 2DRP。


图 10: Kelle 内存子系统。某些词元的输入向量存储在 KV 缓存中，由红色行表示。

刷新与驱逐控制。在每个 MSB 和 LSB 库中，词元根据其注意力分数进一步分为两组，刷新控制器中的一个计数器监控每组的刷新间隔。控制器迭代 eDRAM 条目，通过从寄存器文件中读取其注意力分数来识别每个词元的分组。当特定组的刷新间隔到期时，控制器触发一个刷新信号。然后计算该组中词元的相应 KV 向量地址，并将 KV 向量读出并写回作为刷新过程的一部分。刷新操作在 KV 向量未被模型使用时触发，因此刷新延迟可以被隐藏。当 KV 缓存达到容量且新词元到达时，驱逐控制器从脉动驱逐器接收要驱逐的词元索引，并用新词元替换该词元。

避免 bank 冲突。为了通过并行馈送数据充分利用 32 × 32 RSA 并避免 bank 冲突，Kelle KV 缓存被分为 32 个 bank。具体来说，Key MSB、Key LSB、Value MSB 和 Value LSB 组各分配了 8 个 bank。通过这种设计和流水线化的缓存读取，Kelle eDRAM 提供了足够的带宽来充分利用 RSA 而不产生 bank 冲突。此外，其他 eDRAM 访问操作（如词元读取和词元驱逐）独立运行，这有效地减轻了 bank 冲突。

数据流与片外访存。在 LLM 执行期间，RSA I/O 控制器以最小的开销高效地从不同的 bank 重构数据用于计算。此外，Kelle 加速器将一部分 LLM 层的 KV 向量存储在 eDRAM 中，层数由具体的 LLM 大小和文本长度决定。eDRAM 极大地减少了片外内存访问开销。

5.2 可重构脉动阵列 (RSA)

RSA 架构。脉动阵列核心由一个 32 × 32 的二维阵列组成，它以交错的方式处理输入，并将计算出的部分和发送到累加器和 SFU。它采用权重固定 (weight-stationary) 的数据流，如图 11 (a) 所示。我们采用类似于 FAST【索引 96, FAST: DNN Training Under Variable Precision Block Floating Point with Stochastic Rounding+2021+HPCA】中的可重构策略来执行原地转置矩阵乘法。

重计算的开销。重要的是，第 4.1 节中的重计算在 LLM 解码阶段引入的开销极小。利用脉动阵列在矩阵运算方面的优势，重计算的词元向量可以与当前词元的输入向量结合，高效地创建一个输入矩阵。使用与图 6 (b) 中示例相同的符号，图 11 (a) 显示当前词元的输入向量 $x_5$ 被发送到 RSA 进行 KV 向量计算。为了重新计算第四个词元的 KV 向量，$x_4$ 和 $x_5$ 可以组合成一个矩阵，这在图 11(b) 中导致了最小的延迟和能耗增长。


图 11: (a) 和 (b) 展示了重计算对 RSA 操作的影响。(c) RSA 和脉动驱逐器的集成。(d) 脉动阵列和脉动驱逐器的执行顺序。红色圆圈中的数字表示顺序。

5.3 脉动驱逐器 (Systolic Evictor)

AERP 驱逐流程。AERP 算法中的词元驱逐过程包括：计算注意力分数（如方程 1 所述），根据方程 3 更新重要性分数，识别重要性分数最低的词元，以及执行 KV 缓存更新。

脉动驱逐器设计。为了高效实现驱逐算法，我们提出了一种脉动驱逐器 (Systolic Evictor, SE)，它以脉动方式运行，并与 RSA 集成，以在运行中搜索最小重要性分数。重要性分数是通过将方程 1 中的 $q \cdot k^T$ 结果相加计算的，而无需通过 softmax。这种集成确保了在从 RSA 计算出新词元的注意力分数后，立即找到重要性分数最小的词元。在找到重要性分数最小的词元的索引后，SE 将该索引发送到 eDRAM 控制器中的驱逐控制器，以驱逐相应的词元。

SE 与 RSA 的集成。图 11 (c) 展示了 SE 的设计及其在 RSA 中的集成。SE 包含一列寄存器，在图 11 (c) 中表示为 S，用于预加载先前词元的重要性分数，以及一个寄存器链，表示为 M，周期性地从上到下传播最小重要性分数 (min)。

执行顺序。图 11 (d) 展示了 RSA 和 SE 的执行顺序。在一个周期内，注意力分数从 RSA 的第 i 行计算出来，然后 SE 的第 i 行更新重要性分数和最小重要性分数索引，分别标记为步骤 1 和步骤 2。在下一个周期内，相同的操作在 RSA 和 SE 的下一行执行，分别标记为步骤 3 和步骤 4。脉动驱逐器避免了因最小搜索而产生的额外 LLM 执行延迟。

6 Kelle 调度器

为了进一步减少 eDRAM 的刷新能耗，我们引入了一种新颖的计算模式，该模式缩短了数据生命周期并加速了 LLM 推理，且不影响准确性。


图 12: (a) 和 (b) 分别展示了基线和 Kelle 调度中 SA 模块的计算模式和 eDRAM 数据生命周期。SM 表示 softmax 操作。

SA 模块计算分析。首先，我们对 LLM 解码阶段自注意力 (SA) 架构相关的数据生命周期进行数值分析。如第 2.1 节所述，SA 中的计算首先涉及输入 x 与权重矩阵 $W_Q$、$W_K$ 和 $W_V$ 的矩阵乘法，产生输出 q、k 和 v。这些过程分别表示为 $MM_q$、$MM_k$ 和 $MM_v$。随后，q 和 k 相乘，然后进行 softmax 操作以计算注意力分数 A，这分别称为 $MM_{qk}$ 和 SM。最终，A 与 v 相乘，再与权重矩阵 $W_O$ 相乘以产生 SA 输出，标记为 $MM_{vo}$。矩阵乘法的延迟 ($T_{MM}$) 估计如下：

$$T_{MM} = \frac{N_{MM}}{TOP_{RSA}}$$

延迟建模。其中 $N_{MM}$ 表示矩阵乘法所需的 MAC 操作数。$TOP_{RSA}$ 表示第 5.2 节中描述的 RSA 的吞吐量。与 KV 向量的 eDRAM 访问操作相关的延迟，表示为 $T_{eDRAM}$，建模如下：

$$T_{eDRAM} = \frac{S_{KV}}{B_{eDRAM}}$$

其中 $S_{KV}$ 表示 KV 向量的大小（以字节为单位）。$B_{eDRAM}$ 表示 eDRAM 的带宽。类似地，与权重 SRAM 访问操作相关的延迟，表示为 $T_{SRAM}$，建模如下：

基线调度器的数据生命周期。其中 $S_W$ 表示权重的大小（以字节为单位）。$B_{SRAM}$ 表示 SRAM 的带宽。图 12 (a) 展示了一个基线计算模式，其中矩阵乘法操作 $MM_q$、$MM_k$、$MM_v$ 和 $MM_{qk}$ 是一个接一个进行的，这延长了输入 x、q、k 和 v 的数据生命周期。数据生命周期定义为从数据计算出来到被后续操作使用的时间间隔。例如，在图 12 (a) 中，向量 q 的计算在从 SRAM 访问权重矩阵 $W_Q$ 后于 $t_1$ 开始，而 q 在 $t_4$ 时被 q 和 k 的乘法消耗。在 $t_1$ 和 $t_4$ 之间，$W_K$ 和 $W_V$ 从 SRAM 加载，K 从 eDRAM KV 缓存访问。访问 $W_K$ 和 $W_V$ 的延迟都是 $T_{SRAM}$，访问 K 的延迟是 $T_{eDRAM}$。所以 q 的数据生命周期是 $2 \times T_{SRAM} + T_{eDRAM}$。所有激活值的总数据生命周期是每个激活值数据生命周期的总和，因为它们都存储在 eDRAM 中并需要刷新。我们从方程 4 中省略了计算时间 $T_{MM}$，因为与 $T_{SRAM}$ 和 $T_{eDRAM}$ 相比，它的量级可以忽略不计。这种延长的数据生命周期导致了更高的 eDRAM 刷新成本。基线调度中瞬时数据的总数据生命周期 $L_{bl}$ 建模如下：

Kelle 调度器的数据生命周期。其中 $L_X$、$L_Q$、$L_K$ 和 $L_V$ 分别表示 x、q、k 和 v 的数据生命周期。$T_{SRAM}$ 和 $T_{eDRAM}$ 在方程 6 和方程 5 中定义。相比之下，Kelle 使用的计算模式如图 12 (b) 所示。得益于集成了独立的片上内存，权重和 KV 向量的内存访问是并行的。这种安排减少了激活值的数据生命周期，可以估计如下：

Kelle 调度器的优势。键和值向量立即用于它们各自的计算，消除了长期存储的需要；因此，它们的数据生命周期可以被认为是微不足道的。与基线方案相比，Kelle 调度器显著减少了 eDRAM 中瞬时数据的数据生命周期，从而降低了刷新能耗并增强了系统性能。

A4 实验环境

  • 数据集：

    • 语言生成任务：WikiText-2 (WK2) 【索引 52】，PG19 【索引 66】。WK2 序列长度为数百到数千，PG19 序列长度为数万到数百万。PG19 任务生成长度设置为 8192。
    • 零样本任务：PIQA (PQ) 【索引 10】，Lambada (LA) 【索引 64】，Arc Easy (A-e) 【索引 18】，Arc Challenge (A-c) 【索引 18】，TriviaQA (TQ) 【索引 36】，Qasper (QP) 【索引 19】。
    • 定性评估：CNN/DailyMail (CNN) 【索引 53】（摘要任务，ROUGE-1 分数），TruthfulQA 【索引 43】（事实正确性），BBQ 【索引 58】（偏见评估）。

  • 模型：Llama2 【索引 75】，Llama3 【索引 21】，Llama3.2 【索引 21】，Mistral 【索引 35】，QWEN 【索引 87】，OPT 【索引 95】，模型大小各不相同。

  • 硬件配置：

    • Kelle 加速器：32×32 可重构脉动阵列 (RSA)，2MB SRAM（用于权重），4MB KV 缓存 eDRAM，256KB 激活 eDRAM。
    • 技术节点：45nm NanGate 开放单元库用于综合，65nm 技术节点用于 eDRAM/SRAM 性能评估。
    • 工作频率：1GHz。
    • 片外内存：16GB LPDDR4 DRAM，带宽 64GB/s，类似于 Google Coral 边缘设备【索引 72】。
    • 总片上面积：9.5mm²，功耗：6.52W。

  • 软件配置：

    • 硬件实现：RTL 使用 SystemVerilog。
    • 仿真与综合：Synopsys Design Compiler 【索引 9】，Destiny 【索引 60】 (eDRAM/SRAM)，Cacti 7 【索引 8】 (DRAM)。
    • 评估框架：LM Evaluation Harness 【索引 26】用于零样本任务。
    • 量化：模型权重量化为 8-bit，激活值和 KV 向量为 16-bit（除非另有说明）。

A4 实验结果

7.1 主要准确性结果

与 SOTA 方法的准确性对比。我们将 Kelle 算法与 QuaRot (QR)【索引 6】、StreamLLM (SL)【索引 83】和 H2O【索引 98】等先进方法进行比较。为了保证 KV 缓存预算相同，QR 将 KV 向量量化到 4-bit，而 SL、H2O 和 Kelle 保持 16-bit 未量化。Kelle 根据不同任务动态设置保留的 token 数量 N'（从 128 到 2048 不等），并保留 10 个初始 token 和一定数量的最近 token。如表 2 所示，Kelle 在各种模型和任务上的准确性与使用完整 KV 缓存的 FP16 模型相当，并且优于或达到了与其他方法的相当性能，证明了 AERP 和 2DRP 算法的优越性。

Table 2: 各方法的准确性表现。FP16 表示在 FP16 下无 KV 缓存缩减的 LLM 准确性。

7.2 消融研究

KV 缓存预算的影响。我们调整了 Llama2-7B 模型的缓存预算大小 $N'$。如表 3 所示，随着 $N'$ 减小，准确性持续下降，但当 $N' \ge 128$ 时仍能取得合理的性能。

Table 3: LLaMA2-7B 在不同缓存大小下的准确性。

2DRP 策略的影响。我们将 2DRP 与统一刷新间隔的策略进行比较，同时保持平均保留失败率相同。表 4 显示了在 Llama2-7B 上不同刷新间隔对准确性的影响。结果表明，在所有条件和数据集下，2DRP 相较于统一刷新策略都能提高准确性。

Table 4: LLaMA2-7B 在不同刷新间隔下的准确性。

定性指标和量化兼容性。Kelle 在文本生成质量（如一致性、事实正确性、偏见）方面与 FP16 模型表现相当（表 5）。此外，Kelle 与 Quarot 等量化框架兼容，在进一步量化（W4A8）后，系统性能预计会进一步提升，而准确性影响很小（表 6）。

Table 5: Kelle 定性指标

Table 6: Kelle 与量化结合的准确性

8.1 端到端性能评估

与基线的比较。我们将 Kelle+eDRAM 与四个基线进行比较：Original+SRAM（原始 LLM 在 SRAM 系统上）、Original+eDRAM（原始 LLM 在 Kelle 的 eDRAM 加速器上）、AEP+SRAM（在 SRAM 系统上仅使用注意力驱逐）和 AERP+SRAM（在 SRAM 系统上使用完整的 AERP 算法）。如图 13 所示，与 Original+SRAM 相比，Kelle+eDRAM 平均实现了 3.94 倍的延迟改进和 4.46 倍的能效提升，且随着解码序列变长，性能差距会进一步拉大。

各优化技术的贡献分析：

  1. Original+eDRAM vs. Original+SRAM：速度提升 32%，但能效下降 39%，因 eDRAM 刷新开销大。
  2. AEP+SRAM vs. Original+SRAM：延迟降低 2.39 倍，能效提升 2.41 倍，显示了注意力驱逐策略的有效性。
  3. AERP+SRAM vs. AEP+SRAM：速度和能效分别再提升 1.19 倍和 1.27 倍，体现了重计算策略的价值。
  4. Kelle+eDRAM vs. AERP+SRAM：速度和能效再提升 1.29 倍和 1.45 倍，证明了 2DRP 和 Kelle 调度器对 eDRAM 的优化效果。

开销分析。图 13 中的饼图显示，在 Kelle+eDRAM 系统中，KV 缓存的能耗占比显著降低，表明 Kelle 算法有效缓解了内存访问瓶颈。脉动驱逐器引入的面积和功耗开销极小（分别占片上总量的 0.6% 和 0.4%），但能将系统能效和延迟分别提升 5% 和 7%。


图 13: Kelle 与基线系统的比较。性能以归一化能效和加速比来评估。饼图显示了 Kelle+eDRAM 内主要组件的片上能耗分解。红色虚线描绘了相应设置的加速比。

8.2 与其他加速器的比较

我们将 Kelle+eDRAM 与 LLM.npu【索引 86】、DynaX【索引 85】、COMET【索引 46】以及 NVIDIA Jetson Orin 边缘 GPU【索引 2】进行比较。如图 14 所示，Kelle+eDRAM 在加速比和能效上均优于其他 LLM 加速器。这主要是因为 Kelle 同时解决了计算和 KV 缓存瓶颈，并拥有专用的硬件支持。


图 14: LLM 加速器之间的比较

8.3 硬件消融研究

  • KV 缓存预算的影响：即使在无驱逐的最大缓存预算下（PG19 任务下 $N'=8750$），Kelle 仍比 Original+SRAM 节能约 3 倍，显示了其鲁棒性（表 7）。
  • 重计算的影响：重计算有效降低了 KV 缓存的能耗，而 RSA 能耗增加极小（图 15a）。它通过重叠计算与内存访问，平均将延迟降低了 25%。Roofline 模型（图 16a）显示，适度的重计算可以提升性能，但过多的重计算会使系统从内存受限变为计算受限。
  • 2DRP 和 Kelle 调度器的影响：实验表明，从基础刷新策略（Org）到统一刷新（Uni），再到 2DRP（2D），最后结合 Kelle 调度器（2K），系统能效逐步提升，证明了每个优化步骤的有效性（图 15b）。
  • 输入序列长度的影响：在长输入短输出的计算密集型场景中，Kelle 提升中等（2.1 倍）；在长输入长输出的内存密集型场景中，Kelle 的优势更为显著，平均能效提升 5.6 倍（图 16b）。
  • 批量大小 (Batch Size) 的影响：Kelle 在所有批量大小下均优于基线，且批量越大，优势越明显（表 9）。
  • eDRAM 带宽的影响：即使将 eDRAM 带宽减半，Kelle 依然能获得显著的性能增益，证明了增加 eDRAM 容量本身就能有效减少昂贵的 DRAM 访问。


图 15: (a) Kelle+eDRAM 中 KV 缓存重计算的影响。(b) 对 2DRP 和 Kelle 调度器的评估。


图 16: (a) KV 缓存重计算的影响。(b) 对长输入序列的评估。P 和 D 分别表示预填充和解码阶段。

Table 7: 在多个 KV 缓存预算下的能效。

Table 8: 在不同刷新间隔下的能效。

Table 9: 在不同批量大小下的能效。

A7 补充细节

8.4 讨论

处理长上下文推理。对于长上下文推理，由于 eDRAM 容量有限，多余的 KV 数据会被卸载到 16GB 的 DRAM 中。一个简单的分析显示，对于 LLaMA 2 7B 模型，在没有 AERP 的情况下，Kelle 可以支持多达 19,000 个输入词元。引入 AERP 后，可以在每层执行后立即减少 KV 缓存，从而为后续层释放内存以容纳完整的输入序列。这使得 Kelle 能够支持高达约 60K 词元的输入序列。此外，将 KV 向量量化到 4 位可以支持多达 240K 的词元。虽然存在上限，但这已超过了典型 LLM 数万词元的输入长度。尽管更长的输入序列会增加开销，但方程 1 和 2 的置换不变性允许新的 KV 向量被放置在被驱逐向量的相同位置，极大地简化了分页过程。此外，向量可以顺序预取，无需复杂的查找机制。因此，预取开销随输入长度线性增加，避免了不成比例的增长。

将 Kelle 与 GPU 集成。虽然 Kelle 是用脉动阵列实现的，但 AERP 可以适配到 GPU 上；然而，由于缺少脉动驱逐器，识别注意力分数最低的词元可能会效率低下。2DRP 是 eDRAM 特有的，用于减少刷新能耗。eDRAM 可以与 GPU 现有的内存系统耦合，用于存储 KV 向量。最后，Kelle 调度器可以很容易地使用 CUDA 实现。

A5 结论

KV 缓存技术对于提升 LLM 的效率至关重要。然而，存储大量的 KV 向量会导致巨大的内存占用和增加的数据访问成本。在这项工作中，我们介绍了 Kelle 系统，它利用 eDRAM 作为 KV 向量的主要存储介质。Kelle 的卓越性能凸显了 eDRAM 在实现 KV 缓存机制方面的巨大潜力，为未来的研究铺平了道路。