SuffixDecoding: Extreme Speculative Decoding for Emerging AI Applications

发表时间: 2024-11 · arXiv:2411.04975

原文: https://arxiv.org/abs/2411.04975

  • 文章标题:SuffixDecoding: 面向新兴 AI 应用的极端推测解码
  • 作者/机构:Gabriele Oliaro§,†, Zhihao Jia†, Daniel Campos§, Aurick Qiao§ (§Snowflake AI Research, †Carnegie Mellon University)
速读

一句话结论 本文提出了一种专为重复性 AI Agent 工作负载设计的无模型推测解码方法 SuffixDecoding,通过基于后缀树的模式匹配和自适应推测长度机制,在零 GPU 额外开销的前提下,于 Agent 任务中实现了高达 5.3 倍的推理加速。

要解决什么问题 大型语言模型在标准自回归解码时,每个 token 的生成都依赖前一个输出,这种顺序执行机制限制了吞吐量并造成并行硬件利用率低下。推测解码通过引入轻量级草稿模型快速生成多个候选 token,再交由主模型并行验证来缓解延迟。然而,在自动化编码助手、多 Agent 工作流等新兴应用中,系统会频繁发出包含长且可预测重复序列的推理请求。现有的推测解码方法(无论是基于模型还是无模型)在此类场景下卡在两个关键瓶颈:第一,草稿模型本身仍需消耗 GPU 计算资源并带来编排开销,生成草稿的速度不够快,无法以极低成本最大化利用 Agent 任务中极长的可预测序列;第二,现有方法缺乏自适应性,无法根据当前上下文的接受概率动态调整推测长度,导致在预测置信度低时生成过多无效草稿 token,使主模型的验证计算被大量浪费,验证环节反而成为新的性能瓶颈。

怎么做的 SuffixDecoding 的核心思路是摒弃神经网络草稿模型,转而使用高效的后缀树数据结构缓存和匹配历史 token 序列,实现极速且零 GPU 开销的无模型推测。该方法由四个关键部件构成。首先是双后缀树缓存机制,系统在 CPU 上同时维护一棵记录所有历史请求输出的全局后缀树,以及一棵记录当前请求提示词和已生成内容的独立请求树,节点代表 token,路径代表子序列。其次是基于模式匹配的推测树构建,截取最新生成的长度为 $p$ 的后缀作为模式序列,在两棵树中查找匹配节点并贪心扩展出推测树。扩展时利用历史频率统计对节点评分,节点 N 的得分 $D(N)$ 基于其在语料库中出现的次数计算,代表该节点最终被主模型验证接受的概率,系统优先将高分节点加入推测树。第三是自适应推测长度机制,为避免低置信度下盲目推测导致验证浪费,系统根据模式匹配长度 $p$ 动态限制推测树的最大节点数,定义为 $\text{MAX\_SPEC}(p) = \alpha p$,其中 $\alpha$ 为最大推测因子。匹配到的历史模式越长,推测序列就越长,反之则保守推测。最后是混合推测回退机制,系统计算推测树的期望接受 token 数 $\text{SCORE}(T_{spec}) = \sum_{N \in T_{spec}} D(N)$,若得分低于设定阈值,说明当前处于缺乏重复模式的开放式生成场景,系统会回退到 EAGLE-3 等基于模型的推测解码方法,兼顾 Agent 任务的极致加速与开放任务的通用性。

效果如何 实验在配备 8 块 H100 GPU 和 2TB 内存的硬件上运行,模型选用 Llama-3.1-8B-Instruct 和微调后的 32B 编码模型。对比基线包括代表基于模型路线的 EAGLE-2 和 EAGLE-3,以及代表无模型路线的 Prompt-Lookup Decoding 和 Token Recycling。在 AgenticSQL 多 Agent 工作流任务中,SuffixDecoding 取得 5.3 倍平均加速,比 EAGLE-2/3 快 2.8 倍,比 Token Recycling 快 1.9 倍,平均每步接受 6.3 个 token。在 SWE-Bench 真实代码任务中,EAGLE-2/3 因序列长度限制运行失败,SuffixDecoding 实现 2.5 倍加速,比次优的 Prompt-Lookup Decoding 快 1.7 倍,且作为无损方法保持了原模型 37.2% 的任务解决率。在 Spec-Bench 开放式非 Agent 任务中,单独的 SuffixDecoding 表现不及基线,但结合 EAGLE-3 的混合方法实现 2.5 倍加速,反超所有独立基线。该方法的代价是全局后缀树需占用大量 CPU 内存,长期运行需引入淘汰策略;此外它强依赖输出分布的稳定性,当输入任务分布发生剧烈突变时,加速效果会短暂下降,需等待系统在线收集新分布输出并更新后缀树后才能恢复。

A1 主要贡献

核心问题:大型语言模型(LLMs)是 AI Agent 应用(如自动化编码助手、多 Agent 工作流、检索系统)的基础。与基础聊天机器人不同,这些工作负载会发出重复且可预测的推理请求。例如,多 Agent 系统会重复执行相似任务,而推理循环会重新生成相似的 token 序列。尽管存在这种可预测的重复性,但现有方法未能充分利用这些重复模式,导致延迟成为一个瓶颈。推测解码是缓解推理延迟的常用策略,但现有方法无法满足 Agent 驱动应用中高效处理长重复序列的两个关键要求:1)需要以最小开销快速生成草稿 token,以最大化利用长推测长度;2)必须具备自适应性,只在接受可能性高时生成更多草稿 token,在可能性低时生成较少 token,以避免验证成为瓶颈。

