llama : per-layer KV cache

[ggerganov]

This is an updated version of #3457 by @slaren

Also, adds support for quantum K cache: #4312

It's adapted to the latest code on master. Should work with all model arches

• -ngl now goes up to n_layer + 1 instead of n_layer + 3
• use --no-kv-offload or -nkvo to disable KV cache offloading to GPU

llama.h API Changes

• struct llama_context_params
  • add enum ggml_type type_k; // data type for K cache
  • add enum ggml_type type_v; // data type for V cache
  • add bool offload_kqv; // whether to offload the KQV ops (including the KV cache) to GPU
  • remove bool f16_kv
• LLAMA_SESSION_VERSION bump to 3

Limitations

• Context shift does not work with quantum K cache
• V cache quantization is not yet supported

TODO:

• add command-line args to disable KV offloading (enabled by default)
• extensive testing with different models and different -ngl values
• double-check that the KV shift on context overflow works correctly (initial tests indicate that it does)
• update state save/load to the new KV cache and bump LLAMA_SESSION_VERSION

[slaren]

In cparams it looks like it is still possible to offload K and V separately, but wouldn't this change break that?

[ggerganov]

Good point.

Do we ever want to offload them separately? If not, I will change cparams to just offload_kv

[slaren]

I am not sure, but looking at the difference in performance between 34 and 35 offloaded layers in the graphs in #3457 it is probably not worth supporting offloading V only.

[kalomaze]

On the koboldcpp fork, loading just one of the two KV cache layers (42/43 instead of 43/43 or 41/43 specifically was a problem) caused strange generation issues for me with seemingly random tokens appearing.
I am not sure if this is a bug unique to the koboldcpp fork, or if this is behavior caused by llama.cpp, but it helps support the argument that it's probably not worth supporting

[ggerganov]

Ok, since both ways are equally easy to implement, I will not close the door for now and support separate K and V offloading, even if we don't see a use case at the moment

[kalomaze]

I wonder if it makes a difference specifically in the context of models with/without GQA?

[ggerganov]

Ok, since both ways are equally easy to implement,

It's actually simpler to go with either both K and V offloaded or neither. Added --no-kv-offload cli arg

[oobabooga]

Here is a test on my hardware (RTX 3090, 3200 MHz RAM).

Command:

./main \
  -m ../models/wizardlm-70b-v1.0.Q4_K_S.gguf \
  -e \
  -t 6 \
  -tb 12 \
  -n 256 \
  -c 4096 \
  -ngl 42 \
  --ignore-eos \
  -p "Because of the rapid pace of improvement of large language models, evaluation benchmarks have suffered from short lifespans, with state of the art models quickly saturating existing benchmarks, exceeding the performance of human annotators, leading to efforts to replace or augment the benchmark with more challenging tasks. In addition, there are cases of shortcut learning wherein AIs sometimes cheat on multiplechoice tests by using statistical correlations in superficial test question wording in order to guess the correct responses, without necessarily understanding the actual question being asked.\n\nSome datasets have been constructed adversarially, focusing on particular problems on which extant language models seem to have unusually poor performance compared to humans. One example is the TruthfulQA dataset, a question answering dataset consisting of  questions which language models are susceptible to answering incorrectly by mimicking falsehoods to which they were repeatedly exposed during training. For example, an LLM may answer No to the question Can you teach an old dog new tricks because of its exposure to the English idiom you cant teach an old dog new tricks, even though this is not literally true.\n\nAnother example of an adversarial evaluation dataset is Swag and its successor, HellaSwag, collections of problems in which one of multiple options must be selected to complete a text passage. The incorrect completions were generated by sampling from a language model and filtering with a set of classifiers. The resulting problems are trivial for humans but at the time the datasets were created state of the art language models had poor accuracy on them. For example:"

PR (c80b8a2)

Peak VRAM: 23049MiB

llama_print_timings:        load time =   30315,12 ms
llama_print_timings:      sample time =      66,38 ms /   256 runs   (    0,26 ms per token,  3856,82 tokens per second)
llama_print_timings: prompt eval time =    3365,29 ms /   332 tokens (   10,14 ms per token,    98,65 tokens per second)
llama_print_timings:        eval time =  118578,99 ms /   255 runs   (  465,02 ms per token,     2,15 tokens per second)
llama_print_timings:       total time =  122225,02 ms

Master

Peak VRAM: 22377MiB

