在 ChatGPT 爆火兩年多的時(shí)間里，大語言模型的上下文窗口長度基準(zhǔn)線被拉升，以此為基礎(chǔ)所構(gòu)建的長 CoT 推理、多 Agent 協(xié)作等類型的高級應(yīng)用也逐漸增多。

隨之而來的是，長文本推理速度被提出更高要求，而基于現(xiàn)有 Transformer 架構(gòu)的模型受限于注意力機(jī)制的二次方復(fù)雜度，難以在較短時(shí)延內(nèi)處理超長文本請求。

針對這一痛點(diǎn)，清華大學(xué) NLP 實(shí)驗(yàn)室聯(lián)手中南大學(xué)、北京郵電大學(xué)以及騰訊微信 AI 實(shí)驗(yàn)室取得了突破，共同提出了 APB 框架 —— 其核心是一個(gè)整合了稀疏注意力機(jī)制的序列并行推理框架，通過整合局部 KV 緩存壓縮方式以及精簡的跨 GPU 通信機(jī)制，解決了長上下文遠(yuǎn)距離語義依賴問題，在無性能損失的前提下大幅度提升超長文本預(yù)填充的效率。

在 128K 文本上，APB 能夠出色地平衡性能與速度，達(dá)到相較于傳統(tǒng) Flash Attention 約 10 倍的加速比，在多種任務(wù)上甚至具有超越完整 Attention 計(jì)算的性能；與英偉達(dá)提出的同為分布式設(shè)定下的 Star Attention 相比，APB 也能達(dá)到 1.6 倍加速比，在性能、速度以及整體計(jì)算量上均優(yōu)于 Star Attention。

• 論文鏈接：https://arxiv.org/pdf/2502.12085
• GitHub 鏈接：https://github.com/thunlp/APB

這一方法主要用于降低處理長文本請求的首 token 響應(yīng)時(shí)間。未來，APB 有潛力運(yùn)用在具有低首 token 響應(yīng)時(shí)間要求的模型服務(wù)上，實(shí)現(xiàn)大模型服務(wù)層對長文本請求的高效處理。

瓶頸：加速長文本預(yù)填充效率

長文本預(yù)填充的效率受到計(jì)算的制約。由于注意力機(jī)制的計(jì)算量與序列長度呈二次方關(guān)系，長文本的計(jì)算通常是計(jì)算瓶頸的。主流加速長文本預(yù)填充的路線有兩種，提升并行度和減少計(jì)算：

• 提升并行度：我們可以將注意力機(jī)制的計(jì)算分布在不同設(shè)備上來提升并行度。當(dāng)一個(gè) GPU 的算力被充分的利用時(shí)，簡單的增加 GPU 的數(shù)量就可以增加有效算力?，F(xiàn)存研究中有各種各樣的并行策略，包括張量并行、模型并行、序列并行等。對于長文本推理優(yōu)化，序列并行有很大的優(yōu)化潛力，因?yàn)樗皇苣Ｐ图軜?gòu)的制約，具有很好的可擴(kuò)展性。
• 減少計(jì)算：另一個(gè)加速長文本預(yù)填充的方式是減少計(jì)算，即使用稀疏注意力。我們可以選擇注意力矩陣中計(jì)算的位置，并不計(jì)算其他位置來減少整體的計(jì)算量。此類方法通常會(huì)帶來一定的性能損失。計(jì)算時(shí)忽略重要的上下文會(huì)導(dǎo)致無法處理某些任務(wù)。

然而，簡單地提升并行度和減少計(jì)算并不能在加速長文本預(yù)填充上取得足夠的效果。若要將二者結(jié)合又具有極大挑戰(zhàn)，這是因?yàn)橄∈枳⒁饬C(jī)制中，決定計(jì)算何處注意力通常需要完整輸入序列的信息。在序列并行框架中，每個(gè) GPU 僅持有部分 KV 緩存，無法在不通過大規(guī)模通信的前提下獲得足夠的全局信息來壓縮注意力的計(jì)算。

