优化 TensorRT-LLM 性能的最佳实践

本文档提供了一些优化TensorRT-LLM性能的最佳实践。

如何衡量性能?

TensorRT-LLM 可以使用 C++ 工具进行基准测试。我们正在积极开发 trtllm-bench 命令,这将成为推荐的方式来对 TensorRT-LLM 进行基准测试。

有关详细的性能数据和重现这些结果的步骤,请参阅此文档TensorRT-LLM后端也可用于测量TensorRT-LLM在线服务的性能。

构建选项以优化TensorRT-LLM模型的性能

这部分总结了如何构建引擎以提高运行时的性能。以下选项具有合理的默认值,但对于其中一些选项,可能需要进行调整以获得最佳性能。

请注意,其中一些功能以及如何启用它们可能会在未来发生变化。

max_batch_size, max_seq_lenmax_num_tokens

Explain `max_batch_size`, `max_seq_len` and `max_num_tokens`

关于这三个参数对GPU内存使用的影响,请参考memory.md

max_batch_size

max_batch_size 定义了引擎可以处理的最大请求数。

它控制着在运行时可以调度的最大请求数。

在构建引擎时设置足够高的max_batch_size,以确保它不会成为吞吐量的瓶颈,并且如果您希望获得更好的用户吞吐量或更低的延迟,可以使用运行时的max_batch_size进行调整,而无需重新构建引擎。

max_seq_len

max_seq_len 定义了单个请求的最大序列长度

从 TensorRT-LLM v0.11 开始,当启用 --remove_input_padding--context_fmha 时,max_seq_len 可以替换 max_input_lenmax_output_len,并且默认设置为 max_position_embeddings

使用默认的max_seq_len(即max_position_embeddings),除非你非常确定你的工作负载中的最大序列长度,否则不需要调整它。如果GPU内存非常有限,以至于无法确保即使一个请求达到max_seq_len,你将需要减少它。

max_num_tokens

max_num_tokens 定义了在每个批次中去除填充后的最大批处理输入令牌数。

max_num_tokens 从 v0.11 开始默认设置为 8192,你可以使用运行时 max_num_tokens 进行调整,而无需重新构建引擎。建议调整 --max_num_tokens 以获得更好的性能。

当输入填充未被移除时,最大令牌数等于不会生效。当输入填充被移除时(参见移除输入填充),来自不同序列的令牌会被打包在一起,并且最大令牌数可以设置为不同的(较低的)值,默认值为8192。

有两个方面必须考虑。首先,一些输入序列会比最大输入长度短。其次,当启用飞行中序列批处理时,上下文阶段的请求将与生成阶段的请求一起执行。这些后者的请求产生的标记比max_input_len少得多(最多beam_width个标记)。

使用更现实的max_num_tokens值允许TensorRT-LLM分配更多内存来存储KV缓存并同时执行更多请求。这导致效率提高。

适当增加max_num_tokens将有利于性能。 当将--max_num_tokens增加到一定程度时,GPU利用率将达到平稳状态, 超过该饱和点可能会损害首令牌延迟以及总端到端延迟。

另请参阅 chunked context

多配置文件

--multiple_profiles 在构建的引擎中启用多个TensorRT优化配置文件,这将有利于性能,特别是在GEMM插件被禁用时,因为更多的优化配置文件有助于TensorRT有更多机会选择更好的内核。

注意:此功能会增加引擎构建时间,但预计不会产生其他不良影响。

FP8 上下文融合多头注意力

--use_fp8_context_fmha 启用FP8上下文融合多头注意力。我们建议在使用fp8量化时启用此功能,以提高上下文阶段注意力性能。请注意,仅支持NVIDIA Hopper架构。

GPT 注意力插件和上下文融合多头注意力

GPT注意力插件和融合多头注意力内核默认启用。对于上下文阶段,使用--gpt_attention_plugin--context_fmha参数与trtllm-build来控制。

TensorRT-LLM GPT注意力插件使用高效的内核,并支持KV缓存的就地更新。与使用concat操作符更新KV缓存的实现相比,它减少了内存消耗,并消除了不必要的内存复制操作。

