bug404

bug404

cv、深度学习、模型部署

05-TensorRT-LLM 推测解码与自动部署：让推理再快一步

3 minute read

Published:

背景与目标

大语言模型的自回归生成瓶颈在于串行解码：每个 token 都必须等待前一个 token 完成目标模型的前向计算。推测解码（speculative decoding）将这一过程从严格的“一次一个”改为“一次一批”：用一个轻量级的 draft 模型快速生成多个候选 token，目标模型再以并行方式一次性验证这组候选，接受匹配的前缀并确保分布一致性。相对于基线的自回归，这一范式能成倍提升解码吞吐，但同时也引入了 draft 模型的前向开销与复杂的状态管理。TensorRT-LLM 将 Eagle 系列的 draft 模型与自动部署工具链结合，通过定制 CUDA 内核、动态调度和 torch.compile 深度优化，把推测解码的硅基潜力尽可能释放出来。本文深入解析这套管线在 GPU 上是如何工作的——从 Eagle 树状草稿生成到与 torch.compile 协作的 kernel 自动选择。

核心设计与实现

Eagle-2 多层级树状草稿生成

推测解码的原始形态（如 n-gram 或单个 draft token）接受率低，因为 draft 模型给出的候选过于单一。Eagle 系列利用目标模型自身的隐藏状态来训练 draft 模型，在解码时通过多层级分类头生成一棵候选树，每个叶子都对应一条潜在的前缀。Tree attention 验证机制允许目标模型在单次前向中沿着树结构验证所有叶子，从而同时接受多个正确的分支。

在 TensorRT-LLM 的实现中，EagleDecodeDraftTokensPlugin 将这一过程抽象为一个 CUDA 插件，通过固定索引的输入输出 tensor 在多个层级间传递 logits、路径历史和分数。

// cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp:360-390
// First top-K sampling
TopKSamplingKernelParams<T> params{};
params.logProbsPtrs = logitsPtrs;
params.outputIdsPtrs = firstTopKOutputIdsPtrs;
params.workspace = workspaceSampling;
params.maxTopK = maxTopK;
params.topKs = topKs;
params.batchSize = numInputLogits;
params.returnAllSelectedTokens = true;
params.skipDecode = skipDecode;
params.outputLogProbs = firstTopKOutputLogProbs;
invokeBatchTopKSampling(params, stream);

这段代码对每一层的 logits 做 top-K 采样，参数 maxTopK 控制扩展宽度，skipDecode 掩码标记填充位置（无效 logits），使得 softmax 和采样 kernel 在遇到这些位置时直接跳过计算。采样结果同时包含 token ID 和 log 概率，供后续层级进行分数累积。

跨层级的关键是动态前端索引重建：第一层采样得到 topK 个候选，第二层将这些候选作为父节点再次进行 top-K 采样（每个父节点扩展 topK 个子节点，共 topK^2 个候选），通过第二次 top-K 将候选数压缩回至多 maxTopK 个。之后每个 token 携带指向前驱的索引，使得最终的一次全局 top-K 能重建出树路径，即一组 [batch, max_decoding_tokens, max_path_len] 形状的 paths 张量。所有计算都在单个 CUDA stream 上使用批量 kernel 完成，避免了多次启动和中间 tensor 的全局内存写回。

目标模型验证：树注意力与批量并行

草稿 token 生成后，目标模型按“生成-验证”范式进行检查。由于草稿是一棵树（而非线性序列），验证时需要对树中每个 token 与目标模型的 logits 比较。TensorRT-LLM 利用 tree attention 将不同路径的 token 在 batch 维度上拼接，并构建对应的注意力掩码以阻止跨路径干扰。这样目标模型一次前向即可输出每个叶子 token 的 logits，再通过并行对比（rejection sampling）决定接受哪些 token。因为目标模型的计算密度远高于 draft 模型，这种“一次多检”的策略能有效摊销开销。

torch.compile 与自动部署的协作

推测解码管线若直接使用 eager 模式 PyTorch，会遇到两处瓶颈：一是 draft 模型和目标模型频繁的 kernel launch 与 host-device 同步，二是动态形状（如变长草稿序列）导致 CUDA graph 难以静态编译。TensorRT-LLM 的 auto_deploy 模块通过 MLIR 子图替换、自定义算子调度以及 torch.compile 结合 CUDA graph，将这些路径进行编译时重写。

下面这段代码展示了自动部署在融合子图时的依赖感知插入逻辑：当 MLIR block 中的子图被替换为单个 AdOpaque 融合算子时，编译器会扫描所有输入的生产者位置，确保融合算子插入点在所有生产者之后，避免数据依赖被打破。

# tensorrt_llm/_torch/auto_deploy/mlir/fusion/subgraph_replace.py:73-89
block = subgraph.ops[0].parent
block_ops_order = {op: i for i, op in enumerate(block.ops)}
first_sg_pos = block_ops_order[subgraph.ops[0]]
latest_input_pos = -1
for inp in subgraph.inputs:
    producer = inp.owner
    if producer in block_ops_order:
        latest_input_pos = max(latest_input_pos, block_ops_order[producer])
if latest_input_pos >= first_sg_pos:
    anchor_op = list(block.ops)[latest_input_pos]
    block.insert_op_after(fused_op, anchor_op)
else:
    block.insert_op_before(fused_op, subgraph.ops[0])

在 torch.compile 后端，TorchOptCompiler 通过设置 torch._dynamo.config.recompile_limit 来匹配 CUDA graph 的批量大小个数，防止因 Dynamo 缓存 key 不够而发生重编译回退到 eager。

# tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py:16-22
torch._dynamo.config.recompile_limit = max(
    len(self.cuda_graph_batch_sizes), torch._dynamo.config.recompile_limit
)

