main.cpp

#include <print>
#include <random>
#include "math.hpp"
#include "math_expr.hpp"
#include "tokenizer.hpp"
#include "gguf.hpp"
#include "filesystem.hpp"
#include "cmdline.hpp"

using tensor_1d = misha::math::tensor_1d;
using tensor = misha::math::tensor_view;

struct gemma4_layer_weights_t {
    bool is_full_attention;
    bool is_shared_kv_cache;
    uint16_t kv_cache_index;
    float output_scale;

    tensor pre_attn_norm;
    tensor attn_q_proj;
    tensor attn_q_norm;
    tensor attn_k_proj;
    tensor attn_k_norm;
    tensor attn_v_proj;
    tensor attn_output_proj;
    tensor post_attn_norm;

    tensor pre_ffn_norm;
    tensor ffn_gate_proj;
    tensor ffn_up_proj;
    tensor ffn_down_proj;
    tensor post_ffn_norm;

    tensor per_layer_input_gate;
    tensor per_layer_proj;
    tensor per_layer_norm;
};

struct gemma4_weights_t {
    uint32_t sliding_window_size;
    float layer_rmsnorm_epsilon;
    float final_logit_softcapping;

    tensor embed;
    tensor per_layer_embeds;
    tensor per_layer_inputs_proj;
    tensor per_layer_norm;
    tensor output_norm;
    std::unique_ptr<float[]> rope_freqs;
    std::unique_ptr<float[]> rope_freqs_swa;
    
    std::vector<gemma4_layer_weights_t> layers;
};

struct gemma4_forward_state_t {
    unsigned current_pos = 0;
    tensor_1d x[7] = {};
    std::vector<std::pair<tensor, tensor>> kv_caches = {};

    auto append_kv_cache(const gemma4_layer_weights_t& layer, unsigned sliding_window_size) -> std::pair<tensor_1d, tensor_1d> {
        unsigned pos = current_pos & (layer.is_full_attention ? ~0u : sliding_window_size - 1);

        auto& [k_cache, v_cache] = kv_caches[layer.kv_cache_index];
        for (auto* cache : { &k_cache, &v_cache }) {
            if (pos >= cache->capacity) [[unlikely]] {
                cache->capacity = std::max({64u, pos + 1, cache->capacity * 2});
                float* new_data = (float*)::operator new(cache->dims[0] * cache->capacity * sizeof(float), std::align_val_t(32));
                std::copy_n((float*)cache->data, cache->dims[0] * cache->dims[1], new_data);
                ::operator delete(std::exchange(cache->data, new_data), std::align_val_t(32));
            }
        }

        k_cache.dims[1] = std::max(pos + 1, k_cache.dims[1]);
        v_cache.dims[1] = std::max(pos + 1, v_cache.dims[1]);

        return { k_cache[pos], v_cache[pos] };
    }

    ~gemma4_forward_state_t() {
        std::ranges::for_each(x, &tensor_1d::free);

        for (auto& [k_cache, v_cache] : kv_caches) {
            ::operator delete(k_cache.data, std::align_val_t(32));
            ::operator delete(v_cache.data, std::align_val_t(32));
        }
    }
};


void gemma4_forward(gemma4_forward_state_t& s, const gemma4_weights_t& w, uint32_t token_id) {
    using namespace misha::math;
    unsigned per_layer_input_dim = w.per_layer_inputs_proj.dims[1] / w.layers.size();

    auto& hidden_state = s.x[0];
    auto& per_layer_inputs_buf = s.x[5];
    auto& per_layer_embeds_buf = s.x[6];

    hidden_state = get_embedding(w.embed, token_id, std::sqrt(w.embed.dims[0]));
    per_layer_inputs_buf = hidden_state * w.per_layer_inputs_proj * (1.0f / std::sqrt(per_layer_input_dim));
    per_layer_embeds_buf = get_embedding(w.per_layer_embeds, token_id, std::sqrt(per_layer_input_dim));

    auto per_layer_inputs = tensor(per_layer_inputs_buf.data, { per_layer_input_dim, (uint32_t)w.layers.size() });
    auto per_layer_embeds = tensor(per_layer_embeds_buf.data, { per_layer_input_dim , (uint32_t)w.layers.size() });

    for (auto [layer_idx, layer] : w.layers | std::views::enumerate) {
        // Attention
        {
            unsigned head_dim = layer.attn_q_norm.dims[0];
            unsigned head_count = layer.attn_q_proj.dims[1] / head_dim;
            float* rope_freqs = (layer.is_full_attention ? w.rope_freqs : w.rope_freqs_swa).get();

            auto& input = s.x[1];
            auto& q = s.x[2];
            auto& [k_cache, v_cache] = s.kv_caches[layer.kv_cache_index];
            auto& attn_out = s.x[3].resize(head_dim * head_count);
            auto& scores = s.x[4];
            auto& output = s.x[1];

            input = rmsnorm_to(hidden_state, w.layer_rmsnorm_epsilon, layer.pre_attn_norm);
            q = input * layer.attn_q_proj;

            if (not layer.is_shared_kv_cache) {
                auto [new_key, new_value] = s.append_kv_cache(layer, w.sliding_window_size);

                new_key = input * layer.attn_k_proj;
                rmsnorm(new_key, w.layer_rmsnorm_epsilon, layer.attn_k_norm);
                apply_rope(new_key, s.current_pos, rope_freqs);

                new_value = input * layer.attn_v_proj;
                rmsnorm</*use_weight=*/false>(new_value, w.layer_rmsnorm_epsilon);
            }

            for (unsigned i = 0; i < head_count; i++) {
                tensor_1d q_head = tensor_1d(q.data + i * head_dim, { head_dim }); 
                tensor_1d head_output = tensor_1d(attn_out.data + i * head_dim,{ head_dim });

                rmsnorm(q_head, w.layer_rmsnorm_epsilon, layer.attn_q_norm);
                apply_rope(q_head, s.current_pos, rope_freqs);
                scores = q_head * k_cache;
                softmax(scores);
                head_output = vec_mat_mul(scores, v_cache);
            }

            output = attn_out * layer.attn_output_proj;
            rmsnorm(output, w.layer_rmsnorm_epsilon, layer.post_attn_norm);
            
            hidden_state = (hidden_state + output) * 1.0f;
        }
        
        // FFN
        {
            auto& input = s.x[1];
            auto& ffn_gate = s.x[2];
            auto& ffn_up = s.x[3];
            auto& output = s.x[1];

            input = rmsnorm_to(hidden_state, w.layer_rmsnorm_epsilon, layer.pre_ffn_norm);

            ffn_gate = input * layer.ffn_gate_proj;
            ffn_up = input * layer.ffn_up_proj;
            ffn_gate = gelu(ffn_gate) * ffn_up;

            output = ffn_gate * layer.ffn_down_proj;
            rmsnorm(output, w.layer_rmsnorm_epsilon, layer.post_ffn_norm);

            hidden_state = (hidden_state + output) * 1.0f;
        }

        // Per-layer
        {
            auto per_layer_embed = per_layer_embeds[layer_idx];
            auto& gate = s.x[1];
            auto& output = s.x[2];

            auto input = per_layer_inputs[layer_idx];
            rmsnorm(input, w.layer_rmsnorm_epsilon, w.per_layer_norm);
            input = (input + per_layer_embed) * (1.0f / std::sqrt(2.0f));

            gate = hidden_state * layer.per_layer_input_gate;
            gate = gelu(gate) * input;

            output = gate * layer.per_layer_proj;
            rmsnorm(output, w.layer_rmsnorm_epsilon, layer.per_layer_norm);

            hidden_state = (hidden_state + output) * layer.output_scale;
        }
    }

    s.current_pos++;
}

auto gemma4_output(gemma4_forward_state_t& s, const gemma4_weights_t& w) -> std::span<float> {
    auto& hidden_state = s.x[0];
    auto& logits = s.x[1];

    rmsnorm(hidden_state, w.layer_rmsnorm_epsilon, w.output_norm);
    logits = hidden_state * w.embed;

    for (size_t i = 0; i < w.embed.dims[1]; i++)
        logits[i] = w.final_logit_softcapping * std::tanh(logits[i] * (1.0f / w.final_logit_softcapping));

    return std::span(logits.data, logits.dims[0]);
}


