解耦服务(实验性)#

注意

注意:此功能目前是实验性的,相关 API 在未来版本中可能会发生更改。

目前,TRT-LLM 支持 disaggregated-service,其中请求的上下文和生成阶段可以在不同的执行器上运行。TRT-LLM 的解耦服务依赖于执行器 API,请务必在阅读本文档之前阅读 执行器页面

有关 LLM 推理中解耦服务的更多信息,可以参考诸如 DistServeSplitWise 等论文。

用法#

enum class RequestType
{
    REQUEST_TYPE_CONTEXT_AND_GENERATION = 0,
    REQUEST_TYPE_CONTEXT_ONLY = 1,
    REQUEST_TYPE_GENERATION_ONLY = 2
};

TRT-LLM 执行器可以执行三种类型的请求:REQUEST_TYPE_CONTEXT_AND_GENERATIONREQUEST_TYPE_CONTEXT_ONLYREQUEST_TYPE_GENERATION_ONLY。执行器实例可以执行仅上下文请求的上下文阶段或仅生成请求的生成阶段。当执行器完成仅上下文请求的上下文阶段时,它会维护相应的 KV 缓存,该缓存将由执行器请求用于后续的仅生成请求。

请注意,对于 disaggregated-service,应设置环境变量 TRTLLM_USE_MPI_KVCACHE=1

以下是使用解耦服务的一些关键 API

Request request{...};

request.setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_CONTEXT_ONLY);

auto contextRequestId = contextExecutor.enqueueRequest(request);

auto contextResponses = contextExecutor.awaitResponses(contextRequestId);

auto contextPhaseParams = contextResponses.back().getResult().contextPhaseParams.value();

request.setContextPhaseParams(contextPhaseParams);

request.setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_GENERATION_ONLY);

auto generationRequestId = generationExecutor.enqueueRequest(request);

auto genResponses = generationExecutor.awaitResponses(generationRequestId);

generationExecutor 将需要来自相应 contextExecutor 的数据,例如 KV 缓存,具体取决于附加到请求的 contextPhaseParams,因此请确保在获得 generationExecutor 的响应之前不要关闭相应的 contextExecutor。

在上面的代码示例中,contextExecutor 分配的 contextRequestId 和 generationExecutor 分配的 generationRequestId 是独立的,用户有责任管理从仅上下文请求的 requestId 到仅生成请求的 requestId 的映射。contextResponses 包含上下文阶段生成的第一个输出令牌,而 genResponses 也包含 contextExecutor 生成的第一个输出令牌,因此所有输出令牌都可以从 generationExecutor 的响应中获得。

disaggregated-service usage

disaggregated-service 中需要一个 orchestrator 来管理多个执行器实例并将请求路由到不同的执行器,TRT-LLM 在 cpp/include/tensorrt_llm/executor/disaggServerUtil.h 中提供了类 DisaggExecutorOrchestrator 来启动多个执行器实例,但是,DisaggExecutorOrchestrator 仅以简单的循环策略将请求路由到执行器,用户需要根据其使用场景为解耦服务实现自己的 orchestrator。

示例#

请参考 examples/cpp/executor/executorExampleDisaggregated.cpp

基准测试#

请参考 benchmarks/cpp/disaggServerBenchmark.cppbenchmarks/cpp/README.md

环境变量#

