From dd1ea524333b1e697489067d7a4c39c60d32beee Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Mon, 10 Aug 2026 19:28:56 +0530 Subject: [PATCH 01/10] llama : support multi-output backend sampling (#25532) * Enable backend sampling with token speculation * Clamp the mask sum before converting it into the sampled index * Add a numeric context parameter declaring the maximum outputs one sequence * More fixes * Don't reuse memory for output views. * Match dist between CPU and GPU * Fix CPU and backend sampling mismatches * Simpify some of the changes * Fix tests on Vulkan * More test fixes * Rebase changes * Rebase and address review comments * Address review comments * Address review comments * Update src/llama-sampler.cpp Co-authored-by: Georgi Gerganov --------- Co-authored-by: Georgi Gerganov --- common/common.cpp | 1 + common/common.h | 1 + common/llguidance.cpp | 2 + common/reasoning-budget.cpp | 2 + common/sampling.cpp | 20 + common/sampling.h | 1 + common/speculative.cpp | 12 + common/speculative.h | 9 + docs/speculative.md | 6 + examples/lookup/lookup.cpp | 6 + .../speculative-simple/speculative-simple.cpp | 9 + examples/speculative/speculative.cpp | 8 + include/llama.h | 38 +- src/llama-context.cpp | 309 +++++------ src/llama-cparams.h | 1 + src/llama-graph.cpp | 164 +++--- src/llama-graph.h | 8 +- src/llama-sampler.cpp | 469 +++++++++++++---- src/llama-sampler.h | 5 + tests/test-arg-parser.cpp | 30 ++ tests/test-backend-sampler.cpp | 497 ++++++++++++++++-- tests/test-sampling.cpp | 31 ++ tools/server/server-context.cpp | 36 +- tools/server/tests/unit/test_speculative.py | 31 +- 24 files changed, 1305 insertions(+), 391 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index ffe3e7761..2e3f14cd1 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1639,6 +1639,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 4811345f9..878534dcc 100644 --- a/common/common.h +++ b/common/common.h @@ -447,6 +447,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) 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/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 0ebf9c5ad..2ee1e6b84 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/docs/speculative.md b/docs/speculative.md index 3957db85c..25abef1b6 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -202,6 +202,12 @@ Example Video: If a draft model is combined with a draftless decoding the draftless decoding has higher precedence. +### Backend Sampling + +Use `--backend-sampling` to run supported target-model samplers on the model backend. Draft-model sampling uses the backend by default and can be controlled with `--spec-draft-backend-sampling` and `--no-spec-draft-backend-sampling`. + +Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required. + ### General Speculative Parameters ``` diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp index 2d4c0e528..662105865 100644 --- a/examples/lookup/lookup.cpp +++ b/examples/lookup/lookup.cpp @@ -3,9 +3,11 @@ #include "common.h" #include "ngram-cache.h" #include "sampling.h" +#include "speculative.h" #include "log.h" #include "llama.h" +#include #include #include #include @@ -27,6 +29,10 @@ int main(int argc, char ** argv){ // max. number of additional tokens to draft if match is found const int n_draft = params.speculative.draft.n_max; + const auto output_limits = common_speculative_get_output_limits(params.n_batch, params.n_parallel, n_draft); + params.n_outputs_max = output_limits.total; + params.n_outputs_max_per_seq = output_limits.per_seq; + // init llama.cpp llama_backend_init(); llama_numa_init(params.numa); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index d87ba48be..c727e8139 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -5,6 +5,7 @@ #include "log.h" #include "llama.h" +#include #include #include #include @@ -29,6 +30,11 @@ int main(int argc, char ** argv) { return 1; } + const auto output_limits = common_speculative_get_output_limits( + params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative)); + params.n_outputs_max = output_limits.total; + params.n_outputs_max_per_seq = output_limits.per_seq; + // init llama.cpp llama_backend_init(); llama_numa_init(params.numa); @@ -55,6 +61,9 @@ int main(int argc, char ** argv) { auto params_dft = params; + params_dft.n_outputs_max = params.n_parallel; + params_dft.n_outputs_max_per_seq = 1; + params_dft.devices = params_spec.devices; params_dft.model = params_spec.mparams; params_dft.n_gpu_layers = params_spec.n_gpu_layers; diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index f7fa5e306..17071aa05 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -1,6 +1,7 @@ #include "arg.h" #include "common.h" #include "sampling.h" +#include "speculative.h" #include "log.h" #include "llama.h" @@ -57,6 +58,11 @@ int main(int argc, char ** argv) { // max number of parallel drafting sequences (i.e. tree branches) const int n_seq_dft = params.n_parallel; + const auto output_limits = common_speculative_get_output_limits( + params.n_batch, params.n_parallel, params.speculative.draft.n_max); + params.n_outputs_max = output_limits.total; + params.n_outputs_max_per_seq = output_limits.per_seq; + // probability threshold for splitting a draft branch (only for n_seq_dft > 1) const float p_draft_split = params.speculative.draft.p_split; @@ -83,6 +89,8 @@ int main(int argc, char ** argv) { params.devices = params.speculative.draft.devices; params.model = params.speculative.draft.mparams; params.n_gpu_layers = params.speculative.draft.n_gpu_layers; + params.n_outputs_max = params.n_parallel; + params.n_outputs_max_per_seq = 1; if (params.speculative.draft.cpuparams.n_threads > 0) { params.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads; } diff --git a/include/llama.h b/include/llama.h index a14498925..c14eb6f50 100644 --- a/include/llama.h +++ b/include/llama.h @@ -348,14 +348,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` @@ -1054,6 +1055,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); @@ -1270,9 +1274,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)( @@ -1290,6 +1297,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 { @@ -1310,6 +1324,9 @@ 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); + // copy mutable sampler state without changing dst or its sampling graph bindings + // src and dst must have the same type and configuration + 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); @@ -1499,6 +1516,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 6f2bf1362..0de3a68d1 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 @@ -159,25 +160,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; @@ -265,6 +247,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; @@ -300,18 +303,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", @@ -1231,7 +1235,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; @@ -1576,108 +1580,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); } } @@ -1726,12 +1660,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; } @@ -1740,10 +1674,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; } } @@ -1843,6 +1784,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; @@ -2009,17 +1955,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; @@ -2349,6 +2292,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 || @@ -2357,11 +2301,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; } @@ -2370,6 +2334,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); @@ -2394,14 +2415,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(); @@ -3488,6 +3502,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 2be3b75fb..55d858024 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" @@ -1353,24 +1354,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); } } } @@ -3649,77 +3650,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-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/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 50db29727..ba58f852e 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -2,7 +2,9 @@ #include "common.h" #include "download.h" #include "llama.h" +#include "speculative.h" +#include #include #include #include @@ -14,6 +16,34 @@ static void test(void) { common_params params; + auto assert_output_limits = [](int32_t n_batch, int32_t n_parallel, int32_t n_draft, + int32_t total, int32_t per_seq) { + const auto limits = common_speculative_get_output_limits(n_batch, n_parallel, n_draft); + assert(limits.total == total); + assert(limits.per_seq == per_seq); + }; + + assert_output_limits(16, 2, 3, 8, 4); + assert_output_limits(16, 2, -1, 2, 1); + assert_output_limits( 6, 2, 3, 6, 4); + assert_output_limits( 2, 1, 3, 2, 2); + assert_output_limits( + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max()); + + { + common_params base; + base.n_parallel = 4; + base.n_outputs_max_per_seq = 8; + + const auto draft = common_base_params_to_speculative(base); + assert(draft.n_outputs_max == 4); + assert(draft.n_outputs_max_per_seq == 1); + } + printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n"); for (int ex = 0; ex < LLAMA_EXAMPLE_COUNT; ex++) { try { diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index e5ae634cd..4b3d12635 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -80,7 +81,13 @@ struct test_context { std::unordered_map seq_positions; std::unordered_map last_batch_info; - test_context(const test_params & params, std::vector & configs, int32_t n_seq_max = -1) { + test_context( + const test_params & params, + std::vector & configs, + int32_t n_seq_max = -1, + uint32_t n_outputs_max = 0, + uint32_t n_ubatch = 0, + uint32_t n_outputs_max_per_seq = 1) { auto * model = params.model.get(); GGML_ASSERT(model); @@ -89,6 +96,11 @@ struct test_context { llama_context_params cparams = llama_context_default_params(); cparams.n_ctx = 512; cparams.n_batch = 512; + if (n_ubatch > 0) { + cparams.n_ubatch = n_ubatch; + } + cparams.n_outputs_max = n_outputs_max; + cparams.n_outputs_max_per_seq = n_outputs_max_per_seq; cparams.samplers = configs.data(); cparams.n_samplers = configs.size(); cparams.kv_unified = true; @@ -262,6 +274,66 @@ struct test_context { } }; +struct test_single_output_backend_sampler { + bool backend_initialized = false; + uint32_t backend_outputs_max_per_seq = 0; + int backend_apply_count = 0; + int apply_count = 0; +}; + +static const char * test_single_output_backend_sampler_name(const llama_sampler * /*smpl*/) { + return "single-output-backend"; +} + +static void test_single_output_backend_sampler_apply( + llama_sampler * smpl, llama_token_data_array * /*cur_p*/) { + auto * ctx = (test_single_output_backend_sampler *) smpl->ctx; + ctx->apply_count++; +} + +static void test_single_output_backend_sampler_free(llama_sampler * smpl) { + delete (test_single_output_backend_sampler *) smpl->ctx; +} + +static bool test_single_output_backend_sampler_backend_init( + llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_max_per_seq) { + auto * ctx = (test_single_output_backend_sampler *) smpl->ctx; + ctx->backend_outputs_max_per_seq = n_outputs_max_per_seq; + if (n_outputs_max_per_seq > 1) { + return false; + } + ctx->backend_initialized = true; + return true; +} + +static void test_single_output_backend_sampler_backend_apply( + llama_sampler * smpl, ggml_context * /*ctx*/, ggml_cgraph * /*gf*/, llama_sampler_data * /*data*/) { + auto * ctx = (test_single_output_backend_sampler *) smpl->ctx; + ctx->backend_apply_count++; +} + +static llama_sampler_i test_single_output_backend_sampler_i = { + /* .name = */ test_single_output_backend_sampler_name, + /* .accept = */ nullptr, + /* .apply = */ test_single_output_backend_sampler_apply, + /* .reset = */ nullptr, + /* .clone = */ nullptr, + /* .free = */ test_single_output_backend_sampler_free, + /* .backend_init = */ test_single_output_backend_sampler_backend_init, + /* .backend_accept = */ nullptr, + /* .backend_apply = */ test_single_output_backend_sampler_backend_apply, + /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, +}; + +static llama_sampler * test_single_output_backend_sampler_init( + test_single_output_backend_sampler ** sampler_ctx) { + auto * ctx = new test_single_output_backend_sampler; + *sampler_ctx = ctx; + return llama_sampler_init(&test_single_output_backend_sampler_i, ctx); +} + static void test_backend_greedy_sampling(const test_params & params) { const int seq_id = 0; @@ -661,7 +733,7 @@ static void test_backend_multi_sequence_sampling(const test_params & params) { } static void test_backend_dist_sampling(const test_params & params) { - const int seq_id = 189; + const int seq_id = 0; const int32_t seed = 88; struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params(); @@ -1527,43 +1599,398 @@ static void test_backend_cpu_mixed_batch(const test_params & params) { printf("backend-cpu mixed batch test PASSED\n"); } -static void test_backend_max_outputs(const test_params & params) { - const int seq_id = 0; - const int32_t seed = 88; +static void test_backend_multi_output_limit(const test_params & params) { + const llama_seq_id seq_id = 0; - llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params(); - llama_sampler_ptr backend_sampler_chain(llama_sampler_chain_init(backend_chain_params)); - llama_sampler_chain_add(backend_sampler_chain.get(), llama_sampler_init_dist(seed)); - std::vector backend_sampler_configs = {{ seq_id, backend_sampler_chain.get() }}; + llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88)); + std::vector configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 3, 0, 2); - test_context test_ctx(params, backend_sampler_configs); - - llama_batch batch = llama_batch_init(512, 0, 1); - std::string prompt = "Hello"; - - std::vector tokens; - tokens.push_back(llama_vocab_bos(test_ctx.vocab)); - - std::vector prompt_tokens(32); - int n_tokens = llama_tokenize(test_ctx.vocab, prompt.c_str(), prompt.length(), - prompt_tokens.data(), prompt_tokens.size(), - false, false); - for (int i = 0; i < n_tokens; i++) { - tokens.push_back(prompt_tokens[i]); + llama_batch batch = llama_batch_init(3, 0, 1); + for (int i = 0; i < 3; ++i) { + common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), i, { seq_id }, true); } - for (size_t i = 0; i < tokens.size(); i++) { - // set all tokens as output to trigger error - common_batch_add(batch, tokens[i], i, { seq_id }, true); - } - - printf(">>> test_max_outputs expected error start:\n"); + printf(">>> test_backend_multi_output_limit expected error start:\n"); const int ret = llama_decode(test_ctx.ctx.get(), batch); - GGML_ASSERT(ret != 0 && "llama_decode should not succeed multiple outputs per sequence"); - printf("<<< test_max_outputs expected error end.\n"); + GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit"); + printf("<<< test_backend_multi_output_limit expected error end.\n"); + llama_batch_free(batch); - printf("backend max outputs test PASSED\n"); + printf("backend multi-output limit test PASSED\n"); +} + +static void test_backend_multi_sequence_multi_output_dist(const test_params & params) { + const llama_vocab * vocab = llama_model_get_vocab(params.model.get()); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); + const uint32_t seeds[] = { 88, 1337 }; + // reduce the chance that swapped random inputs select the same token + const float temp = 10.0f; + + llama_sampler_ptr chain_0(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_ptr chain_1(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_chain_add(chain_0.get(), llama_sampler_init_temp(temp)); + llama_sampler_chain_add(chain_0.get(), llama_sampler_init_dist(seeds[0])); + llama_sampler_chain_add(chain_1.get(), llama_sampler_init_temp(temp)); + llama_sampler_chain_add(chain_1.get(), llama_sampler_init_dist(seeds[1])); + std::vector configs = { + { 0, chain_0.get() }, + { 1, chain_1.get() }, + }; + test_context test_ctx(params, configs, 2, 4, 0, 2); + + std::vector reference_configs; + test_context reference_ctx(params, reference_configs, 2, 4); + + const llama_token seq_tokens[2][2] = { + { llama_vocab_bos(vocab), llama_vocab_eos(vocab) }, + { llama_vocab_eos(vocab), llama_vocab_bos(vocab) }, + }; + + llama_batch batch = llama_batch_init(4, 0, 1); + for (int pos = 0; pos < 2; ++pos) { + common_batch_add(batch, seq_tokens[0][pos], pos, { 0 }, true); + common_batch_add(batch, seq_tokens[1][pos], pos, { 1 }, true); + } + + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0); + + std::mt19937 reference_rngs[] = { + std::mt19937(seeds[0]), + std::mt19937(seeds[1]), + }; + std::uniform_real_distribution reference_dist(0.0, 1.0); + + for (int i = 0; i < batch.n_tokens; ++i) { + const llama_seq_id seq_id = batch.seq_id[i][0]; + GGML_ASSERT(seq_id == 0 || seq_id == 1); + + llama_sampler * chain = seq_id == 0 ? chain_0.get() : chain_1.get(); + const llama_token backend_token = llama_sampler_sample(chain, test_ctx.ctx.get(), i); + const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i); + const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i); + const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i); + const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i); + const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i); + + GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab); + GGML_ASSERT(sampled_logits != nullptr); + GGML_ASSERT(sampled_probs != nullptr); + GGML_ASSERT(reference_logits != nullptr); + GGML_ASSERT(n_logits == (uint32_t) n_vocab); + GGML_ASSERT(n_probs == (uint32_t) n_vocab); + + float prob_sum = 0.0f; + float cumsum_before = 0.0f; + for (llama_token token = 0; token < n_vocab; ++token) { + const float expected_logit = reference_logits[token] / temp; + const float tolerance = 1e-4f * std::max(1.0f, std::fabs(expected_logit)); + GGML_ASSERT(std::fabs(sampled_logits[token] - expected_logit) <= tolerance); + GGML_ASSERT(std::isfinite(sampled_probs[token])); + GGML_ASSERT(sampled_probs[token] >= 0.0f); + + prob_sum += sampled_probs[token]; + if (token < backend_token) { + cumsum_before += sampled_probs[token]; + } + } + + GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f); + + const float rnd = reference_dist(reference_rngs[seq_id]); + const float cumsum_sampled = cumsum_before + sampled_probs[backend_token]; + GGML_ASSERT(rnd >= cumsum_before - 1e-4f); + GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f); + } + + llama_batch_free(batch); + + printf("backend multi-sequence multi-output dist test PASSED\n"); +} + +static void test_backend_multi_output_dist_transaction(const test_params & params) { + const llama_seq_id seq_id = 0; + const uint32_t seed = 95; + const llama_vocab * vocab = llama_model_get_vocab(params.model.get()); + + llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_chain_add(chain.get(), llama_sampler_init_temp(10.0f)); + llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed)); + std::vector configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 3, 2, 3); + + auto verify_random = [&](int32_t row, float rnd, bool accept = true) { + const llama_token token = accept ? + llama_sampler_sample(chain.get(), test_ctx.ctx.get(), row) : + llama_get_sampled_token_ith(test_ctx.ctx.get(), row); + const float * probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), row); + + GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab)); + GGML_ASSERT(probs != nullptr); + + float cumsum_before = 0.0f; + for (llama_token i = 0; i < token; ++i) { + cumsum_before += probs[i]; + } + + const float cumsum_sampled = cumsum_before + probs[token]; + GGML_ASSERT(rnd >= cumsum_before - 1e-4f); + GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f); + }; + + std::mt19937 rng(seed); + std::uniform_real_distribution dist(0.0, 1.0); + float randoms[3]; + for (float & rnd : randoms) { + rnd = dist(rng); + } + + int32_t pos = 0; + auto decode = [&]() { + llama_batch batch = llama_batch_init(3, 0, 1); + for (int32_t i = 0; i < 3; ++i) { + common_batch_add(batch, llama_vocab_bos(vocab), pos++, { seq_id }, true); + } + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + return batch; + }; + + llama_batch batch = decode(); + verify_random(0, randoms[0], false); + llama_batch_free(batch); + + batch = decode(); + verify_random(0, randoms[0]); + verify_random(1, randoms[1]); + llama_batch_free(batch); + + batch = decode(); + llama_sampler_ptr saved(llama_sampler_clone(chain.get())); + verify_random(0, randoms[2]); + llama_batch_free(batch); + + llama_sampler_copy(saved.get(), chain.get()); + + batch = decode(); + verify_random(0, randoms[2]); + llama_batch_free(batch); + + printf("backend multi-output dist transaction test PASSED\n"); +} + +static void test_backend_multi_output_sampling_chain(const test_params & params) { + const llama_seq_id seq_id = 0; + const uint32_t seed = 88; + const float p = 0.9f; + const float temp = 0.8f; + const float cdf_epsilon = 1e-4f; + const llama_vocab * vocab = llama_model_get_vocab(params.model.get()); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); + const uint32_t k = std::min(512, n_vocab); + const llama_logit_bias bias = { llama_vocab_bos(vocab), -0.1f }; + + auto make_filter_chain = [&]() { + llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_chain_add(result.get(), llama_sampler_init_logit_bias(n_vocab, 1, &bias)); + llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k)); + llama_sampler_chain_add(result.get(), llama_sampler_init_top_p(p, 1)); + llama_sampler_chain_add(result.get(), llama_sampler_init_min_p(0.01f, 1)); + llama_sampler_chain_add(result.get(), llama_sampler_init_temp(temp)); + return result; + }; + + llama_sampler_ptr chain = make_filter_chain(); + llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed)); + std::vector configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 2, 2, 2); + + std::vector reference_configs; + test_context reference_ctx(params, reference_configs, 1, 2, 2); + + llama_sampler_ptr reference_bias(llama_sampler_init_logit_bias(n_vocab, 1, &bias)); + llama_sampler_ptr reference_top_k(llama_sampler_init_top_k(k)); + llama_sampler_ptr reference_top_p(llama_sampler_init_top_p(p, 1)); + llama_sampler_ptr reference_min_p(llama_sampler_init_min_p(0.01f, 1)); + llama_sampler_ptr reference_temp(llama_sampler_init_temp(temp)); + std::vector reference_data(n_vocab); + + auto make_batch = [&](int32_t pos) { + llama_batch batch = llama_batch_init(2, 0, 1); + for (int i = 0; i < 2; ++i) { + common_batch_add(batch, llama_vocab_bos(vocab), pos + i, { seq_id }, true); + } + return batch; + }; + + llama_batch batch = make_batch(0); + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0); + + for (int i = 0; i < batch.n_tokens; ++i) { + const llama_token backend_token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i); + const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i); + const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i); + const llama_token * sampled_candidates = llama_get_sampled_candidates_ith(test_ctx.ctx.get(), i); + const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i); + const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i); + const uint32_t n_candidates = llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i); + const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i); + + GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab); + GGML_ASSERT(sampled_logits != nullptr); + GGML_ASSERT(sampled_probs != nullptr); + GGML_ASSERT(sampled_candidates != nullptr); + GGML_ASSERT(reference_logits != nullptr); + GGML_ASSERT(n_logits == k); + GGML_ASSERT(n_probs == n_logits); + GGML_ASSERT(n_candidates == n_logits); + + for (llama_token token = 0; token < n_vocab; ++token) { + reference_data[token] = { token, reference_logits[token], 0.0f }; + } + + llama_token_data_array reference = { + /* .data = */ reference_data.data(), + /* .size = */ reference_data.size(), + /* .selected = */ LLAMA_TOKEN_NULL, + /* .sorted = */ false, + }; + + llama_sampler_apply(reference_bias.get(), &reference); + llama_sampler_apply(reference_top_k.get(), &reference); + llama_sampler_apply(reference_top_p.get(), &reference); + GGML_ASSERT(reference.size > 0); + + float cdf = 0.0f; + for (size_t j = 0; j < reference.size; ++j) { + cdf += reference.data[j].p; + } + const float cdf_before = cdf - reference.data[reference.size - 1].p; + const float boundary_distance = std::min(std::fabs(cdf_before - p), std::fabs(cdf - p)); + + llama_sampler_apply(reference_min_p.get(), &reference); + llama_sampler_apply(reference_temp.get(), &reference); + + std::unordered_map reference_by_id; + for (size_t j = 0; j < reference.size; ++j) { + reference_by_id.emplace(reference.data[j].id, reference.data[j].logit); + } + size_t n_backend_only = 0; + int32_t sampled_index = -1; + float prob_sum = 0.0f; + + for (uint32_t j = 0; j < n_logits; ++j) { + GGML_ASSERT(sampled_candidates[j] >= 0 && sampled_candidates[j] < n_vocab); + GGML_ASSERT(std::isfinite(sampled_probs[j])); + GGML_ASSERT(sampled_probs[j] >= 0.0f); + prob_sum += sampled_probs[j]; + + if (sampled_candidates[j] == backend_token) { + sampled_index = j; + } + if (!std::isfinite(sampled_logits[j])) { + GGML_ASSERT(std::isinf(sampled_logits[j]) && sampled_logits[j] < 0.0f); + GGML_ASSERT(sampled_probs[j] == 0.0f); + continue; + } + + const auto match = reference_by_id.find(sampled_candidates[j]); + if (match == reference_by_id.end()) { + ++n_backend_only; + continue; + } + + const float tolerance = 1e-4f * std::max(1.0f, std::fabs(match->second)); + GGML_ASSERT(std::fabs(sampled_logits[j] - match->second) <= tolerance); + reference_by_id.erase(match); + } + + const size_t n_reference_only = reference_by_id.size(); + + if (n_backend_only != 0 || n_reference_only != 0) { + GGML_ASSERT(n_backend_only <= 1); + GGML_ASSERT(n_reference_only <= 1); + GGML_ASSERT(boundary_distance <= cdf_epsilon); + } + + GGML_ASSERT(sampled_index >= 0); + GGML_ASSERT(std::isfinite(sampled_logits[sampled_index])); + GGML_ASSERT(sampled_probs[sampled_index] > 0.0f); + GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f); + } + + llama_batch_free(batch); + + batch = make_batch(2); + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + llama_batch_free(batch); + + printf("backend multi-output sampling chain test PASSED\n"); +} + +static void test_backend_multi_output_cpu_suffix(const test_params & params) { + const llama_seq_id seq_id = 0; + const int32_t k = 8; + const llama_vocab * vocab = llama_model_get_vocab(params.model.get()); + + auto make_chain = [&](test_single_output_backend_sampler ** sampler_ctx) { + llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k)); + llama_sampler_chain_add(result.get(), test_single_output_backend_sampler_init(sampler_ctx)); + llama_sampler_chain_add(result.get(), llama_sampler_init_dist(88)); + return result; + }; + + { + test_single_output_backend_sampler * sampler_ctx = nullptr; + llama_sampler_ptr chain = make_chain(&sampler_ctx); + std::vector configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 1, 0, 4); + + llama_batch batch = llama_batch_init(1, 0, 1); + common_batch_add(batch, llama_vocab_bos(vocab), 0, { seq_id }, true); + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + + GGML_ASSERT(sampler_ctx->backend_initialized); + GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 1); + GGML_ASSERT(sampler_ctx->backend_apply_count > 0); + GGML_ASSERT(sampler_ctx->apply_count == 0); + GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), 0) != LLAMA_TOKEN_NULL); + + llama_batch_free(batch); + } + + { + test_single_output_backend_sampler * sampler_ctx = nullptr; + llama_sampler_ptr chain = make_chain(&sampler_ctx); + std::vector configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 2, 0, 0); + + llama_batch batch = llama_batch_init(2, 0, 1); + for (int i = 0; i < 2; ++i) { + common_batch_add(batch, llama_vocab_bos(vocab), i, { seq_id }, true); + } + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + + GGML_ASSERT(!sampler_ctx->backend_initialized); + GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 2); + GGML_ASSERT(sampler_ctx->backend_apply_count == 0); + for (int i = 0; i < batch.n_tokens; ++i) { + GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), i) == LLAMA_TOKEN_NULL); + GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k); + GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k); + const llama_token token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i); + GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab)); + } + GGML_ASSERT(sampler_ctx->apply_count == batch.n_tokens); + + llama_batch_free(batch); + } + + printf("backend multi-output CPU suffix test PASSED\n"); } struct backend_test_case { @@ -1583,7 +2010,11 @@ static const backend_test_case BACKEND_TESTS[] = { { "dist", test_backend_dist_sampling, true }, { "dist_and_cpu", test_backend_dist_sampling_and_cpu, true }, { "set_sampler", test_backend_set_sampler, true }, - { "max_outputs", test_backend_max_outputs, true }, + { "multi_output_limit", test_backend_multi_output_limit, true }, + { "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true }, + { "multi_output_dist_transaction", test_backend_multi_output_dist_transaction, true }, + { "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true }, + { "multi_output_cpu", test_backend_multi_output_cpu_suffix, true }, { "mixed", test_backend_mixed_sampling, true }, { "min_p", test_backend_min_p_sampling, true }, { "cpu_mixed", test_backend_cpu_mixed_batch, true }, diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index df1eb1a20..d727ab632 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -61,6 +61,35 @@ private: std::vector cur; }; +static llama_token sample_dist(llama_sampler * sampler, const std::vector & logits) { + std::vector cur; + for (llama_token token_id = 0; token_id < (llama_token) logits.size(); ++token_id) { + cur.push_back({ token_id, logits[token_id], 0.0f }); + } + + llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false }; + llama_sampler_apply(sampler, &cur_p); + GGML_ASSERT(cur_p.selected >= 0); + GGML_ASSERT((size_t) cur_p.selected < cur_p.size); + return cur_p.data[cur_p.selected].id; +} + +static void test_dist_singleton_rng() { + llama_sampler * singleton = llama_sampler_init_dist(4242); + llama_sampler * control = llama_sampler_init_dist(4242); + + sample_dist(singleton, { 0.0f }); + sample_dist(control, { 0.0f, 0.0f }); + + const std::vector logits(256, 0.0f); + for (int i = 0; i < 4; ++i) { + GGML_ASSERT(sample_dist(singleton, logits) == sample_dist(control, logits)); + } + + llama_sampler_free(singleton); + llama_sampler_free(control); +} + static void test_temp(const std::vector & probs, const std::vector & probs_expected, float temp) { sampler_tester tester(probs, probs_expected); @@ -308,6 +337,8 @@ static void test_perf() { int main(void) { ggml_time_init(); + test_dist_singleton_rng(); + test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 1.0f); test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.0f, 0.0f, 0.0f, 1.0f}, 0.0f); 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/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(): From 0666ad2b2b2452668733729e8b54234f5964643a Mon Sep 17 00:00:00 2001 From: Mario Limonciello Date: Mon, 10 Aug 2026 12:53:12 -0500 Subject: [PATCH 02/10] ci : target ROCm 7.14 for build and release (#25775) * Switch ROCm from 7.2.1 to 7.14 ROCm 7.14 is the first production release using TheRock build system. It can be installed using multi-arch deliverables from wheels, debs, rpms, tarballs or runfiles. Adjust ROCm targets for Linux and Windows to use this instead. * ci: switch all other Windows ROCm jobs to ROCm 7.14 wheels Move the shared windows-setup-rocm composite action from the HIP SDK PRO Edition installer to the multi-arch ROCm wheels (rocm[libraries,devel]). The wheel-install logic that previously lived inline in release.yml is now in the shared action, and both build-cache.yml and release.yml call it. Also migrate the build-cuda-windows.yml hip job to the same wheel-based layout (cache path/key, rocm-sdk environment setup, llvm/bin compiler paths) so it keeps working after the action's contract changed; drop its now-unused ROCm 7.2.1 rocWMMA download and stale include path. --- .github/actions/windows-setup-rocm/action.yml | 28 +- .github/workflows/build-cache.yml | 10 +- .github/workflows/build-cuda-windows.yml | 76 +++-- .github/workflows/release.yml | 317 +++++++++--------- 4 files changed, 233 insertions(+), 198 deletions(-) diff --git a/.github/actions/windows-setup-rocm/action.yml b/.github/actions/windows-setup-rocm/action.yml index fd9f8e5a4..aecbcf14f 100644 --- a/.github/actions/windows-setup-rocm/action.yml +++ b/.github/actions/windows-setup-rocm/action.yml @@ -8,8 +8,26 @@ inputs: runs: using: "composite" steps: - - name: Setup ROCm - uses: ./.github/actions/install-exe - with: - url: https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ inputs.version }}-Win11-For-HIP.exe - args: -install + - name: Install ROCm with Wheels + shell: pwsh + run: | + $ErrorActionPreference = "Stop" + write-host "Setting up Python virtual environment" + + # Create the venv directly at the cache location to avoid relocation issues + New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null + python -m venv C:\TheRock\build\.venv + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + write-host "Upgrading pip" + python -m pip install --upgrade pip + + write-host "Installing ROCm wheels for multi-arch support" + # Install ROCm wheels for multi-arch support (this may take several minutes) + python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ inputs.version }}" + + # Pre-expand the devel tree so it is included in the cache + write-host "Initializing ROCm devel tree" + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + write-host "Completed ROCm wheel installation to C:\TheRock\build" diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index 327f71978..e15fc5e08 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -123,8 +123,8 @@ jobs: runs-on: windows-2022 env: - # Make sure this is in sync with build.yml - HIPSDK_INSTALLER_VERSION: "26.Q1" + # Make sure this is in sync with release.yml and build-cuda-windows.yml + ROCM_VERSION: "7.14.0" steps: - name: Clone @@ -135,11 +135,11 @@ jobs: uses: actions/cache@v5 id: cache-rocm with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + path: C:\TheRock\build + key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: - version: ${{ env.HIPSDK_INSTALLER_VERSION }} + version: ${{ env.ROCM_VERSION }} diff --git a/.github/workflows/build-cuda-windows.yml b/.github/workflows/build-cuda-windows.yml index 367a3a854..ff900802f 100644 --- a/.github/workflows/build-cuda-windows.yml +++ b/.github/workflows/build-cuda-windows.yml @@ -83,7 +83,7 @@ jobs: env: # Make sure this is in sync with build-cache.yml - HIPSDK_INSTALLER_VERSION: "26.Q1" + ROCM_VERSION: "7.14.0" strategy: matrix: @@ -97,36 +97,53 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: Grab rocWMMA package - id: grab_rocwmma - run: | - curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" - 7z x rocwmma.deb - 7z x data.tar - - - name: Use ROCm Installation Cache + - name: Cache ROCm Installation uses: actions/cache@v5 id: cache-rocm with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + path: C:\TheRock\build + key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: - version: ${{ env.HIPSDK_INSTALLER_VERSION }} + version: ${{ env.ROCM_VERSION }} + + - name: Setup ROCm Environment + run: | + $ErrorActionPreference = "Stop" + + # Activate venv from cache or fresh install + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + # Expand the devel tree (idempotent; no-op if already done during install) + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + + # Get ROCm installation paths using the rocm-sdk CLI tool + $rocmPath = (rocm-sdk path --root) + if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" } + $rocmPath = $rocmPath.Trim() + $cmakePath = (rocm-sdk path --cmake).Trim() + $binPath = (rocm-sdk path --bin).Trim() + write-host "ROCm root: $rocmPath" + + echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV + echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV + echo "$binPath" >> $env:GITHUB_PATH + + # Keep venv in PATH for subsequent steps + echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH - name: Verify ROCm id: verify run: | - # Find and test ROCm installation - $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 - if (-not $clangPath) { - Write-Error "ROCm installation not found" - exit 1 - } - & $clangPath.FullName --version + # Test the ROCm clang shipped in the installed wheel + & "${env:HIP_PATH}\lib\llvm\bin\clang.exe" --version - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -134,28 +151,27 @@ jobs: # TODO: this build does not match the build in release.yml, so we use a different cache key # ideally, the builds should match, similar to the CUDA build above so that we would be able # to populate the ccache for the release with manual runs of this workflow - #key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} + #key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} + key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} - name: Build id: cmake_build run: | - $env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path) - $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}" cmake -G "Unix Makefiles" -B build -S . ` - -DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" ` - -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" ` - -DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/" ` + -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` + -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` + -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DCMAKE_BUILD_TYPE=Release ` -DLLAMA_BUILD_BORINGSSL=ON ` - -DROCM_DIR="${env:HIP_PATH}" ` + -DHIP_PATH="${env:HIP_PATH}" ` -DGGML_HIP=ON ` - -DGPU_TARGETS="gfx1100" ` + -DGPU_TARGETS="gfx1100" ` -DGGML_RPC=ON cmake --build build -j ${env:NUMBER_OF_PROCESSORS} - name: ccache-clear uses: ./.github/actions/ccache-clear with: - #key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} + #key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} + key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c668930b0..3a48a57c1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -748,6 +748,132 @@ jobs: path: llama-bin-win-cpu-${{ matrix.arch }}.zip name: llama-bin-win-cpu-${{ matrix.arch }}.zip + windows-rocm: + runs-on: windows-2022 + + strategy: + matrix: + include: + - ROCM_VERSION: "7.14.0" + gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201" + build: x64 + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: ccache + uses: ggml-org/ccache-action@v1.2.21 + with: + key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + evict-old-files: 1d + + - name: Cache ROCm Installation + id: cache-rocm + uses: actions/cache@v5 + with: + path: C:\TheRock\build + key: rocm-wheels-${{ matrix.ROCM_VERSION }}-multi-arch-${{ runner.os }} + + - name: Setup ROCm + if: steps.cache-rocm.outputs.cache-hit != 'true' + uses: ./.github/actions/windows-setup-rocm + with: + version: ${{ matrix.ROCM_VERSION }} + + - name: Setup ROCm Environment + run: | + $ErrorActionPreference = "Stop" + + # Activate venv from cache or fresh install + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + # Expand the devel tree (idempotent; no-op if already done during install) + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + + # Get ROCm installation paths using the rocm-sdk CLI tool + $rocmPath = (rocm-sdk path --root) + if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" } + $rocmPath = $rocmPath.Trim() + $cmakePath = (rocm-sdk path --cmake).Trim() + $binPath = (rocm-sdk path --bin).Trim() + write-host "ROCm root: $rocmPath" + write-host "CMake path: $cmakePath" + write-host "Bin path: $binPath" + + echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV + echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV + echo "$binPath" >> $env:GITHUB_PATH + + # Keep venv in PATH for subsequent steps + echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH + + - name: Build + run: | + mkdir build + cd build + cmake .. ` + -G "Unix Makefiles" ` + -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` + -DCMAKE_BUILD_TYPE=Release ` + -DGGML_BACKEND_DL=ON ` + -DGGML_NATIVE=OFF ` + -DGGML_CPU=ON ` + -DGGML_CPU_ALL_VARIANTS=ON ` + -DGGML_HIP=ON ` + -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` + -DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" ` + -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DHIP_PATH="${env:HIP_PATH}" ` + -DGGML_HIP_ROCWMMA_FATTN=ON ` + -DAMDGPU_TARGETS="${{ matrix.gpu_targets }}" + cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS} + + - name: ccache-clear + uses: ./.github/actions/ccache-clear + with: + key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + + - name: Verify HIP backend was built + run: | + $hipDll = Get-ChildItem -Path build\bin -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue + if (-not $hipDll) { + Write-Host "##[error]ggml-hip*.dll was NOT produced. The HIP backend silently failed to build." + Write-Host "Contents of build\bin:" + Get-ChildItem build\bin | Format-Table -AutoSize + exit 1 + } + Write-Host "HIP backend artifact found:" + $hipDll | Format-Table FullName, Length -AutoSize + + - name: Determine tag name + id: tag + uses: ./.github/actions/get-tag-name + + - name: Get ROCm short version + run: | + $rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.') + echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV + + - name: Pack artifacts + run: | + cp "LICENSE" "build\bin\" + 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\* + + - name: Upload artifacts + uses: actions/upload-artifact@v6 + with: + path: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip + name: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip + windows: needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1168,8 +1294,8 @@ jobs: strategy: matrix: include: - - ROCM_VERSION: "7.2.1" - gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201" + - ROCM_VERSION: "7.14.0" + gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201" build: 'x64' steps: @@ -1201,38 +1327,36 @@ jobs: run: | sudo apt install -y build-essential git cmake wget - - name: Setup Legacy ROCm - if: matrix.ROCM_VERSION == '7.2.1' - id: legacy_env - run: | - sudo mkdir --parents --mode=0755 /etc/apt/keyrings - wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \ - gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null - - sudo tee /etc/apt/sources.list.d/rocm.list << EOF - deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} jammy main - EOF - - sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF - Package: * - Pin: release o=repo.radeon.com - Pin-Priority: 600 - EOF - - sudo apt update - sudo apt-get install -y libssl-dev rocm-hip-sdk - - - name: Setup TheRock - if: matrix.ROCM_VERSION != '7.2.1' + - name: Setup TheRock with Wheels id: therock_env run: | - wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz - mkdir install - tar -xf *.tar.gz -C install - export ROCM_PATH=$(pwd)/install - echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV - echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV - echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV + # Create Python virtual environment + python3 -m venv .venv + source .venv/bin/activate + + # Install ROCm wheels for build + # libraries = HIP runtime and CMake configs needed for linking + # devel = compilers, headers, static libs + python -m pip install --upgrade pip + python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}" + + # Get ROCm installation paths using the rocm-sdk CLI tool + ROCM_PATH=$(rocm-sdk path --root) + CMAKE_PATH=$(rocm-sdk path --cmake) + BIN_PATH=$(rocm-sdk path --bin) + echo "ROCM_PATH=$ROCM_PATH" + echo "CMAKE_PATH=$CMAKE_PATH" + echo "BIN_PATH=$BIN_PATH" + + # Set environment variables + echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV + echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV + echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV + echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV + + # Keep venv activated for subsequent steps + echo "$(pwd)/.venv/bin" >> $GITHUB_PATH - name: Build with native CMake HIP support id: cmake_build @@ -1276,129 +1400,6 @@ jobs: path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz - windows-hip: - needs: [check-release, get-version] - if: ${{ needs.check-release.outputs.should_release == 'true' }} - - runs-on: windows-2022 - - permissions: - actions: write - - env: - HIPSDK_INSTALLER_VERSION: "26.Q1" - - strategy: - matrix: - include: - - name: "radeon" - gpu_targets: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032" - - steps: - - name: Clone - id: checkout - uses: actions/checkout@v6 - - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - - name: Grab rocWMMA package - id: grab_rocwmma - run: | - curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" - 7z x rocwmma.deb - 7z x data.tar - - - name: Cache ROCm Installation - id: cache-rocm - uses: actions/cache@v5 - with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} - - - name: ccache - uses: ggml-org/ccache-action@v1.2.21 - with: - key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - - - name: Install ROCm - if: steps.cache-rocm.outputs.cache-hit != 'true' - id: depends - run: | - $ErrorActionPreference = "Stop" - write-host "Downloading AMD HIP SDK Installer" - Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe" - write-host "Installing AMD HIP SDK" - $proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru - $completed = $proc.WaitForExit(600000) - if (-not $completed) { - Write-Error "ROCm installation timed out after 10 minutes. Killing the process" - $proc.Kill() - exit 1 - } - if ($proc.ExitCode -ne 0) { - Write-Error "ROCm installation failed with exit code $($proc.ExitCode)" - exit 1 - } - write-host "Completed AMD HIP SDK installation" - - - name: Verify ROCm - id: verify - run: | - # Find and test ROCm installation - $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 - if (-not $clangPath) { - Write-Error "ROCm installation not found" - exit 1 - } - & $clangPath.FullName --version - - - name: Build - id: cmake_build - run: | - $env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path) - $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}" - cmake -G "Unix Makefiles" -B build -S . ` - -DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" ` - -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" ` - -DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/ -Wno-ignored-attributes -Wno-nested-anon-types" ` - -DCMAKE_BUILD_TYPE=Release ` - -DGGML_BACKEND_DL=ON ` - -DGGML_NATIVE=OFF ` - -DGGML_CPU=OFF ` - -DGPU_TARGETS="${{ matrix.gpu_targets }}" ` - -DGGML_HIP=ON ` - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} ` - -DLLAMA_BUILD_BORINGSSL=ON - cmake --build build --target ggml-hip -j ${env:NUMBER_OF_PROCESSORS} - md "build\bin\rocblas\library\" - md "build\bin\hipblaslt\library" - cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\" - cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\" - - - name: ccache-clear - uses: ./.github/actions/ccache-clear - with: - key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - - - name: Pack artifacts - id: pack_artifacts - run: | - 7z a -snl llama-bin-win-hip-${{ matrix.name }}-x64.zip .\build\bin\* - - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - path: llama-bin-win-hip-${{ matrix.name }}-x64.zip - name: llama-bin-win-hip-${{ matrix.name }}-x64.zip - ios-xcode: needs: [check-release, get-version] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1572,7 +1573,7 @@ jobs: - windows-cpu - windows-cuda #- windows-sycl - - windows-hip + - windows-rocm - windows-openvino - ubuntu-22-rocm - ubuntu-cpu @@ -1684,7 +1685,7 @@ jobs: - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz) - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz) - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz) - - [Ubuntu x64 (ROCm 7.2)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.2-x64.tar.gz) + - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz) - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz) - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz) - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz) @@ -1702,7 +1703,7 @@ jobs: - [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip) - [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip) - [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip) - - [Windows x64 (HIP)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-hip-radeon-x64.zip) + - [Windows x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-rocm-7.14-x64.zip) **openEuler:** - [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23705) From 689e227db485c6b33d061555e74034c93a867649 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Mon, 10 Aug 2026 11:09:19 -0700 Subject: [PATCH 03/10] opencl: transpose the K tile in local memory for FA prefill kernels (#26428) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 18 +++++++ .../ggml-opencl/kernels/flash_attn_f32_f16.cl | 51 ++++++++++++++++-- .../kernels/flash_attn_f32_q4_0.cl | 47 ++++++++++++++-- .../kernels/flash_attn_f32_q8_0.cl | 53 +++++++++++++++++-- 4 files changed, 156 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index fc0fce0d7..9874ffe8d 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -73,6 +73,7 @@ typedef const void * (*get_adreno_bin_kernel_func_t)( //------------------------------------------------------------------------------ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor); + static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor); static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor); static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); @@ -4629,6 +4630,23 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) { opts += " -D FA_C8_NO_SG_PIN"; } + // Transposed K tile in local memory: the KV rows the QK loop walks together become + // adjacent, so a group of them is ONE 128-bit local read instead of several narrow + // ones. The QK loop is LDS-read-issue-bound (a wrong-math probe that kept every FMA/dp4a + // but removed the LDS reads ran the kernel ~40% faster), so this is worth up to +26% on + // fa=1 prefill. Output is bit-identical -- only the layout moves. + // + // DK <= 128 only. At DK=256 (gemma-3-4b) it measures 1-2% NEGATIVE and reproduces across + // rounds; padding the row stride does not recover it, so the cause is not a simple bank + // conflict and the wider tile does not want this layout. + // + // Default on within that gate; GGML_OPENCL_FA_K_LDS_T=0 restores the row-major tile. + { + const char * e = getenv("GGML_OPENCL_FA_K_LDS_T"); + if ((e == nullptr || e[0] != '0') && cfg->dk <= 128) { + opts += " -D FA_K_LDS_T"; + } + } return opts; } diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl index 6e43ee81e..bf7695a2c 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl @@ -211,7 +211,30 @@ __kernel void FA_TILE_NAME( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); +#ifdef FA_K_LDS_T + // K tile transposed: [dk vec][kv row] instead of [kv row][dk vec]. + // + // The QK loop walks 2 or 4 KV rows at a time against the same dk element. Row-major + // those are DK_VEC half4s apart, so each is its own 64-bit local read. Transposed they + // are adjacent, so a pair is one 128-bit read -- half the LDS issues for the same bytes, + // no extra registers, arithmetic untouched. + // + // This kernel looked like it should be FMA-bound (a half4 mad does ~4 ALU ops per LDS + // read, unlike the 1:1 of the dp4a loop), but it is NOT: a wrong-math probe that kept + // every FMA and removed the LDS reads ran it 38.6% faster (18.92 -> 11.62 ms/op). + // Explicitly 16-byte aligned: FA_LK_PAIR below reads two adjacent half4 as one float4, + // and the element type only obliges the compiler to align this array to 8. The indices + // are even so the offset is a multiple of 16, but the base has to be too, and relying + // on the compiler to over-align it is relying on luck. + __local KV_DATA_TYPE4 l_k[DK_VEC][BLOCK_N] __attribute__((aligned(16))); +#define FA_LK(ROW, C) l_k[C][ROW] + // Two adjacent KV rows as one 128-bit local read (half4 pair == 16 B). j is even and + // BLOCK_N is even, so &l_k[c][j] is 16 B past a 16 B-aligned base. +#define FA_LK_PAIR(C, J) as_half8(*(__local const float4 *)(&l_k[C][J])) +#else __local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC]; +#define FA_LK(ROW, C) l_k[ROW][C] +#endif __local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC]; #if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE) @@ -254,17 +277,17 @@ __kernel void FA_TILE_NAME( #ifdef FA_K_IMG if (use_kv_pad) { const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; } else { const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row; - l_k[row][col] = read_imageh(k_img, k_row_px + col); + FA_LK(row, col) = read_imageh(k_img, k_row_px + col); } #else const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; #endif } else { - l_k[row][col] = (KV_DATA_TYPE4)(0.0h); + FA_LK(row, col) = (KV_DATA_TYPE4)(0.0h); } } for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) { @@ -292,8 +315,15 @@ __kernel void FA_TILE_NAME( FA_UNROLL for (int k = 0; k < SPLIT_DK_VEC; k++) { const ACC_TYPE4 qk = q_priv[k]; +#if defined(FA_K_LDS_T) + // 2 KV rows adjacent in the transposed tile: one 128-bit local read. + const half8 kk = FA_LK_PAIR(dk_off + k, j); + ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(kk.lo); + ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(kk.hi); +#else ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]); ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]); +#endif partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3; partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3; } @@ -359,7 +389,7 @@ __kernel void FA_TILE_NAME( ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); FA_UNROLL for (int k = 0; k < SPLIT_DK_VEC; k++) { - dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc); + dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(FA_LK(j, dk_off + k)), dot_acc); } local_partial[j][tid] = dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3; @@ -452,10 +482,21 @@ __kernel void FA_TILE_NAME( FA_UNROLL for (int k = 0; k < DK_VEC; k++) { const ACC_TYPE4 qk = q_priv[k]; +#if defined(FA_K_LDS_T) + // 4 KV rows adjacent in the transposed tile: two 128-bit local reads + // instead of four 64-bit ones. + const half8 kk01 = FA_LK_PAIR(k, j); + const half8 kk23 = FA_LK_PAIR(k, j + 2); + dot_acc0 = mad(qk, CONVERT_KV_ACC4(kk01.lo), dot_acc0); + dot_acc1 = mad(qk, CONVERT_KV_ACC4(kk01.hi), dot_acc1); + dot_acc2 = mad(qk, CONVERT_KV_ACC4(kk23.lo), dot_acc2); + dot_acc3 = mad(qk, CONVERT_KV_ACC4(kk23.hi), dot_acc3); +#else dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0); dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1); dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2); dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3); +#endif } ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl index 95d215971..48adba4f7 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl @@ -1631,8 +1631,25 @@ __kernel void flash_attn_f32_q4_0( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); #ifdef FA_HAVE_INT_DOT +// Accessors so the staging code is layout-agnostic. +#ifdef FA_K_LDS_T +#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW] +#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW] +#else +#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX] +#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK] +#endif + +#ifdef FA_K_LDS_T + // K tile transposed: the 4 KV rows the QK loop walks together become adjacent, so each + // (block, group) step is ONE 128-bit local read instead of four 32-bit ones. The QK + // loop is LDS-read-issue-bound. + __local uint l_k_packed[DK_Q4_BLOCKS_PREFILL * 8][BLOCK_N]; + __local float l_k_scale [DK_Q4_BLOCKS_PREFILL][BLOCK_N]; +#else __local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8]; __local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL]; +#endif #else __local half4 l_k[BLOCK_N][DK_VEC]; #endif @@ -1660,17 +1677,17 @@ __kernel void flash_attn_f32_q4_0( const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE; const float df = (float) vload_half(0, (const global half *) blk_ptr); const global uchar * qs = (const global uchar *)(blk_ptr + 2); - l_k_scale[row][blk] = df; + FA_K_SCALE(row, blk) = df; uint k_packed[8]; pack_q4_0_nibbles(qs, k_packed); #pragma unroll for (int j = 0; j < 8; ++j) { - l_k_packed[row][blk * 8 + j] = k_packed[j]; + FA_K_PACKED(row, blk * 8 + j) = k_packed[j]; } } else { - l_k_scale[row][blk] = 0.0f; + FA_K_SCALE(row, blk) = 0.0f; #pragma unroll - for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u; } } #else @@ -1760,6 +1777,19 @@ __kernel void flash_attn_f32_q4_0( for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) { const int b = k_blk_base + b_local; int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; +#ifdef FA_K_LDS_T + // 4 KV rows are adjacent in the transposed tile: one 128-bit local + // read per (block, group) instead of four 32-bit ones. + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]); + sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3); + } +#else #pragma unroll for (int g = 0; g < 8; ++g) { const uint qp = q_packed_pf[b_local * 8 + g]; @@ -1768,12 +1798,21 @@ __kernel void flash_attn_f32_q4_0( sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); } +#endif const float qd = q_d_pf[b_local]; const int q_sum = q_sum_pf[b_local]; +#ifdef FA_K_LDS_T + const float4 ks4 = vload4(0, &l_k_scale[b][j]); + s0 += (float)(sum0 - 8 * q_sum) * qd * ks4.s0; + s1 += (float)(sum1 - 8 * q_sum) * qd * ks4.s1; + s2 += (float)(sum2 - 8 * q_sum) * qd * ks4.s2; + s3 += (float)(sum3 - 8 * q_sum) * qd * ks4.s3; +#else s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b]; s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b]; s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b]; s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b]; +#endif } #else ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl index 7e89ed0bd..f50912d21 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl @@ -1393,8 +1393,31 @@ __kernel void flash_attn_f32_q8_0( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); #ifdef FA_HAVE_INT_DOT +// Accessors so the staging code is layout-agnostic. +#ifdef FA_K_LDS_T +#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW] +#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW] +#else +#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX] +#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK] +#endif + +#ifdef FA_K_LDS_T + // K tile transposed: [block*8 + g][kv row] instead of [kv row][block*8 + g]. + // + // The QK loop walks 4 KV rows at a time against the same (b, g), so in the original + // layout those 4 values are BLOCK_N*8 uints apart and cost 4 separate 32-bit local + // reads. Transposed they are adjacent, so they are one 128-bit read -- 4x fewer LDS + // issues for the same bytes and no extra registers. That matters because the QK loop + // is LDS-read-issue-bound: a wrong-math probe that kept every dp4a but cut the LDS + // reads ran the whole kernel 41% faster (18.51 -> 10.91 ms/op), and deleting QK + // outright only reached 10.88 -- i.e. essentially ALL of QK's cost is these reads. + __local uint l_k_packed[DK_Q8_BLOCKS_PREFILL * 8][BLOCK_N]; + __local float l_k_scale [DK_Q8_BLOCKS_PREFILL][BLOCK_N]; +#else __local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8]; __local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL]; +#endif #else __local half4 l_k[BLOCK_N][DK_VEC]; #endif @@ -1427,7 +1450,7 @@ __kernel void flash_attn_f32_q8_0( const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE; const float df = (float) vload_half(0, (const global half *) blk_ptr); const global uchar * qs = (const global uchar *)(blk_ptr + 2); - l_k_scale[row][blk] = df; + FA_K_SCALE(row, blk) = df; #pragma unroll for (int j = 0; j < 8; ++j) { uint k_packed = @@ -1435,12 +1458,12 @@ __kernel void flash_attn_f32_q8_0( ((uint) qs[j*4 + 1]) << 8 | ((uint) qs[j*4 + 2]) << 16 | ((uint) qs[j*4 + 3]) << 24; - l_k_packed[row][blk * 8 + j] = k_packed; + FA_K_PACKED(row, blk * 8 + j) = k_packed; } } else { - l_k_scale[row][blk] = 0.0f; + FA_K_SCALE(row, blk) = 0.0f; #pragma unroll - for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u; } } #else @@ -1556,6 +1579,19 @@ __kernel void flash_attn_f32_q8_0( for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) { const int b = k_blk_base + b_local; int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; +#if defined(FA_K_LDS_T) + // The 4 KV rows are adjacent in the transposed tile, so each (b, g) + // step is ONE 128-bit local read instead of four 32-bit ones. + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]); + sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3); + } +#else #pragma unroll for (int g = 0; g < 8; ++g) { const uint qp = q_packed_pf[b_local * 8 + g]; @@ -1564,11 +1600,20 @@ __kernel void flash_attn_f32_q8_0( sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); } +#endif const float qd = q_d_pf[b_local]; +#ifdef FA_K_LDS_T + const float4 ks4 = vload4(0, &l_k_scale[b][j]); + s0 += (float)sum0 * qd * ks4.s0; + s1 += (float)sum1 * qd * ks4.s1; + s2 += (float)sum2 * qd * ks4.s2; + s3 += (float)sum3 * qd * ks4.s3; +#else s0 += (float)sum0 * qd * l_k_scale[j ][b]; s1 += (float)sum1 * qd * l_k_scale[j+1][b]; s2 += (float)sum2 * qd * l_k_scale[j+2][b]; s3 += (float)sum3 * qd * l_k_scale[j+3][b]; +#endif } #else ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); From 030ebb558a5820b444a8f836ed5cdd46c9b4bd7a Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Tue, 11 Aug 2026 00:02:25 +0530 Subject: [PATCH 04/10] Address review comment of PR 25532 (#26852) --- include/llama.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/include/llama.h b/include/llama.h index c14eb6f50..5a0b66dc2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1324,8 +1324,6 @@ 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); - // copy mutable sampler state without changing dst or its sampling graph bindings - // src and dst must have the same type and configuration 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); From 84f712946729f8517c972da4eb80db810ffe3210 Mon Sep 17 00:00:00 2001 From: Masashi Yoshimura Date: Tue, 11 Aug 2026 13:10:00 +0900 Subject: [PATCH 05/10] ggml-webgpu: fix CI errors from #25025 and #25262 (#26566) * test new flash_attn test * rebase and fix to disable subgrou matrices when max_kv_tile == 0 * delete log output * Add i32 support to cpy and enables the all ops test * restore the non target ci tests * comment out of TODO of build-cpu.yml * fix format --- .../ggml-webgpu/ggml-webgpu-shader-lib.hpp | 22 +++++++++++++++++-- ggml/src/ggml-webgpu/ggml-webgpu.cpp | 5 ++--- ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl | 2 ++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index 35a55ecaf..0604e1c2b 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -2815,11 +2815,25 @@ class ggml_webgpu_shader_lib { key.common.v_direct &= decisions.use_sg_matrix && key.common.v_type == GGML_TYPE_F16; key.use_sg_matrix = decisions.use_sg_matrix; - const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( + uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u, key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.k_direct || key.common.v_direct); - GGML_ASSERT(max_kv_tile > 0); + + // WorkGroup storage size isn't enough for some params with subgroup matrices path (ref. https://github.com/ggml-org/llama.cpp/pull/26566) + if (max_kv_tile == 0) { + GGML_ASSERT(decisions.use_sg_matrix); + // switch to flash_attn_reg_tile path + decisions.use_sg_matrix = false; + decisions.q_tile = GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; + key.common.k_direct = false; + key.common.v_direct = false; + key.use_sg_matrix = false; + max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( + context.wg_mem_limit_bytes, decisions.q_tile, 1u, key.common.head_dim_qk, key.common.head_dim_v, + key.common.has_mask, key.common.k_direct || key.common.v_direct); + GGML_ASSERT(max_kv_tile > 0); + } decisions.kv_tile = decisions.use_sg_matrix ? std::min(max_kv_tile, context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES) : @@ -2993,6 +3007,10 @@ class ggml_webgpu_shader_lib { defines.push_back("SRC_F16"); variant += "_f16"; break; + case GGML_TYPE_I32: + defines.push_back("SRC_I32"); + variant += "_i32"; + break; default: GGML_ABORT("Unsupported src type for cpy shader"); } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index ba4b91695..98c716247 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4283,9 +4283,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_CPY: case GGML_OP_CONT: - supports_op = ((op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && - (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16)) || - (op->type == GGML_TYPE_I32 && src0->type == GGML_TYPE_F32); + supports_op = (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_I32) && + (src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32); break; case GGML_OP_SET: supports_op = src0->type == src1->type && src0->type == op->type && diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl index 67f1dc092..0d0d81ab6 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl @@ -4,6 +4,8 @@ enable f16; #define SRC_TYPE f32 #elif defined(SRC_F16) #define SRC_TYPE f16 +#elif defined(SRC_I32) +#define SRC_TYPE i32 #endif #ifdef DST_F32 From 48d22e295e2b86b47366c16390794f3e05ba970a Mon Sep 17 00:00:00 2001 From: Aldehir Rojas Date: Mon, 10 Aug 2026 23:10:31 -0500 Subject: [PATCH 06/10] common/peg : suppress incomplete escape sequences (#26780) --- common/peg-parser.cpp | 19 +++++++++++++++---- tests/peg-parser/test-json-parser.cpp | 24 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) 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/tests/peg-parser/test-json-parser.cpp b/tests/peg-parser/test-json-parser.cpp index 5dd00115c..ec7c2e668 100644 --- a/tests/peg-parser/test-json-parser.cpp +++ b/tests/peg-parser/test-json-parser.cpp @@ -77,6 +77,30 @@ void test_json_parser(testing &t) { t.assert_equal("result_is_need_more_input", true, result.need_more_input()); }); + // Test need_more_input() parsing - incomplete escape sequence in a string value + t.test("need_more_input() parsing - incomplete escape sequence", [](testing &t) { + auto json = build_peg_parser([](common_peg_parser_builder & p) { return p.json(); }); + + std::vector inputs { + R"({"text": "hello\)", // dangling backslash + R"({"text": "hello\u)", // incomplete unicode escape sequence + R"({"text": "hello\u00)", + }; + + for (const auto & input : inputs) { + t.test(input, [&](testing &t) { + common_peg_parse_context ctx(input, COMMON_PEG_PARSE_FLAG_LENIENT); + + auto result = json.parse(ctx); + + t.assert_equal("result_is_need_more_input", true, result.need_more_input()); + + // the incomplete escape sequence is not part of the partial value + t.assert_equal("result_end", input.find('\\'), result.end); + }); + } + }); + t.test("object member", [](testing &t) { auto parser = build_peg_parser([](common_peg_parser_builder & p) { return p.json_member("name", "\"" + p.chars("[a-z]") + "\""); From 14e78ddef7a2061e7d5a31dce4eb7ee0bcdbc840 Mon Sep 17 00:00:00 2001 From: Junmo Kim Date: Tue, 11 Aug 2026 13:20:17 +0900 Subject: [PATCH 07/10] model : fix SWA not being enabled for EXAONE 4.5 (#26848) * model : fix SWA not being enabled for EXAONE 4.5 load_arch_hparams tests `hparams.n_layer() == 64` before LLM_KV_NEXTN_PREDICT_LAYERS has been read. n_layer() returns n_layer_all - n_layer_nextn and n_layer_nextn defaults to 0, so a GGUF carrying the MTP head (block_count=65, nextn=1) evaluates to 65 and the whole SWA block is skipped. The model type switch further down in the same function reads 64, because by then the key has been loaded. n_swa is still filled in by the unconditional get_key below the block, so llama_model_n_swa() reports 4096 and the logs look correct while only swa_type stays LLAMA_SWA_TYPE_NONE. This affects the official LGAI-EXAONE GGUF release as well. EXAONE 4.0 has no MTP head, so block_count is 64 there and the check matches. * model-loader : skip TENSOR_SKIP tensors in the metadata-only path create_tensor asserts on a null buffer type when building from metadata alone, but buft_for_tensor returns null by design for tensors marked TENSOR_SKIP, which is how architectures with nextn/MTP layers mark theirs. Those models cannot be constructed by llama_model_init_from_user at all. The file-backed path below already returns nullptr for the same tensors, so callers see the same thing either way. * tests : cover exaone4 hparams ordering Builds a synthetic exaone4 model with the layout the shipped EXAONE 4.5 GGUFs use (block_count 65 + nextn 1). The swa_type check is the one that catches the ordering bug; the n_layer_nextn and n_layer() checks only tell a broken fixture apart from a real regression. Fails before the ordering fix with "swa_type is not STANDARD", passes after. * Revert "tests : cover exaone4 hparams ordering" This reverts commit d2f3bafeee591ad691396b2708de4baef3aaf602. * Revert "model-loader : skip TENSOR_SKIP tensors in the metadata-only path" This reverts commit aecb9bc0c7896b52afbc43921a1f572aa7b5e53c. --- src/models/exaone4.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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; From 4801e3c567d5131dd41b387df5f2d4b1370d92be Mon Sep 17 00:00:00 2001 From: Jim Wu Date: Mon, 10 Aug 2026 21:21:32 -0700 Subject: [PATCH 08/10] tests : disable backend sampler hip multi output (#26878) * test-backend-sampler: skip multi_output_sampling_chain on HIP The new multi_output_sampling_chain test uses top_k, whose backend probs path needs CUB (unavailable on HIP), so sampled_probs is null and the test aborts. Add it to the existing HIP skip list alongside the other TOP_K tests. * ci: keep gpu-rocm logs in a per-run dir keyed by GitHub run id The self-hosted gpu-rocm runner can't upload logs to Azure blob (egress firewalled), so a run's logs were wiped by the next run. Write each run's logs to $OUT/run--/ so an Actions run URL maps to its logs. * test-backend-sampler: also skip multi_output_cpu on HIP Like the other TOP_K-based subtests, multi_output_cpu's backend sampler never initializes on HIP (no CUB TOP_K), so it aborts. Add it to the skip list. --------- Co-authored-by: Jim Wu --- ci/run.sh | 8 ++++++++ tests/test-backend-sampler.cpp | 4 +++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/ci/run.sh b/ci/run.sh index f6c7eb0d5..8046df255 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -49,6 +49,14 @@ mkdir -p "$2" OUT=$(realpath "$1") MNT=$(realpath "$2") +# gpu-rocm self-hosted runner can't upload logs to blob; keep each run's logs in +# their own dir keyed by the GitHub run id so an Actions run URL maps to its logs. +if [ -n "${GG_BUILD_ROCM}" ] && [ -n "${GITHUB_RUN_ID}" ]; then + OUT="$OUT/run-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT:-1}" + mkdir -p "$OUT" + echo "ci results dir: $OUT" +fi + rm -f $OUT/*.log rm -f $OUT/*.exit rm -f $OUT/*.md diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 4b3d12635..c23e7248d 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -2105,7 +2105,9 @@ static std::vector collect_tests_to_run(const std::st #ifdef GGML_USE_HIP // TODO: remove this when https://github.com/ggml-org/llama.cpp/pull/26592 is merged if (test.name == "penalties" || test.name == "set_sampler" || - test.name == "mixed" || test.name == "top_p") { + test.name == "mixed" || test.name == "top_p" || + test.name == "multi_output_sampling_chain" || + test.name == "multi_output_cpu") { fprintf(stderr, "Skipping test '%s' on HIP backend (no backend TOP_K support)\n", test.name.c_str()); continue; } From b3df57286c336255fd02c2162b0b5fe3adc2965e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 11 Aug 2026 09:07:13 +0300 Subject: [PATCH 09/10] tests : clean-up server test, use `tests.sh` in ci (#26886) * tests : remove fetch_server_test_models.py * ci : use tests.sh wrapper of pytest --- .github/workflows/server-sanitize.yml | 4 +- .github/workflows/server-self-hosted.yml | 18 ++-- .github/workflows/server.yml | 12 +-- scripts/fetch_server_test_models.py | 105 ----------------------- tools/server/tests/tests.sh | 5 -- 5 files changed, 17 insertions(+), 127 deletions(-) delete mode 100755 scripts/fetch_server_test_models.py diff --git a/.github/workflows/server-sanitize.yml b/.github/workflows/server-sanitize.yml index 0eeefdf88..5d696282c 100644 --- a/.github/workflows/server-sanitize.yml +++ b/.github/workflows/server-sanitize.yml @@ -110,7 +110,7 @@ jobs: source .venv/bin/activate cd tools/server/tests export ${{ matrix.extra_args }} - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow @@ -119,4 +119,4 @@ jobs: source .venv/bin/activate cd tools/server/tests export ${{ matrix.extra_args }} - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh diff --git a/.github/workflows/server-self-hosted.yml b/.github/workflows/server-self-hosted.yml index 249f389ff..675ddbaaa 100644 --- a/.github/workflows/server-self-hosted.yml +++ b/.github/workflows/server-self-hosted.yml @@ -72,7 +72,7 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx1, backend-sampling) id: server_integration_tests_backend_sampling @@ -81,7 +81,7 @@ jobs: cd tools/server/tests source venv/bin/activate export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2) id: server_integration_tests_gpu2 @@ -90,7 +90,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_METAL_DEVICES=2 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2, backend-sampling) id: server_integration_tests_gpu2_backend_sampling @@ -99,7 +99,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_METAL_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh server-cuda: runs-on: [self-hosted, llama-server, Linux, NVIDIA] @@ -132,7 +132,7 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx1, backend-sampling) id: server_integration_tests_backend_sampling @@ -141,7 +141,7 @@ jobs: cd tools/server/tests source venv/bin/activate export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2) id: server_integration_tests_gpu2 @@ -150,7 +150,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_CUDA_DEVICES=2 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2, backend-sampling) id: server_integration_tests_gpu2_backend_sampling @@ -159,7 +159,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_CUDA_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh server-kleidiai: runs-on: ah-ubuntu_22_04-c8g_8x @@ -219,4 +219,4 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh diff --git a/.github/workflows/server.yml b/.github/workflows/server.yml index 5a02cc15a..d5abf1d23 100644 --- a/.github/workflows/server.yml +++ b/.github/workflows/server.yml @@ -104,21 +104,21 @@ jobs: id: server_integration_tests run: | cd tools/server/tests - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow if: ${{ github.event.schedule || github.event.inputs.slow_tests == 'true' }} run: | cd tools/server/tests - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh - name: Tests (Backend sampling) id: server_integration_tests_backend_sampling run: | cd tools/server/tests export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests (Backend sampling) id: server_integration_tests_slow_backend_sampling @@ -126,7 +126,7 @@ jobs: run: | cd tools/server/tests export LLAMA_ARG_BACKEND_SAMPLING=1 - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh windows: runs-on: windows-2025 @@ -170,7 +170,7 @@ jobs: run: | cd tools/server/tests $env:PYTHONIOENCODING = ":replace" - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow @@ -178,4 +178,4 @@ jobs: run: | cd tools/server/tests $env:SLOW_TESTS = "1" - pytest -v -x + ./tests.sh diff --git a/scripts/fetch_server_test_models.py b/scripts/fetch_server_test_models.py deleted file mode 100755 index f43d1f63c..000000000 --- a/scripts/fetch_server_test_models.py +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env python -''' - This script fetches all the models used in the server tests. - - This is useful for slow tests that use larger models, to avoid them timing out on the model downloads. - - It is meant to be run from the root of the repository. - - Example: - python scripts/fetch_server_test_models.py - ( cd tools/server/tests && ./tests.sh -v -x -m slow ) -''' -import ast -import glob -import logging -import os -from typing import Generator -from pydantic import BaseModel -from typing import Optional -import subprocess - - -class HuggingFaceModel(BaseModel): - hf_repo: str - hf_file: Optional[str] = None - - class Config: - frozen = True - - -def collect_hf_model_test_parameters(test_file) -> Generator[HuggingFaceModel, None, None]: - try: - with open(test_file) as f: - tree = ast.parse(f.read()) - except Exception as e: - logging.error(f'collect_hf_model_test_parameters failed on {test_file}: {e}') - return - - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef): - for dec in node.decorator_list: - if isinstance(dec, ast.Call) and isinstance(dec.func, ast.Attribute) and dec.func.attr == 'parametrize': - param_names = ast.literal_eval(dec.args[0]).split(",") - if "hf_repo" not in param_names: - continue - - raw_param_values = dec.args[1] - if not isinstance(raw_param_values, ast.List): - logging.warning(f'Skipping non-list parametrize entry at {test_file}:{node.lineno}') - continue - - hf_repo_idx = param_names.index("hf_repo") - hf_file_idx = param_names.index("hf_file") if "hf_file" in param_names else None - - for t in raw_param_values.elts: - if not isinstance(t, ast.Tuple): - logging.warning(f'Skipping non-tuple parametrize entry at {test_file}:{node.lineno}') - continue - yield HuggingFaceModel( - hf_repo=ast.literal_eval(t.elts[hf_repo_idx]), - hf_file=ast.literal_eval(t.elts[hf_file_idx]) if hf_file_idx is not None else None) - - -if __name__ == '__main__': - logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') - - models = sorted(list(set([ - model - for test_file in glob.glob('tools/server/tests/unit/test_*.py') - for model in collect_hf_model_test_parameters(test_file) - ])), key=lambda m: (m.hf_repo, m.hf_file)) - - logging.info(f'Found {len(models)} models in parameterized tests:') - for m in models: - logging.info(f' - {m.hf_repo} / {m.hf_file}') - - cli_path = os.environ.get( - 'LLAMA_CLI_BIN_PATH', - os.path.join( - os.path.dirname(__file__), - '../build/bin/Release/llama-cli.exe' if os.name == 'nt' else '../build/bin/llama-cli')) - - for m in models: - if '<' in m.hf_repo or (m.hf_file is not None and '<' in m.hf_file): - continue - if m.hf_file is not None and '-of-' in m.hf_file: - logging.warning(f'Skipping model at {m.hf_repo} / {m.hf_file} because it is a split file') - continue - logging.info(f'Using llama-cli to ensure model {m.hf_repo}/{m.hf_file} was fetched') - cmd = [ - cli_path, - '-hfr', m.hf_repo, - *([] if m.hf_file is None else ['-hff', m.hf_file]), - '-n', '1', - '-p', 'Hey', - '--no-warmup', - '--log-disable', - '-st'] - if m.hf_file != 'tinyllamas/stories260K.gguf' and 'Mistral-Nemo' not in m.hf_repo: - cmd += ('-fa', 'on') - try: - subprocess.check_call(cmd) - except subprocess.CalledProcessError: - logging.error(f'Failed to fetch model at {m.hf_repo} / {m.hf_file} with command:\n {" ".join(cmd)}') - exit(1) 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 From 153d324bcf86d220b235ca010eeb11213f32b5d1 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Tue, 11 Aug 2026 08:20:46 +0200 Subject: [PATCH 10/10] llama: add default load-mode auto, which avoids mmap on iGPUs (#26081) * llama: add new default load-mode auto which picks mmap unless a non-Metal iGPU is used * Update ggml/src/ggml-hexagon/ggml-hexagon.cpp Co-authored-by: Max Krasnyansky * set mmap_support to false on OpenCL backend * fix order of load modes * use -1 for auto * resolve load mode auto earlier to correctly pick gpu host or cpu memory * add load mode auto to llama-bench * bump virtgpu api version, regenerate docs --------- Co-authored-by: Piotr Wilkin (ilintar) Co-authored-by: Max Krasnyansky Co-authored-by: Georgi Gerganov --- common/arg.cpp | 6 ++++-- common/common.h | 2 +- ggml/include/ggml-backend.h | 2 ++ ggml/src/ggml-backend-meta.cpp | 2 ++ ggml/src/ggml-blas/ggml-blas.cpp | 1 + ggml/src/ggml-cann/ggml-cann.cpp | 1 + ggml/src/ggml-cpu/ggml-cpu.cpp | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 1 + ggml/src/ggml-et/ggml-et.cpp | 1 + ggml/src/ggml-hexagon/ggml-hexagon.cpp | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 1 + ggml/src/ggml-opencl/ggml-opencl.cpp | 1 + ggml/src/ggml-openvino/ggml-openvino.cpp | 1 + ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + ggml/src/ggml-sycl/ggml-sycl.cpp | 1 + .../backend/backend-dispatched-device.cpp | 1 + .../backend/shared/api_remoting.h | 2 +- .../ggml-virtgpu/ggml-backend-buffer-type.cpp | 4 ++-- ggml/src/ggml-virtgpu/ggml-backend-device.cpp | 2 +- .../ggml-virtgpu/virtgpu-forward-device.cpp | 4 +++- ggml/src/ggml-virtgpu/virtgpu-forward.gen.h | 3 ++- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1 + ggml/src/ggml-webgpu/ggml-webgpu.cpp | 1 + ggml/src/ggml-zdnn/ggml-zdnn.cpp | 3 ++- ggml/src/ggml-zendnn/ggml-zendnn.cpp | 3 ++- include/llama.h | 11 ++++++----- src/llama-model-loader.cpp | 2 +- src/llama-model.cpp | 19 +++++++++++++++++-- src/llama.cpp | 11 +++++++---- tools/cli/README.md | 2 +- tools/completion/README.md | 2 +- tools/llama-bench/llama-bench.cpp | 8 +++++--- tools/server/README.md | 2 +- 33 files changed, 75 insertions(+), 29 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index c37d5cd0a..cb314eee7 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2605,14 +2605,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.h b/common/common.h index 878534dcc..d485d4fb4 100644 --- a/common/common.h +++ b/common/common.h @@ -473,7 +473,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/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-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5f51ea3bb..ffa361af4 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2815,6 +2815,7 @@ static void ggml_backend_cann_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ true, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 16cc5116c..c0c9aa3cf 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 05e8d7f73..1d4f4dfbd 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4801,6 +4801,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-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b30209095..e8482f734 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1646,6 +1646,7 @@ static void ggml_backend_et_device_get_props(ggml_backend_dev_t dev, struct ggml /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index bdb8af082..f80c60a50 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3930,6 +3930,7 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ (bool) opt_hostbuf, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ false, }; } 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-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 9874ffe8d..19dca4ced 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -10795,6 +10795,7 @@ static void ggml_backend_opencl_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ false, }; } diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index 0e7501fef..dfe80e6c8 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -763,6 +763,7 @@ static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_ /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 17c53a5f0..e9de0d0aa 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1881,6 +1881,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-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 18d58782e..8a19f648b 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5649,6 +5649,7 @@ static void ggml_backend_sycl_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ events, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp index c7acb8b51..87872df1c 100644 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp +++ b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp @@ -111,6 +111,7 @@ uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_ apir_encode_bool_t(enc, &props.caps.host_buffer); apir_encode_bool_t(enc, &props.caps.buffer_from_host_ptr); apir_encode_bool_t(enc, &props.caps.events); + apir_encode_bool_t(enc, &props.caps.mmap_support); return 0; } diff --git a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h index 6bf97e8a3..a5ef3ea47 100644 --- a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h +++ b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h @@ -7,7 +7,7 @@ #include #define APIR_PROTOCOL_MAJOR 0 -#define APIR_PROTOCOL_MINOR 1 +#define APIR_PROTOCOL_MINOR 2 #define APIR_HANDSHAKE_MAGIC 0xab1e diff --git a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp index 8fa20ff43..d5bdc993b 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp @@ -11,9 +11,9 @@ static ggml_backend_buffer_t ggml_backend_remoting_buffer_type_alloc_buffer(ggml context->gpu = gpu; - bool async__unused, host_buffer__unused, events__unused; + bool async__unused, host_buffer__unused, events__unused, mmap_support__unused; bool buffer_from_host_ptr; - apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused); + apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused, &mmap_support__unused); if (buffer_from_host_ptr) { context->apir_context = apir_device_buffer_from_ptr(gpu, size, size); diff --git a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp index a978812cd..987ce9dd1 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp @@ -65,7 +65,7 @@ static void ggml_backend_remoting_device_get_props(ggml_backend_dev_t dev, ggml_ virtgpu * gpu = DEV_TO_GPU(dev); apir_device_get_props(gpu, &props->caps.async, &props->caps.host_buffer, &props->caps.buffer_from_host_ptr, - &props->caps.events); + &props->caps.events, &props->caps.mmap_support); props->caps.buffer_from_host_ptr = false; props->caps.async = false; diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp index 9f513c138..864264f21 100644 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp +++ b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp @@ -144,7 +144,8 @@ void apir_device_get_props(virtgpu * gpu, bool * async, bool * host_buffer, bool * buffer_from_host_ptr, - bool * events) { + bool * events, + bool * mmap_support) { apir_encoder * encoder; apir_decoder * decoder; ApirForwardReturnCode ret; @@ -157,6 +158,7 @@ void apir_device_get_props(virtgpu * gpu, apir_decode_bool_t(decoder, host_buffer); apir_decode_bool_t(decoder, buffer_from_host_ptr); apir_decode_bool_t(decoder, events); + apir_decode_bool_t(decoder, mmap_support); remote_call_finish(gpu, encoder, decoder); diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h index 44b0ad1ff..da28aa5f9 100644 --- a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h +++ b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h @@ -13,7 +13,8 @@ void apir_device_get_props(struct virtgpu * gpu, bool * async, bool * host_buffer, bool * buffer_from_host_ptr, - bool * events); + bool * events, + bool * mmap_support); apir_buffer_context_t apir_device_buffer_from_ptr(struct virtgpu * gpu, size_t size, size_t max_tensor_size); /* buffer-type */ diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a923755f9..45fa97f81 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17891,6 +17891,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 = */ true, + /* .mmap_support = */ !ctx->is_integrated_gpu, }; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 98c716247..6741752b3 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -3942,6 +3942,7 @@ static void ggml_backend_webgpu_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp index 639b818d1..4007ac9df 100644 --- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp +++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp @@ -487,7 +487,8 @@ static void ggml_backend_zdnn_device_get_props(ggml_backend_dev_t dev, ggml_back /* .async = */ false, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, - /* .events = */ false + /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp index e6a9b51b7..ec7ce2331 100644 --- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp +++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp @@ -654,7 +654,8 @@ static void ggml_backend_zendnn_device_get_props(ggml_backend_dev_t dev, struct /* .async = */ false, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, - /* .events = */ false + /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/include/llama.h b/include/llama.h index 5a0b66dc2..bfef0e1d1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -203,11 +203,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); diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 3d50f8a1c..51ba05439 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 3bf3a22f2..cc1917b6d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1271,8 +1271,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); @@ -2452,7 +2467,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.cpp b/src/llama.cpp index d6e0bbfef..94c8f60e0 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -48,6 +48,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: @@ -63,11 +65,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/tools/cli/README.md b/tools/cli/README.md index b42b2e534..880c4a540 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -58,7 +58,7 @@ | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/completion/README.md b/tools/completion/README.md index 552a0c6ab..c2e52ac06 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -141,7 +141,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index c17a27b54..7c495afe2 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -384,7 +384,7 @@ static const cmd_params cmd_params_defaults = { /* n_gpu_layers */ { -1 }, /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, - /* load_mode */ { LLAMA_LOAD_MODE_MMAP }, + /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -459,7 +459,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); - printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); @@ -764,7 +764,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { std::vector modes; for (const auto & m : p) { llama_load_mode mode; - if (m == "none") { + if (m == "auto") { + mode = LLAMA_LOAD_MODE_AUTO; + } else if (m == "none") { mode = LLAMA_LOAD_MODE_NONE; } else if (m == "mmap") { mode = LLAMA_LOAD_MODE_MMAP; diff --git a/tools/server/README.md b/tools/server/README.md index 6927caddb..a2ab872b4 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -75,7 +75,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit |