Skip to content

Studio hands KV preemption to a llama-server that parks slots itself - #10358

Draft
danielhanchen wants to merge 147 commits into
feat/kv-preemptionfrom
feat/kv-preemption-server-swap-c
Draft

Studio hands KV preemption to a llama-server that parks slots itself#10358
danielhanchen wants to merge 147 commits into
feat/kv-preemptionfrom
feat/kv-preemption-server-swap-c

Conversation

@danielhanchen

@danielhanchen danielhanchen commented Sep 5, 2026

Copy link
Copy Markdown
Member

Stacked on #10301; only the last two commits are new. Pairs with unslothai/llama.cpp#184 and #190.

Summary

When the launched llama-server can park a slot in host RAM itself (--preempt-ram, unslothai/llama.cpp#184), Studio stands its own preemption down, lets every chat use the whole context window, relays the server's : preempted and : resumed stream comments (unslothai/llama.cpp#190) to the GUI as the existing paused line, and never mistakes a park for a stall. On the bundled upstream build, or with UNSLOTH_LLAMA_PREEMPT_MODE=studio, behaviour is exactly #10301's.

Why. Studio's own preemption pauses a chat by aborting its upstream stream and resumes by re-prefilling the partial. That is exact at the seam but not byte-identical once the MTP head is involved, and it has to hold back reaction headroom (192 cells per slot plus the drafts, 776 of 8192 at four slots). The server swap restores the sequence in place, byte-identical, at N minus 8 cells. With both active Studio's lower watermark fires first and the better mechanism never runs; with Studio's off, a parked slot was a silent stream that hit the 120 s stall timeout with nothing shown to the user.

Policy

  • Detection: --preempt-ram from the existing --help probe, --kv-unified on, and no hand-typed --preempt-ram 0. UNSLOTH_LLAMA_PREEMPT_MODE=auto|studio|server, default auto (server when the build can). The armed log line says which mode is live.
  • Off by default, on both sides. The fork's llama-server defaults --preempt-ram to 0 (server: preemption notices, asynchronous parks and exact concurrency together llama.cpp#197), and Studio's controller is off unless UNSLOTH_LLAMA_ADMISSION_PREEMPT=1 (studio: KV preemption so parallel chats share one cache instead of killing each other #10301). A default install therefore launches a stock llama-server: no --preempt-ram, no LLAMA_ARG_PREEMPT_RAM, no probe, no participants, share-based admission. With the switch on, a build that carries the flag, a unified cache, no studio override and no budget of the user's own, the launch names --preempt-ram <budget> itself, sized to hold every park's target and draft state (an unsized figure for a pool an auto-fit context leaves unmeasurable, judged after launch against the context the server chose). server_preempts_kv means "this launch turned parking on", re-judged after launch against the argv that spawned, not "this build could": a server that parks nothing never excuses a stalled stream as a park, never tells the controller to stand down for a preemptor nobody is running, and never counts raw relays as pausable. A user's --preempt-ram N or LLAMA_ARG_PREEMPT_RAM still wins, and their zero still stands the server down.
  • Server mode: the controller plans no victims and its buffer is 0 (the server reserves its own drafts and an 8 cell margin, and parks the moment the next batch does not fit). Admission, the tool loop, PARKED_ON_TOOL, context_truncated with fits:false and the ledger keep working. A parked slot is left out of the resident figure, since its cells are in host RAM; counting it made a nearly empty pool read as over the ceiling and erased idle slots for room that was never short.
  • Stream: the three stream consumers (plain chat, tool loop, legacy tool loop) turn : preempted and : resumed into the same pause events the Studio-side preemption yields, so the routes emit : preempt-paused and : preempt-resumed unchanged and the GUI shows "Paused while another chat finishes". The server's 2 s : preempt-keepalive renews the stall deadline. For a swap build without the notices, the read wrapper consults /metrics requests_preempted before treating silence as a stall, capped at 30 minutes.

Results

Qwen3.5-4B with two MTP drafts, -c 8192, four slots, seed 0, temperature 0:

  • Four chats through the API: 4 of 4 in server mode (318 tok/s aggregate, one park), 4 of 4 on the tool-loop path (two parks), 4 of 4 in studio mode on the same build, 4 of 4 on the bundled upstream build. Zero errors in every run.
  • Four Playwright Chromium chats, tools off: 4 of 4 finished, paused label on the three chats that were parked, zero error boxes.

Qwen3.6-35B-A3B with two MTP drafts, -c 8192, four slots:

  • Four chats through the API: 4 of 4 with every chat producing its full 2048 tokens, one natural park, zero errors. Studio mode on the same build and load: 4 of 4 without an HTTP error, but one chat ended empty after the Studio-side controller gave up waiting for room (a studio: KV preemption so parallel chats share one cache instead of killing each other #10301 path, followed up separately).
  • Four Playwright Chromium chats: 4 of 4 finished, paused label on 3 of 4, eight natural parks lasting 59 to 422 s with no stall fired, zero error boxes, zero Context size has been exceeded.

Exactness

Forced parks every 64 tokens (LLAMA_SERVER_PREEMPT_EVERY=64 on the child) against an unforced run, through the Studio API, 1200 tokens, MTP on:

model prompt bytes parks identical
4B 0 5077 18 yes
4B 1 5700 18 yes
35B 0 5276 18 yes
35B 1 4756 18 yes

Cost when it does not fire

Solo 1200-token chats on the 4B, same build, back to back: server mode 209.4 / 192.9 tok/s, studio mode 209.2 / 192.1, server again 208.8 / 193.4. The server-mode path adds one string compare on stream lines that are not data:. The 35B solo numbers were taken on a GPU shared with a foreign job at full utilisation and spread 14x between passes, so they are not quoted.

Tests

31 new tests in test_server_side_preemption_handoff.py (mode and property, probe, controller, stream relay on both surfaces, a park before the first token, an upstream stream without comments is bytewise today's, the read-wrapper grace) plus the parked-slot occupancy test. 737 pass across the handoff, preemption, admission, tool loop, chat generation, slot release and stream cancel suites.

Limitations

  • Server mode trusts the server's --preempt-ram budget. When that budget cannot hold a victim, the paired llama-server parks it by dropping its cells and re-prefilling on resume, so Studio never re-arms its own preemption for that case.
  • A park during prompt processing is announced at park time by the paired llama-server (headers and : preempted go out then), so a build that carries exact_concurrency on /props is never asked the aggregate /metrics reading; only a swap build predating the notices keeps that fallback.
  • The metrics grace for a swap build without the notices is unit-tested only; the notify build never needs it.
  • Non-streaming and passthrough surfaces get no pause signal, and the raw non-streaming relays keep their fixed generation timeout: a park that outlasts it ends the request the way a generation that outlasts it always has. The GUI streams.

Follow-up: slimmer

Five refactor commits, behaviour and coverage unchanged: the scripted backends, recording policies, fake responses and SSE builders that ten test files each carried are one shared tests/preempt_fakes.py; the three pause sites in llama_cpp.py share one charge (the larger of the observed count and the estimate) and one declined-pause helper; the eighteen spellings of the controller key in routes/inference.py are one _preempt_key(); the barrier, recorder and registry-reset fixtures are shared; both no-flash respawns hand exact concurrency back through one helper. 1042 test functions before and after, none removed or renamed.

Two comment sweeps then cut the added comment and docstring lines from 6313 to 1871 (70 percent), the AST gate confirming only comments, docstrings and whitespace changed. Against main the PR is 101 files and 17505 insertions, from 22555.

Verified at every commit by ruff and the 1484 preemption, admission, tool-loop and resume tests, then live on the slimmed Studio with the slimmed llama-server: four Chromium tool chats through the GUI at 8k finished 4 of 4 with 9 parks and 9 resumes, the paused label on three of them and no error; four parallel tool chats through unsloth run at 16k finished 4 of 4 with 10 pauses and no context error.

Follow-up: fixes and fewer tests

Three more commits.

Resume and lease fixes (4b9c38b). A resumed attempt keeps the reasoning that preceded its prose: the assembler folds the attempt's reasoning_content and any reasoning a previous pause left trailing into the partial the next attempt extends, instead of dropping it once prose had started. The final pass replays its own prose rather than the cumulative display text, which carried the thought into the answer. Stop during a pause ends the wait: ControllerPreemptionPolicy.await_resume takes the cancel event and returns as soon as it is set, through one helper at the three pause sites. A park before the first token, and the keepalives that follow it, renew the first-token deadline instead of timing it out. A pause declined at the resume cap keeps the attempt: it is charged, its partial goes back through the assembler, and the final pass extends it with continue_final_message under what is left of the caller's cap, rather than starting afresh with the whole allowance. A long pause no longer expires a durable run: the server's : preempt-keepalive and the Studio-side wait (now a generator that says so every two seconds while the policy waits) both reach the routes as a keepalive event, forwarded as : preempt-keepalive, which the durable producer renews the run lease on at the queue-wait rate and the frontend ignores.

Admission fixes (a31bc71). One predicate, preemption_eligible(), reads the rollout switch together with the admission-control and KV-budget switches for the controller's active and plan_preemptions, for pricing, for arming and for the disarm's cache erase, so an install that switched admission off no longer arms a controller policy over leases charged nothing, and the wire clamp returns None under either opt-out instead of narrowing an unstated cap. LLAMA_ARG_PREEMPT_RAM=0 in the child's environment disables the server hand-off the way --preempt-ram 0 on the launch line does, read in llama.cpp's own order. Exact concurrency is reported on only on the build's own evidence: a build that ignores the variable starts perfectly, so the load reads the --preempt-ram capability the same fork ships and reports unavailable without it, which under the on setting fails the load as designed. The passthrough, Responses and Anthropic raw surfaces reserve as unpausable and pay the honest price, since a stream registered as never chosen has nothing to reclaim.

Fewer tests (915cba3). 44 new test files and 11009 added lines become 10 files and 4494 lines, 604 test functions become 167 plus 21 parametrisations, with every behaviour kept: the incidents folded into their surface files, and every source-level wiring check but two rewritten as behaviour (a real FastAPI app with a recording backend and a patched controller ledger for the arming on every local chat surface, a recording erasures fixture for the disarm, a policy that reads the signal inside on_resumed for the clear-before-resume order). The survivors were checked against five scratch mutations of the production code (the arm made a no-op, the unregister skipped, the one-shot ledger no longer clearing, on_resumed before the clear, settlement never compacting); each was caught. The two source-level checks kept are the frontend's declared comment set, which only the TypeScript can answer, and the respawn retry's argument forwarding, which needs a llama-server that dies and comes back.

Verified: ruff and the pre-push gate clean; the admission and preemption suites 1414 passed and 5 skipped; the whole backend suite 39557 passed and 221 skipped, with every failure in diffusion, audio, video and host-topology files this PR does not touch, the non-diffusion ones failing identically on the merge base. Live on the paired llama-server: four Chromium tool chats through the GUI at 8k finished 4 of 4 with the paused label on three, 9 parks and 9 resumes, no error box and no context error; unsloth run at 16k with MTP and exact concurrency on, four parallel tool chats, finished 4 of 4 at 191 tok/s aggregate with 5 server parks and 5 resumes and no context error, Studio reporting mode=server exact=on for every chat.

Round two of fixes (1e7f37b). The parked-slot metrics probe sends the API key, since /metrics is not in llama-server's public set and every check under --api-key was a 401. Stop is read before every resume grant attempt and passed into the lease's resume, so a grant that succeeds at once no longer carries a stopped chat into the wait. The wire clamp takes the same pausable flag the reservation takes, and the passthrough body is built with it off, so an uncapped passthrough is held to its share while pausable chats keep the window. The settings page shows what the next load resolves to when nothing is stored, so an inherited LLAMA_EXACT_CONCURRENCY reads as on. Three tests added; the admission, preemption and tool-loop suites 1152 passed and 5 skipped.

Round 23

Five review items plus a second round of five isolated reviews of the updated heads. Everything below was reproduced before it was changed; the reproductions are now tests.

Finding Fix
The routed DeferredPreemptionPolicy.await_resume took only a timeout, so the caller's cancel_event keyword raised TypeError and the fallback waited without Stop The wrapper and the protocol carry cancel_event; an older inner policy still works
Exact concurrency was reported on off the --preempt-ram capability, which a build without the mode can carry Read off the running server: /props now carries exact_concurrency (unslothai/llama.cpp#197)
contradicting_args kept every occurrence, so --flash-attn off --flash-attn on read as a contradiction and on killed a valid server Each option resolves to its last value, as llama-server applies argv
The Anthropic streams dropped preempt events; each restarted the stall keepalive, so a long park sent no bytes Both Anthropic streams forward the same : preempt-paused, : preempt-keepalive, : preempt-resumed comments as the OpenAI surface
The exact-concurrency chip read the local server's state next to a hosted model selected beside it Rendered only for the resident local model
A resumed attempt that went on thinking and then called a tool carried only its own thought; the merge replaced the earlier one append_assistant_turn concatenates reasoning_content like it concatenates content
A declined pause's final pass started its snapshots empty, so a non-streaming drain returned the suffix alone, and its cap deducted only the declined attempt The pass extends the partial (no turn boundary, cumulative snapshots) and deducts everything the call spent
A pause with the caller's cap spent reopened the stream for one floored token, past the cap, at all three sites The turn ends with length
A parallel round keyed on the arguments as they arrived, so two spellings healing to one call both ran Keyed on the ledger's healed key (ToolLoopController.call_key)
The Anthropic raw passthrough sent the caller's cap or the whole window while reserved as unpausable Sent the same enforced cap as the OpenAI passthrough
An unstated, unpausable request reserved the flat allowance while being permitted the rest of its share: two beside one stated request reserved 16064 of 16384 and were permitted 22192 Charged what it is sent, the rest of its share
The residency sweep ran with preemption ineligible; a never-configured controller has a zero ceiling, so idle prefix caches were erased on the paths meant to keep their old behaviour Gated on the same predicate arming and pricing use
A policy that raised during the wait was treated as a grant in both tool-loop pause handlers Not a grant: the turn ends with its partial, as a refusal does
A rollback abandoned a provisional tool_start without closing its card The card is closed before the pause is announced
UNSLOTH_LLAMA_ADMISSION_PREEMPT=0 left the child parking on its own default budget One switch: the child gets LLAMA_ARG_PREEMPT_RAM=0 unless something named a budget
A park that outgrows --preempt-ram is re-prefilled, and a re-prefill is not byte-identical on CUDA (it is on CPU) An exact launch sizes --preempt-ram to the whole pool when the server's default would not hold it; the server warns if a re-prefill happens under the mode anyway

| The raw passthrough loops (Responses, chat and completions passthrough, Anthropic passthrough) relayed data: lines only, so a request the server parked sent the client nothing for the whole park | _server_park_sse forwards the server's park notices as the same comments every other surface sends |
| Exact concurrency was reported on with Studio doing the pausing (UNSLOTH_LLAMA_PREEMPT_MODE=studio, or the server's parking off), and a chat Studio resumes is re-prefilled rather than restored | The state needs the server to be the one parking; on fails the load naming the reason, auto loads with a warning |
| A user-named --preempt-ram below the KV pool was kept and reported exact, and the first park to outgrow it is re-prefilled | The shortfall is judged at launch: a warning names the figure to raise it to, the state is unavailable, and on fails the load |
| A swap build predating the stream notices parks in silence, and the read wrapper's /metrics excuse forwarded nothing, so a durable run's lease could expire during a legitimate park | The excuse stamps the backend and the run loop renews the lease from it, with a bounded wait |

| A pause that landed as the caller's cap ran out folded the interrupted attempt into the accumulators and then built the terminal event from the same reading, so a four-token answer reported eight and twice its decode time | Both cap-spent sites hand the builder the attempt's prompt side only |

| An auto-fit context left the KV pool unknown at launch, so the exact launch neither enlarged the parking budget nor judged it while the state read on | With nothing named the child parks without a limit; a named budget is judged after launch off the context the server chose, and a pool that cannot be sized is not certified |
| Under auto, Studio-side pausing or parking switched off were only judged after launch, so the child ran the mode it was then reported not to have | Preflighted: the child is not given the flag, with a warning naming the reason; on still fails the load |
| The raw relays read llama-server without the metrics park grace, so on a swap build predating the stream notices a parked raw request was cut off as a stall | _aiter_llama_stream_items takes the backend's park probe and excuses the silence under the same thirty-minute cap |
| A raw relay was priced at the honest per-slot share even with the server parking slots itself, so a raw stream cost a whole slot of the window that it never used | Both entry points read server_preempts_kv and price a raw stream at the window on a parking server; a Studio-only server keeps the share |
| The raw relays' first-item deadline never asked the park probe, so a raw request parked during its prefill was cut off before its first token | Both first-read timeouts take the same metrics grace as the stall after it, under the same cap |
| UNSLOTH_LLAMA_ADMISSION_PREEMPT=0 with nothing named stood the child's parking down after the auto preflight had passed it | The preflight sees the opt-out too and drops the flag with a warning |
| A diffusion runner was asked to reload for exact concurrency it cannot apply | The reload check returns off for a diffusion runtime |
| The raw relays' park grace retried the httpx iterator after its read timeout, and an async generator that raised is closed, so the retry got StopAsyncIteration and a parked raw request ended as if its answer were complete | _install_park_aware_read waits the park out on the network stream's read, below the iterators, the way the synchronous read wrapper does; a closed iterator is never retried, so a stall is an error rather than a short answer |
| A declined pause that had spent the caller's cap still ran the final pass, floored at one token past the cap | The decline ends the turn with length when the cap is spent, as a granted pause does |
| UNSLOTH_LLAMA_PREEMPT_MODE=studio left the child its default parking budget, so the server parked on its own while Studio reported it did not and no relay excused the silence | Forced Studio mode stands the child's parking down like the global opt-out; a named budget keeps its say |
| The park grace's cap ran from the start of the read, so a request parked at the first-token deadline was cut off after the normal window plus what was left of the cap, not the grace the backstop allows | Both read wrappers start the grace at the deadline they first crossed and bound each retry by what is left of it |
| The park probe scrapes /metrics with a blocking read, and the async relays called it on the event loop, holding every other stream for up to three seconds | The probe runs off the loop, at the network read and above the iterator alike |
| The final pass's refusal to pause again built its terminal event without folding the interrupted attempt, and its not-resumed end folded it and then added the reading again | Both ends fold the attempt once and hand the builder its prompt side |
| Exact concurrency appended --kv-unified after the user's extras, so an explicit --no-kv-unified was reversed by last-arg | The opt-out is kept and reported as the contradiction it is; auto does not start a mode the extras contradict, on fails the load naming the flag |
| A durable run renewed a silent lease only from the stamp the read wrapper writes at its read deadline, and before the first token that deadline is the whole first-token budget, the default lease exactly, so a run parked during prefill on a swap build without the notices was cancelled while healthy | With no stamp in hand the run loop asks /metrics itself, off the event loop and at most once in five seconds across runs, still under the wrapper's cap |
| The exact chip and the field docs said llama-server refused the mode when Studio had withheld it itself | They say the load is not running with it; the load warnings carry the reason |
| The launch generated --preempt-ram -1 on the guess that exact would run, and every abandoned exact attempt kept it, so a fallback load parked on an unlimited budget | No budget is generated for a pool the launch cannot size; the server's default is judged after launch |
| The grace cap test timed real sleeps, which overshoot on a loaded runner and read as a short grace | The test drives a fake clock advanced by exactly each window |
| A parking shortfall found after an auto-fit launch was reported as the server having come up without the mode, while /props said it ran | The warning and the on refusal say the mode runs but cannot hold every park, so a chat re-prefilled after such a park can differ |
| UNSLOTH_LLAMA_PREEMPT_MODE=studio with a named --preempt-ram ran both mechanisms: the stand-down read any named budget as somebody's say, and LLAMA_ARG_PREEMPT_RAM=0 alone could not have won because llama.cpp applies argv after the environment | One owner pauses chats: the stand-down reads only the mode, zeroes the environment and the argv token in place, and the load warning names the flag it overrode; auto no longer starts the mode on the strength of a named budget while preemption is off |
| The parking budget was sized from the target pool plus a margin, while the server saves the target and the MTP draft state per park and holds several parks at once, so a budget the launch certified was short with a drafter | _exact_parking_need_mib prices both pools plus a per-park margin, times the slot count; a draft cache that cannot be sized is judged after launch rather than certified |
| --cpu-moe, --n-cpu-moe and a CPU --override-tensor were not contradictions, though they place per-layer matmuls on a backend whose kernel selection is size dependent; and their last occurrence was taken to decide, while llama.cpp accumulates tensor overrides | All three are Studio's own contradictions (a zero --n-cpu-moe is the default spelled out); an earlier CPU placement is not taken back by a later occurrence |
| A park that fell to recompute under exact concurrency was invisible to the client | The server's : recomputed notice (unslothai/llama.cpp#197) is relayed as : preempt-recomputed on every surface and to durable followers, the final preempt counts ride the turn's metadata, a build that counts without the notice has it synthesised, and the exact chip names the re-prefilled answer on the chat it happened in |
| The stall excuse asked the aggregate requests_preempted, which counts every request, so a neighbour's park excused any stalled stream, including one that had already resumed | ServerParkNotices reads this stream's own : preempted and : resumed off the wire below the parsers; a stream that has heard a notice never falls back to the aggregate probe, which stays for a swap build predating the notices under the same cap |
| The base branch moved under this one: #10120's wire enforcement and the audit fixes on #10301 (the resume grant, the prompt-plus-generated ledger, the reasoning and failed-wait backports, the uncontended token path, the re-cost that prices what it sends and treats a Stop as a Stop) were not here | Merged; one wire bound carries both policies, and this branch's pausable stays on it and on the passthrough retry, so a raw relay keeps its share unless the server parks raw streams itself |
| The import-hoist gate read the shared test fixtures and three unused imports as botched hoists, and missed a name loaded inside an assignment target (app.dependency_overrides[get_current_subject] = ...) | The unused imports go, the fixtures are named once per module, and every assignment target records the names it loads, with a self-test case |

Round 25

Finding Fix
The exact parking budget shared the pool out per slot, but every parked history can be the whole unified pool Sized to the pool plus a per-park margin, times the slot count, with a launch warning when that exceeds the host's free memory
A raw holder registered as measured before its prefill was swallowed by a residency sample from before it Every raw holder registers unmeasured; streams mark themselves at their first data line, non-streaming requests stay counted for their answer
The tool relay cleared the turn's park counters on every final object Summed across the turn's requests
A lease at its floor could expire between two park probes on a swap build without notices The minimum lease covers two probes
A stream on a notice build that heard nothing could still be excused by another stream's park The aggregate reading is never consulted on a build that writes the notices
The exact preflight judged the parent's environment while Manual mode scrubbed the placement twins from the child Judged against the environment the child receives
The recompute note lived only in memory Persisted with the answer and read back on load
Dead longest-wins winner state in the controller Removed

@danielhanchen

Copy link
Copy Markdown
Member Author

Checked the gating in studio/backend/core/inference/llama_cpp.py, where server_preempts_kv needs both the --preempt-ram probe and a unified cache, so the bundled build keeps today's path and the relay only runs on the swap build. Will get this reviewed once the base branch lands.

danielhanchen and others added 3 commits September 6, 2026 06:02
A llama-server built with --preempt-ram parks a slot's sequence in host RAM when
the unified pool fills and restores it in place, byte for byte, at N minus a
handful of cells, and with the stream notices it says so with ": preempted" and
": resumed" SSE comments. Studio's own preemption, which aborts the upstream
stream and re-prefills the partial, stood in its way: its lower watermark fired
first, so the better mechanism never ran, and with it switched off a parked slot
looked like a stalled stream and the user saw nothing.

Detect the flag in the --help probe (supports_preempt_ram) and record it at
launch. On such a build, with a unified cache and the flag not zeroed by hand,
the controller runs in server mode: it keeps its ledger but chooses no victims,
and its buffer is zero, since the server reserves its own drafts and margin and
parks the moment the next batch does not fit, so every chat gets the window.
UNSLOTH_LLAMA_PREEMPT_MODE=auto|studio|server overrides; a build without the
flag keeps every existing path.

The stream consumers turn the two comments into the same pause and resume
events the Studio-side preemption already yields, so the routes emit
": preempt-paused" and ": preempt-resumed" and the GUI shows "Paused while
another chat finishes"; the policy marks the participant PAUSED and DECODING.
For a swap build that predates the notices, the cancel-aware read wrapper asks
/metrics (requests_preempted) before calling a silent stream a stall, bounded
by 30 minutes; it lives in the wrapper because an httpx body iterator that has
raised cannot be waited through.

Durable chat runs, which the GUI uses for a plain chat, kept data lines only and
so never showed a pause of either kind: with nine server-side parks of 5 to 30 s
the label was shown zero times. The worker now relays the two comments as
chunks carrying the frontend's own _admissionStatus field, written at once
rather than batched, since nothing follows a pause for as long as it lasts.

Measured on the 4B with two MTP drafts at -c 8192, four slots: forced parks
every 64 tokens through Studio are byte-identical to the unforced answer on
both prompts (18 parks each); four chats through the API finish 4 of 4 with
one park at 318 tok/s aggregate, and 4 of 4 on the tool-loop path; four GUI
chats finish 4 of 4 with the pause shown on 3 of 4 and no error box; solo
throughput in server mode is 209.4 and 192.9 tok/s against 209.2 and 192.1 in
studio mode on the same build.
llama-server reports a parked slot's logical sequence through /slots, but the
sequence is in host RAM until it is restored. Counting it made a nearly empty
pool read as over the ceiling and erased idle slots for room that was never
short.
@danielhanchen
danielhanchen force-pushed the feat/kv-preemption-server-swap-c branch from d7880b0 to 7044b0e Compare September 6, 2026 06:03
…ve to

A llama-server built with LLAMA_EXACT_CONCURRENCY set (unslothai/llama.cpp#194)
gives a chat the same generated tokens whether it decodes alone, beside three
other chats in one unified KV cache, or across a park and restore that moved its
cells. There is no command-line flag for it, no --help entry and nothing in
/props: the whole interface is an environment variable. So the only way to ask
for it today is to set a llama.cpp variable on the Studio process, which is not
discoverable, is not per load, and hands the mode to every child whether or not
that load wanted it.

UNSLOTH_LLAMA_EXACT_CONCURRENCY=auto|off|on is the switch, shaped like
UNSLOTH_LLAMA_PREEMPT_MODE next to it, with the same three values also on the
load request (exact_concurrency) and in a stored setting the GUI can write
through GET/PUT /api/settings/exact-concurrency. Highest wins: the environment,
then the request, then the store. Default off, because the mode costs about 9
percent of solo decode on a dense model and more on a mixture of experts, and
buys nothing for a chat that never shares its cache. An inherited
LLAMA_EXACT_CONCURRENCY is read as the default rather than ignored, so whoever
set it before this existed keeps it; an explicit off takes it back out of the
child environment, since off is the one answer that has to be obeyed exactly.

The launch sets the variable on the child rather than relying on inheritance,
and squares Studio's own line with the mode: --parallel 1 skips --kv-unified,
which the paged KV pool needs whether or not anything else is decoding, so exact
mode adds it. Studio never emits --cache-reuse and already emits
--no-context-shift; a contradiction in the user's extra args (--cache-reuse,
--context-shift, --no-kv-offload, a quantized KV cache, flash attention off) is
named in a load warning instead of arriving as an error about a launch line the
user did not compose.

Detection is by trying, because there is nothing to ask. Under auto, a child
that dies naming the mode is relaunched once without the variable, ahead of
every other rung in the ladder since the server said which of its own
requirements it could not meet; the attempt bound goes from three to four so the
ROCm correction and the fit recovery do not lose their slot to it. Under on the
refusal is the answer, and it is classified into a message that names the
setting and what the mode requires. The --flash-attn off respawn drops the
variable under auto for the same reason. What is reported is read off the argv
and environment that actually launched, so a respawn that took away flash
attention or the unified cache reads as unavailable rather than still claiming
the guarantee.

The load and status responses gain exact_concurrency (on/off/unavailable) and
requested_exact_concurrency, the duplicate-load check treats a different setting
as a reload since the mode cannot change without a new child, the preemption
snapshot carries it, and the "llama preemption armed" line gains exact=. The
one thing none of this can do is tell a build that predates #194, which ignores
the variable and starts perfectly, from one that granted the mode: that reports
on, and the tests record it as the limit it is.

103 new tests over the resolution order, the child environment, the launch line,
the refusal detection and the auto fallback, the reported state and the armed
line; 2278 pass across the preemption, start-failure and load-mode suites.
@danielhanchen

Copy link
Copy Markdown
Member Author

One more commit, 9956311d6: Studio gets an explicit switch for llama-server's exact concurrency mode (unslothai/llama.cpp#194), so a user does not have to set a llama.cpp variable on the Studio process.

UNSLOTH_LLAMA_EXACT_CONCURRENCY=auto|off|on, also a field on the load request and a stored setting behind GET and PUT /api/settings/exact-concurrency; the environment wins over the request, the request over the store. The default is off, because the mode costs throughput under concurrency and buys nothing for a chat that never shares its cache; an inherited LLAMA_EXACT_CONCURRENCY is read as the default rather than ignored, and an explicit off removes it from the child's environment. When the setting resolves on, Studio sets the variable on the child explicitly, adds --kv-unified if --parallel 1 had dropped it, and names any contradicting extra arg (--cache-reuse, --context-shift, KV offload off, flash attention off, a quantized KV cache). Under auto, a child that dies naming the mode is relaunched once without it and the load reports the mode unavailable; under on, that refusal is the answer, with a message quoting the server's line. The load and status endpoints report exact_concurrency as on, off or unavailable, derived from the argv and environment that actually launched, and the armed log line carries exact=.

Live on the 4B with the integration server: the child's environment carries the variable and the Studio process does not, status reports on, and a 300-token prompt at seed 0 and temperature 0 gives the same 1233 bytes alone twice and beside three neighbours; with the setting off the same cell diverges at byte 78. 103 new tests; the handoff suite and 2278 tests across the preemption, start-failure and load-mode suites pass; the three failures in a wider selection are the same three on the parent commit. No frontend indicator yet; the status field is there for one.

pre-commit-ci Bot and others added 2 commits September 6, 2026 08:36
The backend switch already reports the mode: /api/inference/load and
/api/inference/status carry exact_concurrency (on, off or unavailable) and
requested_exact_concurrency, and GET/PUT /api/settings/exact-concurrency stores
what the next load should ask for. None of it reached the screen, so the one
guarantee a user might want to check, that this chat's tokens do not depend on
the chats sharing its KV cache, was only visible by reading a JSON route.

The chip sits beside the model name, next to the GGUF quant it already shows.
"Exact" when the server came up with the mode, "Exact unavailable" when the load
asked for it under auto and llama-server refused, so it is running without the
guarantee; the title on both says what the mode buys, and the refusal names the
server rather than reading as a Studio fault. Nothing at all for off, which is
the default and every load that never asked, so the header keeps its room until
there is something to report.

The state comes through the ordinary status path: a field on
InferenceStatusResponse, loadedExactConcurrency on the chat runtime store, and
one line in the applier. Read on every status rather than only a seeded load, as
a tab opened onto an already-loaded model performs no load and would otherwise
never learn it. An older backend that does not publish the field normalizes to
off, since claiming the guarantee needs the server to have said so and the
absence of an answer is not one.

Settings gains the three-way select in the GGUF inference engine section, under
the compute backend and the llama.cpp folder, which is where the rest of this
server's launch settings already live. It writes the stored setting, says that
the mode costs throughput while chats decode together and applies on the next
load, and goes read-only while UNSLOTH_LLAMA_EXACT_CONCURRENCY pins the machine,
since a value saved then would not be the one used.

The label and title pair is a plain .ts with no imports: the test runner strips
types but does not transform JSX, so wording inside the component could not be
asserted. Eleven cases in tests/exact-concurrency-indicator.test.ts cover the
mapping, the older backend, both chip texts and the wiring.
@danielhanchen

Copy link
Copy Markdown
Member Author

Added a visible indicator for the exact concurrency switch (fc4ad27).

  • The chat header shows an "Exact" chip next to the model selector while the loaded model is decoding in exact concurrency, and "Exact unavailable" when the setting asks for it but the llama-server build cannot provide it. The chip reads the exact_concurrency field the status endpoint already reports.
  • Settings gains an Exact concurrency row (auto / off / on) under the llama backend section, backed by /api/settings/exact-concurrency.
  • Translations for the new strings in every locale overlay.
  • 11 frontend tests for the chip states and the settings row; typecheck, build and the bundle check pass; the full frontend suite passes (6925 tests). The backend switch and handoff suites pass on the rebased head (135 tests).

@danielhanchen

Copy link
Copy Markdown
Member Author

One measurement worth knowing before turning the switch on. On the llama.cpp side, exact concurrency derives its kernel column bound from the slot count times one plus the draft length, so with the MTP drafts Studio loads by default a four-chat run with the mode on reaches 0.45 of the mode-off aggregate decode rate (2.1x the wall), against 0.91 with speculation off. The switch stays opt-in and off by default, and the setting's help text says it costs throughput, so nothing here changes; a follow-up could disable drafts automatically while the switch is on, which is a separate decision.

@danielhanchen

Copy link
Copy Markdown
Member Author

Follow-up to the note above: the speculative cost was a llama.cpp batch-splitting issue, not the mode itself, and it is fixed on the llama.cpp side (unslothai/llama.cpp#194 at c6c3cb671). With two MTP drafts, four chats in exact concurrency now decode at 0.74 of the mode-off rate instead of 0.43, byte-identical to the solo run on both the 4B and the 35B, with and without parks. Nothing changes in this PR: the switch is still opt-in and off by default, and a llama-server without that commit behaves as before.

@danielhanchen

Copy link
Copy Markdown
Member Author

Merged the base branch forward (a6c99ce): the 23 review fixes from #10301 at d11bd99, no conflicts. The preemption, admission, stream, tool-loop, Anthropic, exact-switch and handoff suites pass on the merge: 1591 tests.

danielhanchen and others added 2 commits September 6, 2026 23:15
…ps its prompt cache

# Conflicts:
#	studio/backend/core/inference/llama_preemption.py
@chatgpt-codex-connector

Copy link
Copy Markdown

🛡️ Codex Security Review · Automatically triggered

Security review completed. No security issues were found in this pull request.

Reviewed commit: 44344f7a22

View security finding report

Only the user who started this review can view the report in Codex.

ℹ️ About Codex security reviews in GitHub

This is an experimental Codex feature. Security reviews are triggered when:

  • You comment "@codex security review"
  • A regular code review gets triggered (for example, "@codex review" or when a PR is opened), and you’re opted in so security review runs alongside code review

Once complete, Codex will leave suggestions, or a comment if no findings are found.

chatgpt-codex-connector[bot]

This comment was marked as resolved.

…r-swap-c for PR #10358

# Conflicts:
#	studio/backend/tests/data/refactor_guard/ast_inventory.json
#	studio/backend/tests/data/refactor_guard/patch_targets.json
…the translated relays, and read the chip's note off the visible branch for PR #10358
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Chef's kiss.

Reviewed commit: 3775ce46bc

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

…#10358

# Conflicts:
#	studio/backend/core/inference/llama_preemption.py
chatgpt-codex-connector[bot]

This comment was marked as resolved.

…#10358

# Conflicts:
#	studio/backend/core/inference/llama_preemption.py
#	studio/backend/tests/test_llama_preemption_wiring.py
chatgpt-codex-connector[bot]

This comment was marked as resolved.

… with the resume wait, name what exact mode needs in settings and spell the overflow retry's raw flag as pausable for PR #10358
…#10358

# Conflicts:
#	studio/backend/core/inference/llama_cpp.py
#	studio/backend/routes/inference.py
#	studio/backend/tests/test_a_pause_charges_what_the_attempt_decoded.py
#	studio/backend/tests/test_llama_preemption_wiring.py
chatgpt-codex-connector[bot]

This comment was marked as resolved.

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. 👍

Reviewed commit: 3489fee96e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@danielhanchen

Copy link
Copy Markdown
Member Author

@codex review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant