Trie Automata for Constrained Decoding over Large Finite Sets

Published: Conference paper at COLM 2026

Large language models increasingly need to generate structured outputs conforming to predefined schemas, with one common constraint being selection from a finite set of valid strings. Current constrained decoding systems handle this through general-purpose grammar compilation, which becomes prohibitively slow as the number of valid values grows into the thousands, a cardinality wall.

The paper identifies a fundamental mismatch: A deeply nested recursive JSON schema and a flat list of 1,000 tool names both undergo the same compilation pipeline, a fundamental mismatch between general-purpose engines and simple constraints. This uniformity creates a bottleneck for one of the most common constraints in production: selecting one string from a known finite set.

The paper documents real-world enum limits across major providers: OpenAI's structured outputs impose a 1,000 enum limit, Google Gemini fails at approximately 120 enum values, and Anthropic's 180-second compilation timeout implies a similar wall at a few hundred values.

These limits are increasingly consequential in agentic workflows where an LLM must select which tool to invoke from a registry that may contain 500–5,000+ APIs, as well as in zero-shot classification over label sets like product taxonomies (1,500+ categories), ICD-10-CM medical codes (74,719 codes), legal case types (10,000+), and entity linking against knowledge bases with tens of thousands of entries.

The core insight is that different constraint types deserve different enforcement mechanisms. A finite set of strings has exploitable structure: shared prefixes, finite depth, and known cardinality.

The paper introduces the trie automaton, a specialized constrained decoding backend for finite-set constraints that:

1. Builds a character-level trie directly from the set

2. Precomputes vocabulary-aware token masks at each trie node using Aho-Corasick multi-pattern matching to align BPE tokens with character-level trie paths

3. Serves masks via O(1) cached lookups at decode time

The two main contributions are:

1. The trie automaton itself: Combining character-level tries, Aho-Corasick multi-pattern matching, and precomputed token masks to achieve O(valid[st]) per-step masking (empirically 10–100 tokens after 3–4 characters of prefix, yielding effectively constant cost versus the O(V·l) cost of general FSM approaches)

2. Empirical characterization of the cardinality wall across seven tokenizer families (32K–262K vocabulary), showing that matching enforcement mechanisms to constraint structure overcomes scaling bottlenecks

The central algorithmic challenge is BPE-trie alignment: a single BPE token can span multiple trie nodes, and a vocabulary of 32K–262K tokens must be matched against every node. The paper notes that to our knowledge, this alignment problem has not been addressed in the constrained decoding literature; prior trie-based work sidesteps it by operating at token granularity.

The paper shows this reduces to multi-pattern string matching, solvable in time linear in the trie size rather than quadratic in the vocabulary. The Aho-Corasick algorithm builds a finite automaton from the pattern set (vocabulary tokens) in O(V·l) time, then processes any input text in a single linear pass. This reduces precomputation from O(Nchars·V·l) to O((Nchars+V)·l), a factor of V improvement.

The AC automaton depends only on the tokenizer (not the enum set), so it can be built once per tokenizer and reused across all enum schemas, reducing per-schema compilation to just the trie traversal: O(Nchars·l).

The paper compares against GENRE's token-level trie approach, which avoids the BPE-alignment problem by building the trie over token IDs. The character-level construction is more expensive per schema, but the reason to prefer it is scaling with K: "A token-level trie shares prefixes only at token boundaries... its size, and hence its compilation cost, grows roughly linearly with K. The character-level trie merges every shared character, and its dominant cost is the K-independent AC build."

The crossover occurs at K ≈ 1,000: the token trie is faster below K ≈ 1,000 (it simply does less work), but the two cross over there and the char trie is 7× faster by K = 10,000. This crossover sits exactly at the cardinality wall we target.

The paper establishes Proposition 1 (Output equivalence): "For any enum E and any prefix y<t, the constrained distributions of the FSM and the trie automaton are identical... Consequently, greedy decoding and fixed-seed sampling produce identical outputs under both methods."

