之前介绍了FlashAttention,现在来看PageAttention的相关内容。
在介绍PageAttention之前,需要对KV Cache有一定的了解。我们知道Attention的计算过程可以使用类似如下的公式进行计算:
$$ \text{Attention}(Q,K,V)=\text{Softmax}(QK^{T})@V $$
在上述过程中我们会计算K与V乘积,而对于序列来说,前面n个的值计算一次就不变了,为了加快计算,于是引入了缓存机制,这就是KV Cache的由来。
需要注意的是,KV Cache只应用于推理阶段,且只存在于Decoder解码器中,其会加大显卡内存占用。由于其需要提前申请过多空间从而造成内部碎片,而过小的空间无法分配给其他请求,从而造成外部碎片。
为了解决KV Cache随着token不断增长而占用的内存也线性增长的问题,于是就引入了PageAttention的概念。
在计算机操作系统中,页(Page)是可以按需加载的,通过动态的方式解决内存占用的问题。同理,对于LLM中的Attention,也可以利用类似的思想按需进行加载,从而解决KV Cache的问题。
PageAttention将连续的键值缓存KV Cache分割为固定大小的物理块,并动态管理,从而显著提升显存利用率和推理吞吐量。其主要分为3个步骤:
- 分块存储:将输入序列的KV Cache划分为固定大小的逻辑块,每个逻辑块通过块表映射到非连续的物理显存块上
- 动态分配:在推理过程,仅按需加载当前计算所需的物理块到显存,避免为每个请求预分配最大长度空间
- 写时复制:在并行采样(如生成多个候选输出)时,多个请求共享相同物理块。当某个请求需修改块内存时,系统自动复制该块并更新映射,确保数据隔离性。
类似地,在SGLang中还有Radix Attention,也是通过对KV Cache进行复用。此时,是不是觉得有些熟悉感,于是我们将大模型的问题在计算机基础课程中找到了解决的方法。
参考视频:
如果喜欢这篇文章或对您有帮助，可以：[☕] 请我喝杯咖啡 | [💓] 小额赞助
