llama : mitigate KV cache fragmentation

[ggerganov]

With the new unified KV cache implementation from #3228 we now support batched decoding.

At runtime, depending on the workload and the length of the decoded sequences, the KV cache can become fragmented. If we think of the cache as an array and each cell being free, or belonging to a certain sequence, then we could end up with many short segments of free cells, instead of one big segment with all the free cells. This hinders the performance since for each new batch, we have to find a free cache segment that can "hold" the entire batch and when we cannot do so, we have to start splitting the batch in smaller batches to be able to fit it.

One possible mitigation is from time to time (based on some logic, or based on user request) to "defragment" the cache. This can be implemented in a very similar way to the existing KV shift functionality, where we add extra graph nodes when there is need to defragment the cache.

Other approaches might be possible. For example, based on the batch size, allocate the KV data either at the start or at the end of the cache. This way we will keep the larger "prompt" segments next to each other at the start of the cache and the small "text" segments at the end of the cache. Not sure if this would be lead to less fragmentation, but similar strategies can be explored

[KerfuffleV2]

When running parallel generation, the kv cache free slot search seems to take up a pretty noticeable amount of time even when it's possible to find an entry. Generation gets slower and as the size shrinks.

I added some debug code to llama_kv_cache_find_slot:

    {
        size_t free_count = 0;
        printf("\033[H\033[2J");
        for (uint32_t i = 0; i < cache.size; i++) {
            if ((i % 80 == 0) && i != 0) putchar('\n');
            if (cache.cells[i].pos == -1) {
                free_count++;
                putchar('.');
            } else {
                putchar('#');
            }
        }
        printf("\n{free:%zu}\n", free_count);
    }

When doing normal parallel generation like the batched example, fragmentation isn't an issue since everything stays in sync. Tests used -np 5 -c 2026 -n 400:

It's terrible with my seqrep sampler using rewind mode though:

Since the items for a sequence id are never contiguous, anything that dynamically removes "clients" or deletes items in a client's history is going to cause terrible fragmentation.

It doesn't help with the fragmentation issues but one way to speed up the normal case is to just save the position just past the slot that was found and start the next search from there (assuming it's not already at the end of the kv cache). The common case will be to append, so it should instantly find a slot that way instead of having to search through the cache from the very beginning every time.

[ggerganov]

It doesn't help with the fragmentation issues but one way to speed up the normal case is to just save the position just past the slot that was found and start the next search from there

Yes, I added this line at some point, but I think it is no longer useful:

https://github.com/ggerganov/llama.cpp/blob/f72f8f22c9cb60465b2e79df2767e4ba9604e576/llama.cpp#L4077-L4081

We should probably just delete this line.
Overall, the KV cache management can be improved with proper data structures. The initial implementation was primarily aiming to provide the functionality

Btw, the correct way to decode in situations with fragmented cache is demonstrated in parallel:

https://github.com/ggerganov/llama.cpp/blob/f72f8f22c9cb60465b2e79df2767e4ba9604e576/examples/parallel/parallel.cpp#L275-L293

This will guarantee full cache utilization, although it would not be optimal since at some point you would start splitting your batches into smaller and smaller batches

[KerfuffleV2]

Commenting out the kv_self.head = 0 line seems to make fragmentation worse:

Not sure why it apparently leaves the first 2 slots empty.

