#include "llama-model.h"

#include "llama-impl.h"

#include "llama-mmap.h"

#include "llama-cparams.h"

#include "llama-model-loader.h"

#include "llama-kv-cache.h"

#include "llama-kv-cache-iswa.h"

#include "llama-memory-hybrid.h"

#include "llama-memory-hybrid-iswa.h"

#include "llama-memory-recurrent.h"

#include "ggml-cpp.h"

#include "models/models.h"

#include <algorithm>

#include <cassert>

#include <cfloat>

#include <cstring>

#include <cmath>

#include <functional>

#include <map>

#include <regex>

#include <sstream>

#include <stdexcept>

const char * llm_type_name(llm_type type) {

    switch (type) {

        case LLM_TYPE_14M:           return "14M";

        case LLM_TYPE_17M:           return "17M";

        case LLM_TYPE_22M:           return "22M";

        case LLM_TYPE_33M:           return "33M";

        case LLM_TYPE_47M:           return "47M";

        case LLM_TYPE_60M:           return "60M";

        case LLM_TYPE_70M:           return "70M";

        case LLM_TYPE_80M:           return "80M";

        case LLM_TYPE_109M:          return "109M";

        case LLM_TYPE_137M:          return "137M";

        case LLM_TYPE_140M:          return "140M";

        case LLM_TYPE_149M:          return "149M";

        case LLM_TYPE_160M:          return "160M";

        case LLM_TYPE_190M:          return "190M";

        case LLM_TYPE_220M:          return "220M";

        case LLM_TYPE_250M:          return "250M";

        case LLM_TYPE_256M:          return "256M";

        case LLM_TYPE_270M:          return "270M";

        case LLM_TYPE_335M:          return "335M";

        case LLM_TYPE_350M:          return "350M";

        case LLM_TYPE_360M:          return "360M";

        case LLM_TYPE_395M:          return "395M";

        case LLM_TYPE_410M:          return "410M";

        case LLM_TYPE_450M:          return "450M";

        case LLM_TYPE_475M:          return "475M";

        case LLM_TYPE_558M:          return "558M";

        case LLM_TYPE_700M:          return "700M";

        case LLM_TYPE_770M:          return "770M";

        case LLM_TYPE_780M:          return "780M";

        case LLM_TYPE_950M:          return "950M";

        case LLM_TYPE_0_3B:          return "0.3B";

        case LLM_TYPE_0_5B:          return "0.5B";

        case LLM_TYPE_0_6B:          return "0.6B";

        case LLM_TYPE_1B:            return "1B";

        case LLM_TYPE_1_2B:          return "1.2B";

        case LLM_TYPE_1_3B:          return "1.3B";

        case LLM_TYPE_1_4B:          return "1.4B";

        case LLM_TYPE_1_5B:          return "1.5B";

        case LLM_TYPE_1_6B:          return "1.6B";

        case LLM_TYPE_1_7B:          return "1.7B";

        case LLM_TYPE_1_8B:          return "1.8B";

        case LLM_TYPE_2B:            return "2B";

        case LLM_TYPE_2_6B:          return "2.6B";

        case LLM_TYPE_2_8B:          return "2.8B";

        case LLM_TYPE_2_9B:          return "2.9B";

        case LLM_TYPE_3B:            return "3B";

        case LLM_TYPE_4B:            return "4B";

        case LLM_TYPE_6B:            return "6B";

        case LLM_TYPE_6_9B:          return "6.9B";

        case LLM_TYPE_7B:            return "7B";

        case LLM_TYPE_8B:            return "8B";

        case LLM_TYPE_9B:            return "9B";

        case LLM_TYPE_11B:           return "11B";

        case LLM_TYPE_12B:           return "12B";

        case LLM_TYPE_13B:           return "13B";

        case LLM_TYPE_14B:           return "14B";

        case LLM_TYPE_15B:           return "15B";

        case LLM_TYPE_16B:           return "16B";

        case LLM_TYPE_20B:           return "20B";

        case LLM_TYPE_26B:           return "26B";

        case LLM_TYPE_27B:           return "27B";

        case LLM_TYPE_30B:           return "30B";

        case LLM_TYPE_32B:           return "32B";

        case LLM_TYPE_34B:           return "34B";

        case LLM_TYPE_35B:           return "35B";

        case LLM_TYPE_36B:           return "36B";

        case LLM_TYPE_40B:           return "40B";

        case LLM_TYPE_65B:           return "65B";

        case LLM_TYPE_70B:           return "70B";

        case LLM_TYPE_120B:          return "120B";

        case LLM_TYPE_142B:          return "142B";

        case LLM_TYPE_236B:          return "236B";

        case LLM_TYPE_290B:          return "290B";

        case LLM_TYPE_314B:          return "314B";

        case LLM_TYPE_405B:          return "405B";

        case LLM_TYPE_671B:          return "671B";

        case LLM_TYPE_SMALL:         return "0.1B";

        case LLM_TYPE_MEDIUM:        return "0.4B";

        case LLM_TYPE_LARGE:         return "0.8B";

        case LLM_TYPE_XL:            return "1.5B";

        case LLM_TYPE_A1_7B:         return "A1.7B";

        case LLM_TYPE_A2_7B:         return "A2.7B";

        case LLM_TYPE_8x7B:          return "8x7B";

        case LLM_TYPE_8x22B:         return "8x22B";

        case LLM_TYPE_16x12B:        return "16x12B";

        case LLM_TYPE_16x3_8B:       return "16x3.8B";

        case LLM_TYPE_10B_128x3_66B: return "10B+128x3.66B";

        case LLM_TYPE_57B_A14B:      return "57B.A14B";

        case LLM_TYPE_17B_16E:       return "17Bx16E (Scout)";

        case LLM_TYPE_17B_128E:      return "17Bx128E (Maverick)";

        case LLM_TYPE_A13B:          return "A13B";

        case LLM_TYPE_7B_A1B:        return "7B.A1B";

        case LLM_TYPE_8B_A1B:        return "8B.A1B";

        case LLM_TYPE_16B_A1B:       return "16B.A1B";

        case LLM_TYPE_21B_A3B:       return "21B.A3B";

        case LLM_TYPE_24B_A2B:       return "24B.A2B";

        case LLM_TYPE_30B_A3B:       return "30B.A3B";

        case LLM_TYPE_31B_A3_5B:     return "31B.A3.5B";

        case LLM_TYPE_35B_A3B:       return "35B.A3B";

        case LLM_TYPE_48B_A3B:       return "48B.A3B";

        case LLM_TYPE_80B_A3B:       return "80B.A3B";

        case LLM_TYPE_100B_A6B:      return "100B.A6B";

        case LLM_TYPE_102B_A12B:     return "102B.A12B";

        case LLM_TYPE_106B_A12B:     return "106B.A12B";

        case LLM_TYPE_196B_A11B:     return "196B.A11B";

        case LLM_TYPE_230B_A10B:     return "230B.A10B";

        case LLM_TYPE_235B_A22B:     return "235B.A22B";

        case LLM_TYPE_300B_A47B:     return "300B.A47B";

        case LLM_TYPE_310B_A15B:     return "310B.A15B";

        case LLM_TYPE_355B_A32B:     return "355B.A32B";

        case LLM_TYPE_744B_A40B:     return "744B.A40B";

        case LLM_TYPE_E2B:           return "E2B";

        case LLM_TYPE_E4B:           return "E4B";

        default:                     return "?B";

    }

}

