大型語(yǔ)言模型 (LLM) 越來(lái)越多地用于需要多個(gè)鏈?zhǔn)缴烧{(diào)用、高級(jí) prompt 技術(shù)、控制流以及與外部環(huán)境交互的復(fù)雜任務(wù)。然而，用于編程和執(zhí)行這些應(yīng)用程序的現(xiàn)有高效系統(tǒng)存在著明顯的缺陷。

現(xiàn)在，開(kāi)源社區(qū)的研究者們面向 LLM 提出了一種結(jié)構(gòu)化生成語(yǔ)言（Structured Generation Language）——SGLang。SGLang 能夠增強(qiáng)與 LLM 的交互，通過(guò)聯(lián)合設(shè)計(jì)后端運(yùn)行時(shí)系統(tǒng)和前端語(yǔ)言，使 LLM 更快、更可控。機(jī)器學(xué)習(xí)領(lǐng)域知名學(xué)者、CMU 助理教授陳天奇還轉(zhuǎn)發(fā)了這項(xiàng)研究。

總的來(lái)說(shuō)，SGLang 的貢獻(xiàn)主要包括：

• 在后端，研究團(tuán)隊(duì)提出了 RadixAttention，這是一種跨多個(gè) LLM 生成調(diào)用的 KV 緩存（KV cache）復(fù)用技術(shù)，自動(dòng)且高效。
• 在前端，研究團(tuán)隊(duì)開(kāi)發(fā)了一種嵌入 Python 的、靈活的域指定（domain-specific）語(yǔ)言來(lái)控制生成過(guò)程。該語(yǔ)言可以在解釋器模式或編譯器模式下執(zhí)行。

后端前端組件協(xié)同工作，可提高復(fù)雜 LLM 程序的執(zhí)行和編程效率。

該研究使用 SGLang 實(shí)現(xiàn)了常見(jiàn)的 LLM 工作負(fù)載，包括智能體、推理、提取、對(duì)話(huà)和小樣本學(xué)習(xí)任務(wù)，并在 NVIDIA A10G GPU 上采用 Llama-7B 和 Mixtral-8x7B 模型。如下圖 1 、圖 2 表明，與現(xiàn)有系統(tǒng)（即 Guidance 和 vLLM）相比，SGLang 的吞吐量提高了 5 倍。

圖 1：不同系統(tǒng)在 LLM 任務(wù)上的吞吐量（A10G、FP16 上的 Llama-7B、張量并行度 = 1）

圖 2：不同系統(tǒng)在 LLM 任務(wù)上的吞吐量（A10G、FP16 上的 Mixtral-8x7B，張量并行度 = 8）

后端：使用 RadixAttention 自動(dòng) KV 緩存復(fù)用

在 SGLang 運(yùn)行時(shí)的開(kāi)發(fā)過(guò)程中，該研究發(fā)現(xiàn)了復(fù)雜 LLM 程序的優(yōu)化關(guān)鍵 ——KV 緩存復(fù)用，當(dāng)前系統(tǒng)對(duì)此處理不佳。KV 緩存復(fù)用意味著具有相同前綴的不同 prompt 可以共享中間 KV 緩存，避免冗余的內(nèi)存和計(jì)算。在涉及多個(gè) LLM 調(diào)用的復(fù)雜程序中，可能存在各種 KV 緩存復(fù)用模式。下圖 3 說(shuō)明了 LLM 工作負(fù)載中常見(jiàn)的四種此類(lèi)模式。雖然某些系統(tǒng)能夠在某些場(chǎng)景下處理 KV 緩存復(fù)用，但通常需要手動(dòng)配置和臨時(shí)調(diào)整。此外，由于可能的復(fù)用模式的多樣性，即使通過(guò)手動(dòng)配置，現(xiàn)有系統(tǒng)也無(wú)法自動(dòng)適應(yīng)所有場(chǎng)景。

圖 3：KV 緩存共享示例。藍(lán)色框是可共享的 prompt 部分，綠色框是不可共享的部分，黃色框是不可共享的模型輸出?？晒蚕淼牟糠职ㄐ颖緦W(xué)習(xí)示例、自洽（self-consistency）問(wèn)題、多輪對(duì)話(huà)中的對(duì)話(huà)歷史以及思維樹(shù)（tree-of-thought）中的搜索歷史。

為了系統(tǒng)地利用這些復(fù)用機(jī)會(huì)，該研究提出了一種在運(yùn)行時(shí)自動(dòng) KV 緩存復(fù)用的新方法 —— RadixAttention。該方法不是在完成生成請(qǐng)求后丟棄 KV 緩存，而是在基數(shù)樹(shù)（radix tree）中保留 prompt 和生成結(jié)果的 KV 緩存。這種數(shù)據(jù)結(jié)構(gòu)可以實(shí)現(xiàn)高效的前綴搜索、插入和驅(qū)逐。該研究實(shí)現(xiàn)了最近最少使用（LRU）驅(qū)逐策略，并輔以緩存感知調(diào)度策略，以提高緩存命中率。

基數(shù)樹(shù)可作為 trie（前綴樹(shù)）節(jié)省空間的替代方案。與典型的樹(shù)不同，基數(shù)樹(shù)的邊緣不僅可以用單個(gè)元素來(lái)標(biāo)記，還可以用不同長(zhǎng)度的元素序列來(lái)標(biāo)記，這提高了基數(shù)樹(shù)的效率。

該研究利用基數(shù)樹(shù)來(lái)管理映射，這種映射是在充當(dāng)鍵的 token 序列和充當(dāng)值的相應(yīng) KV 緩存張量之間進(jìn)行的。這些 KV 緩存張量以分頁(yè)布局存儲(chǔ)在 GPU 上，其中每個(gè)頁(yè)的大小相當(dāng)于一個(gè) token。

考慮到 GPU 內(nèi)存容量有限，無(wú)法重新訓(xùn)練無(wú)限的 KV 緩存張量，這就需要驅(qū)逐策略。該研究采用 LRU 驅(qū)逐策略，遞歸地驅(qū)逐葉節(jié)點(diǎn)。此外，RadixAttention 與連續(xù)批處理和分頁(yè)注意力等現(xiàn)有技術(shù)兼容。對(duì)于多模態(tài)模型，RadixAttention 可以輕松擴(kuò)展以處理圖像 token。

下圖說(shuō)明了在處理多個(gè)傳入請(qǐng)求時(shí)如何維護(hù)基數(shù)樹(shù)。前端總是向運(yùn)行時(shí)發(fā)送完整的 prompt，運(yùn)行時(shí)會(huì)自動(dòng)進(jìn)行前綴匹配、復(fù)用和緩存。樹(shù)形結(jié)構(gòu)存儲(chǔ)在 CPU 上，維護(hù)開(kāi)銷(xiāo)較小。

圖 4. 采用 LRU 驅(qū)逐策略的 RadixAttention 操作示例，分九個(gè)步驟進(jìn)行說(shuō)明。

圖 4 演示了基數(shù)樹(shù)響應(yīng)各種請(qǐng)求的動(dòng)態(tài)演變。這些請(qǐng)求包括兩個(gè)聊天會(huì)話(huà)、一批小樣本學(xué)習(xí)查詢(xún)和自洽性抽樣。每個(gè)樹(shù)邊緣都帶有一個(gè)標(biāo)簽，表示子字符串或 token 序列。節(jié)點(diǎn)采用顏色編碼以反映不同的狀態(tài)：綠色表示新添加的節(jié)點(diǎn)，藍(lán)色表示在該時(shí)間點(diǎn)訪(fǎng)問(wèn)的緩存節(jié)點(diǎn)，紅色表示已被驅(qū)逐的節(jié)點(diǎn)。

前端：使用 SGLang 輕松進(jìn)行 LLM 編程

在前端，該研究提出了 SGLang，一種嵌入在 Python 中的特定于領(lǐng)域的語(yǔ)言，允許表達(dá)高級(jí) prompt 技術(shù)、控制流、多模態(tài)、解碼約束和外部交互。SGLang 函數(shù)可以通過(guò)各種后端運(yùn)行，例如 OpenAI、Anthropic、Gemini 和本地模型。

圖 5. 用 SGLang 實(shí)現(xiàn)多維文章評(píng)分。

圖 5 顯示了一個(gè)具體示例。它利用分支 - 解決 - 合并 prompt 技術(shù)實(shí)現(xiàn)多維文章評(píng)分。該函數(shù)使用 LLM 從多個(gè)維度評(píng)估文章的質(zhì)量，合并判斷，生成摘要，并分配最終等級(jí)。突出顯示的區(qū)域說(shuō)明了 SGLang API 的使用。(1) fork 創(chuàng)建 prompt 的多個(gè)并行副本。(2) gen 調(diào)用 LLM 生成并將結(jié)果存儲(chǔ)在變量中。該調(diào)用是非阻塞的，因此它允許多個(gè)生成調(diào)用在后臺(tái)同時(shí)運(yùn)行。(3) [variable_name] 檢索生成的結(jié)果。(4) 選擇對(duì)生成施加約束。(5) run 使用其參數(shù)執(zhí)行 SGLang 函數(shù)。

給定這樣一個(gè) SGLang 程序，我們可以通過(guò)解釋器執(zhí)行它，也可以將其跟蹤為數(shù)據(jù)流圖并使用圖執(zhí)行器運(yùn)行它。后一種情況為一些潛在的編譯器優(yōu)化開(kāi)辟了空間，例如代碼移動(dòng)、指令選擇和自動(dòng)調(diào)整。

SGLang 的語(yǔ)法很大程度上受到 Guidance 的啟發(fā)，并引入了新的原語(yǔ)，還處理程序內(nèi)并行性和批處理。所有這些新功能都有助于 SGLang 的出色性能。

基準(zhǔn)測(cè)試

研究團(tuán)隊(duì)在常見(jiàn)的 LLM 工作負(fù)載上測(cè)試了其系統(tǒng)，并報(bào)告了所實(shí)現(xiàn)的吞吐量。

具體來(lái)說(shuō)，該研究在 1 個(gè) NVIDIA A10G GPU (24GB) 上測(cè)試了 Llama-7B，在 8 個(gè)具有張量并行性的 NVIDIA A10G GPU 上使用 FP16 精度測(cè)試了 Mixtral-8x7B，并使用 vllm v0.2.5、指導(dǎo) v0.1.8 和 Hugging Face TGI v1.3.0 作為基準(zhǔn)系統(tǒng)。

如圖 1 和圖 2 所示，SGLang 在所有基準(zhǔn)測(cè)試中均優(yōu)于基準(zhǔn)系統(tǒng)，吞吐量提高了 5 倍。它在延遲方面也表現(xiàn)出色，特別是對(duì)于第一個(gè) token 延遲，其中前綴緩存命中可以帶來(lái)顯著的好處。這些改進(jìn)歸功于 RadixAttention 的自動(dòng) KV 緩存復(fù)用、解釋器實(shí)現(xiàn)的程序內(nèi)并行性以及前端和后端系統(tǒng)的協(xié)同設(shè)計(jì)。此外，消融研究表明，即使沒(méi)有緩存命中，也沒(méi)有明顯的開(kāi)銷(xiāo)，這會(huì)導(dǎo)致在運(yùn)行時(shí)始終啟用 RadixAttention。