llama_print_timings:        load time =   31612,69 ms
llama_print_timings:      sample time =      71,63 ms /   256 runs   (    0,28 ms per token,  3573,87 tokens per second)
llama_print_timings: prompt eval time =    4140,64 ms /   332 tokens (   12,47 ms per token,    80,18 tokens per second)
llama_print_timings:        eval time =  119670,77 ms /   255 runs   (  469,30 ms per token,     2,13 tokens per second)
llama_print_timings:       total time =  124069,09 ms

---

The output is stable and coherent in both cases. The conclusion is that the changes in this PR increase the tokens/second for prompt evaluation by +23% while increasing the VRAM usage by +672MB in this particular setup.

[kalomaze]

The prompt processing speed gains are a massive boon it would seem. 13b 4_K_M. Prompt processing was measured in tokens per second for the same exact prompt (~600 tokens) at different layer offload counts, with 4096 tokens ctx.

Generation tps as well, there's some milder gains there. If we focus on 20-35 it looks like:

CPU: i7-6700
GPU: RTX 3060

[kalomaze]

Something that surprised me when testing is that both thread counts seem to matter, perhaps equally now. In the past, it was said that you should focus primarily on batched threads instead of regular threads for CPU offloading... But if I set both to 7, I get better speeds overall. (My CPU has 4 cores and 8 threads)

On 13b, 6_K quantization, 31/41 offloaded layers on RTX 3060 (12GB VRAM):

Threads count 7, Batched threads count 7: 127 t/s prompt processing, 2.5 t/s generation speed
Threads count 2, Batched threads count 7: 137 t/s prompt processing, 1.3t/s generation speed
Threads count 8, Batched threads count 4: 108 t/s prompt processing, 2.09t/s generation speed
Threads count 8, Batched threads count 8: 135 t/s prompt processing, 2.76t/s generation speed

Seems to be optimal to set both to your total CPU thread count, or perhaps 1 layer less for both might more more stable overall.

[ggerganov]

Plan to merge this tomorrow - haven't tested much yet. Hope it does not break something

The CUDA performance with quantum cache is a bit disappointing. Hopefully we will fix this in the future

[kalomaze]

Plan to merge this tomorrow - haven't tested much yet. Hope it does not break something

The CUDA performance with quantum cache is a bit disappointing. Hopefully we will fix this in the future

I have been testing my sampler experiments using it and I don't notice any breaking changes, but that was before you added the K layer quantization stuff apparently if I read the other PR right?

[kalomaze]

The CUDA performance with quantum cache is a bit disappointing.

Thankfully, I'd argue "disappointing" is potentially misleading, looking at the benchmark provided in the other PR. It was my presumption that quantization actually reduces generation throughput because of the overhead of reducing precision, but if it means you can fit a model entirely into VRAM that wouldn't fit before, or more layers of a model into VRAM specifically (thanks to this wonderful spread KV layers PR), that alone is much more important than the actual compute being slower because VRAM is much, much faster.

TLDR; the speed gains from better VRAM utilization probably outpace the speed loss from having to quantize the KV cache if you're VRAM bound. My guess is that high compute low VRAM benefits the most (e.g. most consumer gaming GPUs!!!)

[ggerganov]

I have been testing my sampler experiments using it and I don't notice any breaking changes, but that was before you added the K layer quantization stuff apparently if I read the other PR right?

Yes, thank you for testing. The changes that I expect to break something are related to the "per-layer KV cache". The other changes from #4312 should be good I think.

"breaking change" means that the API has changed - i.e. it is not backwards compatible.

Thankfully, I'd argue "disappointing" is potentially misleading, looking at the benchmark provided in the other PR.

It's disappointing because it could be much better, but the quantum kernels are not well adapted for multi-head tensors (i.e. ne[2] > 1). In that case, we are doing unnecessary copies and launching separate kernels for each head, while it could be done more efficiently in a single pass. I tried to make a quick workaround for the single GPU case, but it's ugly.

[Dampfinchen]

So per-layer KV cache has far exceeded my expectations. Prompt processing went down from 10 ms/t to 5.5 ms/t on a 13B model and 21 layers offloaded. Truly amazing.

However it's a different story with quantum k cache, I can't notice a memory relief when I set -ctk q8_0, like at all. Does it only work with non-k models? Because I have a Q4K_S model.

[x4080]

Is this benefits apple silicon as well ?

[CyborgArmy83]

