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