The result of bool llama_memory_seq_rm() is not checked

[leok7v]

Code in main.cpp passkey.cpp server.cpp speculative.cpp lookup.cpp lookahead.cpp speculative-simple.cpp
do not check boolean result returned by llama_memory_seq_rm() that can be false.

Also see
Discussion:
#16708
and
Usages:
https://github.com/search?q=repo%3Aggml-org%2Fllama.cpp+llama_memory_seq_rm&type=code

[leok7v]

What’s going wrong

1. Hybrid seq_rm is “all-or-nothing”

In llama_memory_hybrid::seq_rm you currently have: oai_citation:0‡GitHub

bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
    // Try removing from the recurrent cache first since it may fail.
    // If it does fail, the cache will not have been mutated.
    if (!mem_recr->seq_rm(seq_id, p0, p1)) {
        return false;
    }
    return mem_attn->seq_rm(seq_id, p0, p1);
}

So: if the recurrent part says “I can’t do that”, hybrid reports failure and the attention cache never gets its range removed.

2. Recurrent seq_rm intentionally refuses partial ranges

llama_memory_recurrent::seq_rm is designed to forbid partial erasures, because Mamba / RWKV style state can’t be meaningfully “cut” in the middle: oai_citation:1‡searchfox.org

bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
    uint32_t new_head = size;

    if (p0 < 0) { p0 = 0; }
    if (p1 < 0) { p1 = std::numeric_limits<llama_pos>::max(); }

    // models like Mamba or RWKV can't have a state partially erased
    if (seq_id >= (int64_t) size) {
        // could be fatal
        return false;
    }

    if (0 <= seq_id) {
        int32_t & tail_id = cells[seq_id].tail;
        if (tail_id >= 0) {
            const auto & cell = cells[tail_id];

            // partial intersection is invalid
            if ((0 < p0 && p0 <= cell.pos) ||
                (0 < p1 && p1 <= cell.pos)) {
                return false;
            }

            // invalidate tails which will be cleared
            if (p0 <= cell.pos && cell.pos < p1) {
                tail_id = -1;
            }
        }
    } else {
        // seq_id is negative, then the range should include everything or nothing
        if (p0 != p1 &&
            (p0 != 0 || p1 != std::numeric_limits<llama_pos>::max())) {
            return false;
        }
    }

    // ... actually clear cells ...

    // If we freed up a slot, set head to it so searching can start there.
    if (new_head != size && new_head < head) {
        head = new_head;
    }

    return true;
}

Key points:

• For seq_id >= 0, it disallows any range [p0, p1) that only partially intersects the tail cell:

  if ((0 < p0 && p0 <= cell.pos) ||
      (0 < p1 && p1 <= cell.pos)) {
      return false;
  }

• For seq_id < 0, it only allows “remove everything” (p0 == 0 and p1 == max) or “remove nothing”.

So the recurrent cache is basically saying:

Either you wipe the whole sequence (or nothing), or I refuse.

3. Main’s context shift / Self-Extend does use partial ranges

In tools/main/main.cpp, the context-management logic is purely written in terms of “sequence positions” and doesn’t know (or care) that the memory implementation might be recurrent, hybrid, whatever. oai_citation:2‡GitHub

• For the simple case (ga_n == 1), the infinite-generation “context shifting” path:

if (ga_n == 1) {
    // infinite text generation via context shifting
    if (n_past + (int) embd.size() >= n_ctx) {
        ...
        const int n_left    = n_past - params.n_keep;
        const int n_discard = n_left/2;

        llama_memory_seq_rm (mem, 0,
                             params.n_keep,
                             params.n_keep + n_discard);
        llama_memory_seq_add(mem, 0,
                             params.n_keep + n_discard,
                             n_past,
                             -n_discard);

        n_past -= n_discard;
        ...
    }
}

This is a partial trimming of the sequence [keep, keep + discard), not a full wipe.

• For Self-Extend (ga_n > 1), it uses partial seq_add/seq_div windows too:

// context extension via Self-Extend
while (n_past >= ga_i + ga_w) {
    const int ib = (ga_n*ga_i)/ga_w;
    const int bd = (ga_w/ga_n)*(ga_n - 1);
    const int dd = (ga_w/ga_n) - ib*bd - ga_w;

    llama_memory_seq_add(mem, 0, ga_i,                n_past,              ib*bd);
    llama_memory_seq_div(mem, 0, ga_i + ib*bd,        ga_i + ib*bd + ga_w, ga_n);
    llama_memory_seq_add(mem, 0, ga_i + ib*bd + ga_w, n_past + ib*bd,      dd);

    n_past -= bd;
    ga_i   += ga_w/ga_n;
}