研究目标与创新点:为了解决上述限制,本文引入了 SuffixDecoding,这是一种专为重复性、Agent 驱动工作负载设计的无模型推测解码方法。其核心创新点如下:
1. 基于后缀树的高效缓存:SuffixDecoding 使用高效的后缀树来缓存来自提示和先前输出的长 token 序列。每个节点代表一个 token,路径编码了先前观察到的子序列,从而能够基于先前出现的情况快速进行模式匹配以识别后续序列。草稿 token 的生成速度极快(每个 token 约 20 微秒),且无 GPU 开销。
2. 自适应推测长度:在每个推理步骤中,SuffixDecoding 根据模式匹配的长度自适应地限制其草稿 token 的数量。较长的模式匹配使其能够自信地推测更长的 token 序列,从而最大化其在 Agent 工作负载上的效果;而较短的模式匹配则触发保守的推测以避免计算浪费。
3. 基于频率的候选择优:利用后缀树中捕获的基于频率的统计数据对最佳推测候选进行评分和选择,从而构建一个更小且可能性更高的推测树。
4. 混合推测能力:SuffixDecoding 可以与现有的基于模型的推测解码方法无缝集成。这种灵活性使得可以采用混合方法,对重复、可预测的 Agent 工作负载利用基于后缀树的推测,同时对开放式对话任务利用基于模型的方法的优势,从而实现两全其美。

主要成果概述:在包括 SWE-Bench 和 Text-to-SQL 在内的 Agent 基准测试上进行的评估表明,SuffixDecoding 实现了高达 5.3 倍的加速,优于最先进的方法——比 EAGLE-2/3 等基于模型的方法快 2.8 倍,比 Token Recycling 等无模型方法快 1.9 倍。

图1:SuffixDecoding 算法概述。两个后缀树分别跟踪正在进行的推理(左上)和以前的输出(左下)。SuffixDecoding 使用这些树根据最近生成的 token 找到匹配的模式。它通过选择最可能的延续来构建一个推测树(中间),并根据频率统计数据对其进行评分。最后,最佳候选由 LLM 在单个前向传递中验证(右侧),接受的 token(以绿色显示)被添加到输出中,并用于下一轮推测。
图1:SuffixDecoding 算法概述。两个后缀树分别跟踪正在进行的推理(左上)和以前的输出(左下)。SuffixDecoding 使用这些树根据最近生成的 token 找到匹配的模式。它通过选择最可能的延续来构建一个推测树(中间),并根据频率统计数据对其进行评分。最后,最佳候选由 LLM 在单个前向传递中验证(右侧),接受的 token(以绿色显示)被添加到输出中,并用于下一轮推测。

A3 背景知识与相关工作

LLM 推理。LLM 推理包括两个阶段:给定一个提示 $x_{prompt} = (x_1, x_2, . . . , x_m)$,LLM 首先并行处理提示(预填充),然后顺序生成新的 token(解码),每个 token $x_{t>m}$ 都以先前生成的 token为条件:

images/ec4ab1ea5a740d59f07e38973707cea141d93378f7184c716592deeb11938fa4.jpg
images/ec4ab1ea5a740d59f07e38973707cea141d93378f7184c716592deeb11938fa4.jpg

在贪心采样中,会迭代选择概率最高的 token,直到满足停止条件。由于每个 token 都依赖于前面的输出,生成过程本质上是顺序的,每个 token 都需要一个单独的前向传播,这限制了吞吐量并导致并行硬件加速器利用不足。

推测解码。推测解码【14,Fast inference from transformers via speculative decoding,2023,ICML】通过使用一个轻量级模型快速生成多个候选 token,然后由主 LLM 并行验证这些 token,从而加速推理。基本方法有两个核心步骤:1. 推测:一个较小的“草稿”模型根据现有的 token 前缀 $x_{<t}$ 快速产生推测的 token $x_{spec} = (x_{t+1}, . . . , x_{t+n})$。2. <strong>验证:LLM 并行验证草稿 token,接受直到第一个不匹配的 token,并丢弃其余部分。这将计算从顺序生成转移到并行验证。然而,草稿模型仍然需要计算资源并增加了编排复杂性。

基于模型的推测解码方法。基于模型的方法包括 Medusa【1,Medusa: Simple llm inference acceleration framework with multiple decoding heads,2024,arXiv】,SpecInfer【23,Specinfer: Accelerating large language model serving with tree-based speculative inference and verification,2024,ASPLOS】(它引入了基于树的推测)、多 token 预测【9,Better & faster large language models via multi-token prediction,2024,ICML】和分块并行解码【27,Blockwise parallel decoding for deep autoregressive models,2018,NIPS;12,Accelerating blockwise parallel language models with draft refinement,2024,NeurIPS】。

无模型的推测解码方法。最近的无模型方法包括 Prompt Lookup Decoding (PLD)【26,Prompt lookup decoding,2023,GitHub】,Token Recycling【20,Turning trash into treasure: Accelerating inference of large language models with token recycling,2024,arXiv】,LLMA【41,Inference with reference: Lossless acceleration of large language models,2023,arXiv】和ANPD【24,Lossless acceleration of large language model via adaptive n-gram parallel decoding,2024,arXiv】。这些方法依赖于小的参考文本并且缺乏自适应推测。SuffixDecoding 是专门为具有长重复序列的 Agent 应用而设计的。

AI Agent 算法。Agent 应用将任务构建为多次 LLM 调用,从而生成长的重复 token 序列。自洽性(Self-consistency)【33,Self-consistency improves chain of thought reasoning in language models,2023a,arXiv】从同一提示中采样多个推理路径,共享相似的思维链步骤。自精炼(Self-refinement)【22,Self-refine: Iterative refinement with self-feedback,2023,arXiv】迭代地修复错误,在保留周围内容的同时修改小部分。多 Agent 工作流【11,Decomposed prompting: A modular approach for solving complex tasks,2023,arXiv】使用具有狭窄功能的专门 Agent,产生重复的输出。这些模式为利用长重复序列创造了机会。

