蚂蚁集团：大模型推理显存优化的深度探索与实践

大模型推理显存挑战

蚂蚁显存优化探索

• 只需导入相关包。

  
  

• 在 Attention kernel 中，替换两行代码的参数即可。

  
  

• 与直接使用 Flash Attention 3 的 kernel 相比，性能基本持平。

  
  

• 与 vLLM 自带的 Flash Attention 2 相比，在某些场景下，性能提升了两倍以上。

  
  

• 消除显存碎片。

  
  

• 提高研发效率，快速适配各种 kernel。

  
  

• Prefill 阶段：这是前向计算阶段，通常是计算密集型的，耗时较长。

  
  

• 排队阶段：排队通常是由于等待显存或等待前一个请求完成（例如前一个 prefill 阶段）。我们主要关注的是显存等待问题。

  
  

• 分层申请显存：在 prefill 阶段，逐层申请 KV cache 显存，而不是一次性申请。

  
  

• 分层 offloading：在不同卡和不同模型长度下，推算出计算时间和传输时间，尝试将两者互相隐藏。

  
  

• 计算完 KV cache 后，立即传输到 CPU。

  
  

• 异步地将 KV cache 从 CPU 拉回到 GPU。

  
  

1. 新请求到达时的显存检查优化：

• 我们放松了对显存检查的要求。原本需要一次性检查整个请求所需的显存是否足够，现在只需检查部分显存（例如 4 层或 8 层）是否足够即可。
• 通过配置参数，系统可以在显存足够的情况下继续推进请求。
• 有两种机会来处理显存不足的情况：
  一种是将显存 Offloading 到其他设备。
  另一种是等待其他请求完成并释放显存后继续推进。

2. 显存分配优化：

• 我们采用了逐层分配显存的方式，而不是一次性分配。
• 在计算完一层的 KV cache 后，立即开启异步的 Uploading 操作，将数据传输回 GPU。

• 为了平衡首次延迟和后续 token 生成的效率，我们设计了调度策略，动态调整显存分配和 Offloading 的优先级。

• 通过配置不同的策略，系统可以在首字生成和后续生成阶段之间进行动态 trade-off。

4. 与现有 vLLM 实现的对比：

• 当前 vLLM 的实现是请求级别的显存管理，即在 prefill 阶段，需要一次性为整个请求的上下文长度预留显存。

• 如果显存不足，vLLM 支持 Offloading 或重新计算（recompute）。

• 首字生成阶段仍然需要确保显存足够，这可能导致排队等待。

1. ShareGPT 负载测试：

• 随着负载压力的增加，优化后的首次延迟（橙色）虽然也会增长，但增长非常缓慢。
• 而基线（蓝色）在负载增加到一定程度后，首次延迟会显著上升，优化后的效果比基线好十几倍。

2. 7B 模型测试：

• 在计算压力较大的情况下，优化后的效果仍然优于基线，大约提升了 8 倍。
• 在 QPS（每秒查询数）方面，优化后相比基线最多降低了 5%（70B 模型），而对于 7B 模型，QPS 仅降低了 1%。

• QPS（吞吐量）：在首次延迟小于 1 秒且生成速度满足体验要求的情况下，优化后的 QPS 相比基线大约提升了 30%。
• 放宽条件：如果放宽一些条件（例如允许首次延迟稍长），QPS 的提升幅度会更大，具体提升多少取决于业务场景的需求。

  
  

1. 给定模型和硬件条件下：

• 如果生成的长度越长，优化后的加速效果会越好。
• 主要原因是，在负载较大的情况下，生成长度越长，基线系统需要等待前一个请求的 decode 阶段结束并释放显存后，才能开始执行新的请求。
• 这种等待会导致新的请求被阻塞，因为显存不足。

2. 测试场景：

