GLRMask GLRMask

Suppose you want to guarantee that an LLM’s output follows a grammar. How do we do this?

The dumb way is to sample repeatedly until we get a token the parser no longer rejects:

while generating:
    logits = model()
    token = sample(logits)
    while not constraint.allows(token):
        token = sample(logits)
    constraint.commit(token)
    model.commit(token)

The problem with rejection sampling is that it blocks the whole decoding pipeline. Nothing can continue until we’ve decided on a token that’s valid, and in the worst case could end up working through almost the entire vocabulary while the GPU sits idle.¹

The alternative is to work out which tokens are legal up front, as a mask over the vocabulary, while the model is busy:

while generating:
    parallel:
        logits = model()
        mask = constraint.get_mask()

    logits[~mask] = -np.inf
    token = sample(logits)

    constraint.commit(token)
    model.commit(token)

get_mask() runs in parallel with the forward pass to determine which model tokens are legal from the current state. The sampler picks one of the legal tokens, then commit(token) advances the lexer and parser using the token that was actually chosen.

If the mask is ready by the time the forward pass finishes, its cost is hidden. If not, it stalls generation. So get_mask has to be fast.

So masked generation works by restricting the sampler to grammatically valid tokens. It literally can’t go wrong²! But what does the mask actually compute?

Suppose the bytes generated so far are $u$, the language we want to generate is $L$ (a language is just a set of strings; here, $L$ is the set allowed by the grammar), and model token $v$ corresponds to the byte string $\beta(v)$. The next-token mask is

$$\mathrm{Mask}_L(u) = \{v\in\mathcal V:\exists w,\;u\beta(v)w\in L\}.$$

Let’s think about what this means. A token is legal if appending its bytes still leaves some way to complete the output into a grammatically valid string. And we need to check that for every token in the vocabulary, at every generation step.

A bruteforce approach might look something like this:

• Build a mask over the entire LLM vocabulary.
• For each LLM token $t$, check whether appending $\beta(t)$ keeps the parse valid.
• If so, set its bit.

With a vocabulary of 100k+ tokens, this probably won’t go too well.

A straightforward way to generate the mask is to explore candidate token bytes online. In practice you put the vocabulary in a byte trie, share work between tokens with common prefixes, cut off branches as soon as they become impossible, and use further shortcuts to settle whole groups of tokens at once. This can already make mask generation very fast, particularly for JSON. llguidance is a strong implementation of this kind of approach.

But how much work remains depends on the grammar and the current parse state, and this can show up most clearly in tail latency (e.g. p99.9).

GLRMask tries to move as much of this work ahead of time as possible. Compilation is heavier, sometimes substantially so, but it only has to happen once, and the compiled constraint can then be reused for every subsequent request.

Weighted automata

To see where we’re headed, it might be useful to understand what we’re trying to build and why it should be possible, at least vaguely.

The central object in GLRMask is the Parser DWA. It eats the current parser stack from the top down, one symbol at a time, and returns the mask directly.

Okay, great: a magical automaton! But why should such a thing exist?

Here’s one way to think about it.

Fix a lexer state $q$ and a model token $v$. Depending on how $v$‘s bytes are lexed, it may produce one or more terminal sequences. Each of those sequences, in turn, makes the parser perform some actions on the stack. For an LR parser, those actions are shifts, reduces, and gotos. But regardless of how the particular parsing paradigm formulates those actions, they must boil down to a sequence of reads and writes on the parser stack.

The lexer, parser, and vocabulary are known ahead of time. So for a given lexer state, by combining the terminal sequences that we know the lexer emits with the stack reads and writes we know the parser performs for them, we should be able to characterise ahead of time exactly how an LLM token acts on the parse state.

Viewed from before the token is consumed, when we just want to know whether or not this token will be accepted, this gives us a set of requirements on the old stack: which parser states need to be sitting at the top of the stack, and in what order. So for each $(q,v)$, there is a set of parser-stack prefixes from which $v$ can be consumed legally.

For reasons we’ll come back to later, it is possible to ensure that an LR parser can produce only a bounded amount of parser activity while processing any one terminal. Furthermore, as long as each terminal consumes at least one byte, we know that each model token can produce only a bounded amount of lexer activity. Putting these together, a model token can induce only a bounded amount of parser activity, and therefore its acceptance depends on only a bounded prefix of the existing parser stack. Thus each $(q,v)$ induces a regular language over stack prefixes, recognizable by a finite automaton.