加速 LLM Agent 的方法。最近的工作旨在解决 Agent 应用中的延迟问题。ALTO【25,Alto: An efficient network orchestrator for compound ai systems,2024,EuroMLSys】通过流水线和调度优化多 Agent 工作流。Dynasor【6,Efficiently serving llm reasoning programs with certaindex,2024,arXiv】提前终止不太可能的推理路径。SuffixDecoding 采用了一种正交的推测解码方法,并且可以与这些方法结合使用。

A2 方法细节

SuffixDecoding 的目标。SuffixDecoding 的目标是实现对长序列的快速、自适应的推测解码,特别适用于 Agent 应用,其中重复的推理调用通常包含高度可预测和重叠的 token 序列。在这种设置下,可以根据先前和正在进行的请求准确预测长的输出片段。

SuffixDecoding 面临的挑战。为了充分利用这些机会,SuffixDecoding 必须解决两个关键挑战。首先,它必须支持快速生成推测序列——包括长延续——而不依赖于草稿模型或昂贵的逐个 token 预测。其次,它必须能适应当前的预测上下文:仅在它们很可能被接受时才积极地推测长延续,而在不确定时推测较短的序列,以避免浪费验证计算。

核心机制:后缀树。为了支持快速推测,SuffixDecoding 在当前和先前请求中的 token 上构建了一个后缀树【38,Linear pattern matching algorithms,1973,SWAT】,并使用该后缀树生成推测 token。树的根节点代表存储在树中任何 token 序列后缀的开始,一个节点的每个子节点代表一个特定的 token,这是从其父节点开始的可能延续,从根到每个节点的路径代表一个不同的子序列。

推测来源。对于每个请求,我们考虑从(1)该请求的提示和输出,以及(2)先前请求的输出中推测 token 序列。这样做可以捕捉到许多 Agent 算法(包括自洽性、自精炼和多 Agent 流水线)中输出 token 重复的来源。

基于模式匹配的推测。SuffixDecoding 利用后缀树执行快速模式匹配,并找到 token 序列的可能延续。假设正在进行的推理的提示和输出 token 为 $x_{1:t}$。考虑一个长度为 $p$ 的后缀 $x_{t-p+1:t}$,我们称之为模式序列。我们从根节点开始遍历后缀树,在每一步选择与 token $x_{t-p+i}$ 对应的子节点。如果不存在这样的子节点,则表示未找到该模式,SuffixDecoding 将退回到标准的非推测解码。否则,经过 $p$ 步后,我们到达一个节点,其下降路径是该模式序列的可能延续。

推测树的构建。虽然这个过程可以快速找到一个(可能很大的)候选序列集合,但在推测解码中验证所有这些序列的成本可能过高。相反,SuffixDecoding 通过贪婪扩展和评分程序构建一个更小且可能性更高的推测树,并在基于树的推测中使用这个更小的树。SuffixDecoding 的整体示意图如图1所示,我们将在本节的其余部分详细介绍。

后缀树的构建过程。构建后缀树并在在线推理服务中更新它涉及两个阶段。首先,之前的推理输出可以在一个离线的处理步骤中被添加到树中(例如,从历史日志中),或者在推理服务期间,在每个推理请求完成后在线添加。其次,当前正在进行的提示和输出 token 会在新请求到达和每个新 token 生成时在线添加。

双后缀树设计。在实践中,我们发现维护两个不同的后缀树很方便:一个用于先前生成的输出的全局树,以及一个用于当前正在进行的推理请求的独立的每个请求的树。这避免了因同步来自多个并发请求的后缀树更新而产生的复杂性和开销。全局树可以离线以 $O(n)$ 时间构建,而每个请求的树可以高效地在线构建和更新【28,On-line construction of suffix trees,1995,Algorithmica】。

内存开销考量。虽然后缀树在空间上是 $O(n)$ 高效的,但当有大量先前输出时,全局树仍然可能变得很大。然而,它们只需要 CPU 内存,这在 LLM 服务场景中通常是充足且未被充分利用的。例如,常用于 LLM 服务的 AWS p5.48xlarge 实例拥有 2TB 的主内存,这足以支持一个包含数百万历史输出和数十亿 token 的后缀树。考虑到典型的服务器配置,SuffixDecoding 在需要缓存淘汰之前大约可以缓存一个月的生成 token(详细的内存开销分析见附录 B)。

推测树的扩展策略。给定一个正在进行的推理 $x_{1:t}$ 的模式序列 $x_{t-p+1:t}$,SuffixDecoding 可以在全局或每个请求的后缀树中快速找到一个节点 $N_p$,该节点的下降路径是该模式序列的可能延续。为了选择一个更小、更可能且大小更适合推测验证的子树,我们从单个节点 $N_p$ 开始,并通过一次扩展一个叶节点来贪婪地生长一个子树。

节点评分与选择。具体来说,我们定义了如下评分函数:

images/50e9aeda3467ae0a4d2f6be40f8ab2f29494b86cad446f46f075d18b04db708c.jpg
images/50e9aeda3467ae0a4d2f6be40f8ab2f29494b86cad446f46f075d18b04db708c.jpg

其中 COUNT(N) 是节点 N 在参考语料库中出现的次数,这可以在构建后缀树时计算。从我们推测子树中的单个节点 $N_p$ 开始,我们考虑其所有叶节点的所有子节点,并添加具有最高 $D(N)$ 值的节点 N。重复此过程,直到子树达到预定的大小限制 MAX_SPEC。直观地看,$C(N)$ 估计了 TOKEN(N) 将是子序列 TOKEN(Np), . . . , TOKEN(PARENT(N)) 中下一个观察到的 token 的概率,而 $D(N)$ 估计了 TOKEN(N) 最终被推测树验证接受的概率,假设输出 token 遵循历史模式。因此,SuffixDecoding 通过贪婪地添加它认为在验证期间最有可能被接受的叶节点来构建推测树。

