模型定义

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

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

此类自由函数的一个简单示例是 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)

为了让用户更容易使用，LLM 中一些最标准的激活函数都由此函数派生而来

# 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——它们涉及将许多操作交织在一起，并且无法自动发现。对于这些情况，您可以在编译时使用插件明确替换图中的部分内容。

模型引擎

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

您可以直接使用该运行时在本地执行模型，也可以使用 NVIDIA Triton Inference Server 的 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），以便在编译后更新权重。TensorRT-LLM 用户可以通过 tensorrt_llm.Builder 类中的 refit_engine 方法使用此功能。

模式匹配和融合

TensorRT 编译网络图时执行的关键步骤之一是操作融合。融合是执行 LLM 时提高效率的常用技术。它有助于减少内存 (DRAM) 与计算核心（位于 GPU 流式多处理器上的 CUDA 核心和 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 提供了一种强大的机制，称为插件。

插件是插入到网络图定义中的节点，映射到用户定义的 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 注意力操作符的更多详细信息，请参阅多头、多查询和分组查询注意力文档。

运行时

TensorRT-LLM 包含一个 API，用于实现 Python 和 C++ 运行时。运行时组件的作用是加载 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 支持：张量并行 (Tensor Parallelism) 和流水线并行 (Pipeline Parallelism)。前一种模式将模型的不同层拆分到多个 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 和张量并行 (TP=4) 运行 Llama 3.1 70B 模型。

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}

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

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

以下示例命令构建了一个引擎，用于在每个节点有 8 个 GPU 的 2 个节点上使用张量并行 (TP=16) 运行 Llama 3.1 405B 模型。虽然模型在多个节点上运行，但您可以在单个节点上构建引擎。

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