server: take the park and restore copies off the decode loop

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"},

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

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);
@@ -125,6 +128,10 @@ 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,
+    // 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);
 
     //
@@ -190,6 +197,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);
 

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
@@ -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 {

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) {

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;
@@ -551,6 +555,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);
@@ -627,6 +643,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);

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

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

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)

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5375,6 +5375,23 @@ 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);
+
+    // 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) {
+        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 +5408,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

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

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

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,
 };
 
 /*

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

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

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,
 };
 }
 

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 {

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

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

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,
 };

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) {

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 */

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,
 };
 
 //

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

include/llama.h

@@ -927,6 +927,75 @@ 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, 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);
+
+    // 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 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. 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,
+                     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
     //