陳天奇團隊LLM結構化生成新引擎XGrammar：百倍加速、近零開銷

机器之心發表於2024-11-26
現在，大語言模型的結構化生成有了一個更加高效、靈活的引擎。

不管是編寫和除錯程式碼，還是透過函式呼叫來使用外部工具，又或是控制機器人，都免不了需要 LLM 生成結構化資料，也就是遵循某個特定格式（如 JSON、SQL 等）的資料。

但使用上下文無關語法（CFG）來進行約束解碼的方法並不高效。針對這個困難，陳天奇團隊提出了一種新的解決方案：XGrammar。
圖片
XGrammar 是一個開源軟體庫，可實現高效、靈活且可移植的結構化生成。該團隊在部落格中表示：「我們毫不妥協地實現了這三個目標，並致力於一個核心使命：將靈活、零開銷的結構化生成帶到任何地方。」
圖片
  • 論文標題：XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models
  • 論文地址：https://arxiv.org/pdf/2411.15100
  • 程式碼地址：https://github.com/mlc-ai/xgrammar

對於結構化生成，一種常用方法是約束解碼。在每個解碼步驟中，約束解碼都會檢查詞表，並透過將無效 token 的機率設定為零來過濾掉違反指定結構的 token。為了支援多種多樣的結構格式，需要一種靈活的機制來指定和檢查這些約束。
圖片
使用 JSON 方案實現約束解碼

上下文無關語法（CFG）就能提供一種通用方法，即透過一組規則來定義結構。其中每條規則都包含一個字元序列或其他規則，並允許遞迴組合來表示複雜的結構。相比於正規表示式等其它格式，CFG 由於支援遞迴結構，因而能提供更大的靈活性，使其適合描述 JSON、SQL 和領域特定語言（DSL）等常見語言。

下圖展示了一個用於陣列和字串的 CFG，可以清楚地看到其中的遞迴結構。
圖片
但是，也正因為 CFG 很靈活，所以直接將其應用於約束解碼的效率並不高。首先，每個解碼步驟都需要對詞表中每個可能的 token 解釋 CFG，在 Llama 3.1 中，這個詞表的大小可能高達 128k。此外，CFG 解釋需要一個堆疊狀態來跟蹤之前匹配的遞迴規則，因此無法提前計算和快取堆疊模式的所有組合。最後，LLM 生成結果中的每個 token 都包含多個字元，這些字元可能會跨越語法元素的邊界，並在執行時執行期間導致進一步的遞迴或堆疊彈出。這種未對齊的邊界問題很棘手，需要在語法執行期間小心處理它們。

XGrammar 便是為解決上述難題而生的，並且效果卓越：相比於之前的 SOTA 方法，XGrammar 可以將上下文無關語法的每 token 延遲減少多達 100 倍！此外，他們還基於 Llama3.1 模型實驗了整合了 XGrammar 的 LLM serving 引擎；在 H100 GPU 上，這能將透過結構化輸出實現端到端 LLM serving 的速度提升 80 倍！

該團隊表示：「我們正在開源 XGrammar 並將其整合到主要的開源 LLM 框架中。」

XGrammar 概覽

如圖 1 所示，Grammar 利用了位元組級下推自動機（byte-level pushdown automaton）來解釋上下文無關語法。
圖片
這種位元組級設計允許每個字元邊緣包含一個或多個位元組，處理不規則的 token 邊界並支援包含 sub-UTF8 字元的 token 。該自動機的結構經過最佳化以加快匹配速度。

在預處理階段，會生成一個自適應 token 掩碼快取，它會透過預先計算與上下文無關的 token 來加快執行時的掩碼生成。上下文擴充套件（context extension）能進一步提升這種快取的有效性。

在執行時，token 掩碼快取會快速生成大部分掩碼，而持續性執行堆疊會高效處理其餘的上下文相關 token。

