很久之前就想复刻一下 vLLM 中的 paged Attention，相关的文章还有论文也基本上都看过一遍，苦于本人很懒，一直处于想做，但又没那么想做的状态，直至上周才下定决心把这个玩意弄完，遂去看了 nano-vLLM 的源码，仔细学习一下相关的架构和设计，然后编写了适配个人推理引擎项目的 cpp 实现。
struct Block;
using block_t = std::shared_ptr<Block>;class BlockManager;
using BlockManager_t = std::shared_ptr<BlockManager>;// kv cache block
struct Block{int _block_id; // 块号long long _hash = -1; // 哈希值 前缀+当前块信息std::vector<int64_t> _token_ids; // 当前块缓存的token信息// std::atomic<size_t> _ref_count = 0;
};// kv cache block manager
class BlockManager{
private:tensor_t _k_cache; // 底层连续缓存张量tensor_t _v_cache;size_t _token_num; // 单个block所包含token缓存个数size_t _block_num; // 维护的总kv cache块数std::vector<block_t> _blocks; // 维护所有块的信息std::unordered_map<long long, int> _hash_2_block; // 哈希值转块号std::queue<int> _free_block_ids; // 可用的block的块号std::unordered_set<int> _used_block_ids; // 已经使用的block的块号std::stack<int> _last_used_block; // TODO: 后续实现为 LRU cache DataType_t _dtype; // 缓存数据类型DeviceType_t _device_type; // 缓存设备类型int _device_id; // 缓存所在设备private:// 自定义哈希函数求解: 计算 token_ids[l: r + 1) 及其前缀信息 prefix 的哈希值long long _compute_hash(const int64_t *token_ids, size_t l, size_t r, long long prefix = -1);// 重新初始化blockvoid _reset(block_t block);// 更新blockvoid _update(block_t block, long long hash, const int64_t *token_ids, size_t l, size_t r);BlockManager(const size_t token_dim, // 单个token的k/v cache占用dtype个数const size_t token_num = 32, // 单个block所包含token总数 const size_t block_num = 8, // block总数DataType_t dtype= LLAISYS_DTYPE_F32, // cache的数据类型DeviceType_t device_type = LLAISYS_DEVICE_NVIDIA, // cache所处设备类型int device_id = 0); // cache所处设备号public:static BlockManager_t create(const size_t token_dim, const size_t token_num = 128, const size_t block_num = 8, DataType_t dtype= DTYPE_F32, DeviceType_t device_type = DEVICE_NVIDIA, int device_id = 0); ~BlockManager() = default;// Prevent copyBlockManager(const BlockManager &) = delete;BlockManager &operator=(const BlockManager &) = delete;// Prevent moveBlockManager(BlockManager &&) = delete;BlockManager &operator=(BlockManager &&) = delete;// 根据token_ids和现有缓存信息分配block块表（对应设备上） 返回匹配前缀块总长size_t allocate(const int64_t *token_ids, // 要分配的序列的信息const size_t ntoken, // 序列总长std::vector<int> &block_ids); // 分配的 block 信息 块表/总块数 (输出)// 根据id获取对应的块的指针block_t block(int block_id) { ASSERT(block_id >= 0 && static_cast<size_t>(block_id) < this->_block_num, "BlockManager_block: block_id out range.");return this->_blocks[block_id]; }// 获取底层缓存的张量tensor_t k_cache() { return this->_k_cache; };tensor_t v_cache() { return this->_v_cache; };// 返回总块数size_t block_num() { return this->_block_num; }// 返回单个block所占token数size_t token_num() { return this->_token_num; }
};
其实还没写完，因为 LRU 缓存我还没写，并且还只是单线程的情况。虽然我没去深究 nano-vLLM 中具体是怎么实现的，但凭借着对整体功能和架构的理解，我基于自己的推理引擎搞了这么个缓存管理的玩意，然后就改进了原有的 flash attention，非常好改，加个块表，计算逻辑不变，就缓存加载时改一下。
由于之前的架构非常不优雅，所以我就对一些部分进行了重构，并且在新的前向链路中增添了 kv cache 管理和 paged attention，再对 paged attention 做过单元测试后，我本以为整个链路会非常顺利的完成测试，当然模型有输出，只是全是乱码（没崩住），成功了一半。
然后就开始分析，正如题目所描述，由于这里我一下引入了分块 kv cache 和 paged attention，再加上整个系统好久没碰了，之前的修改忘记做没做测试了，所以我也不知道哪个部分错了，并且模型前向过程中的张量大的要死，直接打印又看不出什么信息，直接开始控制变量。
首先对整个链路去除 kv cache、paged attention 确保其他算子和流程的正确性，非常 amazing，输出是对的，那么问题就在我们新引入的分块 kv cache 或者 paged attention 上。推理一下由于现在还是单请求，并且 kv cache 就是连续的块，所以分块搬等价于连续搬。所以我构造了一种情况：直接拿现成的 kv cache 构造成可被 flash attention 可用的情况，验证了链路在连续搬运 kv cache 且 无 paged attention 的情况下是正确的。然后替换为 paged attention 发现也可以跑通了，所以 paged attention 也是对的！
// 更新 KV cache: slice 返回视图, rearrange 将新数据写入对应的位置tensor_t k_layer = k_cache->slice(0, layer, layer + 1)->reshape({block_num, token_num, nkvh, dh});tensor_t v_layer = v_cache->slice(0, layer, layer + 1)->reshape({block_num, token_num, nkvh, dh});// 按块加载到 kv cachefor(size_t i = start; i < ntoken; i += token_num){size_t b = i / token_num;size_t l = std::max(start, b * token_num);size_t r = std::min(ntoken - 1, (b + 1) * token_num - 1);// 缓存块切分int block_id = block_ids[b];tensor_t k_block = k_layer->slice(0, block_id, block_id + 1)->reshape({token_num, nkvh, dh});tensor_t v_block = v_layer->slice(0, block_id, block_id + 1)->reshape({token_num, nkvh, dh});k_block = k_block->slice(0, l % token_num, r % token_num + 1)->reshape({r - l + 1, nkvh, dh});v_block = v_block->slice(0, l % token_num, r % token_num + 1)->reshape({r - l + 1, nkvh, dh});// 张量切分tensor_t k_slice = k_rope->slice(0, l - start, r - start + 1);tensor_t v_slice = v_view->slice(0, l - start, r - start + 1);ops::rearrange(k_block, k_slice);ops::rearrange(v_block, v_slice);} // 去掉 layer 维度// k_layer = k_layer->reshape({block_num * token_num, nkvh, dh});// v_layer = v_layer->reshape({block_num * token_num, nkvh, dh});// tensor_t k_slice = k_layer->slice(0, start, ntoken);// tensor_t v_slice = v_layer->slice(0, start, ntoken);// ops::rearrange(k_slice, k_rope);// ops::rearrange(v_slice, v_view);// k_layer = k_layer->slice(0, 0, ntoken);// v_layer = v_layer->slice(0, 0, ntoken);// 自注意力: paged_attention(flash v2)ops::paged_attention(attn_val, q_rope, k_layer, v_layer, dev_block_ids, block_ids.size(), ntoken, scale);// ops::self_attention(attn_val, q_rope, k_layer, v_layer, scale);
所以定位问题在分块搬运的过程中。最后发现 tensor 的 reshape 实现要求张量要连续 isContigous 连续会返回视图，复用内存，否则会调用 contigous 创建临时张量（这里我原本想利用前一种情况），而分块搬运连续两次 slice 导致张量步长不再连续，并且在搬之前还 slice 了一次（slice 只改变了 offset 并没有改变 stride），导致张量直接 reshape 就会调用 contigous 创建临时张量，而不是写入块中。
罪魁祸首：
解决方案：每次 tensor 每次 slice 了，就 reshape 一下，让步长变连续。
非常 amazing！