cwenzi/neuroflow-cpp

1#include <iostream>2#include <exception>3#include "../include/neuroflow/tensor.hpp"4#include "../include/neuroflow/networks.hpp"5 6using namespace neuroflow;7 8int main() {9    try {10        std::cout << "MLA step by step..." << std::endl;11        12        // MLA 参数13        size_t d_model = 64;14        size_t n_heads = 4;15        size_t d_latent = 16;16        size_t head_dim = d_model / n_heads; // 1617        18        std::cout << "d_model=" << d_model << ", n_heads=" << n_heads 19                  << ", d_latent=" << d_latent << ", head_dim=" << head_dim << std::endl;20        21        // 创建 Linear 层22        Linear W_q(d_model, d_model, false);23        Linear W_dkv(d_model, d_latent, false);24        Linear W_uk(d_latent, d_model, false);25        Linear W_uv(d_latent, d_model, false);26        Linear W_o(d_model, d_model, false);27        28        // 输入29        Tensor input({1, d_model});30        for (size_t i = 0; i < input.numel(); ++i) input.as_fp32()[i] = 0.1f * i;31        32        size_t batch = 1;33        size_t seq_len = 1;34        35        std::cout << "1. Q projection..." << std::endl;36        Tensor x_flat = input.reshape({batch * seq_len, d_model});37        std::cout << "   x_flat shape: [" << x_flat.shape_[0] << ", " << x_flat.shape_[1] << "]" << std::endl;38        39        Tensor q = W_q.forward(x_flat);40        std::cout << "   q shape after linear: [" << q.shape_[0] << ", " << q.shape_[1] << "]" << std::endl;41        42        // q 需要 reshape 到 {batch, seq_len, n_heads, head_dim}43        std::cout << "   reshaping q to {batch, seq_len, n_heads, head_dim} = {" << batch << ", " << seq_len << ", " << n_heads << ", " << head_dim << "}" << std::endl;44        std::cout << "   q numel: " << q.numel() << std::endl;45        std::cout << "   expected numel: " << (batch * seq_len * n_heads * head_dim) << std::endl;46        47        q = q.reshape({batch, seq_len, n_heads, head_dim});48        std::cout << "   q reshaped: success" << std::endl;49        50        std::cout << "2. KV compression..." << std::endl;51        Tensor c_kv = W_dkv.forward(x_flat);52        std::cout << "   c_kv shape: [" << c_kv.shape_[0] << ", " << c_kv.shape_[1] << "]" << std::endl;53        c_kv = c_kv.reshape({batch, seq_len, d_latent});54        std::cout << "   c_kv reshaped: success" << std::endl;55        56        std::cout << "3. K/V decompression..." << std::endl;57        size_t total_len = seq_len; // 第一次没有 cache58        59        std::cout << "   creating c_kv_flat {" << (batch * total_len) << ", " << d_latent << "}" << std::endl;60        Tensor c_kv_flat({batch * total_len, d_latent});61        float* ckf = c_kv_flat.as_fp32();62        float* ck = c_kv.as_fp32();63        std::cout << "   copying " << (batch * total_len * d_latent) << " elements" << std::endl;64        for (size_t i = 0; i < batch * total_len * d_latent; ++i) {65            ckf[i] = ck[i];66        }67        std::cout << "   copy done" << std::endl;68        69        Tensor k = W_uk.forward(c_kv_flat);70        Tensor v = W_uv.forward(c_kv_flat);71        std::cout << "   k shape: [" << k.shape_[0] << ", " << k.shape_[1] << "]" << std::endl;72        std::cout << "   v shape: [" << v.shape_[0] << ", " << v.shape_[1] << "]" << std::endl;73        74        k = k.reshape({batch, total_len, n_heads, head_dim});75        v = v.reshape({batch, total_len, n_heads, head_dim});76        std::cout << "   k reshaped, v reshaped" << std::endl;77        78        std::cout << "4. Attention computation..." << std::endl;79        Tensor output({batch, seq_len, d_model});80        float* out = output.as_fp32();81        float* qp = q.as_fp32();82        float* kp = k.as_fp32();83        float* vp = v.as_fp32();84        85        float scale = 1.0f / std::sqrt(static_cast<float>(head_dim));86        87        for (size_t b = 0; b < batch; ++b) {88            for (size_t h = 0; h < n_heads; ++h) {89                for (size_t s = 0; s < seq_len; ++s) {90                    std::vector<float> scores(total_len);91                    for (size_t t = 0; t < total_len; ++t) {92                        float dot = 0;93                        for (size_t d = 0; d < head_dim; ++d) {94                            size_t q_idx = b * seq_len * n_heads * head_dim + s * n_heads * head_dim + h * head_dim + d;95                            size_t k_idx = b * total_len * n_heads * head_dim + t * n_heads * head_dim + h * head_dim + d;96                            std::cout << "   q_idx=" << q_idx << ", k_idx=" << k_idx << std::endl;97                            dot += qp[q_idx] * kp[k_idx];98                        }99                        scores[t] = dot * scale;100                    }101                    // softmax...102                    // output...103                }104            }105        }106        107        std::cout << "Success!" << std::endl;108        return 0;109    } catch (const std::exception& e) {110        std::cout << "Error: " << e.what() << std::endl;111        return 1;112    }113}114