此外，掩碼生成和 LLM 推理是互相重疊的，以最大限度地減少約束解碼的開銷。一旦 LLM 在掩碼約束下生成新 token，就會使用此 token 來更新下推自動機的堆疊狀態，以進行下一次掩碼生成。

具體來說，陳天奇團隊首先得到了一個見解：雖然無法預先計算下推自動機（PDA）無限多個狀態的完整掩碼，但可以預先計算掩碼中相當一部分（通常超過 99%）的 token。因此，可將這些 token 分成兩類：

  • 上下文無關 token：僅透過檢視 PDA 中的當前位置而不是堆疊即可確定其有效性的 token。

  • 上下文相關 token：必須使用整個堆疊來確定其有效性的 token。

下圖展示了一組上下文相關和無關 token 的示例。大多數情況下，上下文無關 token 佔大多數。我們可以預先計算 PDA 中每個位置的上下文無關 token 的有效性，並將它們儲存在自適應 token 掩碼快取中。此過程稱為語法編譯（grammar compilation）。
圖片
下圖則展示了自適應儲存格式。
圖片
在執行時，首先檢索來自快取的上下文無關 token 的有效性。然後，高效地執行 PDA 來檢查其餘的上下文相關 token。透過跳過執行時檢查大多數 token，便可以顯著加快掩碼生成速度。XGrammar 執行時間的整體工作流程見圖 1。

此外，他們還設計了一組額外的演算法和系統最佳化方法，以進一步提高掩碼生成速度並減少預處理時間，包括上下文擴充套件、持續性執行椎棧、下推自動機結構最佳化、並行式語法編譯。

上下文擴充套件

該團隊提出的方法是檢測語法中每個規則的額外上下文資訊，並將其用於減少上下文相關 token 的數量，並進一步加快執行時檢查速度。
圖片
持續性執行堆疊

為了加快由於多種可能的擴充套件路徑而導致的拆分和合並期間多個並行堆疊的維護速度，他們設計了一個基於樹的資料結構，可以有效地同時管理多個堆疊。
圖片
它還可以儲存以前的狀態並實現高效的狀態回滾，從而加快上下文相關 token 的執行時檢查速度。

下推自動機結構最佳化

研究者進行了額外的最佳化，以改進下推自動機的結構，加快最終執行的效率。這些最佳化借鑑了傳統的編譯器最佳化概念，它們對於高效約束解碼特別有用。

一是規則內聯。在指定的上下文無關語法中，可能有許多片段規則，即只有少數元素的規則，然後在下推自動機中將其轉換為小的 FSA（有限狀態自動機）。

為了解決這個問題，研究者為片段規則引入了一種自動內聯策略。他們迭代地選擇不引用其他規則的規則並將它們內聯到父規則中。為了避免自動機大小的爆炸式增長，研究者將內聯規則和內聯結果的大小限制為常量。該內聯過程幾乎消除了片段規則，從而提高了 token 檢查的效率並增強了上下文擴充套件的有效性。

二是下推自動機節點合併。對於下推自動機，在許多情況下，歧義來自具有相同標籤的節點的多個外向邊。在匹配 token 時，如果到達此節點，並且下一個字元恰好與標籤匹配，則匹配堆疊將被拆分為多個堆疊，每個外向邊一個。堆疊數量增多會增加計算量，這是因為需要檢查每個堆疊的上下文相關 token 併合並 token 掩碼。

為了減少這種歧義，節點合併演算法會合並滿足以下兩個條件的後續節點，a）它們由來自同一點的具有相同標籤的邊指向，b）它們沒有被其他邊指向。

以上兩種最佳化保留了自動機的等效性，但減少了節點和邊的數量。執行時，減少了堆疊的數量和 token 檢查所需的計算量，從而加快了掩碼的生成過程。

重疊掩碼生成和 LLM 推理

