How can I swap two KV cache slots?

[KerfuffleV2]

Let's make this simple and say I want to swap the slot at position 10 with the one at position 25, we're running on CPU only and I only care about LLaMA. So this would happen in llm_build_llama around the same place the KV cache shifting stuff happens.

std::swap(kv_self.cells[10], kv_self.cells[25]);

Now something has to happen with kv_self.k and kv_self.v, but what?

Once I know this, I'll probably be able to implement KV cache defragmenting which can make a big difference for parallel generation.

[ggerganov]

Before implementing the defragmentation, we should understand better #3479. KV cache defragmentation would improve things, but not by a large amount.

The performance concerns in #3479 might require reworking the KV cache all together if it turns out that the current way of storing it does not allow to scale the batch size (see #3479 (comment))

[KerfuffleV2]

KV cache defragmentation would improve things, but not by a large amount.

I think it can make a pretty big difference in some cases. Just for example with my testing of the parallel example using -np 10 -cb -ns 50 -c 1200 -ngl 40 --seed 123 (1200 is pretty close to the minimum size necessary):

Total prompt tokens:    782, speed:  4.94 t/s
Total gen tokens:      4177, speed: 26.38 t/s
Total speed (AVG):           speed: 31.32 t/s
Cache misses:          2393

Using -c 1700:

Total prompt tokens:    782, speed: 11.16 t/s
Total gen tokens:      4275, speed: 61.01 t/s
Total speed (AVG):           speed: 72.18 t/s
Cache misses:             0

Those are with the kv slot pull. I couldn't get master pre-merge to even run with a 1200 context. Running the pre-merge version with -c 1700:

Total prompt tokens:    782, speed:  8.71 t/s
Total gen tokens:      4160, speed: 46.32 t/s
Total speed (AVG):           speed: 55.02 t/s
Cache misses:           514

So in the scenario of continuous batching and requests using most of the available context size, it can make a really big difference. Note: like I mentioned I'm not sure the comparison between the version with my kv slot changes and the one without it is really fair, so even though this makes it look good it may not actually be meaningful. The point is that the cache misses really, really hurt performance in that scenario.

Everything else aside, I'd still be interested in finding out how to swap slots with the current KV cache implementation. Just because it's a puzzle I haven't been able to solve so far.

(I've been following #3479 closely!)

[ggerganov]

Without thinking too deeply, I think the defragmentation should happen through the compute graph in a similar way we do the K shift:

https://github.com/ggerganov/llama.cpp/blob/ff3bad83e29e3009010cbc923bebd769055eaa7f/llama.cpp#L4367-L4392

I.e. upon request, we insert extra nodes in the graph to perform the shift. In this case it will perform defragmentation.
The main problem is that we would likely need a dedicated op + GPU kernels for this which seems like too big effort.
I am wondering if we try to somehow defragment the cache using the existing ops: views + copies, etc. But not sure how exactly.

The other option to consider is outside the graph:

• we download the KV cache from device (if it is there)
• perform the defrag on the CPU (because it is easier)
• copy the KV cache back to the device (if needed)

But this seems too custom and might involve extra backend API which is better to avoid.

[KerfuffleV2]

Being able to defrag only if the KV cache isn't offloaded is probably better than nothing.

If you have the time, can we start with the simplest possible example of just swapping two slots? So you can assume:

1. It's a normal LLaMA model
2. Is only on the CPU
3. Doesn't have to be clean/nice (can just use memcpy or whatever on the tensor data)
4. Doesn't even really need to be real code, pseudocode just showing what's needed is probably fine

Right now, unfortunately I don't even really know what a "slot" looks like in terms of structure and the tensors/etc that exist.

[ggerganov]

The K and V buffers are arrays of floats with the following size:

n_elements = n_embd*n_layer*n_ctx;

n_embd is defined by the model. It is equal to n_embd_head*n_head. n_embd_head is typically ~64, 128, 256. n_head is typically ~32

n_layer is defined also by the model (~32)

n_ctx is provided by the user. Can be any number, for most models it does not make sense to be more than the training context of the model

One slot corresponds to n_embd*n_layer numbers. I.e. the slot for a given position goes across the layers (caveat: we might need to update this to support Mistral sliding window per layer, given that my understanding is correct #3581 (comment)).

We store data for n_tokens slots at layer il in cache starting position kv_head like this:

https://github.com/ggerganov/llama.cpp/blob/c9708b4d5ab3b6c21f7884a16ed9f03192ad861b/llama.cpp#L3398-L3419

Note that the K is stored sequentially, but the V is transposed, so a swapping solution has to account for that.

To swap two slots, you would need to make the same views on kv_self.k and kv_self.v with the corresponding src and dst kv_head and n_tokens == 1 and swap the contents of those views.