推理引擎debug记（控制变量法）

很久之前就想复刻一下 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);

    // 重新初始化block
    void _reset(block_t block);

    // 更新block
    void _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 copy
    BlockManager(const BlockManager &) = delete;
    BlockManager &operator=(const BlockManager &) = delete;

    // Prevent move
    BlockManager(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 cache
        for(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！