不管是編寫和調試代碼，還是通過函數調用來使用外部工具，又或是控制機器人，都免不了需要 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 在兩種設置下都幾乎實現了零開銷，在支持未來高性能端側智能體方面具有巨大潛力。

更多技術細節請參閱原論文。