Structured output

Constrained decoding forces the model's output to match a JSON Schema, regular expression, or context-free grammar on every run. At each step, the inferlet computes which tokens are valid next, and the GPU epilogue masks out all other tokens before sampling.

JSON SchemaCopy Anchor Link

grammar.Matcher tracks the output token by token and returns a bitmask of valid next tokens at each step. The inferlet sends that mask into a device Channel, and the epilogue selects the highest-scoring valid token with masked_argmax:

import json

from inferlet import chat, grammar, mask, model
from inferlet.eta import Channel, ForwardPass, Pipeline, StateGeometry, WorkingSet, dtype, intrinsics, masked_argmax

SCHEMA = json.dumps({
    "type": "object",
    "properties": {"city": {"type": "string"}, "country": {"type": "string"}},
    "required": ["city", "country"],
})


async def main(input: dict) -> dict:
    vocab = model.output_vocab_size()
    matcher = grammar.Matcher(grammar.Grammar.from_json_schema(SCHEMA))
    allowed = Channel([vocab], dtype.bool)

    ws, pipe = WorkingSet(), Pipeline()
    tokens = chat.user("Name the capital of France as JSON.") + chat.cue_without_thinking()
    out = []
    while not matcher.is_terminated():
        allowed.put(mask.unpack_mask(matcher.mask(), vocab))     # send valid-token mask to GPU
        fwd = ForwardPass()
        fwd.embed(tokens)
        fwd.state([ws], StateGeometry.folding())
        nxt = fwd.epilogue(lambda: masked_argmax(intrinsics.output(0, "logits"), allowed.take()))
        pipe.submit(fwd)
        token = await nxt
        matcher.accept_tokens([token])                             # advance the grammar state
        out.append(token)
        tokens = [token]
    pipe.close()
    return json.loads(model.decode(out))

pie run main.py
# {"city": "Paris", "country": "France"}

Because invalid tokens are masked out before selection, json.loads always succeeds, even on a 0.8B model.

Regex and EBNF grammarsCopy Anchor Link

You can swap the grammar definition and keep the rest of the loop unchanged:

grammar.Grammar.from_regex(r"\d{3}-\d{4}")
grammar.Grammar.from_ebnf('root ::= "yes" | "no"')
grammar.Grammar.json()               # any valid JSON value

A Matcher also supports fork() when branching a working set, and rollback(n) when discarding rejected draft tokens after speculative decoding.

Token healingCopy Anchor Link

Prompts that end mid-word, such as "https://www.goo", often degrade output quality because the tokenizer splits the trailing characters differently than the full word "google.com". The token-healing inferlet backs up over the last prompt token and constrains the first generated token to start with the removed prefix.

Going beyond per-step maskingCopy Anchor Link

In the loop above, the host waits for each token before computing the next mask, so it runs one pass at a time. Two examples go further:

• constrained-speculative-decoding combines grammar masking with a drafter so one forward pass can verify and commit multiple tokens.
• asap-grammar-aligned-decoding corrects a sampling bias: masking at each step guarantees valid syntax, but skews the probability of complete outputs when sampling. ASAp reweights tokens over time to restore the target distribution.