Stateful Large Language Model Serving with Pensieve

发表时间: 2023-12 · arXiv:2312.05516 (EuroSys 2025)

原文: https://arxiv.org/abs/2312.05516

Lingfan Yu (New York University), Jinkun Lin (New York University), Jinyang Li (New York University)

速读

一句话结论 Pensieve 是一个面向多轮对话的带状态 LLM 推理服务系统，通过在 GPU 和 CPU 间跨请求缓存历史 KV 状态，并设计支持非连续显存的多 token 注意力算子，在真实对话场景下将吞吐量提升了 1.14 到 3.0 倍。

要解决什么问题 现有的 LLM 推理系统在跨请求处理时是无状态的。在目前非常流行的多轮对话场景中，用户和聊天机器人的对话历史会不断累积。为了让模型记住之前的上下文，系统必须在每次收到同一对话的新请求时，将不断增长的对话历史作为提示词重新计算一遍（即 prefill 阶段）。随着对话轮数增加，处理历史记录的计算开销会迅速膨胀，甚至超过逐个生成新 token 的开销。这种机制导致了海量的重复计算，严重拖慢了推理速度。虽然理论上可以通过将处理过的 KV 状态保存在 GPU 中来避免重算，但 GPU 显存容量极小，存不下大量并发对话的历史状态；如果存入磁盘，加载延迟又会严重破坏用户体验。因此，核心卡点在于如何突破显存容量限制，高效地跨请求管理、调度和复用对话历史状态。

怎么做的 核心思路是将 LLM 服务从无状态改为有状态，通过构建 GPU-CPU 双层缓存架构来保存并复用历史请求的 KV 状态。为绕开显存不足的卡点，Pensieve 采取了细粒度的 token 级别（以 32 个 token 为一个数据块）缓存管理策略。当显存吃紧时，系统会将部分数据块驱逐到 CPU 内存，甚至直接丢弃。驱逐策略基于两个偏好：一是优先驱逐闲置时间最长的对话；二是优先驱逐对话历史中最靠前（leading）的 token。这是因为注意力机制具有因果性，越靠前的 token 需要参与计算的上下文越短，重算成本越低。系统为每个数据块计算保留价值：

$$V = \frac{Cost(s, l)}{T}$$


其中 $T$ 是对话闲置时间，$Cost(s, l)$ 是重算大小为 $s$ 的数据块的开销，$l$ 是上下文长度。重算开销被建模为线性增长的注意力开销加上常数级的非注意力开销：

$$Cost(l) = Cost_{attention}(l) + c$$
系统会优先驱逐 $V$ 值最小的数据块。这种驱逐策略导致了一个新问题：被丢弃的头部 token 需要重算，而尾部的新请求 token 也需要计算，中间的 token 则是从 CPU 换入的。这使得 prefill 阶段的上下文在物理显存上是不连续的，且分成了断开的片段。现有的 PagedAttention 算子虽然支持非连续显存，但其本质是矩阵向量乘法，仅限每个请求输入单个 token 的生成阶段使用。为此，Pensieve 开发了一个广义的 Multi-token attention 算子，它执行批处理矩阵乘法，能够计算多个输入 token 的 Query 与分布在非连续显存中的 KV 状态之间的注意力。系统将断开的上下文视为共享部分底层数据的两个子请求来处理，无需进行显存拷贝。借助这个新算子，Pensieve 打破了传统系统将 prefill 和生成阶段分开调度的限制，实现了一个统一的批处理调度器。无论是正在处理长提示词的新请求，还是正在逐个生成 token 的老请求，都可以被拼接到同一个批次中执行。同时，系统通过提前换出和流水线式的层级换入，将 CPU 与 GPU 之间的数据传输延迟与模型计算完全重叠。

