推理优化Survey

算子优化

多头注意力Multi-headed attention

相比于传统注意力只注重于一种分析，多头注意力则是通过将多个注意力头组合使用，来分析不同类型的依赖关系。其基础是传统注意力的基础构建块：缩放点积注意力

核心是Q，K，V三个向量：

Q表示问询内容，即“我正在寻找什么”
K表示所能提供的信息，例如书籍的标题
V则表示实际包含的信息，例如书籍的具体内容

计算过程：

匹配度计算：首先将Q向量和所有K向量相乘，计算出分数score，表示Query和所有Key的相似程度。
缩放： 将分数除以Key向量维度的平方根$\sqrt{dim_k}$，这一步的目的是为了防止维度较高的时候点积过大，导致Softmax梯度消失
Softmax归一化： 对缩放后的分数采用Softmax归一化，将其转为总和为1的注意力全中，这些权重表示输出时，应该从每个Value那取用多少信息
加权求和： 将注意力权重和对应的Value向量相乘并求和，得到最终的输出向量

即$Attention(Q,K,V)=Softmax(\frac{QK^T}{\sqrt{dim_k}})V$

在多头注意力中，则是将上述过程并行执行多次，其计算过程为：

线性投影，创建多个注意力头 对于同一组输入的Q，K，V矩阵，通过h组不同的线性变换来对它们进行投影，其中，h是注意力头的数量，从而得到$Head_i=Attention(QW^{Q_i},KW^{K_i},VW^{V_i})$
并行计算：采用这h个注意力头进行并行计算得到h个输出向量
拼接与线性变换 将h个输出向量拼接起来，最后通过一次可学习的权重矩阵$W^O$进行线性投影，将其维度压缩到可供后续前馈神经网络FFN处理。

对于多头注意力的优化，主要方向包括KV Cache优化(例如，量化)、剪枝(例如，去掉不重要的头)。

分组查询注意力 & 多查询注意力

分组查询注意力Grouped Query Attention是基于多头注意力机制的一种优化，其思想是不同的头划分为不同的组，组内共享相同的K，V投影，但每个头仍然有独立的Q投影。

多查询注意力Multi-Query Attention的优化则更为激进，让所有的头都共享同一组K，V投影，同时每个头有独立的Q投影。

稀疏注意力Sparse Attention

由于传统注意力机制的复杂度是$O(n^2)$，其计算复杂度很大。同时，序列中并非所有token对之间都需要计算注意力，因为很多位置的注意力其实是0。因此，稀疏注意力的基本思想是只计算那些重要的token对之间的注意力。

Band Attention带状注意力： 在注意力矩阵中形成一个宽度为w的滑动窗口，只关注滑动窗口内的token，因此在主对角线上形成了带状。
Dilated Attention扩张注意力： 以固定的间隔关注更远的token，跳跃式的采样
Block Sparse Attention块稀疏注意力： 将注意力矩阵进行分块，只计算某些块，忽略其他部分。

共享注意力 Shared Attention

通过在不同的组件之间共享注意力计算的技术，优化内存使用。

一种是跨层注意力共享，模型的不同层之间共享同样的KV Cache/注意力权重。

一种是跨头注意力共享，例如查询共享：模型的注意力头采用共享的Query，但是有不同的KV向量。

一种是跨时间步注意力共享，一般是用于自回归生成，用于注意力状态复用。

Top-k采样与Nucleus采样

先说一下为什么要采样，如果我们在生成token时永远选择最优的那个，很容易陷入循环，缺乏创造性。而通过引入随机采样，能够打破局部最优，提升模型的创造性。

Top-k采样的思想是，选择概率最高的k个tokens作为候选，而不是从整个词表中选择。其基本逻辑是：

获取概率分布
保留概率最高的k个tokens
将这些tokens的概率归一化
从中随机采样一个token

而nucleus采样的思想则是选择累计概率达到阈值p的最小token集合

推测解码 Speculative Decoding

对于传统的自回归解码，token是串行逐个生成的，缺乏并行性，效率较低。

