Paged Attention&Orca Paper笔记与解读

序

ORCA: A Distributed Serving System for Transformer-Based Generative Models

Orca的目标问题

Orca的Solution

• 早结束；意味着一个batch的请求时长 跟随max_output_length的请求一起结束，延长了所有请求的完成时间

• 晚到等待：意味着下一个batch的请求，其等待时间也是上一个batch中最长请求的完成时间。

• 因此假设是naive的batching思路，“集合一批，处理一批”，latency取每一批中的最长请求为基本单位，且存在计算资源的闲置和内存的空占。

• iterative-level的思路也很简单，自回归本身即循环调度f(x)，因此不必以request为最小颗粒，而是可以以一批对f(x)的调度为batching；那就可以将输出长度为n的请求和长度为m的请求视作(m+n)次调度，当某请求完成时，即退出调度；新的请求可以进入调度。

• 在训练或request-granularity的推理中，计算之所以可以batching，是因为self-attention计算框架约定了[batch_size, max_token_length, hidden_size]这样的tensor shape。此处可以分别做一个讨论，在问题框架上，将推理prefill和decode by step。
• prefill阶段的naive batching可以认为是将token长度简单的padding到max_token_length可以解决
• decode by step阶段，应用简易的kv cache思想，q是当前新增token的一维矩阵，多个请求batching，q的数量可以放置在batch维度中，设置为[n, 1, hidden_size]；K, V为前序token缓存二维矩阵，其shape为 [n, ?, hidden_size] —— 在iterative-level的调度中，K, V的长度不尽相同；若推理时，仍然按照潜在可能的max_token_length对齐，内存空间的使用依然是按照潜在最长输出占用。

• 当然latency上有优化，允许已完成的请求退出，允许等待的请求开始。理想的状况是，一个bacth中的多个请求，内存占用与真实的序列长度线性增长。

• selective-batching：思考batch_size维度在计算上存在的意义，先不考虑multi-head这个因素。同一个batch内，除了scaled-dot attention的QK计算需要行列交互，其他的计算如MLP，softmax，layer norm——更不用说element-wise，都是在行内的操作。因此若假设两个长度不一的请求（两段prompt），假设长度为s1, s2；在原始计算框架下需要合并成两个batch，[2, max(s1, s2), hidden_size]；而在selective-batching方法中，可以合并成[1, s1 + s2, hidden_size]。在non-request-aware operation中，可以当做一个请求处理；在执行attention时，再拆成两个请求（也许是[2, max(s1, s2), hidden_size]）。

• selective-batching不仅可以合并同样在prefill阶段的不同长度请求，在另一个方面也让prefill阶段和decode by step阶段可以合并计算；prefill阶段的输入是长度为任意s的用户promot；decode by step阶段输入是长度为1的last token配合有状态的kv cache manager。在计算时，系统只需要在attention时，能够区分(1) 不同请求的划分 （2）不同请求的阶段 - 决定是否需要查询kv cache参与计算。

• 🤔思考：在selective-batching的框架下，是否需要考虑[X, hidden_size]的形状对GPU计算效率的影响。X可能会变得远大于hidden_size。

• Orca遵循先进先出（文中称FCFS，first come first serve）的原则；对请求到达时间排序，并以内存资源为优先的瓶颈，填充资源
• n_rsrv可以理解为KV cache使用的连续内存空间；每个request的prefill被执行时，会扣除request.max_token这么大的资源预占（这也是paged attention要解决的问题）；而因为按到达序排序，因此prefill之后的请求会稳定地出现在之后批次的serve batch里，直到请求结束，占用的KV cache slots被释放。
• PagedAttention主要优化的问题，即KV cache的生命周期管理与分布方式 & 具体地与decoding algo结合的使用场景。

Efficient Memory Management for Large Language Model Serving with PagedAttention

PagedAttention的目标问题

PagedAttention的方法

On top of it, we build vLLM, an LLM serving system that achieves (1) near-zero waste in KV cache memory and (2) flexible sharing of KV cache within and across requests to further reduce memory usage.

