diff --git a/common/arg.cpp b/common/arg.cpp index a5bfc75c7..6db1a3559 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2606,14 +2606,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_DIO")); add_opt(common_arg( {"-lm", "--load-mode"}, "MODE", - "model loading mode (default: mmap)\n" + "model loading mode (default: auto)\n" + "- auto: mmap, unless a device does not support it\n" "- none: no special loading mode\n" "- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n" "- mlock: force system to keep model in RAM rather than swapping or compressing\n" "- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n" "- dio: use DirectIO if available\n", [](common_params & params, const std::string & value) { - /**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } + /**/ if (value == "auto") { params.load_mode = LLAMA_LOAD_MODE_AUTO; } + else if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; } else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; } else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; } diff --git a/common/common.cpp b/common/common.cpp index a7aba5d5c..b6cfd16f4 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1645,6 +1645,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.n_seq_max = params.n_parallel; cparams.n_rs_seq = params.speculative.need_n_rs_seq(); cparams.n_outputs_max = std::max(params.n_outputs_max, 0); + cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0); cparams.n_batch = params.n_batch; cparams.n_ubatch = params.n_ubatch; cparams.n_threads = params.cpuparams.n_threads; diff --git a/common/common.h b/common/common.h index 2a1927dc1..73fffcf1c 100644 --- a/common/common.h +++ b/common/common.h @@ -448,6 +448,7 @@ struct common_params { int32_t n_parallel = 1; // number of parallel sequences to decode int32_t n_sequences = 1; // number of sequences to decode int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch) + int32_t n_outputs_max_per_seq = 1; // max outputs per sequence int32_t grp_attn_n = 1; // group-attention factor int32_t grp_attn_w = 512; // group-attention width int32_t n_print = -1; // print token count every n tokens (-1 = disabled) @@ -473,7 +474,7 @@ struct common_params { std::vector fit_params_target = std::vector(llama_max_devices(), 1024 * 1024*1024); enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs - enum llama_load_mode load_mode = LLAMA_LOAD_MODE_MMAP; // how to load the model + enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model common_cpu_params cpuparams; common_cpu_params cpuparams_batch; diff --git a/common/llguidance.cpp b/common/llguidance.cpp index d58f147a7..500bb0914 100644 --- a/common/llguidance.cpp +++ b/common/llguidance.cpp @@ -116,6 +116,8 @@ static llama_sampler_i llama_sampler_llg_i = { /* .backend_accept = */ NULL, /* .backend_apply = */ NULL, /* .backend_set_input = */ NULL, + /* .backend_reset = */ NULL, + /* .copy_state = */ NULL, }; static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len, diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index ef290ed7c..4a4be7cf7 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -570,23 +570,34 @@ struct parser_executor { } static common_peg_parse_result handle_escape_sequence(common_peg_parse_context & ctx, size_t start, size_t & pos, const char delimiter) { + auto save = pos; + ++pos; // consume '\' if (pos >= ctx.input.size()) { if (!ctx.is_lenient()) { return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } + pos = save; // suppress unmatched '\' return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); } char c = ctx.input[pos]; + if (c == delimiter || c == '\\' || c == '/' || c == 'b' || c == 'f' || c == 'n' || c == 'r' || c == 't') { ++pos; return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); - } else if (c == 'u') { - return handle_unicode_escape(ctx, start, pos); - } else { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } + + if (c == 'u') { + auto result = handle_unicode_escape(ctx, start, pos); + if (result.need_more_input()) { + pos = save; // suppress incomplete sequence + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); + } + return result; + } + + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } static common_peg_parse_result handle_unicode_escape(common_peg_parse_context & ctx, size_t start, size_t & pos) { diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 1fe242d06..4884299f3 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -217,6 +217,8 @@ static struct llama_sampler_i common_reasoning_budget_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) { diff --git a/common/sampling.cpp b/common/sampling.cpp index ec9c885dd..06dea1e1c 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -518,6 +518,26 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) { }; } +void common_sampler_copy(const common_sampler * src, common_sampler * dst) { + if (!src || !dst || src == dst) { + return; + } + + GGML_ASSERT((src->grmr == nullptr) == (dst->grmr == nullptr)); + GGML_ASSERT((src->rbudget == nullptr) == (dst->rbudget == nullptr)); + + llama_sampler_copy(src->grmr, dst->grmr); + llama_sampler_copy(src->rbudget, dst->rbudget); + llama_sampler_copy(src->chain, dst->chain); + + dst->params = src->params; + dst->prev = src->prev; + dst->cur = src->cur; + dst->cur_p = src->cur_p; + dst->cur_p.data = src->cur_p.data ? dst->cur.data() : nullptr; // re-point to dst's buffer + dst->t_total_us = src->t_total_us; +} + void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) { // TODO: measure grammar performance diff --git a/common/sampling.h b/common/sampling.h index cb90d4ae7..ced3c8364 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -47,6 +47,7 @@ void common_sampler_free(struct common_sampler * gsmpl); void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated); void common_sampler_reset (struct common_sampler * gsmpl); struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl); +void common_sampler_copy (const struct common_sampler * src, struct common_sampler * dst); // arguments can be nullptr to skip printing void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl); diff --git a/common/speculative.cpp b/common/speculative.cpp index 02f4c8054..66fac90a8 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2299,6 +2299,7 @@ common_params common_base_params_to_speculative(const common_params & params) { result.cache_type_k = params_spec.cache_type_k; result.cache_type_v = params_spec.cache_type_v; result.n_outputs_max = params.n_parallel; + result.n_outputs_max_per_seq = 1; return result; } @@ -2384,6 +2385,17 @@ common_speculative_init_result_ptr common_speculative_init_from_params(common_pa return std::make_unique(params, model_tgt, ctx_tgt); } +common_speculative_output_limits common_speculative_get_output_limits( + int32_t n_batch, int32_t n_parallel, int32_t n_draft) { + const int64_t per_seq = 1 + (int64_t) std::max(0, n_draft); + const int64_t total = (int64_t) n_parallel * per_seq; + + return { + /* .total = */ (int32_t) std::min(n_batch, total), + /* .per_seq = */ (int32_t) std::min(n_batch, per_seq), + }; +} + // initialization of the speculative decoding system // common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) { diff --git a/common/speculative.h b/common/speculative.h index 062bf2093..c6986affd 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -25,6 +25,15 @@ int32_t common_speculative_n_max(const common_params_speculative * spec); common_params common_base_params_to_speculative(const common_params & params); +struct common_speculative_output_limits { + int32_t total; + int32_t per_seq; +}; + +// return the output limits needed for speculative decoding +common_speculative_output_limits common_speculative_get_output_limits( + int32_t n_batch, int32_t n_parallel, int32_t n_draft); + common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq); void common_speculative_free(common_speculative * spec); diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 2924fdbe9..cc3f8cd36 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -154,6 +154,8 @@ extern "C" { bool buffer_from_host_ptr; // event synchronization bool events; + // mmap is supported for loading + bool mmap_support; }; // all the device properties diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index a5a3a58ad..7654ea1f3 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -132,6 +132,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ false, // Not implemented. /* .buffer_from_host_ptr = */ false, // Not implemented. /* .events = */ false, // Not implemented. + /* .mmap_support = */ true, }; for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { ggml_backend_dev_props tmp_props; @@ -140,6 +141,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer; props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr; props->caps.events = props->caps.events && tmp_props.caps.events; + props->caps.mmap_support = props->caps.mmap_support && tmp_props.caps.mmap_support; } } diff --git a/ggml/src/ggml-blas/ggml-blas.cpp b/ggml/src/ggml-blas/ggml-blas.cpp index 9745fa29f..e4b5bd254 100644 --- a/ggml/src/ggml-blas/ggml-blas.cpp +++ b/ggml/src/ggml-blas/ggml-blas.cpp @@ -367,6 +367,7 @@ static void ggml_backend_blas_device_get_props(ggml_backend_dev_t dev, struct gg /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 9d26a264a..50b2124c2 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -397,6 +397,7 @@ static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggm /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 0a3cb86b8..1545c12bd 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4814,6 +4814,7 @@ static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ events, + /* .mmap_support = */ props->type != GGML_BACKEND_DEVICE_TYPE_IGPU, }; } diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index a1003b3ac..ef3c92f27 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -681,6 +681,7 @@ static void ggml_backend_metal_device_get_props(ggml_backend_dev_t dev, ggml_bac /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ true, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 3a380229b..f48552f31 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1884,6 +1884,7 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 9f76c1162..d28b79d00 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17924,6 +17924,7 @@ static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml /* .host_buffer = */ true, /* .buffer_from_host_ptr = */ false, /* .events = */ false, //kcpp: triage disable events ref https://github.com/ggml-org/llama.cpp/pull/18397 + /* .mmap_support = */ !ctx->is_integrated_gpu, }; } diff --git a/include/llama.h b/include/llama.h index dfd2ef722..42f464a9a 100644 --- a/include/llama.h +++ b/include/llama.h @@ -206,11 +206,12 @@ extern "C" { }; enum llama_load_mode { - LLAMA_LOAD_MODE_NONE = 0, // no special loading mode - LLAMA_LOAD_MODE_MMAP = 1, // memory map the model - LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available + LLAMA_LOAD_MODE_AUTO = -1, // auto-detect based on device capabilities + LLAMA_LOAD_MODE_NONE = 0, // no special loading mode + LLAMA_LOAD_MODE_MMAP = 1, // memory map the model + LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available }; LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); @@ -351,14 +352,15 @@ extern "C" { // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations // https://github.com/ggml-org/llama.cpp/pull/7544 struct llama_context_params { - uint32_t n_ctx; // text context, 0 = from model - uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode - uint32_t n_ubatch; // physical maximum batch size - uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) - uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] - uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) - int32_t n_threads; // number of threads to use for generation - int32_t n_threads_batch; // number of threads to use for batch processing + uint32_t n_ctx; // text context, 0 = from model + uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode + uint32_t n_ubatch; // physical maximum batch size + uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) + uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] + uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) + uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max) + int32_t n_threads; // number of threads to use for generation + int32_t n_threads_batch; // number of threads to use for batch processing enum llama_context_type ctx_type; // set the context type (e.g. MTP) enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` @@ -1057,6 +1059,9 @@ extern "C" { // // Get the backend sampled token for the ith token. + // With multiple outputs, sampler state advances when the token is accepted, + // not when it is read through this function. + // When accepting multiple outputs, accept a contiguous prefix in output order. // Returns LLAMA_TOKEN_NULL if no token was sampled. LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i); @@ -1273,9 +1278,12 @@ extern "C" { // [EXPERIMENTAL] // backend sampling interface: - // return true if the backend supports all ops needed by the sampler + // return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_max_per_seq outputs per sequence // note: call once per sampler - bool (*backend_init)(struct llama_sampler * smpl, ggml_backend_buffer_type_t buft); + bool (*backend_init)( + struct llama_sampler * smpl, + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq); // call after .backend_apply() void (*backend_accept)( @@ -1293,6 +1301,13 @@ extern "C" { // called before graph execution to set inputs for the current ubatch void (*backend_set_input)(struct llama_sampler * smpl); + + // called before rebuilding a sampling graph to clear any internal sampler state + void (*backend_reset)(struct llama_sampler * smpl); + + // copy mutable state from src into dst while keeping dst's references to the current sampling graph + // src and dst must have the same type and configuration + void (*copy_state)(const struct llama_sampler * src, struct llama_sampler * dst); }; struct llama_sampler { @@ -1313,6 +1328,7 @@ extern "C" { LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p); LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl); LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl); + LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst); // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add) LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl); @@ -1502,6 +1518,7 @@ extern "C" { LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl); /// @details Sample and accept a token from the idx-th output of the last evaluation + // For multiple outputs from one sampler, call this function in output order without gaps. // // Shorthand for: // const auto * logits = llama_get_logits_ith(ctx, idx); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index dc70dd8dc..11366fc7f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -10,6 +10,7 @@ #include "llama-mmap.h" #include "llama-model.h" #include "llama-ext.h" +#include "llama-sampler.h" #include "llama.h" #include @@ -162,25 +163,6 @@ llama_context::llama_context( } } - // Initialize backend samplers here so they are part of the sampling graph - // before the reserve passes run later in this function. This avoids a later - // re-reserve when graph nodes change. - if (params.samplers != nullptr && params.n_samplers > 0) { - for (size_t i = 0; i < params.n_samplers; ++i) { - const auto & config = params.samplers[i]; - - if (llama_sampler_chain_get(config.sampler, -1) == nullptr) { - throw std::runtime_error("the backend samplers must be of type llama_sampler_chain"); - } - - if (set_sampler(config.seq_id, config.sampler)) { - const int n_samplers = llama_sampler_chain_n(config.sampler); - - LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers); - } - } - } - auto rope_scaling_type = params.rope_scaling_type; if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) { rope_scaling_type = hparams.rope_scaling_type_train; @@ -268,6 +250,27 @@ llama_context::llama_context( cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch); cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max; + cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ? + cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max); + + // Initialize backend samplers here so they are part of the sampling graph + // before the reserve passes run later in this function. This avoids a later + // re-reserve when graph nodes change. + if (params.samplers != nullptr && params.n_samplers > 0) { + for (size_t i = 0; i < params.n_samplers; ++i) { + const auto & config = params.samplers[i]; + + if (llama_sampler_chain_get(config.sampler, -1) == nullptr) { + throw std::runtime_error("the backend samplers must be of type llama_sampler_chain"); + } + + if (set_sampler(config.seq_id, config.sampler)) { + const int n_samplers = llama_sampler_chain_n(config.sampler); + + LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers); + } + } + } cparams.op_offload = params.op_offload; cparams.kv_unified = params.kv_unified; @@ -303,18 +306,19 @@ llama_context::llama_context( } } - LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max); - LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); - LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq); - LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); - LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); - LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); - LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); - LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); - LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); - LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); - LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq); - LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); + LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max); + LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); + LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq); + LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); + LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); + LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); + LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); + LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); + LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); + LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); + LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq); + LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); + LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n", __func__, cparams.n_outputs_max_per_seq); if (cparams.n_ctx_seq < hparams.n_ctx_train) { LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", @@ -1241,7 +1245,7 @@ bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) { if (sampler && can_offload) { auto * buft = ggml_backend_dev_buffer_type(model.dev_output()); - sampler->iface->backend_init(sampler, buft); + sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq); sampling.samplers[seq_id] = sampler; @@ -1586,108 +1590,38 @@ int llama_context::encode(const llama_batch & batch_inp) { return 0; } -static std::map build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) { - std::map seq_to_row; - // how many output tokens we have seen so far for this ubatch. - uint32_t local = 0; - for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - // skip tokens that are not output. - if (!ubatch.output[i]) { - continue; - } - - const llama_seq_id seq_id = ubatch.seq_id[i][0]; - // row_offset is the number of output tokens before this ubatch. - seq_to_row[seq_id] = row_offset + local; - ++local; - } - return seq_to_row; -} - -static void copy_tensor_async_ints( - const std::map & tensor_map, - const buffer_view & sampled, - const std::map & seq_to_row, - ggml_backend_sched_t sched) { - if (!sampled.has_data()) { - return; - } - - for (const auto & [seq_id, tensor] : tensor_map) { - auto it = seq_to_row.find(seq_id); - if (it == seq_to_row.end()) { - continue; - } - - const uint32_t row = it->second; - GGML_ASSERT(row < sampled.size); - - GGML_ASSERT(ggml_is_contiguous(tensor) && "sampled tokens tensor must be contiguous for async copy"); - - ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor); - ggml_backend_tensor_get_async(backend, tensor, sampled.data + row, 0, sizeof(sampled.data[row])); - } -} - -static void copy_tensor_async_floats( - const std::map & tensor_map, - const buffer_view & dst, +template +static void copy_tensor_async_rows( + const std::vector & tensors, + const buffer_view & dst, size_t stride, - std::vector & counts, - const std::map & seq_to_row, - ggml_backend_sched_t sched) { + uint32_t row_offset, + ggml_backend_sched_t sched, + std::vector * counts = nullptr) { if (!dst.has_data()) { return; } - for (const auto & [seq_id, tensor] : tensor_map) { - auto it = seq_to_row.find(seq_id); - if (it == seq_to_row.end()) { + for (size_t i = 0; i < tensors.size(); ++i) { + auto * tensor = tensors[i]; + if (tensor == nullptr) { continue; } - const uint32_t row = it->second; - GGML_ASSERT(row < counts.size()); - - GGML_ASSERT(ggml_is_contiguous(tensor) && "logits/probs tensor must be contiguous for async copy"); + const uint32_t row = row_offset + i; + const size_t n_elements = ggml_nelements(tensor); + GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy"); + GGML_ASSERT(n_elements <= stride); + GGML_ASSERT((size_t) row * stride + n_elements <= dst.size); ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor); - float * row_ptr = dst.data + (size_t) row * stride; + T * row_ptr = dst.data + (size_t) row * stride; ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor)); - // Update the actual number of logits/probabilities that were written for this row. - counts[row] = ggml_nelements(tensor); - } -} - -static void copy_tensor_async_candidates( - const std::map & tensor_map, - const buffer_view & dst, - size_t stride, - std::vector & counts, - const std::map & seq_to_row, - ggml_backend_sched_t sched) { - if (!dst.has_data()) { - return; - } - - for (const auto & [seq_id, tensor] : tensor_map) { - auto it = seq_to_row.find(seq_id); - if (it == seq_to_row.end()) { - continue; + if (counts) { + GGML_ASSERT(row < counts->size()); + (*counts)[row] = n_elements; } - - const uint32_t row = it->second; - GGML_ASSERT(row < counts.size()); - - GGML_ASSERT(ggml_is_contiguous(tensor) && "candidates tensor must be contiguous for async copy"); - - ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor); - llama_token * row_ptr = dst.data + (size_t) row * stride; - ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor)); - - // Update the actual number of candidates that were written. - counts[row] = ggml_nelements(tensor); } } @@ -1736,12 +1670,12 @@ int llama_context::decode(const llama_batch & batch_inp) { const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max; - // TODO: avoid this workaround in the future - if (has_samplers && batch_inp.logits) { + // embedding contexts output every token even when batch.logits is not set + if (has_samplers && (output_all || batch_inp.logits)) { std::vector seq_output_count(n_seq_max, 0); for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { - if (batch_inp.logits[i] == 0) { + if (!output_all && batch_inp.logits[i] == 0) { continue; } @@ -1750,10 +1684,17 @@ int llama_context::decode(const llama_batch & batch_inp) { for (int32_t s = 0; s < ns; ++s) { const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0; + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { + continue; + } + seq_output_count[seq_id]++; - if (seq_output_count[seq_id] > 1) { - LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n", - __func__, seq_id, seq_output_count[seq_id]); + auto sampler = sampling.samplers.find(seq_id); + if (sampler != sampling.samplers.end() && + seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) { + LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence " + "(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq, + seq_id, seq_output_count[seq_id]); return -1; } } @@ -1853,6 +1794,11 @@ int llama_context::decode(const llama_batch & batch_inp) { return -2; }; + // start a new sampling transaction for this logical batch + for (const auto & entry : sampling.samplers) { + llama_sampler_backend_begin(entry.second); + } + int64_t n_outputs_prev = 0; int64_t n_tokens_prev = 0; @@ -2019,17 +1965,14 @@ int llama_context::decode(const llama_batch & batch_inp) { } } - // Copy backend sampling output if this ubatch produced any sampling tensors. - if (has_samplers && (!res->t_sampled.empty() || !res->t_sampled_probs.empty() || !res->t_sampled_logits.empty())) { - const auto seq_to_output_row = build_seq_to_output_row(ubatch, n_outputs_prev); + if (has_samplers) { const auto stride = n_vocab; // async copy the sampling data from the backend to the host - copy_tensor_async_ints(res->t_sampled, sampling.sampled, seq_to_output_row, sched.get()); - - copy_tensor_async_floats (res->t_sampled_logits, sampling.logits, stride, sampling.logits_count, seq_to_output_row, sched.get()); - copy_tensor_async_floats (res->t_sampled_probs, sampling.probs, stride, sampling.probs_count, seq_to_output_row, sched.get()); - copy_tensor_async_candidates(res->t_candidates, sampling.candidates, stride, sampling.candidates_count, seq_to_output_row, sched.get()); + copy_tensor_async_rows(res->t_sampled, sampling.sampled, 1, n_outputs_prev, sched.get()); + copy_tensor_async_rows(res->t_sampled_logits, sampling.logits, stride, n_outputs_prev, sched.get(), &sampling.logits_count); + copy_tensor_async_rows(res->t_sampled_probs, sampling.probs, stride, n_outputs_prev, sched.get(), &sampling.probs_count); + copy_tensor_async_rows(res->t_candidates, sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count); } n_outputs_prev += n_outputs; @@ -2359,6 +2302,7 @@ void llama_context::output_reorder() { // uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { + uint32_t res; if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || @@ -2367,11 +2311,31 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_M3) { - return std::max(n_tokens * 40, 32u * model.n_tensors()); + res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else { + res = std::max(1024u, 8u*model.n_tensors()); + for (const auto & lora : model.loras) { + res += lora->get_n_nodes(); + } } - uint32_t res = std::max(1024u, 8u*model.n_tensors()); - for (const auto & lora : model.loras) { - res += lora->get_n_nodes(); + + uint32_t n_sampling_nodes = 0; + uint32_t n_sampling_nodes_max = 0; + for (const auto & [seq_id, sampler] : sampling.samplers) { + const uint32_t n_nodes = llama_sampler_backend_n_nodes(sampler); + n_sampling_nodes += n_nodes; + if (cparams.n_outputs_max_per_seq > 1) { + n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes); + } + } + + const uint32_t n_sampling_outputs_max = std::min( + std::min(n_tokens, cparams.n_outputs_max), + (uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq); + + res += n_sampling_nodes; + if (n_sampling_outputs_max > 1) { + res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max; } return res; } @@ -2380,6 +2344,63 @@ llm_graph_result * llama_context::get_gf_res_reserve() const { return static_cast(gf_res_reserve.get()); } +// pack sampler outputs into as few sequences as possible before using sequences without samplers +static void ubatch_prepare_reserve( + llama_ubatch & ubatch, + uint32_t n_outputs, + const std::map & samplers, + uint32_t n_outputs_max_per_seq) { + const uint32_t n_seqs = ubatch.n_seqs; + const uint32_t n_seq_tokens = ubatch.n_seq_tokens; + + for (uint32_t s = 0; s < n_seqs; ++s) { + for (uint32_t t = 0; t < n_seq_tokens; ++t) { + const uint32_t i = s * n_seq_tokens + t; + ubatch.n_seq_id[i] = 1; + ubatch.seq_id[i] = &ubatch.seq_id_unq[s]; + } + } + + // sequences with a sampler that fit in this ubatch + std::vector sampler_seqs; + std::vector has_sampler(n_seqs, false); + for (const auto & entry : samplers) { + const llama_seq_id seq_id = entry.first; + if (seq_id < 0 || (uint32_t) seq_id >= n_seqs) { + continue; + } + + sampler_seqs.push_back(seq_id); + has_sampler[seq_id] = true; + } + + uint32_t n_outputs_set = 0; + + const uint32_t n_outputs_per_seq = std::min(n_seq_tokens, n_outputs_max_per_seq); + for (uint32_t s : sampler_seqs) { + if (n_outputs_set >= n_outputs) { + break; + } + + for (uint32_t t = 0; t < n_outputs_per_seq && n_outputs_set < n_outputs; ++t) { + ubatch.output[s * n_seq_tokens + t] = true; + ++n_outputs_set; + } + } + + // use sequences without samplers for any remaining outputs + for (uint32_t t = 0; t < n_seq_tokens && n_outputs_set < n_outputs; ++t) { + for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) { + if (has_sampler[s]) { + continue; + } + + ubatch.output[s * n_seq_tokens + t] = true; + ++n_outputs_set; + } + } +} + ggml_cgraph * llama_context::graph_reserve( uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only, size_t * sizes) { //LLAMA_LOG_DEBUG("%s: reserving a graph for ubatch with n_tokens = %4u, n_seqs = %2u, n_outputs = %4u\n", __func__, n_tokens, n_seqs, n_outputs); @@ -2404,14 +2425,7 @@ ggml_cgraph * llama_context::graph_reserve( llama_batch_allocr balloc(model.hparams.n_pos_per_embd()); llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs); - // set one output token per sequence in order to activate all backend samplers - std::vector seq_ids(n_seqs); - for (uint32_t i = 0; i < n_seqs; ++i) { - seq_ids[i] = i; - ubatch.n_seq_id[i] = 1; - ubatch.seq_id[i] = &seq_ids[i]; - ubatch.output[i] = true; - } + ubatch_prepare_reserve(ubatch, n_outputs, sampling.samplers, cparams.n_outputs_max_per_seq); auto * res = gf_res_reserve.get(); @@ -3498,6 +3512,7 @@ llama_context_params llama_context_default_params() { /*.n_seq_max =*/ 1, /*.n_rs_seq =*/ 0, /*.n_outputs_max =*/ 0, + /*.n_outputs_max_per_seq =*/ 1, /*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default /*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS, /*.ctx_type =*/ LLAMA_CONTEXT_TYPE_DEFAULT, diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed..574ce9592 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -15,6 +15,7 @@ struct llama_cparams { uint32_t n_seq_max; uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback uint32_t n_outputs_max; // max outputs supported by the context + uint32_t n_outputs_max_per_seq; int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 7382aa2ad..55f1fff33 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -4,6 +4,7 @@ #include "llama-model.h" #include "llama-batch.h" #include "llama-cparams.h" +#include "llama-sampler.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -1354,24 +1355,24 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) { } } } - for (auto & [seq_id, t] : t_sampled) { - if (t != nullptr) { - ggml_set_output(t); + for (auto * tensor : t_sampled) { + if (tensor != nullptr) { + ggml_set_output(tensor); } } - for (auto & [seq_id, t] : t_sampled_probs) { - if (t != nullptr) { - ggml_set_output(t); + for (auto * tensor : t_sampled_probs) { + if (tensor != nullptr) { + ggml_set_output(tensor); } } - for (auto & [seq_id, t] : t_sampled_logits) { - if (t != nullptr) { - ggml_set_output(t); + for (auto * tensor : t_sampled_logits) { + if (tensor != nullptr) { + ggml_set_output(tensor); } } - for (auto & [seq_id, t] : t_candidates) { - if (t != nullptr) { - ggml_set_output(t); + for (auto * tensor : t_candidates) { + if (tensor != nullptr) { + ggml_set_output(tensor); } } } @@ -3650,77 +3651,102 @@ void llm_graph_context::build_sampling() const { auto inp_sampling = std::make_unique(samplers); res->add_input(std::move(inp_sampling)); - std::map seq_to_logit_row; - int32_t logit_row_idx = 0; - - for (uint32_t i = 0; i < ubatch.n_tokens; i++) { + std::map> sampling_rows; + uint32_t n_rows = 0; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { if (ubatch.output[i]) { - llama_seq_id seq_id = ubatch.seq_id[i][0]; - seq_to_logit_row[seq_id] = logit_row_idx; - logit_row_idx++; + sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++); } } + res->t_sampled.resize(n_rows, nullptr); + res->t_sampled_probs.resize(n_rows, nullptr); + res->t_sampled_logits.resize(n_rows, nullptr); + res->t_candidates.resize(n_rows, nullptr); + // res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1) GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor"); - // add a dummy row of logits - // this trick makes the graph static, regardless of which samplers are activated - // this is important in order to minimize graph reallocations + // add a dummy row to keep the single-output graph static regardless of active samplers + // multi-output graphs can still vary with the number of output rows ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0); - for (const auto & [seq_id, sampler] : samplers) { - const auto it = seq_to_logit_row.find(seq_id); - - // inactive samplers always work on the first row - const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0; - const int i_out = it != seq_to_logit_row.end() ? 1 : 0; - - ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]); - ggml_format_name(logits_seq, "logits_seq_%d", seq_id); - - struct llama_sampler_data data = { - /*.logits =*/ logits_seq, - /*.probs =*/ nullptr, - /*.sampled =*/ nullptr, - /*.candidates =*/ nullptr, - }; - - assert(sampler->iface->backend_apply); - sampler->iface->backend_apply(sampler, ctx0, gf, &data); - - if (data.sampled != nullptr) { - res->t_sampled[seq_id] = data.sampled; - outs[1] = data.sampled; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); - } - - if (data.probs != nullptr) { - res->t_sampled_probs[seq_id] = data.probs; - outs[1] = data.probs; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); - } - - if (data.logits != nullptr) { - res->t_sampled_logits[seq_id] = data.logits; - outs[1] = data.logits; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); - } - - if (data.candidates != nullptr) { - res->t_candidates[seq_id] = data.candidates; - outs[1] = data.candidates; - ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + for (const auto & entry : samplers) { + if (entry.second->iface->backend_reset) { + entry.second->iface->backend_reset(entry.second); } } - // TODO: Call llama_sampler_accept_ggml after all samplers have been applied. + static const std::vector dummy_row = { 0 }; + + for (const auto & [seq_id, sampler] : samplers) { + const auto it = sampling_rows.find(seq_id); + + // inactive samplers always work on the first row + const bool active = it != sampling_rows.end(); + const auto & rows = active ? it->second : dummy_row; + const int i_out = active ? 1 : 0; + + for (uint32_t i = 0; i < rows.size(); ++i) { + ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]); + ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i); + + struct llama_sampler_data data = { + /*.logits =*/ logits_seq, + /*.probs =*/ nullptr, + /*.sampled =*/ nullptr, + /*.candidates =*/ nullptr, + }; + + assert(sampler->iface->backend_apply); + sampler->iface->backend_apply(sampler, ctx0, gf, &data); + + if (data.sampled != nullptr) { + if (active) { + res->t_sampled[rows[i]] = data.sampled; + } + outs[1] = data.sampled; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } + + if (data.probs != nullptr) { + if (active) { + res->t_sampled_probs[rows[i]] = data.probs; + } + outs[1] = data.probs; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } + + if (data.logits != nullptr) { + if (active) { + res->t_sampled_logits[rows[i]] = data.logits; + } + outs[1] = data.logits; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } + + if (data.candidates != nullptr) { + if (active) { + res->t_candidates[rows[i]] = data.candidates; + } + outs[1] = data.candidates; + ggml_build_forward_select(gf, outs.data(), outs.size(), i_out); + } + } + } + + // TODO: Call backend_accept after all samplers have been applied. /* for (const auto & [seq_id, sampler] : samplers) { - if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) { - ggml_tensor * selected_token = it->second; - if (selected_token != nullptr) { - llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token); + const auto it = sampling_rows.find(seq_id); + if (it == sampling_rows.end()) { + continue; + } + + for (uint32_t row : it->second) { + ggml_tensor * selected_token = res->t_sampled[row]; + if (selected_token != nullptr && sampler->iface->backend_accept) { + sampler->iface->backend_accept(sampler, ctx0, gf, selected_token); } } } diff --git a/src/llama-graph.h b/src/llama-graph.h index 32d8d395a..75bc0fe80 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -904,10 +904,10 @@ public: std::vector t_layer_inp; - std::map t_sampled_logits; - std::map t_candidates; - std::map t_sampled; - std::map t_sampled_probs; + std::vector t_sampled; + std::vector t_sampled_probs; + std::vector t_sampled_logits; + std::vector t_candidates; std::vector inputs; std::vector fused_nodes; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 90acde355..7dd140f7c 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -543,7 +543,7 @@ llama_model_loader::llama_model_loader( tensor_buft_overrides = param_tensor_buft_overrides_p; - this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; + this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK || load_mode == LLAMA_LOAD_MODE_AUTO; this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO; if (!fname.empty()) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 78c31ad19..746f7bd7a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1416,8 +1416,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { this->ml = &ml; // to be used by create_tensor() and load_arch_tensors() + if (ml.use_mmap && params.load_mode == LLAMA_LOAD_MODE_AUTO) { + for (const auto & dev : devices) { + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev.dev, &props); + if (!props.caps.mmap_support) { + ml.use_mmap = false; + break; + } + } + } + + const char * load_mode_name = params.load_mode == LLAMA_LOAD_MODE_AUTO + ? llama_load_mode_name(ml.use_mmap ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE) + : llama_load_mode_name(params.load_mode); + LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (load_mode = %s)\n", - __func__, llama_load_mode_name(params.load_mode)); + __func__, load_mode_name); // build a list of buffer types for the CPU and GPU devices pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host); @@ -2597,7 +2612,7 @@ llama_model_params llama_model_default_params() { /*.tensor_buft_overrides =*/ nullptr, /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, - /*.load_mode =*/ LLAMA_LOAD_MODE_MMAP, + /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index e550fbe4a..34a798826 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -467,9 +467,11 @@ static void llama_sampler_empty_free(struct llama_sampler * smpl) { static bool llama_sampler_empty_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { GGML_UNUSED(smpl); GGML_UNUSED(buft); + GGML_UNUSED(n_outputs_max_per_seq); return true; } @@ -511,6 +513,8 @@ static struct llama_sampler_i llama_sampler_empty_i = { /* .backend_accept = */ llama_sampler_empty_backend_accept, /* .backend_apply = */ llama_sampler_empty_backend_apply, /* .backend_set_input = */ llama_sampler_empty_backend_set_input, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_empty(const char * name) { @@ -551,6 +555,12 @@ struct llama_sampler_backend { this->support = support; } + // copy the state that is not tied to the current sampling graph + // samplers that hold only immutable configuration can use this as is + void copy_state(const llama_sampler_backend & src) { + GGML_UNUSED(src); + } + private: std::string name; std::string name_ext; @@ -559,6 +569,71 @@ private: bool support; }; +// .copy_state for samplers deriving from llama_sampler_backend +template +static void llama_sampler_backend_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) { + ((T *) dst->ctx)->copy_state(*(const T *) src->ctx); +} + +struct llama_sampler_backend_probe { + ggml_context_ptr ctx; + ggml_cgraph * gf; +}; + +static llama_sampler_backend_probe llama_sampler_backend_probe_graph( + llama_sampler * sampler, + int64_t n_candidates, + uint32_t max_nodes, + bool with_candidates) { + ggml_init_params params = { + /*.mem_size =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + + ggml_context_ptr ctx_ptr { ggml_init(params) }; + if (!ctx_ptr) { + throw std::runtime_error(format("failed to create ggml context")); + } + + auto * ctx = ctx_ptr.get(); + auto * gf = ggml_new_graph_custom(ctx, max_nodes, false); + + llama_sampler_data data = { + /*.logits =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates), + /*.probs =*/ nullptr, + /*.sampled =*/ nullptr, + /*.candidates =*/ with_candidates ? ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates) : nullptr, + }; + + if (sampler->iface->backend_reset) { + sampler->iface->backend_reset(sampler); + } + sampler->iface->backend_apply(sampler, ctx, gf, &data); + + for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) { + if (output) { + ggml_build_forward_expand(gf, output); + } + } + + if (sampler->iface->backend_reset) { + sampler->iface->backend_reset(sampler); + } + + return { std::move(ctx_ptr), gf }; +} + +static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) { + uint32_t n_tensors = 0; + for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor; + tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) { + ++n_tensors; + } + + return std::max(ggml_graph_n_nodes(probe.gf), n_tensors); +} + // check if all ggml ops used by the sampler are supported by the backend static bool llama_sampler_backend_support( llama_sampler * smpl, @@ -569,50 +644,10 @@ static bool llama_sampler_backend_support( return true; } - ggml_init_params params = { - /*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(), - /*.mem_buffer =*/ NULL, - /*.no_alloc =*/ true, - }; + auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true); - ggml_context_ptr ctx_ptr { ggml_init(params) }; - if (!ctx_ptr) { - throw std::runtime_error(format("failed to create ggml context")); - } - - ggml_context * ctx = ctx_ptr.get(); - - const int64_t n = 1024*1024; - - llama_sampler_data data = { - /*.logits = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n), - /*.probs = */ nullptr, - /*.sampled = */ nullptr, - /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n), - }; - - ggml_cgraph * gf = ggml_new_graph(ctx); - - smpl->iface->backend_apply(smpl, ctx, gf, &data); - - if (data.logits) { - ggml_build_forward_expand(gf, data.logits); - } - - if (data.probs) { - ggml_build_forward_expand(gf, data.probs); - } - - if (data.sampled) { - ggml_build_forward_expand(gf, data.sampled); - } - - if (data.candidates) { - ggml_build_forward_expand(gf, data.candidates); - } - - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - struct ggml_tensor * op = ggml_graph_node(gf, i); + for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) { + struct ggml_tensor * op = ggml_graph_node(probe.gf, i); if (!ggml_backend_dev_supports_op(device, op)) { LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n", @@ -697,7 +732,8 @@ static void llama_sampler_chain_free(struct llama_sampler * smpl) { static bool llama_sampler_chain_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * chain = (llama_sampler_chain *) smpl->ctx; GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice"); @@ -705,26 +741,32 @@ static bool llama_sampler_chain_backend_init( chain->is_init = true; bool res = true; + bool backend_prefix = true; for (auto & smpl : chain->samplers) { - bool res_cur = true; + bool cur_prefix = backend_prefix; // to be able to run a sampler on the backend, it has to: // - have the .backend_init() API implemented // - return true during .backend_init() - if (smpl.ptr->iface->backend_init) { - if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) { - res_cur = false; + // - support the requested per-sequence output limit + if (cur_prefix && smpl.ptr->iface->backend_init) { + if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) { + cur_prefix = false; } } else { - res_cur = false; + cur_prefix = false; } - smpl.is_backend = res_cur; + smpl.is_backend = cur_prefix; + backend_prefix = cur_prefix; - res = res && res_cur; + res = res && cur_prefix; } + auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false); + chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe); + return res; } @@ -780,6 +822,36 @@ static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) { } } +static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) { + auto * chain = (llama_sampler_chain *) smpl->ctx; + + for (auto & entry : chain->samplers) { + if (!entry.is_backend) { + break; + } + if (entry.ptr->iface->backend_reset) { + entry.ptr->iface->backend_reset(entry.ptr); + } + } +} + +static void llama_sampler_chain_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) { + const auto * src_chain = (const llama_sampler_chain *) src->ctx; + auto * dst_chain = (llama_sampler_chain *) dst->ctx; + + GGML_ASSERT(src_chain->samplers.size() == dst_chain->samplers.size()); + + for (size_t i = 0; i < src_chain->samplers.size(); ++i) { + llama_sampler_copy(src_chain->samplers[i].ptr, dst_chain->samplers[i].ptr); + } + + // note: is_init, n_nodes and is_backend belong to the current sampling graph + dst_chain->params = src_chain->params; + dst_chain->cur = src_chain->cur; + dst_chain->t_sample_us = src_chain->t_sample_us; + dst_chain->n_sample = src_chain->n_sample; +} + static struct llama_sampler_i llama_sampler_chain_i = { /* .name = */ llama_sampler_chain_name, /* .accept = */ llama_sampler_chain_accept, @@ -791,22 +863,35 @@ static struct llama_sampler_i llama_sampler_chain_i = { /* .backend_accept = */ llama_sampler_chain_backend_accept, /* .backend_apply = */ llama_sampler_chain_backend_apply, /* .backend_set_input = */ llama_sampler_chain_backend_set_input, + /* .backend_reset = */ llama_sampler_chain_backend_reset, + /* .copy_state = */ llama_sampler_chain_copy_state, }; struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) { return llama_sampler_init( /* .iface = */ &llama_sampler_chain_i, /* .ctx = */ new llama_sampler_chain { - /* .params = */ params, - /* .is_init = */ false, - /* .samplers = */ {}, - /* .cur = */ {}, - /* .t_sample_us = */ 0, - /* .n_sample = */ 0, + /* .params = */ params, + /* .is_init = */ false, + /* .n_nodes = */ 0, + /* .samplers = */ {}, + /* .cur = */ {}, + /* .t_sample_us = */ 0, + /* .n_sample = */ 0, } ); } +uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) { + GGML_ASSERT(sampler != nullptr); + GGML_ASSERT(sampler->iface == &llama_sampler_chain_i); + + const auto * chain = (const llama_sampler_chain *) sampler->ctx; + GGML_ASSERT(chain->is_init); + + return chain->n_nodes; +} + llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) { const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx); const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx); @@ -816,6 +901,7 @@ llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_conte // If a backend sampler has already sampled a token, return it. if (sampled_token != LLAMA_TOKEN_NULL) { LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx); + llama_sampler_accept(smpl, sampled_token); return sampled_token; } @@ -975,8 +1061,10 @@ static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_to static bool llama_sampler_greedy_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_greedy *) smpl->ctx; + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1012,6 +1100,8 @@ static struct llama_sampler_i llama_sampler_greedy_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_greedy_backend_apply, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_greedy() { @@ -1031,7 +1121,25 @@ struct llama_sampler_dist : public llama_sampler_backend { std::mt19937 rng; - ggml_tensor * inp_uniform; + // TODO: refactor + fix naming + // https://github.com/ggml-org/llama.cpp/pull/25532/changes#r3749906719 + // use a temporary RNG for multi-output sampling so rejected tokens do not advance rng + bool backend_transactional; + std::mt19937 rng_backend; + size_t n_backend_draws_generated; + size_t n_backend_draws_committed; + + // inputs for the current sampling graph + std::vector inp_uniforms; + + void copy_state(const llama_sampler_dist & src) { + // note: inp_uniforms and backend_transactional belong to the current sampling graph + seed_cur = src.seed_cur; + rng = src.rng; + rng_backend = src.rng_backend; + n_backend_draws_generated = src.n_backend_draws_generated; + n_backend_draws_committed = src.n_backend_draws_committed; + } }; static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) { @@ -1050,7 +1158,11 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da cur_p->selected = 0; + std::uniform_real_distribution dist(0.0f, 1.0f); + if (cur_p->size == 1) { + // keep the RNG state aligned with backend sampling, which draws once per output + dist(ctx->rng); cur_p->data[0].p = 1.0f; return; } @@ -1075,7 +1187,6 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da // sample from the obtained probabilities and normalize the probs in a single pass // this is ~3x faster on Mac with full gpt-oss vocab than the version below // - std::uniform_real_distribution dist(0.0f, 1.0f); const double rnd = dist(ctx->rng); double sum_run = 0.0f; @@ -1115,6 +1226,9 @@ static void llama_sampler_dist_reset(struct llama_sampler * smpl) { auto * ctx = (llama_sampler_dist *) smpl->ctx; ctx->seed_cur = get_rng_seed(ctx->seed); ctx->rng.seed(ctx->seed_cur); + ctx->rng_backend = ctx->rng; + ctx->n_backend_draws_generated = 0; + ctx->n_backend_draws_committed = 0; } static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) { @@ -1125,7 +1239,12 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample { auto * result_ctx = (llama_sampler_dist *) result->ctx; - result_ctx->rng = ctx->rng; + result_ctx->seed_cur = ctx->seed_cur; + result_ctx->rng = ctx->rng; + result_ctx->backend_transactional = ctx->backend_transactional; + result_ctx->rng_backend = ctx->rng_backend; + result_ctx->n_backend_draws_generated = ctx->n_backend_draws_generated; + result_ctx->n_backend_draws_committed = ctx->n_backend_draws_committed; } return result; @@ -1137,12 +1256,17 @@ static void llama_sampler_dist_free(struct llama_sampler * smpl) { static bool llama_sampler_dist_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_dist *) smpl->ctx; const bool res = llama_sampler_backend_support(smpl, buft); sctx->init(res); + sctx->backend_transactional = n_outputs_max_per_seq > 1; + sctx->rng_backend = sctx->rng; + sctx->n_backend_draws_generated = 0; + sctx->n_backend_draws_committed = 0; return res; } @@ -1156,9 +1280,10 @@ static void llama_sampler_dist_backend_apply( auto * sctx = (llama_sampler_dist *) smpl->ctx; - sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); - ggml_set_name (sctx->inp_uniform, "uniform"); - ggml_set_input(sctx->inp_uniform); + ggml_tensor * inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_format_name(inp_uniform, "uniform_%zu", sctx->inp_uniforms.size()); + ggml_set_input(inp_uniform); + sctx->inp_uniforms.push_back(inp_uniform); // flatten struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits)); @@ -1174,7 +1299,7 @@ static void llama_sampler_dist_backend_apply( // Recall that each entry in cumsum is the cumulative probability up to that // index so values stay negative while the cumulative total is below the // random value, and become zero/positive once the threshold is crossed. - struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform); + struct ggml_tensor * diff = ggml_sub(ctx, cumsum, inp_uniform); ggml_set_name(diff, "dist_cumsum"); // The ggml_step function produces a tensor where entries are 1 if the @@ -1189,6 +1314,9 @@ static void llama_sampler_dist_backend_apply( struct ggml_tensor * idxf = ggml_sum(ctx, mask); ggml_set_name(idxf, "dist_index_f32"); + // Clamp to prevent out-of-bounds access when computing the index. + idxf = ggml_clamp(ctx, idxf, 1.0f, mask->ne[0]); + // Use ggml_scale_bias to scale the index value by -1 and then add the size // of the mask to that value so we get the correct index ((-1 * idxf) + n). struct ggml_tensor * idx = ggml_cast(ctx, ggml_scale_bias(ctx, idxf, -1.0f, mask->ne[0]), GGML_TYPE_I32); @@ -1210,22 +1338,52 @@ static void llama_sampler_dist_backend_apply( static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) { auto * sctx = (llama_sampler_dist *) smpl->ctx; - GGML_ASSERT(sctx->inp_uniform != nullptr); + GGML_ASSERT(!sctx->inp_uniforms.empty()); // We sample in double precision and cast to float to match rnd numbers of - // llama_dampler_dist which uses double precision (sampling from + // llama_sampler_dist which uses double precision (sampling from // std::uniform_real_distribution and // std::uniform_real_distribution with same rng will produce // different sequences). std::uniform_real_distribution dist(0.0f, 1.0f); - const float rnd = dist(sctx->rng); - ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float)); + auto & rng = sctx->backend_transactional ? sctx->rng_backend : sctx->rng; + + for (auto * inp_uniform : sctx->inp_uniforms) { + GGML_ASSERT(inp_uniform != nullptr); + + const float rnd = dist(rng); + ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float)); + + if (sctx->backend_transactional) { + ++sctx->n_backend_draws_generated; + } + } +} + +static void llama_sampler_dist_backend_reset(struct llama_sampler * smpl) { + auto * sctx = (llama_sampler_dist *) smpl->ctx; + sctx->inp_uniforms.clear(); +} + +static void llama_sampler_dist_accept(struct llama_sampler * smpl, llama_token token) { + GGML_UNUSED(token); + + auto * sctx = (llama_sampler_dist *) smpl->ctx; + + if (!sctx->backend_transactional || + sctx->n_backend_draws_committed >= sctx->n_backend_draws_generated) { + return; + } + + std::uniform_real_distribution dist(0.0f, 1.0f); + dist(sctx->rng); + ++sctx->n_backend_draws_committed; } static struct llama_sampler_i llama_sampler_dist_i = { /* .name = */ llama_sampler_dist_name, - /* .accept = */ nullptr, + /* .accept = */ llama_sampler_dist_accept, /* .apply = */ llama_sampler_dist_apply, /* .reset = */ llama_sampler_dist_reset, /* .clone = */ llama_sampler_dist_clone, @@ -1234,6 +1392,8 @@ static struct llama_sampler_i llama_sampler_dist_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_dist_backend_apply, /* .backend_set_input = */ llama_sampler_dist_backend_set_input, + /* .backend_reset = */ llama_sampler_dist_backend_reset, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_dist(uint32_t seed) { @@ -1242,14 +1402,39 @@ struct llama_sampler * llama_sampler_init_dist(uint32_t seed) { /* .iface = */ &llama_sampler_dist_i, /* .ctx = */ new llama_sampler_dist { ("dist"), - /* .seed = */ seed, - /* .seed_cur = */ seed_cur, - /* .rng = */ std::mt19937(seed_cur), - /* .inp_uniform = */ nullptr, + /* .seed = */ seed, + /* .seed_cur = */ seed_cur, + /* .rng = */ std::mt19937(seed_cur), + /* .backend_transactional = */ false, + /* .rng_backend = */ std::mt19937(seed_cur), + /* .n_backend_draws_generated = */ 0, + /* .n_backend_draws_committed = */ 0, + /* .inp_uniforms = */ {}, } ); } +void llama_sampler_backend_begin(llama_sampler * sampler) { + GGML_ASSERT(sampler != nullptr); + + if (sampler->iface == &llama_sampler_chain_i) { + auto * chain = (llama_sampler_chain *) sampler->ctx; + for (auto & entry : chain->samplers) { + if (!entry.is_backend) { + break; + } + llama_sampler_backend_begin(entry.ptr); + } + } else if (sampler->iface == &llama_sampler_dist_i) { + auto * ctx = (llama_sampler_dist *) sampler->ctx; + if (ctx->backend_transactional) { + ctx->rng_backend = ctx->rng; + ctx->n_backend_draws_generated = 0; + ctx->n_backend_draws_committed = 0; + } + } +} + // top-k struct llama_sampler_top_k : public llama_sampler_backend { @@ -1277,8 +1462,10 @@ static void llama_sampler_top_k_free(struct llama_sampler * smpl) { static bool llama_sampler_top_k_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_top_k *) smpl->ctx; + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1325,6 +1512,8 @@ static struct llama_sampler_i llama_sampler_top_k_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_top_k_backend_apply, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_top_k(int32_t k) { @@ -1423,8 +1612,10 @@ static void llama_sampler_top_p_free(struct llama_sampler * smpl) { static bool llama_sampler_top_p_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_top_p *) smpl->ctx; + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1521,6 +1712,8 @@ static struct llama_sampler_i llama_sampler_top_p_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_top_p_backend_apply, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) { @@ -1618,8 +1811,10 @@ static void llama_sampler_min_p_free(struct llama_sampler * smpl) { static bool llama_sampler_min_p_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_min_p *) smpl->ctx; + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1680,6 +1875,8 @@ static struct llama_sampler_i llama_sampler_min_p_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_min_p_backend_apply, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) { @@ -1790,6 +1987,8 @@ static struct llama_sampler_i llama_sampler_typical_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) { @@ -1866,8 +2065,10 @@ static void llama_sampler_backend_temp_sampling( static bool llama_sampler_temp_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_temp *) smpl->ctx; + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1896,6 +2097,8 @@ static struct llama_sampler_i llama_sampler_temp_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_temp_backend_apply, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_temp(float temp) { @@ -2009,8 +2212,10 @@ static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) { static bool llama_sampler_temp_ext_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_temp_ext *) smpl->ctx; + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -2095,6 +2300,8 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_temp_ext_backend_apply, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) { @@ -2202,6 +2409,8 @@ static struct llama_sampler_i llama_sampler_xtc_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) { @@ -2290,7 +2499,7 @@ static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sa // copy the state { - auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx; + auto * result_ctx = (llama_sampler_mirostat *) result->ctx; result_ctx->mu = ctx->mu; result_ctx->rng = ctx->rng; @@ -2321,6 +2530,8 @@ static struct llama_sampler_i llama_sampler_mirostat_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) { @@ -2425,6 +2636,8 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) { @@ -2546,6 +2759,8 @@ static struct llama_sampler_i llama_sampler_grammar_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; static struct llama_sampler * llama_sampler_init_grammar_impl( @@ -2661,6 +2876,12 @@ struct llama_sampler_penalties : public llama_sampler_backend { std::vector host_token_ids; std::vector host_counts; + void copy_state(const llama_sampler_penalties & src) { + // note: inp_token_ids/inp_counts belong to the current sampling graph + prev = src.prev; + token_count = src.token_count; + } + static bool is_disabled( int32_t penalty_last_n, float penalty_repeat, @@ -2790,9 +3011,15 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) { static bool llama_sampler_penalties_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_penalties *) smpl->ctx; + if (n_outputs_max_per_seq > 1) { + sctx->init(false); + return false; + } + const bool res = llama_sampler_backend_support(smpl, buft); sctx->init(res); @@ -2952,6 +3179,12 @@ static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smp ggml_backend_tensor_set(sctx->inp_counts, sctx->host_counts.data(), 0, sctx->n_max * sizeof(int32_t)); } +static void llama_sampler_penalties_backend_reset(struct llama_sampler * smpl) { + auto * sctx = (llama_sampler_penalties *) smpl->ctx; + sctx->inp_token_ids = nullptr; + sctx->inp_counts = nullptr; +} + static struct llama_sampler_i llama_sampler_penalties_i = { /* .name = */ llama_sampler_penalties_name, /* .accept = */ llama_sampler_penalties_accept, @@ -2963,6 +3196,8 @@ static struct llama_sampler_i llama_sampler_penalties_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_penalties_backend_apply, /* .backend_set_input = */ llama_sampler_penalties_backend_set_input, + /* .backend_reset = */ llama_sampler_penalties_backend_reset, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_penalties( @@ -3058,6 +3293,8 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_top_n_sigma(float n) { @@ -3395,6 +3632,8 @@ static struct llama_sampler_i llama_sampler_dry_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) { @@ -3614,6 +3853,8 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_adaptive_p( @@ -3715,13 +3956,17 @@ static void llama_sampler_logit_bias_backend_apply( const size_t n = sctx->logit_bias.size(); - sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n); - ggml_set_name(sctx->inp_logit_bias, "logit_bias"); - ggml_set_input(sctx->inp_logit_bias); + if (sctx->inp_logit_bias == nullptr) { + GGML_ASSERT(sctx->inp_logit_idxs == nullptr); - sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n); - ggml_set_name(sctx->inp_logit_idxs, "logit_idxs"); - ggml_set_input(sctx->inp_logit_idxs); + sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n); + ggml_set_name(sctx->inp_logit_bias, "logit_bias"); + ggml_set_input(sctx->inp_logit_bias); + + sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n); + ggml_set_name(sctx->inp_logit_idxs, "logit_idxs"); + ggml_set_input(sctx->inp_logit_idxs); + } ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f); @@ -3756,10 +4001,18 @@ static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * sm ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs)); } +static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) { + auto * sctx = (llama_sampler_logit_bias *) smpl->ctx; + sctx->inp_logit_bias = nullptr; + sctx->inp_logit_idxs = nullptr; +} + static bool llama_sampler_logit_bias_backend_init( struct llama_sampler * smpl, - ggml_backend_buffer_type_t buft) { + ggml_backend_buffer_type_t buft, + uint32_t n_outputs_max_per_seq) { GGML_UNUSED(buft); + GGML_UNUSED(n_outputs_max_per_seq); auto * sctx = (llama_sampler_logit_bias *) smpl->ctx; @@ -3783,6 +4036,8 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = { /* .backend_accept = */ nullptr, /* .backend_apply = */ llama_sampler_logit_bias_backend_apply, /* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input, + /* .backend_reset = */ llama_sampler_logit_bias_backend_reset, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_logit_bias( @@ -4022,10 +4277,12 @@ static struct llama_sampler_i llama_sampler_infill_i = { /* .reset = */ nullptr, /* .clone = */ llama_sampler_infill_clone, /* .free = */ llama_sampler_infill_free, - /* .backend_apply = */ nullptr, - /* .backend_accept = */ nullptr, - /* .backend_set_input = */ nullptr, /* .backend_init = */ nullptr, + /* .backend_accept = */ nullptr, + /* .backend_apply = */ nullptr, + /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) { @@ -4039,6 +4296,32 @@ struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * voca ); } +void llama_sampler_copy(const struct llama_sampler * src, struct llama_sampler * dst) { + if (!src || !dst || src == dst) { + return; + } + + GGML_ASSERT(src->iface == dst->iface && "llama_sampler_copy: cannot copy between different sampler types"); + + if (dst->iface->copy_state) { + dst->iface->copy_state(src, dst); + return; + } + + // build a temporary sampler carrying src's current state + llama_sampler * tmp = llama_sampler_clone(src); + + // free dst's old state (frees dst->ctx, including children for a chain) + if (dst->iface->free) { + dst->iface->free(dst); + } + + // transplant tmp's state into dst, then destroy the (now empty) temp shell + dst->ctx = tmp->ctx; + tmp->ctx = nullptr; + delete tmp; +} + // utils uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) { diff --git a/src/llama-sampler.h b/src/llama-sampler.h index 929207514..e5db2982b 100644 --- a/src/llama-sampler.h +++ b/src/llama-sampler.h @@ -15,6 +15,8 @@ struct llama_sampler_chain { // has .backend_init() been called? bool is_init = false; + uint32_t n_nodes = 0; + struct info { bool is_backend; @@ -33,6 +35,9 @@ struct llama_sampler_chain { mutable int32_t n_sample; }; +uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler); +void llama_sampler_backend_begin(llama_sampler * sampler); + struct llama_sampler * llama_sampler_init_dry_testing( float dry_multiplier, float dry_base, diff --git a/src/llama.cpp b/src/llama.cpp index f60971292..2d8b893c0 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -69,6 +69,8 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty const char * llama_load_mode_name(enum llama_load_mode load_mode) { switch (load_mode) { + case LLAMA_LOAD_MODE_AUTO: + return "auto"; case LLAMA_LOAD_MODE_NONE: return "none"; case LLAMA_LOAD_MODE_MMAP: @@ -84,11 +86,12 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) { } enum llama_load_mode llama_load_mode_from_str(const char * str) { - if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } - if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } - if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } + if (std::strcmp(str, "auto") == 0) { return LLAMA_LOAD_MODE_AUTO; } + if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } + if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } + if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; } - if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } + if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } throw std::invalid_argument(std::string("unknown load mode: ") + str); } diff --git a/src/models/exaone4.cpp b/src/models/exaone4.cpp index 863268abc..a06819a67 100644 --- a/src/models/exaone4.cpp +++ b/src/models/exaone4.cpp @@ -1,6 +1,9 @@ #include "models.h" void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer"); + if (hparams.n_layer() == 64) { // 32B hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; hparams.n_swa = 4096; @@ -15,9 +18,6 @@ void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); - ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); - - GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer"); switch (hparams.n_layer()) { case 30: type = LLM_TYPE_1_2B; break; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 38d2e5c7a..3b5f6a121 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -39,19 +39,18 @@ using json = nlohmann::ordered_json; constexpr int HTTP_POLLING_SECONDS = 1; -static uint32_t server_n_outputs_max(const common_params & params) { - const uint32_t n_batch = params.n_batch; - +static common_speculative_output_limits server_output_limits(const common_params & params) { if (params.embedding || (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { - return n_batch; + return { params.n_batch, 1 }; } - const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(¶ms.speculative); + auto result = common_speculative_get_output_limits( + params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative)); - const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq; - - return std::max(1, std::min(n_batch, n_outputs)); + result.total = std::max(1, result.total); + result.per_seq = std::max(1, result.per_seq); + return result; } // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 @@ -1063,7 +1062,9 @@ private: const bool is_resume = sleeping; params_base = params; - params_base.n_outputs_max = server_n_outputs_max(params_base); + const auto output_limits = server_output_limits(params_base); + params_base.n_outputs_max = output_limits.total; + params_base.n_outputs_max_per_seq = output_limits.per_seq; const bool has_mmproj = !params.mmproj.path.empty(); const bool has_draft = params.speculative.has_dft(); @@ -1832,18 +1833,13 @@ private: const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs; - bool backend_sampling = true; - - backend_sampling &= task.params.sampling.backend_sampling; - - // TODO: speculative decoding requires multiple samples per batch - not supported yet - backend_sampling &= !(slot.can_speculate()); + bool use_backend_sampling = task.params.sampling.backend_sampling; // TODO: getting pre sampling logits is not yet supported with backend sampling - backend_sampling &= !need_pre_sample_logits; + use_backend_sampling &= !need_pre_sample_logits; // TODO: tmp until backend sampling is fully implemented - if (backend_sampling) { + if (use_backend_sampling) { llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); } else { llama_set_sampler(ctx_tgt, slot.id, nullptr); @@ -3865,7 +3861,8 @@ private: // speculative decoding - main model sample and accept iterate(slots, [&](server_slot & slot) { - if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || slot.spec_draft.empty()) { + if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || + slot.spec_draft.empty() || slot.spec_i_batch.empty()) { return; } @@ -3876,7 +3873,6 @@ private: // verify and try to accept the draft { - // save the sampler sampler state in case we need to restore it common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); @@ -3915,7 +3911,7 @@ private: slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); slot.prompt.tokens.keep_first(ckpt.n_tokens); - slot.smpl = std::move(smpl_save); + common_sampler_copy(smpl_save.get(), slot.smpl.get()); return; } diff --git a/tools/server/tests/tests.sh b/tools/server/tests/tests.sh index 8d6681193..433dc9982 100755 --- a/tools/server/tests/tests.sh +++ b/tools/server/tests/tests.sh @@ -6,11 +6,6 @@ cd $SCRIPT_DIR set -eu -if [[ "${SLOW_TESTS:-0}" == 1 ]]; then - # Slow tests for tool calls need quite a few models ahead of time to avoid timing out. - python $SCRIPT_DIR/../../../scripts/fetch_server_test_models.py -fi - if [ $# -lt 1 ] then if [[ "${SLOW_TESTS:-0}" == 1 ]]; then diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index c6568479c..0184085b4 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -25,33 +25,34 @@ def fixture_create_server(): def test_with_and_without_draft(): global server + request = { + "prompt": "I believe the meaning of life is", + "temperature": 0.8, + "top_k": 40, + "seed": 4242, + "n_predict": 16, + "return_tokens": True, + } + server.model_draft = None # disable draft model server.spec_type = None + server.backend_sampling = True server.start() - res = server.make_request("POST", "/completion", data={ - "prompt": "I believe the meaning of life is", - "temperature": 0.0, - "top_k": 1, - "n_predict": 16, - }) + res = server.make_request("POST", "/completion", data=request) assert res.status_code == 200 - content_no_draft = res.body["content"] + tokens_no_draft = res.body["tokens"] server.stop() # create new server with draft model create_server() + server.backend_sampling = True server.start() - res = server.make_request("POST", "/completion", data={ - "prompt": "I believe the meaning of life is", - "temperature": 0.0, - "top_k": 1, - "n_predict": 16, - }) + res = server.make_request("POST", "/completion", data=request) assert res.status_code == 200 assert res.body["timings"]["draft_n"] > 0 - content_draft = res.body["content"] + tokens_draft = res.body["tokens"] - assert content_no_draft == content_draft + assert tokens_no_draft == tokens_draft def test_different_draft_min_draft_max():