推测解码的思想是通过一个小型草稿模型来快速批量生产候选的tokens，然后再由大型目标模型来进行并行验证。

用小型草稿模型来近似后续$b$个tokens: $\hat{x}{i+1},\dots,\hat{x}{i+b}$
使用大型目标模型并行计算$x_{i+1}=LLM(x_1,\dots,x_{i}),x_{i+2}=LLM(x_1,\dots,x_i,\hat{x}{i+1})$直到$x{i+b+1}=LLM(x_1,\dots,x_i,\hat{x}{i+1},\dots,\hat{x}{i+b})$
找到首个$j$使得$x_{i+j}\ne \hat{x}{i+j}$，返回确定的序列$x_1,\dots,x{i+j}$

Kernels

分块注意力

分块注意力的核心思想是不一次性计算整个注意力矩阵，而是将其切分为多个blocks，更好地利用GPU的并行能力。

比较传统的是Flash Attention技术，其结合了Tiling和Online Softmax的技术。对于传统注意力的计算，我们注意到，为了计算相应的Softmax值$softmax(x_i)=\frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}}$，这导致我们必须先计算出整个$QK^T$矩阵的值才能计算相应的Softmax值。这就导致了内存瓶颈，在缓存中我们无法存储下一个巨大的$QK^T$矩阵，必须要从慢速显存中读取。

