大模型内存优化：PagedAttention移动版

// PagedAttention移动版核心代码

#include <vector>
#include <unordered_map>
#include <algorithm>
#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
#include <cmath>
#include <random>
#include <ctime>
#include <chrono>

using namespace std;

// 定义页面大小（可根据实际情况调整）
const int PAGE_SIZE = 4096;

// 定义内存缓存区大小（可根据设备内存调整）
const int CACHE_SIZE = 1024 * 1024 * 256; // 256MB

// 定义外部存储路径（可根据设备存储情况调整）
const string EXTERNAL_STORAGE_PATH = "/path/to/external/storage";

// 页面类，用于存储模型参数和相关元数据
class Page {
public:
    vector<float> data; // 页面数据
    int page_id; // 页面ID
    int access_count; // 页面访问计数（用于LRU缓存策略）
    chrono::steady_clock::time_point access_time; // 页面最后访问时间

    Page(int id) : page_id(id), access_count(0) {
        data.resize(PAGE_SIZE / sizeof(float));
    }

    // 更新页面访问时间
    void update_access() {
        access_count++;
        access_time = chrono::steady_clock::now();
    }
};

// 内存缓存管理类
class CacheManager {
private:
    vector<Page*> cache; // 内存缓存区
    unordered_map<int, Page*> page_map; // 页面映射表
    list<int> lru_list; // LRU列表，用于缓存淘汰策略
    int cache_hit_count; // 缓存命中次数
    int cache_miss_count; // 缓存未命中次数

public:
    CacheManager() : cache_hit_count(0), cache_miss_count(0) {
        cache.resize(CACHE_SIZE / PAGE_SIZE, nullptr);
    }

    // 将页面加载到内存缓存区
    bool load_page(int page_id, Page* page) {
        // 检查页面是否已在缓存中
        if (page_map.find(page_id) != page_map.end()) {
            cache_hit_count++;
            page_map[page_id]->update_access();
            lru_list.remove(page_id);
            lru_list.push_front(page_id);
            return true;
        }

        // 缓存未命中，需要从外部存储加载页面
        cache_miss_count++;

        // 如果缓存已满，执行LRU淘汰策略
        if (is_cache_full()) {
            evict_lru_page();
        }

        // 将新页面加载到缓存
        int cache_index = get_empty_cache_index();
        if (cache_index != -1) {
            cache[cache_index] = page;
            page_map[page_id] = page;
            lru_list.push_front(page_id);
            return true;
        }

        return false;
    }

    // 从内存缓存区释放页面
    bool release_page(int page_id) {
        if (page_map.find(page_id) != page_map.end()) {
            Page* page = page_map[page_id];
            int cache_index = find_page_in_cache(page);
            if (cache_index != -1) {
                cache[cache_index] = nullptr;
                page_map.erase(page_id);
                lru_list.remove(page_id);
                return true;
            }
        }
        return false;
    }

    // 检查缓存是否已满
    bool is_cache_full() {
        for (auto& page : cache) {
            if (page == nullptr) {
                return false;
            }
        }
        return true;
    }

    // 执行LRU淘汰策略，释放最近最少使用的页面
    void evict_lru_page() {
        if (!lru_list.empty()) {
            int lru_page_id = lru_list.back();
            lru_list.pop_back();
            release_page(lru_page_id);
        }
    }

    // 查找内存缓存区中的空闲位置
    int get_empty_cache_index() {
        for (int i = 0; i < cache.size(); i++) {
            if (cache[i] == nullptr) {
                return i;
            }
        }
        return -1;
    }

    // 查找页面在内存缓存区中的位置
    int find_page_in_cache(Page* page) {
        for (int i = 0; i < cache.size(); i++) {
            if (cache[i] == page) {
                return i;
            }
        }
        return -1;
    }

    // 获取缓存命中率
    float get_cache_hit_ratio() {
        if (cache_hit_count + cache_miss_count == 0) {
            return 0.0f;
        }
        return static_cast<float>(cache_hit_count) / (cache_hit_count + cache_miss_count);
    }
};

// 页面调度管理类
class PageScheduler {
private:
    CacheManager cache_manager; // 内存缓存管理器
    unordered_map<int, string> page_storage_map; // 页面存储映射表
    int total_pages; // 总页面数
    int current_page_id; // 当前页面ID

public:
    PageScheduler() : total_pages(0), current_page_id(0) {
    }

    // 初始化页面存储
    void initialize_storage(const vector<float>& model_parameters) {
        total_pages = (model_parameters.size() + PAGE_SIZE - 1) / PAGE_SIZE;
        current_page_id = 0;

        // 将模型参数分页存储到外部存储
        for (int i = 0; i < total_pages; i++) {
            Page* page = new Page(i);
            for (int j = 0; j < PAGE_SIZE / sizeof(float); j++) {
                int param_index = i * PAGE_SIZE / sizeof(float) + j;
                if (param_index < model_parameters.size()) {
                    page->data[j] = model_parameters[param_index];
                } else {
                    page->data[j] = 0.0f; // 填充剩余位置为0
                }
            }

            // 将页面写入外部存储（这里模拟写入操作）
            string page_path = EXTERNAL_STORAGE_PATH + "/page_" + to_string(i) + ".bin";
            ofstream page_file(page_path, ios::binary);
            if (page_file.is_open()) {
                page_file.write(reinterpret_cast<char*>(page->data.data()), page->data.size() * sizeof(float));
                page_file.close();
                page_storage_map[i] = page_path;
            }

            // 释放页面对象
            delete page;
        }
    }

