在當(dāng)前大模型推理愈發(fā)復(fù)雜的時(shí)代，如何快速、高效地產(chǎn)生超長(zhǎng)文本，成為了模型部署與優(yōu)化中的一大核心挑戰(zhàn)。隨著 GPT-o3, DeepSeek R1 等具備 「超級(jí)上下文窗口」 能力的大模型持續(xù)刷新業(yè)界記錄，百萬(wàn)甚至千萬(wàn) Token 級(jí)別的推理任務(wù)已從研究話題邁入現(xiàn)實(shí)場(chǎng)景。然而，生成這些超長(zhǎng)文本的背后，卻隱藏著令人咋舌的計(jì)算成本 —— 長(zhǎng)時(shí)間的等待、巨大的內(nèi)存負(fù)擔(dān)以及偶爾重復(fù)乏味的輸出，嚴(yán)重制約了這些模型的真正潛力。

面對(duì)這一挑戰(zhàn)，BIGAI NLCo 團(tuán)隊(duì)提出了一項(xiàng)全新的推理加速框架 —— TokenSwift，該工作已成功被 ICML 2025 正式接收！在這項(xiàng)研究中提出了一套可插拔、無(wú)損、高效的生成加速策略，專為 100K Token 級(jí)別的長(zhǎng)文本推理而設(shè)計(jì)。在保持原始模型輸出一致性的前提下，加速比達(dá)到 3 倍以上，極大提升了推理效率。

• 論文標(biāo)題：TokenSwift: Lossless Acceleration of Ultra Long Sequence Generation
• Arxiv: https://arxiv.org/abs/2502.18890
• Github: https://github.com/bigai-nlco/TokenSwift
• Blog: https://bigai-nlco.github.io/TokenSwift/

重新定義超長(zhǎng)生成：為什么傳統(tǒng)方法「慢」？

為了更好地理解 TokenSwift 的意義，我們先看一下目前主流大模型（如 LLaMA、Qwen 等）在長(zhǎng)文本生成中的瓶頸所在。

盡管這些模型具備了強(qiáng)大的生成長(zhǎng)上下文的能力，但大多數(shù)依然采用傳統(tǒng)的自回歸（Autoregressive）生成方式：每次僅生成一個(gè)新的 Token，并以其作為輸入接著生成下一個(gè)。這種方式本身在短文本生成中問(wèn)題不大，但當(dāng)序列長(zhǎng)度擴(kuò)展至 10 萬(wàn)甚至更多時(shí)，性能就會(huì)急劇下降。

主要原因有三：

• 模型重復(fù)重載：每生成一個(gè) Token，都會(huì)觸發(fā)一次完整的前向推理過(guò)程；在多進(jìn)程或流水線執(zhí)行時(shí)，模型需要不斷讀取參數(shù)，造成 I/O 瓶頸。
• KV 緩存無(wú)限膨脹：Transformer 架構(gòu)要求保留所有歷史 Token 的 Key/Value 信息，用于后續(xù) Token 的注意力計(jì)算。隨著生成進(jìn)程推進(jìn)，KV 緩存占用不斷增加，導(dǎo)致計(jì)算與內(nèi)存開(kāi)銷雪上加霜。
• 語(yǔ)義重復(fù)堆疊：生成越長(zhǎng)，模型越容易陷入句式與主題的復(fù)讀循環(huán)，降低輸出多樣性與用戶體驗(yàn)。

尤其在當(dāng)前日益增長(zhǎng)的多輪對(duì)話、大模型代理（Agent）、逐步推理等任務(wù)中，一個(gè) Query 可能會(huì)觸發(fā)幾千甚至上萬(wàn)的推理 Token 輸出。傳統(tǒng)自回歸的效率顯然已經(jīng)難以滿足需求。

TokenSwift：擁抱并行的超長(zhǎng)推理時(shí)代