在上下文阶段启用融合多头注意力将触发一个使用单一内核执行MHA/MQA/GQA块的内核,更多详情请参阅此文档

移除输入填充

默认情况下,移除输入填充功能是启用的,--remove_input_padding参数在trtllm-build中用于控制它。

当输入填充被移除时,不同的标记会被打包在一起。这减少了计算量和内存消耗。更多详情,请参阅此文档

分页键值缓存

默认情况下启用了分页KV缓存,--paged_kv_cache参数在trtllm-build中用于控制它。

分页的KV缓存有助于更高效地管理KV缓存的内存(参见此文档)。它通常会导致批量大小的增加和效率的提高。

减少范数融合

有一个实验性功能称为“Reduce Norm Fusion”,可用于扩展自定义的AllReduce功能。当自定义的AllReduce已经启用时,可以通过使用--reduce_fusion enable参数与trtllm-build来启用此功能。

此功能旨在将ResidualAddLayerNorm内核在AllReduce之后融合为单个内核,从而提高端到端性能。

请注意,目前此功能仅支持llama模型。建议在批量大小较小且生成阶段时间占主导因素时启用此功能。

用户缓冲区

一个名为“用户缓冲区”的实验性功能可用于增强通信性能。可以通过使用--user_buffer enable参数与trtllm-build来启用此功能。 此功能旨在消除通信内核中从本地缓冲区到共享缓冲区的额外复制,从而提高端到端性能。 此功能必须与--reduce_fusion enable一起启用,并且仅支持FP8 LLAMA模型。

嵌入并行性、嵌入共享和查找插件

嵌入并行功能允许将嵌入表分片到多个GPU上,从而减少内存使用并提高吞吐量。嵌入共享功能允许在look_uplm_head层之间共享嵌入表,以减少内存使用。

建议启用嵌入并行性以提高吞吐量,在convert_checkpoint.py中使用--use_parallel_embedding--embedding_sharding_dim

如果满足以下条件,默认情况下会启用嵌入共享:

  1. look_uplm_head 层具有相同的权重。

  2. --gemm_plugin 在构建引擎时未使用。

  3. 对于张量并行的情况,必须设置-embedding_sharding_dim 0。换句话说,我们必须沿词汇维度启用嵌入并行。

详情请参见那些示例

门控MLP中的水平融合

Gated-MLP中的水平融合将两个Matmul操作合并为一个,然后跟随一个独立的SwiGLU内核。这可以有效减少延迟。此功能默认启用。

GEMM 插件

GEMM插件利用NVIDIA cuBLASLt来执行GEMM操作。在FP16和BF16上,建议启用以获得更好的性能和更小的GPU内存使用。在FP8上,建议禁用。

FP8 GEMM 插件用于小批量大小性能优化

FP8 gemm插件是一个实验性功能,旨在提高小批量大小情况下的性能(例如BS<=4),并且可以在构建FP8模型时通过--gemm_plugin fp8启用。虽然可以正确推断出较大批量大小的输入,但随着批量大小的增加,性能可能会下降。因此,目前仅建议在小批量大小场景中使用此功能以减少延迟。

Gated-MLP中的GEMM + SwiGLU融合

Gated-MLP中的GEMM + SwiGLU融合将两个Matmul操作和一个SwiGLU操作合并为一个单一的内核。目前,这仅在Hopper上支持FP8精度。虽然这种融合提高了性能,但由于丢弃了一个量化缩放因子,它可能会在FP8 PTQ中略微降低准确性。

我们建议在Hopper上运行的大型模型使用FP8精度时启用此功能。使用以下trtllm-build参数来启用它:

  • 对于大型模型:--use_fused_mlp=enable --gemm_swiglu_plugin=fp8

  • 对于小批量大小:--use_fused_mlp=enable --low_latency_gemm_swiglu_plugin=fp8 以提高延迟。

我们不建议为非常小的工作负载或如果精度损失不可接受的情况下启用此功能。

BERT 注意力插件和上下文融合多头注意力

