跳到主要内容

AI 流水线中的投机执行:通过预判未来大幅降低延迟

· 阅读需 13 分钟
Tian Pan
Software Engineer

大多数 LLM 流水线在无意中都表现出极其低效的串行化。一个智能体调用天气 API,等待 300 ms,调用日历 API,再等 300 ms,调用交通 API,再次等待——最后才综合出答案。如果这三个调用是并行运行的,那么原本 900 ms 的总延迟本可以缩短到 300 ms。没有人刻意将系统设计成串行;这只是在编写一个又一个异步调用时自然形成的结果。

预测执行(Speculative execution)是一系列技术的统称,通过在确定需要工作之前提前执行工作来降低感知延迟——运行并行的假设、预取可能的下一步、并同时生成多个候选输出。这些技术直接借鉴了 CPU 设计,自 20 世纪 90 年代以来,处理器就已经开始预测执行未来的指令。应用到 AI 流水线中,同样的本能——押注可能的结果、取消失败的部分、接受偶尔的浪费——可以产生显著的提速。但如果你不仔细考虑应用时机,协调开销也可能会抵消掉所有的收益。

预测执行的三个层级

AI 系统中的预测执行在三个不同的层级上运行,每个层级都有不同的权衡。

Token 级预测 (Token-level speculation) 是业界通常所说的“投机采样 (speculative decoding)”。一个小型、快速的草案模型提前生成几个候选 token,然后一个大型目标模型在单次前向传递中验证所有这些 token。由于目标模型可以利用其现有的并行架构同时检查多个 token,因此有效的草案可以低成本地被接受。如果 Google 的基准测试具有代表性——在接受率超过 0.6 的情况下,三个投机 token 可以带来 2–3.4 倍的加速——这是推理服务中延迟优化杠杆最高的手段之一。目前它已成为 vLLM、SGLang 和 TensorRT-LLM 的生产标准。

工具调用级预测 (Tool-call level speculation) 是关于智能体发出的外部调用进行预取或并行化。当智能体的下几个工具调用相互独立时,你可以同时发出所有调用。一个从三个数据源获取信息的搜索智能体不需要源 A 的结果就可以开始获取源 B。天真的实现默认是串行的;正确的实现则会检测独立性并进行分发。

响应级预测 (Response-level speculation) 意味着同时生成多个候选答案——束搜索 (beam search) 就是经典范例。束搜索不是在每一步贪婪地选择概率最高的 token,而是在并行维护 k 个候选序列。这种方法开销更大,但对于初始 token 选择会显著限制后续响应质量的任务,它可以产生更高质量的输出。

这三个层级相互作用。一个在 Token 层级进行投机采样的生产系统也可以在智能体层级并行化工具调用。但每个层级都有自己的失效模式和成本概况,这就是为什么值得将它们分开思考。

投机采样:运作机制与失效原因

投机采样的核心洞察是,大语言模型的瓶颈通常是显存带宽,而不是原始算力。生成一个 token 需要从 GPU 显存中加载整个模型的权重。生成五个 token 则需要加载一次权重并进行五次小型计算传递——以大致相同的显存传输成本获得五倍的输出。草案模型通过一次提出多个 token 来利用这一点;验证器在单个批处理传递中确认或拒绝它们,而无需额外的权重传输。

数学上的保证是明确的:当草案模型的建议与目标模型的预测一致时,它们会被原样接受。当它们偏离时,你会回退到目标模型的输出。由此产生的分布在证明上与单独运行目标模型完全相同。你在不牺牲质量的情况下获得了速度提升。

在实践中破坏这一点的是接受率敏感性。接受率——验证器接受草案 token 的比例——是核心指标。在接受率高于 0.6 时,你会看到明显的提速。低于这个水平,草案模型就是在做无用功。接受率取决于草案模型与目标模型的对齐程度、量化差距以及工作负载特性。在与目标模型不同的数据分布上训练的草案模型接受率会很低。过度量化目标模型也会降低接受率,因为量化模型的 logits 会偏离草案模型训练时的预测。