效果如何 实验在配备 A100-80GB GPU 的硬件上进行。评估使用了 OPT 和 Llama 2 两种模型，分为单 GPU 运行的小模型（OPT-13B、Llama 2-13B）和通过 Megatron-LM 张量并行在 4 张 GPU 上运行的大模型（OPT-66B、Llama 2-70B）。测试数据采用真实的 ShareGPT 和合成的 UltraChat 多轮对话数据集。对比基线是两个代表无状态路线的先进系统：基于 PyTorch 执行的 vLLM 和基于计算图重写与算子融合优化的 TensorRT-LLM。在 ShareGPT 数据集上，对于单卡小模型，在相同的延迟约束下（如每 token 延迟 120 毫秒），Pensieve 的吞吐量达到了 vLLM 的 1.36 倍至 1.70 倍，是 TensorRT-LLM 的 1.14 倍至 1.58 倍。对于 4 卡大模型，由于模型参数和计算量的增长速度快于 KV 状态显存占用的增长速度，Pensieve 的优势被进一步放大，吞吐量达到了 vLLM 的 2.04 倍至 3.0 倍，是 TensorRT-LLM 的 1.64 倍至 2.47 倍。特别是对于采用了分组查询注意力（GQA）的 Llama 2 模型，由于 KV 状态体积大幅减小，Pensieve 能够缓存更多历史记录，性能提升最为显著。作者也指出了方法的局限性：当用户的思考时间（即两次对话请求之间的间隔）过长时（例如达到 600 秒），历史 KV 状态会大量从缓存中流失，导致缓存命中率下降，此时 Pensieve 的性能优势会缩小。此外，由于 PCIe 双向并发传输会导致吞吐量下降约 18% 到 20%，系统在工程实现上做出了妥协，强制优先执行数据换入操作，未能完全榨干 PCIe 的双工带宽。

主要贡献

当前大型语言模型（LLMs）极为流行，高效地为其提供服务至关重要。现有的LLM服务系统在跨请求时是无状态的。因此，当LLM用于常见的多轮对话场景时，服务系统在每一轮都必须将不断增长的对话历史记录与任何新请求一起进行从头处理，这导致了大量的重复计算。

在本文中，作者设计了Pensieve，这是一个专为多轮对话LLM服务优化的系统。Pensieve通过缓存先前处理过的历史记录来在请求之间维护对话状态，从而避免重复处理。Pensieve的多级缓存策略可以同时利用GPU和CPU内存来高效地存储和检索缓存数据。Pensieve还推广了近期的PagedAttention内核，以支持分布在非连续内存上的GPU缓存与多个输入Token之间的注意力计算。评估表明，Pensieve可以实现vLLM和TensorRT-LLM $1.14 - 3.0\times$ 的吞吐量，并显著降低延迟。

背景知识与关键Observation

LLM与注意力机制：流行的按比例放大的语言模型，如GPT-3 【5，Language models are few-shot learners + 2020 + NeurIPS】、OPT 【55，Opt: Open pre-trained transformer language models + 2022 + arXiv】、Llama 【46，Llama: Open and efficient foundation language models + 2023 + arXiv】等，均基于Transformer架构 【48，Attention is all you need + 2017 + NeurIPS】。一个模型由许多Transformer层组成，每一层包含一个注意力模块（如图1虚线框所示）和一个2层的前馈网络。模型接收代表自然语言句子的Token IDs序列作为输入，通过嵌入层为每个Token获取连续表示（即嵌入），然后再输入到Transformer层。LLM具有自回归特性，它基于包含输入Prompt Token及先前迭代生成的输出Token的当前上下文，迭代地预测下一个输出Token。
每层的注意力模块首先对其输入Token嵌入（$X$）执行QKV投影，即线性变换，生成Query（$Q$）、Key（$K$）和Value（$V$）三个新嵌入：

$$Q = X W_{query}$$

$$K = X W_{key}$$
$$V = X W_{value}$$
其中 $W_{query}$、$W_{key}$ 和 $W_{value}$ 是可训练权重。$K$ 和 $V$ 嵌入被称为KV-tokens。然后，该模块使用每对Token的Query和Key的点积计算全互联的注意力分数：
$$A = \frac{Q K^T}{scale}$$
其中 $scale$ 是归一化因子。随后使用softmax对该注意力分数进行归一化，并用于对 $V$ 中的Token嵌入进行加权聚合：
$$O = softmax(A) V$$
上述公式展示了单头注意力的过程，实际应用中使用多头注意力，将公式生成的 $Q, K, V$ 分为多个注意力头独立计算。

