機(jī)器之心編輯部

在大語(yǔ)言模型推理領(lǐng)域，雖然「推測(cè)解碼」（Speculative Decoding，SD）已成為加速生成的標(biāo)準(zhǔn)配置，但它依然存在一個(gè)致命弱點(diǎn)： drafting（草擬）和 verification（驗(yàn)證）之間必須串行進(jìn)行。

近日，來(lái)自斯坦福、普林斯頓大學(xué)和 Together AI 的研究團(tuán)隊(duì)提出SSD 框架及其優(yōu)化算法 SAGUARO，成功實(shí)現(xiàn)了草擬和驗(yàn)證的并行化。

• 論文鏈接：https://arxiv.org/pdf/2603.03251
• GitHub 鏈接：https://github.com/tanishqkumar/ssd

據(jù)介紹，該算法推理速度比世界上最強(qiáng)大的推理引擎都快 2 倍。

「推測(cè)性推測(cè)解碼」（Speculative Speculative Decoding，簡(jiǎn)稱(chēng) SSD），是一種新型的推測(cè)性解碼 (SD)。在傳統(tǒng)的 SD 中，一個(gè)小且快速的模型會(huì)先猜測(cè)大且慢的模型可能生成的下幾個(gè) token，然后大模型通過(guò)一次前向傳播驗(yàn)證這些猜測(cè)，草擬和驗(yàn)證是依次進(jìn)行的。

而在 SSD 中，這兩個(gè)過(guò)程是并行發(fā)生的，完全消除了運(yùn)行小模型的開(kāi)銷(xiāo)。

并行執(zhí)行草擬和驗(yàn)證是很棘手的，因?yàn)槟銦o(wú)法在某事物存在之前對(duì)其進(jìn)行驗(yàn)證，也無(wú)法推測(cè)你不知道的前綴之外的內(nèi)容。在 SSD 中，研究者預(yù)先設(shè)定驗(yàn)證結(jié)果，并在獨(dú)立硬件上進(jìn)行驗(yàn)證的同時(shí)進(jìn)行推測(cè)。這樣一來(lái)，如果其中一種驗(yàn)證結(jié)果出現(xiàn)，推測(cè)就能立即生效。

雖然論文對(duì)算法進(jìn)行了詳盡的理論描述，但在實(shí)踐中，研究者大部分時(shí)間都花在如何讓它與現(xiàn)代推理引擎中各種優(yōu)化技術(shù)（Paged Attention、Prefix Caching、CUDAGraphs 等）協(xié)同工作上。

作者 Tanishq Kumar 表示，「真的，我花在了解 CPU/GPU 同步問(wèn)題上的時(shí)間遠(yuǎn)遠(yuǎn)超過(guò)了我的預(yù)期」。

要使新算法達(dá)到 SOTA 水平，需要精心設(shè)計(jì)系統(tǒng)和算法。總的來(lái)說(shuō)，SSD 推動(dòng)了延遲 - 吞吐量帕累托前沿的發(fā)展，其方式與普通推測(cè)解碼推進(jìn)標(biāo)準(zhǔn)自回歸算法的發(fā)展非常相似。

Tanishq Kumar 稱(chēng)：「我對(duì)快速推理感到興奮，因?yàn)槲曳浅ｊP(guān)注的一個(gè)人工智能工作負(fù)載是超長(zhǎng)時(shí)域推理。想象一下，一個(gè)擁有大量 B200 的數(shù)據(jù)中心完全用于運(yùn)行一個(gè)模型，該模型需要處理數(shù)十億個(gè) token 來(lái)證明 P 與 NP 的區(qū)別。在這種情況下，延遲減半就意味著可以進(jìn)行雙倍深度的思考！」

SSD 如何實(shí)現(xiàn)草擬與驗(yàn)證并行？

現(xiàn)代 AI 對(duì)推理速度有著極高的要求。然而，標(biāo)準(zhǔn)的語(yǔ)言模型解碼是按順序生成單個(gè) token，未能利用現(xiàn)代硬件上可用的大規(guī)模并行計(jì)算。

推測(cè)性解碼（SD）是一種為了解決這個(gè)問(wèn)題而引入的技術(shù)。它使用一個(gè)快速的「草擬模型」來(lái)預(yù)測(cè)目標(biāo)模型可能將生成的下幾個(gè) token，而不是從目標(biāo)模型中進(jìn)行緩慢的自回歸采樣，然后通過(guò)目標(biāo)模型的并行前向傳播驗(yàn)證這些 token。這一驗(yàn)證是按照一個(gè)算法進(jìn)行的，確保生成的 token 是從目標(biāo)模型的分布中采樣的。

