feat: add XGrammar structured output with speculative decoding

[magnetism87]

Problem and scope

Related Issue: #33, specifically its JSON Schema use case. This proposes an implementation of final-response constraints; it does not close the issue's broader GBNF and plain-text grammar requests. The specific XGrammar integration proposed here has not been pre-approved by the maintainer.

NInfer currently rejects JSON response formats, so applications cannot require parseable JSON or a schema-valid final response. This adds native token-constrained generation to the public Engine, CLI, and the three HTTP protocols:

• Chat Completions: response_format with json_object or json_schema.
• Responses: text.format, with the selected format retained in the response object.
• Anthropic Messages: output_config.format.
• CLI: --json and --json-schema FILE.

For example, a request with a constant schema still returns the required JSON value when its prompt asks for unrelated prose. Invalid or unsupported constraints fail explicitly instead of being silently ignored.

Implementation

Grammar and request ownership

Vendor the native C++ portion of XGrammar v0.2.7 at 82505d0d987c36a4209fb3d8571cf6b0f28b5acd, with licenses and a target-scoped CMake build. No configure-time downloads, Python inference dependency, or model conversion are needed. third_party/xgrammar/NINFER.md documents three local correctness patches: JSON control-character exclusion in bounded strings, Unicode/escaped-key handling that prevents additional properties from overwriting constrained declared properties, and rejection of numeric intervals with no value at the supported decimal precision.

The text layer owns compilation, the frontend lazily caches compiled grammars with a bounded cache, and each request owns its matcher state. Output preview uses a fork; only Engine publication commits the matcher. Grammar state is not restored from KV/prefix caches. Serving parses formats into public execution options and rejects unsupported combinations before generation. CLI and serving expose the requested format/schema as a leading instruction before preparation, preserving caller content and cache boundaries. This guides raw JSON generation when tools remain available; the grammar enforces validity independently.

The model frontend composes its reasoning delimiter and native Qwen tool envelope with the final schema into one matcher. A speculative block may cross the reasoning boundary without a separate phase switch. Forced thinking-budget closure advances the matcher through the same preview/commit transaction. Tool names are declared-only, and optional non-strict parameters follow the declaration order shown in the prompt.

Sampling and speculative execution

Apply token masks before penalties and target top-k/top-p/min-p filtering. Each speculative position uses the grammar state conditioned on its preceding draft prefix, including correction and bonus positions. Draft distributions remain unchanged; stochastic acceptance uses the constrained target probability in min(1, p/q), with residual sampling from normalized max(p-q, 0). Rejected drafts do not advance committed grammar state.

Program-owned persistent GPU masks and pinned staging buffers support compact batches, lane reuse, and CUDA graph replay. MTP uses host-known drafts; the shared DFlash/DFlash2 path captures a device-to-host draft copy, a host matcher node, and a host-to-device mask copy before verification. Callback errors surface after synchronization and before publishing output. Existing terminal-prefix commit and model-state handling remain in place.

Contract and tradeoffs

• Supported schema constraints include object properties/required/additional properties, arrays/tuples, length/item bounds, numeric minimum/maximum and exclusive bounds, enums/constants, anyOf, and local recursive references. Unsupported keywords, including multipleOf, patterns/formats, oneOf/allOf, uniqueness, and conditionals, return HTTP 400. Exact restrictions are documented in docs/serving.md.
• Structured requests default to thinking disabled when no explicit reasoning mode or budget is supplied. Explicit thinking, thinking budgets, and active tools are supported: a response may call tools or return constrained final JSON. The schema does not constrain reasoning or tool arguments. Custom stops and assistant-prefill continuation are rejected. Strict tool-argument generation and arbitrary grammar aliases are outside this change. Automatic tool selection remains a model decision.
• Normal completion produces a complete constrained document. Token/context limits, cancellation, transport failures, and generation errors can leave partial JSON. SSE deltas must be concatenated before parsing.
• The compiler produces a documented valid subset: property order follows the schema, bounded strings and additional-property key spellings have restrictions, and whitespace runs are capped at eight characters. Numeric bounds require an explicit numeric type and finite bounds within ±(2^53−1); integer bounds must be integral, and bounded numbers use ordinary decimals with at most six fractional digits.
• DFlash/DFlash2 graph execution adds a CPU matcher synchronization point and transfers per round, including overhead on the common unconstrained DFlash graph path. No throughput improvement or performance-parity claim is made.

