不管是編寫和調(diào)試代碼，還是通過函數(shù)調(diào)用來使用外部工具，又或是控制機(jī)器人，都免不了需要 LLM 生成結(jié)構(gòu)化數(shù)據(jù)，也就是遵循某個(gè)特定格式（如 JSON、SQL 等）的數(shù)據(jù)。

但使用上下文無關(guān)語法（CFG）來進(jìn)行約束解碼的方法并不高效。針對這個(gè)困難，陳天奇團(tuán)隊(duì)提出了一種新的解決方案：XGrammar。

XGrammar 是一個(gè)開源軟件庫，可實(shí)現(xiàn)高效、靈活且可移植的結(jié)構(gòu)化生成。該團(tuán)隊(duì)在博客中表示：「我們毫不妥協(xié)地實(shí)現(xiàn)了這三個(gè)目標(biāo)，并致力于一個(gè)核心使命：將靈活、零開銷的結(jié)構(gòu)化生成帶到任何地方。」

• 論文標(biāo)題：XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models
• 論文地址：https://arxiv.org/pdf/2411.15100
• 代碼地址：https://github.com/mlc-ai/xgrammar

對于結(jié)構(gòu)化生成，一種常用方法是約束解碼。在每個(gè)解碼步驟中，約束解碼都會檢查詞表，并通過將無效 token 的概率設(shè)置為零來過濾掉違反指定結(jié)構(gòu)的 token。為了支持多種多樣的結(jié)構(gòu)格式，需要一種靈活的機(jī)制來指定和檢查這些約束。

使用 JSON 方案實(shí)現(xiàn)約束解碼

上下文無關(guān)語法（CFG）就能提供一種通用方法，即通過一組規(guī)則來定義結(jié)構(gòu)。其中每條規(guī)則都包含一個(gè)字符序列或其他規(guī)則，并允許遞歸組合來表示復(fù)雜的結(jié)構(gòu)。相比于正則表達(dá)式等其它格式，CFG 由于支持遞歸結(jié)構(gòu)，因而能提供更大的靈活性，使其適合描述 JSON、SQL 和領(lǐng)域特定語言（DSL）等常見語言。

下圖展示了一個(gè)用于數(shù)組和字符串的 CFG，可以清楚地看到其中的遞歸結(jié)構(gòu)。

但是，也正因?yàn)?CFG 很靈活，所以直接將其應(yīng)用于約束解碼的效率并不高。首先，每個(gè)解碼步驟都需要對詞表中每個(gè)可能的 token 解釋 CFG，在 Llama 3.1 中，這個(gè)詞表的大小可能高達(dá) 128k。此外，CFG 解釋需要一個(gè)堆棧狀態(tài)來跟蹤之前匹配的遞歸規(guī)則，因此無法提前計(jì)算和緩存堆棧模式的所有組合。最后，LLM 生成結(jié)果中的每個(gè) token 都包含多個(gè)字符，這些字符可能會跨越語法元素的邊界，并在運(yùn)行時(shí)執(zhí)行期間導(dǎo)致進(jìn)一步的遞歸或堆棧彈出。這種未對齊的邊界問題很棘手，需要在語法執(zhí)行期間小心處理它們。

XGrammar 便是為解決上述難題而生的，并且效果卓越：相比于之前的 SOTA 方法，XGrammar 可以將上下文無關(guān)語法的每 token 延遲減少多達(dá) 100 倍！此外，他們還基于 Llama3.1 模型實(shí)驗(yàn)了集成了 XGrammar 的 LLM serving 引擎；在 H100 GPU 上，這能將通過結(jié)構(gòu)化輸出實(shí)現(xiàn)端到端 LLM serving 的速度提升 80 倍！

該團(tuán)隊(duì)表示：「我們正在開源 XGrammar 并將其集成到主要的開源 LLM 框架中?！?/span>

XGrammar 概覽

如圖 1 所示，Grammar 利用了字節(jié)級下推自動機(jī)（byte-level pushdown automaton）來解釋上下文無關(guān)語法。

這種字節(jié)級設(shè)計(jì)允許每個(gè)字符邊緣包含一個(gè)或多個(gè)字節(jié)，處理不規(guī)則的 token 邊界并支持包含 sub-UTF8 字符的 token 。該自動機(jī)的結(jié)構(gòu)經(jīng)過優(yōu)化以加快匹配速度。

在預(yù)處理階段，會生成一個(gè)自適應(yīng) token 掩碼緩存，它會通過預(yù)先計(jì)算與上下文無關(guān)的 token 來加快運(yùn)行時(shí)的掩碼生成。上下文擴(kuò)展（context extension）能進(jìn)一步提升這種緩存的有效性。

在運(yùn)行時(shí)，token 掩碼緩存會快速生成大部分掩碼，而持續(xù)性執(zhí)行堆棧會高效處理其余的上下文相關(guān) token。

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

