LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Control compute buffer sizes for large batches (#1560)

Browse Source 
* Worst graph output tokens

* This is better
This commit is contained in:
Kawrakow
2026-04-01 07:57:52 +02:00
committed by GitHub
parent bc78a1e93b
commit 9660eb4487
7 changed files with 41 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+9
View File
@@ -1517,6 +1517,11 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
}
return true;
}
if (arg == "-wgt" || arg == "--worst-graph-tokens") {
CHECK_ARG;
params.worst_graph_tokens = std::stoi(argv[i]);
return true;
}
if (arg == "--no-mmap") {
params.use_mmap = false;
return true;
@@ -2502,6 +2507,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
options.push_back({ "*", " --cpu-moe", "keep all MoE weights in CPU memory"});
options.push_back({ "*", " --n-cpu-moe N", "keep MoE weights of the first N layers in CPU memory"});
options.push_back({ "*", " --fit-margin N", "safety margin in MiB when auto-fitting model offloading"});
options.push_back({ "*", "-wgt, --worst-graph-tokens N", "number of tokens to use for worst-case graph"});
options.push_back({ "*", " --fit", "automatically determine which tensors to offload to the GPU(s)"});
options.push_back({ "*", " --numa TYPE", "attempt optimizations that help on some NUMA systems\n"
" - distribute: spread execution evenly over all nodes\n"
@@ -3337,6 +3343,7 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params
mparams.ncmoe = params.ncmoe;
mparams.fit = params.fit;
mparams.fit_margin = params.fit_margin;
mparams.worst_graph_tokens = params.worst_graph_tokens;
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;
@@ -3404,6 +3411,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa
cparams.seed = params.seed;
cparams.logits_all = params.logits_all;
cparams.embeddings = params.embedding;
cparams.worst_case_tokens = params.worst_graph_tokens;
cparams.rope_scaling_type = params.rope_scaling_type;
cparams.rope_freq_base = params.rope_freq_base;
cparams.rope_freq_scale = params.rope_freq_scale;
@@ -4386,6 +4394,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l
fprintf(stream, "ncmoe: %d # default: 0\n", params.ncmoe);
fprintf(stream, "fit: %d # default: false\n", params.fit);
fprintf(stream, "fit_margin: %d # default: 0\n", params.fit_margin);
fprintf(stream, "worst_graph_tokens: %d # default: 0\n", params.worst_graph_tokens);
fprintf(stream, "min_keep: %d # default: 0 (disabled)\n", sparams.min_keep);
fprintf(stream, "mirostat: %d # default: 0 (disabled)\n", sparams.mirostat);
fprintf(stream, "mirostat_ent: %f # default: 5.0\n", sparams.mirostat_tau);
common/common.h
+1
View File
@@ -228,6 +228,7 @@ struct gpt_params {
int32_t ncmoe = 0; // number of layers in which MoE tensors are left in VRAM
int32_t fit_margin = 0; // safety margin for auto-fit in MiB
bool fit = false; // automatically fit model (for now just using MoE tensor overrides)
int32_t worst_graph_tokens = 0; // number of tokens to use when reserving the worst graph
float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs
int32_t grp_attn_n = 1; // group-attention factor
int32_t grp_attn_w = 512; // group-attention width
include/llama.h
+2
View File
@@ -380,6 +380,7 @@ extern "C" {
int32_t amb;
int32_t fit_margin;
bool fit;
int32_t worst_graph_tokens;
// proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()
const float * tensor_split;
@@ -426,6 +427,7 @@ extern "C" {
uint32_t n_threads; // number of threads to use for generation
uint32_t n_threads_batch; // number of threads to use for batch processing
int32_t max_extra_alloc; // Max. additional VRAM the scheduler is allowed to allocate
int32_t worst_case_tokens; // number of tokens to use when reserving worst case graphs
enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type`
enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id
src/llama-build-context.cpp
+9 -7
View File
@@ -34,7 +34,8 @@ llm_build_context::llm_build_context(
const llama_batch & batch,
const llm_build_cb & cb,
bool worst_case,
bool warmup) :
bool warmup,
int n_outputs_) :
model (lctx.model),
lctx (lctx),
hparams (model.hparams),
@@ -63,7 +64,7 @@ llm_build_context::llm_build_context(
norm_rms_eps (hparams.f_norm_rms_eps),
n_tokens (batch.n_tokens),
n_kv (worst_case ? kv_self.size : kv_self.n),
n_outputs (worst_case ? n_tokens : lctx.n_outputs),
n_outputs (worst_case ? n_outputs_ > 0 ? n_outputs_ : n_tokens : lctx.n_outputs),
n_outputs_enc (worst_case ? n_tokens : lctx.embd_enc.size() / hparams.n_embd),
kv_head (worst_case ? (kv_self.recurrent ? 0 : kv_self.size - n_tokens) : kv_self.head),
n_ctx_orig (cparams.n_ctx_orig_yarn),
@@ -308,17 +309,17 @@ struct ggml_tensor * llm_build_context::build_inp_embd_mtp(struct ggml_tensor *
if (batch.token) {
lctx.inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch.n_tokens);
cb(lctx.inp_tokens, "inp_tokens", -1);
ggml_set_input(lctx.inp_tokens);
cur = ggml_get_rows(ctx0, mtp_tok_embd, lctx.inp_tokens);
} else {
return nullptr;
return nullptr;
}
cb(cur, "inp_embd", -1);
return cur;
}
@@ -9597,7 +9598,8 @@ struct ggml_cgraph * llm_build_context::llama_build_graph_s_copy(llama_context &
ggml_cgraph * llm_build_context::llama_build_graph(
llama_context & lctx,
const llama_batch & batch,
bool worst_case) {
bool worst_case,
int n_outputs) {
const auto & model = lctx.model;
#if IK_PRINT_TIMING
@@ -9654,7 +9656,7 @@ ggml_cgraph * llm_build_context::llama_build_graph(
llama_token bos = vocab->token_bos();
llama_token eos = vocab->token_eos();
bool is_warming_up = lctx.n_eval == 0 && (batch.n_tokens == 1 && (batch.token[0] == ((bos != -1) ? bos : eos)));
struct llm_build_context llm(lctx, batch, cb, worst_case, is_warming_up);
struct llm_build_context llm(lctx, batch, cb, worst_case, is_warming_up, n_outputs);
llm.init();
src/llama-build-context.h
+5 -4
View File
@@ -102,7 +102,8 @@ struct llm_build_context {
const llama_batch & batch,
const llm_build_cb & cb,
bool worst_case,
bool warmup);
bool warmup,
int n_outputs = 0);
void init();
@@ -428,7 +429,7 @@ llm_expert_gating_func_type gating_op,
static ggml_cgraph * llama_build_graph_s_copy(llama_context & lctx);
static ggml_cgraph * llama_build_graph(llama_context & lctx, const llama_batch & batch, bool worst_case);
static ggml_cgraph * llama_build_graph(llama_context & lctx, const llama_batch & batch, bool worst_case, int n_outputs = 0);
ggml_tensor * build_std_attention(ggml_cgraph * gf, ggml_tensor * attn_norm, ggml_tensor * cur,
ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors,
@@ -439,8 +440,8 @@ llm_expert_gating_func_type gating_op,
static uint32_t llama_kv_qnext_state_slots(const llama_kv_cache & kv_self);
struct ggml_tensor * build_mtp_tail(
const struct llama_layer & mtp_layer,
struct ggml_tensor * prev_embeddings,
const struct llama_layer & mtp_layer,
struct ggml_tensor * prev_embeddings,
int64_t n_embd_head,
struct ggml_cgraph * gf,
struct ggml_tensor * inp_pos,
src/llama-cparams.h
+1
View File
@@ -47,6 +47,7 @@ struct llama_cparams {
int min_experts;
float thresh_experts;
bool mtp;
int worst_graph_tokens;
enum ggml_type reduce_type;
enum llama_pooling_type pooling_type;
src/llama.cpp
+14 -6
View File
@@ -2007,7 +2007,8 @@ struct expert_tensors {
};
static std::pair<std::vector<double>, 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,
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, int worst_case_tokens, bool flash_attn,
std::vector<expert_tensors> & experts) {
int n_layer = model.hparams.n_layer;
std::vector<double> result(n_layer+1, 0);
@@ -2181,7 +2182,8 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
tot_cache += kv_size;
result[il] += kv_size;
}
size_t output_size = model.hparams.n_vocab * n_ubatch * sizeof(float);
int n_output = worst_case_tokens > 0 ? worst_case_tokens : n_ubatch;
size_t output_size = model.hparams.n_vocab * n_output * 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.);
@@ -2209,6 +2211,7 @@ static bool llm_load_tensors(
int n_ubatch,
int amb,
int fit_margin,
int worst_case_tokens,
bool flash_attn,
bool use_mlock,
bool validate_quants,
@@ -2331,7 +2334,8 @@ static bool llm_load_tensors(
std::vector<llama_model_tensor_buft_override> overrides;
if (device_count > 0) {
std::vector<expert_tensors> experts;
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, experts);
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, worst_case_tokens, flash_attn, experts);
size_t required_mem = 0;
for (int i = 0; i <= n_layer; ++i) {
required_mem += layer_sizes[i];
@@ -2894,7 +2898,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam
if (!llm_load_tensors(
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.fit_margin, params.flash_attn,
params.type_k, params.type_v, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.fit_margin,
params.worst_graph_tokens, params.flash_attn,
params.use_mlock, params.validate_quants, params.mtp, params.fit, params.dry_run,
params.progress_callback, params.progress_callback_user_data
)) {
@@ -4651,7 +4656,7 @@ static int32_t llama_kv_cache_update_internal(struct llama_context & lctx) {
int n_tokens = (int)std::min(lctx.cparams.n_ctx, lctx.cparams.n_ubatch);
int n_past = lctx.cparams.n_ctx - n_tokens;
llama_token token = llama_token_bos(&lctx.model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph
ggml_cgraph * gf = llm_build_context::llama_build_graph(lctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true);
ggml_cgraph * gf = llm_build_context::llama_build_graph(lctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true, lctx.cparams.worst_graph_tokens);
// initialize scheduler with the worst-case graph
lctx.reset_scheduler();
@@ -4892,6 +4897,7 @@ struct llama_model_params llama_model_default_params() {
/*.amb =*/ 0,
/*.fit_margin =*/ 0,
/*.fit =*/ false,
/*.worst_graph_tokens =*/ 0,
/*.tensor_split =*/ nullptr,
/*.rpc_servers =*/ nullptr,
/*.progress_callback =*/ nullptr,
@@ -4930,6 +4936,7 @@ struct llama_context_params llama_context_default_params() {
/*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default
/*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS,
/*.max_extra_alloc =*/ 256,
/*.worst_case_tokens =*/ 0,
/*.rope_scaling_type =*/ LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED,
/*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED,
/*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED,
@@ -5350,6 +5357,7 @@ struct llama_context * llama_init_from_model(
cparams.thresh_experts = params.thresh_experts;
cparams.cuda_params = params.cuda_params;
cparams.mtp = params.mtp;
cparams.worst_graph_tokens = params.worst_case_tokens;
cparams.reduce_type = params.type_reduce;
cparams.pooling_type = params.pooling_type;
@@ -5749,7 +5757,7 @@ struct llama_context * llama_init_from_model(
// build worst-case graph
int n_past = cparams.n_ctx - n_tokens;
llama_token token = llama_token_bos(&ctx->model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph
ggml_cgraph * gf = llm_build_context::llama_build_graph(*ctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true);
ggml_cgraph * gf = llm_build_context::llama_build_graph(*ctx, llama_batch_get_one(&token, n_tokens, n_past, 0), true, cparams.worst_graph_tokens);
// initialize scheduler with the worst-case graph
bool gf_success = ggml_backend_sched_reserve(ctx->sched, gf);