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 12 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. |
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.
Tests
tools/server/tests/unit/test_preempt.pygoes from 6 to 10. New: 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.