TokenSwift 的提出，正是為了解決上述超長(zhǎng)生成中的三大瓶頸。它通過(guò)一個(gè)極為輕量且高效的框架，對(duì)傳統(tǒng)自回歸推理進(jìn)行了 「重構(gòu)」，提出了以 「多 Token 草擬 + 并行驗(yàn)證 + 動(dòng)態(tài)緩存更新」 為核心的全新機(jī)制。

讓我們來(lái)逐步拆解 TokenSwift 的關(guān)鍵技術(shù)理念：

 多 Token 并行草擬：告別一次一 Token 的低效時(shí)代

在 TokenSwift 中，不再堅(jiān)持 「一步一 Token」 的生成模式，而是通過(guò)對(duì)已有模型的最小化修改（添加極少量的線性層），實(shí)現(xiàn)了 「一次性草擬多個(gè) Token」。這意味著，模型每一次前向傳播，都可以并行生成 γ 個(gè)候選 Token，大幅降低模型重載頻率，顯著節(jié)省 I/O 時(shí)間。

更關(guān)鍵的是，草擬階段并非 「胡亂猜測(cè)」。引入了上下文引導(dǎo)機(jī)制，使草擬結(jié)果具備較高的語(yǔ)義相關(guān)性與語(yǔ)法一致性，隨后通過(guò)結(jié)構(gòu)化的驗(yàn)證機(jī)制確保其與標(biāo)準(zhǔn) AR 路徑一致。

 n-gram 啟發(fā)式補(bǔ)全：巧用歷史片段，精確草擬結(jié)構(gòu)

為了避免粗略草擬帶來(lái)的語(yǔ)義偏離，TokenSwift 設(shè)計(jì)了基于歷史生成內(nèi)容的 n-gram 片段緩存機(jī)制。會(huì)定期保存頻率較高的 n-gram 序列，并在草擬新 Token 時(shí)，借助這些高頻片段進(jìn)行 「自動(dòng)補(bǔ)全」。

此外，TokenSwift 還引入了 「隨機(jī)篩選器」，在多個(gè) n-gram 片段中選擇語(yǔ)義最優(yōu)的一組進(jìn)行結(jié)構(gòu)構(gòu)建。這不僅提升了草擬的準(zhǔn)確性，也保證了后續(xù)驗(yàn)證的成功率。

 樹結(jié)構(gòu)驗(yàn)證機(jī)制：確保與 AR 一致，輸出 「無(wú)損」 保障

有了草擬機(jī)制，還需要驗(yàn)證其 「正當(dāng)性」。TokenSwift 的另一大亮點(diǎn)在于提出了樹結(jié)構(gòu)的并行驗(yàn)證模塊，通過(guò)構(gòu)建多個(gè)驗(yàn)證路徑，對(duì)草擬的多個(gè) Token 并行進(jìn)行預(yù)測(cè)評(píng)分判斷，并篩選出與標(biāo)準(zhǔn) AR 路徑一致的候選。

換句話說(shuō)，TokenSwift 不僅快，而且不犧牲任何輸出質(zhì)量，生成內(nèi)容與原始模型保持一致，是一套真正 「無(wú)損」 的加速方案。

動(dòng)態(tài) KV 管理 + 重復(fù)懲罰：越長(zhǎng)越快，越長(zhǎng)越優(yōu)

為了解決 KV 緩存膨脹的問(wèn)題，TokenSwift 實(shí)現(xiàn)了動(dòng)態(tài)的 KV 裁剪機(jī)制。模型會(huì)根據(jù) Token 重要度與時(shí)間衰減策略，對(duì) KV 對(duì)進(jìn)行 「主動(dòng)淘汰」，在保證上下文保留質(zhì)量的同時(shí)，顯著減輕緩存負(fù)擔(dān)。

與此同時(shí)，為了緩解長(zhǎng)文本生成過(guò)程中的重復(fù)問(wèn)題，設(shè)計(jì)了重復(fù)懲罰機(jī)制，在生成過(guò)程中動(dòng)態(tài)降低重復(fù) n-gram 的概率，確保最終輸出具備更高的多樣性和可讀性。

實(shí)驗(yàn)評(píng)估：全面剖析 TokenSwift 的性能與質(zhì)量