• 我们测试了不同的输入输出长度组合，例如 1K 输入到 300 输出，以及 1K 输入到 1K 输出。
• 在不同的场景下，你可以根据业务需求约束首字延迟和生成速度，观察 QPS 的变化。
• 通常情况下，首字延迟会有非常明显的提升。

3. 极限情况：

• 在论文中，我们展示了一些极限情况下的测试结果，优化后的效果可能会有 60 多倍的提升。

结语

• 我们的优化主要集中在显存管理和 Offloading 层面，通过分层申请显存、动态调度以及异步传输等技术，显著提升了首次延迟和 QPS。

  
  

• 优化效果与输入输出长度、负载压力等因素密切相关，生成长度越长，优化效果越显著。

  
  

• 可以根据不同的业务场景和需求，灵活调整优化策略，以达到最佳性能。

  
  

1. Virtual Tensor：

• 目标是消除显存碎片，提升适配各种 kernel 的效率。
• 最终效果是仅需 3 行代码即可适配 attention kernel，例如将 Flash Attention      3 集成到 vLLM 中，在 H100 上测试，性能提升了两倍。

2. LayerKV：

• 专注于优化首次生成阶段的显存管理，特别是 prefill 阶段的 KV Cache 管理。
• 通过显存 Offloading 等技术，在 QPS 持平的情况下，优化了平均延迟和 P99 延迟，极限情况下首次延迟提升了 60 多倍。

• GMLake：已开源。

  
  

• Virtual Tensor：部分开源，代码仓库中有相关实现。

  
  

• LayerKV：尚未开源，计划完成后整理报告并开源。

  
  

• 语言模型：例如 LLAMA 和公司内部的百灵大模型，是我们重点优化的目标。

  
  

• 多模态模型：例如文生图、文生视频等。

  
  

• 搜索、推荐、广告类模型：这也是我们需要关注的领域。

  
  

Q&A

• 使用虚拟显存和物理显存的动态管理和映射。

  
  

• 通过动态映射技术，优化显存的使用效率，从而带来显著的收益。

  
  

• 物理显存释放：
  物理显存的头和尾部分被释放回系统，而中间有数据的部分仍然保留，用户可以通过指针继续访问数据。
  我们没有移动数据，数据仍然保持在原来的位置。
• 实现过程：
  假设当前虚拟显存和物理显存对应关系如下：虚拟显存对应四个物理块（四个 block）。
  我们需要通过 mmap 操作将虚拟显存和物理显存解绑（unmap）。
  解绑后，我们只保留中间有数据的物理块，而将前后的物理块通过 mmap 操作解绑。
  解绑后，可以通过 release 操作将物理显存释放回系统，从而让硬件能够看到更多的可用物理显存。
• 后续操作：
  释放后的物理显存已经归还给硬件，用户可以继续通过虚拟显存的方式分配显存。
  数据仍然保留在中间的物理块中，不会被移动或影响。
• 虚拟地址连续性：
  虚拟地址是连续的，用户可以一次性分配一个足够长的虚拟地址空间（例如48 比特，即 2^48），具体长度取决于硬件支持。

• 碎片整理：通过动态映射和释放，实现了显存碎片的整理，而不会影响用户数据的访问。

  
  

• 数据不动：在实现过程中，数据不会被移动，保证了数据的完整性和一致性。

  
  

• 当前放出的代码可能有些混乱，因为我们正在进行重构。

  
  

• 老代码是基于 vLLM 0.4.2 版本的，而我们现在最新的代码是基于 vLLM 0.5.5 版本的。

  
  

• 我们会尽快放出新的代码，请大家稍作等待。

  
  

• 我们主要测试了 L20 硬件，而在 A100 和 H 卡上的效果可能会有所不同。

  
  

• 具体效果可以参考我们的论文，论文中可能包含了 A100 和 L40 的测试结果。

  
  

• 虽然我们没有在 H 卡上进行大量测试，但 A100 和 L40 上仍然会有一定的优化效果。