SSD:并行化草拟与验证,突破大模型推理速度瓶颈

斯坦福等提出SSD框架,并行化草拟和验证,突破大模型推理瓶颈,速度提升2倍!

原文标题:比全球最强推理引擎还快2倍,斯坦福、普林斯顿破解大模型「串行魔咒」

原文作者:机器之心

冷月清谈:

本文介绍了斯坦福、普林斯顿大学和 Together AI 团队提出的新型推测解码框架 SSD 及其优化算法 SAGUARO。该框架旨在解决传统推测解码(SD)中草拟和验证串行进行的瓶颈,通过并行化这两个过程,显著提升大语言模型的推理速度。SSD 的核心思想是预测验证结果,并在独立硬件上并行进行推测和验证。为了优化 SSD 算法,研究者们提出了多种策略,包括将预测验证结果转化为约束优化问题、平衡预测准确性与生成高质量推测,以及设计最优回退策略以应对预测失败的情况。实验结果表明,Saguaro 相比优化的推测性解码实现了高达 2 倍的加速,相比自回归生成实现了高达 5 倍的加速。此外,文章还探讨了未来研究方向,例如将 SSD 与 EAGLE 技术、token-tree 推测相结合,以及在集群层面跨多个目标模型部署共享推测端点。

怜星夜思:

1、SSD框架中提到的“草拟模型预测最可能的验证结果”,这个“最可能”的标准是什么?如果判断失误,会造成多大的性能损失?
2、文章提到SSD可以和EAGLE技术以及token-tree推测结合,这两种技术分别是什么? 结合后能带来哪些好处?
3、文章提到,通过扩展草拟设备的数量以及推测缓存,延迟可以进一步减少。那么,仅仅通过增加更多的GPU来做推测,是不是就能无限提升性能? 瓶颈会在哪里?

原文内容

图片
机器之心编辑部

在大语言模型推理领域,虽然「推测解码」(Speculative Decoding,SD)已成为加速生成的标准配置,但它依然存在一个致命弱点: drafting(草拟)和 verification(验证)之间必须串行进行。


近日,来自斯坦福、普林斯顿大学和 Together AI 的研究团队提出 SSD 框架及其优化算法 SAGUARO,成功实现了草拟和验证的并行化



  • 论文链接:https://arxiv.org/pdf/2603.03251

  • GitHub 链接:https://github.com/tanishqkumar/ssd


据介绍,该算法推理速度比世界上最强大的推理引擎都快 2 倍


图片


「推测性推测解码」(Speculative Speculative Decoding,简称 SSD),是一种新型的推测性解码 (SD)。在传统的 SD 中,一个小且快速的模型会先猜测大且慢的模型可能生成的下几个 token,然后大模型通过一次前向传播验证这些猜测,草拟和验证是依次进行的。


而在 SSD 中,这两个过程是并行发生的,完全消除了运行小模型的开销。



并行执行草拟和验证是很棘手的,因为你无法在某事物存在之前对其进行验证,也无法推测你不知道的前缀之外的内容。在 SSD 中,研究者预先设定验证结果,并在独立硬件上进行验证的同时进行推测。这样一来,如果其中一种验证结果出现,推测就能立即生效。


虽然论文对算法进行了详尽的理论描述,但在实践中,研究者大部分时间都花在如何让它与现代推理引擎中各种优化技术(Paged Attention、Prefix Caching、CUDAGraphs 等)协同工作上。


作者 Tanishq Kumar 表示,「真的,我花在了解 CPU/GPU 同步问题上的时间远远超过了我的预期」。



要使新算法达到 SOTA 水平,需要精心设计系统和算法。总的来说,SSD 推动了延迟 - 吞吐量帕累托前沿的发展,其方式与普通推测解码推进标准自回归算法的发展非常相似。



Tanishq Kumar 称:「我对快速推理感到兴奋,因为我非常关注的一个人工智能工作负载是超长时域推理。想象一下,一个拥有大量 B200 的数据中心完全用于运行一个模型,该模型需要处理数十亿个 token 来证明 P 与 NP 的区别。在这种情况下,延迟减半就意味着可以进行双倍深度的思考!」



SSD 如何实现草拟与验证并行?


现代 AI 对推理速度有着极高的要求。然而,标准的语言模型解码是按顺序生成单个 token,未能利用现代硬件上可用的大规模并行计算。


推测性解码(SD)是一种为了解决这个问题而引入的技术。它使用一个快速的「草拟模型」来预测目标模型可能将生成的下几个 token,而不是从目标模型中进行缓慢的自回归采样,然后通过目标模型的并行前向传播验证这些 token。这一验证是按照一个算法进行的,确保生成的 token 是从目标模型的分布中采样的。


在每次验证中,目标模型决定接受多少个推测的 token,并采样一个额外的奖励 token,该 token 跟随所有已接受的 token。尽管推测性解码有效,但它本身仍受限于串行依赖:必须等待当前验证完成后,才能开始下一轮推测。


那么,我们能否消除草拟和验证之间的顺序依赖呢?


研究者引入推测性推测解码(SSD),这是一个旨在并行化草拟和验证的统一框架。