1. 如何定义一页/一块（page or block）
• OS VMM模块管理的是一一映射的对等概念，虚拟内存与物理内存单位都是byte，只是按一定的page size划分；无论OS层面运行什么样的软件，one page就是4kb/8kb/16kb等整数。而vLLM的管理对象是KV cache，是一个有具体的计算含义的变长实体。具体来说，一个Transformer-based的模型有attention layer count, hidden_size, head_size，其他还有一些算法变体等参数会影响单个token的KV Cache到底多大。此外无论如何，单个token的所有KV cache放在一起过于巨大（100kb~10MB之间），理应有一个“切割”的思考，按类（K or V）切 or 按层切 or 两个维度都切分，也是一个因素。

2. 如何trade off非连续的KV cache存放机制与实际计算时内存连续的要求
• 当KV cache以物理上非连续的方式在逻辑数据结构中被索引时，前序token的KV矩阵就不再是可以直接被GPU矩阵乘法计算的对象；不像Orca中申请max_tokens为基数的连续内存方便。

3. HBM不足时，evict策略是什么
• Orca的连续内存策略能够“悲观”地保证进入bacth的请求，有足够的内存资源能够完成选中的请求；但虚拟内存的方式，只在申请物理内存时，“乐观”地事后知晓物理内存分配不足；（1）若没有evict策略，直接请求失败，则会导致类似惊群或死锁的问题，请求集体失败，集体重试，再集体失败（2）若有evict策略，调度系统逐出一些潜在低效的占用，那么有若干个问题：evict的时机是什么，慢介质是什么，选择策略是什么，数量怎么决定

4. 如何与类似beam search的搜索路径配合，复用的策略

• 引入sequence group的概念，意即完成一个请求所带来的需要KV cache，根据解码算法不同，一个request会产生一到多条推理sequence; 本文采用了all or nothing的策略，在sequence group的维度上，要么逐出sequence group的所有的KV cache，要么不逐出任何一块。因为vLLM假设不存在一种有收益的场景叫做，部分KV cache存在，可以先进行部分的计算。这种比iteration-level更为复杂的计算机制，超出了Orca/vLLM等最细粒度的框架假设。

• 解法一：逐出到CPU内存中；解法二：清除内容，后续重新计算，重新计算的速度是不慢的，重新获得kv cache就是经历一次新的prefill阶段；；需要考量内存显存带宽 pk 计算能力

• 请求调度的额外策略：当逐出发生时，系统就不再接受新的请求，直到先前被抢占的任务完成。这也是一个FCFS的原则（那么在被抢占任务完成前，系统是否有一段时间是跑不满的？）；文中讲到这获得了一个额外的特性，CPU RAM作为缓冲区，容量需求不会超过GPU HBM大小。

• 并行decoding采样。并行decoding采样定义：N个互不相关的随机采样序列以特定的随机参数运行。在这个场景下，request prompt是确定的N个sequence的prefix，token长度为L的prompt会产生至少向下取整(L/block_size)个完整的block，这些block在request生命周期中被每一个sequence的每一次generation复用；最后 < block_size个prompt的KV cache填不满一页，会被copy N份用于后续采样的一部分。

• beam search采样过程中，更复杂的点在于，路径会动态地生成与销毁。prompt也是确定可以被多个sequence复用的；进行过程中对token产生的block进行引用计数；beam search with k的设定下，t时刻下，m（m≤k）条路径依赖来自t-1时刻产生的前缀，那么reference count记为m。在t+1时刻，m条路径会派生出m*k条新路径参与下一轮概率选择，其中有[0, k]个路径会生出，由此也会决定，在t时刻，reference count是[0, k]中的一个数字。若某block的reference count降为0，则释放。

• 这两类采样算法，在PagedAttention中的原理完全是一致的。通过reference count + copy on write实现了复用，区别只是fork和reference count的操作时机不同。前者的fork只存在第一个token生成时（若token length与block size是整数关系，则没有fork过程）；后者的fork在行进中随时会分叉。前者无需管理reference count；后者需要通过reference count管理释放时机。

总结