For plain KV cache, that’s all fine. But for recurrent:

• seq_rm(0, keep, keep + discard) is exactly the kind of partial range seq_rm forbids.
• llama_memory_hybrid::seq_rm forwards that to mem_recr->seq_rm first and bails out if it returns false.

→ So for hybrid models, every attempt to context shift fails at mem_recr->seq_rm, and the attention cache never gets shifted.

4. Why you then see mem_attn->prepare failing

In llama_memory_hybrid::init_batch you prepare recurrent first and attention second: oai_citation:3‡GitHub

// prepare the recurrent batches first
if (!mem_recr->prepare(ubatches)) {
    LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__);
    return std::make_unique<llama_memory_hybrid_context>(
        LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}

// prepare the attention cache
auto heads_attn = mem_attn->prepare(ubatches);
if (heads_attn.empty()) {
    LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__);
    return std::make_unique<llama_memory_hybrid_context>(
        LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}

• mem_recr->prepare(ubatches) succeeds – the recurrent “cell pool” can still accommodate whatever you’re doing.
• mem_attn->prepare(ubatches) fails when KV is full and no space has been freed by a prior seq_rm/seq_add (because that call was aborted by recurrent’s refusal).

Exactly what you described: recurrent prepare succeeds, attention prepare fails once you hit context limits and the shift / Self-Extend that should free KV space never actually runs.

---

What should happen conceptually

For hybrid:

• The attention side is the thing with a finite sliding context where you want:
  • context shifting (trim old tokens, renumber positions), and/or
  • Self-Extend (compress old KV).
• The recurrent side is “infinite” in the sense that its state is already a compressed summary of the whole prefix. For most use-cases:
  • You don’t want to trim it when you trim KV.
  • You only want to clear it when the entire sequence is reset.

So when main does:

llama_memory_seq_rm (mem, 0, keep, keep + discard);
llama_memory_seq_add(mem, 0, keep + discard, n_past, -discard);

…this should, for hybrid:

• Definitely apply to mem_attn (to free KV and renumber positions).
• Be optional / mostly ignored for mem_recr, unless it’s a full-range wipe.

---

Minimal practical fix

The existing comment in hybrid::seq_rm actually hints at the intended semantics: oai_citation:4‡GitHub

// Try removing from the recurrent cache first since it may fail.
// If it does fail, the cache will not have been mutated.

That screams “best-effort, non-fatal” – but the implementation treats failure as fatal.

Patch idea

Change llama_memory_hybrid::seq_rm so that recurrent failure does not abort the attention operation, and optionally only hit recurrent for full-range wipes.

Something like this (semantically clean):

bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
    // Normalize p0/p1 the same way recurrent does
    llama_pos p0n = p0;
    llama_pos p1n = p1;
    if (p0n < 0) {
        p0n = 0;
    }
    if (p1n < 0) {
        p1n = std::numeric_limits<llama_pos>::max();
    }

    // Only try to touch recurrent for "full wipe" ranges.
    // Partial context-shift / Self-Extend is attention-only.
    const bool full_range =
        p0n == 0 && p1n == std::numeric_limits<llama_pos>::max();

    if (full_range) {
        // If this fails, we *do* treat it as real failure.
        if (!mem_recr->seq_rm(seq_id, p0, p1)) {
            return false;
        }
    } else {
        // For partial ranges, best-effort: ignore recurrent failure.
        (void) mem_recr->seq_rm(seq_id, p0, p1);
    }

    // The attention cache must succeed for the call to be considered successful.
    return mem_attn->seq_rm(seq_id, p0, p1);
}

Even simpler (if you’re okay with not distinguishing full vs partial on hybrid level):

bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
    // Best-effort on recurrent; its failure just means "no change".
    (void) mem_recr->seq_rm(seq_id, p0, p1);
    return mem_attn->seq_rm(seq_id, p0, p1);
}

This works because:

• llama_memory_recurrent::seq_rm never partially mutates and then returns false; it either:
  • refuses up front (no changes) or
  • fully applies the operation and returns true. oai_citation:5‡searchfox.org
• So ignoring the boolean is safe: failure ⇒ “no change to recurrent”, success ⇒ “recurrent state updated consistently”.

With that change:

• Context shift (ga_n == 1) will successfully trim & shift attention KV for hybrid models.
• Recurrent state will:
  • stay intact across partial shifts (i.e., it keeps the full prefix summary, which is what you want most of the time), and
  • still be clearable by “full wipe” operations that call seq_rm with a whole-range argument (e.g. reset sequence / clear all).