具體來說，陳天奇團(tuán)隊(duì)首先得到了一個(gè)見解：雖然無法預(yù)先計(jì)算下推自動機(jī)（PDA）無限多個(gè)狀態(tài)的完整掩碼，但可以預(yù)先計(jì)算掩碼中相當(dāng)一部分（通常超過 99%）的 token。因此，可將這些 token 分成兩類：

• 上下文無關(guān) token：僅通過查看 PDA 中的當(dāng)前位置而不是堆棧即可確定其有效性的 token。
• 上下文相關(guān) token：必須使用整個(gè)堆棧來確定其有效性的 token。

下圖展示了一組上下文相關(guān)和無關(guān) token 的示例。大多數(shù)情況下，上下文無關(guān) token 占大多數(shù)。我們可以預(yù)先計(jì)算 PDA 中每個(gè)位置的上下文無關(guān) token 的有效性，并將它們存儲在自適應(yīng) token 掩碼緩存中。此過程稱為語法編譯（grammar compilation）。

下圖則展示了自適應(yīng)存儲格式。

在運(yùn)行時(shí)，首先檢索來自緩存的上下文無關(guān) token 的有效性。然后，高效地執(zhí)行 PDA 來檢查其余的上下文相關(guān) token。通過跳過運(yùn)行時(shí)檢查大多數(shù) token，便可以顯著加快掩碼生成速度。XGrammar 執(zhí)行時(shí)間的整體工作流程見圖 1。

此外，他們還設(shè)計(jì)了一組額外的算法和系統(tǒng)優(yōu)化方法，以進(jìn)一步提高掩碼生成速度并減少預(yù)處理時(shí)間，包括上下文擴(kuò)展、持續(xù)性執(zhí)行椎棧、下推自動機(jī)結(jié)構(gòu)優(yōu)化、并行式語法編譯。

上下文擴(kuò)展

該團(tuán)隊(duì)提出的方法是檢測語法中每個(gè)規(guī)則的額外上下文信息，并將其用于減少上下文相關(guān) token 的數(shù)量，并進(jìn)一步加快運(yùn)行時(shí)檢查速度。

持續(xù)性執(zhí)行堆棧

為了加快由于多種可能的擴(kuò)展路徑而導(dǎo)致的拆分和合并期間多個(gè)并行堆棧的維護(hù)速度，他們設(shè)計(jì)了一個(gè)基于樹的數(shù)據(jù)結(jié)構(gòu)，可以有效地同時(shí)管理多個(gè)堆棧。

它還可以存儲以前的狀態(tài)并實(shí)現(xiàn)高效的狀態(tài)回滾，從而加快上下文相關(guān) token 的運(yùn)行時(shí)檢查速度。

下推自動機(jī)結(jié)構(gòu)優(yōu)化

研究者進(jìn)行了額外的優(yōu)化，以改進(jìn)下推自動機(jī)的結(jié)構(gòu)，加快最終執(zhí)行的效率。這些優(yōu)化借鑒了傳統(tǒng)的編譯器優(yōu)化概念，它們對于高效約束解碼特別有用。

一是規(guī)則內(nèi)聯(lián)。在指定的上下文無關(guān)語法中，可能有許多片段規(guī)則，即只有少數(shù)元素的規(guī)則，然后在下推自動機(jī)中將其轉(zhuǎn)換為小的 FSA（有限狀態(tài)自動機(jī)）。

為了解決這個(gè)問題，研究者為片段規(guī)則引入了一種自動內(nèi)聯(lián)策略。他們迭代地選擇不引用其他規(guī)則的規(guī)則并將它們內(nèi)聯(lián)到父規(guī)則中。為了避免自動機(jī)大小的爆炸式增長，研究者將內(nèi)聯(lián)規(guī)則和內(nèi)聯(lián)結(jié)果的大小限制為常量。該內(nèi)聯(lián)過程幾乎消除了片段規(guī)則，從而提高了 token 檢查的效率并增強(qiáng)了上下文擴(kuò)展的有效性。

二是下推自動機(jī)節(jié)點(diǎn)合并。對于下推自動機(jī)，在許多情況下，歧義來自具有相同標(biāo)簽的節(jié)點(diǎn)的多個(gè)外向邊。在匹配 token 時(shí)，如果到達(dá)此節(jié)點(diǎn)，并且下一個(gè)字符恰好與標(biāo)簽匹配，則匹配堆棧將被拆分為多個(gè)堆棧，每個(gè)外向邊一個(gè)。堆棧數(shù)量增多會增加計(jì)算量，這是因?yàn)樾枰獧z查每個(gè)堆棧的上下文相關(guān) token 并合并 token 掩碼。

為了減少這種歧義，節(jié)點(diǎn)合并算法會合并滿足以下兩個(gè)條件的后續(xù)節(jié)點(diǎn)，a）它們由來自同一點(diǎn)的具有相同標(biāo)簽的邊指向，b）它們沒有被其他邊指向。

