当前位置：首页 > news >正文

TensorRT-LLM中的in-flight batching（IFB）

news 2025/9/21 12:26:21

Ref

官方文档
deepwiki
学习代码

博客阅读

使用IFB时，处在 context 阶段（预填充阶段也即prefill）的序列可以跟处在generation 阶段（生成阶段也即decode）的序列一起被处理。出于效率的原因，in-flight 批处理的支持要求输入张量要被“打包”（packed），不能有填充（no padding）。

顺序上限制的解释

IFB允许 prefill 阶段的 token 处理与 decode 阶段的 token 处理同时存在／并行发生。但是有一个顺序上的限制：所有上下文阶段（prefill/context）的 tokens 在输入 tensor 里必须排在生成阶段的 tokens 之前。也就是说，在一个 batched input 里面，如果有的序列还在 prefill，有的序列已经开始 decode，那么这些在 prefill 的序列在整体输入的先后顺序里必须在前。

这个可能有点难理解，举个比喻：

想象你有两条生产线，一条是做“准备工作”（prefill），一条是做“产出工作”（generate）。顺序约束像是你把所有“准备工作”的物资（对应 token）先排在仓库的入口，然后再把“产出工作”的物资排在后面。但在实际操作中，你可以一边把前面的准备工作物资搬进线，一边启动后面的产出线来处理已经准备好的物资。

具体到 IFB 的情况：

在某个时间点，你有多个请求／序列。有些请求／序列的上下文部分尚未完成，有些已经完成上下文且进入生成阶段。
系统可以构造一个 batch，其中 input tensor 是 packed 的，按顺序把那些尚在上下文阶段的序列的 tokens
放在前面，把那些已经开始生成的序列的 tokens（即生成阶段 token +上下文）放在后面。
GPU kernel 在处理这个 packed tensor的时候，是可以对前部分做上下文阶段的操作，对后部分做生成阶段的操作，中间并不需要等所有 prefill 完全结束才开始后续generate 的步骤，只要顺序满足 packed tensor 的约定。

源码阅读

One of TensorRT-LLM’s key features is it’s inflight batching scheduler
and runtime that handles simultaneously scheduling and executing
requests in both context and generation phases. Max-Num tokens, in
conjunction with Max-Batch size dictates how and when the scheduler
schedules new and current requests, and understanding their role and
how to tune them can provide significant performance benefits.

是TensorRT-LLM中的重要组件，核心代码位于cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp
在我给你的学习代码下运行

cd /LLM-Inference-Notes/third-party/tensorrt-llm/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp

TensorRT-LLM in-flight batching

TensorRT-LLM in-flight batching 源码入口 &说明

以下是几个关键源码文件／目录，以及它们在实现 in-flight batching 中起的作用。
核心代码：

cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp / .h

两个调度器源码：

capacityScheduler.cpp（容量/资源视角，支持 MAX_UTILIZATION 与
GUARANTEED_NO_EVICT）
microBatchScheduler.cpp（把“可放得下”的请求切成当前迭代的
micro-batch）。

运行时/缓冲相关：

runtime/tllmBuffers.h、batch_manager/transformerBuffers.cpp（执行步的张量准备与拷贝常在这些类里）。

trtGptModelInflightBatching.cpp

TrtGptModelInflightBatching

这个文件是TensorRT-LLM中in-flight batching系统的核心实现，包含了 TrtGptModelInflightBatching 类的完整定义

函数定义

// 构造一个“支持 In-Flight Batching（飞行中批处理）”的 GPT 模型实例。
// 这个类在 TRT-LLM 的 C++ 运行时中承担“调度请求 + 执行一步 + 回填结果”的编排角色。
TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr<nvinfer1::ILogger> logger,        // TensorRT 的日志接口：打印 INFO/WARN/ERROR/TRACE，便于排查问题ModelConfig const& modelConfig,                    // 模型结构配置：层数、隐藏维度、注意力类型、是否开启 LoRA/量化 等WorldConfig const& worldConfig,                    // 部署/并行环境配置：设备编号、并行方式（TP/PP）、当前是否是流水线最后一段等RawEngine const& rawEngine,                        // 预构建好的 TensorRT 引擎/权重与算子集：真正“跑起来”的底层执行资源bool ctxGenFusion,                                 // 是否启用“Context(预填充) + Generation(解码) 融合”的执行优化开关executor::ExecutorConfig const& executorConfig,    // 执行层配置：解码策略(top-k/p/温度/beam)、调试与性能旋钮、额外输出等bool isLeaderInOrchMode)                           // 在编排(Orchestration)模式下，本进程/本 rank 是否是“主控/Leader”// --- 基类初始化：把通用的 GPT 运行时（无 IFB 特性）先建好 ---: TrtGptModel(modelConfig, worldConfig, executorConfig)// --- 把关键配置存到成员变量，供后续调度与执行阶段使用 ---, mModelConfig(modelConfig)                        // 保存模型配置（后面会据此决定注意力插件、KV cache 形式、buffer 尺寸等）, mWorldConfig(worldConfig)                        // 保存并行/设备配置（决定哪些 rank 负责产出最终结果、设备选择等）// 初始化设备上下文：根据 WorldConfig 选择 CUDA 设备/流等（多卡/多并行时尤其重要）, mDevice{ runtime::utils::initDevice(worldConfig) }// 解码(生成)配置：若 ExecutorConfig 未提供就落回默认 DecodingConfig{}。// 这里包含温度、top-k/top-p、最大生成长度、beam 等，直接影响生成策略与性能/显存占用。, mDecodingConfig{ executorConfig.getDecodingConfig().value_or(executor::DecodingConfig{}) }// 运行时“性能旋钮”扩展配置：允许按需调优（如批大小上限、token 上限、调度策略细节等），// 便于在不同负载与硬件之间做吞吐/延迟/显存的权衡。, mExtendedRuntimePerfKnobConfig{ executorConfig.getExtendedRuntimePerfKnobConfig() }// 调试配置：是否输出更详细的日志、是否打开额外检查等（生产与开发环境可不同）。, mDebugConfig{ executorConfig.getDebugConfig() }// 额外模型输出：只有在“流水线并行的最后一个 rank”才有意义，才从 ExecutorConfig 中取出；// 否则设为 nullopt（不上报）。额外输出可能是隐藏态、注意力图、分类头等中间张量。, mAdditionalModelOutputs{worldConfig.isLastPipelineParallelRank()? executorConfig.getAdditionalModelOutputs(): std::nullopt}
{// 这里通常会继续做：// 1) 利用 rawEngine 创建/绑定 TensorRT 执行上下文与必需的插件（如 GPT attention + paged KV）。// 2) 根据 ctxGenFusion/DecodingConfig/PerfKnobs 配置调度器与运行时缓冲（RuntimeBuffers）。// 3) 初始化 KV Cache 管理（paged/contiguous、cross-KV 等），为 IFB 的“随时插入/回收”打好基础。// 4) 若启用调试或 TRACE 级日志，可打印当前 rank、最大 batch/token 限额、是否最后 PP rank 等关键信息。
}