Now imagine one of these automata for every $(q,v)$ pair. An LLM vocabulary may contain a hundred thousand or more tokens, and a lexer can easily have tens of thousands of states, so there are a lot of them. But there is also a lot of shared structure among them.

GLRMask exploits this shared structure by storing the whole family in one weighted automaton. Common stack-prefix paths are represented only once, while the weights record which $(q,v)$ pairs use each part of the graph.

A weighted automaton is just an automaton whose transitions carry weights, with rules for how those weights combine. In GLRMask, a weight is a Boolean array

$$w\in\mathbb B^{Q_{\mathrm{lex}}\times\mathcal V},$$

indexed by lexer states and model tokens. Concretely, $w_{q,v}=1$ means that the transition applies to the pair $(q,v)$.

For the current lexer state $q$, each transition’s Boolean weight narrows a running set of allowed tokens as it is taken.

To build the Parser DWA, we handle the lexer and parser separately:

• On the lexer side, the Terminal DWA records which terminal sequences a model token can produce.
• On the parser side, the stack-effect automata describe what each terminal can do to the parser stack.

Terminal DWA

The Terminal DWA records, for every lexer state and model token, the terminal sequences that token can produce.

Fix a lexer state $q$ and a model token $v$, and let $\beta(v)$ be the bytes of that token.

A single LLM token might:

• expand into multiple terminals ("}\n" could lex as RBRACE NEWLINE),
• end partway through a terminal (halfway through a string literal, halfway through a number, halfway through true),
• be lexically ambiguous,
• or produce no terminals yet (e.g. it adds bytes inside a string, but doesn’t close it).

So even for a fixed lexer state and model token, there may be more than one terminal sequence we need to consider.

We store these sequences in a deterministic weighted automaton, or DWA.

A path through the Terminal DWA represents a terminal sequence, and its weights record which lexer-state/model-token pairs can produce it.

Conceptually, for each pair $(q,v)$, run $\beta(v)$ through the lexer from $q$, and add the terminal sequences it can produce as paths from the root, marking those paths with $(q,v)$. If the token ends partway through a lexeme, we also need to account for its possible completions. The model token is only legal if at least one such completion can eventually be accepted by the parser. So branch the path once for each terminal that lexeme could still become, and end each branch with that terminal.

In practice, of course, we traverse a vocabulary trie rather than iterating through each full token individually. But you get the idea. Lots of tries! Just assume, from now on, that any time I mention iterating over the vocabulary, I really mean traversing it as a trie.

Stack effects

For an LR parser, consuming one terminal means performing zero or more reductions followed by a shift.

A shift pushes a state to the stack. Reductions are a bit more involved. But what matters really is that ultimately everything an LR parser does is a sequence of operations that read from (and consume) and write to the top of a stack.

GLRMask abstracts away parser behaviour as sequences of stack effects: $p^-$ for “read and remove parser state $p$,” and $p^+$ for “write parser state $p$.”

A stack effect sequence such as

$$p^-r^+s^+$$

therefore says: the old stack must begin with $p$; consume it; then write $r$ and $s$.

Different terminals can have different stack effects, and a terminal may have more than one possible effect. For each terminal $t$, let $E_t$ denote the set of its possible stack-effect sequences.

Stack effects also compose nicely. Suppose one terminal has the effect

$$p^-r^+s^+$$

and the next has

$$s^-r^-u^+.$$

Composing them gives

$$p^-r^+s^+s^-r^-u^+.$$

Let’s think about what this means. One terminal pushes states $r \, s$ onto the stack, and the next terminal immediately reads $s \, r$ back off. As far as the stack is concerned, those operations cancel:

$$p^-r^+s^+s^-r^-u^+ \;\longrightarrow\; p^-r^+r^-u^+ \;\longrightarrow\; p^-u^+.$$

The two-terminal sequence therefore only requires the original stack to begin with $p$. And if it does, it simply replaces that $p$ with $u$. Everything involving $r$ and $s$ happens internally between the two terminals.

So, the cancellation rule is simply

$$p^+p^-\to\epsilon$$

A mismatched write/read pair, on the other hand, makes the path invalid:

$$p^+q^-\longrightarrow\bot \qquad(p\ne q).$$

For any set of stack effects, repeatedly applying these cancellations takes every surviving effect to a normal form consisting of some reads from the original stack followed by some writes describing the new stack:

$$p_1^-p_2^-\cdots p_m^-q_1^+q_2^+\cdots q_n^+.$$

Finally, for each terminal $t$, let $E_t$ be the set of stack effects by which $t$ can be consumed.