static const char * llama_expert_gating_func_name(llama_expert_gating_func_type type) {

    switch (type) {

        case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX: return "softmax";

        case LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID: return "sigmoid";

        default:                                    return "unknown";

    }

}

static const std::map<llama_rope_scaling_type, const char *> LLAMA_ROPE_SCALING_TYPES = {

    { LLAMA_ROPE_SCALING_TYPE_NONE,       "none"       },

    { LLAMA_ROPE_SCALING_TYPE_LINEAR,     "linear"     },

    { LLAMA_ROPE_SCALING_TYPE_YARN,       "yarn"       },

    { LLAMA_ROPE_SCALING_TYPE_LONGROPE,   "longrope"   },

};

std::string llama_rope_scaling_type_name(llama_rope_scaling_type rope_scaling_type) {

    return LLAMA_ROPE_SCALING_TYPES.at(rope_scaling_type);

}

static llama_rope_scaling_type llama_rope_scaling_type_from_string(const std::string & name) {

    for (const auto & kv : LLAMA_ROPE_SCALING_TYPES) {

        if (kv.second == name) {

            return (llama_rope_scaling_type) kv.first;

        }

    }

    return LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED;

}

// checks if the weight tensor can be used with the specified buffer type and device

static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) {

    GGML_ASSERT(w != nullptr);

    if (op == GGML_OP_NONE) {

        return true;

    }

    ggml_init_params params = {

        /*.mem_size   =*/ ggml_tensor_overhead()*8,

        /*.mem_buffer =*/ NULL,

        /*.no_alloc   =*/ true,

    };

    ggml_context_ptr ctx_ptr { ggml_init(params) };

    if (!ctx_ptr) {

        throw std::runtime_error(format("failed to create ggml context"));

    }

    ggml_context * ctx = ctx_ptr.get();

    ggml_tensor * op_tensor = nullptr;

    switch (op) {

        case GGML_OP_GET_ROWS:

            {

                ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512);

                op_tensor = ggml_get_rows(ctx, w, b);

            } break;

        case GGML_OP_MUL_MAT:

            {

                ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]);

                op_tensor = ggml_mul_mat(ctx, w, b);

            } break;

        case GGML_OP_MUL_MAT_ID:

            {

                int n_expert_used = hparams.n_expert_used;

                ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);

                ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);

                op_tensor = ggml_mul_mat_id(ctx, w, b, ids);

            } break;

        case GGML_OP_ADD:

            {

                ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], w->ne[1], w->ne[2], w->ne[3]);

                op_tensor = ggml_add(ctx, a, w);

            } break;

        case GGML_OP_ADD_ID:

            {

                int n_expert_used = hparams.n_expert_used;

                ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);

                ggml_tensor * c = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);

                op_tensor = ggml_add_id(ctx, a, w, c);

            } break;

        case GGML_OP_MUL:

            {

                ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], w->ne[1], w->ne[2], w->ne[3]);

                op_tensor = ggml_mul(ctx, a, w);

            } break;

        case GGML_OP_DIV:

            {

                ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, w->ne[0]);

                op_tensor = ggml_div(ctx, a, w);

            } break;

        case GGML_OP_ROPE:

            {

                int n_embd_head = hparams.n_embd_head_v;

                int n_head = hparams.n_head();

                ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_head, n_head, 512);

                ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512);

                op_tensor = ggml_rope_ext(

                    ctx, a, b, w,

                    0, 0, 0, 0, 0,

                    0, 0, 0, 0

                );

            } break;

        case GGML_OP_SSM_CONV:

            {

                const int64_t n_seq_tokens = 512;

                const int64_t n_seqs       = 3;

                ggml_tensor * conv_x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0] - 1 + n_seq_tokens, w->ne[1], n_seqs);

                op_tensor = ggml_ssm_conv(ctx, conv_x, w);

            } break;

        case GGML_OP_SSM_SCAN:

            {

                // w is ssm_a, which is used to distinguish Mamba-1 and Mamba-2

                const int64_t d_state      = w->ne[0] == 1 ? hparams.ssm_d_state : w->ne[0];

                const int64_t n_head       = w->ne[1];

                const int64_t head_dim     = hparams.ssm_d_inner / n_head;

                const int64_t n_group      = hparams.ssm_n_group ? hparams.ssm_n_group : 1;

                const int64_t n_seq_tokens = 512;

                const int64_t n_seqs       = 3;

                ggml_tensor * s   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, head_dim, n_head, n_seqs);

                ggml_tensor * x   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_head, n_seq_tokens, n_seqs);

                ggml_tensor * dt  = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_head, n_seq_tokens, n_seqs);

                ggml_tensor * B   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);

                ggml_tensor * C   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);

                ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);

                op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids);

            } break;

        case GGML_OP_RWKV_WKV6:

            {

                // FIXME

                const int64_t S = 123;

                const int64_t H = 123;

                const int64_t n_tokens = 123;

                const int64_t n_seqs = 123;

                ggml_tensor  * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);

                ggml_tensor  * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);

                ggml_tensor  * r = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);

                ggml_tensor  * tf = w;

                ggml_tensor  * td = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);

                ggml_tensor  * state = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, n_seqs, S, H);

                op_tensor = ggml_rwkv_wkv6(ctx, k, v, r, tf, td, state);

            } break;

        case GGML_OP_IM2COL:

            {

                const int n_embd_inp = hparams.n_embd_inp();

                ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd_inp, w->ne[1], 1, 1);

                op_tensor = ggml_im2col(ctx, w, b, 1, 0, 0, 0, 1, 0, false, GGML_TYPE_F16);

            } break;

        case GGML_OP_SCALE:

            {

                op_tensor = ggml_scale(ctx, w, 1.0f);

            } break;

        default:

            GGML_ABORT("%s: missing test for op %s for tensor %s", __func__, ggml_op_name(op), w->name);

    }

    // create a temporary dummy buffer for the weight so that supports_op can check the buffer type

    GGML_ASSERT(w->buffer == nullptr);

    w->buffer = ggml_backend_buft_alloc_buffer(buft, 0);

    bool op_supported = ggml_backend_dev_supports_op(dev, op_tensor);

    ggml_backend_buffer_free(w->buffer);

    w->buffer = nullptr;

    return op_supported;

}