    // 加载页面到内存缓存区
    bool load_page(int page_id) {
        if (page_id >= 0 && page_id < total_pages) {
            // 从外部存储读取页面数据（这里模拟读取操作）
            string page_path = page_storage_map[page_id];
            ifstream page_file(page_path, ios::binary);
            if (page_file.is_open()) {
                vector<float> page_data(PAGE_SIZE / sizeof(float));
                page_file.read(reinterpret_cast<char*>(page_data.data()), page_data.size() * sizeof(float));
                page_file.close();

                // 创建页面对象并加载到缓存
                Page* page = new Page(page_id);
                page->data = page_data;
                bool result = cache_manager.load_page(page_id, page);

                if (result) {
                    cout << "页面 " << page_id << " 加载到内存缓存区成功" << endl;
                } else {
                    cout << "页面 " << page_id << " 加载到内存缓存区失败" << endl;
                    delete page;
                }

                return result;
            }
        }
        return false;
    }

    // 释放页面
    bool release_page(int page_id) {
        return cache_manager.release_page(page_id);
    }

    // 获取缓存命中率
    float get_cache_hit_ratio() {
        return cache_manager.get_cache_hit_ratio();
    }
};

// PagedAttention计算类
class PagedAttention {
private:
    PageScheduler page_scheduler; // 页面调度管理器
    int sequence_length; // 序列长度
    int hidden_size; // 隐藏层大小
    int num_heads; // 注意力头数
    int page_id; // 当前页面ID

public:
    PagedAttention(int seq_len, int hid_size, int heads) 
        : sequence_length(seq_len), hidden_size(hid_size), num_heads(heads), page_id(0) {
    }

    // 初始化PagedAttention
    void initialize(const vector<float>& model_parameters) {
        page_scheduler.initialize_storage(model_parameters);
    }

    // 注意力计算函数
    vector<float> compute_attention(const vector<float>& query, const vector<float>& key, const vector<float>& value) {
        vector<float> attention_scores(sequence_length * sequence_length, 0.0f);
        vector<float> output(sequence_length * hidden_size, 0.0f);

        // 分页计算注意力
        int pages_needed = (sequence_length * sequence_length + PAGE_SIZE - 1) / PAGE_SIZE;
        for (int i = 0; i < pages_needed; i++) {
            // 计算当前页面的起始和结束位置
            int start = i * PAGE_SIZE;
            int end = min((i + 1) * PAGE_SIZE, sequence_length * sequence_length);

            // 加载所需的页面
            if (!page_scheduler.load_page(i)) {
                cout << "加载页面 " << i << " 失败，可能会影响计算结果" << endl;
                continue;
            }

            // 计算当前页面的注意力分数
            for (int j = start; j < end; j++) {
                int row = j / sequence_length;
                int col = j % sequence_length;
                float score = 0.0f;
                for (int k = 0; k < hidden_size; k++) {
                    score += query[row * hidden_size + k] * key[col * hidden_size + k];
                }
                attention_scores[j] = score;
            }

            // 释放页面
            page_scheduler.release_page(i);
        }

        // 计算输出
        for (int i = 0; i < sequence_length; i++) {
            float row_sum = 0.0f;
            for (int j = 0; j < sequence_length; j++) {
                row_sum += exp(attention_scores[i * sequence_length + j]);
            }
            for (int j = 0; j < sequence_length; j++) {
                float weight = exp(attention_scores[i * sequence_length + j]) / row_sum;
                for (int k = 0; k < hidden_size; k++) {
                    output[i * hidden_size + k] += weight * value[j * hidden_size + k];
                }
            }
        }

        return output;
    }
};

// 主函数，用于测试PagedAttention移动版
int main() {
    // 设置随机种子
    srand(time(0));

    // 模型参数配置
    int sequence_length = 512; // 序列长度
    int hidden_size = 1024; // 隐藏层大小
    int num_heads = 16; // 注意力头数

    // 生成随机模型参数（模拟大模型参数）
    int model_size = sequence_length * hidden_size * num_heads;
    vector<float> model_parameters(model_size);
    for (int i = 0; i < model_size; i++) {
        model_parameters[i] = static_cast<float>(rand()) / RAND_MAX;
    }

    // 初始化PagedAttention
    PagedAttention paged_attention(sequence_length, hidden_size, num_heads);
    paged_attention.initialize(model_parameters);

    // 生成随机输入数据
    vector<float> query(sequence_length * hidden_size);
    vector<float> key(sequence_length * hidden_size);
    vector<float> value(sequence_length * hidden_size);
    for (int i = 0; i < sequence_length * hidden_size; i++) {
        query[i] = static_cast<float>(rand()) / RAND_MAX;
        key[i] = static_cast<float>(rand()) / RAND_MAX;
        value[i] = static_cast<float>(rand()) / RAND_MAX;
    }

    // 计算注意力输出
    auto start_time = chrono::steady_clock::now();
    vector<float> attention_output = paged_attention.compute_attention(query, key, value);
    auto end_time = chrono::steady_clock::now();
    chrono::duration<double> duration = end_time - start_time;

    // 输出计算结果和性能指标
    cout << "PagedAttention计算完成，输出大小：" << attention_output.size() << endl;
    cout << "计算时间：" << duration.count() << " 秒" << endl;
    cout << "缓存命中率：" << paged_attention.page_scheduler.get_cache_hit_ratio() << endl;

    return 0;
}