The proof relies on the Myhill-Nerode theorem: the equivalence classes of LE are exactly the distinct enum-value prefixes (plus a dead state), so the trie is isomorphic to the minimal DFA for LE and their transition functions agree.

The paper also establishes Proposition 2 (Hierarchical cardinality reduction): For balanced partitions of an enum into G groups, the effective per-step cardinality is minimized at Ceff* = ⌈√K⌉ when G = ⌈√K⌉.

The trie achieves 0.65µs per-step valid-token computation versus XGrammar's 5.8µs at K=1,000 (approximately 7× faster, with the paper noting a 9× gap in controlled microbenchmarks: 0.08µs vs 4.5µs).

The trie achieves 2–6.5× faster compilation at K ≥ 300. Compilation is nearly flat (30–40ms for Qwen3-8B across K = 10–10,000; 67ms at K = 100,000) because the cost is dominated by the K-independent AC construction over the vocabulary.

The paper reports: end-to-end vLLM throughput reaches 219 req/s vs. XGrammar's 7.5 req/s at batch size 256 (29×). This 29× gap combines the algorithmic speedup with integration-path savings that only precomputed masks make possible.

The 29× gap compounds two effects:

1. Per-step algorithmic advantage: precomputed mask lookup costs 0.65µs vs. XGrammar's 5.9µs dynamic computation (7×)

2. Integration path: "because the trie's masks are precomputed, it integrates as a stateless LogitsProcessor that returns a cached bitmask per step. XGrammar does not currently support this path; its architecture requires dynamic mask computation via vLLM's guided decoding pipeline with per-request grammar compilation, sequential FSM state management, and scheduling overhead."

The paper notes: "In principle, XGrammar could precompute and cache per-state masks for enum constraints, but doing so would effectively reconstruct the trie: the minimal DFA for a finite set is isomorphic to the trie, so caching its per-state masks yields the same data structure. The trie is thus the natural endpoint of optimizing FSM-based decoding for finite sets."

Across seven model families (32K–262K vocabulary), the trie achieves consistent compilation speedups at K ≥ 1,000: 1.2–6.4×, increasing to 3.5–13.7× at K = 5,000–10,000. Per-step cost is unaffected by vocabulary size: Mistral-32K 0.60µs, GPT2-50K 0.62µs, OLMo-100K 0.63µs, Mistral Small-131K 0.64µs, Qwen3-151K 0.65µs, gpt-oss-200K 0.66µs, Gemma3-262K 0.67µs.

The trie's memory scales modestly: 0.9 MB at K = 10,000, 8 MB at K = 100,000, vs. 2 GB for the FSM. The AC automaton requires 150 MB for the largest vocabulary (Gemma3 262K), amortized across all enum schemas sharing that tokenizer.

The trie produces identical outputs to FSM-based constrained decoding (Proposition 1, verified on 1,000 samples). Across 24 model×dataset settings on four public classification benchmarks (TREC, MASSIVE, Banking77, CLINC150), constrained decoding achieves the highest accuracy in 21/24 cases while guaranteeing 100% validity; unconstrained decoding reaches ≥95% validity in only 8/24.

Key findings:

• The trie acts as a no-cost safety net for strong models (Gemma3 12B: matches unconstrained accuracy while eliminating the 0.2–0.5% failure rate)

• It is a critical enabler for weaker ones (Qwen3-1.7B Banking77: 24.8% with trie vs. 4.8% unconstrained at 6.6% validity)

• The think-then-answer strategy is particularly effective with the trie: on Qwen3-8B, think+trie achieves the best accuracy in all four datasets, with gains of up to 21.5 points over single-pass unconstrained (TREC: 63.1 vs. 36.3)

At K = 500–5,000, the trie guarantees 100% validity while unconstrained drops to 84–98%.

