TurboMind的架构

TurboMind的架构#

TurboMind 是一个支持高吞吐量推理的对话式大型语言模型（LLM）推理引擎。它基于 NVIDIA 的 FasterTransformer。TurboMind 的主要特性包括高效的 LLaMa 实现、持久化批量推理模型以及可扩展的 KV 缓存管理器。

TurboMind 的高级概述#

  +--------------------+
  |        API         |
  +--------------------+
          |    ^
  request |    | stream callback
          v    |
  +--------------------+   fetch   +-------------------+
  |  Persistent Batch  | <-------> |  KV Cache Manager |
  +--------------------+   update  +-------------------+
             ^
             |
             v
+------------------------+
|  LLaMA implementation  |
+------------------------+
| FT kernels & utilities |
+------------------------+

持久化批处理#

您可能在其他仓库中认识到这个功能为“连续批处理”。但在功能的并发开发过程中，我们将对话式LLM的推理建模为一个持续运行的批次，其生命周期跨越整个服务过程，因此得名“持久批次”。简单来说

  • 持久化批次有N个预配置的批次槽。

  • 当有空闲插槽可用时，请求会加入批次。一旦请求的令牌生成完成，批次插槽将被释放并可以重复使用。

  • 在缓存命中时（见下文），历史令牌不需要在每一轮对话中解码；响应令牌的生成将立即开始。

  • 批次自动增长或缩小，以最小化不必要的计算。

KV缓存管理器#

TurboMind的KV缓存管理器是一个类似于内存池的对象，它也实现了LRU策略，因此可以将其视为一种KV缓存的缓存形式。它的工作方式如下

  • KV缓存所需的所有设备内存均由管理器分配。预先配置了固定数量的插槽以匹配系统的内存大小。每个插槽对应于单个序列的KV缓存所需的内存。可以配置分配块大小以实现预分配/按需分配策略（或介于两者之间的某种策略）。

  • 当为新序列请求KV缓存空间但池中没有空闲槽位时，最近最少使用的序列将从缓存中逐出，其设备内存将直接由新序列重用。然而，这并不是故事的结束。

  • 获取序列当前位于其中一个槽中类似于缓存命中，历史KV缓存直接返回，无需上下文解码。

  • 被淘汰的序列不会被完全删除，而是转换为最紧凑的形式，即令牌ID。当稍后获取相同的序列ID时（缓存未命中），令牌ID将由FMHA支持的解码器解码并转换回KV缓存。

  • TurboMind内部自动处理驱逐和转换，因此对用户是透明的。从用户的角度来看，使用TurboMind的系统可以访问无限的设备内存。

LLaMa 实现#

我们对LLaMa系列模型的实现是基于FasterTransformer中的Gpt-NeoX模型进行修改的。除了基本的重构和修改以支持LLaMa系列外，我们还进行了一些改进以实现对话模型的高性能推理，最重要的是：

  • 为了支持多轮对话中的快速上下文解码，我们将上下文解码器中的注意力实现替换为基于cutlass的FMHA实现，该实现支持不匹配的Q/K长度。

  • 我们在上下文FMHA和生成FMHA中引入了间接缓冲区指针，以支持批次内KV缓存的不连续性。

  • 为了支持持久批处理的并发推理，设计了新的同步机制来协调在张量并行模式下运行的工作线程。

  • 为了最大化吞吐量，我们实现了INT8 KV缓存支持以增加最大批量大小。这是有效的，因为在现实世界的服务场景中，KV缓存比权重或其他激活消耗更多的内存和内存带宽。

  • 我们解决了在单个进程中以TP模式运行多个模型实例时NCCL挂起的问题，NCCL API现在由主机端同步屏障保护。

API#

TurboMind 支持一个 Python API，可以实现流式输出和张量并行模式。

FasterTransformer与TurboMind之间的区别#

除了上述描述的功能外，还有许多小的差异我们在此文档中未涵盖。值得注意的是，由于目标的不同，FT的许多功能在TurboMind中被舍弃了（例如，前缀提示、束搜索、上下文嵌入、稀疏GEMM、GPT/T5/其他模型系列等）。

常见问题解答#

支持Huggingface模型#

由于历史原因，TurboMind的权重布局基于原始的LLaMa实现（仅通过转置有所不同）。huggingface transformers中的实现使用了不同的布局来处理W_q和W_k，这在deploy.py中进行了处理。