这样一来，无论 draft/target 模型在不同 batch size 间切换，动态编译器都能复用编译缓存，维持端到端的静默执行。

优化技术细节

重叠调度隐藏草稿延迟

草稿模型的前向计算与目标模型的验证之间本质是串行的，但可以在前一个草稿 batch 计算的同时，在 CPU 侧处理上一批次的结果（更新 draft tokens、释放资源），利用双缓冲将计算与调度重叠。

# tensorrt_llm/_torch/speculative/model_drafter.py (generate_draft_tokens_with_overlap)
# Initial forward pass
outputs = self.forward_draft_model(draft_batch, resource_manager, ...)
# Process previous draft results after current forward pass
self._process_previous_draft_results(resource_manager)
# Update target inputs with first iteration results
self._update_draft_tokens_for_target_inputs(...)

这里通过 previous_draft_batch 保存上一批请求，当 GPU 忙于执行 forward_draft_model 时，CPU 可以异步加载下一个批次的 token 和 KV cache 信息，从而将软件延迟隐藏在计算流下。

自适应草稿长度

批次大小直接影响推测加速比：并发请求少时，长草稿序列能提升单请求的有效接受率；并发高时，长序列增加验证开销且 KV cache 压力大，反而降低整体吞吐。Drafter 类维护一个 draft_len_schedule 字典（如 {1: 32, 4: 16, 8: 8}），根据当前 batch size 通过二分查找动态决定草稿长度。

# tensorrt_llm/_torch/speculative/drafter.py (get_draft_len_for_batch_size)
thresholds = list(self.draft_len_schedule.keys())
idx = bisect_right(thresholds, batch_size)
if idx == 0:
    return 0
threshold = thresholds[idx - 1]
return self.draft_len_schedule[threshold]

这种粗粒度的启发式避免了复杂的在线测量，同时能与 KV cache 块管理协调：更短的草稿长度意味着更少的 KV cache 预留，调度器可容纳更多请求。

自动部署中的混合 token 预算

在 MoE 模型中，推理解码时每层需要对激活进行 all-to-all 通信以路由 token 到不同专家。为了支持 CUDA graph 捕获，通常需要固定接收缓冲区大小（max_num_tokens），这会导致图内 kernel 为静态形状浪费计算。trtllm_moe 实现了一个混合策略：当处于 CUDA graph 捕获或 warm-up 阶段时，直接使用实际 token 数量 local_num_tokens 作为预算，前提是 budget * ep_size * 4 <= max_num_tokens，保证 GEMM 的 M 维度处于小 M kernel 的最佳区间。否则回退到静态大预算，避免动态形状破坏图完整性。

# tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py:47-73
if torch.cuda.is_current_stream_capturing() or cuda_graph_state.in_warm_up():
    budget = local_num_tokens
    if budget > 0 and budget * ep_size * 4 <= max_num_tokens:
        return budget
return max_num_tokens

这种混合预算在解码阶段消除了每层的 .item() 主机同步（原本用于读取动态 token 数），使整个 decode step 可以完全被 CUDA graph 捕获而无 host 干预。

依赖感知子图融合与路由 GEMM 调度

自动部署的融合引擎进一步将元素级操作（如激活函数、add）与 matmul 合并为单次 AdOpaque 调用，并通过上述插入策略确保数据流合法。在路由器计算时，根据 token 数量、SM 架构和数据类型，动态选择 torch.ops.trtllm.tinygemm2（小 M 的单 kernel GEMM）或标准的 F.linear（cuBLAS）。此调度检查只在调用点进行，避免了额外的中间表示。

性能与权衡

Eagle 树状生成通过增加 draft 模型的计算量换取了更高的接受率，其性能增益在 batch size 较小时尤为显著（单请求延迟可降低 2-3×）。但其内存包袱不容忽视：需要存储所有层级的分数、前驱索引和 candidate tokens，底层 cudaMalloc 的次数和峰值显存因此上升。重叠调度和自适应草稿长度虽然平衡了吞吐与延迟，却增加了调度器和 drafter 的状态复杂度，使一致性调试变难。

自动部署结合 torch.compile + CUDA graph 极大减少了 kernel launch 和主机同步开销，尤其在小 batch decode 场景下，通常可提升 30-50% 的吞吐。代价是动态性受限：一旦进入 graph 捕获，草稿长度、路由 token 数等必须遵循预定形状，限制了运行时自适应能力。混合 token 预算提供了一个折中，但只有当实际 token 数落在小 M 范围内才生效，否则仍回退到静态 padding。

相关论文

  • EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees. Yuhui Li, et al. arXiv:2406.16858. 首次提出多层级特征缓存和动态树扩展方法。
  • Speculative Decoding for Fast Inference from Large Language Models. Yaniv Leviathan, et al. ICML 2023. 奠定推测解码的理论基础：rejection sampling 保证分布一致性。
  • TensorRT-LLM: A TensorRT-based Toolkit for Optimizing Large Language Models on NVIDIA GPUs. NVIDIA. 提供本文所述自动部署与自定义内核的背景。

总结

TensorRT-LLM 将推测解码落地为一组精心编排的 CUDA 内核与编译优化。Eagle draft 模型利用目标模型的隐藏状态动态构建多层级候选树，通过批量 top-K 和 tree attention 以 GPU 原生并行实现高接受率的验证。而 auto_deploy 工具链则利用 MLIR 融合、torch.compile 与 CUDA graph，将模型定义转化为静态图并消除动态同步，使管线的每一段都在硅基上安静流淌。两者合力，让推理 token 生成从“一步一回头”变成了“一次多步走”的高效流水线。