TRT-LLM 使用一些环境变量来控制解耦服务的行为。

  • TRTLLM_USE_MPI_KVCACHE:是否使用 MPI 传输 KV 缓存。目前,默认值为 0

  • TRTLLM_USE_UCX_KVCACHE:是否使用 UCX 传输 KV 缓存。目前,默认值为 0。要使用解耦服务,需要设置 TRTLLM_USE_MPI_KVCACHE=1TRTLLM_USE_UCX_KVCACHE=1

  • TRTLLM_PARALLEL_CACHE_SEND:如果设置为 1,contextExecutor 将尝试并行发送多个请求的 KV 缓存。默认值为 0

  • TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP:如果设置为 1,generationExecutor 将不会将 KV 缓存传输与模型推理重叠。默认值为 0

  • TRTLLM_ENABLE_KVCACHE_RECEIVE_PARALLEL:当 generation 排名从单个上下文实例中的多个上下文排名接收 KV 缓存时,它将按顺序从每个排名接收 KV 缓存。如果设置为 1,generation 排名将并行地从一个上下文实例中的每个排名接收 KV 缓存。默认值为 0

  • TRTLLM_REQUEST_KV_CACHE_CONCURRENT:如果设置为 1,generationExecutor 会为每个上下文执行器准备独立的资源以接收 KV 缓存,从不同上下文执行器接收到 KV 缓存的请求将被并发处理。如果设置为 0,则 generation 执行器将重用相同的资源来顺序处理每个请求的 KV 缓存传输,从而减少 KV 缓存传输使用的资源,从而降低内存耗尽的风险。默认值为 0

  • TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER:TRT-LLM 通常会将非连续数据复制到临时缓冲区中,然后再发送 KV 缓存。如果设置为 1,TRT-LLM 将尝试直接传输每个 KV 缓存块,从而消除额外的复制。默认值为 0

  • TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE:默认情况下,TRT-LLM 使用 stream-ordered memory allocator 来分配临时缓冲区。如果此环境变量设置为 #Size,则 TRT-LLM 将使用 cudaMalloc 来分配大小为 #Size 的缓冲区以进行 KV 缓存传输。默认值为 0。用户可以设置 TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=1GB 以使用 cudaMalloc 分配 1 GB 的缓冲区以进行 KV 缓存传输。

  • TRTLLM_KVCACHE_TRANSFER_USE_ASYNC_BUFFER:如果设置为 1,TRT-LLM 将使用 cudaMallocAsync 来分配缓冲区以进行 KV 缓存传输。默认值为 0。此环境变量仅在 TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE 大于 0 时才有效。

  • TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM:并发 KV 缓存发送的最大数量。默认值为 4。此环境变量仅在 TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE 大于 0 时才有效。

问题排查和常见问题解答#

常规常见问题解答#

问:TRT-LLM 中分离式服务的局限性是什么?

答:目前,仅支持 decoder-only enginebeamWidth=1,并且模型每一层的 KV 缓存需要是同质的,具有相同的数据类型和相同数量的注意力头。

问:分离式服务使用的引擎与其他引擎不同吗?

答:不。构建引擎的参数没有特殊要求。

问:上下文执行器和生成执行器使用的引擎需要相同吗?

答:不。上下文执行器和生成执行器使用的引擎可以不同,它们的并行度可以是异构的,即 TP、PP 可以不同,TRT-LLM 将处理 KV 缓存的异构性。

问:TRT-LLM 是否支持运行多个上下文执行器实例和生成执行器实例?

答:是。TRT-LLM 支持同时运行多个上下文执行器和生成执行器,并且每个执行器可以使用不同的引擎,但用户有责任将请求路由到不同的执行器并管理 requestId

问:一个执行器可以处理仅上下文请求和仅生成请求吗?

答:是的,但不建议这样做,TRT-LLM 没有针对执行器处理混合的仅上下文请求和仅生成请求的情况实现适当的调度,最好在不同的执行器上运行仅上下文请求和仅生成请求。

问:TRT-LLM 中的分离式服务是否支持多 GPU 和多节点?

答:是的,建议不同的执行器使用不同的 GPU。我们支持仅上下文执行器和仅生成执行器在同一节点或不同节点上运行。 每个执行器使用的 participantIdsdeviceIds 需要由用户显式设置,并且每个执行器的 participantIds 不得相交。

问:TRT-LLM 中分离式服务的要求是什么?

答:TRT-LLM 目前需要 UCX-backend CUDA-aware MPI,TRT-LLM 使用 CUDA-aware MPI 实现 KV 缓存传输,并且未来版本将支持更多用于 KV 缓存传输的通信组件。

调试常见问题解答#

问:如何处理错误 Disaggregated serving is not enabled, please check the configuration?

答:请设置环境变量

export TRTLLM_USE_MPI_KVCACHE=1

export TRTLLM_USE_UCX_KVCACHE=1

