server: preemption notices, asynchronous parks and exact concurrency (pin branch on b10869)

.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:

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

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<llama_adapter_lora_ptr> & 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(&params.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(&params.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));
 

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

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;

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

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

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 {