// lists of buffer types used for each layer

using buft_list_t = std::vector<std::pair<ggml_backend_dev_t, ggml_backend_buffer_type_t>>;

// find the first buffer type in the list that can use the tensor

static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, ggml_op op, const buft_list_t & buft_list) {

    GGML_ASSERT(!buft_list.empty());

    for (const auto & cur : buft_list) {

        ggml_backend_dev_t cur_dev = cur.first;

        ggml_backend_buffer_type_t cur_buft = cur.second;

        if (weight_buft_supported(hparams, tensor, op, cur_buft, cur_dev)) {

            return cur_buft;

        }

    }

    return nullptr;

}

// CPU: ACCEL -> GPU host -> CPU extra -> CPU

static buft_list_t make_cpu_buft_list(const std::vector<ggml_backend_dev_t> & devices, bool use_extra_bufts, bool no_host) {

    buft_list_t buft_list;

    // add ACCEL buffer types

    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {

        ggml_backend_dev_t dev = ggml_backend_dev_get(i);

        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_ACCEL) {

            auto * buft = ggml_backend_dev_buffer_type(dev);

            // skip

            if (buft != ggml_backend_cpu_buffer_type()) {

                buft_list.emplace_back(dev, buft);

            }

        }

    }

    // add a host buffer type

    // storing the tensors in a host buffer is useful when the processing of large batches

    // is offloaded to a GPU device, since it reduces the time spent on data transfers

    // generally, this will be done using the first device in the list

    // a better approach would be to handle this on a weight-by-weight basis using the offload_op

    // function of the device to determine if it would benefit from being stored in a host buffer

    if (!no_host) {

        for (auto * dev : devices) {

            ggml_backend_buffer_type_t buft = ggml_backend_dev_host_buffer_type(dev);

            if (buft) {

                buft_list.emplace_back(dev, buft);

                break;

            }

        }

    }

    // add extra buffer types

    if (use_extra_bufts) {

        auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);

        if (cpu_dev == nullptr) {

            throw std::runtime_error(format("%s: no CPU backend found", __func__));

        }

        auto * cpu_reg = ggml_backend_dev_backend_reg(cpu_dev);

        auto ggml_backend_dev_get_extra_bufts_fn = (ggml_backend_dev_get_extra_bufts_t)

            ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_dev_get_extra_bufts");

        if (ggml_backend_dev_get_extra_bufts_fn) {

            ggml_backend_buffer_type_t * extra_bufts = ggml_backend_dev_get_extra_bufts_fn(cpu_dev);

            while (extra_bufts && *extra_bufts) {

                buft_list.emplace_back(cpu_dev, *extra_bufts);

                ++extra_bufts;

            }

        }

    }

    // add the CPU buffer type

    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {

        ggml_backend_dev_t dev = ggml_backend_dev_get(i);

        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {

            buft_list.emplace_back(dev, ggml_backend_dev_buffer_type(dev));

        }

    }

    return buft_list;

}

