深入解析推测解码：降低AI推理延迟的前沿技术

随着大型语言模型（LLM）的规模和复杂性不断增长，其自回归特性带来的推理延迟成为了一个关键挑战。为了在生产环境中提供实时、流畅的用户体验，学术界和工业界正在积极探索各种优化方法。其中，推测解码（Speculative Decoding） 作为一种创新的解码策略，能够在不牺牲模型准确性的前提下，显著降低LLM的推理延迟。

本文将严格遵循NVIDIA官方博客的核心内容，首先完整复现其对推测解码基本原理的介绍，然后在此基础上，深入探讨如何利用 NVIDIA TensorRT-LLM 和 NVIDIA Triton 推理服务器 等NVIDIA生态系统中的关键技术，在生产环境中高效实现并部署推测解码，从而为开发者提供一套从理论到实践的完整解决方案。

推测解码的基本原理

LLM的自回归推理过程是内存带宽密集型的。在每一步解码中，模型都需要加载完整的模型权重来生成下一个词元（token）。由于大部分时间都消耗在内存访问上，GPU的计算单元往往处于闲置状态，这导致了所谓的“延迟惩罚”。

为了解决这个问题，推测解码引入了一个巧妙的“草稿-验证”机制。它使用一个规模小得多、速度更快的草稿模型（draft model） 来一次性生成一个包含多个词元的草稿序列（K个词元）。然后，使用原始的、更大更精确的目标模型（target model） 对这个草稿序列进行一次并行的前向传播验证。

这个验证步骤是推测解码的核心。目标模型会一次性评估草稿序列中每个词元的生成概率。通过将目标模型的输出与草稿模型的输出进行比较，我们可以确定草稿序列中有多少个词元是“正确”的（即目标模型也会以高概率选择它们）。

假设在前 n 个词元上，两个模型的输出一致，那么我们可以直接接受这 n 个词元，并仅用目标模型在第 n+1 个位置的预测来生成一个新的词元。这样，我们通过一次目标模型的推理，就成功解码了 n+1 个词元，极大地提升了效率。如果草稿序列的第一个词元就未通过验证，系统则退化为传统解码，只接受目标模型生成的第一个词_元_。

这个过程可以用以下伪代码来表示：

# 伪代码：推测解码的基本逻辑
# 译者注：此代码为原文的Python伪代码，并添加了详尽的中文注释

def speculative_decoding(prompt, target_model, draft_model, K):
    """
    执行推测解码的函数。

    参数:
    prompt (list): 输入的初始词元序列。
    target_model (Model): 规模较大、精度较高的目标模型。
    draft_model (Model): 规模较小、速度较快的草稿模型。
    K (int): 草稿模型一次生成的词元数量。
    """
    
    # 初始化生成的词元序列
    generated_tokens = prompt.copy()

    while not is_done(generated_tokens):
        # 1. 草稿阶段：使用草稿模型生成K个候选词元
        # 这是推测执行的核心，快速生成一个“猜测”
        draft_tokens = []
        for _ in range(K):
            # 基于当前已生成的序列，用草稿模型预测下一个词元
            next_token_logits = draft_model(generated_tokens + draft_tokens)
            next_token = sample(next_token_logits)
            draft_tokens.append(next_token)

        # 2. 验证阶段：使用目标模型并行验证所有草稿词元
        # 这是一次性的、高效的验证过程
        target_logits_list = target_model(generated_tokens + draft_tokens)

        # 3. 接受/拒绝阶段：比较草稿和目标模型的输出
        all_draft_tokens_accepted = True
        for i in range(K):
            # 获取草稿模型在当前位置的预测词元
            draft_token = draft_tokens[i]
            # 获取目标模型在相同位置的预测分布
            target_logits = target_logits_list[i]
            
            # 随机接受或拒绝
            # 如果目标模型在当前位置也倾向于生成这个草稿词元，则接受
            # 否则，从目标模型的分布中重新采样一个词元，并终止本次循环
            if draft_token == sample(target_logits):
                # 接受当前草稿词元
                generated_tokens.append(draft_token)
            else:
                # 拒绝草稿，从目标模型分布中重新采样
                corrected_token = sample(target_logits, method="rejection")
                generated_tokens.append(corrected_token)
                all_draft_tokens_accepted = False
                break # 终止验证循环

        # 如果所有草稿词元都被接受了，我们还需要从目标模型中采样最后一个词元
        if all_draft_tokens_accepted:
            final_token = sample(target_logits_list[K])
            generated_tokens.append(final_token)
            
    return generated_tokens