算法1:推测树生成

function EXPANDSPECULATIONTREE(N_p, MAX_SPEC)
    Input: Suffix tree node Np, MAX_SPEC
    Initialize T ← {Np}
    while |T| < MAX_SPEC do
        N ← arg max_{N∈CHILDREN(LEAVES(T))} D(N)
        T ← T ∪ {N}
    end while
    return T
end function

function MATCHPATTERN(S, x_{1:t}, p)
    Input: Suffix tree S, sequence x_{1:t}, length p
    Initialize Np ← ROOT(S)
    for i = 1 to p do
        if NO_CHILD(Np, x_{t−p+i}) then
            return ∅
        end if
        Np ← CHILD(Np, x_{t−p+i})
    end for
    return Np
end function

function GENERATECANDIDATETREE(S_g, S_r, x_{1:t}, α, P)
    Input: Global suffix tree Sg, prompt suffix tree Sr, sequence x_{1:t}, max spec factor α, max pattern size P
    Initialize T_best ← ∅, SCORE_best ← 0
    for S in {Sg, Sr} do
        for p = 1 to P do
            N ← MatchPattern(S, x_{1:t}, p)
            T ← ExpandSpeculationTree(N, αp)
            if SCORE(T) > SCORE_best then
                T_best ← T
                SCORE_best ← SCORE(T)
            end if
        end for
    end for
    return T_best
end function

自适应推测长度。虽然上述过程允许 SuffixDecoding 根据经验概率估计来缓存并快速推测长 token 序列,但它还需要一种机制来自适应地控制其推测的 token 数量。SuffixDecoding 通过动态调整 MAX_SPEC 来实现这一点。较低的值意味着会选择较少但可能性更高的 token 进行推测,而较高的值意味着会选择更多但可能性较低的 token。如果值太低,推测带来的加速可能会受限;如果值太高,则可能会在验证不太可能的 token 上浪费计算资源。

自适应长度的动机。为了指导如何自适应地设置 MAX_SPEC,我们观察到在实践中,接受的 token 数量通常随着模式序列长度 $p$ 的增加而增加(图 2a)。因此,我们将 MAX_SPEC 自适应地定义为:

$$\text{MAX\_SPEC}(p) = \alpha p,$$

其中 $\alpha$ 是用户定义的最大推测因子。图 2b 显示,根据模式长度自适应地设置 MAX_SPEC 可以在接受率和推测加速之间实现更好的权衡。在实践中,我们发现对于 Agent 应用,$\alpha \in [1, 4]$ 的效果很好。

图2:(a) 平均接受的 token 数与模式匹配长度的关系。(b) 使用恒定与自适应 MAX_SPEC 的对比。
图2:(a) 平均接受的 token 数与模式匹配长度的关系。(b) 使用恒定与自适应 MAX_SPEC 的对比。

推测树评分。到目前为止,我们已经讨论了在给定后缀树和模式长度 $p$ 的情况下如何获得一个推测树。然而,SuffixDecoding 维护两个后缀树,即全局后缀树和每个请求的后缀树,每个树都有多种 $p$ 的选择。为了只得到一个推测树,我们为全局后缀树和每个请求的后缀树,以及一系列 $p$ 的值构建推测树。然后,根据一个评分函数选择一个单一的推测树。

评分函数定义。评分函数的定义如下:

$$\text{SCORE}(T_{spec}) = \sum_{N \in T_{spec}} D(N).$$

直观地说,如果 $D(N)$ 估计了推测树 $T_{spec}$ 中的节点 N 将被接受的概率,那么 $SCORE(T_{spec})$ 就估计了期望接受的 token 数。然后,SuffixDecoding 选择具有最高 SCORE 的 $T_{spec}$ 作为最终要验证的推测树。从推测树扩展到评分的端到端候选生成过程在算法1中描述。

混合后缀推测解码。最后,我们发现 $SCORE(T_{spec})$ 可以用来动态地决定是使用 SuffixDecoding 还是回退到基于模型的推测方法,这对于工作负载可能在 Agent 应用和更多样化应用之间混合的实际场景很有用。具体来说,对于每个解码迭代,我们总是首先使用 SuffixDecoding 进行推测。如果 $SCORE(T_{spec}) > \tau$,其中 $\tau$ 是一个可配置的阈值,那么就使用 SuffixDecoding 的草稿 token。否则,我们使用一个回退的推测方法,例如 EAGLE-3【16,Eagle-3: Scaling up inference acceleration of large language models via training-time test,2025a,arXiv】。

混合模式的实践指南。作为一个实践指南,我们发现对于混合工作负载,将 $\tau$ 设置为接近回退推测器的平均接受 token 数效果很好,而对于高度重复的 Agent 任务,仅使用 SuffixDecoding($\tau = 0$)是最佳的。关于 $\tau$ 在不同工作负载类型下的详细敏感性分析在附录 B.1 中提供。对于批处理服务场景,SuffixDecoding 可以与现有的批处理级别推测控制方法集成(见附录 C)。

