From 233225db8f156ae623b33a9cb26d5f4d84c813a8 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Tue, 24 Mar 2026 08:18:28 +0100 Subject: [PATCH] Take into account layer sizes for setting GPU layers (count 2) (#1498) * Also take into account KV cache * Take into account attn_wkv_b and mla = 3 compute buffers * WIP * Minor --- src/llama-arch.h | 4 ++++ src/llama-model.cpp | 23 ++++++++++++++++++++++ src/llama.cpp | 47 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 74 insertions(+) diff --git a/src/llama-arch.h b/src/llama-arch.h index a757a71e..6b56dfe1 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -340,6 +340,8 @@ enum llm_tensor { LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, + + LLM_TENSOR_UNKNOWN, }; llm_arch llm_arch_from_string(const std::string & name); @@ -348,3 +350,5 @@ const char * llama_model_arch_name(llm_arch arch); bool llm_arch_is_recurrent(const llm_arch & arch); bool llm_arch_is_hybrid(const llm_arch & arch); + +llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 8837f271..dd88a4cf 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1837,6 +1837,29 @@ bool llama_model_has_recurrent(const llama_model * model) { return llm_arch_is_hybrid(model->arch) || llm_arch_is_recurrent(model->arch); } +llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il) { + auto it = LLM_TENSOR_NAMES.find(arch); + if (it == LLM_TENSOR_NAMES.end()) { + printf("%s: Oops, did not find arch\n", __func__); + return LLM_TENSOR_UNKNOWN; + } + if (il < 0) { + for (auto & entry : it->second) { + if (tensor_name.find(entry.second) == 0) { + return entry.first; + } + } + return LLM_TENSOR_UNKNOWN; + } + for (auto & entry : it->second) { + auto this_name = ::format(entry.second.c_str(), il); + if (tensor_name.find(this_name) == 0) { + return entry.first; + } + } + return LLM_TENSOR_UNKNOWN; +} + size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const { if (il < 0 || il >= hparams.n_layer) return 0; if (hparams.recurrent_layer_arr[il]) { diff --git a/src/llama.cpp b/src/llama.cpp index 8e7c7803..c28e04f1 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -2013,6 +2013,9 @@ static std::pair, double> get_layer_sizes(const llama_model_ if (has_mla) { mla_tensors.resize(n_layer); } + std::vector ffn_exps(n_layer, 0), ffn_shexp(n_layer, 0), ffn_dense(n_layer, 0); + std::vector attn(n_layer, 0); + std::vector has_layer_norm(n_layer, false); size_t ow_size = 0; size_t embd_size = 0; for (int i = 0; i < ml.n_tensors; ++i) { @@ -2027,6 +2030,9 @@ static std::pair, double> get_layer_sizes(const llama_model_ ow_size = size; continue; } + if (name == "output_norm.weight") { + continue; + } auto pos = name.find("blk."); if (pos != 0) { printf("Oops: tensor with strange name %s\n", name.c_str()); @@ -2049,17 +2055,58 @@ static std::pair, double> get_layer_sizes(const llama_model_ continue; } result[il] += size; + if (auto pos = name.rfind(".bias"); pos < name.size() && name.size() - pos == 4) { + // bias, we don't need to account for those + continue; + } + bool is_mla = false; if (has_mla) { if (name.find("attn_k_b.weight") != std::string::npos) { mla_tensors[il].wk_b = t; + is_mla = true; } else if (name.find("attn_v_b.weight") != std::string::npos) { mla_tensors[il].wv_b = t; + is_mla = true; } else if (name.find("attn_kv_b.weight") != std::string::npos) { mla_tensors[il].wkv_b = t; + is_mla = true; } } + if (!is_mla) { + auto ttype = llm_tensor_type(model.arch, name, il); + if (ttype == LLM_TENSOR_UNKNOWN) printf("Oops: got unknows for tensor %s\n", name.c_str()); + if (ttype == LLM_TENSOR_FFN_GATE_UP_EXPS || ttype == LLM_TENSOR_FFN_GATE_EXPS || ttype == LLM_TENSOR_FFN_UP_EXPS || ttype == LLM_TENSOR_FFN_DOWN_EXPS) { + auto size = t->ne[1] * n_ubatch * model.hparams.n_expert_used * sizeof(float); + ffn_exps[il] += size; + } + else if (ttype == LLM_TENSOR_FFN_UP_SHEXP || ttype == LLM_TENSOR_FFN_GATE_SHEXP || ttype == LLM_TENSOR_FFN_DOWN_SHEXP) { + ffn_shexp[il] += t->ne[1] * n_ubatch * sizeof(float); + } + else if (ttype == LLM_TENSOR_FFN_UP || ttype == LLM_TENSOR_FFN_GATE || ttype == LLM_TENSOR_FFN_DOWN) { + ffn_dense[il] += t->ne[1] * n_ubatch * sizeof(float); + } + else if (ttype == LLM_TENSOR_ATTN_Q || ttype == LLM_TENSOR_ATTN_K || ttype == LLM_TENSOR_ATTN_V) { + attn[il] += t->ne[1] * n_ubatch * sizeof(float); + } + else if (ttype == LLM_TENSOR_ATTN_OUT) { + // This assume FA. With FA, we need to have at the same time Q, K, V and FA(Q, K, V). + // FA(Q, K, V) has the dimension t->ne[0] x n_ubatch + attn[il] += t->ne[0] * n_ubatch * sizeof(float); + } + else if (ttype == LLM_TENSOR_ATTN_NORM || ttype == LLM_TENSOR_ATTN_NORM_2 || ttype == LLM_TENSOR_ATTN_OUT_NORM || ttype == LLM_TENSOR_ATTN_POST_NORM || + ttype == LLM_TENSOR_FFN_NORM || ttype == LLM_TENSOR_FFN_POST_NORM || ttype == LLM_TENSOR_LAYER_OUT_NORM) { + has_layer_norm[il] = true; + } + } + } + for (int il = 0; il < n_layer; ++il) { + auto ffn = std::max(std::max(ffn_exps[il], ffn_shexp[il]), ffn_dense[il]); + auto inp = model.hparams.n_embd * n_ubatch * sizeof(float); + if (has_layer_norm[il]) inp *= 2; + ffn += inp; + compute[il] = std::max(compute[il], ffn); } if (has_mla) { for (int il = 0; il < n_layer; ++il) {