moe-cache : fail the decode instead of aborting on a mailbox timeout

The mailbox's CPU op aborted the process when the GPU did not hand a layer
over within 20 s, taking llama-server down with it. Now it computes no expert
for the step (all -1 ids), sets an error flag, and llama_context::graph_compute
turns the flag into GGML_STATUS_FAILED, so llama_decode returns -3 and rolls
the step back. Later layers of the same step skip the wait once one has
timed out. The GPU-side receive is launched after the CPU split, so giving up
cannot deadlock it, and a late send carries an epoch the next step rejects.

Checked by forcing the timeout: one error line, decode fails, exit code 1
instead of 134.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: voidstackloop

src/llama-context.cpp

@@ -2768,6 +2768,11 @@ ggml_status llama_context::graph_compute(
         LLAMA_LOG_ERROR("%s: ggml_backend_sched_graph_compute_async failed with error %d\n", __func__, status);
     }
 
+    // the CPU splits have run by now, the MoE cache's mailbox ops among them
+    if (moe_cache && moe_cache->take_mailbox_error() && status == GGML_STATUS_SUCCESS) {
+        status = GGML_STATUS_FAILED;
+    }
+
     // fprintf(stderr, "splits: %d\n", ggml_backend_sched_get_n_splits(sched));
 
     return status;

src/llama-moe-cache.cpp

@@ -308,6 +308,7 @@ struct llama_moe_cache::impl {
     float   heat_inc = 1.0f;                 // the weight of a selection now: it grows, which is how older ones decay
     int32_t steps_since_rebalance = 0;
     bool    decode_seen = false;             // a CPU op of the hybrid ran in the last graph
+    std::atomic<bool> mailbox_error{false};  // a hand-over timed out in the last graph, see take_mailbox_error
     size_t  n_swap = 0;
     // mailbox: how long the CPU waited for the GPU's hand-overs (the time the GPU spends before the experts), and how many
     uint64_t wait_ns = 0;
@@ -859,11 +860,25 @@ struct llama_moe_cache::impl {
         impl & self = *h.self;
         const mailbox_layer & m = self.mailbox[h.il];
         const uint32_t epoch = self.epoch;
+        const int64_t k = m.ids->ne[0];
+        int32_t * cold = (int32_t *) dst->data;
+
+        // an earlier layer of this step timed out already: the step fails anyway, do not wait 20 s for every layer
+        if (self.mailbox_error) {
+            std::fill(cold, cold + k, -1);
+            return;
+        }
 
         const auto t0 = std::chrono::steady_clock::now();
         for (uint64_t spins = 0; *m.host_flag != epoch; ++spins) {
             if ((spins & 0xFFFFF) == 0xFFFFF && std::chrono::steady_clock::now() - t0 > std::chrono::seconds(20)) {
-                GGML_ABORT("MoE cache: the GPU did not hand layer %d over within 20 s", h.il);
+                // a stalled GPU fails the step, not the process: the CPU computes no expert (its GPU-side receive is
+                // launched after this split anyway) and graph_compute turns the error into a failed decode. A late send
+                // is stamped with this epoch, which the next step's no longer matches
+                LLAMA_LOG_ERROR("%s: the GPU did not hand layer %d over within 20 s, failing the step\n", __func__, h.il);
+                std::fill(cold, cold + k, -1);
+                self.mailbox_error = true;
+                return;
             }
         }
         std::atomic_thread_fence(std::memory_order_acquire);
@@ -872,8 +887,6 @@ struct llama_moe_cache::impl {
 
         const int32_t * ids   = (const int32_t *) m.ids->data;
         const int32_t * slots = (const int32_t *) dst->src[2]->data;
-        const int64_t k = m.ids->ne[0];
-        int32_t * cold = (int32_t *) dst->data;
 
         // consumed: a graph computed without begin_step() would wait for a new send instead of reading these ids again
         *m.host_flag = 0;
@@ -1461,6 +1474,10 @@ void llama_moe_cache::begin_step() {
     pimpl->begin_step();
 }
 
+bool llama_moe_cache::take_mailbox_error() {
+    return pimpl->mailbox_error.exchange(false);
+}
+
 ggml_tensor * llama_moe_cache::build_hybrid_ids(ggml_context * ctx, ggml_tensor * selected_experts, int32_t il) const {
     return pimpl->build_hybrid_ids(ctx, selected_experts, il);
 }

src/llama-moe-cache.h

@@ -72,6 +72,9 @@ class llama_moe_cache {
     void build_mailbox_recv(ggml_context * ctx, int32_t il, ggml_tensor * e_cold, ggml_tensor *& copy, ggml_tensor *& recv) const;
     // before each graph compute: a new epoch, which the GPU's sends stamp on their flags
     void begin_step();
+    // after each graph compute: true if a hand-over from the GPU timed out in it (the CPU then computed none of its
+    // experts, so the step's results are wrong and it must fail); clears the error
+    bool take_mailbox_error();
 
     // ggml_backend_sched copy callback, returns false if src is not a slot map
     bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph);