fix: surface a failed K-shift instead of decoding over stale K

[yingying0906]

What this fixes

When the K-shift graph fails to allocate or compute, llama_decode still returns 0 and the decode carries on. By that point the cells already hold their shifted positions but K was never rotated to match, so attention runs against stale K. The result is wrong and nothing reports it.

Two layers were swallowing the failure, and both had the same shape: a bool that could not distinguish "nothing to do" from "it failed". A third bool then flattened the failure itself, so the caller could not tell an abort from a real error.

Layer 1: apply() discarded the result

if (ubatches.empty()) {
    kv->update(lctx, do_shift, sc_info);

    return true;
}

llama_kv_cache::update() could not simply be forwarded, because it reported whether it had done any work, not whether it succeeded. updated starts false and is only set once a shift actually runs, so a no-op update returned false. It now reports success instead. Nothing consumed the old meaning: update() is declared on llama_kv_cache rather than the llama_memory_i interface, and apply() is its only caller.

Layer 2: memory_update() logged it and returned true

if (!mctx->apply()) {
    LLAMA_LOG_ERROR("%s: failed to apply memory update\n", __func__);
}

and llama_context::decode calls it for exactly this purpose, discarding the return:

// handle any pending shifts/copies
memory_update(false);

memory_update's bool had the same problem one level up: LLAMA_MEMORY_STATUS_NO_UPDATE returns false, so forwarding it would have failed every decode that had nothing pending. It now returns a three-state result:

enum class memory_update_status { no_update, updated, failed };

decode() stops on failed, and so does the cache-optimization retry, which used to fall through to return 1 and tell the caller to try a smaller batch.

Returning early there also skips the worst-case graph reserve below it. That reserve was what restored the reservation after the memory module reset the scheduler to build its own graph, and it was the only synchronize on the path, since graph_compute dispatches async and ggml_backend_sched_reset does not drain. The failure path drains and asks for a fresh reserve on the next call instead.

Layer 3: the status was flattened into one return code

A false out of the K-shift said nothing about what went wrong. graph_compute() != GGML_STATUS_SUCCESS is true for ALLOC_FAILED, FAILED and ABORTED alike, so all three reached the caller as one -2. decode() does not flatten them for its own ubatch compute:

case GGML_STATUS_ABORTED:      return  2;
case GGML_STATUS_ALLOC_FAILED: return -2;
case GGML_STATUS_FAILED:       return -3;

The abort callback is installed on every backend from the constructor and nothing suspends it for the shift, so a caller cancelling mid-shift got a fatal -2 instead of the non-fatal 2. Stopping is right either way, since K is left half-rotated, but the caller has to be able to tell the two apart. llama_memory_context_i::apply() returns bool and cannot carry the status out, so llama_context records it on the way through and maps it in one place.

Deliberate non-change

has_shift is left set on the failure path. The shift is still owed, and clearing it would silently discard a pending correction rather than let the caller decide. A caller that ignores the decode error will retry the same failing shift, which is preferable to it silently attending over stale K.

There was no way out of it, though. seq_rm resets each cell's own shift but never the cells-level flag, and retrying does not converge because the shift graph is sized off the cache capacity rather than its occupancy, so llama_memory_clear was the only escape and llama-server never calls it. An empty stream owes no shift, so seq_rm clears the flag there.

Context

Found while removing context shifting from reasoning-block compaction in the llm-addon, tetherto/qvac#3938. That change removes the addon's only path into this code, so the addon no longer depends on this fix. It is still reachable by any other caller that shifts, including tools/completion (both the context-shift loop and the self-extend path), examples/passkey, and anything else calling llama_memory_seq_add.

The server side

Two changes are outside src, both pre-existing and both made easier to hit by this one.

update_slots() throws on ret < -1 and nothing catches it, so llama-server goes down with the batch. The other -2 returns in decode() already reach it. It is caught in the callback now, where the slots are, because not every throw site releases them first: common_speculative_process failing, the unsupported has_embd + spec case and an out-of-range speculative batch index all throw mid-batch. A slot left processing is picked up on the next wake-up and throws again, so the request would hang rather than error. The queue loop keeps its own catch as the backstop for a teardown that throws too, since start_loop runs outside any try.

decode() left err empty for ret == 2, so an abort fell through to the "no free space in the KV cache" path and halved n_batch. Retrying cannot help: the abort came from outside, and when it hit a pending K-shift the shift is still owed, so the retry aborts in the same place while n_batch walks down to zero. It stops like the other fatal returns now.

Testing

Full build clean with server on, ctest 67/68 locally on temp-10297. The one failure is test-jinja-py, missing jinja2 in my env. The failure paths are allocation and compute failures inside the K-shift graph, which the test suite cannot synthesize without a memory-constrained backend, so the propagation itself is reviewed by inspection of the call sites above.