Grammar-constrained decoding

A grammar compiled from the tool registry

Capabilities model answers an agent's discovery call with a ranked list of tool pointers. A small model can produce a plausible tool name that does not exist, and an agent that calls a fictional tool fails slowly and confusingly. I decided validity belonged in decoding. I also decided the constraint should make decoding cheaper.

The grammar compiles from the tool registry, through the same frozen contract the evaluation scores against, into a trie over the model's own token ids. One compiled source feeds both the evaluation constraint and the serving constraint, so what was measured is what is served, and a registry change regenerates the grammar along with the training data.

Tool registry → trie over the model's token ids

bottom row: terminals, each mapped to one tool; a separator edge returns to the root and a stop edge ends the answer

One compiled source

the evaluation, through vLLM guided decoding

the serving engine's sampling kernel

the training data, regenerated with it

Where the trie lives on the GPU

When the trie fits in the GPU's constant memory, the decoder keeps it there.

Catalog trie · 42 tools

Constant memory · 64 KiB

every thread in a step reads the same node, and constant memory broadcasts that read

Benchmark trie · 2,000 entries

Global memory

too large for constant memory; the choice is made automatically

Held for each grammar: edge range per node · token and child per edge · terminal-to-tool map · each tool's path of edges

Constant memory is one buffer per program, so a second grammar silently overwrote the first during a health probe. The engine now tracks which grammar owns constant memory and re-binds it before every launch.

Rules that live in the mask

At each step the sampling kernel reads the current node's outgoing edges, drops any that are not allowed, and takes the argmax over the logits of what remains. Ties break toward the lower token id, so the answer never depends on which thread won a race.

Exactly k tools

a tool still owed: the separator is legal and stop is not; once k are out, only stop is

No repeats

edge 1

edge 2

edge 3

each edge counts the unused tools below it; a finished tool decrements its path, and an edge at zero leaves the mask

Fail loudly

!

nothing left standing: the decode ends with an error, never a plausible answer

No-repeat used to be a server-side loop that rebuilt a smaller grammar and asked again, at about 7.5 times the latency on short requests. I replaced it with counts in the mask: one on-device decrement per finished tool. I had the industrial serving stack tested with the real constraint first. Its string-level grammar with token-level state on top could not hold no-repeat soundly, and the workaround did not become correct. That is Trying the industrial path first.

Cheaper than decoding without it

In the serving engine the grammar decides which output-layer rows get computed at all, which is how constrained decoding became cheaper than unconstrained decoding (see GPU inference engine engineering). The wider problem is A hundred tokens instead of forty-two thousand.

picks: always real tools in the 42-tool catalog, never repeated

a real tool that may still be the wrong one: the model's choice

the rest of the catalog

I treat a constraint as information the hardware can use, and I put guarantees where they can be enforced, with their exact scope stated.