Function test_large_window

tests/test_kv_cache.cpp:220–258

}

bool test_large_window() {
    const size_t num_layers = 4;
    const size_t num_kv_heads = 8;
    const size_t head_dim = 64;
    const size_t window_size = 512;

    KVCache cache;
    cache.init(num_layers, 2048, std::vector<size_t>(num_layers, head_dim), std::vector<size_t>(num_layers, num_kv_heads), Precision::FP16);
    cache.set_window_size(window_size, 4);

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

    size_t seq_len = 600;
    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, float(layer + 1)); 
        fill_fp16(v_data, seq_len * num_kv_heads * head_dim, float(layer + 101));

        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() == window_size);
    assert(cache.get_total_seq_len() == seq_len);

    return true;
}

int main() {
    TestUtils::TestRunner runner("KV Cache Sliding Window Tests");