TensorRT LLM自动部署技术架构解析

5 阅读9分钟

使用某机构TensorRT LLM AutoDeploy自动化推理优化

某机构TensorRT LLM使开发者能够为大语言模型(LLM)构建高性能推理引擎,但传统上部署新架构需要大量手动工作。为应对这一挑战,今日宣布在TensorRT LLM中作为测试版功能提供AutoDeploy。

AutoDeploy将开箱即用的PyTorch模型编译为推理优化图。这避免了将推理特定优化直接嵌入模型代码的需要,减少了LLM部署时间。AutoDeploy实现了从手动重新实现和优化每个模型,转向将模型编写与推理优化分离的编译器驱动工作流。

本文介绍AutoDeploy的架构与能力,并展示其如何支持最新的某机构Nemotron模型在发布时即可部署。

什么是AutoDeploy?

每个新的LLM架构都带来自身的推理挑战,从Transformer模型到混合视觉语言模型(VLM),再到状态空间模型(SSM)。将参考实现转变为高性能推理引擎通常需要添加KV缓存管理、跨GPU分片权重、算子融合以及针对特定硬件调优执行图。

AutoDeploy将此工作流转向编译器驱动方法。AutoDeploy不再要求模型作者手动重新实现推理逻辑,而是自动从开箱即用的PyTorch模型中提取计算图,并应用一系列自动转换,生成推理优化的TensorRT LLM图。这使得可以在PyTorch中描述一次模型,将缓存、分片、内核选择和运行时集成等推理相关问题委托给编译器和运行时。

此方法特别适用于长尾模型,包括新的研究架构、内部变体和快速演进的开源模型——这些场景下手动重新实现通常不切实际或得不偿失。AutoDeploy支持在发布时即可部署并达到有竞争力的基准性能,同时为模型成熟后的增量优化保留清晰路径。

AutoDeploy提供:

  • 无缝模型转换:自动将Hugging Face模型转换为TensorRT LLM图,无需手动重写
  • 单一事实来源:保持原始PyTorch模型作为规范定义
  • 推理优化:应用分片、量化、KV缓存插入、注意力融合、CUDA Graphs优化等
  • 发布时部署:支持即时部署,并随时间持续改进性能
  • 开箱即用:作为TensorRT LLM的一部分提供,包含示例和文档

AutoDeploy可用于:

  • 新型或实验性架构:快速部署研究模型、混合设计或新型token混合(注意力)机制
  • 长尾模型支持:服务内部、微调或较少见的模型,无需定制推理实现
  • 快速性能提升:快速达到有竞争力的基准性能,然后增量优化
  • 统一训练到推理工作流:保持PyTorch作为模型定义,同时依赖TensorRT LLM进行运行时集成

AutoDeploy目前支持超过100个文本到文本LLM,并提供对VLM、SSM以及Llama模型家族和某机构Nemotron 3 Nano等性能优化模型的早期支持。

AutoDeploy技术背景

AutoDeploy位于原始Hugging Face模型和TensorRT LLM运行时之间。LLM API接受模型名称或检查点目录,返回一个高级LLM对象。在底层,该对象可以使用AutoDeploy(自动化)或手动后端。

如图1所示,AutoDeploy路径自动提取图、应用优化并生成推理优化图。手动路径需要工程师重写模型(添加KV缓存逻辑、注意力内核、分片、内核融合等),然后通过同一运行时执行。

[图1:AutoDeploy模式概览及其与TensorRT LLM运行时的集成]

图捕获与模式匹配

AutoDeploy使用torch.export API将模型捕获为标准化的Torch图,由核心ATen操作和自定义(用户或AutoDeploy提供)操作组成。导出的图随后经历一系列自动转换,对常见构建块的图表示进行模式匹配和规范化。

在此初始步骤中,AutoDeploy确保专家混合(MoE)、注意力、RoPE或状态空间层等常见构建块,使用表示为自定义操作和图中的单个节点的参考实现来表示。

图2展示了注意力如何跨所有模型表示为PyTorch中单一、易于解释的自定义算子。

[图2:AutoDeploy确保对常见构建块(如注意力)使用规范化表示,以简化下游性能优化,如缓存和内核选择]

这种方法确保了模型支持的无缝接入过程,该过程与性能优化和运行时集成解耦。

此外,模型接入发生在完全自动化(通过模式匹配)和(完全)手动重写之间的滑动尺度上,以确保最终模型图能够完全执行模型。模型作者可以通过将相关操作装饰为PyTorch自定义算子,向模型图中注入自定义内核。AutoDeploy编译器不会修改相关算子(图3)。

[图3:向AutoDeploy模型图中注入自定义算子的示例]

分片、融合与性能优化

