KV Cache Dimension

[rroykuo]

Hi~
I'm figuring out how kv cache be placed into llama_kv_cache (struct defined in llama.cpp)
Now I'm trying llama3 8B model and here are parameters about kv cache:

• n_embd: 4096
• n_head: 32
• n_head_kv: 8
• n_layer: 32
• n_embd_head_k & n_embd_head_v: 128
• n_embd_k_gqa & n_embd_v_gqa: 1024
• n_ctx: 512

As llama3 takes grouped-query attention, so each token would be 8(n_head_kv) * 128(n_embd_head_k(v)) = 1024 (n_embd_k(v)_gqa) tensors for each layer of key or value cache.

Here's a code snippet under llama_kv_cache_init

cache.k_l.reserve(n_layer);
cache.v_l.reserve(n_layer);

    for (int i = 0; i < (int) n_layer; i++) {
        struct ggml_context * ctx = offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();
        ggml_tensor * k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size);
        ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, n_embd_v_gqa*kv_size);
        ggml_format_name(k, "cache_k_l%d", i);
        ggml_format_name(v, "cache_v_l%d", i);
        cache.k_l.push_back(k);
        cache.v_l.push_back(v);
    }

It stores embedding dimension and ith token dimension in 1D dimension for each layer, so I think there should be 1024 tensors for each generated token for a layer of key or value.

I used debugger to trace how kv cache varies (I only observed k_l[0])
As used=4 in kv_self it means there are four tokens be recorded and there should be 1024(n_embd_k_gqa) * 4(tokens) = 4096 elements in k_l[0].
But I only found 512 * 4 = 2048 non zero elements placed in index [0] ~ [2047] ofk_l[0]

For each generated token, it only appended 512 tensors sequentially to k_l[0].

Most weird thing is that it appended elements non-sequentially to later index position

Initially, elements are all zero and it changed to non-zero value as key cache appended to k_l[0]

[ggerganov]

The KV cache is F16 while you are looking at it through F32 pointer

[rroykuo]

Oh! I didn't notice that🤣 Thanks!