大型語言模型 (LLM) 越來越多地用于需要多個鏈式生成調(diào)用、高級 prompt 技術、控制流以及與外部環(huán)境交互的復雜任務。然而，用于編程和執(zhí)行這些應用程序的現(xiàn)有高效系統(tǒng)存在著明顯的缺陷。

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

總的來說，SGLang 的貢獻主要包括：

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

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

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

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

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

后端：使用 RadixAttention 自動 KV 緩存復用

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

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

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

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

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

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

下圖說明了在處理多個傳入請求時如何維護基數(shù)樹。前端總是向運行時發(fā)送完整的 prompt，運行時會自動進行前綴匹配、復用和緩存。樹形結構存儲在 CPU 上，維護開銷較小。

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

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

前端：使用 SGLang 輕松進行 LLM 編程

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

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

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

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

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

基準測試

研究團隊在常見的 LLM 工作負載上測試了其系統(tǒng)，并報告了所實現(xiàn)的吞吐量。

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

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