Rolling buffer cache

[PABannier]

Mistral just released their paper on Mistral 7B.

They use a rolling buffer cache to reduce the memory footprint of the model.

A fixed attention span means that we can limit our cache size using a rolling
buffer cache. The cache has a fixed size of W, and the keys and values for the timestep i are stored
in position i mod W of the cache. As a result, when the position i is larger than W, past values
in the cache are overwritten, and the size of the cache stops increasing. We provide an illustration
in Figure 2 for W = 3. On a sequence length of 32k tokens, this reduces the cache memory usage
by 8x, without impacting the model quality

Do you think this is something we could implement with a ring buffer for instance?

[ggerganov]

The rolling buffer aspect is already supported via the changes in #3228. One has to just manually call llama_kv_cache_seq_rm() to evict tokens from the cache.

The missing part is the support for per-layer positions in the KV cache. I'm still not sure if this is what is being proposed in the paper, but it looks like the KV cache for layer i should hold the cache for positions [i*W, (i+1)*W), but it is not clear how this is supposed to work exactly.

A specific example: let's say we have 2.5*W generated tokens so far.
How does the KV cache for each layer look like?
Is it:

- layer 0: [0,     W)
- layer 1: [W,     2*W)
- layer 2: [1.5*W, 2.5*W)
- layer 3: [1.5*W, 2.5*W)
- layer 4: [1.5*W, 2.5*W)
- layer 5: [1.5*W, 2.5*W)
- layer 6: [1.5*W, 2.5*W)
- layer 7: [1.5*W, 2.5*W)

What happens when we have processed 8.5*W tokens? Maybe this:

- layer 0: [0.5*W, 1.5*W)
- layer 1: [1.5*W, 2.5*W)
- layer 2: [2.5*W, 3.5*W)
- layer 3: [3.5*W, 4.5*W)
- layer 4: [4.5*W, 5.5*W)
- layer 5: [5.5*W, 6.5*W)
- layer 6: [6.5*W, 7.5*W)
- layer 7: [7.5*W, 8.5*W)

For me it is not 100% clear from the paper, but it makes sense to be like this.
If we confirm the above is correct, then we can extend support for per-layer KV cache ops to support this.

[ggerganov]

It's round(2.5*W). I just assume that W is for example 1024 or 2048 or more and 2.5*W is an integer so didn't clarify that

[ggerganov]

I think we can just try the most simple thing and in main we can implement an option that performs:

llama_kv_cache_seq_rm(ctx, 0, params.n_past - W, params.n_past - W + 1);

instead of the existing cache shift. This is trivial to add and we can test if it allows to extend the context up to 32k as proposed. Although I am very doubtful that it will pass the passkey test like this. The SWA strategy is very similar to the existing context shift that we have and we now know that it does not "see" things that leave the context (#3856).

Edit: Also, in the python code, do we see the per-layer strategy implemented? For me it is difficult to read Python as I don't understand the array indexing stuff, so I cannot follow the logic completely.

[KerfuffleV2]

I think it has to be something like this:

def swa(k, i, w = 4096, k_max = 32):
    return i - ((k_max - k) * w)

def go(n_layers = 32, n_tokens = 4096*32, window_size = 4096):
    for i in range(n_tokens):
        if i % 1024 != 0:
            continue
        print()
        for k in range(n_layers):
            min_att = swa(k, i, w = window_size, k_max = n_layers)
            min_att = max(min_att, 0)
            print(f'{i:6}: layer {k:2} attends [{min_att:6}, {i:6}]')

go()

Warning: Running that produces like 5,000 lines of output.

Pregenerated output here: https://gist.github.com/KerfuffleV2/0388066d57918f229ac7e5e376a62cfe

The output of that little Python script seems like it basically matches the paper. edit: Okay, maybe "matches" is an overstatement but I think it's sort-of, kind-of heading in the right direction. It was based on this statement from the site: "Sliding window attention exploits the stacked layers of a transformer to attend in the past beyond the window size: A token i at layer k attends to tokens [i-sliding_window, i] at layer k-1. These tokens attended to tokens [i-2*sliding_window, i]. Higher layers have access to information further in the past than what the attention patterns seems to entail."

Although I am very doubtful that it will pass

Yeah, that doesn't sound like what they were talking about.

Also, in the python code, do we see the per-layer strategy implemented?

From the issue I linked, I don't think so. I was just initially assuming it was there, but apparently that's wrong. I can read the Python, just not the Torch stuff (and maybe not the super fancy array slicing stuff that's common with Torch/numpy code).

[KerfuffleV2]

edit: Oh no, after all this I think it still is wrong. I don't know if I should just delete these posts. I thought I knew English but those two sentences are really hard to parse.

"A token i at layer k attends to tokens [i-sliding_window, i] at layer k-1."

That's talking about the token at i,k attending to tokens from the previous layer at [i-sliding_window, i].

"These tokens" - the tokens from the previous layer at [i-sliding_window, i]?

"attended to tokens [i-2*sliding_window, i]"

So I think my script example is off by one with how it's using k?

---

Actually, at the risk of being annoying with multiple posts, I think the above has to be correct.

Let's say:

>>> i = (4096*30) # 122880
>>> w = 4096 # aka sliding_window

then "A token i at layer k attends to tokens [i-sliding_window, i]"

>>> [i - w, i]
[118784, 122880]

and "at layer k-1 these tokens attended to tokens [i-2*sliding_window, i]"

>>> [i - (2*w), i]
[114688, 122880]

From my script output:

122880: layer  0 attends [     0, 122880]
122880: layer  1 attends [     0, 122880]
122880: layer  2 attends [     0, 122880]
122880: layer  3 attends [  4096, 122880]
122880: layer  4 attends [  8192, 122880]
122880: layer  5 attends [ 12288, 122880]
122880: layer  6 attends [ 16384, 122880]
122880: layer  7 attends [ 20480, 122880]
122880: layer  8 attends [ 24576, 122880]
122880: layer  9 attends [ 28672, 122880]
122880: layer 10 attends [ 32768, 122880]
122880: layer 11 attends [ 36864, 122880]
122880: layer 12 attends [ 40960, 122880]
122880: layer 13 attends [ 45056, 122880]
122880: layer 14 attends [ 49152, 122880]
122880: layer 15 attends [ 53248, 122880]
122880: layer 16 attends [ 57344, 122880]
122880: layer 17 attends [ 61440, 122880]
122880: layer 18 attends [ 65536, 122880]
122880: layer 19 attends [ 69632, 122880]
122880: layer 20 attends [ 73728, 122880]
122880: layer 21 attends [ 77824, 122880]
122880: layer 22 attends [ 81920, 122880]
122880: layer 23 attends [ 86016, 122880]
122880: layer 24 attends [ 90112, 122880]
122880: layer 25 attends [ 94208, 122880]
122880: layer 26 attends [ 98304, 122880]
122880: layer 27 attends [102400, 122880]
122880: layer 28 attends [106496, 122880]
122880: layer 29 attends [110592, 122880]
122880: layer 30 attends [114688, 122880]
122880: layer 31 attends [118784, 122880]

[KerfuffleV2]

@ggerganov It seems like someone finally figured this stuff out. Does this help? mistralai/mistral-inference#40 (comment)