在多個(gè)主流模型上，包括 YaRN-LLaMA2-7b-128k、LLaMA3.1-8b、Qwen2.5-1.5B, 7B, 14B 等，進(jìn)行了大規(guī)模實(shí)驗(yàn)，序列長(zhǎng)度涵蓋從 20K 到 100K，TokenSwift 表現(xiàn)均極其亮眼：

• 加速比普遍在 3 倍以上
• 生成質(zhì)量與原模型一致
• Distinct-n 指標(biāo)顯著優(yōu)于原始 AR 路徑

更令人振奮的是，隨著序列越長(zhǎng)，TokenSwift 的加速效果越顯著。在 100K Token 生成任務(wù)中，LLaMA3.1-8B 的生成時(shí)間從近 5 小時(shí)縮短至 1.5 小時(shí)，極大降低了實(shí)際使用成本。

Token 重用的決定性作用

我們對(duì)比了禁用（k=0）與啟用狀態(tài)下的接受率和加速比。結(jié)果顯示，未啟用重用時(shí)，隨著生成長(zhǎng)度增長(zhǎng)，接受率和加速比均出現(xiàn)明顯下滑；而在 k=20 時(shí)，接受率不僅維持在 70–90% 的高水平，還隨序列越長(zhǎng)而略有提升，加速比亦從～2.1× 提升至～3.1×，凸顯 Token 重用在減少模型重載和保持驗(yàn)證高效中的關(guān)鍵作用 。

動(dòng)態(tài) KV 更新的巧妙平衡

針對(duì) KV 緩存管理，我們進(jìn)一步實(shí)驗(yàn)了「全量緩存」、「僅預(yù)填后一次更新」與 TokenSwift 中的「動(dòng)態(tài)部分緩存」三種策略。在關(guān)閉 Token 重用的前提下，「部分緩存」因接受率低而加速有限，「全量緩存」雖保證了較高接受率，卻因管理成本抵消了速度收益；而 TokenSwift 的動(dòng)態(tài)更新策略恰好在接受率與計(jì)算開(kāi)銷之間取得平衡，使得在 100K Token 任務(wù)上依然能保持近 2× 以上的加速效果 。

上下文懲罰對(duì)多樣性的提升

為了抑制超長(zhǎng)生成中的機(jī)械復(fù)讀，TokenSwift 采用了僅對(duì)最近 W Token 應(yīng)用懲罰值 θ 的「上下文懲罰」機(jī)制。在與 top-p、η-sampling、min-p 等多種采樣方法結(jié)合的實(shí)驗(yàn)中，未啟用懲罰時(shí) Distinct-n 平均值僅約 0.12；引入 θ=1.2、W=1024 后，平均多樣性飆升至 0.43–0.69，同時(shí)加速比僅輕微下降～2–8% 左右，證明了該機(jī)制在保持質(zhì)量一致性與提升文本多樣性上的高效性 。

小結(jié)

TokenSwift 擺脫了傳統(tǒng)自回歸生成的性能枷鎖，讓我們?cè)诿鎸?duì)超長(zhǎng)文本推理任務(wù)時(shí)，不再被時(shí)間與內(nèi)存所拖累。

它并非一個(gè) 「另起爐灶」 的新模型，而是一種可直接嵌入現(xiàn)有主流模型（LLaMA、Qwen 等）的通用加速策略，具備極強(qiáng)的兼容性與部署便利性。

更重要的是，TokenSwift 對(duì)推理質(zhì)量的 「無(wú)損」 保證，讓我們?cè)谙硎芩俣忍嵘耐瑫r(shí)，不犧牲任何精度與語(yǔ)義一致性。我們相信，這一框架將為未來(lái)多輪推理、代碼生成、Agent 計(jì)劃編排等長(zhǎng)文本場(chǎng)景提供堅(jiān)實(shí)的技術(shù)支撐。

歡迎閱讀 ICML 2025 原論文，了解更多技術(shù)細(xì)節(jié)！