前言

影响LLM推理性能的因素有很多，包括但是不限于：模型配置、硬件类型、数据分布、优化算法、推理策略等。本位旨在综述各类技术点，后续会针对核心技术做详细展开。

LLMs计算建模

开源框架分析

开源的LLM推理框架有TensorRT-LLM、FasterTransformer、TGI、VLLM、SGLANG等，以VLLM为例简单介绍下推理框架：

框架优化

Continuous Batching

Paper

Orca: A Distributed Serving System for Transformer-Based Generative Models LLM推理主要氛围Prefill/Decode阶段，其中Prefill阶段为Compute Bound，Decode阶段为Memory-IO Bound，意味着需要充分利用显存的数据带宽。但是由于显存的限制，如A100-40GB而言，以13B模型为例，权重为26GB，其消耗单个Token的显存约为1MB（激活值/Cache等），则当Seq长度为2048时，仅能支持7 batch的推理，此时算力利用率很低。为了提高算力利用率，一方面可以通过量化手段减小显存，也可以通过优化FlashAttention等算子实现减少Memory-IO，另一方面也可以从Schedule层面通过Continuous Batching手段来提高Batching效率。

Motivation

传统的 static batching 如下：

显然，由于不同batch的out_token数目不一致，1，3，4batch存在empty slots，意味着一段时间内，只有batch 2单batch在计算。Continuous batching采用来iteration-level的调度策略：

Key Points

由于Prefill/Decode两个阶段的 computational pattern 差别较大，主要体现在ltency和显存消耗上，所以实际prefill阶段往往没法直接和decode做batching，所以通过 waiting_served_ratio 来控制prefill阶段和decode阶段的调度策略。

Statistics

Prefix prompt cache

Paper

SGLang: Efficient Execution of Structured Language Model Prpgrams

Motivation

多轮对话中，由于Prompts存在大量相同前缀，因此 Prefix prompt cache 的复用就显得很重要：需要考虑：efficient prefix search,reuse,insertion,eviction已经对应的cache aware scheduing。

Key Points

SGLang采用了 prefix tree 的数据结构来进行Cache的数据存储：

另外，Scheduling策略上，以匹配的prefix-length作为最高优先级。

VLLM结合PA和Prefix Cache，通过 hash(prefix tokens + block tokens) 来定义对应的KV Block，通过hash-mapping的方式来实现prefix-cache和physical block的mapping：

Statistics

multi-call 下的性能收益：

MultiLora

Paper

MULTILORA: DEMOCRATIZING LORA FOR BETTER MULTI-TASK LEARNING

S-LoRA: SERVING THOUSANDS OF CONCURRENT LORA ADAPTERS

Motivation

PEFT降低了大模型fine-tuning的硬件算力需求，比较常见如：LoRA；对于多任务PEFT的工作，之前的工作无法像LoRA一样和BaseModel参数做无缝集成。 Multi-LoRA在推理部署的问题的困境在于：

单个LoRA可以通过merge权重的方式，实现无开销的LoRA推理；但是对于Multi-LoRA而言，这种方式的效率是偏低的。
Multi-LoRA的Adapter的数目可能很多，全量加载可能会导致显存占用过高，需要实现Host/Device侧内存的Load/OffLoad
对于2.的问题需要解决内存碎片和动态Load/OffLoad的延迟开销问题。

Key Points

核心流程如下：多个lora独立，但是前向传播过程中共享所以训练数据的前向传播，但loss更新的时候会根据mask矩阵进行单个任务的LoRA权重更新。

推理部署而言，核心为了解决前文的Multi-LoRA的Adator的动态加载的问题，关键方案如下：

Specified-CUDA-Kernels：定制了批处理GEMM内核来计算LoRA。
Memory-Allocator Design：Host内存存储了全量的Adater权重，实际推理过程中需要实时fetch对应的Adater权重到显存
Unified Paging：扩展了PagedAttention的内存管理方式，划分了LoRA权重专门的内存池以减少内存碎片，同时提前fetch权重来实现计算/权重加载并行
张量并行：LoRA侧的并行策略如下：

Statistics

S-LoRA策略下，单个GPU/多个GPU下可以部署数千个LoRA适配器，以下为Llama-7B/13B/30B/70B的部分结果：