当环境变量 TRTLLM_USE_MPI_KVCACHE=1 设置时,TRT-LLM 将使用 CUDA-aware MPI 传输 KV 缓存。 所有涉及的执行器进程必须共享同一个 MPI world 通信器。 因此,使用 TRTLLM_USE_MPI_KVCACHE=1 时,TRT-LLM 仅支持通过 MPI 启动多个执行器。 此外,执行器的 CommunicationMode 必须设置为 kLEADERkORCHESTRATOR,且 SpawnProcesses=false 用于 disaggregated-service。 当设置 TRTLLM_USE_UCX_KVCACHE=1 时,这些限制不适用。

问:为什么一些分析工具显示,即使在配备 NVLink 的设备上,TRT-LLM 的 KV 缓存传输也没有利用 NVLink?

答:确保 TRT-LLM 正在使用 UCX-backend CUDA-aware MPI 运行,并使用 ucx_info -v 检查 UCX 的版本。 如果 UCX 的版本 <=1.17,请设置环境变量 UCX_RNDV_FRAG_MEM_TYPE=cudaUCX_MEMTYPE_CACHE=n 以启用 NVLink。 对于 BlackWell 架构 GPU,需要 UCX 版本 >=1.19 才能启用 NVLink。 如果 UCX 的版本 >=1.18,则有几种方法可以启用 NVLink

  1. 设置环境变量 UCX_CUDA_COPY_ASYNC_MEM_TYPE=cudaUCX_CUDA_COPY_DMABUF=noUCX_MEMTYPE_CACHE=nUCX_RNDV_PIPELINE_ERROR_HANDLING=y

  2. 设置环境变量 TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$SizeUCX_MEMTYPE_CACHE=nUCX_RNDV_PIPELINE_ERROR_HANDLING=y。 $Size 表示 KV 缓存传输的缓冲区大小,建议大于最长请求的 KV 缓存大小。

问:TRT-LLM 是否支持使用 GPU direct RDMA 进行节点间 KV 缓存传输?

答:是的,TRT-LLM 支持使用 GPU direct RDMA 进行节点间 KV 缓存传输,但默认情况下未启用。 有几种方法可以启用 GPU direct RDMA

  1. 设置环境变量 UCX_RNDV_FRAG_MEM_TYPE=cudaUCX_MEMTYPE_CACHE=nUCX_RNDV_PIPELINE_ERROR_HANDLING=y

  2. 设置环境变量 TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$SizeUCX_MEMTYPE_CACHE=nUCX_RNDV_PIPELINE_ERROR_HANDLING=y,$Size 表示 KV 缓存传输的缓冲区大小,建议大于最长请求的 KV 缓存大小。 要在使用 GPU direct RDMA 时获得最佳性能,建议在设置 TRTLLM_USE_MPI_KVCACHE=1 时,在 MPI 初始化之前创建 CUDA 上下文。 一种可能的方法是依靠 MPI 环境变量在 MPI 初始化之前设置正确的设备。

问:KV 缓存传输的性能调优有什么指导原则吗?

答:根据用户的使用情况,某些环境变量集可以帮助避免 KV 缓存传输性能不佳。

环境变量集 A

export UCX_RNDV_FRAG_MEM_TYPES=cuda
export UCX_MEMTYPE_CACHE=n
export UCX_RNDV_PIPELINE_ERROR_HANDLING=y

此集合允许 KV 缓存传输利用节点内的 NVLink 和节点之间的 GDRDMA。

环境变量集 B

export UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda
export UCX_CUDA_COPY_DMABUF=no
export UCX_MEMTYPE_CACHE=n
export UCX_RNDV_PIPELINE_ERROR_HANDLING=y

与集合 A 相比,集合 B 在单个节点上可能提供稍微好一些的性能。但是,当跨多个节点传输 KV 缓存时,可能会导致程序不稳定。

环境变量集 C

export TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$Size
export UCX_MEMTYPE_CACHE=n
export UCX_RNDV_PIPELINE_ERROR_HANDLING=y

集合 C 可以实现比集合 A 和 B 更好的性能,无论是在节点内部还是节点之间。但是,如果 KV 缓存大小超过指定的 $Size,性能可能会下降。