一个简单的 KV Cache 实现

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

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

1
2
3
4
5
6
7
8
9
// 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()

1
2
3
4
5
6
// 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 里基本就是这样调用的：

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