Function test_sliding_window_basic

tests/test_kv_cache.cpp:21–131

}

bool test_sliding_window_basic() {
    const size_t num_layers = 2;
    const size_t max_seq = 2048;
    const size_t num_kv_heads = 8;
    const size_t head_dim = 64;
    const size_t window_size = 16;
    const size_t sink_size = 4;

    KVCache cache;
    cache.init(num_layers, max_seq, std::vector<size_t>(num_layers, head_dim), std::vector<size_t>(num_layers, num_kv_heads), Precision::INT8);
    cache.set_window_size(window_size, sink_size);

    CactusGraph graph;
    graph.input({1}, Precision::FP32);

    {
        size_t seq_len = 10;
        vector<size_t> k_nodes, v_nodes;

        for (size_t layer = 0; layer < num_layers; layer++) {
            size_t k_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16);
            size_t v_node = graph.input({seq_len, num_kv_heads, head_dim}, Precision::FP16);

            vector<uint8_t> k_data, v_data;
            fill_fp16(k_data, seq_len * num_kv_heads * head_dim, layer + 1.0f);
            fill_fp16(v_data, seq_len * num_kv_heads * head_dim, layer + 2.0f);
            graph.set_input(k_node, k_data.data(), Precision::FP16);
            graph.set_input(v_node, v_data.data(), Precision::FP16);

            k_nodes.push_back(k_node);
            v_nodes.push_back(v_node);
        }

        graph.execute();
        cache.update_from_graph(&graph, k_nodes, v_nodes, seq_len, num_layers);

        assert(cache.get_effective_seq_len() == seq_len);
    }

    {
        size_t additional_seq = 6;
        vector<size_t> k_nodes, v_nodes;

        for (size_t layer = 0; layer < num_layers; layer++) {
            size_t k_node = graph.input({additional_seq, num_kv_heads, head_dim}, Precision::FP16);
            size_t v_node = graph.input({additional_seq, num_kv_heads, head_dim}, Precision::FP16);

            vector<uint8_t> k_data, v_data;
            fill_fp16(k_data, additional_seq * num_kv_heads * head_dim, layer + 10.0f);
            fill_fp16(v_data, additional_seq * num_kv_heads * head_dim, layer + 20.0f);
            graph.set_input(k_node, k_data.data(), Precision::FP16);
            graph.set_input(v_node, v_data.data(), Precision::FP16);

            k_nodes.push_back(k_node);
            v_nodes.push_back(v_node);
        }

        graph.execute();