在每次驗(yàn)證中，目標(biāo)模型決定接受多少個(gè)推測(cè)的 token，并采樣一個(gè)額外的獎(jiǎng)勵(lì) token，該 token 跟隨所有已接受的 token。盡管推測(cè)性解碼有效，但它本身仍受限于串行依賴(lài)：必須等待當(dāng)前驗(yàn)證完成后，才能開(kāi)始下一輪推測(cè)。

那么，我們能否消除草擬和驗(yàn)證之間的順序依賴(lài)呢？

研究者引入推測(cè)性推測(cè)解碼（SSD），這是一個(gè)旨在并行化草擬和驗(yàn)證的統(tǒng)一框架。

在 SD 中，草擬模型必須等待驗(yàn)證完成，才能開(kāi)始推測(cè)下一輪，而在 SSD 中，草擬模型會(huì)預(yù)測(cè)最可能的驗(yàn)證結(jié)果，并在驗(yàn)證進(jìn)行的同時(shí)，針對(duì)所有可能的結(jié)果進(jìn)行并行的提前推測(cè)。如果這些預(yù)準(zhǔn)備的結(jié)果中的任何一個(gè)發(fā)生，草擬模型可以立即將預(yù)推測(cè)的 token 發(fā)送給驗(yàn)證器，從而避免草擬階段的開(kāi)銷(xiāo)。與普通的推測(cè)性解碼一樣，SSD 也是無(wú)損的。不同之處在于，SSD 的草擬模型部署在與目標(biāo)模型不同的硬件上。

優(yōu)化 SSD 算法主要面臨三大挑戰(zhàn)。

首先，草擬模型必須準(zhǔn)確預(yù)測(cè)驗(yàn)證結(jié)果，這不僅包括接受了多少個(gè)推測(cè)的 Token，還包括采樣的獎(jiǎng)勵(lì) token。其次，推測(cè)器的接受率與其預(yù)測(cè)驗(yàn)證結(jié)果的能力之間存在微妙的權(quán)衡，必須謹(jǐn)慎處理以最大化加速比。此外，任何 SSD 算法都必須具備處理預(yù)測(cè)失敗的回退策略，因?yàn)樵诖笈幚砹亢透唠S機(jī)性（Temperature）下，預(yù)測(cè)失敗會(huì)頻繁發(fā)生，若處理不當(dāng)，即時(shí)補(bǔ)救的開(kāi)銷(xiāo)將抵消異步帶來(lái)的收益。

為此，他們推出了Saguaro，這是一個(gè)優(yōu)化的 SSD 算法，針對(duì)上述挑戰(zhàn)進(jìn)行了定向優(yōu)化。

• 將預(yù)測(cè)驗(yàn)證結(jié)果的問(wèn)題轉(zhuǎn)化為約束優(yōu)化問(wèn)題，并引入了一種技術(shù)，利用最可能的草擬 logits 來(lái)預(yù)測(cè)獎(jiǎng)勵(lì) token，準(zhǔn)確率最高可達(dá) 90%。
• 識(shí)別了預(yù)測(cè)準(zhǔn)確性與生成高質(zhì)量推測(cè)之間的張力，并開(kāi)發(fā)了一種能夠平衡二者的采樣算法。
• 探討了處理預(yù)測(cè)失敗的多種策略，發(fā)現(xiàn)最優(yōu)回退策略隨批處理大小而異。通過(guò)采用這些優(yōu)化，盡管 Saguaro 在處理每個(gè)批次元素時(shí)進(jìn)行了更多計(jì)算（同時(shí)解碼多種可能的結(jié)果），其表現(xiàn)仍比標(biāo)準(zhǔn) SD 高出 20%。

總的來(lái)看，Saguaro 相比優(yōu)化的推測(cè)性解碼實(shí)現(xiàn)了高達(dá) 2 倍的加速，相比自回歸生成實(shí)現(xiàn)了高達(dá) 5 倍的加速，并在各種批處理規(guī)模下均顯著提升了吞吐量與延遲的帕累托前沿。

不過(guò)，該領(lǐng)域仍有許多值得探索的方向。SSD 可以自然地與 EAGLE 技術(shù)以及token-tree推測(cè)（Token-tree speculation）相結(jié)合，但這種聯(lián)合設(shè)計(jì)及其權(quán)衡空間在很大程度上尚未被發(fā)掘。

此外，通過(guò)擴(kuò)展草擬設(shè)備的數(shù)量以及推測(cè)緩存，延遲可以進(jìn)一步減少，盡管回報(bào)最終會(huì)遞減。最后，在集群層面跨多個(gè)目標(biāo)模型部署共享推測(cè)端點(diǎn)——類(lèi)似于預(yù)填充-解碼分解）——是另一個(gè)自然的研究方向。