// GPU: split if LLAMA_SPLIT_MODE_ROW -> GPU

static buft_list_t make_gpu_buft_list(ggml_backend_dev_t dev, llama_split_mode split_mode, const float * tensor_split) {

    buft_list_t buft_list;

    // add the device split buffer type if requested and available

    if (split_mode == LLAMA_SPLIT_MODE_ROW) {

        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);

        auto ggml_backend_split_buffer_type_fn = (ggml_backend_split_buffer_type_t)

            ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type");

        if (ggml_backend_split_buffer_type_fn) {

            size_t dev_index = [&]() {

                auto * reg = ggml_backend_dev_backend_reg(dev);

                for (size_t i = 0; i < ggml_backend_reg_dev_count(reg); ++i) {

                    if (ggml_backend_reg_dev_get(reg, i) == dev) {

                        return i;

                    }

                }

                throw std::runtime_error(format("device %s not found in its backend reg", ggml_backend_dev_name(dev)));

            }();

            auto * buft = ggml_backend_split_buffer_type_fn(dev_index, tensor_split);

            if (buft != nullptr) {

                buft_list.emplace_back(dev, buft);

            }

        }

    }

    // add the device default buffer type

    buft_list.emplace_back(dev, ggml_backend_dev_buffer_type(dev));

    // add the device extra buffer type (if any)

    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);

    auto ggml_backend_dev_get_extra_bufts_fn = (ggml_backend_dev_get_extra_bufts_t)

        ggml_backend_reg_get_proc_address(reg, "ggml_backend_dev_get_extra_bufts");

    if (ggml_backend_dev_get_extra_bufts_fn) {

        ggml_backend_buffer_type_t * extra_bufts = ggml_backend_dev_get_extra_bufts_fn(dev);

        while (extra_bufts && *extra_bufts) {

            buft_list.emplace_back(dev, *extra_bufts);

            ++extra_bufts;

        }

    }

    return buft_list;

}

