diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d..5bfa4adcdf0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1709,6 +1709,14 @@ 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, -1 - no limit, 0 - disable)", 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( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index de49dac9f63..c99269f9a96 100644 --- a/common/common.h +++ b/common/common.h @@ -614,6 +614,7 @@ struct common_params { int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + int32_t preempt_ram_mib = 8192; // host RAM for parked (preempted) sequences: -1 = no limit, 0 = disable preemption std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b..66940d4fc61 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2617,8 +2617,42 @@ class llama_io_read_host : public llama_io_read_i { ~llama_io_read_host() { // 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++; + } + const size_t tensor_bytes = ggml_nbytes(tensor); + auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + // A fragmented sequence can require thousands of synchronous device + // transfers per layer. For bounded tensors, stage the tensor once and + // preserve every byte belonging to other sequences. Bound scratch RAM + // and leave ordinary contiguous transfers on their original fast path. + if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 && + !ggml_backend_buffer_is_host(buffer)) { + std::vector staging; + try { + staging.resize(tensor_bytes); + } catch (const std::bad_alloc &) { + // Fall back to the individual transfers below. + } + if (!staging.empty()) { + ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes); + 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; + } + } + for (; i < end; ++i) { + const auto & rinfo = rinfos[i]; + ggml_backend_tensor_set(rinfo.tensor, rinfo.ptr, rinfo.offset, rinfo.size); + } } } diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp index d5548afba17..428a9252981 100644 --- a/tests/test-state-restore-fragmented.cpp +++ b/tests/test-state-restore-fragmented.cpp @@ -73,6 +73,14 @@ int main(int argc, char ** argv) { } fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy); + // A fragmented restore may stage a whole device tensor. Check every + // sequence byte-for-byte, including the neighbours that must be preserved. + 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 +104,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/tools/server/README.md b/tools/server/README.md index 93736c3edfa..7b4a0330340 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -164,6 +164,7 @@ 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: 8192, -1 - no limit, 0 - disable)
(env: LLAMA_ARG_PREEMPT_RAM) | | `-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) | @@ -1138,6 +1139,10 @@ 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: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.h b/tools/server/server-common.h index f8ea82ef4cf..f0cf76b8c50 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -467,6 +467,10 @@ struct server_metrics { uint64_t n_decode = 0; uint64_t n_busy_slots = 0; + // [TAG_PREEMPT] slots parked to make room in the unified KV pool, and put back + uint64_t n_preempt = 0; + uint64_t n_resume = 0; + 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 a9edbd7be8b..b18fa4e2f23 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -59,8 +60,37 @@ 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 }; +// [TAG_PREEMPT] server-side request preemption +// +// With --kv-unified the cells are one pool shared by every slot, and each slot believes it +// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves +// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size +// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the +// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and +// continue with the rest". +// +// Nothing is terminated here. The cells of one slot are taken back and given to it again +// later: its sequence is copied to host RAM, its cells are released, and when the pool has +// room the copy goes back and the slot carries on with the same sampler, the same generated +// text and the same open stream. A streaming client sees a pause, not an error. +constexpr int32_t PREEMPT_N_MARGIN = 8; // cells left spare on top of the reservation +constexpr int32_t PREEMPT_N_STARVED = 3; // preemptions after which a slot is protected + +// [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody +// passes a head that does not fit yet: the head keeps the room the pool frees until it fits, so +// its wait is bounded by the slots ahead of it and not by how often a smaller slot can squeeze +// in, grow, and be parked again. Simulated over 60 seeds at eight chats this cuts the longest +// single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often. +// LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest +// parked, and a smaller slot may pass a head that does not fit. +// LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged. +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 + struct server_slot; // forward declaration struct server_batch { @@ -293,6 +323,153 @@ struct server_slot { prompt.clear(); } + // [TAG_PREEMPT] state of a slot whose cells were taken back + // + // Only the KV cells leave. The task, the sampler, the generated text and the position + // the stream has reached stay on the slot, so a resume is a memcpy and not a new + // request: no retokenisation, no replayed prompt, no seam in the output. + slot_state state_before_preempt = SLOT_STATE_IDLE; + std::vector preempt_state_tgt; + std::vector preempt_state_dft; + int32_t n_preempt = 0; // times the CURRENT task has been preempted + int32_t n_ctx_shift = 0; // context shifts the CURRENT task has made: it is at the pool's limit and cycling + int32_t n_preempt_fail = 0; // consecutive failed restores + int64_t t_preempt_us = 0; // when it was parked + bool preempt_rotation_refused = false; // this park has logged a rotation refused for budget + + size_t preempt_state_size() const { + return preempt_state_tgt.size() + preempt_state_dft.size(); + } + + void preempt_state_free() { + preempt_state_tgt.clear(); + preempt_state_tgt.shrink_to_fit(); + preempt_state_dft.clear(); + preempt_state_dft.shrink_to_fit(); + } + + // bytes preempt_save() would need for this slot right now + size_t preempt_state_required() const { + return llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE) + + (ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0); + } + + // copy the sequence out of the cache and release its cells + bool preempt_save() { + const size_t size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + + 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; + } + + // The draft is a prediction, not a result, so it goes with the cells. Preemption + // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already + // holds exactly the tokens the state above covers -- including the rollback done by + // the checkpoint path when a draft was only partially accepted. + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + + // note: prompt.tokens is deliberately kept. It is the mirror of the state just + // copied out, and the resume needs it to know how many cells to ask for. + 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; + } + + // put the sequence back; the slot then continues from the token it was about to decode + bool preempt_restore() { + 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) { + // no room after all: drop the half-written sequence and stay parked + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + if (size_dft > 0 && + llama_state_seq_set_data_ext(ctx_dft, preempt_state_dft.data(), size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) != size_dft) { + mem.seq_rm(id, -1, -1); + n_preempt_fail++; + return false; + } + + preempt_state_free(); + + n_preempt_fail = 0; + + state = state_before_preempt; + + // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks + // that the draft context is where it should be, which the restore above ensures. + // A slot parked while still processing its prompt makes that transition itself + // once the prompt is done. + if (state == SLOT_STATE_GENERATING && can_speculate()) { + common_speculative_begin(spec, id, prompt.tokens.get_text_tokens()); + } + + return true; + } + + // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds for this sequence. + // For a batch that is given up after it was built: the tokens added for this slot + // that were never decoded come off, the sampled token stays in `sampled` and goes into + // the next batch the way it went into this one, and a draft is a prediction that goes + // with them. A chunk that failed to decode left nothing in the cache, so the cache is + // the boundary. + void rewind_to_cache() { + const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1; + + if (n_cached < prompt.n_tokens()) { + prompt.tokens.keep_first(n_cached); + } + + // a prompt whose last chunk was in the batch was marked done when the chunk was + // built; the chunk never ran, so the prompt is not done + if (state == SLOT_STATE_DONE_PROMPT && task && prompt.n_tokens() < task->n_tokens()) { + state = SLOT_STATE_PROCESSING_PROMPT; + } + + spec_draft.clear(); + spec_i_batch.clear(); + spec_ckpt.clear(); + spec_is_replay = false; + + i_batch = -1; + } + std::vector lora; int32_t alora_invocation_start = -1; @@ -351,6 +528,14 @@ struct server_slot { n_predict_max = -1; + // [TAG_PREEMPT] + preempt_state_free(); + state_before_preempt = SLOT_STATE_IDLE; + n_preempt = 0; + n_preempt_fail = 0; + n_ctx_shift = 0; + t_preempt_us = 0; + llama_set_sampler(ctx_tgt, id, nullptr); // clear alora start @@ -503,6 +688,14 @@ struct server_slot { t_last_used = ggml_time_us(); + // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be + // released while parked), so the mirror of them must not outlive them: the next + // task on this slot would otherwise take a prefix match against an empty cache + if (state == SLOT_STATE_PREEMPTED) { + preempt_state_free(); + prompt_clear(); + } + state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough @@ -645,6 +838,8 @@ struct server_slot { {"n_ctx", n_ctx}, {"speculative", can_speculate()}, {"is_processing", is_processing()}, + {"is_preempted", state == SLOT_STATE_PREEMPTED}, + {"n_preempt", n_preempt}, }; const auto & ptask = task ? task : task_prev; @@ -1249,6 +1444,57 @@ struct server_context_impl { } } + { + // read on every load and kept on this context, so a reload after the variable + // changed, or another context loaded in the same process, has an order of its own + preempt_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, for comparing + // policies against each other on the same workload. smallest (the default and the + // shipped one), largest, youngest (the most recent task, as vLLM's scheduler + // preempts), oldest. The leader is kept and the starvation guard applies under all. + 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: the same context reloaded with an attention model after + // a recurrent one gets its preemption back + preempt_recurrent = llama_model_is_recurrent(model_tgt); + + if (preempt_recurrent) { + 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; @@ -2385,11 +2631,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.state == SLOT_STATE_PREEMPTED) { + n_preempted_slots++; + } } SRV_DBG("n_processing_slots = %d\n", n_processing_slots); @@ -2397,6 +2647,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) { @@ -2637,7 +2889,10 @@ struct server_context_impl { void abort_all_slots(const std::string & reason) { for (auto & slot : slots) { - if (slot.is_processing()) { + // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in + // host RAM, not in the cache, and it comes back when there is room, the same as + // in the decode error sweep + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, reason, ERROR_TYPE_SERVER); slot.release(); } @@ -2674,6 +2929,632 @@ struct server_context_impl { }; #endif + // + // [TAG_PREEMPT] server-side request preemption + // + + + // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens, + // whether or not the pool is under pressure. It exists to answer the only question that + // matters about a resume: with one request on an idle server the batch has the same + // shape at every step, so a preempted continuation that is not byte-identical to an + // uninterrupted one is the preemption's fault and nothing else's. + int32_t preempt_test_every = 0; + std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model + + // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so + // the KV-full retry ladder and its last resort are the only thing between a full pool + // and the context error + bool preempt_planner_off = false; + + // LLAMA_SERVER_PREEMPT_RESUME: head (the default) puts parked slots back in the order they + // were parked and only the first until it fits; pass lets a smaller slot pass a head + // that does not fit. Read at load, per context. + bool preempt_resume_head = true; + + // a recurrent cache holds one state per sequence whatever its length: no cell pool, + // nothing to run out of, and the token count the planner measures says nothing about + // it. Preemption is off for those models; a hybrid keeps its attention cache and stays on. + bool preempt_recurrent = false; + + // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop + bool preempt_batch_abandoned = false; + + int32_t preempt_n_spec_max() const { + return spec ? std::max(0, common_speculative_n_max(¶ms_base.speculative)) : 0; + } + + // draft tokens this slot's next step can actually carry: the configured maximum, cut to + // what its context and its prediction budget leave, the way get_n_draft_max() cuts it + 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; + } + + res = std::min(res, slot.n_ctx - slot.prompt.n_tokens() - 2); + + if (slot.n_remaining() > 0) { + res = std::min(res, slot.n_remaining() - 1); + } + + return std::max(0, res); + } + + // host RAM the parked sequences hold right now + size_t preempt_ram_used() const { + size_t res = 0; + + for (const auto & slot : slots) { + res += slot.preempt_state_size(); + } + + return res; + } + + // whether parking this slot stays under --preempt-ram + bool preempt_fits_budget(const server_slot & slot) const { + if (params_base.preempt_ram_mib < 0) { + return true; + } + + const size_t budget = (size_t) params_base.preempt_ram_mib * 1024 * 1024; + + return preempt_ram_used() + slot.preempt_state_required() <= budget; + } + + // cells of the mirrored prompt that a started slot's request keeps, by the rule the batch + // builder applies when it takes the slot: nothing when the request does not cache its + // prompt, otherwise the prefix the two share, cut short of an aLoRA invocation + size_t preempt_n_keep(const server_slot & slot) const { + if (!slot.task->params.cache_prompt) { + return 0; + } + + 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; + } + + // cells of the slot's that its next step keeps: a slot just given a task still mirrors + // the previous request's prompt until the batch builder keeps what preempt_n_keep() + // says and drops the rest, so what it holds, and what it is about to ask for, both + // count from that + int32_t preempt_n_retained(const server_slot & slot) const { + if (slot.state == SLOT_STATE_STARTED && slot.task) { + return (int32_t) preempt_n_keep(slot); + } + + return slot.prompt.n_tokens(); + } + + // cells the slot will ask for on its next step once it is back in the pool + int32_t preempt_n_need(const server_slot & slot) const { + int32_t res = preempt_n_retained(slot); + + if (slot.state_before_preempt == SLOT_STATE_GENERATING) { + res += 1 + preempt_n_spec(slot); + } else { + const int32_t n_left = slot.task ? slot.task->n_tokens() - res : 0; + + res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left)); + } + + return res; + } + + // Cells the pool is holding right now. A released slot keeps its prompt in the cache + // for the next request to reuse as a prefix, so idle slots count too: the first version + // of this counted only the running ones, decided a pool holding 8185 cached cells was + // empty, and every resume failed against a cache that was actually full. + int32_t preempt_kv_used() const { + int32_t res = 0; + + // n_cmpl > 1: the parent and its children share the prompt's cells through seq_cp, so + // the prompt is charged once per family, to whichever resident member comes first; + // the others are charged only what they generated on top of it + 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 + } + + // a child waiting for its parent's prompt does not share anything yet: until + // copy_state_to() runs it still holds whatever the previous request left in its + // cells, so it is charged that on its own, outside the family + if (slot.state == SLOT_STATE_WAIT_OTHER) { + res += 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 += std::max(0, slot.prompt.n_tokens() - slot.task->n_tokens()); + continue; + } + + charged.push_back(family); + } + + // what the pool holds now, the previous request's prompt included for a slot just + // given a task: the batch builder trims that to the prefix the two share, but + // not until the slot is built into a batch, and with continuous batching off that + // can be a long time behind a running generation. Measured by the prefix, a + // restore was found to fit and attempted against cells still occupied. Under + // pressure the planner trims such slots itself, see preempt_normalize_started_all() + res += slot.prompt.n_tokens(); + } + + return res; + } + + // cells those slots are about to ask for on the next decode + int32_t preempt_kv_reserve() const { + const int32_t n_batch = llama_n_batch(ctx_tgt); + + int32_t res = 0; + int32_t res_pmt = 0; + + for (const auto & slot : slots) { + switch (slot.state) { + case SLOT_STATE_GENERATING: + case SLOT_STATE_DONE_PROMPT: + { + res += 1 + preempt_n_spec(slot); + } break; + case SLOT_STATE_STARTED: + case SLOT_STATE_PROCESSING_PROMPT: + { + // from the prefix a started slot keeps, not from the prompt it still + // mirrors: measured by the mirror, a request shorter than the last one + // reserved one cell for a chunk of hundreds + const int32_t n_left = slot.task ? slot.task->n_tokens() - preempt_n_retained(slot) : 0; + + res_pmt += std::max(1, std::min(n_batch, n_left)); + } break; + default: + break; + } + } + + // one batch is all the prompt slots get between them, however many are waiting + return res + std::min(res_pmt, n_batch); + } + + // Keep the slot that is furthest along -- it is the closest to finishing and to giving + // its cells back -- and among the rest prefer one that has not been preempted + // PREEMPT_N_STARVED times already, then the smallest. + // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt + // until the batch builder keeps the prefix the two share and drops the rest (see the + // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out, + // charged and sized by the old prompt, and a short unrelated request could exceed the + // budget or stay parked for room it will never use. Keeping only the shared prefix now + // is what the batch builder does anyway; the chunk reuse it can add on top is given up + // for a slot the planner has to touch, which is rare. + // every started slot, when the pool is short: true when any of them gave cells up + bool preempt_normalize_started_all() { + bool res = false; + + for (auto & slot : slots) { + if (slot.state != SLOT_STATE_STARTED || !slot.task) { + continue; + } + + const int32_t before = slot.prompt.n_tokens(); + + preempt_normalize_started(slot); + + 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 (a recurrent state without rollback + // room for the stale suffix) aborts on a partial removal; for it the whole stale + // sequence goes, and the prompt is processed from the start on resume, as it would be + // without a usable checkpoint + 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); + } + + server_slot * preempt_pick_victim() { + server_slot * leader = nullptr; + int32_t n_running = 0; + + // a slot just given a task is measured by the prefix it keeps, not by the previous + // request's prompt it still mirrors: measured by the mirror, a short request over a + // large stale cache would be the never-parked leader while the longest live + // conversation was parked in its place + for (auto & slot : slots) { + preempt_normalize_started(slot); + } + + for (auto & slot : slots) { + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { + n_running++; + + if (!leader || slot.prompt.n_tokens() > leader->prompt.n_tokens()) { + leader = &slot; + } + } + } + + if (n_running < 2) { + // a single conversation that does not fit the pool on its own is a real context + // overflow and not a scheduling problem - leave it to the existing error path + return nullptr; + } + + server_slot * victim = nullptr; + + for (auto & slot : slots) { + // Before the batch is built every one of these is at a token boundary: a + // generating slot between two sampled tokens, a prompt-processing slot between + // two chunks of its prompt, a started slot with only a cached prefix (or + // nothing) in the pool. A slot holding no cells is still worth parking - it + // is about to ask for a whole batch of them. + 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 + } + + if (!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; + } + + // is a the better victim of the two? the smallest slot under the shipped policy: it + // gives up the least work and its restore is the cheapest (see the PR's simulation); + // the other choices exist for the comparison runs behind LLAMA_SERVER_PREEMPT_POLICY + bool preempt_better_victim(const server_slot & a, const server_slot & b) const { + if (preempt_test_policy == "largest") { + return a.prompt.n_tokens() > b.prompt.n_tokens(); + } + + 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(); + } + + // called once per update_slots(), before the batch is built: at that point every slot is + // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no + // draft is in flight, so a slot can be removed from the picture without unpicking a + // half-decoded batch + 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 + } + + if (params_base.preempt_ram_mib == 0 || preempt_recurrent) { + return; // --preempt-ram 0, or a recurrent cache: the KV-full retry ladder, as before + } + + const int32_t n_cells = n_ctx; + + // Put back what fits, in the order preempt_resume_head describes: by default + // the slot parked longest, and only that one until it fits; under + // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked + // longest, and a smaller slot may pass a head that does not fit. + const bool head_of_line = preempt_resume_head; + + for (;;) { + 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; + + // A parked slot whose sequence plus its next step would not fit an empty pool can + // never be restored, and would otherwise sit at the head of the line for ever + // without a restore ever being attempted: a prompt within n_ctx that was parked + // before it took any cells, but too close to n_ctx to leave room for its first + // batch. That is the single-conversation overflow the KV-full path reports, so + // report it the same way and rescan the line without it. + { + server_slot * impossible = nullptr; + + for (auto * slot : parked) { + if (preempt_n_need(*slot) > n_cells) { + impossible = slot; + break; + } + } + + if (impossible) { + SLT_WRN(*impossible, "parked sequence of %d tokens cannot fit the pool of %d cells even alone, failing it\n", + preempt_n_need(*impossible), n_cells); + send_error(*impossible, "Context size has been exceeded."); + impossible->release(); + continue; + } + } + + // Room for the sequence AND for the next step of everything already running, + // so that a resume cannot immediately trigger the preemption of someone else. + // The margin is headroom for the others; with nothing resident there is nobody + // to keep it for, so a sequence that fits the pool exactly is let back in. + // A cached prompt on an idle slot is worth less than a conversation waiting to + // continue, so give those cells up first - same call the KV-full path makes. + for (;;) { + const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); + const int32_t margin = occupied == 0 ? 0 : PREEMPT_N_MARGIN; + + for (auto * slot : parked) { + if (occupied + preempt_n_need(*slot) + margin <= n_cells) { + best = slot; + break; + } + } + + if (best) { + break; + } + + // a slot just given a task still holds the previous request's prompt until + // the batch builder trims it; trimmed here instead, the cells it will not + // keep are counted out and a parked slot that fits without them comes back + if (preempt_normalize_started_all()) { + continue; + } + + if (!try_clear_idle_slots()) { + break; + } + } + + // Nothing fits. A resident that has reached the pool's limit and is cycling + // through context shifts holds the room for as long as it likes to generate, + // and the head behind it would wait for ever. After the head has waited its + // turn, that resident is parked in its place: it is at a token boundary like + // any other park, and when it comes back it is the one waiting, so the two + // take turns instead of one taking everything. + if (!best) { + server_slot * head = parked.front(); + + if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) { + // the resident whose cells let the head in, the smallest of those; failing + // one that does so alone, the largest, since it makes the most room. Taking + // the first shifting resident in slot order could park one too small to + // matter, spend the park budget on it, and leave the head waiting anyway. + const int32_t occupied = preempt_kv_used() + preempt_kv_reserve(); + const int32_t need = preempt_n_need(*head) + PREEMPT_N_MARGIN; + + server_slot * pick = nullptr; + bool pick_enough = false; + bool budget_refused = false; + + 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, and the cap is a cap on what is held. A budget that + // holds one sequence but not two does not rotate, and the head waits + // for a resident to finish, which is said once per park below. + if (!preempt_fits_budget(slot)) { + budget_refused = true; + continue; + } + + const bool enough = occupied - 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; + } + } + + 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\n", + params_base.preempt_ram_mib); + } + + if (pick && pick->preempt_save()) { + server_slot & slot = *pick; + + metrics.n_preempt++; + + SLT_WRN(slot, "rotated out after %d context shifts: %d cells released, %.1f MiB parked, a head parked %.1f s takes its turn%s, preemptions %d\n", + slot.n_ctx_shift, slot.prompt.n_tokens(), + slot.preempt_state_size() / (1024.0 * 1024.0), + (ggml_time_us() - head->t_preempt_us) / 1e6, + pick_enough ? "" : " (not enough room by itself)", + slot.n_preempt); + + best = head; // re-examined by the loop, which sees the room it just got + } + } + + if (best) { + continue; + } + + break; + } + + const int64_t t_start = ggml_time_us(); + + if (!best->preempt_restore()) { + // update_slots() runs in a tight loop while tasks are pending, so a counter + // alone burns its whole budget in a couple of milliseconds. Give up only on + // a slot that has been failing for a while, and keep the log quiet. + if (best->n_preempt_fail % 64 == 1) { + SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n", + best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6); + } + + 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; + } + + metrics.n_resume++; + + SLT_WRN(*best, "resumed after %.2f s: %d tokens back in the cache in %.2f ms, kv %d/%d, preemptions %d\n", + (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); + } + + // forced preemption, for the determinism test only + if (preempt_test_every > 0) { + for (auto & slot : slots) { + if (slot.state == SLOT_STATE_GENERATING && + (int32_t) slot.stats.n_gen >= (slot.n_preempt + 1) * preempt_test_every && + preempt_fits_budget(slot) && + slot.preempt_save()) { + metrics.n_preempt++; + + SLT_WRN(slot, "preempted on request after %d generated tokens, %.1f MiB parked\n", + (int32_t) slot.stats.n_gen, slot.preempt_state_size() / (1024.0 * 1024.0)); + } + } + } + + if (preempt_planner_off) { + return; // test knob: leave the pool to the retry ladder and its last resort + } + + // and take cells back until the next decode fits + for (;;) { + const int32_t n_used = preempt_kv_used() + preempt_kv_reserve(); + + if (n_used + PREEMPT_N_MARGIN <= n_cells) { + break; + } + + // a prompt cached on an idle slot is the cheapest thing in the pool to give up + if (try_clear_idle_slots()) { + continue; + } + + server_slot * victim = preempt_pick_victim(); + + 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 (!victim->preempt_save()) { + break; // could not park it; the existing retry ladder is still behind us + } + + metrics.n_preempt++; + + 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 update_slots() { #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; @@ -2716,6 +3597,13 @@ struct server_context_impl { } try { + // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured + // after any context shift. Inside the guard with the rest of the step: a shift + // rebuilds a slot's tokens and a park allocates, and either can throw, which the + // slots are told about rather than the loop ending on an uncaught exception + pre_decode_shift(); + update_preemption(); + scoped_timer t(t_pre_decode, n_pre_decode); pre_decode(); batch.render(); @@ -2763,6 +3651,13 @@ struct server_context_impl { llama_synchronize(ctx_tgt); #endif + if (preempt_batch_abandoned) { + // [TAG_PREEMPT] the rest of this batch was never decoded and the slots no + // longer describe it; the next pass builds a new one + preempt_batch_abandoned = false; + break; + } + if (ok) { // move the head of the batch forward with the number of tokens we just processed off_next = off + n_tokens; @@ -2790,9 +3685,11 @@ 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 + // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift, + // not with the cells the shift is about to give back + void pre_decode_shift() { iterate(slots, [&](server_slot & slot) { if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) { if (!params_base.ctx_shift) { @@ -2832,6 +3729,8 @@ 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++; + 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); @@ -2854,7 +3753,9 @@ struct server_context_impl { slot.truncated = true; } }); + } + void pre_decode() { // start populating the batch for this iteration batch.clear(); @@ -3000,7 +3901,9 @@ 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; it takes no part in this pass until it is restored + if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) { return; } @@ -3526,6 +4429,102 @@ struct server_context_impl { } } + // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is + // the context error for every slot in the batch. With a park budget the batch is given + // up instead: every resident slot is rewound to the token boundary the cache is at (a + // batch is applied one chunk at a time, and the chunk that failed left nothing behind), + // the smallest are parked until the planner's own bound holds again, and the next + // update_slots() rebuilds the batch from the survivors. The planner brings the parked + // ones back as cells free up. A multimodal prompt has no boundary the cache can name, + // so it keeps the old path. + 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); + } + + 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) { + continue; + } + + if (slot.prompt.tokens.has_mtmd) { + 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.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; + } + + server_slot * victim = preempt_pick_victim(); + + if (!victim) { + break; + } + + const int32_t n_tokens = victim->prompt.n_tokens(); + const int64_t t_start = ggml_time_us(); + + if (!victim->preempt_save()) { + break; + } + + metrics.n_preempt++; + n_parked++; + + SLT_WRN(*victim, "preempted as a last resort: %d cells released in %.2f ms, %.1f MiB parked, kv %d/%d (wanted %d), preemptions %d\n", + n_tokens, + (ggml_time_us() - t_start) / 1e3, + 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; + } + + // [TAG_PREEMPT] whether a slot in the batch has its sampled token and a draft in it + bool batch_has_spec_groups() const { + for (const auto & slot : slots) { + if (!slot.spec_i_batch.empty()) { + 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) { @@ -3573,7 +4572,27 @@ struct server_context_impl { { std::string err; + // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have + // to stay in one view: a narrower view splits the group and the verify step + // throws for the slot whose tokens straddle it. Halving is no help there, so + // after the idle slots the ladder goes to its last resort straight away. With + // no budget to park into the ladder is what it always was. + 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) { + // [TAG_PREEMPT] park instead of ending everyone, when there is a budget to park into + if (preempt_last_resort(off)) { + preempt_batch_abandoned = true; + return true; + } + // 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."; @@ -3594,7 +4613,9 @@ 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()) { + // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the + // cache; it is not part of this failure and comes back when there is room + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { send_error(slot, err); slot.release(); @@ -3943,7 +4964,8 @@ 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()) { + // [TAG_PREEMPT] a parked slot is processing but took no part in this decode + if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) { metrics.n_busy_slots++; } metrics.n_tokens_max = std::max(metrics.n_tokens_max, (uint64_t) slot.prompt.n_tokens()); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313ce..9afe3c7f06a 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1562,6 +1562,14 @@ 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 }, }; @@ -1586,6 +1594,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 9c99143f8e1..00734924bc6 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -494,6 +494,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 00000000000..0c8f5dc5f29 --- /dev/null +++ b/tools/server/tests/unit/test_preempt.py @@ -0,0 +1,549 @@ +import os +import time +import tempfile +import pytest +from utils import * + +# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked +# (its sequence copied to host RAM, its cells released) instead of every slot being +# terminated. Both tests need more than one slot and --kv-unified, which is the only +# configuration where one slot can take another one's cells. + +server = ServerPreset.tinyllama2() + + +class LogReader: + def __init__(self, path): + self.path = path + self.pos = 0 + + def drain(self): + with open(self.path) as f: + f.seek(self.pos) + content = f.read() + self.pos = f.tell() + return content + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.tinyllama2() + server.n_slots = 2 + server.kv_unified = True + server.server_slots = True + server.temperature = 0.0 + server.seed = 42 + fd, server.log_path = tempfile.mkstemp(suffix=".log") + os.close(fd) + yield + os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None) + os.environ.pop("LLAMA_SERVER_PREEMPT_PLANNER", None) + os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None) + + +def _complete(n_predict: int, prompt: str = "Hi how are you"): + res = server.make_request("POST", "/completion", data={ + "n_predict": n_predict, + "prompt": prompt, + "ignore_eos": True, + "return_tokens": True, + "temperature": 0.0, + "seed": 42, + }) + return res + + +def test_forced_preemption_does_not_change_the_output(): + # Park and restore the only running slot every 8 tokens. With one request the batch + # has the same shape at every step whether or not the slot was parked in between, so + # any difference in the output is the preemption's fault and nothing else's. + global server + server.n_ctx = 512 + server.start() + reference = _complete(64) + assert reference.status_code == 200 + assert reference.body["timings"]["predicted_n"] == 64 + server.stop() + + os.environ["LLAMA_SERVER_PREEMPT_EVERY"] = "8" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_EVERY = 8" in log.drain() + + preempted = _complete(64) + assert preempted.status_code == 200 + assert preempted.body["timings"]["predicted_n"] == 64 + + text = log.drain() + assert text.count("preempted on request") >= 6 + assert text.count("resumed after") >= 6 + + assert preempted.body["content"] == reference.body["content"] + assert preempted.body["tokens"] == reference.body["tokens"] + + +def test_two_slots_that_overflow_the_pool_together_both_finish(): + # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well + # under 256. Together they do not, 336 against 256. Without preemption the retry + # ladder ends with "Context size has been exceeded" on every processing slot; with it + # the smaller slot is parked until the leader finishes and its cells are purged, and + # then it resumes from the token it was parked on. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + + + +_WORDS = ( + "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor " + "incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud " + "exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure " + "dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. " + "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt " + "mollit anim id est laborum. " +) * 4 + + +def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]: + """A prompt whose token count is in [n_tokens - 12, n_tokens], measured on the server.""" + words = (salt + " " + _WORDS).split() + while words: + text = " ".join(words) + res = server.make_request("POST", "/tokenize", data={"content": text}) + assert res.status_code == 200 + n = len(res.body["tokens"]) + if n <= n_tokens: + assert n >= n_tokens - 12, f"could not land near {n_tokens} tokens, got {n}" + return text, n + # about four tokens per word on this model's vocabulary + words = words[: len(words) - max(1, (n - n_tokens) // 8)] + raise AssertionError("empty prompt") + + +def test_two_prompts_that_overflow_the_pool_together_both_finish(): + # Neither slot ever generates before the pool is full: both are still processing their + # prompts. A prompt-processing slot is between two chunks of its prompt, which is as + # clean a boundary as the one between two sampled tokens, so it is parked the same way. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + prompt_a, n_a = _prompt_of_about(150, "Alpha") + prompt_b, n_b = _prompt_of_about(150, "Bravo") + n_predict = 16 + assert n_a + n_predict <= 256 and n_b + n_predict <= 256 + assert n_a + n_b + 2 * n_predict > 256 + + results = parallel_function_calls([ + (_complete, (n_predict, prompt_a)), + (_complete, (n_predict, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert len(res.body["tokens"]) == n_predict + + +def test_a_generating_slot_and_a_large_prompt_both_finish(): + # One slot is generating a long answer to a short prompt when a large prompt arrives + # beside it. Together they need far more than the pool has. The prompt is admitted + # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish. + # This model produces a thousand tokens a second, so the second request is sent right + # behind the first rather than after a delay: its prompt takes several batches to + # process, which is enough for the two to overlap however fast the first one runs. + global server + server.n_ctx = 256 + server.start() + log = LogReader(server.log_path) + + prompt_b, n_b = _prompt_of_about(150, "Charlie") + # b lives long enough for the two to collide: the first run of this used 16 tokens + # and b was finished and purged before a had grown into it + n_predict_a = 230 + n_predict_b = 90 + assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 + assert 8 + n_predict_a + n_b + n_predict_b > 256 + + def _late(n_predict, prompt): + time.sleep(0.02) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_a, "Hi how are you")), + (_late, (n_predict_b, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" in text + + assert results[0].status_code == 200 + assert results[0].body["timings"]["predicted_n"] == n_predict_a + assert results[1].status_code == 200 + assert results[1].body["timings"]["predicted_n"] == n_predict_b + + +def test_preempt_ram_zero_disables_preemption(): + # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the + # KV-full path ends the requests the way it always did. + global server + server.n_ctx = 256 + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0" + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "preempted:" not in text + assert "Context size has been exceeded" in text + assert any(res.status_code != 200 for res in results) + + +def test_metrics_and_slots_report_the_parked_state(): + # A client that wants to tell a parked chat from a slow one reads /slots, and an + # operator reads /metrics. Both must show the preemption happening, and the counters + # must survive the requests finishing. + global server + server.n_ctx = 256 + server.server_metrics = True + server.start() + + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + for slot in res.body: + assert slot["is_preempted"] is False + assert slot["n_preempt"] == 0 + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + for res in results: + assert res.status_code == 200 + + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + metrics = {} + for line in res.body.splitlines(): + if line.startswith("llamacpp:"): + name, value = line.split(" ", 1) + metrics[name[len("llamacpp:"):]] = float(value) + 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 + + res = server.make_request("GET", "/slots") + assert res.status_code == 200 + assert sum(slot["n_preempt"] for slot in res.body) == 0, "n_preempt is per task and resets with the slot" + + +def test_two_prompts_near_the_context_size_both_complete(): + # Two prompts that each fit the context alone but not together. The second one is + # parked before it takes any cells, and it is close enough to n_ctx that its sequence + # plus its first batch would not leave the usual scheduling margin. It must still be + # restored once the first one finishes: with nothing resident there is nobody to keep + # the margin for. Before the fix it was parked for ever, with no restore ever tried. + global server + server.n_ctx = 256 + # the whole prompt in one batch, so the parked slot's first step is the whole prompt + server.n_batch = 256 + server.start() + log = LogReader(server.log_path) + + # sized in tokens, not words: the prompt is the token ids of a short sentence repeated + base = server.make_request("POST", "/tokenize", data={"content": "Once upon a time there was a little girl"}).body["tokens"] + long_prompt = (base * 64)[:250] + n_predict = 4 + together = parallel_function_calls([(_complete, (n_predict, long_prompt)) for _ in range(2)]) + + text = log.drain() + assert "cannot fit the pool" not in text + + for res in together: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + + +def test_the_last_resort_parks_instead_of_ending_everyone(): + # With the planner off nothing is parked ahead of the decode, so two generations that + # fit alone but not together fill the pool until a single token finds no cell. That + # is where upstream ends every slot with the context error. Instead the batch is + # given up, the smaller slot is parked, the larger one finishes, and the parked one + # comes back and finishes too. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + server.start() + log = LogReader(server.log_path) + assert "LLAMA_SERVER_PREEMPT_PLANNER = off" in log.drain() + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted:" not in text, "the planner was off, nothing may be parked ahead of the decode" + assert "preempted as a last resort" in text + assert "last resort: batch given up" in text + assert "resumed after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + assert res.body["truncated"] is False + assert len(res.body["tokens"]) == n_predict + + +def test_the_last_resort_works_with_an_unlimited_budget(): + # --preempt-ram -1 is the documented unlimited setting; it must enable the last resort + # the same as any positive budget does + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "-1" + server.start() + log = LogReader(server.log_path) + + n_predict = 160 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted as a last resort" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + + +def test_the_last_resort_rewinds_a_prompt_in_flight(): + # Same, with a prompt being processed when the pool runs out: the chunk that failed + # is taken back off the slot's tokens and processed again after the resume, so the + # prompt is neither skipped nor fed twice. The prompt is far longer than a batch, so + # the failing chunk is a chunk of it, not its last token. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + server.start() + log = LogReader(server.log_path) + + prompt_b, n_b = _prompt_of_about(150, "Charlie") + n_predict_a = 230 + n_predict_b = 90 + assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256 + assert 8 + n_predict_a + n_b + n_predict_b > 256 + + def _late(n_predict, prompt): + time.sleep(0.02) + return _complete(n_predict, prompt) + + results = parallel_function_calls([ + (_complete, (n_predict_a, "Hi how are you")), + (_late, (n_predict_b, prompt_b)), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "preempted as a last resort" in text + + assert results[0].status_code == 200 + assert results[0].body["timings"]["predicted_n"] == n_predict_a + assert results[1].status_code == 200 + assert results[1].body["timings"]["predicted_n"] == n_predict_b + # the chunk that was in the batch given up is processed once, after the rewind, and + # the count is the prompt plus the BOS the server adds + assert results[1].body["timings"]["prompt_n"] == n_b + 1 + + +def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head(): + # Two generations that each outgrow the pool on their own, with context shift on. The + # resident reaches the limit, shifts, keeps about half the pool and would keep going + # for as long as it has tokens to make, while the parked one never fits beside it. + # After the head has waited its turn the resident is parked in its place, and the two + # take turns until both finish. Long enough that the resident is still going when the + # head's turn comes: this model makes a couple of thousand tokens a second. + global server + server.n_ctx = 256 + server.enable_ctx_shift = True + server.start() + log = LogReader(server.log_path) + + n_predict = 12000 + results = parallel_function_calls([ + (_complete, (n_predict, "Once upon a time there was a brave knight who")), + (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")), + ]) + + text = log.drain() + assert "Context size has been exceeded" not in text + assert "slot context shift" in text + assert "rotated out after" in text + + for res in results: + assert res.status_code == 200 + assert res.body["timings"]["predicted_n"] == n_predict + + +def test_the_rotation_parks_a_resident_that_lets_the_head_in(): + # Three generations with no end in a 256-cell pool with context shift on: two residents + # cycle through shifts while the third waits parked. Every rotation must let the head + # in, so all three keep finishing their tokens and no stream ends short. + global server + server.n_slots = 3 + server.n_ctx = 384 + server.enable_ctx_shift = True + server.start() + n_predict = 9000 + prompts = [ + "Once upon a time there was a brave knight who", + "The quick brown fox jumps over the lazy dog and", + "In a small village by the sea there lived a fisherman who", + ] + results = parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for p in prompts + ]) + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == n_predict + text = open(server.log_path).read() + assert "rotated out after" in text + assert "Context size has been exceeded" not in text + + +def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives(): + # One request asking for two completions is one conversation in two slots: a parent + # and a child sharing the prompt. When the two together do not fit the pool there is + # nobody else to park, since the family is charged once and a member of it is not a + # victim for the other, so the request gets the context error it would get alone, and + # the server carries on serving. + global server + server.n_ctx = 256 + os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off" + server.start() + log = LogReader(server.log_path) + + res = server.make_request("POST", "/completion", data={ + "n_predict": 160, + "n_cmpl": 2, + "prompt": "Once upon a time there was a brave knight who", + "ignore_eos": True, + "return_tokens": True, + "temperature": 0.0, + "seed": 42, + }) + assert res.status_code == 500 + assert "Context size has been exceeded" in res.body["error"]["message"] + + text = log.drain() + assert "preempted as a last resort" not in text, "a family alone in the pool has no victim" + assert "GGML_ASSERT" not in text + + after = _complete(8) + assert after.status_code == 200 + assert after.body["timings"]["predicted_n"] == 8 + + +def test_a_budget_that_holds_one_sequence_does_not_rotate_and_the_head_resumes_when_a_resident_finishes(): + # Three generations with no end in a pool one of them fills, with context shift on, + # under a --preempt-ram that holds the two parked heads but not a head and the resident + # at once. The resident is parked before the head is restored and freed, so a rotation + # holds both states together: under this budget the first one asked for is refused and + # said so, and the heads come back when the resident finishes instead. Every stream + # still finishes its tokens and nothing gets the context error. + global server + server.n_slots = 3 + server.n_ctx = 2048 + server.enable_ctx_shift = True + os.environ["LLAMA_ARG_PREEMPT_RAM"] = "2" + server.start() + # long enough that the resident is still cycling through shifts two seconds after the + # heads were parked, which is when a rotation is first asked for: at 6000 this model + # finished in under three seconds on a fast host and nothing was ever refused + n_predict = 12000 + prompts = [ + "Once upon a time there was a brave knight who", + "The quick brown fox jumps over the lazy dog and", + "In a small village by the sea there lived a fisherman who", + ] + results = parallel_function_calls([ + (server.make_request, ("POST", "/completion", { + "prompt": p, "n_predict": n_predict, "ignore_eos": True, "temperature": 0.0, "seed": 42, + })) for p in prompts + ]) + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == n_predict + text = open(server.log_path).read() + 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_a_recurrent_model_is_served_without_preemption(): + # A recurrent cache holds one state per sequence whatever its length, so the token + # count the planner measures says nothing about it: preemption is off for such a + # model, said so at load, and the forced-park knob parks nothing. + global server + path = os.environ.get("LLAMA_SERVER_TEST_RECURRENT_MODEL") + if path: + server.model_file = path + else: + server.model_file = None + 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) + results = parallel_function_calls([ + (_complete, (64, "Once upon a time")), + (_complete, (64, "The quick brown fox")), + ]) + for res in results: + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 64 + text = open(server.log_path).read() + assert "preemption: off, the recurrent cache holds one state per sequence" in text + assert "preempted" not in text + assert "Context size has been exceeded" not in text