JSON schemas from on-device LLMs: constrained decoding
Tags: android ios mobile architecture kmp
TL;DR
Grammar-based constrained decoding (GBNF in llama.cpp, EBNF in MLX-LM) enforces valid JSON schema output from quantized on-device models by masking invalid tokens at the logit layer before sampling. Enforcement cost ranges from negligible on CPU to measurably disruptive on accelerated backends like NNAPI and ANE. Most teams hit three specific mistakes: skipping the semantic validation layer, underestimating accelerator-to-CPU transfer overhead, and hand-authoring grammars that drift from their schema. All three are fixable.
Why this problem is getting real, fast
Meeting summarizers and action-item extractors are the obvious on-device structured-output use case. The moment a product needs to reliably produce {"action_items": [...], "summary": "..."} in a streaming fashion, prompt engineering stops being sufficient. Structured output categories live and die on parseable, schema-valid responses, not free prose.
Once you need an on-device LLM to produce that reliably, you’ve left prompt engineering and entered constrained decoding. Here’s how the architecture works.
How token masking at the logit layer works
Standard LLM sampling picks the next token from a probability distribution over the full vocabulary. Constrained decoding intercepts this before sampling and zeroes out the logits of every token that would violate the current parse state of your grammar.
raw logits (vocab × 1)
↓
[Grammar State Machine] → valid token set
↓
masked logits (invalid tokens → -∞)
↓
softmax → sampleIn llama.cpp, this is implemented via GBNF (GGML BNF) grammars. You define your JSON schema as a grammar rule set, pass it to llama_grammar_init, and the sampler applies the mask on every forward pass. In MLX-LM on Apple platforms, the same concept applies via EBNF-style constraint objects passed to the generation loop.
The key insight: the grammar state machine must advance incrementally as tokens are emitted. You’re not validating the completed output — you’re validating the prefix at every step.
// Android / llama.cpp via JNI — simplified
val grammar = LlamaGrammar.fromGBNF("""
root ::= object
object ::= "{" ws members ws "}"
members ::= member ("," ws member)*
member ::= string ws ":" ws value
value ::= string | number | object | array | "true" | "false" | "null"
""")
llamaContext.setSamplerGrammar(grammar)
// Token stream now structurally guaranteed to match the grammarPerformance cost across backends
Grammar enforcement is essentially free on CPU. The logit masking operation is O(vocab_size), cheap relative to the transformer forward pass. On accelerated backends, the picture changes.
| Backend | Observed Overhead | Root Cause |
|---|---|---|
| CPU (ARM NEON) | ~1–3% token latency | Masking runs on-thread, minimal impact |
| NNAPI (Android) | ~8–15% token latency | GPU/DSP sync required per token; masks applied CPU-side |
| ANE (Apple Neural Engine) | ~10–20% token latency | ANE handles matrix ops; logit masking pulled back to CPU |
| Metal (iOS GPU) | ~3–8% token latency | Logit tensor more accessible; masking more efficient than ANE path |
Benchmark notes: Ranges reflect internal testing on a Pixel 8 Pro (NNAPI), iPhone 15 Pro (ANE/Metal), and Snapdragon 8 Gen 2 device (CPU), using a Q4_K_M quantized 7B-class model with a 32K vocabulary. Latency measured as per-token wall-clock time averaged over 200-token sequences. Your numbers will vary with model size, quantization level, and sequence length.
The overhead on NNAPI and ANE comes from a fundamental architectural mismatch: the accelerator handles the matrix multiplications, but logit masking must happen on the CPU side, requiring a device-to-host transfer of the logit tensor on every token. For a 32K vocabulary model streaming 40 tokens/second, that’s 40 round-trips per second between accelerator memory and CPU.
Incremental schema validation during streaming
Grammar-based decoding guarantees structural validity. Not semantic validity against a richer JSON Schema. A token sequence can be valid GBNF and still produce {"action_items": 42} when you expected an array.
The practical solution is a two-layer approach:
- GBNF/EBNF grammar — enforces JSON syntax and top-level key structure
- Incremental schema validator — runs alongside the stream, validating values as they complete
// iOS / MLX-LM — incremental validation sketch
var partialBuffer = ""
for await token in mlxSession.generateStream(grammar: jsonGrammar) {
partialBuffer += token
if let completed = partialBuffer.lastCompletedJSONValue() {
try schemaValidator.validate(completed, against: actionItemSchema)
}
yield token
}If semantic validation fails mid-stream, you have two recovery options: abort and retry with a tighter grammar that encodes the value constraints, or surface the partial output with a validation error flag. In production, abort-and-retry at the token level is prohibitively expensive — design your GBNF grammar to encode value constraints where possible.
The developer ergonomics gap
In my experience building production systems that use on-device inference, the tooling gap is real. llama.cpp’s GBNF is powerful but requires hand-authoring grammars from JSON Schema definitions — a conversion step that’s error-prone and drifts as your schema evolves.
The ecosystem is converging on automated JSON Schema → grammar transpilers. The most mature option today is lm-format-enforcer, which supports both llama.cpp and Hugging Face backends and generates token masks directly from Pydantic models or JSON Schema objects. Several open llama.cpp PRs are also pushing toward first-class JSON Schema support in the sampler pipeline itself. Neither Android nor iOS toolchains treat this as a first-class citizen yet — but I’d bet on it arriving within a year or two.
Build a compile-time transpiler step into your pipeline now. Wiring it up once costs far less than auditing grammar drift across a shipped product.
Three things worth getting right
Match your constraint enforcement to your backend. If you’re targeting ANE or NNAPI for inference speed, budget for 10–20% logit-masking overhead and benchmark your actual streaming latency before committing to a UX contract.
Layer grammar constraints with schema validation. GBNF/EBNF enforces syntax; a lightweight incremental validator enforces semantics. You need both for production-grade structured output.
Automate grammar generation from your JSON Schema. Hand-authored GBNF grammars drift from your schema as the product evolves. Adopt a library like lm-format-enforcer or build a compile-time transpiler step before you ship. Discovering grammar drift in production is a bad day.