另一个关键的失效模式是批处理大小 (batch size)。投机采样是一种受限于内存的优化——它在 GPU 利用率不足时有所帮助。在高批处理量(通常是 32 个以上)时,GPU 已经处于计算受限 (compute-bound) 状态。增加草案模型的传递会变成纯粹的开销。vLLM 文档明确警告说:“投机采样在现实系统中表现出脆弱且高度波动的性能”。正确的方法是在低并发时启用它,并随着请求率超过曲线拐点而自动禁用它。

并行工具调用:更易实现的制胜法宝

对于大多数应用开发者来说,投机解码(Speculative Decoding)发生在推理基础设施内部,并不是你可以直接控制的。而并行工具调用则是更具即时操作性的技术。

假设一个智能体(Agent)需要天气数据、日历空闲情况和交通状况来回答调度问题。串行执行:300ms + 300ms + 300ms = 900ms。并行执行:max(300ms, 300ms, 300ms) = 300ms。其中的逻辑很简单。实现障碍通常在于认知——开发者默认采用串行方式,因为这在代码中读起来更自然。

LLMCompiler 框架(ICML 2024)通过将智能体规划视为编译过程,使这一过程正式化。规划器会输出一个工具调用的依赖图;执行器同时调度所有没有前置依赖的调用,然后在调度下一层级之前汇总它们的结果。这实现了规划阶段(需要发生什么)与执行阶段(以什么顺序发生)的分离,这正是正确的抽象方式。

关键的前提条件是独立性。只有当工具调用不共享可变状态,且彼此都不需要对方的输出时,它们才是可以并行的。这听起来显而易见,但在实践中很容易被忽略:如果工具 A 的结果会改变你传递给工具 B 的参数,你就不能并行化它们。依赖图必须准确。对独立性的错误假设会导致错误的行为,而不仅仅是徒劳的计算。

尾部延迟(Tail Latency)问题也值得一提:并行执行意味着你需要等待最慢的那个调用。如果三个调用分别耗时 200ms、300ms 和 800ms,那么并行的延迟就是 800ms——这甚至比你先排列快速调用的优化后的串行执行中位数时间还要糟糕。竞速模式(Racing Patterns)在这里会有所帮助:发起多个等效的调用,取第一个结果,并取消其余的。当你拥有多个返回等效结果的服务提供商时,这种方法效果很好。

预取下一轮工具调用

预取(Pre-fetching)将并行调用的理念扩展到了智能体的推理步骤中。当智能体正在处理当前工具调用的结果时,你是否可以在智能体明确决定调用下一个工具之前,就启动“下一个”工具调用?

PASTE 框架(Pattern-Aware Speculative Tool Execution)通过学习智能体工具日志中的常见序列来实现这一点。如果调用天气 API 的智能体在 80% 的情况下接着会调用日历 API,那么你可以在看到天气调用完成时,就投机性地启动日历调用。如果智能体确认了这一预测,结果就已经在传输中了。如果智能体转向了不同的方向,你就取消这个投机调用。

这比并行调用更复杂,因为它需要一个针对智能体行为的预测模型,而不仅仅是静态依赖分析。PASTE 中的有界浪费模型(Bounded Waste Model)防止了失控的投机:如果接受率低于阈值,投机深度会自动降低。这种自我调节机制至关重要,因为只有当你的预测足够准确时,预取才是有利可图的。一个在每一步都预取错误工具调用的拙劣预测模型,会增加 API 成本而无法降低延迟。

当协调开销淹没收益时

投机技术的失败模式与它的收益同样重要。投机技术有三种常见方式会让情况变得更糟:

伪装成并行任务的串行任务。 并行工具调用示例中的加速效果假设所有调用都是独立的。需要循序渐进推理的任务——每一步都依赖于前一步——无法进行有效的并行化。将并行性强加于本质上串行的推理会破坏正确性。在协调不力的多智能体系统中,错误放大效应可能非常显著:关于多智能体失败模式的研究发现,与正确排序的执行相比,无结构的并行网络会将错误放大 17 倍。