LLM的服务方式：为了使用LLM进行推理，需要在GPU中保留一个KV缓存，以避免在自回归输出生成期间进行重新计算。图1展示了FasterTransformer 【35，FasterTransformer + 2021 + GitHub】、ORCA 【52，Orca: A distributed serving system for Transformer-Based generative models + 2022 + OSDI】、vLLM 【27，Efficient Memory Management for Large Language Model Serving with PagedAttention + 2023 + SOSP】等系统采用的典型LLM推理过程。它分为两个阶段：1）在Prefill（预填充）阶段，所有输入Prompt Token被一起处理，为每一层生成 $K$ 和 $V$（即KV-tokens），并用生成的KV-tokens初始化KV缓存。最后一层最后一个Token的嵌入用于生成第一个输出Token；2）Generation（生成或解码）阶段在多个步骤上迭代工作。在每一步中，上一个生成步骤生成的Token被作为单个新输入Token进行处理。每一层计算新Token的 $Q, K, V$ 嵌入向量，更新KV缓存，并使用新Token的 $Q$ 嵌入与KV缓存中的所有KV-tokens执行注意力计算。
图1 采用Transformer架构的LLM推理过程

迭代级批处理与内存管理：对于输入输出大小可变的LLM，批处理的粒度对系统吞吐量和延迟有巨大影响。ORCA扩展了非Transformer模型的迭代级批处理策略，在Token粒度上执行批处理：每当请求完成一次迭代生成步骤时，调度器会检查其是否到达序列末尾并可以离开批处理，从而腾出空间让新请求立即开始其生成阶段。在内存管理方面，vLLM能够为每个请求动态增加分配的缓存槽，并允许这些槽驻留在非连续的GPU内存中。现有服务系统在请求之间是无状态的，这意味着它们会在请求完成后立即释放该请求使用的所有缓存槽。
图2 现有服务系统在多轮对话中随着每个请求重复处理累积的历史记录

多轮对话中的冗余计算动机：在多轮对话中，用户与聊天机器人进行多轮交谈，因此底层LLM需要了解对话历史记录以生成适当的响应。由于现有服务系统的无状态特性，如图2所示，累积的对话历史记录必须作为原始文本前置到每个新请求中。随着交互的继续，对话历史记录不断增长，使得Prefill阶段的成本掩盖了迭代生成阶段的成本。图3展示了在一个人造工作负载下的沉重提示初始化成本，其中每个请求有200个新提示Token，并具有不同的对话历史记录大小。系统目标是通过在服务系统缓存任何先前处理过的嵌入，并在同一对话的后续请求中重用它们，从而最小化对话历史的冗余计算。
图3 一批32个请求执行Prompt（32个Token）预填充和200步生成的执行时间

GPU内存限制与缓存管理的挑战：LLM模型参数巨大，导致KV-tokens非常大。考虑到有限的GPU内存容量，根据历史记录长度，GPU中只能保留几十或几百个对话历史记录。因此，必须将缓存空间扩展到使用更充足的CPU内存。当使用多级GPU-CPU缓存时，在粗粒度的整个对话历史记录层面进行交换是次优的，因此决定在单个Token粒度上进行交换。当CPU内存压力大时，部分缓存的KV-tokens需要被丢弃并在以后重新计算。图4显示，由于因果注意力计算的性质，出现在上下文序列后面的Token比前面的Token需要更多的计算，因此更倾向于丢弃对话历史前端的Token。
图4 32个Token块在不同上下文大小下的注意力操作执行时间，按Transformer层中非注意力操作的执行时间归一化

