Learning project on how constrained decoding works, starting from an old CUDA loopy belief-propagation project.
Constrained decoding: at every step, mask the model's next-token logits so generation stays on a grammar-valid path. Regex, GBNF, and a subset of JSON Schema, batched on the GPU that holds the logits. The link to belief propagation: reachability in a constraint automaton is boolean message-passing to a fixpoint.
Needs Python 3.10+, torch >= 2.2, and a C++20 toolchain (the install compiles
a small torch extension). CUDA is optional -- everything runs on CPU too.
pip install torch
pip install scikit-build-core cmake ninja
pip install --no-build-isolation -e '.[hf]'from transformers import AutoModelForCausalLM, AutoTokenizer
from bpdecode.hf import RegexLogitsProcessor, GrammarLogitsProcessor
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B")
# regex
lp = RegexLogitsProcessor(r"([0-9]{1,3}\.){3}[0-9]{1,3}", tok)
ids = tok("The router IP is ", return_tensors="pt").input_ids
print(tok.decode(model.generate(ids, logits_processor=[lp], max_new_tokens=20)[0]))
# ... 192.168.1.100
# JSON Schema (or GBNF source via GrammarLogitsProcessor(grammar, tok))
schema = {
"type": "object",
"properties": {"name": {"type": "string"}, "year": {"type": "integer"}},
"required": ["name", "year"],
"additionalProperties": False,
}
gp = GrammarLogitsProcessor.from_json_schema(schema, tok)
# -> {"name": "Rust", "year": 2010}Batched, on tensors (the serving-style path; CPU or CUDA by tensor device):
from bpdecode.batch import GrammarCache, ConstraintBatch
cache = GrammarCache(vocab) # compile once, share across requests
batch = ConstraintBatch(cache.get(pattern), capacity=256, device="cuda")
batch.add(request_id) # add / evict / reset as requests come and go
batch.apply_mask(active_ids, logits) # one call masks the whole batch
batch.commit(active_ids, sampled_tokens) # one call advances itMore in examples/. There's also a vLLM adapter
(bpdecode.vllm) and an experimental lookahead generation loop
(bpdecode.lookahead, see below).
- Regular grammars (regex, non-recursive GBNF) compile to a byte DFA and a
dense
tok_next[state][token]table -- a decode step is a gather. - Context-free grammars run a config-set pushdown automaton; masks are
memoised, and regular sub-loops (
[a-z]+) are spliced out to the dense path. - CUDA kernels in
csrc/cover both, checked against the CPU reference on an RTX 3090.GrammarLogitsProcessorhas a"cpu"backend (per-row, memoised) and a"device"backend (one kernel launch for the whole batch).
Numbers and details are in bench/RESULTS.md.
- Mask speed. Regex: ~8 µs/token on GPU, flat across patterns. JSON Schema: ~1 µs once warm. Caveat: "warm" is a memo hit after a ~0.15 s first-request warmup, while xgrammar and llguidance do their work up front -- so this is not an apples-to-apples win, just how fast a cache can be.
- Soft lookahead: my first idea failed. Nudging the model toward tokens with more valid continuations (counted from the automaton) makes it pad fields and never stop -- more continuations isn't better. Weighting by the model's own probability works (6/6 complete vs 0/6), at the cost of K extra forward passes per step. It's known prior work (Grammar-Aligned Decoding), not a new idea.
- The pushdown kernel doesn't get a GPU speedup. Unlike the branch-free regex gather, it does real per-token simulation, which is warp-divergent, and with no memo its cost grows with batch size.
- Regex: no anchors, backreferences, or lookaround.
- JSON Schema: properties come out in schema order;
minimum/maximumaren't enforced. - The pushdown config-set is bounded (8 alternative stacks, 32 deep); a pathological grammar silently saturates.
- Pre-alpha: not used in production, no independent review.
pip install --no-build-isolation -e '.[dev]'
pytest
ruff check .
cmake -S csrc -B build-cpp && cmake --build build-cpp && ctest --test-dir build-cppMIT