在后续阶段,AutoDeploy通过类似编译器的传递自动应用性能优化,结合融合传递、性能调优配方以及向图表示中插入优化内核。在此阶段,模型还会基于可用启发式或预先指定的分片提示(重用Hugging Face分片提示)进行多GPU推理分片。

灵活的注意力与缓存支持

在图捕获和模式匹配期间,AutoDeploy将token混合(例如注意力)算子表示为简单的仅预填充操作,以AutoDeploy规范化的参考算子表达。这在图3中以softmax注意力为例进行了描绘。

系统随后自动处理切换到性能优化的注意力内核,并自动将token混合算子的缓存机制集成到TensorRT LLM优化缓存管理器系统中。目前,AutoDeploy可以处理由softmax注意力、状态空间层(Mamba2)、线性注意力(DeltaNet)和因果卷积任意组合而成的模型。

为具有缓存的其他算子添加支持遵循严格接口,且易于扩展。

编译工具

AutoDeploy与常见的开箱即用工具集成,用于进一步编译和降低模型,例如torch.compile、与CUDA Graphs集成以处理固定批大小的仅解码批处理、多流优化等。

运行时集成

AutoDeploy处理将模型集成到优化后的TensorRT LLM运行时的所有方面,包括重叠调度器、分块预填充、推测解码或缓存与状态管理等功能,而无需让模型作者承担模型与运行时之间相互交织的依赖关系。

AutoDeploy性能示例:Nemotron 3 Nano

为了评估AutoDeploy的能力,团队接入了某机构Nemotron 3 Nano——一个混合MoE模型。手动调优此类模型进行推理通常需要数周,而AutoDeploy在数天内完成了接入,随后进行的增量优化表现与手动调优的基准相当。

在单个某机构Blackwell DGX B200 GPU上,AutoDeploy的表现与TensorRT LLM中手动优化的基准持平(图4)。对于延迟敏感和高吞吐量应用,分别实现了每用户每秒高达350个token和每秒高达13,000个输出token。

[图4:Nemotron 3 Nano FP8在TensorRT LLM中当前默认PyTorch(手动)后端与AutoDeploy后端的在线性能对比]

数据收集条件:ISL/OSL 1k/1k,TP=1,在某机构DGX B200上使用TensorRT LLM v1.3.0rc1、trtllm-serve和AIPerf基准测试工具。

要自行复现结果,请遵循某机构Nemotron 3 Nano检查点中概述的步骤。

模型接入示例:Nemotron-Flash

Nemotron-Flash是难以使用纯手动推理工作流支持的架构类型的代表性示例。这种混合研究模型结合了多种token混合器——包括状态空间层、softmax注意力和线性注意力——需要大量工程工作才能手动重新实现、优化和维护。

借助AutoDeploy,Nemotron-Flash层的现有优化传递可以开箱即用地重用,无需任何特定于模型的工程工作。新的层类型,如DeltaNet更新规则,作为增量扩展而非完全重写进行集成,并且可以重用于未来的模型接入工作。

结果,Nemotron-Flash在数天内完成了接入和性能优化,现已开箱即用。这凸显了AutoDeploy的核心优势:一旦优化表示为可重用的编译器传递,新的和非常规的架构就可以立即受益于完整的优化堆栈,大幅缩短部署时间,同时保持高推理性能。

团队使用TensorRT LLM AutoDeploy对Nemotron Flash 3B Instruct与Qwen2.5 3B Instruct(一个广泛采用、经过大量手动调优的相似规模模型)进行了基准测试。在图5的基准测试场景(ISL/OSL=8k/16k)中,Nemotron-Flash优于Qwen2.5,凸显了新颖模型架构如何快速接入以实现生产就绪性能。

[图5:AutoDeploy中Nemotron Flash 3B与Qwen2.5 3B的吞吐量-延迟权衡曲线对比]

数据收集条件:ISL/OSL 8k/16k,TP=1,在某机构DGX H100上使用TensorRT LLM v1.3.0rc1、trtllm-serve和AIPerf基准测试工具。

开始使用TensorRT LLM AutoDeploy

TensorRT LLM AutoDeploy标志着向将推理优化视为编译器和运行时责任而非模型作者负担的转变。这种方法支持更快的实验、更广泛的模型覆盖,以及模型设计与部署之间更清晰的分离。

无需手动调优每个模型,只需描述一次架构,让系统应用图转换和优化内核。Nemotron Nano 3和Nemotron-Flash等早期成功案例表明,在不同模型架构上实现模型发布时的峰值性能部署是可行的。

TensorRT LLM AutoDeploy正在快速发展。如果对此功能进行实验或为其开发做贡献感兴趣,请查阅AutoDeploy文档和示例脚本。FINISHED