丢弃前端Token与非连续缓存的挑战：从对话前端丢弃KV-tokens带来了额外的复杂性。图5说明了在Prefill阶段持续对话中典型请求的上下文布局，它被分为四段：被丢弃需重算的、驻留在CPU需取回的、驻留在GPU的、以及新请求的原始Token。两端都需要计算打破了所有现有注意力内核假设输入Token属于连续上下文区域的假设。此外，为了支持KV-tokens在GPU和CPU之间的交换，允许KV-tokens驻留在非连续GPU内存区域更为高效。但vLLM的PagedAttention仅针对生成阶段设计，限制每个请求正好有一个输入Token。为了实现高效的GPU计算，必须解决在Prefill期间支持非连续KV缓存的挑战，这也带来了将处于不同阶段的请求在同一批次中处理的好处（图6）。
图5 典型请求的KV-token上下文布局。阴影区域（出现在上下文两端）标记了必须由预填充阶段处理的Token
图6 使用非连续KV缓存的多Token注意力。批处理包含两个请求，左侧显示它们的Token（Query表示），右侧显示驻留在非连续GPU内存中的上下文（KV-tokens）

方法细节

系统架构概述：图7展示了系统架构。Pensieve由一个调度器和多个工作节点（Worker）组成，每个工作节点管理一个GPU。调度器有两项任务：1）负责批处理请求以执行；2）确保批处理中的请求有充足的GPU内存执行。对于前者，调度器执行细粒度的迭代级批处理，使新请求可以在现有请求执行自回归生成时加入批次。对于后者，调度器管理KV缓存槽的分配，并决定何时在GPU和CPU KV缓存之间进行交换。每个工作节点也有两项任务：1）调用GPU内核处理一批请求；2）根据调度器确定的批处理缓存计划，在GPU和CPU KV缓存之间执行实际的数据移动。
图7 Pensieve的系统架构

统一的批处理调度器：Pensieve执行细粒度的迭代级批处理。与ORCA仅对非注意力操作进行批处理不同，Pensieve和vLLM对非注意力和注意力操作都进行批处理。然而，vLLM仅在同一阶段的请求之间形成批处理并在单独的内核调用中处理，而Pensieve以统一的方式处理Prefill阶段和Generation阶段。这意味着Pensieve调度器将请求组成一个批次，无论它们处于哪个阶段，这种统一批处理由Pensieve的多Token注意力内核设计实现。

调度器的触发与新请求加入：Pensieve的迭代级调度器由生成步骤的完成来触发。具体而言，在每次Token生成的迭代之后，工作节点返回任何已经发出句子结束Token或达到最大解码长度的已完成请求。如果批次中Token的总数少于某个预配置的阈值，调度器会按先到先得的调度策略从等待队列中找到下一个请求加入批次。

统一批处理的拼接：为了形成统一批处理，调度器连接所有请求要处理的Token。对于加入批次的新请求，其Token包括对应于请求新用户Prompt的Token。对于批次中的每个现有请求，其Token包括上一个生成步骤生成的Token。通过将Prefill阶段与Generation阶段结合在一起，系统避免了运行单独的小内核，从而可以提高GPU利用率。

KV缓存管理与检查：传统上，GPU中的KV缓存仅作为计算工作区。在Pensieve中，GPU KV缓存还作为存储空间来缓存最近完成的活跃对话的KV-tokens。Pensieve采用两级分层缓存策略，并使用大得多的CPU内存作为第二级缓存空间。调度器跟踪剩余的空闲GPU KV缓存槽数量。在将一批请求移交给工作节点之前，调度器尝试确保任何新请求过去的KV-tokens将驻留在GPU中，并且有足够的GPU内存用于执行。如果批次中的请求有任何被换出到CPU内存或被丢弃的KV-tokens，调度器会确定需要多少额外的GPU KV缓存槽来换入或重新计算这些缺失的KV-tokens。如果空间充足，调度器指示工作节点执行必要的分配，并开始从CPU交换这些KV-tokens作为批次缓存计划的一部分。

细粒度的缓存驱逐策略：Pensieve执行细粒度的Token级缓存驱逐和丢弃。驱逐策略旨在表达两种偏好：1）优先从较旧的对话中驱逐，即那些已经不活跃较长时间的对话（基于LRU假设）；2）优先从对话历史上下文的前端驱逐Token，这是基于前端Token重算成本较低的观察。

