serve: APIs serving GLM-5.3-Flash
What and why
Draft, in progress. This brings TrnRT's serving layer to SGLang's API and features, for what is needed to serve GLM-5.3-Flash. It is one pull request by the maintainers' choice: the pieces land here as separate commits and are reviewed together.
Out of scope: LoRA, multimodal, embeddings, rerank, classify and audio. Speculative methods other than MTP are also out, as are parallelism and disaggregation methods that do not beat what is served today.
| # | Piece | Changes compiled graphs? | State |
|---|---|---|---|
| 1 | /tokenize, /detokenize (also under /v1), SGLang and vLLM shapes | no | done |
| 2 | SGLang native /generate | no | done |
| 3 | Admin and info: /abort_request, /flush_cache, /get_server_info, /get_model_info, /get_load, /health_generate | no | done |
| 4 | Scheduling policies (fcfs, lpm, spf, priority), per-request priority, queue limits | no | done |
| 5 | Parallel sampling (n > 1), stop regex, kept stop text, skip_special_tokens (min_tokens moved to 8: it needs a device-side mask) | no | done |
| 6 | OpenAI Responses API and Anthropic Messages API | no | done |
| 7 | Richer /v1/stats, dashboard, production log; prompt token cache | no | done |
| 8 | Structured outputs (JSON schema, regex, EBNF), then tool_choice "required" and named tools; min_tokens | yes, behind a flag | done, except tool_choice required / named (below) |
| 9 | Prompt logprobs / echo, penalties, thinking budget | yes, each behind a flag | done, except prompt logprobs (below) |
Every piece that can change a compiled graph is off by default and comes with tools/graph_equiv.py evidence that the default graphs are unchanged.
1. /tokenize and /detokenize
| Piece | Change |
|---|---|
trnrt/serve/api.py | /tokenize takes exactly one of prompt (a string, or a list for a batch) or messages. A chat request is rendered by chat_prompt, the same path /v1/chat/completions uses, so count is the prompt length that request would have. A prompt is encoded as /v1/completions encodes it (add_special_tokens, default true); return_token_strs is vLLM's. /detokenize takes tokens (ids or a batch) checked against the vocabulary, with skip_special_tokens (default true, as SGLang), and answers text (SGLang) and prompt (vLLM). Both run on the prompt thread with its own tokenizer, never on the event loop. |
trnrt/serve/router.py | STATELESS_ROUTES go to the least loaded live engine; a refused connection or a 503 marks it down and the next is asked. They touch no prefix tree, session or load count, so a /tokenize of an agent's context does not steer its next turn. |
docs/API.md, CHANGELOG.md | The endpoints, fields and an example. |
Tests: /tokenize on a chat body equals the prompt ids the engine received for the same body on /v1/chat/completions. Refusals, both detokenize answers, and the router's failover and untouched routing state are covered too. A planted mutant routing these paths through the prefix-tree proxy fails the router test.
2 and 3. /generate and the info and admin endpoints
| Piece | Change |
|---|---|
trnrt/serve/native.py (new) | SGLang's /generate (POST or PUT): text or input_ids, single or batch, SGLang's sampling_params names, SGLang's streaming form, return_logprob with top_logprobs_num and return_text_in_logprobs, client rids, and meta_info with SGLang's finish_reason forms. It samples as /v1/completions does, and refuses what the engine cannot serve (prompt logprobs, n > 1, grammars, penalties, LoRA, multimodal) with a 400 naming the field. Also /model_info, /server_info, /get_load, /health_generate, /abort_request (by rid, or abort_all) and /flush_cache (idle only, as SGLang), each under its get_ name as well. |
trnrt/serve/api.py | Client request ids; an abort delivers the output so far with finish reason abort; abort_all; an idle-only cache flush on the engine thread. |
trnrt/serve/router.py | /generate goes through the prefix-routed proxy, its routing text taken from text/input_ids. Info, load, health, abort and flush fan out to the engines, with results combined; a flush also clears that engine's prefix tree in the router. |
trnrt/__main__.py | Hands the preset name to the server for /server_info (no shape change). |
Tests: the scripted-engine API tests, router tests, and a real-engine test where /generate gives the same ids and logprobs as /v1/completions, and /flush_cache empties both cache tiers (the next turn is cold, the one after hits). Five planted mutants were each caught. Full CPU suite: 1443 passed.
4 and 7. Scheduling, statistics, the dashboard, the production log and the prompt token cache
| Piece | Change |
|---|---|
trnrt/engine/scheduler.py, radix_cache.py, request.py, engine.py | SGLang's --schedule-policy (fcfs, lpm, lof, spf), priority scheduling (--enable-priority-scheduling, the request field priority, preemption of the worst priority, SGLang's threshold), --max-queued-requests and --max-queued-tokens (a full queue answers SGLang's 503 "The request queue is full.", marked by the header x-trnrt-queue-full), and --radix-eviction-policy (lru, lfu, slru, priority, tlru). Host only: the settings travel in a Scheduling object, not EngineConfig, and engine_config() is unchanged with every flag set. With no flag set, the scheduler's plan traces equal the base's in 96 cases (plain, overlapped, blind speculative, KV pressure, sessions). |
trnrt/serve/stats.py, dashboard.py, prodlog.py | GET /v1/stats adds live state and 1/5/15-minute, since-start and daily rates and latency percentiles (the fields the router reads are unchanged); GET /dashboard, one self-contained page; TRNRT_PROD_DIR, an opt-in log of request records (counts and times only, never text) and gauges. |
trnrt/serve/prompt_cache.py | TRNRT_PROMPT_CACHE_MB, opt-in: a chat turn's prompt ids from the previous turn's plus the new tail, which must equal the served render's ids; it checks itself at start, audits its hits and turns off for good on any mismatch. |
trnrt/serve/router.py | Passes a full queue's 503 on without marking the engine down; serves each engine's /v1/stats and /dashboard (the engines listen on 127.0.0.1 only). |
5 and 6. Sampling extras, the Responses API and the Messages API
| Piece | Change |
|---|---|
trnrt/serve/api.py | n choices (up to 128) on both OpenAI endpoints, each its own engine request sharing the prompt through the prefix cache; usage counts the prompt once; a seed gives choice i seed + i (vLLM). SGLang's stop_regex on the host stream; include_stop_str_in_output / no_stop_trim; skip_special_tokens: false; SGLang's matched_stop. The stop token is no longer part of the text when the tokenizer does not mark it special (vLLM, SGLang). min_new_tokens and spaces_between_special_tokens: false are now refused with 400 (they were ignored). |
trnrt/serve/responses.py | POST /v1/responses, stateless: reasoning, message and function-call items, the Responses stream events; store and previous_response_id are refused. |
trnrt/serve/anthropic.py | POST /v1/messages and /v1/messages/count_tokens: thinking, text and tool-use blocks, Anthropic's stream events, usage split and error shape. |
trnrt/serve/router.py | Routes both by session and prefix as it routes chat; count_tokens goes to any engine. |
The three branches were merged together; /generate keeps refusing n, no_stop_trim and skip_special_tokens: false, which only the OpenAI endpoints serve. Mutants planted across pieces 2-7: 28, each caught by a test except one equivalent mutant (a prompt-cache splice that yields the same ids). Full CPU suite at the merged head d00a260: 1641 passed, 22 skipped, 2 xfailed.
8 and 9. Structured outputs, penalties, thinking budgets and min_tokens
All are off by default: each flag adds graph variants that take host inputs, and with both flags off an engine traces exactly the graphs it did before.
| Flag | Serves | New graphs |
|---|---|---|
--structured-outputs | guided decoding through xgrammar (optional extra trnrt[grammar]): response_format json_object / json_schema, vLLM's guided_* and structured_outputs, SGLang's json_schema / regex / ebnf (also on /generate); thinking_token_budget (and the Messages API's budget_tokens); min_tokens | decode and prefill-post graphs that take a packed token mask |
--penalties (--penalty-buckets) | repetition, frequency and presence penalties | decode and prefill-post graphs that take penalty inputs, per distinct-token bucket |
| Piece | Change |
|---|---|
trnrt/engine/grammar.py, guided.py, host_inputs.py (new) | The per-request grammar matcher (the grammar applies after </think>, as in vLLM; a grammar the vocabulary cannot continue ends the request with "stop"), the thinking budget's forced </think>, min_tokens' mask of EOS and the stop ids, and the penalty inputs. |
trnrt/engine/scheduler.py, engine.py | A request needing a mask or penalties waits one read-back while its newest token is in flight; every other row runs as today. Under the dynamic mtp-* presets such a request makes the step plain (it never speculates); mtp-always refuses the flags. |
trnrt/engine/model_runner.py, sampler.py, trnrt/config.py | The variant families' warmup, the in-graph mask unpack and logit shaping, and the two flags. Traced files: edits keep every line the flags-off graphs trace at its line (sample() edited in place; new functions appended; runner edits below __init__). |
Deferred, still refused with a 400: prompt logprobs (prompt_logprobs, echo with logprobs): the model's prefill post path does not compute them yet, a change under trnrt/models/. tool_choice "required" and a named tool: GLM's tool-call tags are special tokens, which an xgrammar grammar over text cannot constrain.
Not yet checked on a Trainium2: that the new variants compile and run (the mask unpack, the penalty gather / scatter at the full vocabulary), guided output at tensor parallel 32, the dynamic presets' plain flip, and the variants' compile time and HBM. Tests are CPU only: an engine on the tiny GLM with the flags off gives the same tokens, logprobs and graph keys as with the flags on and no request using them; six planted mutants were caught. Full CPU suite at d24e30e: 1730 passed, 22 skipped, 2 xfailed.
Graphs
- Pieces 1-7: host-side only.
- Pieces 8-9: add graph variants behind
--structured-outputsand--penalties, off by default.tools/graph_equiv.pyevidence that every default graph equalsdevelopment's, for the default,batch8and all fourmtp-*presets: captures running; results will be added here. - A Trainium2 run of the new variants: not done (see below).
Checklist
-
make checkpasses at the final head - Tests cover each piece as it lands
- Docs and
CHANGELOG.mdupdated per piece
Summary by CodeRabbit
- New Features
- Added SGLang-compatible generation, tokenization, health, and cache-management endpoints, plus OpenAI Responses and Anthropic Messages APIs.
- Added structured-output constraints, minimum-token and thinking-budget controls, sampling penalties, and support for multiple choices.
- Added configurable scheduling priorities, queue limits, and cache eviction policies.
- Added serving statistics, a dashboard, optional production logging, and prompt-token caching.
- Bug Fixes
- Stop sequences now respect minimum-token settings, with options to retain stop text and skip special tokens. Full queues return 503 responses; unsupported parameter values return 400 responses.
- Documentation
- Expanded API, configuration, and architecture guides to cover these capabilities.