PagedAttention实现

/**
 * 简化版PagedAttention实现
 * 
 * 这是一个教学版本，用于理解核心概念
 * 生产环境请使用vLLM
 * 
 * 编译: g++ -std=c++17 -O2 -o paged_attn paged_attn.cpp
 */

#include <iostream>
#include <vector>
#include <unordered_map>
#include <memory>

// ==================== Block定义 ====================
constexpr int BLOCK_SIZE = 16;  // 每个Block存储16个token的KV

struct KVBlock {
    int block_id;
    std::vector<float> keys;    // [BLOCK_SIZE * head_dim]
    std::vector<float> values;
    int ref_count = 0;  // 引用计数（用于CoW）
    
    KVBlock(int id, int head_dim) 
        : block_id(id), keys(BLOCK_SIZE * head_dim), values(BLOCK_SIZE * head_dim) {}
};

// ==================== 页表 ====================
struct PageTable {
    // 逻辑Block → 物理Block的映射
    std::vector<int> logical_to_physical;
    int num_pages = 0;
    
    void add_page(int physical_block_id) {
        logical_to_physical.push_back(physical_block_id);
        num_pages++;
    }
};

// ==================== PagedAttention管理器 ====================
class PagedKVManager {
private:
    int num_blocks_;
    int head_dim_;
    
    // 物理Block池
    std::vector<std::unique_ptr<KVBlock>> blocks_;
    
    // 空闲Block列表
    std::vector<int> free_blocks_;
    
    // 每个序列的页表
    std::unordered_map<int, PageTable> page_tables_;

public:
    PagedKVManager(int num_blocks, int head_dim) 
        : num_blocks_(num_blocks), head_dim_(head_dim) {
        
        // 初始化Block池
        for (int i = 0; i < num_blocks; i++) {
            blocks_.push_back(std::make_unique<KVBlock>(i, head_dim));
            free_blocks_.push_back(i);
        }
    }
    
    // 为序列分配新Block
    bool allocate_block(int seq_id) {
        if (free_blocks_.empty()) {
            std::cerr << "No free blocks available!\n";
            return false;
        }
        
        int block_id = free_blocks_.back();
        free_blocks_.pop_back();
        
        page_tables_[seq_id].add_page(block_id);
        blocks_[block_id]->ref_count = 1;
        
        return true;
    }
    
    // 写入KV数据
    void write_kv(int seq_id, int token_pos, 
                const float* k, const float* v) {
        auto& pt = page_tables_[seq_id];
        int page_idx = token_pos / BLOCK_SIZE;
        int offset = token_pos % BLOCK_SIZE;
        
        int physical_block = pt.logical_to_physical[page_idx];
        auto& block = blocks_[physical_block];
        
        // 复制K,V数据
        for (int i = 0; i < head_dim_; i++) {
            block->keys[offset * head_dim_ + i] = k[i];
            block->values[offset * head_dim_ + i] = v[i];
        }
    }
    
    // Copy-on-Write: 共享Block在写入时复制
    void cow_write(int seq_id, int token_pos,
                 const float* k, const float* v) {
        auto& pt = page_tables_[seq_id];
        int page_idx = token_pos / BLOCK_SIZE;
        int physical_block = pt.logical_to_physical[page_idx];
        
        // 如果引用计数>1，需要复制
        if (blocks_[physical_block]->ref_count > 1) {
            // 分配新Block
            int new_block_id = free_blocks_.back();
            free_blocks_.pop_back();
            
            // 复制数据
            *blocks_[new_block_id] = *blocks_[physical_block];
            blocks_[new_block_id]->block_id = new_block_id;
            blocks_[new_block_id]->ref_count = 1;
            
            // 更新页表
            pt.logical_to_physical[page_idx] = new_block_id;
            blocks_[physical_block]->ref_count--;
            
            physical_block = new_block_id;
        }
        
        // 写入
        write_kv(seq_id, token_pos, k, v);
    }
    
    // 释放序列的所有Block
    void free_sequence(int seq_id) {
        for (int block_id : page_tables_[seq_id].logical_to_physical) {
            blocks_[block_id]->ref_count--;
            if (blocks_[block_id]->ref_count == 0) {
                free_blocks_.push_back(block_id);
            }
        }
        page_tables_.erase(seq_id);
    }
    
    // 打印状态
    void print_status() {
        std::cout << "Total blocks: " << num_blocks_ << std::endl;
        std::cout << "Free blocks: " << free_blocks_.size() << std::endl;
        std::cout << "Active sequences: " << page_tables_.size() << std::endl;
    }
};

// ==================== 测试 ====================
int main() {
    const int num_blocks = 8;
    const int head_dim = 64;
    
    PagedKVManager manager(num_blocks, head_dim);
    
    std::cout << "=== PagedAttention演示 ===\n" << std::endl;
    
    // 序列1: 写入24个token (需要2个Block)
    std::cout << "Sequence 1: writing 24 tokens...";
    float dummy_k[64] = {1.0f};
    float dummy_v[64] = {0.5f};
    
    manager.allocate_block(1);  // Block 0
    manager.allocate_block(1);  // Block 1
    for (int i = 0; i < 24; i++) {
        manager.write_kv(1, i, dummy_k, dummy_v);
    }
    std::cout << " done\n";
    
    // 序列2: 写入16个token (需要1个Block)
    std::cout << "Sequence 2: writing 16 tokens...";
    manager.allocate_block(2);  // Block 2
    for (int i = 0; i < 16; i++) {
        manager.write_kv(2, i, dummy_k, dummy_v);
    }
    std::cout << " done\n" << std::endl;
    
    // 打印状态
    manager.print_status();
    
    // 释放序列1
    std::cout << "\nFreeing sequence 1...";
    manager.free_sequence(1);
    std::cout << " done\n";
    
    manager.print_status();
    
    return 0;
}