驱逐粒度与保留价值：为了减少频繁做出驱逐决策和在PCIe总线上移动少量内存带来的开销，系统将KV-tokens分组为Chunk，并在Chunk粒度上做出驱逐决策。Chunk大小是可配置的，实验发现设置为32个Token效果良好。系统为每个Chunk计算一个保留价值得分 $V = \frac{Cost(s, l)}{T}$，其中 $Cost(s, l)$ 表示重新计算大小为 $s$、上下文大小为 $l$ 的Chunk的成本，分母 $T$ 是对话上次活跃以来的时间。Pensieve按保留价值升序驱逐Chunk。

重算成本估算：重算嵌入的成本被视为重算LLM模型注意力操作和其余非注意力操作的总和：$Cost(s, l) = Cost_{attention}(s, l) + Cost_{other}(s)$。非注意力计算成本独立于上下文大小，而注意力操作需要访问所有 $l$ 个上下文Token。针对固定32个Token的Chunk，成本函数简化为 $Cost(l) = Cost_{attention}(l) + c$。系统执行离线分析以估算 $c$ 以及不同上下文大小下的 $Cost_{attention}(l)$，并使用测量值插值计算其他上下文大小的成本。

提前交换（Ahead-of-the-time swapping）：由于Pensieve试图在GPU中保留KV-tokens供同一对话的后续请求重用，调度器不会像现有系统那样在请求完成后立即释放其GPU缓存槽。如果不等待GPU缓存耗尽，而是在可用GPU缓存槽少于阈值（例如 $25\%$）时，调度器就要求工作节点将选定的KV-tokens复制（即换出）到CPU。对应的GPU内存以惰性方式回收。当CPU缓存空间耗尽时，使用相同的驱逐策略决定丢弃哪些KV-tokens。提前交换允许调度器将缓存驱逐与GPU计算重叠，从而完全隐藏交换的延迟。

流水线KV缓存恢复：调度器在将请求交给工作节点执行之前，不会等待其KV-tokens从CPU完全换入。系统采用流水线方法将计算与数据传输重叠。系统利用LLM模型有许多层且每一层的KV-token仅用于该层自注意力计算的特点，逐层发起传输并同时启动模型计算。工作节点使用GPU事件来保留数据依赖性：只有在该层的KV-tokens完全复制到GPU后，才启动该层的自注意力内核。

处理丢弃的Token：如果请求的部分KV-tokens因CPU内存压力被丢弃，系统将诉诸重算。调度器会从保存在持久存储中的对话历史中获取丢弃Token对应的原始文本Token，将其合并（即前置）到新请求的Prompt中作为批次输入Token的一部分（图8步骤a）。在Prefill阶段，丢弃Token和新Prompt Token的嵌入在连续模型层处理时被连接在一起，计算Query、Key和Value张量（步骤b）。Key和Value存储在KV缓存中，Pensieve维护整个对话上下文（包括先前缓存的Token）的KV位置（步骤c），随后用于执行注意力。
图8 Pensieve如何处理丢弃的Token。 (a) 将丢弃上下文和新Prompt的Token连接作为输入。(b) 计算投影 $Q, K, V$ (c) 使用新计算的 $K, V$ 更新KV缓存；换入之前驱逐到CPU的KV-tokens。(d) 将其作为两个子请求执行注意力：一个让丢弃的上下文关注自身，另一个让新Prompt关注整个上下文。

非连续Query张量的子请求处理：丢弃前端Token带来的挑战是Query张量中的Token对应于上下文中两个不相连的范围。为了解决这个问题，系统将这两个范围视为恰好共享底层上下文部分内容的两个子请求。图8（步骤d）展示了每个子请求的Query张量及其对应的KV上下文位置。由于多Token注意力GPU内核设计支持批处理中不同请求的可变长度Query张量，并且接受非连续的KV-token位置，Pensieve只需更新辅助数据结构，在处理子请求时不会产生内存复制。

生成期间挂起请求：尽管有提前驱逐，由于请求的解码长度未知，调度器在Generation阶段仍可能遇到GPU缓存耗尽的情况。此时，调度器通过将某些请求移出当前批次来暂停其执行，将其KV-tokens换出到CPU，并将其放回等待队列。调度器根据到达时间的降序选择要暂停的请求。为了避免暂停带来的延迟增加，调度器保守地保留 $10\%$ 的GPU缓存槽供处于Generation阶段的现有请求执行。

