投机采样实战：大模型推理提速3倍核心技术解析

一、投机采样核心原理与适用场景

如果你正在被大模型推理延迟拖慢产品落地节奏，花10分钟搞懂投机采样，就能在不损失生成质量的前提下把推理吞吐量提升2-3倍，这笔投入绝对超值。

我们做投机采样的核心思路，本质是把原来大模型逐 token 自回归生成的串行流程，拆成「草稿生成+并行验证」两步。原来

二、主流投机采样方案对比

# 简化版投机采样推理伪代码，基于PyTorch实现
import torch

def speculative_decode(input_ids, draft_model, target_model, max_new_tokens=100):
    """
    输入示例：
    input_ids: torch.tensor([[1, 2, 3, 4, 5]])  # 编码后的输入token序列，shape为(1, seq_len)
    draft_model: 参数量为1B的轻量草稿模型，已和主模型对齐训练
    target_model: 参数量为13B的目标大模型
    max_new_tokens: 最大生成token数
    输出说明：
    返回生成的完整token序列，以及本次推理的加速比、接受的token总数
    """
    generated_tokens = input_ids.clone()
    total_accepted = 0
    total_steps = 0
    
    while generated_tokens.shape[1] < input_ids.shape[1] + max_new_tokens:
        total_steps += 1
        # 1. 草稿模型一次性生成k个候选token
        draft_logits = draft_model(generated_tokens)
        draft_tokens = torch.argmax(draft_logits[:, -1, :], dim=-1, keepdim=True)
        for _ in range(3): # 生成3个候选token
            draft_logits = draft_model(torch.cat([generated_tokens, draft_tokens], dim=-1))
            next_token = torch.argmax(draft_logits[:, -1, :], dim=-1, keepdim=True)
            draft_tokens = torch.cat([draft_tokens, next_token], dim=-1)
        
        # 2. 主模型并行验证所有候选token
        candidate_input = torch.cat([generated_tokens, draft_tokens], dim=-1)
        target_logits = target_model(candidate_input)
        # 逐个比对候选token和主模型输出的匹配度
        accept_len = 0
        for i in range(draft_tokens.shape[1]):
            target_token = torch.argmax(target_logits[:, generated_tokens.shape[1]-1 + i, :], dim=-1)
            if target_token == draft_tokens[:, i]:

三、实战环境与依赖准备

我们首先要搭建支持投机采样的推理框架环境，目前主流的生产级框架vLLM和TensorRT-LLM都已经原生内置了投机采样的完整实现，不需要我们手动魔改核心推理逻辑。vLLM从0.3.0版本开始就把投机采样作为默认支持的功能，我们只需要在pip安装时指定对应版本，启动推理服务时添加 speculative 相关参数即可开启。如果是用TensorRT-LLM的话，官方已经预编译好了支持投机采样的wheel包，我们直接安装对应CUDA版本的包就行，不需要自己编译CUDA内核，大幅降低了环境配置的门槛。

接下来我们需要验证草稿小模型和主大模型的适配性，这是投机采样能生效的核心前提。投机采样的逻辑是先用小模型快速生成多个候选token，再用大模型一次验证这些token是否合法，因此两个模型的词表必须完全一致，架构也要保持同系列，比如我们用Llama 3 70B作为主模型，就必须选择Llama 3架构的8B甚至1B版本作为草稿模型，不能跨架构混用。我们还需要提前用同一组测试prompt验证两个模型的输出重合度，要求前5个候选token的重合度不低于75%，如果重合度太低的话，大模型会频繁拒绝小模型的输出，反而会导致推理速度比原生慢速采样还低。

我们需要提前设定好测试数据集和基准指标，避免后续验证结果出现偏差。测试数据集要覆盖不同难度、不同长度的任务，比如包含短文本问答的MT-Bench、代码生成任务HumanEval、长文本摘要任务CNN/DailyMail，不能只用单一类型的简单数据集。基准指标除了常规的吞吐量、端到端延迟之外，必须重点记录token接受率，这个指标直接决定了投机采样的提速效果，接受率每提升10%，整体吞吐量就能提升15%左右。我们还要设置明确的达标阈值：在相同batch size和输入长度下，开启投机采样后的吞吐量必须比原生采样高2倍以上，否则就是配置或模型适配有问题。

# vLLM开启投机采样的示例代码
from vllm import LLM, SamplingParams

# 初始化推理引擎，指定主模型和草稿模型
llm = LLM(
    model="meta-llama/Llama-3-70B-Instruct",
    speculative_model="meta-llama/Llama-3-8B-Instruct",
    num_speculative_tokens=5,  # 每次小模型生成的候选token数量
    tensor_parallel_size=2  # 主模型的张量并行度
)

# 定义采样参数
sampling_params = SamplingParams(temperature=0.7, top_p=0.95, max_tokens=256)

# 输入测试prompt
prompts = ["解释一下量子纠缠的基本原理，用通俗易懂的语言"]
outputs = llm.generate(prompts, sampling_params)

# 输出说明：每个输出结果会附带acceptance_rate字段，代表大模型接受小模型候选token的比例
for output in outputs:
    print(f"生成内容：{output.outputs[0].text}")
    print(f"Token接受率：{output.metrics.acceptance_rate:.2f}")
    # 当接受率≥0.7时，整体吞吐量可提升2-3倍，达到预期提速目标

对于大部分开发者来说，我们优先推荐使用vLLM内置的投机采样方案，这个方案已经经过了大量生产环境的验证，稳定性有保障，只需要把草稿模型换成和主模型同架构的小版本，调整num_speculative_tokens参数到5-10之间，就能快速验证提速效果。如果后续对性能有极致要求，再切换到TensorRT-LLM的自定义方案，我们可以先拿Llama 3 70B搭配Llama 3 8B跑一遍MT-Bench基准测试，实际验证token接受率和吞吐量的提升情况，确认配置无误后再接入线上

四、推理加速全流程实战

# 输入示例：基于vLLM实现投机采样推理
from vllm import LLM, SamplingParams

# 初始化主模型和辅助小模型
llm = LLM(
    model="meta-llama/Llama-3-7b-instruct",
    speculative_model="meta-llama/Llama-3-1b-instruct",
    num_speculative_tokens=5,  # 候选长度K设为5
    speculative_draft_tensor_parallel_size=1
)

# 配置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=512
)

# 执行推理
prompts = ["请用通俗易懂的话解释什么是投机采样"]
outputs = llm.generate(prompts, sampling_params)

# 输出说明：运行后单张A100的生成速度可达80+ token/s，是原生Llama-3-7B推理速度的3倍左右，输出内容与原生7B模型无显著质量差异