雖然大型語(yǔ)言模型（LLM）的性能表現(xiàn)足夠驚艷，但每次接收用戶請(qǐng)求時(shí)都需要耗費(fèi)大量顯存和計(jì)算資源，一旦請(qǐng)求數(shù)量超出預(yù)期，就極有可能面臨ChatGPT剛發(fā)布時(shí)的宕機(jī)、排隊(duì)、高延遲等窘境。

想要打造一個(gè)高吞吐量的LLM服務(wù)，就需要模型在一個(gè)批次內(nèi)處理盡可能多的請(qǐng)求，不過(guò)現(xiàn)有的系統(tǒng)大多在每次處理請(qǐng)求時(shí)申請(qǐng)大量的key-value（KV）緩存，如果管理效率不高，大量?jī)?nèi)存都會(huì)在碎片和冗余復(fù)制中被浪費(fèi)掉，限制了batch size的增長(zhǎng)。

最近，來(lái)自加州大學(xué)伯克利分校、斯坦福大學(xué)、加州大學(xué)圣迭戈分校的研究人員基于操作系統(tǒng)中經(jīng)典的虛擬內(nèi)存和分頁(yè)技術(shù)，提出了一個(gè)新的注意力算法PagedAttention，并打造了一個(gè)LLM服務(wù)系統(tǒng)vLLM

論文鏈接：https://arxiv.org/pdf/2309.06180.pdf

開(kāi)源鏈接：https://github.com/vllm-project/vllm

vLLM在KV緩存上實(shí)現(xiàn)了幾乎零浪費(fèi)，并且可以在「請(qǐng)求內(nèi)部」和「請(qǐng)求之間」靈活共享KV高速緩存，進(jìn)一步減少了內(nèi)存的使用量。

評(píng)估結(jié)果表明，vLLM可以將常用的LLM吞吐量提高了2-4倍 ，在延遲水平上與最先進(jìn)的系統(tǒng)（如FasterTransformer和Orca）相當(dāng)，并且在更長(zhǎng)序列、更大模型和更復(fù)雜的解碼算法時(shí)，提升更明顯。

PagedAttention

為了解決注意力機(jī)制的內(nèi)存管理問(wèn)題，研究人員開(kāi)發(fā)了一種全新的注意力算法PagedAttention，并構(gòu)建了一個(gè)LLM服務(wù)引擎vLLM，采用集中式調(diào)度器來(lái)協(xié)調(diào)分布式GPU工作線程的執(zhí)行。

1. 算法

受操作系統(tǒng)中分頁(yè)（paging）算法啟發(fā)，PagedAttention將序列中KV緩存劃分為KV塊，其中每個(gè)塊包含固定數(shù)量tokens的鍵（K）和值（V）向量，從而將注意力計(jì)算轉(zhuǎn)換為塊級(jí)運(yùn)算：

在注意力計(jì)算期間，PagedAttention內(nèi)核分別識(shí)別和獲取不同的KV塊，比如下面的例子中，鍵和值向量分布在三個(gè)塊上，并且三個(gè)塊在物理內(nèi)存上是不連續(xù)的，然后將查詢向量與塊中的鍵向量相乘得到部分注意力得分，再乘以塊中的值向量得到最終注意力輸出。

這種設(shè)計(jì)使得KV塊存儲(chǔ)在非連續(xù)物理內(nèi)存中，從而讓vLLM中的分頁(yè)內(nèi)存管理更加靈活。

2. KV緩存管理器

操作系統(tǒng)會(huì)將內(nèi)存劃分為多個(gè)固定大小的頁(yè)，并將用戶程序的邏輯頁(yè)映射到物理頁(yè)，連續(xù)的邏輯頁(yè)可以對(duì)應(yīng)于非連續(xù)的物理內(nèi)存頁(yè)，所以用戶在訪問(wèn)內(nèi)存時(shí)看起來(lái)就像連續(xù)的一樣。

此外，物理內(nèi)存空間不需要提前完全預(yù)留，使操作系統(tǒng)能夠根據(jù)需求動(dòng)態(tài)分配物理頁(yè)。

通過(guò)PageAttention劃分出的KV塊，vLLM利用虛擬內(nèi)存機(jī)制將KV緩存表示為一系列邏輯KV塊，并在生成新token及KV緩存時(shí)，從左到右進(jìn)行填充；最后一個(gè)KV塊的未填充位置預(yù)留給后續(xù)生成操作。

KV塊管理器還負(fù)責(zé)維護(hù)塊表（block table），即每個(gè)請(qǐng)求的邏輯和物理KV塊之間的映射。

將邏輯和物理KV塊分離使得vLLM能夠動(dòng)態(tài)地增長(zhǎng)KV高速緩存存儲(chǔ)器，而無(wú)需預(yù)先將其保留給所有位置，消除了現(xiàn)有系統(tǒng)中的大多數(shù)內(nèi)存浪費(fèi)。

3. 解碼

從下面的例子中可以看出vLLM如何在單個(gè)輸入序列的解碼過(guò)程中執(zhí)行PagedAttention并管理內(nèi)存。

① 與操作系統(tǒng)的虛擬內(nèi)存一樣，vLLM最初不需要為最大可能生成的序列長(zhǎng)度保留內(nèi)存，只保留必要的KV塊，以容納在即時(shí)計(jì)算期間生成的KV緩存。

提示詞中包含7個(gè)tokens，所以vLLM將前兩個(gè)邏輯KV塊（0和1）映射到2個(gè)物理KV塊（7和1）；在預(yù)填充（prefill）步驟中，vLLM使用自注意算法生成提示和首個(gè)輸出token的KV緩存；然后將前4個(gè)token的KV緩存存儲(chǔ)在邏輯塊0中，后面3個(gè)token存儲(chǔ)在邏輯塊1中；剩余的slot被保留用于后續(xù)自回歸生成。

② 在首個(gè)自回歸解碼步驟中，vLLM在物理塊7和1上使用PagedAttention算法生成新token

由于最后一個(gè)邏輯塊中仍有一個(gè)slot可用，所以將新生成的KV緩存存儲(chǔ)在該slot，更新塊表的#filled記錄。

③ 在第二次解碼步驟中，當(dāng)最后一個(gè)邏輯塊已滿時(shí)，vLLM將新生成的KV緩存存儲(chǔ)在新的邏輯塊中，為其分配一個(gè)新的物理塊（物理塊3），并映射存儲(chǔ)在塊表中。

在LLM的計(jì)算過(guò)程中，vLLM使用PagedAttention內(nèi)核訪問(wèn)以前以邏輯KV塊形式存儲(chǔ)的KV緩存，并將新生成的KV緩存保存到物理KV塊中。

在一個(gè)KV塊（塊大小>1）中存儲(chǔ)多個(gè)token使PagedAttention內(nèi)核能夠跨更多位置并行處理KV緩存，從而提高硬件利用率并減少延遲，但較大的塊大小也會(huì)增加內(nèi)存碎片。

隨著生成越來(lái)越多的token及其KV緩存，vLLM會(huì)動(dòng)態(tài)地將新的物理塊分配給邏輯塊，從左到右地填充所有塊，并且只有當(dāng)所有先前的塊都滿時(shí)才分配新的物理塊，即請(qǐng)求的所有內(nèi)存浪費(fèi)限制在一個(gè)塊內(nèi)，可以有效地利用所有內(nèi)存，從而允許更多的請(qǐng)求放入內(nèi)存進(jìn)行批處理，提高了吞吐量；一旦請(qǐng)求完成生成，就可以釋放其KV塊來(lái)存儲(chǔ)其他請(qǐng)求的KV緩存。

4. 通用解碼

除了貪婪解碼和采樣，支持單個(gè)用戶提示輸入生成單個(gè)輸出序列等基本場(chǎng)景外，該算法還可以支持更復(fù)雜的解碼場(chǎng)景，如并行采樣（Parallel sampling）、集束搜索（Beam Search）、共享前綴等。

5. 調(diào)度和搶占（Scheduling and Preemption）

當(dāng)請(qǐng)求流量超過(guò)系統(tǒng)容量時(shí)，vLLM必須對(duì)請(qǐng)求子集進(jìn)行優(yōu)先級(jí)排序，具體采用「先來(lái)先服務(wù)」（FCFS）的調(diào)度策略，可以確保公平性并防止饑餓。

不過(guò)LLM的輸入提示在長(zhǎng)度上可能變化很大，并且輸出長(zhǎng)度是先驗(yàn)未知的，具體取決于輸入提示和模型；隨著請(qǐng)求及其輸出數(shù)量的增長(zhǎng)，vLLM可能會(huì)耗盡GPU的物理塊來(lái)存儲(chǔ)新生成的KV緩存。

交換（Swapping）是大多數(shù)虛擬內(nèi)存算法使用的經(jīng)典技術(shù)，將被釋放的頁(yè)復(fù)制到磁盤(pán)上的交換空間。

除了GPU塊分配器之外，vLLM還包括CPU塊分配器，以管理交換到CPU RAM的物理塊；當(dāng)vLLM耗盡新令牌的空閑物理塊時(shí)，會(huì)選擇一組序列來(lái)釋放KV緩存并將其傳輸?shù)紺PU。

在這種設(shè)計(jì)中，交換到CPU RAM的塊數(shù)永遠(yuǎn)不會(huì)超過(guò)GPU RAM中的物理塊總數(shù)，因此CPU RAM上的交換空間受到分配給KV緩存的GPU內(nèi)存的限制。

