serve stage 2: response_format → per-request grammars → grammar-forced drafts in the multiplexed server (stacked on #148)

[fabio-rovai]

Completes the seam proposed in #146 and started in #148 (stacked on it — the first two commits are that PR + its ws-tolerance follow-up; review the top two commits here): an OpenAI client sends response_format, and the batched server turns it into grammar-forced draft speedup, per request, with no GBNF hand-writing anywhere.

What lands

• Gateway (openai_server.py): response_format {"type":"json_object"} (ships a generic whitespace-tolerant JSON grammar), {"type":"json_schema"} (schema forwarded verbatim; the engine compiles it via schema_gbnf.h), plus a raw-GBNF {"type":"gbnf"} extension. Draft-source semantics end to end: an uncompilable schema costs the speedup, never the request, never the output.
• Protocol: optional 7th SUBMIT field gbytes; grammar text rides the payload after the prompt. 6-field headers unchanged.
• Engine: per-slot GrDraft (a mechanical de-globalization commit precedes this), lifecycle owned by slot reuse, walkers fed on every emitted token. Grammar drafting now runs inside run_serve_mux for greedy requests: a drafting slot steps out of the shared batch for one forward and uses the proven single-sequence verify path (kv_bind + step_all — the exact primitives mux_submit already uses for prefill), then rejoins. No changes to step_decode_batch. Rejected drafts' KV entries are overwritten by the next forward, same as the existing prefix-truncation path. Sampling requests never draft (rejection-resampling under sampling is out of scope and documented as such).

Verified

• make test-c green (incl. new 7-field parse cases); python suite green except the known environmental memory_available failure (resource_plan: memory_available() macOS support (fixes 0 on Darwin) #150 fixes it). The generic JSON grammar is parse+walk validated against grammar.h (46 rules, recursion, escapes).
• Live end-to-end (M3 Max, GLM-5.2 int4, gateway + engine, greedy): plain request → valid JSON; json_schema request → [GRAMMAR] request: 18 rules armed and schema-shaped output; json_object → 46-rule generic grammar armed. The model chose fenced, pretty-printed JSON and the ws-tolerant walker tracked it through the indentation — the SCHEMA=<file.json>: JSON-Schema → GBNF compiler for grammar-forced drafts (follow-up to #70, seam discussed in #146) #148 follow-up doing exactly what it was built for.

Honest caveats

1. Drafted and undrafted runs of the same request can differ in near-tie tokens: the verify forward is S=1+k where the plain path is S=1, and batch-size reduction-order sensitivity is the documented Greedy decoding is not reproducible: MTP (3/5 prompts) and the CUDA expert tier (2/5) both change the output #100 behavior (same as MTP). Verification guarantees each emitted stream is a valid greedy stream of its own forward shapes.
2. Acceptance counters aren't yet surfaced per-request in mux mode (they aggregate in the slot's GrDraft) — a STAT-line extension would make the speedup observable per request; happy to add in this PR or a follow-up, your call.

🤖 Generated with Claude Code

[rajpratham1]

Thanks for the contribution! This is a substantial feature that adds protocol changes, a JSON Schema → GBNF compiler, grammar handling in the OpenAI server, engine integration, and a large new test suite.

Before this can be approved, I think a few things should be clarified:

1. Backward compatibility

   • Please document the protocol change introducing the optional gbytes field in the SUBMIT header and confirm compatibility with older clients/servers.

2. Documentation

   • The new response_format modes (json_object, json_schema, and gbnf) should be documented for users, including examples and any limitations.

3. Testing

   • The new tests cover many success cases, but it would be good to add more negative and edge-case tests (very large schemas, malformed GBNF, invalid grammar payloads, nested depth limits, etc.).

4. Performance

   • Since grammars are compiled/generated per request, it would be helpful to include some discussion or benchmarks on the overhead for large schemas.

Overall the direction looks promising, but given the size of the change I'd prefer these points to be addressed before approval.

[fabio-rovai]

@rajpratham1 all four points addressed in the two pushes just now (plus a rebase of the whole stack onto current main after the merge wave — and since #176 merged underneath, this branch now sits cleanly on top of it):

1. Backward compatibility — the gbytes field is additive in both directions: old client → new server (6-field headers parse unchanged, gbytes=0; covered by the parse tests), new client → old server never occurs through this gateway (it only emits the 7th field when a grammar exists). Now stated explicitly in docs/grammar-draft.md and the decode_batch.h comment.
2. Documentation — docs/grammar-draft.md (the canonical page from docs: grammar-forced drafts reference page (requested in #146) #170) gained a "Server usage: response_format" section: all three modes with examples, the draft-source-not-constraint semantics, the greedy-only drafting rule, the 1 MiB cap, and the Greedy decoding is not reproducible: MTP (3/5 prompts) and the CUDA expert tier (2/5) both change the output #100-class near-tie caveat.
3. Testing — added: non-dict response_format, empty GBNF, oversized grammar (>1 MiB, now pre-checked at the gateway to match the engine's gbytes bound), and an explicit test documenting that malformed GBNF passes the gateway by design — the engine fail-softs it, because a bad grammar must cost the speedup, never the request. Depth limits are enforced and tested engine-side (SGB_MAX_DEPTH in the SCHEMA=<file.json>: JSON-Schema → GBNF compiler for grammar-forced drafts (follow-up to #70, seam discussed in #146) #148 suite).
4. Performance — measured (M3 Max, 10k iterations): schema→GBNF compile + gr_parse = 7.8 µs/request for a typical 4-field schema, 17.9 µs at the 32-level nesting cap. Against generation at ~1 token/s on this hardware class, per-request compilation is 5–6 orders of magnitude below one token; caching compiled grammars would be complexity without a payoff.

🤖 Generated with Claude Code

[JustVugg]

#111 (GPU resident pipeline) just landed on dev and overlaps this on decode_batch.h/glm.c — one rebase onto current dev and I'll merge straight away (everything else in today's wave is in, so nothing further will move under you). The stage-2 design itself looks right on review: fail-soft grammar setup, per-request state in one struct, no unsafe constructs. Sorry for the churn — your stack landed 3-for-3 today (#148, #170, #176).

[fabio-rovai]

Rebased onto current main (e21318e, post-#111/#195/#201/#202) — the stack is now the three stage-2 commits only (the #148 base deduped away after your merge). One conflict resolved in run_serve_mux init: kept the #195 Windows _setmode binary-mode block alongside the per-slot GrDraft array. make glm clean, full make test-c + python suite green on Apple Silicon. Ready when you are — and thanks for the fast reviews today.

[JustVugg]

Heads-up: c/glm.c was just split into c/colibri.c + header modules (#391, now merged into dev). This PR touches the old glm.c, so it will need a rebase onto current dev with its changes moved into colibri.c (or the relevant extracted header: quant.h, sample.h, kv_persist.h, telemetry.h, grammar.h). The make glm target still works (alias of make colibri), so no build scripts break. Apologies for the churn — ping me if the rebase gets thorny and I can help place the hunks.

[JustVugg]

@fabio-rovai rebase heads-up: I rebased most of the way onto current dev for you (colibri.c + decode_batch.h merged clean after the #391 refactor), but the openai_server.py conflict is semantic and I don't want to guess it: #437 (now merged) refactored the same request path — it moved tools/tool_choice to generation() parameters (filtered in chat_completion()) and changed generation_options's return, while your PR adds grammar to generation_options's return and threads it through generation(). Both touch the exact same two signatures. The right merge is to keep #437's tools-param flow and add your grammar return/param alongside it, but that's a serve-path change I can't validate without the model, so it's your call. Should be a quick reconcile now that the C side is clean. Thanks for the grammar-forced-drafts work.

[fabio-rovai]

Rebased onto current dev (e809171). The openai_server.py reconcile went exactly as you sketched: #437's tools-param flow is kept as-is (tools/tool_choice stay generation() parameters, filtered in chat_completion()), and generation_options now returns grammar as a 4th value threaded through both stream and non-stream paths. The C side landed in colibri.c/decode_batch.h with no manual surgery — your partial rebase read was right.

Verified locally on macOS: make check fully green, python suite 103 passed / 8 skipped.

One drive-by you should know about: make check is currently broken on dev itself — #270 merged after the #391 split and still builds tests/test_pipe_block against the deleted glm.c (both the Makefile rule and the #include). I included a one-commit fix here (2f4dca0, rule + include moved to colibri.c with the post-split header deps) because CI on this PR can't go green without it. Happy to split it into its own PR if you'd rather land it independently — it affects every open PR's CI, not just this one.

[fabio-rovai]

@rajpratham1 status check on your 14 Jul review: all four points were addressed the same day (point-by-point response — gbytes backward-compat documented, response_format modes documented in docs/grammar-draft.md, negative/edge tests added incl. malformed GBNF and the 1 MiB payload gate, compile overhead measured at 7.8µs typical / 17.9µs at depth cap), and the branch has since been rebased onto current dev with the #437 serve-path reconcile. All 8 CI jobs are green. #168 was in the same position and merged earlier today after re-review — would you take another look at this one?