serve stage 2: per-request grammars end-to-end — response_format -> SUBMIT -> grammar-forced drafts in the multiplexed server
The OpenAI gateway previously 400'd every response_format; the mux engine path
ran with speculation disabled entirely. Now:
- openai_server.py: response_format {"type":"json_object"} (generic ws-tolerant
JSON grammar), {"type":"json_schema"} (schema forwarded as-is, compiled
engine-side by schema_gbnf.h), and a raw-GBNF {"type":"gbnf"} extension.
Draft-source semantics throughout: a schema the engine cannot compile costs
the speedup, never the request and never the output.
- SUBMIT protocol: optional 7th field gbytes; grammar text appended to the
payload after the prompt. 6-field headers unchanged (back-compatible).
- Engine: per-slot GrDraft (grammar_setup_text/grammar_teardown split out of
the env-driven setup); walkers fed on every emitted token. Grammar-forced
drafting in run_serve_mux for greedy requests: a drafting slot leaves the
shared batch for one forward and runs the proven single-sequence verify path
(kv_bind + step_all) — the same primitives prefill already uses per
submission — then rejoins; rejected drafts' KV entries are overwritten by the
next forward exactly like the existing prefix-truncation path. Sampling
requests never draft (verification under sampling needs rejection resampling;
out of scope).
Tests: 7-field SUBMIT parse cases; response_format->grammar plumbing incl.
fail cases; test doubles updated; generic JSON grammar parse+walk validated
against grammar.h. make test-c green; python suite green except the known
environmental memory_available failure (#150 fixes it).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -44,6 +44,12 @@ static void test_submit_header(void)
|
||||
assert(coli_submit_parse("SUBMIT 1 0 16777216 3 1 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 16777217 3 1 1", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 1 0 2 3 1 1 trailing", &sub));
|
||||
/* optional 7th field: per-request grammar length (0 when absent) */
|
||||
assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95", &sub) && sub.gbytes == 0);
|
||||
assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 512", &sub) && sub.gbytes == 512);
|
||||
assert(coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 1048576", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 1048577", &sub));
|
||||
assert(!coli_submit_parse("SUBMIT 42 3 17 64 0.7 0.95 512 extra", &sub));
|
||||
}
|
||||
|
||||
int main(void)
|
||||
|
||||
@@ -20,8 +20,8 @@ class FakeEngine:
|
||||
self.calls = []
|
||||
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
|
||||
cancelled=None):
|
||||
self.calls.append((prompt, maximum, temperature, top_p, cache_slot))
|
||||
cancelled=None, grammar=None):
|
||||
self.calls.append((prompt, maximum, temperature, top_p, cache_slot, grammar))
|
||||
on_text("Hé")
|
||||
on_text("llo")
|
||||
return {"prompt_tokens": 7, "completion_tokens": 2, "length_limited": False}
|
||||
@@ -34,7 +34,7 @@ class BlockingEngine(FakeEngine):
|
||||
self.release = threading.Event()
|
||||
|
||||
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
|
||||
cancelled=None):
|
||||
cancelled=None, grammar=None):
|
||||
self.entered.set()
|
||||
self.release.wait(2)
|
||||
return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot,
|
||||
@@ -71,11 +71,11 @@ class TemplateTest(unittest.TestCase):
|
||||
|
||||
def test_validates_generation_limits(self):
|
||||
self.assertEqual(generation_options({"max_tokens": 4, "temperature": 0, "top_p": 1}, 8),
|
||||
(4, 0.0, 1.0))
|
||||
(4, 0.0, 1.0, None))
|
||||
# max_tokens above the server cap is clamped, not rejected (#260): OpenAI
|
||||
# clients default to large values; erroring breaks them.
|
||||
self.assertEqual(generation_options({"max_tokens": 9, "temperature": 0, "top_p": 1}, 8),
|
||||
(8, 0.0, 1.0))
|
||||
(8, 0.0, 1.0, None))
|
||||
# non-positive / non-int max_tokens is still a hard error
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"max_tokens": 0}, 8)
|
||||
@@ -84,7 +84,21 @@ class TemplateTest(unittest.TestCase):
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"top_p": math.inf}, 8)
|
||||
self.assertEqual(generation_options({"temperature": None, "top_p": None}, 8),
|
||||
(8, 0.7, 0.9))
|
||||
(8, 0.7, 0.9, None))
|
||||
# response_format -> grammar plumbing (draft source, never a constraint)
|
||||
opts = generation_options({"max_tokens": 4, "response_format": {"type": "json_object"}}, 8)
|
||||
self.assertIn("root ::=", opts[3])
|
||||
schema = {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]}
|
||||
opts = generation_options({"max_tokens": 4, "response_format":
|
||||
{"type": "json_schema", "json_schema": {"schema": schema}}}, 8)
|
||||
self.assertEqual(json.loads(opts[3]), schema)
|
||||
opts = generation_options({"max_tokens": 4, "response_format":
|
||||
{"type": "gbnf", "grammar": 'root ::= "x"'}}, 8)
|
||||
self.assertEqual(opts[3], 'root ::= "x"')
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"response_format": {"type": "yaml"}}, 8)
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"response_format": {"type": "json_schema", "json_schema": {}}}, 8)
|
||||
|
||||
|
||||
class ProtocolTest(unittest.TestCase):
|
||||
|
||||
Reference in New Issue
Block a user