echodict/llama.cpp

1#include "llama-hparams.h"2 3#include "ggml.h"4 5#include <algorithm>6#include <cassert>7 8void llama_hparams::set_swa_pattern(uint32_t n_pattern, bool dense_first) {9    if (dense_first) {10        for (uint32_t il = 0; il < n_layer; ++il) {11            swa_layers[il] = n_pattern == 0 || (il % n_pattern != 0);12        }13    } else {14        for (uint32_t il = 0; il < n_layer; ++il) {15            swa_layers[il] = n_pattern == 0 || (il % n_pattern < (n_pattern - 1));16        }17    }18}19 20bool llama_hparams::is_swa_any() const {21    for (uint32_t il = 0; il < n_layer; ++il) {22        if (swa_layers[il]) {23            return true;24        }25    }26 27    return false;28}29 30uint32_t llama_hparams::n_head(uint32_t il) const {31    if (il < n_layer) {32        return n_head_arr[il];33    }34 35    GGML_ABORT("fatal error");36}37 38uint32_t llama_hparams::n_head_kv(uint32_t il) const {39    if (il < n_layer) {40        return n_head_kv_arr[il];41    }42 43    GGML_ABORT("fatal error");44}45 46uint32_t llama_hparams::n_ff(uint32_t il) const {47    if (il < n_layer) {48        return n_ff_arr[il];49    }50 51    GGML_ABORT("fatal error");52}53 54uint32_t llama_hparams::n_gqa(uint32_t il) const {55    const uint32_t n_head    = this->n_head(il);56    const uint32_t n_head_kv = this->n_head_kv(il);57 58    if (n_head_kv == 0) {59        return 0;60    }61 62    return n_head/n_head_kv;63}64 65uint32_t llama_hparams::n_rot(uint32_t il) const {66    if (il < n_layer) {67        return is_swa(il) ? n_rot_swa : n_rot_full;68    }69 70    GGML_ABORT("fatal error");71}72 73uint32_t llama_hparams::n_embd_inp() const {74    uint32_t n_embd_inp = n_embd;75 76    if (n_deepstack_layers > 0) {77        n_embd_inp += n_embd * n_deepstack_layers;78    }79 80    return n_embd_inp;81}82 83uint32_t llama_hparams::n_embd_out() const {84    return n_embd_out_impl > 0 ? n_embd_out_impl : n_embd;85}86 87uint32_t llama_hparams::n_embd_head_k(uint32_t il) const {88    if (il < n_layer) {89        return is_swa(il) ? n_embd_head_k_swa : n_embd_head_k_full;90    }91 92    GGML_ABORT("fatal error");93}94 95uint32_t llama_hparams::n_embd_head_v(uint32_t il) const {96    if (il < n_layer) {97        return is_swa(il) ? n_embd_head_v_swa : n_embd_head_v_full;98    }99 100    GGML_ABORT("fatal error");101}102 103uint32_t llama_hparams::n_embd_k_gqa(uint32_t il) const {104    const uint32_t n_head_kv = this->n_head_kv(il);105 106    return n_embd_head_k(il) * n_head_kv;107}108 109uint32_t llama_hparams::n_embd_v_gqa(uint32_t il) const {110    const uint32_t n_head_kv = this->n_head_kv(il);111 112    return n_embd_head_v(il) * n_head_kv;113}114 115bool llama_hparams::is_n_embd_k_gqa_variable() const {116    const uint32_t val = n_embd_k_gqa();117    for (uint32_t il = 0; il < n_layer; ++il) {118        if (val != n_embd_k_gqa(il)) {119            return true;120        }121    }122 123    return false;124}125 126bool llama_hparams::is_n_embd_v_gqa_variable() const {127    const uint32_t val = n_embd_v_gqa();128    for (uint32_t il = 0; il < n_layer; ++il) {129        if (val != n_embd_v_gqa(il)) {130            return true;131        }132    }133 134    return false;135}136 137uint32_t llama_hparams::n_embd_k_gqa_max() const {138    uint32_t val = n_embd_k_gqa();139    for (uint32_t il = 0; il < n_layer; ++il) {140        val = std::max(val, n_embd_k_gqa(il));141    }142 143    return val;144}145 146uint32_t llama_hparams::n_embd_v_gqa_max() const {147    uint32_t val = n_embd_v_gqa();148    for (uint32_t il = 0; il < n_layer; ++il) {149        val = std::max(val, n_embd_v_gqa(il));150    }151 152    return val;153}154 155uint32_t llama_hparams::n_embd_r() const {156    if (wkv_head_size != 0) {157        // for RWKV models158        return token_shift_count * n_embd;159    }160 161    if (n_shortconv_l_cache != 0) {162        // for LFM2 models163        return n_embd * (n_shortconv_l_cache - 1);164    }165 166    if (n_embd_head_kda != 0) {167        // for Kimi KDA layers168        // Conv state for Q, K, V: 3 * (d_conv - 1) * n_head * head_dim169        const uint32_t d_inner = n_head() * n_embd_head_kda;  // 32 * 128 = 4096170        return 3 * (ssm_d_conv > 0 ? ssm_d_conv - 1 : 3) * d_inner;171    }172 173    // TODO: maybe support other convolution strides than 1174    // NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed175    // Corresponds to Mamba's conv_states size176    return (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state);177}178 179uint32_t llama_hparams::n_embd_s() const {180    if (wkv_head_size != 0) {181        // corresponds to RWKV's wkv_states size182        return n_embd * wkv_head_size;183    }184 185    if (n_embd_head_kda != 0) {186        // for Kimi KDA layers187        // Full recurrent state: head_dim * head_dim * n_head188        // h tensor shape for delta attention: [head_dim, head_dim, n_head]189        return n_embd_head_kda * n_embd_head_kda * n_head();  // 128 * 128 * 32 = 524288190    }191 192    // corresponds to Mamba's ssm_states size193    return ssm_d_state * ssm_d_inner;194}195 196bool llama_hparams::is_recurrent(uint32_t il) const {197    if (il < n_layer) {198        return recurrent_layer_arr[il];199    }200 201    GGML_ABORT("%s: il (%u) out of bounds (n_layer: %u)\n", __func__, il, n_layer);202}203 204uint32_t llama_hparams::n_pos_per_embd() const {205    return rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1;206}207 208bool llama_hparams::is_swa(uint32_t il) const {209    if (il < n_layer) {210        return swa_layers[il];211    }212 213    GGML_ABORT("fatal error");214}215 216bool llama_hparams::is_mla() const {217    assert((n_embd_head_k_mla_impl == 0 && n_embd_head_v_mla_impl == 0) ||218           (n_embd_head_k_mla_impl != 0 && n_embd_head_v_mla_impl != 0));219 220    return n_embd_head_k_mla_impl != 0 && n_embd_head_v_mla_impl != 0;221}222 223uint32_t llama_hparams::n_embd_head_k_mla() const {224    return is_mla() ? n_embd_head_k_mla_impl : n_embd_head_k();225}226 227uint32_t llama_hparams::n_embd_head_v_mla() const {228    return is_mla() ? n_embd_head_v_mla_impl : n_embd_head_v();229}230 231bool llama_hparams::has_kv(uint32_t il) const {232    if (n_layer_kv_from_start >= 0) {233        if (il < (uint32_t) n_layer_kv_from_start) {234            return true;235        }236 237        return false;238    }239 240    // by default, all layers have kv241    return true;242}243 244uint32_t llama_hparams::n_layer_kv() const {245    uint32_t res = 0;246 247    for (uint32_t il = 0; il < n_layer; ++il) {248        if (has_kv(il)) {249            res++;250        }251    }252 253    return res;254}255 256bool llama_hparams::use_mrope() const {257    return rope_sections[0] > 0 && rope_sections[1] > 0;258}259