diff --git a/README.md b/README.md index 76ee63c..966b83e 100644 --- a/README.md +++ b/README.md @@ -265,7 +265,7 @@ gmlx builds on [llama.cpp and ggml](https://github.com/ggml-org/llama.cpp) for the GGUF format and the K-quant reference implementations, [MLX and mlx-lm](https://github.com/ml-explore/mlx-lm) for the runtime and model implementations, [mlx-vlm](https://github.com/Blaizzy/mlx-vlm) for the -batching server engine and the vision towers, +server app, generation step loop and vision towers, [mlx-whisper](https://pypi.org/project/mlx-whisper/) for speech-to-text, and [mlx-audio](https://pypi.org/project/mlx-audio/) for text-to-speech. diff --git a/docs/api.md b/docs/api.md index 5db1867..943af6f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -153,10 +153,11 @@ profile. `requests[]` is high-cardinality and contributes only its count. ## API capabilities -The protocol surface is inherited from mlx-vlm, since gmlx swaps the model -layer and not the handlers, so its request features work unchanged on GGUF -models. Context windows come from the GGUF's metadata, with no server-side -override or request-level context setting. +The protocol surface follows mlx-vlm's, with gmlx patches on the handlers +(around twenty route-level patches, including the added `/v1/completions` +route), so upstream request features work on GGUF models. Context windows +come from the GGUF's metadata, with no server-side override or +request-level context setting. ### Tool calling diff --git a/docs/internals/README.md b/docs/internals/README.md index ec820b2..c32a837 100644 --- a/docs/internals/README.md +++ b/docs/internals/README.md @@ -5,7 +5,7 @@ How gmlx works internally, for contributors. Users start at | Page | Contains | |------|----------| -| [serving-architecture.md](serving-architecture.md) | how the loader, engine, batching and HTTP layers compose, plus the config-server call graph | +| [serving-architecture.md](serving-architecture.md) | the upstream mechanism and the gmlx scheduling policy, from GGUF bytes to streamed response | | [speculative-batching.md](speculative-batching.md) | how speculative decoding and continuous batching run together | | [prompt-cache.md](prompt-cache.md) | prompt cache tiers per architecture, reuse counters, environment switches | | [adding-architectures.md](adding-architectures.md) | what adding a model family involves and the acceptance gate | diff --git a/docs/internals/serving-architecture.md b/docs/internals/serving-architecture.md index 4c4aab0..8672ce5 100644 --- a/docs/internals/serving-architecture.md +++ b/docs/internals/serving-architecture.md @@ -5,80 +5,104 @@ server, for contributors. This page covers the implementation, while the config surface is documented in [server-config.md](../server-config.md) and the endpoints in [api.md](../api.md). -The server is stock mlx-vlm, with its app, batching engine and protocol -handlers untouched, and gmlx reaches into it through patched seams. Loads -route to the gmlx loader, which reads GGUF bytes through mlx-kquant's C++ -reader and swaps model leaves for K-quant kernels, and the stock engine then -executes those kernels in its own forward pass. There is no engine fork. -The seam inventory, and why each one is fragile, is in -[upstream-upgrades.md](upstream-upgrades.md). +The mechanism is stock mlx-vlm and the policy is gmlx. Upstream owns the +FastAPI app object, the protocol handlers and SSE formatters, the engine's +step loop, and the fp16 `BatchKVCache` layout. gmlx owns the scheduling +policy around that loop: admission (`admit_gate`), prioritization +(`batch_sched` with `auto_ratio`), and resource arbitration (`governor`, +`capacity`, `queue_cap`). All of it installs through the patch layer, whose +seam inventory, well over a hundred entries, is in +[upstream-upgrades.md](upstream-upgrades.md). Loads route to the gmlx +loader, which reads GGUF bytes through mlx-kquant's C++ reader and swaps +model leaves for K-quant kernels, and the stock step loop executes those +kernels in its own forward pass. There is no engine fork. ## From file to response ```mermaid flowchart TD subgraph DISK["GGUF on disk"] + direction LR LLM["text LLM GGUF
K-quant: Q4_K / Q6_K / MXFP4 ..."] MM["mmproj GGUF
float or Q8_0, VLM only"] + LLM ~~~ MM end subgraph LOAD["Loader: gmlx.load_model"] - direction TB - PARSE["parse file bytes + GGUF->HF name remap"] + direction LR + PARSE["parse file bytes
GGUF->HF name remap"] SYNTH["config + tokenizer synth
including the chat template"] - BUILD["build stock model class"] - KQ["install K-quant leaves
KQuantLinear, gather_qmm, KQuantMultiLinear"] + BUILD["build stock
model class"] + KQ["install K-quant leaves
KQuantLinear, gather_qmm"] PARSE --> SYNTH --> BUILD --> KQ end - LLM --> PARSE - MM -. VLM .-> PARSE subgraph ADAPT["Model adapter + residency"] - direction TB - WRAP["text -> gmlx vendored text_only.Model
(VLM -> mlx_vlm vision/audio model class)"] - STOP["attach StoppingCriteria to tokenizer"] - REG["multi-model residency pool
pinned + LRU, one process wired_limit"] + direction LR + WRAP["text -> gmlx vendored text_only.Model
VLM -> mlx_vlm vision/audio class"] + STOP["attach StoppingCriteria
to tokenizer"] + REG["residency pool
pinned + LRU, one wired_limit"] WRAP --> STOP --> REG end - KQ --> WRAP - subgraph ENGINE["Engine: mlx_vlm.generate.ar.BatchGenerator"] - direction TB + subgraph TICK["gmlx tick policy: wrappers on BatchGenerator._next, install order"] + direction LR + TG["tick guard
OOM / GPU-fault"] + GOV["governor
collision bands"] + MT["memtrace (opt)
GMLX_SERVE_MEMSTATS"] + QCC["queue-cap
census"] + FG["fresh gate
cache-freshness hold"] + AG["admit gate
headroom projection"] + PACE["pacer
prefill pacing"] + ST["step timing (opt)
GMLX_STEP_LOG"] + TG --> GOV --> MT --> QCC --> FG --> AG --> PACE --> ST + end + + subgraph ENGINE["Upstream mechanism: mlx_vlm.generate.ar"] + direction LR EMB["precompute inputs_embeds
get_input_embeddings(input_ids)"] - BG["continuous batching (embeds-in)
insert(inputs_embeds) -> next()"] - KVC["BatchKVCache (ragged, left-pad)"] - APC["prompt cache
exact / checkpoint / block tiers"] + BG["step loop (embeds-in)
one decode step + one prefill chunk"] + KVC["BatchKVCache (ragged, left-pad)
or gmlx kvarn KV, per-row ends"] + APC["prompt cache (gmlx APC)
exact / checkpoint / block tiers"] SAMP["sampler / logits procs / stop"] - MTP["speculative draft / MTP"] + MTP["owned MTP round loop
gmlx.spec.engine"] EMB --> BG BG --- KVC BG --- APC BG --- SAMP - BG -.- MTP + BG --- MTP end - REG --> EMB - subgraph PROTO["HTTP: mlx-vlm FastAPI"] - direction TB + subgraph PROTO["HTTP: mlx-vlm FastAPI, ~20 gmlx route patches incl. added /v1/completions"] + direction LR TOOLS["tool-call extractor
mlx_lm.tool_parsers (from chat template)"] ANTH["/v1/messages (Anthropic)"] OAI["/v1/chat/completions (OpenAI)"] RESP["/v1/responses (OpenAI Responses)"] SSE["streaming SSE formatters"] TOOLS --> ANTH & OAI & RESP - ANTH --- SSE - OAI --- SSE - RESP --- SSE + ANTH & OAI & RESP --> SSE + end + + subgraph CLIENTS["Clients"] + direction LR + CC["Anthropic-API client on /v1/messages
such as Claude Code via ANTHROPIC_BASE_URL"] + SDK["OpenAI SDK / curl / apps
on /v1/chat/completions, /v1/responses"] + CC ~~~ SDK end - BG --> TOOLS - CC["Anthropic-API client
such as Claude Code via ANTHROPIC_BASE_URL"] - SDK["OpenAI SDK / curl / apps"] - ANTH --> CC - OAI --> SDK - RESP --> SDK + LR["live requests: per-request rows,
wraps ResponseGenerator._step, never steps the engine"] + + DISK --> LOAD --> ADAPT --> TICK --> ENGINE --> PROTO --> CLIENTS + TICK -.- LR ``` +Eight wrappers assign `BatchGenerator._next`, shown in install order in the +tick-policy box. Six install by default; memtrace and step timing are +env-gated. The live-requests publisher sits beside the stack, not in it: it +wraps `ResponseGenerator._step` to publish per-request rows and never steps +the engine. + ## What the diagram leaves out The loader's output is a model, config and tokenizer triple with no @@ -101,6 +125,66 @@ assistant id never reaches the HTTP layer as itself: the tool loop runs on a worker thread and each round re-enters the server as an ordinary loopback client ([served assistants](../assistant.md#served-assistants)). +## Scheduling policy + +Admission, prioritization and resource arbitration run in the gmlx +wrappers around the stock step loop. + +| Module | Policy | +|---|---| +| `batch_sched.py` | decode-priority prefill pacing: a chunk runs only after decode has banked ratio x last chunk time | +| `auto_ratio.py` | derives the pacing ratio from a retention floor, with hysteresis, dwell and a pacing-attributable deadline | +| `admit_gate.py` | projects committed bytes before a prompt batch forms and defers the join instead of failing | +| `governor.py` | ticks-to-collision banding with separate rate and one-shot accounting | +| `queue_cap.py` | rejects over-cap requests with a 503 and a computed Retry-After instead of holding sockets | +| `capacity.py` | derives the depth-width frontier at boot and sets decode concurrency from it | + +Stock `_next` runs one decode step then, unconditionally, one 2048-token +prefill chunk per tick. At depth that chunk head-of-line blocks decode: +measured at d50k, 80 to 84 percent of decode wall is stall, and about 56 +percent at d14k. The pacer admits a chunk only once decode has accumulated +ratio x last_chunk_time since the previous chunk. Prefill runs at full +speed whenever no decode batch is live, so single-stream TTFT is untouched. + +`auto_ratio` resolves `decode_prefill_ratio: auto` per tick. A retention +floor rho, default 0.5, fixes the paced ratio at rho/(1-rho), and the +resolver selects that ratio or zero using an incumbency rule, a chunk-cost +threshold with hysteresis and dwell, and a deadline that ages +pacing-attributable seconds only, so capacity-blocked waits accrue nothing. +There is no queue-depth term. + +The admit gate prices the bytes a candidate join would commit against +measured headroom before the stock admission arm forms a prompt batch, and +hides the pending list for the tick while the projection does not fit. Two +anti-deadlock rules bound it: an idle server is never declined, and past +the defer ceiling it admits one row per tick, loudly. + +The governor computes ticks-to-collision from one shared accounting and +walks a band ladder from green to red. Bands are rates, not levels: a deep +batch at flat headroom is green, a shallow one growing fast is not. Rate +and one-shot costs are accounted separately, and dwell minimums plus a cap +on sheds per minute prevent thrash. + +The queue cap rejects before enqueue with an HTTP 503, a body naming the +cap and depth, and a Retry-After set to the estimated drain time clamped +between 2 and 60 seconds, instead of holding sockets until the queue +timeout. + +Capacity derives a table at model build time from the same admit-side cost +model requests are priced with: max context at width 1, max width at +representative depths, and the depth-width frontier. Decode concurrency is +min(`GMLX_DECODE_BATCH`, frontier width) and the queue cap default follows +it. A configuration that cannot fit at width 1 is refused at boot with +numbers; `GMLX_OVERCOMMIT=1` disables the refusal and the derived ceilings. + +The modules interlock. The governor's band is the admit gate's hard hold. +The admit gate's deferred set keeps `auto_ratio` from charging capacity +waits to pacing. The pacer's observed chunk cost feeds `auto_ratio`'s +threshold and the prefill chunk sizing. The capacity frontier bounds +decode width, which sets the queue cap default. All of it surfaces on +`/v1/metrics`: the capacity table, per-request rows and rate views +([api.md](../api.md#capacity-and-live-request-metrics)). + ## The request path through the seams ```mermaid @@ -110,8 +194,9 @@ sequenceDiagram participant R as residency pool participant S as serving (resolver + bridge) participant L as loader + kquant swap - participant E as BatchGenerator (kq.* kernels) + participant E as gmlx tick stack + stock step loop C->>A: POST /v1/chat (model "id@profile") + A->>A: queue depth cap, over-cap gets 503 + Retry-After [patched] A->>R: get_cached_model(id) [patched] R->>S: resolve_request_model(id@profile) S-->>R: abspath + ResolvedModel, sets _active_spec @@ -126,8 +211,15 @@ sequenceDiagram end A->>A: _build_gen_args seeds sampling from the active profile [patched] A->>E: generate(...) + E->>E: admit gate holds the join until projected bytes fit [patched] + E->>E: paced ticks, governor bands, contained faults [patched] E-->>C: stream tokens ``` -On this path the patched seams are the residency lookup, the load call and -the generation argument builder. Everything between them is stock. +On this path the patched seams shown are the queue cap, the residency +lookup, the load call, the generation argument builder, the admit gate and +the wrapped tick. The registry in `gmlx/upstream/seams.py` declares 124 +seams, 71 of them critical; the largest clusters sit on +`mlx_vlm.generate.ar` (23), `mlx_vlm.apc` (19) and `mlx_vlm.models.cache` +(18). The step loop inside the wrappers is stock. The policy around it is +not. diff --git a/gmlx/upstream/seams.py b/gmlx/upstream/seams.py index 0570ab0..1d7925d 100644 --- a/gmlx/upstream/seams.py +++ b/gmlx/upstream/seams.py @@ -1,7 +1,7 @@ """Contract registry for every upstream symbol gmlx patches or deep-imports. -gmlx monkeypatches ~30 private symbols across mlx-vlm and mlx-lm and -deep-imports model internals. Those seams are guarded structurally (try/except +gmlx monkeypatches well over a hundred symbols across mlx-vlm and mlx-lm +and deep-imports model internals. Those seams are guarded structurally (try/except ImportError + idempotence flags), so upstream renames or rewrites fail silently - stock behavior quietly returns, or a confusing error surfaces far downstream (mlx-vlm 0.6.4 vendoring switch_layers turned into a gather_mm