server: a failed prompt fill leaves the batch and its prompt clean (Windows bad_alloc abort)

.github/workflows/server-sanitize.yml

@@ -79,6 +79,7 @@ jobs:
           cmake -B build \
             -DLLAMA_BUILD_BORINGSSL=ON \
             -DGGML_SCHED_NO_REALLOC=ON \
+            -DLLAMA_SERVER_TEST_FAULTS=ON \
             -DGGML_SANITIZE_ADDRESS=${{ matrix.sanitizer == 'ADDRESS' }} \
             -DGGML_SANITIZE_THREAD=${{ matrix.sanitizer == 'THREAD' }} \
             -DGGML_SANITIZE_UNDEFINED=${{ matrix.sanitizer == 'UNDEFINED' }} \

.github/workflows/server-self-hosted.yml

@@ -55,7 +55,7 @@ jobs:
       - name: Build
         id: cmake_build
         run: |
-          cmake -B build -DGGML_SCHED_NO_REALLOC=ON
+          cmake -B build -DGGML_SCHED_NO_REALLOC=ON -DLLAMA_SERVER_TEST_FAULTS=ON
           cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) --target llama-server
 
       - name: Python setup
@@ -115,7 +115,7 @@ jobs:
       - name: Build
         id: cmake_build
         run: |
-          cmake -B build -DGGML_CUDA=ON -DGGML_SCHED_NO_REALLOC=ON
+          cmake -B build -DGGML_CUDA=ON -DGGML_SCHED_NO_REALLOC=ON -DLLAMA_SERVER_TEST_FAULTS=ON
           cmake --build build --config Release -j $(nproc) --target llama-server
 
       - name: Python setup
@@ -202,7 +202,7 @@ jobs:
       - name: Build
         id: cmake_build
         run: |
-          cmake -B build -DGGML_SCHED_NO_REALLOC=ON -DGGML_CPU_KLEIDIAI=ON
+          cmake -B build -DGGML_SCHED_NO_REALLOC=ON -DGGML_CPU_KLEIDIAI=ON -DLLAMA_SERVER_TEST_FAULTS=ON
           cmake --build build --config Release -j $(nproc) --target llama-server
 
       - name: Python setup

.github/workflows/server.yml

@@ -90,7 +90,8 @@ jobs:
         id: cmake_build
         run: |
           cmake -B build \
-            -DGGML_SCHED_NO_REALLOC=ON
+            -DGGML_SCHED_NO_REALLOC=ON \
+            -DLLAMA_SERVER_TEST_FAULTS=ON
           cmake --build build --config Release -j $(nproc) --target llama-server test-server-prefill
 
       - name: Prefill scheduling test
@@ -167,7 +168,8 @@ jobs:
             -DCMAKE_TOOLCHAIN_FILE=cmake/x64-windows-llvm.cmake ^
             -DCMAKE_BUILD_TYPE=Release ^
             -DLLAMA_BUILD_BORINGSSL=ON ^
-            -DGGML_SCHED_NO_REALLOC=ON
+            -DGGML_SCHED_NO_REALLOC=ON ^
+            -DLLAMA_SERVER_TEST_FAULTS=ON
           set /A NINJA_JOBS=%NUMBER_OF_PROCESSORS%-1
           cmake --build build --config Release -j %NINJA_JOBS% --target llama-server test-server-prefill
 

tools/server/CMakeLists.txt

@@ -34,6 +34,12 @@ target_include_directories(${TARGET} PRIVATE ../mtmd)
 target_include_directories(${TARGET} PRIVATE ${CMAKE_SOURCE_DIR})
 target_link_libraries(${TARGET} PUBLIC llama-common mtmd ${CMAKE_THREAD_LIBS_INIT})
 
+# the faults tools/server/tests inject (server-test-faults.h); only the server test build sets it
+option(LLAMA_SERVER_TEST_FAULTS "llama-server: compile in the test fault injection (tests only, never a release)" OFF)
+if (LLAMA_SERVER_TEST_FAULTS)
+    target_compile_definitions(${TARGET} PRIVATE LLAMA_SERVER_TEST_FAULTS)
+endif()
+
 # llama-server-impl: server logic, reusable by app
 
 set(TARGET llama-server-impl)

tools/server/server-context.cpp

@@ -11,6 +11,7 @@
 #include "server-stream.h"
 #include "server-page-adapters.h"
 #include "server-prefill.h"
+#include "server-test-faults.h"
 
 #include "build-info.h"
 #include "common.h"
