XGrammar 2: Dynamic and Efficient Structured Generation Engine for Agentic LLMs

TL;DR

XGrammar 2, a highly optimized structured generation engine for agentic LLMs, accelerates the mask generation for these dynamic structured generation tasks through a new dynamic dispatching semantics: TagDispatch, and extends the previous PDA-based mask generation algorithm to the Earley-parser-based one.

Abstract

Modern LLM agents are required to handle increasingly complex structured generation tasks, such as tool calling and conditional structured generation. These tasks are significantly more dynamic than predefined structures, posing new challenges to the current structured generation engines. In this paper, we propose XGrammar 2, a highly optimized structured generation engine for agentic LLMs. XGrammar 2 accelerates the mask generation for these dynamic structured generation tasks through a new dynamic dispatching semantics: TagDispatch. We further introduce a just-in-time (JIT) compilation method to reduce compilation time and a cross-grammar caching mechanism to leverage the common sub-structures across different grammars. Additionally, we extend the previous PDA-based mask generation algorithm to the Earley-parser-based one and design a repetition compression algorithm to handle repetition structures in grammars. Evaluation results show that XGrammar 2 can achieve more than 6x speedup over the existing structured generation engines. Integrated with an LLM inference engine, XGrammar 2 can handle dynamic structured generation tasks with near-zero overhead.

Figures (11)

• Figure 1: Some examples of tool calling and conditional structured generation.
• Figure 2: Overview of our approach. We design a new dynamic dispatching semantics, TagDispatch (§\ref{['method_tagdispatch']}), to efficiently support dynamic structured generation tasks. we design a JIT compilation method (§\ref{['method_jit']}) to reduce the preprocessing time. To leverage the sub-structures across different grammars, we designed a cross-grammar caching algorithm (§\ref{['method_cgc']}) based on the Earley parser (§\ref{['method_earley']}. Besides, we also introduce a repetition compression algorithm (§\ref{['method_repetition']}) to handle some long-tail cases of the grammars.
• Figure 3: The basic structure of TagDispatch.
• Figure 4: The perfect cache hit and the partial cache hit. In cross-grammar caching, with a given rule with the FSM $A$, we will check if the token mask caches for the same FSM have been computed. If there is, then it is a cache hit. If the rules share the same lookahead assertion, then it is a perfect cache hit, and we can reuse the token mask cache directly. Otherwise, it is a partial cache hit, and we need to re-check all the uncertain tokens with $A$'s lookahead assertion, and add the new cache into the global cache pool.
• Figure 5: Repetition State Compression. Most states within a repetition structure share highly similar token mask caches. We compress these states into a single representative state and generate only one token mask cache for it. The final few states in the repetition structure are left uncompressed. On the one hand, their token mask caches are often not sufficiently similar to the preceding ones; on the other hand, these states provide an invariant subsequence for the repetition expression, helping to reduce the number of uncertain tokens.