模型定义

TensorRT-LLM 提供了一个模型定义API，可用于定义大型语言模型。该API建立在强大的TensorRT Python API之上，用于在TensorRT中创建深度神经网络的图形表示。为了熟悉TensorRT API的核心概念，请先参考TensorRT文档中的核心概念部分。

在TensorRT-LLM中，tensorrt_llm.Builder类包含一个tensorrt.Builder对象。该实例在tensorrt_llm.Builder.create_network方法中用于创建tensorrt.INetworkDefinition类的实例。然后可以使用tensorrt_llm.functional中定义的函数来填充INetworkDefinition对象。

这种自由函数的一个简单例子是 tensorrt_llm.activation，它在模型的图中插入一个 tensorrt.IActivationLayer 节点：

# In tensorrt_llm.functional:

def activation(input: Tensor, act_type: trt.ActivationType) -> Tensor:
    layer = default_trtnet().add_activation(input.trt_tensor, act_type)   # default_trtnet() -> INetworkDefinition
    return _create_tensor(layer.get_output(0), layer)

为了让用户更加方便，LLMs中一些最标准的激活函数是从该函数派生的：

# In tensorrt_llm.functional:

relu    = partial(activation, act_type=trt.ActivationType.RELU)
sigmoid = partial(activation, act_type=trt.ActivationType.SIGMOID)

可以使用专门的激活函数来组装更高级的函数，例如silu激活函数：

# In tensorrt_llm.functional:

def silu(input: Tensor) -> Tensor:
    return input * sigmoid(input)

当使用TensorRT-LLM的模型定义API时，会组装一个网络图。该图稍后可以使用由tensorrt.ILayer类暴露的图遍历API进行遍历或转换。该图还将在引擎编译期间由TensorRT进行优化，如下一节所述。

编译

一旦填充完成，tensorrt.INetworkDefinition的实例可以通过tensorrt.Builder编译成一个高效的引擎。在TensorRT-LLM中，这是通过tensorrt_llm.Builder类的build_engine成员函数完成的，该函数调用了tensorrt.Builder对象的build_serialized_network方法。如果一切顺利，该调用将生成一个tensorrt.IHostMemory类的实例。该对象是一个优化的TensorRT引擎，可以存储为二进制文件。

TensorRT 编译器

TensorRT编译器可以遍历图，为每个操作和可用的GPU选择最佳的内核。关键的是，它还可以识别图中的模式，其中多个操作是融合到单个内核中的良好候选者。这减少了所需的内存移动量和启动多个GPU内核的开销。

TensorRT 还将操作图编译成一个可以一次性启动的CUDA Graph，进一步减少了内核启动的开销。

TensorRT 编译器在融合层和提高执行速度方面非常强大，但有一些复杂的层融合——比如 FlashAttention ——涉及将许多操作交织在一起，这些操作无法自动发现。对于这些情况，你可以在编译时显式地用 plugins 替换图的部分内容。

模型引擎

引擎文件包含执行模型所需的信息，但在实际使用LLM时，需要的远不止通过模型的单次前向传递。TensorRT-LLM包含一个高度优化的C++运行时，用于执行构建的LLM引擎并管理诸如从模型输出中采样令牌、管理KV缓存以及将请求批量处理等过程。

您可以直接使用该运行时在本地执行模型，或者可以使用NVIDIA Triton推理服务器的TensorRT-LLM运行时后端为多个用户提供服务。

权重绑定

TensorRT引擎嵌入了网络权重，这些权重在编译时必须已知。 因此，在调用tensorrt_llm.Builder.build_engine之前，必须将权重绑定到模型定义中的参数。这导致了如下代码：

# The Linear operator exposes two parameters (see tensorrt_llm/layers/linear.py):
class Linear(Module):
    def __init__(self, ...):
        self.weight = Parameter(shape=(self.out_features, self.in_features), dtype=dtype)
        self.bias   = Parameter(shape=(self.out_features, ), dtype=dtype)

# The parameters are bound to the weights before compiling the model. See examples/gpt/weight.py:
tensorrt_llm_gpt.layers[i].mlp.fc.weight.value = fromfile(...)
tensorrt_llm_gpt.layers[i].mlp.fc.bias.value   = fromfile(...)

请注意，TensorRT 还可以 refit 引擎以在编译后更新权重。此功能可通过 refit_engine 方法在 tensorrt_llm.Builder 类中提供给 TensorRT-LLM 用户。

模式匹配与融合

TensorRT 在编译网络图时执行的关键步骤之一是操作融合。融合是一种众所周知的技术，用于提高执行 LLMs 时的效率。它有助于减少内存（DRAM）和计算核心（CUDA 核心以及位于 GPU 的 Streaming Multiprocessors 上的 Tensor 核心）之间传输的数据量。它还消除了内核启动开销（每次在 GPU 上启动内核时，都会有一个小的额外 CPU 成本，称为启动开销）。一个经典的例子是将激活函数与通常在网络中位于其前面的矩阵乘法（matmul）进行融合。

