使用 Huggingface 编译 LLM 模型¶

本教程将引导您完成使用 Torch-TensorRT 从 Huggingface 编译 LLM 模型的过程。我们还将介绍 Torch-TensorRT 中的 KV 缓存，它可以显著提高 LLM 推理的性能。代码可在 tools/llm 目录中找到。我们使用 run_llm.py 脚本来编译模型、生成输出并测量性能。

注意

这是一个实验性发布，API 在未来版本中可能会发生变化。

注意

Llama-2-7b-chat-hf 和 gpt2 模型的编译脚本和教程已合并到位于 tools/llm 目录中的统一 run_llm.py 脚本中。

tools/llm 目录概述¶

tools/llm 目录提供了以下工具来从 Huggingface 编译 LLM 模型：

run_llm.py：模型编译、生成输出和基准测试的主要入口点
run_vlm.py：编译和基准测试视觉语言模型 (VLM) 的入口点
静态缓存实用程序：用于 KV 缓存优化的 static_cache_v1.py 和 static_cache_v2.py
SDPA Attention：用于注册缩放点积注意力转换器和降低传递的 sdpa_converter.py 和 register_sdpa.py。
测试组件：模型特定的验证测试文件
实用函数：用于公共操作的 utils.py 和 cache_utils.py

支持的模型¶

我们已正式验证支持以下 LLM 系列：

模型系列	HuggingFace 模型卡	精度	KV 缓存支持？
GPT-2	gpt2	FP16, FP32	是
LLaMA 2	meta-llama/Llama-2-7b-chat-hf	FP16, FP32	是
LLaMA 3.1	meta-llama/Llama-3.1-8B-Instruct	FP16, FP32	是
LLaMA 3.2	meta-llama/Llama-3.2-1B-Instruct meta-llama/Llama-3.2-3B-Instruct	FP16, FP32	是
Qwen 2.5	Qwen/Qwen2.5-0.5B-Instruct Qwen/Qwen2.5-1.5B-Instruct Qwen/Qwen2.5-3B-Instruct Qwen/Qwen2.5-7B-Instruct	FP16, FP32	是
Gemma 3	google/gemma-3-1b-it	FP16, FP32	是

支持的 VLM 模型¶

我们已正式验证支持以下视觉语言模型 (VLM)：

模型系列	HuggingFace 模型卡	精度	KV 缓存支持？	组件支持
Qwen 2.5 VL	Qwen/Qwen2.5-VL-3B-Instruct	FP16, FP32	是 (仅限 static_v1)	仅限语言模型 (不支持图像编码器)
Eagle2	nvidia/Eagle2-2B	FP16, FP32	是 (仅限 static_v1)	语言模型和图像编码器均受支持

使用 run_llm.py 入门¶

主要入口点是 run_llm.py，它提供了完整的模型编译和基准测试工作流程。

基本用法¶

python tools/llm/run_llm.py \
  --model meta-llama/Llama-3.2-1B-Instruct \
  --prompt "What is parallel programming?" \
  --precision FP16 \
  --num_tokens 128 \
  --cache static_v2 \
  --benchmark

关键参数¶

--model：HuggingFace LLM 的名称或路径
--tokenizer：(可选) 分词器名称；默认为模型名称
--prompt：用于文本生成的输入提示
--precision：精度模式 (FP16, FP32)
--num_tokens：要生成的输出 token 数量
--cache：KV 缓存类型 (static_v1, static_v2, 或留空表示无 KV 缓存)
--benchmark：启用基准测试模式以进行性能比较
--enable_pytorch_run：也运行并比较 PyTorch 基线

其他用法示例¶

# Compare different models performance
python tools/llm/run_llm.py --model gpt2 --benchmark --enable_pytorch_run
python tools/llm/run_llm.py --model meta-llama/Llama-3.2-1B-Instruct --benchmark --enable_pytorch_run