Is this benefits apple silicon as well ?

Curious about this as well. Would be nice if it could speed up prompt processing on Apple Silicon.

[LostRuins]

@Dampfinchen hmm I've tried the per-layer KV offloading, and I've had a the opposite experience from you, I think that it feels like an overall downgrade (to me), at least for the 6GB RTX 2060.

I tested both approaches with the max number of GPU layers I could offload before going OOM (at 2048 ctx).

main: build = 1620 (fe680e3)

max 19 layers, per layer kv offload

.\main.exe -m e:\LLaMA\models\LLaMA2-13B-Tiefighter.Q4_K_M.gguf -ngl 19 -n 32 -f c:\test\test.txt -c 2048
llama_print_timings:        load time =    3418.61 ms
llama_print_timings:      sample time =       8.17 ms /    32 runs   (    0.26 ms per token,  3916.77 tokens per second)
llama_print_timings: prompt eval time =    1797.20 ms /   202 tokens (    8.90 ms per token,   112.40 tokens per second)
llama_print_timings:        eval time =    6757.10 ms /    31 runs   (  217.97 ms per token,     4.59 tokens per second)
llama_print_timings:       total time =    8597.05 ms

max 23 layers, no kv offload

.\main.exe -m e:\LLaMA\models\LLaMA2-13B-Tiefighter.Q4_K_M.gguf -ngl 23 -nkvo -n 32 -f c:\test\test.txt -c 2048
llama_print_timings:        load time =    3705.89 ms
llama_print_timings:      sample time =       8.25 ms /    32 runs   (    0.26 ms per token,  3880.20 tokens per second)
llama_print_timings: prompt eval time =    1993.56 ms /   202 tokens (    9.87 ms per token,   101.33 tokens per second)
llama_print_timings:        eval time =    6076.94 ms /    31 runs   (  196.03 ms per token,     5.10 tokens per second)
llama_print_timings:       total time =    8107.60 ms

I then switched to my KCPP build which sets GGML_CUDA_FORCE_MMQ to true, allowing me to offload even more layers without going OOM.

max 25 layers, mmq, no kv offload

ContextLimit: 234/2048, Processing:2.58s (12.8ms/T), Generation:5.67s (177.2ms/T), Total:8.25s (3.88T/s)

Overall, trading a ~40% increase PP speed for a 22% reduced generation speed does not really feel like a worthwhile trade to me, especially if we take into account the reduced need for rapid prompt processing while using context shifting. I think perhaps you were testing at -c 512 which has a much smaller difference in number of layers offloaded. At 4096 ctx, the difference will probably be even greater.

Edit: @Pyroserenus suggested that most of the gains are only apparently when comparing between nearly offloaded cases, so here's some 7B tests, which I am able to offload all layers for.

per-layer KV offload, 7B, 28 layers, cublas
ContextLimit: 1606/2048, Processing:3.36s (2.2ms/T), Generation:3.62s (56.6ms/T), Total:6.98s (9.17T/s)

per-layer KV offload, 7B, 28 layers, MMQ
ContextLimit: 1606/2048, Processing:5.02s (3.3ms/T), Generation:3.72s (58.1ms/T), Total:8.74s (7.32T/s)

no KV offload, 7B, 28 layers, cublas
ContextLimit: 1606/2048, Processing:6.84s (4.4ms/T), Generation:5.43s (84.8ms/T), Total:12.27s (5.22T/s)

no KV offload, 7B, 28 layers, MMQ
ContextLimit: 1606/2048, Processing:8.17s (5.3ms/T), Generation:5.06s (79.0ms/T), Total:13.22s (4.84T/s)

full offload, 7B, cublas
ContextLimit: 1606/2048, Processing:2.38s (1.5ms/T), Generation:2.18s (34.0ms/T), Total:4.56s (14.03T/s)

indeed for the last few layers of a model, the PP speed difference between no KV offload and per-layer KV offload does differ much more.

[Ph0rk0z]

So for full offload this will not make much difference? I think I have to test this via the built in server rather than python bindings. I would hope that 8bit KV would have helped something.

[Dampfinchen]

@LostRuins That's because you are just processing 200 tokens, while I am using a full 2K prompt and also generate 180 Tokens instead of 31. The difference offloaded KV cache makes, is huge, both in generation and prompt processing speed.

22 Layers, per-layer KV cache offloaded, 6 GB 2060

25 Layers, per layer KV Cache disabled