A4 实验环境

  • 基准比较方法:使用 Spec-Bench【39,Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding,2024,arXiv】与基于模型和无模型的推测解码方法进行比较。

    • 基于模型:EAGLE-2 和 EAGLE-3【15,Eagle-2: Faster inference of language models with dynamic draft trees,2024,arXiv】。
    • 无模型:Prompt-Lookup Decoding (PLD)【26,Prompt lookup decoding,2023,GitHub】和 Token Recycling【20,Turning trash into treasure: Accelerating inference of large language models with token recycling,2024,arXiv】。
  • 数据集与应用

    • Agent 应用

      1. SWE-Bench【10,Swe-bench: Can language models resolve real-world github issues?,2024,arXiv】:一个解决 GitHub 问题的基准,使用 OpenHands【31,OpenHands: An Open Platform for AI Software Developers as Generalist Agents,2024c,arXiv】框架进行追踪。
      2. AgenticSQL:一个专有的多 Agent SQL 生成工作流(如图3所示),包含结构化生成、非结构化生成和检索增强生成等步骤。
    • 非 Agent 工作负载

      1. Spec-Bench:包含13个类别的开放式单轮任务,其中包括8个 MT-Bench【46,Judging llm-as-a-judge with mt-bench and chatbot arena,2023,arXiv】类别(写作、角色扮演、推理、数学、编码、提取、STEM、人文)。
  • 模型架构

    • Spec-Bench 实验使用 Llama-3.1-8B-Instruct。
    • 端到端 SWE-Bench 实验使用为该任务微调的 all-hands/openhands-lm-32b-v0.1-ep3 模型。
  • 硬件配置

    • 实验在一台 AWS p5.48xlarge 实例上进行,配备 8 块 NVIDIA H100 80G GPU 和 2TB 主内存。
  • 软件配置

    • SuffixDecoding 在 vLLM【13,Efficient memory management for large language model serving with pagedattention,2023,arXiv】中实现。
    • 端到端实验中,vLLM 配置了 4 路张量并行和前缀缓存。
  • 模拟实验:除了在真实硬件上的评估,还在第 4.4 节和附录 A.2 中使用模拟验证器进行了额外的消融实验,通过将推测的 token 与真实的响应进行比较来模拟验证过程。

图3:AgenticSQL 是一个多 Agent 工作流,由多个不同 LLM 的结构化生成、非结构化生成和检索增强生成步骤组成。从用户问题中提取有用特征(分类和提取),并用检索到的上下文进行补充(丰富)。多个文本到 SQL 的步骤并行提出用户问题的解决方案(SQL 1...N),并伴有错误校正器的反馈。最后的合并步骤将提出的 SQL 候选合成为最终的 SQL 查询和文本响应。
图3:AgenticSQL 是一个多 Agent 工作流,由多个不同 LLM 的结构化生成、非结构化生成和检索增强生成步骤组成。从用户问题中提取有用特征(分类和提取),并用检索到的上下文进行补充(丰富)。多个文本到 SQL 的步骤并行提出用户问题的解决方案(SQL 1...N),并伴有错误校正器的反馈。最后的合并步骤将提出的 SQL 候选合成为最终的 SQL 查询和文本响应。

A5 实验结果

基准比较(图 4)
- Agent 工作负载:在 AgenticSQL 和 SWE-Bench 上,SuffixDecoding 均优于所有基线。
- 在 AgenticSQL 上,SuffixDecoding 实现了 5.3 倍的平均加速,比 EAGLE-2/3 快 2.8 倍,比 Token Recycling 快 1.9 倍。其平均每步接受 token 数为 6.3,远高于 EAGLE-3(3.6)和 Token Recycling(3.2)。
- 在 SWE-Bench 上,EAGLE-2/3 因序列长度限制而失败。SuffixDecoding 实现了 2.5 倍的平均加速,比次优的 PLD 快 1.7 倍。其平均每步接受 token 数为 7.8,而 PLD 仅为 3.2。

  • 非 Agent 工作负载(Spec-Bench):在这种场景下,单独的 SuffixDecoding 表现不如 EAGLE-2/3 和 Token Recycling。然而,采用 Hybrid 混合方法(SuffixDecoding + EAGLE-3)可以实现两全其美,获得了 2.5 倍的平均加速,优于独立的 EAGLE-3(2.4 倍)和 Token Recycling(2.2 倍)。
  • 混合方法在 Agent 任务上的表现:在 AgenticSQL 上,Hybrid 方法也表现出色,树变体实现了 4.1 倍的加速,远优于独立的 EAGLE-2/3(1.9 倍)。尽管其平均接受长度(7.5)高于 SuffixDecoding(6.3),但 SuffixDecoding 凭借其极低的推测成本和高接受率,在 Agent 任务中仍是最终的赢家(5.3 倍 vs 4.1 倍加速)。
图4:与 Vanilla 解码相比,SuffixDecoding 和基线方法在三个基准测试(Spec-Bench、AgenticSQL 和 SWE-Bench)上的推测加速(上图)和每步平均接受的 token 数(下图)。实验使用 Llama-3.1-8B-Instruct 在单张 H100 GPU 上进行,批处理大小为 1。加速比是相对于 Vanilla 解码的每输出 token 壁钟时间的比率。Suffix (tree) 和 Hybrid (tree) 使用 SuffixDecoding 的树推测算法。Suffix (linear) 和 Hybrid (linear) 使用更简单的线性推测方法。混合变体将 SuffixDecoding 与 EAGLE-3 结合。注意 EAGLE-2/3 和 Token Recycling 因上下文过长在几个 SWE-Bench 任务上运行失败,由缺失的条形图表示。
图4:与 Vanilla 解码相比,SuffixDecoding 和基线方法在三个基准测试(Spec-Bench、AgenticSQL 和 SWE-Bench)上的推测加速(上图)和每步平均接受的 token 数(下图)。实验使用 Llama-3.1-8B-Instruct 在单张 H100 GPU 上进行,批处理大小为 1。加速比是相对于 Vanilla 解码的每输出 token 壁钟时间的比率。Suffix (tree) 和 Hybrid (tree) 使用 SuffixDecoding 的树推测算法。Suffix (linear) 和 Hybrid (linear) 使用更简单的线性推测方法。混合变体将 SuffixDecoding 与 EAGLE-3 结合。注意 EAGLE-2/3 和 Token Recycling 因上下文过长在几个 SWE-Bench 任务上运行失败,由缺失的条形图表示。