# Generate the outputs (disable benchmarking) by specifying the number of tokens to generate. Default = 128
python tools/llm/run_llm.py --model gpt2 --prompt "What is parallel programming?" --num_tokens 128
python tools/llm/run_llm.py --model meta-llama/Llama-3.2-1B-Instruct --prompt "What is parallel programming?" --num_tokens 128

# Test different caching approaches
python tools/llm/run_llm.py --model meta-llama/Llama-3.2-1B-Instruct --cache static_v1
python tools/llm/run_llm.py --model meta-llama/Llama-3.2-1B-Instruct --cache static_v2

# Compare FP16 vs FP32 performance
python tools/llm/run_llm.py --model Qwen/Qwen2.5-1.5B-Instruct --precision FP16 --benchmark
python tools/llm/run_llm.py --model Qwen/Qwen2.5-1.5B-Instruct --precision FP32 --benchmark

使用 run_vlm.py 入门¶

对于视觉语言模型 (VLM)，请使用 run_vlm.py 来编译和基准测试同时处理文本和图像的模型。

基本用法¶

python tools/llm/run_vlm.py \
  --model Qwen/Qwen2.5-VL-3B-Instruct \
  --precision FP16 \
  --num_tokens 128 \
  --cache static_v1 \
  --enable_pytorch_run \
  --benchmark

关键参数¶

--model：HuggingFace VLM 的名称或路径
--prompt：用于生成的输入提示
--image_path：(可选) 输入图像文件的路径。如果未提供，将使用示例图像
--precision：精度模式 (FP16, FP32)
--num_tokens：要生成的输出 token 数量
--cache：KV 缓存类型 (static_v1 或留空表示无 KV 缓存)
--benchmark：启用基准测试模式
--enable_pytorch_run：也运行并比较 PyTorch 基线

Torch-TensorRT 中的 KV 缓存¶

我们提供了两个版本的静态 KV 缓存：static_cache_v1 和 static_cache_v2。在这两个实现中，我们将静态 KV 缓存张量添加为模型的输入/输出，而不将它们存储为外部内存。KV 缓存的长度 = 输入序列长度 + 输出序列长度 (由 --num_tokens 指定)。头部的数量和头部的维度由模型配置确定。

静态缓存 v1¶

static_cache_v1.py 在模型图中实现了 KV 缓存，如下所示：

class StaticCacheV1Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q, k, v, key_cache, value_cache, start_idx, end_idx, is_causal=True):
        # Concatenate new key/value pairs with existing cache
        new_key_cache = torch.cat((key_cache[:, :, :start_idx, :], k, key_cache[:, :, end_idx:, :]), dim=2)
        new_value_cache = torch.cat((value_cache[:, :, :start_idx, :], v, value_cache[:, :, end_idx:, :]), dim=2)

        # Compute attention using the updated cache
        attn_output = torch._C._nn.scaled_dot_product_attention(
            q,
            new_key_cache[:, :, :end_idx, :],
            new_value_cache[:, :, :end_idx, :],
            dropout_p=0.0,
            is_causal=is_causal
        )

        return attn_output, new_key_cache, new_value_cache

在上面的代码中，我们将新的键/值对与现有缓存连接起来并进行更新。为了计算注意力，我们使用更新后的缓存并从缓存中收集相应的键/值，直到包含当前 token 索引。上面的代码实际上是作为一个 FX 图转换来实现的。当我们导入 static_cache_v1.py 模块时，我们使用装饰器 @_aten_lowering_pass 将其注册为一个 Torch-TensorRT 降低传递。

注意

start_idx 和 end_idx 是当前 token 在缓存中的开始和结束索引。对于预填充阶段，start_idx 为 0，end_idx 为输入序列长度。对于解码阶段，start_idx 从输入序列长度开始，end_idx 等于 start_idx + 1。 start_idx 以 1 为增量，直到序列结束或达到要生成的 token 的最大数量。

静态缓存 v2¶

