PagedAttention

smartcat2010

已于 2024-09-10 20:23:12 修改

阅读量44

点赞数

文章标签：计算框架

于 2024-09-10 08:57:47 首次发布

原文链接：https://zhuanlan.zhihu.com/p/691038809

版权

用于推理阶段，节约显存；多request并发处理；

中文讲解：

https://zhuanlan.zhihu.com/p/691038809

13B模型的显存占用量：

Others: 主要是activation；推理阶段不需要保存下来；

Prefill阶段：

Decode阶段：

随着prompt数量变多和序列变长，KV cache也变大，对gpu显存造成压力

由于输出的序列长度无法预先知道，所以我们很难提前为KV cache量身定制存储空间

在常规的推理框架中，当我们的服务接收到一条请求时，它会为这条请求中的prompts分配gpu显存空间，其中就包括对KV cache的分配。由于推理所生成的序列长度大小是无法事先预知的，所以大部分框架会按照(batch_size, max_seq_len)这样的固定尺寸，在gpu显存上预先为一条请求开辟一块连续的矩形存储空间。然而，这样的分配方法很容易引起“gpu显存利用不足”的问题，进而影响模型推理时的吞吐量。

PagedAttention

和操作系统的虚拟内存分页机制一样。

请求（request）可理解为操作系统中的一个进程（可并发处理多个request)
逻辑内存（logical KV blocks）可理解为操作系统中的虚拟内存，每个block类比于虚拟内存中的一个page。每个block的大小是固定的，在vLLM中默认大小为16，即可装16个token的K/V值
块表（block table）可理解为操作系统中的虚拟内存到物理内存的映射表
物理内存（physical KV blocks）可理解为操作系统中的物理内存，物理块在gpu显存上，每个block类比于虚拟内存中的一个page

同时处理多个request:

默认：16个token的KV，作为1个block;

Parallel Sampling:

例如：给模型发送一个请求，希望它对prompt做续写，并给出三种不同的回答。

这个场景叫parallel sampling。在这个场景中，我们可以将prompt复制3次后拼接成1个batch喂给模型，让它做推理。但这种方式会产生prompt部分KV cache的重复存储。

解决：复用公共文本的block；一半公共一半不同的，copy-on-write;

(TensorRT-LLM的--enable_kv_cache_reuse选项，就是干这个的）