TurboMind：高效推理引擎引领大语言模型新时代-CSDN博客

本文链接：https://blog.csdn.net/weixin_36829761/article/details/140746200

在人工智能蓬勃发展的当下，大语言模型（LLM）已成为 AI 领域的明星技术。然而，随着模型规模的不断扩大，如何高效地进行推理成为了一个棘手的问题。面对这一挑战，一款名为 TurboMind 的高效推理引擎应运而生，为 LLM 的应用带来了新的可能。

TurboMind：FasterTransformer 的进化版

TurboMind 是基于英伟达的 FasterTransformer 研发而成的 LLM 推理引擎。它不仅继承了 FasterTransformer 的高效性，还在此基础上进行了多项创新和优化。其核心功能包括对 LLaMa 结构模型的支持、persistent batch 推理模式以及可扩展的 KV 缓存管理器。

架构创新：持久化批处理

TurboMind 引入了一种名为 “persistent batch” 的创新机制，这一机制在某些项目中也被称为 “continuous batching”。这种方法将对话式 LLM 的推理过程建模为一个持续运行的批处理，其生命周期贯穿整个服务过程。

具体实现方式如下：

预先准备 N 个批处理槽位（batch slots）。
当有空闲槽位时，新的请求会被加入到批处理中。
请求对应的 tokens 生成完毕后，相应的槽位立即释放，准备接收新的请求。
当一个序列命中缓存时，其历史 token 无需在每轮中重新解码，从而实现即时的 token 生成。
整个批处理会自动扩缩容，以避免不必要的计算资源浪费。

这种机制大大提高了模型的推理效率，使得 TurboMind 能够更好地应对高并发的推理需求。

内存管理的突破：KV 缓存管理器

TurboMind 的另一个亮点是其独特的 KV 缓存管理器。这个管理器本质上是一个内存池类型的对象，巧妙地结合了 LRU（最近最少使用）算法的实现。这使得整个管理器可以被视为一个 “KV 缓存的缓存”。

KV 缓存管理器的工作原理如下：

管理器根据预先配置的槽位数量分配 KV 缓存空间。
每个槽位对应一个序列所需的 KV 缓存。
内存块的分配可以通过配置实现预分配或按需分配。
当缓存池中没有空闲槽位时，管理器会根据 LRU 机制踢除最近使用最少的序列，将其占用的槽位分配给新的请求。
被踢除的序列不会被完全删除，而是转换为最简洁的形式（如 token IDs）。
当再次获取到相同的序列 ID 时，这些 token IDs 将被 FMHA 的 context decoder 解码并转回 KV 缓存。

这种设计使得 TurboMind 能够高效地管理内存资源，从用户的角度来看，使用 TurboMind 的系统就像是可以访问无限的设备内存一样。

LLaMa 模型的优化实现

TurboMind 对 LLaMa 系列模型的实现是在 FasterTransformer 的 Gpt-NeX 模型基础上进行修改和优化的。除了基本的重构和修改外，TurboMind 还引入了多项改进，以实现会话模型的高性能推理：

支持多轮对话中的快速文本解码：使用基于 cutlass 的 FMHA 实现替代了 context decoder 中的注意力机制，支持 Q/K 长度不匹配的情况。
引入间接缓冲指针：在 context FMHA 和 generation FMHA 中加入间接缓冲指针，支持批处理中不连续的 KV 缓存。
新的同步机制：为支持 persistent batch 的并发推理，设计了新的同步机制，协调张量并行模式下的工作线程。
INT8 KV cache 实现：降低内存开销，提高批处理大小和系统吞吐量。这在实际应用中尤其有用，因为 KV cache 通常比权重和其他激活消耗更多的内存和内存带宽。
解决 NCCL 卡住问题：解决了单个进程内多个模型实例在张量并行（TP）模式下运行时 NCCL 卡住的问题。NCCL APIs 现由主机端的同步障碍保护。

这些优化使得 TurboMind 在处理 LLaMa 系列模型时表现出色，大大提升了推理效率和性能。

API 设计：灵活与高效并重

TurboMind 的 Python API 支持流式结果返回和张量并行模式，为开发者提供了灵活而高效的接口。这使得开发者可以更方便地集成 TurboMind 到自己的应用中，充分发挥其高性能推理的优势。

TurboMind vs FasterTransformer：精简与专注

虽然 TurboMind 是基于 FasterTransformer 开发的，但两者还是存在不少差异。TurboMind 去除了 FasterTransformer 中的一些功能，包括前缀提示词、beam search、上下文嵌入、稀疏化 GEMM 操作以及对 GPT 或 T5 等结构模型的支持。这种"减法"使得 TurboMind 更加专注于 LLaMa 结构模型的高效推理，从而在特定场景下实现更优的性能。

兼容性探讨：Hugging Face 模型的支持

值得注意的是，TurboMind 的权重设计是基于 LLaMa 的官方实现完成的，两者仅相差一个转置操作。然而，Hugging Face 版本的实现采用了另一种形式。为了解决这一兼容性问题，TurboMind 在 deploy.py 文件中进行了适配处理，主要涉及 W_q 和 W_k 的差异处理。这一设计使得 TurboMind 能够更好地支持来自不同来源的模型权重。