KV Cache 技术

KV Cache

我们提到过，模型的推理阶段的實質就是不断将新生成的 token append 到 input token lists 中，然后将这个新的 input token lists 重新走一遍模型的 forward() 最后取出 last row of logits 作为下一个 token 的概率分布．

所以，我们发现既然我们只需要 last row of logits，那么我們就不需要對很長的 input token lists 算各種 Attention/Position Embedding，相反，我們只需要對那個新 token 進行 Position Embedding 等即可。而對於 Attention 中需要計算的 $S=QK^\top, O=PV$ ，我們也只需要保存 previous $K,V$ 的值即可，而不需要保存所有的 tokens 再重下標計算其 Keys and Values．這樣的技術就是 KV Cache，我們保存之前 tokens 的 $K,V$ Cache，這樣我們只需要對新的 token 進行 $Q,K,V$ Projection 即可，大大降低了計算量．

于是，我們可以進一步區分 Prefill 和 Decode：

Prefill 就是用戶最一開始輸入一個 prompt 交給 LLM 推理，初始化每一層 Attention 的 KV Cache 並推理出第一個 token
Decode 就是 LLM 開始自回歸生成後續 tokens 的步驟

由於 Prefill 和 Decode 對於計算資源的需求是不一樣的，在 KV Cache 這一優化技術上，我們又開發了 P/D 分離架構．

一个简单的 KV Cache 实现

以下代码的来源都是 InfiniTensor 的 llaisys 作业，做了一些简化

一个简单的 KV Cache 实现思路是，我们先开辟出足够的空间给 Key Cache 和 Value Cache 并记录当前的 cache_size

// struct KVCache ...
tensor_t keys;
tensor_t values;
size_t cache_size;

// init()
keys = Tensor::create({max_seq_len, num_kv_head, head_dim}, dtype);
values = Tensor::create({max_seq_len, num_kv_head, vdim}, dtype);
cache_len = 0;

当我们想要 append KV Cache，我们直接执行数据拷贝。

要注意的是，keys->data() 由于是 std::shared_ptr<> 实现，故返回的其实是 std::byte*，需要再乘以 elementSize()

// append() for key
auto begin = keys->data() + cache_size * num_kv_head * head_dim * elementSize();
auto numel = new_len * num_kv_head * head_dim;
std::memcpy(begin, new_keys->data(), numel * elementSize());

// similar for value

然后在我们需要 Key cache 和 Value cache 的时候，返回一个切片，其第一维的范围设在 [0, cache_size)

1
2
3

// slice(dim, start, end); end is exclusive
tensor_t getKeys() { return keys->slice(0, 0, cache_size); }
tensor_t getValues() { return values->slice(0, 0, cache_size); }

所以在 Attention 里基本就是这样调用的：

// ...... (inside attention)
auto kcache = model->kvcaches[layer]->getKeysSlice();
auto vcache = model->kvcaches[layer]->getValuesSlice();
ops::self_attention(attn_out, pos_q, kcache, vcache, scale);