Dumping n_tested does show that it immediately finds an entry without having to search though. I actually don't notice a difference in speed though, so maybe I was wrong about kv cache slot searching being the reason it gets slower over time. Maybe that's just something related to decoding in parallel mode? (It doesn't seem to happen when -np 1.)

Btw, the correct way to decode in situations with fragmented cache is demonstrated

Thanks, I actually saw that. Just kind of playing around with the parallel generation stuff (unfortunately I didn't get a chance to while this was in testing).

Do the items in a batch have to be laid out like:

0123012301230123....

so if you delete 0, you get:

.123.123.123.123....

or could it be:

0000111122223333....

Basically if you knew advance how many sequences you were going to use, you could dedicate a section of the kv cache to each sequence id. There would be limitations to that approach though, so I'm not really saying it would be the way to go even if possible. Just curious.

[ggerganov]

The items in a batch can be laid out in any order. The current limitation is that the batch will have to be put sequentially inside the cache. I.e. we don't have means to distribute the batch items into arbitrary positions in the cache. To achieve this, the only way currently is to split the batch into smaller batches

[KerfuffleV2]

@ggerganov

Does something like this look reasonably sane? edit: Maybe not. The fact that it produced pretty reasonable output made me think it was close, but now I really don't know.

static uint32_t llama_kv_cache_compact(struct llama_kv_cache & cache) {
    uint32_t counter = 0;
    int32_t free_slot = -1;
    for (int32_t i = 0; i < cache.size; i++) {
        if (cache.cells[i].pos < 0 || cache.cells[i].seq_id.empty()) {
            if (free_slot == -1) {
                free_slot = i;
            }
            // Probably not necessary but just to be safe.
            cache.cells[i].pos = -1;
            cache.cells[i].seq_id.clear();
            continue;
        }
        if (free_slot == -1) {
            continue;
        }
        std::swap(cache.cells[i], cache.cells[free_slot]);
        // Or maybe the free cell delta needs to be preserved?
        cache.cells[i].delta = 0;
        cache.cells[free_slot].delta += i - free_slot;
        cache.has_shift = true;
        counter++;
        if (cache.cells[free_slot + 1].pos < 0 || cache.cells[free_slot + 1].seq_id.empty()) {
            free_slot++;
        } else {
            free_slot = -1;
        }
    }
    cache.head = free_slot != -1 ? free_slot : 0;
    return counter;
}

It just moves through the kv cache keeping track of the first available free slot. Then if it hits an occupied slot, it swaps that with the free slot and moves the free slot index forward if the slot past it is also free.

I think the logic for that part is fine but I'm much less sure about the kv cache delta manipulation part. It actually seems to work, however, some kind of weird stuff:

1. Even without running the compact function, commenting/uncommenting the kv_self.head = 0 line causes different generation even with the same seed (and seems like it happens even on CPU though I haven't tested that extensively).
2. Defragging the cache also causes different generation.

If it seems like this is on the right track, I can make a pull. Also even running it once per llama_decode doesn't seem to have much of a speed impact.

---

The current limitation is that the batch will have to be put sequentially inside the cache.

Thanks, that's actually what I was doing a really bad job of asking. So it's a current limitation, but potentially could be changed in the future.

---

edit: I figured out why without cache.head = 0 the layout looked like ..#### initially. It's because while the warmup run in common.cpp deletes the tokens it doesn't reset head, so head starts out at 2. Something like this fixes it:

static void llama_kv_cache_tokens_rm(struct llama_kv_cache & cache, int32_t c0, int32_t c1) {
    if (c0 < 0) c0 = 0;
    if (c1 < 0) c1 = cache.size;
    if (c0 == 0 && c1 == int32_t(cache.size)) {
        cache.head = 0;
    }

[ggerganov]

You are swapping just the cell info that lives in the CPU memory. You have to also swap the actual KV data stored in llama_kv_cache.buf and since it can be on the GPU, it becomes more tricky to implement.

Yup, I forgot about the warmup leaving a trace in the cache. Your suggested fix should be fine.

[KerfuffleV2]

Is it possibly time to revisit this, since it doesn't seem like more major KV cache changes happening right now? I'm still interested in finding the answer to #3507

[MrJackSpade]

On the subject of "Splitting Batches" what I've done locally to help mitigate the issue, is to identify a block of the cache that requires moving the fewest number of cells and then prepending those cells to the batch I'm attempting to evaluate and clearing the cells.

It increases the size of the batch, but since the bottleneck (For me) is number of batches specifically, it still ends up being substantially faster than trying to split/search for free space repeatedly. The difference between a 21 token batch and an 18 token batch is negligible, but the difference between an 18 token batch and 3x6 token batches is huge

This is kind of a lazy way to go about it because I'm not actually moving the data in memory, I'm recalculating it by sending it back through decode, so its definitely not an "ideal" solution. I'm also not even sure its a functional solution outside of "it appears to work so far"

[ggerganov]

That's an interesting approach. Definitely worth looking into it. Likely moving the data is not the right approach and techniques like the one you describe could yield better results.

Additionally, I think if we change the "meta information" data structure to not be per-cell, but per-batch we can also gain. Currently, we have meta data for each separate cell in cell[i] that straight up corresponds to the KV buffer at index i - i.e. one-to-one mapping between meta cells and data cells. Instead of cell[i] we can be storing metadata for entire batches (batch[k]) and have a "pointer index i" to the KV buffer where this batch is located. I.e. "decouple" the meta info from the actual data via an indirection (the pointer index). Searches could become more efficient I think

[MrJackSpade]

Just messing around with the algorithm I managed to get a 10x speed up on a fairly limited scope test, by breaking up the cache into pages and then tracking the number of free cells on each page.

When a new batch is requested, the number of pages required is calculated and then the head moves across the cache += page_size until page[head / page_count] => page[(head / page_count) + required_pages] sum > n_tokens. Though having a total free some exceeding the batch size doesn't guarantee they would be contiguous, this has a fairly high benefit once the cache becomes highly fragmented since it allows skipping the majority of the cache when looking for free blocks.

During the actual final mapping step I just said page[(head + i / page_size)]-- to update the page free space tracking.

The up-side is that a change like this would require very little in the way of modifying the cache and the additional memory usage is negligible, the downsides are that the cache clearing functions would be responsible for updating the pages which could lead to bugs with synchronization, and the actual page size is largely arbitrary as performance would depend on the batch sizes being requested.

My test involved a max batch of 64 and a page size of 32, and requested random batch sizes (with the / 2 split on failure) until the cache was full.

I did >> instead of / for obvious reasons which locks the page size to 2^n but avoids issues with division.