diff --git a/common/arg.cpp b/common/arg.cpp index 5bfa4adcdf0..af32812b930 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1304,6 +1304,11 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e exit(0); } params.lr.init(); + + if (!common_exact_concurrency_init(ctx_arg.params)) { + ctx_arg.params = params_org; + return false; + } } catch (const std::invalid_argument & ex) { fprintf(stderr, "%s\n", ex.what()); ctx_arg.params = params_org; @@ -1717,6 +1722,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.cpp b/common/common.cpp index 3d54bd6002d..18477a91ffb 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1,4 +1,5 @@ #include "ggml.h" +#include "ggml-backend.h" #include "gguf.h" #include "build-info.h" @@ -1289,6 +1290,12 @@ struct common_init_result::impl { common_init_result::common_init_result(common_params & params, bool model_only) : pimpl(new impl{}) { + // [TAG_EXACT_CONCURRENCY] before any context exists, so one is never created under a figure the explicit bound does not cover + if (!model_only && !common_exact_concurrency_init(params)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n"); + return; + } + auto mparams = common_model_params_to_llama(params); auto cparams = common_context_params_to_llama(params); @@ -1433,6 +1440,71 @@ std::vector & common_init_result::lora() { return pimpl->lora; } +// [TAG_EXACT_CONCURRENCY] +bool common_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + +int common_exact_decode_width(const common_params & params) { + const int64_t n_slots = std::max(1, params.n_parallel); + + const int64_t n_draft = std::max(0, (int) common_speculative_n_max(¶ms.speculative)); + + // the product is handed to a backend as an int; one that overflows is reported, not wrapped + const int64_t n_cols = n_slots*(1 + n_draft); + + return n_cols > INT32_MAX ? -1 : (int) n_cols; +} + +bool common_exact_concurrency_init(const common_params & params) { + if (!common_exact_concurrency()) { + return true; + } + + // DFlash drafting turns causal attention off on its draft context, which the paged attention needs; say so instead of asserting in the graph. DSpark is the same. + for (const auto type : params.speculative.types) { + if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n"); + return false; + } + } + + const int n_cols = common_exact_decode_width(params); + + if (n_cols < 0) { + COM_ERR("LLAMA_EXACT_CONCURRENCY: a decode step of %d slots with %d draft tokens each is too wide to report\n", + std::max(1, params.n_parallel), std::max(0, (int) common_speculative_n_max(¶ms.speculative))); + return false; + } + + const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + if (bound) { + const int max_cols = atoi(bound); + if (max_cols > 0 && max_cols < n_cols) { + COM_ERR("GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at " + "least %d to cover a decode step of %d slots, above which a matmul is left " + "batched and its rows depend on the other rows in the ubatch. Raise it to %d, " + "set it to 0 for no bound, or unset it to let it default to %d.\n", + max_cols, n_cols, std::max(1, params.n_parallel), n_cols, n_cols); + return false; + } + } + + // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is; this also covers a caller that decodes before creating a context + if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) || + !llama_set_exact_decode_width((uint32_t) n_cols)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n"); + return false; + } + + return true; +} + common_init_result_ptr common_init_from_params(common_params & params, bool model_only) { common_init_result_ptr res(new common_init_result(params, model_only)); diff --git a/common/common.h b/common/common.h index c99269f9a96..e22c07bad12 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 @@ -931,6 +932,14 @@ using common_init_result_ptr = std::unique_ptr; common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false); +// [TAG_EXACT_CONCURRENCY] true when LLAMA_EXACT_CONCURRENCY is set for this process +bool common_exact_concurrency(); + +int common_exact_decode_width(const common_params & params); + +// report that width to the CUDA backend, refusing a smaller explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS; false if the configuration must not run +bool common_exact_concurrency_init(const common_params & params); + struct llama_model_params common_model_params_to_llama ( common_params & params); struct llama_context_params common_context_params_to_llama(const common_params & params); diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e3..84a2f8458ea 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -62,6 +62,8 @@ 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 +127,8 @@ 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. + 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 +194,8 @@ 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, 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/include/ggml-cuda.h b/ggml/include/ggml-cuda.h index 1cd81eeaebc..897da6ca5f8 100644 --- a/ggml/include/ggml-cuda.h +++ b/ggml/include/ggml-cuda.h @@ -38,6 +38,9 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); + +// [TAG_EXACT_CONCURRENCY] report the widest ubatch a decode step of this process can build, so the column policy covers it; call before the first graph is computed +GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3d..e417b3cc448 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,9 @@ 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: a missing entry is NULL and ggml_backend_event_query() then blocks instead. + 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..d29ec88bde9 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2656,6 +2656,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten return true; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return false; + } #ifdef ASCEND_310P // FA not support on 310p device return false; @@ -2948,6 +2952,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..7ea548bcbce 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -474,6 +474,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return ggml_is_contiguous(op->src[0]); case GGML_OP_SSM_SCAN: return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; + // [TAG_EXACT_CONCURRENCY] note: FLASH_ATTN_EXT with src[5], the page table, is deliberately still accepted: the CPU ignores it, but it is the reference test-backend-ops uses default: return true; } @@ -500,6 +501,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/common.cuh b/ggml/src/ggml-cuda/common.cuh index 14dd1098c97..47ecee48c2f 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -49,6 +49,11 @@ #define GGML_CUDA_CC_PASCAL 600 #define GGML_CUDA_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products +// [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits +int ggml_cuda_batch_invariant(); +// widest batch the split applies to, 0 = no bound; bounding it gives up prompt-phase invariance only +int ggml_cuda_batch_invariant_max_cols(); + #define GGML_CUDA_CC_VOLTA 700 #define GGML_CUDA_CC_TURING 750 #define GGML_CUDA_CC_AMPERE 800 diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index e67cc7fdf78..61aa0dc376a 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1091,7 +1091,9 @@ void launch_fattn( // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped. // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or // multiple sequences of possibly different lengths. - if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) { + // [TAG_BATCH_INVARIANT] without this scan the KV loop runs to K->ne[1], which grows with the other sequences; the mask bounds it by the sequence's own extent + const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0; + if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) { const int64_t s31 = mask->nb[1] / sizeof(half2); const int64_t s33 = mask->nb[3] / sizeof(half2); @@ -1148,6 +1150,13 @@ void launch_fattn( if (ntiles_dst % blocks_num.x != 0) { // Fixup is only needed if the SMs work on fractional tiles. dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2))); } + } else if (dst->src[5] || ggml_cuda_batch_invariant()) { + // [TAG_BATCH_INVARIANT] the KV split between blocks, and so the order the partials combine in, follows K->ne[1]: pin it to one block per tile + parallel_blocks = 1; + + blocks_num.x = ntiles_x; + blocks_num.y = parallel_blocks; + blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3]; } else { // parallel_blocks must not be larger than what the tensor size allows: parallel_blocks = std::min(parallel_blocks, ntiles_KV); @@ -1214,7 +1223,7 @@ void launch_fattn( V_data, mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, - KV_max.ptr, + dst->src[5] ? (const int *) dst->src[5]->data : KV_max.ptr, !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3], diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 69dd9368624..0dba4b80b9c 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -16,7 +16,7 @@ static constexpr __device__ int ggml_cuda_fattn_vec_get_nthreads_device() { #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wpass-failed" #endif // __clang__ -template // D == head size +template // D == head size __launch_bounds__(ggml_cuda_fattn_vec_get_nthreads_device(), 1) static __global__ void flash_attn_ext_vec( const char * Q_ptr, @@ -247,13 +247,24 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } - const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11; + // in the paged specialization KV_max carries [count, physical page IDs...] per query; the loop and each warp's recurrence follow logical positions, never physical addresses + static_assert(!paged || ncols == 1, "paged attention has one query per block"); + const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr; + const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11); + const char * K_base = K; + const char * V_base = V; + const half * mask_base = maskh; K += blockIdx.y*nthreads * nb11; V += blockIdx.y*nthreads * nb21; maskh += blockIdx.y*nthreads; for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads, - // Increment pointers after each loop: K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) { + if constexpr (paged) { + const int physical = pages[1 + k_VKQ_0/FATTN_KQ_STRIDE]*FATTN_KQ_STRIDE + k_VKQ_0%FATTN_KQ_STRIDE; + K = K_base + int64_t(physical)*nb11; + V = V_base + int64_t(physical)*nb21; + maskh = mask_base + physical; + } // Calculate KQ tile and keep track of new maximum KQ values: float KQ_reg[ncols]; // KQ in registers. diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c0..5a1303401d0 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -457,6 +457,11 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; + // [TAG_BATCH_INVARIANT] every choice below switches on Q->ne[1] or K->ne[1], both of which grow with the other sequences, so pin the kernel a batch of one would use + if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) { + return BEST_FATTN_KERNEL_VEC; + } + // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel) { @@ -569,6 +574,52 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + + if (dst->src[5]) { + GGML_ASSERT(dst->src[0]->ne[0] == 256 && dst->src[2]->ne[0] == 256); + GGML_ASSERT(dst->src[1]->type == GGML_TYPE_F16 && dst->src[2]->type == GGML_TYPE_F16); + GGML_ASSERT(dst->src[3] && dst->src[0]->ne[3] == 1); + GGML_ASSERT(dst->src[5]->type == GGML_TYPE_I32 && ggml_is_contiguous(dst->src[5])); + GGML_ASSERT(dst->src[5]->ne[0] == 1 + dst->src[1]->ne[1]/FATTN_KQ_STRIDE); + GGML_ASSERT(dst->src[5]->ne[1] == dst->src[0]->ne[1]); + float softcap; + memcpy(&softcap, (const float *) dst->op_params + 2, sizeof(softcap)); + GGML_ASSERT(softcap == 0.0f); + fattn_kernel_t kernel = flash_attn_ext_vec<256, 1, GGML_TYPE_F16, GGML_TYPE_F16, false, true>; + launch_fattn<256, 1, 1>(ctx, dst, kernel, 4, 0, 128, false, false, false); + return; + } + + // [TAG_BATCH_INVARIANT] attend one query row at a time, as a batch of one would + const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols(); + if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 && + (fattn_max_cols <= 0 || dst->src[0]->ne[1] <= fattn_max_cols)) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * mask = dst->src[3]; + + for (int64_t i = 0; i < Q->ne[1]; ++i) { + ggml_tensor Q_row = *Q; + Q_row.ne[1] = 1; + Q_row.data = (char *) Q->data + i*Q->nb[1]; + + ggml_tensor mask_row; + ggml_tensor dst_row = *dst; + // ne[2] runs to the end of dst so the F16 K/V scratch behind dst stays in place + dst_row.ne[2] = dst->ne[2] - i; + dst_row.data = (char *) dst->data + i*dst->nb[2]; + dst_row.src[0] = &Q_row; + if (mask) { + mask_row = *mask; + mask_row.ne[1] = 1; + mask_row.data = (char *) mask->data + i*mask->nb[1]; + dst_row.src[3] = &mask_row; + } + + ggml_cuda_flash_attn_ext(ctx, &dst_row); + } + return; + } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc62..b2bc1934795 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1758,6 +1758,11 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { + // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving it on would give a solo request a different code path from a batched one + if (ggml_cuda_batch_invariant()) { + return false; + } + ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1785,6 +1790,10 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { } static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { + if (ggml_cuda_batch_invariant()) { + return false; + } + ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; const ggml_tensor * dst = tensor; @@ -1813,60 +1822,281 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - GGML_TENSOR_BINARY_OP_LOCALS +// [TAG_BATCH_INVARIANT] the token count picks the matmul and how its K loop is split, so the same request produces different bits. GGML_CUDA_BATCH_INVARIANT: +// 1 - compute every destination column on its own, exactly as a batch of one would +// 2 - split off only the columns whose batch-of-one configuration differs from the batched one +static bool ggml_cuda_exact_concurrency() { + static const bool exact = []() { + const char * value = getenv("LLAMA_EXACT_CONCURRENCY"); + return value && atoi(value) != 0; + }(); + return exact; +} - const int32_t hint = ggml_get_op_params_i32(dst, 1); - if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { +int ggml_cuda_batch_invariant() { + static const int mode = []() { + if (ggml_cuda_exact_concurrency()) { return 2; } + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT"); + return val ? atoi(val) : 0; + }(); + return mode; +} + +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch the caller says it can build, 0 if it never said +static std::atomic g_exact_decode_width{0}; + +void ggml_backend_cuda_set_exact_decode_width(int n_cols) { + // monotonic: the widest figure ever reported stays, whatever order the reports arrive in + int cur = g_exact_decode_width.load(std::memory_order_relaxed); + + while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + } +} + +int ggml_cuda_batch_invariant_max_cols() { + // [TAG_EXACT_CONCURRENCY] prompt ubatches hold one sequence, so a prefill already matches its solo run; an explicit bound always wins + static const int explicit_cols = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + return val ? atoi(val) : -1; + }(); + + if (explicit_cols >= 0) { + return explicit_cols; + } + + if (!ggml_cuda_exact_concurrency()) { + return 0; + } + + const int width = g_exact_decode_width.load(std::memory_order_relaxed); + + return width > 0 ? width : 16; +} + +// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so say so once. Only when nothing reported a decode width: with one, wider batches are single-sequence prefills. +static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) { + if (!ggml_cuda_exact_concurrency()) { + return; + } + + if (g_exact_decode_width.load(std::memory_order_relaxed) > 0) { return; } + static std::atomic_flag warned = ATOMIC_FLAG_INIT; + if (warned.test_and_set(std::memory_order_relaxed)) { + return; + } + + GGML_LOG_WARN("%s: LLAMA_EXACT_CONCURRENCY is set, but this %s is %d columns wide while " + "GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d, so it is left batched and its result depends " + "on the other columns in the ubatch. Raise the bound, set it to 0 for no bound, or call " + "ggml_backend_cuda_set_exact_decode_width() with the widest decode this process builds. " + "Reported once.\n", __func__, op, (int) ncols, max_cols); +} + +enum ggml_cuda_mm_path { + GGML_CUDA_MM_CUBLAS_UNSUPPORTED, + GGML_CUDA_MM_MMVF, + GGML_CUDA_MM_MMVF_TRANSPOSED, + GGML_CUDA_MM_MMF, + GGML_CUDA_MM_MMVQ, + GGML_CUDA_MM_MMQ, + GGML_CUDA_MM_CUBLAS, +}; + +static ggml_cuda_mm_path ggml_cuda_mul_mat_path( + int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) { // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. // Therefore, in such cases use cuBLAS. const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); - return; + return GGML_CUDA_MM_CUBLAS_UNSUPPORTED; } - - const int cc = ggml_cuda_info().devices[ctx.device].cc; - const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; - if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { // The custom F16 vector kernel can be used over batched cuBLAS GEMM. // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVF; } // A transposed vector can still use MMVQ (i.e. ne01 == 1) - if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1 + if (src0->ne[1] == 1 && ne11 > MMVF_MAX_BATCH_SIZE && dst->ne[2] == 1 && dst->ne[3] == 1 && src0->type == GGML_TYPE_F32 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && ggml_cuda_should_use_mmvf(src1->type, cc, src1->ne, src1->nb, /*ne11 =*/ 1)) { - ggml_tensor dst_vec = *dst; - dst_vec.ne[0] = ne11; - dst_vec.ne[1] = 1; - dst_vec.nb[1] = dst_vec.nb[0]*ne11; - dst_vec.nb[2] = dst_vec.nb[1]; - dst_vec.nb[3] = dst_vec.nb[1]; - ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); - return; + return GGML_CUDA_MM_MMVF_TRANSPOSED; } if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { - ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMF; } if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MM_MMVQ; } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return GGML_CUDA_MM_MMQ; + } + return GGML_CUDA_MM_CUBLAS; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); + +// [TAG_BATCH_INVARIANT] the widest slice of columns that can be recomputed in one launch while every column still sums as a batch of one; always below ncols_dst, so the recursion ends +static int64_t ggml_cuda_mul_mat_invariant_width( + int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, + ggml_cuda_mm_path path_one, int64_t ncols_dst) { + if (path_one != GGML_CUDA_MM_MMVF && path_one != GGML_CUDA_MM_MMVQ) { + return 1; + } + const int64_t widest = path_one == GGML_CUDA_MM_MMVF ? MMVF_MAX_BATCH_SIZE : MMVQ_MAX_BATCH_SIZE; + for (int64_t w = std::min(ncols_dst - 1, widest); w > 1; --w) { + if (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, w) != path_one) { + continue; + } + if (path_one == GGML_CUDA_MM_MMVQ && !ggml_cuda_mmvq_matches_single_column(src0->type, cc, w)) { + continue; + } + return w; + } + return 1; +} + +static bool ggml_cuda_mul_mat_split_columns( + ggml_backend_cuda_context & ctx, int cc, int warp_size, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + // recurrent output projections broadcast one weight matrix over sequence planes, so normalize each plane before applying the column policy + if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 && + (dst->ne[2] > 1 || dst->ne[3] > 1) && + src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) { + for (int64_t i3 = 0; i3 < dst->ne[3]; ++i3) { + for (int64_t i2 = 0; i2 < dst->ne[2]; ++i2) { + ggml_tensor src_plane = *src1; + ggml_tensor dst_plane = *dst; + src_plane.ne[2] = src_plane.ne[3] = 1; + dst_plane.ne[2] = dst_plane.ne[3] = 1; + src_plane.data = (char *) src1->data + i2*src1->nb[2] + i3*src1->nb[3]; + dst_plane.data = (char *) dst->data + i2*dst->nb[2] + i3*dst->nb[3]; + ggml_cuda_mul_mat(ctx, src0, &src_plane, &dst_plane); + } + } + return true; + } + + const int64_t ncols_dst = dst->ne[1]; + if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) { + return false; + } + if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { + return false; + } + const int max_cols = ggml_cuda_batch_invariant_max_cols(); + if (max_cols > 0 && ncols_dst > max_cols) { + ggml_cuda_warn_above_exact_bound("MUL_MAT", ncols_dst, max_cols); + return false; + } + + // mode 1 recomputes one column at a time; mode 2, which exact concurrency runs under, uses the widest slices that keep the batch-of-one arithmetic + int64_t width = 1; + if (ggml_cuda_batch_invariant() >= 2) { + const ggml_cuda_mm_path path_one = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1); + const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst); + if (path_one == path_batched) { + // same implementation, but it still has to sum in the same order + if (path_batched == GGML_CUDA_MM_MMVF) { + return false; // the block size follows K alone + } + if (path_batched == GGML_CUDA_MM_MMVQ && + ggml_cuda_mmvq_matches_single_column(src0->type, cc, ncols_dst)) { + return false; + } + } + width = ggml_cuda_mul_mat_invariant_width(cc, warp_size, src0, src1, dst, path_one, ncols_dst); + if (width >= ncols_dst) { + width = 1; + } + } + + for (int64_t i = 0; i < ncols_dst; i += width) { + const int64_t n = std::min(width, ncols_dst - i); + + ggml_tensor src1_col = *src1; + ggml_tensor dst_col = *dst; + + src1_col.ne[1] = n; + src1_col.nb[2] = n*src1_col.nb[1]; + src1_col.nb[3] = n*src1_col.nb[1]; + src1_col.data = (char *) src1->data + i*src1->nb[1]; + + dst_col.ne[1] = n; + dst_col.nb[2] = n*dst_col.nb[1]; + dst_col.nb[3] = n*dst_col.nb[1]; + dst_col.data = (char *) dst->data + i*dst->nb[1]; + + ggml_cuda_mul_mat(ctx, src0, &src1_col, &dst_col); + } + return true; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS + + const int32_t hint = ggml_get_op_params_i32(dst, 1); + if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { return; } - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + + if (ggml_cuda_batch_invariant() && ggml_cuda_mul_mat_split_columns(ctx, cc, warp_size, src0, src1, dst)) { + return; + } + + switch (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ne11)) { + case GGML_CUDA_MM_CUBLAS_UNSUPPORTED: + case GGML_CUDA_MM_CUBLAS: + ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + return; + case GGML_CUDA_MM_MMVF: + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMVF_TRANSPOSED: { + ggml_tensor dst_vec = *dst; + dst_vec.ne[0] = ne11; + dst_vec.ne[1] = 1; + dst_vec.nb[1] = dst_vec.nb[0]*ne11; + dst_vec.nb[2] = dst_vec.nb[1]; + dst_vec.nb[3] = dst_vec.nb[1]; + ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); + return; + } + case GGML_CUDA_MM_MMF: + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMVQ: + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MM_MMQ: + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return; + } + GGML_ABORT("fatal error"); +} + +static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) { + if (!ggml_cuda_batch_invariant()) { + return false; + } + const int64_t ntokens = dst->ne[2]; + if (ntokens <= 1) { + return false; + } + const int max_cols = ggml_cuda_batch_invariant_max_cols(); + if (max_cols > 0 && ntokens > max_cols) { + ggml_cuda_warn_above_exact_bound("MUL_MAT_ID", ntokens, max_cols); + return false; + } + return true; } // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization @@ -1879,9 +2109,12 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c return true; } - if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) { + // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path that decides whether the stream is synchronized is the single-token one + const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2]; + + if (ntokens <= MMVQ_MAX_BATCH_SIZE) { if (ggml_is_quantized(src0->type)) { - if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) { + if (ntokens <= get_mmvq_mmid_max_batch(src0->type, cc)) { return false; } } else if (GGML_CUDA_CC_IS_AMD(cc)) { @@ -1889,17 +2122,51 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c } } - if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) { + if (ggml_cuda_should_use_mmq(src0->type, cc, ntokens, /*n_experts=*/src0->ne[2])) { return false; } - if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, ntokens, /*mul_mat_id=*/true)) { return false; } return true; } +static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// [TAG_BATCH_INVARIANT] recompute dst one token at a time: every implementation below groups the ubatch's tokens by the expert they routed to, so shapes depend on the other tokens +static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * ids = dst->src[2]; + + const int64_t ntokens = dst->ne[2]; + + for (int64_t i = 0; i < ntokens; ++i) { + ggml_tensor src1_token = *src1; + ggml_tensor ids_token = *ids; + ggml_tensor dst_token = *dst; + + src1_token.ne[2] = 1; + src1_token.nb[3] = src1_token.nb[2]; + src1_token.data = (char *) src1->data + i*src1->nb[2]; + + ids_token.ne[1] = 1; + ids_token.nb[2] = ids_token.nb[1]; + ids_token.nb[3] = ids_token.nb[1]; + ids_token.data = (char *) ids->data + i*ids->nb[1]; + + dst_token.ne[2] = 1; + dst_token.nb[3] = dst_token.nb[2]; + dst_token.data = (char *) dst->data + i*dst->nb[2]; + + dst_token.src[1] = &src1_token; + dst_token.src[2] = &ids_token; + + ggml_cuda_mul_mat_id(ctx, &dst_token); + } +} + static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -1912,6 +2179,18 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + // [TAG_BATCH_INVARIANT] + if (ggml_cuda_mul_mat_id_splits_tokens(dst)) { + GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1); + // a quantized expert matrix takes the single-token MMVQ path at every token count and can put the tokens on its sample axis in one launch; anything else goes token by token + if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); + return; + } + ggml_cuda_mul_mat_id_split_tokens(ctx, dst); + return; + } + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); @@ -3297,9 +3576,10 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } - //topk-moe - if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || - cgraph->nodes[i]->op == GGML_OP_ARGSORT) { + // [TAG_BATCH_INVARIANT] the routing fusion passes its memory-range check only for a one-token ubatch, so a solo request takes the fused top-k kernel and a batched one the long chain + if (!ggml_cuda_batch_invariant() && + (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || + cgraph->nodes[i]->op == GGML_OP_ARGSORT)) { ggml_cuda_topk_moe_args args; const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); std::vector ops; @@ -5375,6 +5655,21 @@ 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, so collecting one here would consume somebody else's + 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 +5686,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 @@ -5492,6 +5788,10 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con if (strcmp(name, "ggml_backend_get_features") == 0) { return (void *)ggml_backend_cuda_get_features; } + // [TAG_EXACT_CONCURRENCY] + if (strcmp(name, "ggml_backend_cuda_set_exact_decode_width") == 0) { + return (void *)ggml_backend_cuda_set_exact_decode_width; + } return nullptr; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 97053480980..1fe68cebfa4 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -541,6 +541,20 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } +// [TAG_BATCH_INVARIANT] +bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst) { + if (ncols_dst < 1 || ncols_dst > MMVQ_MAX_BATCH_SIZE) { + return false; + } + const mmvq_parameter_table_id table_id = get_device_table_id(cc); + if (table_id == MMVQ_PARAMETERS_GB10) { + // There nwarps also depends on the K loop trip count, which the caller does not pass in. + return ncols_dst == 1; + } + // blocks_per_iter, which assigns K blocks to threads, is proportional to nwarps; rows_per_cuda_block only changes which rows a block owns, not the order within a row + return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id); +} + template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( @@ -577,9 +591,10 @@ static __global__ void mul_mat_vec_q( uint32_t sample_dst; ggml_cuda_pdl_sync(); - channel_x = ncols_dst == 1 && ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); - channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; sample_dst = blockIdx.z; + // [TAG_BATCH_INVARIANT] with ids, a sample is a token: every token goes on the z axis of one single-column launch, so each (token, expert slot) block runs the single-token configuration + channel_x = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio); + channel_y = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); const uint32_t sample_y = sample_dst; @@ -1266,7 +1281,11 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); - GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE); + // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token configuration with the tokens on the sample axis, so the count is not bounded by the column templates + const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant(); + + GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples); + GGML_ASSERT(!tokens_as_samples || !fusion); const float * src1_d = (const float *) src1->data; const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; @@ -1354,6 +1373,17 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; + if (tokens_as_samples) { + GGML_ASSERT(ne03 == 1 && ne13 == 1 && ne3 == 1); + // one column, one sample per token: y advances by s12 per token, dst by s2, x not at all + mul_mat_vec_q_switch_type( + src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, + ne01, 1, s01, stride_col_y, stride_col_dst, + ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, + 1, ne2, s03, s12, s2, ids_stride, stream); + return; + } + mul_mat_vec_q_switch_type( src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4e6..688c944c1f9 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -4,6 +4,9 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11); +// [TAG_BATCH_INVARIANT] true when an MMVQ launch of ncols_dst columns sums each destination element in the same order as a single-column launch, i.e. when nwarps is unchanged +bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst); + // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID, // based on the quantization type and GPU architecture (compute capability). int get_mmvq_mmid_max_batch(ggml_type type, int cc); 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..e80c8110265 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1266,6 +1266,11 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm (op->src[1]->ne[1] % op->src[4]->ne[1] == 0); break; case GGML_OP_FLASH_ATTN_EXT: + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + supported = false; + break; + } if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] && (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] && (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr && @@ -1684,6 +1689,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..0eac93570c4 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4157,7 +4157,8 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons break; case GGML_OP_FLASH_ATTN_EXT: - supp = ggml_hexagon_supported_flash_attn_ext(sess, op); + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + supp = op->src[5] == nullptr && ggml_hexagon_supported_flash_attn_ext(sess, op); break; case GGML_OP_SET_ROWS: @@ -4274,6 +4275,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-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 19c57820e85..bc831a5afac 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1592,6 +1592,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_ROLL: return ggml_is_contiguous(op->src[0]); case GGML_OP_FLASH_ATTN_EXT: + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; walking the pool in physical order here would be silently wrong + if (op->src[5] != NULL) { + return false; + } // for new head sizes, add checks here if (op->src[0]->ne[0] != 32 && op->src[0]->ne[0] != 40 && 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..d27b6212b4a 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7842,6 +7842,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_MEAN: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return false; + } // The E17 compilers segfault while building FA kernels, skip E17 for now if (adreno_e17_compiler_quirks(backend_ctx)) { return false; @@ -11332,6 +11336,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..22b51a306ff 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1128,6 +1128,10 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { break; } case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return true; + } float scale = 1.0f; float max_bias = 0.0f; float logit_softcap = 0.0f; @@ -1454,6 +1458,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..7c51ddd8206 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1915,7 +1915,10 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { GGML_UNUSED(dev); - GGML_UNUSED(op); + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads; the remote end is not asked, so it is not claimed here + if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) { + return false; + } //TODO: call the remote backend and cache the results return true; } @@ -1945,6 +1948,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..82aa1d7e273 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6342,7 +6342,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K; case GGML_OP_FLASH_ATTN_EXT: - return ggml_sycl_flash_attn_ext_supported(device, op); + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op); default: return false; } @@ -6448,6 +6449,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..c14e3bc8864 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18192,6 +18192,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + return false; + } bool coopmat2 = device->coopmat2; uint32_t HSK = op->src[1]->ne[0]; uint32_t HSV = op->src[2]->ne[0]; @@ -18798,6 +18802,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..98a2298964d 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4408,6 +4408,12 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_FLASH_ATTN_EXT: { + // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads + if (op->src[5]) { + supports_op = false; + break; + } + // conservative support checks for whether the more resource-intensive shader paths // can be used, to avoid cases where flash_attn is assigned to the CPU later on supports_op = src0->type == GGML_TYPE_F32 && @@ -4661,6 +4667,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..13362a7f098 100644 --- a/include/llama.h +++ b/include/llama.h @@ -795,6 +795,17 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); + // [TAG_EXACT_CONCURRENCY] cells the memory allocates in one indivisible unit: 1 ordinarily, larger where a mode places cells in blocks, when n tokens occupy round_up(n, granularity) + LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); + + // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or 1 plus the draft length. Never lowered; false when a column bound cannot cover it. + LLAMA_API bool llama_set_exact_decode_tokens(uint32_t n_tokens); + LLAMA_API uint32_t llama_exact_decode_tokens(void); + + // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns; never lowered, and false when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is below it + LLAMA_API bool llama_set_exact_decode_width(uint32_t n_cols); + LLAMA_API uint32_t llama_exact_decode_width(void); + // // State / sessions // @@ -927,6 +938,46 @@ extern "C" { llama_seq_id dest_seq_id, llama_state_seq_flags flags); + // [TAG_STATE_ASYNC] asynchronous per-sequence state transfer, polled with llama_state_seq_copy_done(). + // Until it completes the caller must not touch the buffer, free the cells read, or decode what is written. + struct llama_state_seq_copy; + + // NULL when the backends cannot copy asynchronously, or cannot say whether a copy has finished without waiting for it; the caller then uses the synchronous 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 is far too slow to redo per transfer, so only llama_state_seq_copy_buf_free() frees it. + 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); + 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 held right now is page-locked. False while no buffer is held: ask llama_state_seq_copy_buf_can_pin() instead. + LLAMA_API bool llama_state_seq_copy_buf_is_pinned(struct llama_state_seq_copy * cpy); + + LLAMA_API bool llama_state_seq_copy_buf_can_pin(struct llama_state_seq_copy * cpy); + + // issue the copies; the bytes covered, 0 on failure. size must be within llama_state_seq_copy_buf_size(), and LLAMA_STATE_SEQ_FLAGS_ON_DEVICE is refused. + 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); + + LLAMA_API size_t llama_state_seq_copy_n_copies(struct llama_state_seq_copy * cpy); + + LLAMA_API int64_t llama_state_seq_copy_sync_us(struct llama_state_seq_copy * cpy); + + 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-batch.cpp b/src/llama-batch.cpp index 2b98a552f48..561adccd530 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -507,7 +507,36 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) { +bool llama_batch_allocr::has_shared_tokens() const { + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] > 1) { + return true; + } + } + + return false; +} + +bool llama_batch_allocr::has_seq_wider_than(uint32_t n_tokens) const { + std::vector n_per_seq(n_seq_max, 0); + + for (int32_t i = 0; i < batch.n_tokens; ++i) { + // tokens already placed in an earlier ubatch do not make the rest of the batch a prompt + if (used[i]) { + continue; + } + + for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { + if (++n_per_seq[batch.seq_id[i][s]] > n_tokens) { + return true; + } + } + } + + return false; +} + +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -518,6 +547,9 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, llama_seq_id last_seq_id = -1; + // [TAG_EXACT_CONCURRENCY] tokens left in the first set taken, when isolating: only sets with the same count join it, so every set in the ubatch finishes in it + uint32_t n_left_first = 0; + // determine the non-overlapping sequence sets participating in this ubatch for (int32_t i = 0; i < batch.n_tokens; ++i) { if (used[i]) { @@ -540,6 +572,38 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, } if (add) { + // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a prompt and gets a ubatch of its own; grouped sets need equal tokens left, or the expansion below changes their sum order + if (isolate_seqs_above > 0) { + uint32_t n_left = 0; + + for (const auto idx : seq_set_map[seq_set[i]]) { + if (!used[idx]) { + ++n_left; + } + } + + if (n_left > isolate_seqs_above) { + if (!cur_seq_set.empty()) { + // let the sets already taken have this ubatch; the prompt gets the next one + break; + } + + cur_seq_set.push_back(seq_set[i]); + + last_seq_id = batch.seq_id[i][0]; + + break; + } + + if (cur_seq_set.empty()) { + n_left_first = n_left; + } else if (n_left != n_left_first) { + continue; + } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) { + break; + } + } + cur_seq_set.push_back(seq_set[i]); last_seq_id = batch.seq_id[i][0]; diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a0..edb01045b10 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,7 +105,13 @@ class llama_batch_allocr { // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail); + // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this many tokens left is a prompt and gets a ubatch of its own + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0); + + // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place, i.e. what remains of the batch holds a prompt + bool has_seq_wider_than(uint32_t n_tokens) const; + + bool has_shared_tokens() const; // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 66940d4fc61..a62d2dbd534 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 @@ -101,6 +102,13 @@ llama_context::llama_context( throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ)); } + // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build, reported so a backend that splits columns covers it; reported at the end of the constructor + if (llama_exact_concurrency()) { + if (!llama_exact_check_n_seq(cparams.n_seq_max)) { + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } + } + cparams.n_rs_seq = params.n_rs_seq; if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) { LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n", @@ -393,6 +401,12 @@ llama_context::llama_context( }; memory.reset(model.create_memory(params_mem, cparams)); + + // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a non-causal context with a cache would assert on its first graph + if (llama_exact_concurrency() && memory && !cparams.causal_attn) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__); + throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache"); + } } // init backends @@ -476,12 +490,24 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + + // [TAG_EXACT_CONCURRENCY] nothing above can fail now, so publish the width; a refusal here means the bound moved + if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) { + throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width"); + } } 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 its KV buffers are about to go + 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]; @@ -1188,6 +1214,11 @@ void llama_context::set_causal_attn(bool value) { return; } + if (!value && memory && llama_exact_concurrency()) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__); + return; + } + cparams.causal_attn = value; sched_need_reserve = true; @@ -1577,6 +1608,10 @@ int llama_context::encode(const llama_batch & batch_inp) { } } + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -2022,6 +2057,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 +2597,132 @@ 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; the restore side asks for one per cell, and the transposed V layout repeats every run once per 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] a comb of equal runs at a constant stride is one strided copy: sequences sharing a unified cache take their cells in turn +template +static void llama_io_emit(const std::vector & infos, size_t first, size_t last, emit_t emit) { + 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); + if (deferred) { + return; // [TAG_STATE_ASYNC] the derived class posts the copies itself } + + 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 { @@ -2597,10 +2752,8 @@ class llama_io_write_host : public llama_io_write_i { return size_written; } -private: - uint8_t * ptr; - size_t buf_size = 0; - size_t size_written = 0; +protected: + llama_io_write_host(uint8_t * p, size_t len, bool deferred) : ptr(p), buf_size(len), deferred(deferred) {} struct write_info { ggml_tensor * tensor; @@ -2609,6 +2762,12 @@ class llama_io_write_host : public llama_io_write_i { size_t offset; }; std::vector winfos; + +private: + uint8_t * ptr; + size_t buf_size = 0; + size_t size_written = 0; + const bool deferred = false; }; class llama_io_read_host : public llama_io_read_i { @@ -2616,6 +2775,10 @@ class llama_io_read_host : public llama_io_read_i { llama_io_read_host(const uint8_t * p, size_t len) : ptr(p), buf_size(len) {} ~llama_io_read_host() { + if (deferred) { + return; // [TAG_STATE_ASYNC] the derived class posts the copies itself + } + // flush the reads for (size_t i = 0; i < rinfos.size();) { auto * tensor = rinfos[i].tensor; @@ -2623,19 +2786,23 @@ class llama_io_read_host : public llama_io_read_i { while (end < rinfos.size() && rinfos[end].tensor == tensor) { end++; } + // [TAG_STATE_COALESCE] the restore emits one fragment per cell, but the cost is the number of runs of adjacent cells, so count runs before falling back to staging const size_t tensor_bytes = ggml_nbytes(tensor); auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - // A fragmented sequence can require thousands of synchronous device - // transfers per layer. For bounded tensors, stage the tensor once and - // preserve every byte belonging to other sequences. Bound scratch RAM - // and leave ordinary contiguous transfers on their original fast path. - if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && + + 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; + }); + if (n_runs >= 64 && tensor_bytes <= 64 * 1024 * 1024 && !ggml_backend_buffer_is_host(buffer)) { std::vector staging; try { staging.resize(tensor_bytes); } catch (const std::bad_alloc &) { - // Fall back to the individual transfers below. } if (!staging.empty()) { ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); @@ -2649,10 +2816,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; } } @@ -2683,10 +2853,8 @@ class llama_io_read_host : public llama_io_read_i { return size_read; } -private: - const uint8_t * ptr; - size_t buf_size = 0; - size_t size_read = 0; +protected: + llama_io_read_host(const uint8_t * p, size_t len, bool deferred) : ptr(p), buf_size(len), deferred(deferred) {} struct read_info { ggml_tensor * tensor; @@ -2695,6 +2863,12 @@ class llama_io_read_host : public llama_io_read_i { size_t offset; }; std::vector rinfos; + +private: + const uint8_t * ptr; + size_t buf_size = 0; + size_t size_read = 0; + const bool deferred = false; }; class llama_io_write_file : public llama_io_write_i { @@ -2998,6 +3172,273 @@ 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: the host buffer, one backend per device, each with its own stream, and one event per device +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; + + size_t n_copies = 0; + 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 none: a host tensor is a memcpy, and a split buffer fails every backend's async copy assert + 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(); + } + + void record() { + + for (auto & it : devs) { + if (it.second.pending) { + ggml_backend_event_record(it.second.event, it.second.backend.get()); + } + } + } + + // order the copies behind the compute already queued on each device: the copy stream waits for the context's fence, recorded at the end of every decode + 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, for a restore only: its copies write KV cells while other sequences read every cell up to n_kv + 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 costs about as long as the copy it is for, and a caller parking the same sequence asks for a slightly different size each time + 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 off, so believe the buffer that came back rather than the type + 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; + } + + 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(); + } +}; + +// [TAG_STATE_ASYNC] the buffer walk of llama_io_write_host, with the copies posted on the transfer's stream instead of made here +class llama_io_write_host_async : public llama_io_write_host { +public: + llama_io_write_host_async(uint8_t * p, size_t len, llama_state_seq_copy & cpy) : + llama_io_write_host(p, len, true), cpy(cpy) {} + + // posted from the destructor, and only once serialisation reached the end: a caller told of a partial failure by a zero return is free to reuse the buffer at once + 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(); + } + +private: + llama_state_seq_copy & cpy; + + bool committed = false; +}; + +// [TAG_STATE_ASYNC] the read half of the same, without llama_io_read_host's whole-tensor staging: a write-back would undo whatever the sequences sharing the tensor wrote while these copies ran +class llama_io_read_host_async : public llama_io_read_host { +public: + llama_io_read_host_async(const uint8_t * p, size_t len, llama_state_seq_copy & cpy) : + llama_io_read_host(p, len, true), cpy(cpy) {} + + // see llama_io_write_host_async::commit(): a restore that failed part way has dropped the sequence, and copies posted for it would write cells that are no longer its own + void commit() { + committed = true; + } + + ~llama_io_read_host_async() { + if (!committed) { + return; + } + + 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(); + } + +private: + 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 +3512,237 @@ 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 makes the first poll wait for the whole copy, so leave it out and let state_seq_copy_init() return NULL + 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 using, so a transfer could end up ordered behind a graph + 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 the ones the state lives on: with most layers on the CPU every copy takes the synchronous branch of backend_for() + 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 and recorded after every decode; installed only after the checks above, so a refused transfer leaves nothing behind + 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); + } + + 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) { + // the library owns this buffer, so the extent can be checked instead of believed: every bounds check validates against it, so an oversized one agrees and the copy overruns + 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 has nowhere to leave the data here, and get_size_ext() with that flag reports a metadata-sized state, so the two cannot be paired + 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; + } + + 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) { + 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; + } + + 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, masked, by a graph in flight, so the copy stream waits for the compute stream on the device, see order_after() + 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; + } + } + + 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"); @@ -3226,6 +3898,11 @@ size_t llama_context::state_write_data(llama_io_write_i & io) { } size_t llama_context::state_read_data(llama_io_read_i & io) { + // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical index, which the paged pool owns; refused before anything is parsed + if (memory && memory->alloc_granularity() > 1) { + throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence"); + } + LLAMA_LOG_DEBUG("%s: reading state\n", __func__); // read model info @@ -4030,6 +4707,14 @@ bool llama_memory_can_shift(llama_memory_t mem) { return mem->get_can_shift(); } +uint32_t llama_memory_alloc_granularity(llama_memory_t mem) { + if (!mem) { + return 1; + } + + return mem->alloc_granularity(); +} + // llama state API // deprecated @@ -4125,6 +4810,66 @@ 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); } +llama_state_seq_copy * llama_state_seq_copy_init(llama_context * ctx) { + return ctx->state_seq_copy_init(); +} + +void llama_state_seq_copy_free(llama_state_seq_copy * cpy) { + delete cpy; // waits for anything still in flight +} + +uint8_t * llama_state_seq_copy_buf_resize(llama_state_seq_copy * cpy, size_t size) { + return cpy->buf_resize(size); +} + +uint8_t * llama_state_seq_copy_buf(llama_state_seq_copy * cpy) { + return cpy->data; +} + +size_t llama_state_seq_copy_buf_size(llama_state_seq_copy * cpy) { + return cpy->size; +} + +size_t llama_state_seq_copy_buf_capacity(llama_state_seq_copy * cpy) { + return cpy->capacity; +} + +size_t llama_state_seq_copy_n_copies(llama_state_seq_copy * cpy) { + return cpy->n_copies; +} + +int64_t llama_state_seq_copy_sync_us(llama_state_seq_copy * cpy) { + return cpy->t_sync_us; +} + +void llama_state_seq_copy_buf_free(llama_state_seq_copy * cpy) { + cpy->buf_free(); +} + +bool llama_state_seq_copy_buf_is_pinned(llama_state_seq_copy * cpy) { + return cpy->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..f44f505a05f 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,8 @@ struct llama_memory_buffer { using llama_memory_buffers = std::map; +struct llama_state_seq_copy; + struct llama_context { // init scheduler and compute buffers, reserve worst-case graphs llama_context( @@ -156,6 +159,19 @@ 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(); + + void state_seq_copy_release(llama_state_seq_copy * cpy); + + void state_seq_copies_drain(); + bool state_load_file( const char * filepath, llama_token * tokens_out, @@ -348,6 +364,12 @@ 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, 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/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..e9fd4ecd6e8 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -21,11 +21,25 @@ #include #include #include +#include #include #include // dedup helpers +// [TAG_EXACT_CONCURRENCY] the page table is wired into llm_graph_input_attn_kv only, so a V-less layout would attend in physical order with the mode reporting itself on +static void llm_graph_reject_exact_concurrency(const char * layout) { + if (!llama_exact_concurrency()) { + return; + } + + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, but this model uses the %s attention " + "layout, which carries no page table and would attend in physical cell order\n", + __func__, layout); + + throw std::runtime_error("exact concurrency: unsupported attention layout"); +} + static ggml_tensor * build_attn_inp_kq_mask( ggml_context * ctx, const llama_kv_cache_context * mctx, @@ -468,6 +482,7 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { + if (self_pages && self_pages->buffer) { mctx->set_input_pages(self_pages, ubatch); } mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch); @@ -1084,6 +1099,7 @@ void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { + if (inp_attn->self_pages) { mctx->get_attn()->set_input_pages(inp_attn->self_pages, ubatch); } mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); @@ -2547,7 +2563,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, - int il) const { + int il, + ggml_tensor * pages) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -2580,6 +2597,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + cur->src[5] = pages; res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); @@ -2769,6 +2787,8 @@ static std::unique_ptr build_attn_inp_kv_impl( inp->self_kq_mask_cnv = inp->self_kq_mask; } + inp->self_pages = mctx_cur->build_input_pages(ctx0, ubatch); + GGML_ASSERT(!inp->self_pages || (cparams.flash_attn && cparams.causal_attn)); inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0); inp->self_v_rot = mctx_cur->build_input_v_rot(ctx0); @@ -2831,7 +2851,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = mctx_cur->get_v(ctx0, il); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il, inp->self_pages); cb(cur, "kqv_out", il); if (inp->self_v_rot) { @@ -2865,6 +2885,8 @@ static std::unique_ptr build_attn_inp_k_impl( const llama_cparams & cparams, const llama_kv_cache_context * mctx_cur) { + llm_graph_reject_exact_concurrency("V-less KV (attn_k)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { @@ -3241,6 +3263,8 @@ static std::unique_ptr build_attn_inp_k_dsa_impl( const llama_cparams & cparams, const llama_kv_cache_dsa_context * mctx_cur) { + llm_graph_reject_exact_concurrency("sparse V-less KV (attn_k_dsa)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { @@ -3358,6 +3382,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { const auto * mctx_cur = static_cast(mctx); + llm_graph_reject_exact_concurrency("V-less sliding window KV (attn_k_iswa)"); + auto inp = std::make_unique(hparams, cparams, mctx_cur); { diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb5..26f2169532d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -319,6 +319,7 @@ class llm_graph_input_attn_no_cache : public llm_graph_input_i { class llm_graph_input_attn_kv : public llm_graph_input_i { public: + ggml_tensor * self_pages = nullptr; // I32 [1 + physical pages, n_tokens] llm_graph_input_attn_kv( const llama_hparams & hparams, const llama_cparams & cparams, @@ -1172,7 +1173,8 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, - int il) const; + int il, + ggml_tensor * pages = nullptr) const; llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index b3a94b946d2..50037a46b84 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -1,11 +1,15 @@ #include "llama-impl.h" +#include "ggml-backend.h" #include "gguf.h" #include "llama.h" #include #include +#include +#include #include +#include #include #include #include @@ -169,3 +173,142 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) { return gguf_data_to_str(type, gguf_get_val_data(ctx_gguf, i), 0); } } + +bool llama_exact_concurrency() { + static const bool enabled = []() { + const char * val = getenv("LLAMA_EXACT_CONCURRENCY"); + return val && atoi(val) != 0; + }(); + + return enabled; +} + +// [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h +static std::atomic g_exact_decode_tokens{1}; + +// one lock for the token figure, the sequence count and the width: a report interleaved with a change of figure could leave the backend with a width that covers neither +static std::recursive_mutex g_exact_mutex; + +// the most sequences any context was created with; the tokens figure is process wide, so raising it re-reports every context's width +static std::atomic g_exact_max_n_seq{0}; + +static bool llama_exact_width_within_explicit_bound(uint32_t n_cols); + +static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) { + const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens; + + if (w > (uint64_t) INT32_MAX) { + LLAMA_LOG_ERROR("%s: a decode step of %u sequences with %u tokens each is too wide to report\n", __func__, n_seq, n_tokens); + return false; + } + + n_cols = (uint32_t) w; + + return true; +} + +bool llama_exact_check_n_seq(uint32_t n_seq) { + std::lock_guard lock(g_exact_mutex); + + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); + + uint32_t n_cols = 0; + + return llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) && llama_exact_width_within_explicit_bound(n_cols); +} + +bool llama_exact_report_n_seq(uint32_t n_seq) { + std::lock_guard lock(g_exact_mutex); + + const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed)); + + uint32_t n_cols = 0; + + if (!llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) || !llama_set_exact_decode_width(n_cols)) { + return false; + } + + uint32_t cur = g_exact_max_n_seq.load(std::memory_order_relaxed); + + while (n_seq > cur && !g_exact_max_n_seq.compare_exchange_weak(cur, n_seq, std::memory_order_relaxed)) { + } + + return true; +} + +bool llama_set_exact_decode_tokens(uint32_t n_tokens) { + n_tokens = n_tokens > 0 ? n_tokens : 1; + + std::lock_guard lock(g_exact_mutex); + + // never lowered: a narrower context set up later would turn an existing speculative context's verify steps into prompts + if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) { + return true; + } + + // every context widens with the figure, so report the width first; one the explicit bound cannot cover leaves the old figure in place + const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed); + + uint32_t n_cols = 0; + + if (n_seq > 0 && (!llama_exact_width_of(n_seq, n_tokens, n_cols) || !llama_set_exact_decode_width(n_cols))) { + return false; + } + + g_exact_decode_tokens.store(n_tokens, std::memory_order_relaxed); + + return true; +} + +uint32_t llama_exact_decode_tokens(void) { + return g_exact_decode_tokens.load(std::memory_order_relaxed); +} + +// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h; reached through the registry so an absent or late-loaded backend costs nothing +static std::atomic g_exact_decode_width{0}; + +// an explicit column bound wins in the CUDA backend, so a width above it would leave decodes batched past the bound +static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) { + static const int explicit_cols = []() { + const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS"); + return val ? atoi(val) : -1; + }(); + + if (explicit_cols > 0 && (uint32_t) explicit_cols < n_cols) { + LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u columns for the decode step just requested; raise it, set it to 0 for no bound, or unset it\n", + __func__, explicit_cols, n_cols); + return false; + } + + return true; +} + +bool llama_set_exact_decode_width(uint32_t n_cols) { + if (!llama_exact_width_within_explicit_bound(n_cols)) { + return false; + } + + std::lock_guard lock(g_exact_mutex); + + uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed); + + while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) { + } + + const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed); + + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + ggml_backend_reg_t reg = ggml_backend_reg_get(i); + + auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width"); + if (fn) { + fn((int) widest); + } + } + + return true; +} + +uint32_t llama_exact_decode_width(void) { + return g_exact_decode_width.load(std::memory_order_relaxed); +} diff --git a/src/llama-impl.h b/src/llama-impl.h index 4988b06d2ca..01970f2a88c 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -103,3 +103,11 @@ std::string llama_format_tensor_shape(const std::vector & ne); std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); + +// [TAG_EXACT_CONCURRENCY] opt-in mode under which a sequence's attention depends only on its own cells, so its output does not change when others share the KV cache +bool llama_exact_concurrency(); + +// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so the backend knows the width every context needs +bool llama_exact_report_n_seq(uint32_t n_seq); + +bool llama_exact_check_n_seq(uint32_t n_seq); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a75314..8bbed8264b2 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include static bool ggml_is_power_of_2(int n) { @@ -61,6 +62,73 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { // llama_kv_cache // +// [TAG_EXACT_CONCURRENCY] the paged specialization lives in the CUDA sources; every other backend ignores src[5] and walks the pool in physical cell order +static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) { + if (!dev) { + return false; + } + + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (!reg) { + return false; + } + + const char * name = ggml_backend_reg_name(reg); + if (!name) { + return false; + } + + return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0; +} + +// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a layer of this shape: the registry name only says which backends carry the kernels +static bool llama_dev_supports_paged_attn( + ggml_backend_dev_t dev, + ggml_type type_k, ggml_type type_v, + uint32_t n_embd_head_k, uint32_t n_embd_head_v, + uint32_t n_head, uint32_t n_head_kv, + uint32_t n_cells, uint32_t page_size) { + if (!llama_dev_has_paged_attn(dev)) { + return false; + } + + ggml_init_params ip = { + /*.mem_size =*/ ggml_tensor_overhead()*16 + ggml_graph_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + return false; + } + + bool res = true; + + const int64_t n_kv = page_size; + + for (const int64_t n_tokens : { (int64_t) 1, (int64_t) 4, (int64_t) 16, (int64_t) 512 }) { + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd_head_k, n_tokens, n_head, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, n_embd_head_k, n_kv, n_head_kv, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, n_embd_head_v, n_kv, n_head_kv, 1); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_tokens, 1, 1); + + ggml_tensor * op = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf((float) n_embd_head_k), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(op, GGML_PREC_F32); + + op->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + n_cells/page_size, n_tokens); + + if (!ggml_backend_dev_supports_op(dev, op)) { + res = false; + break; + } + } + + ggml_free(ctx); + + return res; +} + llama_kv_cache::llama_kv_cache( const llama_model & model, const llama_hparams & hparams, @@ -84,6 +152,9 @@ llama_kv_cache::llama_kv_cache( v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { + // [TAG_EXACT_CONCURRENCY] read the knob through the same cached reader the graph and the CUDA dispatcher use, so a mid-process change cannot leave them disagreeing + exact_pages = llama_exact_concurrency(); + // shared cells view the source cache's K/V tensors, so the cell count // follows the source allocation: a fitted target can be smaller than the // draft default and oversized views would overflow the source tensors @@ -97,6 +168,27 @@ llama_kv_cache::llama_kv_cache( GGML_ASSERT(kv_size % n_pad == 0); + if (exact_pages) { + const char * unsupported = nullptr; + + if (!unified) { + unsupported = "it needs a unified KV cache (pass --kv-unified)"; + } else if (v_trans) { + unsupported = "it needs a non-transposed V cache (pass --flash-attn on)"; + } else if (n_swa != 0) { + unsupported = "the paged pool does not support sliding window attention"; + } else if (type_k != GGML_TYPE_F16 || type_v != GGML_TYPE_F16) { + unsupported = "it needs an F16 KV cache (do not pass --cache-type-k or --cache-type-v)"; + } else if (kv_size % exact_page_size != 0) { + unsupported = "the context size must be a multiple of 256 (pass -c as a multiple of 256)"; + } + + if (unsupported) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s\n", __func__, unsupported); + throw std::runtime_error("exact concurrency: unsupported KV cache configuration"); + } + } + const uint32_t n_layer = hparams.n_layer_all; // define a comparator for the buft -> ctx map to ensure that the order is well-defined: @@ -220,6 +312,39 @@ llama_kv_cache::llama_kv_cache( LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); + // [TAG_EXACT_CONCURRENCY] the paged kernel handles 256-wide K and V heads only; any other width would run unpaged while the mode reports itself as on + if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, " + "and the paged attention kernel supports 256-wide K and V heads only\n", + __func__, il, hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), is_mla ? " (MLA)" : ""); + throw std::runtime_error("exact concurrency: unsupported attention head size"); + } + + if (exact_pages && hparams.attn_soft_cap) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), " + "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping); + throw std::runtime_error("exact concurrency: attention soft cap is not supported"); + } + + if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, " + "which has no paged attention: every layer must be offloaded to the CUDA backend " + "(pass -ngl to offload all layers and do not pass --no-kv-offload)\n", + __func__, il, dev_name); + throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend"); + } + + // [TAG_EXACT_CONCURRENCY] right backend; ask whether this layer's attention, with the page table attached, lands on one of its kernels at all + if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v, + hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), + hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s cannot run the paged attention for layer %d " + "(K %s, V %s, %u-wide heads): the build or the device has no flash attention kernel for it, " + "and the op would fall to the CPU, which ignores the page table\n", + __func__, dev_name, il, ggml_type_name(type_k), ggml_type_name(type_v), hparams.n_embd_head_k(il)); + throw std::runtime_error("exact concurrency: the device cannot run the paged attention"); + } + ggml_context * ctx = ctx_for_buft(buft); if (!ctx) { throw std::runtime_error("failed to create ggml context for kv cache"); @@ -364,7 +489,71 @@ llama_kv_cache::llama_kv_cache( debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0; } +void llama_kv_cache::exact_pages_rebuild() const { + const auto & cells = v_cells[0]; + + exact_page_owner.assign(cells.size()/exact_page_size, exact_page{}); + + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { + continue; + } + + GGML_ASSERT(cells.seq_count(i) == 1); + + const auto pos = cells.pos_get(i); + + GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size); + + const exact_page cur { cells.seq_get(i), llama_pos(pos/(llama_pos) exact_page_size) }; + + auto & owner = exact_page_owner[i/exact_page_size]; + + GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg)); + + owner = cur; + } + + exact_page_owner_dirty = false; +} + +void llama_kv_cache::exact_pages_sync() const { + if (exact_page_owner_dirty) { + exact_pages_rebuild(); + + return; + } + + if (debug > 0) { + const auto kept = exact_page_owner; + + exact_pages_rebuild(); + + GGML_ASSERT(kept.size() == exact_page_owner.size()); + + for (size_t p = 0; p < kept.size(); ++p) { + GGML_ASSERT(kept[p].seq == exact_page_owner[p].seq && kept[p].lpg == exact_page_owner[p].lpg); + } + } +} + +void llama_kv_cache::exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos) { + if (exact_page_owner_dirty || exact_page_owner.empty()) { + return; + } + + const exact_page cur { seq, llama_pos(pos/(llama_pos) exact_page_size) }; + + auto & owner = exact_page_owner[idx/exact_page_size]; + + GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg)); + + owner = cur; +} + void llama_kv_cache::clear(bool data) { + exact_page_owner_dirty = true; + for (uint32_t s = 0; s < n_stream; ++s) { v_cells[s].reset(); v_heads[s] = 0; @@ -383,6 +572,9 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { return true; } + // [TAG_EXACT_CONCURRENCY] a removal can empty a page, which only the cells know + exact_page_owner_dirty = true; + // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); @@ -452,6 +644,14 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll return; } + // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so refuse a copy that would share cells rather than abort. After the shared-cells return, so a draft cache is unaffected. + if (exact_pages && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between " + "sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size()); GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size()); @@ -544,6 +744,8 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { return; } + exact_page_owner_dirty = true; + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); auto & cells = v_cells[seq_to_stream[seq_id]]; @@ -571,6 +773,14 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll return; } + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo the page size, so shifting positions would misplace every cell + if (exact_pages && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions " + "(seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1"); @@ -621,6 +831,13 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in return; } + if (exact_pages && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions " + "(seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1"); @@ -704,11 +921,23 @@ llama_memory_context_ptr llama_kv_cache::init_batch( GGML_UNUSED(embd_all); do { + // [TAG_EXACT_CONCURRENCY] a token shared by several sequences would be one cell in a page that belongs to one sequence, so refuse it here rather than assert at placement + if (exact_pages && balloc.has_shared_tokens()) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " + "give every token exactly one sequence id\n", __func__); + break; + } + balloc.split_reset(); std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); + // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt into one ubatch, so a prefill would run at a width its solo run never sees; the set split gives each its own + const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0; + + auto ubatch = n_stream == 1 && !isolate + ? balloc.split_simple(n_ubatch) + : balloc.split_equal(n_ubatch, n_stream > 1, 0, isolate); if (ubatch.n_tokens == 0) { break; @@ -755,6 +984,9 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector v_heads_old; // old positions of the heads, before placing the ubatch std::vector v_cells; // copy of the old cells, before placing the ubatch + + // [TAG_EXACT_CONCURRENCY] page ownership before the ubatch, so undoing a speculative placement does not force a rebuild from every cell + std::vector exact_page_owner_old; }; // remember the old state of the cells so we can restore it in the end @@ -775,7 +1007,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vectorv_cells[s]); head = it->v_heads_old[s]; } + + // [TAG_EXACT_CONCURRENCY] put back what the allocator knew, unless the placement also removed cells, when only the cells can say what is left + if (!exact_page_owner_dirty) { + exact_page_owner = it->exact_page_owner_old; + } } if (!success) { @@ -961,6 +1198,45 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, } } + if (exact_pages) { + const auto & cells = v_cells[0]; + + exact_pages_sync(); + + using page_key = std::pair; + + std::vector owner = exact_page_owner; + std::map pages; + + for (uint32_t p = 0; p < owner.size(); ++p) { + if (owner[p].seq >= 0) { + pages.emplace(page_key {owner[p].seq, owner[p].lpg}, p); + } + } + + std::set assigned; + slot_info res {0, 0, {0}, {{}}}; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1 && ubatch.pos[i] >= 0); + const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size}; + auto it = pages.find(key); + if (it == pages.end()) { + uint32_t page = v_heads[0]/exact_page_size; + uint32_t tested = 0; + while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; } + if (tested == owner.size()) { return {}; } + page %= owner.size(); + owner[page] = exact_page {key.first, key.second}; + it = pages.emplace(key, page).first; + } + const uint32_t idx = it->second*exact_page_size + ubatch.pos[i]%exact_page_size; + if (!cells.is_empty(idx) || !assigned.insert(idx).second) { return {}; } + res.idxs[0].push_back(idx); + } + if (cont && !res.is_contiguous()) { return {}; } + return res; + } + uint32_t n_tokens = ubatch.n_tokens; uint32_t n_seqs = 1; @@ -1139,6 +1415,12 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) { cells.seq_add(idx, ubatch.seq_id[i][s]); } + + if (exact_pages) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1); + + exact_pages_claim(idx, ubatch.seq_id[i][0], ubatch.pos[i]); + } } } @@ -1170,7 +1452,16 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } } +uint32_t llama_kv_cache::alloc_granularity() const { + // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so n tokens hold round_up(n, exact_page_size) cells: the tail page is charged in full + return exact_pages ? exact_page_size : 1; +} + bool llama_kv_cache::get_can_shift() const { + // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo 256, so the pool cannot shift positions; reporting it disables --context-shift and --cache-reuse at load + if (exact_pages) { + return false; + } // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot. if (model.arch == LLM_ARCH_STEP35) { return false; @@ -1232,7 +1523,50 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const { return v_cells[seq_to_stream[seq_id]]; } +ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const { + if (!exact_pages) { return nullptr; } + auto * pages = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + get_size()/exact_page_size, ubatch.n_tokens); + ggml_set_input(pages); + ggml_set_name(pages, "attn_logical_pages"); + return pages; +} + +void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { + GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens); + + exact_pages_sync(); + + std::map> pages; + for (uint32_t p = 0; p < exact_page_owner.size(); ++p) { + const auto & owner = exact_page_owner[p]; + if (owner.seq >= 0) { + pages[owner.seq][owner.lpg] = p; + } + } + std::vector data(ggml_nelements(dst), -1); + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + GGML_ASSERT(ubatch->n_seq_id[i] == 1); + auto * row = data.data() + i*dst->ne[0]; + row[0] = 0; + for (const auto & page : pages[ubatch->seq_id[i][0]]) { + if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; } + row[++row[0]] = page.second; + } + } + ggml_backend_tensor_set(dst, data.data(), 0, data.size()*sizeof(int32_t)); +} + +ggml_tensor * llama_kv_cache_context::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const { + return kv->build_input_pages(ctx, ubatch); +} + +void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const { + kv->set_input_pages(dst, ubatch); +} + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { + // the per-query page map is the only loop bound for exact attention, so neighbours cannot extend it + if (exact_pages) { return get_size(); } uint32_t result = 0; // pad the n_kv value so that the graph remains constant across batches and can be reused @@ -2037,6 +2371,12 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical index, which the paged pool owns; refused before a byte is read + if (exact_pages && seq_id == -1) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__); + throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY"); + } + // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -2285,6 +2625,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 } else { // whole KV cache restore + GGML_ASSERT(!exact_pages); + if (cell_count > cells.size()) { LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__); return false; diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f98..c36be8c0154 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -131,6 +131,9 @@ class llama_kv_cache : public llama_memory_i { bool get_can_shift() const override; + // [TAG_EXACT_CONCURRENCY] the page size under exact mode, 1 otherwise + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; @@ -171,6 +174,8 @@ class llama_kv_cache : public llama_memory_i { // uint32_t get_n_kv(const slot_info & sinfo) const; + ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const; + void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const; // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; @@ -235,6 +240,24 @@ class llama_kv_cache : public llama_memory_i { std::vector v_stream; }; + static constexpr uint32_t exact_page_size = 256; + bool exact_pages = false; + + // [TAG_EXACT_CONCURRENCY] which (sequence, logical page) owns each physical page; seq < 0 means free, and it is kept current as cells are placed and dirtied by removals + struct exact_page { + llama_seq_id seq = -1; + llama_pos lpg = -1; + }; + + mutable std::vector exact_page_owner; + mutable bool exact_page_owner_dirty = true; + + void exact_pages_sync() const; + + void exact_pages_rebuild() const; + + void exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos); + bool v_trans = true; // the value tensor is transposed const uint32_t n_seq_max = 1; @@ -365,6 +388,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // uint32_t get_n_kv() const; + ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const; + void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const; ggml_type type_k() const; ggml_type type_v() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..a596f35bd8d 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -66,6 +66,13 @@ llama_memory_hybrid::llama_memory_hybrid( llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { + // [TAG_EXACT_CONCURRENCY] refused before the attention half asserts on it, see llama_kv_cache::init_batch + if (llama_exact_concurrency() && balloc.has_shared_tokens()) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; " + "give every token exactly one sequence id\n", __func__); + break; + } + balloc.split_reset(); // follow the recurrent pattern for creating the ubatch splits @@ -86,7 +93,10 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // so that the rollback snapshots remain valid const uint32_t n_rs_seq = mem_recr->n_rs_seq; - ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + // [TAG_EXACT_CONCURRENCY] the recurrent half is not invariant to the ubatch shape, so a prompt gets a ubatch of its own + const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } if (ubatch.n_tokens == 0) { @@ -135,6 +145,11 @@ bool llama_memory_hybrid::get_can_shift() const { return mem_attn->get_can_shift(); } +uint32_t llama_memory_hybrid::alloc_granularity() const { + // the recurrent half holds one state per sequence, so the attention half is the one whose cells a caller is planning capacity for + return mem_attn->alloc_granularity(); +} + void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); mem_recr->clear(data); @@ -150,6 +165,13 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + // [TAG_EXACT_CONCURRENCY] the attention half refuses this, so refuse before either half is touched or the two could end up describing different states + if (llama_exact_concurrency() && seq_id_src != seq_id_dst) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n", + __func__, seq_id_src, seq_id_dst); + return; + } + mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); } @@ -160,11 +182,23 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + if (llama_exact_concurrency() && shift != 0) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions (seq %d, shift %d); ignoring the shift\n", + __func__, seq_id, shift); + return; + } + mem_attn->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + if (llama_exact_concurrency() && d != 1) { + LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions (seq %d, d %d); ignoring the division\n", + __func__, seq_id, d); + return; + } + mem_attn->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); } diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb7499..70ba19ca323 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -58,6 +58,8 @@ class llama_memory_hybrid : public llama_memory_i { bool get_can_shift() const override; + uint32_t alloc_granularity() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7..d6ac2e55358 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -431,7 +431,10 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // [TAG_RECURRENT_ROLLBACK_SPLITS] // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch // so that the rollback snapshots remain valid - ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: the state a prompt leaves behind depends on what shared its ubatch, so isolate prompts + const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0; + + ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate); } if (ubatch.n_tokens == 0) { diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645..61cd348f2dd 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,6 +100,9 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; + // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode allocates in larger blocks, when n tokens occupy round_up(n, granularity) cells. Not pure, so old modules inherit 1. + virtual uint32_t alloc_granularity() const { return 1; } + // // ops // 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-backend-ops.cpp b/tests/test-backend-ops.cpp index 53e93a1448d..b86f2703aea 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7193,6 +7193,40 @@ struct test_flash_attn_ext : public test_case { } }; +// same attention as the CPU mask reference, but visiting nonadjacent pages in a different order +struct test_flash_attn_ext_pages : public test_flash_attn_ext { + test_flash_attn_ext_pages(int64_t batch) : + test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {} + + std::string vars() override { return test_flash_attn_ext::vars() + ",exact_pages=1"; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + auto * out = test_flash_attn_ext::build_graph(ctx); + out->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 5, nb); + ggml_set_name(out->src[5], "pages"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + test_flash_attn_ext::initialize_tensors(ctx); + auto * pages = ggml_get_tensor(ctx, "pages"); + auto * mask = ggml_get_tensor(ctx, "m"); + std::vector ids(5*nb, -1); + std::vector values(1024*nb, ggml_fp32_to_fp16(-INFINITY)); + for (int64_t q = 0; q < nb; ++q) { + ids[5*q] = q%2 ? 1 : 2; + ids[5*q + 1] = 2; + ids[5*q + 2] = 0; + for (int j = 0; j < 256; ++j) { values[1024*q + 512 + j] = ggml_fp32_to_fp16(0.0f); } + if (q%2 == 0) { + for (int j = 0; j < 17; ++j) { values[1024*q + j] = ggml_fp32_to_fp16(0.0f); } + } + } + ggml_backend_tensor_set(pages, ids.data(), 0, ids.size()*sizeof(int32_t)); + ggml_backend_tensor_set(mask, values.data(), 0, values.size()*sizeof(ggml_fp16_t)); + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; @@ -9170,6 +9204,21 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) { + for (int n : {1, 17, 307}) { + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1})); + test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {1, 4})); + } + } + + // MoE projections at the token counts a decode ubatch forms; 17 tokens is past the width exact concurrency pins + for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) { + for (int n : {1, 2, 4, 8, 17}) { + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true, 512, n, 2048)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, false, 2048, n, 512)); + } + } + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); @@ -9937,6 +9986,9 @@ static std::vector> make_test_cases_eval() { } // mixed quant and Q1_0 test cases + for (int64_t batch : {1, 4, 12}) { + test_cases.emplace_back(new test_flash_attn_ext_pages(batch)); + } test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0)); diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index 428a9252981..5a1502f747f 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,8 +73,7 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); - // A fragmented restore may stage a whole device tensor. Check every - // sequence byte-for-byte, including the neighbours that must be preserved. + // a fragmented restore may stage a whole device tensor, so check every sequence byte-for-byte, neighbours included std::vector> before(params.n_parallel); for (int s = 0; s < params.n_parallel; ++s) { before[s].resize(llama_state_seq_get_size(ctx, s)); diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp new file mode 100644 index 00000000000..bd3cd9cc30a --- /dev/null +++ b/tests/test-state-seq-copy.cpp @@ -0,0 +1,140 @@ +// [TAG_STATE_ASYNC] guards on the asynchronous state transfer: the buffer belongs to the transfer, so an oversized size is refused, and ON_DEVICE is refused as these go via the host + +#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; + } + + // 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"); + + 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); + + 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); + + 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); + + 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 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__); + + std::vector before(llama_state_seq_get_size(ctx, seq_id)); + CHECK(llama_state_seq_get_data(ctx, before.data(), before.size(), seq_id) == before.size()); + + CHECK(llama_state_seq_copy_get(cpy, size, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) == size); + llama_state_seq_copy_wait(cpy); + + llama_memory_seq_rm(llama_get_memory(ctx), seq_id, -1, -1); + + CHECK(llama_state_seq_copy_set(cpy, size - 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); + + 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-common.h b/tools/server/server-common.h index f0cf76b8c50..423e9982f1f 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,7 +467,6 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; - // [TAG_PREEMPT] slots parked to make room in the unified KV pool, and put back uint64_t n_preempt = 0; uint64_t n_resume = 0; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b18fa4e2f23..c624c2cee57 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -18,6 +18,7 @@ #include "mtmd-helper.h" #include +#include #include #include #include @@ -61,36 +62,63 @@ 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 -// -// With --kv-unified the cells are one pool shared by every slot, and each slot believes it -// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves -// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size -// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the -// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and -// continue with the rest". -// -// Nothing is terminated here. The cells of one slot are taken back and given to it again -// later: its sequence is copied to host RAM, its cells are released, and when the pool has -// room the copy goes back and the slot carries on with the same sampler, the same generated -// text and the same open stream. A streaming client sees a pause, not an error. +// [TAG_PREEMPT] server-side request preemption: instead of ending every conversation in flight with a context error, one slot's sequence is copied to host RAM and back when there is room constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation +constexpr int64_t PREEMPT_KEEPALIVE_MS = 2000; // SSE keepalive period while a streaming slot is parked constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected -// [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody -// passes a head that does not fit yet: the head keeps the room the pool frees until it fits, so -// its wait is bounded by the slots ahead of it and not by how often a smaller slot can squeeze -// in, grow, and be parked again. Simulated over 60 seeds at eight chats this cuts the longest -// single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often. -// LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest -// parked, and a smaller slot may pass a head that does not fit. -// LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. +static std::string preempt_notice_comment(const server_task_result_preempt_notice & notice) { + std::string res = notice.parked ? ": preempted" : ": resumed"; + + if (notice.index > 0) { + res += " " + std::to_string(notice.index); + } + + return res + "\n\n"; +} constexpr int32_t PREEMPT_N_FAIL_MAX = 8; // failed restores before the slot is given up on constexpr int64_t PREEMPT_FAIL_US = 60ll * 1000 * 1000; // ... and only after this long parked 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] an asynchronous park only releases its cells when its copy lands, so it must fire this many decode steps before the pool would run out +constexpr int32_t PREEMPT_N_ASYNC_STEPS = 8; + +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_free) : llama_state_seq_copy_ptr(); +} + +// [TAG_EXACT_CONCURRENCY] the planner counts cells, not tokens: a page belongs to one sequence, so a token count sees room find_slot cannot find and nobody is ever parked + +static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) { + return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g; +} + +static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) { + return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g); +} + +static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 && + preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3, + "at a granularity of 1 a run of n tokens has to cost exactly n cells"); +static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(8191, 1, 1) == 1 && + preempt_n_cells_step_g(1000, 512, 1) == 512, + "at a granularity of 1 a step of n tokens has to cost exactly n cells"); + +static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 && + preempt_n_cells_g(257, 256) == 512, + "a tail page is charged in full"); +static_assert(preempt_n_cells_step_g(255, 1, 256) == 0 && preempt_n_cells_step_g(256, 1, 256) == 256 && + preempt_n_cells_step_g(256, 257, 256) == 512, + "a step is free until it crosses a page boundary and costs whole pages when it does"); + struct server_slot; // forward declaration struct server_batch { @@ -323,42 +351,191 @@ struct server_slot { prompt.clear(); } - // [TAG_PREEMPT] state of a slot whose cells were taken back - // - // Only the KV cells leave. The task, the sampler, the generated text and the position - // the stream has reached stay on the slot, so a resume is a memcpy and not a new - // request: no retokenisation, no replayed prompt, no seam in the output. 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, and are shared_ptr only so a slot survives the vector's reallocation + 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; + } + + template + auto preempt_sum(F f) const -> decltype(f(preempt_cpy_tgt.get())) { + if (!preempt_is_async()) { + return 0; + } + + return f(preempt_cpy_tgt.get()) + (preempt_cpy_dft ? f(preempt_cpy_dft.get()) : 0); + } + + template + void preempt_each(F f) const { + if (preempt_cpy_tgt) { + f(preempt_cpy_tgt.get()); + } + + if (preempt_cpy_dft) { + f(preempt_cpy_dft.get()); + } + } + + int64_t preempt_sync_us() const { + return preempt_sum(llama_state_seq_copy_sync_us); + } + + size_t preempt_n_copies() const { + return preempt_sum(llama_state_seq_copy_n_copies); + } + + // [TAG_PREEMPT_ASYNC] a copy is running: the slot must not be scheduled but still owns cells, so it is neither running nor parked + bool preempt_in_flight() const { + return state == SLOT_STATE_PREEMPTING || state == SLOT_STATE_RESTORING; + } + + 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_ctx_shift = 0; // context shifts it 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 { - return preempt_state_tgt.size() + preempt_state_dft.size(); + // for a transfer the capacity, not the live size: the pinned buffers are kept between parks, so --preempt-ram has to bound what is held + return preempt_is_async() ? preempt_sum(llama_state_seq_copy_buf_capacity) + : preempt_state_tgt.size() + preempt_state_dft.size(); } void preempt_state_free() { + // waits for anything in flight first: release() is reached with a copy possibly still using the buffer + preempt_each(llama_state_seq_copy_buf_free); + preempt_state_tgt.clear(); preempt_state_tgt.shrink_to_fit(); preempt_state_dft.clear(); preempt_state_dft.shrink_to_fit(); } - // bytes preempt_save() would need for this slot right now + void preempt_copy_wait() { + preempt_each(llama_state_seq_copy_wait); + } + 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); } - // copy the sequence out of the cache and release its cells + // 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 + void preempt_detach() { + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + } + + bool preempt_copy_done() { + return llama_state_seq_copy_done(preempt_cpy_tgt.get()) && + (!preempt_cpy_dft || llama_state_seq_copy_done(preempt_cpy_dft.get())); + } + + bool preempt_resumed() { + n_preempt_fail = 0; + + state = state_before_preempt; + + if (state == SLOT_STATE_GENERATING && can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + + return true; + } + + bool preempt_save_poll() { + if (!preempt_copy_done()) { + return false; + } + + mem.seq_rm(id, -1, -1); + + state = SLOT_STATE_PREEMPTED; + + return true; + } + + bool preempt_restore_poll() { + if (!preempt_copy_done()) { + return false; + } + + 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); + } + + return preempt_resumed(); + } + + // [TAG_PREEMPT_ASYNC] copy the sequence out and release its cells; with a transfer this returns once the copy is issued and the cells stay the slot's until preempt_save_poll() sees it land 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 larger buffer can still come back pageable, and a copy into pageable memory blocks; such a slot parks synchronously from now on + 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; preempt_save_poll() releases 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,19 +559,9 @@ 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. + // note: prompt.tokens is deliberately kept - the resume sizes its request from it mem.seq_rm(id, -1, -1); state_before_preempt = state; @@ -407,20 +574,33 @@ struct server_slot { return true; } - // put the sequence back; the slot then continues from the token it was about to decode + // [TAG_PREEMPT_ASYNC] put the sequence back; with a transfer this returns once the copy is issued, leaving the slot RESTORING: it owns the cells, but they hold no state until the copy lands bool preempt_restore() { - const size_t size_tgt = preempt_state_tgt.size(); - const size_t size_dft = preempt_state_dft.size(); + 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 what was issued finish before the half-written sequence is dropped, or cells go while a copy still writes them + preempt_copy_wait(); + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } - if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { - // no room after all: drop the half-written sequence and stay parked - mem.seq_rm(id, -1, -1); - n_preempt_fail++; - return false; + state = SLOT_STATE_RESTORING; + + return true; } - if (size_dft > 0 && - llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + const size_t size_tgt = preempt_state_tgt.size(); + const size_t size_dft = preempt_state_dft.size(); + + if (llama_state_seq_set_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt || + (size_dft > 0 && + llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft)) { mem.seq_rm(id, -1, -1); n_preempt_fail++; return false; @@ -428,27 +608,10 @@ struct server_slot { preempt_state_free(); - n_preempt_fail = 0; - - state = state_before_preempt; - - // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks - // that the draft context is where it should be, which the restore above ensures. - // A slot parked while still processing its prompt makes that transition itself - // once the prompt is done. - if (state == SLOT_STATE_GENERATING && can_speculate()) { - common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); - } - - return true; + return preempt_resumed(); } - // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds for this sequence. - // For a batch that is given up after it was built: the tokens added for this slot - // that were never decoded come off, the sampled token stays in `sampled` and goes into - // the next batch the way it went into this one, and a draft is a prediction that goes - // with them. A chunk that failed to decode left nothing in the cache, so the cache is - // the boundary. + // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up after it was built: never-decoded tokens and the draft come off, `sampled` is kept void rewind_to_cache() { const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; @@ -456,18 +619,12 @@ struct server_slot { prompt.tokens.keep_first(n_cached); } - // a prompt whose last chunk was in the batch was marked done when the chunk was - // built; the chunk never ran, so the prompt is not done + // the last chunk was marked done when it was built but never ran, so it is not done if (state == SLOT_STATE_DONE_PROMPT && task && prompt.n_tokens() < task->n_tokens()) { state = SLOT_STATE_PROCESSING_PROMPT; } - spec_draft.clear(); - spec_i_batch.clear(); - spec_ckpt.clear(); - spec_is_replay = false; - - i_batch = -1; + preempt_detach(); } std::vector lora; @@ -528,7 +685,6 @@ struct server_slot { n_predict_max = -1; - // [TAG_PREEMPT] preempt_state_free(); state_before_preempt = SLOT_STATE_IDLE; n_preempt = 0; @@ -688,10 +844,9 @@ struct server_slot { t_last_used = ggml_time_us(); - // [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] [TAG_PREEMPT_ASYNC] a parked slot's cells are already gone, so the mirror must not outlive them or the next task prefix-matches an empty cache; wait for any copy first, its buffer and its cells are about to be handed on + if (preempt_is_out()) { + preempt_copy_wait(); preempt_state_free(); prompt_clear(); } @@ -838,7 +993,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 +1238,16 @@ struct server_context_impl { int64_t t_last_load_progress_ms = 0; void destroy() { + // [TAG_PREEMPT_ASYNC] the slots outlive this call and may hold a copy reading or writing KV tensors of the contexts about to be freed; release() makes the same wait for one slot + for (auto & slot : slots) { + slot.preempt_copy_wait(); + + slot.preempt_cpy_tgt.reset(); + slot.preempt_cpy_dft.reset(); + } + + preempt_ram_kind_logged = false; + spec.reset(); spec_init.reset(); @@ -1423,6 +1588,20 @@ struct server_context_impl { } }; + // [TAG_PREEMPT_ASYNC] one transfer per context, reused for every park and resume, because each owns a backend and installs the fences the context records after every decode + 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 mid-park, so the whole slot stays synchronous + slot.preempt_cpy_tgt.reset(); + } + } + } + slot.reset(); } @@ -1445,8 +1624,63 @@ struct server_context_impl { } { - // 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 + 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) { + // a copy into pageable memory is staged by the driver and blocks the thread that issued it, and a host buffer type is free to hand back pageable memory rather than fail + 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(); + } + } + } + + { + preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt))); + + // test knob: the paged kernel needs a head size of 256, so a harness model cannot reach the paged arithmetic otherwise + const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY"); + + if (LLAMA_SERVER_PREEMPT_GRANULARITY) { + preempt_alloc_granularity = std::max(1, atoi(LLAMA_SERVER_PREEMPT_GRANULARITY)); + + SRV_WRN("LLAMA_SERVER_PREEMPT_GRANULARITY = %d (test knob: planning the kv pool in blocks of %d cells)\n", + preempt_alloc_granularity, preempt_alloc_granularity); + } else if (preempt_alloc_granularity > 1) { + SRV_INF("preemption: the kv pool allocates %d cells at a time, planning in pages\n", + preempt_alloc_granularity); + } + } + + { preempt_resume_head = true; const char * LLAMA_SERVER_PREEMPT_RESUME = getenv("LLAMA_SERVER_PREEMPT_RESUME"); @@ -1463,10 +1697,7 @@ struct server_context_impl { const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY"); preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0; - // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks, for comparing - // policies against each other on the same workload. smallest (the default and the - // shipped one), largest, youngest (the most recent task, as vLLM's scheduler - // preempts), oldest. The leader is kept and the starvation guard applies under all. + // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks; smallest (default), largest, youngest, oldest const char * LLAMA_SERVER_PREEMPT_POLICY = getenv("LLAMA_SERVER_PREEMPT_POLICY"); preempt_test_policy = LLAMA_SERVER_PREEMPT_POLICY ? LLAMA_SERVER_PREEMPT_POLICY : "smallest"; @@ -1486,8 +1717,7 @@ struct server_context_impl { SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_PLANNER = off (test knob: nothing is parked ahead of the decode, only as a last resort)\n"); } - // assigned, not only set: the same context reloaded with an attention model after - // a recurrent one gets its preemption back + // assigned, not only set: a context reloaded with an attention model after a recurrent one gets preemption back preempt_recurrent = llama_model_is_recurrent(model_tgt); if (preempt_recurrent) { @@ -2189,6 +2419,22 @@ struct server_context_impl { queue_results.send(std::move(res)); } + void send_preempt_notice(server_slot & slot, bool parked) { + if (!slot.task || !slot.task->params.stream) { + return; + } + + auto res = std::make_unique(); + + res->id = slot.task->id; + res->index = slot.task->index; + res->id_slot = slot.id; + res->parked = parked; + res->n_preempt = slot.n_preempt; + + queue_results.send(std::move(res)); + } + void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) { auto res = std::make_unique(); @@ -2637,7 +2883,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++; } } @@ -2889,10 +3135,8 @@ struct server_context_impl { void abort_all_slots(const std::string & reason) { for (auto & slot : slots) { - // [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] a parked slot, or one with a copy in flight, took no part in what failed and comes back when there is room + if (slot.is_processing() && !slot.preempt_is_out()) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); } @@ -2929,43 +3173,59 @@ struct server_context_impl { }; #endif - // - // [TAG_PREEMPT] server-side request preemption - // - - // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, - // whether or not the pool is under pressure. It exists to answer the only question that - // matters about a resume: with one request on an idle server the batch has the same - // shape at every step, so a preempted continuation that is not byte-identical to an - // uninterrupted one is the preemption's fault and nothing else's. + // LLAMA_SERVER_PREEMPT_EVERY=N: preempt every generating slot every N tokens, pressure or not, so the determinism test can blame any difference on the preemption int32_t preempt_test_every = 0; std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model - // 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 + // [TAG_PREEMPT_ASYNC] whether parks go through a transfer; false with --no-preempt-async or a backend that cannot copy asynchronously + bool preempt_async_ok = false; + + // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load: 1 ordinarily, the page size under exact concurrency. LLAMA_SERVER_PREEMPT_GRANULARITY overrides it. + int32_t preempt_alloc_granularity = 1; + + int32_t preempt_n_cells(int32_t n_tokens) const { + return preempt_n_cells_g(n_tokens, preempt_alloc_granularity); + } + + int32_t preempt_n_cells_step(int32_t n_tokens, int32_t n_step) const { + return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity); + } + + // LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, leaving only the KV-full retry ladder bool preempt_planner_off = false; - // LLAMA_SERVER_PREEMPT_RESUME: head (the default) puts parked slots back in the order they - // were parked and only the first until it fits; pass lets a smaller slot pass a head - // that does not fit. Read at load, per context. + // LLAMA_SERVER_PREEMPT_RESUME=head or pass, read at load, per context bool preempt_resume_head = true; - // a recurrent cache holds one state per sequence whatever its length: no cell pool, - // nothing to run out of, and the token count the planner measures says nothing about - // it. Preemption is off for those models; a hybrid keeps its attention cache and stays on. bool preempt_recurrent = false; - // 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 in place inside the next llama_decode + 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; } - // 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 + 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"); + } + int32_t preempt_n_spec(const server_slot & slot) const { int32_t res = preempt_n_spec_max(); @@ -2982,7 +3242,6 @@ struct server_context_impl { return std::max(0, res); } - // host RAM the parked sequences hold right now size_t preempt_ram_used() const { size_t res = 0; @@ -2993,20 +3252,68 @@ struct server_context_impl { return res; } - // whether parking this slot stays under --preempt-ram - bool preempt_fits_budget(const server_slot & slot) const { - if (params_base.preempt_ram_mib < 0) { - return true; + // [TAG_PREEMPT_ASYNC] a restored slot keeps its pinned buffer, so that idle capacity is given back largest first when a park does not fit under --preempt-ram + 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(); } + } - const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + size_t preempt_ram_budget() const { + return params_base.preempt_ram_mib < 0 ? SIZE_MAX : (size_t) params_base.preempt_ram_mib * 1024 * 1024; + } + + bool preempt_fits_budget(const server_slot & slot) { + const size_t budget = preempt_ram_budget(); + + // what this slot already holds is counted by preempt_ram_used() and reused, so a park costs only the rest + 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() + slot.preempt_state_required() <= budget; + return preempt_ram_used() + extra <= budget; + } + + void preempt_trim_ram(server_slot & slot) { + if (preempt_ram_used() > preempt_ram_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 - // builder applies when it takes the slot: nothing when the request does not cache its - // prompt, otherwise the prefix the two share, cut short of an aLoRA invocation size_t preempt_n_keep(const server_slot & slot) const { if (!slot.task->params.cache_prompt) { return 0; @@ -3021,10 +3328,7 @@ struct server_context_impl { return n_keep; } - // cells of the slot's that its next step keeps: a slot just given a task still mirrors - // the previous request's prompt until the batch builder keeps what preempt_n_keep() - // says and drops the rest, so what it holds, and what it is about to ask for, both - // count from that + // a slot just given a task still mirrors the previous request's prompt, so what it holds and what it asks for both count from preempt_n_keep() int32_t preempt_n_retained(const server_slot & slot) const { if (slot.state == SLOT_STATE_STARTED && slot.task) { return (int32_t) preempt_n_keep(slot); @@ -3033,7 +3337,6 @@ struct server_context_impl { return slot.prompt.n_tokens(); } - // cells the slot will ask for on its next step once it is back in the pool int32_t preempt_n_need(const server_slot & slot) const { int32_t res = preempt_n_retained(slot); @@ -3045,19 +3348,14 @@ struct server_context_impl { res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); } - return res; + // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in full; undercounting admits a resume find_slot cannot satisfy + return preempt_n_cells(res); } - // Cells the pool is holding right now. A released slot keeps its prompt in the cache - // for the next request to reuse as a prefix, so idle slots count too: the first version - // of this counted only the running ones, decided a pool holding 8185 cached cells was - // empty, and every resume failed against a cache that was actually full. int32_t preempt_kv_used() const { int32_t res = 0; - // n_cmpl > 1: the parent and its children share the prompt's cells through seq_cp, so - // the prompt is charged once per family, to whichever resident member comes first; - // the others are charged only what they generated on top of it + // n_cmpl > 1: a family shares the prompt's cells through seq_cp, so it is charged once, to the first resident member std::vector charged; for (const auto & slot : slots) { @@ -3065,11 +3363,10 @@ struct server_context_impl { continue; // parked: its cells are in host RAM, not in the pool } - // 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 + // [TAG_PREEMPT_ASYNC] deliberately not skipped: a slot with a copy in flight holds cells either way, and skipping it would hand the same cells out twice + if (slot.state == SLOT_STATE_WAIT_OTHER) { - res += slot.prompt.n_tokens(); + res += preempt_n_cells(slot.prompt.n_tokens()); continue; } @@ -3077,77 +3374,81 @@ struct server_context_impl { const int family = slot.task->is_parent() ? slot.task->id : slot.task->id_parent; if (std::find(charged.begin(), charged.end(), family) != charged.end()) { - res += std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens()); + res += preempt_n_cells(std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens())); continue; } charged.push_back(family); } - // what the pool holds now, the previous request's prompt included for a slot just - // given a task: the batch builder trims that to the prefix the two share, but - // not until the slot is built into a batch, and with continuous batching off that - // can be a long time behind a running generation. Measured by the prefix, a - // restore was found to fit and attempted against cells still occupied. Under - // pressure the planner trims such slots itself, see preempt_normalize_started_all() - res += slot.prompt.n_tokens(); + res += preempt_n_cells(slot.prompt.n_tokens()); } return res; } - // cells those slots are about to ask for on the next decode + // [TAG_PREEMPT_ASYNC] the room the pool is kept clear of, so everything still decoding has somewhere to put its tokens until a park lands; a resume candidate is charged the same runway + int32_t preempt_n_margin(int32_t n_additional_running = 0) const { + if (!preempt_async_ok) { + // [TAG_EXACT_CONCURRENCY] a margin of eight cells is no margin where a step can cost a whole page + return preempt_n_cells(PREEMPT_N_MARGIN); + } + + int32_t n_running = n_additional_running; + + for (const auto & slot : slots) { + if (slot.is_processing() && (!slot.preempt_is_out() || slot.state == SLOT_STATE_RESTORING)) { + n_running++; + } + } + + // [TAG_EXACT_CONCURRENCY] round the runway up to a page, once and not per slot, which would keep a page per slot out of the users' reach + return preempt_n_cells( + PREEMPT_N_MARGIN + n_running * (1 + preempt_n_spec_max()) * PREEMPT_N_ASYNC_STEPS); + } + int32_t preempt_kv_reserve() const { const int32_t n_batch = llama_n_batch(ctx_tgt); int32_t res = 0; int32_t res_pmt = 0; + int32_t n_pmt = 0; + // [TAG_EXACT_CONCURRENCY] reserve the cells the next step ADDS, not its tokens: the used figure already rounds every tail page up, and only a page crossing can empty the pool for (const auto & slot : slots) { - switch (slot.state) { + const int32_t n_cur = slot.prompt.n_tokens(); + + // [TAG_PREEMPT_ASYNC] a restoring slot decodes as soon as its copy lands, so it is charged the step of the state it goes back to, or that first step preempts somebody else + const slot_state state = slot.state == SLOT_STATE_RESTORING ? slot.state_before_preempt : slot.state; + + switch (state) { case SLOT_STATE_GENERATING: case SLOT_STATE_DONE_PROMPT: { - res += 1 + preempt_n_spec(slot); + res += preempt_n_cells_step(n_cur, 1 + preempt_n_spec(slot)); } break; case SLOT_STATE_STARTED: case SLOT_STATE_PROCESSING_PROMPT: { - // from the prefix a started slot keeps, not from the prompt it still - // mirrors: measured by the mirror, a request shorter than the last one - // reserved one cell for a chunk of hundreds - const int32_t n_left = slot.task ? slot.task->n_tokens() - preempt_n_retained(slot) : 0; + const int32_t n_have = preempt_n_retained(slot); + const int32_t n_left = slot.task ? slot.task->n_tokens() - n_have : 0; - res_pmt += std::max(1, std::min(n_batch, n_left)); + res_pmt += preempt_n_cells_step(n_have, std::max(1, std::min(n_batch, n_left))); + n_pmt++; } break; default: break; } } - // one batch is all the prompt slots get between them, however many are waiting - return res + std::min(res_pmt, n_batch); + return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); } - // Keep the slot that is furthest along -- it is the closest to finishing and to giving - // its cells back -- and among the rest prefer one that has not been preempted - // PREEMPT_N_STARVED times already, then the smallest. - // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt - // until the batch builder keeps the prefix the two share and drops the rest (see the - // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out, - // charged and sized by the old prompt, and a short unrelated request could exceed the - // budget or stay parked for room it will never use. Keeping only the shared prefix now - // is what the batch builder does anyway; the chunk reuse it can add on top is given up - // for a slot the planner has to touch, which is rare. - // every started slot, when the pool is short: true when any of them gave cells up + // [TAG_PREEMPT] trim a just-started slot to the prefix it keeps first, or it is copied out, charged and sized by the previous request's prompt bool preempt_normalize_started_all() { bool res = false; for (auto & slot : slots) { - if (slot.state != SLOT_STATE_STARTED || !slot.task) { - continue; - } - const int32_t before = slot.prompt.n_tokens(); preempt_normalize_started(slot); @@ -3173,10 +3474,7 @@ struct server_context_impl { return; } - // a memory that cannot remove part of a sequence (a recurrent state without rollback - // room for the stale suffix) aborts on a partial removal; for it the whole stale - // sequence goes, and the prompt is processed from the start on resume, as it would be - // without a usable checkpoint + // a memory that cannot remove part of a sequence aborts on a partial removal, so drop the whole stale sequence const bool partial_ok = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART && (!ctx_dft || ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART); @@ -3194,16 +3492,12 @@ struct server_context_impl { server_slot * leader = nullptr; int32_t n_running = 0; - // a slot just given a task is measured by the prefix it keeps, not by the previous - // request's prompt it still mirrors: measured by the mirror, a short request over a - // large stale cache would be the never-parked leader while the longest live - // conversation was parked in its place for (auto & slot : slots) { preempt_normalize_started(slot); } 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()) { @@ -3213,19 +3507,14 @@ struct server_context_impl { } if (n_running < 2) { - // a single conversation that does not fit the pool on its own is a real context - // overflow and not a scheduling problem - leave it to the existing error path + // one conversation that does not fit alone is a real overflow, not a scheduling problem return nullptr; } server_slot * victim = nullptr; for (auto & slot : slots) { - // Before the batch is built every one of these is at a token boundary: a - // generating slot between two sampled tokens, a prompt-processing slot between - // two chunks of its prompt, a started slot with only a cached prefix (or - // nothing) in the pool. A slot holding no cells is still worth parking - it - // is about to ask for a whole batch of them. + // before the batch is built every slot is at a token boundary; one holding no cells is still worth parking if (slot.state != SLOT_STATE_GENERATING && slot.state != SLOT_STATE_PROCESSING_PROMPT && slot.state != SLOT_STATE_STARTED) { @@ -3240,6 +3529,16 @@ struct server_context_impl { continue; // n_cmpl > 1 slots share one sequence, out of scope here } + // a started slot the STARTED block is about to reject gets its error on its own pass: a park notice would open the stream and turn that 4xx into 200 plus an in-stream error + if (slot.state == SLOT_STATE_STARTED) { + std::string msg; + error_type type = ERROR_TYPE_SERVER; + + if (slot_prompt_rejected(slot, msg, type)) { + continue; + } + } + if (!preempt_fits_budget(slot)) { continue; } @@ -3257,9 +3556,6 @@ struct server_context_impl { return victim; } - // is a the better victim of the two? the smallest slot under the shipped policy: it - // gives up the least work and its restore is the cheapest (see the PR's simulation); - // the other choices exist for the comparison runs behind LLAMA_SERVER_PREEMPT_POLICY bool preempt_better_victim(const server_slot & a, const server_slot & b) const { if (preempt_test_policy == "largest") { return a.prompt.n_tokens() > b.prompt.n_tokens(); @@ -3276,10 +3572,118 @@ struct server_context_impl { return a.prompt.n_tokens() < b.prompt.n_tokens(); } - // called once per update_slots(), before the batch is built: at that point every slot is - // 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] park a slot: a synchronous park is finished here, an asynchronous one only issued, and update_preempt_copies() counts it when its copy lands. The notice goes with the save, not the cell release: preempt_save() has already detached the slot, so a release-time notice would leave the copy's silence unexplained. + bool preempt_park(server_slot & slot, int64_t t_start) { + slot.t_preempt_copy_us = t_start; + + if (!slot.preempt_save()) { + return false; + } + + preempt_log_ram_kind(slot); + + if (slot.state == SLOT_STATE_PREEMPTED) { + metrics.n_preempt++; + } + + send_preempt_notice(slot, true); + + return true; + } + + void preempt_parked(server_slot & slot, const char * note) { + metrics.n_preempt++; + + SLT_WRN(slot, "park completed after %.2f ms%s: %d cells released, %.1f MiB parked, kv %d/%d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, note, + slot.prompt.n_tokens(), + slot.preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_ctx); + } + + // [TAG_PREEMPT] a resume whose copy has landed; announced here rather than where the restore was issued, this being the first moment the slot can be scheduled again + void preempt_restored(server_slot & slot, const char * note) { + metrics.n_resume++; + + preempt_trim_ram(slot); + + send_preempt_notice(slot, false); + + SLT_WRN(slot, "restore completed after %.2f ms%s: %d tokens back in the cache, kv %d/%d, preemptions %d\n", + (ggml_time_us() - slot.t_preempt_copy_us) / 1e3, note, + slot.prompt.n_tokens(), + preempt_kv_used(), n_ctx, + slot.n_preempt); + } + + void update_preempt_copies() { + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTING) { + if (slot.preempt_save_poll()) { + preempt_parked(slot, ""); + } + } else if (slot.state == SLOT_STATE_RESTORING) { + if (slot.preempt_restore_poll()) { + preempt_restored(slot, ""); + } + } + } + } + + // [TAG_PREEMPT_ASYNC] wait for every copy in flight before a shift: the shift is one in-place graph over the whole K cache, so a copy beside it reads or writes half-shifted cells + 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()) { + preempt_restored(slot, " (waited for, a context shift is due)"); + } + } + } + + bool preempt_copies_in_flight() const { + for (const auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTING) { + return true; + } + } + + return false; + } + + 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; + } + + preempt_parked(slot, " (waited for)"); + + 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,16 +3693,14 @@ 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 } const int32_t n_cells = n_ctx; - // Put back what fits, in the order preempt_resume_head describes: by default - // the slot parked longest, and only that one until it fits; under - // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked - // longest, and a smaller slot may pass a head that does not fit. const bool head_of_line = preempt_resume_head; for (;;) { @@ -3328,40 +3730,21 @@ struct server_context_impl { server_slot * best = nullptr; - // A parked slot whose sequence plus its next step would not fit an empty pool can - // never be restored, and would otherwise sit at the head of the line for ever - // without a restore ever being attempted: a prompt within n_ctx that was parked - // before it took any cells, but too close to n_ctx to leave room for its first - // batch. That is the single-conversation overflow the KV-full path reports, so - // report it the same way and rescan the line without it. - { - server_slot * impossible = nullptr; - - for (auto * slot : parked) { - if (preempt_n_need(*slot) > n_cells) { - impossible = slot; - break; - } - } + const auto impossible = std::find_if(parked.begin(), parked.end(), + [this, n_cells](const server_slot * slot) { return preempt_n_need(*slot) > n_cells; }); - if (impossible) { - SLT_WRN(*impossible, "parked sequence of %d tokens cannot fit the pool of %d cells even alone, failing it\n", - preempt_n_need(*impossible), n_cells); - send_error(*impossible, "Context size has been exceeded."); - impossible->release(); - continue; - } + if (impossible != parked.end()) { + SLT_WRN(**impossible, "parked sequence of %d tokens cannot fit the pool of %d cells even alone, failing it\n", + preempt_n_need(**impossible), n_cells); + send_error(**impossible, "Context size has been exceeded."); + (*impossible)->release(); + continue; } - // 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. - // 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. + // room for the sequence and for the next step of everything running, the candidate included, or a resume immediately preempts somebody; with nobody resident an exact fit is let in 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) { @@ -3374,9 +3757,6 @@ struct server_context_impl { break; } - // a slot just given a task still holds the previous request's prompt until - // the batch builder trims it; trimmed here instead, the cells it will not - // keep are counted out and a parked slot that fits without them comes back if (preempt_normalize_started_all()) { continue; } @@ -3386,22 +3766,14 @@ struct server_context_impl { } } - // Nothing fits. A resident that has reached the pool's limit and is cycling - // through context shifts holds the room for as long as it likes to generate, - // and the head behind it would wait for ever. After the head has waited its - // turn, that resident is parked in its place: it is at a token boundary like - // any other park, and when it comes back it is the one waiting, so the two - // take turns instead of one taking everything. + // nothing fits: a resident cycling through context shifts holds the room for as long as it generates, so it is parked once the head has waited its turn if (!best) { server_slot * head = parked.front(); - if (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 - // matter, spend the park budget on it, and leave the head waiting anyway. + // [TAG_PREEMPT_ASYNC] a park still copying holds its cells, so a rotation now would only park another resident on top + if (!preempt_copies_in_flight() && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); - const int32_t need = preempt_n_need(*head) + PREEMPT_N_MARGIN; + const int32_t need = preempt_n_need(*head) + preempt_n_margin(1); server_slot * pick = nullptr; bool pick_enough = false; @@ -3416,17 +3788,13 @@ struct server_context_impl { continue; } - // The head's own bytes are not credited as leaving: the resident is - // parked before the head is restored and freed, so both states are - // held at once, and the cap is a cap on what is held. A budget that - // holds one sequence but not two does not rotate, and the head waits - // for a resident to finish, which is said once per park below. + // the head's own bytes are not credited as leaving: the resident is parked before the head is restored and freed, so both states are held at once if (!preempt_fits_budget(slot)) { budget_refused = true; continue; } - const bool enough = occupied - slot.prompt.n_tokens() + need <= n_cells; + const bool enough = occupied - preempt_n_cells(slot.prompt.n_tokens()) + need <= n_cells; if (!pick || (enough && !pick_enough) || @@ -3437,6 +3805,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; @@ -3444,11 +3814,9 @@ struct server_context_impl { params_base.preempt_ram_mib); } - if (pick && pick->preempt_save()) { + if (pick && preempt_park(*pick, t_start)) { server_slot & slot = *pick; - 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(), slot.preempt_state_size() / (1024.0 * 1024.0), @@ -3456,7 +3824,10 @@ 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 on the pass that sees the copy land + if (slot.state == SLOT_STATE_PREEMPTED) { + best = head; + } } } @@ -3469,10 +3840,9 @@ 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 - // a slot that has been failing for a while, and keep the log quiet. if (best->n_preempt_fail % 64 == 1) { SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); @@ -3487,8 +3857,23 @@ struct server_context_impl { break; } + 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); + + continue; + } + metrics.n_resume++; + // [TAG_PREEMPT] the synchronous restore returns with the slot already back in its old state, so issue and landing are the same moment here + send_preempt_notice(*best, false); + SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", (ggml_time_us() - best->t_preempt_us) / 1e6, best->prompt.n_tokens(), @@ -3497,15 +3882,11 @@ struct server_context_impl { best->n_preempt); } - // forced preemption, for the determinism test only if (preempt_test_every > 0) { 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) && preempt_park(slot, 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)); } @@ -3516,19 +3897,28 @@ struct server_context_impl { return; // test knob: leave the pool to the retry ladder and its last resort } - // and take cells back until the next decode fits 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; } - // a prompt cached on an idle slot is the cheapest thing in the pool to give up if (try_clear_idle_slots()) { continue; } + // [TAG_PREEMPT_ASYNC] a park issued and not landed holds cells that are already spoken for, so waiting for it is quicker than parking somebody else + if (preempt_copies_in_flight()) { + if (n_used > n_cells) { + if (preempt_wait_in_flight()) { + continue; + } + } else { + break; + } + } + server_slot * victim = preempt_pick_victim(); if (!victim) { @@ -3540,11 +3930,26 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); - if (!victim->preempt_save()) { + if (!preempt_park(*victim, t_start)) { break; // could not park it; the existing retry ladder is still behind us } - metrics.n_preempt++; + 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] short of the lookahead only, the step still fits and leaving is the point; out of room for it the cells are held until the copy lands, so the retry ladder ends every request instead of waiting + if (n_used + preempt_n_margin() > n_cells) { + continue; + } + + break; + } SLT_WRN(*victim, "preempted: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", n_tokens, @@ -3555,6 +3960,53 @@ struct server_context_impl { } } + // the checks a request has to pass before its prompt is processed; true when it is rejected. An empty prompt is not here: it is a final response, not an error. + bool slot_prompt_rejected(const server_slot & slot, std::string & msg, error_type & type) const { + if (!slot.task) { + return false; + } + + // TODO: support memory-less logits computation + if (slot.task->need_logits() && !llama_get_memory(ctx_tgt)) { + msg = "the current context does not logits computation. skipping"; + type = ERROR_TYPE_SERVER; + return true; + } + + if (!slot.can_split()) { + const int32_t n_ubatch = llama_n_ubatch(ctx_tgt); + + if (slot.task->n_tokens() > n_ubatch) { + msg = string_format( + "input (%d tokens) is too large to process. increase the physical batch " + "size (current batch size: %d)", + slot.task->n_tokens(), n_ubatch); + type = ERROR_TYPE_SERVER; + return true; + } + + if (slot.task->n_tokens() > slot.n_ctx) { + msg = string_format( + "input (%d tokens) is larger than the max context size (%d tokens). skipping", + slot.task->n_tokens(), slot.n_ctx); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; + } + + return false; + } + + if (slot.task->n_tokens() >= slot.n_ctx) { + msg = string_format( + "request (%d tokens) exceeds the available context size (%d tokens), try increasing it", + slot.task->n_tokens(), slot.n_ctx); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; + } + + return false; + } + void update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -3569,7 +4021,6 @@ struct server_context_impl { } #endif - // check if all slots are idle { bool all_idle = true; @@ -3597,13 +4048,13 @@ struct server_context_impl { } try { - // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured - // after any context shift. Inside the guard with the rest of the step: a shift - // rebuilds a slot's tokens and a park allocates, and either can throw, which the - // slots are told about rather than the loop ending on an uncaught exception + // [TAG_PREEMPT] make the pool fit the step about to be built, measured after any context shift; inside the guard because a shift or a park can throw pre_decode_shift(); update_preemption(); + // [TAG_PREEMPT_ASYNC] before pre_decode(), not only the target decode: the draft it asks for applies the draft cache's pending shift in place + preempt_wait_for_shift(); + scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); batch.render(); @@ -3639,6 +4090,10 @@ 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 + 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 { @@ -3652,8 +4107,6 @@ struct server_context_impl { #endif if (preempt_batch_abandoned) { - // [TAG_PREEMPT] the rest of this batch was never decoded and the slots no - // longer describe it; the next pass builds a new one preempt_batch_abandoned = false; break; } @@ -3687,8 +4140,6 @@ struct server_context_impl { // apply context-shift if needed // TODO: simplify and improve - // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift, - // not with the cells the shift is about to give back void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { @@ -3730,6 +4181,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); @@ -3901,9 +4353,8 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - // [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) { + // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to batch until it is restored + if (!slot.is_processing() || slot.preempt_is_out()) { return; } @@ -3961,46 +4412,18 @@ struct server_context_impl { return; } - // TODO: support memory-less logits computation - if (slot.task->need_logits() && !llama_get_memory(ctx_tgt)) { - send_error(slot, "the current context does not logits computation. skipping", ERROR_TYPE_SERVER); - slot.release(); - return; - } - - if (!slot.can_split()) { - if (slot.task->n_tokens() > n_ubatch) { - send_error(slot, - string_format( - "input (%d tokens) is too large to process. increase the physical batch " - "size (current batch size: %d)", - slot.task->n_tokens(), n_ubatch), - ERROR_TYPE_SERVER); - slot.release(); - return; - } + { + std::string msg; + error_type type = ERROR_TYPE_SERVER; - if (slot.task->n_tokens() > slot.n_ctx) { - send_error( - slot, - string_format( - "input (%d tokens) is larger than the max context size (%d tokens). skipping", - slot.task->n_tokens(), slot.n_ctx), - ERROR_TYPE_EXCEED_CONTEXT_SIZE); - slot.release(); - return; - } - } else { - if (slot.task->n_tokens() >= slot.n_ctx) { - send_error(slot, - string_format("request (%d tokens) exceeds the available context size (%d " - "tokens), try increasing it", - slot.task->n_tokens(), slot.n_ctx), - ERROR_TYPE_EXCEED_CONTEXT_SIZE); + if (slot_prompt_rejected(slot, msg, type)) { + send_error(slot, msg, type); slot.release(); return; } + } + if (slot.can_split()) { if (slot.task->params.cache_prompt) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); @@ -4056,6 +4479,8 @@ 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); + 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++; @@ -4429,18 +4854,17 @@ struct server_context_impl { } } - // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is - // the context error for every slot in the batch. With a park budget the batch is given - // up instead: every resident slot is rewound to the token boundary the cache is at (a - // batch is applied one chunk at a time, and the chunk that failed left nothing behind), - // the smallest are parked until the planner's own bound holds again, and the next - // 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 at all and go asynchronously + 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); } + // [TAG_PREEMPT] the retry ladder ran out: give the batch up, rewind every resident to the token boundary the cache is at and park the smallest. Multimodal keeps the old path. bool preempt_last_resort(int32_t off) { if (!preempt_last_resort_possible()) { return false; @@ -4449,7 +4873,7 @@ struct server_context_impl { int32_t n_running = 0; for (auto & slot : slots) { - if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED || slot.preempt_in_flight()) { continue; } @@ -4465,7 +4889,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 +4901,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 +4914,22 @@ struct server_context_impl { const int32_t n_tokens = victim->prompt.n_tokens(); const int64_t t_start = ggml_time_us(); - if (!victim->preempt_save()) { + if (!preempt_park(*victim, t_start)) { break; } - metrics.n_preempt++; n_parked++; + // [TAG_PREEMPT_ASYNC] the cells are wanted now, not next iteration: wait for the copy, which releases them + 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; + } + 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, @@ -4514,7 +4948,6 @@ struct server_context_impl { return true; } - // [TAG_PREEMPT] whether a slot in the batch has its sampled token and a draft in it bool batch_has_spec_groups() const { for (const auto & slot : slots) { if (!slot.spec_i_batch.empty()) { @@ -4569,14 +5002,16 @@ struct server_context_impl { }); if (ret != 0) { + // [TAG_PREEMPT_ASYNC] halving the batch returns no cells, so wait for an issued park first, or the ladder runs down to n_batch == 1 and ends every request + 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; - // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have - // to stay in one view: a narrower view splits the group and the verify step - // throws for the slot whose tokens straddle it. Halving is no help there, so - // after the idle slots the ladder goes to its last resort straight away. With - // no budget to park into the ladder is what it always was. + // [TAG_PREEMPT] a slot's sampled token and its draft have to stay in one view, so halving would split the group and make the verify step throw if (ret == 1 && n_batch > 1 && preempt_last_resort_possible() && batch_has_spec_groups()) { if (try_clear_idle_slots()) { SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n"); @@ -4587,7 +5022,6 @@ struct server_context_impl { } if (n_batch == 1 && ret == 1) { - // [TAG_PREEMPT] park instead of ending everyone, when there is a budget to park into if (preempt_last_resort(off)) { preempt_batch_abandoned = true; return true; @@ -4613,9 +5047,7 @@ struct server_context_impl { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); 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(); @@ -4964,8 +5396,7 @@ struct server_context_impl { void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { 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()); @@ -5222,6 +5653,14 @@ std::unique_ptr server_routes::handle_completions_impl( task.params.oaicompat_cmpl_id = completion_id; task.params.oaicompat_model = meta->model_name; + // [TAG_EXACT_CONCURRENCY] exact mode gives a page to a single sequence, so refuse an n_cmpl > 1 child here, where it becomes a 400 rather than at seq_cp + if (task.params.n_cmpl > 1 && common_exact_concurrency()) { + throw std::runtime_error( + "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each " + "completion needs its own sequence, and in exact mode a KV page belongs " + "to a single sequence. Send n separate requests, or unset the variable."); + } + // prepare child tasks if (task.params.n_cmpl > 1) { int n_children = task.params.n_cmpl - 1; @@ -5275,37 +5714,51 @@ std::unique_ptr server_routes::handle_completions_impl( // in streaming mode, the first error must be treated as non-stream response // this is to match the OAI API behavior // ref: https://github.com/ggml-org/llama.cpp/pull/16486#discussion_r2419657309 + // [TAG_PREEMPT] a slot can be parked before any token exists, so those notices are kept and sent in front of the first real result + std::string preempt_prefix; + std::set parked_idx; // prompts of this request that are parked right now auto first_result = rd.next(req.should_stop); - if (first_result == nullptr) { - GGML_ASSERT(req.should_stop()); - return res; // connection is closed - } + if (first_result != nullptr && dynamic_cast(first_result.get()) != nullptr) { + const auto * notice = static_cast(first_result.get()); + preempt_prefix = preempt_notice_comment(*notice); + if (notice->parked) { + parked_idx.insert(notice->index); + } else { + parked_idx.erase(notice->index); + } + first_result.reset(); + } else { + if (first_result == nullptr) { + GGML_ASSERT(req.should_stop()); + return res; // connection is closed + } - if (first_result->is_error()) { - res->error(first_result->to_json()); - return res; - } + if (first_result->is_error()) { + res->error(first_result->to_json()); + return res; + } - GGML_ASSERT( - dynamic_cast(first_result.get()) != nullptr || - dynamic_cast (first_result.get()) != nullptr - ); + GGML_ASSERT( + dynamic_cast(first_result.get()) != nullptr || + dynamic_cast (first_result.get()) != nullptr + ); + } - // next responses are streamed - // to be sent immediately - json first_result_json = first_result->to_json(); + json first_result_json = first_result ? first_result->to_json() : json(nullptr); if (first_result_json == nullptr) { - res->data = ""; // simply send HTTP headers and status code + res->data = preempt_prefix; // simply send HTTP headers and status code } else if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { - res->data = format_anthropic_sse(first_result_json); + res->data = preempt_prefix + format_anthropic_sse(first_result_json); } else if (res_type == TASK_RESPONSE_TYPE_OAI_RESP) { - res->data = format_oai_resp_sse(first_result_json); + res->data = preempt_prefix + format_oai_resp_sse(first_result_json); } else { - res->data = format_oai_sse(first_result_json); + res->data = preempt_prefix + format_oai_sse(first_result_json); } res->status = 200; res->content_type = "text/event-stream"; - res->set_next([res_this = res.get(), res_type, sse_ping_interval](std::string & output) -> bool { + res->set_next([res_this = res.get(), res_type, sse_ping_interval, parked_idx](std::string & output) mutable -> bool { + const bool parked = !parked_idx.empty(); + static auto format_error = [](task_response_type res_type, const json & res_json) { if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { return format_anthropic_sse({ @@ -5356,10 +5809,13 @@ std::unique_ptr server_routes::handle_completions_impl( // receive subsequent results bool timeout = false; int64_t start_time = ggml_time_ms(); - auto result = rd.next([&timeout, &start_time, sse_ping_interval, &effective_should_stop]() { + // [TAG_PREEMPT] a parked slot produces nothing, so ping at least every 2 s whether or not --sse-ping asked for one, and name it; a shorter interval asked for is kept + const int64_t ping_cfg = sse_ping_interval > 0 ? (int64_t) sse_ping_interval * 1000 : -1; + const int64_t ping_ms = parked ? (ping_cfg > 0 ? std::min(ping_cfg, PREEMPT_KEEPALIVE_MS) : PREEMPT_KEEPALIVE_MS) : ping_cfg; + auto result = rd.next([&timeout, &start_time, ping_ms, &effective_should_stop]() { if (effective_should_stop()) { return true; // should_stop condition met - } else if (sse_ping_interval > 0 && ggml_time_ms() - start_time > (int64_t)sse_ping_interval * 1000) { + } else if (ping_ms > 0 && ggml_time_ms() - start_time > ping_ms) { timeout = true; return true; // timeout } @@ -5369,7 +5825,7 @@ std::unique_ptr server_routes::handle_completions_impl( if (timeout) { // some clients may time out (e.g. undici) will time out if no data is received for a while, so we need to send a ping to keep the connection alive SRV_DBG("%s", "sending SSE ping\n"); - output = ":\n\n"; + output = parked ? ": preempt-keepalive\n\n" : ":\n\n"; return true; } @@ -5385,12 +5841,23 @@ std::unique_ptr server_routes::handle_completions_impl( output = format_error(res_type, res_json); SRV_DBG("%s", "error received during streaming, terminating stream\n"); return false; // terminate on error + } else if (const auto * notice = dynamic_cast(result.get())) { + if (notice->parked) { + parked_idx.insert(notice->index); + } else { + parked_idx.erase(notice->index); + } + output = preempt_notice_comment(*notice); } else { GGML_ASSERT( dynamic_cast(result.get()) != nullptr || dynamic_cast(result.get()) != nullptr ); json res_json = result->to_json(); + if (res_json.is_null()) { + // [TAG_PREEMPT] the empty signal a prompt sends before its first token has nothing to add once a notice has opened the stream + return true; + } if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) { output = format_anthropic_sse(res_json); } else if (res_type == TASK_RESPONSE_TYPE_OAI_RESP) { diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d8..b5c8ab4a8ac 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -448,6 +448,9 @@ server_task_result_ptr server_response::recv(const std::unordered_set & id_ } server_task_result_ptr server_response::recv_with_timeout(const std::unordered_set & id_tasks, int timeout) { + // [TAG_PREEMPT] the timeout is a deadline, not a per-wait duration: send() notify_all()s for every result of every task, and with wait_for() each wakeup restarted the wait + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout); + while (true) { std::unique_lock lock(mutex_results); @@ -459,7 +462,7 @@ server_task_result_ptr server_response::recv_with_timeout(const std::unordered_s } } - std::cv_status cr_res = condition_results.wait_for(lock, std::chrono::seconds(timeout)); + std::cv_status cr_res = condition_results.wait_until(lock, deadline); if (!running) { RES_DBG("%s : queue result stop\n", __func__); std::terminate(); // we cannot return here since the caller is HTTP code diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 9afe3c7f06a..48fc77c960f 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1023,6 +1023,13 @@ void server_task_result_cmpl_partial::update(task_result_state & state) { } } +json server_task_result_preempt_notice::to_json() { + return json { + {"preempted", parked}, + {"n_preempt", n_preempt}, + }; +} + json server_task_result_cmpl_partial::to_json() { GGML_ASSERT(is_updated && "update() must be called before to_json()"); if (is_begin) { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 00734924bc6..f731378b1f6 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -392,6 +392,17 @@ struct server_task_result_cmpl_final : server_task_result { json to_json_anthropic_stream(); }; +// [TAG_PREEMPT] out-of-band notice for a streaming task whose slot was parked or restored, sent as an SSE comment (": preempted", ": resumed") every existing client ignores +struct server_task_result_preempt_notice : server_task_result { + bool parked = false; // true when the slot was just parked, false when restored + int32_t n_preempt = 0; // how many times this task has been parked so far + + virtual bool is_stop() override { + return false; + } + virtual json to_json() override; +}; + struct server_task_result_cmpl_partial : server_task_result { std::string content; llama_tokens tokens; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py index 0c8f5dc5f29..4e7aa6d851c 100644 --- a/tools/server/tests/unit/test_preempt.py +++ b/tools/server/tests/unit/test_preempt.py @@ -1,13 +1,11 @@ import os +import re import time import tempfile import pytest from utils import * -# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked -# (its sequence copied to host RAM, its cells released) instead of every slot being -# terminated. Both tests need more than one slot and --kv-unified, which is the only -# configuration where one slot can take another one's cells. +# Preemption on a unified KV pool: one slot is parked, its sequence copied to host RAM and its cells released, instead of every slot being terminated. Needs --kv-unified. server = ServerPreset.tinyllama2() @@ -38,8 +36,10 @@ def create_server(): os.close(fd) yield os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_SERVER_PREEMPT_GRANULARITY", 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"): @@ -54,21 +54,56 @@ def _complete(n_predict: int, prompt: str = "Hi how are you"): return res -def test_forced_preemption_does_not_change_the_output(): - # Park and restore the only running slot every 8 tokens. With one request the batch - # has the same shape at every step whether or not the slot was parked in between, so - # any difference in the output is the preemption's fault and nothing else's. - global server - server.n_ctx = 512 +_PROMPT_A = "Once upon a time there was a brave knight who" +_PROMPT_B = "The quick brown fox jumps over the lazy dog and" +_PROMPT_C = "In a small village by the sea there lived a fisherman who" + + +def _start(**kwargs) -> LogReader: + """Start the server with these settings, and read its log from the first line.""" + for key, value in kwargs.items(): + setattr(server, key, value) server.start() + return LogReader(server.log_path) + + +def _late(n_predict: int, prompt: str, delay: float = 0.02): + time.sleep(delay) + return _complete(n_predict, prompt) + + +def _complete_all(n_predict: int, prompts=(_PROMPT_A, _PROMPT_B)): + return parallel_function_calls([(_complete, (n_predict, prompt)) for prompt in prompts]) + + +def _complete_all_raw(n_predict: int, prompts): + """As _complete_all, without return_tokens: these ask for thousands of tokens.""" + return parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": prompt, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for prompt in prompts + ]) + + +def _assert_completed(results, n_predict: int, whole: bool = False): + """Every request generated what it asked for; `whole` also pins the untruncated body.""" + for res in results: + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_predict + if whole: + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + + +def test_forced_preemption_does_not_change_the_output(): + _start(n_ctx=512) reference = _complete(64) assert reference.status_code == 200 assert reference.body["timings"]["predicted_n"] == 64 server.stop() os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" - server.start() - log = LogReader(server.log_path) + log = _start() assert "LLAMA_SERVER_PREEMPT_EVERY = 8" in log.drain() preempted = _complete(64) @@ -84,33 +119,40 @@ def test_forced_preemption_does_not_change_the_output(): def test_two_slots_that_overflow_the_pool_together_both_finish(): - # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well - # under 256. Together they do not, 336 against 256. Without preemption the retry - # ladder ends with "Context size has been exceeded" on every processing slot; with it - # the smaller slot is parked until the leader finishes and its cells are purged, and - # then it resumes from the token it was parked on. - global server - server.n_ctx = 256 - server.start() - log = LogReader(server.log_path) + # each request fits the pool alone (168 of 256 cells) but not together; without preemption both end with "Context size has been exceeded" + log = _start(n_ctx=256) n_predict = 160 - 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")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "preempted:" in text assert "resumed after" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict - assert res.body["truncated"] is False - assert len(res.body["tokens"]) == n_predict + _assert_completed(results, n_predict, whole=True) + + +def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages(): + # a block allocator gives a whole block to one sequence, so the planner has to count cells: counting tokens it sees room the allocator cannot find. GRANULARITY injects the size. + os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" + log = _start(n_ctx=256) + assert "LLAMA_SERVER_PREEMPT_GRANULARITY = 64" in log.drain() + + n_predict = 160 + results = _complete_all(n_predict) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + held = [int(n) for n in re.findall(r"kv (\d+)/256", text)] + wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)] + assert held and wanted, f"the planner logged no figures:\n{text}" + assert all(n % 64 == 0 for n in held + wanted), f"not whole blocks: {held} {wanted}" + _assert_completed(results, n_predict, whole=True) _WORDS = ( @@ -134,19 +176,12 @@ def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: if n <= n_tokens: assert n >= n_tokens - 12, f"could not land near {n_tokens} tokens, got {n}" return text, n - # about four tokens per word on this model's vocabulary words = words[: len(words) - max(1, (n - n_tokens) // 8)] raise AssertionError("empty prompt") def test_two_prompts_that_overflow_the_pool_together_both_finish(): - # Neither slot ever generates before the pool is full: both are still processing their - # prompts. A prompt-processing slot is between two chunks of its prompt, which is as - # clean a boundary as the one between two sampled tokens, so it is parked the same way. - global server - server.n_ctx = 256 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) prompt_a, n_a = _prompt_of_about(150, "Alpha") prompt_b, n_b = _prompt_of_about(150, "Bravo") @@ -154,46 +189,27 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish(): assert n_a + n_predict <= 256 and n_b + n_predict <= 256 assert n_a + n_b + 2 * n_predict > 256 - results = parallel_function_calls([ - (_complete, (n_predict, prompt_a)), - (_complete, (n_predict, prompt_b)), - ]) + results = _complete_all(n_predict, [prompt_a, prompt_b]) text = log.drain() assert "Context size has been exceeded" not in text assert "preempted:" in text assert "resumed after" in text + _assert_completed(results, n_predict) for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict assert len(res.body["tokens"]) == n_predict def test_a_generating_slot_and_a_large_prompt_both_finish(): - # One slot is generating a long answer to a short prompt when a large prompt arrives - # beside it. Together they need far more than the pool has. The prompt is admitted - # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish. - # This model produces a thousand tokens a second, so the second request is sent right - # behind the first rather than after a delay: its prompt takes several batches to - # process, which is enough for the two to overlap however fast the first one runs. - global server - server.n_ctx = 256 - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) prompt_b, n_b = _prompt_of_about(150, "Charlie") - # b lives long enough for the two to collide: the first run of this used 16 tokens - # and b was finished and purged before a had grown into it n_predict_a = 230 n_predict_b = 90 assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 assert 8 + n_predict_a + n_b + n_predict_b > 256 - def _late(n_predict, prompt): - time.sleep(0.02) - return _complete(n_predict, prompt) - results = parallel_function_calls([ (_complete, (n_predict_a, "Hi how are you")), (_late, (n_predict_b, prompt_b)), @@ -210,19 +226,11 @@ def _late(n_predict, prompt): def test_preempt_ram_zero_disables_preemption(): - # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the - # KV-full path ends the requests the way it always did. - global server - server.n_ctx = 256 os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) n_predict = 160 - 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")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "preempted:" not in text @@ -231,13 +239,7 @@ def test_preempt_ram_zero_disables_preemption(): def test_metrics_and_slots_report_the_parked_state(): - # A client that wants to tell a parked chat from a slow one reads /slots, and an - # operator reads /metrics. Both must show the preemption happening, and the counters - # must survive the requests finishing. - global server - server.n_ctx = 256 - server.server_metrics = True - server.start() + _start(n_ctx=256, server_metrics=True) res = server.make_request("GET", "/slots") assert res.status_code == 200 @@ -246,10 +248,7 @@ def test_metrics_and_slots_report_the_parked_state(): assert slot["n_preempt"] == 0 n_predict = 160 - 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")), - ]) + results = _complete_all(n_predict) for res in results: assert res.status_code == 200 @@ -270,51 +269,191 @@ 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, only on a backend that can copy asynchronously and signal an event; a CPU-only build falls back and these 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" + return _start(**kwargs).drain() + + +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(): + text = _start_async(n_ctx=512, n_gpu_layer=99) + _require_async(text) + + res_plain = _complete(64) + assert res_plain.status_code == 200 + + server.stop() + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + log = _start() + + 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 + 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(): + text = _start_async(n_ctx=256, n_gpu_layer=99) + _require_async(text) + + n_predict = 160 + + alone = [_complete(n_predict, prompt) for prompt in (_PROMPT_A, _PROMPT_B)] + for res in alone: + assert res.status_code == 200 + + server.stop() + log = _start() + + together = _complete_all(n_predict) + + text = log.drain() + _require_async(text) + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + _assert_completed(together, n_predict) + for res, ref in zip(together, alone): + 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, where the host buffer is freed and the cells handed on, so both have to wait for the copy + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + text = _start_async(n_ctx=512, n_gpu_layer=99) + _require_async(text) + + for i in range(4): + _cancel_soon(96, _PROMPT_A, 0.05 + 0.1 * i) + + 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" + + 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(): + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "0" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + log = _start(n_ctx=512, n_gpu_layer=99) + + 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 + 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: an asynchronous park does not return the cells before update_slots() carries on + log = _start(n_ctx=512, n_gpu_layer=99, n_slots=4) + + 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") + + 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 + + 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, 0.25)), + ]) + + 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 - # plus its first batch would not leave the usual scheduling margin. It must still be - # restored once the first one finishes: with nothing resident there is nobody to keep - # the margin for. Before the fix it was parked for ever, with no restore ever tried. - global server - server.n_ctx = 256 - # the whole prompt in one batch, so the parked slot's first step is the whole prompt - server.n_batch = 256 - server.start() - log = LogReader(server.log_path) + # the second prompt is parked before it takes any cells and is too close to n_ctx to leave the usual margin, but must still be restored once the first finishes + log = _start(n_ctx=256, n_batch=256) - # sized in tokens, not words: the prompt is the token ids of a short sentence repeated base = server.make_request("POST", "/tokenize", data={"content": "Once upon a time there was a little girl"}).body["tokens"] - long_prompt = (base * 64)[:250] + long_prompt = (base * 64)[:240] n_predict = 4 - together = parallel_function_calls([(_complete, (n_predict, long_prompt)) for _ in range(2)]) + together = _complete_all(n_predict, [long_prompt, long_prompt]) text = log.drain() assert "cannot fit the pool" not in text - for res in together: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict + _assert_completed(together, n_predict) def test_the_last_resort_parks_instead_of_ending_everyone(): - # With the planner off nothing is parked ahead of the decode, so two generations that - # fit alone but not together fill the pool until a single token finds no cell. That - # is where upstream ends every slot with the context error. Instead the batch is - # given up, the smaller slot is parked, the larger one finishes, and the parked one - # comes back and finishes too. - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) assert "LLAMA_SERVER_PREEMPT_PLANNER = off" in log.drain() n_predict = 160 - 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")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text @@ -323,48 +462,27 @@ def test_the_last_resort_parks_instead_of_ending_everyone(): assert "last resort: batch given up" in text assert "resumed after" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict - assert res.body["truncated"] is False - assert len(res.body["tokens"]) == n_predict + _assert_completed(results, n_predict, whole=True) def test_the_last_resort_works_with_an_unlimited_budget(): - # --preempt-ram -1 is the documented unlimited setting; it must enable the last resort - # the same as any positive budget does - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" os.environ["LLAMA_ARG_PREEMPT_RAM"] = "-1" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) n_predict = 160 - 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")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "preempted as a last resort" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict + _assert_completed(results, n_predict) def test_the_last_resort_rewinds_a_prompt_in_flight(): - # Same, with a prompt being processed when the pool runs out: the chunk that failed - # is taken back off the slot's tokens and processed again after the resume, so the - # prompt is neither skipped nor fed twice. The prompt is far longer than a batch, so - # the failing chunk is a chunk of it, not its last token. - global server - server.n_ctx = 256 os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) prompt_b, n_b = _prompt_of_about(150, "Charlie") n_predict_a = 230 @@ -372,10 +490,6 @@ def test_the_last_resort_rewinds_a_prompt_in_flight(): assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 assert 8 + n_predict_a + n_b + n_predict_b > 256 - def _late(n_predict, prompt): - time.sleep(0.02) - return _complete(n_predict, prompt) - results = parallel_function_calls([ (_complete, (n_predict_a, "Hi how are you")), (_late, (n_predict_b, prompt_b)), @@ -389,60 +503,29 @@ def _late(n_predict, prompt): assert results[0].body["timings"]["predicted_n"] == n_predict_a assert results[1].status_code == 200 assert results[1].body["timings"]["predicted_n"] == n_predict_b - # the chunk that was in the batch given up is processed once, after the rewind, and - # the count is the prompt plus the BOS the server adds + # the chunk in the batch given up is processed once after the rewind; the count is the prompt plus the BOS the server adds assert results[1].body["timings"]["prompt_n"] == n_b + 1 def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): - # Two generations that each outgrow the pool on their own, with context shift on. The - # resident reaches the limit, shifts, keeps about half the pool and would keep going - # for as long as it has tokens to make, while the parked one never fits beside it. - # After the head has waited its turn the resident is parked in its place, and the two - # take turns until both finish. Long enough that the resident is still going when the - # head's turn comes: this model makes a couple of thousand tokens a second. - global server - server.n_ctx = 256 - server.enable_ctx_shift = True - server.start() - log = LogReader(server.log_path) + # with context shift on the resident would hold half the pool for as long as it generates, so once the head has waited its turn the resident is parked and the two take turns + log = _start(n_ctx=256, enable_ctx_shift=True) n_predict = 12000 - 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")), - ]) + results = _complete_all(n_predict) text = log.drain() assert "Context size has been exceeded" not in text assert "slot context shift" in text assert "rotated out after" in text - for res in results: - assert res.status_code == 200 - assert res.body["timings"]["predicted_n"] == n_predict + _assert_completed(results, n_predict) def test_the_rotation_parks_a_resident_that_lets_the_head_in(): - # Three generations with no end in a 256-cell pool with context shift on: two residents - # cycle through shifts while the third waits parked. Every rotation must let the head - # in, so all three keep finishing their tokens and no stream ends short. - global server - server.n_slots = 3 - server.n_ctx = 384 - server.enable_ctx_shift = True - server.start() + _start(n_slots=3, n_ctx=384, enable_ctx_shift=True) n_predict = 9000 - prompts = [ - "Once upon a time there was a brave knight who", - "The quick brown fox jumps over the lazy dog and", - "In a small village by the sea there lived a fisherman who", - ] - results = parallel_function_calls([ - (server.make_request, ("POST", "/completion", { - "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, - })) for p in prompts - ]) + results = _complete_all_raw(n_predict, (_PROMPT_A, _PROMPT_B, _PROMPT_C)) for res in results: assert res.status_code == 200, res.body assert res.body["tokens_predicted"] == n_predict @@ -452,21 +535,14 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in(): def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): - # One request asking for two completions is one conversation in two slots: a parent - # and a child sharing the prompt. When the two together do not fit the pool there is - # nobody else to park, since the family is charged once and a member of it is not a - # victim for the other, so the request gets the context error it would get alone, and - # the server carries on serving. - global server - server.n_ctx = 256 + # a family member is not a victim for the other, so a two-completion request gets the context error it would get alone and the server carries on os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" - server.start() - log = LogReader(server.log_path) + log = _start(n_ctx=256) res = server.make_request("POST", "/completion", data={ "n_predict": 160, "n_cmpl": 2, - "prompt": "Once upon a time there was a brave knight who", + "prompt": _PROMPT_A, "ignore_eos": True, "return_tokens": True, "temperature": 0.0, @@ -484,33 +560,34 @@ 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(): + # an asynchronous slot keeps its pinned buffer after a restore, and that idle capacity counts against --preempt-ram: unless it is given back, the first restore spends the budget + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "256" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" + text = _start_async(n_ctx=8192, n_gpu_layer=99) + _require_async(text) + log = LogReader(server.log_path) + + n_predict = 1800 + results = _complete_all(n_predict) + + 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" + 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" + + _assert_completed(results, n_predict) + for res in results: + 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 - # at once. The resident is parked before the head is restored and freed, so a rotation - # holds both states together: under this budget the first one asked for is refused and - # said so, and the heads come back when the resident finishes instead. Every stream - # still finishes its tokens and nothing gets the context error. - global server - server.n_slots = 3 - server.n_ctx = 2048 - server.enable_ctx_shift = True + # a rotation holds both states at once, since the resident is parked before the head is restored and freed, so a budget for two heads but not a head plus the resident must refuse os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" - server.start() - # long enough that the resident is still cycling through shifts two seconds after the - # heads were parked, which is when a rotation is first asked for: at 6000 this model - # finished in under three seconds on a fast host and nothing was ever refused + _start(n_slots=3, n_ctx=2048, enable_ctx_shift=True) n_predict = 12000 - prompts = [ - "Once upon a time there was a brave knight who", - "The quick brown fox jumps over the lazy dog and", - "In a small village by the sea there lived a fisherman who", - ] - results = parallel_function_calls([ - (server.make_request, ("POST", "/completion", { - "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, - })) for p in prompts - ]) + results = _complete_all_raw(n_predict, (_PROMPT_A, _PROMPT_B, _PROMPT_C)) for res in results: assert res.status_code == 200, res.body assert res.body["tokens_predicted"] == n_predict @@ -521,10 +598,6 @@ def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_w def test_a_recurrent_model_is_served_without_preemption(): - # A recurrent cache holds one state per sequence whatever its length, so the token - # count the planner measures says nothing about it: preemption is off for such a - # model, said so at load, and the forced-park knob parks nothing. - global server path = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") if path: server.model_file = path @@ -536,10 +609,7 @@ def test_a_recurrent_model_is_served_without_preemption(): server.n_ctx = 1024 os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" server.start(timeout_seconds=300) - results = parallel_function_calls([ - (_complete, (64, "Once upon a time")), - (_complete, (64, "The quick brown fox")), - ]) + results = _complete_all(64, ["Once upon a time", "The quick brown fox"]) for res in results: assert res.status_code == 200, res.body assert res.body["tokens_predicted"] == 64 diff --git a/tools/server/tests/unit/test_preempt_notify.py b/tools/server/tests/unit/test_preempt_notify.py new file mode 100644 index 00000000000..fa6a956e878 --- /dev/null +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -0,0 +1,255 @@ +import os +import tempfile +import threading +import time +import pytest +import requests +from utils import * + +# [TAG_PREEMPT] a streaming client is told when its slot is parked and restored, as SSE comments every existing client ignores; a keepalive every 2 s keeps proxies from giving up + +server = ServerPreset.tinyllama2() + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + server.temperature = 0.0 + server.seed = 42 + # A build without libcurl cannot fetch the model itself; point it at a local copy. + local = os.environ.get("LLAMA_SERVER_TEST_MODEL") + if local: + server.model_hf_repo = None + server.model_hf_file = None + server.model_file = local + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + + +def _stream_raw(path: str, data: dict) -> tuple[list[str], list[str]]: + """The SSE lines of one streaming request: (comment lines, data lines).""" + url = f"http://{server.server_host}:{server.server_port}{path}" + res = requests.post(url, json=data, stream=True) + assert res.status_code == 200 + comments, datas = [], [] + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line.startswith(":"): + comments.append(line) + elif line.startswith("data: "): + datas.append(line[6:]) + return comments, datas + + +def _content(datas: list[str]) -> str: + out = "" + for d in datas: + if d == "[DONE]": + break + j = json.loads(d) + if "content" in j: + out += j["content"] + for ch in j.get("choices", []) or []: + delta = ch.get("delta") or {} + out += delta.get("content") or "" + return out + + +def _completion_payload(n_predict: int) -> dict: + return { + "n_predict": n_predict, + "prompt": "Hi how are you", + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + "stream": True, + } + + +def _chat_payload(n_predict: int) -> dict: + return { + "max_tokens": n_predict, + "messages": [{"role": "user", "content": "Hi how are you"}], + "temperature": 0.0, + "seed": 42, + "stream": True, + } + + +_PROMPT_A = "Once upon a time there was a brave knight who" +_PROMPT_B = "The quick brown fox jumps over the lazy dog and" + + +def _start(**kwargs): + """Start the server with these settings.""" + for key, value in kwargs.items(): + setattr(server, key, value) + server.start() + + +def _final(datas: list[str]) -> dict: + """The last response object of a finished stream, past the [DONE] marker.""" + return json.loads([d for d in datas if d != "[DONE]"][-1]) + + +def _stream_both(n_predict: int): + """One streaming completion per prompt, both at once.""" + return parallel_function_calls([ + (_stream_raw, ("/completion", _completion_payload(n_predict) | {"prompt": prompt})) + for prompt in (_PROMPT_A, _PROMPT_B) + ]) + + +def test_a_stream_announces_its_parks_and_the_body_is_unchanged(): + _start(n_ctx=512) + ref_comments, ref_datas = _stream_raw("/completion", _completion_payload(64)) + assert not any(c.startswith(": preempted") or c.startswith(": resumed") for c in ref_comments) + assert _content(ref_datas) + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + comments, datas = _stream_raw("/completion", _completion_payload(64)) + + parked = [c for c in comments if c == ": preempted"] + resumed = [c for c in comments if c == ": resumed"] + assert len(parked) >= 6, comments + assert len(resumed) == len(parked), comments + seq = [c for c in comments if c in (": preempted", ": resumed")] + assert seq == [": preempted", ": resumed"] * len(parked), seq + def _pieces(ds): + return [json.loads(d).get("content") for d in ds if d != "[DONE]"] + + assert _pieces(datas) == _pieces(ref_datas) + assert _content(datas) == _content(ref_datas) + final, ref_final = json.loads(datas[-1]), json.loads(ref_datas[-1]) + assert final["tokens_predicted"] == ref_final["tokens_predicted"] == 64 + + +def test_the_oai_chat_stream_carries_the_same_comments(): + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=512) + comments, datas = _stream_raw("/v1/chat/completions", _chat_payload(48)) + assert ": preempted" in comments and ": resumed" in comments + assert datas[-1] == "[DONE]" + assert _content(datas) + + +def test_non_streaming_requests_see_nothing(): + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=512) + res = server.make_request("POST", "/completion", data={ + "n_predict": 32, + "prompt": "Hi how are you", + "ignore_eos": True, + "temperature": 0.0, + "seed": 42, + }) + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == 32 + assert "preempted" not in res.body + + +def test_two_overflowing_streams_both_finish_and_the_parked_one_says_so(): + _start(n_ctx=256) + + n_predict = 160 + results = _stream_both(n_predict) + announced = 0 + for comments, datas in results: + final = _final(datas) + assert final["timings"]["predicted_n"] == n_predict + assert final["truncated"] is False + if ": preempted" in comments: + announced += 1 + assert ": resumed" in comments + assert announced >= 1, [r[0] for r in results] + + +def test_a_stream_parked_before_its_first_token_starts_with_the_notice(): + # n_batch: the whole prompt in one batch, so the planner sees its size at once + _start(n_ctx=512, n_batch=512) + url = f"http://{server.server_host}:{server.server_port}/completion" + first = _completion_payload(390) | {"prompt": " ".join([_PROMPT_A] * 6)} + second = _completion_payload(32) | {"prompt": " ".join([_PROMPT_B] * 14)} + + timeline = [] + lock = threading.Lock() + + def _run(name, payload, started=None): + res = requests.post(url, json=payload, stream=True) + assert res.status_code == 200 + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if not line: + continue + with lock: + timeline.append((time.monotonic(), name, line)) + if started is not None and line.startswith("data: "): + started.set() + + started = threading.Event() + t = threading.Thread(target=_run, args=("first", first, started)) + t.start() + assert started.wait(30) + _run("second", second) + t.join(60) + + second_lines = [(ts, line) for ts, name, line in timeline if name == "second"] + first_end = max(ts for ts, name, _ in timeline if name == "first") + assert second_lines[0][1] == ": preempted", second_lines[:3] + assert second_lines[0][0] < first_end + events = [line for _, line in second_lines if line in (": preempted", ": resumed") or line.startswith("data: ")] + assert events[0] == ": preempted" and events[1] == ": resumed" and events[2].startswith("data: "), events[:3] + datas = [line[6:] for _, line in second_lines if line.startswith("data: ")] + assert _content(datas) + assert _final(datas)["tokens_predicted"] == 32 + + +def test_a_resident_rotated_out_for_a_parked_head_is_told_so(): + _start(n_ctx=256, enable_ctx_shift=True) + n_predict = 12000 + results = _stream_both(n_predict) + n_parked = 0 + for comments, datas in results: + final = _final(datas) + assert final["tokens_predicted"] == n_predict + seq = [c for c in comments if c in (": preempted", ": resumed")] + assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq + n_parked += len(seq) // 2 + assert n_parked >= 2, [r[0] for r in results] + + +def test_an_oversized_prompt_is_errored_instead_of_parked(): + # A slot just given a task has not passed the prompt checks yet, and a notice opens the stream, so parking it would turn a plain error response into 200 plus an in-stream one. + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=512, n_batch=512) + url = f"http://{server.server_host}:{server.server_port}/completion" + resident = _completion_payload(390) | {"prompt": " ".join([_PROMPT_A] * 6)} + oversized = _completion_payload(16) | {"prompt": " ".join([_PROMPT_B] * 80)} + + started = threading.Event() + + def _run_resident(): + res = requests.post(url, json=resident, stream=True) + assert res.status_code == 200 + for raw in res.iter_lines(): + if raw.decode("utf-8").startswith("data: "): + started.set() + + t = threading.Thread(target=_run_resident) + t.start() + try: + assert started.wait(60) + res = requests.post(url, json=oversized, stream=True) + body = res.text + assert res.status_code != 200, body + assert not body.lstrip().startswith(":"), body + assert "error" in json.loads(body), body + finally: + t.join(120)