作者 | 趙驍勇 阿里巴巴智能引擎事業部
審校 | 劉侃,Kitty
RTP-LLM 是阿里巴巴大模型預測團隊開發的高性能 LLM 推理加速引擎。它在阿里巴巴集團內廣泛應用,支撐著淘寶、天貓、高德、餓了么等核心業務部門的大模型推理需求。在 RTP-LLM 上,我們實現了一個通用的投機采樣框架,支持多種投機采樣方法,能夠幫助業務有效降低推理延遲以及提升吞吐。本文重點介紹其中的一種投機編輯技術,該技術已被應用于阿里內部智能化產品,該技術在這一特定業務場景下可實現每秒推理 1000 個 token 的能力。
背 景
相比于其他 AI 應用,LLM 的推理階段最大的特點是它是一個自回歸過程,絕大部分推理框架的優化工作都是圍繞這個自回歸過程展開的。
具體來講,在下面這個例子中,用戶的輸入是“你是誰”,LLM 的返回是“我是 Qwen ”,假設這些“你是誰”這些字符串在分詞后對應的 token id 時 5121,5122,5123,“我是 Qwen ”對應的 token id 時 5124,5122,5125 和 65536。在這個過程中,處理用戶輸入的 prompt 那個階段稱為 Prefill 階段,它會對用戶輸入的所有 prompt(包含 5121,5122,5123) 進行并行計算,將中間計算產生的 KV Cache 保存起來,并且輸出第一個 token 5124。之后生成新 token id(5122, 5125, 65536) 的階段稱為 Decode 階段,與 Prefill 階段相比,Decode 階段每次只會對當前的 token 進行計算,比如在第一次 Decode 時,模型的輸入只有 5124 這個 token,我們會將中間計算得到的 KV Cache 保存下來,并產生下一個 token 5122,之后這個過程會自回歸的進行下去,直到遇到結束 token 65536。通常情況下,對于一個推理請求,處理 prompt 的 Prefill 只會進行一次,而生成新 token 的 Decode 會自回歸的進行非常多次,取決于用戶的輸出有多長。
LLM 的推理離不開 GPU,對于任何一張 GPU,有四個硬件特性非常關鍵,它們分別是算力,顯存帶寬,通信和顯存。以 A100 80GB SXM 為例,它的 FP16 Tensor Core 的算力 312 TFLOPS,顯存帶寬為 2,039 GB/s,NVLINK 卡間通信帶寬為 600 GB/s,PCI-E 卡間通信 64 GB/s,顯存為 80GB。不管是推理還是訓練系統,最大化算力利用率都是至關重要的優化目標,這通常被稱為提高 MFU(Model FLOPs Utilization)。在實現這一目標的過程中,開發人員往往會面臨多種挑戰,比如帶寬瓶頸、通信瓶頸以及顯存限制等。為了突破這些障礙,開發人員需要采用各種優化策略,力求將算力發揮到極致。
對于推理來說,絕大部分情況下都是單機部署(可能會有單機多卡),我們這里只分析其中最為樸素的一種情況,單機單卡部署,并且顯存是充裕的。這時候影響推理性能的硬件因素的只有兩個,算力和顯存帶寬。
對于 Decode 來說,它每次輸入的 token 只有一個,但是在計算的過程中需要將 LLM 全部的 weights 加載到寄存器中,在這個過程中,首先會達到硬件的顯存帶寬瓶頸,導致 GPU 算力無法充分發揮。
對于 Prefill,它每次輸入的 token 來自完整的 prompt,假設這個 prompt 包含的 token 比較多,比如幾百個以上,那么在計算過程中,加載一次 weights 就可以同時計算幾百個 token,這時候硬件會達到算力瓶頸。
上圖這張圖 (來自于 FlashInfer) 生動的闡述了這一過程,這張圖通常也被稱為 Roofline Model,橫軸是計算強度,即每進行 1FLOPS 計算所需要的訪存量,縱軸是 GPU 每秒的浮點數計算次數,在計算強度降低時,比如 Decode 階段,我們會達到顯存帶寬瓶頸,導致只有一部分 GPU 算力被使用,隨著計算強度拉大,在 Verification 和 Append 過程,越來越多的 GPU 算力被使用,最后,到了 Prefill 階段,我們達到了算力瓶頸,每秒進行的浮點數計算次數達到 312 TFLOPs,這也是推理最理想的情況。注意:上面的 Verication 對應著投機采樣中大模型的 Score 過程,Append 過程對應著使用 reuse cache 的情況下,Prefill 過程中被 reuse 的那部分 KV Cache 不需要重新計算,只需要計算后面的 KV Cache。
上述Roofline Model展示的是請求并發度為1的情況,對于主流的推理引擎,通常可以在一次迭代過程中可以同時對多個請求進行并行處理。計算強度跟batch_size成線性關系, batch_size的計算公式如下:batch_size = seq_len1 + seq_len2+ ... + seq_lenn
這里的 n 是并發的請求數,seq_len~i~ 是每一請求當前輪迭代要處理的 token 數量,對于 Decode,seq_len 為 1,對于 Prefill,seq_len 為 prompt 的 token 數量(不考慮 reuse cache)。
根據上述 batch_size 的計算公式可以發現,提升計算強度的方式有兩種,一是提升請求的并發度,讓同一時刻有多個請求在 Decode,二是提升每一個請求的 seq_len。投機采樣正是這樣一種提升計算強度的方式,它在保持原有請求并發度的情況下,通過小模型來猜測未來大模型可能會生成的 k 個 token,進而讓大模型一次性對這 k 個 token 進行并行計算 (每一個請求的 seq_len 從 1 增加到了 k),之后通過驗證算法來決定接受多少 token。大模型一次性對 k 個 token 進行并行計算的過程稱為 Score,對于 Score 而言,通過調整小模型提議的 token 數量 (即 k 值),我們可以控制投機采樣的計算強度,當 seq_len 長到一定程度之后,Score 也會跟 Prefill 一樣達到算力瓶頸。
投機采樣簡介及性能分析
簡介
投機采樣會有兩個以上的模型(圖來自于pearl):
小模型Mq(后面稱為Propose Model),Mq既可以是概率模型,比如神經網絡,也可以是規則模型,比如n-gram token匹配,RAG匹配等
大模型Mp(后面稱為Score Model),Mp是需要加速的GPT模型本身
投機采樣每一步的迭代過程都會包含下面三個階段
Propose Stage: 小模型 M~q~ 提議 k 個 token
Score Stage: 大模型 M~p~ 對這 k 個 token 進行打分,這個打分過程是并行計算的
Verification Stage: 根據小模型和大模型的輸出的 token 和 probs 來決定本輪迭代接收多少個 token
通過上面這種方式,投機采樣把原來逐個 token 逐個 token 的 Decode 過程轉換成一次性對 k 個 token 進行 Score 并行計算的過程,從而更高效的利用了 GPU 的算力資源。另外,相比于其他的推理加速手段,比如量化,投機采樣在某種特定的驗證算法可以保證精度不丟失,這是投機采樣相比于其他加速手段一個重要優勢。
延遲分析
請求并發度為 1
在本小節,我們通過數學公式來分析下投機采樣在延遲方面的理論降低上限,為了簡化分析,我們做如下
假設:
請求并發度為 1
不考慮框架的額外開銷,并且 Verification Stage 的開銷為 0
LLM 的輸出足夠長,可以忽略 Prefill 階段的影響
小模型為 GPT 模型
此處假設小模型也為GPT模型,每輪迭代小模型提議k步(實際情況中,k通常為10以內的較小值),小模型單輪Decode的耗時為Tq,大模型單輪Decode和對k個token進行Score的耗時均為Tp(k比較小), 此時投機采樣單輪的耗時為 Tq* k + Tp,假設平均接收長度為L,那么每個token的平均接收耗時為 (Tq* k + Tp) / L ,除以大模型單次Decode的耗時 Tp即可得到理論加速比:(Tq* k + Tp) / L * Tp
舉個例子,小模型單次 Decode 耗時 7ms, 大模型單次 Decode 和對 5 個 token 進行 Score 的耗時均為 18ms,小模型每次提議 5 個 token,那么每輪投機采樣的耗時為 5 * 7 + 18 = 53ms,假設每輪迭代平均接收的 token 為 3 個,那么生成每個 token 的平均耗時是 53/3=17.66ms,加速比為 18/17.66=1.02,同理,如果平均接收的 token 為 4 個,那么生成每個 token 的平均耗時是 53/4=13.25ms,加速比為 18/13.25=1.35。
由此可見,決定投機采樣模型加速比的主要因素有兩個:
小模型單輪 Decode 的延遲,越低越好,對于某些規則模型,該項延遲可以忽略不計
小模型提議 token 的接受率,越高越好
小模型為規則模型
在某些投機采樣方法中,小模型可以通過簡單的規則(比如n-gram token匹配)就可以提議token,這種通過規則匹配提議token的過程耗時可以幾乎忽略不計,因此我們可以假設小模型提議token的耗時Tq為0。對于這一類模型,我們可以把k設置的很大,但是當k比較大(比如1024)情況下,大模型單次Decode和單次對k個token進行Score的耗時的差異不可以忽略,此時上述理論的加速比會變成(Tp-score-k) / L * Tp-decode
這里Tp-score-k表示大模型對k個token進行Score的耗時,Tp-decode 表示大模型單次Decode的耗時考慮其極限情況,假設小模型預測的token都能被大模型完美接收,此時k就等于L,在這種情況下,我們可以得到所有投機采樣方法加速比的極限:max Tp-score-k/(Tp-decode * k), k=1,2,...+∞
拿 A100 卡舉例,當模型為 qwen72b 時, 開啟 TP=2 的情況下,在 RTP-LLM 框架下:
對 128 個 token 進行逐個 token 逐個 token 的 Decode 的耗時為 51*128=6528ms,對 128 個 token 進行 Score 的耗時為 65ms,此時投機采樣加速比為 100。
對 2048 個 token 進行逐個 token 逐個 token 的 Decode 的耗時為 54*2048=104448ms,對 2048 個 token 進行 Score 的耗時為 701ms,此時投機采樣加速比為 148。
總結下,在請求并發度為 1 時,投機采樣在我看來 GPU 資源充裕的情況下追求極致低延遲的技術,在使用規則模型并且接收率非常高的情況下,投機采樣可以把一個原本是訪存瓶頸的 Decode 任務轉化計算瓶頸 Score 的任務,后面 Aone Copilot 能夠實現 1000 tokens 一秒的輸出也是基于這個理論。
請求并發度>1
在上面對投機采樣延遲分析過程中,我們假設請求并發度為 1,但這個假設太嚴苛了,在很多時候都是不成立的。目前主流的推理框架都已經實現了 continous batching,在實際運行的過程中請求并發度經常會大于 1 的情況,當請求并發度大于 1 時,投機采樣框架對延遲的提升就會快速下降,甚至可能會出現比不開投機采樣還要差的情況,在下圖中可以看到,隨著 QPS 上升,投機采樣的的延遲快速上升,在 QPS 為 10 的時候,它比不開投機采樣還要慢 1.8 倍 (下圖來自于 vllm)。
但是這也并不意味著投機采樣只要在請求并發度>1 的時候就一定會帶來負面收益,在上圖中可以看到,在低 QPS 的時候,投機采樣還是更好的利用到 GPU 空閑算力,從而達到比非投機采樣更低的延遲。具體到我們的實際業務場景中,會有以下兩種場景:
在線任務中,用戶流量確定,GPU 資源不太充裕,每個請求的輸入比較長,想要在保證在一定延遲的情況下,盡可能提升提升吞吐。由于用戶的資源不太充裕,不能夠部署 PD 分離實例。此時為了保證用戶服務的延遲,每張 GPU 卡上分配的用戶請求的并發度就不能開的太高,否則會出現大量 a 請求的 Decode 被 b 請求的 Prefill 拖慢的情況,為了避免這種情況,我們可以開啟 chunked_prefill,但是模型比較大的情況下,chunked_prefill 增加的 batch_size 還是太大了,它還是會拖慢 Decode 的執行效率,最終,用戶還是需要更多的 GPU 才能滿足他們的服務需求。投機采樣的好處是,它可以在每張 GPU 卡保持較低的請求并發度的情況下,通過設置比較小的 k 值,在不對延遲產生大影響的情況下來提升吞吐,從而滿足這部分用戶的需求。
在線任務中,用戶流量確定且比較小,GPU 資源充裕,用戶想要追求極致且穩定的低延遲。由于用戶的 GPU 資源比較充裕,我們可以通過部署 PD 分離來獲得穩定的低延遲。用戶想在此基礎上追求更加極致的低延遲,并且用戶的流量又比較小,無法充分利用起 Decode 實例的算力,此時就可以在 Decode 實例上通過投機采樣來進一步降低用戶的延遲。
總結下,在流量確定的情況下,通過調節投機采樣的 k 值,我們可以有效地平衡用戶對于延遲和吞吐量的需求。在流量不確定的情況下,我們也可以讓投機采樣能夠依據實時流量負載動態調整 k 值來讓它的延遲在最壞情況下也比非投機采樣好。在這個角度上來看,投機采樣是一項能賦予單機計算靈活性的關鍵技術。當請求的并發度很大,算力已接近飽和時,投機采樣的增益較為有限,此時應適當減小 k 值,以免影響延遲表現。相反,若請求的并發度較小,GPU 算力尚未完全利用,通過調整 k 值,投機采樣可提議出更多 token 供大模型并行 Score,從而大幅提升算力利用率。
吞吐分析
以下對吞吐影響的討論主要是在小模型為規則模型(即 Prompt Lookup)的基礎上進行討論的,它的提議過程耗時可以幾乎忽略不計,我們只需要考慮大模型 Score 的耗時。
在不開投機采樣的情況,通過加大請求的并發數可以讓算力使用達到極限情況,此時就達到了最大吞吐。如果此時我們開啟投機采樣后,按照同樣的請求并發數,吞吐會出現一定程度的下降,這主要是因為投機采樣不能保證每一個提議 token 會被成功接受,但是無論接受與否,每一個提議的 token 都會通過大模型的 Score 驗證,對于那些沒有被接受,計算是浪費的。上文中 QPS 很高的時候延遲下降也是同理。
通過上述簡單的分析,我們可以得到一個結論,投機采樣對于 LLM 推理的吞吐是有害的,這個結論是大部分人的共識,我最初的認知也是如此,但事實真的是這樣嘛?投機采樣對于吞吐是有害的這一結論在長期極限壓力并且顯存是完全充裕的情況下基本上是沒問題的,但是實際使用的場景中,這一場景的假設并不總是成立的:
顯存有時候比算力更快達到瓶頸。為了讓吞吐盡可能提高,我們會把并發調到最大,直至顯存的上限。但是在內部業務場景中,有各種各樣類型的 GPU 卡,有些卡的顯存很小,比如 L20,它的顯存是 48G,當我們把一個 Qwen-32B-int8 模型加載上去后,留給 KV Cache 的顯存就不是特別多了,在這種情況下,如果請求的輸出特別長,請求并發度就不能打的特別高(運行期大部分的顯存都會被每個請求的 KV Cache 占據,它跟輸入輸出的長度呈線性關系),此時如果再把請求并發度打更高,會讓部分請求在運行過程中因為分配不到 KV Cache 而被淘汰,導致吞吐降低。在這種情況下,由于存在顯存瓶頸,GPU 的算力不能被充分使用,投機采樣在這種情況下就可以在不增加請求并發度的情況進一步提高算力的利用率。
由于推理引擎框架層的限制,通過提升請求的并發度來讓 GPU 卡打滿算力有時候并不是一件容易的事情,特別是對于尺寸比較小的模型。如下圖所示,以 RTP-LLM 在 Qwen 1.8b 模型為例,要想讓 A10 完全打滿算力,需要請求的并發度達到 1024(此后 seq_len 和 prefill latency 的關系是系數為 1 的線性關系,對于算力更強的卡比如 A100,需要的請求并發度更高),這樣 Decode 階段湊的 batch size 才能夠達到算力瓶頸,但在真實場景下達到這樣的并發度會對推理框架本身的實現提出非常高要求,比如湊批流程,tokenizer,PB 序列化,Python 的 GIL 鎖等任何一個環節都可能成為瓶頸。投機采樣可以在請求并發度較低的情況下,通過調整 k 值來進一步提升 batch size,更好的利用顯卡算力。
對于超長 sequence,通過投機采樣增大 batch_size 對于 attention 的歷史 KV Cache 訪存更加友好一些。非投機采樣為了湊大 batch_size,只能加大請求的并發度,考慮 128 個請求湊批的情況,128 個請求在 attention 階段都要訪問一下自己歷史所有的 KV Cache,這會嚴重降低 attention 計算階段的計算強度,sequence 特別長的情況下甚至可能在這一階段打不到計算瓶頸。但是,如果用投機采樣來增大 batch_size 的話,可以不增加請求的并發度,比如對于 1 個請求,我們可以直接將 k 設置為 128,此時也會湊滿 128 的 batch_size,但是在 attention 階段只需要訪問這一個請求的歷史 KV Cache,可以獲得更大的計算強度,進而帶來吞吐的提升。
總結下,對于提升吞吐而言,在許多情況下受限于顯存大小,推理框架本身實現,訪問歷史 KV Cache 等的限制,我們不能將 GPU 算力充分利用起來。此時投機采樣可以通過調整 k 值,從另一個維度提升 batch_size,進而充分利用起 GPU 的算力。
RTP-LLM 通用投機采樣框架實現
投機采樣一直是大語言模型領域學術界和工業界關注的熱點,有各種各樣的投機模型,比較主流的有:
樸素投機采樣:拿小尺寸的 GPT 模型作為小模型 M~q~
Medusa:在 last hidden state 上面添加幾個額外的 lm head 來預測未來若干個 token
Eagle:訓練一個新的 Auto-Regression Head 來預測未來的 hidden state,根據未來的 hidden state 去預測未來的 token
Prompt Lookup: 去歷史 prompt 中進行 n-grams 匹配,將 n-grams 匹配成功后面的若干個 token 作為提議 token
可以看到,當前階段有非常多的投機采樣模型需要支持,同時未來可能還有樹采樣,動態投機采樣等需求。為此,我們在 RTP-LLM 中實現一套通用的投機采樣框架,它能夠很好應對未來可能增加的各種投機采樣需求的拓展,同時這套投機采樣框架盡可能復用了原來的 GenerateStream,BatchStreamProcessor,NormalExecutor 等模塊組件,以減少開發的工作量。
我們根據投機采樣的三個階段 (Propose, Score, Verification),在實現時將投機采樣分為四大類模塊組件:
ProposeExecutor:小模型提議 token,根據不同投機采樣算法有多種實現,比如樸素投機采樣,PromptLookup,Eagle 和 Medusa 等
ScoreExecutor:大模型對小模型提議的 token 進行打分
SpeculativeSampler:對小模型大模型的輸出的 token 進行驗證來決定接收多少個 token
SpeculativeUpdater:將接收的 token 更新到最初的 strema 上面
四個模塊組件都有明確的輸入輸出并且各自都是無狀態的,從而實現功能的接觸耦合。
RTP-LLM 投機編輯實現與業務提升
Prompt Lookup 投機采樣簡介
在正式介紹投機編輯前,我們首先介紹下 Prompt Lookup 投機采樣。
Prompt Lookup 投機采樣每一輪的基本流程如下:
根據最近幾個生成的 token 在 prompt 中進行 n-gram token 匹配
如果匹配命中,則把命中位置之后的 k 個 token 作為 Propose token(上圖中叫做 candidate tokens),到此完成小模型提議 token 階段
將小模型提議的 tokens 交給大模型打分
驗證算法根據小模型和大模型的輸出來決定當輪迭代接受多少個 token
對于抽取式場景,Prompt Lookup 的加速效果非常明顯,主要是由于小模型提議的 token 接收率特別高。
基于 Prompt Lookup 的投機編輯實現
代碼編輯也是這樣一類抽取式場景,比如下圖中,我們實現了一個快速排序算法,想要通過大模型把其中的 arr 變量的 long 類型改成 int 類型,并且要求模型輸出修改完后的完整代碼,此時大部分的模型輸出都可以直接從輸入中摘抄即可,僅在發生差異的地方特殊處理一下即可,Aone Copilot 的 Fast Apply 實現一鍵應用修改功能也是類似的場景。
除此之外,編輯這一類特殊的抽取式場景還有一個特點,大模型在做摘抄的時候會滿足一個前后順序關系,比如上圖中假設大模型當前階段的摘抄到了第 7 行的 quickSort(arr, low, pi - 1); ,下次摘抄的起始位置一定是從之后的位置(第 8 行)開始的,基于這個特性,我們對 prompt lookup 做了如下改進來提高它在投機編輯場景的性能:
額外維護一個游標來記錄上次 prompt lookup 成功的位置,每次匹配都以游標作為起始位置,從之后的位置開始進行尋找 n-grams token 匹配,在每輪投機采樣迭代后也會把游標的位置往后面更新。
第一輪迭代時跳過 n-grams token 匹配,直接將最開始的 k 個 token 作為提議 token。
投機編輯對業務的性能提升
延遲提升
Aone Copilot 基于投機編輯技術實現了 Fast Apply 功能。該功能的輸入是原始代碼及需修改的部分,輸出則是重新編寫的代碼。由于需要重寫全部代碼,因此對延遲提出了很大的挑戰。借助投機編輯技術,我們在 32B 規模的模型上實現了單個請求每秒可以推理 1000 條 token,相比未使用投機編輯之前提升了 10 倍性能,滿足了 Aone Copilot 上線業務的需求。
公司內部的大模型業務廣泛采用了結構化輸出方式,這些業務的 LLM 輸出內容大部分都是比較固定的,對于固定的輸出部分可以直接通過投機編輯去進行“摘抄”即可。在使用投機編輯技術之后,這部分業務的延遲得到了顯著的降低。具體而言,在一項使用 72B 規模模型的結構化輸出業務中,輸出長度約為 100 個字符,在應用投機編輯后,Decode 部分的延遲從原來的 3 秒降至 1.2 秒,實現了大約 60% 的 Decode 延遲提升。
吞吐提升
公司內部有一個訓練數據清洗的業務,相比傳統的規則匹配方法,使用大模型能夠顯著提高數據清洗的效果,從而更好地支持后續的訓練任務。大語言模型的輸入是未經處理的原始語料,輸出則是經過清洗后的干凈語料。輸入語料的平均長度約為 3000 個 token,而輸出語料的平均長度約為 1000 個 token。通過使用投機編輯技術進行加速后,單張 GPU 卡的每秒請求率(RPS)從 3 提升到了 10。
未來方向
投機采樣的優化跟場景密切相關,在使用時自身也有許多參數需要配置,目前在阿里的內部業務中只有一部分業務使用上了投機采樣,對于大多數業務方來說投機采樣使用起來還是有一定的門檻。在未來, 我們將增強 Prompt Lookup 投機采樣并將其設置為一種默認開啟的功能,讓它能夠開箱即用且能夠根據線上流量做自適應調整,使所有業務方都能享受其性能優勢。
最后,Prompt Lookup 投機采樣方法能夠在很大程度上加速 LLM 輸出的前提是大部分內容都能被準確預測。然而在不少場景下 (比如對話場景),該假設并不總是成立。此時需要借助一個額外的小型神經網絡來預測未來的 token 以進一步提升接收率。得益于 RTP-LLM 通用的投機采樣框架的支持,我們將很輕松地集成其他先進的投機采樣方法,例如近期備受關注的 DeepSeek 的 MTP。在未來,我們也會在 RTP-LLM 中將 Prompt Lookup 投機采樣和 Deepseek MTP 的結合起來,對較為簡單的預測采用 Prompt Lookup 直接完成,而對于更為復雜的預測,則利用 MTP 進行預測,從而進一步加快 Deepseek 的生成速度。
參考文獻
https://flashinfer.ai/2024/02/02/introduce-flashinfer.html
https://pearl-code.github.io/
https://docs.google.com/presentation/d/1wUoLmhfX6B7CfXy3o4m-MdodRL26WvY3/
https://github.com/apoorvumang/prompt-lookup-decoding.git
AICon 2025 強勢來襲,5 月上海站、6 月北京站,雙城聯動,全覽 AI 技術前沿和行業落地。大會聚焦技術與應用深度融合,匯聚 AI Agent、多模態、場景應用、大模型架構創新、智能數據基建、AI 產品設計和出海策略等話題。即刻掃碼購票,一同探索 AI 應用邊界!
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.