极智AI | 大模型优化技术PagedAttention

1,362 阅读3分钟

欢迎关注我的公众号 [极智视界],获取我的更多经验分享

大家好,我是极智视界,本文来介绍一下 大模型优化技术PagedAttention。

邀您加入我的知识星球「极智视界」,星球内有超多好玩的项目实战源码下载,链接:t.zsxq.com/0aiNxERDq

之前写过一篇文章是关于 KV Cache 优化的,是《大模型优化之KV Cache》,这里这个 PagedAttention 里面做的事情其实就是在优化 KV Cache,而且跟之前文章中提到的方法类似。那为什么要单独再拎出来说呢?主要还是因为大模型部署领域有很多专业名词,还是有必要拿出来说说它们的来龙去脉的。

首先需要明确的是,在 KV Cache 内存管理中,一般会存在三种形式的浪费:reserved fragmentation(预留碎片)、internal fragmentation(内部碎片)、external fragmentation(外部碎片),如下,

而说到 Attention,最熟悉的应该还是那一句 "Attention Is All You Need.",确实,回过头来看,这句话引发了人工智能的新革命。Transformer 中最重要的结构应该就是 Attention 注意力机制了。对于 PagedAttention,其实从字面意思来理解,它的概念也会大差不差。Paged 是受计算机操作系统中内存分页的思想,把这个思想应用到大模型中 Attention 机制其中 KV Cache 优化,这两者一结合就是 PagedAttention。

具体来说,PagedAttention 将 Sequence 中 KV Cache 划分为 KV 块,就像内存分页那样,每个块 (每个页) 的容量是固定的。对于 PagedAttention 中的 KV 块来说,每个块可以存固定数量 tokens 的 key (K) 和 value (V) Tensor,从而将注意力机制的计算转换为块级的计算,如下,

现在的大语言模型主要是预测下一个 token,PagedAttention 是将 KV Cache 以块的方式存储起来,每个块就是一个 Block,每个 BlockSize 一般为 16。如下 ,这个例子中,KV 向量分布在三个 blocks 上,但是这三个 blocks 在物理内存上并不一定是连续的。比如在某个时刻,Query 向量 qi "forth" 和 blocks 中的 Key 向量 Kj "Four score and seven" 去计算 attention score Aij,然后和 blocks 中的 Vj 去计算输出的 token Oj,中间涉及的计算方式就是采用上面的公式。

基于上面的 KV Cache 管理,形成了如下的 LLM 内存管理系统,采用一个 centralized scheduler 来统筹、分配不通 GPU 工作站的内存管理。中间核心的 KV Cache Manager 在每个 paged fashion 中高效地管理 KV Cache,而具体的管理策略就是使能于 PagedAttention。

在 PagedAttention 中会将 KV blocks 从概念上分为 逻辑KV blocks (Logical KV blocks) 和 物理KV blocks (Physical KV blocks),如下,

KV blocks 实际存储的物理内存可能并不是连续的,但是 PagedAttention 会将它们的逻辑 KV blocks 组织为连续的,这样便于上层的使用,这中间会涉及到 KV blocks 的内存映射 (mapping),这种内存映射有点像链表的数据结构。这样做的好处有挺多,一方面可以高效地管理动态增长的 KV Cache 内存,因为增加的内存是不定的,如果是连续开辟内存难免会产生多余的浪费;另外还可以很好地将琐碎的内存片段利用起来,这也是为什么 BlockSize 设定为 16 的一部分原因。

好了,以上分享了 大模型优化技术PagedAttention,希望我的分享能对你的学习有一点帮助。



 【公众号传送】

《极智AI | 算法部署中需要注意的Lazy Loading》


畅享人工智能的科技魅力,让好玩的AI项目不难玩。邀请您加入我的知识星球, 星球内我精心整备了大量好玩的AI项目,皆以工程源码形式开放使用,涵盖人脸、检测、分割、多模态、AIGC、自动驾驶、工业等。不敢说会对你学习有所帮助,但一定非常好玩,并持续更新更加有趣的项目。 t.zsxq.com/0aiNxERDq

logo_show.gif