失败投机的成本。 你运行的每一个投机分支都会产生实际成本。Token 级的投机会生成草稿 Token,如果接受失败,这些 Token 将被丢弃。预取的工具调用会返回结果,如果智能体选择了不同的路径,这些结果将被忽略。在分支因子为 2、深度为 4 的基于树的投机解码设置中,会生成一个包含 15 个候选 Token 的树;如果只有一条包含 4 个 Token 的路径被接受,则浪费了 11 个 Token。在大规模应用中,除非仔细限制,否则这种浪费可能会超过节省的开销。

速率限制耗尽。 并行工具调用和预取会同时成倍增加你的 API 请求速率。一个以 300ms 间隔进行 4 次串行调用的智能体,每秒发出大约 3–4 个请求。并行运行这些调用意味着同时发出 4 个请求。在此基础上再加上预取的投机调用,可能会让你迅速触及速率限制(Rate Limit)红线,从而创造一个新的瓶颈,抵消掉延迟的节省。

决策框架

并非所有的延迟都值得通过投机来优化。以下是一个实用的启发式方法:

针对 Token 级投机(即投机解码):当你控制或影响推理基础设施、并发量较低、输出较长(100+ Token),且你使用的模型有良好的草稿模型支持时应用。当 Batch Size 超过 32 或请求速率导致推理框架显示吞吐量下降时,自动禁用它。

针对 并行工具调用:当你能够绘制工具调用的依赖图,且某些调用没有前置依赖时应用。这几乎总是值得做的——认知开销低,在任何异步框架中实现都很直接,且收益与独立调用的数量成正比。

针对 预取:只有当你拥有历史数据来训练预测模型、能够监测接受率,并且在预测效果不佳时能够自动回退到非投机执行时才应用。这是一种高投入的优化,在特定场景下(如具有可预测工具序列的高流量智能体)会有显著收益。

针对 响应级并行(如束搜索 Beam Search、生成多个候选结果):当输出质量比成本更重要,而不是以降低延迟为主要目标时应用。束搜索会随着束宽线性增加计算成本。质量的瓶颈通常在成本收益达到平衡之前就已经出现了。

协作成本

这一族群中的每一项技术都增加了顺序代码所不具备的协作复杂性。你需要处理失败分支的取消工作。你需要透传并行调用的错误,且不能丢失上下文。你需要确保并行变更(mutations)不会产生竞争条件。随着流量模式的变化,你需要监控采纳率并调整推测参数。

那些真正从推测中获益的团队,往往是首先在可观测性上投入的团队。在启用推测前后,他们能够衡量实际的采纳率、实际的浪费率以及实际的尾部延迟分布。如果没有这些监控埋点,推测就只是一个带有性能噱头的额外开销。

CPU 的类比在这里也很有启发意义:CPU 中的推测执行也引入了新的失效模式(如 Spectre 和 Meltdown)。性能提升是实打实的,但安全隐患也同样存在。AI 流水线推测也有类似的情况:一个推测性地执行写入或变更(而不仅仅是读取)的预取系统,可能会因为本该被丢弃的推测分支而产生真实的副作用。工具调用预取对于幂等的读取操作通常是安全的,但对于任何改变状态的操作都是危险的。

趋势走向

目前的趋势是让推理服务基础设施透明地处理推测,这样应用开发者就不必为此操心。vLLM 的推测解码、LLMCompiler 的并行分发以及 PASTE 式的模式感知调度,都在朝着根据工作负载特征自动做出正确选择的方向迈进。

不会变得自动化的是应用层关于哪些工具调用实际上是独立的决策。这需要理解你的业务领域,并正确构建依赖图。这是仍然需要工程判断的部分,也是大多数团队获得最大剩余收益的地方。

最简单的第一步:审视你当前智能体的工具调用,在纸上画出依赖图,并统计有多少调用没有前置节点。这个数字告诉了你理论上的并行度预算。如果该数字大于 1,那么你可能正在浪费降低延迟的机会。

References:Let's stay in touch and Follow me for more thoughts and updates