TensorRT LLM--In-Flight Batching
原创TensorRT LLM--In-Flight Batching
原创
TensorRT LLM依赖于一个名为Batch Manager的组件来支持请求的即时批处理，该技术旨在减少队列中的等待时间，达到更高的GPU利用率。
更详细地说，该功能允许当一个请求再处理中，同时开始处理另一个新请求。
批处理管理器API
客户端可以使用两个主要的回调与批处理管理器交互，它们的签名在callbacks.h文件中定义。
添加新请求、获取处理结果
向批处理管理器传递新请求的入口点是GetInferenceRequestsCallback，该回调的实现返回需要批处理管理器处理的请求列表（std:：list＜std:：shared_ptr＜InferenceRequest＞），有一个参数表示可以接受的最大请求数（负值表示可以接受无限数量的请求）。
传递新请求的实现如下：
using GetInferenceRequestsCallback = std::function<std::list<std::shared_ptr<InferenceRequest>>(int32_t)>; |
|---|
对于每个新请求，客户端必须向批处理管理器提供输入张量和一个64位无符号数字（uint64_t），该数字将唯一标识该请求。该标识符称为请求ID。输入张量被收集在一个map（std:：map＜std:：string，Tensor＞）中，该map将输入名称与张量相关联。
结果通过SendResponseCallback传递到客户端，一致回调必须接受唯一标识请求的64位请求ID、输出张量列表、布尔值（设置为true时标识请求的最后一个响应）和潜在的非空错误消息。非空错误消息表示遇到错误。在这种情况下，指示这是最后一个响应的布尔值将设置为true，回调必须正确处理错误。
获取请求结果的实现如下：
using SendResponseCallback = std::function<void(uint64_t, std::list<std::shared_ptr<Tensor>> const&, bool, const std::string&)>; |
|---|
需要注意的是，如果客户端传递的请求ID与批次管理器正在处理的请求的请求ID相同，则批次管理器将拒绝使用GetInferenceRequestsCallback回调发送的任何请求。请求ID出现在对标记为final（第三个参数设置为true）的SendResponseCallback回调的调用中后，可以重用。
中断请求
批处理管理器允许用户停止执行当前正在运行的请求，需要停止的一组请求ID可以通过回调传递给批处理管理器：
using PollStopSignalCallback = std::function<std::unordered_set<uint64_t>()>; |
|---|
当一个正在处理中的请求出现在要中断的请求集中时，批处理管理器将确保它被正确停止。
处理中请求统计
当调用以下接口时，批处理管理器可以返回统计信息：
using ReturnBatchManagerStatsCallback = std::function<void(const std::string&)>; |
|---|
统计信息被打包为JSON字符串。该字符串包含三个字段： 时间戳，请求的时间戳（使用std:：put_time（&tm，“%m-%d-%Y%H:%m:%S”）获得）， 迭代计数器（Iteration Counter）是与给定请求的执行相对应的计数器值， Active Request Count，正在处理中的请求数。
GptManager设计
GptManager管理活跃工作项池，活跃工作项池由服务器主动接收的请求流组成。GptManager在其构造函数中生成一个工作线程，然后一直循环获取新请求。工作线程在每个循环迭代开始时调用GetInferenceRequestsCallback，用于读取新请求。当一个或多个请求处理结束了要返回时，它会在每次迭代结束时调用SendResponseCallback。在启用流模式的请求的情况下，此响应可以是单个令牌，或者在禁用流模式时，此响应是完整响应。PollStopSignalCallback和ReturnBatchManagerStatsCallback（如果提供）都会在每次迭代循环结束时调用。当系统没有活动请求时，不调用ReturnBatchManagerStatsCallback。当批处理管理器在SendResponseCallback中通知（通过final_response布尔参数）完成时，服务器可以安全地从其工作项池中撤回请求。在此之前，与该请求相关的所有TensorRTLLM内部状态都将被释放。可以如下创建批处理管理器的实例以服务于像GPT这样的自回归模型：
#include <tensorrt_llm/batch_manager/GptManager.h> using namespace tensorrt_llm::batch_manager; GptManager batchManager(pathToTrtEngine, // Path to the TensorRT engine of the model, TrtGptModelType::InflightBatching, // Use in-flight batching, maxBeamWidth, // Maximum beam width (must be >= 1), schedulerPolicy, // Scheduling policy (see below), maxNumRequests, // Maximum number of requests, getInferenceRequestsCb, // The Get callback (see above), sendResponseCb); // The Send callback (see above). |
|---|
调度程序策略帮助批处理管理器调整请求的执行调度方式。批处理管理器可以尝试通过积极地调度请求（schedulerPolicy设置为MAX_utilization）来最大限度地提高GPU的利用率，如果KV缓存的内存不足，则可能不得不暂停请求。请注意，任何暂停的请求都将自动恢复，用户唯一可见的影响可能是延迟增加。它还可以采用更保守的方法，只有在知道内存分配足以处理所有活动请求时才调度请求，即使在KV缓存消耗的最坏情况下也是如此。该模式对应于设置为GUARANTED_NO_EVICT的schedulerPolicy。
多GPU计算
当使用张量并行或流水线并行在多个GPU上运行时，需要服务器启动的进程数量与GPU排列的进程数量一样多，并且每个进程都运行自己的GptManager副本。
给定节点上可见的GPU数量可以使用CUDA_visible_DEVICES环境变量进行控制。
必须注意确保所有列在生成循环的每次迭代中都能看到相同的输入，在TensorRT LLM Triton后端，在GetInferenceRequestsCallback中执行MPI广播，以确保每个MPI列都能看到相同的请求集。ReturnBatchManagerStatsCallback只需要从单个列组中调用；所有等级都持有最终结果的相同副本。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。
原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。
如有侵权，请联系 cloudcommunity@tencent.com 删除。