推测树案例分析(图 5)
- 为了直观理解 SuffixDecoding 的高性能,本文分析了 AgenticSQL 中 Extract 任务的一个推测树。该任务的输出是格式和键名相同的 JSON 文档,其中包含许多离散值(尤其是布尔值 true/false)。
- SuffixDecoding 的全局后缀树记录了这些模式。如图 5 所示,它构建了一个在连续特征的 true/false 值处分支的大型推测树。这个树总有一个分支具有很高的接受率,能够一步生成数十个 token,这展示了 SuffixDecoding 发现和利用复杂模式的能力,尤其是在结构化生成任务中。

图5:一个为 AgenticSQL Extract 任务构建的包含 66 个 token 的 SuffixDecoding 推测树。
图5:一个为 AgenticSQL Extract 任务构建的包含 66 个 token 的 SuffixDecoding 推测树。

端到端 SWE-Bench 性能(图 6)
- 将 SuffixDecoding 集成到 vLLM 中,并实时运行 OpenHands Agent 解决 SWE-Bench 问题。
- 实验表明,解码时间(即输出生成)在所有 SWE-Bench 任务中占主导地位。
- 在此端到端场景中,SuffixDecoding 的性能比 vLLM 自带的 PLD 高出 1.3-3 倍,相对于 Vanilla 解码实现了 1.8-4.5 倍的推测加速。
- 由于 SuffixDecoding 是无损的,它保持了原始模型在 SWE-Bench Verified 上的 37.2% 的得分。

图6:OpenHands Agent 在 SWE-Bench Verified 上的端到端任务完成时间。基准测试以 8 个任务并发运行。vLLM 部署在 4 个 H100 GPU 上,配置为 4 路张量并行并启用前缀缓存。结果按 SWE-Bench 中的不同代码库进行了分解。
图6:OpenHands Agent 在 SWE-Bench Verified 上的端到端任务完成时间。基准测试以 8 个任务并发运行。vLLM 部署在 4 个 H100 GPU 上,配置为 4 路张量并行并启用前缀缓存。结果按 SWE-Bench 中的不同代码库进行了分解。

消融实验
- 全局 vs. 请求级后缀树(图 7)
- 在 AgenticSQL 上的实验表明,同时使用两个后缀树(全局树和请求级树)通常比只使用一个效果更好。
- 在大多数任务中,全局树的贡献大于请求级树。然而,对于 Combine 任务,请求级树更优,因为它严重依赖于当前上下文中的 token。这表明 SuffixDecoding 通过结合两种树,能够适应不同特性的任务。

  • 开放式场景下的表现(图 8)
    • 在 WildChat(开放式聊天)和 Magicoder(代码导向聊天)等重复性较低的工作负载上评估 SuffixDecoding。
    • 结果显示,随着后缀树规模(即缓存的输出示例数量)的增长,加速比持续提高。这表明即使在重复性较低的场景中,SuffixDecoding 也能学习到有用的模式。
    • 接受率在不同树规模下保持稳定,这归功于自适应的推测长度机制 MAX_SPEC = αp,它会根据模式匹配的长度自动调整推测的 token 数量。
图7:AgenticSQL 中各任务的加速因子和推测 token 数。SuffixDecoding 分别在仅使用全局后缀树、仅使用请求级后缀树和同时使用两者(基线)的情况下运行。
图7:AgenticSQL 中各任务的加速因子和推测 token 数。SuffixDecoding 分别在仅使用全局后缀树、仅使用请求级后缀树和同时使用两者(基线)的情况下运行。
图8:在 Magicoder 和 Wildchat 上,加速比(左)和接受率(右)与全局后缀树大小的关系 (α = 1)。随着先前输出示例的增多,SuffixDecoding 的加速比持续增加,而接受率保持稳定。
图8:在 Magicoder 和 Wildchat 上,加速比(左)和接受率(右)与全局后缀树大小的关系 (α = 1)。随着先前输出示例的增多,SuffixDecoding 的加速比持续增加,而接受率保持稳定。

A6 结论

本文介绍了 SuffixDecoding,一种为新兴 Agent 应用设计的无模型推测解码方法。通过使用高效的后缀树数据结构,SuffixDecoding 有效地利用了许多 Agent 算法(如自洽性、自精炼和多 Agent 流水线)中发现的长而重复的 token 序列。通过在两个实际的 Agent 应用(OpenHands 和 AgenticSQL)上的实验,我们证明了 SuffixDecoding 显著加速了它们的解码延迟和任务完成时间,并且比其他基于模型和无模型的推测解码基线快得多。

A7 附录

A.1 主要实验细节

Spec-Bench 实验设置 (4.2节)。我们通过运行原始仓库中的 Spec-Bench 代码库进行了 Spec-Bench 实验,并做了以下修改。首先,我们更新了代码以适配最新版本的 transformers 库,这是运行 Llama-3.1 等最新开源 LLM 所必需的。我们还添加了对任意数据集(如 SWE-Bench 和 AgenticSQL)的支持,并在框架内实现了 SuffixDecoding。我们在一个 8xH100 80GB GPU 集群(1TB RAM)上运行实验。我们为每个基线使用一个 GPU,批处理大小为 1,与原始 SpecBench 代码一致。

