从 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: 模型编译、生成输出和基准测试的主要入口点
静态缓存实用程序：用于 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	是

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

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 模型时存在问题。

要求¶

Torch-TensorRT 2.8.0 或更高版本
Transformers v4.52.3