Executor API#

TensorRT-LLM 包含一个高级 C++ API,称为 Executor API,它允许您异步执行请求,支持飞行中批处理,并且无需定义回调函数。

软件组件(在下文中称为“客户端”)可以使用 executor.h 文件中定义的 API 与执行器交互。有关 API 的详细信息,请参阅 _cpp_gen/executor.rst

以下章节概述了 Executor API 中定义的主要类。

API#

Executor 类#

Executor 负责接收来自客户端的请求,并为这些请求提供响应。执行器的构造需要提供一个包含 TensorRT-LLM 引擎的目录路径,或者包含引擎和模型 JSON 配置的缓冲区。客户端可以使用 Executor 类的 enqueueRequestenqueueRequests 方法创建请求并将它们加入执行队列。加入队列的请求将由执行器调度执行,并且在主执行循环的每次迭代中,多个独立请求可以被批量处理(这个过程通常被称为连续批处理或迭代级批处理)。通过调用 awaitResponses 方法并提供请求 ID,可以等待特定请求的响应。另外,调用 awaitResponses 时省略请求 ID,可以等待任何请求的响应。Executor 类还允许使用 cancelRequest 方法取消请求,并使用 getLatestIterationStats 获取每次迭代和每个请求的统计信息。

Request 类#

Request 用于定义请求的属性,例如输入令牌 ID 和要生成的最大令牌数。streaming 参数可用于指示请求是应为每个新生成的令牌生成响应(streaming = true)还是仅在所有令牌生成后生成响应(streaming = false)。请求的其他强制参数包括采样配置(由 SamplingConfig 类定义),其中包含控制解码过程的参数,以及输出配置(由 OutputConfig 类定义),该配置控制特定响应的 Result 中应包含哪些信息。

构造请求时还可以提供可选参数,例如违禁词列表、停止词列表、客户端 ID,或用于 Prompt Tuning、LoRA 或推测解码的配置对象,或者要生成的序列数等。

Response 类#

Executor 类的 awaitResponses 方法返回一个响应向量。每个响应都包含与该响应关联的请求 ID,并且包含一个错误或一个 Result。在使用 getResult 方法尝试获取与此响应关联的 Result 之前,请使用 hasError 方法检查响应是否包含错误。

Result 类#

Result 保存给定请求的结果。它包含一个名为 isFinal 的布尔参数,指示这是否是给定请求 ID 将返回的最后一个 Result。它还包含生成的令牌。如果请求配置为 streaming = falsenumReturnSequences = 1,则将返回单个响应,isFinal 布尔值将设置为 true,并且所有生成的令牌都将包含在 outputTokenIds 中。如果使用 streaming = truenumReturnSequences = 1,则除了最后一个结果外,Result 将包含一个或多个令牌(取决于请求的 returnAllGeneratedTokens 参数),并且与此请求关联的最后一个结果的 isFinal 标志将设置为 true

请求的 numReturnSequences 参数控制为每个提示生成的输出序列数。使用此选项时,Executor 将为每个请求返回至少 numReturnSequences 个响应,每个响应包含一个 Result。在 Beam Search(beamWidth > 1)中,要返回的 Beam 数将受到 numReturnSequences 的限制,并且 Result 类的 sequenceIndex 属性将始终为零。否则,在 Sampling(beamWidth = 1)中,sequenceIndex 属性指示结果中生成序列的索引(0 <= sequenceIndex < numReturnSequences)。它包含一个名为 isSequenceFinal 的布尔参数,指示这是否是该序列的最后一个结果,还包含一个布尔参数 isFinal,指示何时为该请求生成了所有序列。当 numReturnSequences = 1 时,isFinalisSequenceFinal 相同。

这是一个示例,展示了当 numReturnSequences = 3 时,部分 3 个响应的样子

Response 1: requestId = 1, Result with sequenceIndex = 0, isSequenceFinal = false, isFinal = false
Response 2: requestId = 1, Result with sequenceIndex = 1, isSequenceFinal = true,  isFinal = false
Response 3: requestId = 1, Result with sequenceIndex = 2, isSequenceFinal = false, isFinal = false

在此示例中,每个响应包含针对不同序列的一个结果。第二个 Result 的 isSequenceFinal 标志设置为 true,表示它是 sequenceIndex = 1 的最后一个结果,但是,每个 Response 的 isFinal 标志设置为 false,因为序列 0 和 2 尚未完成。

发送具有不同 Beam Width 的请求#

如果满足以下条件,执行器可以处理具有不同 Beam Width 的请求

  • 模型构建时 max_beam_width > 1

  • 执行器配置时 maxBeamWidth > 1(配置的 maxBeamWidth 必须小于或等于模型的 max_beam_width)。

  • 请求的 Beam Width 小于或等于配置的 maxBeamWidth

执行器可以同时调度具有相同 Beam Width 的连续请求。对于具有两个不同 Beam Width xy 的连续请求,只有在所有具有 Beam Width x 的请求都处理完毕后,才会调度具有 Beam Width y 的请求。这使得运行时可以在没有正在处理的请求时重新配置自身以适应新的 Beam Width。每次检测到请求的 Beam Width 与当前配置不同时,都会自动发生重新配置。等待先前的请求完成并重新配置运行时可能会导致显著的开销并降低整体吞吐量。

使用 Logits 后处理器控制输出#

或者,您可以通过提供 Executor::LogitsPostProcessorConfig 的实例来修改网络产生的 logits。例如,此功能可用于生成 JSON 格式的输出。Executor::LogitsPostProcessorConfig 以以下形式指定了一个命名回调映射

std::unordered_map<std::string, function<Tensor(IdType, Tensor&, BeamTokens const&, StreamPtr const&, std::optional<IdType>)>>