针对非连续缓存的多Token注意力设计：为了将GPU缓存中现有的KV-tokens与刚从CPU换入的KV-tokens结合，分配连续内存会产生昂贵的内存复制。vLLM的PagedAttention内核可以处理Generation阶段的非连续KV缓存，但不能用于Prefill阶段，因为它假设批处理中的每个请求正好有一个输入Token。系统构建了一个多Token注意力内核，支持在每个请求的多个输入Token的Query表示（Q）和分布在非连续内存上的KV缓存之间执行注意力计算。

多Token注意力内核运算原理：新内核与vLLM内核的主要区别如图9所示。vLLM的内核执行单个输入Token与现有KV-tokens之间的注意力计算，底层计算可描述为两个矩阵-向量乘法。相反，新内核同时处理每个请求的多个输入Token，计算所有输入Token与对话现有KV-tokens之间的注意力分数，底层计算可描述为两个矩阵-矩阵乘法。由于Q张量中存在额外的维度，该内核在GPU上有更多的并行化和分块机会，但也必须处理Q是一个参差不齐大小的张量（不同请求输入Token数不同）的挑战。
图9 单Token注意力与多Token注意力的对比，省略了softmax和重缩放。$d$ 是嵌入向量的维度。多Token注意力分数中的红色阴影区域代表因果掩码。$K$ 和 $V$ 在逻辑上是连续的，但在使用分页内存管理时物理上放置在随机位置。

融合因果掩码与内核实现：当请求的多个输入Token在内核中一起处理时，需要应用因果掩码，以便早期的输入Token不会“关注”后期的Token。这要求将注意力分数矩阵的相应上三角部分设置为0（图9红色阴影区域）。系统将因果掩码操作融合在多Token注意力内核内部，以避免物化中间的注意力分数矩阵 【10，FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness + 2022 + NeurIPS】。系统基于PyTorch现有的多Token注意力内核进行扩展以处理非连续KV缓存，并使用NVIDIA Cutlass模板库提供的块级矩阵-矩阵乘法 【33，CUTLASS + 2017 + GitHub】。

统一Prefill和Generation阶段的执行过程：在Pensieve的批处理形成中，处于Prefill阶段的新请求可以与处于Generation阶段的现有请求分组。这是因为Generation阶段执行的单Token注意力可以被视为Query大小等于1的多Token注意力的特例。Pensieve调度器连接所有请求要处理的输入Token，并维护辅助数据结构来跟踪每个请求对应的区域。在工作节点执行期间，这些批处理的输入Token通过线性层生成每个Token的QKV表示。新生成的KV-tokens存储在GPU KV缓存分配的槽中。然后工作节点应用多Token注意力内核为所有请求生成输出Token。

多GPU支持机制：对于无法装入单个GPU的大型LLM模型，系统采用在Megatron-LM 【45，Megatron-lm: Training multi-billion parameter language models using model parallelism + 2019 + arXiv】中实现的张量并行化，将大型模型划分到多个GPU上。Pensieve中的KV缓存也相应地被划分。每个模型分区由一个工作节点进程管理，该进程拥有自己的GPU和CPU缓存分区来存储分配给它的注意力状态分区。由于模型划分发生在KV缓存的特征维度上，它不影响决定迁移或丢弃哪个Token的驱逐策略。因此，每个工作节点遵循相同的迁移计划。

实现细节与优化：Pensieve原型使用约7000行C++/CUDA代码实现。系统依赖PyTorch (v2.0.0, CUDA 11.8) C++前端API来执行LLM中的GPU算子。系统开发了融合多Token注意力内核。在优化方面，系统发现当CPU到GPU的数据传输与GPU到CPU的数据传输并发进行时，两个方向的吞吐量都会显著下降（$18-20\%$）。为了防止驱逐操作减慢过去KV-token的换入速度，系统设置了等待机制：如果工作节点有任何正在进行的换入任务，它会等待执行GPU到CPU的复制，直到换入任务完成。