針對這一問題，有兩個(gè)先驅(qū)方法：一是英偉達(dá)提出的 Star Attention ，直接去除了序列并行中的所有通信，并只計(jì)算每個(gè) GPU 上局部上下文的注意力，但這樣計(jì)算也導(dǎo)致了很大程度的性能損失；二是卡內(nèi)基梅隆大學(xué)提出的 APE，關(guān)注 RAG 場景下長文本預(yù)填充加速，通過將上下文均勻分割、對注意力進(jìn)行放縮和調(diào)整 softmax 溫度，實(shí)現(xiàn)并行編碼，同樣在需要遠(yuǎn)距離依賴的場景上有一定的性能損失。

區(qū)別于上述方法，APB 通過設(shè)計(jì)面向序列并行場景的低通信稀疏注意力機(jī)制，構(gòu)建了一個(gè)更快、性能更好，且適配通用長文本任務(wù)的長文本加速方法。

APB：面相序列并行框架的稀疏注意力機(jī)制

相較于之前的研究，APB 通過如下方法提出了一種面相序列并行框架的稀疏注意力機(jī)制：

• 增加較小的 Anchor block：Star Attention 中引入的 Anchor  block（輸入序列開始的若干 token）能夠極大恢復(fù)性能，然而其尺寸需要和局部上下文塊一樣大。過大的 anchor block 會(huì)在 FFN 中引入過多的額外開銷。APB 通過減少 anchor  block 的大小，使其和上下文塊的 1/4 或 1/8 一樣大。
• 解決長距離語義依賴問題：先前研究某些任務(wù)上性能下降的原因是它們無法處理長距離語義依賴，后序 GPU 分塊無法看到前序上下文塊中的信息，導(dǎo)致無法處理特定任務(wù)。APB 通過構(gòu)建 passing  block 的方式來解決這一問題。Passing  block 由前面設(shè)備上的重要 KV 對組成。每個(gè)上下文塊先被壓縮，然后將被壓縮的上下文塊通信到后續(xù) GPU 上來構(gòu)建 passing block。
• 壓縮上下文塊：在不進(jìn)行大規(guī)模通信的前提下，每個(gè)設(shè)備只對自己持有的上下文有訪問權(quán)限。因此，現(xiàn)存的 KV Cache 壓縮算法（例如 H2O 和 SnapKV）不適用于這一場景，因?yàn)樗鼈円蕾嚾蛄械男畔ⅰＨ欢?，該特點(diǎn)與 Locret 一致，KV Cache 重要性分?jǐn)?shù)僅僅與對應(yīng) KV 對的 Q, K, V 相關(guān)。APB 使用 Locret 中引入的 retaining heads 作為上下文壓縮器。
• 查詢感知的上下文壓縮：APB 在 anchor  block 的開頭嵌入查詢。當(dāng)預(yù)填充結(jié)束時(shí)，這些查詢可以隨著 anchor  block 一同被丟棄，不會(huì)影響整體計(jì)算的同時(shí)還能讓上下文壓縮器看到查詢的內(nèi)容。通過這種方式，保留頭能夠更精準(zhǔn)地識別出查詢相關(guān)的 KV 對，并通過通信機(jī)制傳給后續(xù)設(shè)備。

以此機(jī)制為基礎(chǔ)，APB 的推理過程如下：

• 上下文分割：長文本被均勻的分到每個(gè)設(shè)備上，開頭拼接一個(gè) anchor block，其中包含了查詢問題。
• 上下文壓縮：我們用 Locret 引入的保留頭來壓縮 KV Cache。
• 通信：我們對壓縮過的 KV Cache 施加一個(gè) AllGather 算子。每個(gè)設(shè)備會(huì)拿到前序設(shè)備傳來的壓縮緩存，并構(gòu)建 passing block。
• 計(jì)算：我們使用一個(gè)特殊的 Flash Attention Kernel 來實(shí)現(xiàn)這個(gè)特殊的注意力機(jī)制。我們更改了注意力掩碼的形狀。Passing block 在注意力計(jì)算結(jié)束后就被刪除，不參與后續(xù)計(jì)算。