Online Softmax的提出就是为了解决这一问题，其思想是将序列分块后，通过记录处理到第t个块的最大值和调整过后的指数和作为全局统计量，再用于计算最终的Softmax分数，更多详细的信息可以查看这篇专栏。这里只解释其核心思想，online softmax利用了同底指数的乘除等价于其幂次的加减的性质。用$m_j$表示到第$j$个元素的最大值，用$d_j$表示到第$j$个元素的指数和，在遍历过程中即可更新指数和为 $\begin{aligned} d_j &= \sum_{i=1}^j e^{x_i - m_j} \quad \text{//safe softmax,avoid overflow}\\ &=\sum_{i=1}^{j-1}e^{x_i -m_j}+e^{x_j-m_j}\\ &=\frac{\sum_{i=1}^{j-1}e^{x_i}}{e^{m_j}}+\frac{e^{x_j}}{e^{m_j}}\\ &=\frac{\sum_{i=1}^{j-1}e^{x_i}}{e^{m_{j-1}}}\times \frac{e^{m_{j-1}}}{e^{m_j}}+\frac{e^{x_j}}{e^{m_j}}\\ &=\sum_{i=1}^{j-1}e^{x_i-m_{j-1}}\times \frac{e^{m_{j-1}}}{e^{m_j}}+\frac{e^{x_j}}{e^{m_j}}\\ &=d_{j-1} \times \frac{e^{m_{j-1}}}{e^{m_j}}+\frac{e^{x_j}}{e^{m_j}}\\ &=d_{j-1} \times e^{m_{j-1} - m_j}+e^{x_j-m_j}\\ \end{aligned}$ 通过上述方式，一次遍历就可以得到所有的指数和和最大值，其对应的代码实现如下

void onlineSoftmax(float *src, float *dst, int data_len) 
{
  float old_max = -FLT_MAX;
  float d_j = 0.0f;
  for (int i = 0; i < data_len; i++) {
    float new_max = std::max(old_max, src[i]);
    d_j = d_j * std::expf(old_max - new_max) + std::expf(src[i] - new_max)
      old_max = new_max
  }
  
  // calculate the softmax
  for (int i = 0; i < data_len; i++) {
    dst[i] = std::expf(src[i] - old_max) / d_j;
  }
}

FlashAttention就是利用了OnlineSoftmax的技术，将softmax计算过程和与V矩阵的计算融合在了一起，避免了中间结果的存储。在外循环里，遍历$Q$的块；在内循环，遍历$K,V$的块，具体流程如下

计算当前块的原始分数$S_{ij}=Q_iK_{j}^{T}$存储在缓存中，不写回显存
对于$Q_i$块中的每个查询$q$，使用上述Online Softmax的规则，更新其查询的$m[q]$和$d[q]$，形成新的向量$\vec{m},\vec{d}$
根据新的统计量来修正之前的输出 $O_i^{new}=diag(e^{\vec{m}^{old}-\vec{m}^{new}})\cdot O_i^{old}+(e^{S_{ij}-\vec{m}^{new}})\cdot V_j$ 这里$diag()$的作用是将原本的向量转化为一个对角矩阵来和输出矩阵相乘。该式子中前半部分代表了对于旧输出的缩放，后半部分代表了新块计算的分数的贡献。内循环遍历完所有的$K,V$块后，所得结果就是$Q_i$对应的正确答案。

需要注意的是，FlashAttention是按照Q块进行分配，如果说序列长度不规则，容易出现负载不均衡的问题。Stream-K技术针对这一问题提出了更加细粒度的划分方式，将整个$Q K^T$的计算看作一个任务池，细分为多个子块的计算，所有的处理单元从中动态领取计算任务，最后通过精细同步来保持一致性。

Scheduling

经典的调度方式包括First-Come First-Serve, Shortest-Job First, Multi-Level Queuing(队列间具有不同优先级，队列内可设置不同策略)。

对于调度，其主要目的是实现负载均衡，一种简单的贪心策略就是每次都把新的工作负载交给当前负载最低的worker(这里的负载指的是makespan)。要执行这种策略，不可避免的需要负载预测来预估新负载所需的开销。除了减少makespan，还可以充分利用缓存共享来提高缓存命中率。

Cache Aware Scheduling

SGLang

首先是SGLang，其采用了Radix Tree的数据结构来实现cache-aware的调度方法，详情可见我的另一篇关于SGLang的博客。

但SGLang这种基于公共前缀来复用Cache资源的调度存在一个问题，那就是饥饿问题：恶意客户端可以通过发送大量带有相同长前缀的请求来垄断共享资源，影响正常请求的处理，负载均衡会受到影响，导致公平性问题。

Preble

由于实际生产环境必然是分布式的，现有的分布式服务系统则没有平衡好公平性和局部性复用的关系，例如vLLM集群采用轮询负载均衡，SGLang的调度则是以Locality复用为优先。Preble研究的问题就是：如何在分布式（多GPU）的LLM服务系统中，协同优化“提示KV缓存的复用”和“跨GPU的计算负载均衡”？

如果只追求缓存复用：把所有共享同一提示的请求都发送到同一个GPU上，那么这个GPU会成为热点，负载极高，而其他GPU空闲，导致系统整体吞吐量下降和尾延迟飙升。
如果只追求负载均衡：（像现有系统那样）均匀地把请求分到所有GPU，那么缓存复用率会很低，大量计算被浪费在重复计算相同的提示前缀上。

Preble提出了E2分布式调度算法和两级调度架构。

Preble设计的两层调度架构：

全局调度器：一个中心化的调度器，掌握整个集群的全局视图。它负责请求级的调度，即为每个新到来的请求决定应该由哪个GPU（或模型实例）来处理。这是E2算法发挥作用的地方。
本地调度器：每个GPU上都有一个本地调度器，负责迭代级的调度。它管理分配给该GPU的请求队列，决定每轮迭代具体执行哪些请求的哪个步骤（Prefill/Decode），并管理本地的KV缓存。

在全局调度器上，Preble应用了其提出的E2全局调度算法，这里的E2指的是Exploitation & Exploration，其核心思想是，当一个新的请求到达时，调度器要决定是Exploit(利用)现有缓存，还是Explore(探索)新的GPU来实现负载均衡。其流程如下

第一步：匹配全局前缀树

全局调度器维护一个全局前缀树，所有request的prompt都被插入到这棵树上。共享的前缀是树中的内部节点，独特的后缀是叶子节点。每个节点记录了哪些GPU缓存了它的KV状态。
当一个新请求到达时，将其提示与这棵前缀树进行匹配，找到最长的匹配前缀。假设匹配的token数为 cached_len，剩余未匹配的token数为 missed_len。

第二步：Exploit 还是 Explore？

决策规则：比较 cached_len 和 missed_len。
- 如果 cached_len > missed_len → 选择 Exploitation：
  - 理由：复用缓存所节省的计算量（cached_len）大于需要新计算的计算量（missed_len）。此时，复用的收益很高，应该优先利用。
  - 将请求发送到缓存了最长匹配前缀的GPU。如果多个GPU都缓存了，则选择其中“负载”最轻的一个。
- 否则 → 选择 Exploration：
  - 理由：需要新计算的部分占了大头，复用的收益相对较小。此时，更应该关注负载均衡，为这个请求寻找一个“好”的GPU，避免给缓存了前缀但已经很忙的GPU增加负担。
  - 基于Prefill和Decode的计算成本与prompt的token数成比例的insight，提出了prompt-aware负载成本模型，为请求选择一个总成本最低的GPU。

第三步：prompt-aware负载成本计算（用于Exploration）

当需要探索时，E2不是简单地看哪个GPU空闲，而是计算将一个请求分配给某个GPU 的总成本 Cost_i，它由三部分组成：

历史负载成本 $L_i$：估算 $GPU_i$ 在最近一个时间窗口内处理所有请求所花费的总计算时间（包括预填充和解码）。关键点：在计算预填充时间时，只计算未命中缓存的令牌，命中部分视为0成本。
驱逐成本 $M_i$：为了给新请求腾出内存空间，$GPU_i$可能需要驱逐一些已缓存的KV状态。驱逐成本不是简单的“驱逐了多少令牌”，而是 “驱逐这些令牌的预期未来损失” 。计算公式为：被驱逐节点的Prefill时间$\times$该节点的历史访问频率，保证不轻易驱逐那些热门且计算代价高的缓存。
新请求成本 $P_i$：估算在$GPU_i$处理这个新请求中未命中缓存的令牌所需的预填充时间。

E2选择$L_i +M_i+ P_i$总和最小的那个GPU。

在本地调度器上，Preble设计了一套基于优先级的队列配置，按照prompt中tokens的缓存比例进行分级，比例越高，其优先级就越高。然后在组建running batch的时候，再基于优先级进行比例采样，从而保证各个优先级的请求都能获得相应比例的服务，避免低优先级的饥饿问题。

Deficit Longest Prefix Match

虽然Preble一定程度上平衡了负载均衡和局部性复用，但其是从请求的复用性角度来进行公平性评估的，服务资源仍然容易被那些恶意请求垄断，于是DLPM针对这一问题，采用了基于服务量的评估方式。DLPM通过结合了SGLang中的Longest Prefix Match的思想和Virtual Token Counter(VTC)两种方法，平衡cache复用与公平性。

Virtual Token Counter方法以一个时间段内针对一个用户请求所处理的input&output tokens数量作为服务成本的度量。VTC针对请求队列采用了Monitoring Stream和Execution Stream来进行管理。

Monitoring Stream：用于处理请求入队的情况。为了防止新的客户端因为其历史数据较低而透支未来服务资源的情况，需要对其进行counter lift。如果队列为空，该客户端的计数器就设置为最后离开队列的那个客户端计数器的大小；如果队列不为空，则设置为最小的计数器数值；
Execution Stream：从等待队列中选择计数器最小且最早加入队列的请求，根据prefill阶段的tokens量其计数器直接增加相应token长度，然后将其弹出进入到running batch中；decode阶段，随着tokens输出实时增加其计数器数值。注意，请求的处理过程是非抢占的，只有当前请求执行完，才会释放相应的KV Cache内存然后触发新的请求进入running batch。

DLPM中一个核心概念是backlog，如果说一个客户端backlogged了，意思是说处理来自该客户端的请求无法进一步增加吞吐量，只会造成额外的排队延迟。基于这一概念，DLPM提出了三个公平性原则：

一个时间窗口内，两个backlogged的客户端之间接受的服务量之差不应该高于一个特定常数
一个没有backlogged的客户端所接受的服务量相对于一个backlogged的客户端接受的服务量的差值不应该超过特定常数
只要队列里还有未处理的请求，worker就不能空闲

第一个原则保证了那些一直在以高速率发送请求的客户端不会得到更多的服务资源；第二个原则保证了那些一开始以低速率发送请求的客户端无法通过累计未使用的服务从而在后续垄断资源。

DLPM在不打破LPM固有的本地排序的前提下(i.e., 请求的优先级按照其重复前缀长度降序排列)，借鉴了Deficit Round Robin技术，迫使调度器偶尔优先处理来自服务较少客户端的请求，从而来维持公平性。

其核心组件和工作流程如下：

初始化：为每个客户$i$维护一个计数器$q_i$，初始值为0。同时设定一个$Q^u$，这是一个可调的超参数。
每个连续批处理步骤：
- a. 排序：像LPM一样，将等待队列中的请求按照其匹配前缀长度进行降序排序。这一步是为了最大化局部性。
- b. 资格检查与批次填充：遍历排序后的请求列表，尝试将请求加入当前运行的批次$B$中。一个请求能被加入的前提条件是：其所属客户的$q_i > 0$
  - 如果$q_i > 0$，则将该请求加入批次$B$。
  - 如果$q_i \le 0 $，则跳过该请求，即使它有很长的匹配前缀。
- c. 量子补充：如果遍历完所有请求后，发现所有活跃客户的$q_i \le 0$，并且批次$B$还未填满，则系统执行补充：为所有活跃客户的赤字计数器增加$Q^u$。然后，算法重新从步骤a开始，用更新后的$q_i$再次尝试填充批次。
- d. 服务扣除：
  - 当一个请求被加入批次$B$后，立即从其客户的$q_i$中扣除该请求的extend tokens（即不共享的、需要新计算的那部分input token）的成本。
  - 在模型完成一个decode步骤（生成一个output token）后，从相应客户的$q_i$中扣除成本。

该方法中，超参数$Q^u$起到了控制局部性-公平性权衡的作用，越小则越接近VTC方法，越大则越接近LPM方法。

在分布式环境下，论文提出了$D^2LPM$的方法。因为集中式调度需要保证所有实例的缓存信息、前缀树需要实时同步，开销十分巨大，因此$D^2LPM$采用了去中心化的调度。在局部性上，尽量将相同前缀的请求发往已有该前缀缓存的GPU；在公平性上，通过双重赤字计数器机制，防止某些客户端过度占用某个GPU。其主要的组件包括：

全局调度器：负责接收请求，并根据前缀匹配结果和赤字计数器决定将请求发往哪个GPU。
本地调度器：每个GPU上运行一个DLPM调度器，负责本地的公平调度和前缀缓存管理。
全局前缀树（Radix Tree）：记录所有GPU上的前缀缓存情况，用于匹配请求的前缀。

$D^2LPM$在两层维护赤字计数器：

a) Global Layer

