ref #3479

Description

I wanted to find out what is currently llama.cpp's main limitation when running batched decoding, so I ran a few tests on different hardware to profile mainly the self attention overhead when using the existing unified KV cache implementation (#3228).

Below are the results on 3 different hardware:

• M2 Ultra
• RTX 4080
• A100

I'm using the batched-bench tool to run PP + TG for different number of batches:

• Shared prompt of 512 tokens
• Each batch generates 128 tokens
• Batch nums: 1,2,3,4,5,6,7,8,16,32
• Models: LLaMA 1B and 7B, F16 precision

This PR adds a hack to allow for conveniently turning on and off some of the attention ops via environment variables:

• Set env SKIP_KQ_KQV=1 to skip the 2 matrix multiplications KQ and KQV
• Set env SKIP_KQ_ALL=1 to skip all attention ops (KQ, KQ_scaled, KQ_masked, KQ_soft_max, KQV)
• Without setting the env this branch runs normally as master

I've also performed 2 custom diffs for Metal and CUDA to run the full computation but force only 1 KV head to be computed during matrix-multiplication:

diff --git a/ggml-cuda.cu b/ggml-cuda.cu
index 654d3632..28ed98de 100644
--- a/ggml-cuda.cu
+++ b/ggml-cuda.cu
@@ -6653,13 +6653,13 @@ static void ggml_cuda_op_mul_mat(
 
     const int64_t ne00 = src0->ne[0];
     const int64_t ne01 = src0->ne[1];
-    const int64_t ne02 = src0->ne[2];
+    const int64_t ne02 = 1;
     const int64_t ne03 = src0->ne[3];
     const int64_t nrows0 = ggml_nrows(src0);
 
     const int64_t ne10 = src1->ne[0];
     const int64_t ne11 = src1->ne[1];
-    const int64_t ne12 = src1->ne[2];
+    const int64_t ne12 = 1;
     const int64_t ne13 = src1->ne[3];
     const int64_t nrows1 = ggml_nrows(src1);

diff --git a/ggml-metal.m b/ggml-metal.m
index c908106b..4b7a5226 100644
--- a/ggml-metal.m
+++ b/ggml-metal.m
@@ -736,7 +736,7 @@ void ggml_metal_graph_compute(

                 const int64_t  ne00 = src0 ? src0->ne[0] : 0;
                 const int64_t  ne01 = src0 ? src0->ne[1] : 0;
-                const int64_t  ne02 = src0 ? src0->ne[2] : 0;
+                      int64_t  ne02 = src0 ? src0->ne[2] : 0;
                 const int64_t  ne03 = src0 ? src0->ne[3] : 0;

                 const uint64_t nb00 = src0 ? src0->nb[0] : 0;
@@ -746,7 +746,7 @@ void ggml_metal_graph_compute(

                 const int64_t  ne10 = src1 ? src1->ne[0] : 0;
                 const int64_t  ne11 = src1 ? src1->ne[1] : 0;
-                const int64_t  ne12 = src1 ? src1->ne[2] : 0;
+                      int64_t  ne12 = src1 ? src1->ne[2] : 0;
                 const int64_t  ne13 = src1 ? src1->ne[3] : 0; UNUSED(ne13);

                 const uint64_t nb10 = src1 ? src1->nb[0] : 0;
@@ -786,6 +786,11 @@ void ggml_metal_graph_compute(
                 //            dst->name);
                 //}

+                if (dst->op == GGML_OP_MUL_MAT) {
+                    ne02 = 1;
+                    ne12 = 1;
+                }
+
                 switch (dst->op) {
                     case GGML_OP_NONE:
                     case GGML_OP_RESHAPE:

All these options allow us to measure the overhead from the following computations individually:

• KQ and KQV matrix multiplications for all heads
• KQ and KQV matrix multiplications per attention head
• KQ_scale + KQ_masked + KQ_soft_max

Results

These are the raw numbers that I measured. In each file, first are the 7B runs, followed by the 1B runs:

• M2 Ultra perf-m2-ultra.txt
• RTX 4080 perf-4080.txt
• A100 perf-A100.txt

Here I'll inline part of the A100 results for convenience. For the rest of the results, checkout the text files above:

normal

LLAMA_CUBLAS=1 make -j && ./batched-bench /workspace/openllama-7b/ggml-model-f16.gguf 4608 1 99 0 512 128 1,2,3,4,5,6,7,8,16,32

SKIP_KQ_ALL=1

LLAMA_CUBLAS=1 make -j && SKIP_KQ_ALL=1 ./batched-bench /workspace/openllama-7b/ggml-model-f16.gguf 4608 1 99 0 512 128 1,2,3,4,5,6,7,8,16,32

normal + force 1 KV head

LLAMA_CUBLAS=1 make -j && ./batched-bench /workspace/openllama-7b/ggml-model-f16.gguf 4608 1 99 0 512 128 1,2,3,4,5,6,7,8,16,32

Observations

• Using the text generation times (T_TG) for 8 batches, with Metal the entire self attention calculation amounts to ~12% of the text generation time for 1B and ~8% for 7B. This means that even with infinitely fast KV cache and self attention, the Metal performance cannot improve by more than 1.13x for 1B and 1.09x for 7B
• The same calculations for CUDA show that the self attention amounts to (64% 1B, 44% 7B) for RTX 4080 and (76% 1B, 77% 7B) for A100 of the text generation time. So the maximum theoretical improvement in this case is (2.6x 1B, 1.8x 7B) for RTX 4080 and (3.7x 1B, 4.3x 7B)
• Of course, we cannot ever have an infinitely fast self attention that takes 0 time, but from the above observations we can see that the KQ and KQV matrix multiplications on CUDA take a much larger toll compared to Metal, both for prompt processing and for more than 1 batches
• Another observation is that on CUDA, the KQ and KQV processing time scales linearly with the number of KV heads for more than 1 batch, while on Metal where we have a custom matrix-matrix multiplication kernel, the computation scales much better

If my analysis above are correct, there is a significant speedup to be gained for CUDA - both for batched decoding and for prompt processing. I'm not familiar with the best practices for CUDA, but I think we should either:

• Utilize CUDA streams for each separate KV head (currently we run n_head CUBLAS GEMMs in a single CUDA stream). If I remember correctly, we have tried utilizing CUDA streams, but only for single-batch decoding. Probably we have to revisit?
• Implement a custom 3D matrix-matrix multiplication kernel similar to the one that we have in Metal. Use it for the KQ and KQV ops where ne02 > 1 and ne12 > 1

These observations could also explain the poor performance observed for speculative decoding on A100 reported here: #3649

Reproducing

If anyone is interested in re-running the CUDA tests above, I used the following script:

#!/bin/bash

# setup deps
apt-get update
apt-get install git-lfs cmake cmake-curses-gui vim ruby
git-lfs install

# this is useful to git clone repos without doubling the disk size due to .git
git clone https://github.com/iboB/git-lfs-download
ln -sfn /git-lfs-download/git-lfs-download /usr/local/bin/git-lfs-download

# download data
cd workspace

git-lfs-download https://huggingface.co/PY007/TinyLlama-1.1B-Chat-v0.3
git-lfs-download https://huggingface.co/openlm-research/open_llama_7b

# llama.cpp
cd /
git clone https://github.com/ggerganov/llama.cpp

cd llama.cpp

ln -sfn /workspace/open_llama_7b            /workspace/openllama-7b
ln -sfn /workspace/TinyLlama-1.1B-Chat-v0.3 /workspace/tinyllama-1b

pip install -r requirements.txt
python3 convert.py /workspace/tinyllama-1b --outfile /workspace/tinyllama-1b/ggml-model-f16.gguf --outtype f16
python3 convert.py /workspace/openllama-7b --outfile /workspace/openllama-7b/ggml-model-f16.gguf --outtype f16

LLAMA_CUBLAS=1 make -j && ./batched-bench /workspace/tinyllama-1b/ggml-model-f16.gguf 4608 1 99 0 512 128 1,2,3,4,5,6,7,8,16,32
LLAMA_CUBLAS=1 make -j && ./batched-bench /workspace/openllama-7b/ggml-model-f16.gguf 4608 1 99 0 512 128 1,2,3,4,5,6,7,8,16,32

On runpod, the above RTX 4080 and A100 tests cost me a total of ~$1.11 to perform. You would need ~40 GB storage.
Alternatively, you can run them locally - the tests require 16GB VRAM

cc @slaren @JohannesGaessler for any comments and insights