[Paper Reading] InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

Paper Source: InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management(OSDI’2024)

Overview

上图给出了KV Cache存储在不同位置时的开销。

1. KV Cache可由GPU完全存放，没有数据在CPU和GPU之间传输，所需的只是GPU对KV Cache的读开销，可以忽略不记（GPU的显存高带宽），这种情况受限于GPU显存大小
2. KV Cache 被 offload 到 CPU 内存存储，所需为CPU与GPU之间的传输开销（可达上百GB）
3. 采用预取技术来提前传输完整KV Cache，部分传输开销可以和计算过程overlap
4. 只预取部分关键的KV Cache，使得传输开销可以被完全overlap

InfiniGen, a KV Cache management framework for offloading-based inference systems.

本文场景着眼于 offloading-based inference system 中的 KV Cache 占用问题，基于此，提出一种动态 KV Cache 管理机制，创新点在于利用前一层的 attention 输入来预取当前层的重要 tokens 的 KV Cache 并推测当前层的 attention score。

这个系统设计就主要包括预取模块、权重索引生成控制器、KV选择控制器、推理控制器和倾斜控制器等组件，下面依次介绍

Prefetching Module

Observation 1： The attention inputs of consecutive attention layers are highly similar in LLMs

可以看到，由于离群值的影响，$Tblock\_in_{i-1}$ 在x轴上的投影值较大，而由于做了 LayerNorm， $Attn\_out_{i-1}$ 和 $FFN\_out_{i-1}$ 相比之下在两类维度的数值都很小。

因此，相比之下，$Tblock\_in_i$ 主要受到 $Tblock\_in_{i-1}$​ 的影响，这也就是相邻block间输入的相似性，而又由于attention input 是 LayerNorm 之后的 block input，这也导致了相邻 attention layer 输入的相似性。

以上，我们解释了 observation 1 的由来，而基于这个 observation 1，文章提出，利用 Layer i attention input speculate Layer i 的 attention pattern。

需要注意的是 $Tblock\_in$​​ 是会一层层逐渐改变的，因此相距较远的 input 差异会比较明显，也因此，最终所推测出的 token 会呈现出一种好像 select 过的现象，这也是标题中Dynamic的来源。

Observation 2: the attention score highly depends on a few columns in the query and key matrices.

也就是说，attention 分数主要是由Q和K矩阵中的少数几列决定的，也就是所谓的 important tokens。

那么，我们找到 Layer i-1 的 important tokens，就可以据此推测出 Layer i 的 attention score。

这里，文章利用SVD（奇异值分解）进一步将 Q 和 K 矩阵加以倾斜（Skewed），让某些列的幅值更大，目的是为了让大幅值的列更少，同时这些列的表达能力更强。

具体而言，就是将 Q 和 K 与一个正交矩阵相乘，正交矩阵的转置就是自己的逆矩阵，故最终计算结果不变。

Prefetching Scheme

具体流程见上图，

1. Prefill之前，对Q和K矩阵进行前文提到的倾斜操作

2. Prefill阶段，从Q权重矩阵（$W_Q$）和 Key Cache 矩阵中选择若干重要列（token Ids），得到部分Q权重矩阵(Partial $W_Q$​)与部分Key Cache 矩阵（Partial Key Cache）

   

3. Decode阶段，使用Step 2 中预取的 Partial $W_Q$​ 和 Partial Key Cache计算得到推测的Attention Score