Speculative decoding

Paper

MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads A Hitchhiker’s Guide to Speculative Decoding

Motivation

LLM的decoding阶段通常会受限于显存带宽(memory-bandwidth-bound)，因为decoding阶段单次仅生成单个token，导致计算效率偏低。提高decoding阶段的计算效率，其中一种方式便是：提高单次inference的计算强度/减少deocding的迭代次数。其设计思路是：采用一个轻量的draft model来预测多个token，然后通过original model来refine生成的token，从而减少总的计算量，同时又不损失精度。

Key Points

MEDUSA策略是：添加了多个head在最后一层的hidden states后，在实际的推理的过程中，连续的heads能生成连续多个token，最后再通过原始模型对多个Candidates进行确认。

对于多个tokens的计算，采用了Tree-Attention的方式来进行Attention的计算。

实际部署过程中的一些细节设计：

HEAD的效率问题：更多的HEAD，能够增加单次推理的token数目，如果完全被original model接受的话，效率更高；但是更多的HAD，对训练的难度更高，且如果draft model的预测精度不高的话，会造成一定的计算资源/显存损失，实验上看3-4 heads收益比较大，部分常见6-8heads也有比较大的收益。
KV-Cache的设计：original model对多个token的验证，Lookahead scheduling设计能够减少对proposal token的KV Cache的重新计算/显存拷贝

Statistics

Llama2 13B的测试结果如下，在保持TTFT/ITL一致的情况下，能提升约一倍的吞吐。

Pipeline/Tensor parallism

训练场景，并行策略有DP/TP/PP/EP等策略，也有不同并行策略的组合，以及Zero等OffLoad的策略。对于推理而言，最常用的是TP/PP策略。 TP最常用的策略主要在MLP/Self-Attention（num_head切分）进行切分，整体为：col-spllit->row-split->all-reduce。从参数量上看，单卡的权重为1/TP，计算量上为1/TP，激活值为1/TP，但是由于输入为完整的seq_len，所以KV_Cache为完整的，通信量上看是 2*b*s*h + 2*b*s*h ，通信量比较大，所以一般需要较大的通信带宽。

PP的切分策略比较简单，通常按照Transformer-Block进行切分(num_layers/pp_world_size)，不同Block之间通过send/reiv通信，且由于推理过程中仅涉及forward阶段，所以PP的Bubble现象也不显著：需要保证不同卡间TransforBlock计算的均衡。从参数量上看，单卡的权重为1/PP，计算量为1/PP，激活值为1/PP，KV_Cache为1/PP；通信量上是 b*s*h 。

TODO PD Seperate

TODO ChunkPrefill

模型优化

Graph Compile

图编译优化主要是围绕AI编译器来实现的，整体框架可参考：The Deep Learning Compiler: A Comprehensive Survey:

对于LLM来说，本身模型结构的变化不大，意味这图层面的优化策略是比较通用的，如：Quant算子的融合/PointWise算子融合/自定义融合算子等，所以常用的策略是复用PyTorch等框架的Compiler，自定义LLM的Compiler，如[RFC] A Graph Optimization System in vLLM using torch.compile中介绍的：基于TorchDynamo的FX Graph添加Used-defined Compiler来实现LLM的自定义Compiler优化。其优势是能够不开发专用Compiler的前提下，完成high-level的IR优化，同时复用TorchDynoma Low-IR侧的优化能力，但是缺点是其输入的图属于”raw FX Graph”，对于例如code elimination/topo sort,实现上会有些困难。

量化

llm int8

weight混合精度量化，channel-wise，outlier采用FP16.

GPTQ

W4A16，W8A16属于PTQ（训练后量化），一般需要校准数据集：通过分解海森矩阵来迭代量化，对权重矩阵进行block切分，同样为混合精度量化。

AWQ

W4A16，W8A16，激活感知的量化，为完全的int4/int8权重量化；核心思路：对关键权重先乘一个方法系数再量化进行一个保护。首先，识别关键权重的方法是分析Activation分布。显著权重乘一个放大系数scale后，量化误差会较之前变小，并且对于显著权重也给予不同的保护粒度。论文通过最小化layer量化前后的差值来在搜索空间寻找最优的 scaling。