struct llama_model::impl {

    impl() = default;

    ~impl() = default;

    uint64_t n_elements = 0;

    size_t n_bytes = 0;

    std::string desc_str;

    // model memory mapped files

    llama_mmaps mappings;

    // objects representing data potentially being locked in memory

    llama_mlocks mlock_bufs;

    llama_mlocks mlock_mmaps;

    // contexts where the model tensors metadata is stored as well as the corresponding buffers:

    std::vector<std::pair<ggml_context_ptr, std::vector<ggml_backend_buffer_ptr>>> ctxs_bufs;

    buft_list_t cpu_buft_list;

    std::map<ggml_backend_dev_t, buft_list_t> gpu_buft_list;

    struct layer_dev {

        ggml_backend_dev_t dev;

        buft_list_t * buft_list;

    };

    layer_dev dev_input = {};

    layer_dev dev_output = {};

    std::vector<layer_dev> dev_layer;

    bool has_tensor_overrides;

};

llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique<impl>()) {

    pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern;

}

llama_model::~llama_model() {

    for (auto * lora : loras) {

        delete lora;

    }

}

void llama_model::load_stats(llama_model_loader & ml) {

    pimpl->n_elements = ml.n_elements;

    pimpl->n_bytes = ml.n_bytes;

}

void llama_model::load_arch(llama_model_loader & ml) {

    arch = ml.get_arch();

    if (arch == LLM_ARCH_UNKNOWN) {

        throw std::runtime_error("unknown model architecture: '" + ml.get_arch_name() + "'");

    }

}