vLLM SWE-Bench 实验设置 (4.3节)。我们在一个 8xH100 80GB GPU 集群(1TB RAM)上进行了端到端的 SWE-Bench 实验。我们使用 vLLM 在本地部署了 all-hands/openhands-lm-32b-v0.1-ep3 模型,张量并行度为 4,并启用了前缀缓存。我们使用 flashinfer 内核进行采样。我们对 vLLM 做了一些微小修改,以记录感兴趣的每个请求和每步的统计数据(每个 token 的延迟、吞吐量、接受长度、接受率)。我们对所有基线使用了相同的设置。我们在同一台机器上运行 OpenHands 守护进程,并使用 OpenAI API 与 vLLM 服务器交互。我们使用 CodeActAgent【30,Executable code actions elicit better llm agents,2024b,arXiv】运行 OpenHands,设置 ITERATIVE_EVAL_MODE=true,最大迭代次数为 100,这是 OpenHands 作者推荐的。我们使用最多 16 个并发工作者来运行 SWE-Bench 任务。

详细的子任务结果。以下表格展示了所有评估基准的详细性能指标。

SWE-Bench: 平均接受 token 数 (tokens/step)
SWE-Bench: 平均接受 token 数 (tokens/step)

SWE-Bench: 平均接受率
SWE-Bench: 平均接受率

SWE-Bench: 每个输出 token 的时间 (ms)
SWE-Bench: 每个输出 token 的时间 (ms)

SWE-Bench: 每个生成 token 的推测时间 (ms)
SWE-Bench: 每个生成 token 的推测时间 (ms)

SWE-Bench: 相对于 vanilla 解码的加速比
SWE-Bench: 相对于 vanilla 解码的加速比

AgenticSQL: 平均接受 token 数 (tokens/step)
AgenticSQL: 平均接受 token 数 (tokens/step)

AgenticSQL: 平均接受率
AgenticSQL: 平均接受率

AgenticSQL: 每个输出 token 的时间 (ms)
AgenticSQL: 每个输出 token 的时间 (ms)

AgenticSQL: 每个生成 token 的推测时间 (ms)
AgenticSQL: 每个生成 token 的推测时间 (ms)

AgenticSQL: 相对于 vanilla 解码的加速比
AgenticSQL: 相对于 vanilla 解码的加速比

Spec-Bench: 平均接受 token 数 (tokens/step)
Spec-Bench: 平均接受 token 数 (tokens/step)

Spec-Bench: 平均接受率
Spec-Bench: 平均接受率

Spec-Bench: 每个输出 token 的时间 (ms)
Spec-Bench: 每个输出 token 的时间 (ms)

Spec-Bench: 每个生成 token 的推测时间 (ms)
Spec-Bench: 每个生成 token 的推测时间 (ms)

Spec-Bench: 相对于 vanilla 解码的加速比
Spec-Bench: 相对于 vanilla 解码的加速比

A.2 额外的消融实验

额外数据集详情

  1. WildChat【45,Wildchat: 1m chatGPT interaction logs in the wild,2024,ICLR】:使用来自 WildChat 数据集的指令,该数据集包含用户与 ChatGPT 服务的真实交互。
  2. Magicoder【37,Magicoder: Source code is all you need,2023,arXiv】:使用 Magicoder-EvolInstruct-110K 数据集的指令,其中包含通过 Self-Instruct【2,Code alpaca: An instruction-following llama model for code generation,2023,GitHub;34,Self-instruct: Aligning language models with self-generated instructions,2023b,ACL】生成的与代码相关的问题和指令。
  3. SpiderSQL:Spider【43,Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-SQL task,2018,EMNLP】是一个包含手动标注的问题和 SQL 响应的数据集。我们使用 DAIL-SQL【7,Text-to-sql empowered by large language models: A benchmark evaluation,2024a,VLDB】中的指令。

输入分布变化的影响。在实际的 LLM 服务中,请求的输入特征可能会随时间变化,并且可能与 SuffixDecoding 训练所用的输出示例分布不同。为了评估这种情况,我们运行了在 WildChat 输出上训练的 SuffixDecoding,并开始向其发送来自 SpiderSQL 的输入,这代表了一个非常突然的分布变化。如图 9 所示,SuffixDecoding 从拥有 4000 个 WildChat 输出示例开始,并开始接收 SpiderSQL 推理请求。没有任何适应的情况下,SuffixDecoding 仍然实现了 1.5 倍的加速和 8% 的接受率,但远未达到如果它在 500 个 SpiderSQL 示例上训练所能达到的 2.6 倍加速和 20% 接受率。在处理每个 SpiderSQL 推理请求后,SuffixDecoding 可以将其输出插入其全局后缀树中,这意味着它可以以在线方式适应新的输入分布。如图 9 所示,SuffixDecoding 的性能随着处理的 SpiderSQL 推理请求数量的增加而提高。令人惊讶的是,在观察并在线适应 500 个 SpiderSQL 示例后,SuffixDecoding 的性能与仅在 500 个 SpiderSQL 示例上离线训练的性能几乎没有区别。这表明 SuffixDecoding 能够快速适应输入分布的变化,且性能没有损失。

图9:SuffixDecoding 在输入分布变化下的性能。SuffixDecoding 在 WildChat 的输出上训练,但在 SpiderSQL 上进行评估。X 轴:SpiderSQL 输入的数量,它们在处理后被添加到全局后缀树中。红线:如果 SuffixDecoding 仅在 500 个 SpiderSQL 的输出示例上离线训练的性能。
图9:SuffixDecoding 在输入分布变化下的性能。SuffixDecoding 在 WildChat 的输出上训练,但在 SpiderSQL 上进行评估。X 轴:SpiderSQL 输入的数量,它们在处理后被添加到全局后缀树中。红线:如果 SuffixDecoding 仅在 500 个 SpiderSQL 的输出示例上离线训练的性能。