The paper benchmarks LLGuidance (v1.6.1) and finds: "LLGuidance achieves near-zero compilation (0.6–24ms) but per-step masking costs 73–141µs. XGrammar balances both (3–239ms compilation, 5–10µs masking). The trie minimizes per-step cost (0.65µs) through precomputation, with nearly-flat compilation (30–40ms)."

The paper notes: The two are complementary: LLGuidance excels at schema diversity with negligible startup, while the trie exploits finite-set structure for per-step speedups that compound in batch serving.

In batch serving at B=128: LLGuidance masking (3.7ms) consumes 37% of the GPU forward pass, becoming the throughput bottleneck; XGrammar's 783µs is 7.8%; the trie's 10µs is negligible (0.1%).

To isolate the algorithmic contribution, the paper implements a controlled comparison: precomputing XGrammar's per-state bitmasks for all DFA states and serving them via cached lookup. Results on Qwen3-8B (K=1,000): both have 6,786 states/nodes (identical, confirming Proposition 1), but precomputation takes 33ms for the trie vs. 6.5s for XGrammar (196× slower). This confirms that the trie is the efficient algorithm for precomputing per-state masks over finite sets.

The paper sketches two preliminary extensions for K > 50,000:

1. Hierarchical schema rewriting: partitioning the enum into G ≈ √K groups, reducing per-step cardinality from K to √K

2. Speculative short-circuiting: using a lightweight scoring function to identify a top-k shortlist before applying trie-constrained decoding

The paper also validates that the dispatch principle generalizes beyond finite sets, showing constraint-aware dispatch, matching specialized engines to constraint structure, yields large speedups whenever the constraint has exploitable regularity, not just for finite sets. For fixed-format strings (e.g., YYYY-MM-DD), a character-position mask engine achieves compilation speedups ranging from 2× (UUID on 32K vocab) to 7,939× (date on 262K vocab) versus xgrammar regex compilation.

The paper acknowledges: The trie automaton is specialized for flat finite-set constraints; production schemas that also contain nested objects or arrays still require a general-purpose backend for the structural portions. Additional limitations include: the per-step speedup matters most in batch serving; for dynamic one-shot schemas at small K, LLGuidance's near-zero compilation may outweigh the trie's per-step advantage; at extreme K, validity is guaranteed but accuracy is bounded by the model, not the decoder; and end-to-end throughput is measured only on vLLM, with the 29× figure being a vLLM-specific figure that should not be read as engine-independent.

The paper concludes: "This work addresses a mismatch in constrained decoding: applying general-purpose automata to constraints with exploitable structure. We solve it for finite-set constraints using Aho-Corasick precomputation, expanding the practical enum limit from hundreds to tens of thousands. The trie achieves 7× faster per-step masking through precomputed lookups; because precomputed masks also enable a stateless serving path that bypasses the guided decoding pipeline, the advantage compounds to 29× higher end-to-end batch throughput. Across seven tokenizer families and six models, the trie maintains sub-100ms compilation and flat per-step cost while guaranteeing 100% output validity."

The paper's broader principle: "matching enforcement mechanisms to constraint structure, generalizes beyond enums to any constraint with exploitable regularity. As structured generation becomes central to agentic systems, we expect constraint-specialized backends to become the norm rather than the exception."

Based on the paper, here are the specific improvements I can make to AI systems:

• What I can do: When an AI system needs to output one value from a known finite set (e.g., tool names, classification labels, medical codes), I can build a character-level trie automaton with precomputed token masks using Aho-Corasick multi-pattern matching.

• Improved capability: Handle enum constraints up to 100,000+ values with sub-100ms compilation and 0.65µs per-step masking, versus current systems that fail or become prohibitively slow beyond 1,000 values. This enables:

• Tool selection from registries with 5,000+ APIs

• Zero-shot classification over 74,719 ICD-10-CM medical codes

• Entity linking against knowledge bases with tens of thousands of entries

• What I can do: Automatically detect constraint type (finite set, regex, fixed-format string, nested schema) and route to the optimal enforcement mechanism rather than using one general-purpose grammar compiler for everything.

