GLRMask

GLRMask is a Rust and Python library for grammar-constrained LLM generation. It accepts JSON Schema and general context-free grammars, then restricts each decoding step to model tokens that can still lead to a valid result.

The runtime is built for low next-token mask latency, including the slow end of the distribution on complex grammars. A compiled constraint is immutable and reusable, while each generation request gets its own small mutable state.

(archived) Inside GLRMask: constrained decoding with weighted automata starts with pathological model tokens and follows the problem through to the automaton GLRMask runs over the parser stacks.

python -m pip install glrmask

cargo add glrmask

Using it

This example constrains a Llama model to one of three JSON strings:

import numpy as np
from llama_cpp import Llama
from torch import from_numpy
from torch.distributions import Categorical

import glrmask


llm = Llama(model_path="model.gguf", logits_all=True)
vocab = glrmask.Vocab.from_llama_cpp(llm)
end_token_ids = vocab.llama_cpp_end_token_ids

schema = '{"type":"string","enum":["positive","negative","neutral"]}'
constraint = glrmask.Constraint.from_json_schema(
    schema,
    vocab,
    end_token_ids=end_token_ids,
)

prompt = "Classify this review: The story dragged badly. Sentiment: "
llm.eval(llm.tokenize(prompt.encode()))

state = constraint.start()
generated = []

for _ in range(64):
    logits = llm.scores[llm.n_tokens - 1]
    mask = state.mask(llm.n_vocab())
    logits[~mask] = -np.inf

    token = Categorical(logits=from_numpy(logits)).sample().item()
    llm.eval([token])
    state.commit_token(token)
    generated.append(token)

    if token in end_token_ids:
        break

print(llm.detokenize(generated).decode())

mask() depends only on the current constraint state, so it can run alongside the model forward pass. commit_token() advances the lexer and GLR parser after sampling.

The compiled Constraint can be serialized and cached. DynamicConstraint has the same runtime interface when a cached compilation is unavailable. It starts much faster, with higher mask latency, while the normal constraint is compiled for later requests.

Grammars

JSON Schema is the main practical use case. GLRMask also accepts its own EBNF-like GLRM syntax, along with Lark and EBNF grammars. GLRM can refer to exact model-token IDs as terminals, which is useful for end tokens and other tokenizer-level structure that has no byte representation.

General context-free grammars bring lexer ambiguity, reductions, and multiple parser stacks into the constraint state. GLRMask keeps those alternatives in a graph-structured stack rather than choosing one parse early.

Mask generation

Committing a sampled token uses an incremental GLR parser. Generating the next mask uses a different path: a deterministic weighted automaton reads the current parser stacks.

Its transition weights are sets of model tokens. Traversing a path intersects those sets; merging parser alternatives unions them. The result is the vocabulary mask for the current state. Mask generation therefore works over the compact parse configuration instead of speculatively advancing the parser for each candidate token.

The persistent weighted graph-structured stack used by the runtime is also available separately as weighted-gss.

Repository

The repository contains the Rust crate, Python extension, JSON Schema importer, grammar compiler, runtime, examples, tests, and benchmark harnesses. Published wheels include the native extension.

The implementation has been exercised across the JSONSchemaBench corpus and larger programming-language grammars. Compilation is substantially heavier than mask generation, which is why constraints are serializable and why the dynamic fallback exists.

The library is under active development. The repository contains the current API documentation, benchmark reports, and examples.