cwenzi/neuroflow-cpp

1#include <iostream>2#include <exception>3#include "../include/neuroflow/memory.hpp"4 5using namespace neuroflow;6 7int main() {8    try {9        std::cout << "MLA debug test..." << std::endl;10        11        LatentKVCache mla(64, 4, 16, 128);12        13        std::cout << "mla.d_model: " << mla.d_model << std::endl;14        std::cout << "mla.n_heads: " << mla.n_heads << std::endl;15        std::cout << "mla.d_latent: " << mla.d_latent << std::endl;16        std::cout << "mla.head_dim: " << mla.head_dim << std::endl;17        std::cout << "mla.cache_len: " << mla.cache_len << std::endl;18        19        Tensor input({1, 64});20        std::cout << "input shape: [" << input.shape_[0] << ", " << input.shape_[1] << "]" << std::endl;21        std::cout << "input numel: " << input.numel() << std::endl;22        23        for (size_t i = 0; i < input.numel(); ++i) input.as_fp32()[i] = 0.1f * i;24        25        std::cout << "Calling mla.forward(input, true)..." << std::endl;26        Tensor output1 = mla.forward(input, true);27        28        std::cout << "output1 shape size: " << output1.shape_.size() << std::endl;29        for (size_t i = 0; i < output1.shape_.size(); ++i) std::cout << "  dim " << i << ": " << output1.shape_[i] << std::endl;30        std::cout << "mla.cache_len after 1st forward: " << mla.cache_len << std::endl;31        32        std::cout << "Second forward..." << std::endl;33        Tensor input2({1, 64});34        for (size_t i = 0; i < input2.numel(); ++i) input2.as_fp32()[i] = 0.2f * i;35        Tensor output2 = mla.forward(input2, true);36        37        std::cout << "output2 shape size: " << output2.shape_.size() << std::endl;38        for (size_t i = 0; i < output2.shape_.size(); ++i) std::cout << "  dim " << i << ": " << output2.shape_[i] << std::endl;39        std::cout << "mla.cache_len after 2nd forward: " << mla.cache_len << std::endl;40        41        std::cout << "Success!" << std::endl;42        return 0;43    } catch (const std::exception& e) {44        std::cout << "Error: " << e.what() << std::endl;45        return 1;46    }47}48