记为$q_{i,w}$表示客户端 $i$ 在GPU $w$ 上的“服务额度”。
作用：控制客户端对某个GPU的资源占有，防止某个客户端的所有请求都发往同一个GPU。

b) Local Layer

每个GPU本地运行DLPM，用于保证本地GPU上各客户端之间的公平性。

$D^2LPM$调度流程如下：

请求到达全局调度器，提取其前缀。
前缀匹配：查询全局前缀树，找到所有缓存了该前缀的GPU集合$G$。
选择GPU：
- 优先选择$G$中 $q_{i,w}>0$ 的GPU（即该客户端在该GPU上还有额度）。
- 如果有多个可选，选择队列最短的GPU。
- 如果$G$中所有GPU的$q_{i,w} \le 0$，则选择所有GPU中队列最短的。
更新全局赤字计数器：
- 请求被发往GPU后，更新$q_{i,w}$（扣除该请求的“服务成本”）。
本地调度：
- GPU收到请求后，将其加入本地队列，由本地DLPM调度器进行调度。
定期更新：
- 当请求完成时，全局调度器异步更新前缀树和赤字计数器。

除了上述这类启发式的调度方法，还有一种是通过预估工作负载来进行调度的方法，例如vLLM-LTR就是通过对请求按照预测的输出长度进行排序，然后采用Shortest-Job-First的排序思想来降低延迟。

针对推理系统的优化调研综述