• Improved capability:

• For finite sets: use trie automaton (7× faster per-step, 29× end-to-end throughput)

• For fixed-format strings (dates, UUIDs): use character-position mask engines (up to 7,939× faster compilation)

• For complex nested schemas: fall back to general-purpose FSM

• This matches the paper's principle: different constraint types deserve different enforcement mechanisms

• What I can do: Because trie masks are precomputed, I can serve constrained decoding as a stateless LogitsProcessor that returns cached bitmasks, bypassing the guided decoding pipeline entirely.

• Improved capability: Achieve 219 req/s vs. 7.5 req/s (29×) in batch serving at batch size 256, eliminating per-request grammar compilation, sequential FSM state management, and scheduling overhead. This is critical for production agentic systems serving many concurrent requests.

• What I can do: For enums with K > 50,000, partition into G ≈ ⌈√K⌉ groups, reducing effective per-step cardinality from K to √K.

• Improved capability: Handle extremely large label sets (e.g., 1 million+ product categories) while maintaining constant-time masking, by first selecting the group then the specific value within it.

• What I can do: Use a lightweight scoring function to identify a top-k shortlist before applying trie-constrained decoding.

• Improved capability: For very large enums, reduce the effective search space to the most likely candidates (e.g., top-100), then apply exact trie constraints only to that subset, combining speed with guaranteed validity.

• What I can do: Integrate the trie automaton with a two-phase generation strategy: first allow free-form reasoning, then constrain the final answer to the enum set.

• Improved capability: On Qwen3-8B, this achieves up to 21.5 points accuracy improvement over single-pass unconstrained decoding (TREC: 63.1 vs. 36.3), while guaranteeing 100% output validity. This is particularly valuable for weaker models that otherwise produce invalid outputs 93% of the time on classification tasks.

• What I can do: Build the Aho-Corasick automaton once per tokenizer (not per enum schema) and reuse it across all enum constraints sharing that tokenizer.

• Improved capability: For a production system serving many enum schemas with the same model, amortize the 150MB AC automaton cost across all schemas, reducing per-schema compilation to just trie traversal: O(Nchars·l) instead of O(Nchars·V·l), a factor of V (vocabulary size) improvement.

• What I can do: Provide formal guarantees (Proposition 1) that constrained decoding with the trie produces identical outputs to FSM-based methods for greedy decoding and fixed-seed sampling.

• Improved capability: AI systems can safely switch to the faster trie backend without any behavioral change, verified on 1,000 samples across 24 model×dataset settings. This eliminates the 0.2–0.5% failure rate of strong models and enables weaker models to achieve 100% validity where they previously failed 93% of the time.

Large language models increasingly need to generate structured outputs that conform to predefined schemas, with one common constraint being selection from a finite set of valid strings. Current constrained decoding systems handle this through general-purpose grammar compilation, which becomes prohibitively slow as the number of valid values grows into the thousands, a cardinality wall. We introduce the trie automaton, a specialized mechanism that exploits finite-set structure (shared prefixes, bounded depth, known cardinality) via Aho-Corasick multi-pattern matching to precompute per-node token masks. The trie achieves 7X faster per-step valid-token computation (0.65 us vs. 5.8 us) compared to XGrammar, one of the primary backends in vLLM and SGLang, and 2--6.5X faster compilation at K >= 300. Because precomputed masks enable a stateless serving path that bypasses the guided decoding pipeline, this advantage compounds in batch serving: end-to-end vLLM throughput reaches 219 req/s vs. XGrammar's 7.5 req/s at batch size 256 (29X). The 29X combines the algorithmic speedup with integration-path savings that only precomputed masks can unlock. Across seven tokenizer families (32K--262K vocabulary), the trie maintains sub-100ms compilation up to K = 10,000 and flat per-step cost regardless of set size, while guaranteeing 100% output validity.