auto load_gemma4_weights(const misha::gguf::gguf_model& model) -> gemma4_weights_t {
    using enum misha::gguf::gguf_type;

    uint32_t layers_count = model.get_metadata("gemma4.block_count", GGUF_TYPE_UINT32).value;
    uint32_t layers_count_share_kv_cache = model.get_metadata("gemma4.attention.shared_kv_layers", GGUF_TYPE_UINT32).value;
    uint32_t rope_dims_count = model.get_metadata("gemma4.rope.dimension_count", GGUF_TYPE_UINT32).value;
    uint32_t rope_dims_count_swa = model.get_metadata("gemma4.rope.dimension_count_swa", GGUF_TYPE_UINT32).value;
    float rope_freq_base = model.get_metadata("gemma4.rope.freq_base", GGUF_TYPE_FLOAT32).as<float>();
    float rope_freq_base_swa = model.get_metadata("gemma4.rope.freq_base_swa", GGUF_TYPE_FLOAT32).as<float>();

    gemma4_weights_t weights;
    weights.final_logit_softcapping = model.get_metadata("gemma4.final_logit_softcapping", GGUF_TYPE_FLOAT32).as<float>();
    weights.layer_rmsnorm_epsilon = model.get_metadata("gemma4.attention.layer_norm_rms_epsilon", GGUF_TYPE_FLOAT32).as<float>();
    weights.sliding_window_size = model.get_metadata("gemma4.attention.sliding_window", GGUF_TYPE_UINT32).value;
    weights.embed = model.get_tensor("token_embd.weight");
    weights.per_layer_embeds = model.get_tensor("per_layer_token_embd.weight");
    weights.per_layer_inputs_proj = model.get_tensor("per_layer_model_proj.weight");
    weights.per_layer_norm = model.get_tensor("per_layer_proj_norm.weight");
    weights.output_norm = model.get_tensor("output_norm.weight");
    weights.rope_freqs = misha::math::compute_rope_freqs(rope_dims_count, rope_freq_base);
    weights.rope_freqs_swa = misha::math::compute_rope_freqs(rope_dims_count_swa, rope_freq_base_swa);

    float* rope_freq_factors = (float*)model.get_tensor("rope_freqs.weight").data;
    for (uint32_t i = 0; i < rope_dims_count / 2; i++)
        weights.rope_freqs[i] = weights.rope_freqs[i] / rope_freq_factors[i];
    
    auto get_layer_tensor = [&](int layer_id, std::string_view name) -> tensor {
        char full_name_buffer[128];
        auto res = std::format_to_n(full_name_buffer, sizeof full_name_buffer, "blk.{}.{}", layer_id, name);
        auto full_tensor_name = std::string_view(full_name_buffer, std::min<size_t>(res.size, sizeof full_name_buffer));
        return model.get_tensor(full_tensor_name);
    };
    
    unsigned last_kv_cache_idx = 0;
    unsigned last_sliding_window_kv_cache_idx = 0;
    unsigned last_full_attention_kv_cache_idx = 0;

    for (unsigned i = 0; i < layers_count; i++) {
        gemma4_layer_weights_t layer;

        layer.is_full_attention = i % 5 == 4;
        layer.is_shared_kv_cache = i >= layers_count - layers_count_share_kv_cache;
        
        if (not layer.is_shared_kv_cache) {
            layer.attn_k_proj = get_layer_tensor(i, "attn_k.weight");
            layer.attn_k_norm = get_layer_tensor(i, "attn_k_norm.weight");
            layer.attn_v_proj = get_layer_tensor(i, "attn_v.weight");
            layer.kv_cache_index = last_kv_cache_idx++;
            (layer.is_full_attention ? last_full_attention_kv_cache_idx : last_sliding_window_kv_cache_idx) = layer.kv_cache_index;
        } else {
            layer.kv_cache_index = layer.is_full_attention ? last_full_attention_kv_cache_idx : last_sliding_window_kv_cache_idx;
        }

        layer.pre_attn_norm = get_layer_tensor(i, "attn_norm.weight");
        layer.attn_q_proj = get_layer_tensor(i, "attn_q.weight");
        layer.attn_q_norm = get_layer_tensor(i, "attn_q_norm.weight");
        layer.attn_output_proj = get_layer_tensor(i, "attn_output.weight");
        layer.post_attn_norm = get_layer_tensor(i, "post_attention_norm.weight");

        layer.pre_ffn_norm = get_layer_tensor(i, "ffn_norm.weight");
        layer.ffn_gate_proj = get_layer_tensor(i, "ffn_gate.weight");
        layer.ffn_up_proj = get_layer_tensor(i, "ffn_up.weight");
        layer.ffn_down_proj = get_layer_tensor(i, "ffn_down.weight");
        layer.post_ffn_norm = get_layer_tensor(i, "post_ffw_norm.weight");

        layer.per_layer_input_gate = get_layer_tensor(i, "inp_gate.weight");
        layer.per_layer_proj = get_layer_tensor(i, "proj.weight");
        layer.per_layer_norm = get_layer_tensor(i, "post_norm.weight");
        layer.output_scale = ((float*)get_layer_tensor(i, "layer_output_scale.weight").data)[0];

        weights.layers.emplace_back(std::move(layer));
    }

    return weights;
}

