基于KV Cache卸载的推理系统设计
《Stateful Large Language Model Serving with Pensieve》基于vLLM的卸载设计。
《Cost-Efficient Large Language Model Serving for Multi-turn Conversations with CachedAttention》提出的策略比较基础，是基本的流水线重叠和预取思想。
《Accelerating LLM Serving for Multi-turn Dialogues with Efficient Resource Management》提出重排执行和长请求抢占的机制来解决队头阻塞。
Pensieve
Pensieve针对的是以下几个问题
- 对于KV Cache，显存空间是有限的，因此要利用起来CPU内存
- Token级别的缓存管理和恢复问题
- 对于CPU与GPU之间的KV缓存交换，需要非连续的KV缓存管理
Pensieve scheduler的作用主要有两个，一个是要构建running batch，一个是确保batch中的请求能够得到充足的显存来运行。
支持混合阶段Batch
Pensieve scheduler支持在同一个batch中同时包含Prefill阶段和Decode阶段的请求，通过这种组合，避免了分开运行较小的Kernels，提高了GPU的利用率。这一组合依赖于其实现的matrix-matrix注意力计算Kernel。
复合倾向的驱逐策略
对于Cache的驱逐策略，Pensieve结合了两种倾向，首先是最基本的LRU驱逐策略，另一个是优先移除上下文开头的词元缓存(因为这一部分的重计算成本更低)。在Pensieve中，按照32 tokens的大小构建chunk, 在进行驱逐时，会估计每个chunk的保留价值$V$，其计算公式为$V = \frac{cost(s, l)}{T}$，其中$T$表示该chunk自从上一次活跃过后的时间，$cost(s, l)$则表示对于上下文大小为$l$、本身chunk大小为$s$(这里文中固定选取为32)的重计算量。
提前换出显存
此外，Pensieve还设计了提前换出的策略，当GPU空闲的显存低于一个百分比阈值时，会进行一次向CPU中的拷贝，同时GPU中的部分是采用了懒收回的方式，只有当scheduler要使用这部分显存的时候，才会进行释放。另一种情况是在decode阶段会出现GPU显存不足的情况，此时Pensieve会选择换出那些到达时间戳更大的请求的KV Cache到CPU，并将这些请求放回等待队列。
流水线式换入
Pensieve不会完全等待KV Cache换入，而是采用流水线的方式，用计算来掩盖传输成本。这是因为每一层的KV Cache只有在进行这一层的self-attention计算时才会使用，因此可以先换入浅层的KV Cache，利用其计算过程来掩盖传输。
Dropped部分重计算
对于那些被drop掉的KV Cache，调度器会从持久化存储的对话历史中拿到其原始文本，并将其加入到该对话的新的请求中，重新进行计算。但由于dropped文本和new request存在上下文不连续的问题，Pensieve是将其划分为两个子请求，dropped部分之和自己相应的KV Cache进行注意力计算，而New Prompt部分则是和Dropper + Cached Token + New Prompt三部分联合组成的上下文的KV Cache进行注意力计算。
CachedAttention
CachedAttention的架构是采用了三级架构，包括GPU的HBM显存、主机的内存和磁盘存储。
对于层与层之间的传输成本，CachedAttention同样采用通过计算与传输重叠的方式。考虑到磁盘较慢的传输速度，CachedAttention设计了一种scheduler-aware的预取方式。
此外，由于模型对于单次输入令牌数有限制，随着对话的不断进行，需要进行截断(Token truncation)，将一部分输入直接丢弃，并重新拼接。这导致输入中token的位置会发生变化，其旧的嵌入位置编码所对应的那些被存储的KV Cache失效。因此CachedAttention设计了与位置编码解耦的截断方法。
Layer-wise Pre-loading
大语言模型是由多个Transformer层构成，每一层都有自己的KV Cache，CachedAttention基于此特性，在计算当前层的self-attention时，提前从Host memory中加载下一层的KV Cache到HBM。
Asynchronous Saving
在经过Prefill和Decode后，需要从HBM中向Host Memory写入新的KV Cache。CachedAttention提出了一种异步的写入方式，当前层的计算完成后，在下一个层的计算过程中向Host Memory中写入新产生的KV cache。
Scheduler-aware Fetching
在CachedAttention，大部分的KV Cache被卸载到磁盘中。因此CachedAttention会从请求的等待队列中进行look-ahead prefetching，如果说这些等待的请求在Disk上发生了命中，那么就会将这些请求对应的KV Cache移动到host memory里面。
Scheduler-aware Eviction
当host memory和disk中的KV Cache存满之后，需要进行驱逐时，Cached Attention会根据等待队列中从头到尾的顺序进行评估。首先，基于内存和磁盘的存储大小计算出相应的驱逐窗口，对于等待队列中的请求，越靠近队尾位置的请求，其KV cache越容易被从host memory中驱逐到disk中，当disk中也满的时候，会优先驱逐那些靠近等待队列队尾的请求对应的KV Cache。
Decoupled KV Cache Truncation
CachedAttention这里借鉴了相对位置编码的思想，即注意力层在评估token之间的相关性时，不应该基于绝对位置编码，而应该是考虑词与词之间的相对位置关系。其中，RoPE是一个主流的相对位置编码方案，其简化理解就是将Q、K向量根据其绝对位置，使用旋转矩阵R进行旋转，这样旋转后Q、K的点积结果，仅依赖于其相对位置，其简化的公式理解就是$Score(i, j)=(Q_iR_i)\cdot(R_j K_j)$。
CachedAttention则是利用了这一点，其存储了那些未经过旋转矩阵操作的原始KV cache，当进行截断时，其选取截断后的原始KV Cache部分，并进行旋转操作得到新的位置编码后的结果用于推理。
FlashGen
FlashGen针对的是两个问题，一个是重计算成本，一个是FCFS调度带来的长请求队头阻塞问题。
FlashGen同样采用GPU HBM – Host Memory – Disk Storage三级架构。
GPU-only Design
在GPU上，一个请求完成后的KV Cache会被保留，但同时其也会被标记为Reclaimable，允许正在运行的请求占用这部分空间。
GPU-CPU Design
在引入CPU的Host Memory后，FlashGen在推理进行时会将新产生的KV Cache异步地逐渐拷贝到主机内存上，这样做法的好处是，在对GPU上的KV Cache进行reclaim的时候无需等待，因为其备份已经在Host Memory上完成了。
针对队头阻塞问题
FlashGen提出了一种重排执行的方法，因为多个请求占据同一个GPU时，会出现碎片化问题，导致一部分显存没有办法运行当前的长请求，必须等待显存释放，因此会出现阻塞情况。此时FlashGen会对排队中的其他请求进行扫描，找到一个最近的能够执行的请求，先将其放到GPU上去执行。这样就保证了设备的利用率，当新的显存被释放出来之后，FlashGen会让长请求抢占原本promoted的请求，让长请求到设备上运行，避免长请求饥饿。
