From b7a2bde4cc4ce7772594a65569c2f659a3420583 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 23 Mar 2026 17:46:53 +0100 Subject: [PATCH] Take into account layer sizes for setting GPU layers (cont) (#1476) * Also take into account KV cache * Take into account attn_wkv_b and mla = 3 compute buffers --- common/common.cpp | 119 ++++++++++++++++++++++++++------------------ include/llama.h | 8 +++ src/llama-model.cpp | 27 ++++++++++ src/llama-model.h | 4 ++ src/llama.cpp | 105 +++++++++++++++++++++++++++++++++++--- 5 files changed, 207 insertions(+), 56 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 1cd51fef..ddde7502 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -3235,46 +3235,6 @@ void llama_lora_adapters_apply(struct llama_context * ctx, std::vector get_batch_ubatch(const gpt_params & params) { + int n_batch = params.n_batch; + int n_ubatch = params.n_ubatch; + if (params.n_ctx > 0) { + n_batch = std::min(n_batch, params.n_ctx); + } + if (!params.mmproj.path.empty()) { + // temporary fix for qwen mtmd + n_batch = std::max(n_batch, n_ubatch); + n_ubatch = n_batch; + fprintf(stdout, "Adjust batch size for mtmd: u_batch = %d, batch = %d\n", n_ubatch, n_batch); + } else { + n_ubatch = std::min(n_batch, n_ubatch); + } + return {n_batch, n_ubatch}; +} + +struct llama_model_params common_model_params_to_llama(const gpt_params & params) { + auto mparams = llama_model_default_params(); + mparams.devices = params.devices.c_str(); + + if (params.n_gpu_layers != -1) { + mparams.n_gpu_layers = params.n_gpu_layers; + } + mparams.mla = params.mla_attn; + mparams.dry_run = params.dry_run; + mparams.rpc_servers = params.rpc_servers.c_str(); + mparams.main_gpu = params.main_gpu; + mparams.max_gpu = params.max_gpu; + mparams.ncmoe = params.ncmoe; + mparams.type_k = kv_cache_type_from_str(params.cache_type_k); + mparams.type_v = kv_cache_type_from_str(params.cache_type_v); + mparams.max_ctx_size = params.n_ctx; + mparams.n_seq_max = params.n_parallel; + mparams.n_ubatch = get_batch_ubatch(params).second; + mparams.amb = params.attn_max_batch; + mparams.split_mode = params.split_mode; + mparams.tensor_split = params.tensor_split; + mparams.use_mmap = params.use_mmap; + mparams.use_mlock = params.use_mlock; + mparams.check_tensors = params.check_tensors; + mparams.repack_tensors = params.repack_tensors; + mparams.use_thp = params.use_thp; + mparams.validate_quants = params.validate_quants; + mparams.merge_qkv = params.merge_qkv; + mparams.merge_up_gate_exps = params.merge_up_gate_exps; + mparams.mtp = params.has_mtp; + mparams.flash_attn = params.flash_attn; + if (params.kv_overrides.empty()) { + mparams.kv_overrides = NULL; + } else { + GGML_ASSERT(params.kv_overrides.back().key[0] == 0 && "KV overrides not terminated with empty key"); + mparams.kv_overrides = params.kv_overrides.data(); + } + if (params.tensor_buft_overrides.empty()) { + mparams.tensor_buft_overrides = NULL; + } else { + GGML_ASSERT(params.tensor_buft_overrides.back().pattern == nullptr && "Tensor buffer overrides not terminated with empty pattern"); + mparams.tensor_buft_overrides = params.tensor_buft_overrides.data(); + } + if (!mparams.flash_attn && ggml_is_quantized(mparams.type_v)) { + throw std::runtime_error("Quantized V cache cannot be used without flash attention"); + } + + return mparams; +} + static ggml_type ggml_type_from_str(const std::string & s) { if (s == "f32") { return GGML_TYPE_F32; @@ -3331,15 +3358,8 @@ static ggml_type ggml_type_from_str(const std::string & s) { struct llama_context_params common_context_params_to_llama(const gpt_params & params) { auto cparams = llama_context_default_params(); - int n_batch = params.n_batch; - int n_ubatch = params.n_ubatch; - // temporary fix for qwen mtmd - if (!params.mmproj.path.empty()) { - n_batch = std::max(params.n_batch, params.n_ubatch); - n_ubatch = params.n_batch; - fprintf(stdout, "Adjust batch size for mtmd: u_batch = %d, batch = %d\n", n_ubatch, n_batch); - } + auto [n_batch, n_ubatch] = get_batch_ubatch(params); cparams.n_ctx = params.n_ctx; cparams.n_seq_max = params.n_parallel; @@ -3387,6 +3407,9 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa cparams.type_k = kv_cache_type_from_str(params.cache_type_k); cparams.type_v = kv_cache_type_from_str(params.cache_type_v); cparams.type_reduce = ggml_type_from_str(params.reduce_type); + if (!cparams.flash_attn && ggml_is_quantized(cparams.type_v)) { + throw std::runtime_error("Quantized V cache cannot be used without flash attention"); + } if (!params.offload_policy.empty()) cparams.offload_policy = (void *)¶ms.offload_policy; if (!params.cuda_params.empty()) cparams.cuda_params = (void *)params.cuda_params.data(); diff --git a/include/llama.h b/include/llama.h index 9fb7e77d..186b16fb 100644 --- a/include/llama.h +++ b/include/llama.h @@ -372,6 +372,13 @@ extern "C" { int32_t max_gpu; int32_t ncmoe; + enum ggml_type type_k; + enum ggml_type type_v; + uint32_t max_ctx_size; + int32_t n_seq_max; + int32_t n_ubatch; + int32_t amb; + // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; @@ -403,6 +410,7 @@ extern "C" { bool merge_up_gate_exps; // if true, merge ffn_up_exps and ffn_gate_exps tensors into a single, contiguous tensor bool mtp; // if true, load MTP layers if present bool dry_run; // skip loading tensors + bool flash_attn; }; // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 328220d7..8837f271 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1,4 +1,5 @@ #include "llama-model.h" +#include "llama-cparams.h" #include @@ -1835,3 +1836,29 @@ bool llama_model_is_hybrid(const llama_model * model) { bool llama_model_has_recurrent(const llama_model * model) { return llm_arch_is_hybrid(model->arch) || llm_arch_is_recurrent(model->arch); } + +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]) { + auto state_sots = std::min(std::max(1, n_seq_max), kv_size); + return hparams.n_embd_v_s() * state_sots * sizeof(float); + } + bool is_mla_attn = arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4; + if (is_mla_attn && mla_attn) { + auto n_embd_head_qk_rope = hparams.n_rot; + auto kv_lora_rank = hparams.n_lora_kv; + if (flash_attn) { + return ggml_row_size(type_k, kv_lora_rank + n_embd_head_qk_rope) * kv_size; + } + auto kv_type = mla_attn == 1 ? type_k : type_v; + auto size = ggml_row_size(kv_type, kv_lora_rank + n_embd_head_qk_rope) * kv_size; + if (mla_attn == 1) { + size += ggml_row_size(type_v, kv_lora_rank*kv_size); + } + return size; + } + auto n_head_kv = hparams.n_head_kv(il); + auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k) * n_head_kv*kv_size; + auto v_size = ggml_row_size(type_v, hparams.n_embd_v_gqa(il)) * kv_size; + return k_size + v_size; +} diff --git a/src/llama-model.h b/src/llama-model.h index 2e81dbe8..a89aac59 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -362,6 +362,8 @@ struct rpc_device { uint32_t device; }; +struct llama_cparams; + struct llama_model { e_model type = MODEL_UNKNOWN; llm_arch arch = LLM_ARCH_UNKNOWN; @@ -457,6 +459,8 @@ struct llama_model { return tensor_overrides; } + size_t 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; + void set_tensor_overrides(const llama_model_params& params); int device_count() const; diff --git a/src/llama.cpp b/src/llama.cpp index e9ce56ef..c6c8c89c 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1997,9 +1997,22 @@ static bool is_model_split_supported(const llama_model & model) { return it != k_supported.end(); } -static std::vector get_layer_sizes(const llama_model_loader & ml, const llama_model & model) { + +static std::pair, double> get_layer_sizes(const llama_model_loader & ml, const llama_model & model, + ggml_type cache_type_k, ggml_type cache_type_v, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, int amb, bool flash_attn) { int n_layer = model.hparams.n_layer; std::vector result(n_layer+1, 0); + std::vector compute(n_layer+1, 0); + struct mla_tensors { + ggml_tensor * wk_b = nullptr; + ggml_tensor * wv_b = nullptr; + ggml_tensor * wkv_b = nullptr; + }; + std::vector mla_tensors; + bool has_mla = model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4; + if (has_mla) { + mla_tensors.resize(n_layer); + } size_t ow_size = 0; size_t embd_size = 0; for (int i = 0; i < ml.n_tensors; ++i) { @@ -2036,13 +2049,70 @@ static std::vector get_layer_sizes(const llama_model_loader & ml, const continue; } result[il] += size; + if (has_mla) { + if (name.find("attn_k_b.weight") != std::string::npos) { + mla_tensors[il].wk_b = t; + } + else if (name.find("attn_v_b.weight") != std::string::npos) { + mla_tensors[il].wv_b = t; + } + else if (name.find("attn_kv_b.weight") != std::string::npos) { + mla_tensors[il].wkv_b = t; + } + } + } + if (has_mla) { + for (int il = 0; il < n_layer; ++il) { + auto & mla = mla_tensors[il]; + if (mla.wk_b && mla.wv_b && !mla.wkv_b) { + auto type = ggml_is_quantized(mla.wk_b->type) ? GGML_TYPE_Q8_0 : mla.wk_b->type; + auto wkv_b_size = ggml_row_size(type, mla.wv_b->ne[0]) * (mla.wv_b->ne[1] + mla.wk_b->ne[1]) * mla.wv_b->ne[2] * 2; + result[il] += wkv_b_size; + } + else if (mla.wkv_b) { + if (!mla.wk_b) result[il] += ggml_nbytes(mla.wkv_b)/2; + if (!mla.wv_b) result[il] += ggml_nbytes(mla.wkv_b)/2; + } + if (mla_attn == 3 && mla.wv_b) { + auto kv_f32_size = (2 * mla.wv_b->ne[1] * mla.wv_b->ne[2] * max_ctx_size * sizeof(float))/(1024.*1024.); + //printf("wv_b: %ld x %ld x %ld -> %g\n", mla.wv_b->ne[0], mla.wv_b->ne[1], mla.wv_b->ne[2], kv_f32_size); + int n_head = mla.wv_b->ne[2]; + int n_max_head = n_head; + if (amb > 0 && kv_f32_size > amb) { + n_max_head = 1; + for (int niter = 2; niter < n_head; ++niter) { + if (n_head % niter == 0 && kv_f32_size/niter <= amb) { + n_max_head = n_head/niter; + break; + } + } + } + kv_f32_size = 2. * mla.wv_b->ne[1] * n_max_head * max_ctx_size * sizeof(float); + compute[il] = std::max(compute[il], kv_f32_size); + } + } } if (!ow_size) ow_size = embd_size; result[n_layer] = ow_size; LLAMA_LOG_INFO("------------------- Layer sizes:\n"); - for (int il = 0; il < n_layer; ++il) LLAMA_LOG_INFO("Layer %2d: %g MiB\n", il, result[il]/1024./1024.); - LLAMA_LOG_INFO("Layer %2d: %g MiB (output layer)\n", n_layer, result[n_layer]/1024./1024.); - return result; + double tot_model = 0, tot_cache = 0, max_compute = 0; + for (int il = 0; il < n_layer; ++il) { + auto kv_size = model.cache_size(il, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, flash_attn); + LLAMA_LOG_INFO("Layer %2d: %9.2f, %9.2f, %9.2f %9.2f MiB\n", il, result[il]/1024./1024., kv_size/1024./1024., (result[il] + kv_size)/1024./1024., compute[il]/1024./1024.); + max_compute = std::max(max_compute, compute[il]); + tot_model += result[il]; + tot_cache += kv_size; + result[il] += kv_size; + } + size_t output_size = model.hparams.n_vocab * n_ubatch * sizeof(float); + if (output_size < max_compute) output_size = max_compute; + output_size -= max_compute; + LLAMA_LOG_INFO("Layer %2d: %9.2f, %9.2f, %9.2f MiB (output layer)\n", n_layer, result[n_layer]/1024./1024., output_size/1024./1024., (result[n_layer] + output_size)/1024./1024.); + result[n_layer] += output_size; + tot_cache += output_size; + LLAMA_LOG_INFO("--------------------------------------------------------------------------\n"); + LLAMA_LOG_INFO("Total : %9.2f, %9.2f, %9.2f MiB\n", tot_model/1024./1024., tot_cache/1024./1024., (tot_model + tot_cache)/1024./1024.); + return std::make_pair(std::move(result), max_compute); } // Returns false if cancelled by progress_callback @@ -2055,6 +2125,13 @@ static bool llm_load_tensors( int main_gpu, int max_gpu, const float * tensor_split, + ggml_type cache_type_k, + ggml_type cache_type_v, + uint32_t max_ctx_size, + int n_seq_max, + int n_ubatch, + int amb, + bool flash_attn, bool use_mlock, bool validate_quants, bool mtp, @@ -2088,6 +2165,9 @@ static bool llm_load_tensors( } else { LLAMA_LOG_INFO("======================================= HAVE_FANCY_SIMD is NOT defined\n"); } + if (max_ctx_size == 0) { + max_ctx_size = model.hparams.n_ctx_train; + } model.split_mode = split_mode; model.main_gpu = main_gpu; @@ -2139,14 +2219,15 @@ static bool llm_load_tensors( model.default_layer_device = std::vector(hparams.n_layer+1, device_count-1); int act_gpu_layers = std::min(n_gpu_layers, (int)n_layer + 1); if (device_count > 1) { - auto layer_sizes = get_layer_sizes(ml, model); + auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, n_ubatch, amb, flash_attn); int n_last = n_layer; if (n_gpu_layers > n_layer) ++n_last; - double sum = 0; + double sum = max_compute * device_count; for (int i = i_gpu_start; i < n_last; ++i) sum += layer_sizes[i]; int last = i_gpu_start; float loaded_sum = 0; - for (int id = 0; id < int(model.splits.size()); ++id) { + for (int id = 0; id < device_count; ++id) { + loaded_sum += max_compute; float split_size = model.splits[id]*sum; int il = last; for (; il < n_last; ++il) { @@ -2513,7 +2594,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam #endif if (!llm_load_tensors( - ml, model, params.n_gpu_layers, params.mla, params.split_mode, params.main_gpu, params.max_gpu, params.tensor_split, + ml, model, params.n_gpu_layers, params.mla, params.split_mode, params.main_gpu, params.max_gpu, params.tensor_split, + params.type_k, params.type_v, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.flash_attn, params.use_mlock, params.validate_quants, params.mtp, params.dry_run, params.progress_callback, params.progress_callback_user_data )) { @@ -4477,6 +4559,12 @@ struct llama_model_params llama_model_default_params() { /*.main_gpu =*/ 0, /*.max_gpu =*/ 0, /*.ncmoe =*/ 0, + /*.type_k =*/ GGML_TYPE_F16, + /*.type_v =*/ GGML_TYPE_F16, + /*.max_ctx_size =*/ 0, + /*.n_seq_max =*/ 1, + /*.n_ubatch =*/ 512, + /*.amb =*/ 0, /*.tensor_split =*/ nullptr, /*.rpc_servers =*/ nullptr, /*.progress_callback =*/ nullptr, @@ -4494,6 +4582,7 @@ struct llama_model_params llama_model_default_params() { /*.merge_up_gate_exps =*/ false, /*.mtp =*/ false, /*.dry_run =*/ false, + /*.flash_attn =*/ true, }; #ifdef GGML_USE_METAL