重新計(jì)算（Recomputation），當(dāng)被搶占的序列被重新調(diào)度時(shí)，可以簡(jiǎn)單地重新計(jì)算KV緩存，其延遲可以顯著低于原始延遲，因?yàn)榻獯a時(shí)生成的token可以與原始用戶提示連接起來(lái)作為新的提示，所有位置的KV緩存可以在一次提示階段迭代中生成。

交換和重計(jì)算的性能取決于CPU、RAM和GPU內(nèi)存之間的帶寬以及GPU的計(jì)算能力。

6. 分布式執(zhí)行（Distributed Execution）

vLLM支持Megatron-LM風(fēng)格的張量模型并行策略，遵循SPMD（單程序多數(shù)據(jù)）執(zhí)行調(diào)度，其中線性層被劃分以執(zhí)行逐塊矩陣乘法，并且GPU通過(guò)allreduce操作不斷同步中間結(jié)果。

具體來(lái)說(shuō)，注意算子在注意頭維度上被分割，每個(gè)SPMD過(guò)程負(fù)責(zé)多頭注意中的注意頭子集，不過(guò)每個(gè)模型分片仍然處理相同的輸入token集合，即在同一位置需要KV緩存。

不同的GPU worker共享管理器，以及從邏輯塊到物理塊的映射，使用調(diào)度程序?yàn)槊總€(gè)輸入請(qǐng)求提供的物理塊來(lái)執(zhí)行模型；盡管每個(gè)GPU工作線程具有相同的物理塊id，但是一個(gè)工作線程僅為其相應(yīng)的注意頭存儲(chǔ)KV緩存的一部分。

在每一步中，調(diào)度程序首先為批處理中的每個(gè)請(qǐng)求準(zhǔn)備帶有輸入token id的消息，以及每個(gè)請(qǐng)求的塊表；

然后調(diào)度程序?qū)⒃摽刂葡V播給GPU worker，使用輸入token id執(zhí)行模型；在注意力層，根據(jù)控制消息中的塊表讀取KV緩存；在執(zhí)行過(guò)程中，將中間結(jié)果與all-reduce通信原語(yǔ)同步，而無(wú)需調(diào)度程序的協(xié)調(diào)。

最后，GPU worker將該迭代的采樣token發(fā)送回調(diào)度器。

評(píng)估結(jié)果

基礎(chǔ)采樣

在ShareGPT數(shù)據(jù)集上，隨著請(qǐng)求速率的增加，延遲最初緩慢增加，之后會(huì)突然激增，可能是因?yàn)楫?dāng)請(qǐng)求速率超過(guò)服務(wù)系統(tǒng)的容量時(shí)，導(dǎo)致隊(duì)列長(zhǎng)度無(wú)限增長(zhǎng)。

vLLM可以承受比Orca高1.7倍-2.7倍的請(qǐng)求速率，比Orca（Max）高2.7倍-8倍的請(qǐng)求速率，同時(shí)保持相似的延遲，因?yàn)镻agedAttention可以有效地管理內(nèi)存使用，從而能夠比Orca在一個(gè)批次內(nèi)處理更多的請(qǐng)求。

對(duì)于OPT-13B模型，vLLM同時(shí)處理的請(qǐng)求比Orca多2.2倍，比Orca（Max）多4.3倍。

與FasterTransformer相比，vLLM實(shí)現(xiàn)高達(dá)22倍的請(qǐng)求速率，可能是因?yàn)闆](méi)有利用細(xì)粒度的調(diào)度機(jī)制，并且與Orca（Max）一樣在內(nèi)存管理方面很低效。

多序列

在并行采樣中，請(qǐng)求中的所有并行序列可以共享提示符的KV緩存，隨著采樣序列數(shù)量的增加，vLLM實(shí)現(xiàn)了比Orca基線更大的提升。

由于集束搜索中共享內(nèi)容更多，vLLM展示出了更大的性能優(yōu)勢(shì)。

在OPT-13B和Alpaca數(shù)據(jù)集上，vLLM相對(duì)于Orca（Oracle）的改進(jìn)從基本采樣的1.3倍增加到寬度為6的集束搜索的2.3倍。

通過(guò)計(jì)算共享保存的塊數(shù)除以未共享的總塊數(shù)計(jì)算的存儲(chǔ)器節(jié)省量，結(jié)果顯示并行采樣節(jié)省了6.1%-9.8%的內(nèi)存，集束搜索節(jié)省了37.6%-55.2%的內(nèi)存。

在使用ShareGPT數(shù)據(jù)集的相同實(shí)驗(yàn)中，可以看到并行采樣節(jié)省了16.2%-30.5%的內(nèi)存，集束搜索節(jié)省了44.3%-66.3%的內(nèi)存。