透過上述最佳化，token 掩碼生成過程顯著加快，但仍需要 CPU 計算。為了進一步消除約束解碼的開銷，研究者將 mask 生成計算與 LLM 推理過程重疊，如下圖 8 所示。
圖片
研究者觀察到，mask 生成過程和 LLM 推理過程可以重疊，原因在於 mask 生成只需要 CPU，並且只依賴於之前生成的 token。LLM 推理過程（除取樣階段外）只需要 GPU，並且也只依賴於之前生成的 token。因此可以將 CPU 上的 mask 生成過程與 GPU 上的 LLM 推理過程並行化。

評估結果

研究者利用 12,000 行核心 C++ 程式碼來實現 XGrammar，並提供了 Python 捆綁包以方便與 LLM 推理框架無縫整合。他們在評估 XGrammar 過程中回答以下幾個問題：

  • XGrammar 能否高效支援約束解碼的每個步驟？

  • XGrammar 能否在 LLM serving 中實現端到端結構化生成的最小開銷？

  • XGrammar 能否部署在更廣泛的平臺上？

語法引擎效率

本節中評估了語法引擎的效能。研究者在 Llama-3.1-8B Instruct 上評估了他們的方法和基線，該模型能夠遵循人類的指令。

結果如下圖 9 所示，在 JSON 模式設定中，XGrammar 可以實現高達 3 倍的加速；在 JSON 語法用例下，可以實現超過 100 倍的加速。與 JSON 模式（更受限制）相比，JSON 的上下文無關語法包含更復雜的規則，因為它可以包含遞迴列表和字典，導致語法引擎更難有效地執行它。

在這兩種情況下，XGrammar 都可以在不到 40 微秒的時間內生成每個 token 的掩碼，使其成為低延遲 LLM 推理的理想選擇。
圖片
端到端 LLM 引擎評估

本節在 LLM serving 設定下來評估 XGrammar。研究者將 XGrammar 整合到端到端 LLM 推理框架中，並與其他 LLM serving 框架進行效率比較。同時，他們還與其他支援結構化生成的 LLM 引擎進行效率比較，包括整合 Outlines 的 vLLM (v0.6.3) 和內建語法引擎的 llama.cpp。

實驗結果如下圖 10 所示，XGrammar 在 CFG 和 JSON 模式的所有基線中實現了最佳的 TTFT 和 TPOT。vLLM 和 llama.cpp 的計算受到其語法引擎更長預處理和每個 token 處理時長的阻礙。

在批次較大的情況下，vLLM 中 TPOT 速度的下降尤為明顯。與現有解決方案相比，XGrammar 引擎總體上可以將輸出 token 的速度提高 80 倍。這種加速來自 XGrammar 帶來的效能最佳化。
圖片
研究者還在下表 1 中研究了語法處理的開銷問題。由於 token 掩碼生成效率和語法 GPU 重疊，語法過程在 TPOT 中幾乎不產生任何開銷。
圖片
跨平臺部署

本節探討如何將 XGrammar 引入各種平臺。研究者利用 Emscripten 將 XGrammar 編譯成 WebAssembly 並構建 JavaScript 捆綁包。他們進一步將 web-binding 與瀏覽器內 LLM 推理框架 WebLLM 整合，以實現結構化生成。

研究者使用 JSON-mode-eval 資料集評估端到端效能，在裝有 Google Chrome 的 MacBook Pro M3 Max（macOS 14.5）上使用 4 位量化模型 Llama-3.1-8B-Instruct，並在裝有 Safari 的 iPhone 14 Pro Max（iOS 18）上使用 Qwen2.5-0.5B-Instruct。

結果如下圖 11 所示，研究者比較了使用 XGrammar 進行結構化生成和非結構化生成時的第一個 token 時間 (TTFT) 和每個輸出 token 時間 (TPOT)，同時確保生成的 token 數量相同。結果表明，XGrammar 在兩種設定下都幾乎實現了零開銷，在支援未來高效能端側智慧體方面具有巨大潛力。
圖片
更多技術細節請參閱原論文。

相關文章