推测解码的成功关键在于草稿模型和目标模型之间的一致性。如果草稿模型能够很好地“预测”目标模型的行为，那么每次迭代接受的词元数量就会增加，从而获得更高的加速比。

借助 NVIDIA TensorRT-LLM 高效实现推测解码

理论虽好，但在实际生产中高效实现推测解码却充满挑战，需要深度优化内存访问、并行计算和模型调度。NVIDIA TensorRT-LLM 作为一个端到端的LLM推理优化库，为推测解码提供了原生、开箱即用的支持，极大地简化了开发和部署过程。

TensorRT-LLM 的推测解码实现经过了高度优化，能够充分利用NVIDIA GPU（尤其是 Hopper 和 Ampere 架构）的硬件特性。它不仅包含了经过优化的MHA/MQA/GQA Kernel，还通过高效的K/V Cache管理和模型并行技术，确保了草稿模型和目标模型的推理流程无缝衔接。

在 TensorRT-LLM 中启用推测解码

在 TensorRT-LLM 中，启用推测解码非常直观。开发者在构建和运行模型时，只需指定目标模型和草稿模型即可。TensorRT-LLM 会自动处理两者之间的同步、验证和采样逻辑。

以下是一个简化的示例，展示了如何在 TensorRT-LLM Python API 中使用推测解码：

# 译者注：此代码为基于TensorRT-LLM API思想的示例代码，非官方完整代码，用于说明API使用方法

import tensorrt_llm
from tensorrt_llm.runtime import SamplingConfig, ModelRunner

# 1. 分别为目标模型和草稿模型创建TensorRT-LLM引擎
#    这通常通过 tensorrt_llm.build 命令完成，这里假设引擎文件已存在
target_engine_dir = "path/to/target_model_engine"
draft_engine_dir = "path/to/draft_model_engine"

# 2. 创建采样配置 (SamplingConfig)
#    可以定义温度、Top-K、Top-P等参数
sampling_config = SamplingConfig(end_id=2, pad_id=2, num_beams=1)

# 3. 创建模型运行器 (ModelRunner)
#    在创建时同时传入目标模型和草稿模型的路径
#    TensorRT-LLM会自动识别并启用推测解码模式
runner = ModelRunner.from_dir(
    engine_dir=target_engine_dir,
    rank=tensorrt_llm.mpi_rank(),
    speculative_decoding_draft_engine_dir=draft_engine_dir
)

# 4. 准备输入数据
input_text = ["NVIDIA TensorRT-LLM is a powerful tool for"]
input_ids, input_lengths = runner.prepare_inputs([input_text])

# 5. 执行生成
#    TensorRT-LLM的运行时会自动处理推测解码的复杂逻辑
#    开发者无需关心内部的“草稿-验证”循环
output_ids = runner.generate(
    input_ids,
    sampling_config,
    max_new_tokens=100,
)

# 6. 处理输出结果
output_text = runner.decode(output_ids, input_lengths)
print(f"Input: {input_text[0]}")
print(f"Output: {output_text[0]}")

通过这种方式，TensorRT-LLM 将底层的复杂性完全封装，让开发者可以专注于业务逻辑，同时享受到推测解码带来的显著性能提升。在NVIDIA H100 GPU上，使用TensorRT-LLM实现的推测解码通常可以为Llama 2 70B等模型带来 2倍到3倍 的吞吐量提升和延迟降低。

性能对比与选择合适的草稿模型

推测解码的性能增益与草稿模型的选择密切相关。一个理想的草稿模型应该具备以下特点：

1. 速度快：模型参数量小，结构简单，推理延迟极低。
2. 一致性高：其输出的概率分布与目标模型尽可能相似。

下表展示了不同大小的草稿模型与Llama 2 70B目标模型搭配使用时，在NVIDIA H100 GPU上的性能表现（数据来源于NVIDIA官方博客）。