在TensorRT-LLM中，定义模型时，这样的序列可以写成：

c = tensorrt_llm.functional.matmul(a, b)
c = tensorrt_llm.functional.relu(c)

在推理过程中，如果上述序列在没有融合的情况下执行，c张量必须在matmul结束时写入全局内存，在relu中从同一内存中读取，并在relu之后再次写入。如果没有其他操作使用matmul和relu之间的中间值，这是次优的。这就是为什么在编译期间，TensorRT会识别该模式并自动生成一个GPU内核，该内核在matmul结束时应用relu，而无需通过全局内存的中间步骤。通过这种优化，c张量只写入一次（在relu之后），而不是两次，并且在两个操作之间不进行读取。

识别可以融合的操作序列的过程被称为模式匹配。TensorRT 拥有一个强大的模式匹配算法，能够识别许多可能的融合。所有识别的模式都会通过一个高级内核编译器转换为更高效的内核。

插件

可能的融合数量几乎是无限的，一些有用的融合涉及对图进行非常高级的修改。一个众所周知的例子是Flash-Attention技术，用于优化在许多LLM中找到的Multihead-Attention块。Flash-Attention需要对序列BMM-Softmax-BMM（其中BMM代表批量矩阵-矩阵乘积）中执行的算术进行修改，并且需要对两个批量矩阵乘积的for循环进行交错。这是非平凡的，并且不一定是你可以期望编译器“自己发现”的东西（或者它可能需要支持多面体模型）。

因此，即使TensorRT拥有强大的模式匹配算法并支持许多可能的融合，仍然存在无法识别不常见和/或非常高级模式的风险。为了克服这一不可避免的限制，TensorRT提供了一种强大的机制，称为plugins。

插件是插入到网络图定义中的节点，这些节点映射到用户定义的GPU内核。TensorRT-LLM使用了多个这样的插件。它们可以在cpp/tensorrt_llm/plugins目录中找到。

插件是用C++编写的，并遵循TensorRT开发者指南中使用自定义层扩展TensorRT部分描述的明确定义的接口。当在TensorRT引擎中执行时，插件会触发其封装的GPU内核的执行。一个相当简单的插件示例是QuantizeTensorPlugin，它在QuantizeTensorPlugin::enqueue成员函数中触发CUDA内核：

// In cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp:

int QuantizeTensorPlugin::enqueue(...) {
    if (inputDesc[0].type == DataType::kFLOAT) {
        invokeQuantization<float>(...);
    } else {
        invokeQuantization<half>(...);
    }
    return 0;
}

// In cpp/tensorrt_llm/kernels/quantization.cu:

template <typename T>
void invokeQuantization(...) {
    // The standard <<< >>> construct to launch CUDA kernels
    quantizedKernel<<<grid, block, 0, stream>>>(...);
}

有关TensorRT-LLM如何实现GPT Attention操作符的更多详细信息，请参阅多头、多查询和组查询注意力文档。

运行时

TensorRT-LLM 包含一个用于实现 Python 和 C++ 运行时的 API。运行时组件的作用是加载 TensorRT 引擎并驱动它们的执行。通常，对于像 GPT 这样的自回归模型，运行时负责加载实现输入序列处理以及生成循环主体的引擎。有关 C++ 运行时的详细信息，请参阅 GPT C++ 运行时 文档。

多GPU和多节点支持

即使TensorRT是为单GPU系统设计的，TensorRT-LLM增加了对多GPU和多节点系统的支持。这是通过使用TensorRT插件实现的，这些插件封装了来自NCCL库的通信原语，以及一个自定义插件，该插件在GPU之间存在All-to-all连接（通过DGX系统中的NVSwitch）时优化了All-Reduce原语。

通信插件可以在 cpp/tensorrt_llm/plugins/ncclPlugin 找到，多GPU功能在TensorRT-LLM模型定义API中暴露为：

# In tensorrt_llm/functional.py:

# Collectives.
def allreduce(tensor: Tensor, group: List[int]) -> Tensor
def allgather(tensor: Tensor, group: List[int], gather_dim: int = 0) -> Tensor

# Point-to-point communication primitives.
def send(tensor: Tensor, tgt: int) -> Tensor
def recv(tensor: Tensor, src: int) -> Tensor