Verification

Tested on an RTX 5090 with CUDA 13.1, a Release sm_120a build, and Python 3.11. Live inference used an existing Qwen3.8-27B NVFP4 Swift artifact containing MTP and DFlash2 weights.

• Full build passed. The follow-up passed all nine affected tests: grammar, CLI options, prompt input/cache boundaries, Chat Completions, Responses, Anthropic Messages, frontend thinking-budget control, native tool parsing, and captured structured-mask rounds. The original sampling/speculative numerical qualification and additional frontend/runtime/state-layout checks also passed.
• Sampling/speculative tests compare production kernels against independent FP64 probability oracles, including forbidden high-logit tokens, masks before filters, first/middle rejection, all-accepted/bonus cases, draft lengths 1/5/15, batch sizes 1/8, production vocabulary size, and graph replay.
• Dedicated graph-stage tests cover physical-lane remapping, matcher failure propagation, and descriptor/error reset. Grammar tests cover EOS, preview isolation, UTF-8, control characters, escaped-key aliases, recursive references, unsupported constraints, and whitespace bounds. Follow-up cases cover numeric endpoints/empty intervals, schema enforcement with multiple tool alternatives, split and mixed-token reasoning boundaries, draft-preview isolation, and injected thinking-budget closure.
• The final live suite passed ordinary decoding, MTP with five draft tokens, and DFlash2 with seven draft tokens in graph and eager execution. It checks greedy/stochastic output, contradictory prompts, nested schemas, independent schema validation, prefix reuse, concurrent/mixed requests, streaming, output limits, cancellation followed by a fresh request, invalid requests, and all three HTTP protocols. The final suite also checks a real tool-result-to-bounded-JSON exchange with explicit reasoning and multiple tools still active, including an irrelevant code tool and a Markdown-formatted example.
• DFlash2 also passed at concurrency eight with fifteen draft tokens, including eight simultaneous arrays of 32 constrained objects. That run preceded the final whitespace bound; the final four-mode suite was rerun after the bound was added.
• A real CLI --json invocation produced parseable JSON. git diff --check passed.
• Unmodified tool-eval-bench 2.6.1.dev65+g6be685f0e, TC-64 through TC-69, using Qwen3.8-27B with MTP-5, temperature 1.0, seed 42, top-p 0.95, top-k 20, thinking/preserve-thinking enabled, medium reasoning, and a 16,384-token limit: TC-64/65/66/67/69 passed 2/2 each. TC-68 failed 0/2; it deliberately omits response_format and tests prompt-only adherence. Total: 10/12 (83%). This is a six-scenario rerun, not a new full-benchmark score. All 224 installed benchmark Python files were verified unchanged.
• The same unchanged benchmark initially exposed repeated tool calls after reasoning. An unconstrained replay produced Markdown-fenced JSON; exposing the requested raw-JSON contract in the server prompt and preserving tool parameter declaration order resolved the affected final-turn scenarios without disabling automatic tools or changing benchmark requests.

Build and focused-test reproduction:

cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=ON
cmake --build build -j
ctest --test-dir build --output-on-failure \
  -R '^ninfer_(structured_output|cli_options|openai_schema|openai_responses|anthropic_schema|qwen3_5_structured_round|sampling|speculative_round|qwen3_5_frontend|tool_call_parser|prompt_input)_test$'

Live-test reproduction, with Python 3.11 plus requests and jsonschema, and the explicit local artifact path substituted as needed:

python3.11 tests/test_structured_output_live.py \
  --server build/apps/ninfer-serve \
  --artifact models/qwen3_8_27b_nvfp4swift.ninfer \
  --output-dir /tmp/ninfer-structured-live \
  --modes none mtp dflash2 dflash2-eager

python3.11 tests/test_structured_output_live.py \
  --server build/apps/ninfer-serve \
  --artifact models/qwen3_8_27b_nvfp4swift.ninfer \
  --output-dir /tmp/ninfer-structured-c8-d15 \
  --modes dflash2 --concurrency 8 --draft-tokens 15

Unverified: end-to-end legacy DFlash inference, because no compatible artifact was available. Its shared mask-stage integration and operator-level acceptance/correction semantics are covered, but that is not a substitute for a real legacy DFlash run. No end-to-end performance comparison was conducted. origin/dev was inspected and has no commits ahead of the target master at submission.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.