预测 SuffixDecoding 的有效性。SuffixDecoding 在更结构化的任务上(如 AgenticSQL)往往比在非常开放式的任务上(如 WildChat)表现更好。我们可以使用经验熵来衡量这种“结构化”程度。步骤如下:(1)从示例模型输出中创建一个后缀树(通常 100 个示例就足够了),(2)通过确定每个子节点被访问的频率来计算每个节点输出分布的熵,以及(3)计算所有节点熵的加权平均值。较低的平均熵表明输出 token 基于其前缀更可预测,这通常意味着 SuffixDecoding 会表现得更好。表 1 显示了在我们各个评估数据集样本上测得的经验熵。我们发现平均熵与对每个数据集“结构化”程度的直观理解密切相关。此外,它与 SuffixDecoding 在这些数据集上的性能也很好地相关。因此,从业者可以使用少量输出示例来计算此值,以评估 SuffixDecoding 是否适合他们的特定任务。

表1:我们各个评估数据集的实测经验熵。
表1:我们各个评估数据集的实测经验熵。

B 可扩展性和内存开销

性能指标。SuffixDecoding 使用的后缀树在时间和内存上都非常高效。表 2 显示了不同树大小下(使用 SWE-Bench 数据集的请求)的每个 token 的查找时间、更新时间和内存消耗。总体而言,总内存消耗随树的大小线性扩展,并且即使树变大,更新和查找时间也保持快速。

表2:不同条目数量的性能指标
表2:不同条目数量的性能指标

内存容量分析。考虑到典型的服务器配置,SuffixDecoding 在达到 CPU 内存限制之前可以缓存数周的生成 token。例如,在 A100 GPU 上运行 Llama-8B 的优化推理引擎通常每秒可生成约 5000 个 token【42,Achieving faster open-source llama3 serving with sglang runtime (vs. tensorrt-llm, vllm),2024,LMSYS Org Blog】,即每天约 4.32 亿个 token。典型的 A100 系统(例如 AWS p4d.24xlarge)每个 A100 GPU 有 144GB 的 CPU 内存。这意味着 SuffixDecoding 在达到 CPU 内存限制之前,每个 A100 GPU 可能可以缓存 31 天的生成 token。超出此限制后,也可以直接在 SuffixDecoding 中集成缓存淘汰策略(例如 LRU)以避免内存不足问题。

B.1 混合回退阈值敏感性分析

阈值 τ 的作用。混合 SuffixDecoding 方法使用一个回退阈值 $\tau$ 来决定何时使用基于后缀树的推测,何时回退到基于模型的推测器(例如 EAGLE-3)。本节分析了在开放式和 Agent 工作负载中,性能对不同阈值的敏感性。

开放式工作负载(Spec-Bench)。表 3 显示了 Spec-Bench 在不同阈值下的壁钟加速结果。对于开放式生成任务,我们发现将 $\tau$ 设置为接近或略高于模型基推测器的平均接受 token 数(MAT)会产生最佳性能。例如,EAGLE-3 在 Spec-Bench 上的 MAT 约为 4.65 tokens/step。结果显示,$\tau \in [5, 7]$ 产生了最高的整体加速(2.5 倍),且在该范围内性能相对稳定。

表3:混合 SuffixDecoding 在不同阈值 τ 下于 Spec-Bench 上的加速比。
表3:混合 SuffixDecoding 在不同阈值 τ 下于 Spec-Bench 上的加速比。

Agent 工作负载(AgenticSQL)。相反,表 4 显示了 AgenticSQL(一个高度重复的 Agent 工作负载)的结果。在这里,独立的 SuffixDecoding($\tau = 0$)实现了最佳的整体加速(5.35 倍),因为它能够自信地推测比基于模型的方法长得多的序列。较高的阈值会逐渐降低性能,因为系统越来越多地回退到效果较差的基于模型的推测。

表4:混合 SuffixDecoding 在不同阈值 τ 下于 AgenticSQL 上的加速比。
表4:混合 SuffixDecoding 在不同阈值 τ 下于 AgenticSQL 上的加速比。

实践指南。总而言之,从业者应为混合或开放式工作负载设置 $\tau \approx MAT_{fallback}$,其中 $MAT_{fallback}$ 是回退的基于模型推测器的平均接受 token 数。对于高度重复的 Agent 工作负载,单独使用 SuffixDecoding($\tau = 0$ 或非常低的值)会产生最佳性能。

C 批处理级别的推测控制

兼容性。对于批处理服务场景,优化批处理中每个请求的推测对于许多实际的在线部署至关重要。虽然 SuffixDecoding 专注于推测哪些 token,但它也与探索每个请求应推测多少 token 的现有工作兼容。例如,TurboSpec【19,Turbospec: Closed-loop speculation control system for optimizing llm serving goodput,2024,arXiv】使用“goodput”指标来平衡推测长度和批处理大小,并指出最佳推测随着批处理大小的增加而减少。AdaServe【17,Adaserve: Accelerating multi-slo llm serving with slo-customized speculative decoding,2025b,arXiv】将多 SLO 服务构建为一个约束优化问题,并引入了 SLO 定制的推测解码,该解码构建了针对每个请求的延迟目标的推测树。

未来工作。通过动态调整 SuffixTree 为每个请求生成的推测 token 数量,SuffixDecoding 可以与 TurboSpec 集成以最大化批处理的 goodput 指标。此外,SuffixDecoding 使用的基于统计的评分可以潜在地为像 TurboSpec 这样的方法提供更好的信息。例如,选择从具有更高边际分数的序列中进行更多推测。这些都是未来工作中有趣且重要的方向。