SmoothQuant

W8A8，支持weight/activation联合量化，其动机是发现activations需要channel-wise才能保持精度，但是channel-wise的activation量化效率很低，需要采用：把Activation量化的难度平滑转移到weight的量化上。SmoothQuant 提出了一种数学上等价的逐通道缩放变换（per-channel scaling transformation），可显著平滑通道间的幅度，从而使模型易于量化。

TODO 稀疏

算子优化

硬件简介

当前深度学习应用硬件主要是GPU和各类XPU（AISC），GPU架构设计上的最大的特点是：

内存架构：GPU的片上Cache相对较小，且通常采用HBM等高带宽的存储，会设计多级存储方案
计算单元：GPU有强大的计算单元，且超配大量的多线程，专注大规模并行的计算任务。计算单元以GPU SM为单位，GPU单个时钟周期能支持多个Warp，A100为例，每个SM包含64个Warp，包含2048个线程。
GPU独特的硬件设计对应其SIMT编程特性，而传统的CPU或者XPU都是SIMD编程模型，其主要区别在于：
- SIMD采用一个线程处理一个指令，指令是向量化处理的，意味可以支持单指令多地址，且维度是固定的，如32bit的4维向量vec4，需要4的整数倍数据。其优势是，对于大尺寸连续数据的处理上，SIMD性能更有优势，但是对于小数据尤其是没法保持为整数倍的数据，会有一定程度计算浪费。
- SIMT采用的思路是单核多线程，即单个计算单元（SM）是支持多线程的，即多线程同指令，通过多线程的方式来降低IO的时延。SIMT对开发人员的要求是更低的，因为SIMT始终是统一指令的，只要处理好不同线程之间的数据排布即可。

GEMM

GEMM优化的核心在于：

提高Cache命中率，设计更好的数据排布(Tiling)
提高并行度，充分利用指令向量化和多核并行

先看一个native的实现：

void naive_row_major_sgemm(const float* A, const float* B, float* C, const int M,
                           const int N, const int K) {
    for (int m = 0; m < M; ++m) {
        for (int n = 0; n < N; ++n) {
            for (int k = 0; k < K; ++k) {
                C[m * N + n] += A[m * K + k] * B[k * N + n];
            }
        }
    }
}

其浮点运算量为2xMxNxK，缓存读次数为2xMxNxK,写次数为MxN；但是这种写法会造成大量的Cache miss：

简单的循环重排就可以得到加速：

void optimize_row_major_sgemm(const float* A, const float* B, float* C, const int M,
                              const int N, const int K) {
    for (int m = 0; m < M; ++m) {
        for (int k = 0; k < K; ++k) {
            for (int n = 0; n < N; ++n) {
                C[m * N + n] += A[m * K + k] * B[k * N + n];
            }
        }
    }
}

进一步提高Cache命中率，则需要做矩阵A/B的分块计算，以CUDA而言：需要设计从Global Memory->Shared Memory->Register的分块逻辑。

FlashAttention

FlashAttentionV2算是FlashAttention的集大成者，V3更多的是基于GPU硬件定制优化。本身FlashAttentionV1缓解了IO显存的时延问题，FlashAttentionV2在其基础上进一步做了优化：

减少了non-matmul FLOPS数量
基于不同thread block实现序列并行化，优化了warp的工作分配。

前置补充FlashAttentionV1的核心点：主要通过tiling设计/Online-Softmax来优化中间结果的缓存处处，可以将内存开销降低到线性级别，提高了2-4倍加速。

v2 python实现：

import torch

torch.manual_seed(456)

N, d = 16, 8
Q_mat = torch.rand((N, d))
K_mat = torch.rand((N, d))
V_mat = torch.rand((N, d))

expected_softmax = torch.softmax(Q_mat @ K_mat.T, dim=1)
expected_attention = expected_softmax @ V_mat

# 分块（tiling）尺寸，以SRAM的大小计算得到
Br = 4
Bc = d

O = torch.zeros((N, d))