方法细节引用的参考文献汇总：
* 【5，Language models are few-shot learners + 2020 + NeurIPS】：在“LLM与注意力机制”段落中提及GPT-3模型时引用。
* 【55，Opt: Open pre-trained transformer language models + 2022 + arXiv】：在“LLM与注意力机制”段落中提及OPT模型时引用。
* 【46，Llama: Open and efficient foundation language models + 2023 + arXiv】：在“LLM与注意力机制”段落中提及Llama模型时引用。
* 【48，Attention is all you need + 2017 + NeurIPS】：在“LLM与注意力机制”段落中说明基础模型架构时引用，指出流行大模型基于Transformer架构。
* 【35，FasterTransformer + 2021 + GitHub】：在“LLM的服务方式”段落中引用，作为采用典型LLM推理过程的系统代表。
* 【52，Orca: A distributed serving system for Transformer-Based generative models + 2022 + OSDI】：在“LLM的服务方式”段落中引用，作为采用典型推理过程及迭代级批处理策略的系统代表。
* 【27，Efficient Memory Management for Large Language Model Serving with PagedAttention + 2023 + SOSP】：在“LLM的服务方式”段落中引用，作为支持动态分配非连续缓存及PagedAttention内核的系统代表。
* 【10，FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness + 2022 + NeurIPS】：在“融合因果掩码与内核实现”段落中引用，说明将因果掩码操作融合在内核内部以避免物化中间矩阵的技术来源。
* 【33，CUTLASS + 2017 + GitHub】：在“融合因果掩码与内核实现”段落中引用，说明内核底层使用了NVIDIA Cutlass模板库。
* 【45，Megatron-lm: Training multi-billion parameter language models using model parallelism + 2019 + arXiv】：在“多GPU支持机制”段落中引用，说明张量并行化方案的实现参考。

实验环境

  • 硬件配置：评估在Azure NC A100 v4系列上进行，配备多达4个A100-80GB GPU，24核AMD EPYC 7003处理器，每个GPU 220 GB CPU内存。为公平比较，每个系统配置在每个GPU上分配40 GB内存用于KV缓存。
  • 模型架构关键参数：使用OPT和Llama 2开源模型。评估了单卡小型模型（OPT-13B，Llama 2-13B）和使用张量并行化在4个GPU上分区的大型模型（OPT-66B，Llama 2-70B）。为展示Pensieve对分组查询注意力（GQA）的有效性，将Llama 2-13B的KV头数从40修改为10。所有实验中，模型参数和中间隐藏表示均采用16位半精度浮点格式。
  • 数据集名称、规模及用途：使用ShareGPT（真实ChatGPT对话，48,159个对话，平均5.56轮，用于模拟真实多轮场景）和UltraChat（合成数据集，1,468,352个对话，平均3.86轮）。最大上下文限制为16384个Token。
  • 软件配置：基线系统为vLLM (v0.2.0) 和 TensorRT-LLM (v0.12.0)。vLLM使用PyTorch作为执行后端，TensorRT-LLM使用图重写优化并使用TensorRT Runtime执行。还测试了仅使用GPU缓存的变体Pensieve (GPU cache)。