在 SD 中,草拟模型必须等待验证完成,才能开始推测下一轮,而在 SSD 中,草拟模型会预测最可能的验证结果,并在验证进行的同时,针对所有可能的结果进行并行的提前推测。如果这些预准备的结果中的任何一个发生,草拟模型可以立即将预推测的 token 发送给验证器,从而避免草拟阶段的开销。与普通的推测性解码一样,SSD 也是无损的。不同之处在于,SSD 的草拟模型部署在与目标模型不同的硬件上。


优化 SSD 算法主要面临三大挑战


首先,草拟模型必须准确预测验证结果,这不仅包括接受了多少个推测的 Token,还包括采样的奖励 token。其次,推测器的接受率与其预测验证结果的能力之间存在微妙的权衡,必须谨慎处理以最大化加速比。此外,任何 SSD 算法都必须具备处理预测失败的回退策略,因为在大批处理量和高随机性(Temperature)下,预测失败会频繁发生,若处理不当,即时补救的开销将抵消异步带来的收益。


为此,他们推出了 Saguaro,这是一个优化的 SSD 算法,针对上述挑战进行了定向优化。


  • 将预测验证结果的问题转化为约束优化问题,并引入了一种技术,利用最可能的草拟 logits 来预测奖励 token,准确率最高可达 90%。

  • 识别了预测准确性与生成高质量推测之间的张力,并开发了一种能够平衡二者的采样算法。

  • 探讨了处理预测失败的多种策略,发现最优回退策略随批处理大小而异。通过采用这些优化,尽管 Saguaro 在处理每个批次元素时进行了更多计算(同时解码多种可能的结果),其表现仍比标准 SD 高出 20%。


总的来看,Saguaro 相比优化的推测性解码实现了高达 2 倍的加速,相比自回归生成实现了高达 5 倍的加速,并在各种批处理规模下均显著提升了吞吐量与延迟的帕累托前沿


不过,该领域仍有许多值得探索的方向。SSD 可以自然地与 EAGLE 技术以及 token-tree 推测(Token-tree speculation)相结合,但这种联合设计及其权衡空间在很大程度上尚未被发掘。


此外,通过扩展草拟设备的数量以及推测缓存,延迟可以进一步减少,尽管回报最终会递减。最后,在集群层面跨多个目标模型部署共享推测端点——类似于预填充-解码分解)——是另一个自然的研究方向。




© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:liyazhou@jiqizhixin.com

我觉得瓶颈在于“草拟”这个动作本身。毕竟草拟模型也是需要训练的,而且要尽可能拟合目标模型,不然预测不准就白搭。所以,草拟模型的质量决定了能提升多少性能。无限增加GPU数量,如果草拟模型拉胯,那就是浪费资源。

EAGLE (Early Aggregated Gradient Layer Ensemble) 技术,我理解是利用多个模型进行ensemble,提高预测的准确性,从而减少回退操作。如果SSD和EAGLE结合,可以提高草拟模型预测验证结果的准确性,降低预测失败的概率,进一步提升整体性能。

这个问题很有意思!我理解的“最可能”,应该是指基于草拟模型对目标模型行为的最佳估计,可以理解为概率最高的那个结果。如果判断失误,性能损失主要体现在回退策略的开销上。论文里也提到,要仔细设计回退策略才能弥补异步带来的收益,具体损失大小取决于回退策略的效率,以及预测失败的频率。

这俩技术听起来就很fancy!token-tree感觉像是在赌哪个分支会中,EAGLE像是集思广益,提高胜率。SSD本身就在搞并行,再结合这俩,不得起飞? 感觉以后炼丹又多了好多调参的新姿势了 (痛苦面具)。

从硬件角度来说,GPU数量增加会导致互联带宽成为瓶颈,数据传输速度跟不上。从软件角度来说,任务调度和资源管理的复杂度会大大增加,需要更精细的优化。而且,别忘了还有成本问题,无限增加GPU,钞能力也遭不住啊!

我认为这个“最可能”的标准肯定和loss function有关!草拟模型训练的时候就是要尽可能拟合验证模型的结果,loss越小,预测就越准,翻车的可能性就越低。不过就算翻车了,也就是多算一步的事情,毕竟人家都说了快2倍了,底子好!

学术一点说,这里的“最可能”大概率是指草拟模型输出的logits中概率最高的那个token序列。判断失误的性能损失主要来源于两方面:一是需要回退到串行验证模式,二是错误的推测可能导致缓存失效,需要重新计算。具体损失可以用公式量化,但涉及到模型参数和硬件性能,比较复杂。

EAGLE我不太熟,但token-tree推测我了解一些。简单说,就是不再只推测一个token序列,而是推测多个可能的token分支,形成一个树状结构。这样可以提高推测的覆盖范围,降低回退的概率。如果SSD和token-tree结合,感觉能进一步减少草拟和验证之间的依赖,提高并行度。

理论上无限增加GPU肯定不现实。瓶颈可能出现在以下几个方面:一是GPU之间的通信开销会随着数量增加而上升,二是目标模型的验证速度可能跟不上,三是推测缓存的大小有限,频繁的缓存未命中会降低效率。最终会达到一个收益递减点。