void llama_model::load_hparams(llama_model_loader & ml) {

    const gguf_context * ctx = ml.meta.get();

    // get metadata as string

    for (int i = 0; i < gguf_get_n_kv(ctx); i++) {

        gguf_type type = gguf_get_kv_type(ctx, i);

        if (type == GGUF_TYPE_ARRAY) {

            continue;

        }

        const char * name = gguf_get_key(ctx, i);

        const std::string value = gguf_kv_to_str(ctx, i);

        gguf_kv.emplace(name, value);

    }

    // get general kv

    ml.get_key(LLM_KV_GENERAL_NAME, name, false);

    // everything past this point is not vocab-related

    // for CLIP models, we only need to load tensors, no hparams

    if (hparams.vocab_only || ml.get_arch() == LLM_ARCH_CLIP) {

        return;

    }

    ml.get_key(LLM_KV_CONTEXT_LENGTH,          hparams.n_ctx_train);

    ml.get_key(LLM_KV_EMBEDDING_LENGTH,        hparams.n_embd);

    ml.get_key(LLM_KV_EMBEDDING_LENGTH_OUT,    hparams.n_embd_out_impl, false);

    ml.get_key(LLM_KV_BLOCK_COUNT,             hparams.n_layer);

    ml.get_key(LLM_KV_EXPERT_COUNT,            hparams.n_expert,        false);

    ml.get_key(LLM_KV_EXPERT_USED_COUNT,       hparams.n_expert_used,   false);

    ml.get_key(LLM_KV_EXPERT_GROUP_COUNT,      hparams.n_expert_groups, false);

    ml.get_key(LLM_KV_EXPERT_GROUP_USED_COUNT, hparams.n_group_used,    false);

    if (arch == LLM_ARCH_WAVTOKENIZER_DEC) {

        ml.get_key(LLM_KV_FEATURES_LENGTH,  hparams.n_embd);

        ml.get_key(LLM_KV_EMBEDDING_LENGTH, hparams.n_embd_out_impl);

        ml.get_key(LLM_KV_POSNET_EMBEDDING_LENGTH, hparams.posnet.n_embd);

        ml.get_key(LLM_KV_POSNET_BLOCK_COUNT,      hparams.posnet.n_layer);

        ml.get_key(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, hparams.convnext.n_embd);

        ml.get_key(LLM_KV_CONVNEXT_BLOCK_COUNT,      hparams.convnext.n_layer);

    }

    GGML_ASSERT(hparams.n_expert <= LLAMA_MAX_EXPERTS);

    GGML_ASSERT(hparams.n_expert_used <= hparams.n_expert);

    if (hparams.n_expert > 0) {

        GGML_ASSERT(hparams.n_expert_used > 0);

        GGML_ASSERT(hparams.n_expert_groups < hparams.n_expert);

        if (hparams.n_expert_groups > 1) {

            GGML_ASSERT(hparams.n_expert % hparams.n_expert_groups == 0);

            GGML_ASSERT(hparams.n_group_used > 0);

            GGML_ASSERT(hparams.n_group_used < hparams.n_expert_groups);

        }

    } else {

        GGML_ASSERT(hparams.n_expert_used == 0);

        GGML_ASSERT(hparams.n_expert_groups == 0);

    }

    std::fill(hparams.n_head_arr.begin(),    hparams.n_head_arr.end(),    0);

    std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0);

    std::fill(hparams.n_ff_arr.begin(),      hparams.n_ff_arr.end(),      0);

    std::fill(

        hparams.recurrent_layer_arr.begin(),

        hparams.recurrent_layer_arr.end(),

        llm_arch_is_recurrent(ml.get_arch()));

    std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0);

    std::fill(hparams.swa_layers.begin(), hparams.swa_layers.end(), 0);

    std::fill(hparams.xielu_alpha_n.begin(), hparams.xielu_alpha_n.end(), 0.0f);

    std::fill(hparams.xielu_alpha_p.begin(), hparams.xielu_alpha_p.end(), 0.0f);

    std::fill(hparams.xielu_beta.begin(), hparams.xielu_beta.end(), 0.0f);

    std::fill(hparams.xielu_eps.begin(), hparams.xielu_eps.end(), 0.0f);

    std::fill(hparams.swiglu_clamp_exp.begin(),   hparams.swiglu_clamp_exp.end(),   0.0f);

    std::fill(hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.end(), 0.0f);

    ml.get_key_or_arr(LLM_KV_FEED_FORWARD_LENGTH,  hparams.n_ff_arr,   hparams.n_layer, false);

    ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, hparams.n_layer, false);

    // n_head_kv is optional, default to n_head

    hparams.n_head_kv_arr = hparams.n_head_arr;

    ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT_KV, hparams.n_head_kv_arr, hparams.n_layer, false);

    bool rope_finetuned = false;

    ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);

    hparams.rope_finetuned = rope_finetuned;

    hparams.n_ctx_orig_yarn = hparams.n_ctx_train;

    ml.get_key(LLM_KV_ROPE_SCALING_ORIG_CTX_LEN, hparams.n_ctx_orig_yarn, false);

    // rope_freq_base (optional)

    hparams.rope_freq_base_train = 10000.0f;

    ml.get_key(LLM_KV_ROPE_FREQ_BASE, hparams.rope_freq_base_train, false);

    std::string rope_scaling("linear");

    ml.get_key(LLM_KV_ROPE_SCALING_TYPE, rope_scaling, false);

    hparams.rope_scaling_type_train = llama_rope_scaling_type_from_string(rope_scaling);

    GGML_ASSERT(hparams.rope_scaling_type_train != LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED);

    // TODO: Handle SWA metadata similarly when models start implementing it

    // rope_freq_scale (inverse of the kv) is optional

    float ropescale = 0.0f;

    if (!ml.get_key(LLM_KV_ROPE_SCALING_FACTOR, ropescale, false)) {

        // try the old key name

        ml.get_key(LLM_KV_ROPE_SCALE_LINEAR, ropescale, false);

    }

    hparams.rope_freq_scale_train = ropescale == 0.0f ? 1.0f : 1.0f/ropescale;

    ml.get_key(LLM_KV_ROPE_SCALING_ATTN_FACTOR, hparams.rope_attn_factor, false);

    // non-transformer models do not have attention heads

    if (hparams.n_head() > 0) {

        // gpt-neox n_rot = rotary_pct * (n_embd / n_head)

        // gpt-j n_rot = rotary_dim

        hparams.n_embd_head_k = hparams.n_embd / hparams.n_head();

        ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH, hparams.n_embd_head_k, false);

        hparams.n_embd_head_v = hparams.n_embd / hparams.n_head();

        ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH, hparams.n_embd_head_v, false);

        // sanity check for n_rot (optional)

        hparams.n_rot = hparams.n_embd_head_k;

        ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT, hparams.n_rot, false);

        if (arch == LLM_ARCH_LLAMA || arch == LLM_ARCH_DECI || arch == LLM_ARCH_FALCON || arch == LLM_ARCH_LLAMA_EMBED) {

            if (hparams.n_rot != hparams.n_embd_head_k) {

                throw std::runtime_error(format("invalid n_rot: %u, expected %u", hparams.n_rot, hparams.n_embd_head_k));

            }

        }

    } else {

        hparams.n_rot = 0;

        hparams.n_embd_head_k = 0;

        hparams.n_embd_head_v = 0;

    }

    // for differentiating model types

    uint32_t n_vocab = 0;

    ml.get_key(LLM_KV_VOCAB_SIZE, n_vocab, false) || ml.get_arr_n(LLM_KV_TOKENIZER_LIST, n_vocab, false);

    // for classifier models

    ml.get_arr(LLM_KV_CLASSIFIER_OUTPUT_LABELS, classifier_labels, false);

    if (!classifier_labels.empty()) {

        hparams.n_cls_out = classifier_labels.size();

    }

    // arch-specific KVs

    switch (arch) {

        case LLM_ARCH_LLAMA:

        case LLM_ARCH_LLAMA_EMBED:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                if (hparams.n_expert == 8) {

                    switch (hparams.n_layer) {

                        case 32: type = LLM_TYPE_8x7B; break;

                        case 56: type = LLM_TYPE_8x22B; break;

                        default: type = LLM_TYPE_UNKNOWN;

                    }

                } else {

                    switch (hparams.n_layer) {

                        case 16: type = LLM_TYPE_1B; break; // Llama 3.2 1B

                        case 22: type = LLM_TYPE_1B; break;

                        case 26: type = LLM_TYPE_3B; break;

                        case 28: type = LLM_TYPE_3B; break; // Llama 3.2 3B

                        case 30: type = LLM_TYPE_256M; break; // smoldocling 256M

                        // granite uses a vocab with len 49152

                        case 32: type = n_vocab == 49152 ? LLM_TYPE_3B : (n_vocab < 40000 ? LLM_TYPE_7B : LLM_TYPE_8B); break;

                        case 36: type = LLM_TYPE_8B; break; // granite

                        case 40: type = LLM_TYPE_13B; break;

                        case 48: type = LLM_TYPE_34B; break;

                        case 60: type = LLM_TYPE_30B; break;

                        case 80: type = hparams.n_head() == hparams.n_head_kv() ? LLM_TYPE_65B : LLM_TYPE_70B; break;

                        default: type = LLM_TYPE_UNKNOWN;

                    }

                }

            } break;

        case LLM_ARCH_LLAMA4:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,  hparams.n_ff_exp);

                ml.get_key(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,   hparams.n_moe_layer_step);

                const bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);

                if (found_swa && hparams.n_swa == 0) {

                    hparams.swa_type             = LLAMA_SWA_TYPE_NONE;

                    hparams.n_no_rope_layer_step = hparams.n_layer; // always use rope

                } else {

                    hparams.swa_type                = LLAMA_SWA_TYPE_CHUNKED;

                    hparams.n_swa                   = 8192;

                    hparams.n_attn_temp_floor_scale = 8192;

                    hparams.f_attn_temp_scale       = 0.1f;

                    hparams.f_attn_temp_offset      = 1.0f;

                    hparams.set_swa_pattern(4);   // pattern: 3 chunked - 1 full

                    hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;

                    hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;

                    ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);

                }

                switch (hparams.n_expert) {

                    case 0: {

                        // MobileLLM (no MoE)

                        switch (hparams.n_embd) {

                            case 2048: type = LLM_TYPE_140M; break;

                            case 4096: type = LLM_TYPE_360M; break;

                            case 6144: type = LLM_TYPE_950M; break;

                            default:   type = LLM_TYPE_UNKNOWN;

                        }

                    } break;

                    case 16:  type = LLM_TYPE_17B_16E; break;

                    case 128: type = LLM_TYPE_17B_128E; break;

                    default:  type = LLM_TYPE_UNKNOWN;

                }

                hparams.use_kq_norm = type != LLM_TYPE_17B_128E;

            } break;

        case LLM_ARCH_ARCEE:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                // Arcee uses the same structure as Llama

                switch (hparams.n_layer) {

                    case 36: type = LLM_TYPE_4B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_AFMOE:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead);

                ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,  hparams.n_ff_exp);

                ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);

                ml.get_key(LLM_KV_EXPERT_GATING_FUNC,          hparams.expert_gating_func, false);

                ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);

                ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,         hparams.expert_weights_norm, false);

                ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW,    hparams.n_swa, false);

                // Set up interleaved sliding window attention (ISWA)

                // Pattern: 3 sliding - 1 full (global_attn_every_n_layers = 4)

                if (hparams.n_swa > 0) {

                    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;

                    hparams.set_swa_pattern(4);

                    hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;

                    hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;

                    ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);

                } else {

                    hparams.swa_type = LLAMA_SWA_TYPE_NONE;

                }

                // Default to sigmoid if not set

                if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {

                    hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;

                }

                switch (hparams.n_layer) {

                    case 56: type = LLM_TYPE_6B; break;

                    case 32: type = LLM_TYPE_26B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_DECI:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                switch (hparams.n_layer) {

                    case 32: type = LLM_TYPE_7B; break;

                    case 80: type = LLM_TYPE_70B; break;

                    case 162: type = LLM_TYPE_405B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_MINICPM:

            {

                // Backward-compatible defaults for older MiniCPM GGUFs

                hparams.f_embedding_scale = 12.0f;

                hparams.f_residual_scale  = 1.4f / sqrtf(float(hparams.n_layer));

                hparams.f_logit_scale     = hparams.n_embd ? (256.0f / float(hparams.n_embd)) : 1.0f;

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                // Optional KV reads, override defaults if present in newer GGUF exports

                ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, /*required=*/false);

                ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, /*required=*/false);

                ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, /*required=*/false);

                // MiniCPM uses rope by default, unlike Granite which uses it as a switch

                hparams.rope_finetuned = true;

                switch (hparams.n_layer) {

                    case 52: type = LLM_TYPE_1B; break;

                    case 40: type = LLM_TYPE_2B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_MINICPM3:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,       hparams.n_lora_q);

                ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK,      hparams.n_lora_kv);

                switch (hparams.n_layer) {

                    case 62: type = LLM_TYPE_4B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_GROK:

            {

                // defaults for old GGUFs

                hparams.yarn_beta_fast = 8.0f;

                hparams.f_logit_scale = 0.5773502691896257f;

                hparams.f_embedding_scale = 78.38367176906169f;

                hparams.f_attn_out_scale = 0.08838834764831845f;

                hparams.f_attn_logit_softcapping = 30.0f;

                hparams.f_router_logit_softcapping = 30.0f;

                // no final_logit_softcapping in grok-1

                hparams.f_final_logit_softcapping = 0.0f;

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,  hparams.f_norm_rms_eps);

                ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,   hparams.n_ff_exp, false);

                ml.get_key(LLM_KV_LOGIT_SCALE,                  hparams.f_logit_scale, false);

                ml.get_key(LLM_KV_EMBEDDING_SCALE,              hparams.f_embedding_scale, false);

                ml.get_key(LLM_KV_ATTENTION_OUTPUT_SCALE,       hparams.f_attn_out_scale, false);

                ml.get_key(LLM_KV_ATTN_LOGIT_SOFTCAPPING,       hparams.f_attn_logit_softcapping, false);

                ml.get_key(LLM_KV_ROUTER_LOGIT_SOFTCAPPING,     hparams.f_router_logit_softcapping, false);

                ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING,      hparams.f_final_logit_softcapping, false);

                ml.get_key(LLM_KV_ATTENTION_TEMPERATURE_LENGTH,  hparams.attn_temp_length, false);

                ml.get_key(LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR,  hparams.yarn_ext_factor, false);

                ml.get_key(LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, hparams.yarn_attn_factor, false);

                ml.get_key(LLM_KV_ROPE_SCALING_YARN_BETA_FAST,   hparams.yarn_beta_fast, false);

                ml.get_key(LLM_KV_ROPE_SCALING_YARN_BETA_SLOW,   hparams.yarn_beta_slow, false);

                switch (hparams.n_layer) {

                    case 64: type = LLM_TYPE_314B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_FALCON:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);

                switch (hparams.n_layer) {

                    case 32: type = LLM_TYPE_7B; break;

                    case 60: type = LLM_TYPE_40B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_BAICHUAN:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                switch (hparams.n_layer) {

                    case 32: type = LLM_TYPE_7B; break;

                    case 40: type = LLM_TYPE_13B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

                if (type == LLM_TYPE_13B) {

                    // TODO: become GGUF KV parameter

                    hparams.f_max_alibi_bias = 8.0f;

                }

            } break;

        case LLM_ARCH_STARCODER:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);

                switch (hparams.n_layer) {

                    case 24: type = LLM_TYPE_1B; break;

                    case 36: type = LLM_TYPE_3B; break;

                    case 42: type = LLM_TYPE_7B; break;

                    case 40: type = LLM_TYPE_15B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

            } break;

        case LLM_ARCH_REFACT:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);

                switch (hparams.n_layer) {

                    case 32: type = LLM_TYPE_1B; break;

                    default: type = LLM_TYPE_UNKNOWN;

                }

                // TODO: become GGUF KV parameter

                hparams.f_max_alibi_bias = 8.0f;

            } break;

        case LLM_ARCH_BERT:

            {

                ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS,    hparams.f_norm_eps);

                ml.get_key(LLM_KV_ATTENTION_CAUSAL,           hparams.causal_attn);

                ml.get_key(LLM_KV_POOLING_TYPE,               hparams.pooling_type, false);

                switch (hparams.n_layer) {

                    case 3: