diff --git a/.github/workflows/bench.yml.disabled b/.github/workflows/bench.yml.disabled index f2d7e16e981a..5829359ad10e 100644 --- a/.github/workflows/bench.yml.disabled +++ b/.github/workflows/bench.yml.disabled @@ -162,7 +162,7 @@ jobs: tools/server/bench/*.log - name: Commit status - uses: Sibz/github-status-action@v1 + uses: Sibz/github-status-action@faaa4d96fecf273bd762985e0e7f9f933c774918 # v1 with: authToken: ${{secrets.GITHUB_TOKEN}} sha: ${{ inputs.sha || github.event.pull_request.head.sha || github.sha }} @@ -172,7 +172,7 @@ jobs: state: 'success' - name: Upload benchmark images - uses: devicons/public-upload-to-imgur@v2.2.2 + uses: devicons/public-upload-to-imgur@352cf5f2805c692539a96cfe49a09669e6fca88e # v2.2.2 continue-on-error: true # Important as it looks unstable: 503 id: imgur_step with: @@ -221,7 +221,7 @@ jobs: echo "IMAGE_3=${{ fromJSON(steps.imgur_step.outputs.imgur_urls)[3] }}" >> $GITHUB_ENV - name: Comment PR - uses: mshick/add-pr-comment@v2 + uses: mshick/add-pr-comment@b8f338c590a895d50bcbfa6c5859251edc8952fc # v2 id: comment_pr if: ${{ github.event.pull_request != '' && matrix.pr_comment_enabled == 'true' }} with: diff --git a/common/arg.cpp b/common/arg.cpp index 74241f931285..1b7e477c72c3 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1315,6 +1315,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; @@ -1728,6 +1733,25 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.cache_ram_mib = value; } ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--preempt-ram"}, "N", + string_format("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: %d - disabled, -1 - no limit)", params.preempt_ram_mib), + [](common_params & params, int value) { + 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: the copy out " + "overlaps with the slots that keep decoding, while a copy back in, and a kv-full retry behind a " + "copy out that has not landed, 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 d162a38800e0..944028da7180 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); @@ -1337,6 +1344,11 @@ common_init_result::common_init_result(common_params & params, bool model_only) return; } + if (!common_exact_concurrency_model(params, model)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to create a context, see the error above\n"); + return; + } + const llama_vocab * vocab = llama_model_get_vocab(model); // load and optionally apply lora adapters @@ -1403,6 +1415,12 @@ common_init_result::common_init_result(common_params & params, bool model_only) pimpl->context.reset(lctx); + if (!common_exact_concurrency_context(params, lctx)) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to serve this context, see the error above\n"); + pimpl->context.reset(); + return; + } + set_process_priority(params.cpuparams.priority); pimpl->threadpools.init(lctx, params); @@ -1433,6 +1451,148 @@ 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_batch_geometry(int n_batch, int n_ubatch, int n_decode_width, int * n_batch_min) { + // an unset ubatch is the whole batch, and a ubatch never exceeds it + const int n_ub = std::min(n_batch, n_ubatch <= 0 ? n_batch : n_ubatch); + + const int n_min = n_ub + std::max(0, n_decode_width); + + if (n_batch_min) { + *n_batch_min = n_min; + } + + return n_batch >= n_min; +} + +// [TAG_EXACT_CONCURRENCY] the refusals that need the loaded model, run before a context exists +bool common_exact_concurrency_model(const common_params & params, const llama_model * model) { + if (!common_exact_concurrency() || params.mmproj.path.empty()) { + return true; + } + + // the paged pool places a cell from the sequence and the position alone, and M-RoPE gives every token of one image the same temporal position, so the second of them lands on the first one's cell and the batch is refused at the first image + const llama_rope_type rope_type = llama_model_rope_type(model); + + if (rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE) { + COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support M-RoPE together with a projector: the tokens of one image share a temporal position and the paged pool would give them one cell\n"); + return false; + } + + return true; +} + +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; + } + } + + // a prompt is added to a batch in whole ubatches, so a batch that cannot hold one beside a decode step of every slot would leave a prefill shorter ubatches than it gets alone, and the mode would report itself as on while a shared step changed the prompt's arithmetic + // a causal context clamps the batch to the context size, so that is the batch a prefill really gets; an unset -c is only known once the context exists, which common_exact_concurrency_context() checks + const int n_batch_eff = params.n_ctx > 0 ? std::min(params.n_ctx, params.n_batch) : params.n_batch; + + int n_batch_min = 0; + + if (!common_exact_batch_geometry(n_batch_eff, params.n_ubatch, n_cols, &n_batch_min)) { + COM_ERR("LLAMA_EXACT_CONCURRENCY needs a batch of at least %d tokens for a %d-token ubatch " + "and a decode step of %d slots (%d columns), but the batch is %d: a prefill beside " + "a running slot would be split into shorter ubatches than the same prompt gets alone. " + "Raise -b to %d (and -c to at least that), or lower -ub.\n", + n_batch_min, std::min(n_batch_eff, params.n_ubatch <= 0 ? n_batch_eff : params.n_ubatch), + std::max(1, params.n_parallel), n_cols, n_batch_eff, n_batch_min); + 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; +} + +bool common_exact_concurrency_context(const common_params & params, const llama_context * ctx) { + if (!common_exact_concurrency()) { + return true; + } + + const int n_cols = common_exact_decode_width(params); + + if (n_cols < 0) { + return false; // already reported by common_exact_concurrency_init() + } + + // the context clamps the batch to the context size and the ubatch to the batch, and an unset -c takes its size from the model or from the fit to device memory, so this is the geometry a prefill really gets + const int n_batch = (int) llama_n_batch(ctx); + const int n_ubatch = (int) llama_n_ubatch(ctx); + + int n_batch_min = 0; + + if (!common_exact_batch_geometry(n_batch, n_ubatch, n_cols, &n_batch_min)) { + COM_ERR("LLAMA_EXACT_CONCURRENCY needs a batch of at least %d tokens for a %d-token ubatch " + "and a decode step of %d slots (%d columns), but the context was created with a batch " + "of %d: a context of %d tokens clamps it, so a prefill beside a running slot would be " + "split into shorter ubatches than the same prompt gets alone. Raise -c to at least %d " + "(-fitc as well when the context was fitted to device memory), or lower -ub.\n", + n_batch_min, n_ubatch, std::max(1, params.n_parallel), n_cols, n_batch, + (int) llama_n_ctx(ctx), n_batch_min); + 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 63d0badd0f74..60bda08d74fd 100644 --- a/common/common.h +++ b/common/common.h @@ -630,6 +630,8 @@ struct common_params { int32_t kv_unified_per_slot = 0; // max context per parallel slot; 0 = unset 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 = 0; // host RAM for parked (preempted) sequences: 0 = preemption off (the default), -1 = no limit + 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 @@ -947,6 +949,23 @@ 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); + +// [TAG_EXACT_CONCURRENCY] whether a batch of this shape holds a whole prompt ubatch beside a decode step of every slot, which a prefill needs to be split into the ubatches it would get alone; n_batch_min reports the batch size that would +bool common_exact_batch_geometry(int n_batch, int n_ubatch, int n_decode_width, int * n_batch_min = nullptr); + +// 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); + +// the same for what only the loaded model tells: false if the model must not be served in exact mode +bool common_exact_concurrency_model(const common_params & params, const struct llama_model * model); + +// the same for the geometry the created context settled on, which the context size may have clamped below what -b and -ub asked for +bool common_exact_concurrency_context(const common_params & params, const struct llama_context * ctx); + 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/common/speculative.cpp b/common/speculative.cpp index 2db381d58086..82fc175281e7 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -810,7 +810,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + // the per-call bound comes from the caller's remaining context, so it stops the loop as well as the configured maximum + if ((params.n_max <= (int) result.size()) || + (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; continue; @@ -1193,7 +1195,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n = (int32_t) dp.n_past; - const int32_t n_draft = params.n_max; + // the caller's remaining context bounds the block as well as the configured maximum: the whole block is decoded before any truncation + const int32_t n_draft = dp.n_max > 0 ? std::min(params.n_max, dp.n_max) : params.n_max; const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1); i_block_beg[seq_id] = batch.n_tokens; @@ -1691,7 +1694,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + // the per-call bound comes from the caller's remaining context, so it stops the loop as well as the configured maximum + if ((params.n_max <= (int) result.size()) || + (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; continue; diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e35..84a2f8458ea2 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 1cd81eeaebcd..897da6ca5f82 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 ef05905cf9ab..241f5bfb1dd7 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 @@ -215,6 +215,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 9f484c4eb86f..5e979c5eebfc 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 4feaf77a5ece..da3838a52fa8 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -184,6 +184,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; @@ -560,6 +564,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); @@ -636,6 +652,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 e4b5bd254747..7271b6b632b6 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 c2745014a192..20c0e59df711 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2660,6 +2660,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; @@ -2952,6 +2956,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 8cece71f186f..7ea548bcbce8 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 7d14ce9067ee..b3006642ad48 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 ea7b46d5708a..5ab83a1f7c89 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -10,6 +10,12 @@ #define HALF_MAX_HALF __float2half(65504.0f/2) // Use neg. of this instead of -INFINITY to initialize KQ max vals to avoid NaN upon subtraction. #define SOFTMAX_FTZ_THRESHOLD -20.0f // Softmax exp. of values smaller than this are flushed to zero to avoid NaNs. +// [TAG_EXACT_CONCURRENCY] the page table of the paged path, which only the ordinary flash +// attention op carries: another op is free to keep a tensor of its own in the same slot +static __forceinline__ const ggml_tensor * ggml_cuda_fattn_pages(const ggml_tensor * dst) { + return dst->op == GGML_OP_FLASH_ATTN_EXT ? dst->src[5] : nullptr; +} + // log(2) = 0.6931, by adding this to the KQ maximum used for the softmax the numerical range representable // by the VKQ accumulators is effectively being shifted up by a factor of 2. // This reduces issues with numerical overflow but also causes larger values to be flushed to zero. @@ -1106,7 +1112,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 (!use_sparse && 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 (!use_sparse && !ggml_cuda_fattn_pages(dst) && 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); @@ -1164,6 +1172,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 (ggml_cuda_fattn_pages(dst) || 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); @@ -1238,7 +1253,7 @@ void launch_fattn( V_data, mask ? ((const char *) mask->data) : nullptr, sinks ? ((const char *) sinks->data) : nullptr, - KV_max.ptr, + ggml_cuda_fattn_pages(dst) ? (const int *) ggml_cuda_fattn_pages(dst)->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 57a285565913..fc35c0f6358d 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 ab262156fafc..9b12ef629655 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -626,6 +626,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) { @@ -738,6 +743,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 (const ggml_tensor * pages = ggml_cuda_fattn_pages(dst)) { + 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(pages->type == GGML_TYPE_I32 && ggml_is_contiguous(pages)); + GGML_ASSERT(pages->ne[0] == 1 + dst->src[1]->ne[1]/FATTN_KQ_STRIDE); + GGML_ASSERT(pages->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, /*use_sparse =*/ 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 56e398a66d09..ff5ce6c5fbf0 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1791,6 +1791,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; @@ -1818,6 +1823,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; @@ -1846,63 +1855,286 @@ 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; - const bool f32_pedantic = src0->type == GGML_TYPE_F32 && - ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; - 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 (!f32_pedantic && 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; + // a pedantic F32 request never takes the mmf path; checked here so the batch-invariant + // width search and the column split below pick the same implementation the switch does + const bool f32_pedantic = src0->type == GGML_TYPE_F32 && + ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; + if (!f32_pedantic && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { + 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 + // every mode owes the caller the batch-of-one column policy, and a plane the policy never sees is left batched + if (ggml_cuda_batch_invariant() && 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; + } + + 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; } - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + + 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 @@ -1919,9 +2151,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)) { @@ -1929,17 +2164,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 (!f32_pedantic && ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + if (!f32_pedantic && 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]; @@ -1954,6 +2223,18 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * const bool f32_pedantic = src0->type == GGML_TYPE_F32 && ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; + // [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); @@ -1977,8 +2258,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * return; } - if (!f32_pedantic && ggml_cuda_should_use_mmf( - src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + if (!f32_pedantic && ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); return; } @@ -3508,9 +3788,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; @@ -5624,6 +5905,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, @@ -5640,6 +5936,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 @@ -5741,6 +6038,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 c96cf64f67c7..c9218d48c7b8 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -580,6 +580,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( @@ -616,9 +630,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; @@ -1422,7 +1437,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; @@ -1512,6 +1531,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 5605bf7a4e60..688c944c1f90 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 2fc0fe9fdbb7..1a09deffd56a 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 6d725c7ec196..ebecf679950e 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 61c31d6f2912..b6305eb2a240 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1267,6 +1267,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 && @@ -1685,6 +1690,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 a39df2a878c5..48325cad4f2b 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -5930,7 +5930,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: @@ -6036,6 +6037,7 @@ static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .event_new = */ ggml_backend_hexagon_device_event_new, /* .event_free = */ ggml_backend_hexagon_device_event_free, /* .event_synchronize = */ ggml_backend_hexagon_device_event_synchronize, + /* .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 afd6f521011e..590df1bd24bd 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1703,6 +1703,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 3bd6abd06fdc..69af3678fa06 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -821,6 +821,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 3002835e8aea..0ed3e8fc5c77 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -8745,6 +8745,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; @@ -12410,6 +12414,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 a7956227830b..51cc2168ce72 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1171,6 +1171,10 @@ static ggml_openvino_op_support is_op_supported_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 {false, "FLASH_ATTN_EXT with a page table is CUDA only"}; + } float scale = 1.0f; float max_bias = 0.0f; float logit_softcap = 0.0f; @@ -1495,6 +1499,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 cc7d7206933f..fa5f1cfb6ede 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -2185,7 +2185,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; } @@ -2233,6 +2236,7 @@ static const struct ggml_backend_device_i ggml_backend_rpc_device_i = { /* .event_new = */ ggml_backend_rpc_device_event_new, /* .event_free = */ ggml_backend_rpc_device_event_free, /* .event_synchronize = */ ggml_backend_rpc_device_event_synchronize, + /* .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 4091f73a4674..05b66a3a9e58 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6649,7 +6649,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; } @@ -6755,6 +6756,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 987ce9dd110c..13a70c594df6 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 42e79233eb68..75b16cd016ed 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -19491,6 +19491,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]; @@ -20168,6 +20172,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 f06a9c872db9..38f2cfca7bc5 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4425,6 +4425,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 && @@ -4679,6 +4685,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 4007ac9dfc7d..bbd74fb9d5aa 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 ec7ce233145a..89c6c36a0f19 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 ef7a012c43a1..fc5c2a2abac3 100644 --- a/include/llama.h +++ b/include/llama.h @@ -804,6 +804,22 @@ 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 block: 1 ordinarily, larger where a mode places cells in blocks + // n contiguous tokens then occupy round_up(n, granularity) cells; a sequence left with holes still holds every block one live cell is in + LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem); + + // [TAG_PREEMPT] run the in-place update a seq_add() recorded, which llama_decode() would otherwise run at the start of the next batch + // takes the context because the update is a graph; returns true when one was run + LLAMA_API bool llama_memory_update(struct llama_context * ctx); + + // [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 // @@ -936,6 +952,47 @@ 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 + // Also NULL for a recurrent or hybrid model: its states move between rows on every decode, so a transfer beside a decode can read another sequence + 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-adapter.cpp b/src/llama-adapter.cpp index e6678a66d2a9..309840167d19 100644 --- a/src/llama-adapter.cpp +++ b/src/llama-adapter.cpp @@ -351,6 +351,21 @@ static void llama_adapter_lora_init_impl(llama_model & model, const char * path_ LLAMA_LOG_DEBUG("%s: lora for '%s' -> '%s'\n", __func__, model_tensor->name, ggml_backend_buft_name(buft)); + // [TAG_EXACT_CONCURRENCY] the adapter follows the weight it adapts, so a weight the mode + // leaves on the host puts the adapted matmul there too. token_embd is exempt from the + // context's weight check because get_rows is not offloaded by width, but its adapter is + // applied with a mul_mat, which is: see llm_graph_context::build_inp_embd(). + if (llama_exact_concurrency() && !llama_exact_buft_invariant(buft)) { + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but the lora for '%s' would sit in a %s " + "buffer, which has no batch-invariant kernels: the adapted matmul's result would " + "depend on how many sequences share the step (move the tensor to the device that " + "holds the layers, for example --override-tensor %s=CUDA0, or serve this adapter " + "without the mode)\n", + __func__, model_tensor->name, ggml_backend_buft_name(buft), model_tensor->name); + + throw std::runtime_error("exact concurrency: a lora weight is not on the CUDA backend"); + } + ggml_context * dev_ctx = ctx_for_buft(buft); // validate tensor shape if (is_token_embd) { diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48f..5d52f5bac0b2 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -507,7 +507,54 @@ 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_repeated_positions() const { + std::vector n_per_seq(n_seq_max, 0); + + for (int32_t i = 0; i < batch.n_tokens; ++i) { + for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { + n_per_seq[batch.seq_id[i][s]]++; + } + } + + for (uint32_t s = 0; s < n_seq_max; ++s) { + if (n_per_seq[s] > seq_pos[s].size()) { + 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 +565,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 +590,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 a3d1889d4a04..b70987864503 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -105,7 +105,16 @@ 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; + + // [TAG_EXACT_CONCURRENCY] true if a sequence has several tokens at one position, which the paged pool would give one cell + bool has_repeated_positions() 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 dcccf005d690..34d3db0ccea4 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 @@ -32,6 +33,52 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) { throw std::runtime_error("Unsupported ctx type"); } +// [TAG_EXACT_CONCURRENCY] the caches check where the KV lives; this checks the weights. Every per-layer weight and the output head must sit on a backend with the mode's kernels, else a sequence's own matmuls change with the width of the step it shares. +// token_embd is exempt: it feeds GET_ROWS, a per-row copy that reports a batch size of 0 to the offload test, so it stays put at every width. A tied head is that same tensor and model.output points at it, so the head check covers it; a lora on it runs as MUL_MAT, which llama_adapter_lora_init_impl() refuses on a host buffer. +static void llama_exact_check_weights(const llama_model & model) { + auto host_buft = [](const ggml_tensor * t) -> ggml_backend_buffer_type_t { + if (!t || !t->buffer) { + return nullptr; + } + + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(t->buffer); + + return llama_exact_buft_invariant(buft) ? nullptr : buft; + }; + + auto refuse = [&model](const char * name, ggml_backend_buffer_type_t buft, const char * what) { + const std::string tname = name; + + const char * fix = "pass -ngl to offload every layer, and no --override-tensor that keeps one on the host"; + + if (tname.find("_exps") != std::string::npos) { + fix = "do not pass --cpu-moe or --n-cpu-moe, and no --override-tensor that keeps an expert on the host"; + } else if (model.has_tensor_overrides()) { + fix = "drop the --override-tensor that placed it there, and pass -ngl to offload every layer"; + } + + LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s %s is in a %s buffer, which has no " + "batch-invariant kernels: its result would depend on how many sequences share the step (%s)\n", + __func__, what, tname.c_str(), ggml_backend_buft_name(buft), fix); + + throw std::runtime_error("exact concurrency: a weight is not on the CUDA backend"); + }; + + for (const auto & [name, t] : model.tensors_by_name) { + if (name.rfind("blk.", 0) != 0) { + continue; + } + + if (auto * buft = host_buft(t)) { + refuse(name.c_str(), buft, "layer weight"); + } + } + + if (auto * buft = host_buft(model.output)) { + refuse(ggml_get_name(model.output), buft, "the output head"); + } +} + struct llm_fused_op_probe { llm_fused_op op; const char * name; @@ -101,6 +148,17 @@ 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"); + } + + if (!hparams.vocab_only) { + llama_exact_check_weights(model); + } + } + 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 +451,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 +540,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); + } + // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation if (!model.hparams.no_alloc && !opt_ctx) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { @@ -1197,6 +1273,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; @@ -1586,6 +1667,10 @@ int llama_context::encode(const llama_batch & batch_inp) { } } + if (!state_copy_fences.empty()) { + state_seq_copy_fence(); + } + return 0; } @@ -1707,6 +1792,12 @@ int llama_context::decode(const llama_batch & batch_inp) { return -1; } + // [TAG_EXACT_CONCURRENCY] an invalid batch, not a full cache: left to the memory it came back as 1, which callers retry + if (llama_exact_concurrency() && (balloc->has_shared_tokens() || balloc->has_repeated_positions())) { + LLAMA_LOG_ERROR("%s: exact concurrency needs every token at one sequence id and one position of its own\n", __func__); + return -1; + } + const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs(); @@ -2033,6 +2124,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; } @@ -2575,16 +2670,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 { @@ -2614,10 +2825,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; @@ -2626,6 +2835,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 { @@ -2633,9 +2848,54 @@ 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 (const auto & rinfo : rinfos) { - ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); + for (size_t i = 0; i < rinfos.size();) { + auto * tensor = rinfos[i].tensor; + size_t end = i + 1; + 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; + + 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 &) { + } + if (!staging.empty()) { + ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); + for (size_t j = i; j < end; ++j) { + const auto & rinfo = rinfos[j]; + GGML_ASSERT(rinfo.offset <= tensor_bytes && rinfo.size <= tensor_bytes - rinfo.offset); + memcpy(staging.data() + rinfo.offset, rinfo.ptr, rinfo.size); + } + ggml_backend_tensor_set(tensor, staging.data(), 0, tensor_bytes); + i = end; + continue; + } + } + 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; } } @@ -2666,10 +2926,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; @@ -2678,6 +2936,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 { @@ -3062,6 +3326,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) { @@ -3135,6 +3666,243 @@ 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() { + // [TAG_STATE_ASYNC] a recurrent state keeps no fixed row: find_slot() gathers the live rows together, so a decode beside a transfer moves or overwrites the row the transfer reads. A hybrid carries that half too + if (llm_arch_is_recurrent(model.arch) || llm_arch_is_hybrid(model.arch)) { + LLAMA_LOG_INFO("%s: this model moves sequence states between rows, so they are copied synchronously\n", __func__); + return nullptr; + } + + 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"); @@ -3290,6 +4058,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 @@ -4121,6 +4894,22 @@ 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(); +} + +bool llama_memory_update(llama_context * ctx) { + if (!ctx) { + return false; + } + + return ctx->memory_update(false); +} + // llama state API // deprecated @@ -4216,6 +5005,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 bf91daa8b562..f44f505a05f3 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 3909546e10e3..c475d97b5249 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); @@ -1087,6 +1102,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); @@ -2604,7 +2620,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * v_mla, int64_t n_kv_max, 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 @@ -2637,6 +2654,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); @@ -2828,6 +2846,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); @@ -2890,7 +2910,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, 0, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il, inp->self_pages); cb(cur, "kqv_out", il); if (inp->self_v_rot) { @@ -2924,6 +2944,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); { @@ -3300,6 +3322,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); { @@ -3417,6 +3441,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 f79a274c9112..7b926522d85d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -320,6 +320,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, @@ -1188,7 +1189,8 @@ struct llm_graph_context { ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] int64_t n_kv_max, 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 b3a94b946d28..923845e1ecd6 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,160 @@ 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); } } + +// [TAG_EXACT_CONCURRENCY] +bool llama_exact_backend_name(const char * reg_name) { + return reg_name && (strcmp(reg_name, "CUDA") == 0 || strcmp(reg_name, "ROCm") == 0 || strcmp(reg_name, "MUSA") == 0); +} + +// [TAG_EXACT_CONCURRENCY] a host buffer is the case that matters: the scheduler runs an op on the backend that holds its weight, and moves a host weight's op to the GPU only once the batch is wide enough (ggml_backend_cuda_device_offload_op). +// The CPU matmul picks between its SGEMM and its vector dot by the batch width too. +bool llama_exact_buft_invariant(ggml_backend_buffer_type_t buft) { + if (!buft || ggml_backend_buft_is_host(buft)) { + return false; + } + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; + + return reg && llama_exact_backend_name(ggml_backend_reg_name(reg)); +} + +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 4988b06d2ca0..dc4e21eb8114 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -1,6 +1,7 @@ #pragma once #include "ggml.h" // for ggml_log_level +#include "ggml-backend.h" #include #include @@ -103,3 +104,17 @@ 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] whether a backend registry carries the mode's batch-invariant kernels +bool llama_exact_backend_name(const char * reg_name); + +// [TAG_EXACT_CONCURRENCY] whether a tensor placed in this buffer type is computed by such a backend +bool llama_exact_buft_invariant(ggml_backend_buffer_type_t buft); + +// [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 120e45732af2..3628082a7dc2 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -62,6 +63,68 @@ 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; + } + + return llama_exact_backend_name(ggml_backend_reg_name(reg)); +} + +// [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_prec_set_acc(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, @@ -86,6 +149,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 @@ -99,6 +165,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: @@ -222,6 +309,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"); @@ -366,7 +486,99 @@ 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{}); + exact_page_live .assign(cells.size()/exact_page_size, 0); + + 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_live[i/exact_page_size]; + } + + 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) { + // what was maintained has to say what the cells say + const auto kept = exact_page_owner; + const auto kept_live = exact_page_live; + + 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); + GGML_ASSERT(kept_live[p] == exact_page_live[p]); + } + } +} + +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; + + ++exact_page_live[idx/exact_page_size]; +} + +// [TAG_EXACT_CONCURRENCY] a page stays with its sequence for as long as one of its cells is live, +// so a removal frees it only when it takes the last one. Counting per page is what keeps a removal +// that empties nothing, such as the rejected tail of every accepted speculative step, from costing +// a rescan of the pool. +void llama_kv_cache::exact_pages_release(uint32_t idx) { + ++exact_page_n_release; + + if (exact_page_owner_dirty || exact_page_owner.empty()) { + return; + } + + const uint32_t page = idx/exact_page_size; + + GGML_ASSERT(exact_page_live[page] > 0); + + if (--exact_page_live[page] == 0) { + exact_page_owner[page] = exact_page{}; + } +} + 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; @@ -408,6 +620,11 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { } if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) { + // [TAG_EXACT_CONCURRENCY] the cell is gone; the page goes with the last of them + if (exact_pages) { + exact_pages_release(i); + } + if (new_head == cells.size()) { new_head = i; } @@ -431,6 +648,10 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { continue; } + if (exact_pages) { + exact_pages_release(i); + } + cells.rm(i); if (new_head == cells.size()) { @@ -454,6 +675,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()); @@ -555,6 +784,11 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) { for (uint32_t i = 0; i < cells.size(); ++i) { if (cells.seq_keep(i, seq_id)) { + // [TAG_EXACT_CONCURRENCY] as in seq_rm, the cell emptied here + if (exact_pages) { + exact_pages_release(i); + } + if (new_head == cells.size()) { new_head = i; } @@ -573,6 +807,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"); @@ -623,6 +865,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"); @@ -706,11 +955,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; @@ -757,11 +1018,19 @@ 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 and occupancy before the ubatch, so undoing a speculative placement does not force a rebuild from every cell + std::vector exact_page_owner_old; + std::vector exact_page_live_old; }; // remember the old state of the cells so we can restore it in the end std::vector states; + // [TAG_EXACT_CONCURRENCY] a placement can purge positions outside the cells it restores below, + // and those are not undone; count removals to notice + const uint64_t n_release_before = exact_page_n_release; + bool success = true; for (const auto & ubatch : ubatches) { @@ -777,7 +1046,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vectorsinfo; @@ -805,6 +1078,16 @@ 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_rebuild) { + exact_page_owner = it->exact_page_owner_old; + exact_page_live = it->exact_page_live_old; + } + } + + if (exact_rebuild) { + exact_page_owner_dirty = true; } if (!success) { @@ -963,6 +1246,48 @@ 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; + + exact_page_owner_tmp = exact_page_owner; + + auto & owner = exact_page_owner_tmp; + + 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; @@ -1125,6 +1450,10 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & seq_pos_max_rm[seq_id] = std::max(seq_pos_max_rm[seq_id], pos); + if (exact_pages) { + exact_pages_release(idx); + } + cells.rm(idx); } @@ -1154,6 +1483,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]); + } } } @@ -1185,7 +1520,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; @@ -1247,7 +1591,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 @@ -2208,6 +2595,12 @@ void llama_kv_cache::state_read_sinfo( llama_state_seq_flags flags, slot_info_vec_t * sinfos_out, const slot_info_vec_t * sinfos_in) { + // [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; @@ -2527,6 +2920,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 8c62300f8356..f120a0c2ed17 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -133,6 +133,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; @@ -184,6 +187,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; // active cell count when position p lives in physical cell p; 0 for any non-contiguous layout uint32_t get_n_kv_pos_contiguous(const slot_info & sinfo, const llama_ubatch & ubatch) const; @@ -266,6 +271,38 @@ 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; + + // live cells in each physical page, so a removal can free the page it emptied without rescanning the pool + mutable std::vector exact_page_live; + + // with LLAMA_KV_CACHE_DEBUG set this also rebuilds, to check what was maintained + void exact_pages_sync() const; + + void exact_pages_rebuild() const; + + void exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos); + + // record that the cell at physical index idx has just become empty + void exact_pages_release(uint32_t idx); + + // how many cells have been released, so prepare() can tell whether a placement removed cells + // it is not going to restore + uint64_t exact_page_n_release = 0; + + // scratch for find_slot(), which must not touch the ownership it reads + mutable std::vector exact_page_owner_tmp; + bool v_trans = true; // the value tensor is transposed const uint32_t n_seq_max = 1; @@ -398,6 +435,8 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; uint32_t get_n_kv_pos_contiguous() 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 42c7381a9e6f..a596f35bd8d5 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 484eafb74991..70ba19ca3239 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 57919accf095..88b14faa8e4b 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -442,7 +442,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 db825396645e..61cd348f2dd9 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 bbbdcfddff2b..4508c0c002d4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -160,6 +160,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-grammar-integration.cpp) llama_build_and_test(test-llama-grammar.cpp) llama_build_and_test(test-batch-alloc.cpp) + # [TAG_EXACT_CONCURRENCY] the buffer types the mode treats as invariant, through llama-impl.h + llama_build_and_test(test-exact-buft.cpp) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) @@ -333,6 +335,15 @@ 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) + +# [TAG_EXACT_CONCURRENCY] page bookkeeping of the paged KV pool; skips itself where the mode cannot run +llama_build_and_test(test-exact-pages.cpp LABEL "model" ARGS -m "${MODEL_DEST}") +set_tests_properties(test-exact-pages PROPERTIES FIXTURES_REQUIRED test-download-model) + if (APPLE) llama_build(test-rset-release.cpp) endif() @@ -355,6 +366,16 @@ unset(LLAMA_TEST_NAME) llama_build_and_test(test-mtmd-impl.cpp) target_link_libraries(test-mtmd-impl PRIVATE mtmd) +# [TAG_EXACT_CONCURRENCY] the batch shape the mode requires, checked without a model +llama_build_and_test(test-exact-geometry.cpp) + +# server helpers that need no model +if (LLAMA_BUILD_TOOLS) + llama_build_and_test(test-server-tokens.cpp) + target_link_libraries(test-server-tokens PRIVATE server-context mtmd) + target_include_directories(test-server-tokens PRIVATE ${PROJECT_SOURCE_DIR}/tools/server ${PROJECT_SOURCE_DIR}/tools/mtmd) +endif() + # GGUF model data fetcher library for tests that need real model metadata # Only compile when cpp-httplib has SSL support (CPPHTTPLIB_OPENSSL_SUPPORT) if (TARGET cpp-httplib) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f4e85f2b47ea..f06748c79b06 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7715,6 +7715,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_FLASH_ATTN_EXT_BANDED struct test_flash_attn_ext_banded : public test_case { const int64_t d; @@ -9866,6 +9900,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, m, 2, 1024, { 1, 1 }, { 1, 1 })); } + 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})); @@ -10682,6 +10731,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-exact-buft.cpp b/tests/test-exact-buft.cpp new file mode 100644 index 000000000000..c7f93c857d7a --- /dev/null +++ b/tests/test-exact-buft.cpp @@ -0,0 +1,60 @@ +// [TAG_EXACT_CONCURRENCY] which buffer types the mode accepts a weight in. A host buffer is not one +// of them: the scheduler runs an operation on the backend holding its weight, and moves a host +// weight's operation to the GPU only once the batch is wide enough, so its result would depend on +// how many sequences share the step. This is the predicate behind both the context's weight check +// and the refusal of a lora that would inherit such a buffer. + +#include "ggml-backend.h" + +#include "../src/llama-impl.h" + +#include + +#undef NDEBUG +#include + +int main() { + ggml_backend_load_all(); + + // nothing placed anywhere is nothing to trust + assert(!llama_exact_buft_invariant(nullptr)); + + // the plain CPU buffer, and the pinned host buffer a GPU backend offers, are both host memory + assert(!llama_exact_buft_invariant(ggml_backend_cpu_buffer_type())); + + bool checked_gpu = false; + + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + + const bool invariant = reg && llama_exact_backend_name(ggml_backend_reg_name(reg)); + + // a device's own buffer follows its backend, and its host buffer never does + assert(llama_exact_buft_invariant(ggml_backend_dev_buffer_type(dev)) == invariant); + + if (auto * host = ggml_backend_dev_host_buffer_type(dev)) { + assert(!llama_exact_buft_invariant(host)); + } + + checked_gpu = checked_gpu || invariant; + } + + // the registry names the mode trusts, whatever this build has + assert(llama_exact_backend_name("CUDA")); + assert(llama_exact_backend_name("ROCm")); + assert(llama_exact_backend_name("MUSA")); + assert(!llama_exact_backend_name("CPU")); + assert(!llama_exact_backend_name("BLAS")); + assert(!llama_exact_backend_name(nullptr)); + + printf("%s: all tests passed%s\n", __func__, + checked_gpu ? "" : " (no batch-invariant device here, the positive case was not exercised)"); + + return 0; +} diff --git a/tests/test-exact-geometry.cpp b/tests/test-exact-geometry.cpp new file mode 100644 index 000000000000..7e6e80e07f6b --- /dev/null +++ b/tests/test-exact-geometry.cpp @@ -0,0 +1,55 @@ +// [TAG_EXACT_CONCURRENCY] the batch shape a prefill needs to be split into the ubatches it would +// get alone: the server adds a prompt in whole ubatches, so the batch has to hold one of those +// beside a decode step of every slot, or the prompt is left the shorter remainder. + +#include "common.h" + +#include + +#undef NDEBUG +#include + +int main() { + int n_min = 0; + + // the reported minimum is the ubatch plus the decode step, whether or not the batch reaches it + assert(common_exact_batch_geometry(2048, 512, 4, &n_min)); + assert(n_min == 516); + + // the case that used to warn and carry on: one decoder beside the prompt leaves it 511 tokens + assert(!common_exact_batch_geometry(512, 512, 1, &n_min)); + assert(n_min == 513); + + assert(!common_exact_batch_geometry(512, 512, 2, &n_min)); + assert(n_min == 514); + + // exactly enough, and one short of it + assert(common_exact_batch_geometry(514, 512, 2, &n_min) && n_min == 514); + assert(!common_exact_batch_geometry(513, 512, 2, &n_min) && n_min == 514); + + // a single slot with no draft still needs room for its own decoded token + assert(!common_exact_batch_geometry(512, 512, 1)); + assert(common_exact_batch_geometry(1024, 512, 1)); + + // an unset ubatch is the whole batch, which then cannot hold a decode step as well + assert(!common_exact_batch_geometry(2048, 0, 4, &n_min)); + assert(n_min == 2052); + + // a ubatch larger than the batch is clamped to it, so it cannot pass either + assert(!common_exact_batch_geometry(512, 4096, 1, &n_min)); + assert(n_min == 513); + + // the shape a context settles on when its size clamps the batch: n_batch becomes min(n_ctx, -b) + // and n_ubatch min(n_batch, -ub), so a context of 256 cells leaves the two equal and no column + // for a decode step, whatever -b and -ub asked for + assert(!common_exact_batch_geometry(256, 256, 2, &n_min)); + assert(n_min == 258); + + // no slot decoding at all: the prompt has the batch to itself + assert(common_exact_batch_geometry(512, 512, 0, &n_min)); + assert(n_min == 512); + + printf("%s: all tests passed\n", __func__); + + return 0; +} diff --git a/tests/test-exact-pages.cpp b/tests/test-exact-pages.cpp new file mode 100644 index 000000000000..e9cdef82ffe7 --- /dev/null +++ b/tests/test-exact-pages.cpp @@ -0,0 +1,166 @@ +// [TAG_EXACT_CONCURRENCY] drives the removal paths of the paged KV pool - a removal that empties nothing, one that leaves holes, and the pages those holes keep reserved - while LLAMA_KV_CACHE_DEBUG=1 makes the pool cross-check page ownership against the live cells every ubatch. +// Needs a CUDA build with 256-wide heads and a fully offloaded F16 cache; without one the test reports what it skipped and passes. + +#include "arg.h" +#include "common.h" +#include "llama.h" + +#include +#include +#include + +static const uint32_t PAGE = 256; + +static bool decode_range(llama_context * ctx, llama_seq_id seq, llama_pos first, llama_pos last) { + llama_batch batch = llama_batch_init(64, 0, 1); + + bool ok = true; + + for (llama_pos p = first; p <= last && ok; ) { + common_batch_clear(batch); + + for (int i = 0; i < 64 && p <= last; ++i, ++p) { + common_batch_add(batch, 1, p, {seq}, false); + } + + // every decode asks for one set of logits, so none of them is a batch with no output + batch.logits[batch.n_tokens - 1] = true; + + ok = llama_decode(ctx, batch) == 0; + } + + llama_batch_free(batch); + + return ok; +} + +// Windows has no setenv +static void set_env_default(const char * name, const char * value) { + if (getenv(name)) { + return; + } +#ifdef _WIN32 + _putenv_s(name, value); +#else + setenv(name, value, 0); +#endif +} + +int main(int argc, char ** argv) { + // read before the model is loaded: both are latched on first use + set_env_default("LLAMA_EXACT_CONCURRENCY", "1"); + set_env_default("LLAMA_KV_CACHE_DEBUG", "1"); + + common_params params; + + params.sampling.seed = 1234; + params.kv_unified = true; + params.n_parallel = 2; + params.n_ctx = 2*4*PAGE; + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + + common_init(); + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + // after the parser, which requires the default here + params.n_gpu_layers = 999; + + 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) { + printf("%s : skipped, this build and model cannot run exact concurrency\n", __func__); + return 0; + } + + llama_memory_t mem = llama_get_memory(ctx); + + const uint32_t gran = llama_memory_alloc_granularity(mem); + if (gran != PAGE) { + fprintf(stderr, "%s : allocation granularity is %u, expected %u\n", __func__, gran, PAGE); + return 1; + } + + // positions 0..599 of sequence 0: three pages, the last one part full + if (!decode_range(ctx, 0, 0, 599)) { + fprintf(stderr, "%s : failed to fill sequence 0\n", __func__); + return 1; + } + + // a page belongs to one sequence, so a cross-sequence copy is refused whole rather than half + // applied: the pool logs the refusal and leaves the destination empty and the source as it was + llama_memory_seq_cp(mem, 0, 1, -1, -1); + + if (llama_memory_seq_pos_max(mem, 1) != -1 || llama_memory_seq_pos_max(mem, 0) != 599) { + fprintf(stderr, "%s : a refused copy left sequence 1 at %d and sequence 0 at %d\n", __func__, + llama_memory_seq_pos_max(mem, 1), llama_memory_seq_pos_max(mem, 0)); + return 1; + } + + // the removal every accepted speculative step makes: a rejected tail that is not there. It + // must leave the pool alone, ownership included + if (!llama_memory_seq_rm(mem, 0, 600, -1) || llama_memory_seq_pos_max(mem, 0) != 599) { + fprintf(stderr, "%s : a removal past the tail changed the sequence, its end is %d\n", + __func__, llama_memory_seq_pos_max(mem, 0)); + return 1; + } + + if (!decode_range(ctx, 0, 600, 655)) { + fprintf(stderr, "%s : failed to continue sequence 0 after a removal that removed nothing\n", __func__); + return 1; + } + + // holes: positions 1 to 510 go, 0 and 511 to 655 stay, so the first two pages each keep a live + // cell and neither is free for another sequence. A hybrid memory refuses to remove the middle + // of a sequence, and then there is nothing to check here + const bool holes = llama_memory_seq_rm(mem, 0, 1, 511); + + if (holes && llama_memory_seq_pos_max(mem, 0) != 655) { + fprintf(stderr, "%s : a partial removal changed the end of the sequence: %d\n", __func__, + llama_memory_seq_pos_max(mem, 0)); + return 1; + } + + printf("%s : interior removal %s\n", __func__, holes ? "left holes" : "was refused, skipping the hole case"); + + // sequence 1 fills what is left of the pool. The pool holds 8 pages and sequence 0 holds 3 of + // them, holes and a part full tail page included, so 5 remain + if (!decode_range(ctx, 1, 0, 5*PAGE - 1)) { + fprintf(stderr, "%s : failed to fill the pages sequence 0 does not hold\n", __func__); + return 1; + } + + // one page more than the pool has left: it has to refuse rather than take a page that still + // has a live cell in it + if (decode_range(ctx, 1, 5*PAGE, 5*PAGE)) { + fprintf(stderr, "%s : the pool allocated a page that sequence 0 still holds\n", __func__); + return 1; + } + + // a whole sequence goes back to the pool as whole pages, holes included + if (!llama_memory_seq_rm(mem, 0, -1, -1) || llama_memory_seq_pos_max(mem, 0) != -1) { + fprintf(stderr, "%s : sequence 0 is still in the pool after a full removal\n", __func__); + return 1; + } + + if (!decode_range(ctx, 1, 5*PAGE, 8*PAGE - 1)) { + fprintf(stderr, "%s : the three pages of the removed sequence were not reusable\n", __func__); + return 1; + } + + // the pool is full again + if (decode_range(ctx, 1, 8*PAGE, 8*PAGE)) { + fprintf(stderr, "%s : the pool allocated a ninth page\n", __func__); + return 1; + } + + printf("%s : ok, page ownership survived a no-op removal, holes and a full removal\n", __func__); + + return 0; +} diff --git a/tests/test-server-tokens.cpp b/tests/test-server-tokens.cpp new file mode 100644 index 000000000000..051e01a15f0d --- /dev/null +++ b/tests/test-server-tokens.cpp @@ -0,0 +1,184 @@ +// [TAG_PREEMPT] the server converts between a KV position and a token count when it rewinds a slot to +// what the cache holds. With M-RoPE media the two differ, so the conversion is exercised here on a +// hand-built image chunk, without a model. + +#include "server-common.h" + +#include "mtmd.h" + +#include +#include +#include +#include +#include +#include + +#undef NDEBUG +#include + +// the wire format of mtmd_input_chunk_save(), which needs a context to produce a chunk; written here so +// that a chunk of a known shape can be loaded without one +struct chunk_writer { + std::vector buf; + + template void put(T v) { + const char * p = reinterpret_cast(&v); + buf.insert(buf.end(), p, p + sizeof(T)); + } + + void put_str(const std::string & s) { + put(s.size()); + buf.insert(buf.end(), s.begin(), s.end()); + } +}; + +// nx*ny tokens of one image, max(nx, ny) positions under M-RoPE +static mtmd::input_chunk_ptr make_image_chunk(uint32_t nx, uint32_t ny) { + chunk_writer w; + + w.put(2); // MTMD_SERIALIZATION_VERSION + w.put(MTMD_INPUT_CHUNK_TYPE_IMAGE); + w.put(0); // tokens_text + w.put(1); // tokens_image follows + w.put(nx); + w.put(ny); + w.put(1); // MTMD_POS_TYPE_MROPE + w.put(0); // image_idx + w.put(1); // n_temporal_merge + w.put_str("test-image"); // id + w.put(0); // batch_f32.is_audio + w.put(1); // one entry + w.put(0); // entry.add_viewsep + w.put(0); // entry.add_newline + w.put(0); // entry.lead_pad + w.put(1); // entry.nx + w.put(1); // entry.ny + w.put(0); // no tokens_audio + + mtmd::input_chunk_ptr chunk(mtmd_input_chunk_load(w.buf.data(), w.buf.size())); + + assert(chunk && "the serialized image chunk was rejected"); + + return chunk; +} + +// 10 text tokens, an image of 256 tokens and 16 positions, 20 text tokens, then n_gen generated tokens +static server_tokens make_prompt(size_t n_gen, const mtmd_input_chunk * chunk) { + server_tokens res; + + res.has_mtmd = true; + + for (size_t i = 0; i < 10; ++i) { + res.push_back((llama_token) (100 + i)); + } + + res.push_back(chunk); + + for (size_t i = 0; i < 20; ++i) { + res.push_back((llama_token) (200 + i)); + } + + for (size_t i = 0; i < n_gen; ++i) { + res.push_back((llama_token) (300 + i)); + } + + return res; +} + +int main() { + const mtmd::input_chunk_ptr chunk = make_image_chunk(16, 16); + + assert(mtmd_input_chunk_get_n_tokens(chunk.get()) == 256); + assert(mtmd_input_chunk_get_n_pos (chunk.get()) == 16); + + // a cut in the generated tail: the cache reports 86 positions, which is 326 tokens + { + server_tokens prompt = make_prompt(40, chunk.get()); + + assert(prompt.size() == 326); + assert(prompt.pos_next() == 86); + + const llama_pos pos_cached = 86; + const size_t n_cached = prompt.size_up_to_pos(pos_cached); + + assert(n_cached == 326); + assert(prompt.pos_next(n_cached) == pos_cached); + + // the same number taken for a token count falls inside the image + bool threw = false; + + try { + prompt.keep_first((size_t) pos_cached); + } catch (const std::exception &) { + threw = true; + } + + assert(threw && "a position used as a token count cuts the image in half"); + } + + // the same prompt with a longer tail, cut inside the generated tokens + { + server_tokens prompt = make_prompt(300, chunk.get()); + + assert(prompt.size() == 586); + assert(prompt.pos_next() == 346); + + const llama_pos pos_cached = 106; // 10 text + 16 image + 20 text + 60 generated + const size_t n_cached = prompt.size_up_to_pos(pos_cached); + + assert(n_cached == 346); + assert(prompt.pos_next(n_cached) == pos_cached); + + prompt.keep_first(n_cached); + + assert(prompt.size() == 346); + assert(prompt.pos_next() == pos_cached); + } + + // a cut before the image, and one at its first token: both are token boundaries + { + server_tokens prompt = make_prompt(0, chunk.get()); + + assert(prompt.size_up_to_pos(10) == 10); + assert(prompt.pos_next(10) == 10); + + // the image ends at position 26 and token 266 + assert(prompt.size_up_to_pos(26) == 266); + assert(prompt.pos_next(266) == 26); + } + + // a cut inside the image: the conversion cannot land there, and stepping back reaches the chunk's first token + { + server_tokens prompt = make_prompt(0, chunk.get()); + + const llama_pos pos_cached = 20; // inside the image, which spans positions 10..25 + + size_t n_cached = prompt.size_up_to_pos(pos_cached); + + assert(n_cached == 266); // rounded up to the whole chunk + + while (n_cached > 0 && prompt.pos_next(n_cached) > pos_cached) { + n_cached--; + } + + assert(n_cached == 10); + assert(prompt.pos_next(n_cached) == 10); + + prompt.keep_first(n_cached); // would throw if it cut the image in half + assert(prompt.size() == 10); + } + + // an empty cache has to be handled by the caller: the walk always consumes its first token + { + server_tokens prompt = make_prompt(4, chunk.get()); + + const llama_pos pos_cached = 0; + + assert(prompt.size_up_to_pos(pos_cached) == 1); + assert((pos_cached > 0 ? prompt.size_up_to_pos(pos_cached) : 0) == 0); + } + + printf("%s: all tests passed\n", __func__); + + return 0; +} diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index d5548afba179..5a1502f747fd 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,6 +73,13 @@ 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, 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)); + GGML_ASSERT(llama_state_seq_get_data(ctx, before[s].data(), before[s].size(), s) == before[s].size()); + } + // clear seq 1 to create a "hole" in the KV cache (fragmentation) // 0.20.20.20.2.... llama_memory_t mem = llama_get_memory(ctx); @@ -96,6 +103,13 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : restored state into seq 1, %zu bytes\n", __func__, nset); + for (int s = 0; s < params.n_parallel; ++s) { + std::vector after(llama_state_seq_get_size(ctx, s)); + GGML_ASSERT(llama_state_seq_get_data(ctx, after.data(), after.size(), s) == after.size()); + GGML_ASSERT(before[s] == after); + } + fprintf(stderr, "%s : all %d sequence snapshots are byte-identical after restore\n", __func__, params.n_parallel); + // Verify we can decode with the restored state // Generate one token to verify the restored state is usable auto sparams = llama_sampler_chain_default_params(); diff --git a/tests/test-state-seq-copy.cpp b/tests/test-state-seq-copy.cpp new file mode 100644 index 000000000000..7a322fd20543 --- /dev/null +++ b/tests/test-state-seq-copy.cpp @@ -0,0 +1,147 @@ +// [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); + + // a recurrent state does not stay in one row, so these models are refused a transfer whatever the backend can do + if (llama_model_is_recurrent(llama_init->model()) || llama_model_is_hybrid(llama_init->model())) { + CHECK(cpy == nullptr); + fprintf(stderr, "%s : a recurrent or hybrid model is refused a transfer, as it must be\n", __func__); + return 0; + } + + 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 71ebb95434e4..383aafbfd739 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -168,6 +168,8 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-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: 0 - disabled, -1 - no limit)
(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: the copy out overlaps with the slots that keep decoding, while a copy back in, and a kv-full retry behind a copy out that has not landed, 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) | @@ -675,6 +677,18 @@ These words will not be included in the completion, so make sure to add them to - `tokens_cached`: Number of tokens from the prompt which could be re-used from previous completion - `tokens_evaluated`: Number of tokens evaluated in total from the prompt - `truncated`: Boolean indicating if the context size was exceeded during generation, i.e. the number of tokens provided in the prompt (`tokens_evaluated`) plus tokens generated (`tokens predicted`) exceeded the context size (`n_ctx`) +- `preempt`: How the request was served while the unified KV cache was full (see `--preempt-ram`). `parks` is how often the request was parked to make room for another, and `recomputes` is how many of those parks dropped the sequence's cells because `--preempt-ram` was spent, so that the resume re-prefilled its tokens instead of restoring the bytes that were saved. A re-prefilled sequence continues from the same tokens, but its numerics are not guaranteed identical to the sequence that left, `LLAMA_EXACT_CONCURRENCY` included: raise `--preempt-ram` until `recomputes` stays 0 where that matters. Both fields are present in the final response of a streamed completion as well, and on the OpenAI-compatible endpoints: the final chunk of a streamed `/v1/chat/completions`, the `message_delta` event of `/v1/messages`, and the response object of `/v1/responses` streamed or not, which in a stream is the `response` of the `response.completed` event. + +While a request is streaming, the server sends SSE comment lines that a client reading raw lines can act on and every SSE event consumer ignores: + +- `: preempted` - the slot was parked and the stream is silent until it comes back. A parked stream is kept alive with the same comment about every two seconds. +- `: resumed` - the slot is running again. +- `: recomputed` - sent right after `: resumed` when that resume re-prefilled the sequence rather than restoring its saved bytes, i.e. what follows is the continuation `preempt.recomputes` counts. +- `: preempt-keepalive` - sent while the slot stays parked, at most every 2 seconds, or at the request's `sse_ping_interval` when that is shorter. + +A park can happen while the prompt is still being processed, before the request has produced a token. The notice is not held back for the first chunk in that case: the response headers and the `: preempted` line go out at the moment the slot is parked, on every streaming surface (`/completion`, `/v1/chat/completions`, `/v1/responses`, `/v1/messages`), so a client never has to tell that silence from a stall. `: resumed`, and `: recomputed` where it applies, follow when the slot runs again. + +When the request asks for more than one completion, either several prompts or `n` above one, the index of the completion follows the word, for example `: resumed 1`, including index `0`. A request with a single completion carries no index. ### POST `/tokenize`: Tokenize a given text @@ -980,6 +994,8 @@ This endpoint is enabled by default and can be disabled with `--no-slots`. It ca If query param `?fail_on_no_slot=1` is set, this endpoint will respond with status code 503 if there is no available slots. +Every entry also reports how its request has been served under preemption (see `--preempt-ram`): `is_preempted` and `is_transferring` say whether the slot's cells have been released or a copy is in flight, `n_preempt` counts the parks of the current task, and `n_recompute` counts how many of those dropped the cells, so that the resume re-prefilled the sequence instead of restoring its saved bytes. + **Response format**
@@ -1147,6 +1163,11 @@ In *router mode* the query param `?model={model_id}` has to be set. This endpoin | `llamacpp:spec_decode_num_accepted_tokens_total` | Counter | Total draft tokens accepted by the target model (0 when spec-decode is off). | | `llamacpp:spec_decode_num_drafts_total` | Counter | Total speculative decoding verification steps (0 when spec-decode is off). | | `llamacpp:spec_decode_num_accepted_tokens_per_pos_total` | Counter | Accepted tokens per draft position (labeled `position="N"`; absent when spec-decode is off or before the first completed speculative request). | +| `llamacpp:n_preempt_total` | Counter | Slots parked to make room in the unified KV cache (0 unless `--kv-unified` with more than one slot). | +| `llamacpp:n_resume_total` | Counter | Parked slots put back. | +| `llamacpp:preempt_recompute_total` | Counter | Parks that dropped their cells because `--preempt-ram` was spent, so the resume re-prefills the sequence instead of restoring its saved bytes. | +| `llamacpp:requests_preempted` | Gauge | Requests currently parked, waiting for room in the unified KV cache. | +| `llamacpp:preempt_ram_bytes` | Gauge | Host RAM held by parked sequences. | ### POST `/slots/{id_slot}?action=save`: Save the prompt cache of the specified slot to a file. diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 2ac98b6fddbc..76e348dab1cc 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -470,6 +470,12 @@ const mtmd::input_chunk_ptr & server_tokens::find_chunk(size_t idx) const { throw std::runtime_error("Chunk not found"); } +size_t server_tokens::chunk_n_tokens_at(size_t idx) const { + auto it = map_idx_to_media.find(idx); + + return it == map_idx_to_media.end() ? 0 : mtmd_input_chunk_get_n_tokens(it->second.get()); +} + std::pair server_tokens::find_next_media_chunk(size_t idx) const { auto it = map_idx_to_media.upper_bound(idx); if (it != map_idx_to_media.end()) { diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6c681a2cf56d..2c3e4a26352b 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -186,6 +186,9 @@ struct server_tokens { const mtmd::input_chunk_ptr & find_chunk(size_t idx) const; + // tokens of the media chunk that starts at idx, 0 if none starts there + size_t chunk_n_tokens_at(size_t idx) const; + // find next media chunk after idx // returns a pair of pointer to the chunk (nullptr if not found) and its start index in tokens std::pair find_next_media_chunk(size_t idx) const; @@ -474,6 +477,12 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; + uint64_t n_preempt = 0; + uint64_t n_resume = 0; + + // [TAG_PREEMPT] parks that dropped their cells: those resumes re-prefill, and a re-prefill is not bit-for-bit the state that left + uint64_t n_preempt_recompute = 0; + uint64_t n_draft_tokens = 0; // Total draft tokens generated uint64_t n_draft_accepted = 0; // Draft tokens actually accepted uint64_t n_draft_verif_steps = 0; // Total draft token verification steps by the target model diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index fe068d3e9104..9d13145138b5 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -18,7 +18,9 @@ #include "mtmd-helper.h" #include +#include #include +#include #include #include #include @@ -104,8 +106,79 @@ enum slot_state { SLOT_STATE_PROCESSING_PROMPT, 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: 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 + +static std::string preempt_notice_comment(const server_task_result_preempt_notice & notice) { + const std::string suffix = (notice.batched ? " " + std::to_string(notice.index) : "") + "\n\n"; + + std::string res = (notice.parked ? ": preempted" : ": resumed") + suffix; + + // [TAG_PREEMPT] the resume rebuilt the sequence from its tokens, so what follows is not the continuation the saved bytes would have given + if (!notice.parked && notice.recomputed) { + res += ": recomputed" + suffix; + } + + return res; +} +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 +constexpr int64_t PREEMPT_ROTATE_RECOMPUTE_US = 30ll * 1000 * 1000; // ... and after this long it gives way even where that costs the resident a re-prefill + +// [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; + +// [TAG_PREEMPT] test knob: LLAMA_SERVER_PREEMPT_FAIL_SAVE=N fails the host allocation of the Nth park, which no budget check can rule out, so that the fall back to a recompute park is exercised +static bool preempt_fail_save() { + static int32_t n_left = []() { + const char * val = getenv("LLAMA_SERVER_PREEMPT_FAIL_SAVE"); + + return val ? std::max(0, atoi(val)) : 0; + }(); + + return n_left > 0 && --n_left == 0; +} + +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 { @@ -339,6 +412,384 @@ struct server_slot { prompt.clear(); } + 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_recompute = 0; // ... of which parked by dropping the cells, so the resume re-prefilled + 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 + + // [TAG_PREEMPT] a park with no room left under --preempt-ram: the cells are dropped instead of copied out, and the resume re-prefills the tokens + bool preempt_recompute = false; // parked by dropping its cells + bool preempt_reprefill = false; // putting back, as a prompt, what such a park dropped + server_tokens preempt_tokens; // what the re-prefill decodes: the prompt and everything generated so far + + size_t preempt_state_size() const { + // 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(); + } + + 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); + } + + // 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(); + } + + // the tokens the prompt step works through: its own list while re-prefilling, the request's otherwise + const server_tokens & preempt_input() const { + return preempt_tokens.empty() ? task->tokens : preempt_tokens; + } + + int32_t preempt_n_input() const { + return preempt_tokens.empty() ? (task ? task->n_tokens() : 0) : (int32_t) preempt_tokens.size(); + } + + // [TAG_PREEMPT] park with the host budget spent: drop the cells, keep the tokens, re-prefill them on resume. The sampler and the counters are untouched, so the stream carries on from the same token; the resume is bit-exact only as far as prefill numerics match decode numerics. A media chunk comes back the way it went in: the prompt step reads the chunk's data off the task and reserves its cells whole, so the placeholder the re-prefill list carries is all it needs + bool preempt_save_recompute() { + preempt_state_free(); + preempt_detach(); + + if (state == SLOT_STATE_GENERATING) { + preempt_tokens = std::move(prompt.tokens); + prompt.tokens = server_tokens(); + + prompt.tokens.has_mtmd = preempt_tokens.has_mtmd; // the re-prefill pushes the chunk's placeholder back into it + } + + prompt_clear(); + + state_before_preempt = state; + state = SLOT_STATE_PREEMPTED; + t_preempt_us = ggml_time_us(); + preempt_recompute = true; + preempt_rotation_refused = false; + + n_preempt++; + n_recompute++; + + return true; + } + + bool preempt_restore_recompute() { + preempt_recompute = false; + n_preempt_fail = 0; + + if (preempt_tokens.empty()) { + state = state_before_preempt; // its prompt had not been processed yet, so it is processed again from the start + + // the park dropped the cells the prompt step had already filled, and the restart does not pass through SLOT_STATE_STARTED, where these two are set: left alone they would count the dropped prefix a second time + stats.n_prompt_cached = 0; + stats.n_prompt_processed = 0; + + return true; + } + + preempt_reprefill = true; + state = SLOT_STATE_PROCESSING_PROMPT; + + return true; + } + + // every token is back in the cache: the slot goes on generating from the token it had already sampled + void preempt_reprefill_done() { + preempt_reprefill = false; + preempt_tokens.clear(); + + i_batch = -1; + state = SLOT_STATE_GENERATING; + + if (can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + } + + // [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() { + if (preempt_fail_save()) { + SLT_ERR(*this, "%s", "failed to allocate the host memory for the preemption state (test knob)\n"); + preempt_state_free(); + return false; + } + + 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); + } catch (const std::bad_alloc & e) { + SLT_ERR(*this, "failed to allocate %.3f MiB for the preemption state: %s\n", + (size_tgt + size_dft) / (1024.0 * 1024.0), e.what()); + preempt_state_free(); + return false; + } + + if (llama_state_seq_get_data_ext(ctx_tgt, preempt_state_tgt.data(), size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_tgt) { + SLT_ERR(*this, "%s", "failed to copy the target sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + if (size_dft > 0 && + llama_state_seq_get_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + SLT_ERR(*this, "%s", "failed to copy the draft sequence out of the KV cache\n"); + preempt_state_free(); + return false; + } + + preempt_detach(); + + // note: prompt.tokens is deliberately kept - the resume sizes its request from it + mem.seq_rm(id, -1, -1); + + state_before_preempt = state; + state = SLOT_STATE_PREEMPTED; + t_preempt_us = ggml_time_us(); + preempt_rotation_refused = false; + + n_preempt++; + + return true; + } + + // [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() { + if (preempt_recompute) { + return preempt_restore_recompute(); + } + + 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; + } + + state = SLOT_STATE_RESTORING; + + return true; + } + + 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; + } + + preempt_state_free(); + + return preempt_resumed(); + } + + // [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() { + // the memory counts positions, and with M-RoPE media a position is not a token, so convert before truncating + const llama_pos pos_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; + + size_t n_cached = pos_cached > 0 ? prompt.tokens.size_up_to_pos(pos_cached) : 0; + + // a cut inside a media chunk is not a token boundary: keep what precedes the chunk, and the cells of its head go too + bool split_chunk = false; + + while (n_cached > 0 && prompt.tokens.pos_next(n_cached) > pos_cached) { + n_cached--; + split_chunk = true; + } + + if (n_cached < (size_t) prompt.n_tokens()) { + prompt.tokens.keep_first(n_cached); + } + + if (split_chunk) { + mem.seq_rm(id, prompt.tokens.pos_next(), -1); + } + + // what is kept ends where the cache does, or the next decode is positioned from the wrong count + GGML_ASSERT(prompt.tokens.pos_next() <= pos_cached); + + // 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() < preempt_n_input()) { + state = SLOT_STATE_PROCESSING_PROMPT; + } + + preempt_detach(); + } + std::vector lora; int32_t alora_invocation_start = -1; @@ -397,6 +848,17 @@ struct server_slot { n_predict_max = -1; + preempt_state_free(); + preempt_tokens.clear(); + preempt_recompute = false; + preempt_reprefill = false; + state_before_preempt = SLOT_STATE_IDLE; + n_preempt = 0; + n_recompute = 0; + n_preempt_fail = 0; + n_ctx_shift = 0; + t_preempt_us = 0; + llama_set_sampler(ctx_tgt, id, nullptr); // clear alora start @@ -549,6 +1011,13 @@ struct server_slot { t_last_used = ggml_time_us(); + // [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(); + } + state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough @@ -687,10 +1156,15 @@ struct server_slot { json res; res = { - {"id", id}, - {"n_ctx", n_ctx}, - {"speculative", can_speculate()}, - {"is_processing", is_processing()}, + {"id", id}, + {"n_ctx", n_ctx}, + {"speculative", can_speculate()}, + {"is_processing", is_processing()}, + // [TAG_PREEMPT] parked means the cells are gone; a copy out still owns them and a restore has already taken them back, so a scraper counting residency has to keep counting those two + {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"is_transferring", preempt_in_flight()}, + {"n_preempt", n_preempt}, + {"n_recompute", n_recompute}, }; const auto & ptask = task ? task : task_prev; @@ -872,6 +1346,30 @@ struct server_context_impl { metrics.reset_bucket(); } + // [TAG_PREEMPT] the first prompt of a request that cannot be served, with the error response it gets; false when every one of them passes. + // A park notice opens the stream of the member it belongs to, so a member rejected after that could only be told inside a stream that has already answered 200. + bool tasks_prompt_rejected(const std::vector & tasks, json & error) const { + std::string msg; + error_type type = ERROR_TYPE_SERVER; + + for (const auto & task : tasks) { + if (!task_prompt_rejected(task, msg, type)) { + continue; + } + + error = format_error_response(msg, type); + + if (type == ERROR_TYPE_EXCEED_CONTEXT_SIZE) { + error["n_prompt_tokens"] = task.n_tokens(); + error["n_ctx"] = n_ctx_slot(); + } + + return true; + } + + return false; + } + private: // note: accessing these fields outside of this class is not thread-safe // use server_context methods instead @@ -936,6 +1434,17 @@ 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; + preempt_recompute_logged = false; + spec.reset(); spec_init.reset(); @@ -1310,6 +1819,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(); } @@ -1331,6 +1854,111 @@ struct server_context_impl { } } + { + 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"); + } + } else if (params_base.preempt_ram_mib != 0 && params_base.preempt_async && + !llama_model_is_recurrent(model_tgt) && preempt_state_relocates()) { + SRV_WRN("%s", "preemption: a recurrent state does not stay in one row, so a copy running beside the decode could read another sequence; 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 && params_base.preempt_ram_mib != 0) { + 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"); + if (LLAMA_SERVER_PREEMPT_RESUME && strcmp(LLAMA_SERVER_PREEMPT_RESUME, "head") != 0) { + if (strcmp(LLAMA_SERVER_PREEMPT_RESUME, "pass") != 0) { + SRV_ERR("LLAMA_SERVER_PREEMPT_RESUME = %s is not a resume order; use head (the default) or pass\n", + LLAMA_SERVER_PREEMPT_RESUME); + return false; + } + preempt_resume_head = false; + SRV_WRN("%s", "LLAMA_SERVER_PREEMPT_RESUME = pass (parked slots come back most-preempted first, and a smaller slot may pass a head that does not fit)\n"); + } + + 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; 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"; + + if (preempt_test_policy != "smallest") { + SRV_WRN("LLAMA_SERVER_PREEMPT_POLICY = %s (test knob: victim choice for comparison only)\n", preempt_test_policy.c_str()); + } + + if (preempt_test_every > 0) { + SRV_WRN("LLAMA_SERVER_PREEMPT_EVERY = %d (test knob: preempting every slot every %d tokens)\n", + preempt_test_every, preempt_test_every); + } + + const char * LLAMA_SERVER_PREEMPT_PLANNER = getenv("LLAMA_SERVER_PREEMPT_PLANNER"); + preempt_planner_off = LLAMA_SERVER_PREEMPT_PLANNER && strcmp(LLAMA_SERVER_PREEMPT_PLANNER, "off") == 0; + + if (preempt_planner_off) { + 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: 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 && params_base.preempt_ram_mib != 0) { + SRV_WRN("%s", "preemption: off, the recurrent cache holds one state per sequence whatever its length, so there is no cell pool to run out of\n"); + } + } + { const char * LLAMA_SERVER_SLOTS_N_DIFF = getenv("LLAMA_SERVER_SLOTS_N_DIFF"); slots_n_diff = LLAMA_SERVER_SLOTS_N_DIFF ? atoi(LLAMA_SERVER_SLOTS_N_DIFF) : 0; @@ -2043,7 +2671,25 @@ struct server_context_impl { 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) { + void send_preempt_notice(server_slot & slot, bool parked, bool recomputed = false) { + 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->recomputed = recomputed; + res->n_preempt = slot.n_preempt; + res->batched = slot.task->batched; + + 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(); res->id = slot.task->id; @@ -2120,6 +2766,8 @@ struct server_context_impl { res->stopping_word = slot.stopping_word; res->stop = slot.stop; res->post_sampling_probs = slot.task->params.post_sampling_probs; + res->n_preempt = slot.n_preempt; + res->n_recompute = slot.n_recompute; res->verbose = slot.task->params.verbose; res->stream = slot.task->params.stream; @@ -2501,11 +3149,15 @@ struct server_context_impl { case SERVER_TASK_TYPE_METRICS: { int n_processing_slots = 0; + int n_preempted_slots = 0; for (server_slot & slot : slots) { if (slot.is_processing()) { n_processing_slots++; } + if (slot.preempt_is_out()) { + n_preempted_slots++; + } } SRV_DBG("n_processing_slots = %d\n", n_processing_slots); @@ -2513,6 +3165,8 @@ struct server_context_impl { res->id = task.id; res->n_processing_slots = n_processing_slots; res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size(); + res->n_preempted_slots = n_preempted_slots; + res->preempt_ram_bytes = preempt_ram_used(); res->metrics = metrics; if (task.metrics_reset_bucket) { @@ -2724,71 +3378,1068 @@ struct server_context_impl { } break; } - return true; - } + return true; + } + + void iterate(std::vector & slots, std::function callback) { + for (auto & slot : slots) { + try { + callback(slot); + } catch (const std::exception & e) { + SLT_ERR(slot, "got exception: %s\n", e.what()); + send_error(slot, std::string("got exception: ") + e.what(), ERROR_TYPE_SERVER); + slot.release(); + } + } + } + + void iterate(std::vector & slots, std::function callback) { + for (auto & slot : slots) { + try { + callback(*slot); + } catch (const std::exception & e) { + SLT_ERR(*slot, "got exception: %s\n", e.what()); + send_error(*slot, std::string("got exception: ") + e.what(), ERROR_TYPE_SERVER); + slot->release(); + } + } + } + + void abort_all_slots(const std::string & reason) { + for (auto & slot : slots) { + // [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(); + } + } + } + + // @ngxson : for debugging only + int64_t t_pre_decode = 0; + int64_t t_decode = 0; + int64_t t_post_decode = 0; + int64_t t_sampl = 0; + int64_t n_pre_decode = 0; + int64_t n_decode = 0; + int64_t n_post_decode = 0; + int64_t n_sampl = 0; +// #define DEBUG_TIMINGS +#ifdef DEBUG_TIMINGS + struct scoped_timer { + int64_t & t; + int64_t & n; + int64_t t_start; + scoped_timer(int64_t & t_, int64_t & n_) : t(t_), n(n_) { + t_start = ggml_time_us(); + } + ~scoped_timer() { + t += ggml_time_us() - t_start; + n++; + } + }; +#else + struct scoped_timer { + scoped_timer(int64_t &, int64_t &) {} + ~scoped_timer() {} + }; +#endif + + + // 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 + + // [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 or pass, read at load, per context + bool preempt_resume_head = true; + + bool preempt_recurrent = false; + + 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; + } + + bool preempt_ram_kind_logged = false; + + // [TAG_PREEMPT] the fall back to a recompute park is logged once, not per park + bool preempt_recompute_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(); + + if (res == 0 || !slot.task || !slot.can_speculate()) { + return 0; + } + + // a recompute park moved the prompt out of the slot, so the tokens it comes back with bound the draft, not the empty prompt: read as empty, a 2000-token sequence in a 2048-cell pool was charged a whole draft and failed as impossible + const int32_t n_tokens = std::max(slot.prompt.n_tokens(), slot.preempt_n_input()); + + res = std::min(res, slot.n_ctx - n_tokens - 2); + + if (slot.n_remaining() > 0) { + res = std::min(res, slot.n_remaining() - 1); + } + + return std::max(0, res); + } + + size_t preempt_ram_used() const { + size_t res = 0; + + for (const auto & slot : slots) { + res += slot.preempt_state_size(); + } + + return res; + } + + // [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(); + } + } + + 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() + 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(); + } + } + + size_t preempt_n_keep(const server_slot & slot) const { + if (!slot.task->params.cache_prompt) { + return 0; + } + + size_t n_keep = slot.prompt.tokens.get_common_prefix(slot.task->tokens); + + if (slot.alora_invocation_start > 0) { + n_keep = std::min(n_keep, (size_t) (slot.alora_invocation_start - 1)); + } + + return n_keep; + } + + // 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.preempt_recompute) { + return (int32_t) slot.preempt_tokens.size(); // parked by dropping its cells: it comes back needing all of them at once + } + + if (slot.state == SLOT_STATE_STARTED && slot.task) { + return (int32_t) preempt_n_keep(slot); + } + + return slot.prompt.n_tokens(); + } + + // [TAG_PREEMPT] the cells the media chunks pending at n_have take: pre_decode() runs a whole chunk through llama_decode() calls of its own, which no kv-full retry covers, so the planner reserves the lot before it is decoded + int32_t preempt_n_mtmd_pending(const server_slot & slot, int32_t n_have) const { + if (!slot.task || !slot.task->tokens.has_mtmd) { + return 0; + } + + const auto & tokens = slot.task->tokens; + + int32_t res = 0; + + for (int32_t i = n_have; i >= 0 && i < (int32_t) tokens.size(); ) { + const int32_t n = (int32_t) tokens.chunk_n_tokens_at(i); + + if (n <= 0) { + break; + } + + res += n; + i += n; + } + + return res; + } + + int32_t preempt_n_need(const server_slot & slot) const { + int32_t res = preempt_n_retained(slot); + + if (slot.state_before_preempt == SLOT_STATE_GENERATING) { + res += 1 + preempt_n_spec(slot); + } else { + const int32_t n_mtmd = preempt_n_mtmd_pending(slot, res); + const int32_t n_left = slot.preempt_n_input() - res; + + res += n_mtmd > 0 ? n_mtmd : std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); + } + + // [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); + } + + int32_t preempt_kv_used() const { + int32_t res = 0; + + // 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) { + if (slot.state == SLOT_STATE_PREEMPTED) { + continue; // parked: its cells are in host RAM, not in the pool + } + + // [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 += preempt_n_cells(slot.prompt.n_tokens()); + continue; + } + + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + 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 += preempt_n_cells(std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens())); + continue; + } + + charged.push_back(family); + } + + res += preempt_n_cells(slot.prompt.n_tokens()); + } + + return res; + } + + // [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 res_mm = 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) { + 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 += preempt_n_cells_step(n_cur, 1 + preempt_n_spec(slot)); + } break; + case SLOT_STATE_STARTED: + case SLOT_STATE_PROCESSING_PROMPT: + { + const int32_t n_have = preempt_n_retained(slot); + const int32_t n_mtmd = preempt_n_mtmd_pending(slot, n_have); + + // a media chunk is decoded whole, past the batch cap below and past the kv-full retry + if (n_mtmd > 0) { + res_mm += preempt_n_cells_step(n_have, n_mtmd); + break; + } + + const int32_t n_left = slot.preempt_n_input() - n_have; + + res_pmt += preempt_n_cells_step(n_have, std::max(1, std::min(n_batch, n_left))); + n_pmt++; + } break; + default: + break; + } + } + + return res + res_mm + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1)); + } + + // [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) { + const int32_t before = slot.prompt.n_tokens(); + + preempt_normalize_started(slot); + + if (slot.prompt.n_tokens() < before) { + SLT_INF(slot, "trimmed to the %d cells its request keeps ahead of the batch builder, %d released\n", + slot.prompt.n_tokens(), before - slot.prompt.n_tokens()); + res = true; + } + } + + return res; + } + + void preempt_normalize_started(server_slot & slot) { + if (slot.state != SLOT_STATE_STARTED || !slot.task) { + return; + } + + const size_t n_keep = preempt_n_keep(slot); + + if (n_keep >= slot.prompt.tokens.size()) { + return; + } + + // 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); + + if (!partial_ok) { + slot.prompt.tokens.clear(); + slot.mem.seq_rm(slot.id, -1, -1); + return; + } + + slot.prompt.tokens.keep_first(n_keep); + slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); + } + + // `recompute` asks for a fallback: with no victim the host budget can hold, the same policy picks again and the park drops the cells instead of copying them + server_slot * preempt_pick_victim(bool * recompute = nullptr) { + server_slot * leader = nullptr; + int32_t n_running = 0; + + for (auto & slot : slots) { + preempt_normalize_started(slot); + } + + for (auto & slot : slots) { + if (slot.is_processing() && !slot.preempt_is_out()) { + n_running++; + + if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) { + leader = &slot; + } + } + } + + if (n_running < 2) { + // one conversation that does not fit alone is a real overflow, not a scheduling problem + return nullptr; + } + + server_slot * victim = preempt_pick_victim_pass(leader, false); + + if (!victim && recompute) { + victim = preempt_pick_victim_pass(leader, true); + + *recompute = victim != nullptr; + } + + return victim; + } + + server_slot * preempt_pick_victim_pass(const server_slot * leader, bool recompute) { + server_slot * victim = nullptr; + + for (auto & slot : slots) { + // 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) { + continue; + } + + if (&slot == leader) { + continue; + } + + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + 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 (!recompute && !preempt_fits_budget(slot)) { + continue; + } + + const bool starved = slot.n_preempt >= PREEMPT_N_STARVED; + const bool starved_cur = victim && victim->n_preempt >= PREEMPT_N_STARVED; + + if (!victim || + (starved_cur && !starved) || + (starved_cur == starved && preempt_better_victim(slot, *victim))) { + victim = &slot; + } + } + + return victim; + } + + 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(); + } + + if (preempt_test_policy == "youngest") { + return a.task->id > b.task->id; + } + + if (preempt_test_policy == "oldest") { + return a.task->id < b.task->id; + } + + return a.prompt.n_tokens() < b.prompt.n_tokens(); + } + + // [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, bool recompute = false) { + slot.t_preempt_copy_us = t_start; + + // [TAG_PREEMPT] a budget check grants permission to allocate, not a successful allocation: preempt_save() unwinds and waits for whatever it issued, so the same victim can still be parked by dropping its cells + if (!recompute && !slot.preempt_save()) { + SLT_WRN(slot, "%s", "the park could not take the host memory the budget allowed, so it drops its cells instead and the resume re-prefills its tokens\n"); + + recompute = true; + } + + if (recompute) { + if (!slot.preempt_save_recompute()) { + return false; + } + + if (!preempt_recompute_logged) { + preempt_recompute_logged = true; + + // [TAG_EXACT_CONCURRENCY] a state that comes back from host memory is the state that left; one rebuilt by re-prefilling is the same on CPU and differs in the last bits on CUDA, where a prefill of a token and a decode of it take different kernels + if (common_exact_concurrency()) { + SRV_WRN("%s", "exact concurrency: a re-prefilled sequence is not guaranteed byte-identical to one that was never parked; raise --preempt-ram until every parked sequence fits it\n"); + } + + SRV_WRN("preemption: --preempt-ram %d MiB holds no further parked sequence, so a park drops its cells and the resume re-prefills its tokens\n", + params_base.preempt_ram_mib); + } + } + + preempt_log_ram_kind(slot); + + if (slot.state == SLOT_STATE_PREEMPTED) { + metrics.n_preempt++; + } + + if (slot.preempt_recompute) { + metrics.n_preempt_recompute++; + } + + 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)"); + } + } + } + + // [TAG_PREEMPT] run the recorded shift here rather than leave it to the next llama_decode(): a park in between would serialize the positions the shift has already moved together with the K values it has not, and removing the sequence would drop the pending deltas with it + void preempt_apply_shift() { + if (!preempt_shift_pending) { + return; + } + + preempt_wait_for_shift(); + + llama_memory_update(ctx_tgt); + + if (ctx_dft) { + llama_memory_update(ctx_dft); + } + } + + 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 + } + + if (!llama_get_memory(ctx_tgt)) { + 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; + + const bool head_of_line = preempt_resume_head; + + for (;;) { + std::vector parked; + + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_PREEMPTED) { + parked.push_back(&slot); + } + } + + if (parked.empty()) { + break; + } + + std::sort(parked.begin(), parked.end(), [head_of_line](const server_slot * a, const server_slot * b) { + if (!head_of_line && a->n_preempt != b->n_preempt) { + return a->n_preempt > b->n_preempt; + } + + return a->t_preempt_us < b->t_preempt_us; + }); + + if (head_of_line) { + parked.resize(1); + } + + server_slot * best = nullptr; + + 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 != 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 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(1); + + for (auto * slot : parked) { + if (occupied + preempt_n_need(*slot) + margin <= n_cells) { + best = slot; + break; + } + } + + if (best) { + break; + } + + if (preempt_normalize_started_all()) { + continue; + } + + if (!try_clear_idle_slots()) { + break; + } + } + + // 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(); + + // [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(1); + + server_slot * pick = nullptr; + bool pick_enough = false; + bool budget_refused = false; + bool recompute = false; + + auto rotate_pick = [&](bool with_recompute) { + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_GENERATING || slot.n_ctx_shift == 0) { + continue; + } + + if (slot.task && (slot.task->is_parent() || slot.task->is_child())) { + 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 + if (!with_recompute && !preempt_fits_budget(slot)) { + budget_refused = true; + continue; + } + + const bool enough = occupied - preempt_n_cells(slot.prompt.n_tokens()) + need <= n_cells; + + if (!pick || + (enough && !pick_enough) || + (enough == pick_enough && (enough ? slot.prompt.n_tokens() < pick->prompt.n_tokens() + : slot.prompt.n_tokens() > pick->prompt.n_tokens()))) { + pick = &slot; + pick_enough = enough; + } + } + }; + + rotate_pick(false); + + // [TAG_PREEMPT] a resident the budget cannot swap out is rotated by dropping its cells, as ordinary victim selection does: waiting instead has no bound, a resident that keeps shifting need never finish. The head waits longer for this than for a swap: the rotated resident pays a whole re-prefill, and under exact concurrency it stops being the sequence that was parked + if (!pick && budget_refused && ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_RECOMPUTE_US) { + rotate_pick(true); + + recompute = pick != nullptr; + } + + const int64_t t_start = ggml_time_us(); + const int32_t n_rotated = pick ? pick->prompt.n_tokens() : 0; + + if (!pick && budget_refused && !head->preempt_rotation_refused) { + head->preempt_rotation_refused = true; + + SLT_WRN(*head, "no rotation: --preempt-ram %d MiB does not hold this parked state and a resident's at once, and the two are held together while the resident is parked and the head restored; the head waits for a resident to finish, or %.0f s for one to be rotated out by dropping its cells\n", + params_base.preempt_ram_mib, PREEMPT_ROTATE_RECOMPUTE_US / 1e6); + } + + if (pick && preempt_park(*pick, t_start, recompute)) { + server_slot & slot = *pick; + + if (slot.preempt_recompute) { + SLT_WRN(slot, "rotated out after %d context shifts: %d cells dropped, %d tokens to re-prefill on resume, a head parked %.1f s takes its turn%s, preemptions %d\n", + slot.n_ctx_shift, n_rotated, + slot.preempt_n_input(), + (ggml_time_us() - head->t_preempt_us) / 1e6, + pick_enough ? "" : " (not enough room by itself)", + slot.n_preempt); + } else { + 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), + (ggml_time_us() - head->t_preempt_us) / 1e6, + pick_enough ? "" : " (not enough room by itself)", + slot.n_preempt); + } + + // [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; + } + } + } + + if (best) { + continue; + } + + break; + } + + const int64_t t_start = ggml_time_us(); + + const bool recompute = best->preempt_recompute; + + best->t_preempt_copy_us = t_start; + + if (!best->preempt_restore()) { + 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); + } + + if (best->n_preempt_fail >= PREEMPT_N_FAIL_MAX && + ggml_time_us() - best->t_preempt_us > PREEMPT_FAIL_US) { + send_error(*best, "failed to restore the preempted sequence"); + best->release(); + } + + 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, recompute); + + if (recompute) { + SLT_WRN(*best, "resumed after %.2f s: %d tokens to re-prefill, kv %d/%d, preemptions %d\n", + (ggml_time_us() - best->t_preempt_us) / 1e6, + best->preempt_n_input(), + preempt_kv_used(), n_cells, + best->n_preempt); + } else { + 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(), + (ggml_time_us() - t_start) / 1e3, + preempt_kv_used(), n_cells, + best->n_preempt); + } + } + + 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) { + continue; + } - void iterate(std::vector & slots, std::function callback) { - for (auto & slot : slots) { - try { - callback(slot); - } catch (const std::exception & e) { - SLT_ERR(slot, "got exception: %s\n", e.what()); - send_error(slot, std::string("got exception: ") + e.what(), ERROR_TYPE_SERVER); - slot.release(); + // the budget refusing is the recompute park's case, so the knob reaches it too + const bool recompute = !preempt_fits_budget(slot); + + if (preempt_park(slot, ggml_time_us(), recompute)) { + 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)); + } } } - } - void iterate(std::vector & slots, std::function callback) { - for (auto & slot : slots) { - try { - callback(*slot); - } catch (const std::exception & e) { - SLT_ERR(*slot, "got exception: %s\n", e.what()); - send_error(*slot, std::string("got exception: ") + e.what(), ERROR_TYPE_SERVER); - slot->release(); + if (preempt_planner_off) { + return; // test knob: leave the pool to the retry ladder and its last resort + } + + for (;;) { + const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); + + if (n_used + preempt_n_margin() <= n_cells) { + break; + } + + 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; + } + } + + bool recompute = false; + + server_slot * victim = preempt_pick_victim(&recompute); + + if (!victim) { + SRV_DBG("the kv pool needs %d of %d cells and nothing can be preempted (parked %.1f MiB of the %d MiB --preempt-ram budget)\n", + n_used, n_cells, preempt_ram_used() / (1024.0 * 1024.0), params_base.preempt_ram_mib); + break; + } + + const int32_t n_tokens = victim->prompt.n_tokens(); + const int64_t t_start = ggml_time_us(); + + if (!preempt_park(*victim, t_start, recompute)) { + break; // could not park it; the existing retry ladder is still behind us + } + + 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; + } + + if (victim->preempt_recompute) { + SLT_WRN(*victim, "preempted: %d cells dropped in %.2f ms, %d tokens to re-prefill on resume, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_n_input(), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + } else { + SLT_WRN(*victim, "preempted: %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, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); } } } - void abort_all_slots(const std::string & reason) { - for (auto & slot : slots) { - if (slot.is_processing()) { - send_error(slot, reason, ERROR_TYPE_SERVER); - slot.release(); + // 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 task_prompt_rejected(const server_task & task, std::string & msg, error_type & type) const { + // TODO: support memory-less logits computation + if (task.need_logits() && !llama_get_memory(ctx_tgt)) { + msg = "the current context does not logits computation. skipping"; + type = ERROR_TYPE_SERVER; + return true; + } + + // as launch_slot_with_task(), ahead of it: a sibling parked behind a running one used to fail inside a stream that had already opened 200 + if (!task.tokens.validate(ctx_tgt)) { + msg = "Prompt contains invalid tokens"; + type = ERROR_TYPE_INVALID_REQUEST; + return true; + } + + // as server_slot::can_split(), from the task alone + const bool can_split = + !task.need_embd() || + (llama_get_memory(ctx_tgt) && llama_pooling_type(ctx_tgt) == LLAMA_POOLING_TYPE_LAST); + + if (!can_split) { + const int32_t n_ubatch = llama_n_ubatch(ctx_tgt); + + if (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)", + task.n_tokens(), n_ubatch); + type = ERROR_TYPE_SERVER; + return true; + } + + if (task.n_tokens() > n_ctx_slot()) { + msg = string_format( + "input (%d tokens) is larger than the max context size (%d tokens). skipping", + task.n_tokens(), n_ctx_slot()); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; } + + return false; } - } - // @ngxson : for debugging only - int64_t t_pre_decode = 0; - int64_t t_decode = 0; - int64_t t_post_decode = 0; - int64_t t_sampl = 0; - int64_t n_pre_decode = 0; - int64_t n_decode = 0; - int64_t n_post_decode = 0; - int64_t n_sampl = 0; -// #define DEBUG_TIMINGS -#ifdef DEBUG_TIMINGS - struct scoped_timer { - int64_t & t; - int64_t & n; - int64_t t_start; - scoped_timer(int64_t & t_, int64_t & n_) : t(t_), n(n_) { - t_start = ggml_time_us(); + if (task.n_tokens() >= n_ctx_slot()) { + msg = string_format( + "request (%d tokens) exceeds the available context size (%d tokens), try increasing it", + task.n_tokens(), n_ctx_slot()); + type = ERROR_TYPE_EXCEED_CONTEXT_SIZE; + return true; } - ~scoped_timer() { - t += ggml_time_us() - t_start; - n++; + + return false; + } + + bool slot_prompt_rejected(const server_slot & slot, std::string & msg, error_type & type) const { + if (!slot.task) { + return false; } - }; -#else - struct scoped_timer { - scoped_timer(int64_t &, int64_t &) {} - ~scoped_timer() {} - }; -#endif + + return task_prompt_rejected(*slot.task, msg, type); + } void update_slots() { #ifdef DEBUG_TIMINGS @@ -2804,7 +4455,6 @@ struct server_context_impl { } #endif - // check if all slots are idle { bool all_idle = true; @@ -2832,6 +4482,14 @@ struct server_context_impl { } try { + // [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(); + + // before update_preemption(), and not only before the decode: a slot must never be parked with a shift still pending on its cells + preempt_apply_shift(); + + update_preemption(); + scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); batch.render(); @@ -2867,6 +4525,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 { @@ -2879,6 +4541,11 @@ struct server_context_impl { llama_synchronize(ctx_tgt); #endif + if (preempt_batch_abandoned) { + preempt_batch_abandoned = false; + break; + } + if (ok) { // move the head of the batch forward with the number of tokens we just processed off_next = off + n_tokens; @@ -2906,9 +4573,9 @@ struct server_context_impl { } } - void pre_decode() { - // apply context-shift if needed - // TODO: simplify and improve + // apply context-shift if needed + // TODO: simplify and improve + void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { if (!params_base.ctx_shift) { @@ -2948,6 +4615,9 @@ 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); @@ -2970,7 +4640,9 @@ struct server_context_impl { slot.truncated = true; } }); + } + void pre_decode() { // start populating the batch for this iteration batch.clear(); @@ -3116,7 +4788,8 @@ struct server_context_impl { return; // batch is full, skip remaining slots } - if (!slot.is_processing()) { + // [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; } @@ -3133,7 +4806,10 @@ struct server_context_impl { // this slot still has a prompt to be processed if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) { - const auto & input_tokens = slot.task->tokens; + const auto & input_tokens = slot.preempt_input(); + + // [TAG_PREEMPT] what the prompt step works towards: the re-prefill list of a park that dropped its cells, the request otherwise + const int32_t n_input_tokens = slot.preempt_n_input(); // used to determine the number of tokens added to the batch for the current slot const auto n_tokens_prev = batch.size(); @@ -3174,46 +4850,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); @@ -3269,6 +4917,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++; @@ -3427,11 +5077,32 @@ struct server_context_impl { if (!slot.can_split()) { // cannot fit the prompt in the current batch - will try next iter - if (batch.size() + slot.task->n_tokens() > n_batch) { + if (batch.size() + n_input_tokens > n_batch) { return; } } + // [TAG_EXACT_CONCURRENCY] a prompt is isolated into ubatches of its own, so it gets the shapes it would get alone only if what it adds here is a whole number of ubatches: otherwise a neighbour's decoded token shortens the last one, and 512,512,512,509 is not 512,512,512,512 + int32_t n_batch_cur = n_batch; + + if (common_exact_concurrency() && slot.can_split() && !slot.prompt.tokens.has_mtmd) { + const int32_t n_avail = n_batch - (int32_t) batch.size(); + const int32_t n_left = n_input_tokens - slot.prompt.n_tokens(); + + if (n_left > n_avail) { + const int32_t n_take = n_avail - n_avail % n_ubatch; + + if (n_take > 0) { + n_batch_cur = (int32_t) batch.size() + n_take; + } else { + // the waiting ends: common_exact_batch_geometry() refused a batch that cannot hold a whole ubatch beside a decode step of every slot, and the prompts ahead of this one in the same batch are finite + SLT_DBG(slot, "exact concurrency: %d of %d batch tokens left, short of a %d-token ubatch: the prefill waits\n", + n_avail, n_batch, n_ubatch); + return; + } + } + } + // note: the prompt timing is advanced in post_decode(), so it does not cover // the tokens added to the batch below slot.print_timings_pp(); @@ -3462,6 +5133,9 @@ struct server_context_impl { // make checkpoints only for completion tasks do_checkpoint = do_checkpoint && slot.task->type == SERVER_TASK_TYPE_COMPLETION; + // a re-prefill walks its own list, which the request's message spans do not index + do_checkpoint = do_checkpoint && !slot.preempt_reprefill; + // make a checkpoint of the parts of the memory that cannot be rolled back. // checkpoints are created only if: // - the model does not support partial sequence removal @@ -3478,12 +5152,16 @@ struct server_context_impl { while (true) { auto cur_token_idx = slot.prompt.n_tokens(); if ( - cur_token_idx >= slot.task->n_tokens() || + cur_token_idx >= n_input_tokens || input_tokens[cur_token_idx] != LLAMA_TOKEN_NULL // encountered a text token ) { break; } + // [TAG_PREEMPT_ASYNC] the chunk decodes whole, past the kv-full retry, so a park the planner issued for it has to land first + while (preempt_wait_in_flight()) { + } + // process the mtmd chunk // note: it submits its own decode, potentially be async // so the timing is queued and flushed on the next sync @@ -3504,8 +5182,12 @@ struct server_context_impl { } metrics_queue_prompt(n_tokens_out); - slot.stats.n_prompt_processed += n_tokens_out; - slot.stats.update_prompt_last(); + + // [TAG_PREEMPT] a re-prefill puts back what a park dropped: real compute, counted above, but it is not the request's prompt + if (!slot.preempt_reprefill) { + slot.stats.n_prompt_processed += n_tokens_out; + slot.stats.update_prompt_last(); + } // add the mtmd chunk to cache { @@ -3521,7 +5203,7 @@ struct server_context_impl { const auto last_user_pos = spans.last_user_message_pos(); // add prompt tokens for processing in the current batch - while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) { + while (slot.prompt.n_tokens() < n_input_tokens && batch.size() < n_batch_cur) { // get next token to process llama_token cur_tok = input_tokens[slot.prompt.n_tokens()]; if (cur_tok == LLAMA_TOKEN_NULL) { @@ -3546,6 +5228,11 @@ struct server_context_impl { /* is_prompt = */ true); slot.prompt.tokens.push_back(cur_tok); + // [TAG_EXACT_CONCURRENCY] a token that was decoded goes back through the arithmetic that decoded it: one per step, in the narrow set beside the other decodes. Re-prefilled wide it went through batched arithmetic, and the output diverged at the second park + if (slot.preempt_reprefill && common_exact_concurrency() && slot.prompt.n_tokens() >= slot.task->n_tokens()) { + break; + } + // break at the last user message, or at user messages at least min step past the last checkpoint if (do_checkpoint && spans.is_user_start(slot.prompt.n_tokens())) { const auto pos = slot.prompt.n_tokens(); @@ -3567,7 +5254,7 @@ struct server_context_impl { bool should_break = false; for (int offset : checkpoint_offsets) { const int n_last = std::min(n_batch, offset); - if (slot.task->n_tokens() == slot.prompt.n_tokens() + n_last) { + if (n_input_tokens == slot.prompt.n_tokens() + n_last) { should_break = true; break; } @@ -3583,13 +5270,13 @@ struct server_context_impl { const auto n_tokens_start = slot.prompt.n_tokens() - n_tokens_cur; - const bool near_prompt_end = slot.task->n_tokens() < slot.prompt.n_tokens() + n_ubatch; + const bool near_prompt_end = n_input_tokens < slot.prompt.n_tokens() + n_ubatch; const bool is_user_start = spans.is_user_start(n_tokens_start); const bool is_last_user_message = n_tokens_start == last_user_pos; // entire prompt has been processed - if (slot.prompt.n_tokens() == slot.task->n_tokens()) { + if (slot.prompt.n_tokens() == n_input_tokens) { slot.state = SLOT_STATE_DONE_PROMPT; GGML_ASSERT(batch.size() > 0); @@ -3597,10 +5284,14 @@ struct server_context_impl { // extract the logits only for the last token batch.set_output(batch.size() - 1, true); - slot.stats.n_gen = 0; - slot.i_batch = batch.size() - 1; + slot.i_batch = batch.size() - 1; - slot.init_sampler(); + // [TAG_PREEMPT] a re-prefill only puts back what the park dropped: the sampler and the counters carry on from where the park found them + if (!slot.preempt_reprefill) { + slot.stats.n_gen = 0; + + slot.init_sampler(); + } } else { // skip ordinary mid-prompt checkpoints, unless the batch starts a user // message or we are near the end of the prompt @@ -3642,6 +5333,123 @@ struct server_context_impl { } } + // [TAG_PREEMPT_ASYNC] a recurrent memory keeps no fixed row per sequence: find_slot() gathers the active ones into contiguous rows, so a decode beside a park moves or overwrites the row the copy is still reading. A hybrid carries that half too, and so can the draft. + static bool preempt_state_relocates(const llama_model * model) { + return model && (llama_model_is_recurrent(model) || llama_model_is_hybrid(model)); + } + + bool preempt_state_relocates() const { + return preempt_state_relocates(model_tgt) || preempt_state_relocates(model_dft); + } + + // [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) && !preempt_state_relocates(); + } + + 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. A media chunk mid-prompt keeps the old path. + bool preempt_last_resort(int32_t off) { + if (!preempt_last_resort_possible()) { + return false; + } + + int32_t n_running = 0; + + for (auto & slot : slots) { + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED || slot.preempt_in_flight()) { + continue; + } + + // a media chunk decodes whole through calls of its own, so a resident still inside its prompt cannot be rewound to a token boundary; one that is generating can + if (slot.prompt.tokens.has_mtmd && slot.state != SLOT_STATE_GENERATING) { + return false; + } + + n_running++; + } + + if (n_running < 2) { + return false; // one conversation that does not fit alone is a real overflow + } + + for (auto & slot : slots) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && slot.state != SLOT_STATE_WAIT_OTHER && + !slot.preempt_in_flight()) { + slot.rewind_to_cache(); + } + } + + const int32_t n_cells = n_ctx; + int32_t n_parked = 0; + + for (;;) { + const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); + + if (n_parked > 0 && n_used + preempt_n_margin() <= n_cells) { + break; + } + + bool recompute = false; + + server_slot * victim = preempt_pick_victim(&recompute); + + if (!victim) { + break; + } + + const int32_t n_tokens = victim->prompt.n_tokens(); + const int64_t t_start = ggml_time_us(); + + if (!preempt_park(*victim, t_start, recompute)) { + break; + } + + 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%s: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + victim->preempt_recompute ? " by dropping its cells" : "", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + victim->preempt_state_size() / (1024.0 * 1024.0), + preempt_kv_used(), n_cells, n_used, + victim->n_preempt); + } + + if (n_parked == 0) { + return false; // nothing could be parked: the error path clears what the rewind left + } + + SRV_WRN("last resort: batch given up at off = %d, %d slot(s) parked, kv %d/%d resident\n", + off, n_parked, preempt_kv_used(), n_cells); + + return true; + } + + bool batch_has_spec_groups() const { + for (const auto & slot : slots) { + if (!slot.spec_i_batch.empty()) { + return true; + } + } + + return false; + } + // returns true = success ; false = retry with smaller batch size // throw std::runtime_error on fatal error bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) { @@ -3686,10 +5494,31 @@ 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] 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"); + return false; // retry at the same width + } + + n_batch = 1; + } + if (n_batch == 1 && ret == 1) { + if (preempt_last_resort(off)) { + preempt_batch_abandoned = true; + return true; + } + // TODO: try to terminate only the largest active slot/sequence and continue with the rest // need to remove the tokens from the current batch too err = "Context size has been exceeded."; @@ -3710,7 +5539,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) { - if (slot.is_processing()) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED && !slot.preempt_in_flight()) { send_error(slot, err); slot.release(); @@ -3805,7 +5634,7 @@ struct server_context_impl { iterate(slots, [&](server_slot & slot) { // optionally send prompt processing progress if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_DONE_PROMPT) { - if (slot.task->params.stream && slot.task->params.return_progress) { + if (slot.task->params.stream && slot.task->params.return_progress && !slot.preempt_reprefill) { send_partial_response(slot, {}, true); } } @@ -3816,6 +5645,12 @@ struct server_context_impl { } if (slot.state == SLOT_STATE_DONE_PROMPT) { + // [TAG_PREEMPT] the re-prefill is back in the cache; the token this slot had already sampled is decoded next, so nothing is sampled here + if (slot.preempt_reprefill) { + slot.preempt_reprefill_done(); + return; + } + if (slot.task->type == SERVER_TASK_TYPE_EMBEDDING) { // prompt evaluated for embedding send_embedding(slot, batch_view); @@ -4071,7 +5906,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) { - if (slot.is_processing()) { + 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()); @@ -4092,7 +5927,8 @@ struct server_context_impl { n_prompt_tokens++; auto & slot = slots[t.id_slot]; - if (slot.stats.is_set()) { + // [TAG_PREEMPT] replayed tokens stay out of the slot's prompt count, they were counted when the request first processed its prompt + if (slot.stats.is_set() && !slot.preempt_reprefill) { slot.stats.n_prompt_processed++; } } @@ -4110,7 +5946,8 @@ struct server_context_impl { for (int i = off; i < off + n_tokens; ++i) { const auto & t = batch.tokens[i]; auto & slot = slots[t.id_slot]; - if (t.is_prompt && slot.stats.is_set()) { + // [TAG_PREEMPT] a re-prefill must not move the prompt/generation boundary: n_gen carries across the park, so the generation time would then cover only the tokens after it + if (t.is_prompt && slot.stats.is_set() && !slot.preempt_reprefill) { slot.stats.set_prompt_last(t_now); } } @@ -4328,6 +6165,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; @@ -4339,6 +6184,16 @@ std::unique_ptr server_routes::handle_completions_impl( tasks.push_back(std::move(task)); } + // [TAG_PREEMPT] every prompt of the request, before any of them is queued: one member can be parked, and its notice opens the stream, before another member is rejected + { + json error; + + if (ctx_server.tasks_prompt_rejected(tasks, error)) { + res->error(error); + return res; + } + } + rd.post_tasks(std::move(tasks)); } catch (const std::exception & e) { res->error(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST)); @@ -4381,37 +6236,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({ @@ -4462,10 +6331,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 } @@ -4475,7 +6347,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; } @@ -4491,12 +6363,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) { @@ -4624,6 +6507,8 @@ static json get_res_props(const server_context_meta & meta, const common_params { "endpoint_slots", params.endpoint_slots }, { "endpoint_props", params.endpoint_props }, { "endpoint_metrics", params.endpoint_metrics }, + // [TAG_EXACT_CONCURRENCY] a client that asked for the mode reads here whether this process runs it: a build that ignores the variable starts all the same + { "exact_concurrency", common_exact_concurrency() }, { "ui", params.ui }, { "ui_settings", meta.json_ui_settings }, { "chat_template", tmpl_default }, diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 78169e9a5d86..c555e1856c2c 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 @@ -527,12 +530,16 @@ void server_response_reader::post_tasks(std::vector && tasks, bool id_tasks = server_task::get_list_id(tasks); states.reserve(tasks.size()); size_t index = 0; + // [TAG_PREEMPT] several prompts, or several completions of one prompt, all number their results, and their preempt notices have to say which one they belong to + const bool batched = id_tasks.size() > 1; for (auto & task : tasks) { - task.index = index++; + task.index = index++; + task.batched = batched; states.push_back(task.create_state()); // for child tasks for (auto & child_task : task.child_tasks) { - child_task.index = index++; + child_task.index = index++; + child_task.batched = batched; states.push_back(child_task.create_state()); } } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..5c2376b1fc2d 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -355,6 +355,7 @@ json server_task_result_cmpl_final::to_json_non_oaicompat() { {"stopping_word", stopping_word}, {"tokens_cached", n_tokens_cached}, {"timings", stats.to_json()}, + {"preempt", preempt_to_json()}, }; if (!stream && !probs_output.empty()) { res["completion_probabilities"] = completion_token_output::probs_vector_to_json(probs_output, post_sampling_probs); @@ -362,6 +363,14 @@ json server_task_result_cmpl_final::to_json_non_oaicompat() { return response_fields.empty() ? res : json_get_nested_values(response_fields, res); } +// [TAG_PREEMPT] how the request was served: a recompute resume re-prefilled its tokens, so its continuation is not the bytes that were parked +json server_task_result_cmpl_final::preempt_to_json() const { + return json { + {"parks", n_preempt}, + {"recomputes", n_recompute}, + }; +} + json server_task_result_cmpl_final::usage_json_oaicompat() { return json { {"completion_tokens", n_decoded}, @@ -406,6 +415,7 @@ json server_task_result_cmpl_final::to_json_oaicompat() { } if (stats.is_set()) { res["timings"] = stats.to_json(); + res["preempt"] = preempt_to_json(); } return res; @@ -454,6 +464,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat() { } if (stats.is_set()) { res["timings"] = stats.to_json(); + res["preempt"] = preempt_to_json(); } return res; @@ -515,6 +526,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat_stream() { if (stats.is_set()) { deltas.back()["timings"] = stats.to_json(); + deltas.back()["preempt"] = preempt_to_json(); } // extra fields for debugging purposes @@ -591,6 +603,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp() { {"total_tokens", n_decoded + n_prompt_tokens}, {"input_tokens_details", json { {"cached_tokens", n_prompt_tokens_cache} }}, }}, + {"preempt", preempt_to_json()}, }; return res; @@ -701,7 +714,9 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() { {"output_tokens", n_decoded}, {"total_tokens", n_decoded + n_prompt_tokens}, {"input_tokens_details", json { {"cached_tokens", n_prompt_tokens_cache} }}, - }} + }}, + // [TAG_PREEMPT] inside the response object, where the non-streaming body carries it: that object is what a client keeps from the stream + {"preempt", preempt_to_json()}, }}, }} }); @@ -724,6 +739,7 @@ json server_task_result_cmpl_final::to_json_oaicompat_asr() { {"total_tokens", n_decoded + n_prompt_tokens}, {"input_tokens_details", json { {"cached_tokens", n_prompt_tokens_cache} }}, }}, + {"preempt", preempt_to_json()}, }; return event; } @@ -788,7 +804,8 @@ json server_task_result_cmpl_final::to_json_anthropic() { {"cache_read_input_tokens", n_prompt_tokens_cache}, {"input_tokens", n_prompt_tokens - n_prompt_tokens_cache}, {"output_tokens", n_decoded} - }} + }}, + {"preempt", preempt_to_json()} }; return res; @@ -968,7 +985,8 @@ json server_task_result_cmpl_final::to_json_anthropic_stream() { }}, {"usage", { {"output_tokens", n_decoded} - }} + }}, + {"preempt", preempt_to_json()} }} }); @@ -1023,6 +1041,14 @@ void server_task_result_cmpl_partial::update(task_result_state & state) { } } +json server_task_result_preempt_notice::to_json() { + return json { + {"preempted", parked}, + {"recomputed", recomputed}, + {"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) { @@ -1562,6 +1588,18 @@ std::string server_task_result_metrics::to_metrics() { "spec_decode_num_drafts_total", "Speculative: Total speculative decoding verification steps", (double) metrics.n_draft_verif_steps + }, { + "n_preempt_total", + "Preemption: Total slots parked to make room in the unified KV cache", + (double) metrics.n_preempt + }, { + "n_resume_total", + "Preemption: Total parked slots put back", + (double) metrics.n_resume + }, { + "preempt_recompute_total", + "Preemption: Total parks that dropped their cells, whose resume re-prefills instead of restoring the saved bytes", + (double) metrics.n_preempt_recompute }, }; @@ -1586,6 +1624,14 @@ std::string server_task_result_metrics::to_metrics() { "n_busy_slots_per_decode", "Average number of busy slots per llama_decode() call", (double) metrics.n_busy_slots / std::max((double) metrics.n_decode, 1.0) + }, { + "requests_preempted", + "Preemption: Number of requests currently parked, waiting for room in the unified KV cache", + (double) n_preempted_slots + }, { + "preempt_ram_bytes", + "Preemption: Host RAM held by parked sequences", + (double) preempt_ram_bytes }, }; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..c5dd2206108e 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -139,6 +139,9 @@ struct server_task { // TODO @ngxson : remove this field and implement a mapping task_id -> idx in the response_reader size_t index = 0; // used when there are multiple prompts (batch request) + // [TAG_PREEMPT] this request yielded more than one task, so index tells its results apart and the preempt notices carry it + bool batched = false; + // used by SERVER_TASK_TYPE_CANCEL int id_target = -1; int id_slot = -1; @@ -339,6 +342,12 @@ struct server_task_result_cmpl_final : server_task_result { std::vector probs_output; std::vector response_fields; + // [TAG_PREEMPT] how the request was served: how often it was parked, and how many of those parks re-prefilled instead of restoring saved bytes + int32_t n_preempt = 0; + int32_t n_recompute = 0; + + json preempt_to_json() const; + task_params generation_params; // response formatting @@ -392,6 +401,19 @@ 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 + bool recomputed = false; // this resume re-prefilled its tokens instead of restoring saved bytes + int32_t n_preempt = 0; // how many times this task has been parked so far + bool batched = false; // one of several tasks of its request, so the notice names which one by index + + 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; @@ -494,6 +516,8 @@ struct server_task_result_metrics : server_task_result { // these are immediate stats, not accumulated (server_metrics is cumulative) int n_processing_slots = 0; int n_tasks_deferred = 0; + int n_preempted_slots = 0; // [TAG_PREEMPT] processing slots currently parked + size_t preempt_ram_bytes = 0; // [TAG_PREEMPT] host RAM their parked sequences hold server_metrics metrics; diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py new file mode 100644 index 000000000000..c47cda62b17f --- /dev/null +++ b/tools/server/tests/unit/test_preempt.py @@ -0,0 +1,1094 @@ +import base64 +import json +import os +import re +import struct +import subprocess +import threading +from concurrent.futures import ThreadPoolExecutor +import time +import tempfile +import pytest +import requests +from utils import * + +# 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() + +_ASYNC_BANNER = "parking and resuming asynchronously" + +_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" + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + # the server parks only when asked: --preempt-ram defaults to 0, and this suite is about parking + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "8192" + server.server_slots = True + server.server_metrics = True + server.temperature = 0.0 + server.seed = 42 + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + for name in ("LLAMA_SERVER_PREEMPT_EVERY", "LLAMA_SERVER_PREEMPT_GRANULARITY", + "LLAMA_SERVER_PREEMPT_PLANNER", "LLAMA_ARG_PREEMPT_RAM", "LLAMA_ARG_PREEMPT_ASYNC", + "LLAMA_SERVER_PREEMPT_FAIL_SAVE", "LLAMA_ARG_SPEC_DRAFT_P_MIN", "LLAMA_ARG_LOG_VERBOSITY", + "LLAMA_BATCH_DEBUG", "LLAMA_ARG_CTX_CHECKPOINTS", + "LLAMA_MEDIA_MARKER", "LLAMA_EXACT_CONCURRENCY"): + os.environ.pop(name, None) + + +def _start(**kwargs): + """Start the server with these settings; its log starts empty again on every start.""" + for key, value in kwargs.items(): + setattr(server, key, value) + server.start() + + +def _start_async(**kwargs): + """As _start, on the asynchronous park path; a backend that cannot copy off-thread skips the test.""" + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "1" + _start(n_gpu_layer=99, **kwargs) + _require_async(_log()) + + +def _log() -> str: + """The server log once its writer thread has stopped growing it: on a loaded host the log lags the response that came from it.""" + deadline = time.time() + 5.0 + last = -1 + while time.time() < deadline: + size = os.path.getsize(server.log_path) + if size == last: + break + last = size + time.sleep(0.1) + return open(server.log_path, errors="replace").read() + + +def _require_async(text: str): + if _ASYNC_BANNER not in text: + pytest.skip("this backend cannot copy asynchronously, the async park path is not exercised") + + +def _complete(n_predict: int, prompt="Hi how are you", id_slot: int = -1, delay: float = 0.0, after_slot_busy=None, + timeout: float = DEFAULT_REQUEST_TIMEOUT): + time.sleep(delay) + if after_slot_busy is not None: + # sent once that slot is processing, so the request queues behind it whatever the host's speed + for _ in range(200): + slots = server.make_request("GET", "/slots").body + if any(s["id"] == after_slot_busy and s["is_processing"] for s in slots): + break + time.sleep(0.02) + return server.make_request("POST", "/completion", data={ + "n_predict": n_predict, "prompt": prompt, "id_slot": id_slot, + "ignore_eos": True, "return_tokens": True, "temperature": 0.0, "seed": 42, + }, timeout=timeout) + + +def _complete_all(n_predict: int, prompts=(_PROMPT_A, _PROMPT_B), timeout: float = DEFAULT_REQUEST_TIMEOUT): + return parallel_function_calls([(_complete, (n_predict, prompt, -1, 0.0, None, timeout)) for prompt in prompts]) + + +def _wait_processing(slot_ids, timeout: float = 30.0): + """Return once every one of these slots is processing; a request that ended first is a failure, not a hang.""" + deadline = time.time() + timeout + while time.time() < deadline: + slots = server.make_request("GET", "/slots").body + if all(any(s["id"] == i and s["is_processing"] for s in slots) for i in slot_ids): + return + time.sleep(0.005) + pytest.fail(f"slots {slot_ids} never showed as processing") + + +def _complete_overlapping(n_predict, n_prompt, timeout: float = DEFAULT_REQUEST_TIMEOUT): + """A leader on slot 0 and a follower on slot 1 that certainly overlap: the follower is sent once the leader is seen processing, so the lengths and not the client's speed decide what the pool has to hold.""" + leader = _prompt_of(n_prompt[0], _PROMPT_A) + other = _prompt_of(n_prompt[1], _PROMPT_B) + with ThreadPoolExecutor(1) as pool: + first = pool.submit(_complete, n_predict[0], leader, 0, 0.0, None, timeout) + _wait_processing([0]) + second = _complete(n_predict[1], other, 1, 0.0, None, timeout) + return [first.result(), second] + + +def _wait_preempted(timeout: float = 30.0) -> bool: + """True once some slot is parked: its cells are in host RAM and it wants them back.""" + deadline = time.time() + timeout + while time.time() < deadline: + slots = server.make_request("GET", "/slots").body + if any(s["is_preempted"] for s in slots): + return True + time.sleep(0.005) + return False + + +def _prompt_of(n_tokens: int, text: str) -> list: + """A prompt of exactly n_tokens tokens, as ids: no BOS is added to one of those.""" + base = server.make_request("POST", "/tokenize", data={"content": text}).body["tokens"] + assert base + return (base * (n_tokens // len(base) + 1))[:n_tokens] + + +def _assert_completed(results, n_predict: int): + for res in results: + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_predict + + +_PARK_MARKERS = ("preempted:", "preempted as a last resort", "preempted on request") + + +def _assert_nothing_parked(text: str): + # the park log lines, not the bare word: a verbose log prints every slot's "is_preempted" + assert not any(m in text for m in _PARK_MARKERS), "nothing could be parked here" + + +def _assert_recovered(text: str, parked: str = "preempted:"): + """Nothing was ended for want of cells: a slot was parked and came back.""" + assert "Context size has been exceeded" not in text + assert parked in text + assert "resumed after" in text + + +def _metrics() -> dict: + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + return { + name[len("llamacpp:"):]: float(value) + for name, value in (line.split(" ", 1) for line in res.body.splitlines() if line.startswith("llamacpp:")) + } + + +@pytest.mark.parametrize("mode", ["sync", "async", "no-async"]) +def test_forced_parks_do_not_change_the_output(mode): + if mode != "sync": + server.n_gpu_layer = 99 + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "1" if mode == "async" else "0" + _start(n_ctx=512) + if mode == "async": + _require_async(_log()) + reference = _complete(64) + assert reference.status_code == 200 + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start() + preempted = _complete(64) + assert preempted.status_code == 200 + assert preempted.body["timings"]["predicted_n"] == 64 + assert preempted.body["content"] == reference.body["content"] + assert preempted.body["tokens"] == reference.body["tokens"] + + text = _log() + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 + if mode == "async": + assert "park completed after" in text + assert "restore issued in" in text + assert "restore completed after" in text + if mode == "no-async": + assert _ASYNC_BANNER not in text + assert "park issued in" not in text + + +@pytest.mark.parametrize("knob", ["planner", "pages", "async", "last-resort", "last-resort-unlimited"]) +def test_two_generations_that_do_not_fit_together_both_finish(knob): + # each request fits the pool alone (960 and 600 of 1024 cells) but not together; without preemption both end with "Context size has been exceeded" + if knob == "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 + os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64" + if knob.startswith("last-resort"): + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + if knob == "last-resort-unlimited": + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "-1" + n_ctx = 1024 + (_start_async if knob == "async" else _start)(n_ctx=n_ctx) + + # the lengths, not the client's speed, decide the overlap: two equal requests fired together did not overlap on a Windows runner, the first finished before the second arrived, and the last resort never saw the two residents it needs. + # the follower is sent once the leader is seen processing, so it holds its cells while the leader grows into the rest of the pool + n_predict = (460, 400) + results = _complete_overlapping(n_predict, (500, 200)) + + text = _log() + _assert_recovered(text, "preempted as a last resort" if knob.startswith("last-resort") else "preempted:") + for res, n_wanted in zip(results, n_predict): + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_wanted + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_wanted + + if knob == "pages": + held = [int(n) for n in re.findall(rf"kv (\d+)/{n_ctx}", 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}" + if knob.startswith("last-resort"): + assert "preempted:" not in text, "the planner was off, nothing may be parked ahead of the decode" + assert "last resort: batch given up" in text + if knob == "planner": + metrics = _metrics() + assert metrics["n_preempt_total"] >= 1 + assert metrics["n_resume_total"] == metrics["n_preempt_total"] + assert metrics["requests_preempted"] == 0 + assert metrics["preempt_ram_bytes"] == 0 + + +@pytest.mark.parametrize("knob", ["ram-0", "family"]) +def test_a_request_that_cannot_be_helped_gets_the_context_error_and_the_server_lives(knob): + if knob == "ram-0": + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + else: + # a family member is not a victim for the other, so a two-completion request gets the error it would get alone + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + _start(n_ctx=256) + + if knob == "ram-0": + # the overflow has to be a matter of lengths: two equal requests fired together did not overlap on a Windows runner, and each one fits the pool alone + assert any(res.status_code != 200 for res in _complete_overlapping((110, 100), (120, 60))) + else: + res = server.make_request("POST", "/completion", data={ + "n_predict": 160, "n_cmpl": 2, "prompt": _PROMPT_A, + "ignore_eos": True, "temperature": 0.0, "seed": 42, + }) + assert res.status_code == 500 + assert "Context size has been exceeded" in res.body["error"]["message"] + + text = _log() + assert "Context size has been exceeded" in text + _assert_nothing_parked(text) + assert "GGML_ASSERT" not in text + after = _complete(8) + assert after.status_code == 200 + assert after.body["timings"]["predicted_n"] == 8 + + +def test_a_server_that_never_asked_for_parking_behaves_as_upstream(): + """--preempt-ram defaults to 0, so a unified-cache server started without it parks nothing.""" + os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) + _start(n_ctx=256) + + text = _log() + assert "preemption:" not in text, "a server that did not ask for parking announced it" + assert _ASYNC_BANNER not in text, "the async park path was set up without being asked for" + + # as above, the two have to be resident together for the pool to overflow at all + assert any(res.status_code != 200 for res in _complete_overlapping((110, 100), (120, 60))) + + text = _log() + assert "Context size has been exceeded" in text + _assert_nothing_parked(text) + assert "last resort" not in text, "the retry ladder consulted the planner" + assert "GGML_ASSERT" not in text + + metrics = _metrics() + assert metrics["n_preempt_total"] == 0 + assert metrics["preempt_ram_bytes"] == 0 + + after = _complete(8) + assert after.status_code == 200 + assert after.body["timings"]["predicted_n"] == 8 + + +@pytest.mark.parametrize("planner", ["on", "off"]) +def test_a_late_prompt_and_a_generating_slot_both_finish(planner): + if planner == "off": + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + _start(n_ctx=256) + + n_b = 150 + n_predict_a = 230 + n_predict_b = 90 + assert 8 + n_predict_a + n_b + n_predict_b > 256 + results = parallel_function_calls([ + (_complete, (n_predict_a, "Hi how are you")), + (_complete, (n_predict_b, _prompt_of(n_b, _PROMPT_C), -1, 0.02)), + ]) + + text = _log() + assert "Context size has been exceeded" not in text + assert ("preempted as a last resort" if planner == "off" else "preempted:") in text + for res, n_predict in zip(results, (n_predict_a, n_predict_b)): + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_predict + # the chunk in the batch given up is processed once after the rewind, never twice + assert results[1].body["timings"]["prompt_n"] == n_b + + +def test_a_prompt_parked_before_its_first_token_is_issued_whole(): + # both prompts are too close to n_ctx to leave the usual margin, so the second is parked before it takes a cell and has to come back once the first has finished + _start(n_ctx=256, n_batch=256) + + n_prompt = 240 + n_predict = 4 + long_prompt = _prompt_of(n_prompt, "Once upon a time there was a little girl") + together = _complete_all(n_predict, [long_prompt, long_prompt]) + + assert "cannot fit the pool" not in _log() + _assert_completed(together, n_predict) + for res in together: + assert res.body["timings"]["prompt_n"] == n_prompt, "the prompt was not issued once and whole" + + +def test_a_resident_cycling_through_context_shifts_is_rotated_out_for_a_parked_head(): + # with context shift on a resident would hold its cells for as long as it generates, so once the head has waited its turn the resident is parked and the two take turns + _start(n_slots=3, n_ctx=384, enable_ctx_shift=True) + + # the rotation waits on a clock, not on a token count, so the generation has to be long enough on a fast host; that is a lot of tokens for a slow one, and it takes turns with two others, so it is given more than the usual wait + n_predict = 9000 + _assert_completed(_complete_all(n_predict, (_PROMPT_A, _PROMPT_B, _PROMPT_C), timeout=1800), n_predict) + + text = _log() + _assert_recovered(text, "rotated out after") + assert "slot context shift" in text + + +def test_a_park_whose_host_allocation_fails_is_parked_by_recompute(): + # the budget grants permission to allocate, not a successful allocation: a failed save used to stop the planner and leave the pool to overflow, although the same victim could be parked by dropping its cells + os.environ["LLAMA_SERVER_PREEMPT_FAIL_SAVE"] = "1" + _start(n_ctx=1024, n_slots=2) + + # lengths decide the overlap, not the host's speed: two 171-cell requests fired together did not overlap on a Windows runner, so nothing was parked. The leader ends at 960 of 1024 cells, so the second is parked whatever the client's lag + leader = _prompt_of(500, _PROMPT_A) + other = _prompt_of(200, _PROMPT_B) + with ThreadPoolExecutor(1) as pool: + first = pool.submit(_complete, 460, leader, 0) + _wait_processing([0]) + second = _complete(400, other, 1) + results = [first.result(), second] + + text = _log() + assert "could not take the host memory" in text, "the injected allocation failure never fired" + assert "tokens to re-prefill" in text, "the failed save did not fall back to recompute" + assert "Context size has been exceeded" not in text + for res, n_predict in zip(results, (460, 400)): + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_predict + + +def test_a_resident_that_cannot_be_swapped_out_is_rotated_by_recompute(): + # 1 MiB holds no snapshot, so rotation refused every resident and only logged that the head waits: a resident that keeps context-shifting need never finish, and the head waited behind it for good + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + _start(n_ctx=2048, n_slots=2, n_batch=2048, enable_ctx_shift=True) + + # a small n_discard keeps the resident near the end of the pool, so its state never fits the budget + def unending_request(): + return server.make_request("POST", "/completion", data={ + "n_predict": 100000, "prompt": _PROMPT_A, "n_keep": 1, "n_discard": 64, + "ignore_eos": True, "temperature": 0.0, "seed": 42, + }, timeout=600) + + # the server is stopped below with this request still in flight, so the disconnect it then sees is expected and must not surface as an unhandled thread exception + unending = [] + + def run_unending(): + try: + unending.append(unending_request()) + except requests.exceptions.RequestException: + pass + + t = threading.Thread(target=run_unending) + t.start() + + try: + # the resident has to be at the pool's limit and cycling before a second prompt cannot fit beside it + for _ in range(3000): + if "slot context shift" in _log(): + break + time.sleep(0.05) + else: + pytest.fail("the resident never reached the end of the pool") + + waiting = server.make_request("POST", "/completion", data={ + "n_predict": 8, "prompt": _prompt_of(1800, _PROMPT_C), + "ignore_eos": True, "temperature": 0.0, "seed": 42, + }, timeout=300) + + assert waiting.status_code == 200, waiting.body + assert waiting.body["timings"]["predicted_n"] == 8, "the second request never made progress" + assert not unending, "the first request ended before the second made progress" + finally: + server.stop() + t.join(60) + + text = _log() + assert "slot context shift" in text + assert re.search(r"rotated out after .* cells dropped", text), "the rotation did not fall back to recompute" + assert "Context size has been exceeded" not in text + + +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" + _start_async(n_ctx=512) + + for i in range(4): + try: + server.make_request("POST", "/completion", data={ + "n_predict": 96, "prompt": _PROMPT_A, "ignore_eos": True, "temperature": 0.0, "seed": 42, + }, timeout=0.05 + 0.1 * i) + except Exception: + pass # the point is the drop, not the response + + for _ in range(600): + 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 + assert slot["is_transferring"] is False + assert _metrics()["preempt_ram_bytes"] == 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_a_started_slot_is_counted_by_the_cells_it_holds_not_by_the_prompt_it_keeps(): + # the last request waits for slot 0 and is started on it holding the first request's cells; counted by the prompt it keeps instead, the pool looks free and a parked slot is restored into cells that are still taken + _start(n_ctx=1024, n_slots=3) + + # the lengths, not the host's speed, decide who is parked: the three prompts (500 + 200 + 200) fit the 1024 cells, so both of the others are parked holding at least their whole prompt once slot 0 grows into the rest, and slot 0 is the largest slot throughout, which the planner never picks as a victim. Slot 0 ends holding 960 of the 1024 cells, too few left for either parked slot to come back + ids = _prompt_of(500, _PROMPT_C) + + with ThreadPoolExecutor(4) as pool: + first = pool.submit(_complete, 460, ids, 0) + _wait_processing([0]) + # queued behind slot 0 whatever the host's speed, and a real prefix of what slot 0 holds: it starts on 960 cells while keeping 8 of them + follower = pool.submit(_complete, 8, ids[:8], 0) + long_ones = [pool.submit(_complete, 400, _prompt_of(200, _PROMPT_A), 1), + pool.submit(_complete, 400, _prompt_of(200, _PROMPT_B), 2)] + parked = _wait_preempted() + results = [first.result(), follower.result(), long_ones[0].result(), long_ones[1].result()] + + text = _log() + assert parked, "the pool never came under pressure, so no slot was waiting for the cells slot 0 keeps" + assert "trimmed to the" in text, "the started slot kept the cells of the request before it" + assert "resume failed" not in text + assert "Context size has been exceeded" not in text + for res, n_predict in zip(results, (460, 8, 400, 400)): + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == n_predict + + +def test_a_recurrent_model_is_served_without_preemption(): + server.model_file = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") + if server.model_file: + server.model_hf_repo = server.model_hf_file = None + else: + server.model_hf_repo = "Felladrin/gguf-mamba-130m-hf" + server.model_hf_file = "mamba-130m-hf.Q2_K.gguf" + server.offline = False + server.n_ctx = 1024 + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start(timeout_seconds=300) + + # not "Once upon a time": what this Q2_K model decodes from it on CUDA carries bytes the content parser refuses, master included, which is not what this test measures + results = _complete_all(64, ["The quick brown fox", "Hello world"]) + _assert_completed(results, 64) + + text = _log() + assert "preemption: off, the recurrent cache holds one state per sequence" in text + _assert_nothing_parked(text) + assert "Context size has been exceeded" not in text + + +def _stream_completion(n_predict: int, prompt: str) -> tuple[list[str], dict]: + """One streaming completion: its SSE comment lines and the last response object.""" + url = f"http://{server.server_host}:{server.server_port}/completion" + res = requests.post(url, json={ + "prompt": prompt, "n_predict": n_predict, "ignore_eos": True, + "temperature": 0.0, "seed": 42, "stream": True, + }, stream=True, timeout=600) + assert res.status_code == 200, res.text + comments, datas = [], [] + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line.startswith(":"): + comments.append(line) + elif line.startswith("data: ") and line[6:] != "[DONE]": + datas.append(json.loads(line[6:])) + return comments, datas[-1] + + +@pytest.mark.parametrize("planner", ["on", "off"]) +def test_a_budget_that_holds_no_sequence_parks_by_dropping_the_cells(planner): + # 1 MiB holds neither sequence, so no victim fits the budget: the park drops the cells and the resume re-prefills the tokens, instead of the pool overflowing and ending both; the last resort falls back the same way + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + if planner == "off": + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + # one batch for the re-prefill: a prompt of this length in 32-token batches hangs the CUDA build with graphs on, master included, so that is not what this test measures + _start(n_ctx=3840, n_batch=2048) + + n_predict = 2000 + if planner == "on": + results = parallel_function_calls([(_stream_completion, (n_predict, p)) for p in (_PROMPT_A, _PROMPT_B)]) + for comments, final in results: + assert "error" not in final, final + assert final["tokens_predicted"] == n_predict + comments = [c for cs, _ in results for c in cs] + assert ": preempted" in comments and ": resumed" in comments, comments + else: + _assert_completed(_complete_all(n_predict), n_predict) + + text = _log() + assert "Context size has been exceeded" not in text + assert "tokens to re-prefill" in text, "no park fell back to recompute" + if planner == "off": + assert "preempted as a last resort by dropping its cells" in text + + +_IMG_URL = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/11_truck.png" + + +def test_a_media_chunk_is_reserved_whole_before_it_is_decoded(): + # a chunk is decoded whole inside one iteration, through decodes the kv-full retry does not cover: unless the planner reserves every cell it takes, the second of two image requests that each fit alone fails part way through its chunk + os.environ["LLAMA_MEDIA_MARKER"] = "<__media__>" + server.model_hf_repo = "ggml-org/tinygemma3-GGUF:Q8_0" + server.model_hf_file = None + server.model_alias = "tinygemma3" + _start(n_ctx=400, n_batch=64, n_ubatch=64) + + image = base64.b64encode(requests.get(_IMG_URL, timeout=60).content).decode() + prompt = {"prompt_string": "<__media__>\nWhat is in this image?", "multimodal_data": [image]} + results = parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": prompt, "n_predict": 4, "temperature": 0.0, "seed": 42, + })) for _ in range(2) + ]) + + text = _log() + assert "failed to process mtmd chunk" not in text + assert "preempted:" in text, "nothing was parked to make room for a chunk" + for res in results: + assert res.status_code == 200, res.body + assert res.body["timings"]["prompt_n"] > 64, "the chunk fits one batch, so it never spans several decodes" + + +def test_an_mtp_draft_stays_inside_the_reservation_it_was_priced_for(): + # near the end of the pool the planner prices one draft token and one sampled token, but the draft loop stopped against the configured window and could attempt positions past the reservation + path = os.environ.get("LLAMA_SERVER_TEST_MTP_MODEL") + if not path: + pytest.skip("set LLAMA_SERVER_TEST_MTP_MODEL to a gguf carrying an MTP head") + server.model_file = path + server.model_hf_repo = server.model_hf_file = None + server.spec_type = "draft-mtp" + os.environ["LLAMA_ARG_SPEC_DRAFT_P_MIN"] = "0.0" # nothing but the bounds stops the draft loop + os.environ["LLAMA_ARG_LOG_VERBOSITY"] = "5" # the wrapper says so when it truncates what an implementation returned + _start(n_ctx=2048, n_slots=2, n_batch=2048, n_gpu_layer=99, spec_draft_n_max=128, spec_draft_n_min=1) + + n_prompt = 2045 + res = server.make_request("POST", "/completion", data={ + "prompt": _prompt_of(n_prompt, _PROMPT_C), "n_predict": 2, "ignore_eos": True, + "temperature": 0.0, "seed": 42, "cache_prompt": False, + }, timeout=600) + + assert res.status_code == 200, res.body + assert res.body["timings"]["prompt_n"] == n_prompt + assert res.body["tokens_predicted"] == 2 + + text = _log() + assert "truncating draft to" not in text, "the draft was scheduled past the tokens the planner reserved" + assert "llama_decode[" not in text + assert "Context size has been exceeded" not in text + + +def test_a_hybrid_model_parks_synchronously(): + # the recurrent half of a hybrid gathers the active sequences into contiguous rows on every batch, so a copy running beside the decode could read a row another sequence has been moved into + path = os.environ.get("LLAMA_SERVER_TEST_HYBRID_MODEL") + if not path: + pytest.skip("set LLAMA_SERVER_TEST_HYBRID_MODEL to a hybrid attention/recurrent gguf") + server.model_file = path + server.model_hf_repo = server.model_hf_file = None + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "1" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=1024, n_gpu_layer=99) + + res = _complete(24, "Once upon a time") + assert res.status_code == 200, res.body + assert res.body["timings"]["predicted_n"] == 24 + + text = _log() + assert "a recurrent state does not stay in one row" in text + assert _ASYNC_BANNER not in text + assert "park issued in" not in text + _assert_recovered(text, "preempted on request") + + +# [TAG_EXACT_CONCURRENCY] the paged pool places a cell from the sequence and the position alone, so a layout that gives several tokens one position cannot be served + +def _mrope_model() -> str: + path = os.environ.get("LLAMA_SERVER_TEST_MROPE_MODEL") + if not path: + pytest.skip("set LLAMA_SERVER_TEST_MROPE_MODEL to an M-RoPE gguf") + return path + + +def test_exact_concurrency_refuses_an_mrope_model_with_a_projector(): + # every token of one image shares a temporal position under M-RoPE, so the pool would give the second one the first one's cell + path = _mrope_model() + with tempfile.TemporaryDirectory() as tmp: + mmproj = os.path.join(tmp, "mmproj.gguf") + with open(mmproj, "wb") as f: + f.write(b"GGUF" + struct.pack(" list: + """The tokens of every ubatch a split produced, in order; needs LLAMA_BATCH_DEBUG.""" + res = [] + pending = False + for line in text.splitlines(): + if "added ubatch to split" in line: + pending = True + elif pending and "n_tokens" in line: + res.append(int(line.split("=")[-1])) + pending = False + return res + + +def test_exact_concurrency_prefills_a_prompt_in_the_ubatches_it_would_get_alone(): + # generated tokens enter the batch first and a prompt took what was left, so its ubatches were 512,512,512,509 beside three decoders and 512,512,512,512 alone: isolating the sequences does not make the shapes equal by itself + # geometry: -b 2048 -ub 512 -np 4, so the batch holds a whole ubatch beside a decode step of every slot (516 tokens), which is what common_exact_batch_geometry() requires of a start + path = os.environ.get("LLAMA_SERVER_TEST_EXACT_MODEL") + if not path: + pytest.skip("set LLAMA_SERVER_TEST_EXACT_MODEL to a gguf exact concurrency accepts") + server.model_file = path + server.model_hf_repo = server.model_hf_file = None + os.environ["LLAMA_EXACT_CONCURRENCY"] = "1" + os.environ["LLAMA_BATCH_DEBUG"] = "1" + os.environ["LLAMA_ARG_LOG_VERBOSITY"] = "5" + os.environ["LLAMA_ARG_CTX_CHECKPOINTS"] = "0" + _start(n_ctx=16384, n_slots=4, n_batch=2048, n_ubatch=512, fa="on", n_gpu_layer=99, cache_ram=0) + + def prefill(first_token: int) -> list: + mark = len(_log()) + res = server.make_request("POST", "/completion", data={ + "prompt": list(range(first_token, first_token + 3500)), "n_predict": 1, + "cache_prompt": False, "temperature": 0.0, "seed": 42, + }, timeout=600) + assert res.status_code == 200, res.body + assert res.body["timings"]["prompt_n"] == 3500 + text = _log()[mark:] + # a decode step is one token per slot, so the prompt's own ubatches are the wide ones + return [w for w in _ubatch_widths(text) if w > 3] + + alone = prefill(1000) + assert alone, "no ubatch was recorded, LLAMA_BATCH_DEBUG did not reach the log" + + def decoder(i): + server.make_request("POST", "/completion", data={ + "prompt": list(range(20000 + 100 * i, 20000 + 100 * i + 8)), "n_predict": 100000, + "cache_prompt": False, "ignore_eos": True, "temperature": 0.0, "seed": 42, + }, timeout=600) + + threads = [threading.Thread(target=decoder, args=(i,), daemon=True) for i in range(3)] + for t in threads: + t.start() + + try: + for _ in range(600): + slots = server.make_request("GET", "/slots").body + if sum(1 for s in slots if s["is_processing"] and s["n_prompt_tokens"] > 0) >= 3: + break + time.sleep(0.1) + else: + pytest.fail("the decoders never started") + time.sleep(1.0) + + beside = prefill(50000) + finally: + server.stop() + for t in threads: + t.join(30) + + assert beside == alone, f"alone {alone}, beside three decoders {beside}" + + +def test_slots_reports_a_transferring_slot_apart_from_a_parked_one(): + # a copy out still owns its cells and a restore has already taken them back, so a reader counting residency has to keep counting both; only a fully parked slot holds nothing + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "1" + _start_async(n_ctx=256) + + done = [] + t = threading.Thread(target=lambda: done.extend(_complete_all(900))) + t.start() + seen_parked = seen_transferring = False + try: + deadline = time.time() + 120 + while time.time() < deadline and not (seen_parked and seen_transferring): + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + for slot in res.body: + assert not (slot["is_preempted"] and slot["is_transferring"]), slot + if slot["is_transferring"]: + seen_transferring = True + assert slot["n_prompt_tokens"] > 0, "a slot with a copy in flight still holds its cells" + seen_parked = seen_parked or slot["is_preempted"] + finally: + t.join(180) + + assert len(done) == 2, done + for res in done: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] > 0 + assert seen_parked, "no parked slot was ever reported" + assert seen_transferring, "no slot with a copy in flight was ever reported" + + +def _shift_completion(n_predict: int): + """A completion whose context shifts, on a token prompt so its length is exact.""" + return server.make_request("POST", "/completion", data={ + "prompt": [1] + list(range(10, 70)), "n_predict": n_predict, "n_keep": 16, "n_discard": 64, + "ignore_eos": True, "return_tokens": True, "cache_prompt": False, "temperature": 0.0, "seed": 42, + }) + + +def test_a_park_right_after_a_context_shift_does_not_change_the_output(): + # the shift moves the positions and leaves the K transformation for the next decode, so a park in between used to save the new positions with the old K + os.environ["LLAMA_ARG_PREEMPT_ASYNC"] = "0" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + _start(n_ctx=256, n_batch=32, n_ubatch=32, enable_ctx_shift=True, cache_ram=0) + + n_predict = 320 + reference = _shift_completion(n_predict) + assert reference.status_code == 200, reference.body + assert reference.body["timings"]["predicted_n"] == n_predict + n_prompt = reference.body["timings"]["prompt_n"] + server.stop() + + # park on the step the shift lands on: the pool holds n_ctx cells, so the first shift is that many tokens in + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "8192" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = str(256 - n_prompt) + _start() + + parked = _shift_completion(n_predict) + assert parked.status_code == 200, parked.body + assert parked.body["timings"]["predicted_n"] == n_predict + + text = _log() + assert "slot context shift" in text + _assert_recovered(text, "preempted on request") + first_diff = next((i for i, (a, b) in enumerate(zip(reference.body["tokens"], parked.body["tokens"])) if a != b), None) + assert first_diff is None, f"the parked run diverged at token {first_diff}" + + +def test_a_sibling_prompt_with_an_invalid_token_is_refused_before_anything_streams(): + # validated with the others ahead of posting: parked behind a running sibling, it used to fail inside a stream that had already opened 200 + _start(n_ctx=256, n_slots=2, n_batch=256) + + res = server.make_request("POST", "/completion", data={ + "prompt": [[1] * 240, [1] * 240, [9999999]], "n_predict": 4, "temperature": 0.0, "seed": 42, + }) + assert res.status_code == 400, res.body + assert "invalid tokens" in str(res.body) + + text = _log() + _assert_nothing_parked(text) + + +def test_a_recompute_park_bounds_its_draft_by_the_tokens_it_comes_back_with(): + # a recompute park moves the prompt out of the slot, and the draft was bounded by the empty prompt: 2000 tokens and a whole draft could not fit a 2048-cell pool "even alone", failing a request that fits + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "1" + server.spec_type = "ngram-mod" + _start(n_ctx=2048, n_slots=2, n_batch=2048, n_ubatch=512, spec_ngram_mod_n_max=128, spec_ngram_mod_n_min=1) + + prompt = [1] + list(range(10, 110)) * 19 + list(range(10, 109)) + assert len(prompt) == 2000 + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, "n_predict": 24, "ignore_eos": True, "temperature": 0.0, "seed": 42, "cache_prompt": False, + }) + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 24 + + text = _log() + assert "tokens to re-prefill" in text + assert "cannot fit the pool" not in text + assert "Context size has been exceeded" not in text + + +def test_props_says_whether_exact_concurrency_is_running(): + # a client that asked for the mode reads the answer here: a build that ignores the variable starts all the same + _start(n_ctx=256) + res = server.make_request("GET", "/props") + assert res.status_code == 200 + assert res.body["exact_concurrency"] is False + + +def test_props_reports_exact_concurrency_on(): + server.model_file = _mrope_model() + server.model_hf_repo = server.model_hf_file = None + os.environ["LLAMA_EXACT_CONCURRENCY"] = "1" + _start(n_ctx=512, n_slots=2, n_batch=512, n_ubatch=128, fa="on", n_gpu_layer=99) + res = server.make_request("GET", "/props") + assert res.status_code == 200 + assert res.body["exact_concurrency"] is True + + +def test_a_recompute_park_under_exact_concurrency_says_it_is_not_byte_identical(): + # a state that comes back from host memory is the state that left; one rebuilt by re-prefilling differs in the last bits on CUDA, so the mode says so the first time it happens + server.model_file = _mrope_model() + server.model_hf_repo = server.model_hf_file = None + os.environ["LLAMA_EXACT_CONCURRENCY"] = "1" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "4" + _start(n_ctx=512, n_slots=2, n_batch=512, n_ubatch=128, fa="on", n_gpu_layer=99) + + res = _complete(16, "Once upon a time") + assert res.status_code == 200, res.body + text = _log() + assert "tokens to re-prefill" in text + assert "not guaranteed byte-identical" in text + + +def test_a_recompute_park_is_reported_to_the_client_and_to_metrics(): + # a recompute resume re-prefills instead of restoring the saved bytes, so it is not the continuation the parked state would have given: the request, /slots and /metrics all say how often that happened + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + # a sequence this long holds more than the 1 MiB budget, so every park drops its cells + _start(n_ctx=2048, n_batch=2048) + + n_predict = 24 + prompt = _prompt_of(1950, _PROMPT_C) + comments, final = _stream_completion(n_predict, prompt) + + assert "error" not in final, final + assert final["tokens_predicted"] == n_predict + assert final["preempt"]["parks"] >= 2, final["preempt"] + assert final["preempt"]["recomputes"] == final["preempt"]["parks"], final["preempt"] + + # the notice comes right after the resume it belongs to + resumed = [i for i, c in enumerate(comments) if c.startswith(": resumed")] + assert len(resumed) == final["preempt"]["recomputes"], comments + for i in resumed: + assert comments[i + 1].startswith(": recomputed"), comments + + metrics = _metrics() + assert metrics["preempt_recompute_total"] == final["preempt"]["recomputes"] + assert metrics["preempt_recompute_total"] == metrics["n_preempt_total"] + + # a request that is never parked says so + plain = server.make_request("POST", "/completion", data={ + "n_predict": 4, "prompt": _PROMPT_B, "temperature": 0.0, "seed": 42, + }) + assert plain.status_code == 200, plain.body + assert plain.body["preempt"] == {"parks": 0, "recomputes": 0} + + +def test_a_recompute_restore_does_not_count_its_replay_as_prompt(): + # the re-prefill puts back what the park dropped: counted as prompt it would move the prompt/generation boundary, and since n_gen carries across the park the generation time would then cover only the tokens after the last re-prefill + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=2048, n_batch=2048) + + n_prompt = 1950 + n_predict = 24 + res = _complete(n_predict, _prompt_of(n_prompt, _PROMPT_C)) + + assert res.status_code == 200, res.body + assert res.body["preempt"]["recomputes"] >= 1, res.body["preempt"] + + timings = res.body["timings"] + assert timings["prompt_n"] == n_prompt, timings + assert timings["predicted_n"] == n_predict, timings + # every re-prefill happens inside the generation, so the generation holds the longer time of the two + assert timings["predicted_ms"] > timings["prompt_ms"], timings + + +def test_slots_reports_the_recomputes_of_the_current_task(): + # a reader watching the slots sees the same count the request is given at the end + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=2048, n_batch=2048) + + done = [] + prompt = _prompt_of(1950, _PROMPT_C) + t = threading.Thread(target=lambda: done.append(_complete(64, prompt))) + t.start() + + seen = 0 + try: + deadline = time.time() + 180 + while time.time() < deadline and seen == 0 and t.is_alive(): + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + for slot in res.body: + assert slot["n_recompute"] <= slot["n_preempt"] + seen = max(seen, slot["n_recompute"]) + time.sleep(0.02) + finally: + t.join(180) + + assert seen > 0, "no slot ever reported a recompute park" + assert done and done[0].status_code == 200, done + assert done[0].body["preempt"]["recomputes"] >= seen + + +def test_a_swap_park_is_not_reported_as_a_recompute(): + # the same park with room for its bytes keeps the sequence it saved, and the client is told so + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=512) + + comments, final = _stream_completion(24, _PROMPT_A) + + assert final["preempt"]["parks"] >= 1, final["preempt"] + assert final["preempt"]["recomputes"] == 0, final["preempt"] + assert ": resumed" in comments and not any(c.startswith(": recomputed") for c in comments), comments + assert _metrics()["preempt_recompute_total"] == 0 + + +def _stream_responses(n_predict: int, prompt: str) -> dict: + """One streaming /v1/responses request: the data of its response.completed event.""" + url = f"http://{server.server_host}:{server.server_port}/v1/responses" + res = requests.post(url, json={ + "model": "test", "input": prompt, "max_output_tokens": n_predict, + "temperature": 0.0, "stream": True, + }, stream=True, timeout=600) + assert res.status_code == 200, res.text + completed = None + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line.startswith("data: "): + data = json.loads(line[6:]) + if data.get("type") == "response.completed": + completed = data + assert completed is not None, "the stream never reached response.completed" + return completed + + +def test_a_streamed_response_carries_the_preempt_record_where_a_plain_one_does(): + # what a client keeps from a streamed /v1/responses is data["response"], so the record has to be in that object, the same place the non-streamed body carries it + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=512) + + completed = _stream_responses(24, _PROMPT_A) + + assert completed["response"]["preempt"]["parks"] >= 1, completed["response"] + assert completed["response"]["preempt"]["recomputes"] == 0, completed["response"] + + plain = server.make_request("POST", "/v1/responses", data={ + "model": "test", "input": _PROMPT_B, "max_output_tokens": 4, "temperature": 0.0, + }) + assert plain.status_code == 200, plain.body + assert sorted(plain.body["preempt"]) == sorted(completed["response"]["preempt"]) == ["parks", "recomputes"] + + +def test_two_image_chats_that_outgrow_the_parking_budget_both_finish(): + # a media chunk could not be parked by recompute, so with the host budget spent nothing could be parked at all and the pool overflowing ended both chats. The chunk comes back the way it went in: re-encoded off the task, its cells reserved whole + os.environ["LLAMA_MEDIA_MARKER"] = "<__media__>" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "1" + server.model_hf_repo = "ggml-org/tinygemma3-GGUF:Q8_0" + server.model_hf_file = None + server.model_alias = "tinygemma3" + _start(n_ctx=1024, n_slots=2, n_batch=64, n_ubatch=64) + + image = base64.b64encode(requests.get(_IMG_URL, timeout=60).content).decode() + prompt = {"prompt_string": "<__media__>\nWhat is in this image?", "multimodal_data": [image]} + n_predict = 700 + results = parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": prompt, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for _ in range(2) + ]) + + text = _log() + assert "Context size has been exceeded" not in text + assert "failed to process mtmd chunk" not in text + assert "tokens to re-prefill" in text, "no park fell back to recompute" + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == n_predict + + +def test_a_park_during_the_prefill_does_not_count_the_restarted_prompt_twice(): + # a park with no budget for the state drops the cells of a prompt that is still being processed, and the resume starts that prefill again from nothing: what it had counted before the park has to go with the cells, or the request reports more prompt tokens than it has + # the park is made to fail its host allocation, so it drops the cells instead: the state of a half processed prompt is small enough to fit any budget + os.environ["LLAMA_SERVER_PREEMPT_FAIL_SAVE"] = "1" + _start(n_ctx=2048, n_batch=256) + + resident = [] + t = threading.Thread(target=lambda: resident.append(_complete(1900, _PROMPT_A)), daemon=True) + t.start() + + # the pool has to be nearly full before the second prompt starts, so that it is that prefill which runs out of cells + deadline = time.time() + 90 + while t.is_alive() and time.time() < deadline: + slots = server.make_request("GET", "/slots").body + if any(slot.get("n_prompt_tokens", 0) >= 1600 for slot in slots): + break + time.sleep(0.005) + else: + pytest.fail("the resident never grew into the pool") + + n_prompt = 500 + n_predict = 8 + comments, final = _stream_completion(n_predict, _prompt_of(n_prompt, _PROMPT_B)) + t.join(120) + + assert resident and resident[0].status_code == 200, resident + assert "error" not in final, final + assert comments and comments[0] == ": preempted", comments + + # a park that dropped fewer cells than the resume has tokens to put back is a park taken mid-prefill, which is the case this test is about + parks = [(int(cells), int(again)) for cells, again in re.findall( + r"preempted: (\d+) cells dropped .*? (\d+) tokens to re-prefill", _log())] + assert any(0 < cells < again for cells, again in parks), parks + + timings = final["timings"] + assert timings["prompt_n"] == n_prompt, timings + assert timings["cache_n"] == 0, timings + assert timings["predicted_n"] == n_predict, timings 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 000000000000..420f557fa009 --- /dev/null +++ b/tools/server/tests/unit/test_preempt_notify.py @@ -0,0 +1,298 @@ +import json +import os +import re +import tempfile +import time +import threading +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() + +_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" + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + # the server parks only when asked: --preempt-ram defaults to 0, and this suite is about parking + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "8192" + server.temperature = 0.0 + server.seed = 42 + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + for name in ("LLAMA_SERVER_PREEMPT_EVERY", "LLAMA_ARG_PREEMPT_RAM"): + os.environ.pop(name, None) + + +def _start(**kwargs): + for key, value in kwargs.items(): + setattr(server, key, value) + server.start() + + +def _completion_payload(n_predict: int, prompt: str = "Hi how are you", **extra) -> dict: + return {"n_predict": n_predict, "prompt": prompt, "ignore_eos": True, + "temperature": 0.0, "seed": 42, "stream": True, **extra} + + +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} + + +def _log() -> str: + """The server log once its writer thread has stopped growing it: on a loaded host the log lags the response that came from it.""" + deadline = time.time() + 5.0 + last = -1 + while time.time() < deadline: + size = os.path.getsize(server.log_path) + if size == last: + break + last = size + time.sleep(0.1) + return open(server.log_path, errors="replace").read() + + +def _post(path: str, data: dict): + return requests.post(f"http://{server.server_host}:{server.server_port}{path}", json=data, stream=True) + + +def _stream_raw(path: str, data: dict) -> tuple[list[str], list[str]]: + """The SSE lines of one streaming request: (comment lines, data lines).""" + res = _post(path, data) + 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 _stream_all(n_predict: int, prompts, **extra): + return parallel_function_calls([ + (_stream_raw, ("/completion", _completion_payload(n_predict, prompt, **extra))) for prompt in prompts + ]) + + +def _behind_a_resident(payload: dict) -> tuple[int, str, list[str]]: + """Run this request behind a resident holding the pool: its status, its body, and its SSE lines.""" + started = threading.Event() + + def _resident(): + res = _post("/completion", _completion_payload(390, " ".join([_PROMPT_A] * 6))) + assert res.status_code == 200 + for raw in res.iter_lines(): + if raw.decode("utf-8").startswith("data: "): + started.set() + + t = threading.Thread(target=_resident) + t.start() + try: + assert started.wait(60) + res = _post("/completion", payload) + if res.status_code != 200: + return res.status_code, res.text, [] + lines, alive = [], None + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line: + alive = t.is_alive() if alive is None else alive + lines.append(line) + assert alive, "the resident had finished before this request was told anything" + return 200, "", lines + finally: + t.join(120) + + +def _content(datas: list[str]) -> str: + out = "" + for d in datas: + if d == "[DONE]": + break + j = json.loads(d) + out += j.get("content") or "" + for ch in j.get("choices", []) or []: + out += (ch.get("delta") or {}).get("content") or "" + return out + + +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 _notices(comments: list[str]) -> list[str]: + return [c for c in comments if c in (": preempted", ": resumed")] + + +@pytest.mark.parametrize("path,payload", [ + ("/completion", _completion_payload(64)), + ("/v1/chat/completions", _chat_payload(64)), +]) +def test_every_park_in_a_stream_is_announced_paired_with_a_resume_and_changes_nothing(path, payload): + _start(n_ctx=512) + ref_comments, ref_datas = _stream_raw(path, payload) + assert _notices(ref_comments) == [] + assert _content(ref_datas) + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + comments, datas = _stream_raw(path, payload) + seq = _notices(comments) + assert len(seq) >= 12, comments + assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq + assert _content(datas) == _content(ref_datas) + + +def _prefill_payload(path: str, prompt: str, n_predict: int) -> dict: + """The same request on each streaming surface.""" + if path == "/completion": + return {"prompt": prompt, "n_predict": n_predict, "ignore_eos": True, + "temperature": 0.0, "seed": 42, "stream": True} + if path == "/v1/responses": + return {"model": "test", "input": prompt, "max_output_tokens": n_predict, + "temperature": 0.0, "stream": True} + return {"model": "test", "messages": [{"role": "user", "content": prompt}], + "max_tokens": n_predict, "temperature": 0.0, "stream": True} + + +@pytest.mark.parametrize("path", ["/completion", "/v1/chat/completions", "/v1/responses", "/v1/messages"]) +def test_a_park_during_prompt_processing_opens_the_stream_with_the_notice(path): + # a park before the first token is the case a client cannot tell from a stall, so the notice goes out with the response headers rather than waiting for a chunk that is not coming + server.server_slots = True + _start(n_ctx=2048, n_batch=256) + + def _resident(): + res = _post("/completion", _completion_payload(1900, _PROMPT_A)) + for _ in res.iter_lines(): + pass + + t = threading.Thread(target=_resident, daemon=True) + t.start() + + # the pool has to be nearly full before the second prompt starts, so that its prefill is what runs out of cells + # the resident grows by decoding: 1400 tokens took over 12 s on a loaded CI runner + deadline = time.time() + 90 + while t.is_alive() and time.time() < deadline: + slots = requests.get(f"http://{server.server_host}:{server.server_port}/slots").json() + if any(slot.get("n_prompt_tokens", 0) >= 1400 for slot in slots): + break + time.sleep(0.02) + else: + pytest.fail("the resident never grew into the pool") + + res = _post(path, _prefill_payload(path, " ".join([_PROMPT_B] * 31), 8)) + assert res.status_code == 200 + + t0 = time.time() + seen = [] + for raw in res.iter_lines(): + line = raw.decode("utf-8") + if line: + seen.append((time.time() - t0, line)) + t.join(120) + + text = _log() + assert "preempted:" in text, "nothing was parked while the prompt was being processed" + + comments = [(at, line) for at, line in seen if line.startswith(":")] + datas = [(at, line) for at, line in seen if line.startswith("data:")] + + assert comments and comments[0][1] == ": preempted", [line for _, line in seen[:4]] + assert datas, "the request never produced a chunk" + + # sent when the slot was parked, not batched with the chunk that came later + assert comments[0][0] + 0.05 < datas[0][0], [(round(at, 3), line[:24]) for at, line in seen[:4]] + assert any(line == ": resumed" for _, line in comments), [line for _, line in comments[:4]] + + +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) + + status, _, lines = _behind_a_resident(_completion_payload(32, " ".join([_PROMPT_B] * 14))) + assert status == 200 + events = [l for l in lines if l in (": preempted", ": resumed") or l.startswith("data: ")] + assert events[:2] == [": preempted", ": resumed"], events[:3] + assert events[2].startswith("data: "), events[:3] + datas = [l[6:] for l in lines if l.startswith("data: ")] + assert _content(datas) + assert _final(datas)["tokens_predicted"] == 32 + + +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) + + status, body, _ = _behind_a_resident(_completion_payload(16, " ".join([_PROMPT_B] * 80))) + assert status != 200, body + assert not body.lstrip().startswith(":"), body + assert "error" in json.loads(body), body + + +def test_a_rotation_tells_both_streams_and_a_head_parked_past_the_budget_is_kept_alive(): + # --preempt-ram 2 MiB holds one parked state but not a resident's and the head's at once, so that rotation is refused and the head waits parked for longer than the 2 s keepalive + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" + _start(n_slots=3, n_ctx=2048, enable_ctx_shift=True) + + n_predict = 12000 + # the default parked keepalive is 2 s, which is also when a resident is rotated out for the head: + # a park that ends with that rotation could beat its own keepalive. Ask for a 1 s ping instead, so + # any park that outlasts one rotation window is still required to say so + results = _stream_all(n_predict, (_PROMPT_A, _PROMPT_B, _PROMPT_C), sse_ping_interval=1) + n_parked = n_keepalive = 0 + for comments, datas in results: + assert _final(datas)["tokens_predicted"] == n_predict + seq = _notices(comments) + assert seq == [": preempted", ": resumed"] * (len(seq) // 2), seq + n_parked += len(seq) // 2 + n_keepalive += comments.count(": preempt-keepalive") + assert n_parked >= 2, [r[0] for r in results] + assert n_keepalive >= 1, "a parked stream was left silent past its keepalive interval" + + text = _log() + assert "rotated out after" in text + assert "no rotation: --preempt-ram 2 MiB" in text + assert "resumed after" in text + assert "Context size has been exceeded" not in text + + +def test_every_notice_of_a_multi_prompt_stream_names_the_prompt_it_is_about(): + # one request, two prompts: a client reading the shared stream can only tell the notices apart by their index, so index 0 has to be spelled out like any other + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + _start(n_ctx=512) + + comments, datas = _stream_raw("/completion", _completion_payload(32) | {"prompt": [_PROMPT_A, _PROMPT_B]}) + notices = [c for c in comments if c.startswith(": preempted") or c.startswith(": resumed")] + assert notices, comments + assert all(re.fullmatch(r": (preempted|resumed) [01]", c) for c in notices), notices + for index in (0, 1): + assert f": preempted {index}" in notices, notices + assert f": resumed {index}" in notices, notices + + +def test_an_oversized_sibling_prompt_is_errored_before_a_valid_one_is_parked(): + # a request can carry several prompts; a valid one can be parked and its notice opens the stream, so the sibling that does not fit has to be found before any of them is queued + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "8192" + _start(n_ctx=256, n_slots=3, n_batch=512) + + status, body, _ = _behind_a_resident(_completion_payload(8) | {"prompt": [[1] * 120, [1] * 300]}) + assert status == 400, (status, body) + assert not body.lstrip().startswith(":"), body + assert "error" in json.loads(body), body +