xgrammar-based constrained decoding with per-model reasoning parser

Summary

This PR adds grammar-constrained structured output generation to oMLX using xgrammar. Users can request structured outputs (JSON schema, regex, EBNF grammar, choice lists) via the OpenAI-compatible structured_outputs field, and oMLX enforces them at the logit level during generation.

A key design decision is the per-model reasoning_parser setting, which maps to xgrammar's builtin structural tags. This lets grammar constraints work transparently with different model protocols — Qwen's <think>...</think>, Harmony's <|channel|>analysis/final system, DeepSeek, Llama, etc. — without any client-side changes.

How it works

1. Client sends a request with structured_outputs: {"json": {...schema...}}
2. Server reads the model's reasoning_parser setting (configured via Admin UI)
3. If set, calls xgrammar.get_builtin_structural_tag(parser_name, reasoning=...) to get the model's protocol structure, then patches the user's grammar into the output slot
4. If not set, compiles a bare grammar (no structural wrapping)
5. During generation, GrammarConstraintProcessor applies token-level bitmasks to enforce the grammar
6. Bitmask computation runs in parallel with the model forward pass via mx.async_eval

Client request ──► reasoning_parser setting ──► get_builtin_structural_tag()
                                                        │
                                               Patch user grammar into
                                               output slot of tag
                                                        │
                                               compile_structural_tag()
                                                        │
                                               GrammarConstraintProcessor
                                               (bitmask per token, batched)

API

Uses the vLLM-compatible structured_outputs field in extra_body:

response = client.chat.completions.create(
    model="Qwen3.5-4B-4bit",
    messages=[{"role": "user", "content": "Give me a person"}],
    extra_body={
        "structured_outputs": {
            "json": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "age": {"type": "integer"}
                },
                "required": ["name", "age"]
            }
        }
    }
)

Supported grammar types:

• json — JSON schema (dict or string)
• regex — regular expression pattern
• grammar — EBNF/GBNF grammar string
• choice — list of allowed string values

Per-model reasoning_parser

Configured via the Admin UI dropdown (or PUT /admin/api/models/{id}/settings). Valid values: qwen, qwen_coder, harmony, llama, deepseek_r1, deepseek_v3_2, kimi, minimax, or empty (no structural tag).

When set, the auto-budget logic also kicks in: if the user doesn't specify a thinking_budget, one is auto-set so the model exits reasoning before the grammar activates.

Performance

Benchmarked with 128 max_tokens, 60-second time-boxed runs per configuration:

• Grammar overhead: 1.09x–1.24x (9–24% slower per request)
• TTFT unaffected by grammar constraints
• Overhead inversely correlates with model size — larger models have longer forward passes that better overlap with bitmask computation
• Full CSV with all concurrency levels in tests/bench_results.csv

Files changed

Testing

• 55 unit tests covering _build_format_element, _patch_output_format, _compile_with_structural_tag, _compile_bare_grammar, _compile_grammar_for_request, GrammarConstraintProcessor, scheduler grammar path, batched grammar, and _get_model_vocab_size
• Live integration tests verifying JSON schema, regex, and choice grammars produce correct output across Qwen, Gemma, and Harmony/OSS models
• Time-boxed performance benchmarks at concurrency 1/2/4 with thinking on/off

Dependencies

• xgrammar — added as optional: pip install 'omlx[grammar]'. When not installed, structured output requests via response_format fall back to prompt injection (existing behavior); requests via structured_outputs return a 400 error explaining how to install.