以上兩種優(yōu)化保留了自動機(jī)的等效性，但減少了節(jié)點(diǎn)和邊的數(shù)量。運(yùn)行時(shí)，減少了堆棧的數(shù)量和 token 檢查所需的計(jì)算量，從而加快了掩碼的生成過程。

重疊掩碼生成和 LLM 推理

通過上述優(yōu)化，token 掩碼生成過程顯著加快，但仍需要 CPU 計(jì)算。為了進(jìn)一步消除約束解碼的開銷，研究者將 mask 生成計(jì)算與 LLM 推理過程重疊，如下圖 8 所示。

研究者觀察到，mask 生成過程和 LLM 推理過程可以重疊，原因在于 mask 生成只需要 CPU，并且只依賴于之前生成的 token。LLM 推理過程（除采樣階段外）只需要 GPU，并且也只依賴于之前生成的 token。因此可以將 CPU 上的 mask 生成過程與 GPU 上的 LLM 推理過程并行化。

評估結(jié)果

研究者利用 12,000 行核心 C++ 代碼來實(shí)現(xiàn) XGrammar，并提供了 Python 捆綁包以方便與 LLM 推理框架無縫集成。他們在評估 XGrammar 過程中回答以下幾個(gè)問題：

• XGrammar 能否高效支持約束解碼的每個(gè)步驟？
• XGrammar 能否在 LLM serving 中實(shí)現(xiàn)端到端結(jié)構(gòu)化生成的最小開銷？
• XGrammar 能否部署在更廣泛的平臺上？

語法引擎效率

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

結(jié)果如下圖 9 所示，在 JSON 模式設(shè)置中，XGrammar 可以實(shí)現(xiàn)高達(dá) 3 倍的加速；在 JSON 語法用例下，可以實(shí)現(xiàn)超過 100 倍的加速。與 JSON 模式（更受限制）相比，JSON 的上下文無關(guān)語法包含更復(fù)雜的規(guī)則，因?yàn)樗梢园f歸列表和字典，導(dǎo)致語法引擎更難有效地執(zhí)行它。

在這兩種情況下，XGrammar 都可以在不到 40 微秒的時(shí)間內(nèi)生成每個(gè) token 的掩碼，使其成為低延遲 LLM 推理的理想選擇。

端到端 LLM 引擎評估

本節(jié)在 LLM serving 設(shè)置下來評估 XGrammar。研究者將 XGrammar 集成到端到端 LLM 推理框架中，并與其他 LLM serving 框架進(jìn)行效率比較。同時(shí)，他們還與其他支持結(jié)構(gòu)化生成的 LLM 引擎進(jìn)行效率比較，包括集成 Outlines 的 vLLM (v0.6.3) 和內(nèi)置語法引擎的 llama.cpp。

實(shí)驗(yàn)結(jié)果如下圖 10 所示，XGrammar 在 CFG 和 JSON 模式的所有基線中實(shí)現(xiàn)了最佳的 TTFT 和 TPOT。vLLM 和 llama.cpp 的計(jì)算受到其語法引擎更長預(yù)處理和每個(gè) token 處理時(shí)長的阻礙。

在批量較大的情況下，vLLM 中 TPOT 速度的下降尤為明顯。與現(xiàn)有解決方案相比，XGrammar 引擎總體上可以將輸出 token 的速度提高 80 倍。這種加速來自 XGrammar 帶來的性能優(yōu)化。

研究者還在下表 1 中研究了語法處理的開銷問題。由于 token 掩碼生成效率和語法 GPU 重疊，語法過程在 TPOT 中幾乎不產(chǎn)生任何開銷。

跨平臺部署

本節(jié)探討如何將 XGrammar 引入各種平臺。研究者利用 Emscripten 將 XGrammar 編譯成 WebAssembly 并構(gòu)建 JavaScript 捆綁包。他們進(jìn)一步將 web-binding 與瀏覽器內(nèi) LLM 推理框架 WebLLM 集成，以實(shí)現(xiàn)結(jié)構(gòu)化生成。

研究者使用 JSON-mode-eval 數(shù)據(jù)集評估端到端性能，在裝有 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。

結(jié)果如下圖 11 所示，研究者比較了使用 XGrammar 進(jìn)行結(jié)構(gòu)化生成和非結(jié)構(gòu)化生成時(shí)的第一個(gè) token 時(shí)間 (TTFT) 和每個(gè)輸出 token 時(shí)間 (TPOT)，同時(shí)確保生成的 token 數(shù)量相同。結(jié)果表明，XGrammar 在兩種設(shè)置下都幾乎實(shí)現(xiàn)了零開銷，在支持未來高性能端側(cè)智能體方面具有巨大潛力。

更多技術(shù)細(xì)節(jié)請參閱原論文。