@@ -171,6 +172,32 @@ struct server_batch {
         return (int32_t)tokens.size();
     }
 
+    // what a slot finds in the batch before it adds to it
+    struct mark {
+        int32_t       n_tokens;
+        size_t        n_embd;
+        bool          has_embd;
+        server_slot * slot_batched;
+        float         alora_scale;
+        size_t        alora_disabled_id;
+    };
+
+    mark get_mark() const {
+        return { size(), embd.size(), has_embd, slot_batched, alora_scale, alora_disabled_id };
+    }
+
+    // undo everything added since the mark, so a slot that fails leaves the batch as it found it
+    void rollback(const mark & m) {
+        GGML_ASSERT(!batch_rendered);
+        GGML_ASSERT(m.n_tokens <= size() && m.n_embd <= embd.size());
+        tokens.resize(m.n_tokens);
+        embd.resize(m.n_embd);
+        has_embd          = m.has_embd;
+        slot_batched      = m.slot_batched;
+        alora_scale       = m.alora_scale;
+        alora_disabled_id = m.alora_disabled_id;
+    }
+
     void set_output(int32_t idx, bool output) {
         GGML_ASSERT(idx >= 0 && idx < (int32_t)tokens.size());
         tokens[idx].output = output;
@@ -943,6 +970,7 @@ struct server_context_impl {
     int trace = 0;        // env: LLAMA_TRACE
     int slots_debug = 0;  // env: LLAMA_SERVER_SLOTS_DEBUG
     int slots_n_diff = 0; // env: LLAMA_SERVER_SLOTS_N_DIFF
+    server_test_faults test_faults; // the server test build only (LLAMA_SERVER_TEST_FAULTS)
 
     int n_empty_consecutive = 0;
 
@@ -1340,6 +1368,14 @@ struct server_context_impl {
             }
         }
 
+        if constexpr (server_test_faults::enabled) {
+            test_faults = server_test_faults::from_env();
+            if (test_faults.any()) {
+                SRV_WRN("test faults: prompt checkpoints = %d, speculative checkpoints = %d, prompt fills = %d\n",
+                        test_faults.prompt_checkpoints, test_faults.spec_checkpoints, test_faults.prompt_fills);
+            }
+        }
+
         // the update_slots() logic will always submit a maximum of n_batch or n_parallel tokens
         // note that n_batch can be > n_ctx (e.g. for non-causal attention models such as BERT where the KV cache is not used)
         {
@@ -2482,6 +2518,46 @@ struct server_context_impl {
         return (slot.task->n_tokens() + params_base.n_ctx_checkpoints - 1) / params_base.n_ctx_checkpoints;
     }
 
+    enum class checkpoint_kind { prompt, speculative };
+
+    // which parts of a sequence's state a checkpoint copies
+    struct checkpoint_parts {
+        bool tgt;
+        bool dft;
+        bool spec; // the speculative implementation's own state
+    };
+
+    // Continuum fork: the one place a checkpoint's state copy is allocated. A copy that cannot be
+    // allocated is refused, not a slot failure: the checkpoint is cleared, false is returned and
+    // the caller goes on without it (a restore from a prompt checkpoint is cold; a draft that a
+    // speculative checkpoint would protect is dropped, and the step decodes without it).
+    bool copy_checkpoint(const server_slot & slot, common_prompt_checkpoint & ckpt, checkpoint_kind kind, checkpoint_parts parts) {
+        const bool is_prompt = kind == checkpoint_kind::prompt;
+        try {
+            if constexpr (server_test_faults::enabled) {
+                if (server_test_faults::take(is_prompt ? test_faults.prompt_checkpoints : test_faults.spec_checkpoints)) {
+                    throw std::bad_alloc(); // where the copies below allocate
+                }
+            }
+            if (parts.tgt) {
+                ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+            }
+            if (parts.dft) {
+                ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+            }
+            if (parts.spec) {
+                common_speculative_get_state(spec.get(), slot.id, ckpt.data_spec);
+            }
+            return true;
+        } catch (const std::bad_alloc &) {
+            SLT_WRN(slot, "%s checkpoint refused: its state copy could not be allocated (n_tokens = %" PRId64 "), %s\n",
+                    is_prompt ? "context" : "speculative", ckpt.n_tokens,
+                    is_prompt ? "a restore from here will be cold" : "the draft is dropped");
+            ckpt.clear();
+            return false;
+        }
+    }
+
     // n_tokens_cur: the number of tokens added to the batch for the current slot
     void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) {
         const int id_task = slot.task->id;
@@ -2513,7 +2589,8 @@ struct server_context_impl {
             slot.prompt.checkpoints.erase(slot.prompt.checkpoints.begin());
         }
 
-        auto & cur = slot.prompt.checkpoints.emplace_back();
+        // the copy is made aside and kept only once complete
+        common_prompt_checkpoint cur;
 
         cur.id_task = id_task;
 
@@ -2522,15 +2599,17 @@ struct server_context_impl {
         //       this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225
         cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max);
 
-        cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-        cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-        // stash the draft's speculative state with the checkpoint
-        common_speculative_get_state(spec.get(), slot.id, cur.data_spec);
+        // with the draft's speculative state stashed beside it
+        if (!copy_checkpoint(slot, cur, checkpoint_kind::prompt, { /* tgt */ true, /* dft */ true, /* spec */ true })) {
+            return;
+        }
+
+        const auto & kept = slot.prompt.checkpoints.emplace_back(std::move(cur));
 
         SLT_TRC(slot,
                 "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n",
-                (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min,
-                cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
+                (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, kept.pos_min,
+                kept.pos_max, kept.n_tokens, (float) kept.size() / 1024 / 1024);
     }
 
     // returns false to decline the task, it is offered again after the decode is done
@@ -3621,8 +3700,8 @@ struct server_context_impl {
                                 llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id),
                                 llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id));
 
-                        if (use_ckpt_dft) {
-                            slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                        if (use_ckpt_dft && !copy_checkpoint(slot, slot.spec_ckpt, checkpoint_kind::speculative, { false, true, false })) {
+                            return; // no draft this step
                         }
 
                         slot.spec_prompt = slot.prompt.tokens.get_text_tokens();
@@ -3654,8 +3733,6 @@ struct server_context_impl {
             auto & draft = slot.spec_draft;
             auto & ckpt  = slot.spec_ckpt;
 
-            slot.stats.n_draft_tokens += draft.size();
-
             // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL]
             const bool use_ckpt_dft = ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
 
@@ -3677,24 +3754,18 @@ struct server_context_impl {
                 const bool use_ckpt_dft =
                    (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft));
 
-                if (use_ckpt_tgt) {
-                    //const int64_t t_start = ggml_time_us();
-
-                    ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-
-                    //const int64_t t_total = ggml_time_us() - t_start;
-                    //printf("checkpoint total: %f ms\n", t_total / 1000.0);
-
+                if ((use_ckpt_tgt || use_ckpt_dft) &&
+                    !copy_checkpoint(slot, ckpt, checkpoint_kind::speculative, { use_ckpt_tgt, use_ckpt_dft, false })) {
+                    draft.clear(); // a draft that could not be rolled back is not verified
+                } else if (use_ckpt_tgt) {
                     SLT_DBG(slot, "created speculative checkpoint (pos_min = %d, pos_max = %d, n_tokens = %d, size = %.3f MiB, draft = %.3f MiB)\n",
                             ckpt.pos_min, ckpt.pos_max, slot.prompt.n_tokens(),
                             (float) ckpt.size() / 1024 / 1024,
                             (float) ckpt.data_dft.size() / 1024 / 1024);
                 }
-
-                if (use_ckpt_dft) {
-                    ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
-                }
             }
+
+            slot.stats.n_draft_tokens += draft.size();
         });
 
         // update the batch with the sampled/drafted tokens
@@ -3723,7 +3794,9 @@ struct server_context_impl {
             for (const size_t index : prefill_order.order(slots.size(), params_base.kv_unified)) {
                 prompt_slots.push_back(&slots[index]);
             }
-            iterate(prompt_slots, [&](server_slot & slot) {
+            // n_written: the slot's prompt tokens whose memory is written, set once the fill
+            // starts adding tokens that are only in the batch (not decoded yet)
+            const auto fill_prompt = [&](server_slot & slot, int32_t & n_written) {
                 const auto batch_before = batch.size();
                 if (!add_ok || batch.size() >= n_batch) {
                     return; // batch is full, skip remaining slots
@@ -4176,6 +4249,7 @@ struct server_context_impl {
                     slot.prompt.end_from = std::max(0, slot.task->n_tokens() - std::min(n_batch, 4 + n_ubatch));
 
                     // add prompt tokens for processing in the current batch
+                    n_written = slot.prompt.n_tokens();
                     while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) {
                         // get next token to process
                         llama_token cur_tok = input_tokens[slot.prompt.n_tokens()];
@@ -4245,6 +4319,12 @@ struct server_context_impl {
                     // the number of tokens added to the batch for the current slot
                     const auto n_tokens_cur = batch.size() - n_tokens_prev;
 
+                    if constexpr (server_test_faults::enabled) {
+                        if (n_tokens_cur > 0 && server_test_faults::take(test_faults.prompt_fills)) {
+                            throw std::runtime_error("test fault: the prompt fill failed after adding its tokens");
+                        }
+                    }
+
                     const auto n_tokens_start = slot.prompt.n_tokens() - n_tokens_cur;
 
                     const bool near_prompt_end = slot.task->n_tokens() < slot.prompt.n_tokens() + n_ubatch;
@@ -4304,6 +4384,22 @@ struct server_context_impl {
                 if (!slot_batched) {
                     slot_batched = &slot;
                 }
+            };
+
+            // a slot whose fill fails leaves the batch and its prompt as it found them: its
+            // tokens were never decoded, so no batched slot owns them and its memory lacks them
+            iterate(prompt_slots, [&](server_slot & slot) {
+                const auto batch_mark = batch.get_mark();
+                int32_t n_written = -1;
+                try {
+                    fill_prompt(slot, n_written);
+                } catch (...) {
+                    batch.rollback(batch_mark);
+                    if (n_written >= 0) {
+                        slot.prompt.tokens.keep_first(n_written);
+                    }
+                    throw; // iterate() reports the error and releases the slot
+                }
             });
         }
     }

tools/server/server-test-faults.h

@@ -0,0 +1,54 @@
+#pragma once
+
+#include <cstdint>
+#ifdef LLAMA_SERVER_TEST_FAULTS
+#include <cstdlib>
+#endif
+
+// Continuum fork: failures the server tests inject through the environment, so a test can show
+// that a failure leaves the server consistent. Compiled in only by the server test build
+// (cmake -DLLAMA_SERVER_TEST_FAULTS=ON). In any other build `enabled` is false, every use sits
+// under `if constexpr (server_test_faults::enabled)` and is discarded, and no environment is
+// read, so a stray variable can never fail a real slot.
+struct server_test_faults {
+#ifdef LLAMA_SERVER_TEST_FAULTS
+    static constexpr bool enabled = true;
+#else
+    static constexpr bool enabled = false;
+#endif
+
+    int32_t prompt_checkpoints = 0; // LLAMA_SERVER_TEST_FAIL_CHECKPOINTS: the next N prompt checkpoint copies fail to allocate
+    int32_t spec_checkpoints   = 0; // LLAMA_SERVER_TEST_FAIL_SPEC_CHECKPOINTS: the next N speculative checkpoint copies fail to allocate
+    int32_t prompt_fills       = 0; // LLAMA_SERVER_TEST_FAIL_PROMPT_FILLS: the next N prompt fills throw after adding their tokens
+
+    static server_test_faults from_env() {
+        server_test_faults faults;
+#ifdef LLAMA_SERVER_TEST_FAULTS
+        faults.prompt_checkpoints = read("LLAMA_SERVER_TEST_FAIL_CHECKPOINTS");
+        faults.spec_checkpoints   = read("LLAMA_SERVER_TEST_FAIL_SPEC_CHECKPOINTS");
+        faults.prompt_fills       = read("LLAMA_SERVER_TEST_FAIL_PROMPT_FILLS");
+#endif
+        return faults;
+    }
+
+    bool any() const {
+        return prompt_checkpoints || spec_checkpoints || prompt_fills;
+    }
+
+    // true when this occurrence is one to fail, counting it
+    static bool take(int32_t & remaining) {
+        if (remaining <= 0) {
+            return false;
+        }
+        --remaining;
+        return true;
+    }
+
+#ifdef LLAMA_SERVER_TEST_FAULTS
+private:
+    static int32_t read(const char * name) {
+        const char * value = std::getenv(name);
+        return value ? std::atoi(value) : 0;
+    }
+#endif
+};

tools/server/tests/README.md

@@ -17,10 +17,12 @@ To mitigate it, you can increase values in `n_predict`, `kv_size`.
 
 ```shell
 cd ../../..
-cmake -B build
+cmake -B build -DLLAMA_SERVER_TEST_FAULTS=ON
 cmake --build build --target llama-server
 ```
 
+`LLAMA_SERVER_TEST_FAULTS` compiles in the faults `unit/test_failed_fill.py` injects; without it those tests skip. Never set it for a build that serves.
+
 2. Start the test: `./tests.sh`
 
 It's possible to override some scenario steps values with environment variables: