server: take the park and restore copies off the decode loop - #192
server: take the park and restore copies off the decode loop#192danielhanchen wants to merge 48 commits into
Conversation
ggml_backend_event_synchronize() is the only way to find out whether the work recorded before an event has finished, and it answers by waiting for it. A caller that issued an asynchronous copy so that it could get on with something else has no way to ask "is it done yet" without giving that up again. ggml_backend_event_query() is that question. It is optional, and it is the last field of ggml_backend_device_i so that a backend which does not implement it needs no change: a missing entry is NULL and the generic implementation falls back to a blocking synchronize and returns true, which is correct, just no better than what a caller could do already. CUDA implements it with cudaEventQuery, treating cudaErrorNotReady as the answer "not yet" rather than as a failure, and clearing it so it is not reported against the next call. The other sixteen device interfaces get an explicit NULL. Trailing initializers could have been left off, since these are positional aggregate initializers and the new member would be value-initialized, but -Wmissing-field-initializers is part of -Wextra and becomes an error under LLAMA_FATAL_WARNINGS.
Two changes to how a sequence's state is copied out of and back into the cache, the first of which the second one needs. Coalescing. The save side works out which cells belong to the sequence, merges them into ranges and emits one write per range per tensor. The restore side does not: it emits one read per cell, thousands of them, even when the cells it was given are a handful of long runs. Merging fragments that are adjacent in both the tensor and the buffer fixes both sides at once, and covers the transposed V layout where the same runs are emitted once per embedding row. Sequences sharing a unified cache take their cells in turn, so what is left after merging is a regular comb rather than one block; a comb is what a strided copy describes, so runs of one length at a constant stride become a single 2d transfer. Measured on a 4B at -c 8192 with four chats, a 1989-cell sequence goes from 1989 transfers per tensor to about 160, and a sequence that has the cache to itself to one. Asynchronous transfers. llama_state_seq_copy is a transfer that can be issued and left running: it owns the host buffer, a backend per device holding part of the cache so the copies get a stream of their own rather than queueing behind the graphs, and an event per device to say when its half is done. The buffer is pinned where the backend offers pinned memory, which is what makes the copies overlap at all, and grow-only, because page-locking a hundred MiB costs about as long as the copy it is for and a caller parking the same sequence repeatedly asks for a slightly different size each time. The restore side of the asynchronous path deliberately does not use the whole-tensor staging the synchronous one does. Staging reads a tensor, patches the sequence's bytes into the host copy and writes the tensor back, which keeps the neighbours only while nothing else is touching the cache. These copies exist so that decoding can carry on beside them, so the write-back would undo whatever the sequences sharing the tensor wrote to their own cells in the meantime. Writing only this sequence's runs cannot, and coalescing is what makes that affordable. llama_state_seq_copy_init() returns NULL when no backend can copy asynchronously, so a caller keeps the synchronous calls on those.
preempt_save() and preempt_restore() run inside update_slots(), so while one sequence is copied out of or back into the KV pool every other slot stops. On a 4B at -c 8192 with four chats that is a 250 ms freeze at a park and 149 ms at a restore, seen by chats that had nothing to do with either, against a p99 inter-token gap of 19 ms when nothing is being parked. The park was cheap for the slot it saved; it was the other three that paid for it. A park now has two halves. preempt_save() issues the copy and leaves the slot PREEMPTING: the cells are still its own, because the copy is still reading them, and nobody may take them. update_slots() polls the event each iteration and only then releases the cells and marks the slot PREEMPTED. A restore is the mirror, RESTORING: the cells are allocated and owned by the sequence, so nobody else can take them, but they do not hold its state until the copy lands, which is why the slot is not scheduled and its drafter not rearmed until it does. An asynchronous park does not hand its cells back before update_slots() carries on, so it has to fire earlier than a synchronous one, or the slots that keep decoding have nowhere to put their tokens and end up waiting for the copy after all. preempt_n_margin() keeps eight decode steps of every running slot clear ahead of the pool filling, which is about the tenth of a second a copy of one sequence takes. The same figure gates a resume, so that a slot is not put back into a pool it would immediately have to be taken out of again. When that lookahead is not enough the loop waits for the outstanding park rather than let the KV-full path end every request, which is no worse than the synchronous path and is the last thing tried before giving up. Everything that reads a slot's state had to learn the two new ones. is_processing() is deliberately left as "not idle", because it is what keeps NEXT_RESPONSE posted and the loop polling; narrowing it would deadlock a server whose only slots are mid-copy. preempt_kv_used() deliberately still counts them, since a slot on its way out has not released its cells and one on its way back in has already been given them. release() waits for an outstanding copy before freeing the buffer and handing the cells on, which is the path a cancelled request and every error path take, and where a transfer would otherwise outlive the memory on both ends. --preempt-async (LLAMA_ARG_PREEMPT_ASYNC) is on by default and falls back to the synchronous path on a backend that cannot copy asynchronously, saying so once at load. --no-preempt-async keeps the old behaviour, so both can be compared on one binary. The pinned buffers are held for as long as the task that parked owns the slot rather than freed between two of its parks, so --preempt-ram now bounds the host memory actually held; it still reads zero once the slots are released. Measured on the same four chats, survivors now see 38 to 43 ms at a park and 19 ms at a restore under LLAMA_SERVER_PREEMPT_EVERY=64, against 127 to 158 ms and 112 to 115 ms before, and four-chat throughput goes from 179-184 to 243-267 tok/s. What is left is issue cost: about 11500 transfers at 4 us each, because four chats interleaving in one pool leave a sequence in roughly 160 runs per tensor. A sequence that has the pool to itself is 66 transfers and 0.26 ms. Tests: the asynchronous path is byte-identical to an uninterrupted run and to the synchronous path, two slots that overflow the pool together finish with the tokens they produce alone, cancelling while a copy is in flight leaves no slot stuck and no parked memory held, and --no-preempt-async really does switch it off.
ggml-cuda.cu is compiled for ROCm and for MUSA through the vendor headers, which rename every cuda* name it uses. The non-blocking event query added cudaEventQuery and cudaErrorNotReady, and neither header maps them, so both builds stop at an undeclared identifier while the adjacent cudaEventSynchronize has been mapped all along. hipEventQuery and musaEventQuery have the same signature and the same convention: success when everything recorded before the event has finished, hipErrorNotReady or musaErrorNotReady while it has not, which is exactly what the query reads them as.
…ously state_seq_copy_init() took any device advertising async and events, but ggml_backend_event_query() is optional: a device that does not implement it gets the generic fallback, which answers "is it done" by waiting for it. Metal, Vulkan and SYCL all advertise both capabilities and all leave event_query null, so they were handed a transfer object, told the caller the copies were asynchronous, and then blocked it for the whole copy on its first poll. That is the stall the transfer exists to remove, made worse by the caller no longer expecting it. ggml_backend_dev_supports_event_query() is the question the fallback hides, and state_seq_copy_init() now asks it. A device without a query is left out, so those backends get NULL and keep the synchronous llama_state_seq_*_data_ext calls they always used, which is the documented behaviour and is what the server already falls back to. The reason is logged once.
…s down destroy() resets llama_init and nulls ctx_tgt and ctx_dft, but the slots are declared after llama_init and are still alive at that point, and one of them can be holding a park or a resume that is still reading or writing KV tensors of the context being freed. release() already makes that wait for a single slot, on the path a cancelled request takes; nothing made it for all of them. The sleeping-state path is where it shows: /sleep calls destroy() and the server carries on running, so a copy issued an iteration earlier is left pointing at freed tensors and load_model() then clears the slots, running the transfer destructor's own wait against the same memory. Shutdown has the same hole with less time to notice it. destroy() now waits for every slot's outstanding copy and lets go of the transfers before anything is freed, which also means the next context does not inherit a backend and a host buffer belonging to the previous one.
preempt_n_margin() keeps eight decode steps of every running slot clear ahead of the pool filling, and the resume gate uses the same figure so that a slot is not put back into a pool it would immediately have to leave. It was not doing that for the slot being resumed. The candidate is still PREEMPTED while it is being considered, so the loop that counts running slots skips it, and the runway it needs appears only after it has been let in, at which point the pool is short by exactly that much and somebody gets parked. At -c 256 with a 1 + n_spec step and the eight-step runway, totals from 233 to 240 cells admit a restore that then cannot take its first step, and under load the same slot was seen restored and parked again five times over. preempt_n_margin() takes the number of slots that are about to be running as well as those that already are, and the resume gate passes one for the candidate. Everything else keeps the count it had. preempt_kv_reserve() already reserves a restoring slot's next step; this is the eight-step runway behind it.
… for llama_state_seq_copy_buf_is_pinned() returned can_pin, which is worked out from the buffer type the backend offers and is fixed for the life of the transfer. The header promises the buffer is page-locked. Those are different questions: the CUDA host buffer type is handed out whether or not pinning is available, and under GGML_CUDA_NO_PINNED its allocation falls back to an ordinary CPU buffer, so the server logged "pinned host memory" while every park ran through pageable memory. It was also true before any buffer existed and after buf_free(). buf_resize() already records which it got, by comparing the buffer that came back against the type that was asked for, so is_pinned() now returns that. llama_state_seq_copy_buf_can_pin() is the capability question, for a caller that wants to know before allocating anything. The load banner asked the capability question at a point where no buffer exists and printed the answer as though one did. It now says what the backend offers, in those words, and the first park reports what the buffer it allocated actually turned out to be.
Both issue functions validated only that a buffer existed. The caller's size was handed straight to the io object, which then validated every fragment against that number rather than against the allocation, so a save issued with a size larger than the buffer wrote past the end of it and a restore read whatever was next on the heap and sent it to the device. Unlike the legacy API the library owns this buffer, so it can simply check: a size of zero, or one beyond llama_state_seq_copy_buf_size(), is refused with a log line. The flags word had the same problem from the other end. LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to stay in device buffers, and both functions built the host serializers regardless, while llama_state_seq_get_size_ext() with that flag reports a state without the tensor bytes in it. A caller pairing the documented size call with these ones sized a buffer for the metadata and then tried to fill it with the whole sequence. The flag is refused here and the restriction is written down in llama.h; the synchronous calls still serve it. tests/test-state-seq-copy.cpp covers both refusals in both directions, checks that a refused call posts nothing, that the same call at the buffer's own size still round-trips the sequence byte-for-byte, and that a transfer reports itself as pinned only while it holds memory that is. It skips itself where no backend can copy asynchronously.
The cudaErrorNotReady branch of the CUDA event query called cudaGetLastError() on the belief that the result had to be cleared. It does not: cudaEventQuery() returns cudaErrorNotReady as its return value without recording it in the thread's last-error state, so the only thing that call can collect is an error somebody else planted and has not looked at yet. Checked on a B200 with CUDA 13.1. An unrelated cudaSetDevice(99) leaves 101 pending; cudaEventQuery() on an outstanding event returns 600 and cudaPeekAtLastError() still reads 101 afterwards, so the cudaGetLastError() returned 101 and left the state clean. A real launch failure would have been thrown away the same way, and its owner would never have seen it.
ggml_backend_device_i gained event_query, so a device interface built against the previous header is one member shorter than the one ggml now reads. Every in-tree initializer was updated, but a backend loaded from a shared library is not: ggml_backend_reg_load_backend() accepts it on api_version alone, and a prebuilt .so still reporting 2 would have been let in and its iface.event_query read past the end of the object. Rejecting it is what the version is for.
… needs A review of #192 pointed at the victim loop in update_preemption(). When the pool has no room for the step about to be built and no park is in flight, the loop issues the victim's asynchronous park and breaks. The cells are held until the copy lands, so the batch is built into a pool that has not got smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in microseconds without ever polling the copy, ending in "Context size has been exceeded" for every slot. The synchronous path freed the cells before returning, so it could not do this. I could not reproduce it. Six live rounds on the 4B at -c 8192, four chats with 1000-token prompts and 2048 tokens each, with the fourth chat's prompt held back 20 s so it arrives into a pool the other three have filled, exact mode on and off, on a binary without this change: 4 of 4 every round, no context errors, and the retry ladder was not entered once ("failed to find free space" appears zero times in both server logs). The reason is that preempt_kv_reserve() counts an incoming prompt chunk before it is allocated, so the planner crosses the lookahead threshold an iteration before the pool actually fills, and every park in those runs was issued with the 80 cells of asynchronous runway still ahead of it, never at the hard threshold this is about. Committing it anyway, because the described state is real even if these workloads do not reach it, and the change is inert unless it is reached: * The victim loop goes round again instead of leaving, but only when n_used + PREEMPT_N_MARGIN > n_cells, that is when there is no room for the step itself rather than merely less than the asynchronous lookahead wants. The next pass reaches preempt_wait_in_flight() and waits for the park just issued, which is what that function was written for and no worse than the synchronous path. Short of the lookahead only, it still breaks, because parking early and letting the copy run beside the decode is the entire point of #192. * On llama_decode returning 1, an outstanding park is waited for before any batch width is given up. Halving n_batch returns no cells, so without this the ladder can walk to n_batch == 1 and end every request while the room it needed was one event query away. Safe at that point because the slot was detached before the batch was built, so completing its park cannot change what is about to be retried; that is also why update_preemption() itself is not called from here. New test, test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_ everything: three slots generating near the ceiling and a fourth request whose prompt does not fit in what is left, which is the shape the existing tests miss because their victim holds almost no cells. The three are sized to oversubscribe the pool between them so the pressure does not depend on when the fourth arrives. It is kept for the shape it covers rather than as an attribution: it passes either way, and the attribution above was done at live scale.
|
Nine more commits. Eight are the review items, and the ninth,
I could not reproduce the failure. Six live rounds on the 4B at Of the eight review commits the one with a real measured effect is The rest. Evidence on the new head. |
|
Merged #184's 2a7e277 as 4e49a05: a parked sequence that cannot fit the pool alone is failed with the context error instead of parked for ever, and the scheduling margin is waived when nothing is resident. Here the margin is the async lookahead one, One thing recorded, not changed: the lookahead margin counts running slots and excludes slots in RESTORING while their cells are already counted in |
|
Merged #184's 84311fd and 6dbc4e7 as 01dfecc: when the KV-full retry ladder runs out, resident slots are rewound to the cache boundary and the smallest parked instead of every slot getting the context error, and a batch holding a draft is never narrowed through the group. Here the parks issued as a last resort are waited for before the batch is given up, since the cells are wanted now rather than next pass, and slots in transfer are left out of the rewind and the error sweep. Suite 14 passed. Exact-mode accuracy with the planner off (last resort only), 4B and 35B, with and without MTP drafts: identical to solo in every round, 0 context errors. |
A slot on its way back in already holds its cells and starts decoding the moment its copy lands, so it needs the same runway as the slots already running. Leaving it out let two back-to-back restores land into a pool that then had to park someone again at once.
|
cd1cd4e: the lookahead margin now counts a slot being restored. It already holds its cells and starts decoding the moment its copy lands, so it needs the same runway as the slots already running; leaving it out let two back-to-back restores land into a pool that then had to park someone again at once (the round 2 item recorded as a rare cost). Suite 14 passed; merged into #197. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…code, so restores in one pass do not wait for each other order_after() recorded its fence on the compute stream at the time of the copy, which put it behind the waits order_before() had queued for the restores issued earlier in the same pass: restore B's copies then waited for restore A's to land, and restores issued together ran one after the other on the device, though the host no longer blocked. The fence is now one event per device owned by the context, recorded on the compute stream at the end of every decode and encode once a transfer exists (and once when the first transfer is created), and every park or restore waits for that point instead of recording its own.
…s from the prefix it keeps
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a56d49e1f1
ℹ️ 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".
| llama_io_get(cpy.backend_for(tensor), tensor, ptr, offset, size, | ||
| n_copies, stride_tensor, stride_data); |
There was a problem hiding this comment.
When the KV cache spans multiple CUDA-family devices, order_after() visits all copy backends and can leave a different device current before this callback posts a tensor transfer. The CUDA async set/get callbacks use cuda_ctx->stream() without selecting cuda_ctx->device, and once that stream already exists, stream() does not select it either; submitting the copy with a stream from the non-current CUDA context can therefore abort with an invalid resource handle when the first asynchronous park runs. Explicitly select the copy backend's device before each transfer (including both scalar and 2-D callbacks).
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not a defect. The CUDA programming guide's stream and event rules (Programming Systems with Multiple GPUs) say a memory copy succeeds even when issued to a stream not associated with the current device; only kernel launches require the current device to match, cudaEventRecord needs the event and the stream on the same device (they are, both are created for the copy device), and cudaStreamWaitEvent works across devices. Every transfer here is a memcpy or a 2-D memcpy on the copy backend's own stream, so the current device does not matter to it, and ggml's own CUDA set_tensor_async and event_wait select no device either.
…r the layout check; transfers only where a park can happen The decision to stage a fragmented restore counted the calls the emitter makes, and on a buffer without 2-D copies one strided call expands into one synchronous transfer per row, so a regularly interleaved sequence that had been staged at 64 runs was no longer staged and paid for every row. ggml_backend_buffer_supports_2d() says which kind of buffer it is, and the count is by rows where a row is what a call costs. The per-device fences were installed before the check that refuses a transfer for a state not all in device memory, so a server that then fell back to synchronous copies recorded them after every decode for nobody. They are installed after the checks, and an install that fails is undone. The server made a transfer for every slot, fences included, even where no park can happen: one slot, no memory, a recurrent cache. The transfers and the banner are now gated on the same conditions as the planner.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…ansfer takes the fences with it The wait for copies in flight before a context shift is applied ran just before the target decode, but pre_decode() had already asked the draft context for its draft, a decode that applies that cache's pending shift in place while a park or restore may still be copying draft cells. The wait now follows update_preemption() and precedes pre_decode(), so both caches shift after the copies have landed. When the pinning probe found pageable memory the server gave its transfers up, but the fences a transfer installs stayed with the context and were recorded after every decode for nobody. The context counts its live transfers now and the last one to go frees the fences.
… resident cycling past the rotation delay
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…he shift wait runs before the decode as well synchronize() covers the graph backends, not the copy backend a transfer owns, so a context freed while a transfer was still copying could free the KV buffers under it, and freeing the transfer afterwards touched the dead context. The context keeps the set of its live transfers now: at teardown each is waited for and disowned, and its own free then touches nothing of the context. The wait for copies in flight before a shift is applied runs before pre_decode(), for the draft, and again before the decode: a shift that --cache-reuse asks for is found inside pre_decode(), after the first wait, and the decode applies it in place like any other.
…ntext; the recurrent flag assigned on every load
|
Codex Review: Didn't find any major issues. Keep them coming! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
…at/server-side-preemption-async
|
@codex review |
|
Codex Review: Didn't find any major issues. 🚀 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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". |
1 participant
Stacked on #184; the last three commits are new.
Summary
preempt_save()andpreempt_restore()run insideupdate_slots(), so while one sequence iscopied out of or back into the unified KV pool, every other slot stops. On a 4B at
-c 8192with four chats that is a 250 ms freeze at a park and 149 ms at a restore, paid by chats that
had nothing to do with either, against a p99 inter-token gap of 19 ms when nothing is parked.
This makes both copies asynchronous and takes the waiting off the decode loop. A park becomes
two halves:
preempt_save()issues the copy and leaves the slotPREEMPTING, still owningits cells because the copy is still reading them;
update_slots()polls the event eachiteration and only then releases them. A restore is the mirror,
RESTORING: cells allocatedand owned, but the slot is not scheduled and its drafter not rearmed until the copy lands.
Three supporting pieces.
ggml_backend_event_queryis new -- the existing event API could onlyanswer "is it done" by waiting for it; it is the last field of
ggml_backend_device_isobackends that do not implement it need no change and get a blocking fallback. The park buffers
move from pageable
std::vectorto the backend's pinned host buffer type, grow-only, becausepage-locking 120 MiB costs about as long as the copy it is for. And the state transfers now
coalesce adjacent cells into runs and regular runs into strided 2d copies, which fixes a
long-standing asymmetry: the save side already merged its cells into ranges, the restore side
emitted one transfer per cell.
Policy
The park fires earlier than a synchronous one would, because it does not return its cells
before
update_slots()carries on:preempt_n_margin()keeps 8 decode steps of every runningslot clear ahead of the pool filling, roughly the tenth of a second a sequence copy takes. The
same margin gates a resume, so a slot is not put back into a pool it would immediately have to
leave. When the lookahead is not enough, the loop waits for the outstanding park rather than
letting the KV-full path end every request; that is the last thing tried before giving up, and
it is no worse than the synchronous path.
--preempt-async/--no-preempt-async(LLAMA_ARG_PREEMPT_ASYNC), on by default, fallingback to the synchronous path on a backend that cannot copy asynchronously and saying so once
at load, so both can be compared on one binary.
Results
Four chats, 1000-token prompts, 2048 tokens each with
ignore_eos,LLAMA_SERVER_PREEMPT_EVERY=64so both sides do identical work (124 parks). Three runs each,back to back.
Every "after" figure is outside the range of all three "before" figures. With natural
(unforced) parks the survivor stall at a park goes from 250 / 273 / 233 ms to 72 / 52 / 36 ms.
--no-preempt-asyncon the same binary lands between the two, keeping the synchronous pathbut still getting the coalescing.
Exactness
One prompt, seed 0, temperature 0, 1000 tokens, streamed: forced parks through transfers,
unforced, and forced with
--no-preempt-asyncall give the identical sha256126193ef...4a45d3c8, which is also the value recorded for this prompt on an earlier build.test-state-restore-fragmentedpasses on CUDA with the 4B, all three sequence snapshotsbyte-identical after a fragmented restore.
Cost when it does not fire
None on the decode path.
llama_state_seq_copy_init()returns NULL where the backend cannotcopy asynchronously and every park and resume is the synchronous one it was.
--preempt-ram 0and a non-unified cache return before any of this is reached. The pinned buffers are allocated
on a slot's first park, not at startup, and released when the task gives up the slot;
llamacpp:preempt_ram_bytesnow reports the memory actually held and still reads zero whennothing is parked. The coalescing and strided-copy changes apply to the synchronous path too
and only reduce the number of transfers.
A round that recorded a context shift waits for every copy in flight before it decodes. The shift is applied inside the next
llama_decodeas one graph over the whole K cache, in place; a restore still writing its cells on another stream would be read half done and written back stale, and a park still reading would read through the rewrite. Shifts are rare, so the wait is too.No rotation runs while a park is still copying, and an asynchronous rotation park is re-examined on the pass that sees its copy land rather than on the same pass, so one rotation cannot park every shifting resident. The rotation's budget counts nothing of an asynchronous head as leaving (its pinned buffer is kept through the restore by design) and charges the resident only what it does not hold yet; a restored slot's buffer is returned when the pool is over its budget. The
--cache-reuseshift is applied by the same in-place graph as a context shift and makes the round wait for copies in flight in the same way.Tests
tools/server/tests/unit/test_preempt.py: 16 pass on this branch, the eleven of #184 plus five here. New here: forced async parks givebyte-identical output; two slots that overflow the pool together finish with the tokens they
produce alone; cancelling while a copy is in flight leaves no slot stuck and no parked memory
held;
--no-preempt-asyncreally switches it off. The async ones skip on a backend withoutasynchronous copies.
test-state-restore-fragmentedcovers the coalescing and stridedgrouping preserving neighbouring sequences.
Limitations
What is left of the stall is issue cost: ~11500 transfers at ~4 us each, because four chats
interleaving in one pool leave a sequence in ~160 runs per tensor. A sequence with the pool to
itself is 66 transfers and 0.26 ms. Cutting that further needs a batched scatter/gather in the
backend or a less interleaved allocator; whole-tensor staging is not available, because
read-modify-write is exactly what is unsafe while decoding continues. The first park of a task
pays ~45 ms to page-lock its buffer. The lookahead margin is a constant tuned to these decode
rates.
ggml_backend_event_queryis implemented for CUDA only; the other backends compileagainst the new field and take the blocking fallback. Multi-GPU KV and
v_trans = true(without
--flash-attn) are written for but not exercised here.Follow-up commits after review
Nine more commits, one per item. The event query is mapped for HIP and MUSA, which did not compile without it. A sequence is copied asynchronously only when the device implements a real non-blocking event query, so Metal, Vulkan and SYCL take the documented synchronous path instead of blocking the decode loop on every poll. Teardown drains in-flight copies before the contexts are freed. The resume gate charges the candidate its own lookahead: four chats at a very tight
-c 4096went from 38 and 31 parks to 8 and 8, with the worst slot going from 14 and 16 preemptions to 4, at 26 to 27 percent lower throughput at that size and no difference at-c 8192. The pinned-memory report and the startup banner say what was allocated rather than what could be. The two issue functions refuse a size beyond the owned buffer and the on-device flag, with a new unit test. The error-collection call in the CUDA event query is deleted, since an experiment showed it could not clear the not-ready state and did consume an unrelated error. The backend API version is bumped for the new interface member. Finally, the victim loop no longer leaves an issued asynchronous park holding the room the current step needs, and a decode that fails for lack of room waits for an outstanding park before halving the batch; that path did not reproduce in six staggered live rounds, and the fix is inert unless it is reached.Harness 11 of 11, forced parks byte-identical to unforced with the same sha256 as the first build, survivor stalls unchanged by the fixes.
Three more after a second pass. The asynchronous adapters posted their queued copies from the destructor whether or not serialisation had got to the end, so a buffer one byte short returned 0 while 64 copies were still reading it; both adapters now commit only after the serialisation succeeds, and the unit test asserts a short transfer posts nothing on either side. A host buffer type may hand back pageable memory instead of failing, and a copy into or out of it blocks the issuing thread, so the server takes a one MiB buffer at load, looks at what it got, and parks synchronously with a warning if it is pageable.
llama_state_seq_copy_initreturned a transfer whenever a device could copy asynchronously, even with most layers on the CPU where every state tensor is host memory and every copy took the synchronous branch; it returns NULL unless every non-empty memory buffer lives on one of its devices, and the server's banner follows.Two more from the next pass. A park buffer many times larger than the one MiB startup probe can still come back pageable (a host-locking limit), so the park now looks at the buffer it actually got, and a pageable one is given back with the slot's transfers so that slot parks synchronously from then on. And the copies of a restore and the graphs that follow it shared no ordering: the restore writes cells on the copy stream while other sequences decode, and an attention that is not paged reads every cell up to n_kv, masked ones included. After the copies are posted and their events recorded, every compute backend on that device now waits for the event before its next graph, a stream wait rather than a host wait, so the thread carries on and the next decode starts the moment the copy lands. A park needs nothing of the kind: it reads cells nobody writes until it has landed.
Two more. Both copy functions drained the host to order their copies behind the compute already queued, and with the compute stream waiting on the previous restore that drain blocked the decode loop until the previous copy had landed, so two restores issued in one pass ran one after the other. Each copy device now carries a second event, recorded on the compute stream and waited for on the copy stream: the ordering is on the device and the host drains nothing. And a restored slot keeps its pinned buffer for its next park while that capacity is charged against
--preempt-ram, so a budget that held one sequence was spent for good by the first restore and every later park was refused. Both budget checks give idle buffers back before deciding, largest first, never one that still holds a parked sequence, has a copy in flight, is the candidate's own or is the head of a rotation. New test: two sequences parked in turn under a budget that holds one, both preempted, no context error. Harness 19 of 19.Carried from #184 since: a pure recurrent cache is served without preemption, and the rotation asks the plain budget check (both states are held at once while the resident is parked and the head restored), so a budget that holds one sequence but not two does not rotate and says so once per park; here the idle-buffer return runs first, the parked head is never among the buffers given back. Harness 21 of 21.
One more from the next pass. Recorded at the time of the copy, the fence landed on the compute stream behind the waits queued for the restores issued earlier in the same pass, so restore B's copies waited for restore A's to land and restores issued together ran one after the other on the device. The fence is now one event per device owned by the context, recorded on the compute stream at the end of every decode and encode once a transfer exists, and every park or restore waits for that point instead of recording its own. Also carried from #184: a started slot's reservation counts from the prefix it keeps, as its used count already did.
Three more from the pass after that. The decision to stage a fragmented restore counted the emitter's calls, and on a buffer without 2-D copies one strided call is one synchronous transfer per row, so a regularly interleaved sequence that had been staged at 64 runs was no longer staged;
ggml_backend_buffer_supports_2d()says which kind of buffer it is and the count is by rows on one that copies row by row. The per-device fences were installed before the layout check could refuse the transfer, so a server that fell back to synchronous copies recorded them after every decode for nobody; they follow the checks now and a failed install is undone. And a transfer was made for every slot even where no park can happen (one slot, no memory, a recurrent cache); the transfers and the banner are gated on the planner's own conditions.Two more. The wait for copies in flight before a context shift is applied ran just before the target decode, after pre_decode() had already asked the draft context for its draft, a decode that applies that cache's pending shift in place; the wait now follows update_preemption() and precedes pre_decode(), so both caches shift after the copies have landed. And giving the transfers up (the pageable fallback) left the fences with the context, recorded after every decode for nobody; the context counts its live transfers and the last one to go frees the fences.
And two from the pass after: the wait for copies in flight before a shift is applied runs before pre_decode() for the draft and again before the decode, since a shift that --cache-reuse asks for is found inside pre_decode(); and a context freed with a transfer still alive drains it and disowns it first, synchronize() covering the graph backends and not the copy backend a transfer owns, so the transfer's own free then touches nothing of the context.
From #184: a started slot is counted by the cells it holds, and trimmed to the prefix its request keeps when nothing fits, so a resume is never attempted against cells the batch builder has not yet released.