架构概览#

TensorRT-LLM 是一个旨在为大型语言模型 (LLM) 推理创建优化解决方案的工具包。除了 TensorRT，PyTorch 也可以作为 TensorRT-LLM 的后端。本文档概述了 PyTorch 后端的架构。

顶层 API#

PyTorch 后端的接口是 tensorrt._torch.LLM。

from tensorrt_llm._torch import LLM
llm = LLM(model=<path_to_llama_from_hf>)

LLM 还管理输入的 Token 化和反 Token 化过程。

PyExecutor#

与使用 执行器 API 的 TensorRT 后端类似，PyTorch 后端采用了一个 PyExecutor 类。该类具有与 Executor 类似的接口，可以作为替代后端集成到 LLM 中。PyExecutor 的关键组件包括

  • 模型引擎 (Model Engine)：持有语言模型并高效支持单步模型前向传播。

  • 解码器 (Decoder)：基于模型引擎的输出生成输出 Token。目前仅支持贪婪搜索。

  • 调度器 (Scheduler)：决定是否为请求分配资源（如 KV Cache），以及在当前步骤是否为每个请求运行前向传播。

PyExecutor 的单步流程包括

  • 从请求队列中获取新请求（如果有）。

  • 调度一些请求。

  • 对调度的请求运行模型前向传播。

  • 使用调度的请求的模型前向输出运行解码器。

  • 为每个请求添加输出 Token 并处理已完成的请求。

模型引擎 (Model Engine)#

PyExecutor 的核心组件是 ModelEngine，负责在 GPU 上高效执行模型的前向传播。ModelEngine 的关键方法是处理前向传播计算的 forward。对于 PyTorch 后端，派生类是 PyTorchModelEngine，在 pytorch_model_engine.py 中声明。

解码器 (Decoder)#

解码器根据模型引擎的输出生成输出 Token，并支持贪婪搜索解码。

调度器 (Scheduler)#

调度器分两步操作

  1. CapacityScheduler：确定是否有足够的资源来满足请求。

  2. MicroBatchScheduler：选择一些请求供模型运行前向传播。

CapacityScheduler 和 MicroBatchScheduler 目前都使用 C++ 绑定。但是，由于接口是用 Python 实现的，因此可以进行自定义。文档 scheduler.md 解释了如何实现自定义调度逻辑。

资源管理器 (ResourceManager)#

ResourceManager 帮助分配和管理运行单个请求的推理可能需要的资源。它是一个包含继承自 BaseResourceManager 的对象的容器，每个对象管理特定类型的资源。BaseResourceManager 有三个重要的接口

  • prepare_resources：在 PyExecutor 中当前批处理的模型前向传播之前的每一步调用。

  • update_resources：在当前批处理的每一步完成时调用。

  • free_resources：在每个请求完成时调用。

一个关键资源是 transformer 模型的 KV Cache。BaseResourceManager 用于 KV Cache 的是 KVCacheManager。

KVCacheManager#

目前，KVCacheManager 使用 C++ 绑定。但是，由于其接口是用 Python 实现的，因此可以在 Python 中进行自定义。文档 kv_cache_manager.md 详细说明了如何实现自定义 KVCacheManager。