目标模型	草稿模型	加速比 (Speedup)
Llama 2 70B	Llama 2 7B	2.1x
Llama 2 70B	Llama 2 13B	1.8x
Llama 2 70B	TinyLlama 1.1B	2.4x

（注：加速比数据仅为示例，实际性能取决于具体任务、硬件和配置）

从表格中可以看出，并非草稿模型越大越好。虽然更大的草稿模型（如13B）可能与目标模型（70B）的输出更一致，但其自身的推理延迟也更高，反而可能拖累整体性能。相反，一个极小但高效的模型（如TinyLlama 1.1B）由于其极快的速度，即使接受率稍低，也可能带来最佳的整体加速效果。

因此，选择合适的草稿模型是一个需要在“速度”和“准确性”之间进行权衡的工程问题。开发者需要根据自己的应用场景和性能目标，通过实验来确定最佳的模型组合。

使用 NVIDIA Triton 推理服务器进行生产级部署

在实验室中验证了推测解码的性能后，下一步就是将其稳定、高效地部署到生产环境中。NVIDIA Triton 推理服务器 是专为大规模AI推理设计的标准化解决方案，它与TensorRT-LLM紧密集成，为部署推测解码模型提供了强大的支持。

Triton 提供了动态批处理（Dynamic Batching）、模型并发执行、多模型编排等企业级功能，这些功能与TensorRT-LLM的推测解码优化相结合，可以实现最大化的GPU利用率和系统吞吐量。

Triton 部署配置

使用Triton部署一个启用了推测解码的TensorRT-LLM模型，主要涉及到模型的目录结构和 config.pbtxt 文件的配置。

1. 模型仓库目录结构:

您需要将转换好的TensorRT-LLM目标模型引擎和草稿模型引擎放置在Triton的模型仓库中。

/model_repository
└── /my_llm_model
    ├── /1
    │   ├── target_model_engine.engine  # 目标模型引擎
    │   └── draft_model_engine.engine   # 草稿模型引擎
    └── config.pbtxt                    # Triton配置文件

2. config.pbtxt 配置文件:

在配置文件中，您需要指定使用 tensorrtllm 后端，并提供必要的参数来启用推测解码。

# config.pbtxt 示例
# 译者注：此为Triton配置文件示例，用于说明如何启用推测解码

name: "my_llm_model"
backend: "tensorrtllm"
max_batch_size: 64 # 可根据需要调整

# ... 其他输入输出配置 ...

parameters {
  key: "model_type"
  value: { string_value: "inflight_batching" }
}

parameters {
  key: "engine_dir"
  value: { string_value: "/model_repository/my_llm_model/1" } # 引擎文件所在目录
}

parameters {
  key: "speculative_decoding_mode"
  value: { string_value: "tensorrt_llm" } # 启用推测解码
}

parameters {
  key: "draft_model_name"
  value: { string_value: "draft_model_engine.engine" } # 指定草稿模型文件名
}

# 动态批处理配置
dynamic_batching {
  max_queue_delay_microseconds: 100
}

通过这样的配置，Triton服务器在加载模型时，会自动识别并启动推测解码模式。当客户端发送推理请求时，Triton和TensorRT-LLM后端会协同工作，高效地执行从草稿生成到目标验证的整个流程，并将最终结果返回给用户。

结论

推测解码为解决大型语言模型推理延迟问题提供了一个非常有效的途径。通过引入一个轻量级的草稿模型来辅助解码，它能够在不牺牲输出质量的前提下，显著提升推理速度和吞吐量。

然而，要将这一技术的潜力完全发挥出来，离不开强大的软件和硬件生态系统的支持。NVIDIA提供的 TensorRT-LLM 通过其深度优化的内核和简洁的API，极大地降低了实现推测解码的门槛。而 NVIDIA Triton 推理服务器 则为将其部署到复杂的生产环境提供了稳定、高效和可扩展的标准化路径。

作为开发者，理解推测解码的原理，并学会利用NVIDIA的全栈AI平台来实践和部署它，将是您在构建下一代低延迟、高性能AI应用过程中的一项核心竞争力。我们鼓励您亲自尝试使用TensorRT-LLM，探索不同模型组合带来的性能惊喜。