auto init_gemma4_state(const gemma4_weights_t& weights) -> gemma4_forward_state_t {
    using enum misha::gguf::ggml_type;

    auto state = gemma4_forward_state_t();
    std::ranges::fill(state.x, tensor_1d { .capacity = 0 });

    for (auto& layer : weights.layers)
        if (not layer.is_shared_kv_cache)
            state.kv_caches.emplace_back(
                tensor { .dims = { layer.attn_k_proj.dims[1], 0 }, .ggml_type = GGML_TYPE_F32 },
                tensor { .dims = { layer.attn_v_proj.dims[1], 0 }, .ggml_type = GGML_TYPE_F32 }
            );

    return state;
}

auto main(int argc, char* argv[]) -> int try {
    auto args = misha::cmdline::parse(argc, argv);

    std::string gguf_file_path = args["--model"] | args["-m"] | "gemma-4-E2B_q4_0-it.gguf";
    float temperature = args["--temp"] | args["--temperature"] | 1.0f;
    float top_p = args["--top-p"] | 0.9f;
    unsigned max_tokens = args["--max-tokens"] | 4096;
    unsigned random_seed = args["--seed"] | std::random_device()();
    std::string system_prompt = args["--system-prompt"] | "";
    std::string prompt = args["--prompt"] | args["-p"] | "Hello, World!";
    bool enable_thinking = args["--thinking"] || false;

    std::string chat_template = std::string() +
        "<|turn>system\n" + (enable_thinking ? "<|think|>" : "") + "{}<turn|>\n" +
        "<|turn>user\n{}<turn|>\n" +
        "<|turn>model\n";

    auto mapped_file = misha::filesystem::mmap(gguf_file_path.c_str());
    if (not mapped_file)
        throw std::runtime_error(std::format("failed to mmap file: {}", gguf_file_path));

    auto model = misha::gguf::gguf_model(mapped_file.bytes());
    auto tokenizer = misha::gemma4::load_tokenizer_from_gguf(model);
    auto gemma4_weights = load_gemma4_weights(model);
    auto gemma4_state = init_gemma4_state(gemma4_weights);
    auto random_generator = std::mt19937(random_seed);

    std::string fmt_prompt = std::vformat(chat_template, std::make_format_args(system_prompt, prompt));
    std::vector<int> tokens = tokenizer.encode(fmt_prompt);

    std::print("tokens: ");
    for (uint32_t token : tokens) {
        gemma4_forward(gemma4_state, gemma4_weights, token);
        std::print("{:?}, ", tokenizer.decode(token));
        std::fflush(stdout);
    }

    std::print("\ngeneration: ");
    for (unsigned step = 0; step < max_tokens; step++) {
        std::span<float> logits = gemma4_output(gemma4_state, gemma4_weights);
        unsigned token = misha::math::sample(logits, random_generator, temperature, top_p);
        std::string_view decoded_token = tokenizer.decode(token);
        
        if (token == tokenizer.eos_token_id or decoded_token == "<turn|>")
            break;

        std::print("{}",decoded_token);
        std::fflush(stdout);

        gemma4_forward(gemma4_state, gemma4_weights, token);
    }
    
    std::println();
    return 0;
}
catch (const std::exception& e) {
    std::fprintf(stderr, "[fatal error] %s\n", e.what());
    return -1;
}