Rather than enumerate $E_t$ explicitly, we represent it with a small automaton over stack operations. Each path through the automaton corresponds to one stack effect in $E_t$.

Parser DWA

Let’s recap. We’ve built two kinds of automata.

The Terminal DWA compactly represents the terminal sequences that model tokens can produce, with each path weighted by the lexer states and model tokens for which it applies.

Then for each terminal, we have a small automaton describing its effect on the parser stack.

So the Terminal DWA connects model tokens to terminals, while the stack-effect automata connect terminals to stack effects.

$$\begin{array}{rcl} \text{model tokens} & \xrightarrow{\text{Terminal DWA}} \text{terminals} \xrightarrow{\text{stack-effect automata}} & \text{stack effects} \end{array}$$

It’s time to cut out the middleman.

$$\begin{array}{rcl} \text{model tokens} & \xrightarrow{\hspace{7em}\text{Parser DWA}\hspace{7em}} & \text{stack effects} \end{array}$$

Wherever the Terminal DWA has an edge labelled $t$, splice in the stack-effect automaton for $t$. A path that previously represented a terminal sequence now expands into paths describing the corresponding parser-stack operations.

Writes from one terminal can run into reads from the next. Adjacent write/read symbols cancel when they match, and kill the path when they don’t. We keep cancelling until all of these internal write/read pairs are gone.

Any remaining writes sit at the end of the path, describing the net additions to the stack after the required stack prefix has been read.

For mask generation, the trailing writes aren’t needed. We could keep them and use the same automaton to update the parser state during commit, but that adds more complexity than it is worth. So instead we just drop them.

What we are left with is a weighted automaton over parser-stack symbols, whose paths read prefixes of the parser stack directly.

And that’s it: the Parser DWA!

I think it’s helpful to see the Parser DWA as basically a compressed representation of the family $L_{q,v}$. For any particular $(q,v)$, the corresponding read-only automaton for $L_{q,v}$ can in principle be recovered from the Parser DWA by keeping only transitions whose weights contain $(q,v)$. So, effectively, slapping Boolean-array weights onto transitions lets us pack a whole family of automata into one, with the weights encoding the distinctions among them.

Finally, one nice property of the Parser DWA is that it only ever needs to read a bounded prefix of the parser stack. First, it should be fairly clear that each terminal sequence emitted by the lexer for a model token has bounded length. Every terminal emitted during the scan consumes at least one byte, and a model token contains only finitely many bytes. And because the vocabulary is finite, there is a fixed maximum token length. For each terminal, the stack effect is bounded too. After grammar normalization, the number of consecutive reductions between shifts is bounded by a constant depending only on the grammar [Aycock et al., 2001]. The Parser DWA is just the composition of these bounded terminal sequences and bounded stack effects. So it too must be bounded. Nice!

Appendix: The construction in equations

A model token may produce one or more terminal sequences. Each sequence induces a stack effect. The Parser DWA keeps only the read prefix of that effect: the part that must be present on the parser stack.

The weights are Boolean masks indexed by lexer state and model token:

$$K=\mathbb B^{Q_{\mathrm{lex}}\times\mathcal V}, \qquad W_{qv}\in\mathbb B.$$

Weights combine pointwise by union and intersection:

