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.