---

Why your observed behavior matches this

llama_memory_seq_rm() returns false for hybrid models because …

Yep:

1. Main calls llama_memory_seq_rm(mem, 0, keep, keep + discard) during context shift / Self-Extend. main.cpp
2. Hybrid forwards that to mem_recr->seq_rm(0, keep, keep + discard).
3. Recurrent sees this is a partial range intersecting its sequence tail and returns false. llama-memory-recurrent.cpp
4. Hybrid returns false and never calls mem_attn->seq_rm.
5. KV stays full; then mem_attn->prepare(ubatches) can’t find space and fails on the next batch build. llama-memory-hybrid.cpp

After the patch, step 4 changes to:

• Hybrid ignores recurrent failure and does call mem_attn->seq_rm, freeing space in KV.
• mem_attn->prepare then succeeds again.

---

TL;DR

• The bug is in llama_memory_hybrid::seq_rm treating recurrent refusal as fatal.
• Recurrent seq_rm is deliberately strict and disallows partial erasures.
• Context shifting / Self-Extend are partial erasures, so hybrid currently can’t use them.
• Fix: make hybrid’s seq_rm treat recurrent as best-effort (or only apply it for full-range wipes), and always execute mem_attn->seq_rm for these context-management operations.

[gabe-l-hart]

Hi @leok7v, thanks for identifying this issue. It looks like you've got the situation pretty well in hand, so you likely have a better grip on this by now than I did when I wrote the initial hybrid cache implementation. My understanding is that seq_rm is used with offsets in order to do a partial erasure which (as you note) is not possible (currently) for the recurrent cache due to the lack of retained state from previous iterations. The reason for treating this as a fatal issue is that the bool return semantics would be out-of-sync if the attn cache were partially cleared but the recurrent portion were not. If this is used to determine when to re-perform prefill on some tokens, it would either result in the wrong starting state for recurrent, or lost context for the attn cache. This is the case with the one place where the return value actually is checked (https://github.com/ggml-org/llama.cpp/blob/master/tools/server/server.cpp#L3963).

I think what you're identifying is that seq_rm is actually used for multiple purposes and that for context shifting, this best-effort behavior is actually desired. I'm not really sure what the right answer is there since the current function signature is not able to distinguish these two cases. A simple flag for best_effort could be used, but I'm not sure if there's appetite for changing the public API that way.

[leok7v]

Actually recurrent memory does not need shift because it is keeping the cumulative final state but the kv memory may benefit from shifting or self expand (at least my understanding is that it can).
And there are other reasons why seq_rm() may fail on both recr-> and attn-> subsystems.
I am not advocating any public API changes - no-one wants that unless we have a really compelling reason.
But we do need to check and at least log seq_rm() failures and pass it up (especially in server.cpp)
Also some extra tests for this situation would be good.

Alternative is simple - do NOT use group attention shift/self expand but I don't think it is supported right now, though server.cpp is in major reimplementation of slots and memory right as we speak.

There is also this:
#16983

And after all of that there are still issues with find_slot()...

But again - I am just humble user of the code - I don't have deep understanding of all models architectures and especially hybrid models.

[gabe-l-hart]

That all makes sense. I think the issue is that there are two scenarios where seq_rm is used:

• Context Shifting: this removes tokens at the beginning of the context (AFAIK) which, as you state, is a no-op for the recurrent cache.
• Common Prefix Caching: this removes tokens at the end of the context which is where the recurrent cache is unable to "reverse time" without sequence checkpoints.

Now that I think about it, this should be pretty easy to differentiate based on the starting position in seq_rm. If it starts at 0 or <0, it's a no-op and just return true. If it starts >0, it's impossible and return false.

[gabe-l-hart]

Ok, I really needed to read this more carefully! This piece of context shifting and self-extension is pretty nuanced (thanks for spelling it out so carefully in the issue). I think I was close to correct before, but I wasn't quite understanding the Context Shifting portion. The logic you outline in the issue is all spot-on as long as the partial erasure does not include the tail of the sequence, so I think the logic in the recurrent cache needs to be updated to only refuse partial clears that include the tail-end of the sequence. This should percolate up to not failing in the hybrid cache and fix the core issue of the attention cache never getting cleared.

[gabe-l-hart]

Ok, I think I have the right fix in place for this: #17009

I haven't tested this with llama-server. I believe the shifting logic there is similar, though the calculation of the upper-bound seems to be different.