Felipe97/llama-cpp-compiled
01.1k
1#include "llama-context.h"2 3#include "ggml.h"4#include "llama-arch.h"5#include "llama-graph.h"6#include "llama-impl.h"7#include "llama-batch.h"8#include "llama-io.h"9#include "llama-memory.h"10#include "llama-mmap.h"11#include "llama-model.h"12#include "llama-ext.h"13#include "llama-sampler.h"14#include "llama.h"15 16#include <cinttypes>17#include <cmath>18#include <cstring>19#include <limits>20#include <stdexcept>21#include <string>22 23//24// llama_context25//26 27static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) {28 switch (ctx_type) {29 case LLAMA_CONTEXT_TYPE_DEFAULT: return LLM_GRAPH_TYPE_DEFAULT;30 case LLAMA_CONTEXT_TYPE_MTP : return LLM_GRAPH_TYPE_DECODER_MTP;31 }32 throw std::runtime_error("Unsupported ctx type");33}34 35struct llm_fused_op_probe {36 llm_fused_op op;37 const char * name;38 uint32_t n_tokens_per_seq;39};40 41static const llm_fused_op_probe llm_fused_op_flash_attn_probe = {42 /*.op =*/ LLM_FUSED_OP_FLASH_ATTN,43 /*.name =*/ "Flash Attention",44 /*.n_tokens_per_seq =*/ 1,45};46 47static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = {48 /*.op =*/ LLM_FUSED_OP_GDN_AR,49 /*.name =*/ "fused Gated Delta Net (autoregressive)",50 /*.n_tokens_per_seq =*/ 1,51};52 53static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = {54 /*.op =*/ LLM_FUSED_OP_GDN_CH,55 /*.name =*/ "fused Gated Delta Net (chunked)",56 /*.n_tokens_per_seq =*/ 16,57};58 59static const llm_fused_op_probe llm_fused_op_lid_probe = {60 /*.op =*/ LLM_FUSED_OP_LIGHTNING_INDEXER,61 /*.name =*/ "Lightning Indexer",62 /*.n_tokens_per_seq =*/ 1,63};64 65static const llm_fused_op_probe llm_fused_op_dsv4_hc_pre_probe = {66 /*.op =*/ LLM_FUSED_OP_DSV4_HC_PRE,67 /*.name =*/ "fused DeepSeek V4 HC pre",68 /*.n_tokens_per_seq =*/ 1,69};70 71static const llm_fused_op_probe llm_fused_op_dsv4_hc_comb_probe = {72 /*.op =*/ LLM_FUSED_OP_DSV4_HC_COMB,73 /*.name =*/ "fused DeepSeek V4 HC comb",74 /*.n_tokens_per_seq =*/ 1,75};76 77static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = {78 /*.op =*/ LLM_FUSED_OP_DSV4_HC_POST,79 /*.name =*/ "fused DeepSeek V4 HC post",80 /*.n_tokens_per_seq =*/ 1,81};82 83llama_context::llama_context(84 const llama_model & model,85 llama_context_params params) :86 model(model),87 cvec(std::make_unique<llama_adapter_cvec>()),88 loras(std::make_unique<llama_adapter_loras>()),89 balloc(std::make_unique<llama_batch_allocr>(model.hparams.n_pos_per_embd())) {90 // TODO warning when creating llama_context with awkward ctx size that is not a power of 2,91 // may need to be backend-dependent92 LLAMA_LOG_INFO("%s: constructing llama_context\n", __func__);93 94 t_start_us = model.t_start_us;95 t_load_us = model.t_load_us;96 97 const auto & hparams = model.hparams;98 99 cparams.n_seq_max = std::max(1u, params.n_seq_max);100 if (cparams.n_seq_max > LLAMA_MAX_SEQ) {101 throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ));102 }103 104 cparams.n_rs_seq = params.n_rs_seq;105 if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {106 LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",107 __func__, cparams.n_rs_seq);108 cparams.n_rs_seq = 0;109 }110 111 cparams.n_threads = params.n_threads;112 cparams.n_threads_batch = params.n_threads_batch;113 cparams.yarn_ext_factor = params.yarn_ext_factor >= 0.0f ? params.yarn_ext_factor : hparams.yarn_ext_factor;114 cparams.yarn_attn_factor = params.yarn_attn_factor >= 0.0f ? params.yarn_attn_factor : hparams.yarn_attn_factor;115 cparams.yarn_beta_fast = params.yarn_beta_fast >= 0.0f ? params.yarn_beta_fast : hparams.yarn_beta_fast;116 cparams.yarn_beta_slow = params.yarn_beta_slow >= 0.0f ? params.yarn_beta_slow : hparams.yarn_beta_slow;117 cparams.embeddings = params.embeddings;118 cparams.embeddings_nextn = false;119 cparams.embeddings_nextn_masked = false;120 cparams.offload_kqv = params.offload_kqv;121 cparams.no_perf = params.no_perf;122 cparams.warmup = false;123 124 // +1: id n_layer() taps the output of the last layer ("input" of the head)125 cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false);126 embd_layer_inp.resize(hparams.n_layer() + 1);127 128 cparams.ctx_type = params.ctx_type;129 cparams.rope_scaling_type = params.rope_scaling_type;130 cparams.pooling_type = params.pooling_type;131 132 cparams.n_ctx = params.n_ctx == 0 ? hparams.n_ctx_train : params.n_ctx;133 cparams.rope_freq_base = params.rope_freq_base == 0.0f ? hparams.rope_freq_base_train : params.rope_freq_base;134 cparams.rope_freq_scale = params.rope_freq_scale == 0.0f ? hparams.rope_freq_scale_train : params.rope_freq_scale;135 136 cparams.n_ctx_orig_yarn = params.yarn_orig_ctx != 0 ? params.yarn_orig_ctx :137 hparams.n_ctx_orig_yarn != 0 ? hparams.n_ctx_orig_yarn :138 hparams.n_ctx_train;139 140 cparams.cb_eval = params.cb_eval;141 cparams.cb_eval_user_data = params.cb_eval_user_data;142 143 cparams.ctx_other = nullptr;144 145 // TODO: more generic146 if (model.arch == LLM_ARCH_GEMMA4_ASSISTANT) {147 if (params.ctx_other == nullptr) {148 // TODO: change from runtime_error to llama_exception to avoid printing error message149 throw std::runtime_error("Gemma4Assistant requires ctx_other to be set (this warning is normal during memory fitting)");150 }151 152 cparams.ctx_other = params.ctx_other;153 }154 155 if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) {156 if (model.tok_embd == nullptr || model.output == nullptr) {157 if (params.ctx_other == nullptr) {158 throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)");159 }160 cparams.ctx_other = params.ctx_other;161 }162 }163 164 if (cparams.rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) {165 cparams.rope_scaling_type = hparams.rope_scaling_type_train;166 }167 168 if (cparams.rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_NONE) {169 cparams.rope_freq_scale = 1.0f; // never scale if scaling type is none170 }171 172 if (cparams.yarn_ext_factor < 0.0f) { // negative indicates 'not set'173 cparams.yarn_ext_factor = cparams.rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_YARN ? 1.0f : 0.0f;174 }175 176 if (cparams.yarn_ext_factor != 0) {177 static auto get_mscale = [](float scale, float mscale) {178 return scale <= 1.0f ? 1.0f : (0.1f * mscale * logf(scale) + 1.0f);179 };180 181 const float factor = 1.0f / cparams.rope_freq_scale;182 183 // ref: https://github.com/huggingface/transformers/blob/6d00f6b0a5679c36510f203e4226e36f517c3032/src/transformers/modeling_rope_utils.py#L336-L348184 if (hparams.rope_yarn_log_mul != 0.0f) {185 // note: here we assume `mscale == 1.0f`186 // TODO: start reading the actual value of mscale and handle the case where it is not 1.0f187 float mscale = 1.0f;188 const float mscale_all_dims = hparams.rope_yarn_log_mul;189 190 // [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]191 // special-case DEEPSEEK v2:192 // https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite-Chat/blob/main/config.json#L42-L43193 if (model.arch == LLM_ARCH_DEEPSEEK2 && mscale_all_dims != 1.0f) {194 mscale = mscale_all_dims;195 }196 197 cparams.yarn_attn_factor = get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dims);198 199 LLAMA_LOG_WARN("%s: setting new yarn_attn_factor = %.4f (mscale == %.1f, mscale_all_dim = %.1f)\n",200 __func__, cparams.yarn_attn_factor, mscale, mscale_all_dims);201 } else {202 cparams.yarn_attn_factor = get_mscale(factor, 1.0f);203 }204 205 // when YARN is applied with yarn_ext_factor != 0.0f, we need to cancel this factor:206 // https://github.com/ggml-org/llama.cpp/blob/a81a569577cc38b32558958b048228150be63eae/ggml/src/ggml-cpu/ops.cpp#L5541-L5544207 //208 // ref: https://github.com/ggml-org/llama.cpp/discussions/7416209 // https://github.com/ggml-org/llama.cpp/pull/17945210 cparams.yarn_attn_factor *= 1.0f / (1.0f + 0.1f * logf(factor));211 }212 213 cparams.yarn_attn_factor *= hparams.rope_attn_factor;214 215 if (cparams.pooling_type == LLAMA_POOLING_TYPE_UNSPECIFIED) {216 if (hparams.pooling_type == LLAMA_POOLING_TYPE_UNSPECIFIED) {217 cparams.pooling_type = LLAMA_POOLING_TYPE_NONE;218 } else {219 cparams.pooling_type = hparams.pooling_type;220 }221 }222 223 if (params.attention_type == LLAMA_ATTENTION_TYPE_UNSPECIFIED) {224 cparams.causal_attn = hparams.causal_attn;225 } else {226 cparams.causal_attn = params.attention_type == LLAMA_ATTENTION_TYPE_CAUSAL;227 }228 229 cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED;230 cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO;231 232 cparams.fused_gdn_ar = true;233 cparams.fused_gdn_ch = true;234 cparams.auto_fgdn = false;235 236 cparams.fused_lid = true;237 cparams.auto_flid = false;238 239 cparams.fused_dsv4_hc_pre = true;240 cparams.fused_dsv4_hc_comb = true;241 cparams.fused_dsv4_hc_post = true;242 cparams.auto_fhc = true;243 244 // with causal attention, the batch size is limited by the context size245 cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;246 247 cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch);248 249 cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max;250 cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ?251 cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max);252 253 // Initialize backend samplers here so they are part of the sampling graph254 // before the reserve passes run later in this function. This avoids a later255 // re-reserve when graph nodes change.256 if (params.samplers != nullptr && params.n_samplers > 0) {257 for (size_t i = 0; i < params.n_samplers; ++i) {258 const auto & config = params.samplers[i];259 260 if (llama_sampler_chain_get(config.sampler, -1) == nullptr) {261 throw std::runtime_error("the backend samplers must be of type llama_sampler_chain");262 }263 264 if (set_sampler(config.seq_id, config.sampler)) {265 const int n_samplers = llama_sampler_chain_n(config.sampler);266 267 LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers);268 }269 }270 }271 272 cparams.op_offload = params.op_offload;273 cparams.kv_unified = params.kv_unified;274 275 // initialized later276 cparams.pipeline_parallel = false;277 278 {279 const char * LLAMA_GRAPH_REUSE_DISABLE = getenv("LLAMA_GRAPH_REUSE_DISABLE");280 graph_reuse_disable = LLAMA_GRAPH_REUSE_DISABLE ? (atoi(LLAMA_GRAPH_REUSE_DISABLE) != 0) : graph_reuse_disable;281 282 if (graph_reuse_disable) {283 LLAMA_LOG_WARN("%s: graph reuse disabled\n", __func__);284 }285 }286 287 // ref: https://github.com/ggml-org/llama.cpp/pull/17046#discussion_r2503085732288 cparams.n_ctx = GGML_PAD(cparams.n_ctx, 256);289 290 if (cparams.kv_unified) {291 cparams.n_ctx_seq = cparams.n_ctx;292 } else {293 cparams.n_ctx_seq = cparams.n_ctx / cparams.n_seq_max;294 cparams.n_ctx_seq = GGML_PAD(cparams.n_ctx_seq, 256);295 296 if (cparams.n_ctx_seq == 0) {297 throw std::runtime_error("n_ctx_seq == 0");298 }299 300 if (cparams.n_ctx != cparams.n_ctx_seq * cparams.n_seq_max) {301 cparams.n_ctx = cparams.n_ctx_seq * cparams.n_seq_max;302 LLAMA_LOG_WARN("%s: n_ctx is not divisible by n_seq_max - rounding down to %u\n", __func__, cparams.n_ctx);303 }304 }305 306 LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max);307 LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx);308 LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq);309 LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch);310 LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch);311 LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn);312 LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type));313 LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false");314 LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base);315 LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale);316 LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq);317 LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max);318 LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n", __func__, cparams.n_outputs_max_per_seq);319 320 if (cparams.n_ctx_seq < hparams.n_ctx_train) {321 LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n",322 __func__, cparams.n_ctx_seq, hparams.n_ctx_train);323 }324 325 if (cparams.n_ctx_seq > hparams.n_ctx_train) {326 LLAMA_LOG_WARN("%s: n_ctx_seq (%u) > n_ctx_train (%u) -- possible training context overflow\n",327 __func__, cparams.n_ctx_seq, hparams.n_ctx_train);328 }329 330 if (!hparams.vocab_only) {331 // GPU backends332 for (const auto & dev : model.devices) {333 ggml_backend_t backend = ggml_backend_dev_init(dev.dev, nullptr);334 if (backend == nullptr) {335 throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev.dev)));336 }337 backends.emplace_back(backend);338 }339 340 // add ACCEL backends (such as BLAS)341 for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {342 ggml_backend_dev_t dev = ggml_backend_dev_get(i);343 if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_ACCEL) {344 ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);345 if (backend == nullptr) {346 throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev)));347 }348 backends.emplace_back(backend);349 }350 }351 352 // add CPU backend353 backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);354 if (backend_cpu == nullptr) {355 throw std::runtime_error("failed to initialize CPU backend");356 }357 backends.emplace_back(backend_cpu);358 359 // create a list of the set_n_threads functions in the backends360 for (auto & backend : backends) {361 ggml_backend_dev_t dev = ggml_backend_get_device(backend.get());362 ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr;363 if (reg) {364 auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");365 if (ggml_backend_set_n_threads_fn) {366 set_n_threads_fns.emplace_back(backend.get(), ggml_backend_set_n_threads_fn);367 }368 }369 }370 371 llama_set_abort_callback(this, params.abort_callback, params.abort_callback_data);372 373 // graph outputs buffer374 {375 if (output_reserve(params.n_seq_max) < params.n_seq_max) {376 throw std::runtime_error("failed to reserve initial output buffer");377 }378 379 LLAMA_LOG_INFO("%s: %10s output buffer size = %8.2f MiB\n", __func__,380 ggml_backend_buffer_name (buf_output.get()),381 ggml_backend_buffer_get_size(buf_output.get()) / 1024.0 / 1024.0);382 }383 }384 385 // init the memory module386 if (!hparams.vocab_only) {387 llama_memory_params params_mem = {388 /*.type_k =*/ params.type_k,389 /*.type_v =*/ params.type_v,390 /*.swa_full =*/ params.swa_full,391 /*.ctx_type =*/ cparams.ctx_type,392 /*.mem_other =*/ llama_get_memory(cparams.ctx_other),393 };394 395 memory.reset(model.create_memory(params_mem, cparams));396 }397 398 // init backends399 if (!hparams.vocab_only) {400 LLAMA_LOG_DEBUG("%s: enumerating backends\n", __func__);401 402 backend_buft.clear();403 backend_ptrs.clear();404 backend_buf_exp_size.clear();405 406 for (auto & backend : backends) {407 auto * buft = ggml_backend_get_default_buffer_type(backend.get());408 auto backend_type = ggml_backend_dev_type(ggml_backend_get_device(backend.get()));409 410 if (backend_type == GGML_BACKEND_DEVICE_TYPE_CPU && !model.devices.empty()) {411 // use the host buffer of the first device CPU for faster transfer of the intermediate state412 const auto & dev = model.devices[0];413 auto * host_buft = ggml_backend_dev_host_buffer_type(dev.dev);414 if (host_buft) {415 buft = host_buft;416 }417 }418 419 backend_buft.push_back(buft);420 backend_ptrs.push_back(backend.get());421 backend_buf_exp_size.push_back(0);422 }423 424 LLAMA_LOG_DEBUG("%s: backend_ptrs.size() = %zu\n", __func__, backend_ptrs.size());425 426 // TODO: move these checks to ggml_backend_sched427 // enabling pipeline parallelism in the scheduler increases memory usage, so it is only done when necessary428 bool pipeline_parallel =429 model.n_devices() > 1 &&430 model.n_gpu_layers() > model.hparams.n_layer_all &&431 model.split_mode() == LLAMA_SPLIT_MODE_LAYER &&432 cparams.offload_kqv &&433 !model.has_tensor_overrides();434 435 // pipeline parallelism requires support for async compute and events in all devices436 if (pipeline_parallel) {437 for (auto & backend : backends) {438 auto dev_type = ggml_backend_dev_type(ggml_backend_get_device(backend.get()));439 if (dev_type == GGML_BACKEND_DEVICE_TYPE_CPU) {440 // ignore CPU backend441 // TODO: should we ignore ACCEL types too?442 continue;443 }444 auto * dev = ggml_backend_get_device(backend.get());445 ggml_backend_dev_props props;446 ggml_backend_dev_get_props(dev, &props);447 if (!props.caps.async || !props.caps.events) {448 // device does not support async compute or events449 pipeline_parallel = false;450 break;451 }452 }453 }454 455 cparams.pipeline_parallel = pipeline_parallel;456 457 if (cparams.pipeline_parallel) {458 LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__);459 }460 461 sched_reserve();462 463 if (!cparams.flash_attn) {464 if (ggml_is_quantized(params.type_v)) {465 throw std::runtime_error("quantized V cache was requested, but this requires Flash Attention");466 }467 }468 }469 470 // Initialize the full vocabulary token ids for backend samplers.471 {472 const int n_vocab = model.vocab.n_tokens();473 474 sampling.token_ids_full_vocab.resize(n_vocab);475 for (int i = 0; i < n_vocab; ++i) {476 sampling.token_ids_full_vocab[i] = i;477 }478 }479}480 481llama_context::~llama_context() {482 // wait for any pending asynchronous copies into the output buffers before they are freed483 synchronize();484 485 // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation486 if (!model.hparams.no_alloc && !opt_ctx) {487 for (size_t i = 0; i < backend_ptrs.size(); ++i) {488 ggml_backend_t backend = backend_ptrs[i];489 ggml_backend_buffer_type_t buft = backend_buft[i];490 491 const size_t size_exp = backend_buf_exp_size[i];492 const size_t size_act = ggml_backend_sched_get_buffer_size(sched.get(), backend);493 if (size_exp == size_act) {494 LLAMA_LOG_DEBUG("%s: %10s compute buffer size is %8.4f MiB, matches expectation of %8.4f MiB\n",495 __func__, ggml_backend_buft_name(buft), size_act / (1024.0*1024.0), size_exp / (1024.0*1024.0));496 } else {497 LLAMA_LOG_WARN("%s: %10s compute buffer size of %8.4f MiB, does not match expectation of %8.4f MiB\n",498 __func__, ggml_backend_buft_name(buft), size_act / (1024.0*1024.0), size_exp / (1024.0*1024.0));499 }500 }501 }502 ggml_opt_free(opt_ctx);503}504 505void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) {506 const char * func = __func__;507 auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) {508 if (!enabled) {509 return;510 }511 512 const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs;513 514 auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true);515 if (!gf) {516 throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check");517 }518 519 bool device_mismatch = false;520 for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) {521 if (node.op != probe.op) {522 continue;523 }524 525 GGML_ASSERT(node.il >= 0);526 527 ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor);528 ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr;529 530 // TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior,531 // but is still wrong for cases like --no-kv-offload.532 ggml_backend_dev_t device_layer = model.dev_layer(node.il);533 534 if (device_fused != device_layer) {535 LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s "536 "is assigned to device %s (usually due to missing support)\n",537 func, node.il,538 device_layer ? ggml_backend_dev_name(device_layer) : "none",539 probe.name,540 device_fused ? ggml_backend_dev_name(device_fused) : "none");541 device_mismatch = true;542 break;543 }544 }545 546 if (device_mismatch) {547 enabled = false;548 LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name);549 } else {550 enabled = true;551 LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name);552 }553 };554 555 if (cparams.auto_fa) {556 resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn);557 cparams.auto_fa = false;558 }559 560 if (cparams.auto_fgdn) {561 LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func);562 resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar);563 resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);564 cparams.auto_fgdn = false;565 }566 567 if (cparams.auto_flid) {568 LLAMA_LOG_INFO("%s: resolving fused Lightning Indexer support:\n", func);569 resolve(llm_fused_op_lid_probe, cparams.fused_lid);570 cparams.auto_flid = false;571 }572 573 if (cparams.auto_fhc) {574 LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 HC support:\n", func);575 resolve(llm_fused_op_dsv4_hc_pre_probe, cparams.fused_dsv4_hc_pre);576 resolve(llm_fused_op_dsv4_hc_comb_probe, cparams.fused_dsv4_hc_comb);577 resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post);578 cparams.auto_fhc = false;579 }580}581 582void llama_context::sched_reserve() {583 if (!sched_need_reserve) {584 return;585 }586 587 sched_need_reserve = false;588 589 LLAMA_LOG_INFO("%s: reserving ...\n", __func__);590 591 synchronize();592 593 const int64_t t_start_us = ggml_time_us();594 595 const uint32_t n_seqs = cparams.n_seq_max;596 const uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch);597 598 const size_t max_nodes = this->graph_max_nodes(n_tokens);599 600 LLAMA_LOG_DEBUG("%s: max_nodes = %zu\n", __func__, max_nodes);601 602 for (auto & res : gf_res_prev) {603 res.reset();604 }605 gf_res_reserve.reset(new llm_graph_result(max_nodes));606 gf_res_prev_active = nullptr;607 608 sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));609 610 llama_memory_context_ptr mctx;611 if (memory) {612 LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__);613 mctx = memory->init_full();614 if (!mctx) {615 throw std::runtime_error("failed to initialize memory module");616 }617 }618 619 // avoid reserving graphs with zero outputs - assume one output per sequence620 const int n_outputs = n_seqs;621 622 LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs);623 624 resolve_fused_ops(mctx.get(), n_seqs);625 626 // reserve worst-case graph627 int n_splits_pp = -1;628 int n_nodes_pp = -1;629 630 int n_splits_tg = -1;631 int n_nodes_tg = -1;632 633 const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max);634 635 // reserve pp (prompt processing) graph first so that buffers are only allocated once636 {637 auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(),638 model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr);639 if (!gf) {640 if (cparams.pipeline_parallel) {641 LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);642 cparams.pipeline_parallel = false;643 sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload));644 gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());645 }646 if (!gf) {647 throw std::runtime_error("failed to allocate compute pp buffers");648 }649 }650 651 n_splits_pp = ggml_backend_sched_get_n_splits(sched.get());652 n_nodes_pp = ggml_graph_n_nodes(gf);653 }654 655 // reserve with tg (token generation) graph to get the number of splits and nodes656 {657 auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.hparams.no_alloc);658 if (!gf) {659 throw std::runtime_error("failed to allocate compute tg buffers");660 }661 662 n_splits_tg = ggml_backend_sched_get_n_splits(sched.get());663 n_nodes_tg = ggml_graph_n_nodes(gf);664 }665 666 // reserve again with pp graph to avoid ggml-alloc reallocations during inference667 {668 // TODO: the worst case graph is not always reached for `n_seqs > 1`669 // need to implement a more robust mechanism that tries a few different inputs and analyzes the results670 ggml_cgraph * gf = nullptr;671 switch (model.arch) {672 case LLM_ARCH_KIMI_LINEAR:673 case LLM_ARCH_MINIMAX_01:674 // [TAG_RESERVE_DIAG_DECAY]675 // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which676 // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1`677 gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.hparams.no_alloc);678 break;679 default:680 gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc);681 };682 683 if (!gf) {684 throw std::runtime_error("failed to allocate compute pp buffers");685 }686 }687 688 for (size_t i = 0; i < backend_ptrs.size(); ++i) {689 ggml_backend_t backend = backend_ptrs[i];690 ggml_backend_buffer_type_t buft = backend_buft[i];691 if (!model.hparams.no_alloc) {692 backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend);693 }694 if (backend_buf_exp_size[i] > 1) {695 LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__,696 ggml_backend_buft_name(buft),697 backend_buf_exp_size[i] / 1024.0 / 1024.0);698 }699 }700 701 if (n_nodes_pp == n_nodes_tg) {702 LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp);703 } else {704 LLAMA_LOG_INFO("%s: graph nodes = %d (with bs=%d), %d (with bs=1)\n", __func__, n_nodes_pp, n_tokens, n_nodes_tg);705 }706 707 if (n_splits_pp == n_splits_tg) {708 LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits_pp);709 } else {710 LLAMA_LOG_INFO("%s: graph splits = %d (with bs=%d), %d (with bs=1)\n", __func__, n_splits_pp, n_tokens, n_splits_tg);711 }712 713 const int64_t t_end_us = ggml_time_us();714 715 LLAMA_LOG_INFO("%s: reserve took %.2f ms, sched copies = %d\n",716 __func__, (t_end_us - t_start_us)/1000.0, ggml_backend_sched_get_n_copies(sched.get()));717}718 719void llama_context::synchronize() {720 if (!sched) {721 return;722 }723 724 ggml_backend_sched_synchronize(sched.get());725 726 // FIXME: if multiple single tokens are evaluated without a synchronization,727 // the stats will be added to the prompt evaluation stats728 // this should only happen when using batch size 1 to evaluate a batch729 730 // add the evaluation to the stats731 if (n_queued_tokens == 1) {732 if (!cparams.no_perf) {733 t_eval_us += ggml_time_us() - t_compute_start_us;734 }735 n_eval++;736 } else if (n_queued_tokens > 1) {737 if (!cparams.no_perf) {738 t_p_eval_us += ggml_time_us() - t_compute_start_us;739 }740 n_p_eval += n_queued_tokens;741 }742 743 // get a more accurate load time, upon first eval744 if (n_queued_tokens > 0 && !has_evaluated_once) {745 t_load_us = ggml_time_us() - t_start_us;746 has_evaluated_once = true;747 }748 749 n_queued_tokens = 0;750 t_compute_start_us = 0;751}752 753const llama_model & llama_context::get_model() const {754 return model;755}756 757const llama_cparams & llama_context::get_cparams() const {758 return cparams;759}760 761ggml_backend_sched_t llama_context::get_sched() const {762 return sched.get();763}764 765uint32_t llama_context::n_ctx() const {766 return cparams.n_ctx;767}768 769uint32_t llama_context::n_ctx_seq() const {770 return cparams.n_ctx_seq;771}772 773uint32_t llama_context::n_batch() const {774 return cparams.n_batch;775}776 777uint32_t llama_context::n_ubatch() const {778 return cparams.n_ubatch;779}780 781uint32_t llama_context::n_seq_max() const {782 return cparams.n_seq_max;783}784 785uint32_t llama_context::n_threads() const {786 return cparams.n_threads;787}788 789uint32_t llama_context::n_threads_batch() const {790 return cparams.n_threads_batch;791}792 793llama_memory_t llama_context::get_memory() const {794 return memory.get();795}796 797bool llama_context::memory_update(bool optimize) {798 if (!memory) {799 return false;800 }801 802 {803 const auto mctx = memory->init_update(this, optimize);804 switch (mctx->get_status()) {805 case LLAMA_MEMORY_STATUS_SUCCESS:806 {807 // noop808 } break;809 case LLAMA_MEMORY_STATUS_NO_UPDATE:810 {811 // no updates need to be performed812 return false;813 }814 case LLAMA_MEMORY_STATUS_FAILED_PREPARE:815 case LLAMA_MEMORY_STATUS_FAILED_COMPUTE:816 {817 LLAMA_LOG_ERROR("%s: failed to prepare memory update\n", __func__);818 return false;819 }820 }821 822 // reset the previous graph results to make sure that they won't be reused823 // TODO: make mctx->apply() report if a graph reserve is needed, then reset graph results only if the memory module reset the scheduler824 for (auto & res : gf_res_prev) {825 if (res) {826 res->reset();827 }828 }829 gf_res_prev_active = nullptr;830 831 if (!mctx->apply()) {832 LLAMA_LOG_ERROR("%s: failed to apply memory update\n", __func__);833 }834 }835 836 // if the memory module did any computation, we have to reserve a new worst-case graph837 {838 const auto mctx = memory->init_full();839 if (!mctx) {840 throw std::runtime_error("failed to initialize memory context");841 }842 843 const uint32_t n_seqs = cparams.n_seq_max;844 const uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch);845 846 const uint32_t n_outputs_max = std::min(n_tokens, cparams.n_outputs_max);847 848 auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_max, mctx.get());849 if (!gf) {850 LLAMA_LOG_ERROR("%s: failed to reserve graph after the memory update\n", __func__);851 }852 }853 854 return true;855}856 857enum llama_pooling_type llama_context::pooling_type() const {858 return cparams.pooling_type;859}860 861float * llama_context::get_logits() {862 output_reorder();863 864 return logits.data;865}866 867int64_t llama_context::output_resolve_row(int32_t i) const {868 int64_t j = -1;869 870 // support negative indices (last output row)871 if (i < 0) {872 j = n_outputs + i;873 if (j < 0) {874 throw std::runtime_error(format("negative index out of range [0, %d)", n_outputs));875 }876 } else if ((size_t) i >= output_ids.size()) {877 throw std::runtime_error(format("out of range [0, %zu)", output_ids.size()));878 } else {879 // use output_ids to translate the batch token index into a row number880 // that holds this token's data.881 j = output_ids[i];882 }883 884 if (j < 0) {885 // the batch token was not configured to output anything886 throw std::runtime_error(format("batch.logits[%d] != true", i));887 }888 889 if (j >= n_outputs) {890 throw std::runtime_error(format("corrupt output buffer (j=%" PRId64 ", n_outputs=%d)", j, n_outputs));891 }892 893 return j;894}895 896float * llama_context::get_logits_ith(int32_t i) {897 output_reorder();898 899 try {900 if (logits.data == nullptr) {901 throw std::runtime_error("no logits");902 }903 904 const int64_t j = output_resolve_row(i);905 return logits.data + j*model.vocab.n_tokens();906 } catch (const std::exception & err) {907 LLAMA_LOG_ERROR("%s: invalid logits id %d, reason: %s\n", __func__, i, err.what());908#ifndef NDEBUG909 GGML_ABORT("fatal error");910#else911 return nullptr;912#endif913 }914}915 916float * llama_context::get_embeddings() {917 output_reorder();918 919 return embd.data;920}921 922llama_token * llama_context::get_sampled_tokens() const{923 return sampling.sampled.data;924}925 926float * llama_context::get_embeddings_ith(int32_t i) {927 output_reorder();928 929 try {930 if (embd.data == nullptr) {931 throw std::runtime_error("no embeddings");932 }933 934 const int64_t j = output_resolve_row(i);935 const uint32_t n_embd_out = model.hparams.n_embd_out();936 return embd.data + j*n_embd_out;937 } catch (const std::exception & err) {938 LLAMA_LOG_ERROR("%s: invalid embeddings id %d, reason: %s\n", __func__, i, err.what());939#ifndef NDEBUG940 GGML_ABORT("fatal error");941#else942 return nullptr;943#endif944 }945}946 947float * llama_context::get_embeddings_seq(llama_seq_id seq_id) {948 auto it = embd_seq.find(seq_id);949 if (it == embd_seq.end()) {950 return nullptr;951 }952 953 return it->second.data();954}955 956float * llama_context::get_embeddings_nextn() {957 output_reorder();958 959 return embd_nextn.data;960}961 962float * llama_context::get_embeddings_nextn_ith(int32_t i) {963 output_reorder();964 965 try {966 if (embd_nextn.data == nullptr) {967 throw std::runtime_error("no nextn embeddings");968 }969 970 const uint32_t n_embd = model.hparams.n_embd_out();971 972 if (!cparams.embeddings_nextn_masked) {973 // unmasked: nextn rows are stored densely, indexed by raw token position.974 if (i < 0 || (size_t)(i + 1) * n_embd > embd_nextn.size) {975 throw std::runtime_error(format("out of range [0, %zu)", embd_nextn.size / n_embd));976 }977 return embd_nextn.data + (size_t) i * n_embd;978 }979 980 const int64_t j = output_resolve_row(i);981 return embd_nextn.data + j*n_embd;982 } catch (const std::exception & err) {983 LLAMA_LOG_ERROR("%s: invalid nextn embeddings id %d, reason: %s\n", __func__, i, err.what());984#ifndef NDEBUG985 GGML_ABORT("fatal error");986#else987 return nullptr;988#endif989 }990}991 992float * llama_context::get_embeddings_layer_inp(uint32_t lid) {993 output_reorder();994 995 GGML_ASSERT(lid < embd_layer_inp.size() && embd_layer_inp[lid].has_data());996 997 return embd_layer_inp[lid].data;998}999 1000llama_token llama_context::get_sampled_token_ith(int32_t idx) {1001 output_reorder();1002 1003 if (!sampling.sampled.has_data()) {1004 return LLAMA_TOKEN_NULL;1005 }1006 1007 try {1008 const int64_t row = output_resolve_row(idx);1009 GGML_ASSERT(row < (int64_t) sampling.sampled.size);1010 return sampling.sampled.data[row];1011 } catch (const std::exception & err) {1012 LLAMA_LOG_ERROR("%s: invalid backend sampled token id %d, reason: %s\n", __func__, idx, err.what());1013 return LLAMA_TOKEN_NULL;1014 }1015}1016 1017float * llama_context::get_sampled_probs_ith(int32_t idx) {1018 output_reorder();1019 1020 if (!sampling.probs.has_data()) {1021 return nullptr;1022 }1023 1024 try {1025 const int64_t row = output_resolve_row(idx);1026 if ((size_t) row >= sampling.probs_count.size() || sampling.probs_count[row] == 0) {1027 return nullptr;1028 }1029 return sampling.probs.data + row*model.vocab.n_tokens();1030 } catch (const std::exception & err) {1031 LLAMA_LOG_ERROR("%s: invalid backend sampled probs id %d, reason: %s\n", __func__, idx, err.what());1032 return nullptr;1033 }1034}1035 1036float * llama_context::get_sampled_logits_ith(int32_t idx) {1037 output_reorder();1038 1039 if (!sampling.logits.has_data()) {1040 return nullptr;1041 }1042 1043 try {1044 const int64_t row = output_resolve_row(idx);1045 if ((size_t) row >= sampling.logits_count.size() || sampling.logits_count[row] == 0) {1046 return nullptr;1047 }1048 return sampling.logits.data + row*model.vocab.n_tokens();1049 } catch (const std::exception & err) {1050 LLAMA_LOG_ERROR("%s: invalid backend sampled logits id %d, reason: %s\n", __func__, idx, err.what());1051 return nullptr;1052 }1053}1054 1055const llama_token * llama_context::get_sampled_candidates_ith(int32_t idx) {1056 output_reorder();1057 1058 try {1059 const int64_t row = output_resolve_row(idx);1060 if (sampling.candidates.has_data() &&1061 (size_t) row < sampling.candidates_count.size() &&1062 sampling.candidates_count[row] > 0) {1063 return sampling.candidates.data + row*model.vocab.n_tokens();1064 }1065 } catch (const std::exception & err) {1066 // fallback to full vocab list1067 GGML_UNUSED(err);1068 }1069 1070 return sampling.token_ids_full_vocab.data();1071}1072 1073size_t llama_context::get_sampled_candidates_count(int32_t idx) {1074 output_reorder();1075 1076 if (!sampling.candidates.has_data()) {1077 return 0;1078 }1079 1080 try {1081 const int64_t row = output_resolve_row(idx);1082 if ((size_t) row >= sampling.candidates_count.size()) {1083 return 0;1084 }1085 return sampling.candidates_count[row];1086 } catch (const std::exception & err) {1087 LLAMA_LOG_ERROR("%s: invalid backend sampled candidates count id %d, reason: %s\n", __func__, idx, err.what());1088 return 0;1089 }1090}1091 1092size_t llama_context::get_sampled_logits_count(int32_t idx) {1093 output_reorder();1094 1095 if (!sampling.logits.has_data()) {1096 return model.vocab.n_tokens();1097 }1098 1099 try {1100 const int64_t row = output_resolve_row(idx);1101 if ((size_t) row >= sampling.logits_count.size()) {1102 return 0;1103 }1104 return sampling.logits_count[row];1105 } catch (const std::exception & err) {1106 LLAMA_LOG_ERROR("%s: invalid backend sampled logits count id %d, reason: %s\n", __func__, idx, err.what());1107 return 0;1108 }1109}1110 1111size_t llama_context::get_sampled_probs_count(int32_t idx) {1112 output_reorder();1113 1114 if (!sampling.probs.has_data()) {1115 return 0;1116 }1117 1118 try {1119 const int64_t row = output_resolve_row(idx);1120 if ((size_t) row >= sampling.probs_count.size()) {1121 return 0;1122 }1123 return sampling.probs_count[row];1124 } catch (const std::exception & err) {1125 LLAMA_LOG_ERROR("%s: invalid backend sampled probs count id %d, reason: %s\n", __func__, idx, err.what());1126 return 0;1127 }1128}1129 1130 1131void llama_context::attach_threadpool(1132 ggml_threadpool_t threadpool,1133 ggml_threadpool_t threadpool_batch) {1134 LLAMA_LOG_DEBUG("%s: call\n", __func__);1135 1136 this->threadpool = threadpool;1137 this->threadpool_batch = threadpool_batch ? threadpool_batch : threadpool;1138}1139 1140void llama_context::detach_threadpool() {1141 LLAMA_LOG_DEBUG("%s: call\n", __func__);1142 1143 this->threadpool = nullptr;1144 this->threadpool_batch = nullptr;1145}1146 1147void llama_context::set_n_threads(int32_t n_threads, int32_t n_threads_batch) {1148 LLAMA_LOG_DEBUG("%s: n_threads = %d, n_threads_batch = %d\n", __func__, n_threads, n_threads_batch);1149 1150 cparams.n_threads = n_threads;1151 cparams.n_threads_batch = n_threads_batch;1152}1153 1154void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data) {1155 LLAMA_LOG_DEBUG("%s: call\n", __func__);1156 1157 this->abort_callback = abort_callback;1158 this->abort_callback_data = abort_callback_data;1159 1160 for (auto & backend : backends) {1161 auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get()));1162 if (reg) {1163 auto * set_abort_callback_fn = (ggml_backend_set_abort_callback_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_abort_callback");1164 if (set_abort_callback_fn) {1165 set_abort_callback_fn(backend.get(), this->abort_callback, this->abort_callback_data);1166 }1167 }1168 }1169}1170 1171void llama_context::set_embeddings(bool value) {1172 LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);1173 1174 cparams.embeddings = value;1175 1176 // TODO: not sure yet if we want to reserve here1177 //sched_need_reserve = true;1178}1179 1180void llama_context::set_embeddings_nextn(bool value, bool masked) {1181 LLAMA_LOG_DEBUG("%s: value = %d, masked = %d\n", __func__, value, masked);1182 1183 cparams.embeddings_nextn = value;1184 cparams.embeddings_nextn_masked = masked;1185}1186 1187void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) {1188 LLAMA_LOG_DEBUG("%s: lid = %d, enable = %d\n", __func__, lid, enable);1189 1190 GGML_ASSERT(lid <= model.hparams.n_layer());1191 1192 cparams.embeddings_layer_inp[lid] = enable;1193 1194 // note: without this reserve, the draft acceptance drops to zero. not sure why - this is unexpected1195 sched_need_reserve = true;1196}1197 1198void llama_context::set_nextn_layer_offset(int32_t offset) {1199 cparams.nextn_layer_offset = offset;1200}