Source
Dong et al. (2024), LIT-552 — XGrammar. Read as NOTE-294.
The practice
If downstream code parses the output, constrain decoding to the grammar, and do not rely on prompting and retrying. At each step, set the logits of tokens that would violate the grammar to −∞ before sampling. Unconstrained, Llama-3.1 8B produced valid function-call JSON 62% of the time and valid XML 80%. The failures were extra prose and wrong types. With the mask, both are 100% by construction.
Make the mask cheap enough not to matter:
- Precompute per automaton node. Most tokens’ validity depends only on the current position in the current rule (the stack top), not on the whole parse stack. Cache those verdicts per node, stored as whichever of the accept or reject lists is smaller. Check only the context-dependent remainder (1,134 of 128k tokens for JSON) against the full stack
- Build the mask on the CPU while the GPU computes logits, and synchronize only before sampling. Build the per-grammar cache during prefill
- Use a full context-free grammar when the structure nests. Regex engines cannot express arbitrarily nested JSON
| optimization (cumulative) | mask latency per token |
|---|---|
| PDA over the full vocabulary | 65.8 ms |
| + node merging | 38.3 ms |
| + per-node mask cache | 0.154 ms |
| + rule inlining | 0.035 ms |
| + context expansion | 0.018 ms |
End to end on Llama-3.1 8B (MLC-LLM), TPOT goes from 6.2 to 6.3 ms at batch 1 and from 9.0 to 9.2 ms at batch 16.
The design and the table both come from LIT-552, XGrammar, which introduced the split into context-independent tokens (all but 1,134 of Llama-3.1’s 128k for JSON) cached per node, ablated each optimization in its Table 3, and measured the overlap’s end-to-end cost in MLC-LLM in its Table 2.
Conditions
- It guarantees syntax, not correctness. The paper’s “accuracy” is syntactic validity. Masking tokens and renormalizing the rest can change the distribution the model samples from, and whether task accuracy holds up under the constraint is not measured here. Check the end task (SOTA-308 makes the same point for retrieval)
- A spare CPU core per stream, for the overlap to hide the cost
- Measured on one model family against library versions from late 2024
Known implementations
- XGrammar (as a library), used as the backend in SGLang and MLC-LLM