$$(W\oplus W')_{qv}=W_{qv}\lor W'_{qv}, \qquad (W\otimes W')_{qv}=W_{qv}\land W'_{qv}.$$

A deterministic weighted automaton is

$$\mathcal A=(S,s_0,\tau,\varphi), \qquad \tau:S\times\Sigma\rightharpoonup S\times K.$$

If its run on $a_1\cdots a_m$ is

$$s_0\xrightarrow{a_1/W_1}s_1 \xrightarrow{a_2/W_2}\cdots \xrightarrow{a_m/W_m}s_m,$$

its value is the intersection of the transition weights and the final weight:

$$\big[\!\big[\mathcal A\big]\!\big](a_1\cdots a_m) =W_1\otimes\cdots\otimes W_m\otimes\varphi(s_m).$$

The Terminal DWA stores which lexer-state/token pairs produce each grammar terminal sequence. If $\beta(v)$ is the byte string of model token $v$, then

$$\big[\!\big[\mathcal T_{\mathrm{DWA}}\big]\!\big](r)_{qv}=1 \quad\Longleftrightarrow\quad r\in\mathrm{Lex}(q,\beta(v)).$$

Here $r=t_1\cdots t_k$ is a sequence of grammar terminals.

Each terminal $t$ has a language $E_t$ of parser-stack effects. For parser stack alphabet $\Gamma$, introduce a read and a write symbol for every parser state:

$$\Gamma^\pm=\{p^-,p^+:p\in\Gamma\}.$$

An effect $u^-z^+\in E_t$ means that consuming $t$ can replace the top prefix $u$ with $z$:

$$u\alpha\overset{t}{\mapsto}z\alpha \qquad(\alpha\in\Gamma^*).$$

Adjacent writes and reads cancel when their parser states agree. A disagreement kills the path:

$$p^+p^-\to\epsilon, \qquad p^+q^-\to\bot\quad(p\ne q).$$

Writing $\preceq$ for prefix order, composition of two stack effects is

$$(u^-x^+)\odot(y^-z^+) := \begin{cases} u^-a^+z^+, & x=ya,\\ u^-a^-z^+, & y=xa,\\ \bot, & x\not\preceq y\ \land\ y\not\preceq x. \end{cases}$$

Composition extends from individual effects to their languages:

$$E\odot F =\{x\odot y:x\in E,\ y\in F,\ x\odot y\ne\bot\}.$$

So the net effects of a terminal sequence $r=t_1\cdots t_k$ are

$$E_r=E_{t_1}\odot\cdots\odot E_{t_k}, \qquad E_\epsilon=\{\epsilon\}.$$

For mask generation, only the read prefix matters: the part that had to exist before the token began. Project each net effect onto that prefix:

$$\widehat E_r = \{u\in\Gamma^*: \exists z\in\Gamma^*,\ u^-z^+\in E_r\}.$$

A terminal sequence is parseable from stack $\sigma$ exactly when one of these read prefixes occurs at the top of $\sigma$:

$$\mathrm{Parseable}(r,\sigma) \quad\Longleftrightarrow\quad \exists u\in\widehat E_r,\ \exists\alpha:\ \sigma=u\alpha.$$

A model token may produce several terminal sequences, so its read-prefix language is their union:

$$\widehat E_{q,v} = \bigcup_{r\in\mathrm{Lex}(q,\beta(v))}\widehat E_r.$$

The Parser DWA stores all of these languages in one weighted automaton:

$$\big[\!\big[\mathcal P_{\mathrm{DWA}}\big]\!\big](\rho)_{qv}=1 \quad\Longleftrightarrow\quad \rho\in\widehat E_{q,v}.$$

The Parser DWA therefore satisfies:

$$\begin{aligned} \big[\!\big[\mathcal P_{\mathrm{DWA}}\big]\!\big](\rho)_{qv}=1 &\quad\Longleftrightarrow\quad \rho\in\widehat E_{q,v}\\ &\quad\Longleftrightarrow\quad \exists r:\ \big[\!\big[\mathcal T_{\mathrm{DWA}}\big]\!\big](r)_{qv}=1 \ \land\ \rho\in\widehat E_r. \end{aligned}$$

Ignoring longest-match exclusions for now, $\mathrm{ParserOK}(q,\sigma,v)$ holds exactly when some top prefix $\rho$ of $\sigma$ has Parser-DWA value containing $(q,v)$:

$$\mathrm{ParserOK}(q,\sigma,v) \quad\Longleftrightarrow\quad \exists\rho,\alpha:\ \sigma=\rho\alpha \ \land\ \big[\!\big[\mathcal P_{\mathrm{DWA}}\big]\!\big](\rho)_{qv}=1.$$

Suppose one token emits the two terminals x and ], with effects

$$b^-b^+p^+\in E_{\texttt{x}}, \qquad p^-b^-c^+\in E_{\texttt{]}}.$$

Their internal parser work disappears:

$$(b^-b^+p^+)(p^-b^-c^+) \longrightarrow b^-b^+b^-c^+ \longrightarrow b^-c^+.$$

The token may perform several reductions and shifts internally, but the mask query only has to check that the old stack begins with $b$.

In one line:

$$\big[\!\big[\mathcal P_{\mathrm{DWA}}\big]\!\big](\rho)_{qv} = \bigvee_{r\in\mathcal T^*} \left( \big[\!\big[\mathcal T_{\mathrm{DWA}}\big]\!\big](r)_{qv} \land [\rho\in\widehat E_r] \right).$$

Footnotes

1. Assuming sampling without replacement. ↩

2. Constrained decoding can hurt task performance even when it guarantees syntactic validity. Greedily masking and renormalizing over the tokens that are legal at the current step can distort the model’s original distribution, while token/grammar misalignment can force unnatural tokenizations and reduce accuracy. See Grammar-Aligned Decoding and DOMINO. ↩