2024年，我们已经进入大模型全面爆发的时代，作为大模型很重要的工程实践: 推理服务，则成为熟悉和了解大模型工程的关键一环。vLLM是23年开始出现的一款较为优秀的大模型推理框架，很值得学习和研究，我将发布一系列的Blog，针对近期学习vLLM的主要内容，通过图解的方式从工程和算法角度进行总结。本篇文章主要关注于vLLM的原理与整体架构，构建一个全貌，后面的文章会分层逐步细化代码具体实现。
1. Transformer架构模型与KVCache
目前大模型都是基于Transformer架构进行设计，transformer的核心能力是其对输入上下文进行SelfAttention，从而获得跟输入信息最相关的信息量，作为下一层step的输入，最终通过一些forward和normalization，得到最终的概率分布。
1.1 SelfAttention
自注意力可以用如下公式1进行表达
$Attention(Q, K, V)=softmax(\frac{QK^T}{\sqrt{b_k}})\quad\quad\quad(1)$
以上公式可进一步进行分解：
- 将输入的序列 映射到三个向量空间中，分别得到三个向量$q, k, v$, 对应公式2
- 计算$q, k$向量匹配度，得到注意力得分$a_{ij}$，表示第 i 个位置的 token 和第 j 个位置的 token 的匹配度得分。各注意力得分组成一个向量$a_i$，表示第$i$个位置的 token 和以前其余位置所有token的得分向量
- 注意力得分向量$a_i$与向量$v$做内积，得到SelfAttention的最终输出$o_i$, 对应公式3
$q_i = W_qx_i,\quad k_i = W_kx_i,\quad v_i = W_vx_i,\quad\quad\quad(2)$
$a_{ij} = {exp(q^T_i/\sqrt{d}) \over \sum_{t=1}^iexp(q^T_t/\sqrt{d})},\quad o_i = \sum_{j=1}^i{a_{ij}v_j}\quad\quad\quad(3)$,
上图就是Transformer中，公式2和公式3的具体求解过程图解
1.2 KVCache
通过公式3和图解可知，对于SelfAttention来说，在计算$a_{ij}$时，除了需要公式2计算得到的$q_i, k_i, v_i$，还需要$k_i, v_i(j < i)$。由于之前迭代生成第$j$个token时已经计算过了$k_i, v_i(j < i)$, 并且后续迭代生成可以直接复用这部分$k_i, v_i$，所以Transformer相关模型在所计算时，都使用缓存的方式，将计算过的$k, v$向量缓存到显存中，避免每次迭代重复计算这两个向量, 因此向量的缓存就是大家熟知的 KVCache了。
2. 推理服务优化目标
根据上文，如何利用的KVCache原理，用最少的资源获得最大的吞吐量，就是推理服务突破瓶颈的关键目标。
在推理过程中，主要分为两个阶段: Prefill 和 Decode 阶段
- Prefill: 即用户输出prompt阶段，推理服务器会根据prompt解析出token序列，然后使用token序列去初始化GPU显存中的KVCache数据，并预测下一个生成的token。
- Decode: 接下来针对产生的新token，进行循环迭代，直到达到最大生成长度或者终止符为止；迭代生成每个 token。在生成第$t+1$个token时，需要将 prompt token、已生成的 token 及当前第$t$个token的KVCache拼接起来，与第$t$个token的query vector完成SelfAttention等计算。完成当前迭代的token生成时，若生成的token为终止符或者生成的长度到达最大长度，则停止迭代。
在 LLM 服务中，由于Decode阶段生成token时依赖之前生成的token，无法并行生成多个token，只能串行迭代生成。因此会影响整体推理服务吞吐。那么如何提升整体吞吐量，我们先了解一下GPU内存分布。显存占用主要由三部分组成：模型权重、KVCache 以及其他激活值，KVCach 和激活值显存开销又与推理时的batch size成正比。
在论文中，以A100 40GB为例，LLaMA-13B BF16 模型的显存占比。模型权重有26GB，其余显存均用来存储KVCache和激活值。设置最大batch size使KVCache和激活值显存占满剩余显存空间，那么KVCache的显存占比大于30%，激活值显存占比不及5%
所以想提高吞吐量，就是要提高KVCache的内存占用大小，我们来看一下内存占用的公式:
KVCacheMemory = 2 * batch_size * max_seq_len * num_layers * hidden_size * num_heads * dtype_size
从公式中可以看出，KVCache显存开销与batch_size、num_layers、max_seq_len、hidden_size、num_heads以及数值类型大小六个维度相关。在KVCache 显存开销上限一定的情况下，为了提升单次推理最大合并请求数，也就是batch_size的值，可通过降低其他几个维度的值来提升。因此我们可以尝试如下几个方面进行优化:
2.1 数值类型量化
训练时通常使用FP16数值类型，在推理阶段可使用int8量化，降低数值类型大小。如使用KVCache int8量化, 这会使batch_size增加一倍。
2.2 模型精简
num_layers、hidden_size、num_heads：理论上可以通过剪枝等稀疏化的方式，降低这两个维度的值, 这个并不是本文的重点内容，暂时不深入探讨。
2.3 Batching
我可以在单次请求中尽量多的放入prompt request序列，同时也在某些request完成输出后，及时的插入新的reqeust，填满batch size，具体可以通过下面图示理解
- 延时方面：可以大幅度降低请求端到端延时。由于队列的请求在Decode阶段动态地插入到 Batch，所以请求的队列等待时间从原来的秒级别下降到毫秒级别（仅需等待一轮迭代即可插入到Batch中）。
- 资源利用率方面：提升GPU资源有效利用率。由于队列的请求在Decode阶段每一轮结束后动态地插入到Batch，而Batch也在迭代过程中动态地删除已完成推理的请求，这样避免了GPU空转，提升资源利用率。
2.4 PageAttention
根据上文了解到，max_seq_len是token的最大长度，是一个固定值，gpt2中的大小是2048。在Prefill阶段，KV向量初始化操作一次性的。因此如果实际按照max_seq_len的大小进行KVCache的实际分配，必然会提前占用很多还未使用的空间，显存利用率会明显降低。
vLLM框架论文中提出了PageAttention，原理是借鉴操作系统对虚拟内存设计，也对KVCache进行划分，每个KVCache块包含若干个固定tokens的 key、value向量, 称之为一个Block。vLLM会维护一个block_table, 这是一个逻辑Block和物理Block的映射表。逻辑上看，多组Block在虚拟空间上是连续的，但从实际分配上看，这样我们可以对物理显存上的Block添加管理策略，每个Block在物理显存上可以动态的、不连续的实时分配；当物理内存不足时，又可以通过和CPU 内存进行灵活交换物理显存内容，从而高效利用gpu显存，增加并发的prompt数量，提升整体吞吐量。