映射键是与该 logits 后处理回调关联的名称。然后,每个请求可以指定用于该特定请求的 logits 后处理器的名称(如果有)。

回调的第一个参数是请求 ID,第二个是 logits 张量,第三个是请求到目前为止产生的令牌,第四个是 logits 张量使用的操作流,最后一个是可选的客户端 ID。回调返回修改后的 logits 张量。多个请求可以共享相同的客户端 ID,回调可以根据客户端 ID 使用不同的逻辑。

您必须使用流来访问 logits 张量。例如,要执行与 bias 张量的加法运算,加法运算在该流上入队。另外,您可以调用 stream->synchronize(),但这会减慢整个执行管道。

执行器还包含一个 LogitsPostProcessorBatched 方法,该方法可以批量更改多个请求的 logits。批量方法允许进一步优化并减少回调开销。

std::function<void(std::vector<IdType> const&, std::vector<Tensor>&, std::vector<std::reference_wrapper<BeamTokens const>> const&, StreamPtr const&, std::vector<std::optional<IdType>> const&)>

可以在 LogitsPostProcessorConfig 中指定单个批量回调。每个请求可以通过将 logits 后处理器的名称指定为 Request::kBatchedPostProcessorName 来选择应用此回调。

注意:目前,两种回调变体都不支持 STATIC 批处理类型。

在多 GPU 运行时,回调默认在第一个张量并行组的所有 rank 上调用。为确保正确执行,请在这些 rank 上复制回调访问的客户端状态。如果复制开销大或不可行,请使用 LogitsPostProcessorConfig::setReplicate(false) 使回调仅在 rank 0 上调用。执行器会在内部广播采样令牌以确保正确执行。

使用引导解码生成结构化输出#

引导解码控制生成输出,使其适合预定义的结构化格式,例如 JSON 或 XML。目前,引导解码支持 XGrammar 后端。

要启用引导解码,构造 Executor 时必须提供 GuidedDecodingConfig 的有效实例。GuidedDecodingConfig 应包含一些分词器信息,包括 encodedVocabtokenizerStr(可选)和 stopTokenIds(可选)。对于 Hugging Face 分词器,可以通过以下方式提取这些信息

encoded_vocab = tokenizer.get_vocab()
encoded_vocab = [token for token, _ in sorted(encoded_vocab.items(), key=lambda x: x[1])]
tokenizer_str = tokenizer.backend_tokenizer.to_str()
stop_token_ids = [tokenizer.eos_token_id]

有关更多详细信息,请参阅 tensorrt_llm/llmapi/tokenizer.py。您可以将这些材料转储到磁盘,然后重新加载到 C++ 运行时中使用。

每个请求可以选择性地指定一个 GuidedDecodingParams,该参数定义了所需的结构化格式。目前,它支持四种类型

  • GuidedDecodingParams::GuideType::kJSON:生成的文本符合 JSON 格式;

  • GuidedDecodingParams::GuideType::kJSON_SCHEMA:生成的文本符合带有附加限制的 JSON 格式;

  • GuidedDecodingParams::GuideType::kREGEX:生成的文本符合正则表达式;

  • GuidedDecodingParams::GuideType::kEBNF_GRAMMAR:生成的文本符合扩展巴科斯-瑙尔范式(EBNF)语法。

后三种类型应与提供给 GuidedDecodingParams 的 schema/regex/grammar 一起使用。

获取任意输出张量#

Executor API 允许用户读取模型的任意输出。例如,可以获取隐藏状态或 logits。

将张量标记为输出#

要使用此功能获取张量,在构建 TRT 引擎之前,需要在模型定义中将其标记为输出(例如,添加 topk_logits.mark_output("TopKLogits"))。

配置 Executor#

假设您计划使用的 TensorRT 引擎有一个名为 TopKLogits 的张量被标记为输出,您应该通过将其名称传递给 ExecutorConfig 配置对象来配置 Executor 从此输出张量读取

auto const executorConfig = ExecutorConfig{};

std::vector<executor::AdditionalModelOutput> additionalOutputs{
    executor::AdditionalModelOutput{"TopKLogits", /*whether or not to get the output for the context too */ true}};
executorConfig.setAdditionalModelOutputs(additionalOutputs);

// ... set more configuration options if needed
// ... create the `Executor` instance

请求附加输出#

构造一个请求并将其加入执行器队列,以查询此张量输出

executor::Request request{requestTokens, parameters.maxOutputLength, true, executor::SamplingConfig{},
    executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}};
executor.enqueueRequest(request);

输出可以在每个响应的 additionalOutputs 属性中找到。

关于上下文输出的注意事项#

如果启用了 KV cache 重用,上下文输出将不包含已重用部分的上下文输出。这部分的输出只能从生成此部分 KV cache 的具有相同前缀的先前请求中获取。

C++ Executor API 示例#

提供了两个 C++ 示例,展示了如何使用 Executor API,这些示例位于 examples/cpp/executor 文件夹中。

Executor API 的 Python 绑定#

还提供了 Executor API 的 Python 绑定,以便从 Python 使用 Executor API。Python 绑定定义在 bindings.cpp 中,构建后可在 tensorrt_llm.bindings.executor 包中使用。在 Python 解释器中运行 'help('tensorrt_llm.bindings.executor') 将提供可用类的概述。

此外,还提供了三个 Python 示例,演示如何将 Python 绑定用于单 GPU 和多 GPU 模型的 Executor API。它们位于 examples/bindings 中。

使用 Triton 推理服务器进行飞行中批处理#

TensorRT-LLM 提供了一个 Triton 推理服务器 C++ 后端,其中包含使用飞行中批处理服务模型所需的机制。该后端也是一个很好的入门示例,展示了如何使用 TensorRT-LLM C++ Executor API 实现飞行中批处理。