聲明:本文來(lái)自于（ID:ykqsd.com）授權(quán)轉(zhuǎn)載發(fā)布。

機(jī)器之心由

機(jī)器之心編輯

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

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

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

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

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

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

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

針對(duì)這一問(wèn)題，現(xiàn)有研究提出了兩種關(guān)鍵方法：英偉達(dá)提出的 Star Attention 通過(guò)去除序列并行中的所有通信，僅計(jì)算 GPU 局部區(qū)域的注意力，但這種簡(jiǎn)化策略導(dǎo)致顯著的計(jì)算效率損失。另一項(xiàng)研究由卡內(nèi)基梅隆大學(xué)提出，即 APE 方法，專注于 RAG 桅ma場(chǎng)景下的長(zhǎng)文本預(yù)填充加速，通過(guò)均勻劃分上下文、縮放注意力計(jì)算并調(diào)整 Softmax 溫度，實(shí)現(xiàn)了并行編碼。然而，該方法在處理需要遠(yuǎn)距離依賴的文本時(shí)仍存在一定效率損失。

與上述方法相比，APB 通過(guò)創(chuàng)新設(shè)計(jì)面向序列并行的低通信稀疏注意力機(jī)制，成功構(gòu)建了一個(gè)在計(jì)算效率和性能表現(xiàn)上均更優(yōu)的長(zhǎng)文本加速框架，該框架能夠很好地適應(yīng)通用長(zhǎng)文本任務(wù)。

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

相較于現(xiàn)有研究，APB 通過(guò)以下方式提出了一種創(chuàng)新的面向序列并行框架的稀疏注意力機(jī)制：

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

該方法顯著提升了長(zhǎng)文本推理速度的同時(shí)保持了較高的計(jì)算效率。

研究團(tuán)隊(duì)采用了 Llama-3.1-8B-instruct、Qwen-2.5-14B-instruct 以及 Yi-34B-200K 模型，在 InfiniteBench 和 RULER 平臺(tái)上進(jìn)行了系統(tǒng)評(píng)估，全面考察了任務(wù)性能（%）和推理速度（tok/s）指標(biāo)。

從實(shí)驗(yàn)結(jié)果來(lái)看，F(xiàn)lash Attention 作為無(wú)序列并行的精準(zhǔn)注意力算法，在任務(wù)性能方面表現(xiàn)優(yōu)異，但推理速度最慢；而 Ring Attention 和 Ulysses 作為序列并行的精準(zhǔn)注意力算法，通過(guò)優(yōu)化并行度顯著提升了推理速度；MInference 作為一種無(wú)序列并行的稀疏注意力機(jī)制，在性能上也存在一定局限性；Star Attention 則作為序列并行與稀疏注意力結(jié)合的初始嘗試，在推理速度上表現(xiàn)尚可，但計(jì)算效率明顯下降。

相較于上述基線算法，APB 在多種模型和任務(wù)場(chǎng)景中展現(xiàn)出更優(yōu)的性能表現(xiàn)和更快的推理速度。這表明，APB 方法實(shí)現(xiàn)了計(jì)算效率與推理速度的最佳平衡。

此外，研究團(tuán)隊(duì)在不同數(shù)據(jù)長(zhǎng)度的場(chǎng)景下對(duì) APB 與基線算法的性能表現(xiàn)和推理速度進(jìn)行了全面評(píng)估，并提供了整體計(jì)算量的數(shù)據(jù)支持，結(jié)果表明 APB 在各類場(chǎng)景中均展現(xiàn)出顯著的優(yōu)勢(shì)。

From the figure, it can be observed that APB demonstrates superior task performance and inference speed across various input lengths. The speed advantage becomes more pronounced as the input sequence lengthens. APB achieves this improvement because it requires fewer computations, with the computational efficiency gap widening as the sequence becomes longer.

Additionally, researchers conducted a pre-filled time decomposition analysis on both APB and baseline algorithms and found that leveraging sequence parallelism can significantly reduce the time spent on attention and feed-forward (FFN) operations.

Furthermore, the sparse attention mechanism in APB allows for further reduction in attention time. The Star Attention mechanism, however, suffers from significant extra overhead due to its large anchor block, while APB mitigates this issue by utilizing passing blocks to transmit long-range semantic dependencies, thereby enabling a much smaller anchor block size and reducing the extra overhead at the FFN stage.

APB exhibits excellent compatibility, being adaptable to different distributed settings (e.g., number of GPUs) and various model sizes. It achieves impressive performance and inference speed across multiple model configurations and distributed setups.

**Biography of Core Authors**

- **Yuxiang Huang**: An undergraduate student in the fourth year at Tsinghua University, currently a Ph.D. candidate candidate in the 2025 intake at the THUNLP Lab, advised by Prof. Liu Zhiyuan. He has previously participated in projects such as MiniCPM, model efficient fine-tuning, and speculative sampling. His primary research interests revolve around constructing efficient large-model inference systems, with a focus on model compression, speculative sampling, and fast inference techniques for long-text sparse retrieval.

- **Mingyue Li**: An undergraduate student in their third year at Central South University, who joined the THUNLP Lab as an intern in June 2024. He has been involved in the speculative sampling project. His main research interests include accelerating large-model inference, particularly techniques such as speculative sampling and acceleration methods for long-text retrieval.