From 5c6d79e1a8587c168eafb3c8d40cdc87eceb7fad Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:34:40 +0000 Subject: [PATCH 01/25] ggml: add a non-blocking query for backend events 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. --- ggml/include/ggml-backend.h | 3 +++ ggml/src/ggml-backend-impl.h | 5 +++++ ggml/src/ggml-backend-meta.cpp | 1 + ggml/src/ggml-backend.cpp | 12 ++++++++++++ ggml/src/ggml-blas/ggml-blas.cpp | 1 + ggml/src/ggml-cann/ggml-cann.cpp | 1 + ggml/src/ggml-cpu/ggml-cpu.cpp | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 17 +++++++++++++++++ ggml/src/ggml-et/ggml-et.cpp | 1 + ggml/src/ggml-hexagon/ggml-hexagon.cpp | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-opencl/ggml-opencl.cpp | 1 + ggml/src/ggml-openvino/ggml-openvino.cpp | 1 + ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + ggml/src/ggml-sycl/ggml-sycl.cpp | 1 + ggml/src/ggml-virtgpu/ggml-backend-device.cpp | 1 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1 + ggml/src/ggml-webgpu/ggml-webgpu.cpp | 1 + ggml/src/ggml-zdnn/ggml-zdnn.cpp | 1 + ggml/src/ggml-zendnn/ggml-zendnn.cpp | 1 + 20 files changed, 53 insertions(+) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e3..30d8304d492 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -125,6 +125,9 @@ extern "C" { GGML_API void ggml_backend_event_free(ggml_backend_event_t event); GGML_API void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend); GGML_API void ggml_backend_event_synchronize(ggml_backend_event_t event); + // non-blocking: true once everything recorded before the event has completed. + // backends without a query implementation fall back to a blocking synchronize and return true. + GGML_API bool ggml_backend_event_query(ggml_backend_event_t event); GGML_API void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event); // diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3d..902b0963afa 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -200,6 +200,11 @@ extern "C" { ggml_backend_event_t (*event_new) (ggml_backend_dev_t dev); void (*event_free) (ggml_backend_dev_t dev, ggml_backend_event_t event); void (*event_synchronize) (ggml_backend_dev_t dev, ggml_backend_event_t event); + + // (optional) non-blocking completion test for an event. + // kept last so that backends that do not implement it need no change: a missing entry + // is NULL, and ggml_backend_event_query() then falls back to a blocking synchronize. + bool (*event_query) (ggml_backend_dev_t dev, ggml_backend_event_t event); }; struct ggml_backend_device { diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 3ec40fb1af7..d531ae4b5fa 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -193,6 +193,7 @@ static const ggml_backend_device_i ggml_backend_meta_device_iface = { /* .event_new = */ nullptr, /* .event_free = */ nullptr, /* .event_synchronize = */ nullptr, + /* .event_query = */ NULL, }; static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e519bdf50a1..a56ab30862a 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -551,6 +551,18 @@ void ggml_backend_event_synchronize(ggml_backend_event_t event) { event->device->iface.event_synchronize(event->device, event); } +bool ggml_backend_event_query(ggml_backend_event_t event) { + GGML_ASSERT(event); + + if (event->device->iface.event_query == NULL) { + // no way to ask: the honest answer is to wait for it and then say yes + ggml_backend_event_synchronize(event); + return true; + } + + return event->device->iface.event_query(event->device, event); +} + void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { GGML_ASSERT(backend); GGML_ASSERT(backend->iface.event_wait != NULL); diff --git a/ggml/src/ggml-blas/ggml-blas.cpp b/ggml/src/ggml-blas/ggml-blas.cpp index e4b5bd25474..7271b6b632b 100644 --- a/ggml/src/ggml-blas/ggml-blas.cpp +++ b/ggml/src/ggml-blas/ggml-blas.cpp @@ -469,6 +469,7 @@ static const struct ggml_backend_device_i ggml_backend_blas_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // backend reg interface diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5e5541aac94..f0bc9205bfa 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2948,6 +2948,7 @@ static const ggml_backend_device_i ggml_backend_cann_device_interface = { /* .event_new = */ ggml_backend_cann_device_event_new, /* .event_free = */ ggml_backend_cann_device_event_free, /* .event_synchronize = */ ggml_backend_cann_device_event_synchronize, + /* .event_query = */ NULL, }; // backend reg diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 8cece71f186..a11d707aa14 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -500,6 +500,7 @@ static const struct ggml_backend_device_i ggml_backend_cpu_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // CPU backend - backend (reg) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc62..91902099011 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5375,6 +5375,22 @@ static void ggml_backend_cuda_device_event_synchronize(ggml_backend_dev_t dev, g CUDA_CHECK(cudaEventSynchronize((cudaEvent_t)event->context)); } +static bool ggml_backend_cuda_device_event_query(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + const cudaError_t err = cudaEventQuery((cudaEvent_t)event->context); + + if (err == cudaErrorNotReady) { + // not an error: clear it so it is not reported against the next call + (void) cudaGetLastError(); + return false; + } + + CUDA_CHECK(err); + + return true; +} + static const ggml_backend_device_i ggml_backend_cuda_device_interface = { /* .get_name = */ ggml_backend_cuda_device_get_name, /* .get_description = */ ggml_backend_cuda_device_get_description, @@ -5391,6 +5407,7 @@ static const ggml_backend_device_i ggml_backend_cuda_device_interface = { /* .event_new = */ ggml_backend_cuda_device_event_new, /* .event_free = */ ggml_backend_cuda_device_event_free, /* .event_synchronize = */ ggml_backend_cuda_device_event_synchronize, + /* .event_query = */ ggml_backend_cuda_device_event_query, }; // backend reg diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b87b189a57a..ace1cdedad1 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1684,6 +1684,7 @@ static const struct ggml_backend_device_i ggml_backend_et_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; /* diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e8a5009b381..497ae043a4e 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4274,6 +4274,7 @@ static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; //** backend registry diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9756d47050c..8962e2f7f00 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -818,6 +818,7 @@ static ggml_backend_device_i ggml_backend_metal_device_i = { /* .event_new = */ ggml_backend_metal_device_event_new, /* .event_free = */ ggml_backend_metal_device_event_free, /* .event_synchronize = */ ggml_backend_metal_device_event_synchronize, + /* .event_query = */ NULL, }; // backend registry diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 64f3325b2a5..fc8dbbac80a 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -11332,6 +11332,7 @@ struct ggml_backend_device_i ggml_backend_opencl_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; } diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index e299e16c778..995f1328442 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1454,6 +1454,7 @@ static const struct ggml_backend_device_i ggml_backend_openvino_device_interface /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; struct ggml_backend_openvino_reg_context { diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..4a814d3f534 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1945,6 +1945,7 @@ static const struct ggml_backend_device_i ggml_backend_rpc_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // backend reg interface diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834..6c0cc073117 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6448,6 +6448,7 @@ static const ggml_backend_device_i ggml_backend_sycl_device_interface = { /* .event_new = */ ggml_backend_sycl_device_event_new, /* .event_free = */ ggml_backend_sycl_device_event_free, /* .event_synchronize = */ ggml_backend_sycl_device_event_synchronize, + /* .event_query = */ NULL, }; // backend reg diff --git a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp index 987ce9dd110..13a70c594df 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp @@ -157,4 +157,5 @@ const ggml_backend_device_i ggml_backend_remoting_device_interface = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1d86aaac5c..b94f520185a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18798,6 +18798,7 @@ static const struct ggml_backend_device_i ggml_backend_vk_device_i = { /* .event_new = */ ggml_backend_vk_device_event_new, /* .event_free = */ ggml_backend_vk_device_event_free, /* .event_synchronize = */ ggml_backend_vk_device_event_synchronize, + /* .event_query = */ NULL, }; static const char * ggml_backend_vk_reg_get_name(ggml_backend_reg_t reg) { diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 2434848a55a..4de48d7fa69 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4661,6 +4661,7 @@ static struct ggml_backend_device_i ggml_backend_webgpu_device_i = { /* .event_new = */ ggml_backend_webgpu_device_event_new, /* .event_free = */ ggml_backend_webgpu_device_event_free, /* .event_synchronize = */ ggml_backend_webgpu_device_event_synchronize, + /* .event_query = */ NULL, }; /* End GGML Backend Device Interface */ diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp index 4007ac9dfc7..bbd74fb9d5a 100644 --- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp +++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp @@ -547,6 +547,7 @@ static ggml_backend_device_i ggml_backend_zdnn_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp index ec7ce233145..89c6c36a0f1 100644 --- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp +++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp @@ -781,6 +781,7 @@ static const struct ggml_backend_device_i ggml_backend_zendnn_device_i = { /* .event_new = */ NULL, /* .event_free = */ NULL, /* .event_synchronize = */ NULL, + /* .event_query = */ NULL, }; // backend reg interface From b0487839371a7e395596d40c9dbc9c2a1eea19d5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:34:40 +0000 Subject: [PATCH 02/25] llama: coalesce sequence state transfers, and issue them asynchronously 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. --- include/llama.h | 58 ++++ src/llama-context.cpp | 644 +++++++++++++++++++++++++++++++++++++++++- src/llama-context.h | 9 + 3 files changed, 702 insertions(+), 9 deletions(-) diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..2c61b3d2a71 100644 --- a/include/llama.h +++ b/include/llama.h @@ -927,6 +927,64 @@ extern "C" { llama_seq_id dest_seq_id, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] asynchronous per-sequence state transfer + // + // llama_state_seq_get_data_ext / set_data_ext do not return until every byte has moved, + // so a caller that copies a sequence out of the cache to make room stops doing anything + // else for as long as the copy takes. A transfer object issues the same copies on a + // stream of its own and hands back control immediately; the caller polls + // llama_state_seq_copy_done() and gets on with its other work in between. + // + // The transfer owns the host buffer it reads from or writes into. That buffer is pinned + // when the backend offers pinned memory, which is what makes the copy fast, and it + // cannot be freed while a copy is still using it. + // + // Between issuing and completion the caller must not touch the buffer, must not free or + // reuse the cells of a sequence being read, and must not decode a sequence being + // written. llama_state_seq_copy_free() waits for an outstanding copy first. + struct llama_state_seq_copy; + + // NULL if the context's backends cannot copy asynchronously; the caller then uses the + // synchronous llama_state_seq_*_data_ext calls + LLAMA_API struct llama_state_seq_copy * llama_state_seq_copy_init(struct llama_context * ctx); + LLAMA_API void llama_state_seq_copy_free(struct llama_state_seq_copy * cpy); + + // Size the transfer's host buffer, keeping no contents; NULL on failure. Grow-only: + // page-locking host memory is far too slow to do once per transfer, so the memory is + // kept between them and only given back by llama_state_seq_copy_buf_free(). + LLAMA_API uint8_t * llama_state_seq_copy_buf_resize (struct llama_state_seq_copy * cpy, size_t size); + LLAMA_API uint8_t * llama_state_seq_copy_buf (struct llama_state_seq_copy * cpy); + LLAMA_API size_t llama_state_seq_copy_buf_size (struct llama_state_seq_copy * cpy); + // host memory actually held, which is what a caller budgeting host RAM has to count + LLAMA_API size_t llama_state_seq_copy_buf_capacity(struct llama_state_seq_copy * cpy); + LLAMA_API void llama_state_seq_copy_buf_free (struct llama_state_seq_copy * cpy); + + // true when the buffer is page-locked, i.e. when the copies can really overlap + LLAMA_API bool llama_state_seq_copy_buf_is_pinned(struct llama_state_seq_copy * cpy); + + // issue the copies; return the number of bytes covered, 0 on failure + LLAMA_API size_t llama_state_seq_copy_get( + struct llama_state_seq_copy * cpy, + size_t size, + llama_seq_id seq_id, + llama_state_seq_flags flags); + + LLAMA_API size_t llama_state_seq_copy_set( + struct llama_state_seq_copy * cpy, + size_t size, + llama_seq_id dest_seq_id, + llama_state_seq_flags flags); + + // transfers the last issue posted: one per run of adjacent cells, per tensor + LLAMA_API size_t llama_state_seq_copy_n_copies(struct llama_state_seq_copy * cpy); + + // microseconds the last issue spent waiting for the compute streams before it could start + LLAMA_API int64_t llama_state_seq_copy_sync_us(struct llama_state_seq_copy * cpy); + + // non-blocking completion test, and the blocking wait behind it + LLAMA_API bool llama_state_seq_copy_done(struct llama_state_seq_copy * cpy); + LLAMA_API void llama_state_seq_copy_wait(struct llama_state_seq_copy * cpy); + // // Decoding // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 66940d4fc61..0747d1c6366 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2558,16 +2558,145 @@ class llama_io_write_dummy : public llama_io_write_i { size_t size_written = 0; }; +// [TAG_STATE_COALESCE] one transfer per run of cells, not one per cell +// +// A sequence's state is emitted in cell order, so a run of cells that is contiguous in the +// cache is contiguous both in the tensor and in the host buffer, and the fragments covering +// it are one transfer. The save side already coalesces its cells into ranges before it emits +// them; the restore side does not, and asks for one transfer per cell even when the cells it +// was given are a handful of long runs. Merging here fixes both sides at once, and covers +// the transposed V layout, where the same runs are emitted once per embedding row. +template +static size_t llama_io_run_end(const std::vector & infos, size_t i) { + size_t end = i + 1; + + while (end < infos.size() && + infos[end].tensor == infos[end - 1].tensor && + infos[end].offset == infos[end - 1].offset + infos[end - 1].size && + infos[end].ptr == infos[end - 1].ptr + infos[end - 1].size) { + end++; + } + + return end; +} + +template +static size_t llama_io_run_size(const std::vector & infos, size_t i, size_t end) { + size_t size = 0; + + for (size_t j = i; j < end; ++j) { + size += infos[j].size; + } + + return size; +} + +// [TAG_STATE_COALESCE] runs of one length at a constant stride are a single strided copy +// +// Sequences sharing a unified cache take their cells in turn, so a sequence's cells are not +// one block but a regular comb: a few cells, a gap, a few cells, for as long as the sequence +// is. Merging adjacent cells still leaves hundreds of runs per tensor, and at a few +// microseconds to post each one that is tens of milliseconds spent issuing copies. A comb is +// exactly what a strided copy describes, so one call replaces a whole group of runs. +// +// emit(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); n_copies == 1 means +// an ordinary contiguous transfer and the strides are not meaningful. +template +static void llama_io_emit(const std::vector & infos, size_t first, size_t last, emit_t emit) { + // the runs of adjacent cells, as index ranges into infos + std::vector> runs; + + for (size_t i = first; i < last; ) { + const size_t end = llama_io_run_end(infos, i); + + runs.emplace_back(i, end); + + i = end; + } + + for (size_t r = 0; r < runs.size(); ) { + const auto & head = infos[runs[r].first]; + + const size_t size = llama_io_run_size(infos, runs[r].first, runs[r].second); + + size_t n_copies = 1; + size_t stride_tensor = 0; + size_t stride_data = 0; + + if (r + 1 < runs.size()) { + const auto & next = infos[runs[r + 1].first]; + + if (next.tensor == head.tensor && next.offset > head.offset && next.ptr > head.ptr && + llama_io_run_size(infos, runs[r + 1].first, runs[r + 1].second) == size) { + stride_tensor = next.offset - head.offset; + stride_data = (size_t) (next.ptr - head.ptr); + + // a strided copy may not have its rows overlap, on either side + if (stride_tensor >= size && stride_data >= size) { + while (r + n_copies < runs.size()) { + const auto & cur = infos[runs[r + n_copies].first]; + + if (cur.tensor != head.tensor || + cur.offset != head.offset + n_copies * stride_tensor || + cur.ptr != head.ptr + n_copies * stride_data || + llama_io_run_size(infos, runs[r + n_copies].first, runs[r + n_copies].second) != size) { + break; + } + + n_copies++; + } + } + } + } + + emit(head.tensor, head.ptr, head.offset, size, n_copies, stride_tensor, stride_data); + + r += n_copies; + } +} + +// a null backend means the caller wants the copy to have happened by the time this returns +static void llama_io_get(ggml_backend_t backend, ggml_tensor * tensor, void * ptr, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + if (n_copies > 1) { + if (backend) { + ggml_backend_tensor_get_2d_async(backend, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } else { + ggml_backend_tensor_get_2d(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } + } else if (backend) { + ggml_backend_tensor_get_async(backend, tensor, ptr, offset, size); + } else { + ggml_backend_tensor_get(tensor, ptr, offset, size); + } +} + +static void llama_io_set(ggml_backend_t backend, ggml_tensor * tensor, const void * ptr, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + if (n_copies > 1) { + if (backend) { + ggml_backend_tensor_set_2d_async(backend, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } else { + ggml_backend_tensor_set_2d(tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + } + } else if (backend) { + ggml_backend_tensor_set_async(backend, tensor, ptr, offset, size); + } else { + ggml_backend_tensor_set(tensor, ptr, offset, size); + } +} + class llama_io_write_host : public llama_io_write_i { public: llama_io_write_host( uint8_t * p, size_t len) : ptr(p), buf_size(len) {} ~llama_io_write_host() { - // TODO: add backend support to batch tensor_get? or some other way to speed this up - for (const auto & winfo : winfos) { - ggml_backend_tensor_get(winfo.tensor, winfo.ptr, winfo.offset, winfo.size); - } + llama_io_emit(winfos, 0, winfos.size(), + [](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_get(nullptr, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + }); } void write(const void * src, size_t size) override { @@ -2623,13 +2752,23 @@ class llama_io_read_host : public llama_io_read_i { while (end < rinfos.size() && rinfos[end].tensor == tensor) { end++; } + // [TAG_STATE_COALESCE] the fragments the restore emits are one per cell; what + // matters is how many runs of adjacent cells they form, because that is how many + // transfers they actually cost. Count the runs first, and only fall back to + // staging the whole tensor when even the runs are too many. + size_t n_runs = 0; + llama_io_emit(rinfos, i, end, + [&n_runs](ggml_tensor *, const uint8_t *, size_t, size_t, size_t, size_t, size_t) { + n_runs++; + }); + const size_t tensor_bytes = ggml_nbytes(tensor); auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; // A fragmented sequence can require thousands of synchronous device // transfers per layer. For bounded tensors, stage the tensor once and // preserve every byte belonging to other sequences. Bound scratch RAM // and leave ordinary contiguous transfers on their original fast path. - if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && + if (n_runs >= 64 && tensor_bytes <= 64 * 1024 * 1024 && !ggml_backend_buffer_is_host(buffer)) { std::vector staging; try { @@ -2649,10 +2788,13 @@ class llama_io_read_host : public llama_io_read_i { continue; } } - for (; i < end; ++i) { - const auto & rinfo = rinfos[i]; - ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); - } + llama_io_emit(rinfos, i, end, + [](ggml_tensor * tensor, const uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_set(nullptr, tensor, ptr, offset, size, n_copies, stride_tensor, stride_data); + }); + + i = end; } } @@ -2998,6 +3140,321 @@ size_t llama_context::state_set_data(const uint8_t * src, size_t size) { } } +// [TAG_STATE_ASYNC] a sequence state transfer that runs beside the decode instead of in it +// +// Everything the transfer needs to outlive the call that issued it lives here: the host +// buffer the bytes land in or come from, one backend per device holding part of the cache +// (each with a stream of its own, so the copies never queue behind the graphs), and one +// event per device to tell the caller when its half is finished. +struct llama_state_seq_copy { + llama_context * ctx = nullptr; + + struct dev_copy { + ggml_backend_ptr backend; + ggml_backend_event_t event = nullptr; + bool pending = false; + }; + + std::map devs; + + ggml_backend_buffer_ptr host_buf; + + uint8_t * data = nullptr; + size_t size = 0; // bytes the current transfer covers + size_t capacity = 0; // bytes actually held, kept across transfers + bool pinned = false; + bool can_pin = false; + + // transfers the last issue actually posted, i.e. runs of adjacent cells over all tensors + size_t n_copies = 0; + // microseconds the last issue spent draining the compute streams before it could start + int64_t t_sync_us = 0; + + ~llama_state_seq_copy() { + wait(); + + for (auto & it : devs) { + if (it.second.event) { + ggml_backend_event_free(it.second.event); + } + } + } + + // The stream this tensor is copied on, or null when it needs no stream: tensors already + // in host memory are a memcpy, and a tensor in a split or otherwise non-default buffer + // fails the buffer check every backend's async copy asserts, so both take the plain + // synchronous path. Handing a backend out marks it, so record() knows which ones ran. + ggml_backend_t backend_for(const ggml_tensor * t) { + ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; + + if (!buf || ggml_backend_buffer_is_host(buf)) { + return nullptr; + } + + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf); + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + + if (!dev || buft != ggml_backend_dev_buffer_type(dev)) { + return nullptr; + } + + auto it = devs.find(dev); + + if (it == devs.end()) { + return nullptr; + } + + it->second.pending = true; + + return it->second.backend.get(); + } + + // close every stream the transfer just used + void record() { + + for (auto & it : devs) { + if (it.second.pending) { + ggml_backend_event_record(it.second.event, it.second.backend.get()); + } + } + } + + bool done() { + bool res = true; + + for (auto & it : devs) { + if (!it.second.pending) { + continue; + } + + if (ggml_backend_event_query(it.second.event)) { + it.second.pending = false; + } else { + res = false; + } + } + + return res; + } + + void wait() { + for (auto & it : devs) { + if (!it.second.pending) { + continue; + } + + ggml_backend_event_synchronize(it.second.event); + + it.second.pending = false; + } + } + + // Grow-only. Pinning host memory is expensive -- a hundred MiB of it costs about as long + // as the copy it is for -- and a caller that parks the same sequence over and over asks + // for a slightly different size every time, so freeing between transfers would put that + // cost back on the very loop this is keeping clear. The memory is given back by + // buf_free() when the caller is finished with the slot, not between two of its parks. + uint8_t * buf_resize(size_t size_new) { + if (size_new <= capacity) { + size = size_new; + + return size_new == 0 ? nullptr : data; + } + + // never move memory a copy could still be reading or writing + wait(); + + host_buf.reset(); + + data = nullptr; + size = 0; + capacity = 0; + pinned = false; + + ggml_backend_buffer_type_t host_buft = host_buffer_type(); + + ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(host_buft, size_new); + + if (!buf) { + return nullptr; + } + + uint8_t * base = (uint8_t *) ggml_backend_buffer_get_base(buf); + + if (!base) { + ggml_backend_buffer_free(buf); + return nullptr; + } + + host_buf.reset(buf); + + data = base; + size = size_new; + capacity = size_new; + // a host buffer type may quietly hand back ordinary memory when pinning is turned + // off, so believe the buffer that came back rather than the type that was asked + pinned = can_pin && ggml_backend_buffer_get_type(buf) == host_buft; + + return data; + } + + void buf_free() { + wait(); + + host_buf.reset(); + + data = nullptr; + size = 0; + capacity = 0; + pinned = false; + } + + // Pinned host memory is the point of allocating through the backend at all: a copy in or + // out of pageable memory is staged through a pinned bounce buffer by the driver and + // blocks, which is exactly the stall being removed here. + ggml_backend_buffer_type_t host_buffer_type() { + for (auto & it : devs) { + ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(it.first); + + if (buft) { + return buft; + } + } + + return ggml_backend_cpu_buffer_type(); + } +}; + +class llama_io_write_host_async : public llama_io_write_i { +public: + llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : + ptr(p), buf_size(len), cpy(cpy) {} + + ~llama_io_write_host_async() { + llama_io_emit(winfos, 0, winfos.size(), + [this](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_get(cpy.backend_for(tensor), tensor, ptr, offset, size, + n_copies, stride_tensor, stride_data); + + cpy.n_copies++; + }); + + cpy.record(); + } + + void write(const void * src, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + memcpy(ptr, src, size); + ptr += size; + size_written += size; + buf_size -= size; + } + + void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + + winfos.push_back({tensor, ptr, size, offset}); + + ptr += size; + size_written += size; + buf_size -= size; + } + + size_t n_bytes() override { + return size_written; + } + +private: + uint8_t * ptr; + size_t buf_size = 0; + size_t size_written = 0; + + struct write_info { + ggml_tensor * tensor; + uint8_t * ptr; + size_t size; + size_t offset; + }; + std::vector winfos; + + llama_state_seq_copy & cpy; +}; + +class llama_io_read_host_async : public llama_io_read_i { +public: + llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : + ptr(p), buf_size(len), cpy(cpy) {} + + ~llama_io_read_host_async() { + // No whole-tensor staging here, unlike the synchronous path above. Staging reads a + // tensor, patches this sequence's bytes into the host copy and writes the whole + // tensor back, which preserves the neighbours only while nothing else is touching + // the cache. These copies are issued precisely so that decoding can carry on beside + // them, so a 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: + // every byte in them belongs to the sequence being restored. That is affordable + // because the runs have been coalesced -- one transfer per run of adjacent cells, + // which is what staging was working around in the first place. + llama_io_emit(rinfos, 0, rinfos.size(), + [this](ggml_tensor * tensor, const uint8_t * ptr, size_t offset, size_t size, + size_t n_copies, size_t stride_tensor, size_t stride_data) { + llama_io_set(cpy.backend_for(tensor), tensor, ptr, offset, size, + n_copies, stride_tensor, stride_data); + + cpy.n_copies++; + }); + + cpy.record(); + } + + void read(void * dst, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + memcpy(dst, ptr, size); + ptr += size; + size_read += size; + buf_size -= size; + } + + void read_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + if (size > buf_size) { + throw std::runtime_error("unexpectedly reached end of buffer"); + } + + rinfos.push_back({tensor, ptr, size, offset}); + + ptr += size; + size_read += size; + buf_size -= size; + } + + size_t n_bytes() override { + return size_read; + } + +private: + const uint8_t * ptr; + size_t buf_size = 0; + size_t size_read = 0; + + struct read_info { + ggml_tensor * tensor; + const uint8_t * ptr; + size_t size; + size_t offset; + }; + std::vector rinfos; + + llama_state_seq_copy & cpy; +}; + static constexpr uint32_t io_magic = 0xaf143cd8; size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_flags flags) { @@ -3071,6 +3528,117 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr } } +// [TAG_STATE_ASYNC] + +llama_state_seq_copy * llama_context::state_seq_copy_init() { + std::unique_ptr cpy(new llama_state_seq_copy()); + + cpy->ctx = this; + + for (auto & backend : backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); + + if (!dev || cpy->devs.find(dev) != cpy->devs.end()) { + continue; + } + + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev, &props); + + if (!props.caps.async || !props.caps.events) { + continue; + } + + // a backend of its own, not the one the graphs are computed on: that one moves its + // copies to whichever stream it is currently using, so a transfer posted to it could + // end up ordered behind a graph -- which is the stall this exists to avoid + ggml_backend_t backend_cpy = ggml_backend_dev_init(dev, nullptr); + + if (!backend_cpy) { + continue; + } + + ggml_backend_event_t event = ggml_backend_event_new(dev); + + if (!event) { + ggml_backend_free(backend_cpy); + continue; + } + + auto & dc = cpy->devs[dev]; + + dc.backend.reset(backend_cpy); + dc.event = event; + } + + if (cpy->devs.empty()) { + return nullptr; + } + + cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); + + return cpy.release(); +} + +size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { + if (!cpy.data) { + return 0; + } + + // The copies run on their own stream and are ordered against nothing, so the decode that + // produced these cells has to be finished before they are read. This is the one part of + // the transfer that stays on the caller's thread, and it costs nothing where it is used: + // a caller preempting a sequence does it between two decodes, with the previous one + // already drained by the sampling that followed it. + const int64_t t_sync = ggml_time_us(); + synchronize(); + cpy.t_sync_us = ggml_time_us() - t_sync; + + cpy.n_copies = 0; + + llama_io_write_host_async io(cpy.data, size, cpy); + + try { + io.write(&io_magic, sizeof(io_magic)); + io.write(&seq_id, sizeof(seq_id)); + + return state_seq_write_data(io, seq_id, flags); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); + return 0; + } +} + +size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { + if (!cpy.data) { + return 0; + } + + const int64_t t_sync = ggml_time_us(); + synchronize(); + cpy.t_sync_us = ggml_time_us() - t_sync; + + cpy.n_copies = 0; + + llama_io_read_host_async io(cpy.data, size, cpy); + + try { + uint32_t magic_read; + io.read(&magic_read, sizeof(magic_read)); + if (io_magic != magic_read) { + throw std::runtime_error("wrong sequence state magic"); + } + + llama_seq_id seq_id_read; + io.read(&seq_id_read, sizeof(seq_id_read)); + + return state_seq_read_data(io, seq_id, flags); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); + return 0; + } +} + bool llama_context::state_load_file(const char * filepath, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { llama_file file(filepath, "rb"); @@ -4125,6 +4693,64 @@ size_t llama_state_seq_set_data_ext(llama_context * ctx, const uint8_t * src, si return ctx->state_seq_set_data(seq_id, src, size, flags); } +// [TAG_STATE_ASYNC] + +llama_state_seq_copy * llama_state_seq_copy_init(llama_context * ctx) { + return ctx->state_seq_copy_init(); +} + +void llama_state_seq_copy_free(llama_state_seq_copy * cpy) { + delete cpy; // waits for anything still in flight +} + +uint8_t * llama_state_seq_copy_buf_resize(llama_state_seq_copy * cpy, size_t size) { + return cpy->buf_resize(size); +} + +uint8_t * llama_state_seq_copy_buf(llama_state_seq_copy * cpy) { + return cpy->data; +} + +size_t llama_state_seq_copy_buf_size(llama_state_seq_copy * cpy) { + return cpy->size; +} + +size_t llama_state_seq_copy_buf_capacity(llama_state_seq_copy * cpy) { + return cpy->capacity; +} + +size_t llama_state_seq_copy_n_copies(llama_state_seq_copy * cpy) { + return cpy->n_copies; +} + +int64_t llama_state_seq_copy_sync_us(llama_state_seq_copy * cpy) { + return cpy->t_sync_us; +} + +void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { + cpy->buf_free(); +} + +bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { + return cpy->can_pin; +} + +size_t llama_state_seq_copy_get(llama_state_seq_copy * cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { + return cpy->ctx->state_seq_copy_get(*cpy, size, seq_id, flags); +} + +size_t llama_state_seq_copy_set(llama_state_seq_copy * cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags) { + return cpy->ctx->state_seq_copy_set(*cpy, size, dest_seq_id, flags); +} + +bool llama_state_seq_copy_done(llama_state_seq_copy * cpy) { + return cpy->done(); +} + +void llama_state_seq_copy_wait(llama_state_seq_copy * cpy) { + cpy->wait(); +} + size_t llama_state_seq_save_file(llama_context * ctx, const char * filepath, llama_seq_id seq_id, const llama_token * tokens, size_t n_token_count) { ctx->synchronize(); diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b56..3d66f2a948a 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -39,6 +39,9 @@ struct llama_memory_buffer { using llama_memory_buffers = std::map; +// [TAG_STATE_ASYNC] defined in llama-context.cpp +struct llama_state_seq_copy; + struct llama_context { // init scheduler and compute buffers, reserve worst-case graphs llama_context( @@ -156,6 +159,12 @@ struct llama_context { size_t state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size, llama_state_seq_flags flags); size_t state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] the same two transfers, issued on a stream of their own and left running + llama_state_seq_copy * state_seq_copy_init(); + + size_t state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags); + size_t state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags); + bool state_load_file( const char * filepath, llama_token * tokens_out, From 1d98f93cbfe4d4ed97cdea5b2aa9ed1c41302d1a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 5 Sep 2026 23:35:19 +0000 Subject: [PATCH 03/25] server: take the park and restore copies off the decode loop 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. --- common/arg.cpp | 10 + common/common.h | 1 + tools/server/README.md | 1 + tools/server/server-context.cpp | 489 ++++++++++++++++++++++-- tools/server/tests/unit/test_preempt.py | 180 +++++++++ 5 files changed, 657 insertions(+), 24 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 5bfa4adcdf0..30a208e707f 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1717,6 +1717,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.preempt_ram_mib = value; } ).set_env("LLAMA_ARG_PREEMPT_RAM").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--preempt-async"}, + {"--no-preempt-async"}, + "copy a parked sequence out of and back into the KV cache on a stream of its own, so the " + "slots that keep running do not wait for it (default: enabled, needs a backend that can " + "copy asynchronously, otherwise the copies are synchronous as before)", + [](common_params & params, bool value) { + params.preempt_async = value; + } + ).set_env("LLAMA_ARG_PREEMPT_ASYNC").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index c99269f9a96..1c220e6b785 100644 --- a/common/common.h +++ b/common/common.h @@ -615,6 +615,7 @@ struct common_params { int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = disable preemption + bool preempt_async = true; // park and restore on a stream of their own, off the decode loop std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/tools/server/README.md b/tools/server/README.md index 7b4a0330340..bec3f09ddff 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -165,6 +165,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | | `--preempt-ram N` | with a unified KV cache, park a slot in host RAM instead of failing every slot when the cache fills; N is the maximum host RAM for parked sequences in MiB (default: 8192, -1 - no limit, 0 - disable)
(env: LLAMA_ARG_PREEMPT_RAM) | +| `--preempt-async`, `--no-preempt-async` | copy a parked sequence out of and back into the KV cache on a stream of its own, so the slots that keep running do not wait for it (default: enabled, needs a backend that can copy asynchronously, otherwise the copies are synchronous as before)
(env: LLAMA_ARG_PREEMPT_ASYNC) | | `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)
(env: LLAMA_ARG_KV_UNIFIED) | | `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)
(env: LLAMA_ARG_CACHE_IDLE_SLOTS) | | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6723c51397e..f15f78f8201 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -60,6 +60,8 @@ enum slot_state { SLOT_STATE_DONE_PROMPT, SLOT_STATE_GENERATING, SLOT_STATE_PREEMPTED, // [TAG_PREEMPT] cells released, everything needed to resume is in host RAM + SLOT_STATE_PREEMPTING, // [TAG_PREEMPT_ASYNC] the copy out is running; the cells are still this slot's + SLOT_STATE_RESTORING, // [TAG_PREEMPT_ASYNC] the copy back in is running; the cells are allocated but not yet filled }; // [TAG_PREEMPT] server-side request preemption @@ -80,6 +82,34 @@ constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked +// [TAG_PREEMPT_ASYNC] how far ahead of the pool filling an asynchronous park is triggered +// +// A synchronous park hands the cells back before update_slots() goes on, so it only has to +// fire once the next step would not fit. An asynchronous one does not: the copy is still +// reading the cells, and they are only released when it lands. The slots that keep decoding +// in the meantime need somewhere to put their tokens, so the park has to be triggered this +// many decode steps before the pool would actually have run out. Too small and the pool +// fills while the copy is still running, and the decode ends up waiting for it after all, +// which is no worse than not doing this at all but no better either. Too large and slots +// are parked, and so parked again, earlier and more often than they need to be, which costs +// more in total than the one late park it avoided. +// +// Eight steps of every running slot is about a tenth of a second of runway at the speeds a +// handful of parallel chats decode at, which is the order a copy of one sequence takes. +constexpr int32_t PREEMPT_N_ASYNC_STEPS = 8; + +struct llama_state_seq_copy_deleter { + void operator()(llama_state_seq_copy * cpy) const { llama_state_seq_copy_free(cpy); } +}; + +using llama_state_seq_copy_ptr = std::shared_ptr; + +static llama_state_seq_copy_ptr llama_state_seq_copy_make(llama_context * ctx) { + llama_state_seq_copy * cpy = ctx ? llama_state_seq_copy_init(ctx) : nullptr; + + return cpy ? llama_state_seq_copy_ptr(cpy, llama_state_seq_copy_deleter{}) : llama_state_seq_copy_ptr(); +} + struct server_slot; // forward declaration struct server_batch { @@ -320,32 +350,210 @@ struct server_slot { slot_state state_before_preempt = SLOT_STATE_IDLE; std::vector preempt_state_tgt; std::vector preempt_state_dft; + + // [TAG_PREEMPT_ASYNC] the two transfers this slot parks and resumes through + // + // They own the pinned host buffers the sequence lives in while it is parked, so when + // they exist the std::vectors above stay empty and the state is in the transfers. Held + // by shared_ptr only because the slots are built with emplace_back into a vector that + // reallocates as it grows, and a slot must survive being moved. + llama_state_seq_copy_ptr preempt_cpy_tgt; + llama_state_seq_copy_ptr preempt_cpy_dft; + + bool preempt_is_async() const { + return (bool) preempt_cpy_tgt; + } + + // microseconds the last park or resume spent draining the compute streams + int64_t preempt_sync_us() const { + if (!preempt_is_async()) { + return 0; + } + + return llama_state_seq_copy_sync_us(preempt_cpy_tgt.get()) + + (preempt_cpy_dft ? llama_state_seq_copy_sync_us(preempt_cpy_dft.get()) : 0); + } + + // transfers the last park or resume posted, which is what its issue cost is made of + size_t preempt_n_copies() const { + if (!preempt_is_async()) { + return 0; + } + + return llama_state_seq_copy_n_copies(preempt_cpy_tgt.get()) + + (preempt_cpy_dft ? llama_state_seq_copy_n_copies(preempt_cpy_dft.get()) : 0); + } + + // [TAG_PREEMPT_ASYNC] a copy is running for this slot: it is not decoding and must not be + // scheduled, but it still owns cells, so it is neither running nor parked + bool preempt_in_flight() const { + return state == SLOT_STATE_PREEMPTING || state == SLOT_STATE_RESTORING; + } + + // parked, or on its way out or back in: in none of these does the slot take part in a decode + bool preempt_is_out() const { + return state == SLOT_STATE_PREEMPTED || preempt_in_flight(); + } int32_t n_preempt = 0; // times the CURRENT task has been preempted int32_t n_preempt_fail = 0; // consecutive failed restores int64_t t_preempt_us = 0; // when it was parked + int64_t t_preempt_copy_us = 0; // [TAG_PREEMPT_ASYNC] when the current copy was issued size_t preempt_state_size() const { + if (preempt_is_async()) { + // the capacity, not the live size: the pinned buffers are kept between two parks + // of the same task because page-locking them again would cost as much as the + // copy, so what --preempt-ram has to bound is what is held, not what is in use + return llama_state_seq_copy_buf_capacity(preempt_cpy_tgt.get()) + + (preempt_cpy_dft ? llama_state_seq_copy_buf_capacity(preempt_cpy_dft.get()) : 0); + } + return preempt_state_tgt.size() + preempt_state_dft.size(); } void preempt_state_free() { + // resizing waits for anything still in flight first: this is called from release(), + // which a cancelled request reaches while its copy may still be reading or writing + // the buffer, and freeing it underneath a running transfer would be a use-after-free + if (preempt_cpy_tgt) { + llama_state_seq_copy_buf_free(preempt_cpy_tgt.get()); + } + + if (preempt_cpy_dft) { + llama_state_seq_copy_buf_free(preempt_cpy_dft.get()); + } + preempt_state_tgt.clear(); preempt_state_tgt.shrink_to_fit(); preempt_state_dft.clear(); preempt_state_dft.shrink_to_fit(); } + // [TAG_PREEMPT_ASYNC] give up on an outstanding copy without using its result + void preempt_copy_wait() { + if (preempt_cpy_tgt) { + llama_state_seq_copy_wait(preempt_cpy_tgt.get()); + } + + if (preempt_cpy_dft) { + llama_state_seq_copy_wait(preempt_cpy_dft.get()); + } + } + // bytes preempt_save() would need for this slot right now size_t preempt_state_required() const { return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0); } + // take the slot out of the step that is about to be built + // + // The draft is a prediction, not a result, so it goes with the cells. Preemption runs + // before the batch is built, so spec_i_batch is empty and prompt.tokens already holds + // exactly the tokens the state covers -- including the rollback done by the checkpoint + // path when a draft was only partially accepted. + void preempt_detach() { + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + } + + // [TAG_PREEMPT_ASYNC] has the copy out finished? if so, the cells can finally go + bool preempt_save_poll() { + if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { + return false; + } + + if (preempt_cpy_dft && !llama_state_seq_copy_done(preempt_cpy_dft.get())) { + return false; + } + + mem.seq_rm(id, -1, -1); + + state = SLOT_STATE_PREEMPTED; + + return true; + } + + // [TAG_PREEMPT_ASYNC] has the copy back in finished? if so, the slot can decode again + bool preempt_restore_poll() { + if (!llama_state_seq_copy_done(preempt_cpy_tgt.get())) { + return false; + } + + if (preempt_cpy_dft && !llama_state_seq_copy_done(preempt_cpy_dft.get())) { + return false; + } + + // the state is back in the cache, so the buffers hold nothing that matters; the + // memory itself is kept for the next park of this task and handed back by release() + llama_state_seq_copy_buf_resize(preempt_cpy_tgt.get(), 0); + + if (preempt_cpy_dft) { + llama_state_seq_copy_buf_resize(preempt_cpy_dft.get(), 0); + } + + n_preempt_fail = 0; + + state = state_before_preempt; + + // same call the DONE_PROMPT -> GENERATING transition makes; it reads the restored + // sequence, so it has to wait for the copy like everything else + if (state == SLOT_STATE_GENERATING && can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + + return true; + } + // copy the sequence out of the cache and release its cells + // + // [TAG_PREEMPT_ASYNC] With a transfer this returns as soon as the copy has been issued, + // leaving the slot PREEMPTING: the cells are still its own, because the copy is still + // reading them, and nobody may take them until preempt_save_poll() says the copy landed. bool preempt_save() { const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + if (preempt_is_async()) { + if (!llama_state_seq_copy_buf_resize(preempt_cpy_tgt.get(), size_tgt) || + (size_dft > 0 && (!preempt_cpy_dft || + !llama_state_seq_copy_buf_resize(preempt_cpy_dft.get(), size_dft)))) { + SLT_ERR(*this, "failed to allocate %.3f MiB of pinned host memory for the preemption state\n", + (size_tgt + size_dft) / (1024.0 * 1024.0)); + preempt_state_free(); + return false; + } + + if (llama_state_seq_copy_get(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to issue the copy of the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_copy_get(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to issue the copy of the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + preempt_detach(); + + // note: no mem.seq_rm() here. The copy is still reading these cells, so they are + // released in preempt_save_poll() once it has finished with them. + state_before_preempt = state; + state = SLOT_STATE_PREEMPTING; + t_preempt_us = ggml_time_us(); + + n_preempt++; + + return true; + } + try { preempt_state_tgt.resize(size_tgt); preempt_state_dft.resize(size_dft); @@ -369,16 +577,7 @@ struct server_slot { return false; } - // The draft is a prediction, not a result, so it goes with the cells. Preemption - // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already - // holds exactly the tokens the state above covers -- including the rollback done by - // the checkpoint path when a draft was only partially accepted. - spec_draft.clear(); - spec_i_batch.clear(); - spec_ckpt.clear(); - spec_is_replay = false; - - i_batch = -1; + preempt_detach(); // note: prompt.tokens is deliberately kept. It is the mirror of the state just // copied out, and the resume needs it to know how many cells to ask for. @@ -394,7 +593,33 @@ struct server_slot { } // put the sequence back; the slot then continues from the token it was about to decode + // + // [TAG_PREEMPT_ASYNC] With a transfer this returns as soon as the copy has been issued, + // leaving the slot RESTORING: the cells are allocated and owned by this 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 until preempt_restore_poll() says so. bool preempt_restore() { + if (preempt_is_async()) { + const size_t size_tgt = llama_state_seq_copy_buf_size(preempt_cpy_tgt.get()); + const size_t size_dft = preempt_cpy_dft ? llama_state_seq_copy_buf_size(preempt_cpy_dft.get()) : 0; + + if (llama_state_seq_copy_set(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt || + (size_dft > 0 && + llama_state_seq_copy_set(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft)) { + // no room after all: let whatever was already issued finish before the + // half-written sequence is dropped, or the cells would go while a copy is + // still writing into them + preempt_copy_wait(); + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + state = SLOT_STATE_RESTORING; + + return true; + } + const size_t size_tgt = preempt_state_tgt.size(); const size_t size_dft = preempt_state_dft.size(); @@ -649,7 +874,13 @@ struct server_slot { // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be // released while parked), so the mirror of them must not outlive them: the next // task on this slot would otherwise take a prefix match against an empty cache - if (state == SLOT_STATE_PREEMPTED) { + // + // [TAG_PREEMPT_ASYNC] a slot can also be released with a copy still running, by + // a cancelled request or by the error paths. Wait for it before anything else: + // the host buffer is about to be freed and the cells about to be handed to the + // next task, and a transfer still reading or writing either would outlive both. + if (preempt_is_out()) { + preempt_copy_wait(); preempt_state_free(); prompt_clear(); } @@ -796,7 +1027,7 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, - {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"is_preempted", preempt_is_out()}, {"n_preempt", n_preempt}, }; @@ -1381,6 +1612,22 @@ struct server_context_impl { } }; + // [TAG_PREEMPT_ASYNC] one transfer per context, made once and reused for every + // park and resume this slot ever does, because each owns a backend and a stream + if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + slot.preempt_cpy_tgt = llama_state_seq_copy_make(ctx_tgt); + + if (slot.preempt_cpy_tgt && ctx_dft) { + slot.preempt_cpy_dft = llama_state_seq_copy_make(ctx_dft); + + if (!slot.preempt_cpy_dft) { + // a draft that cannot go asynchronously would have to be waited for + // in the middle of the park, so the whole slot stays synchronous + slot.preempt_cpy_tgt.reset(); + } + } + } + slot.reset(); } @@ -1402,6 +1649,31 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] the slots either all park through a transfer or none do + { + preempt_async_ok = !slots.empty(); + + for (const auto & slot : slots) { + preempt_async_ok = preempt_async_ok && slot.preempt_is_async(); + } + + if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + if (preempt_async_ok) { + SRV_INF("preemption: parking and resuming asynchronously, %s host memory\n", + llama_state_seq_copy_buf_is_pinned(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + } else { + SRV_WRN("%s", "preemption: this backend cannot copy asynchronously, parking and resuming synchronously\n"); + } + } + + if (!preempt_async_ok) { + for (auto & slot : slots) { + slot.preempt_cpy_tgt.reset(); + slot.preempt_cpy_dft.reset(); + } + } + } + { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; @@ -2554,7 +2826,7 @@ struct server_context_impl { if (slot.is_processing()) { n_processing_slots++; } - if (slot.state == SLOT_STATE_PREEMPTED) { + if (slot.preempt_is_out()) { n_preempted_slots++; } } @@ -2855,6 +3127,11 @@ struct server_context_impl { // uninterrupted one is the preemption's fault and nothing else's. int32_t preempt_test_every = 0; + // [TAG_PREEMPT_ASYNC] whether the slots park and resume through a transfer. False when + // --no-preempt-async was given, or when the backend cannot copy asynchronously, in which + // case every park and resume is the synchronous one it always was. + bool preempt_async_ok = false; + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -2878,7 +3155,13 @@ struct server_context_impl { const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - return preempt_ram_used() + slot.preempt_state_required() <= budget; + // whatever this slot already holds is counted by preempt_ram_used() and will be + // reused, so parking it again only costs what it does not have yet + const size_t held = slot.preempt_state_size(); + const size_t need = slot.preempt_state_required(); + const size_t extra = need > held ? need - held : 0; + + return preempt_ram_used() + extra <= budget; } // cells the slot will ask for on its next step once it is back in the pool @@ -2908,12 +3191,40 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } + // [TAG_PREEMPT_ASYNC] deliberately not skipped: a slot whose copy is still + // running holds cells either way. One on its way out has not released them yet + // because the copy is still reading them, and one on its way back in has already + // been given them. Skipping either would hand the same cells out twice. + res += slot.prompt.n_tokens(); } return res; } + // [TAG_PREEMPT_ASYNC] the room the pool is kept clear of + // + // A synchronous park releases the cells before update_slots() carries on, so it only has + // to fire once the next step would not fit. An asynchronous one leaves them held until + // its copy lands, so it has to fire early enough that everything still decoding has + // somewhere to put its tokens until then. 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. + int32_t preempt_n_margin() const { + if (!preempt_async_active()) { + return PREEMPT_N_MARGIN; + } + + int32_t n_running = 0; + + for (const auto & slot : slots) { + if (slot.is_processing() && !slot.preempt_is_out()) { + n_running++; + } + } + + return PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS; + } + // cells those slots are about to ask for on the next decode int32_t preempt_kv_reserve() const { const int32_t n_spec = preempt_n_spec_max(); @@ -2936,6 +3247,21 @@ struct server_context_impl { res_pmt += std::max(1, std::min(n_batch, n_left)); } break; + case SLOT_STATE_RESTORING: + { + // [TAG_PREEMPT_ASYNC] its cells are already counted by preempt_kv_used(), + // but it starts decoding as soon as its copy lands, so the step it will + // take has to be reserved now -- otherwise the pool is handed out from + // under it and its first step preempts somebody else straight away + if (slot.state_before_preempt == SLOT_STATE_GENERATING) { + res += 1 + n_spec; + } else { + const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0; + + res_pmt += std::max(1, std::min(n_batch, n_left)); + } + } break; + // a preempting slot is on its way out and will not decode: nothing to reserve default: break; } @@ -2953,7 +3279,7 @@ struct server_context_impl { int32_t n_running = 0; for (auto & slot : slots) { - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + if (slot.is_processing() && !slot.preempt_is_out()) { n_running++; if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) { @@ -3011,11 +3337,75 @@ struct server_context_impl { // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no // draft is in flight, so a slot can be removed from the picture without unpicking a // half-decoded batch + // [TAG_PREEMPT_ASYNC] is any slot parking or resuming through a transfer right now + bool preempt_async_active() const { + return preempt_async_ok; + } + + // Pick up the copies that have landed since the last iteration. This runs before + // anything reads preempt_kv_used(), so a park whose cells came back is seen as free + // room straight away and a resume that landed can be scheduled in the same iteration. + void update_preempt_copies() { + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTING) { + if (slot.preempt_save_poll()) { + metrics.n_preempt++; + + SLT_WRN(slot, "park completed after %.2f ms: %d cells released, %.1f MiB parked, kv %d/%d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + slot.preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_ctx); + } + } else if (slot.state == SLOT_STATE_RESTORING) { + if (slot.preempt_restore_poll()) { + metrics.n_resume++; + + SLT_WRN(slot, "restore completed after %.2f ms: %d tokens back in the cache, kv %d/%d, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + } + } + } + + // Wait for one outstanding park, the last thing tried before giving up on finding room. + // It is what keeps a pool that fills faster than the copies drain no worse than the + // synchronous path: the decode waits for the copy exactly as it used to. + bool preempt_wait_in_flight() { + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_PREEMPTING) { + continue; + } + + slot.preempt_copy_wait(); + + if (!slot.preempt_save_poll()) { + continue; + } + + metrics.n_preempt++; + + SLT_WRN(slot, "park completed after %.2f ms (waited for): %d cells released, kv %d/%d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx); + + return true; + } + + return false; + } + void update_preemption() { if (!params_base.kv_unified || slots.size() < 2) { return; // with a cache per slot, no slot can take another one's cells } + update_preempt_copies(); + if (params_base.preempt_ram_mib == 0) { return; // --preempt-ram 0: the KV-full retry ladder, as before } @@ -3055,7 +3445,7 @@ struct server_context_impl { // continue, so give those cells up first - same call the KV-full path makes. for (;;) { for (auto * slot : parked) { - if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin() <= n_cells) { best = slot; break; } @@ -3072,6 +3462,8 @@ struct server_context_impl { const int64_t t_start = ggml_time_us(); + best->t_preempt_copy_us = t_start; + if (!best->preempt_restore()) { // update_slots() runs in a tight loop while tasks are pending, so a counter // alone burns its whole budget in a couple of milliseconds. Give up only on @@ -3090,6 +3482,22 @@ struct server_context_impl { break; } + // [TAG_PREEMPT_ASYNC] with a transfer the copy has only been issued; the slot is + // RESTORING and update_preempt_copies() counts it and logs it when it lands + if (best->state == SLOT_STATE_RESTORING) { + SLT_WRN(*best, "resumed after %.2f s: %d tokens, restore issued in %.2f ms (%zu transfers, %.2f ms sync), kv %d/%d, preemptions %d\n", + (ggml_time_us() - best->t_preempt_us) / 1e6, + best->prompt.n_tokens(), + (ggml_time_us() - t_start) / 1e3, + best->preempt_n_copies(), best->preempt_sync_us() / 1e3, + preempt_kv_used(), n_cells, + best->n_preempt); + + // it holds cells now but is not decoding yet, so there is nothing more to + // decide about it this iteration + continue; + } + metrics.n_resume++; SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", @@ -3105,12 +3513,19 @@ struct server_context_impl { for (auto & slot : slots) { if (slot.state == SLOT_STATE_GENERATING && (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && - preempt_fits_budget(slot) && - slot.preempt_save()) { - metrics.n_preempt++; + preempt_fits_budget(slot)) { + slot.t_preempt_copy_us = ggml_time_us(); + + if (slot.preempt_save()) { + // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by + // update_preempt_copies() when its copy lands, not here + if (slot.state == SLOT_STATE_PREEMPTED) { + metrics.n_preempt++; + } - SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", - (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", + (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + } } } } @@ -3119,7 +3534,7 @@ struct server_context_impl { for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); - if (n_used + PREEMPT_N_MARGIN <= n_cells) { + if (n_used + preempt_n_margin() <= n_cells) { break; } @@ -3128,6 +3543,15 @@ struct server_context_impl { continue; } + // [TAG_PREEMPT_ASYNC] Out of room for the step about to be built, rather than + // merely short of the lookahead the asynchronous path keeps. A park that has + // been issued but not landed is holding cells that are already spoken for, and + // waiting for it is both quicker and more useful than parking somebody else, + // whose cells would not come back this iteration either. + if (n_used + PREEMPT_N_MARGIN > n_cells && preempt_wait_in_flight()) { + continue; + } + server_slot * victim = preempt_pick_victim(); if (!victim) { @@ -3139,10 +3563,27 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); + victim->t_preempt_copy_us = t_start; + if (!victim->preempt_save()) { break; // could not park it; the existing retry ladder is still behind us } + // [TAG_PREEMPT_ASYNC] the copy has only been issued; the cells are still the + // victim's until it lands, so nothing further can be decided about the pool this + // iteration. update_preempt_copies() picks it up on the next one, and the step + // that wanted the room is built from whatever is free right now. + if (victim->state == SLOT_STATE_PREEMPTING) { + SLT_WRN(*victim, "preempted: %d cells, park issued in %.2f ms (%zu transfers, %.2f ms sync), %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_n_copies(), victim->preempt_sync_us() / 1e3, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + break; + } + metrics.n_preempt++; SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", @@ -3485,7 +3926,7 @@ struct server_context_impl { // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to // batch; it takes no part in this pass until it is restored - if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { + if (!slot.is_processing() || slot.preempt_is_out()) { return; } @@ -4429,7 +4870,7 @@ struct server_context_impl { metrics.n_decode++; for (const auto & slot : slots) { // [TAG_PREEMPT] a parked slot is processing but took no part in this decode - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + if (slot.is_processing() && !slot.preempt_is_out()) { metrics.n_busy_slots++; } metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens()); diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 0da885bcafd..1aba01751bd 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -39,6 +39,7 @@ def create_server(): yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) + os.environ.pop("LLAMA_ARG_PREEMPT_ASYNC", None) def _complete(n_predict: int, prompt: str = "Hi how are you"): @@ -267,3 +268,182 @@ def test_metrics_and_slots_report_the_parked_state(): res = server.make_request("GET", "/slots") assert res.status_code == 200 assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot" + + +# [TAG_PREEMPT_ASYNC] parking and resuming on a stream of their own +# +# The copies are only asynchronous on a backend that can copy asynchronously and signal an +# event, which today means a GPU one. On a CPU-only build the server says so and falls back +# to the synchronous path, and the tests below that need the asynchronous one skip. + +_ASYNC_BANNER = "parking and resuming asynchronously" + + +def _start_async(**kwargs) -> str: + """Start the server with the asynchronous path asked for, and return its log so far.""" + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "1" + for key, value in kwargs.items(): + setattr(server, key, value) + server.start() + with open(server.log_path) as f: + return f.read() + + +def _require_async(text: str): + if _ASYNC_BANNER not in text: + pytest.skip("this backend cannot copy asynchronously, the async park path is not exercised") + + +def test_async_preemption_does_not_change_the_output(): + # The same question the synchronous determinism test asks, of the asynchronous path: + # with one request the batch has the same shape at every step, so a continuation that + # was parked and resumed through a transfer and is not byte-identical to an + # uninterrupted one is the transfer's fault and nothing else's. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + text = _start_async() + _require_async(text) + + res_plain = _complete(64) + assert res_plain.status_code == 200 + + server.stop() + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + + res_preempted = _complete(64) + assert res_preempted.status_code == 200 + + text = log.drain() + _require_async(text) + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 + # the asynchronous path is the one that ran, not the synchronous fallback: only it + # splits a park and a resume into an issue and a completion + assert "park completed after" in text + assert "restore issued in" in text + assert "restore completed after" in text + + assert res_preempted.body["content"] == res_plain.body["content"] + assert res_preempted.body["tokens"] == res_plain.body["tokens"] + + +def test_async_preemption_under_load_keeps_every_slot_and_its_output(): + # Two requests that do not fit the pool together, parked and resumed asynchronously + # while the other one keeps decoding. Every slot must finish, and finish with exactly + # the tokens it produces when it has the pool to itself. + global server + server.n_ctx = 256 + server.n_gpu_layer = 99 + text = _start_async() + _require_async(text) + + n_predict = 160 + prompts = [ + "Once upon a time there was a brave knight who", + "The quick brown fox jumps over the lazy dog and", + ] + + # each one alone, for the reference tokens + alone = [_complete(n_predict, prompt) for prompt in prompts] + for res in alone: + assert res.status_code == 200 + + server.stop() + server.start() + log = LogReader(server.log_path) + + together = parallel_function_calls([(_complete, (n_predict, prompt)) for prompt in prompts]) + + text = log.drain() + _require_async(text) + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res, ref in zip(together, alone): + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert res.body["tokens"] == ref.body["tokens"] + + +def _cancel_soon(n_predict: int, prompt: str, timeout: float): + try: + server.make_request("POST", "/completion", data={ + "n_predict": n_predict, + "prompt": prompt, + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + }, timeout=timeout) + except Exception: + pass # the point is the drop, not the response + + +def test_cancel_while_a_copy_is_in_flight_frees_the_slot(): + # A cancelled request can reach release() with a park or a resume still running, which + # is where the host buffer is freed and the cells are handed on. Both have to wait for + # the copy first. LLAMA_SERVER_PREEMPT_EVERY keeps every slot cycling between the two + # states, so cancelling at a spread of moments lands in both; what is asserted is that + # the server survives it, the slots come back, and it still answers correctly. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + # every 8 tokens, so a slot spends most of its life in one of the two copy states, but + # not so often that the abandoned requests take minutes to drain + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + text = _start_async() + _require_async(text) + + for i in range(4): + _cancel_soon(96, "Once upon a time there was a brave knight who", 0.05 + 0.1 * i) + + # every slot back, and none of them still holding a parked sequence + deadline = time.time() + 120 + while time.time() < deadline: + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + if all(not slot["is_processing"] for slot in res.body): + break + time.sleep(0.2) + else: + pytest.fail("a slot never came back after a cancel during a copy") + + for slot in res.body: + assert slot["is_preempted"] is False + + if server.server_metrics: + res = server.make_request("GET", "/metrics") + for line in res.body.splitlines(): + if line.startswith("llamacpp:preempt_ram_bytes"): + assert float(line.split(" ", 1)[1]) == 0, "a cancelled slot kept its parked memory" + + # and the server still works + res = _complete(16) + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == 16 + + +def test_no_preempt_async_falls_back_to_the_synchronous_path(): + # The flag has to really switch it off, so that the two can be compared on one binary. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "0" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + + res = _complete(64) + assert res.status_code == 200 + + text = log.drain() + assert _ASYNC_BANNER not in text + assert "park issued in" not in text + assert "restore issued in" not in text + # the synchronous path still parks and resumes + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 From e7e88e9de8bbb7d559bddc8969795480f5674edd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:30:36 +0000 Subject: [PATCH 04/25] ggml: map the event query onto HIP and MUSA 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. --- ggml/src/ggml-cuda/vendors/hip.h | 2 ++ ggml/src/ggml-cuda/vendors/musa.h | 2 ++ 2 files changed, 4 insertions(+) diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index 9aa558f3f4c..83d37aeeb2a 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -58,10 +58,12 @@ #define cudaDeviceSynchronize hipDeviceSynchronize #define cudaError_t hipError_t #define cudaErrorMemoryAllocation hipErrorOutOfMemory +#define cudaErrorNotReady hipErrorNotReady #define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled #define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled #define cudaEventCreateWithFlags hipEventCreateWithFlags #define cudaEventDisableTiming hipEventDisableTiming +#define cudaEventQuery hipEventQuery #define cudaEventRecord hipEventRecord #define cudaEventSynchronize hipEventSynchronize #define cudaEvent_t hipEvent_t diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h index 6d725c7ec19..ebecf679950 100644 --- a/ggml/src/ggml-cuda/vendors/musa.h +++ b/ggml/src/ggml-cuda/vendors/musa.h @@ -46,10 +46,12 @@ #define cudaDeviceSynchronize musaDeviceSynchronize #define cudaError_t musaError_t #define cudaErrorMemoryAllocation musaErrorMemoryAllocation +#define cudaErrorNotReady musaErrorNotReady #define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled #define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled #define cudaEventCreateWithFlags musaEventCreateWithFlags #define cudaEventDisableTiming musaEventDisableTiming +#define cudaEventQuery musaEventQuery #define cudaEventRecord musaEventRecord #define cudaEventSynchronize musaEventSynchronize #define cudaEvent_t musaEvent_t From d33f04fd639b8209a0d6b94d97803b4efc391591 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:33:35 +0000 Subject: [PATCH 05/25] llama: require a real event query before copying a sequence asynchronously 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. --- ggml/include/ggml-backend.h | 6 +++++- ggml/src/ggml-backend.cpp | 5 +++++ include/llama.h | 5 +++-- src/llama-context.cpp | 19 +++++++++++++++++++ 4 files changed, 32 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 30d8304d492..d21bf40dd58 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -126,7 +126,8 @@ extern "C" { GGML_API void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backend); GGML_API void ggml_backend_event_synchronize(ggml_backend_event_t event); // non-blocking: true once everything recorded before the event has completed. - // backends without a query implementation fall back to a blocking synchronize and return true. + // backends without a query implementation fall back to a blocking synchronize and return true, + // which ggml_backend_dev_supports_event_query() tells apart from a real non-blocking query. GGML_API bool ggml_backend_event_query(ggml_backend_event_t event); GGML_API void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event); @@ -193,6 +194,9 @@ extern "C" { GGML_API ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size); GGML_API bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op); + // whether ggml_backend_event_query() on this device really is non-blocking, i.e. whether + // the device implements it rather than falling back to a blocking synchronize + GGML_API bool ggml_backend_dev_supports_event_query(ggml_backend_dev_t device); GGML_API bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft); GGML_API bool ggml_backend_dev_offload_op(ggml_backend_dev_t device, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index a56ab30862a..b13d9c811c4 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -639,6 +639,11 @@ bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_t return device->iface.supports_op(device, op); } +bool ggml_backend_dev_supports_event_query(ggml_backend_dev_t device) { + GGML_ASSERT(device); + return device->iface.event_query != NULL; +} + bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { GGML_ASSERT(device); return device->iface.supports_buft(device, buft); diff --git a/include/llama.h b/include/llama.h index 2c61b3d2a71..26db483e6c2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -944,8 +944,9 @@ extern "C" { // written. llama_state_seq_copy_free() waits for an outstanding copy first. struct llama_state_seq_copy; - // NULL if the context's backends cannot copy asynchronously; the caller then uses the - // synchronous llama_state_seq_*_data_ext calls + // NULL if the context's backends cannot copy asynchronously, or cannot say whether a + // copy has finished without waiting for it, which would put the stall straight back; the + // caller then uses the synchronous llama_state_seq_*_data_ext calls LLAMA_API struct llama_state_seq_copy * llama_state_seq_copy_init(struct llama_context * ctx); LLAMA_API void llama_state_seq_copy_free(struct llama_state_seq_copy * cpy); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0747d1c6366..d2e929a1429 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -13,6 +13,7 @@ #include "llama-sampler.h" #include "llama.h" +#include #include #include #include @@ -3549,6 +3550,24 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } + // A device that advertises events but does not implement event_query is no use + // here. ggml_backend_event_query() then answers the only way it can, by waiting for + // the event, so the first poll of a transfer blocks the caller for the whole copy -- + // the very stall this exists to remove, except that the caller has been told the + // copy is asynchronous and has stopped looking for it. Such a device is left out, so + // that state_seq_copy_init() returns NULL and the caller keeps the synchronous calls + // it already had. + if (!ggml_backend_dev_supports_event_query(dev)) { + static std::atomic warned(false); + + if (!warned.exchange(true)) { + LLAMA_LOG_INFO("%s: %s cannot test an event without waiting for it, so sequence " + "states are copied synchronously\n", __func__, ggml_backend_dev_name(dev)); + } + + continue; + } + // a backend of its own, not the one the graphs are computed on: that one moves its // copies to whichever stream it is currently using, so a transfer posted to it could // end up ordered behind a graph -- which is the stall this exists to avoid From 0160ea46748249c54be27bc6c0d4cd4a2421d4cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:36:42 +0000 Subject: [PATCH 06/25] server: wait for the parks still in flight before tearing the contexts 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. --- tools/server/server-context.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f15f78f8201..e31a96f427f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1272,6 +1272,23 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { + // [TAG_PREEMPT_ASYNC] the slots outlive this call -- they are declared after + // llama_init, so they are still there when it is reset here, and load_model() clears + // them only after the next context exists -- and any one of them may be holding a + // park or a resume that is still reading or writing KV tensors of the contexts about + // to be freed. release() makes the same wait for a single slot; this is the one that + // covers all of them, and on the sleeping-state path it is the only one there is, + // because the server carries on running afterwards. + for (auto & slot : slots) { + slot.preempt_copy_wait(); + + // the transfer holds a backend and an event of its own, and its host buffer is + // no use to the context that comes back: let go of both before that context's + // successor makes new ones + slot.preempt_cpy_tgt.reset(); + slot.preempt_cpy_dft.reset(); + } + spec.reset(); spec_init.reset(); From c0d92970dd208439ccd97fad226e20250c5c04c8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:36:57 +0000 Subject: [PATCH 07/25] server: charge a resume candidate its own lookahead before admitting it 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. --- tools/server/server-context.cpp | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e31a96f427f..1effa101417 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3226,12 +3226,19 @@ struct server_context_impl { // its copy lands, so it has to fire early enough that everything still decoding has // somewhere to put its tokens until then. 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. - int32_t preempt_n_margin() const { + // + // n_additional_running is for slots that are not running yet but are about to be: a + // resume candidate is still PREEMPTED while it is being considered, so it is not counted + // by the loop below, yet the moment it is admitted it starts decoding and needs the same + // runway as everybody else. Admitting it without charging it that runway is what the + // margin exists to prevent, and it showed up as a slot restored and parked again a few + // iterations later, over and over. + int32_t preempt_n_margin(int32_t n_additional_running = 0) const { if (!preempt_async_active()) { return PREEMPT_N_MARGIN; } - int32_t n_running = 0; + int32_t n_running = n_additional_running; for (const auto & slot : slots) { if (slot.is_processing() && !slot.preempt_is_out()) { @@ -3457,12 +3464,14 @@ struct server_context_impl { server_slot * best = nullptr; // Room for the sequence AND for the next step of everything already running, - // so that a resume cannot immediately trigger the preemption of someone else. + // and for the lookahead of the candidate itself, which is about to become one of + // them: a resume must not immediately trigger the preemption of someone else, or + // of itself. // A cached prompt on an idle slot is worth less than a conversation waiting to // continue, so give those cells up first - same call the KV-full path makes. for (;;) { for (auto * slot : parked) { - if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin() <= n_cells) { + if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin(1) <= n_cells) { best = slot; break; } From cd54f6089240237cce65c030d3555dc7848d91c4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:37:18 +0000 Subject: [PATCH 08/25] llama: report the host memory a transfer holds, not the kind it asked 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. --- include/llama.h | 8 +++++++- src/llama-context.cpp | 7 +++++++ tools/server/server-context.cpp | 36 +++++++++++++++++++++++++++++++-- 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/include/llama.h b/include/llama.h index 26db483e6c2..a5be59cb8f8 100644 --- a/include/llama.h +++ b/include/llama.h @@ -960,9 +960,15 @@ extern "C" { LLAMA_API size_t llama_state_seq_copy_buf_capacity(struct llama_state_seq_copy * cpy); LLAMA_API void llama_state_seq_copy_buf_free (struct llama_state_seq_copy * cpy); - // true when the buffer is page-locked, i.e. when the copies can really overlap + // true when the buffer that is held right now is page-locked, i.e. when the copies can + // really overlap. False while no buffer is held, since none is page-locked then: a + // caller asking before the first resize wants llama_state_seq_copy_buf_can_pin(). LLAMA_API bool llama_state_seq_copy_buf_is_pinned(struct llama_state_seq_copy * cpy); + // true when the backend offers pinned host memory at all. It is what the next resize + // will ask for, not what any buffer is: an allocation can still come back pageable. + LLAMA_API bool llama_state_seq_copy_buf_can_pin(struct llama_state_seq_copy * cpy); + // issue the copies; return the number of bytes covered, 0 on failure LLAMA_API size_t llama_state_seq_copy_get( struct llama_state_seq_copy * cpy, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index d2e929a1429..a807944b4b5 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -4751,6 +4751,13 @@ void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { } bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { + // what was allocated, not what could be: a host buffer type is free to hand back + // ordinary memory, which is what CUDA does under GGML_CUDA_NO_PINNED, and there is + // nothing page-locked before the first resize or after buf_free() + return cpy->pinned; +} + +bool llama_state_seq_copy_buf_can_pin(llama_state_seq_copy * cpy) { return cpy->can_pin; } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1effa101417..007f3298cb0 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1289,6 +1289,10 @@ struct server_context_impl { slot.preempt_cpy_dft.reset(); } + // the next context allocates its own host buffers, so say again what they turn out + // to be + preempt_ram_kind_logged = false; + spec.reset(); spec_init.reset(); @@ -1676,8 +1680,11 @@ struct server_context_impl { if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { if (preempt_async_ok) { - SRV_INF("preemption: parking and resuming asynchronously, %s host memory\n", - llama_state_seq_copy_buf_is_pinned(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + // no buffer has been allocated yet, so this is what the backend offers, + // not what is held. What was actually got is reported by the first park, + // because a host buffer type may still hand back ordinary memory. + SRV_INF("preemption: parking and resuming asynchronously, backend offers %s host memory\n", + llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); } else { SRV_WRN("%s", "preemption: this backend cannot copy asynchronously, parking and resuming synchronously\n"); } @@ -3153,6 +3160,27 @@ struct server_context_impl { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } + // [TAG_PREEMPT_ASYNC] whether the kind of host memory the parks actually got has been + // reported. It is only knowable once a buffer exists, and it is worth knowing: pinned + // memory is what lets the copies overlap, and a host buffer type is free to hand back + // ordinary memory instead of failing. + bool preempt_ram_kind_logged = false; + + void preempt_log_ram_kind(const server_slot & slot) { + if (preempt_ram_kind_logged || !slot.preempt_is_async()) { + return; + } + + if (llama_state_seq_copy_buf_capacity(slot.preempt_cpy_tgt.get()) == 0) { + return; // nothing held, so nothing to report yet + } + + preempt_ram_kind_logged = true; + + SRV_INF("preemption: parking into %s host memory\n", + llama_state_seq_copy_buf_is_pinned(slot.preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + } + // host RAM the parked sequences hold right now size_t preempt_ram_used() const { size_t res = 0; @@ -3543,6 +3571,8 @@ struct server_context_impl { slot.t_preempt_copy_us = ggml_time_us(); if (slot.preempt_save()) { + preempt_log_ram_kind(slot); + // [TAG_PREEMPT_ASYNC] a slot left PREEMPTING is counted by // update_preempt_copies() when its copy lands, not here if (slot.state == SLOT_STATE_PREEMPTED) { @@ -3595,6 +3625,8 @@ struct server_context_impl { break; // could not park it; the existing retry ladder is still behind us } + preempt_log_ram_kind(*victim); + // [TAG_PREEMPT_ASYNC] the copy has only been issued; the cells are still the // victim's until it lands, so nothing further can be decided about the pool this // iteration. update_preempt_copies() picks it up on the next one, and the step From 22c90bd9ffcebf70f66bb9f865711e330fa77a8c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:40:27 +0000 Subject: [PATCH 09/25] llama: check the size and the flags a sequence transfer is issued with 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. --- include/llama.h | 6 +- src/llama-context.cpp | 34 +++++++- tests/CMakeLists.txt | 5 ++ tests/test-state-seq-copy.cpp | 146 ++++++++++++++++++++++++++++++++++ 4 files changed, 188 insertions(+), 3 deletions(-) create mode 100644 tests/test-state-seq-copy.cpp diff --git a/include/llama.h b/include/llama.h index a5be59cb8f8..3c64888d25f 100644 --- a/include/llama.h +++ b/include/llama.h @@ -969,7 +969,11 @@ extern "C" { // will ask for, not what any buffer is: an allocation can still come back pageable. LLAMA_API bool llama_state_seq_copy_buf_can_pin(struct llama_state_seq_copy * cpy); - // issue the copies; return the number of bytes covered, 0 on failure + // Issue the copies; return the number of bytes covered, 0 on failure. size must be + // between 1 and llama_state_seq_copy_buf_size(): the buffer belongs to the transfer, and + // a size beyond it is refused rather than believed. LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is + // refused too, since these copies serialise through host memory; use + // llama_state_seq_get_data_ext / set_data_ext for that flag. LLAMA_API size_t llama_state_seq_copy_get( struct llama_state_seq_copy * cpy, size_t size, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a807944b4b5..09e53123da6 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3600,7 +3600,22 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { } size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { - if (!cpy.data) { + // Unlike the legacy API the library owns this buffer, so the extent the io object is + // built with can be checked instead of believed. Every bounds check inside that object + // validates against the extent it was given, so a size larger than the allocation makes + // all of them agree with the caller and the copy runs past the buffer. + if (!cpy.data || size == 0 || size > cpy.size) { + LLAMA_LOG_ERROR("%s: cannot cover %zu bytes, the transfer's buffer holds %zu\n", __func__, size, cpy.size); + return 0; + } + + // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to be left in device buffers, + // and this path has nowhere to leave it: it serialises through the host buffer it owns, + // which is the whole point of it. llama_state_seq_get_size_ext() with that flag reports + // a metadata-sized state, so a caller pairing the two would size a buffer for one thing + // and fill it with another; the synchronous calls serve that flag. + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + LLAMA_LOG_ERROR("%s: LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is not supported here, the copies go through host memory\n", __func__); return 0; } @@ -3629,7 +3644,22 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size } size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags) { - if (!cpy.data) { + // Unlike the legacy API the library owns this buffer, so the extent the io object is + // built with can be checked instead of believed. Every bounds check inside that object + // validates against the extent it was given, so a size larger than the allocation makes + // all of them agree with the caller and the copy runs past the buffer. + if (!cpy.data || size == 0 || size > cpy.size) { + LLAMA_LOG_ERROR("%s: cannot cover %zu bytes, the transfer's buffer holds %zu\n", __func__, size, cpy.size); + return 0; + } + + // LLAMA_STATE_SEQ_FLAGS_ON_DEVICE asks for the tensor data to be left in device buffers, + // and this path has nowhere to leave it: it serialises through the host buffer it owns, + // which is the whole point of it. llama_state_seq_get_size_ext() with that flag reports + // a metadata-sized state, so a caller pairing the two would size a buffer for one thing + // and fill it with another; the synchronous calls serve that flag. + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + LLAMA_LOG_ERROR("%s: LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is not supported here, the copies go through host memory\n", __func__); return 0; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78af..55f46d7c6b1 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -299,6 +299,11 @@ llama_build_and_test(test-backend-sampler.cpp LABEL "model") llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}") set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model) +# Guards on the asynchronous per-sequence state transfer +# Skips itself on a backend that cannot copy asynchronously +llama_build_and_test(test-state-seq-copy.cpp LABEL "model" ARGS -m "${MODEL_DEST}") +set_tests_properties(test-state-seq-copy PROPERTIES FIXTURES_REQUIRED test-download-model) + # Test state save/load functionality llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}") set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model) diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp new file mode 100644 index 00000000000..a4633954cee --- /dev/null +++ b/tests/test-state-seq-copy.cpp @@ -0,0 +1,146 @@ +// [TAG_STATE_ASYNC] guards on the asynchronous per-sequence state transfer +// +// llama_state_seq_copy_get / _set take a size and a flags word from the caller and hand both +// to an io object that validates everything else against them. The buffer belongs to the +// transfer, so a size larger than it is refused rather than believed, and +// LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is refused because these copies serialise through host +// memory. This also checks that the buffer reports itself as page-locked only while it holds +// memory that is. +// +// Skipped, not failed, on a backend that cannot copy asynchronously: there is no transfer to +// make and the synchronous calls are what a caller uses there. + +#include "arg.h" +#include "common.h" +#include "llama.h" + +#include +#include +#include + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + fprintf(stderr, "%s : FAILED at line %d: %s\n", __func__, \ + __LINE__, #cond); \ + return 1; \ + } \ + } while (0) + +int main(int argc, char ** argv) { + common_params params; + + params.sampling.seed = 1234; + params.kv_unified = true; + params.n_parallel = 2; + params.n_ctx = 256; + + common_init(); + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + ggml_backend_load_all(); + + common_init_result_ptr llama_init = common_init_from_params(params); + + llama_context * ctx = llama_init->context(); + + if (llama_init->model() == nullptr || ctx == nullptr) { + fprintf(stderr, "%s : failed to init\n", __func__); + return 1; + } + + // put something in the cache to copy: two sequences interleaved, so the cells of each + // are a comb rather than one block, which is what the transfer is built for + std::vector tokens(60, 1); + + llama_batch batch = llama_batch_init(params.n_parallel*tokens.size(), 0, 1); + for (size_t i = 0; i < tokens.size(); i++) { + for (int s = 0; s < params.n_parallel; ++s) { + common_batch_add(batch, tokens[i], i, {s}, false); + } + } + batch.logits[batch.n_tokens - 1] = true; + + if (llama_decode(ctx, batch)) { + fprintf(stderr, "%s : failed to decode\n", __func__); + llama_batch_free(batch); + return 1; + } + + llama_batch_free(batch); + + llama_state_seq_copy * cpy = llama_state_seq_copy_init(ctx); + + if (cpy == nullptr) { + fprintf(stderr, "%s : this backend cannot copy sequence states asynchronously, skipping\n", __func__); + return 0; + } + + const int seq_id = 1; + const size_t size = llama_state_seq_get_size_ext(ctx, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); + + CHECK(size > 0); + + // nothing is allocated yet, so nothing is page-locked yet, whatever the backend offers + CHECK(llama_state_seq_copy_buf_is_pinned(cpy) == false); + CHECK(llama_state_seq_copy_buf(cpy) == nullptr); + + CHECK(llama_state_seq_copy_buf_resize(cpy, size) != nullptr); + CHECK(llama_state_seq_copy_buf_size(cpy) == size); + + fprintf(stderr, "%s : seq %d state is %zu bytes, %s host memory (backend offers %s)\n", + __func__, seq_id, size, + llama_state_seq_copy_buf_is_pinned(cpy) ? "pinned" : "pageable", + llama_state_seq_copy_buf_can_pin(cpy) ? "pinned" : "pageable"); + + // a size beyond the buffer the transfer owns is refused, on both directions + CHECK(llama_state_seq_copy_get(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_set(cpy, size + 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + + // and so is an empty one, which cannot even hold the header + CHECK(llama_state_seq_copy_get(cpy, 0, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_set(cpy, 0, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + + // ON_DEVICE keeps the tensor data off the host, which is where these copies go + CHECK(llama_state_seq_copy_get(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) == 0); + CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) == 0); + + // none of that may have posted anything + CHECK(llama_state_seq_copy_done(cpy)); + + fprintf(stderr, "%s : oversized, empty and ON_DEVICE transfers are all refused\n", __func__); + + // the same call at the size the transfer does own still works, and round-trips + std::vector before(llama_state_seq_get_size(ctx, seq_id)); + CHECK(llama_state_seq_get_data(ctx, before.data(), before.size(), seq_id) == before.size()); + + CHECK(llama_state_seq_copy_get(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); + llama_state_seq_copy_wait(cpy); + + llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1); + + CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); + llama_state_seq_copy_wait(cpy); + + std::vector after(llama_state_seq_get_size(ctx, seq_id)); + CHECK(after.size() == before.size()); + CHECK(llama_state_seq_get_data(ctx, after.data(), after.size(), seq_id) == after.size()); + CHECK(before == after); + + fprintf(stderr, "%s : a transfer at the buffer's own size round-trips seq %d byte-for-byte\n", + __func__, seq_id); + + // giving the memory back leaves nothing page-locked to report + llama_state_seq_copy_buf_free(cpy); + CHECK(llama_state_seq_copy_buf_is_pinned(cpy) == false); + CHECK(llama_state_seq_copy_buf_capacity(cpy) == 0); + + llama_state_seq_copy_free(cpy); + + fprintf(stderr, "%s : SUCCESS\n", __func__); + + return 0; +} From e8f8b2fcf1c85ba16777c3ccd1a851f0fffcee1e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:41:15 +0000 Subject: [PATCH 10/25] ggml: stop collecting an error the CUDA event query never sets 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/src/ggml-cuda/ggml-cuda.cu | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 91902099011..1fca6352403 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5380,9 +5380,10 @@ static bool ggml_backend_cuda_device_event_query(ggml_backend_dev_t dev, ggml_ba const cudaError_t err = cudaEventQuery((cudaEvent_t)event->context); + // not an error, and nothing to clear: cudaEventQuery() returns cudaErrorNotReady + // without recording it as the thread's last error, so collecting one here would only + // consume somebody else's, and a real launch failure would be swallowed if (err == cudaErrorNotReady) { - // not an error: clear it so it is not reported against the next call - (void) cudaGetLastError(); return false; } From 86315eac6bc5a110a2df4e7f927b39dc7a9a6c25 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:41:15 +0000 Subject: [PATCH 11/25] ggml: bump the backend API version for the new device interface member 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. --- ggml/src/ggml-backend-impl.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 902b0963afa..bb3be31217b 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -8,7 +8,7 @@ extern "C" { #endif - #define GGML_BACKEND_API_VERSION 2 + #define GGML_BACKEND_API_VERSION 3 // // Backend buffer type From 888603d03578749219b190fd1dd0ed9554496fdb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 06:21:36 +0000 Subject: [PATCH 12/25] server: do not leave an issued async park holding the room the decode 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. --- tools/server/server-context.cpp | 33 ++++++++++++++ tools/server/tests/unit/test_preempt.py | 57 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 007f3298cb0..5509d4ea06c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3639,6 +3639,27 @@ struct server_context_impl { victim->preempt_state_size() / (1024.0 * 1024.0), preempt_kv_used(), n_cells, n_used, victim->n_preempt); + + // [TAG_PREEMPT_ASYNC] Whether we may leave now depends on which of the two + // thresholds we are under. + // + // Short of the lookahead only: there is still room for the step about to be + // built, the park is early by design, and leaving is the whole point -- the + // copy runs beside the decode and update_preempt_copies() collects it next + // iteration. + // + // Out of room for the step itself: the cells are held until the copy lands, + // so leaving now builds a batch into a pool that has not got smaller. The + // decode fails, and the retry ladder halves n_batch to 1 without ever + // polling the copy, ending in "Context size has been exceeded" for every + // slot. The synchronous path did not have this problem because it returned + // the cells before it returned. Go round instead: the next pass reaches + // preempt_wait_in_flight() and waits for the park just issued, which is no + // worse than the synchronous path and is what it was written for. + if (n_used + PREEMPT_N_MARGIN > n_cells) { + continue; + } + break; } @@ -4593,6 +4614,18 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] Before giving up any batch width: a park that has been + // issued and not yet landed is holding cells that are already spoken for, and + // waiting for it returns them. Halving the batch returns nothing, so without + // this the ladder can run all the way down to n_batch == 1 and end every + // request while the room it needed was moments from arriving. Safe from here + // because the slot was detached before this batch was built, so completing its + // park cannot change what the batch about to be retried contains. + if (ret == 1 && preempt_wait_in_flight()) { + SRV_WRN("%s", "waited for an in-flight park before retrying the decode\n"); + return false; // retry at the same batch size, with the cells it freed + } + // retry with half the batch size to try to find a free slot in the KV cache if (!try_clear_idle_slots()) { n_batch /= 2; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 1aba01751bd..4c3e8c465d8 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -447,3 +447,60 @@ def test_no_preempt_async_falls_back_to_the_synchronous_path(): # the synchronous path still parks and resumes assert text.count("preempted on request") >= 6 assert text.count("resumed after") >= 6 + + +def test_a_prompt_arriving_into_a_nearly_full_pool_parks_rather_than_ends_everything(): + # [TAG_PREEMPT_ASYNC] The case the async path made worse than the synchronous one, and + # that the existing tests miss because their victim holds almost no cells. + # + # Three slots are well into generating when a fourth request arrives whose prompt does + # not fit in what is left. update_preemption() picks a victim and issues its park, but + # an asynchronous park does not return the cells before update_slots() carries on. If + # the loop leaves at that point, the batch is built into a pool that has not got any + # smaller, llama_decode returns 1, and the retry ladder halves n_batch to 1 in + # microseconds without ever polling the copy -- ending every request with "Context size + # has been exceeded" while the room it wanted was one event query away. + # + # Pass is what the synchronous path gave: a park, and all four requests finish. + global server + server.n_ctx = 512 + server.n_gpu_layer = 99 + server.n_slots = 4 + server.start() + log = LogReader(server.log_path) + + prompt_a, n_a = _prompt_of_about(100, "Alpha") + prompt_b, n_b = _prompt_of_about(100, "Bravo") + prompt_c, n_c = _prompt_of_about(100, "Charlie") + prompt_d, n_d = _prompt_of_about(150, "Delta") + + # A, B and C oversubscribe the pool between them, so the pressure does not depend on + # when D arrives, and every occupant is holding real cells rather than the handful the + # other tests park. Each of the four still fits on its own. + n_predict_abc = 130 + n_predict_d = 40 + assert max(n_a, n_b, n_c) + n_predict_abc < 512 and n_d + n_predict_d < 512 + assert n_a + n_b + n_c + 3 * n_predict_abc > 512 + + def _late(n_predict, prompt): + # D's prompt arrives into a pool the other three have already grown into; this + # model decodes about 120 tokens a second, so they are all still running + time.sleep(0.25) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_abc, prompt_a)), + (_complete, (n_predict_abc, prompt_b)), + (_complete, (n_predict_abc, prompt_c)), + (_late, (n_predict_d, prompt_d)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted" in text + + for i, res in enumerate(results): + assert res.status_code == 200, (i, res.body) + for i in range(3): + assert results[i].body["timings"]["predicted_n"] == n_predict_abc + assert results[3].body["timings"]["predicted_n"] == n_predict_d From cd1cd4e6d37474b9ee6157179e44e351f2aeed57 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 16:36:25 +0000 Subject: [PATCH 13/25] server: count a slot being restored in the asynchronous lookahead margin 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. --- tools/server/server-context.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index cea3665774d..0395a3e22f6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3337,7 +3337,9 @@ struct server_context_impl { int32_t n_running = n_additional_running; for (const auto & slot : slots) { - if (slot.is_processing() && !slot.preempt_is_out()) { + // A slot on its way back in already holds its cells and starts decoding the + // moment its copy lands, so it needs the runway now; one on its way out does not. + if (slot.is_processing() && (!slot.preempt_is_out() || slot.state == SLOT_STATE_RESTORING)) { n_running++; } } From 42e536970c7d2348f0e7cb8a7fe67bbefe362e11 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 18:33:13 +0000 Subject: [PATCH 14/25] server: a round with a context shift waits for every park and restore copy first A context shift is recorded by pre_decode_shift() and applied inside the next llama_decode as one graph over the whole K cache, in place. A restore still writing its cells on its own stream could be read half done and written back stale, and a park still reading its cells would read through the rewrite. The round that recorded a shift now waits for every copy in flight before it decodes; shifts are rare, so the wait is too. The rotation park also stamps its issue time, so its completion line measures the copy rather than the previous park. --- tools/server/server-context.cpp | 48 +++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8c77bac7adc..13ca47798e9 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3227,6 +3227,10 @@ struct server_context_impl { // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop bool preempt_batch_abandoned = false; + // [TAG_PREEMPT_ASYNC] a context shift was recorded this round; it is applied inside the + // next llama_decode as one graph over the whole K cache, in place + bool preempt_shift_pending = false; + int32_t preempt_n_spec_max() const { return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; } @@ -3537,6 +3541,42 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] wait for every copy in flight, parks and restores alike. The + // context shift a slot recorded this round is applied inside the next llama_decode as one + // graph over the whole K cache, in place: a restore still writing its cells on its own + // stream could be read half done and written back stale, and a park still reading its + // cells would read through the rewrite. Shifts are rare, so this round waits. + void preempt_wait_for_shift() { + if (!preempt_shift_pending) { + return; + } + + preempt_shift_pending = false; + + while (preempt_wait_in_flight()) { + } + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_RESTORING) { + continue; + } + + slot.preempt_copy_wait(); + + if (!slot.preempt_restore_poll()) { + continue; + } + + metrics.n_resume++; + + SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + } + // Wait for one outstanding park, the last thing tried before giving up on finding room. // It is what keeps a pool that fills faster than the copies drain no worse than the // synchronous path: the decode waits for the copy exactly as it used to. @@ -3682,10 +3722,14 @@ struct server_context_impl { continue; } + const int64_t t_start = ggml_time_us(); + if (!preempt_fits_budget(slot) || !slot.preempt_save()) { continue; } + slot.t_preempt_copy_us = t_start; + // [TAG_PREEMPT_ASYNC] an asynchronous park is counted when its copy // lands, and the head is re-examined on the pass that sees the room if (slot.state != SLOT_STATE_PREEMPTING) { @@ -3956,6 +4000,9 @@ struct server_context_impl { llama_batch batch_view; int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); + + preempt_wait_for_shift(); + for (int32_t off = 0; off < batch.size(); off = off_next) { const int32_t n_tokens = std::min(n_batch, batch.size() - off); try { @@ -4047,6 +4094,7 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); slot.n_ctx_shift++; + preempt_shift_pending = true; slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); From 02a3e11d43efaa6b94101c6c16ba2e9af731238c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 19:57:19 +0000 Subject: [PATCH 15/25] server: the abort sweep also leaves a slot whose copy is in flight alone --- tools/server/server-context.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d4961788d93..88b6d7f7e20 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3179,7 +3179,8 @@ struct server_context_impl { // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in // host RAM, not in the cache, and it comes back when there is room, the same as // in the decode error sweep - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + // [TAG_PREEMPT_ASYNC] a slot whose copy is in flight is out of the round as well + if (slot.is_processing() && !slot.preempt_is_out()) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); } From 5a13c675840827b6e3902ac43e01b77953b8e606 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 20:58:26 +0000 Subject: [PATCH 16/25] server: asynchronous copies and the rotation, the budget and the cache-reuse shift No rotation runs while a park is still copying: its cells are still held, the head would not fit yet, and the rotation would only park another resident on top. An asynchronous rotation park is not re-examined on the same pass either; the head is re-examined when the copy lands. The rotation's budget no longer counts an asynchronous head's bytes as leaving, since 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, so a buffer held by a running slot cannot keep every other slot from being parked. The cache-reuse shift is applied by the same in-place graph as a context shift, so it sets the flag that makes the round wait for copies in flight. --- tools/server/server-context.cpp | 53 +++++++++++++++++++++++++++++---- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 88b6d7f7e20..7a13443b4d5 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3322,16 +3322,38 @@ struct server_context_impl { // resident, so its bytes are on their way out and are not held against the resident. // A budget that holds one sequence but not two would otherwise refuse every rotation // and leave the head parked for as long as the resident cares to generate. + // [TAG_PREEMPT_ASYNC] an asynchronous head keeps its pinned buffer through the restore + // (see preempt_state_size), so nothing of it leaves; what the resident already holds is + // reused, as in preempt_fits_budget, and only the rest is charged. bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { if (params_base.preempt_ram_mib < 0) { return true; } + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + const size_t used = preempt_ram_used(); + const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); + const size_t held = slot.preempt_state_size(); + const size_t need = slot.preempt_state_required(); + const size_t extra = need > held ? need - held : 0; + + return used - leaving + extra <= budget; + } + + // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer for its next park, which is + // worth it while the budget has room for it and not otherwise: over budget, a buffer held + // by a slot that is running would keep every other slot from being parked at all + void preempt_trim_ram(server_slot & slot) { + if (params_base.preempt_ram_mib < 0) { + return; + } + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - const size_t used = preempt_ram_used(); - const size_t leaving = std::min(used, head.preempt_state_size()); - return used - leaving + slot.preempt_state_required() <= budget; + if (preempt_ram_used() > budget && slot.preempt_state_size() > 0) { + SLT_INF(slot, "%.1f MiB of parked RAM returned: the pool is over its budget\n", slot.preempt_state_size() / (1024.0 * 1024.0)); + slot.preempt_state_free(); + } } // cells the slot will ask for on its next step once it is back in the pool @@ -3634,6 +3656,8 @@ struct server_context_impl { if (slot.preempt_restore_poll()) { metrics.n_resume++; + preempt_trim_ram(slot); + SLT_WRN(slot, "restore completed after %.2f ms: %d tokens back in the cache, kv %d/%d, preemptions %d\n", (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, slot.prompt.n_tokens(), @@ -3672,6 +3696,8 @@ struct server_context_impl { metrics.n_resume++; + preempt_trim_ram(slot); + SLT_WRN(slot, "restore completed after %.2f ms (waited for, a context shift is due): %d tokens back in the cache, kv %d/%d, preemptions %d\n", (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, slot.prompt.n_tokens(), @@ -3815,7 +3841,15 @@ struct server_context_impl { if (!best) { server_slot * head = parked.front(); - if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { + // [TAG_PREEMPT_ASYNC] a park still copying holds its cells, so the head would + // not fit yet and a rotation now would only park another resident on top + bool parking = false; + + for (const auto & slot : slots) { + parking = parking || slot.state == SLOT_STATE_PREEMPTING; + } + + if (!parking && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { // the resident whose cells let the head in, the smallest of those; failing // one that does so alone, the largest, since it makes the most room. Taking // the first shifting resident in slot order could park one too small to @@ -3870,7 +3904,12 @@ struct server_context_impl { pick_enough ? "" : " (not enough room by itself)", slot.n_preempt); - best = head; // re-examined by the loop, which sees the room it just got + // [TAG_PREEMPT_ASYNC] a synchronous park has released its cells, so the + // head is re-examined now; an asynchronous one has not, and the head is + // re-examined on the pass that sees the copy land + if (slot.state == SLOT_STATE_PREEMPTED) { + best = head; + } } } @@ -4548,6 +4587,10 @@ struct server_context_impl { slot.mem.seq_rm (slot.id, head_p, head_c); slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); + // [TAG_PREEMPT_ASYNC] applied inside the next llama_decode by the + // same in-place graph as a context shift, see preempt_wait_for_shift + preempt_shift_pending = true; + for (size_t i = 0; i < n_match; i++) { slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]); n_past++; From 82f40df769c8153b5d2c602d973267c5b2807805 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:35:50 +0000 Subject: [PATCH 17/25] preempt: post no copies for a transfer that failed part way, park synchronously into pageable memory, offer no transfer for state that is not on a device The asynchronous state adapters posted their queued copies from the destructor whether or not serialisation had got to the end, so a buffer one byte short made llama_state_seq_copy_get() return 0 while 64 copies were still reading it, and the restore counterpart wrote into cells the failed restore had already given up. Both adapters commit only after the serialisation succeeds; a failed one posts nothing. A host buffer type may hand back pageable memory instead of failing, and a copy into or out of pageable memory blocks the thread that issued it. The server takes a one MiB buffer at load and looks at what it got; if it is pageable the slots park synchronously and say so. state_seq_copy_init() returned a transfer whenever a device could copy asynchronously, even when every state tensor lived in host memory (most layers on the CPU) and every copy took the synchronous branch. It returns NULL unless every non-empty memory buffer lives on one of its devices. --- src/llama-context.cpp | 69 ++++++++++++++++++++++++++++++++- tools/server/server-context.cpp | 28 ++++++++++--- 2 files changed, 90 insertions(+), 7 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 09e53123da6..c73ba5ddacf 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3332,7 +3332,19 @@ class llama_io_write_host_async : public llama_io_write_i { llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : ptr(p), buf_size(len), cpy(cpy) {} + // The transfers are posted from the destructor, and only once serialisation has got to + // the end: a failure part way, a buffer one byte short say, is reported to the caller as + // a zero return, and a caller told that is free to reuse the buffer at once. Copies + // posted regardless would still be reading it. + void commit() { + committed = true; + } + ~llama_io_write_host_async() { + if (!committed) { + return; + } + llama_io_emit(winfos, 0, winfos.size(), [this](ggml_tensor * tensor, uint8_t * ptr, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { @@ -3385,6 +3397,8 @@ class llama_io_write_host_async : public llama_io_write_i { std::vector winfos; llama_state_seq_copy & cpy; + + bool committed = false; }; class llama_io_read_host_async : public llama_io_read_i { @@ -3392,7 +3406,18 @@ class llama_io_read_host_async : public llama_io_read_i { llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : ptr(p), buf_size(len), cpy(cpy) {} + // see llama_io_write_host_async::commit(): the restore that failed part way has already + // dropped the sequence, and copies posted for it would write into cells that are no + // longer its own + void commit() { + committed = true; + } + ~llama_io_read_host_async() { + if (!committed) { + return; + } + // No whole-tensor staging here, unlike the synchronous path above. Staging reads a // tensor, patches this sequence's bytes into the host copy and writes the whole // tensor back, which preserves the neighbours only while nothing else is touching @@ -3454,6 +3479,8 @@ class llama_io_read_host_async : public llama_io_read_i { std::vector rinfos; llama_state_seq_copy & cpy; + + bool committed = false; }; static constexpr uint32_t io_magic = 0xaf143cd8; @@ -3594,6 +3621,36 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { return nullptr; } + // The devices above are the ones the graphs run on, not necessarily the ones the state + // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor + // there takes the synchronous branch of backend_for(). A transfer whose every copy would + // do that is not asynchronous, whatever it is called, and the caller is better served by + // the synchronous calls it already has and a log line that says so. + if (memory) { + bool on_device = false; + + for (const auto & [buft, size] : memory->memory_breakdown()) { + if (size == 0) { + continue; + } + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + + if (ggml_backend_buft_is_host(buft) || !dev || buft != ggml_backend_dev_buffer_type(dev) || + cpy->devs.find(dev) == cpy->devs.end()) { + LLAMA_LOG_INFO("%s: the sequence state is not all in device memory (%s), so it is copied synchronously\n", + __func__, ggml_backend_buft_name(buft)); + return nullptr; + } + + on_device = true; + } + + if (!on_device) { + return nullptr; + } + } + cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); return cpy.release(); @@ -3636,7 +3693,11 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size io.write(&io_magic, sizeof(io_magic)); io.write(&seq_id, sizeof(seq_id)); - return state_seq_write_data(io, seq_id, flags); + const size_t n = state_seq_write_data(io, seq_id, flags); + + io.commit(); + + return n; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; @@ -3681,7 +3742,11 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size llama_seq_id seq_id_read; io.read(&seq_id_read, sizeof(seq_id_read)); - return state_seq_read_data(io, seq_id, flags); + const size_t n = state_seq_read_data(io, seq_id, flags); + + io.commit(); + + return n; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); return 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a0d1564405c..a4b40961842 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1725,11 +1725,29 @@ struct server_context_impl { if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { if (preempt_async_ok) { - // no buffer has been allocated yet, so this is what the backend offers, - // not what is held. What was actually got is reported by the first park, - // because a host buffer type may still hand back ordinary memory. - SRV_INF("preemption: parking and resuming asynchronously, backend offers %s host memory\n", - llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()) ? "pinned" : "pageable"); + // Pinned host memory is what lets a copy run beside the decode: one into or + // out of pageable memory is staged by the driver and blocks the thread that + // issued it, which is the stall the asynchronous path exists to remove. A + // host buffer type is free to hand back ordinary memory instead of failing + // (GGML_CUDA_NO_PINNED, or a pinning limit), and that is only knowable from + // a buffer, so a small one is taken and looked at before the first park. + bool pinned = llama_state_seq_copy_buf_can_pin(slots[0].preempt_cpy_tgt.get()); + + if (pinned) { + auto * cpy = slots[0].preempt_cpy_tgt.get(); + + pinned = llama_state_seq_copy_buf_resize(cpy, 1u << 20) != nullptr && + llama_state_seq_copy_buf_is_pinned(cpy); + + llama_state_seq_copy_buf_free(cpy); + } + + if (pinned) { + SRV_INF("%s", "preemption: parking and resuming asynchronously through pinned host memory\n"); + } else { + SRV_WRN("%s", "preemption: the host memory on offer is pageable, so a copy would block the decode; parking and resuming synchronously\n"); + preempt_async_ok = false; + } } else { SRV_WRN("%s", "preemption: this backend cannot copy asynchronously, parking and resuming synchronously\n"); } From 2b4a6912299776ad5e335b2f22adaaf20f0ef823 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 21:36:26 +0000 Subject: [PATCH 18/25] tests: a state transfer that fails one byte short posts no copies, on both sides --- tests/test-state-seq-copy.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp index a4633954cee..dabc4db50c3 100644 --- a/tests/test-state-seq-copy.cpp +++ b/tests/test-state-seq-copy.cpp @@ -113,6 +113,14 @@ int main(int argc, char ** argv) { fprintf(stderr, "%s : oversized, empty and ON_DEVICE transfers are all refused\n", __func__); + // a transfer that fails part way, one byte short of the state, must post nothing: the + // caller is told it failed and is free to reuse the buffer at once + CHECK(llama_state_seq_copy_get(cpy, size - 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_n_copies(cpy) == 0); + CHECK(llama_state_seq_copy_done(cpy)); + + fprintf(stderr, "%s : a transfer one byte short is refused and posts no copies\n", __func__); + // the same call at the size the transfer does own still works, and round-trips std::vector before(llama_state_seq_get_size(ctx, seq_id)); CHECK(llama_state_seq_get_data(ctx, before.data(), before.size(), seq_id) == before.size()); @@ -122,6 +130,11 @@ int main(int argc, char ** argv) { llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1); + // the restore side too: a buffer claimed one byte short is refused before a copy is posted + CHECK(llama_state_seq_copy_set(cpy, size - 1, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == 0); + CHECK(llama_state_seq_copy_n_copies(cpy) == 0); + CHECK(llama_state_seq_copy_done(cpy)); + CHECK(llama_state_seq_copy_set(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); llama_state_seq_copy_wait(cpy); From 9e827b53d98aad78ac65994c4a4340039e623faf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 23:26:59 +0000 Subject: [PATCH 19/25] preempt: a slot whose park buffer comes back pageable parks synchronously from then on The load-time probe takes one MiB and looks at it, but a buffer many times larger can still come back pageable (a host-locking limit, say), and a copy into pageable memory blocks the thread that issued it. The park now looks at the buffer it actually got: a pageable one is given back with the slot's transfers, and the slot takes the synchronous path for this park and every later one. --- tools/server/server-context.cpp | 55 +++++++++++++++++++++------------ 1 file changed, 36 insertions(+), 19 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a4b40961842..2d79a259684 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -545,30 +545,47 @@ struct server_slot { return false; } - if (llama_state_seq_copy_get(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { - SLT_ERR(*this, "%s", "failed to issue the copy of the target sequence out of the KV cache\n"); - preempt_state_free(); - return false; - } + // [TAG_PREEMPT_ASYNC] the load-time probe saw pinned memory, but a buffer this much + // larger can still come back pageable (a host-locking limit, say): the host buffer + // type hands back ordinary memory rather than failing, and a copy into pageable + // memory blocks the thread that issued it, which is the stall this path exists to + // remove. Such a slot parks synchronously from now on: its transfers are given + // back and the plain path below takes over, for this park and every later one. + const bool pageable = !llama_state_seq_copy_buf_is_pinned(preempt_cpy_tgt.get()) || + (size_dft > 0 && !llama_state_seq_copy_buf_is_pinned(preempt_cpy_dft.get())); + + if (pageable) { + SLT_WRN(*this, "the host memory for a %.3f MiB park is pageable, so this slot parks synchronously from now on\n", + (size_tgt + size_dft) / (1024.0 * 1024.0)); - if (size_dft > 0 && - llama_state_seq_copy_get(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { - SLT_ERR(*this, "%s", "failed to issue the copy of the draft sequence out of the KV cache\n"); - preempt_state_free(); - return false; - } + preempt_cpy_tgt.reset(); + preempt_cpy_dft.reset(); + } else { + if (llama_state_seq_copy_get(preempt_cpy_tgt.get(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to issue the copy of the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_copy_get(preempt_cpy_dft.get(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to issue the copy of the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } - preempt_detach(); + preempt_detach(); - // note: no mem.seq_rm() here. The copy is still reading these cells, so they are - // released in preempt_save_poll() once it has finished with them. - state_before_preempt = state; - state = SLOT_STATE_PREEMPTING; - t_preempt_us = ggml_time_us(); + // note: no mem.seq_rm() here. The copy is still reading these cells, so they are + // released in preempt_save_poll() once it has finished with them. + state_before_preempt = state; + state = SLOT_STATE_PREEMPTING; + t_preempt_us = ggml_time_us(); - n_preempt++; + n_preempt++; - return true; + return true; + } } try { From a4b62f2be5dea11c887e24ea491e10c0d85d5c36 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 23:36:16 +0000 Subject: [PATCH 20/25] preempt: the graphs that follow a restore wait for its copies, on the device A restore writes cells of the KV cache on the copy stream while the other sequences keep decoding on the compute stream, and an attention that is not paged reads every cell up to n_kv, masked ones included, so the reads and the writes were unordered. After the copies are posted and their events recorded, every backend the graphs run on waits for the event of the transfer on its device: a stream wait, so the thread that issued the restore carries on and the next decode starts the moment the copy lands. A park needs none of this: it reads cells nobody writes until it has landed, behind the synchronize at the top of the issue. --- src/llama-context.cpp | 62 ++++++++++++++++++++++++++++++++----------- 1 file changed, 47 insertions(+), 15 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c73ba5ddacf..34d2f51399a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3221,6 +3221,30 @@ struct llama_state_seq_copy { } } + // Order the context's compute behind the copies just recorded, on the device: every + // backend the graphs run on waits for the event of the transfer on its device before + // the next graph it is given. This is a stream wait, not a host wait, so the caller's + // thread carries on and the decode it issues next starts the moment the copy lands. + // + // Needed for a restore and only a restore: its copies write cells of the KV cache + // while other sequences keep decoding, and an attention that is not paged reads every + // cell up to n_kv, masked ones included, so without this the reads and the writes are + // unordered. A park reads cells nobody writes until it has landed, and the decode that + // produced them has been drained by the synchronize() at the top of the issue. + void order_before(const std::vector & compute) { + for (auto & it : devs) { + if (!it.second.pending) { + continue; + } + + for (const auto & backend : compute) { + if (ggml_backend_get_device(backend.get()) == it.first) { + ggml_backend_event_wait(backend.get(), it.second.event); + } + } + } + } + bool done() { bool res = true; @@ -3730,27 +3754,35 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size cpy.n_copies = 0; - llama_io_read_host_async io(cpy.data, size, cpy); + size_t n = 0; - try { - uint32_t magic_read; - io.read(&magic_read, sizeof(magic_read)); - if (io_magic != magic_read) { - throw std::runtime_error("wrong sequence state magic"); - } + { + llama_io_read_host_async io(cpy.data, size, cpy); - llama_seq_id seq_id_read; - io.read(&seq_id_read, sizeof(seq_id_read)); + try { + uint32_t magic_read; + io.read(&magic_read, sizeof(magic_read)); + if (io_magic != magic_read) { + throw std::runtime_error("wrong sequence state magic"); + } - const size_t n = state_seq_read_data(io, seq_id, flags); + llama_seq_id seq_id_read; + io.read(&seq_id_read, sizeof(seq_id_read)); - io.commit(); + n = state_seq_read_data(io, seq_id, flags); - return n; - } catch (const std::exception & err) { - LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); - return 0; + io.commit(); + } catch (const std::exception & err) { + LLAMA_LOG_ERROR("%s: error loading state: %s\n", __func__, err.what()); + return 0; + } } + + // the adapter has posted the copies and recorded the events on its way out; the + // graphs that follow on these devices wait for them, see order_before() + cpy.order_before(backends); + + return n; } bool llama_context::state_load_file(const char * filepath, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { From 294d2a9120a7e63c1188be2eeca2920b3cae4c38 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 00:23:45 +0000 Subject: [PATCH 21/25] preempt: idle parked RAM is given back when another slot needs to park; the copies wait for the compute stream on the device instead of draining it A restored slot keeps its pinned buffer for its next park, and that capacity was charged against --preempt-ram while it held nothing, so a budget that held one sequence was spent for good by the first restore: every later park was refused, and once the slot holding the buffer was the leader nothing could be parked at all. Both budget checks now give idle buffers back before deciding: largest first, never a buffer that still holds a parked sequence or has a copy in flight, never the candidate's own, and never the head of a rotation. state_seq_copy_get and state_seq_copy_set drained the host with synchronize() so that the copies were ordered behind the compute already queued. With the wait that order_before() puts on the compute stream for the previous restore, that drain blocked the calling thread until the previous copy had landed, and two restores issued in one pass ran one after the other with the whole transfer back on the decode loop. Each copy device now carries a second event, recorded on the compute stream and waited for on the copy stream, so the ordering is on the device and the host drains nothing. New server test: two sequences parked in turn under a budget that holds only one, both preempted, no context error, the idle-return line logged. --- src/llama-context.cpp | 48 ++++++++++++++++--- tools/server/server-context.cpp | 63 +++++++++++++++++++++++-- tools/server/tests/unit/test_preempt.py | 37 +++++++++++++++ 3 files changed, 137 insertions(+), 11 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 34d2f51399a..20e548220db 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3153,6 +3153,8 @@ struct llama_state_seq_copy { struct dev_copy { ggml_backend_ptr backend; ggml_backend_event_t event = nullptr; + // recorded on the compute stream and waited for on the copy stream, see order_after() + ggml_backend_event_t fence = nullptr; bool pending = false; }; @@ -3178,6 +3180,9 @@ struct llama_state_seq_copy { if (it.second.event) { ggml_backend_event_free(it.second.event); } + if (it.second.fence) { + ggml_backend_event_free(it.second.fence); + } } } @@ -3221,6 +3226,25 @@ struct llama_state_seq_copy { } } + // Order the copies about to be posted behind the compute already queued on each device: + // the decode that produced the cells a park reads, or that a restore's cells were + // carved out of, has to be finished before the copy touches them. Recorded on the + // compute backend's stream and waited for on the copy stream, so the host drains + // nothing. Draining it (synchronize()) is what this replaces: with the wait that + // order_before() queues on the compute stream for the previous restore, a host drain + // blocked this thread until that copy had landed, and two restores issued in one pass + // ran one after the other with the whole transfer back on the decode loop. + void order_after(const std::vector & compute) { + for (auto & it : devs) { + for (const auto & backend : compute) { + if (ggml_backend_get_device(backend.get()) == it.first) { + ggml_backend_event_record(it.second.fence, backend.get()); + ggml_backend_event_wait(it.second.backend.get(), it.second.fence); + } + } + } + } + // Order the context's compute behind the copies just recorded, on the device: every // backend the graphs run on waits for the event of the transfer on its device before // the next graph it is given. This is a stream wait, not a host wait, so the caller's @@ -3635,10 +3659,19 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } + ggml_backend_event_t fence = ggml_backend_event_new(dev); + + if (!fence) { + ggml_backend_event_free(event); + ggml_backend_free(backend_cpy); + continue; + } + auto & dc = cpy->devs[dev]; dc.backend.reset(backend_cpy); dc.event = event; + dc.fence = fence; } if (cpy->devs.empty()) { @@ -3700,13 +3733,11 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size return 0; } - // The copies run on their own stream and are ordered against nothing, so the decode that - // produced these cells has to be finished before they are read. This is the one part of - // the transfer that stays on the caller's thread, and it costs nothing where it is used: - // a caller preempting a sequence does it between two decodes, with the previous one - // already drained by the sampling that followed it. + // The copies run on their own stream, so the decode that produced these cells has to be + // finished before they are read: the copy stream waits for the compute stream, on the + // device, see order_after(). Nothing stays on the caller's thread. const int64_t t_sync = ggml_time_us(); - synchronize(); + cpy.order_after(backends); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; @@ -3748,8 +3779,11 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size return 0; } + // the cells this restore was given may still be read by a graph in flight (masked, but + // read), so the copy stream waits for the compute stream before it writes them: on the + // device, see order_after(), rather than by draining the compute stream on this thread const int64_t t_sync = ggml_time_us(); - synchronize(); + cpy.order_after(backends); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 2d79a259684..3745762a725 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3337,7 +3337,51 @@ struct server_context_impl { } // whether parking this slot stays under --preempt-ram - bool preempt_fits_budget(const server_slot & slot) const { + // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer for its next park, and + // that capacity counts against the budget while it holds no state. When a park does + // not fit, that idle capacity is what to give back first: largest first, never a + // buffer that still holds a parked sequence or has a copy in flight, and never the + // candidate's own, which it reuses. Without this a budget that holds one sequence was + // spent for good by the first restore: every later park was refused, and once the + // slot holding the buffer was the leader nothing could be parked at all. + void preempt_reclaim_idle_ram(size_t budget, size_t extra, const server_slot & keep) { + for (;;) { + if (preempt_ram_used() + extra <= budget) { + return; + } + + server_slot * best = nullptr; + + for (auto & other : slots) { + if (&other == &keep) { + continue; + } + + if (other.state == SLOT_STATE_PREEMPTED || other.state == SLOT_STATE_PREEMPTING || other.state == SLOT_STATE_RESTORING) { + continue; + } + + if (other.preempt_state_size() == 0) { + continue; + } + + if (!best || other.preempt_state_size() > best->preempt_state_size()) { + best = &other; + } + } + + if (!best) { + return; + } + + SLT_INF(*best, "%.1f MiB of idle parked RAM returned so that another slot can park\n", + best->preempt_state_size() / (1024.0 * 1024.0)); + + best->preempt_state_free(); + } + } + + bool preempt_fits_budget(const server_slot & slot) { if (params_base.preempt_ram_mib < 0) { return true; } @@ -3350,6 +3394,8 @@ struct server_context_impl { const size_t need = slot.preempt_state_required(); const size_t extra = need > held ? need - held : 0; + preempt_reclaim_idle_ram(budget, extra, slot); + return preempt_ram_used() + extra <= budget; } @@ -3360,18 +3406,27 @@ struct server_context_impl { // [TAG_PREEMPT_ASYNC] an asynchronous head keeps its pinned buffer through the restore // (see preempt_state_size), so nothing of it leaves; what the resident already holds is // reused, as in preempt_fits_budget, and only the rest is charged. - bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const { + bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) { if (params_base.preempt_ram_mib < 0) { return true; } const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; - const size_t used = preempt_ram_used(); - const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); const size_t held = slot.preempt_state_size(); const size_t need = slot.preempt_state_required(); const size_t extra = need > held ? need - held : 0; + // the head is parked, so it is never among the idle buffers given back here + { + const size_t used = preempt_ram_used(); + const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); + + preempt_reclaim_idle_ram(budget + leaving, extra, slot); + } + + const size_t used = preempt_ram_used(); + const size_t leaving = head.preempt_is_async() ? 0 : std::min(used, head.preempt_state_size()); + return used - leaving + extra <= budget; } diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 26c6e7579ac..de9da44c4ed 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -719,3 +719,40 @@ def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_ after = _complete(8) assert after.status_code == 200 assert after.body["timings"]["predicted_n"] == 8 + + +def test_a_restored_slot_gives_its_idle_buffer_back_when_another_slot_needs_to_park(): + # Under a finite --preempt-ram an asynchronous slot keeps its pinned buffer after a + # restore, for its next park, and that idle capacity counted against the budget. With + # a budget that holds one sequence, the first restore spent it for good: every later + # park of the other slot was refused. The idle buffer is given back when another slot + # needs the room, and both slots go on being parked. + global server + # a pool of 8192 cells, but the model's own window is 2048, so each generation stays + # under that; 1800 tokens of this model's state is about 1.1 MiB, so a budget of + # 2 MiB holds one sequence and not two + server.n_ctx = 8192 + server.n_gpu_layer = 99 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "256" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" + text = _start_async() + _require_async(text) + log = LogReader(server.log_path) + + n_predict = 1800 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "idle parked RAM returned" in text, "the idle buffer of a restored slot was never given back" + import re + parked = re.findall(r"id\s+(\d+) \| task \d+ \| preempted on request", text) + assert {"0", "1"} <= set(parked), f"only slots {sorted(set(parked))} were ever parked" + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False From 012ef75479ea0fd9d9c9e6032cd32794a53b5767 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:07:14 +0000 Subject: [PATCH 22/25] llama: the copies wait for a fence the context records after every decode, 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. --- src/llama-context.cpp | 74 ++++++++++++++++++++++++++++--------------- src/llama-context.h | 8 +++++ 2 files changed, 57 insertions(+), 25 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 20e548220db..21776af9afd 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,6 +483,10 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + for (auto & it : state_copy_fences) { + ggml_backend_event_free(it.second); + } + if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; @@ -1578,6 +1582,10 @@ int llama_context::encode(const llama_batch & batch_inp) { } } + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -2023,6 +2031,10 @@ int llama_context::decode(const llama_batch & batch_inp) { // wait for the computation to finish (automatically done when obtaining the model output) //synchronize(); + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -3153,8 +3165,6 @@ struct llama_state_seq_copy { struct dev_copy { ggml_backend_ptr backend; ggml_backend_event_t event = nullptr; - // recorded on the compute stream and waited for on the copy stream, see order_after() - ggml_backend_event_t fence = nullptr; bool pending = false; }; @@ -3180,9 +3190,6 @@ struct llama_state_seq_copy { if (it.second.event) { ggml_backend_event_free(it.second.event); } - if (it.second.fence) { - ggml_backend_event_free(it.second.fence); - } } } @@ -3228,19 +3235,20 @@ struct llama_state_seq_copy { // Order the copies about to be posted behind the compute already queued on each device: // the decode that produced the cells a park reads, or that a restore's cells were - // carved out of, has to be finished before the copy touches them. Recorded on the - // compute backend's stream and waited for on the copy stream, so the host drains - // nothing. Draining it (synchronize()) is what this replaces: with the wait that - // order_before() queues on the compute stream for the previous restore, a host drain - // blocked this thread until that copy had landed, and two restores issued in one pass - // ran one after the other with the whole transfer back on the decode loop. - void order_after(const std::vector & compute) { + // carved out of, has to be finished before the copy touches them. The copy stream waits + // for the context's fence on its device, an event the context records on the compute + // stream at the end of every decode, so the host drains nothing. The fence is recorded + // there and not here: recorded here, it would land behind the waits that order_before() + // queued for the restores issued earlier in the same pass, and each restore would then + // wait for the previous one's copies. Draining the host (synchronize()) is what this + // replaced: with those same waits on the compute stream, a host drain blocked this thread + // until the previous restore had landed. + void order_after(const std::map & fences) { for (auto & it : devs) { - for (const auto & backend : compute) { - if (ggml_backend_get_device(backend.get()) == it.first) { - ggml_backend_event_record(it.second.fence, backend.get()); - ggml_backend_event_wait(it.second.backend.get(), it.second.fence); - } + const auto fence = fences.find(it.first); + + if (fence != fences.end()) { + ggml_backend_event_wait(it.second.backend.get(), fence->second); } } } @@ -3606,6 +3614,16 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr // [TAG_STATE_ASYNC] +void llama_context::state_seq_copy_fence() { + for (const auto & it : state_copy_fences) { + for (const auto & backend : backends) { + if (ggml_backend_get_device(backend.get()) == it.first) { + ggml_backend_event_record(it.second, backend.get()); + } + } + } +} + llama_state_seq_copy * llama_context::state_seq_copy_init() { std::unique_ptr cpy(new llama_state_seq_copy()); @@ -3659,25 +3677,31 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } - ggml_backend_event_t fence = ggml_backend_event_new(dev); + if (state_copy_fences.find(dev) == state_copy_fences.end()) { + ggml_backend_event_t fence = ggml_backend_event_new(dev); - if (!fence) { - ggml_backend_event_free(event); - ggml_backend_free(backend_cpy); - continue; + if (!fence) { + ggml_backend_event_free(event); + ggml_backend_free(backend_cpy); + continue; + } + + state_copy_fences[dev] = fence; } auto & dc = cpy->devs[dev]; dc.backend.reset(backend_cpy); dc.event = event; - dc.fence = fence; } if (cpy->devs.empty()) { return nullptr; } + // the fences say where the compute streams are now, before any transfer asks + state_seq_copy_fence(); + // The devices above are the ones the graphs run on, not necessarily the ones the state // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor // there takes the synchronous branch of backend_for(). A transfer whose every copy would @@ -3737,7 +3761,7 @@ size_t llama_context::state_seq_copy_get(llama_state_seq_copy & cpy, size_t size // finished before they are read: the copy stream waits for the compute stream, on the // device, see order_after(). Nothing stays on the caller's thread. const int64_t t_sync = ggml_time_us(); - cpy.order_after(backends); + cpy.order_after(state_copy_fences); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; @@ -3783,7 +3807,7 @@ size_t llama_context::state_seq_copy_set(llama_state_seq_copy & cpy, size_t size // read), so the copy stream waits for the compute stream before it writes them: on the // device, see order_after(), rather than by draining the compute stream on this thread const int64_t t_sync = ggml_time_us(); - cpy.order_after(backends); + cpy.order_after(state_copy_fences); cpy.t_sync_us = ggml_time_us() - t_sync; cpy.n_copies = 0; diff --git a/src/llama-context.h b/src/llama-context.h index 3d66f2a948a..efa69f33cdc 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -165,6 +165,10 @@ struct llama_context { size_t state_seq_copy_get(llama_state_seq_copy & cpy, size_t size, llama_seq_id seq_id, llama_state_seq_flags flags); size_t state_seq_copy_set(llama_state_seq_copy & cpy, size_t size, llama_seq_id dest_seq_id, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] mark the point the compute streams have reached, for the copies to + // wait for; recorded after every decode and encode once a transfer exists + void state_seq_copy_fence(); + bool state_load_file( const char * filepath, llama_token * tokens_out, @@ -357,6 +361,10 @@ struct llama_context { ggml_backend_t backend_cpu = nullptr; std::vector backends; + // [TAG_STATE_ASYNC] one event per device that copies asynchronously, recorded on the + // compute stream at the end of every decode; see state_seq_copy_fence() + std::map state_copy_fences; + // training ggml_opt_context_t opt_ctx = nullptr; From 7eff72650d2e90edb45812805f6d00879b14d50b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 01:48:16 +0000 Subject: [PATCH 23/25] llama: staging counts by what a buffer charges; fences installed after 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. --- ggml/include/ggml-backend.h | 3 ++ ggml/src/ggml-backend.cpp | 4 +++ src/llama-context.cpp | 59 ++++++++++++++++++++++----------- tools/server/server-context.cpp | 16 +++++++-- 4 files changed, 59 insertions(+), 23 deletions(-) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index d21bf40dd58..09ee64a6561 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -62,6 +62,9 @@ extern "C" { GGML_API size_t ggml_backend_buffer_get_alloc_size(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor); GGML_API void ggml_backend_buffer_clear (ggml_backend_buffer_t buffer, uint8_t value); GGML_API bool ggml_backend_buffer_is_host (ggml_backend_buffer_t buffer); + // whether the buffer copies a strided set of rows in one call (see ggml_backend_tensor_set_2d); + // without it the generic path issues one transfer per row + GGML_API bool ggml_backend_buffer_supports_2d (ggml_backend_buffer_t buffer); GGML_API void ggml_backend_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); GGML_API enum ggml_backend_buffer_usage ggml_backend_buffer_get_usage (ggml_backend_buffer_t buffer); GGML_API ggml_backend_buffer_type_t ggml_backend_buffer_get_type (ggml_backend_buffer_t buffer); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index b13d9c811c4..1ac8ecad9f6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -175,6 +175,10 @@ bool ggml_backend_buffer_is_host(ggml_backend_buffer_t buffer) { return ggml_backend_buft_is_host(ggml_backend_buffer_get_type(buffer)); } +bool ggml_backend_buffer_supports_2d(ggml_backend_buffer_t buffer) { + return buffer->iface.set_tensor_2d != NULL && buffer->iface.get_tensor_2d != NULL; +} + void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { GGML_ASSERT(buffer); buffer->usage = usage; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 21776af9afd..4c63f0bec2a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2769,14 +2769,19 @@ class llama_io_read_host : public llama_io_read_i { // matters is how many runs of adjacent cells they form, because that is how many // transfers they actually cost. Count the runs first, and only fall back to // staging the whole tensor when even the runs are too many. + const size_t tensor_bytes = ggml_nbytes(tensor); + auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + // A strided set of rows is one transfer on a buffer that copies 2-D, and one + // per row on one that does not (the generic path expands it), so it is counted + // by what it costs on this buffer, not by the calls it makes. + const bool has_2d = ggml_backend_buffer_supports_2d(buffer); + size_t n_runs = 0; llama_io_emit(rinfos, i, end, - [&n_runs](ggml_tensor *, const uint8_t *, size_t, size_t, size_t, size_t, size_t) { - n_runs++; + [&n_runs, has_2d](ggml_tensor *, const uint8_t *, size_t, size_t, size_t n_copies, size_t, size_t) { + n_runs += has_2d ? 1 : n_copies; }); - - const size_t tensor_bytes = ggml_nbytes(tensor); - auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; // A fragmented sequence can require thousands of synchronous device // transfers per layer. For bounded tensors, stage the tensor once and // preserve every byte belonging to other sequences. Bound scratch RAM @@ -3677,18 +3682,6 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { continue; } - if (state_copy_fences.find(dev) == state_copy_fences.end()) { - ggml_backend_event_t fence = ggml_backend_event_new(dev); - - if (!fence) { - ggml_backend_event_free(event); - ggml_backend_free(backend_cpy); - continue; - } - - state_copy_fences[dev] = fence; - } - auto & dc = cpy->devs[dev]; dc.backend.reset(backend_cpy); @@ -3699,9 +3692,6 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { return nullptr; } - // the fences say where the compute streams are now, before any transfer asks - state_seq_copy_fence(); - // The devices above are the ones the graphs run on, not necessarily the ones the state // lives on: with most layers left on the CPU the KV cache is host memory, and a tensor // there takes the synchronous branch of backend_for(). A transfer whose every copy would @@ -3734,6 +3724,35 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { cpy->can_pin = cpy->host_buffer_type() != ggml_backend_cpu_buffer_type(); + // One fence per device, shared by every transfer on this context and recorded after + // every decode from now on. Installed only here, after the checks above: a transfer + // refused for its layout must leave nothing behind that every later decode would keep + // recording for nobody. + std::vector fences_new; + + for (const auto & it : cpy->devs) { + if (state_copy_fences.find(it.first) != state_copy_fences.end()) { + continue; + } + + ggml_backend_event_t fence = ggml_backend_event_new(it.first); + + if (!fence) { + for (auto dev : fences_new) { + ggml_backend_event_free(state_copy_fences[dev]); + state_copy_fences.erase(dev); + } + + return nullptr; + } + + state_copy_fences[it.first] = fence; + fences_new.push_back(it.first); + } + + // the fences say where the compute streams are now, before any transfer asks + state_seq_copy_fence(); + return cpy.release(); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index dcba6d0408e..5991a2ba3a3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1698,8 +1698,11 @@ struct server_context_impl { }; // [TAG_PREEMPT_ASYNC] one transfer per context, made once and reused for every - // park and resume this slot ever does, because each owns a backend and a stream - if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + // park and resume this slot ever does, because each owns a backend and a stream. + // Only where a park can happen at all (see update_preemption): a transfer also + // installs the fences the context records after every decode, which a server + // that will never park has no use for. + if (preempt_async_possible()) { slot.preempt_cpy_tgt = llama_state_seq_copy_make(ctx_tgt); if (slot.preempt_cpy_tgt && ctx_dft) { @@ -1742,7 +1745,7 @@ struct server_context_impl { preempt_async_ok = preempt_async_ok && slot.preempt_is_async(); } - if (params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0) { + if (preempt_async_possible()) { if (preempt_async_ok) { // Pinned host memory is what lets a copy run beside the decode: one into or // out of pageable memory is staged by the driver and blocks the thread that @@ -5069,6 +5072,13 @@ struct server_context_impl { // update_slots() rebuilds the batch from the survivors. The planner brings the parked // ones back as cells free up. A multimodal prompt has no boundary the cache can name, // so it keeps the old path. + // [TAG_PREEMPT_ASYNC] whether a park can happen and go asynchronously: the conditions + // update_preemption() gates on, and the asynchronous switch + bool preempt_async_possible() const { + return params_base.preempt_async && params_base.kv_unified && params_base.preempt_ram_mib != 0 && + slots.size() >= 2 && llama_get_memory(ctx_tgt) && !llama_model_is_recurrent(model_tgt); + } + bool preempt_last_resort_possible() const { return params_base.kv_unified && params_base.preempt_ram_mib != 0 && !preempt_recurrent && slots.size() >= 2 && llama_get_memory(ctx_tgt); } From 3475eb0708382b60ecaac1a13268b6e5b98ef02a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:30:12 +0000 Subject: [PATCH 24/25] llama: the shift wait runs before the draft is asked for; the last transfer 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. --- src/llama-context.cpp | 22 ++++++++++++++++++++++ src/llama-context.h | 7 +++++++ tools/server/server-context.cpp | 8 ++++++-- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4c63f0bec2a..4e6132bb0a8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,6 +483,7 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + // transfers outlive nothing: the server frees its slots before the contexts for (auto & it : state_copy_fences) { ggml_backend_event_free(it.second); } @@ -3177,6 +3178,8 @@ struct llama_state_seq_copy { ggml_backend_buffer_ptr host_buf; + bool counted = false; // held in the context's count of live transfers + uint8_t * data = nullptr; size_t size = 0; // bytes the current transfer covers size_t capacity = 0; // bytes actually held, kept across transfers @@ -3189,6 +3192,10 @@ struct llama_state_seq_copy { int64_t t_sync_us = 0; ~llama_state_seq_copy() { + if (counted) { + ctx->state_seq_copy_release(); + } + wait(); for (auto & it : devs) { @@ -3619,6 +3626,18 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr // [TAG_STATE_ASYNC] +void llama_context::state_seq_copy_release() { + GGML_ASSERT(state_copy_live > 0); + + if (--state_copy_live == 0) { + for (auto & it : state_copy_fences) { + ggml_backend_event_free(it.second); + } + + state_copy_fences.clear(); + } +} + void llama_context::state_seq_copy_fence() { for (const auto & it : state_copy_fences) { for (const auto & backend : backends) { @@ -3753,6 +3772,9 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { // the fences say where the compute streams are now, before any transfer asks state_seq_copy_fence(); + state_copy_live++; + cpy->counted = true; + return cpy.release(); } diff --git a/src/llama-context.h b/src/llama-context.h index efa69f33cdc..db7ffc6b4c0 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -169,6 +169,9 @@ struct llama_context { // wait for; recorded after every decode and encode once a transfer exists void state_seq_copy_fence(); + // a transfer letting go of this context: the last one takes the fences with it + void state_seq_copy_release(); + bool state_load_file( const char * filepath, llama_token * tokens_out, @@ -365,6 +368,10 @@ struct llama_context { // compute stream at the end of every decode; see state_seq_copy_fence() std::map state_copy_fences; + // transfers alive on this context; the fences go when the last one does, so a server + // that made transfers and then gave them up records nothing after its decodes + int32_t state_copy_live = 0; + // training ggml_opt_context_t opt_ctx = nullptr; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 418d17580db..c3a8195bf0a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4244,6 +4244,12 @@ struct server_context_impl { pre_decode_shift(); update_preemption(); + // [TAG_PREEMPT_ASYNC] before pre_decode(), not only before the target decode: the + // draft it asks for is a decode on the draft context, which applies that cache's + // pending shift in place, and a park or restore still copying draft cells would + // read through it or be overwritten by it just the same + preempt_wait_for_shift(); + scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); batch.render(); @@ -4280,8 +4286,6 @@ struct server_context_impl { int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); - preempt_wait_for_shift(); - for (int32_t off = 0; off < batch.size(); off = off_next) { const int32_t n_tokens = std::min(n_batch, batch.size() - off); try { From ebfa47b1678a2525840f82ad249f13cad19eda13 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 7 Sep 2026 02:59:15 +0000 Subject: [PATCH 25/25] llama: a context freed with live transfers drains and disowns them; the 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. --- src/llama-context.cpp | 27 +++++++++++++++++++++------ src/llama-context.h | 11 ++++++++--- tools/server/server-context.cpp | 5 +++++ 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4e6132bb0a8..a5e9f269233 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -483,7 +483,12 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); - // transfers outlive nothing: the server frees its slots before the contexts + // A transfer still alive is drained first: synchronize() covers the graph backends, + // not the copy backend a transfer owns, and the KV buffers it may still be reading or + // writing are about to go. It is then let go of, so freeing it later touches nothing + // of this context. + state_seq_copies_drain(); + for (auto & it : state_copy_fences) { ggml_backend_event_free(it.second); } @@ -3193,7 +3198,7 @@ struct llama_state_seq_copy { ~llama_state_seq_copy() { if (counted) { - ctx->state_seq_copy_release(); + ctx->state_seq_copy_release(this); } wait(); @@ -3626,10 +3631,20 @@ size_t llama_context::state_seq_set_data(llama_seq_id seq_id, const uint8_t * sr // [TAG_STATE_ASYNC] -void llama_context::state_seq_copy_release() { - GGML_ASSERT(state_copy_live > 0); +void llama_context::state_seq_copies_drain() { + for (auto * cpy : state_copies) { + cpy->wait(); + cpy->ctx = nullptr; + cpy->counted = false; + } + + state_copies.clear(); +} + +void llama_context::state_seq_copy_release(llama_state_seq_copy * cpy) { + GGML_ASSERT(state_copies.erase(cpy) == 1); - if (--state_copy_live == 0) { + if (state_copies.empty()) { for (auto & it : state_copy_fences) { ggml_backend_event_free(it.second); } @@ -3772,7 +3787,7 @@ llama_state_seq_copy * llama_context::state_seq_copy_init() { // the fences say where the compute streams are now, before any transfer asks state_seq_copy_fence(); - state_copy_live++; + state_copies.insert(cpy.get()); cpy->counted = true; return cpy.release(); diff --git a/src/llama-context.h b/src/llama-context.h index db7ffc6b4c0..106dc3922e6 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -12,6 +12,7 @@ #include "ggml-opt.h" #include +#include #include struct llama_model; @@ -170,7 +171,10 @@ struct llama_context { void state_seq_copy_fence(); // a transfer letting go of this context: the last one takes the fences with it - void state_seq_copy_release(); + void state_seq_copy_release(llama_state_seq_copy * cpy); + + // at teardown: wait for every live transfer and let it go + void state_seq_copies_drain(); bool state_load_file( const char * filepath, @@ -369,8 +373,9 @@ struct llama_context { std::map state_copy_fences; // transfers alive on this context; the fences go when the last one does, so a server - // that made transfers and then gave them up records nothing after its decodes - int32_t state_copy_live = 0; + // that made transfers and then gave them up records nothing after its decodes, and a + // context freed with transfers still alive drains them and lets them go first + std::set state_copies; // training ggml_opt_context_t opt_ctx = nullptr; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 6b24b528b2a..99751e77e46 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4290,6 +4290,11 @@ struct server_context_impl { int32_t off_next = 0; int32_t n_batch = llama_n_batch(ctx_tgt); + // [TAG_PREEMPT_ASYNC] and once more here: a shift --cache-reuse asks for is found + // inside pre_decode(), after the wait above, and the decode below applies it in + // place like any other + preempt_wait_for_shift(); + for (int32_t off = 0; off < batch.size(); off = off_next) { const int32_t n_tokens = std::min(n_batch, batch.size() - off); try {