多GPU支持可以通过两种不同的模型并行模式启用：张量并行和管道并行。前一种模式将模型的不同层拆分到多个GPU上。每个GPU运行整个网络，并在需要时与其兄弟GPU同步。管道并行将不同的层分配到不同的GPU上。每个GPU运行整个模型的一个子集，通信发生在这些层子集的边界上。张量并行通常会导致更平衡的执行，但需要GPU之间更多的内存带宽。管道并行减少了对高带宽通信的需求，但可能会导致负载均衡问题，并且在GPU利用率方面可能效率较低。

示例

以下是Llama 3.1 70B和Llama 3.1 405B的示例，展示了如何在TensorRT-LLM中执行多GPU和多节点推理。Llama 3.1 70B的示例在单个节点上执行多GPU推理，而Llama 3.1 405B的示例执行多节点推理。

Llama 3.1 70B

以下示例命令用于在单个节点上使用4个GPU构建一个运行Llama 3.1 70B模型的引擎，并采用张量并行（TP=4）。

folder_trt_llm=../TensorRT-LLM
model_dir=Llama-3.1-70B
ckpt_dir=ckpt_llama_3.1_70b
engine_dir=engine_llama_3.1_70b
dtype=bfloat16
tp_size=4
pp_size=1
kv_cache_type=paged
max_input_len=128
max_output_len=128
max_batch_size=4
workers=$(( tp_size * pp_size ))

python ${folder_trt_llm}/examples/llama/convert_checkpoint.py \
    --output_dir ${ckpt_dir} \
    --model_dir ${model_dir} \
    --dtype ${dtype} \
    --tp_size ${tp_size} \
    --pp_size ${pp_size} \
    --workers ${workers} \
    --use_parallel_embedding

trtllm-build \
    --output_dir ${engine_dir} \
    --checkpoint_dir ${ckpt_dir} \
    --gemm_plugin ${dtype} \
    --gpt_attention_plugin ${dtype} \
    --kv_cache_type ${kv_cache_type} \
    --max_input_len ${max_input_len} \
    --max_seq_len $(( max_input_len + max_output_len )) \
    --max_batch_size ${max_batch_size} \
    --workers ${workers}

以下示例命令通过在单个节点上使用4个GPU运行examples/run.py来执行推理。

input_text="Born in north-east France, Soyer trained as a"

mpirun -n $(( tp_size * pp_size )) \
    python ${folder_trt_llm}/examples/run.py \
        --engine_dir ${engine_dir} \
        --tokenizer_dir ${model_dir} \
        --input_text "${input_text}" \
        --max_output_len ${max_output_len}

Llama 3.1 405B

以下示例命令用于在2个节点上构建一个引擎，每个节点有8个GPU，以运行Llama 3.1 405B模型，并使用张量并行（TP=16）。尽管模型在多个节点上运行，但你可以在单个节点上构建引擎。

folder_trt_llm=../TensorRT-LLM
model_dir=Llama-3.1-405B
ckpt_dir=ckpt_llama_3.1_405b
engine_dir=engine_llama_3.1_405b
dtype=bfloat16
tp_size=16
pp_size=1
kv_cache_type=paged
max_input_len=128
max_output_len=128
max_batch_size=4
workers=8

python ${folder_trt_llm}/examples/llama/convert_checkpoint.py \
    --output_dir ${ckpt_dir} \
    --model_dir ${model_dir} \
    --dtype ${dtype} \
    --tp_size ${tp_size} \
    --pp_size ${pp_size} \
    --workers ${workers} \
    --use_parallel_embedding

trtllm-build \
    --output_dir ${engine_dir} \
    --checkpoint_dir ${ckpt_dir} \
    --gemm_plugin ${dtype} \
    --gpt_attention_plugin ${dtype} \
    --kv_cache_type ${kv_cache_type} \
    --max_input_len ${max_input_len} \
    --max_seq_len $(( max_input_len + max_output_len )) \
    --max_batch_size ${max_batch_size} \
    --workers ${workers}

以下示例脚本，launch_llama_3.1_405b.sh，展示了如何在每个节点有8个GPU的2个节点上使用Slurm进行推理。如果您使用不同的工作负载管理软件，关键是要运行examples/run.py命令。

#!/bin/bash
#SBATCH --account account
#SBATCH --partition partition
#SBATCH --job-name job-name
#SBATCH --time 1:00:00
#SBATCH --nodes 2

folder_trt_llm=../TensorRT-LLM
engine_dir=engine_llama_3.1_405b
model_dir=Llama-3.1-405B
max_output_len=128

input_text="Born in north-east France, Soyer trained as a"

srun \
    --ntasks-per-node 8 \
    --mpi pmix \
    python ${folder_trt_llm}/examples/run.py \
        --engine_dir ${engine_dir} \
        --tokenizer_dir ${model_dir} \
        --input_text "${input_text}" \
        --max_output_len ${max_output_len}

您可以通过在Slurm集群上运行脚本来执行推理。

sbatch launch_llama_3.1_405b.sh