# 算法流程第3步，执行外循环
for block_start_Br in range(0, N, Br):
    block_end_Br = block_start_Br + Br
    # 算法流程第4步，从HBM中load Qi 的一个block到SRAM
    Qi = Q_mat[block_start_Br:block_end_Br, :]
    # 算法流程第5步，初始化每个block的值
    Oi = torch.zeros((Br, d))  # shape Br x d
    li = torch.zeros((Br, 1))  # shape Br x 1
    mi = torch.full((Br, 1), -torch.inf)  # shape Br x 1

    # 算法流程第6步，执行内循环
    for block_start_Bc in range(0, N, Bc):
        block_end_Bc = block_start_Bc + Bc

        # 算法流程第7步，load Kj, Vj到SRAM
        Kj = K_mat[block_start_Bc:block_end_Bc, :]
        Vj = V_mat[block_start_Bc:block_end_Bc, :]

        # 算法流程第8步
        Sij = Qi @ Kj.T
        # 算法流程第9步
        mi_new = torch.max(torch.column_stack([mi, torch.max(Sij, dim=1).values[:, None]]), dim=1).values[:, None]
        Pij_hat = torch.exp(Sij - mi_new)
        li = torch.exp(mi - mi_new) * li + torch.sum(Pij_hat, dim=1)[:, None]
        # 算法流程第10步
        Oi = Oi * torch.exp(mi - mi_new) + Pij_hat @ Vj

        mi = mi_new

    # 第12步
    Oi = Oi / li

    # 第14步
    O[block_start_Br:block_end_Br, :] = Oi
assert torch.allclose(O, expected_attention)

TODO PageAttention

场景优化

MOEs(Mixed Expert Models)

MoEs最早由Switch Transformer提出,核心结构由稀疏MoE层和门控网络或者路由组成：

稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构。
门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如，在下图中，“More”这个令牌可能被发送到第二个专家，而“Parameters”这个令牌被发送到第一个专家。有时，一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。

MoEs加速的问题在于：其本身的分支结构会导致计算效率低下，且由于Experts数目的增多会导致需要多卡/多机之间的分布式部署，从而引入额外的通信开销。对于MoEs模型，通常需要采用多种并行计算策略：DP/PP(MP)/TP/EP，且多种策略往往需要组合使用。

DeepSpeed的并行策略组合如下：下图是以单个 MoE 层为例解释了如何使用混合并行策略，假设总共有16 个 GPU，专家数量为 8 个，可以看到并行模式是针对专家和非专家模块分别设计的，具体如下：

对于非 expert 参数模块：
- 使用4 路数据并行，也就是说该部分参数复制了 4 遍
- 每一路数据并行的内部采用 4 路的tensor 并行
对于 expert 参数模块：
- 使用 8 路专家并行
- 每个专家通过 2 路的 tensor 并行进行参数拆分

同时，TP/EP间也可以做通信优化：核心思路是TP的输入/输出数据是一致的，不需要做all-to-all通信，能够节省EP的通信开销：

长文本

长文本对LLM带来的挑战主要在于显存上的，核心原因是：原始Attention的计算过程中，中间变量的大小是O(n^2)增长的，所以核心思路是优化Attention的中间计算结果。其中，最为典型的思路便是：SP(序列并行，Sequence Parallism)，而序列并行最典型的思路便是：ULYSESS和Ring Attention。 Ulysess的核心思路是：Q/K/V矩阵按照N维度做切分，然后workers之间通过All2All转换为d维度切分（前置知识：All2All等价于分布式Transpose），然后可以按照Softmax(QK^T)计算,然后再做一次All2All通信转置回来，单次通信量为O(N*d)，与卡数目无关，总通信量为3xO(N*d)限制条件是d/P可整除，对GQA/MQA不友好。

RingAttention可以认为是分布式的FlashAttention，其核心修改点是：内层循环K、V的时候需要通过P2P通信来得到不同Device间的K、V中间结果，RingAttention的单次通信量是O(Nxd)*P，和卡量相关，总通信量为N/P/c*(P-1)*O(N*d)，其中P为卡量，c为单次计算的block大小，但是计算和通信理论上可以流水，且当计算时间大于通信时间的时候，能够overlap。Ring-Attention的优点在于泛化性好，即插即用，P2P对通信要求低，但是通信量更大，且存在负载均衡的问题，具体表现为越到后面的seq实际的计算量越大。