实验结果

  • 端到端服务性能 (单GPU)：测试了OPT-13B和Llama 2-13B在ShareGPT和UltraChat数据集上的归一化延迟与吞吐量。结果显示，在ShareGPT上，Pensieve在OPT-13B上的吞吐量是vLLM的 $1.36\times$，TensorRT-LLM的 $1.14\times$；在Llama 2-13B上的吞吐量是vLLM的 $1.70\times$，TensorRT-LLM的 $1.58\times$。Pensieve在ShareGPT上的增益大于UltraChat，因为真实数据集包含更多对话轮次。对于Llama 2-13B的性能增益更显著，因为它使用了GQA，减少了KV-tokens的内存占用，使Pensieve能存储更多历史记录。图表引用：图10。
  • 多GPU服务性能：在ShareGPT数据集上测试了四卡运行的OPT-66B和Llama 2-70B。结果表明，大模型放大了Pensieve的优势，因为计算量的增长快于KV缓存内存的增长。Pensieve在OPT-66B上的吞吐量达到vLLM的 $2.04\times$ 和TensorRT-LLM的 $1.64\times$；在Llama 2-70B上达到vLLM的 $3.0\times$ 和TensorRT-LLM的 $2.47\times$。图表引用：图11。
  • 多Token注意力内核性能：对比了理想情况（连续内存）、分配连续内存复制（CopyOut + Attention）、多次调用vLLM单Token内核（Multi-round PagedAttention）以及Pensieve内核的执行时间。结果显示，两种基线方案引入了显著开销，而Pensieve内核匹配了理想基线的性能，甚至因将辅助数据计算卸载到CPU而表现略好。图表引用：图12。
  • 统一调度的影响：评估了Llama 2-13B上分离处理与统一处理Prefill和Generation阶段的差异。结果表明，统一调度避免了以少量请求执行Prefill阶段的情况，从而实现了更好的吞吐量和延迟。图表引用：图13。
  • 驱逐策略的影响：在OPT-13B上对比了Pensieve缓存策略与经典LRU策略。结果显示，当工作负载超过3 req/s时，Pensieve策略优于LRU。Pensieve策略的CPU缓存命中率比LRU高出最多4.4个百分点，重算的KV-tokens数量平均减少了最多 $14.6\%$。图表引用：图14。
  • 用户思考时间的影响：使用Llama 2-13B评估了不同平均用户思考时间（30s至600s）对性能的影响。结果表明，随着思考时间增加，吞吐量下降，因为过去的KV-tokens以更高速率从缓存中丢弃。但即使思考时间增加到600秒，Pensieve的延迟和吞吐量仍优于vLLM。图表引用：图15。

补充细节

相关工作：
* LLM推理与服务系统：近期开发了许多优化系统（如vLLM, ORCA, TensorRT-LLM, DeepSpeed等），它们研究了与Pensieve不同的优化机会，如增量解码、迭代级批处理、内核融合、推测解码和量化。DistServe主张在不同GPU上分离Prefill和Generation阶段以满足SLA，而Pensieve执行统一批处理以优化吞吐量。
* 缓存LLM注意力状态的系统：与Pensieve同期的系统（PromptCache, SGLang, RAGCache, CachedAttention）也缓存注意力状态。PromptCache需要用户预先定义Schema；SGLang和RAGCache采用基于树的缓存以共享前缀，且优先从树底部（即尾端）驱逐。CachedAttention在整个对话粒度上驱逐且不重算被截断的上下文。Pensieve专为多轮对话设计，其驱逐策略在前端的细粒度Token组层面进行以降低重算成本，并通过其多Token注意力内核解决非连续区域计算的挑战。CacheGen优化远程网络存储的上下文加载延迟，与Pensieve使用本地缓存的技术正交。
* 解决GPU内存限制的技术：针对DNN训练的GPU-CPU交换技术（如SwapAdvisor, ZeroOffload）或针对弱GPU推理的系统（DeepSpeed-ZeroInference, FlexGen）主要面向延迟不敏感的应用，且不在请求间持久化KV缓存。DNN训练中的重算技术通常在层或算子粒度进行，而Pensieve在Token(Chunk)级别进行细粒度重算以利用因果注意力特性。Pensieve未采用统一内存（Unified Memory）或直接主机访问（Direct Host Access），因为系统需要显式管理数据移动以预取KV缓存，且直接主机访问在生成步骤重复使用缓存时效率显著降低。

结论

在为多轮对话提供大型语言模型服务时，一个主要的低效源于对话累积历史上下文的重复计算。论文开发了Pensieve，这是一个有状态的LLM服务系统，它将历史嵌入保存在多级GPU-CPU缓存中。它使用新的GPU注意力内核，在请求的新多Token输入与其存储在非连续GPU内存中的已保存上下文之间执行注意力计算。实验表明，对于小型模型（OPT-13B, Llama 2-13B），Pensieve实现了基线系统 $1.14 - 1.70\times$ 的吞吐量，对于大型模型（OPT-66B, Llama 2-70B），则达到了 $1.64 - 3.0\times$ 的吞吐量。