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/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e3..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); @@ -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); diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3d..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 @@ -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..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; @@ -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); 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..1fca6352403 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/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 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 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 diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..3c64888d25f 100644 --- a/include/llama.h +++ b/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 // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 66940d4fc61..a5e9f269233 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 @@ -482,6 +483,16 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + // 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); + } + if (!model.hparams.no_alloc) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; @@ -1577,6 +1588,10 @@ int llama_context::encode(const llama_batch & batch_inp) { } } + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -2022,6 +2037,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; } @@ -2558,16 +2577,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 +2771,28 @@ 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. 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, 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; + }); // 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 +2812,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 +3164,398 @@ 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; + + 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 + 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() { + if (counted) { + ctx->state_seq_copy_release(this); + } + + 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()); + } + } + } + + // 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. 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) { + const auto fence = fences.find(it.first); + + if (fence != fences.end()) { + ggml_backend_event_wait(it.second.backend.get(), fence->second); + } + } + } + + // 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; + + 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) {} + + // 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) { + 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; + + bool committed = false; +}; + +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) {} + + // 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 + // 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; + + bool committed = false; +}; + 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 +3629,276 @@ 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_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_copies.empty()) { + 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) { + 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()); + + 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 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 + 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; + } + + // 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(); + + // 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(); + + state_copies.insert(cpy.get()); + cpy->counted = true; + + 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) { + // 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; + } + + // 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(); + cpy.order_after(state_copy_fences); + 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)); + + 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; + } +} + +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) { + // 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; + } + + // 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(); + cpy.order_after(state_copy_fences); + cpy.t_sync_us = ggml_time_us() - t_sync; + + cpy.n_copies = 0; + + size_t n = 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)); + + n = state_seq_read_data(io, seq_id, flags); + + 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) { llama_file file(filepath, "rb"); @@ -4125,6 +4953,71 @@ 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) { + // 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; +} + +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..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; @@ -39,6 +40,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 +160,22 @@ 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); + + // [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(); + + // a transfer letting go of this context: the last one takes the fences with it + 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, llama_token * tokens_out, @@ -348,6 +368,15 @@ 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; + + // 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, 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/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..dabc4db50c3 --- /dev/null +++ b/tests/test-state-seq-copy.cpp @@ -0,0 +1,159 @@ +// [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__); + + // 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()); + + 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); + + // 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); + + 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; +} 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 b18fa4e2f23..6ecd33ce827 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -61,6 +61,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 @@ -91,6 +93,34 @@ constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked constexpr int64_t PREEMPT_ROTATE_US = 2ll * 1000 * 1000; // a resident cycling through context shifts gives way to a parked head that has waited this long +// [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 { @@ -331,34 +361,229 @@ 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_ctx_shift = 0; // context shifts the CURRENT task has made: it is at the pool's limit and cycling 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 bool preempt_rotation_refused = false; // this park has logged a rotation refused for budget 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; + } + + // [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)); + + 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(); + + // 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); @@ -382,16 +607,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. @@ -408,7 +624,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(); @@ -691,7 +933,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(); } @@ -838,7 +1086,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}, }; @@ -1083,6 +1331,27 @@ 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(); + } + + // 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(); @@ -1423,6 +1692,25 @@ 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. + // 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) { + 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(); } @@ -1444,6 +1732,52 @@ 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 (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 + // 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"); + } + } + + if (!preempt_async_ok) { + for (auto & slot : slots) { + slot.preempt_cpy_tgt.reset(); + slot.preempt_cpy_dft.reset(); + } + } + } + { // read on every load and kept on this context, so a reload after the variable // changed, or another context loaded in the same process, has an order of its own @@ -2637,7 +2971,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++; } } @@ -2892,7 +3226,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(); } @@ -2942,6 +3277,11 @@ struct server_context_impl { int32_t preempt_test_every = 0; std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model + // [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; + // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so // the KV-full retry ladder and its last resort are the only thing between a full pool // and the context error @@ -2960,10 +3300,35 @@ 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; } + // [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"); + } + // draft tokens this slot's next step can actually carry: the configured maximum, cut to // what its context and its prediction budget leave, the way get_n_draft_max() cuts it int32_t preempt_n_spec(const server_slot & slot) const { @@ -2994,14 +3359,82 @@ 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; } 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; + + preempt_reclaim_idle_ram(budget, extra, slot); + + return preempt_ram_used() + 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; + + 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 of the mirrored prompt that a started slot's request keeps, by the rule the batch @@ -3065,6 +3498,11 @@ 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. + // a child waiting for its parent's prompt does not share anything yet: until // copy_state_to() runs it still holds whatever the previous request left in its // cells, so it is charged that on its own, outside the family @@ -3096,6 +3534,38 @@ struct server_context_impl { 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. + // + // 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 = n_additional_running; + + for (const auto & slot : slots) { + // 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++; + } + } + + 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_batch = llama_n_batch(ctx_tgt); @@ -3120,6 +3590,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 + preempt_n_spec(slot); + } 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; } @@ -3203,7 +3688,7 @@ struct server_context_impl { } 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()) { @@ -3280,6 +3765,108 @@ 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++; + + 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(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + } + } + } + + // [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++; + + 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(), + 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 @@ -3289,6 +3876,8 @@ struct server_context_impl { return; // no cache at all (an embedding model): nothing to run out of, nothing to park } + update_preempt_copies(); + if (params_base.preempt_ram_mib == 0 || preempt_recurrent) { return; // --preempt-ram 0, or a recurrent cache: the KV-full retry ladder, as before } @@ -3354,14 +3943,15 @@ struct server_context_impl { } // 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. - // The margin is headroom for the others; with nothing resident there is nobody - // to keep it for, so a sequence that fits the pool exactly is let back in. + // 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. The margin is headroom for the others; with nothing resident there + // is nobody to keep it for, so a sequence that fits the pool exactly is let back in. // 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 (;;) { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); - const int32_t margin = occupied == 0 ? 0 : PREEMPT_N_MARGIN; + const int32_t margin = occupied == 0 ? 0 : preempt_n_margin(1); for (auto * slot : parked) { if (occupied + preempt_n_need(*slot) + margin <= n_cells) { @@ -3395,7 +3985,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 @@ -3437,6 +4035,8 @@ struct server_context_impl { } } + const int64_t t_start = ggml_time_us(); + if (!pick && budget_refused && !head->preempt_rotation_refused) { head->preempt_rotation_refused = true; @@ -3447,7 +4047,13 @@ struct server_context_impl { if (pick && pick->preempt_save()) { server_slot & slot = *pick; - metrics.n_preempt++; + 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) { + metrics.n_preempt++; + } SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn%s, preemptions %d\n", slot.n_ctx_shift, slot.prompt.n_tokens(), @@ -3456,7 +4062,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; + } } } @@ -3469,6 +4080,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 @@ -3487,6 +4100,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", @@ -3502,12 +4131,21 @@ 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(); - 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)); + 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) { + 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)); + } } } } @@ -3520,7 +4158,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; } @@ -3529,6 +4167,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) { @@ -3540,10 +4187,50 @@ 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 } + 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 + // 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); + + // [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; + } + metrics.n_preempt++; SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", @@ -3604,6 +4291,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(); @@ -3639,6 +4332,12 @@ struct server_context_impl { llama_batch batch_view; 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 { @@ -3730,6 +4429,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); @@ -3903,7 +4603,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; } @@ -4056,6 +4756,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++; @@ -4437,6 +5141,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); } @@ -4449,7 +5160,8 @@ struct server_context_impl { int32_t n_running = 0; for (auto & slot : slots) { - if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { + // [TAG_PREEMPT_ASYNC] a slot in transfer is not in this batch either way + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED || slot.preempt_in_flight()) { continue; } @@ -4465,7 +5177,8 @@ struct server_context_impl { } for (auto & slot : slots) { - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && slot.state != SLOT_STATE_WAIT_OTHER) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && slot.state != SLOT_STATE_WAIT_OTHER && + !slot.preempt_in_flight()) { slot.rewind_to_cache(); } } @@ -4476,7 +5189,7 @@ struct server_context_impl { for (;;) { const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); - if (n_parked > 0 && n_used + PREEMPT_N_MARGIN <= n_cells) { + if (n_parked > 0 && n_used + preempt_n_margin() <= n_cells) { break; } @@ -4489,13 +5202,29 @@ 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; } - metrics.n_preempt++; + preempt_log_ram_kind(*victim); + n_parked++; + // [TAG_PREEMPT_ASYNC] the cells are wanted now, not next iteration: wait for the + // copy to land, which releases them and logs the park the way the planner does + if (victim->state == SLOT_STATE_PREEMPTING) { + while (preempt_wait_in_flight()) { + } + + SLT_WRN(*victim, "preempted as a last resort: %d cells released, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, preempt_kv_used(), n_cells, n_used, victim->n_preempt); + continue; + } + + metrics.n_preempt++; + SLT_WRN(*victim, "preempted as a last resort: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, (ggml_time_us() - t_start) / 1e3, @@ -4569,6 +5298,18 @@ struct server_context_impl { }); if (ret != 0) { + // [TAG_PREEMPT_ASYNC] Before giving up any batch width, and before the last resort: + // 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 + } + { std::string err; @@ -4615,7 +5356,7 @@ struct server_context_impl { for (auto & slot : slots) { // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the // cache; it is not part of this failure and comes back when there is room - if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && !slot.preempt_in_flight()) { send_error(slot, err); slot.release(); @@ -4965,7 +5706,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 0c8f5dc5f29..34846487ed3 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -40,6 +40,7 @@ def create_server(): os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) os.environ.pop("LLAMA_SERVER_PREEMPT_PLANNER", 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"): @@ -270,6 +271,242 @@ def test_metrics_and_slots_report_the_parked_state(): 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 + + +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 + + def test_two_prompts_near_the_context_size_both_complete(): # Two prompts that each fit the context alone but not together. The second one is # parked before it takes any cells, and it is close enough to n_ctx that its sequence @@ -484,6 +721,43 @@ def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_ 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 + + def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_when_a_resident_finishes(): # Three generations with no end in a pool one of them fills, with context shift on, # under a --preempt-ram that holds the two parked heads but not a head and the resident