llama_kv_cache_seq_shift delta does not appear to be calculated properly

[MrJackSpade]

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• [Y] I am running the latest code. Development is very rapid so there are no tagged versions as of now.
• [Y] I carefully followed the README.md.
• [Y] I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• [Y] I reviewed the Discussions, and have a new bug or useful enhancement to share.

Not 100% certain if this is a bug or not, but I was playing with the kv cache shifting functionality and I was getting some weird results so I figured I'd step through it and see what was going on.

I noticed that after performing a double shift on a chunk of the kv cache, that the cell delta only reflected the second shift. The cell positions are properly updated, however. It looks like the shift code treats the cell pos differently than the cell delta.

See cache.cells[i].pos += delta vs. cache.cells[i].delta = delta

The position is cumulative, but the delta maintains the value of the last shift, which throws the position out of sync with the delta once its been shifted more than once.

Of course, the example in main.cpp discards half the disposable context on every shift, so if my understanding is correct, this isn't something that would be noticed through normal usage of the "main.cpp" application. A block of the KV cache wouldn't be shifted twice using that code in the first place, since the second shift would dispose of the block that was moved during the first shift.

So at least superficially it looks like this might be an oversight that wouldn't have come up without directly calling the API. I cant think of why the delta would be fixed to the last shift operation instead of cumulative like the position either. I figured I would make a note of this here in case this was an actual issue and not intentional

[TikaToka]

not sure with this, but due to this issue, I think i have a problem like below

{"timestamp":1698484388,"level":"INFO","function":"log_server_request","line":2156,"message":"request","remote_addr":"127.0.0.1","remote_port":37522,"status":200,"method":"POST","path":"/completion","params":{}}
system prompt updated
slot 0 is processing [task id: 2]
slot 1 released (360 tokens in cache)
slot 0 : kv cache rm - [124, end)
slot 0: context shift - n_keep = 0, n_left = 510, n_discard = 255
slot 0: context shift - n_keep = 0, n_left = 510, n_discard = 255
slot 0: context shift - n_keep = 0, n_left = 510, n_discard = 255
slot 0: context shift - n_keep = 0, n_left = 510, n_discard = 255
slot 0: context shift - n_keep = 0, n_left = 510, n_discard = 255
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 256
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 128
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 64
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 32
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 16
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 8
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 4
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 2
update_slots : failed to find free space in the KV cache, retrying with smaller n_batch = 1
update_slots : failed to decode the batch, n_batch = 1, ret = 1

looking for solutions...

[ggerganov]

@MrJackSpade Yes, this is an oversight. But I wonder why would you want to do 2 shifts instead of one?

@TikaToka I think you are running out of context - increase it through the -c arg. But your issue is not related to this one, so if the problem persist, you should open a separate issue with more details of what you are doing.

[MrJackSpade]

@ggerganov When rolling the context for the "infinite window" I started using 1/8 to 1/4 batches instead of 1/2 batches, because the amount of context data lost on 1/2 can make the transition more jarring. Rolling a smaller chunk at a time smooths out the memory loss and slightly increases the length of time that certain facts are remembered. It increases the amount of work, but also increases the quality quite a bit so it's worth it.

The result of this of course, is that with 1/4 size blocks, each cell would be shifted 3 times within the cache, with the 4th shift clearing it out.

[ggerganov]

@MrJackSpade Could you please review and test #3835

[MrJackSpade]

@ggerganov This does appear to resolve the immediate issue of the deltas falling out of sync through subsequent shifts.

Just to ensure that I'm doing my DD though, I should ask, would there be any benefit in using cache.cells[i].delta = cache.cells[i].pos - i instead of cache.cells[i].delta += delta? If the delta is supposed to represent the difference between the cell index and the cell position, it seems sensible that it might be calculated based on those properties for clarity even if the ideal end result of the calculation is the same. It would also potentially account for any "drift" that may potentially be introduced elsewhere, though I am generally cautious about introducing "self correcting" code as a rule.