static_cache_v2.py 与 static_cache_v1.py 类似，但它使用的切片操作更少。它在模型图中实现了 KV 缓存，如下所示：

class StaticCacheV2Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q, k, v, key_cache, value_cache, start_idx, end_idx, is_causal=True):
        concat_keys = torch.cat((key_cache[:, :, :start_idx, :], k), dim=2)
        concat_values = torch.cat((value_cache[:, :, :start_idx, :], v), dim=2)
        new_key_cache = torch.cat((concat_keys, key_cache[:, :, end_idx:, :]), dim=2)
        new_value_cache = torch.cat((concat_values, value_cache[:, :, end_idx:, :]), dim=2)
        attn_output = torch._C._nn.scaled_dot_product_attention(
              q, concat_keys, concat_values, dropout_p=0.0, is_causal=is_causal
        )

        return attn_output, new_key_cache, new_value_cache

在上面的代码中，我们将现有的键/值缓存与当前 token 的键/值连接起来。我们使用它来直接计算注意力和更新键/值缓存，插入当前的键/值。上面的代码实际上是作为一个 FX 图转换来实现的。当我们导入 static_cache_v1.py 模块时，我们使用装饰器 @_aten_lowering_pass 将其注册为一个 Torch-TensorRT 降低传递。 start_idx 和 end_idx 的定义与 static_cache_v1.py 相同。

模型使用静态 KV 缓存编译后，模型的输入签名会发生变化。新的输入签名是 (input_ids, position_ids, key_cache_0, value_cache_0, ..., start_idx, end_idx)。键/值缓存张量的数量等于模型中注意力头的数量。我们可以使用 generate_with_static_cache 函数来生成输出。

生成输出¶

我们使用自定义的 generate 函数来生成输出。此函数执行标准的自回归解码，不使用 KV 缓存。还有一个 generate_with_static_cache 函数，它执行带有 KV 缓存的自回归解码。

generate_with_static_cache 函数负责准备编译了静态 KV 缓存的模型输入。模型输入是 input_ids, position_ids, key_cache_0, value_cache_0, …., start_idx, end_idx。我们用零初始化键/值缓存张量，并且对于生成的每个 token，新的键/值缓存张量就是模型的输出。

SDPA 转换器 (sdpa_converter.py)¶

使用 TRT Python API 转换缩放点积注意力操作。
支持因果和标准自注意力。

SDPA 注册 (register_sdpa.py)¶

这是一个 Torch-TensorRT 降低传递，它将 SDPA 的变体替换为 torch.nn.functional.scaled_dot_product_attention。
注册用于转换 torch.nn.functional.scaled_dot_product_attention 操作的 SDPA 转换器。

局限性和已知问题¶

滑动窗口注意力 (在 Gemma3 和 Qwen 3 模型中使用) 尚不支持。
某些模型架构 (例如 Phi-4) 在导出 torch 模型时存在问题。
对于 VLM，由于与 torch.export 不兼容的动态操作，Qwen2.5-VL 图像编码器的编译不被支持。

要求¶

Torch-TensorRT 2.8.0 或更高版本
Transformers v4.52.3
对于 VLM 模型 (run_vlm.py)：- pip install qwen-vl-utils (用于 Qwen2.5-VL-3B-Instruct 模型) - pip install flash-attn --no-build-isolation -v (用于 Eagle2-2B 模型)

使用 Huggingface 编译 LLM 模型¶

tools/llm 目录概述¶

支持的模型¶

支持的 VLM 模型¶

使用 run_llm.py 入门¶

基本用法¶

关键参数¶

其他用法示例¶

使用 run_vlm.py 入门¶

基本用法¶

关键参数¶

Torch-TensorRT 中的 KV 缓存¶

静态缓存 v1¶

静态缓存 v2¶

生成输出¶

SDPA 转换器 (sdpa_converter.py)¶

SDPA 注册 (register_sdpa.py)¶

局限性和已知问题¶

要求¶

文档

教程

资源