BERT注意力插件和上下文融合多头注意力都推荐用于BERT模型。它们默认通过使用--bert_attention_plugin--context_fmha参数与trtllm-build一起启用。

优化TensorRT-LLM模型性能的运行时选项

这部分总结了可以调整的运行时配置旋钮,以增强已构建引擎的性能。请注意,目前可以使用Executor API以及TensorRT-LLM backend来修改配置。

容量调度策略

目前有三种批处理调度策略:GUARANTEED_NO_EVICT(默认), MAX_UTILIZATIONSTATIC_BATCH

调度策略可以设置为MAX_UTILIZATION,以便在启用飞行中序列批处理时,在每次前向循环迭代中尽可能多地打包请求。它通过积极调度请求来最大化GPU的利用率,但存在达到KV缓存大小限制时不得不暂停请求的风险。

为了在内存分配方面对KV缓存限制采取更保守的方法,应将CapacitySchedulerPolicy设置为GUARANTEED_NO_EVICT,以确保已启动的请求永远不会被暂停。

如果目标是最大化吞吐量,用户应该尝试MAX_UTILIZATION。 然而,他们需要记住,如果请求必须暂停,这可能会对延迟产生负面影响。

STATIC_BATCH 是一种遗留模式,不建议在生产环境中使用。

上下文分块策略

上下文分块将增加上下文和生成阶段之间批处理的机会,从而平衡每次迭代的计算量并增加吞吐量。

目前有两种上下文分块策略:FIRST_COME_FIRST_SERVED(默认) 和EQUAL_PROGRESS

FIRST_COME_FIRST_SERVED 应该能够实现整体更好的性能,而 EQUAL_PROGRESS 在理论上可以帮助确保大多数请求的首个令牌时间(TTFT) 相对相似。

批处理类型

批处理类型可以设置为INFLIGHT(默认)和STATIC。 建议使用INFLIGHT以提高吞吐量并减少延迟。

分页KV缓存中的最大令牌数和KV缓存空闲GPU内存比例

max_tokens_in_paged_kv_cachekv_cache_free_gpu_mem_fraction 参数可用于控制KV缓存管理器处理的最大令牌数。正确设置这些参数有助于更好地控制在推理过程中KV缓存管理器的可用内存量。请记住,增加KV缓存管理器的可用内存量往往会转化为更高的可达到吞吐量。

max_tokens_in_paged_kv_cache 标志直接设置 KV 缓存管理器中的最大令牌数。如果未设置,该值将根据 kv_cache_free_gpu_mem_fraction 设置计算。

kv_cache_free_gpu_mem_fraction 是一个介于 0.01.0 之间的浮点数,表示在加载模型后,GPU内存中用于KV缓存的最大比例。默认值为 0.90,意味着90%的可用GPU内存将用于保存KV缓存中的令牌。基于该值,TensorRT-LLM可以确定KV缓存管理器中的最大令牌数。

当两个参数都设置时,KV缓存管理器中的最大令牌数将设置为max_tokens_in_paged_kv_cache和根据KV缓存可用内存量计算出的值之间的较小值。

除非用户清楚地知道模型所需的KV缓存中的最大令牌数,否则建议保持max_tokens_in_paged_kv_cache未设置。对于kv_cache_free_gpu_mem_fraction,如果在同一GPU上没有执行其他程序,建议测试使用高达0.95的值以追求高吞吐量。请注意,kv_cache_free_gpu_mem_fraction参数不能设置为1.0,因为必须保留一些内存用于输入和输出。

最大注意力窗口大小

max_attention_window_size 标志设置了在使用滑动窗口注意力等技术时,生成一个令牌所关注的最大令牌数。详情请参阅此 文档。 它默认为最大序列长度(在构建引擎时为max_seq_len),这意味着该功能默认是禁用的。

当设置为比max_seq_len更小的值(在引擎构建期间),只会存储最后max_attention_window_size个令牌的KV缓存。如果运行时的输入序列长度超过max_attention_window_size值,准确性可能会开始下降,但运行时性能会更好(由于计算和GPU内存分配的减少)。用户可以修改该值以提高运行时性能,但会牺牲准确性。