APB 實(shí)現(xiàn)更快、性能更好的長文本推理

團(tuán)隊(duì)使用 Llama-3.1-8B-instruct, Qwen-2.5-14B-instruct 以及 Yi-34B-200K 模型在 InfiniteBench 和 RULER 上進(jìn)行了測試，測量任務(wù)性能（%）以及處理完整長文本請求的推理速度（tok /s）。研究人員選擇 Flash Attention, Ulysses, Ring Attention, MInference 以及 Star Attention 作為基線算法，實(shí)驗(yàn)結(jié)果如下：

從上圖可見，F(xiàn)lash Attention 作為無序列并行的精準(zhǔn)注意力算法，具有較好的任務(wù)性能，但推理速度最慢；Ring Attention 和 Ulysses 作為序列并行的精準(zhǔn)注意力算法，通過增加并行度的方式提升了推理速度；MInference 是一種無序列并行的稀疏注意力機(jī)制，表現(xiàn)出了一定的性能損失；Star Attention 作為序列并行與稀疏注意力結(jié)合的最初嘗試，具有較好的推理速度，然而表現(xiàn)出了顯著的性能下降。

相較于基線算法，APB 在多種模型和任務(wù)上表現(xiàn)出更優(yōu)的性能和更快的推理速度。這意味著，APB 方法能夠?qū)崿F(xiàn)最好的任務(wù)性能與推理速度的均衡。

除此之外，研究人員在不同長度的數(shù)據(jù)上測量了 APB 與基線算法的性能、速度，并給出了整體計(jì)算量，結(jié)果如下：

可以從上圖中看到，APB 在各種輸入長度下均表現(xiàn)出更優(yōu)的任務(wù)性能與推理速度。速度優(yōu)勢隨著輸入序列變長而變得更加明顯。APB 相較于其他方法更快的原因是它需要更少的計(jì)算，且計(jì)算量差異隨著序列變長而加大。

并且，研究人員還對 APB 及基線算法進(jìn)行了預(yù)填充時(shí)間拆解分析，發(fā)現(xiàn)序列并行可以大幅度縮減注意力和 FFN 時(shí)間。

通過稀疏注意力機(jī)制，APB 能進(jìn)一步縮減注意力時(shí)間。Star Attention 由于使用了過大的 anchor block，其 FFN 的額外開銷十分明顯，而 APB 由于使用了 passing block 來傳遞遠(yuǎn)距離語義依賴，能夠大幅度縮小 anchor block 大小，從而降低 FFN 處的額外開銷。

APB 支持具有卓越的兼容性，能適應(yīng)不同分布式設(shè)定（顯卡數(shù)目）以及不同模型大小，在多種模型和分布式設(shè)定下均在性能與推理速度上取得了優(yōu)異的效果。

核心作者簡介

黃宇翔，清華大學(xué)四年級本科生，THUNLP 實(shí)驗(yàn)室 2025 年準(zhǔn)入學(xué)博士生，導(dǎo)師為劉知遠(yuǎn)副教授。曾參與過 MiniCPM、模型高效微調(diào)、以及投機(jī)采樣研究項(xiàng)目。主要研究興趣集中在構(gòu)建高效的大模型推理系統(tǒng)，關(guān)注模型壓縮、投機(jī)采樣、長文本稀疏等推理加速技術(shù)。

李明業(yè)，中南大學(xué)三年級本科生，2024 年 6 月份加入 THUNLP 實(shí)驗(yàn)室實(shí)習(xí)，參與過投機(jī)采樣研究項(xiàng)目。主要研究興趣集中在大模型的推理加速，例如投機(jī)采樣以及長文本推理加速等。