[ggerganov]

This wouldn't work in the general case. Think of the cache just as a container and the positions of the cells in that container are irrelevant to their contents. You should be able to shuffle the cells in random way and the evaluation should still work.

In the simple case of a single sequence generation, it so happens that the index in the container matches the position of the tokens in the cache. But in general this is not the case.

[MrJackSpade]

Should I expect to see all "pos" accounted for depending on how far into the context I've written?

I've just noticed that post-shift there's now a huge chunk of numbers missing. The "pos" property of the cells appears to be contiguous between 0 - ~3K, then the entire block between 3K and 5K is missing, and new data is being written in the 5K range.

The issue that I'm currently working through is that after I've shifted the cells, the model seems to forget almost everything after the prompt. I'm trying to pin down if this is a problem with how I'm calling the API, or a bug within the API itself

Just to clarify, I'm not talking about the index of the cell data, that is expected to have gaps, but the actual "pos" property on the cells has fairly massive gaps.

[ggerganov]

Let say you have 8 tokens in the cache and it's size is 16:

idx:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
pos:  0  1  2  3  4  5  6  7 -1 -1 -1 -1 -1 -1 -1 -1

You want to keep 2-7 (i.e. evict oldest 2 tokens):
You shift the cache with -2:

idx:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
pos: -1 -1  0  1  2  3  4  5 -1 -1 -1 -1 -1 -1 -1 -1

A pos equal to -1 indicates that the cell is unoccupied so new data can be written in it.
If you want to evaluate a new token, you now have to give it position 6 instead of 8.

The cache might look like this:

idx:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
pos:  6 -1  0  1  2  3  4  5 -1 -1 -1 -1 -1 -1 -1 -1

Or this:

idx:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
pos: -1 -1  0  1  2  3  4  5  6 -1 -1 -1 -1 -1 -1 -1

Depending on where the logic finds a "free slot" for the new batch (i.e. token).
Hope this helps. Let me know if there are more questions - I'm looking forward to people using the KV cache API.

[MrJackSpade]

Ok, thanks for the explanation. That's what I was expecting.

What I'm seeing is

idx:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
pos: -1 -1  0  1  2  3 -1 -1  6 -1 -1 -1 -1 -1 -1 -1

Where the "2 tokens" appear to be trimmed from both the beginning, AND the end.

What happening is that I'm writing ~7500 tokens before shifting, and then removing ~2000 of them. That should logically mean that the next positions being written are ~5,500 post shift, which does seem to the case. The problem is that when I check the cache, the cells 3K-5K are missing.

However given the tricky method of tracking and shifting I've implemented on my client side, its possible (likely even) that this is the result of an error in my calculation. I just have to confirm that for sure because superficially all of my numbers appear correct.

I'm going to keep stepping through the code and see if I can track this one down, it seems like it should be relatively easy.

The only other potential issue I see in relation to this specific ticket, is that some cells within the cache have a pos == -1 and a delta != 0. I haven't nailed down the exact cause of that one yet because I'm more focused on the disappearing tokens right now. I don't have a full view of the code base at the moment, so I cant say whether or not initializing a cell with an already < 0 delta would cause an issue, or if that cell delta is always reset at the point of initialization. If that's a problem we can leave this issue open and I can see if I can track that down. If not, could close this issue and if the "missing tokens" issue ends up being on the Llama.cpp side, then I can just open a new bug for that.

[ggerganov]

Hm, not sure what is causing this. Check the example code here:

https://github.com/ggerganov/llama.cpp/blob/8925cf9ef8216621de37d41b787d25305c9d556d/examples/main/main.cpp#L496-L501

This is all you need to shift the cache and keep n_keep tokens at the start (i.e. prompt, a.k.a. sink cache).

• First it removes all tokens with positions [n_keep, n_keep + n_discard)
• Then it shifts all tokens with positions [n_keep + n_discard, n_past) with -n_discard
• +1 is to account for BOS token