总结：大模型推理优化策略

原文地址：大模型推理优化策略

• 7.1 显存优化
  • Qunatized KV Cache
  • MQA/GQA
    • 核心思想是减少kv-cache的数量，以少量kv-cache对应多个query

      ​

    • ​

  • PagedAttention
    • KV cache，其具有以下特点:1. 显存占用大，14b级别的模型，每个token需要约0.7M-1M的显存；2. 动态变化：KV 缓存的大小取决于序列长度，这是高度可变和不可预测的。因此，这对有效管理 KV cache 挑战较大。该研究发现，由于碎片化和过度保留，现有系统浪费了 60% - 80% 的显存。
    • 为了解决这个问题，该研究引入了 PagedAttention，这是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 允许在非连续的内存空间中存储连续的 key 和 value 。具体来说，PagedAttention 将每个序列的 KV cache 划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。因为块在内存中不需要连续，因而可以用一种更加灵活的方式管理 key 和 value ，就像在操作系统的虚拟内存中一样：可以将块视为页面，将 token 视为字节，将序列视为进程。序列的连续逻辑块通过块表映射到非连续物理块中。物理块在生成新 token 时按需分配。在 PagedAttention 中，内存浪费只会发生在序列的最后一个块中。这使得在实践中可以实现接近最佳的内存使用，仅浪费不到 4%。
    • PagedAttention 还有另一个关键优势 —— 高效的内存共享。例如在并行采样中，多个输出序列是由同一个 prompt 生成的。在这种情况下，prompt 的计算和内存可以在输出序列中共享。PagedAttention 自然地通过其块表格来启动内存共享。与进程共享物理页面的方式类似，PagedAttention 中的不同序列可以通过将它们的逻辑块映射到同一个物理块的方式来共享块。为了确保安全共享，PagedAttention 会对物理块的引用计数进行跟踪，并实现写时复制（Copy-on-Write）机制。PageAttention 的内存共享大大减少了复杂采样算法的内存开销，例如并行采样和集束搜索的内存使用量降低了 55%。这可以转化为高达 2.2 倍的吞吐量提升。
    • vllm的其他优化：continuous batching,CUDA kernel优化
  • FlashAttention
    • 解释1：记I为模型的计算强度I，单位FLOP/byte代表模型进行单位byte数据交互可实现的操作数,则I*带宽beta即模型的计算性能，单位为FLOP/s。令I_max=计算平台算力/计算平台带宽，当模型的计算强度I小于平台的理论计算强度I_max，模型的计算能力P即I*beta,当模型的计算强度大于I_max,则模型的计算性能P等于平台的算力。故若模型的计算强度小，则瓶颈在带宽，若模型的计算强度大，则瓶颈在算力。为提高计算性能，需提高计算强度，即每bytes数据交互的操作数。
    • 解释2：记N=每次操作要求的FLOP,单位FLOP/OP；pi=平台的算力，单位FLOP/s;beta=内存带宽，单位byte/s；P=实际实现计算速度，单位为FLOP/s；优化目标为O=P/N每秒钟实现的运算次数，单位为OP/s.由于N固定，故优化目标转而为P，P=min{beta_r*I_max=beta_r*pi/beta,pi}，故优化目标转而为beta，即改变内存访问策略，实现beta最大化。
    • 出发点：注意力操作中，S和P的计算空间复杂度都是O(N^2),此外,scale,mask,softmax,dropout都是带宽约束操作。

      ​

    • 思路：输入数据K、Q、V是存储在HBM上的，中间结果S、A都不需要存储到HBM上。通过这种方式，FlashAttention可以将内存开销降低到线性级别，并实现了2-4倍的加速，同时避免了对中间结果的频繁读写，从而提高了计算效率。

      ​

    • 方法：O(N^2)空间复杂度的矩阵计算对HBM的读写是主要的内存瓶颈，因此主要优化点是：1. 在不访问整个输入的情况下计算softmax;2. 不为反向传播存储大的中间attention矩阵。FlashAttention提出两种方法来分步解决上述问题：tiling,recomputation。tiling - 注意力计算被重新构造，将输入分割成块，并通过在输入块上进行多次传递来递增地执行softmax操作。recomputation - 存储来自前向的 softmax 归一化因子，以便在反向中快速重新计算芯片上的 attention，这比从HBM读取中间矩阵的标准注意力方法更快。由于重新计算，这虽然导致FLOPS增加，但是由于大量减少HBM访问，FlashAttention运行速度更快。该算法背后的主要思想是分割输入，将它们从慢速HBM加载到快速SRAM，然后计算这些块的 attention 输出。在将每个块的输出相加之前，将其按正确的归一化因子进行缩放，从而得到正确的结果。kernel融合：tiling分块计算使得我们可以用一个CUDA kernel来执行注意力的所有操作。从HBM中加载输入数据，在SRAM中执行所有的计算操作（矩阵乘法，mask，softmax，dropout，矩阵乘法），再将计算结果写回到HBM中。通过kernel融合将多个操作融合为一个操作，避免了反复地从HBM中读写数据。基于以上思路，flash-attention天然在batch和heads两个维度上进行了并行化：使用一个thread block来处理一个attention head，总共需要thread block的数量等于batch size × number of heads。每个block被调到到一个SM上运行（SM可以调度多个thread block），当线程块数目较大时（例如80个），该调度方法能高效利用GPU上接近全部的计算资源。

      ​

      • ​

      • ​

      • 此处P=A/l

        ​

  • FlashAttentionv2
    • 出发点：FlashAttention利用了GPU内存的非对称层次结构，将内存消耗降至线性（而非二次方），并相较于优化基线实现了2到4倍的运行速度提升，然而，该技术的速度依然没有达到优化矩阵乘法（GEMM）操作的速度，前向传播的计算吞吐量仅达到理论最大浮点运算速率(FLOPs/s)的30-50%，而反向传播只能达到25-35%。这种低效率是由于GPU上不同线程块之间的负载分配不佳，导致低占用率或不必要的共享内存读/写。
    • 方法：(1)调整算法，减少中间缩放的次数。通常实现Softmax算子为了数值稳定性（因为指数增长太快，数值会过大甚至溢出），会减去最大值，这样带来的代价就是要对token遍历3次。而如果不保存中间最大值和指数和而只保留对数指数和，可以减少非矩阵乘法操作的浮点计算次数；(2)序列长度维度的并行。v1在batch_size和num_heads维度使用了并行化，但在处理长序列输入时，由于内存限制，通常会减小batch size和head数量，这样并行化程度就降低了。因此，FlashAttention-2还在序列长度这一维度上进行并行化，显著提升了计算速度。而序列并行要求我们交换循环的层次，FlashAttention-v2将Q移到外循环，将KV移到内循环，从而分块计算的注意力值矩阵块只存储在SRAM上，避免了在HBM上的频繁读写。(b*s*d->b*s*k*m)。(3)在每个线程块内部，分散warps之间的工作负载，减少通过共享内存的通信。v1使用是split-K策略，这种策略会使所有warp将中间结果写入共享内存进行同步，然后将中间结果相加，这些共享内存读取会拖慢前向传播的计算。v2使用的是split-Q策略，这样在每个warp计算QK^\top后，结果只需要对应的V分片即可得到O的对应分片，而无需进行warps间的通信，减少了中间共享内存读写。

      ​

      • 此处e^{S-m}即为A

        ​

      • ？每个token的计算都需要加载K,V,O，得到S，故M大小的SRAM容纳的最大token数为[M/4d]，取该数与d间的最小值作为K、V的token上限数。

        ​

    • 结果：在前向传播中可以达到理论最大吞吐量的73%，反向传播中则可达到63%。当用于端到端训练GPT风格的模型时，每个A100 GPU的训练速度可以达到225 TFLOPs/s。
    • 细节1：GPU由计算元素（如浮点运算单元）和内存层次结构组成。现代GPU中通常包含专门的单元来加速低精度（如FP16/BF16）的矩阵乘法运算，例如Nvidia GPU上的Tensor Core。内存层次结构包括高带宽存储器（HBM上的存储都属于DRAM）和静态随机存取存储器（GPU上所有的on-chip memory都是SRAM）。以A100 GPU为例，它拥有40-80GB的高带宽存储（HBM），带宽达到1.5-2.0TB/s，以及每个流式多处理器108个中的192KB芯片上SRAM，其带宽估计约为19TB/s。Nvidia的A100 GPU具有最大理论吞吐量达到312 TFLOPs/s（每秒万亿次浮点运算）的FP16/BF16格式矩阵乘法性能，然而对于非矩阵乘法的FP32格式计算，这一数字仅为19.5 TFLOPs/s。
    • 细节2：GPU有大量的threads用于执行操作（an operation，也称为a kernel）。这些thread组成了thread block，接着这些blocks被调度在SMs上运行。在每个thread block中，threads被组成了warps（32个threads为一组）。一个warp内的threads可以通过快速shuffle指令进行通信或者合作执行矩阵乘法。在每个thread block内部，warps可以通过读取/写入共享内存进行通信。每个kernel从HBM加载数据到寄存器和SRAM中，进行计算，最后将结果写回HBM中。
  • Flash-Decoding: kv序列长度方向的并行
    • 出发点：FlashAttention优化不适合直接应用于推理过程。因为在训练过程中，FlashAttention对batch size和query length进行了并行化加速。而在推理过程中，query length通常为1，这意味着如果batch size小于GPU上的SM数量（例如A100上有108个SMs），那么整个计算过程只使用了GPU的一小部分！特别是当上下文较长时，通常会减小batch size来适应GPU内存。例如batch size = 1时，FlashAttention对GPU利用率小于1%！
    • FlashAttention对batch size和query length进行了并行化加速，Flash-Decoding在此基础上增加了一个新的并行化维度：keys/values的序列长度。即使batch size很小，但只要上下文足够长，它就可以充分利用GPU。与FlashAttention类似，Flash-Decoding几乎不用额外存储大量数据到全局内存中，从而减少了内存开销。
    • Flash Decoding主要包含以下三个步骤:1将keys和values分成较小的block;2 使用FlashAttention并行计算query与每个block的注意力（这是和FlashAttention最大的区别）。对于每个block的每行（因为一行是一个特征维度），Flash Decoding会额外记录attention values的log-sum-exp（标量值，用于第3步进行rescale）;3 对所有output blocks进行reduction得到最终的output，需要用log-sum-exp值来重新调整每个块的贡献。第1步中的数据分块不涉及GPU操作（因为不需要在物理上分开），只需要对第2步和第3步执行单独的kernels。虽然最终的reduction操作会引入一些额外的计算，但在总体上，Flash-Decoding通过增加并行化的方式取得了更高的效率。

      ​

  • SparQ Attention：基于稀疏注意力思想减少kv的数量
    • 出发点：1. 注意力机制中的softmax函数输出主要由少数分量主导，大部分分量接近于0。因此，只需从内存中获取具有最高注意力分数的tokens的key和value对，就可以在不影响任务性能的情况下显著减少内存传输。2. 通过稀疏化查询向量q，仅保留r个最大幅度分量，可以有效预测具有最大注意力分数的索引，而无需获取完整的K矩阵。步骤：1. 找到输入查询向量q中r个最大分量的索引，并仅沿着这些索引获取key缓存K。使用切片查询和键来计算近似注意力分数。2. 在近似注意力分数中找到前k个位置，并获取相应