diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 36f1e0cd50f..1e87312d207 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -100,6 +100,8 @@ add_library(${TARGET} reasoning-budget.h sampling.cpp sampling.h + speculative-prefill.cpp + speculative-prefill.h speculative.cpp speculative.h subproc.cpp diff --git a/common/arg.cpp b/common/arg.cpp index 4469612cd5b..6b95cadd809 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -361,6 +361,7 @@ static bool spec_types_is_default(const common_params & params) { common_models_handler common_models_handler_init(const common_params & params, llama_example curr_ex) { common_download_hf_plan plan; common_download_hf_plan plan_spec; + common_download_hf_plan plan_prefill; common_download_opts opts; const bool spec_type_draft_mtp = std::find(params.speculative.types.begin(), @@ -413,7 +414,11 @@ common_models_handler common_models_handler_init(const common_params & params, l plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec); } - return common_models_handler{plan, plan_spec, opts}; + if (!params.speculative.prefill.model.hf_repo.empty()) { + plan_prefill = common_download_get_hf_plan(params.speculative.prefill.model, opts); + } + + return common_models_handler{plan, plan_spec, plan_prefill, opts}; } bool common_models_handler_is_preset_repo(const common_models_handler & handler) { @@ -461,8 +466,9 @@ static std::vector build_url_tasks(const common_params_mod void common_models_handler_apply(common_models_handler & handler, common_params & params, common_download_callback * callback) { std::vector tasks; - auto & plan = handler.plan; - auto & plan_spec = handler.plan_spec; + auto & plan = handler.plan; + auto & plan_spec = handler.plan_spec; + auto & plan_prefill = handler.plan_prefill; auto opts = handler.opts; // copy opts.callback = callback; @@ -478,6 +484,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params handle_url(params.model); handle_url(params.mmproj); handle_url(params.speculative.draft.mparams); + handle_url(params.speculative.prefill.model); // optionally, if docker repo is set, resolve it if (!params.model.docker_repo.empty()) { @@ -513,6 +520,13 @@ void common_models_handler_apply(common_models_handler & handler, common_params tasks.push_back(task); had_spec_url = true; } + if (!params.speculative.prefill.model.url.empty()) { + common_download_task task; + task.url = params.speculative.prefill.model.url; + task.local_path = params.speculative.prefill.model.path; + task.opts = opts; + tasks.push_back(task); + } // handle hf_plan tasks auto add_tasks = [&opts, &tasks](const hf_cache::hf_files & model_files, @@ -626,6 +640,11 @@ void common_models_handler_apply(common_models_handler & handler, common_params had_spec_url = true; } + // handle plan_prefill (e.g. --spec-prefill-hf) + if (!plan_prefill.model_files.empty()) { + add_tasks(plan_prefill.model_files, plan_prefill.primary, params.speculative.prefill.model); + } + if (!plan.model_files.empty()) { add_tasks(plan.model_files, plan.primary, params.model); } @@ -4245,6 +4264,111 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.mparams.hf_file = value; // will be used if --spec-draft-hf is set } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL")); + add_opt(common_arg( + {"--spec-prefill", "--speculative-prefill"}, + "enable speculative prefill using draft model to filter prompt tokens", + [](common_params & params) { + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL")); + add_opt(common_arg( + {"--spec-prefill-draft-model", "--spec-prefill-model", "-mpd", "--speculative-prefill-model", "--speculative-prefill-draft-model"}, "FNAME", + "draft model for speculative prefill (default: unused)", + [](common_params & params, const std::string & value) { + params.speculative.prefill.model.path = value; + params.speculative.prefill.model.hf_file = value; // will be used if --spec-prefill-hf is set + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_MODEL")); + add_opt(common_arg( + {"--spec-prefill-draft-hf", "--spec-prefill-hf", "-hfpd", "--speculative-prefill-hf", "--speculative-prefill-draft-hf"}, "/[:quant]", + "Hugging Face model repository for speculative prefill draft model (default: unused)", + [](common_params & params, const std::string & value) { + params.speculative.prefill.model.hf_repo = value; + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_HF")); + add_opt(common_arg( + {"--spec-prefill-draft-ngl", "--spec-prefill-ngl", "-nglpd", "--speculative-prefill-ngl", "--speculative-prefill-draft-ngl"}, "N", + string_format("max. number of speculative prefill draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)", + params.speculative.prefill.n_gpu_layers == -1 ? "auto" : "all"), + [](common_params & params, const std::string & value) { + if (value == "auto") { + params.speculative.prefill.n_gpu_layers = -1; + } else if (value == "all") { + params.speculative.prefill.n_gpu_layers = -2; + } else { + params.speculative.prefill.n_gpu_layers = std::stoi(value); + } + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: no usable GPU found, --spec-prefill-ngl option will be ignored\n"); + fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n"); + fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_N_GPU_LAYERS")); + add_opt(common_arg( + {"--spec-prefill-draft-device", "--spec-prefill-device", "-devpd", "--speculative-prefill-device", "--speculative-prefill-draft-device"}, "", + "comma-separated list of devices to use for offloading the speculative prefill draft model (none = don't offload)\n" + "use --list-devices to see a list of available devices", + [](common_params & params, const std::string & value) { + params.speculative.prefill.devices = parse_device_list(value); + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_DEVICE")); + add_opt(common_arg( + {"--spec-prefill-draft-ctx", "--spec-prefill-ctx", "--spec-prefill-ctx-size", "--spec-prefill-max-ctx", "-cpd", "--speculative-prefill-ctx", "--speculative-prefill-max-ctx"}, "N", + string_format("context size for speculative prefill draft model (default: %d, 0 = draft training limit or main context)", params.speculative.prefill.n_ctx), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("spec-prefill context size must be >= 0"); + } + params.speculative.prefill.n_ctx = value; + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_CTX_SIZE")); + add_opt(common_arg( + {"--spec-prefill-p", "--spec-prefill-percentage"}, "P", + string_format("fraction of prompt tokens to retain during speculative prefill (default: %.2f)", (double) params.speculative.prefill.percentage), + [](common_params & params, const std::string & value) { + const float val = std::stof(value); + if (val <= 0.0f || val > 1.0f) { + throw std::invalid_argument("spec-prefill percentage must be between 0.0 (exclusive) and 1.0 (inclusive)"); + } + params.speculative.prefill.enabled = true; + params.speculative.prefill.percentage = val; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_P")); + add_opt(common_arg( + {"--spec-prefill-chunk", "--spec-prefill-chunk-size"}, "N", + string_format("chunk grouping size for speculative prefill (default: %d, 0 to disable)", params.speculative.prefill.chunk_size), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("spec-prefill chunk size must be >= 0"); + } + params.speculative.prefill.chunk_size = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-prefill-lookahead", "--spec-prefill-lah"}, "N", + string_format("number of lookahead decode steps on draft model for attention estimation (default: %d)", params.speculative.prefill.look_ahead_cnt), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("spec-prefill lookahead count must be >= 1"); + } + params.speculative.prefill.look_ahead_cnt = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-prefill-pool-kernel"}, "N", + string_format("1D average pooling kernel size for attention smoothing (default: %d, 0 to disable)", params.speculative.prefill.pool_kernel_size), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("spec-prefill pool kernel size must be >= 0"); + } + params.speculative.prefill.pool_kernel_size = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"--spec-type"}, common_speculative_all_types_str(), string_format("comma-separated list of types of speculative decoding to use (default: %s)\n", diff --git a/common/arg.h b/common/arg.h index 421bc295fc2..caa7e9f5219 100644 --- a/common/arg.h +++ b/common/arg.h @@ -138,6 +138,7 @@ void common_params_add_preset_options(std::vector & args); struct common_models_handler { common_download_hf_plan plan; common_download_hf_plan plan_spec; + common_download_hf_plan plan_prefill; common_download_opts opts; }; diff --git a/common/common.h b/common/common.h index 4e9448bb106..298c5bab7ed 100644 --- a/common/common.h +++ b/common/common.h @@ -367,6 +367,20 @@ struct common_params_speculative_ngram_cache { std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding }; +struct common_params_speculative_prefill { + bool enabled = false; // enable speculative prefill + common_params_model model; // draft model for speculative prefill + int32_t n_ctx = 0; // context size for draft model (0 = default/target context size) + int32_t n_gpu_layers = -1; // max draft model layers to store in VRAM (-1 - use default) + std::vector devices; // devices to use for offloading the draft model + float percentage = 0.3f; // fraction of prompt tokens to retain (0.0 < p <= 1.0) + int32_t chunk_size = 32; // chunk grouping size (0 to disable chunking) + int32_t look_ahead_cnt = 8; // lookahead decode steps on draft model + int32_t pool_kernel_size = 13; // 1D average pooling kernel size for smoothing + bool keep_bos = true; // preserve first token (BOS) + bool keep_last = true; // preserve last token / tail chunk +}; + struct common_params_speculative { std::vector types = { COMMON_SPECULATIVE_TYPE_NONE }; @@ -383,6 +397,8 @@ struct common_params_speculative { common_params_speculative_ngram_cache ngram_cache; + common_params_speculative_prefill prefill; + bool has_dft() const { return !draft.mparams.empty(); } @@ -391,6 +407,9 @@ struct common_params_speculative { return synth_len != -1.0 || !synth_rates.empty(); } + bool has_prefill() const { + return prefill.enabled && !prefill.model.empty(); + } uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; diff --git a/common/speculative-prefill.cpp b/common/speculative-prefill.cpp new file mode 100644 index 00000000000..ab20ad08f1e --- /dev/null +++ b/common/speculative-prefill.cpp @@ -0,0 +1,337 @@ +#include "speculative-prefill.h" + +#include "common.h" +#include "ggml.h" +#include "ggml-backend.h" +#include "llama.h" +#include "log.h" +#include "sampling.h" + +#include +#include +#include +#include +#include + +#define SPF_DBG(fmt, ...) LOG_DBG("spec-prefill: " fmt, __VA_ARGS__) +#define SPF_INF(fmt, ...) LOG_INF("spec-prefill: " fmt, __VA_ARGS__) +#define SPF_ERR(fmt, ...) LOG_ERR("spec-prefill: " fmt, __VA_ARGS__) + +std::vector common_speculative_prefill_compute_importance( + const std::vector> & attn_layers, + int32_t n_prompt, + int32_t n_heads, + int32_t pool_kernel_size) { + if (attn_layers.empty() || n_prompt <= 0 || n_heads <= 0) { + return std::vector(std::max(0, n_prompt), 1.0f); + } + + std::vector max_importance(n_prompt, 0.0f); + std::vector smoothed(n_prompt, 0.0f); + + const int32_t pad = pool_kernel_size > 0 ? (pool_kernel_size / 2) : 0; + + for (const auto & layer_attn : attn_layers) { + if ((int32_t) layer_attn.size() < n_heads * n_prompt) { + continue; + } + + for (int32_t h = 0; h < n_heads; ++h) { + const float * head_attn = layer_attn.data() + (size_t) h * n_prompt; + + if (pool_kernel_size > 1) { + for (int32_t i = 0; i < n_prompt; ++i) { + const int32_t start = std::max(0, i - pad); + const int32_t end = std::min(n_prompt - 1, i + pad); + + float sum = 0.0f; + for (int32_t j = start; j <= end; ++j) { + sum += head_attn[j]; + } + smoothed[i] = sum / (float) (end - start + 1); + } + } else { + std::copy(head_attn, head_attn + n_prompt, smoothed.begin()); + } + + for (int32_t i = 0; i < n_prompt; ++i) { + if (smoothed[i] > max_importance[i]) { + max_importance[i] = smoothed[i]; + } + } + } + } + + return max_importance; +} + +std::vector common_speculative_prefill_select_indices( + const std::vector & importance_scores, + const common_params_speculative_prefill & params) { + const int32_t n_prompt = (int32_t) importance_scores.size(); + if (n_prompt <= 0) { + return {}; + } + + const float percentage = std::max(0.01f, std::min(1.0f, params.percentage)); + std::vector selected; + + if (params.chunk_size > 0 && n_prompt > params.chunk_size) { + const int32_t chunk_size = params.chunk_size; + const int32_t n_chunks = (n_prompt + chunk_size - 1) / chunk_size; + + std::vector> chunk_scores(n_chunks); + + for (int32_t c = 0; c < n_chunks; ++c) { + const int32_t start = c * chunk_size; + const int32_t end = std::min(n_prompt, (c + 1) * chunk_size); + + float sum = 0.0f; + for (int32_t i = start; i < end; ++i) { + sum += importance_scores[i]; + } + chunk_scores[c] = { sum / (float) (end - start), c }; + } + + const int32_t n_keep_chunks = std::max(1, (int32_t) std::ceil(n_chunks * percentage)); + + std::sort(chunk_scores.begin(), chunk_scores.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + + for (int32_t k = 0; k < n_keep_chunks && k < n_chunks; ++k) { + const int32_t c = chunk_scores[k].second; + const int32_t start = c * chunk_size; + const int32_t end = std::min(n_prompt, (c + 1) * chunk_size); + + for (int32_t i = start; i < end; ++i) { + selected.push_back(i); + } + } + } else { + const int32_t n_keep = std::max(1, (int32_t) std::ceil(n_prompt * percentage)); + + std::vector> token_scores(n_prompt); + for (int32_t i = 0; i < n_prompt; ++i) { + token_scores[i] = { importance_scores[i], i }; + } + + std::sort(token_scores.begin(), token_scores.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + + for (int32_t k = 0; k < n_keep; ++k) { + selected.push_back(token_scores[k].second); + } + } + + if (params.keep_bos) { + selected.push_back(0); + } + + if (params.keep_last) { + if (params.chunk_size > 0 && n_prompt > params.chunk_size) { + const int32_t start = ((n_prompt - 1) / params.chunk_size) * params.chunk_size; + for (int32_t i = start; i < n_prompt; ++i) { + selected.push_back(i); + } + } else { + selected.push_back(n_prompt - 1); + } + } + + std::sort(selected.begin(), selected.end()); + selected.erase(std::unique(selected.begin(), selected.end()), selected.end()); + + return selected; +} + +struct cb_attn_collector_data { + int32_t n_prompt = 0; + std::vector> layer_attns; +}; + +static bool cb_collect_attn(ggml_tensor * t, bool ask, void * user_data) { + if (ask) { + return strncmp(t->name, "kq_soft_max", 11) == 0; + } + + if (strncmp(t->name, "kq_soft_max", 11) != 0) { + return true; + } + + auto * data = (cb_attn_collector_data *) user_data; + if (data == nullptr || data->n_prompt <= 0) { + return true; + } + + if (t->type != GGML_TYPE_F32) { + return true; + } + + const int32_t n_past_total = (int32_t) t->ne[0]; + const int32_t n_heads = (int32_t) t->ne[2]; + const int32_t n_prompt = data->n_prompt; + + if (t->nb[0] != sizeof(float) || n_heads <= 0 || n_past_total < n_prompt) { + return true; + } + + std::vector layer_data((size_t) n_heads * n_prompt); + + for (int32_t h = 0; h < n_heads; ++h) { + float * dst = layer_data.data() + (size_t) h * n_prompt; + ggml_backend_tensor_get(t, dst, (size_t) h * t->nb[2], (size_t) n_prompt * sizeof(float)); + } + + data->layer_attns.push_back(std::move(layer_data)); + + return true; +} + +common_speculative_prefill_result common_speculative_prefill_execute( + llama_context * ctx_dft, + common_sampler * smpl_dft, + const std::vector & prompt, + llama_seq_id seq_id, + const common_params_speculative_prefill & params) { + common_speculative_prefill_result res; + res.n_prompt_orig = (int32_t) prompt.size(); + + if (prompt.empty()) { + return res; + } + + if (!params.enabled || params.percentage >= 1.0f || (int32_t) prompt.size() <= params.chunk_size) { + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.importance_scores.assign(prompt.size(), 1.0f); + return res; + } + + const int32_t n_ctx_dft = llama_n_ctx(ctx_dft); + const int32_t lookahead = std::max(1, params.look_ahead_cnt); + + if ((int32_t) prompt.size() + lookahead > n_ctx_dft) { + SPF_INF("prompt size (%d) + lookahead (%d) exceeds draft context (%d); skipping speculative prefill\n", + (int32_t) prompt.size(), lookahead, n_ctx_dft); + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.importance_scores.assign(prompt.size(), 1.0f); + return res; + } + + const auto t_start = ggml_time_us(); + + const llama_model * model_dft = llama_get_model(ctx_dft); + const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); + const int32_t n_heads_dft = llama_model_n_head(model_dft); + + // 1. evaluate full prompt on draft model + { + const int32_t n_batch_dft = llama_n_batch(ctx_dft); + llama_batch batch_prompt = llama_batch_init(std::min((int32_t) prompt.size(), n_batch_dft), 0, 1); + + for (int32_t i = 0; i < (int32_t) prompt.size(); i += n_batch_dft) { + const int32_t n_eval = std::min((int32_t) prompt.size() - i, n_batch_dft); + common_batch_clear(batch_prompt); + + for (int32_t j = 0; j < n_eval; ++j) { + const int32_t idx = i + j; + const bool is_last = (idx == (int32_t) prompt.size() - 1); + common_batch_add(batch_prompt, prompt[idx], (llama_pos) idx, { seq_id }, is_last); + } + + const int ret = llama_decode(ctx_dft, batch_prompt); + if (ret != 0) { + SPF_ERR("failed to decode prompt on draft model, ret = %d\n", ret); + llama_batch_free(batch_prompt); + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + return res; + } + } + llama_batch_free(batch_prompt); + } + + const auto t_prefill_end = ggml_time_us(); + res.t_draft_eval_us = t_prefill_end - t_start; + + // 2. lookahead decode steps with attention extraction + std::vector total_importance(prompt.size(), 0.0f); + + cb_attn_collector_data cb_data; + cb_data.n_prompt = (int32_t) prompt.size(); + + // attach callback + llama_set_eval_callback(ctx_dft, cb_collect_attn, &cb_data); + + llama_batch batch_decode = llama_batch_init(1, 0, 1); + + int32_t actual_steps = 0; + int32_t cur_pos = (int32_t) prompt.size(); + + for (int32_t k = 0; k < lookahead; ++k) { + cb_data.layer_attns.clear(); + + const llama_token token_id = common_sampler_sample(smpl_dft, ctx_dft, -1); + common_sampler_accept(smpl_dft, token_id, true); + + if (llama_vocab_is_eog(vocab_dft, token_id)) { + break; + } + + common_batch_clear(batch_decode); + common_batch_add(batch_decode, token_id, cur_pos++, { seq_id }, true); + + const int ret = llama_decode(ctx_dft, batch_decode); + if (ret != 0) { + SPF_ERR("failed lookahead decode step %d, ret = %d\n", k, ret); + break; + } + + if (!cb_data.layer_attns.empty()) { + const auto step_importance = common_speculative_prefill_compute_importance( + cb_data.layer_attns, + (int32_t) prompt.size(), + n_heads_dft, + params.pool_kernel_size); + + for (size_t i = 0; i < prompt.size(); ++i) { + total_importance[i] += step_importance[i]; + } + actual_steps++; + } + } + + llama_batch_free(batch_decode); + + // detach callback + llama_set_eval_callback(ctx_dft, nullptr, nullptr); + + if (actual_steps == 0) { + LOG_WRN("spec-prefill: attention was not captured (flash attn on?); keeping full prompt\n"); + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.importance_scores.assign(prompt.size(), 1.0f); + return res; + } + + for (size_t i = 0; i < prompt.size(); ++i) { + total_importance[i] /= (float) actual_steps; + } + + const auto t_est_start = ggml_time_us(); + res.importance_scores = total_importance; + res.kept_indices = common_speculative_prefill_select_indices(total_importance, params); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.t_estimate_us = ggml_time_us() - t_est_start; + + return res; +} diff --git a/common/speculative-prefill.h b/common/speculative-prefill.h new file mode 100644 index 00000000000..0752449133d --- /dev/null +++ b/common/speculative-prefill.h @@ -0,0 +1,37 @@ +#pragma once + +#include "llama.h" +#include "common.h" + +#include +#include + + +struct common_speculative_prefill_result { + std::vector kept_indices; // sorted original token positions kept for target model + std::vector importance_scores; // per-token aggregated importance scores (size = n_prompt) + int64_t t_draft_eval_us = 0; + int64_t t_estimate_us = 0; + int32_t n_prompt_orig = 0; + int32_t n_prompt_kept = 0; +}; + +// calculate 1D average-pooled smoothed token importance from raw attention matrices +std::vector common_speculative_prefill_compute_importance( + const std::vector> & attn_layers, + int32_t n_prompt, + int32_t n_heads, + int32_t pool_kernel_size); + +// select kept token indices using chunk-based top-k or token-level ranking +std::vector common_speculative_prefill_select_indices( + const std::vector & importance_scores, + const common_params_speculative_prefill & params); + +// execute draft prefill, lookahead decoding, attention collection, and index selection +common_speculative_prefill_result common_speculative_prefill_execute( + llama_context * ctx_dft, + common_sampler * smpl_dft, + const std::vector & prompt, + llama_seq_id seq_id, + const common_params_speculative_prefill & params); diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 39f802d250e..f023b4bfb24 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -30,6 +30,7 @@ else() add_subdirectory(simple) add_subdirectory(simple-chat) add_subdirectory(speculative) + add_subdirectory(speculative-prefill) add_subdirectory(speculative-simple) add_subdirectory(gen-docs) add_subdirectory(training) diff --git a/examples/speculative-prefill/CMakeLists.txt b/examples/speculative-prefill/CMakeLists.txt new file mode 100644 index 00000000000..9bfdec3148b --- /dev/null +++ b/examples/speculative-prefill/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET llama-speculative-prefill) +add_executable(${TARGET} speculative-prefill.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp new file mode 100644 index 00000000000..1a7ae656046 --- /dev/null +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -0,0 +1,270 @@ +#include "arg.h" +#include "common.h" +#include "log.h" +#include "llama.h" +#include "sampling.h" +#include "speculative.h" +#include "speculative-prefill.h" +#include "../../src/llama-ext.h" + +#include +#include +#include +#include +#include +#include +#include + +#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + + common_init(); + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) { + return 1; + } + + common_params_model dft_model = params.speculative.prefill.model; + if (dft_model.empty()) { + dft_model = params.speculative.draft.mparams; + } + + if (dft_model.empty()) { + LOG_ERR("%s: draft model is required for speculative prefill (specify with -mpd, --spec-prefill-model, or -md)\n", __func__); + return 1; + } + + // force speculative prefill enabled + params.speculative.prefill.enabled = true; + + llama_backend_init(); + llama_numa_init(params.numa); + + // load target model + LOG_INF("%s: loading target model...\n", __func__); + auto init_tgt = common_init_from_params(params); + if (!init_tgt) { + LOG_ERR("%s: failed to load target model\n", __func__); + return 1; + } + + llama_model * model_tgt = init_tgt->model(); + llama_context * ctx_tgt = init_tgt->context(); + + if (llama_model_is_recurrent(model_tgt)) { + LOG_ERR("%s: speculative prefill is not supported for recurrent models\n", __func__); + return 1; + } + + // load draft model with standard attention to allow attention extraction + LOG_INF("%s: loading draft model...\n", __func__); + common_params params_dft = common_base_params_to_speculative(params); + params_dft.model = dft_model; + if (params.speculative.prefill.n_ctx > 0) { + params_dft.n_ctx = params.speculative.prefill.n_ctx; + } + if (params.speculative.prefill.n_gpu_layers != -1) { + params_dft.n_gpu_layers = params.speculative.prefill.n_gpu_layers; + } + if (!params.speculative.prefill.devices.empty()) { + params_dft.devices = params.speculative.prefill.devices; + } else if (!params.speculative.draft.devices.empty()) { + params_dft.devices = params.speculative.draft.devices; + } + params_dft.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; // standard attention needed to capture kq_soft_max weights + + auto init_dft = common_init_from_params(params_dft, /*model_only=*/true); + if (!init_dft) { + LOG_ERR("%s: failed to load draft model\n", __func__); + return 1; + } + + llama_model * model_dft = init_dft->model(); + if (!model_dft) { + LOG_ERR("%s: failed to load draft model\n", __func__); + return 1; + } + + if (llama_model_is_recurrent(model_dft)) { + LOG_ERR("%s: draft model is recurrent and not supported for speculative prefill\n", __func__); + return 1; + } + + if (llama_model_target_layer_ids_n(model_dft) > 0) { + LOG_ERR("%s: draft model '%s' is target-dependent and cannot be used for speculative prefill\n", __func__, dft_model.path.c_str()); + return 1; + } + + if (params.speculative.prefill.n_ctx <= 0 && params_dft.n_ctx > (int32_t) llama_model_n_ctx_train(model_dft)) { + params_dft.n_ctx = llama_model_n_ctx_train(model_dft); + LOG_INF("%s: capping speculative prefill draft context to training limit (%d tokens)\n", __func__, params_dft.n_ctx); + } + + llama_context_params cparams_dft = common_context_params_to_llama(params_dft); + llama_context_ptr ctx_dft_own(llama_init_from_model(model_dft, cparams_dft)); + llama_context * ctx_dft = ctx_dft_own.get(); + if (!ctx_dft) { + LOG_ERR("%s: failed to create draft context\n", __func__); + return 1; + } + + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); + const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); + + if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || + (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) { + LOG_ERR("%s: draft model bos tokens must match target model. add: %d - %d, id: %d - %d\n", + __func__, + llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft), + llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft)); + return 1; + } + + { + const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); + const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); + const int vocab_diff = n_vocab_tgt > n_vocab_dft ? n_vocab_tgt - n_vocab_dft : n_vocab_dft - n_vocab_tgt; + if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { + LOG_ERR("%s: target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + __func__, n_vocab_tgt, n_vocab_dft, vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); + return 1; + } + } + + // tokenize prompt + std::vector prompt_tokens = common_tokenize(ctx_tgt, params.prompt, true, true); + const int32_t n_prompt = (int32_t) prompt_tokens.size(); + + if (n_prompt == 0) { + LOG_ERR("%s: empty prompt provided\n", __func__); + return 1; + } + + LOG_INF("%s: prompt tokens = %d\n", __func__, n_prompt); + LOG_INF("%s: speculative prefill config: percentage = %.2f, chunk_size = %d, lookahead = %d, pool_kernel = %d\n", + __func__, + (double) params.speculative.prefill.percentage, + params.speculative.prefill.chunk_size, + params.speculative.prefill.look_ahead_cnt, + params.speculative.prefill.pool_kernel_size); + + llama_seq_id seq_id = 0; + + // initialize draft sampler for lookahead steps + common_params_sampling sparams_dft = params.sampling; + sparams_dft.temp = 0.0f; + common_sampler_ptr smpl_dft(common_sampler_init(model_dft, sparams_dft)); + + // step 1: run speculative prefill on draft model + const auto t_spec_prefill_start = ggml_time_us(); + + common_speculative_prefill_result spec_res = common_speculative_prefill_execute( + ctx_dft, + smpl_dft.get(), + prompt_tokens, + seq_id, + params.speculative.prefill); + + LOG_INF("%s: speculative prefill kept %d / %d tokens (%.1f%%)\n", + __func__, + spec_res.n_prompt_kept, + spec_res.n_prompt_orig, + 100.0f * (float) spec_res.n_prompt_kept / (float) spec_res.n_prompt_orig); + LOG_INF("%s: draft eval time: %.2f ms, importance estimation time: %.2f ms\n", + __func__, + spec_res.t_draft_eval_us / 1000.0f, + spec_res.t_estimate_us / 1000.0f); + + // step 2: sparse target model prefill + const auto t_tgt_prefill_start = ggml_time_us(); + + const int32_t n_batch_tgt = llama_n_batch(ctx_tgt); + const int32_t n_kept_total = (int32_t) spec_res.kept_indices.size(); + llama_batch batch_tgt = llama_batch_init(std::min(n_kept_total, n_batch_tgt), 0, 1); + + int ret = 0; + for (int32_t i = 0; i < n_kept_total; i += n_batch_tgt) { + const int32_t n_eval = std::min(n_kept_total - i, n_batch_tgt); + common_batch_clear(batch_tgt); + + for (int32_t j = 0; j < n_eval; ++j) { + const int32_t k = i + j; + const int32_t orig_idx = spec_res.kept_indices[k]; + const bool is_last = (k == n_kept_total - 1); + common_batch_add(batch_tgt, prompt_tokens[orig_idx], (llama_pos) k, { seq_id }, is_last); + } + + ret = llama_decode(ctx_tgt, batch_tgt); + if (ret != 0) { + LOG_ERR("%s: failed to decode sparse prompt on target model, ret = %d\n", __func__, ret); + llama_batch_free(batch_tgt); + return 1; + } + } + llama_batch_free(batch_tgt); + + llama_synchronize(ctx_tgt); + llama_synchronize(ctx_dft); + + const auto t_tgt_prefill_end = ggml_time_us(); + const double ttft_ms = (t_tgt_prefill_end - t_spec_prefill_start) / 1000.0; + const double tgt_prefill_ms = (t_tgt_prefill_end - t_tgt_prefill_start) / 1000.0; + + LOG_INF("%s: sparse target prefill time = %.2f ms\n", __func__, tgt_prefill_ms); + LOG_INF("%s: total Time-To-First-Token (TTFT) = %.2f ms (%.2f effective prompt tokens/s)\n", + __func__, ttft_ms, ttft_ms > 0.0 ? (1000.0 * (double) n_prompt / ttft_ms) : 0.0); + + // step 3: autoregressive generation + common_sampler_ptr smpl_tgt(common_sampler_init(model_tgt, params.sampling)); + + LOG("\n--- Generation Start ---\n"); + + int32_t n_predict = params.n_predict > 0 ? params.n_predict : 32; + int32_t cur_pos = n_kept_total; + int32_t n_generated = 0; + + llama_batch batch_gen = llama_batch_init(1, 0, 1); + + const auto t_gen_start = ggml_time_us(); + + for (int32_t i = 0; i < n_predict; ++i) { + const llama_token token_id = common_sampler_sample(smpl_tgt.get(), ctx_tgt, -1); + common_sampler_accept(smpl_tgt.get(), token_id, true); + + if (llama_vocab_is_eog(vocab_tgt, token_id)) { + break; + } + + const std::string piece = common_token_to_piece(ctx_tgt, token_id); + LOG("%s", piece.c_str()); + fflush(stdout); + + common_batch_clear(batch_gen); + common_batch_add(batch_gen, token_id, (llama_pos) cur_pos++, { seq_id }, true); + + ret = llama_decode(ctx_tgt, batch_gen); + if (ret != 0) { + LOG_ERR("%s: failed to decode generated token %d, ret = %d\n", __func__, i, ret); + break; + } + + n_generated++; + } + + const auto t_gen_end = ggml_time_us(); + const double gen_ms = (t_gen_end - t_gen_start) / 1000.0; + + LOG("\n--- Generation End ---\n\n"); + + LOG_INF("generated %d tokens in %.2f ms (%.2f tokens/s)\n", + n_generated, gen_ms, n_generated > 0 ? (n_generated / (gen_ms / 1000.0)) : 0.0); + + llama_batch_free(batch_gen); + llama_backend_free(); + + return 0; +} diff --git a/include/llama.h b/include/llama.h index ef7a012c43a..51f59e2b1b6 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1021,6 +1021,9 @@ extern "C" { // Set abort callback LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data); + // Set evaluation callback + LLAMA_API void llama_set_eval_callback(struct llama_context * ctx, ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data); + // Wait until all computations are finished // This is automatically done when using one of the functions below to obtain the computation results // and is not necessary to call it explicitly in most cases diff --git a/scripts/compare-llama-bench.py b/scripts/compare-llama-bench.py index e5f26b5a41f..2fbcf555583 100755 --- a/scripts/compare-llama-bench.py +++ b/scripts/compare-llama-bench.py @@ -30,7 +30,9 @@ "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "tensor_split", "tensor_buft_overrides", "load_mode", "embeddings", "no_op_offload", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts", "n_cpu_moe", - "fit_target", "fit_min_ctx" + "fit_target", "fit_min_ctx", + "spec_prefill_model", "spec_prefill_n_gpu_layers", "spec_prefill_percentage", + "spec_prefill_chunk_size", "spec_prefill_lookahead", "spec_prefill_pool_kernel" ] LLAMA_BENCH_DB_TYPES = [ @@ -38,9 +40,11 @@ "TEXT", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "TEXT", "INTEGER", "INTEGER", "TEXT", "TEXT", "INTEGER", "TEXT", "INTEGER", "INTEGER", "INTEGER", "TEXT", "TEXT", - "TEXT", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER", + "TEXT", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "INTEGER", "TEXT", "INTEGER", "INTEGER", "REAL", "REAL", "INTEGER", - "INTEGER", "INTEGER" + "INTEGER", "INTEGER", + "TEXT", "INTEGER", "REAL", + "INTEGER", "INTEGER", "INTEGER" ] # All test-backend-ops SQL fields @@ -64,7 +68,9 @@ "cpu_info", "gpu_info", "backends", "n_gpu_layers", "n_cpu_moe", "tensor_buft_overrides", "model_filename", "model_type", "n_batch", "n_ubatch", "embeddings", "cpu_mask", "cpu_strict", "poll", "n_threads", "type_k", "type_v", "load_mode", "no_kv_offload", "split_mode", "main_gpu", "tensor_split", "flash_attn", "n_prompt", "n_gen", "n_depth", - "fit_target", "fit_min_ctx" + "fit_target", "fit_min_ctx", + "spec_prefill_model", "spec_prefill_n_gpu_layers", "spec_prefill_percentage", + "spec_prefill_chunk_size", "spec_prefill_lookahead", "spec_prefill_pool_kernel" ] # Properties by which to differentiate results per commit for test-backend-ops: @@ -84,6 +90,8 @@ "cpu_mask": "CPU mask", "cpu_strict": "CPU strict", "poll": "Poll", "n_threads": "Threads", "type_k": "K type", "type_v": "V type", "load_mode": "Load mode", "no_kv_offload": "NKVO", "split_mode": "Split mode", "main_gpu": "Main GPU", "tensor_split": "Tensor split", "flash_attn": "FlashAttention", + "spec_prefill_model": "Draft Model", "spec_prefill_n_gpu_layers": "Draft ngl", "spec_prefill_percentage": "SPF %", + "spec_prefill_chunk_size": "SPF chunk", "spec_prefill_lookahead": "SPF lookahead", "spec_prefill_pool_kernel": "SPF pool", } # Header names for the table (test-backend-ops): diff --git a/scripts/compare_spec_prefill.py b/scripts/compare_spec_prefill.py new file mode 100755 index 00000000000..bf4a5d6dbe3 --- /dev/null +++ b/scripts/compare_spec_prefill.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +import argparse +import os +import re +import statistics +import subprocess +import sys +import time + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + + parser = argparse.ArgumentParser(description="Compare Speculative Prefill ON vs OFF (Baseline)") + parser.add_argument("-m", "--target-model", help="Path to target model GGUF") + parser.add_argument("-hf", "--hf-repo", help="Hugging Face repo for target model (/[:quant])") + parser.add_argument("-hff", "--hf-file", help="Hugging Face model file for target model") + parser.add_argument("-hft", "--hf-token", help="Hugging Face access token") + parser.add_argument("-mpd", "-md", "--draft-model", "--spec-prefill-model", help="Path to draft model GGUF") + parser.add_argument("-hfpd", "-hfd", "--hf-repo-draft", "--draft-hf", "--spec-prefill-hf", help="Hugging Face repo for draft model (/[:quant])") + parser.add_argument("-hffd", "--hf-file-draft", help="Hugging Face model file for draft model") + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99, help="Number of GPU layers for target model (default: 99)") + parser.add_argument("-nglpd", "-ngld", "--n-gpu-layers-draft", "--spec-prefill-ngl", type=int, default=99, help="Number of GPU layers for draft model (default: 99)") + parser.add_argument("-p", "--prompt", help="Prompt text") + parser.add_argument("-f", "--file", help="File containing prompt text") + parser.add_argument("--n-prompt", type=int, help="Synthesize prompt of approx N tokens") + parser.add_argument("-n", "--n-predict", type=int, default=32, help="Number of tokens to generate (default: 32)") + parser.add_argument("--percentages", default="0.2,0.3,0.5", help="Comma-separated keep percentages (default: 0.2,0.3,0.5)") + parser.add_argument("--chunk-size", type=int, default=32, help="Chunk size (default: 32)") + parser.add_argument("--lookahead", type=int, default=4, help="Lookahead count (default: 4)") + parser.add_argument("--reps", type=int, default=1, help="Repetitions per test (default: 1)") + parser.add_argument("--bin", default=default_bin, help="Path to binary (default: auto-detected)") + parser.add_argument("--threads", type=int, default=0, help="Number of threads (0 for auto)") + return parser.parse_args() + +def generate_synthetic_prompt(n_tokens): + base = ( + "In an ancient kingdom surrounded by mist-covered mountains, scholars gathered in the great library " + "to study the ancient manuscripts of astronomy, mathematics, and philosophy. They observed the stars " + "every evening through brass telescopes, recording every shift in planetary alignment with meticulous " + "precision. Among them was an eager apprentice named Nicholas, who discovered an enigmatic cipher " + "hidden within the margins of an age-old star chart. " + ) + repeats = max(1, n_tokens // 75) + return (base * repeats).strip() + +def run_single(bin_path, model_args, prompt, percentage, chunk_size, lookahead, n_predict, threads, ngl, ngld): + cmd = [ + bin_path, + *model_args, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(n_predict), + "--spec-prefill-percentage", str(percentage), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + if threads > 0: + cmd.extend(["-t", str(threads)]) + + start_t = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + wall_ms = (time.perf_counter() - start_t) * 1000.0 + + output = proc.stdout + "\n" + proc.stderr + + if proc.returncode != 0: + print("Error running command (exit %d):\n%s" % (proc.returncode, output)) + return None + + res = { + "wall_ms": wall_ms, + "n_prompt_orig": 0, + "n_prompt_kept": 0, + "kept_pct": 100.0, + "draft_eval_ms": 0.0, + "estimate_ms": 0.0, + "sparse_prefill_ms": 0.0, + "ttft_ms": 0.0, + "gen_tokens": 0, + "gen_speed_ts": 0.0, + "generation": "", + } + + m = re.search(r"prompt tokens = (\d+)", output) + if m: + res["n_prompt_orig"] = int(m.group(1)) + + m = re.search(r"speculative prefill kept (\d+) / (\d+) tokens \(([\d\.]+)%\)", output) + if m: + res["n_prompt_kept"] = int(m.group(1)) + res["kept_pct"] = float(m.group(3)) + else: + res["n_prompt_kept"] = res["n_prompt_orig"] + res["kept_pct"] = 100.0 + + m = re.search(r"draft eval time: ([\d\.]+) ms, importance estimation time: ([\d\.]+) ms", output) + if m: + res["draft_eval_ms"] = float(m.group(1)) + res["estimate_ms"] = float(m.group(2)) + + m = re.search(r"sparse target prefill time = ([\d\.]+) ms", output) + if m: + res["sparse_prefill_ms"] = float(m.group(1)) + + m = re.search(r"total Time-To-First-Token \(TTFT\) = ([\d\.]+) ms", output) + if m: + res["ttft_ms"] = float(m.group(1)) + elif res["sparse_prefill_ms"] > 0: + res["ttft_ms"] = res["sparse_prefill_ms"] + + m = re.search(r"generated (\d+) tokens in ([\d\.]+) ms \(([\d\.]+) tokens/s\)", output) + if m: + res["gen_tokens"] = int(m.group(1)) + res["gen_speed_ts"] = float(m.group(3)) + + gen_m = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if gen_m: + res["generation"] = gen_m.group(1).strip() + + return res + +def main(): + args = parse_args() + + if not os.path.exists(args.bin): + print("Error: binary %s not found. Please build llama.cpp first." % args.bin) + sys.exit(1) + + # Validate target model arguments + if not args.target_model and not args.hf_repo: + print("Error: either -m/--target-model or -hf/--hf-repo must be specified for target model.") + sys.exit(1) + + model_args = [] + target_desc = "" + draft_desc = "" + + if args.target_model: + if not os.path.exists(args.target_model): + print("Error: target model file %s not found." % args.target_model) + sys.exit(1) + model_args.extend(["-m", args.target_model]) + target_desc = args.target_model + elif args.hf_repo: + model_args.extend(["-hf", args.hf_repo]) + target_desc = args.hf_repo + if args.hf_file: + model_args.extend(["-hff", args.hf_file]) + target_desc += " (" + args.hf_file + ")" + + if args.hf_token: + model_args.extend(["-hft", args.hf_token]) + + # Validate draft model arguments + if args.draft_model: + if not os.path.exists(args.draft_model): + print("Error: draft model file %s not found." % args.draft_model) + sys.exit(1) + model_args.extend(["-md", args.draft_model]) + draft_desc = args.draft_model + elif args.hf_repo_draft: + model_args.extend(["-hfd", args.hf_repo_draft]) + draft_desc = args.hf_repo_draft + if args.hf_file_draft: + model_args.extend(["-hffd", args.hf_file_draft]) + draft_desc += " (" + args.hf_file_draft + ")" + else: + # Default draft model to target model if unspecified + if args.target_model: + model_args.extend(["-md", args.target_model]) + draft_desc = args.target_model + " (self-draft)" + elif args.hf_repo: + model_args.extend(["-hfd", args.hf_repo]) + if args.hf_file: + model_args.extend(["-hffd", args.hf_file]) + draft_desc = args.hf_repo + " (self-draft)" + + if args.prompt: + prompt = args.prompt + elif args.file: + with open(args.file, "r") as f: + prompt = f.read().strip() + elif args.n_prompt: + prompt = generate_synthetic_prompt(args.n_prompt) + else: + prompt = ( + "Once upon a time in a bustling mountain village, there was a master clockmaker named Jonathan. " + "Every morning, he wound the town clock high upon the clocktower, watching the villagers below begin " + "their daily routines. One foggy autumn morning, an enigmatic traveler arrived bearing a broken mechanical " + "device covered in strange celestial engravings." + ) + + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + print("=" * 80) + print("SPECULATIVE PREFILL COMPARISON: ON vs OFF (Baseline)") + print("=" * 80) + print("Target Model : %s" % target_desc) + print("Draft Model : %s" % draft_desc) + print("Backend : Vulkan (GPU layers: %d tgt, %d draft)" % (args.n_gpu_layers, args.n_gpu_layers_draft)) + print("Binary Path : %s" % args.bin) + print("Chunk Size : %d" % args.chunk_size) + print("Lookahead : %d" % args.lookahead) + print("Repetitions : %d" % args.reps) + print("Gen Tokens : %d" % args.n_predict) + print("-" * 80) + + # 1. Warmup + print("Running warmup...") + run_single(args.bin, model_args, prompt, 1.0, args.chunk_size, args.lookahead, 4, args.threads, args.n_gpu_layers, args.n_gpu_layers_draft) + + # 2. Run Baseline (OFF: percentage = 1.0) + print("\n[1/2] Benchmarking Baseline (Speculative Prefill: OFF, Keep: 100%)...") + base_runs = [] + base_gen = "" + for r in range(args.reps): + res = run_single(args.bin, model_args, prompt, 1.0, args.chunk_size, args.lookahead, args.n_predict, args.threads, args.n_gpu_layers, args.n_gpu_layers_draft) + if res: + base_runs.append(res) + base_gen = res["generation"] + print(" Rep %d: TTFT = %.2f ms | Prefill = %.2f ms" % (r+1, res["ttft_ms"], res["sparse_prefill_ms"])) + + if not base_runs: + print("Error: baseline runs failed.") + sys.exit(1) + + base_ttft_avg = statistics.mean([r["ttft_ms"] for r in base_runs]) + base_ttft_std = statistics.stdev([r["ttft_ms"] for r in base_runs]) if len(base_runs) > 1 else 0.0 + n_prompt_total = base_runs[0]["n_prompt_orig"] + + # 3. Run Speculative Prefill (ON: percentage = p) + spec_results = {} + for p in percentages: + print("\n[2/2] Benchmarking Speculative Prefill ON (Keep: %.0f%%)..." % (p * 100.0)) + runs = [] + spec_gen = "" + for r in range(args.reps): + res = run_single(args.bin, model_args, prompt, p, args.chunk_size, args.lookahead, args.n_predict, args.threads, args.n_gpu_layers, args.n_gpu_layers_draft) + if res: + runs.append(res) + spec_gen = res["generation"] + print(" Rep %d: TTFT = %.2f ms | Draft = %.2f ms | Target Prefill = %.2f ms | Kept = %d/%d" % ( + r+1, res["ttft_ms"], res["draft_eval_ms"], res["sparse_prefill_ms"], res["n_prompt_kept"], res["n_prompt_orig"])) + + if runs: + ttft_avg = statistics.mean([r["ttft_ms"] for r in runs]) + ttft_std = statistics.stdev([r["ttft_ms"] for r in runs]) if len(runs) > 1 else 0.0 + speedup = base_ttft_avg / ttft_avg if ttft_avg > 0 else 0.0 + spec_results[p] = { + "runs": runs, + "ttft_avg": ttft_avg, + "ttft_std": ttft_std, + "speedup": speedup, + "n_kept": runs[0]["n_prompt_kept"], + "draft_eval_avg": statistics.mean([r["draft_eval_ms"] for r in runs]), + "target_prefill_avg": statistics.mean([r["sparse_prefill_ms"] for r in runs]), + "gen_speed_avg": statistics.mean([r["gen_speed_ts"] for r in runs]), + "gen_text": spec_gen, + } + + # Print Summary Table + print("\n" + "=" * 90) + print("BENCHMARK SUMMARY (Prompt Length N = %d tokens)" % n_prompt_total) + print("=" * 90) + header = "%-24s | %-13s | %-16s | %-9s | %-14s" % ("Configuration", "Tokens Kept", "TTFT (ms)", "Speedup", "Throughput") + print(header) + print("-" * 90) + + base_tp = (n_prompt_total / (base_ttft_avg / 1000.0)) if base_ttft_avg > 0 else 0.0 + print("%-24s | %5d/%-5d (100%%) | %6.2f ± %-5.2f | %-9s | %8.1f t/s" % ( + "Baseline (OFF, 100%)", n_prompt_total, n_prompt_total, base_ttft_avg, base_ttft_std, "1.00x", base_tp)) + + for p, data in spec_results.items(): + kept_str = "%5d/%d (%4.1f%%)" % (data["n_kept"], n_prompt_total, p*100.0) + ttft_str = "%6.2f ± %-5.2f" % (data["ttft_avg"], data["ttft_std"]) + speedup_str = "%5.2fx" % data["speedup"] + eff_tp = (n_prompt_total / (data["ttft_avg"] / 1000.0)) if data["ttft_avg"] > 0 else 0.0 + cfg_str = "SpecPrefill (ON, p=%.2f)" % p + print("%-24s | %-13s | %-16s | %-9s | %8.1f t/s" % ( + cfg_str, kept_str, ttft_str, speedup_str, eff_tp)) + + print("=" * 90) + + # Detailed Latency Breakdown + print("\nLATENCY BREAKDOWN (Speculative Prefill Phases):") + print("-" * 90) + print("%-10s | %-15s | %-16s | %-22s | %-12s" % ("Keep %", "Draft Prefill", "Lookahead/Attn", "Target Sparse Prefill", "Total TTFT")) + print("-" * 90) + for p, data in spec_results.items(): + lah_time = max(0.0, data["ttft_avg"] - data["draft_eval_avg"] - data["target_prefill_avg"]) + print("%5.1f%% | %8.2f ms | %9.2f ms | %12.2f ms | %7.2f ms" % ( + p*100.0, data["draft_eval_avg"], lah_time, data["target_prefill_avg"], data["ttft_avg"])) + print("-" * 90) + + # Output Sample Comparison + print("\nSAMPLE CONTINUATION OUTPUTS:") + print("-" * 90) + print("[Baseline (OFF)]:\n\"%s\"\n" % base_gen) + for p, data in spec_results.items(): + print("[SpecPrefill (ON, p=%.2f)]:\n\"%s\"\n" % (p, data["gen_text"])) + print("-" * 90) + +if __name__ == "__main__": + main() diff --git a/scripts/eval_spec_prefill_entropy.py b/scripts/eval_spec_prefill_entropy.py new file mode 100644 index 00000000000..d931bbfad6d --- /dev/null +++ b/scripts/eval_spec_prefill_entropy.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +""" +Empirical demonstration of mathematical failure boundaries for speculative prefill: +1. High-entropy multi-key retrieval (non-sequential random IDs across distributed chunks) +2. Whole-document comprehensive extraction (aggregating 8 distinct entity-value pairs) +3. Conflicting multi-version overrides with adversarial distractor anchors +""" + +import argparse +import os +import re +import subprocess +import time + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + default_tgt = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/4ca720788d1e01f1bff70c033e0d0028fd02e502/Qwen3.8-27B-UD-Q6_K_XL.gguf" + default_dft = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.5-2B-GGUF/snapshots/f6d5376be1edb4d416d56da11e5397a961aca8ae/Qwen3.5-2B-UD-Q4_K_XL.gguf" + + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--target-model", default=default_tgt) + parser.add_argument("-md", "--draft-model", default=default_dft) + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99) + parser.add_argument("-ngld", "--n-gpu-layers-draft", type=int, default=99) + parser.add_argument("--bin", default=default_bin) + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.15,0.08") + return parser.parse_args() + +def run_test(bin_path, target_model, draft_model, prompt, p, ngl, ngld, max_gen=48, chunk_size=32, lookahead=4): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(max_gen), + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept_tokens = int(m_kept.group(1)) if m_kept else 0 + total_tokens = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "kept": kept_tokens, + "total": total_tokens, + } + +def get_distractors(n=25): + text = [ + "The geological survey team conducted detailed mineralogical analyses in the western foothills. Core samples revealed substantial granite and quartz strata with trace deposits of copper ore.", + "Municipal water engineers redesigned the gravity-fed aqueduct conduits. Pressure regulation cisterns and basalt masonry arches ensured uninterrupted flow into central metropolitan fountains.", + "Agricultural registers recorded seasonal barley harvests across forty alluvial valleys. Yields were cataloged in standard bushel units and distributed to state granaries via river barges.", + "Archival preservationists treated ancient parchment scrolls using mild cedar oil emulsions. Manuscript codices from the third century were rebound in pigskin covers and stored in dry vaults.", + "Textile manufacturing facilities calibrated loom tension for fine linen production. Natural madder dye vats produced crimson vestments exported to neighboring maritime principalities.", + "Harbor authorities logged the arrival of merchant frigates carrying spices, porcelain, and timber. Port tariffs were collected in silver bullion and recorded in double-entry ledgers.", + "Astronomers observed the transit of celestial satellites across equatorial constellations. Quadrant measurements were compiled into seasonal navigational ephemerides for mariners." + ] + return [text[i % len(text)] for i in range(n)] + +def test_random_kv_pairs(): + paras = get_distractors(24) + # High-entropy non-sequential codes in 5 separated chunks + items = [ + ("Database Node Alpha is assigned security token 49152.", 2), + ("Database Node Beta is assigned security token 19842.", 7), + ("Database Node Gamma is assigned security token 33419.", 12), + ("Database Node Delta is assigned security token 58201.", 17), + ("Database Node Epsilon is assigned security token 27182.", 22), + ] + for text, idx in sorted(items, key=lambda x: x[1], reverse=True): + paras.insert(idx, text) + + prompt = ( + "%s\n\n" + "Question: List the exact 5-digit security tokens for all five database nodes (Alpha, Beta, Gamma, Delta, Epsilon):\n" + "Answer: The security tokens for Alpha, Beta, Gamma, Delta, Epsilon are:" + ) % ("\n\n".join(paras)) + + tokens = ["49152", "19842", "33419", "58201", "27182"] + def check(g): + found = [t for t in tokens if t in g] + return len(found) == len(tokens), len(found) + + return { + "name": "High-Entropy Distributed Multi-Key Retrieval", + "prompt": prompt, + "check": check, + "total_items": 5 + } + +def main(): + args = parse_args() + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + task = test_random_kv_pairs() + + print("=" * 85) + print("SPECULATIVE PREFILL: ENTROPY & PRUNING LIMIT EVALUATION") + print("=" * 85) + print("Prompt Length: ~1,100 tokens, 5 non-redundant targets spread across context") + print("-" * 85) + + for p in percentages: + res = run_test(args.bin, args.target_model, args.draft_model, task["prompt"], p, args.n_gpu_layers, args.n_gpu_layers_draft) + all_passed, count = task["check"](res["gen_text"]) + status = "PASS (5/5)" if all_passed else f"FAIL ({count}/5 retrieved)" + print("Keep %3d%% | Kept %3d/%4d tokens | Status: %-18s | Output: %s" % ( + int(p*100), res["kept"], res["total"], status, res["gen_text"].replace("\n", " ")[:40])) + print("=" * 85) + +if __name__ == "__main__": + main() diff --git a/scripts/eval_spec_prefill_failure_modes.py b/scripts/eval_spec_prefill_failure_modes.py new file mode 100644 index 00000000000..3f7d2836467 --- /dev/null +++ b/scripts/eval_spec_prefill_failure_modes.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python3 +""" +Stress test / Failure Mode benchmark for Speculative Prefill in llama.cpp. +Identifies challenging scenarios where speculative prefill drops below 100% accuracy: +1. Multi-hop indirect reasoning (bridging without keyword overlap) +2. Global aggregation / counting across non-redundant distributed context +3. Temporal revision / conflicting overrides (adversarial distractors) +4. Multi-step distributed variable tracking across distant chunks +5. Extreme compression (low keep ratio p <= 0.05 / 0.10) +""" + +import argparse +import os +import re +import subprocess +import sys +import time + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + default_tgt = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/4ca720788d1e01f1bff70c033e0d0028fd02e502/Qwen3.8-27B-UD-Q6_K_XL.gguf" + default_dft = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.5-2B-GGUF/snapshots/f6d5376be1edb4d416d56da11e5397a961aca8ae/Qwen3.5-2B-UD-Q4_K_XL.gguf" + + parser = argparse.ArgumentParser(description="Speculative Prefill Failure Mode / Stress Benchmark") + parser.add_argument("-m", "--target-model", default=default_tgt, help="Target model GGUF") + parser.add_argument("-md", "--draft-model", default=default_dft, help="Draft model GGUF") + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99) + parser.add_argument("-ngld", "--n-gpu-layers-draft", type=int, default=99) + parser.add_argument("--bin", default=default_bin) + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.15,0.08,0.05") + return parser.parse_args() + +def run_test(bin_path, target_model, draft_model, prompt, p, ngl, ngld, max_gen=40, chunk_size=32, lookahead=4): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(max_gen), + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + t0 = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + t_ms = (time.perf_counter() - t0) * 1000.0 + + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_ttft = re.search(r"total Time-To-First-Token \(TTFT\) = ([\d\.]+) ms", output) + ttft_ms = float(m_ttft.group(1)) if m_ttft else 0.0 + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept_tokens = int(m_kept.group(1)) if m_kept else 0 + total_tokens = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "ttft_ms": ttft_ms, + "kept": kept_tokens, + "total": total_tokens, + "raw": output + } + +def build_distractor_background(num_paras=25): + topics = [ + "In northern agricultural districts, crop rotation schedules were synchronized with solar calendars. Farmers monitored soil moisture and planted legumes every third season to restore nitrogen balance. Yield estimates were cataloged by regional grain administrators.", + "Geological survey teams mapped mineral deposits along the western mountain ranges. They noted extensive veins of quartz, granite, and copper ore embedded within metamorphic formations. Laboratory assays confirmed high purity in several core samples.", + "Maritime navigational treatises detailed the seasonal reversal of coastal monsoons. Vessel captains adjusted their trade voyages to harness prevailing trade winds and avoid shallow reefs. Harbor logs recorded arrival dates and cargo manifests.", + "Urban sanitation works expanded rapidly during the middle administrative period. Brick-lined drainage channels and elevated aqueducts delivered fresh spring water across residential sectors, drastically reducing waterborne illnesses in the central wards.", + "Archival historians indexed thousands of legal decrees enacted over three dynasties. Scribes preserved fragile scrolls in temperature-controlled stone vaults to prevent moisture decay. Cross-referencing genealogical records resolved land boundary disputes.", + "Textile manufacturing guilds established standardized dye recipes using natural pigments. Indigo, madder root, and crushed ochre were blended in copper vats under precise temperature controls. Master weavers inspected linen quality before export.", + "Hydrological engineers constructed earthen levees along the delta tributaries. Seasonal flooding patterns were regulated through sluice gates to irrigate terraced rice paddies. Silt deposits replenished topsoil fertility across the lower basin.", + "Astronomical observatories recorded lunar occultations and planetary conjunctions with calibrated astrolabes. Star charts were published biannually for maritime navigation and calendar harmonization across all inland provinces." + ] + paras = [] + for i in range(num_paras): + paras.append(topics[i % len(topics)]) + return paras + +def create_multi_hop_latent_task(): + paras = build_distractor_background(25) + paras.insert(4, "Project Leadership Log: The chief architect and primary designer of the Helios Project is Dr. Elena Rostova.") + paras.insert(18, "Personnel Biography Note: In her personal leisure time, Dr. Elena Rostova is an accomplished player of the harpsichord.") + + prompt = ( + "%s\n\n" + "Question: What musical instrument is played by the chief architect of the Helios Project?\n" + "Answer: The musical instrument played by the chief architect of the Helios Project is the" + ) % ("\n\n".join(paras)) + + return { + "name": "1. Multi-Hop Latent (No Keyword Overlap in Hop 2)", + "prompt": prompt, + "expected": "harpsichord", + "check": lambda g: "harpsichord" in g.lower() + } + +def create_distributed_aggregation_task(): + paras = build_distractor_background(20) + incidents = [ + ("Facility Alpha logged exactly 4 security incidents this quarter.", 2), + ("Facility Beta logged exactly 7 security incidents this quarter.", 6), + ("Facility Gamma logged exactly 3 security incidents this quarter.", 10), + ("Facility Delta logged exactly 5 security incidents this quarter.", 14), + ("Facility Epsilon logged exactly 2 security incidents this quarter.", 18), + ("Facility Zeta logged exactly 6 security incidents this quarter.", 22), + ] # Total = 4 + 7 + 3 + 5 + 2 + 6 = 27 + + for text, idx in sorted(incidents, key=lambda x: x[1], reverse=True): + paras.insert(idx, text) + + prompt = ( + "%s\n\n" + "Question: What is the total sum of security incidents reported across Facility Alpha, Beta, Gamma, Delta, Epsilon, and Zeta?\n" + "Answer: The total sum of security incidents across all six facilities is" + ) % ("\n\n".join(paras)) + + return { + "name": "2. Global Aggregation (6 Distributed Items)", + "prompt": prompt, + "expected": "27", + "check": lambda g: "27" in g or "twenty-seven" in g.lower() + } + +def create_temporal_override_task(): + paras = build_distractor_background(24) + paras.insert(3, "Security Notice A: The current primary server access code is ALPHA-1111.") + paras.insert(12, "Security Notice B: Update regarding server configuration: server access code changed to BETA-2222.") + paras.insert(22, "FINAL OVERRIDE NOTICE: Under emergency security protocol, the active server access code is ZETA-9999. All earlier codes are revoked and obsolete.") + + prompt = ( + "%s\n\n" + "Question: According to the final active override notice, what is the valid server access code?\n" + "Answer: The active valid server access code is" + ) % ("\n\n".join(paras)) + + return { + "name": "3. Temporal Revision / Final Override (Adversarial)", + "prompt": prompt, + "expected": "ZETA-9999", + "check": lambda g: "zeta" in g.lower() or "9999" in g + } + +def create_distributed_var_tracking_task(): + paras = build_distractor_background(22) + paras.insert(2, "Code Section 1: initialize register_x = 100;") + paras.insert(8, "Code Section 2: register_x = register_x + 50;") + paras.insert(14, "Code Section 3: register_x = register_x * 2;") + paras.insert(20, "Code Section 4: register_x = register_x - 30;") + + prompt = ( + "%s\n\n" + "Question: Following all four sequential code sections from 1 to 4, what is the final numeric value of register_x?\n" + "Answer: The final value of register_x is" + ) % ("\n\n".join(paras)) + + return { + "name": "4. Distributed Variable Tracking", + "prompt": prompt, + "expected": "270", + "check": lambda g: "270" in g + } + +def main(): + args = parse_args() + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + print("=" * 95) + print("SPECULATIVE PREFILL FAILURE MODE & STRESS TEST BENCHMARK") + print("=" * 95) + print("Target Model : %s" % args.target_model) + print("Draft Model : %s" % args.draft_model) + print("Keep Ratios : %s" % percentages) + print("=" * 95) + + tasks = [ + create_multi_hop_latent_task(), + create_distributed_aggregation_task(), + create_temporal_override_task(), + create_distributed_var_tracking_task(), + ] + + results = {p: {} for p in percentages} + + for p in percentages: + print(f"\n>>> Running evaluation for Keep Ratio: {int(p*100)}% ({p:.2f}) <<<") + for task in tasks: + res = run_test(args.bin, args.target_model, args.draft_model, task["prompt"], p, args.n_gpu_layers, args.n_gpu_layers_draft) + passed = task["check"](res["gen_text"]) + status = "PASS" if passed else "FAIL" + results[p][task["name"]] = { + "passed": passed, + "ttft": res["ttft_ms"], + "kept": res["kept"], + "total": res["total"], + "gen": res["gen_text"].replace("\n", " ")[:50], + } + print(" %-45s : [%s] (Kept: %d/%d) | Output: \"%s\"" % ( + task["name"], status, res["kept"], res["total"], results[p][task["name"]]["gen"])) + + print("\n" + "=" * 95) + print("FAILURE MODE BENCHMARK SUMMARY TABLE") + print("=" * 95) + header = "%-8s | %-11s | " + " | ".join([f"Task {i+1}" for i in range(len(tasks))]) + " | %-8s" + print(header % ("Keep %", "Tokens Kept", "Accuracy")) + print("-" * 95) + + for p in percentages: + passes = [results[p][t["name"]]["passed"] for t in tasks] + kept_tokens = results[p][tasks[0]["name"]]["kept"] + total_tokens = results[p][tasks[0]["name"]]["total"] + acc = (sum(passes) / len(passes)) * 100.0 + pass_strs = ["PASS" if s else "FAIL" for s in passes] + row_str = "%-8s | %4d/%-6d | " + " | ".join(["%-6s" for _ in passes]) + " | %5.1f%%" + print(row_str % tuple([f"{int(p*100)}%", kept_tokens, total_tokens] + pass_strs + [acc])) + print("=" * 95) + +if __name__ == "__main__": + main() diff --git a/scripts/eval_spec_prefill_longbench_nonparametric.py b/scripts/eval_spec_prefill_longbench_nonparametric.py new file mode 100644 index 00000000000..4bacc2fdbca --- /dev/null +++ b/scripts/eval_spec_prefill_longbench_nonparametric.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +""" +Fictitious / Non-Parametric LongBench Multi-Hop QA: +Tests if LongBench multi-hop fails when models cannot cheat using pre-trained parametric memory. +""" + +import argparse +import os +import re +import subprocess +import time + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + default_tgt = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/4ca720788d1e01f1bff70c033e0d0028fd02e502/Qwen3.8-27B-UD-Q6_K_XL.gguf" + default_dft = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.5-2B-GGUF/snapshots/f6d5376be1edb4d416d56da11e5397a961aca8ae/Qwen3.5-2B-UD-Q4_K_XL.gguf" + + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--target-model", default=default_tgt) + parser.add_argument("-md", "--draft-model", default=default_dft) + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99) + parser.add_argument("-ngld", "--n-gpu-layers-draft", type=int, default=99) + parser.add_argument("--bin", default=default_bin) + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.15,0.08") + return parser.parse_args() + +def run_test(bin_path, target_model, draft_model, prompt, p, ngl, ngld, max_gen=32, chunk_size=32, lookahead=4): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(max_gen), + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + t0 = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept_tokens = int(m_kept.group(1)) if m_kept else 0 + total_tokens = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "kept": kept_tokens, + "total": total_tokens, + } + +def main(): + args = parse_args() + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + distractors = [ + "The deep-space research vessel Hyperion-1 was commissioned for extrasolar navigation under Captain Robert Vance, who completed flight certification at the Orbital Institute founded in 2185.", + "Planetary survey vessel Solaris-4 mapped asteroid belts in the Gliese system under Commander Maya Lin, who graduated from the Ceres Technical Academy founded in 2240.", + "Cargo transport vessel Atlas-7 transported hydrogen fuel cells across outer lunar outposts under Navigator Eric Zhao, trained at the Lunar Flight School established in 2199.", + "Atmospheric probing vessel Zephyr-1 investigated gas giant storm dynamics in the Kepler perimeter under Pilot Laura Gomez, an alumna of the Jovian Navigation College founded in 2275.", + "Deep reconnaissance cruiser Vanguard-6 conducted quantum radar calibration under Captain Julian Ross, who attended the Sirius Space Academy founded in 2310.", + "Mining support cruiser Titan-3 escorted mineral transport convoys along the Kuiper barrier under Officer Daniel Park, who studied at the Neptune Defense Institute founded in 2225.", + "Hydroponics research station Demeter-2 developed closed-loop atmospheric life support systems under Director Sophia Kim, educated at the Bio-Engineering Academy founded in 2260.", + "Orbital telescope array Copernicus-8 tracked gravitational microlensing events under Astrophysicist Marcus Silva, certified by the Stellar Observation College founded in 2305.", + "Planetary defense frigate Aegis-5 patrolled inner asteroid orbital lanes under Commander Nathan Drake, trained at the Martian Aerospace Center established in 2170.", + "Communication relay platform Hermes-9 maintained laser relay arrays across the Sol grid under Technician Elena Cruz, certified by the Quantum Network Academy founded in 2280." + ] + + doc_target_1 = "The exploratory starship Zephyr-9 was commissioned for deep-space cartography under the command of Captain Alyssa Thorne." + doc_target_2 = "Captain Alyssa Thorne completed her advanced astrogation degree at the Pioneer Space Academy, which was founded in 2348 on New Horizon." + + context_paras = list(distractors[:5]) + [doc_target_1] + list(distractors[5:]) + [doc_target_2] + context = "\n\n".join(context_paras) + + prompt = ( + f"{context}\n\n" + "Question: In what year was the academy where the commander of starship Zephyr-9 completed her degree founded?\n" + "Answer: The academy where the commander of starship Zephyr-9 completed her degree was founded in" + ) + + print("=" * 85) + print("NON-PARAMETRIC LONGBENCH MULTI-HOP QA BENCHMARK") + print("=" * 85) + print("Target Year : 2348") + print("Prompt Size : ~1,000 tokens across 12 distinct non-redundant documents") + print("-" * 85) + + for p in percentages: + res = run_test(args.bin, args.target_model, args.draft_model, prompt, p, args.n_gpu_layers, args.n_gpu_layers_draft) + passed = "2348" in res["gen_text"] + status = "PASS" if passed else "FAIL" + print("Keep %3d%% | Kept %4d/%4d tokens | Status: [%-4s] | Output: \"%s\"" % ( + int(p*100), res["kept"], res["total"], status, res["gen_text"].replace("\n", " ")[:45])) + print("=" * 85) + +if __name__ == "__main__": + main() diff --git a/scripts/eval_spec_prefill_longbench_real.py b/scripts/eval_spec_prefill_longbench_real.py new file mode 100644 index 00000000000..0d048f27499 --- /dev/null +++ b/scripts/eval_spec_prefill_longbench_real.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +""" +Evaluate LongBench on non-redundant, realistic multi-page documents: +1. Multi-hop QA across unique distinct documents (no 10x paragraph repetitions!) +2. Long-document comprehension with scattered distractors +""" + +import argparse +import os +import re +import string +import subprocess +import time +from collections import Counter + +def normalize_answer(s): + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + def white_space_fix(text): + return " ".join(text.split()) + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + def lower(text): + return text.lower() + return white_space_fix(remove_articles(remove_punc(lower(s)))) + +def f1_score(prediction, ground_truth): + pred_tokens = normalize_answer(prediction).split() + gt_tokens = normalize_answer(ground_truth).split() + common = Counter(pred_tokens) & Counter(gt_tokens) + num_same = sum(common.values()) + if num_same == 0: + return 0.0 + precision = 1.0 * num_same / len(pred_tokens) + recall = 1.0 * num_same / len(gt_tokens) + return (2 * precision * recall) / (precision + recall) + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + default_tgt = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/4ca720788d1e01f1bff70c033e0d0028fd02e502/Qwen3.8-27B-UD-Q6_K_XL.gguf" + default_dft = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.5-2B-GGUF/snapshots/f6d5376be1edb4d416d56da11e5397a961aca8ae/Qwen3.5-2B-UD-Q4_K_XL.gguf" + + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--target-model", default=default_tgt) + parser.add_argument("-md", "--draft-model", default=default_dft) + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99) + parser.add_argument("-ngld", "--n-gpu-layers-draft", type=int, default=99) + parser.add_argument("--bin", default=default_bin) + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.15,0.08") + return parser.parse_args() + +def run_test(bin_path, target_model, draft_model, prompt, p, ngl, ngld, max_gen=32, chunk_size=32, lookahead=4): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(max_gen), + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + t0 = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + t_ms = (time.perf_counter() - t0) * 1000.0 + + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_ttft = re.search(r"total Time-To-First-Token \(TTFT\) = ([\d\.]+) ms", output) + ttft_ms = float(m_ttft.group(1)) if m_ttft else 0.0 + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept_tokens = int(m_kept.group(1)) if m_kept else 0 + total_tokens = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "ttft_ms": ttft_ms, + "kept": kept_tokens, + "total": total_tokens, + } + +def main(): + args = parse_args() + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + # Realistic non-redundant multi-document LongBench tasks: + # Task 1: HotpotQA multi-hop without duplicate text + doc_apollo = "The Apollo program was the human spaceflight program by NASA. Apollo 11 successfully landed the first humans on the Moon in July 1969, led by Commander Neil Armstrong." + doc_armstrong = "Neil Armstrong was an aeronautical engineer and naval aviator who served in the Korean War. He completed his undergraduate engineering degree at an institution established in Indiana in 1869." + doc_purdue = "Purdue University is a land-grant university in West Lafayette, Indiana, founded in 1869 by benefactor John Purdue. Its engineering school educated Neil Armstrong and Gene Cernan." + + # Non-redundant distractor articles: + distractors = [ + "Project Gemini was NASA's second human spaceflight program, conducting ten crewed flights in 1965 and 1966 to develop space rendezvous and docking techniques.", + "The Saturn V was an American super heavy-lift launch vehicle developed by NASA under Wernher von Braun for the Apollo lunar exploration missions.", + "The Lunar Roving Vehicle was an electric vehicle designed to operate in the low-gravity vacuum of the Moon during Apollo 15, 16, and 17 missions.", + "The Command Module Columbia was the only spacecraft of the Apollo 11 mission to return safely to Earth after splashing down in the Pacific Ocean.", + "Mission Control Center at Lyndon B. Johnson Space Center in Houston managed flight control for all American crewed spaceflights starting from Gemini 4.", + "The Skylab space station orbited Earth from 1973 to 1979, supporting three crewed missions that conducted solar astronomy and biomedical experiments.", + "The Space Shuttle program was NASA's reusable spacecraft system, flying 135 missions between 1981 and 2011 to construct the International Space Station.", + "Alan Shepard became the first American in space during the Mercury-Redstone 3 flight in 1961, piloting the Freedom 7 capsule into a sub-orbital trajectory.", + "John Glenn orbited the Earth three times aboard Friendship 7 in 1962, becoming the first American astronaut to enter Earth orbit.", + "The International Space Station is a modular space station in low Earth orbit, developed through a multinational collaboration including NASA, ESA, and JAXA.", + ] + + # Assemble non-redundant long context: + context_paras = list(distractors[:5]) + [doc_apollo] + list(distractors[5:]) + [doc_armstrong] + context = "\n\n".join(context_paras) + + prompt = ( + f"{context}\n\n" + "Question: In what year was the university where the commander of Apollo 11 completed his degree founded?\n" + "Answer: The university was founded in" + ) + + print("=" * 85) + print("REALISTIC NON-REDUNDANT LONGBENCH MULTI-HOP QA BENCHMARK") + print("=" * 85) + print(f"Context Length: ~1,150 tokens (all non-redundant distinct articles)") + print(f"Ground Truth : 1869") + print("-" * 85) + + for p in percentages: + res = run_test(args.bin, args.target_model, args.draft_model, prompt, p, args.n_gpu_layers, args.n_gpu_layers_draft) + f1 = f1_score(res["gen_text"], "1869") + passed = "1869" in res["gen_text"] + status = "PASS" if passed else "FAIL" + print("Keep %3d%% | Kept %4d/%4d tokens | Status: [%-4s] | F1 = %5.1f%% | Gen: \"%s\"" % ( + int(p*100), res["kept"], res["total"], status, f1 * 100.0, res["gen_text"].replace("\n", " ")[:40])) + print("=" * 85) + +if __name__ == "__main__": + main() diff --git a/scripts/eval_spec_prefill_quality.py b/scripts/eval_spec_prefill_quality.py new file mode 100644 index 00000000000..299bc0214cb --- /dev/null +++ b/scripts/eval_spec_prefill_quality.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +import argparse +import os +import re +import subprocess +import sys +import time + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + + parser = argparse.ArgumentParser(description="Evaluate Speculative Prefill Quality Impact (Needle-in-a-Haystack & QA)") + parser.add_argument("-m", "--target-model", required=True, help="Path to target model GGUF") + parser.add_argument("-mpd", "-md", "--draft-model", "--spec-prefill-model", help="Path to prefill draft model GGUF") + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99, help="GPU layers for target model") + parser.add_argument("-nglpd", "-ngld", "--n-gpu-layers-draft", "--spec-prefill-ngl", type=int, default=99, help="GPU layers for draft model") + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.20,0.15", help="Comma-separated keep percentages") + parser.add_argument("--chunk-size", type=int, default=32, help="Chunk size (default: 32)") + parser.add_argument("--lookahead", type=int, default=4, help="Lookahead count (default: 4)") + parser.add_argument("--context-size", type=int, default=1500, help="Approximate prompt context size in words") + parser.add_argument("--bin", default=default_bin, help="Path to binary") + return parser.parse_args() + +def generate_haystack(target_words): + paragraphs = [ + "In the quiet valleys of the northern province, astronomers built observatories to chart the movement of celestial bodies. " + "Every evening, scholars logged coordinates of distant nebulae and recorded fluctuations in stellar brightness with brass instruments. " + "The archive of records grew into hundreds of bound volumes filled with geometrical proofs and astronomical charts.", + + "Trade routes crisscrossed the continent, carrying spices, textiles, and precious metals between bustling harbor cities and inland markets. " + "Caravans traveled along well-guarded mountain passes, stopping at desert oases to trade horses and exchange news of foreign realms. " + "Merchant guilds maintained detailed ledgers of commerce, recording tariffs and grain prices across maritime hubs.", + + "Architects designed aqueducts and arched bridges to bring clean mountain water into the centers of expanding metropolises. " + "Engineers perfected the composition of volcanic mortar, enabling the construction of domes and bathhouses that endured for centuries. " + "Civic planners organized city blocks into grids surrounding central plazas, facilitating administration and commerce.", + + "The great library of Alexandria preserved parchment rolls of ancient philosophy, drama, geometry, and medicinal botany. " + "Translators worked continuously to transcribe foreign manuscripts into standardized scripts, adding scholarly annotations. " + "Philosophers met under columned porticos to debate the nature of knowledge, virtue, and physical cosmology.", + ] + + text = [] + total = 0 + while total < target_words: + for p in paragraphs: + text.append(p) + total += len(p.split()) + if total >= target_words: + break + return text + +def run_test(bin_path, target_model, draft_model, prompt, p, chunk_size, lookahead, ngl, ngld): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", "32", + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + t0 = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + t_ms = (time.perf_counter() - t0) * 1000.0 + + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_ttft = re.search(r"total Time-To-First-Token \(TTFT\) = ([\d\.]+) ms", output) + ttft_ms = float(m_ttft.group(1)) if m_ttft else 0.0 + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept_tokens = int(m_kept.group(1)) if m_kept else 0 + total_tokens = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "ttft_ms": ttft_ms, + "kept": kept_tokens, + "total": total_tokens, + "raw": output + } + +def main(): + args = parse_args() + draft_model = args.draft_model if args.draft_model else args.target_model + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + needle_secret = "DELTA-7492-OMEGA" + needle_sentence = "IMPORTANT NOTE: The secret project verification passcode is DELTA-7492-OMEGA." + + depths = [0.20, 0.50, 0.80] # 20% (early), 50% (middle), 80% (late in document) + + print("=" * 85) + print("SPECULATIVE PREFILL QUALITY EVALUATION: NEEDLE-IN-A-HAYSTACK (NIAH)") + print("=" * 85) + print("Target Model : %s" % args.target_model) + print("Draft Model : %s" % draft_model) + print("Target Needle: \"%s\"" % needle_secret) + print("Depths Tested: %s" % [("%d%%" % int(d*100)) for d in depths]) + print("-" * 85) + + haystack_paras = generate_haystack(args.context_size) + + results_matrix = {} # p -> depth -> dict + + for p in percentages: + results_matrix[p] = {} + print("\nEvaluating Keep Percentage: %d%%..." % int(p*100)) + + for depth in depths: + insert_idx = int(len(haystack_paras) * depth) + test_paras = list(haystack_paras) + test_paras.insert(insert_idx, needle_sentence) + full_context = "\n\n".join(test_paras) + + prompt = ( + "%s\n\n" + "Question: What is the secret project verification passcode mentioned in the text above?\n" + "Answer: The secret project verification passcode is" + ) % full_context + + res = run_test(args.bin, args.target_model, draft_model, prompt, p, args.chunk_size, args.lookahead, args.n_gpu_layers, args.n_gpu_layers_draft) + + passed = ("DELTA" in res["gen_text"]) or ("7492" in res["gen_text"]) or ("OMEGA" in res["gen_text"]) + results_matrix[p][depth] = { + "passed": passed, + "gen": res["gen_text"], + "ttft": res["ttft_ms"], + "kept": res["kept"], + "total": res["total"] + } + status_str = "PASS [FOUND]" if passed else "FAIL [LOST]" + print(" Depth %2d%%: %s (TTFT: %6.1f ms) | Kept: %d/%d | Gen: \"%s\"" % ( + int(depth*100), status_str, res["ttft_ms"], res["kept"], res["total"], res["gen_text"].replace("\n", " ")[:45])) + + print("\n" + "=" * 85) + print("QUALITY ACCURACY SUMMARY TABLE (Needle Retrieval vs Keep Ratio)") + print("=" * 85) + depth_headers = " | ".join([("Depth %2d%%" % int(d*100)) for d in depths]) + print("Keep %% | Tokens Kept | Avg TTFT | %s | Accuracy" % depth_headers) + print("-" * 85) + + for p in percentages: + scores = [results_matrix[p][d]["passed"] for d in depths] + avg_ttft = sum([results_matrix[p][d]["ttft"] for d in depths]) / len(depths) + n_kept = results_matrix[p][depths[0]]["kept"] + n_total = results_matrix[p][depths[0]]["total"] + score_pct = (sum(scores) / len(scores)) * 100.0 + + depth_results = " | ".join(["PASS" if s else "FAIL" for s in scores]) + cfg_name = "%d%%" % int(p*100) + print("%-8s | %4d/%-4d | %7.1f ms | %-26s | %5.1f%%" % ( + cfg_name, n_kept, n_total, avg_ttft, depth_results, score_pct)) + + print("=" * 85) + +if __name__ == "__main__": + main() diff --git a/scripts/eval_spec_prefill_targeted_failures.py b/scripts/eval_spec_prefill_targeted_failures.py new file mode 100644 index 00000000000..2db7a08dc64 --- /dev/null +++ b/scripts/eval_spec_prefill_targeted_failures.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +""" +Diagnostic test script demonstrating hard failure modes for speculative prefill: +1. True Semantic Disconnection Multi-Hop (Hop 2 has zero semantic overlap with the prompt) +2. Multi-Entity / Multi-Fact High-Recall Aggregation (5 distinct facts scattered across 5 chunks) +3. Conflicting Override / Negation +4. Distributed Key-Value Retrieval (10 KV pairs where question queries 3 random pairs) +""" + +import argparse +import os +import re +import subprocess +import time + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + default_tgt = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.8-27B-GGUF/snapshots/4ca720788d1e01f1bff70c033e0d0028fd02e502/Qwen3.8-27B-UD-Q6_K_XL.gguf" + default_dft = "/home/rocko/.cache/huggingface/hub/models--unsloth--Qwen3.5-2B-GGUF/snapshots/f6d5376be1edb4d416d56da11e5397a961aca8ae/Qwen3.5-2B-UD-Q4_K_XL.gguf" + + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--target-model", default=default_tgt) + parser.add_argument("-md", "--draft-model", default=default_dft) + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99) + parser.add_argument("-ngld", "--n-gpu-layers-draft", type=int, default=99) + parser.add_argument("--bin", default=default_bin) + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.15,0.10") + return parser.parse_args() + +def run_test(bin_path, target_model, draft_model, prompt, p, ngl, ngld, max_gen=64, chunk_size=32, lookahead=4): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(max_gen), + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + t0 = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + t_ms = (time.perf_counter() - t0) * 1000.0 + + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_ttft = re.search(r"total Time-To-First-Token \(TTFT\) = ([\d\.]+) ms", output) + ttft_ms = float(m_ttft.group(1)) if m_ttft else 0.0 + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept_tokens = int(m_kept.group(1)) if m_kept else 0 + total_tokens = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "ttft_ms": ttft_ms, + "kept": kept_tokens, + "total": total_tokens, + "raw": output + } + +def get_distractors(n=20): + text = [ + "The geological survey team conducted detailed mineralogical analyses in the western foothills. Core samples revealed substantial granite and quartz strata with trace deposits of copper ore.", + "Municipal water engineers redesigned the gravity-fed aqueduct conduits. Pressure regulation cisterns and basalt masonry arches ensured uninterrupted flow into central metropolitan fountains.", + "Agricultural registers recorded seasonal barley harvests across forty alluvial valleys. Yields were cataloged in standard bushel units and distributed to state granaries via river barges.", + "Archival preservationists treated ancient parchment scrolls using mild cedar oil emulsions. Manuscript codices from the third century were rebound in pigskin covers and stored in dry vaults.", + "Textile manufacturing facilities calibrated loom tension for fine linen production. Natural madder dye vats produced crimson vestments exported to neighboring maritime principalities.", + "Harbor authorities logged the arrival of merchant frigates carrying spices, porcelain, and timber. Port tariffs were collected in silver bullion and recorded in double-entry ledgers.", + "Astronomers observed the transit of celestial satellites across equatorial constellations. Quadrant measurements were compiled into seasonal navigational ephemerides for mariners." + ] + return [text[i % len(text)] for i in range(n)] + +def test_pure_latent_multihop(): + """ + Question asks: What clearance tier is assigned to Director Vance's facility? + Clue 1 (Paragraph 3): 'Director Vance is the sole administrator of Complex Theta.' (Matches Vance, administrator) + Clue 2 (Paragraph 16): 'Complex Theta is designated as Clearance Tier Cobalt.' (Matches Complex Theta, but NO Vance, NO administrator, NO question terms) + """ + paras = get_distractors(22) + paras.insert(3, "Facility Administration Records: Director Vance is the sole administrator of Complex Theta.") + paras.insert(16, "Security Protocol Index: Complex Theta is designated as Clearance Tier Cobalt.") + + prompt = ( + "%s\n\n" + "Question: What clearance tier is assigned to the complex managed by Director Vance?\n" + "Answer: The clearance tier assigned to the complex managed by Director Vance is Clearance Tier" + ) % ("\n\n".join(paras)) + + return { + "name": "1. True Latent Multi-Hop (Cobalt)", + "prompt": prompt, + "check": lambda g: "cobalt" in g.lower() + } + +def test_multi_fact_recall(): + """ + Question asks to list ALL 5 server port numbers. + Each port is in a completely separate chunk with 0 redundancy. + """ + paras = get_distractors(20) + items = [ + ("Network Configuration A: Server Alpha listens on port 8081.", 2), + ("Network Configuration B: Server Beta listens on port 8082.", 6), + ("Network Configuration C: Server Gamma listens on port 8083.", 10), + ("Network Configuration D: Server Delta listens on port 8084.", 14), + ("Network Configuration E: Server Epsilon listens on port 8085.", 18), + ] + for text, idx in sorted(items, key=lambda x: x[1], reverse=True): + paras.insert(idx, text) + + prompt = ( + "%s\n\n" + "Question: List the port numbers for all five servers (Alpha, Beta, Gamma, Delta, Epsilon) in order:\n" + "Answer: Alpha, Beta, Gamma, Delta, Epsilon port numbers:" + ) % ("\n\n".join(paras)) + + def check(g): + ports = ["8081", "8082", "8083", "8084", "8085"] + return all(p in g for p in ports) + + return { + "name": "2. Multi-Fact Full Recall (All 5 Ports)", + "prompt": prompt, + "check": check + } + +def test_indirect_math(): + """ + Math with separated values: + Val A = 15 + Val B = 25 + Val C = 60 + Sum = 100 + """ + paras = get_distractors(20) + paras.insert(2, "Inventory Section 1: Item Box A contains exactly 15 units.") + paras.insert(9, "Inventory Section 2: Item Box B contains exactly 25 units.") + paras.insert(17, "Inventory Section 3: Item Box C contains exactly 60 units.") + + prompt = ( + "%s\n\n" + "Question: What is the total sum of units contained in Box A, Box B, and Box C combined?\n" + "Answer: The total sum of units in Box A, Box B, and Box C is exactly" + ) % ("\n\n".join(paras)) + + return { + "name": "3. Distributed Summation (15 + 25 + 60 = 100)", + "prompt": prompt, + "check": lambda g: "100" in g or "one hundred" in g.lower() + } + +def main(): + args = parse_args() + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + tasks = [ + test_pure_latent_multihop(), + test_multi_fact_recall(), + test_indirect_math(), + ] + + print("=" * 90) + print("SPECULATIVE PREFILL: TARGETED FAILURE CASE INVESTIGATION") + print("=" * 90) + print("Target Model : %s" % args.target_model) + print("Draft Model : %s" % args.draft_model) + print("Percentages : %s" % percentages) + print("=" * 90) + + results = {p: {} for p in percentages} + + for p in percentages: + print(f"\nEvaluating Keep Ratio: {int(p*100)}% ({p:.2f})...") + for task in tasks: + res = run_test(args.bin, args.target_model, args.draft_model, task["prompt"], p, args.n_gpu_layers, args.n_gpu_layers_draft, max_gen=32) + passed = task["check"](res["gen_text"]) + status = "PASS" if passed else "FAIL" + results[p][task["name"]] = { + "passed": passed, + "ttft": res["ttft_ms"], + "kept": res["kept"], + "total": res["total"], + "gen": res["gen_text"].replace("\n", " ")[:45] + } + print(" %-45s : [%s] (Kept: %4d/%4d) | Gen: \"%s\"" % ( + task["name"], status, res["kept"], res["total"], results[p][task["name"]]["gen"])) + + print("\n" + "=" * 90) + print("SUMMARY COMPARISON MATRIX") + print("=" * 90) + header = "%-10s | %-11s | " + " | ".join([f"Task {i+1}" for i in range(len(tasks))]) + " | %-8s" + print(header % ("Keep %", "Tokens Kept", "Accuracy")) + print("-" * 90) + for p in percentages: + passes = [results[p][t["name"]]["passed"] for t in tasks] + kept = results[p][tasks[0]["name"]]["kept"] + total = results[p][tasks[0]["name"]]["total"] + acc = (sum(passes) / len(passes)) * 100.0 + pass_strs = ["PASS" if s else "FAIL" for s in passes] + row_str = "%-10s | %4d/%-6d | " + " | ".join(["%-6s" for _ in passes]) + " | %5.1f%%" + print(row_str % tuple([f"{int(p*100)}%", kept, total] + pass_strs + [acc])) + print("=" * 90) + +if __name__ == "__main__": + main() diff --git a/scripts/run_paper_benchmarks.py b/scripts/run_paper_benchmarks.py new file mode 100644 index 00000000000..cda17a751ab --- /dev/null +++ b/scripts/run_paper_benchmarks.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +import re +import string +import subprocess +import sys +import time +from collections import Counter + +def normalize_answer(s): + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + def white_space_fix(text): + return " ".join(text.split()) + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + def lower(text): + return text.lower() + return white_space_fix(remove_articles(remove_punc(lower(s)))) + +def f1_score(prediction, ground_truth): + pred_tokens = normalize_answer(prediction).split() + gt_tokens = normalize_answer(ground_truth).split() + common = Counter(pred_tokens) & Counter(gt_tokens) + num_same = sum(common.values()) + if num_same == 0: + return 0.0 + precision = 1.0 * num_same / len(pred_tokens) + recall = 1.0 * num_same / len(gt_tokens) + return (2 * precision * recall) / (precision + recall) + +def parse_args(): + default_bin = "./build-vulkan/bin/llama-speculative-prefill" if os.path.exists("./build-vulkan/bin/llama-speculative-prefill") else "./build/bin/llama-speculative-prefill" + + parser = argparse.ArgumentParser(description="Run ICML 2025 Speculative Prefill Paper Benchmarks (RULER, LongBench QA, and Scaling)") + parser.add_argument("-m", "--target-model", required=True, help="Target model path") + parser.add_argument("-mpd", "-md", "--draft-model", "--spec-prefill-model", required=True, help="Draft model path") + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99) + parser.add_argument("-nglpd", "-ngld", "--n-gpu-layers-draft", "--spec-prefill-ngl", type=int, default=99) + parser.add_argument("--percentages", default="1.0,0.50,0.30,0.15", help="Keep percentages") + parser.add_argument("--chunk-size", type=int, default=32) + parser.add_argument("--lookahead", type=int, default=4) + parser.add_argument("--bin", default=default_bin) + parser.add_argument("--suite", choices=["all", "ruler", "longbench", "scaling"], default="all") + return parser.parse_args() + +def run_inference(bin_path, target_model, draft_model, prompt, p, chunk_size, lookahead, ngl, ngld, max_gen=32): + cmd = [ + bin_path, + "-m", target_model, + "-md", draft_model, + "-ngl", str(ngl), + "-ngld", str(ngld), + "-p", prompt, + "-n", str(max_gen), + "--spec-prefill-percentage", str(p), + "--spec-prefill-chunk-size", str(chunk_size), + "--spec-prefill-lookahead", str(lookahead), + ] + + t0 = time.perf_counter() + proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + t_ms = (time.perf_counter() - t0) * 1000.0 + + output = proc.stdout + "\n" + proc.stderr + + gen_text = "" + m_gen = re.search(r"--- Generation Start ---\n(.*?)\n--- Generation End ---", output, re.DOTALL) + if m_gen: + gen_text = m_gen.group(1).strip() + + m_ttft = re.search(r"total Time-To-First-Token \(TTFT\) = ([\d\.]+) ms", output) + ttft_ms = float(m_ttft.group(1)) if m_ttft else 0.0 + + m_kept = re.search(r"speculative prefill kept (\d+) / (\d+) tokens", output) + kept = int(m_kept.group(1)) if m_kept else 0 + total = int(m_kept.group(2)) if m_kept else 0 + + return { + "gen_text": gen_text, + "ttft_ms": ttft_ms, + "kept": kept, + "total": total, + } + +def run_ruler_suite(args, percentages): + print("\n" + "=" * 90) + print("1. RULER BENCHMARK (Synthetic Context Probing: Single-Needle, Multi-Key, & Variable Tracking)") + print("=" * 90) + + # Generate synthetic distractors + noise = ( + "In the tranquil coastal cities, fishermen departed at dawn to cast broad nets into the deep shoals. " + "They monitored changing tides and ocean currents using lunar navigation tables developed by mariners. " + "The daily catch was unloaded at lively dockside markets where merchants bid for fresh mackerel and tuna. " + ) * 40 # ~1500 tokens + + tasks = [ + { + "name": "NIAH_Single_Needle (Depth 25%)", + "prompt": noise[:500] + "\n\nIMPORTANT: The secret passphrase is ZEPHYR-9042.\n\n" + noise[500:] + "\n\nQuestion: What is the secret passphrase?\nAnswer: The secret passphrase is", + "ground_truth": "ZEPHYR-9042", + "check": lambda g: "ZEPHYR" in g or "9042" in g + }, + { + "name": "NIAH_Single_Needle (Depth 75%)", + "prompt": noise[:1500] + "\n\nIMPORTANT: The access token for vault B is OMEGA-3184.\n\n" + noise[1500:] + "\n\nQuestion: What is the access token for vault B?\nAnswer: The access token for vault B is", + "ground_truth": "OMEGA-3184", + "check": lambda g: "OMEGA" in g or "3184" in g + }, + { + "name": "NIAH_Multi_Key (Key-Value Pairs)", + "prompt": noise[:400] + "\n\nRecord A: city is Prague, code is 101.\n\n" + noise[400:1000] + "\n\nRecord B: city is Kyoto, code is 505.\n\n" + noise[1000:] + "\n\nQuestion: What is the code for the city of Kyoto?\nAnswer: The code for Kyoto is", + "ground_truth": "505", + "check": lambda g: "505" in g + }, + { + "name": "Variable_Tracking (Chain Assignment)", + "prompt": noise[:500] + "\n\nvar_x = 42; var_y = var_x; var_z = var_y; result = var_z + 10;\n\n" + noise[500:] + "\n\nQuestion: What is the final value of result?\nAnswer: The value of result is", + "ground_truth": "52", + "check": lambda g: "52" in g + } + ] + + ruler_results = {p: {} for p in percentages} + + for p in percentages: + print("\nEvaluating RULER at Keep Percentage: %.0f%%..." % (p * 100)) + for task in tasks: + res = run_inference(args.bin, args.target_model, args.draft_model, task["prompt"], p, args.chunk_size, args.lookahead, args.n_gpu_layers, args.n_gpu_layers_draft) + passed = task["check"](res["gen_text"]) + ruler_results[p][task["name"]] = {"passed": passed, "ttft": res["ttft_ms"], "gen": res["gen_text"]} + status = "PASS" if passed else "FAIL" + print(" %-35s : %s | TTFT: %6.1f ms | Output: \"%s\"" % ( + task["name"], status, res["ttft_ms"], res["gen_text"].replace("\n", " ")[:35])) + + # Print RULER Table + print("\n" + "-" * 90) + print("RULER TASK ACCURACY & SPEEDUP SUMMARY:") + print("-" * 90) + task_names = [t["name"] for t in tasks] + header = "%-10s | %-12s | " + " | ".join(["%-12s" for _ in tasks]) + " | %-8s" + print(header % tuple(["Keep %", "Avg TTFT"] + [f"Task {i+1}" for i in range(len(tasks))] + ["Accuracy"])) + print("-" * 90) + base_ttft = sum([ruler_results[1.0][t["name"]]["ttft"] for t in tasks]) / len(tasks) + for p in percentages: + passes = [ruler_results[p][t["name"]]["passed"] for t in tasks] + avg_ttft = sum([ruler_results[p][t["name"]]["ttft"] for t in tasks]) / len(tasks) + acc = (sum(passes) / len(passes)) * 100.0 + pass_strs = ["PASS" if s else "FAIL" for s in passes] + row_str = "%-10s | %7.1f ms | " + " | ".join(["%-12s" for _ in passes]) + " | %5.1f%%" + print(row_str % tuple([f"{int(p*100)}%", avg_ttft] + pass_strs + [acc])) + print("-" * 90) + +def run_longbench_suite(args, percentages): + print("\n" + "=" * 90) + print("2. LONGBENCH MULTI-DOMAIN QA & SUMMARIZATION BENCHMARK") + print("=" * 90) + + # Multi-hop QA sample from HotpotQA & Multi-Doc QA + doc1 = ( + "The Apollo program was the third United States human spaceflight program carried out by NASA. " + "It succeeded in landing the first humans on the Moon in 1969. The mission that accomplished this " + "historic landing was Apollo 11, commanded by Neil Armstrong alongside lunar module pilot Buzz Aldrin. " + ) * 10 + + doc2 = ( + "Neil Armstrong was an American astronaut and aeronautical engineer who became the first person to walk " + "on the Moon on July 20, 1969. Before becoming an astronaut, Armstrong served as a naval aviator in the " + "United States Navy and flew combat missions during the Korean War. He graduated from Purdue University. " + ) * 10 + + doc3 = ( + "Purdue University is a public land-grant research university in West Lafayette, Indiana. " + "Founded in 1869 after benefactor John Purdue donated land and money to establish a college of science, " + "technology, and agriculture, Purdue has educated numerous prominent engineers and twenty-five astronauts. " + ) * 10 + + qa_tasks = [ + { + "dataset": "HotpotQA (Multi-hop QA)", + "context": doc1 + "\n\n" + doc2 + "\n\n" + doc3, + "question": "Which university did the commander of Apollo 11 graduate from?", + "ground_truth": "Purdue University", + }, + { + "dataset": "Qasper (Single-doc QA)", + "context": doc2 + "\n\n" + doc1, + "question": "What military branch did Neil Armstrong serve in before becoming an astronaut?", + "ground_truth": "United States Navy", + }, + { + "dataset": "2WikiMQA (Multi-hop Reasoning)", + "context": doc3 + "\n\n" + doc1, + "question": "In what year was the university that educated the commander of Apollo 11 founded?", + "ground_truth": "1869", + } + ] + + lb_results = {p: {} for p in percentages} + + for p in percentages: + print("\nEvaluating LongBench Tasks at Keep Percentage: %.0f%%..." % (p * 100)) + for task in qa_tasks: + prompt = f"{task['context']}\n\nQuestion: {task['question']}\nAnswer:" + res = run_inference(args.bin, args.target_model, args.draft_model, prompt, p, args.chunk_size, args.lookahead, args.n_gpu_layers, args.n_gpu_layers_draft, max_gen=24) + f1 = f1_score(res["gen_text"], task["ground_truth"]) + lb_results[p][task["dataset"]] = {"f1": f1, "ttft": res["ttft_ms"], "gen": res["gen_text"]} + print(" %-30s : F1 = %5.1f%% | TTFT: %6.1f ms | Gen: \"%s\"" % ( + task["dataset"], f1 * 100.0, res["ttft_ms"], res["gen_text"].replace("\n", " ")[:35])) + + # Print LongBench Table + print("\n" + "-" * 90) + print("LONGBENCH TASK QA F1 SCORE & SPEEDUP SUMMARY:") + print("-" * 90) + header = "%-10s | %-12s | " + " | ".join(["%-20s" for _ in qa_tasks]) + " | %-10s" + print(header % tuple(["Keep %", "Avg TTFT"] + [t["dataset"] for t in qa_tasks] + ["Average F1"])) + print("-" * 90) + for p in percentages: + f1s = [lb_results[p][t["dataset"]]["f1"] * 100.0 for t in qa_tasks] + avg_ttft = sum([lb_results[p][t["dataset"]]["ttft"] for t in qa_tasks]) / len(qa_tasks) + avg_f1 = sum(f1s) / len(f1s) + f1_strs = [f"{score:5.1f}%" for score in f1s] + row_str = "%-10s | %7.1f ms | " + " | ".join(["%-20s" for _ in qa_tasks]) + " | %5.1f%%" + print(row_str % tuple([f"{int(p*100)}%", avg_ttft] + f1_strs + [avg_f1])) + print("-" * 90) + +def run_scaling_suite(args, percentages): + print("\n" + "=" * 90) + print("3. CONTEXT-LENGTH SCALING & EFFICIENCY BENCHMARK (N = 1024 to 8192 Tokens)") + print("=" * 90) + + base_block = ( + "The Renaissance was a fervent period of European cultural, artistic, political and economic rebirth " + "following the Middle Ages. Generally described as taking place from the 14th century to the 17th century, " + "the Renaissance promoted the rediscovery of classical philosophy, literature and art. Some of the greatest " + "thinkers, authors, statesmen, scientists and artists in human history thrived during this era. " + ) + + context_lengths = [1024, 2048, 4096] + scaling_results = {} + + for n_ctx in context_lengths: + repeats = max(1, n_ctx // 50) + prompt = (base_block * repeats)[: n_ctx * 4] + scaling_results[n_ctx] = {} + print(f"\nBenchmarking Prompt Length N ~= {n_ctx} tokens...") + for p in percentages: + res = run_inference(args.bin, args.target_model, args.draft_model, prompt, p, args.chunk_size, args.lookahead, args.n_gpu_layers, args.n_gpu_layers_draft, max_gen=16) + scaling_results[n_ctx][p] = res + speedup = scaling_results[n_ctx][1.0]["ttft_ms"] / res["ttft_ms"] if (1.0 in scaling_results[n_ctx] and res["ttft_ms"] > 0) else 1.0 + print(" Keep %3d%% : TTFT = %6.1f ms | Kept %4d/%-4d | Speedup = %4.2fx" % ( + int(p*100), res["ttft_ms"], res["kept"], res["total"], speedup)) + + print("\n" + "=" * 90) + print("SCALING BENCHMARK SUMMARY (TTFT in milliseconds):") + print("=" * 90) + header = "%-12s | " + " | ".join([f"N = {n:4d} tokens" for n in context_lengths]) + print(header) + print("-" * 90) + for p in percentages: + row = [f"{int(p*100):3d}% Keep"] + for n_ctx in context_lengths: + ttft = scaling_results[n_ctx][p]["ttft_ms"] + base_ttft = scaling_results[n_ctx][1.0]["ttft_ms"] + sp = base_ttft / ttft if ttft > 0 else 1.0 + row.append(f"{ttft:6.1f} ms ({sp:4.2f}x)") + print("%-12s | " % row[0] + " | ".join(["%-17s" % cell for cell in row[1:]])) + print("=" * 90) + +def main(): + args = parse_args() + percentages = [float(p.strip()) for p in args.percentages.split(",") if p.strip()] + + print("*" * 90) + print("REPRODUCING ICML 2025 SPECULATIVE PREFILL PAPER BENCHMARKS ON VULKAN") + print("Target Model : %s" % args.target_model) + print("Draft Model : %s" % args.draft_model) + print("Backend : Vulkan (AMD Radeon 8060S Graphics)") + print("Percentages : %s" % percentages) + print("*" * 90) + + if args.suite in ["all", "ruler"]: + run_ruler_suite(args, percentages) + + if args.suite in ["all", "longbench"]: + run_longbench_suite(args, percentages) + + if args.suite in ["all", "scaling"]: + run_scaling_suite(args, percentages) + +if __name__ == "__main__": + main() diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9aed8013327..0842d2fd1e8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1158,6 +1158,17 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void } } +void llama_context::set_eval_callback(ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data) { + LLAMA_LOG_DEBUG("%s: call\n", __func__); + + cparams.cb_eval = cb_eval; + cparams.cb_eval_user_data = cb_eval_user_data; + + if (sched) { + ggml_backend_sched_set_eval_callback(sched.get(), cb_eval, cb_eval_user_data); + } +} + void llama_context::set_embeddings(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); @@ -3799,6 +3810,14 @@ void llama_set_abort_callback(llama_context * ctx, bool (*abort_callback)(void * ctx->set_abort_callback(abort_callback, abort_callback_data); } +void llama_set_eval_callback(llama_context * ctx, ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data) { + if (!ctx) { + return; + } + + ctx->set_eval_callback(cb_eval, cb_eval_user_data); +} + void llama_set_embeddings(llama_context * ctx, bool embeddings) { ctx->set_embeddings(embeddings); } diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b56..b089e1267aa 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -111,6 +111,7 @@ struct llama_context { void set_n_threads(int32_t n_threads, int32_t n_threads_batch); void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data); + void set_eval_callback(ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data); void set_embeddings (bool value); void set_embeddings_nextn(bool value, bool masked); diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index e0907631abd..b823858405e 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -254,6 +254,19 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); assert(params.speculative.draft.n_max == 123); + argv = {"binary_name", "-mpd", "prefill-draft.gguf", "-nglpd", "24", "-devpd", "none", "-cpd", "4096"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); + assert(params.speculative.prefill.model.path == "prefill-draft.gguf"); + assert(params.speculative.prefill.n_gpu_layers == 24); + assert(params.speculative.prefill.devices.size() == 1 && params.speculative.prefill.devices[0] == nullptr); + assert(params.speculative.prefill.n_ctx == 4096); + assert(params.speculative.prefill.enabled == true); + + argv = {"binary_name", "--spec-prefill-device", "none"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); + assert(params.speculative.prefill.devices.size() == 1 && params.speculative.prefill.devices[0] == nullptr); + assert(params.speculative.prefill.enabled == true); + { common_params synth_params; argv = {"binary_name", "--spec-synth-len", "3.4"}; @@ -273,7 +286,6 @@ static void test(void) { argv = {"binary_name", "--spec-synth-len", "3.4x"}; assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); } - argv = {"binary_name", "-lm", "none"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_NONE); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 1b4bbde4d4f..11d1e0a353b 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -27,6 +27,8 @@ #include "ggml.h" #include "llama.h" #include "log.h" +#include "sampling.h" +#include "speculative-prefill.h" #ifdef _WIN32 # define WIN32_LEAN_AND_MEAN @@ -332,12 +334,34 @@ static std::vector parse_int_range(const std::string & s, bool allow_negati return result; } +static std::vector parse_float_list(const std::string & s) { + auto p = string_split(s, ','); + std::vector result; + for (const auto & item : p) { + float val = std::stof(item); + if (val <= 0.0f || val > 1.0f) { + throw std::invalid_argument("speculative prefill percentage must be between 0.0 (exclusive) and 1.0 (inclusive)"); + } + result.push_back(val); + } + return result; +} + struct cmd_params { std::vector model; std::vector hf_repo; std::vector hf_file; std::string hf_token; bool offline; + std::vector spec_prefill_model; + std::vector spec_prefill_hf_repo; + std::vector spec_prefill_hf_file; + std::vector spec_prefill_n_ctx; + std::vector spec_prefill_n_gpu_layers; + std::vector spec_prefill_percentage; + std::vector spec_prefill_chunk_size; + std::vector spec_prefill_lookahead; + std::vector spec_prefill_pool_kernel; std::vector n_prompt; std::vector n_gen; std::vector> n_pg; @@ -378,48 +402,57 @@ struct cmd_params { }; static const cmd_params cmd_params_defaults = { - /* model */ { "models/7B/ggml-model-q4_0.gguf" }, - /* hf_repo */ {}, - /* hf_file */ {}, - /* hf_token */ "", - /* offline */ false, - /* n_prompt */ { 512 }, - /* n_gen */ { 128 }, - /* n_pg */ {}, - /* n_depth */ { 0 }, - /* n_batch */ { 2048 }, - /* n_ubatch */ { 512 }, - /* type_k */ { GGML_TYPE_F16 }, - /* type_v */ { GGML_TYPE_F16 }, - /* n_threads */ { common_cpu_get_num_math() }, - /* cpu_mask */ { "0x0" }, - /* cpu_strict */ { false }, - /* poll */ { 50 }, - /* n_gpu_layers */ { -1 }, - /* n_cpu_moe */ { 0 }, - /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, - /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, - /* lazy_mode */ { LLAMA_LAZY_MODE_AUTO }, - /* main_gpu */ { 0 }, - /* no_kv_offload */ { false }, - /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, - /* devices */ { {} }, - /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, - /* tensor_buft_overrides*/ { std::vector{ { nullptr, nullptr } } }, - /* embeddings */ { false }, - /* no_op_offload */ { false }, - /* no_host */ { false }, - /* fit_params_target */ { 0 }, - /* fit_params_min_ctx */ { 0 }, - /* numa */ GGML_NUMA_STRATEGY_DISABLED, - /* reps */ 5, - /* prio */ GGML_SCHED_PRIO_NORMAL, - /* delay */ 0, - /* verbose */ false, - /* progress */ false, - /* no_warmup */ false, - /* output_format */ MARKDOWN, - /* output_format_stderr */ NONE, + /* model */ { "models/7B/ggml-model-q4_0.gguf" }, + /* hf_repo */ {}, + /* hf_file */ {}, + /* hf_token */ "", + /* offline */ false, + /* spec_prefill_model */ {}, + /* spec_prefill_hf_repo */ {}, + /* spec_prefill_hf_file */ {}, + /* spec_prefill_n_ctx */ { 0 }, + /* spec_prefill_n_gpu_layers */ { -1 }, + /* spec_prefill_percentage */ { 0.30f }, + /* spec_prefill_chunk_size */ { 32 }, + /* spec_prefill_lookahead */ { 8 }, + /* spec_prefill_pool_kernel */ { 13 }, + /* n_prompt */ { 512 }, + /* n_gen */ { 128 }, + /* n_pg */ {}, + /* n_depth */ { 0 }, + /* n_batch */ { 2048 }, + /* n_ubatch */ { 512 }, + /* type_k */ { GGML_TYPE_F16 }, + /* type_v */ { GGML_TYPE_F16 }, + /* n_threads */ { common_cpu_get_num_math() }, + /* cpu_mask */ { "0x0" }, + /* cpu_strict */ { false }, + /* poll */ { 50 }, + /* n_gpu_layers */ { -1 }, + /* n_cpu_moe */ { 0 }, + /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, + /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, + /* lazy_mode */ { LLAMA_LAZY_MODE_AUTO }, + /* main_gpu */ { 0 }, + /* no_kv_offload */ { false }, + /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, + /* devices */ { {} }, + /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, + /* tensor_buft_overrides */ { std::vector{ { nullptr, nullptr } } }, + /* embeddings */ { false }, + /* no_op_offload */ { false }, + /* no_host */ { false }, + /* fit_params_target */ { 0 }, + /* fit_params_min_ctx */ { 0 }, + /* numa */ GGML_NUMA_STRATEGY_DISABLED, + /* reps */ 5, + /* prio */ GGML_SCHED_PRIO_NORMAL, + /* delay */ 0, + /* verbose */ false, + /* progress */ false, + /* no_warmup */ false, + /* output_format */ MARKDOWN, + /* output_format_stderr */ NONE, }; static void print_usage(int /* argc */, char ** argv) { @@ -455,6 +488,24 @@ static void print_usage(int /* argc */, char ** argv) { printf(" (default: value from HF_TOKEN environment variable)\n"); printf(" --offline Offline mode: forces use of cache, prevents network access\n"); printf(" (default: disabled)\n"); + printf(" -spf, --spec-prefill, --speculative-prefill <0|1|on|off>\n"); + printf(" enable speculative prefill (default: disabled)\n"); + printf(" -mpd, --spec-prefill-model, --spec-prefill-draft-model \n"); + printf(" draft model for speculative prefill (default: unused)\n"); + printf(" -hfpd, --spec-prefill-hf, --spec-prefill-draft-hf /[:quant]\n"); + printf(" Hugging Face draft model repo for speculative prefill (default: unused)\n"); + printf(" -hffpd, --spec-prefill-hf-file Hugging Face draft model file for speculative prefill (default: unused)\n"); + printf(" -cpd, --spec-prefill-ctx, --spec-prefill-draft-ctx \n"); + printf(" draft context size for speculative prefill (default: %s)\n", join(cmd_params_defaults.spec_prefill_n_ctx, ",").c_str()); + printf(" -nglpd, --spec-prefill-ngl, --spec-prefill-draft-ngl \n"); + printf(" GPU layers for speculative prefill draft model (default: %s)\n", join(cmd_params_defaults.spec_prefill_n_gpu_layers, ",").c_str()); + printf(" -spfp, --spec-prefill-p, --spec-prefill-percentage

\n"); + printf(" token retention percentage for speculative prefill (default: %.2f)\n", cmd_params_defaults.spec_prefill_percentage[0]); + printf(" -spfcs, --spec-prefill-chunk, --spec-prefill-chunk-size \n"); + printf(" chunk size for speculative prefill (default: %s)\n", join(cmd_params_defaults.spec_prefill_chunk_size, ",").c_str()); + printf(" -spflah, --spec-prefill-lookahead, --spec-prefill-lah \n"); + printf(" lookahead steps for speculative prefill (default: %s)\n", join(cmd_params_defaults.spec_prefill_lookahead, ",").c_str()); + printf(" -spfpool, --spec-prefill-pool-kernel attention pooling kernel size for speculative prefill (default: %s)\n", join(cmd_params_defaults.spec_prefill_pool_kernel, ",").c_str()); printf(" -p, --n-prompt (default: %s)\n", join(cmd_params_defaults.n_prompt, ",").c_str()); printf(" -n, --n-gen (default: %s)\n", join(cmd_params_defaults.n_gen, ",").c_str()); printf(" -pg (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str()); @@ -581,6 +632,90 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { params.hf_token = argv[i]; } else if (arg == "--offline") { params.offline = true; + } else if (arg == "-mpd" || arg == "--spec-prefill-model" || arg == "--spec-prefill-draft-model" || + arg == "--speculative-prefill-model" || arg == "--speculative-prefill-draft-model") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + params.spec_prefill_model.insert(params.spec_prefill_model.end(), p.begin(), p.end()); + } else if (arg == "-hfpd" || arg == "--spec-prefill-hf" || arg == "--spec-prefill-draft-hf" || + arg == "--speculative-prefill-hf" || arg == "--speculative-prefill-draft-hf") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + params.spec_prefill_hf_repo.insert(params.spec_prefill_hf_repo.end(), p.begin(), p.end()); + } else if (arg == "-hffpd" || arg == "--spec-prefill-hf-file" || arg == "--spec-prefill-draft-hf-file") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + params.spec_prefill_hf_file.insert(params.spec_prefill_hf_file.end(), p.begin(), p.end()); + } else if (arg == "-cpd" || arg == "--spec-prefill-ctx" || arg == "--spec-prefill-draft-ctx" || + arg == "--spec-prefill-ctx-size" || arg == "--spec-prefill-max-ctx" || + arg == "--speculative-prefill-ctx" || arg == "--speculative-prefill-max-ctx") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + params.spec_prefill_n_ctx.insert(params.spec_prefill_n_ctx.end(), p.begin(), p.end()); + } else if (arg == "-nglpd" || arg == "--spec-prefill-ngl" || arg == "--spec-prefill-draft-ngl" || + arg == "--speculative-prefill-ngl" || arg == "--speculative-prefill-draft-ngl") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i], /*allow_negative=*/true); + params.spec_prefill_n_gpu_layers.insert(params.spec_prefill_n_gpu_layers.end(), p.begin(), p.end()); + } else if (arg == "-spfp" || arg == "--spec-prefill-p" || arg == "--spec-prefill-percentage") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_float_list(argv[i]); + params.spec_prefill_percentage.insert(params.spec_prefill_percentage.end(), p.begin(), p.end()); + } else if (arg == "-spfcs" || arg == "--spec-prefill-chunk" || arg == "--spec-prefill-chunk-size") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + params.spec_prefill_chunk_size.insert(params.spec_prefill_chunk_size.end(), p.begin(), p.end()); + } else if (arg == "-spflah" || arg == "--spec-prefill-lookahead" || arg == "--spec-prefill-lah") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + params.spec_prefill_lookahead.insert(params.spec_prefill_lookahead.end(), p.begin(), p.end()); + } else if (arg == "-spfpool" || arg == "--spec-prefill-pool-kernel") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = parse_int_range(argv[i]); + params.spec_prefill_pool_kernel.insert(params.spec_prefill_pool_kernel.end(), p.begin(), p.end()); + } else if (arg == "-spf" || arg == "--spec-prefill" || arg == "--speculative-prefill") { + if (i + 1 < argc && argv[i + 1][0] != '-') { + i++; + auto p = string_split(argv[i], split_delim); + for (const auto & item : p) { + if (item == "0" || item == "off" || item == "false" || item == "none" || item == "disable" || item == "disabled") { + params.spec_prefill_model.push_back(""); + } else if (item == "1" || item == "on" || item == "true" || item == "enable" || item == "enabled" || item == "self") { + params.spec_prefill_model.push_back("self"); + } else { + params.spec_prefill_model.push_back(item); + } + } + } else { + params.spec_prefill_model.push_back("self"); + } } else if (arg == "-p" || arg == "--n-prompt") { if (++i >= argc) { invalid_param = true; @@ -1139,10 +1274,56 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } } + if (!params.spec_prefill_hf_repo.empty()) { + for (size_t i = 0; i < params.spec_prefill_hf_repo.size(); i++) { + common_params p; + p.hf_token = params.hf_token; + p.offline = params.offline; + p.model.hf_repo = params.spec_prefill_hf_repo[i]; + if (!params.spec_prefill_hf_file.empty() && !params.spec_prefill_hf_file[i].empty()) { + p.model.hf_file = params.spec_prefill_hf_file[i]; + } + + // only the text model file is needed + common_models_handler models_handler = common_models_handler_init(p, LLAMA_EXAMPLE_BENCH); + common_models_handler_apply(models_handler, p); + if (p.model.path.empty()) { + fprintf(stderr, "error: failed to download speculative prefill draft model from HuggingFace\n"); + exit(1); + } + + params.spec_prefill_model.push_back(p.model.path); + } + } + // set defaults if (params.model.empty()) { params.model = cmd_params_defaults.model; } + if (params.spec_prefill_model.empty()) { + params.spec_prefill_model = { "" }; + } + if (params.spec_prefill_hf_file.empty()) { + params.spec_prefill_hf_file = cmd_params_defaults.spec_prefill_hf_file; + } + if (params.spec_prefill_n_ctx.empty()) { + params.spec_prefill_n_ctx = cmd_params_defaults.spec_prefill_n_ctx; + } + if (params.spec_prefill_n_gpu_layers.empty()) { + params.spec_prefill_n_gpu_layers = cmd_params_defaults.spec_prefill_n_gpu_layers; + } + if (params.spec_prefill_percentage.empty()) { + params.spec_prefill_percentage = cmd_params_defaults.spec_prefill_percentage; + } + if (params.spec_prefill_chunk_size.empty()) { + params.spec_prefill_chunk_size = cmd_params_defaults.spec_prefill_chunk_size; + } + if (params.spec_prefill_lookahead.empty()) { + params.spec_prefill_lookahead = cmd_params_defaults.spec_prefill_lookahead; + } + if (params.spec_prefill_pool_kernel.empty()) { + params.spec_prefill_pool_kernel = cmd_params_defaults.spec_prefill_pool_kernel; + } if (params.n_prompt.empty()) { params.n_prompt = cmd_params_defaults.n_prompt; } @@ -1233,6 +1414,13 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { struct cmd_params_instance { std::string model; + std::string spec_prefill_model; + int spec_prefill_n_ctx; + int spec_prefill_n_gpu_layers; + float spec_prefill_percentage; + int spec_prefill_chunk_size; + int spec_prefill_lookahead; + int spec_prefill_pool_kernel; int n_prompt; int n_gen; int n_depth; @@ -1347,6 +1535,13 @@ static std::vector get_cmd_params_instances(const cmd_param // this ordering minimizes the number of times that each model needs to be reloaded // clang-format off for (const auto & m : params.model) + for (const auto & mpd : params.spec_prefill_model) + for (const auto & spfc : (mpd.empty() || mpd == "none" ? std::vector{0} : params.spec_prefill_n_ctx)) + for (const auto & nglpd : (mpd.empty() || mpd == "none" ? std::vector{-1} : params.spec_prefill_n_gpu_layers)) + for (const auto & spfp : (mpd.empty() || mpd == "none" ? std::vector{0.30f} : params.spec_prefill_percentage)) + for (const auto & spfcs : (mpd.empty() || mpd == "none" ? std::vector{32} : params.spec_prefill_chunk_size)) + for (const auto & spflah : (mpd.empty() || mpd == "none" ? std::vector{8} : params.spec_prefill_lookahead)) + for (const auto & spfpool : (mpd.empty() || mpd == "none" ? std::vector{13} : params.spec_prefill_pool_kernel)) for (const auto & fpt : params.fit_params_target) for (const auto & fpc : params.fit_params_min_ctx) for (const auto & nl : params.n_gpu_layers) @@ -1377,34 +1572,41 @@ static std::vector get_cmd_params_instances(const cmd_param continue; } cmd_params_instance instance = { - /* .model = */ m, - /* .n_prompt = */ n_prompt, - /* .n_gen = */ 0, - /* .n_depth = */ nd, - /* .n_batch = */ nb, - /* .n_ubatch = */ nub, - /* .type_k = */ tk, - /* .type_v = */ tv, - /* .n_threads = */ nt, - /* .cpu_mask = */ cm, - /* .cpu_strict = */ cs, - /* .poll = */ pl, - /* .n_gpu_layers = */ nl, - /* .n_cpu_moe = */ ncmoe, - /* .split_mode = */ sm, - /* .load_mode = */ lm, - /* .lazy_mode = */ lzm, - /* .main_gpu = */ mg, - /* .no_kv_offload = */ nkvo, - /* .flash_attn = */ fa, - /* .devices = */ devs, - /* .tensor_split = */ ts, - /* .tensor_buft_overrides = */ ot, - /* .embeddings = */ embd, - /* .no_op_offload = */ nopo, - /* .no_host = */ noh, - /* .fit_target = */ fpt, - /* .fit_min_ctx = */ fpc, + /* .model = */ m, + /* .spec_prefill_model = */ mpd, + /* .spec_prefill_n_ctx = */ spfc, + /* .spec_prefill_n_gpu_layers = */ nglpd, + /* .spec_prefill_percentage = */ spfp, + /* .spec_prefill_chunk_size = */ spfcs, + /* .spec_prefill_lookahead = */ spflah, + /* .spec_prefill_pool_kernel = */ spfpool, + /* .n_prompt = */ n_prompt, + /* .n_gen = */ 0, + /* .n_depth = */ nd, + /* .n_batch = */ nb, + /* .n_ubatch = */ nub, + /* .type_k = */ tk, + /* .type_v = */ tv, + /* .n_threads = */ nt, + /* .cpu_mask = */ cm, + /* .cpu_strict = */ cs, + /* .poll = */ pl, + /* .n_gpu_layers = */ nl, + /* .n_cpu_moe = */ ncmoe, + /* .split_mode = */ sm, + /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, + /* .main_gpu = */ mg, + /* .no_kv_offload = */ nkvo, + /* .flash_attn = */ fa, + /* .devices = */ devs, + /* .tensor_split = */ ts, + /* .tensor_buft_overrides = */ ot, + /* .embeddings = */ embd, + /* .no_op_offload = */ nopo, + /* .no_host = */ noh, + /* .fit_target = */ fpt, + /* .fit_min_ctx = */ fpc, }; instances.push_back(instance); } @@ -1414,34 +1616,41 @@ static std::vector get_cmd_params_instances(const cmd_param continue; } cmd_params_instance instance = { - /* .model = */ m, - /* .n_prompt = */ 0, - /* .n_gen = */ n_gen, - /* .n_depth = */ nd, - /* .n_batch = */ nb, - /* .n_ubatch = */ nub, - /* .type_k = */ tk, - /* .type_v = */ tv, - /* .n_threads = */ nt, - /* .cpu_mask = */ cm, - /* .cpu_strict = */ cs, - /* .poll = */ pl, - /* .n_gpu_layers = */ nl, - /* .n_cpu_moe = */ ncmoe, - /* .split_mode = */ sm, - /* .load_mode = */ lm, - /* .lazy_mode = */ lzm, - /* .main_gpu = */ mg, - /* .no_kv_offload = */ nkvo, - /* .flash_attn = */ fa, - /* .devices = */ devs, - /* .tensor_split = */ ts, - /* .tensor_buft_overrides = */ ot, - /* .embeddings = */ embd, - /* .no_op_offload = */ nopo, - /* .no_host = */ noh, - /* .fit_target = */ fpt, - /* .fit_min_ctx = */ fpc, + /* .model = */ m, + /* .spec_prefill_model = */ mpd, + /* .spec_prefill_n_ctx = */ spfc, + /* .spec_prefill_n_gpu_layers = */ nglpd, + /* .spec_prefill_percentage = */ spfp, + /* .spec_prefill_chunk_size = */ spfcs, + /* .spec_prefill_lookahead = */ spflah, + /* .spec_prefill_pool_kernel = */ spfpool, + /* .n_prompt = */ 0, + /* .n_gen = */ n_gen, + /* .n_depth = */ nd, + /* .n_batch = */ nb, + /* .n_ubatch = */ nub, + /* .type_k = */ tk, + /* .type_v = */ tv, + /* .n_threads = */ nt, + /* .cpu_mask = */ cm, + /* .cpu_strict = */ cs, + /* .poll = */ pl, + /* .n_gpu_layers = */ nl, + /* .n_cpu_moe = */ ncmoe, + /* .split_mode = */ sm, + /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, + /* .main_gpu = */ mg, + /* .no_kv_offload = */ nkvo, + /* .flash_attn = */ fa, + /* .devices = */ devs, + /* .tensor_split = */ ts, + /* .tensor_buft_overrides = */ ot, + /* .embeddings = */ embd, + /* .no_op_offload = */ nopo, + /* .no_host = */ noh, + /* .fit_target = */ fpt, + /* .fit_min_ctx = */ fpc, }; instances.push_back(instance); } @@ -1451,34 +1660,41 @@ static std::vector get_cmd_params_instances(const cmd_param continue; } cmd_params_instance instance = { - /* .model = */ m, - /* .n_prompt = */ n_pg.first, - /* .n_gen = */ n_pg.second, - /* .n_depth = */ nd, - /* .n_batch = */ nb, - /* .n_ubatch = */ nub, - /* .type_k = */ tk, - /* .type_v = */ tv, - /* .n_threads = */ nt, - /* .cpu_mask = */ cm, - /* .cpu_strict = */ cs, - /* .poll = */ pl, - /* .n_gpu_layers = */ nl, - /* .n_cpu_moe = */ ncmoe, - /* .split_mode = */ sm, - /* .load_mode = */ lm, - /* .lazy_mode = */ lzm, - /* .main_gpu = */ mg, - /* .no_kv_offload = */ nkvo, - /* .flash_attn = */ fa, - /* .devices = */ devs, - /* .tensor_split = */ ts, - /* .tensor_buft_overrides = */ ot, - /* .embeddings = */ embd, - /* .no_op_offload = */ nopo, - /* .no_host = */ noh, - /* .fit_target = */ fpt, - /* .fit_min_ctx = */ fpc, + /* .model = */ m, + /* .spec_prefill_model = */ mpd, + /* .spec_prefill_n_ctx = */ spfc, + /* .spec_prefill_n_gpu_layers = */ nglpd, + /* .spec_prefill_percentage = */ spfp, + /* .spec_prefill_chunk_size = */ spfcs, + /* .spec_prefill_lookahead = */ spflah, + /* .spec_prefill_pool_kernel = */ spfpool, + /* .n_prompt = */ n_pg.first, + /* .n_gen = */ n_pg.second, + /* .n_depth = */ nd, + /* .n_batch = */ nb, + /* .n_ubatch = */ nub, + /* .type_k = */ tk, + /* .type_v = */ tv, + /* .n_threads = */ nt, + /* .cpu_mask = */ cm, + /* .cpu_strict = */ cs, + /* .poll = */ pl, + /* .n_gpu_layers = */ nl, + /* .n_cpu_moe = */ ncmoe, + /* .split_mode = */ sm, + /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, + /* .main_gpu = */ mg, + /* .no_kv_offload = */ nkvo, + /* .flash_attn = */ fa, + /* .devices = */ devs, + /* .tensor_split = */ ts, + /* .tensor_buft_overrides = */ ot, + /* .embeddings = */ embd, + /* .no_op_offload = */ nopo, + /* .no_host = */ noh, + /* .fit_target = */ fpt, + /* .fit_min_ctx = */ fpc, }; instances.push_back(instance); } @@ -1521,6 +1737,12 @@ struct test { bool no_host; size_t fit_target; uint32_t fit_min_ctx; + std::string spec_prefill_model; + int spec_prefill_n_gpu_layers; + float spec_prefill_percentage; + int spec_prefill_chunk_size; + int spec_prefill_lookahead; + int spec_prefill_pool_kernel; int n_prompt; int n_gen; int n_depth; @@ -1555,12 +1777,18 @@ struct test { flash_attn = inst.flash_attn; devices = inst.devices; tensor_split = inst.tensor_split; - tensor_buft_overrides = inst.tensor_buft_overrides; - embeddings = inst.embeddings; - no_op_offload = inst.no_op_offload; - no_host = inst.no_host; - fit_target = inst.fit_target; - fit_min_ctx = inst.fit_min_ctx; + tensor_buft_overrides = inst.tensor_buft_overrides; + embeddings = inst.embeddings; + no_op_offload = inst.no_op_offload; + no_host = inst.no_host; + fit_target = inst.fit_target; + fit_min_ctx = inst.fit_min_ctx; + spec_prefill_model = inst.spec_prefill_model; + spec_prefill_n_gpu_layers = inst.spec_prefill_n_gpu_layers; + spec_prefill_percentage = inst.spec_prefill_percentage; + spec_prefill_chunk_size = inst.spec_prefill_chunk_size; + spec_prefill_lookahead = inst.spec_prefill_lookahead; + spec_prefill_pool_kernel = inst.spec_prefill_pool_kernel; n_prompt = inst.n_prompt; n_gen = inst.n_gen; n_depth = inst.n_depth; @@ -1620,6 +1848,8 @@ struct test { "tensor_buft_overrides", "load_mode", "lazy_mode", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", + "spec_prefill_model", "spec_prefill_n_gpu_layers", "spec_prefill_percentage", + "spec_prefill_chunk_size", "spec_prefill_lookahead", "spec_prefill_pool_kernel", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" }; @@ -1633,17 +1863,19 @@ struct test { field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" || field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" || - field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") { + field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn" || + field == "spec_prefill_n_gpu_layers" || field == "spec_prefill_chunk_size" || + field == "spec_prefill_lookahead" || field == "spec_prefill_pool_kernel") { return INT; } if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" || field == "embeddings" || field == "no_host") { return BOOL; } - if (field == "avg_ts" || field == "stddev_ts") { + if (field == "avg_ts" || field == "stddev_ts" || field == "spec_prefill_percentage") { return FLOAT; } - if (field == "load_mode" || field == "lazy_mode") { + if (field == "load_mode" || field == "lazy_mode" || field == "spec_prefill_model") { return STRING; } return STRING; @@ -1719,6 +1951,12 @@ struct test { std::to_string(no_host), std::to_string(fit_target), std::to_string(fit_min_ctx), + spec_prefill_model, + std::to_string(spec_prefill_n_gpu_layers), + std::to_string(spec_prefill_percentage), + std::to_string(spec_prefill_chunk_size), + std::to_string(spec_prefill_lookahead), + std::to_string(spec_prefill_pool_kernel), std::to_string(n_prompt), std::to_string(n_gen), std::to_string(n_depth), @@ -1910,6 +2148,16 @@ struct markdown_printer : public printer { if (field == "no_host") { return 4; } + if (field == "spec_prefill_model") { + return -20; + } + if (field == "spec_prefill_n_gpu_layers" || field == "spec_prefill_chunk_size" || + field == "spec_prefill_lookahead" || field == "spec_prefill_pool_kernel") { + return 5; + } + if (field == "spec_prefill_percentage") { + return 6; + } int width = std::max((int) field.length(), 10); @@ -1962,6 +2210,24 @@ struct markdown_printer : public printer { if (field == "fit_min_ctx") { return "fitc"; } + if (field == "spec_prefill_model") { + return "draft_model"; + } + if (field == "spec_prefill_n_gpu_layers") { + return "nglpd"; + } + if (field == "spec_prefill_percentage") { + return "spf_p"; + } + if (field == "spec_prefill_chunk_size") { + return "spf_chunk"; + } + if (field == "spec_prefill_lookahead") { + return "spf_lah"; + } + if (field == "spec_prefill_pool_kernel") { + return "spf_pool"; + } return field; } @@ -2046,6 +2312,31 @@ struct markdown_printer : public printer { if (params.fit_params_min_ctx.size() > 1 || params.fit_params_min_ctx != cmd_params_defaults.fit_params_min_ctx) { fields.emplace_back("fit_min_ctx"); } + bool has_any_spec_prefill = false; + for (const auto & mpd : params.spec_prefill_model) { + if (!mpd.empty() && mpd != "none") { + has_any_spec_prefill = true; + break; + } + } + if (has_any_spec_prefill) { + fields.emplace_back("spec_prefill_model"); + if (params.spec_prefill_n_gpu_layers.size() > 1 || params.spec_prefill_n_gpu_layers != cmd_params_defaults.spec_prefill_n_gpu_layers) { + fields.emplace_back("spec_prefill_n_gpu_layers"); + } + if (params.spec_prefill_percentage.size() > 1 || params.spec_prefill_percentage != cmd_params_defaults.spec_prefill_percentage) { + fields.emplace_back("spec_prefill_percentage"); + } + if (params.spec_prefill_chunk_size.size() > 1 || params.spec_prefill_chunk_size != cmd_params_defaults.spec_prefill_chunk_size) { + fields.emplace_back("spec_prefill_chunk_size"); + } + if (params.spec_prefill_lookahead.size() > 1 || params.spec_prefill_lookahead != cmd_params_defaults.spec_prefill_lookahead) { + fields.emplace_back("spec_prefill_lookahead"); + } + if (params.spec_prefill_pool_kernel.size() > 1 || params.spec_prefill_pool_kernel != cmd_params_defaults.spec_prefill_pool_kernel) { + fields.emplace_back("spec_prefill_pool_kernel"); + } + } fields.emplace_back("test"); fields.emplace_back("t/s"); @@ -2100,6 +2391,22 @@ struct markdown_printer : public printer { snprintf(buf + len, sizeof(buf) - len, " @ d%d", t.n_depth); } value = buf; + } else if (field == "spec_prefill_model") { + value = t.spec_prefill_model.empty() || t.spec_prefill_model == "none" ? "-" : t.spec_prefill_model; + } else if (field == "spec_prefill_percentage") { + if (t.spec_prefill_model.empty() || t.spec_prefill_model == "none") { + value = "-"; + } else { + snprintf(buf, sizeof(buf), "%.2f", (double) t.spec_prefill_percentage); + value = buf; + } + } else if (field == "spec_prefill_n_gpu_layers" || field == "spec_prefill_chunk_size" || + field == "spec_prefill_lookahead" || field == "spec_prefill_pool_kernel") { + if (t.spec_prefill_model.empty() || t.spec_prefill_model == "none") { + value = "-"; + } else if (vmap.find(field) != vmap.end()) { + value = vmap.at(field); + } } else if (field == "t/s") { snprintf(buf, sizeof(buf), "%.2f ± %.2f", t.avg_ts(), t.stdev_ts()); value = buf; @@ -2199,6 +2506,80 @@ static bool test_prompt(llama_context * ctx, int n_prompt, int n_batch, int n_th return true; } +static bool test_speculative_prefill_prompt( + llama_context * ctx_tgt, + llama_context * ctx_dft, + common_sampler * smpl_dft, + const common_params_speculative_prefill & spf_params, + int n_prompt, + int n_batch_tgt, + int n_threads) { + + llama_set_n_threads(ctx_tgt, n_threads, n_threads); + llama_set_n_threads(ctx_dft, n_threads, n_threads); + + const llama_model * model_tgt = llama_get_model(ctx_tgt); + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); + const int32_t n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); + + if (llama_model_is_recurrent(model_tgt)) { + fprintf(stderr, "%s: speculative prefill is not supported for recurrent models, skipping instance\n", __func__); + return false; + } + + const llama_model * model_dft = llama_get_model(ctx_dft); + const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); + const int32_t n_vocab_dft = llama_vocab_n_tokens(vocab_dft); + + const int32_t n_vocab_min = std::min(n_vocab_tgt, n_vocab_dft); + const int32_t vocab_diff = n_vocab_tgt > n_vocab_dft ? n_vocab_tgt - n_vocab_dft : n_vocab_dft - n_vocab_tgt; + if (vocab_diff > 128) { + fprintf(stderr, "%s: vocab size difference %d exceeds 128, skipping instance\n", __func__, vocab_diff); + return false; + } + + std::vector prompt_tokens(n_prompt); + prompt_tokens[0] = llama_vocab_get_add_bos(vocab_tgt) ? llama_vocab_bos(vocab_tgt) : (std::rand() % n_vocab_min); + for (int i = 1; i < n_prompt; i++) { + prompt_tokens[i] = std::rand() % n_vocab_min; + } + + llama_seq_id seq_id = 0; + + common_speculative_prefill_result spec_res = common_speculative_prefill_execute( + ctx_dft, + smpl_dft, + prompt_tokens, + seq_id, + spf_params); + + const int32_t n_kept_total = (int32_t) spec_res.kept_indices.size(); + llama_batch batch_tgt = llama_batch_init(std::max(1, std::min(n_kept_total, n_batch_tgt)), 0, 1); + + for (int32_t i = 0; i < n_kept_total; i += n_batch_tgt) { + const int32_t n_eval = std::min(n_kept_total - i, n_batch_tgt); + common_batch_clear(batch_tgt); + + for (int32_t j = 0; j < n_eval; ++j) { + const int32_t k = i + j; + const int32_t orig_idx = spec_res.kept_indices[k]; + const bool is_last = (k == n_kept_total - 1); + common_batch_add(batch_tgt, prompt_tokens[orig_idx], (llama_pos) k, { seq_id }, is_last); + } + + const int ret = llama_decode(ctx_tgt, batch_tgt); + if (ret != 0) { + fprintf(stderr, "%s: failed to decode sparse prompt on target model, ret = %d\n", __func__, ret); + llama_batch_free(batch_tgt); + return false; + } + } + llama_batch_free(batch_tgt); + + llama_synchronize(ctx_tgt); + return true; +} + static bool test_gen(llama_context * ctx, int n_gen, int n_threads) { llama_set_n_threads(ctx, n_threads, n_threads); @@ -2306,8 +2687,10 @@ int llama_bench(int argc, char ** argv) { std::vector params_instances = get_cmd_params_instances(params); - llama_model * lmodel = nullptr; - const cmd_params_instance * prev_inst = nullptr; + llama_model * lmodel = nullptr; + const cmd_params_instance * prev_inst = nullptr; + llama_model * lmodel_dft = nullptr; + const cmd_params_instance * prev_inst_dft = nullptr; // store the llama_context state at the previous depth that we performed a test // ref: https://github.com/ggml-org/llama.cpp/pull/16944#issuecomment-3478151721 @@ -2366,21 +2749,128 @@ int llama_bench(int argc, char ** argv) { lmodel = llama_model_load_from_file(inst.model.c_str(), mparams); if (lmodel == NULL) { fprintf(stderr, "%s: error: failed to load model '%s'\n", __func__, inst.model.c_str()); + if (lmodel_dft) { + llama_model_free(lmodel_dft); + } return 1; } prev_inst = &inst; } + bool has_spec_prefill = !inst.spec_prefill_model.empty() && inst.spec_prefill_model != "none"; + std::string dft_model_path = inst.spec_prefill_model == "self" ? inst.model : inst.spec_prefill_model; + + if (has_spec_prefill) { + bool dft_model_changed = !lmodel_dft || !prev_inst_dft || + prev_inst_dft->spec_prefill_model != inst.spec_prefill_model || + prev_inst_dft->spec_prefill_n_gpu_layers != inst.spec_prefill_n_gpu_layers || + prev_inst_dft->split_mode != inst.split_mode || + prev_inst_dft->load_mode != inst.load_mode || + prev_inst_dft->main_gpu != inst.main_gpu || + prev_inst_dft->devices != inst.devices || + prev_inst_dft->no_host != inst.no_host; + + if (dft_model_changed) { + if (lmodel_dft) { + llama_model_free(lmodel_dft); + } + + llama_model_params dft_mparams = llama_model_default_params(); + dft_mparams.n_gpu_layers = inst.spec_prefill_n_gpu_layers; + if (!inst.devices.empty()) { + dft_mparams.devices = const_cast(inst.devices.data()); + } + dft_mparams.split_mode = inst.split_mode; + dft_mparams.load_mode = inst.load_mode; + dft_mparams.main_gpu = inst.main_gpu; + dft_mparams.tensor_split = inst.tensor_split.data(); + dft_mparams.no_host = inst.no_host; + + lmodel_dft = llama_model_load_from_file(dft_model_path.c_str(), dft_mparams); + if (lmodel_dft == NULL) { + fprintf(stderr, "%s: error: failed to load speculative prefill draft model '%s'\n", __func__, dft_model_path.c_str()); + llama_model_free(lmodel); + return 1; + } + prev_inst_dft = &inst; + } + } else { + if (lmodel_dft) { + llama_model_free(lmodel_dft); + lmodel_dft = nullptr; + prev_inst_dft = nullptr; + } + } + + if (has_spec_prefill) { + const llama_vocab * vocab_tgt = llama_model_get_vocab(lmodel); + const llama_vocab * vocab_dft = llama_model_get_vocab(lmodel_dft); + const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); + const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); + const int vocab_diff = n_vocab_tgt > n_vocab_dft ? n_vocab_tgt - n_vocab_dft : n_vocab_dft - n_vocab_tgt; + if (vocab_diff > 128) { + fprintf(stderr, "%s: vocab size difference %d exceeds 128, skipping instance\n", __func__, vocab_diff); + continue; + } + } + llama_context * ctx = llama_init_from_model(lmodel, cparams); if (ctx == NULL) { fprintf(stderr, "%s: error: failed to create context with model '%s'\n", __func__, inst.model.c_str()); + if (lmodel_dft) { + llama_model_free(lmodel_dft); + } llama_model_free(lmodel); return 1; } + llama_context * ctx_dft = nullptr; + common_sampler_ptr smpl_dft; + common_params_speculative_prefill spf_params; + + if (has_spec_prefill) { + llama_context_params dft_cparams = llama_context_default_params(); + const uint32_t needed_dft_ctx = (uint32_t) (inst.n_prompt + inst.spec_prefill_lookahead + 16); + dft_cparams.n_ctx = inst.spec_prefill_n_ctx > 0 ? inst.spec_prefill_n_ctx : std::min(needed_dft_ctx, (uint32_t) llama_model_n_ctx_train(lmodel_dft)); + dft_cparams.n_batch = inst.n_batch; + dft_cparams.n_ubatch = inst.n_ubatch; + dft_cparams.type_k = inst.type_k; + dft_cparams.type_v = inst.type_v; + dft_cparams.offload_kqv = !inst.no_kv_offload; + dft_cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + dft_cparams.embeddings = false; + dft_cparams.op_offload = !inst.no_op_offload; + dft_cparams.no_perf = true; + + ctx_dft = llama_init_from_model(lmodel_dft, dft_cparams); + if (ctx_dft == NULL) { + fprintf(stderr, "%s: error: failed to create draft context with model '%s'\n", __func__, dft_model_path.c_str()); + llama_free(ctx); + llama_model_free(lmodel_dft); + llama_model_free(lmodel); + return 1; + } + + common_params_sampling sparams_dft; + sparams_dft.temp = 0.0f; + smpl_dft.reset(common_sampler_init(lmodel_dft, sparams_dft)); + + spf_params.enabled = true; + spf_params.n_ctx = inst.spec_prefill_n_ctx; + spf_params.percentage = inst.spec_prefill_percentage; + spf_params.chunk_size = inst.spec_prefill_chunk_size; + spf_params.look_ahead_cnt = inst.spec_prefill_lookahead; + spf_params.pool_kernel_size = inst.spec_prefill_pool_kernel; + spf_params.keep_bos = true; + spf_params.keep_last = true; + } + test t(inst, lmodel, ctx); llama_memory_clear(llama_get_memory(ctx), false); + if (ctx_dft) { + llama_memory_clear(llama_get_memory(ctx_dft), false); + } // cool off before the test if (params.delay) { @@ -2390,7 +2880,9 @@ int llama_bench(int argc, char ** argv) { struct ggml_threadpool_params tpp = ggml_threadpool_params_default(t.n_threads); if (!parse_cpu_mask(t.cpu_mask, tpp.cpumask)) { fprintf(stderr, "%s: failed to parse cpu-mask: %s\n", __func__, t.cpu_mask.c_str()); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } @@ -2401,12 +2893,17 @@ int llama_bench(int argc, char ** argv) { struct ggml_threadpool * threadpool = ggml_threadpool_new_fn(&tpp); if (!threadpool) { fprintf(stderr, "%s: threadpool create failed : n_threads %d\n", __func__, tpp.n_threads); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } llama_attach_threadpool(ctx, threadpool, NULL); + if (ctx_dft) { + llama_attach_threadpool(ctx_dft, threadpool, NULL); + } // warmup run if (!params.no_warmup) { @@ -2414,11 +2911,14 @@ int llama_bench(int argc, char ** argv) { if (params.progress) { fprintf(stderr, "llama-bench: benchmark %d/%zu: warmup prompt run\n", params_idx, params_count); } - //test_prompt(ctx, std::min(t.n_batch, std::min(t.n_prompt, 32)), 0, t.n_batch, t.n_threads); - bool res = test_prompt(ctx, t.n_prompt, t.n_batch, t.n_threads); + bool res = has_spec_prefill + ? test_speculative_prefill_prompt(ctx, ctx_dft, smpl_dft.get(), spf_params, t.n_prompt, t.n_batch, t.n_threads) + : test_prompt(ctx, t.n_prompt, t.n_batch, t.n_threads); if (!res) { fprintf(stderr, "%s: error: failed to run prompt warmup\n", __func__); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } @@ -2430,7 +2930,9 @@ int llama_bench(int argc, char ** argv) { bool res = test_gen(ctx, 1, t.n_threads); if (!res) { fprintf(stderr, "%s: error: failed to run gen warmup\n", __func__); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } @@ -2439,6 +2941,10 @@ int llama_bench(int argc, char ** argv) { for (int i = 0; i < params.reps; i++) { llama_memory_clear(llama_get_memory(ctx), false); + if (ctx_dft) { + llama_memory_clear(llama_get_memory(ctx_dft), false); + common_sampler_reset(smpl_dft.get()); + } if (t.n_depth > 0) { bool is_cached = t.n_depth == cstate.depth; @@ -2460,7 +2966,9 @@ int llama_bench(int argc, char ** argv) { bool res = test_prompt(ctx, t.n_depth, t.n_batch, t.n_threads); if (!res) { fprintf(stderr, "%s: error: failed to run depth\n", __func__); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } @@ -2484,10 +2992,14 @@ int llama_bench(int argc, char ** argv) { fprintf(stderr, "llama-bench: benchmark %d/%zu: prompt run %d/%d\n", params_idx, params_count, i + 1, params.reps); } - bool res = test_prompt(ctx, t.n_prompt, t.n_batch, t.n_threads); + bool res = has_spec_prefill + ? test_speculative_prefill_prompt(ctx, ctx_dft, smpl_dft.get(), spf_params, t.n_prompt, t.n_batch, t.n_threads) + : test_prompt(ctx, t.n_prompt, t.n_batch, t.n_threads); if (!res) { fprintf(stderr, "%s: error: failed to run prompt\n", __func__); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } @@ -2500,7 +3012,9 @@ int llama_bench(int argc, char ** argv) { bool res = test_gen(ctx, t.n_gen, t.n_threads); if (!res) { fprintf(stderr, "%s: error: failed to run gen\n", __func__); + if (ctx_dft) llama_free(ctx_dft); llama_free(ctx); + if (lmodel_dft) llama_model_free(lmodel_dft); llama_model_free(lmodel); exit(1); } @@ -2522,11 +3036,17 @@ int llama_bench(int argc, char ** argv) { llama_perf_context_print(ctx); + if (ctx_dft) { + llama_free(ctx_dft); + } llama_free(ctx); ggml_threadpool_free_fn(threadpool); } + if (lmodel_dft) { + llama_model_free(lmodel_dft); + } llama_model_free(lmodel); if (p) { diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index c30955e89f0..d2db0f92c4c 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -546,7 +546,7 @@ void server_tokens::insert(const llama_tokens & inp_tokens) { } const llama_tokens & server_tokens::get_tokens() const { - GGML_ASSERT(!has_mtmd); + GGML_ASSERT(!has_media()); return tokens; } @@ -629,7 +629,7 @@ llama_tokens server_tokens::get_text_tokens() const { } void server_tokens::set_token(llama_pos pos, llama_token id) { - GGML_ASSERT(!has_mtmd); // only allow this if mtmd is disabled + GGML_ASSERT(!has_media()); tokens[pos] = id; } diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6c681a2cf56..3e37a6a4c0d 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -220,6 +220,8 @@ struct server_tokens { bool empty() const { return tokens.empty(); } + bool has_media() const { return !map_idx_to_media.empty(); } + void clear() { map_idx_to_media.clear(); tokens.clear(); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f5477356d61..f41650ecb65 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -14,6 +14,8 @@ #include "log.h" #include "sampling.h" #include "speculative.h" +#include "speculative-prefill.h" +#include "src/llama-ext.h" #include "mtmd.h" #include "mtmd-helper.h" @@ -241,6 +243,7 @@ struct server_slot { llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; + llama_context * ctx_spf = nullptr; common_memory mem; @@ -258,6 +261,9 @@ struct server_slot { bool spec_is_replay = false; std::mt19937 spec_synth_rng; + bool spec_prefill_active = false; + server_tokens spec_prefill_tokens; + // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 std::unique_ptr task; @@ -335,6 +341,9 @@ struct server_slot { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); mem.seq_rm(id, -1, -1); + if (ctx_spf) { + llama_memory_seq_rm(llama_get_memory(ctx_spf), id, -1, -1); + } prompt.clear(); } @@ -404,6 +413,9 @@ struct server_slot { // clear multimodal state mbatch.reset(); + + spec_prefill_active = false; + spec_prefill_tokens.clear(); } void init_sampler() const { @@ -471,6 +483,14 @@ struct server_slot { return !!spec; } + const server_tokens & prompt_src() const { + return spec_prefill_active ? spec_prefill_tokens : task->tokens; + } + + int32_t n_prompt_src() const { + return spec_prefill_active ? (int32_t) spec_prefill_tokens.size() : task->n_tokens(); + } + void add_token(const completion_token_output & token) { if (!is_processing()) { SLT_WRN(*this, "%s", "slot is not processing\n"); @@ -508,10 +528,12 @@ struct server_slot { // no speculative decoding i_batch = batch.size(); + llama_pos pos = prompt.tokens.pos_next(); + if (!inp_embd.empty()) { - add_ok &= batch.add(id, inp_embd, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, inp_embd, pos, true, false); } else { - add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, sampled, pos, true, false); } SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n", @@ -527,7 +549,7 @@ struct server_slot { spec_i_batch.push_back(batch.size() + i + 1); } - auto pos0 = prompt.tokens.pos_next(); + llama_pos pos0 = prompt.tokens.pos_next(); add_ok &= batch.add(id, sampled, pos0++, true, false); for (auto token : spec_draft) { @@ -551,6 +573,10 @@ struct server_slot { state = SLOT_STATE_IDLE; + if (ctx_spf) { + llama_memory_seq_rm(llama_get_memory(ctx_spf), id, -1, -1); + } + // do not keep context of the child slots - the parent's context is enough if (task->is_child()) { prompt_clear(); @@ -622,11 +648,18 @@ struct server_slot { return; } - const double n_prompt_second = stats.n_prompt_tps(); - const double f_progress = task->n_tokens() > 0 ? (double) prompt.n_tokens() / task->n_tokens() : 0.0; + const double f_progress = n_prompt_src() > 0 ? (double) prompt.n_tokens() / n_prompt_src() : 0.0; + int n_tokens_disp = (int) stats.n_prompt_processed; + double n_prompt_second = stats.n_prompt_tps(); + + if (spec_prefill_active && task && n_prompt_src() > 0) { + const int n_uncached = (int) (task->n_tokens() - stats.n_prompt_cached); + n_tokens_disp = (int) (f_progress * n_uncached); + n_prompt_second = t_prompt_total > 0.0 ? (1e3 / t_prompt_total) * n_tokens_disp : 0.0; + } SLT_INF(*this, "prompt processing, n_tokens = %6d, progress = %.2f, t = %6.2f s / %.2f tokens per second\n", - (int) stats.n_prompt_processed, f_progress, t_prompt_total / 1e3, n_prompt_second); + n_tokens_disp, f_progress, t_prompt_total / 1e3, n_prompt_second); } void print_timings() const { @@ -697,7 +730,7 @@ struct server_slot { if (ptask) { res["id_task"] = ptask->id; - res["n_prompt_tokens"] = (int32_t) prompt.tokens.size(); + res["n_prompt_tokens"] = ptask->n_tokens(); res["n_prompt_tokens_processed"] = stats.n_prompt_processed; res["n_prompt_tokens_cache"] = stats.n_prompt_cached; res["params"] = ptask->params.to_json(only_metrics); @@ -725,6 +758,10 @@ struct server_slot { mem.seq_rm(other.id, -1, -1); mem.seq_cp(id, other.id, -1, -1); + if (ctx_spf) { + llama_memory_seq_rm(llama_get_memory(ctx_spf), other.id, -1, -1); + } + other.i_batch = i_batch; other.stats = stats; @@ -890,6 +927,12 @@ struct server_context_impl { common_speculative_init_result_ptr spec_init; + common_init_result_ptr llama_init_spf; + llama_context_ptr ctx_spf_own; + llama_model * model_spf = nullptr; + llama_context * ctx_spf = nullptr; + common_sampler_ptr smpl_spf; + common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -937,6 +980,13 @@ struct server_context_impl { void destroy() { spec.reset(); + + smpl_spf.reset(); + ctx_spf_own.reset(); + llama_init_spf.reset(); + ctx_spf = nullptr; + model_spf = nullptr; + spec_init.reset(); ctx_dft = nullptr; @@ -1001,6 +1051,53 @@ struct server_context_impl { return true; } + void apply_spec_prefill(server_slot & slot) { + slot.spec_prefill_active = false; + slot.spec_prefill_tokens.clear(); + + if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { + return; + } + if (!slot.task || !slot.task->need_logits() || slot.task->tokens.has_media()) { + return; + } + + const llama_tokens & prompt = slot.task->tokens.get_tokens(); + if (prompt.empty()) { + return; + } + + llama_memory_seq_rm(llama_get_memory(ctx_spf), slot.id, -1, -1); + common_sampler_reset(smpl_spf.get()); + + const common_speculative_prefill_result res = common_speculative_prefill_execute( + ctx_spf, smpl_spf.get(), prompt, slot.id, params_base.speculative.prefill); + + llama_memory_seq_rm(llama_get_memory(ctx_spf), slot.id, -1, -1); + + if (res.kept_indices.empty() || (int32_t) res.kept_indices.size() >= (int32_t) prompt.size()) { + return; + } + + llama_tokens kept; + kept.reserve(res.kept_indices.size()); + for (const int32_t idx : res.kept_indices) { + if (idx >= 0 && idx < (int32_t) prompt.size()) { + kept.push_back(prompt[idx]); + } + } + if (kept.size() < 2) { + return; + } + + slot.spec_prefill_tokens = server_tokens(kept, slot.task->tokens.has_mtmd); + slot.spec_prefill_active = true; + + SLT_INF(slot, "speculative prefill kept %d / %d tokens (%.1f%%)\n", + (int) kept.size(), (int) prompt.size(), + 100.0f * (float) kept.size() / (float) prompt.size()); + } + // load the model and initialize llama_context // this may also be called to resume from sleeping state bool load_model(common_params & params) { @@ -1150,6 +1247,107 @@ struct server_context_impl { load_progress_callback(1.0f, &load_progress_spec); } + if (params_base.speculative.prefill.enabled) { + if (llama_model_is_recurrent(model_tgt)) { + SRV_ERR("%s", "speculative prefill is not supported for recurrent models\n"); + return false; + } + + const bool has_target_dependent_dft = std::any_of( + params_base.speculative.types.begin(), + params_base.speculative.types.end(), + [](common_speculative_type t) { + return t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || + t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK || + t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || + t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP; + }); + + common_params_model spf_model = params_base.speculative.prefill.model; + if (spf_model.empty() && !has_target_dependent_dft) { + spf_model = params_base.speculative.draft.mparams; + } + + if (spf_model.empty()) { + SRV_ERR("%s", "speculative prefill enabled but no draft model was provided\n"); + return false; + } + + common_params params_spf = common_base_params_to_speculative(params_base); + params_spf.model = spf_model; + if (params_base.speculative.prefill.n_ctx > 0) { + params_spf.n_ctx = params_base.speculative.prefill.n_ctx; + } + if (params_base.speculative.prefill.n_gpu_layers != -1) { + params_spf.n_gpu_layers = params_base.speculative.prefill.n_gpu_layers; + } + if (!params_base.speculative.prefill.devices.empty()) { + params_spf.devices = params_base.speculative.prefill.devices; + } + params_spf.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + params_spf.fit_params = false; + params_spf.n_parallel = params_base.n_parallel; + params_spf.speculative.prefill.enabled = false; + params_spf.speculative.draft.mparams = {}; + + const bool reuse_dft = model_dft && !has_target_dependent_dft && ( + params_base.speculative.prefill.model.empty() || + spf_model.path == params_base.speculative.draft.mparams.path); + + if (reuse_dft) { + SRV_INF("reusing draft model for speculative prefill '%s'\n", params_spf.model.get_name().c_str()); + model_spf = model_dft; + } else { + SRV_INF("loading speculative prefill draft model '%s'\n", params_spf.model.get_name().c_str()); + llama_init_spf = common_init_from_params(params_spf, /*model_only=*/true); + model_spf = llama_init_spf ? llama_init_spf->model() : nullptr; + } + if (model_spf == nullptr) { + SRV_ERR("failed to load speculative prefill draft model, '%s'\n", params_spf.model.path.c_str()); + return false; + } + + if (params_base.speculative.prefill.n_ctx <= 0 && params_spf.n_ctx > (int32_t) llama_model_n_ctx_train(model_spf)) { + params_spf.n_ctx = llama_model_n_ctx_train(model_spf); + SRV_INF("capping speculative prefill draft context to training limit (%d tokens)\n", params_spf.n_ctx); + } + + llama_context_params cparams = common_context_params_to_llama(params_spf); + ctx_spf_own.reset(llama_init_from_model(model_spf, cparams)); + ctx_spf = ctx_spf_own.get(); + if (ctx_spf == nullptr) { + SRV_ERR("failed to create speculative prefill context for '%s'\n", params_spf.model.path.c_str()); + return false; + } + + if (llama_model_is_recurrent(model_spf)) { + SRV_ERR("speculative prefill draft model '%s' is recurrent and not supported\n", params_spf.model.get_name().c_str()); + return false; + } + + if (llama_model_target_layer_ids_n(model_spf) > 0) { + SRV_ERR("speculative prefill draft model '%s' is target-dependent and cannot be used for speculative prefill\n", params_spf.model.get_name().c_str()); + return false; + } + + const llama_vocab * vocab_spf = llama_model_get_vocab(model_spf); + const int n_vocab_tgt = llama_vocab_n_tokens(vocab); + const int n_vocab_spf = llama_vocab_n_tokens(vocab_spf); + const int vocab_diff = n_vocab_tgt > n_vocab_spf ? n_vocab_tgt - n_vocab_spf : n_vocab_spf - n_vocab_tgt; + if (vocab_diff > 128) { + SRV_ERR("speculative prefill vocab size difference %d exceeds 128\n", vocab_diff); + return false; + } + + common_params_sampling sparams_spf; + sparams_spf.temp = 0.0f; + smpl_spf.reset(common_sampler_init(model_spf, sparams_spf)); + if (!smpl_spf) { + SRV_ERR("%s", "failed to init speculative prefill sampler\n"); + return false; + } + } + if (has_mmproj) { if (callback_state) { callback_state(SERVER_STATE_LOADING, {{"stage", "mmproj_model"}}); @@ -1290,6 +1488,7 @@ struct server_context_impl { slot.id = i; slot.ctx_tgt = ctx_tgt; slot.ctx_dft = ctx_dft; + slot.ctx_spf = ctx_spf; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); slot.n_ctx = n_ctx_slot(); @@ -3106,19 +3305,23 @@ struct server_context_impl { // this slot still has a prompt to be processed if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) { - const auto & input_tokens = slot.task->tokens; + if (slot.state == SLOT_STATE_STARTED) { + slot.stats.update_prompt_start(); + apply_spec_prefill(slot); + } + + const auto & input_tokens = slot.prompt_src(); + const int32_t n_prompt_in = slot.n_prompt_src(); // used to determine the number of tokens added to the batch for the current slot const auto n_tokens_prev = batch.size(); // TODO: maybe move branch to outside of this loop in the future if (slot.state == SLOT_STATE_STARTED) { - slot.stats.update_prompt_start(); - slot.state = SLOT_STATE_PROCESSING_PROMPT; SLT_TRC(slot, "new prompt, n_ctx_slot = %d, n_keep = %d, task.n_tokens = %d\n", - slot.n_ctx, slot.task->params.n_keep, slot.task->n_tokens()); + slot.n_ctx, slot.task->params.n_keep, n_prompt_in); // print prompt tokens (for debugging) /*if (1) { @@ -3155,39 +3358,39 @@ struct server_context_impl { } if (!slot.can_split()) { - if (slot.task->n_tokens() > n_ubatch) { + if (n_prompt_in > n_ubatch) { send_error(slot, string_format( "input (%d tokens) is too large to process. increase the physical batch " "size (current batch size: %d)", - slot.task->n_tokens(), n_ubatch), + n_prompt_in, n_ubatch), ERROR_TYPE_SERVER); slot.release(); return; } - if (slot.task->n_tokens() > slot.n_ctx) { + if (n_prompt_in > slot.n_ctx) { send_error( slot, string_format( "input (%d tokens) is larger than the max context size (%d tokens). skipping", - slot.task->n_tokens(), slot.n_ctx), + n_prompt_in, slot.n_ctx), ERROR_TYPE_EXCEED_CONTEXT_SIZE); slot.release(); return; } } else { - if (slot.task->n_tokens() >= slot.n_ctx) { + if (n_prompt_in >= slot.n_ctx) { send_error(slot, string_format("request (%d tokens) exceeds the available context size (%d " "tokens), try increasing it", - slot.task->n_tokens(), slot.n_ctx), + n_prompt_in, slot.n_ctx), ERROR_TYPE_EXCEED_CONTEXT_SIZE); slot.release(); return; } - if (slot.task->params.cache_prompt) { + if (slot.task->params.cache_prompt && !slot.spec_prefill_active) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); @@ -3264,7 +3467,7 @@ struct server_context_impl { llama_pos pos_next = slot.prompt.tokens.pos_next(n_past); // ref: https://github.com/ggml-org/llama.cpp/pull/24110 - const bool has_new_tokens = (n_past < slot.task->n_tokens()); + const bool has_new_tokens = (n_past < n_prompt_in); // the largest pos_min required for a checkpoint to be useful const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); @@ -3280,7 +3483,7 @@ struct server_context_impl { // this is useful for debugging prompt caching if (slots_debug) { const int np0 = std::max(n_past - slots_n_diff, 0); - const int np1 = std::min(n_past + slots_n_diff + 2, std::min(slot.prompt.tokens.size(), slot.task->tokens.size())); + const int np1 = std::min(n_past + slots_n_diff + 2, std::min(slot.prompt.tokens.size(), input_tokens.size())); std::stringstream ss0; std::stringstream ss1; @@ -3305,7 +3508,7 @@ struct server_context_impl { } { - const auto token = slot.task->tokens[i]; + const auto token = input_tokens[i]; const auto piece = token != LLAMA_TOKEN_NULL ? common_token_to_piece(ctx_tgt, token) : "[mtmd]"; ss1 << piece; st1 << std::setw(8) << token; @@ -3373,8 +3576,8 @@ struct server_context_impl { } // [TAG_PROMPT_LOGITS] - if (n_past == slot.task->n_tokens() && n_past > 0) { - SLT_WRN(slot, "need to evaluate at least 1 token for each active slot (n_past = %d, task.n_tokens() = %d)\n", n_past, slot.task->n_tokens()); + if (n_past == n_prompt_in && n_past > 0) { + SLT_WRN(slot, "need to evaluate at least 1 token for each active slot (n_past = %d, task.n_tokens() = %d)\n", n_past, n_prompt_in); n_past--; SLT_WRN(slot, "n_past was set to %d\n", n_past); } @@ -3400,7 +3603,7 @@ struct server_context_impl { if (!slot.can_split()) { // cannot fit the prompt in the current batch - will try next iter - if (batch.size() + slot.task->n_tokens() > n_batch) { + if (batch.size() + n_prompt_in > n_batch) { return; } } @@ -3451,7 +3654,7 @@ struct server_context_impl { while (true) { auto cur_token_idx = slot.prompt.n_tokens(); if ( - cur_token_idx >= slot.task->n_tokens() || + cur_token_idx >= n_prompt_in || input_tokens[cur_token_idx] != LLAMA_TOKEN_NULL // encountered a text token ) { break; @@ -3494,7 +3697,7 @@ struct server_context_impl { const auto last_user_pos = spans.last_user_message_pos(); // add prompt tokens for processing in the current batch - while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) { + while (slot.prompt.n_tokens() < n_prompt_in && batch.size() < n_batch) { // get next token to process llama_token cur_tok = input_tokens[slot.prompt.n_tokens()]; if (cur_tok == LLAMA_TOKEN_NULL) { @@ -3512,9 +3715,10 @@ struct server_context_impl { // embedding requires all tokens in the batch to be output; // MTP also wants logits at every prompt position so the // streaming hook can mirror t_h_nextn into ctx_dft. + llama_pos pos = slot.prompt.tokens.pos_next(); add_ok &= batch.add(slot.id, cur_tok, - /* pos = */ slot.prompt.tokens.pos_next(), + /* pos = */ pos, /* output = */ slot.need_embd(), /* is_prompt = */ true); slot.prompt.tokens.push_back(cur_tok); @@ -3540,7 +3744,7 @@ struct server_context_impl { bool should_break = false; for (int offset : checkpoint_offsets) { const int n_last = std::min(n_batch, offset); - if (slot.task->n_tokens() == slot.prompt.n_tokens() + n_last) { + if (n_prompt_in == slot.prompt.n_tokens() + n_last) { should_break = true; break; } @@ -3556,15 +3760,19 @@ struct server_context_impl { const auto n_tokens_start = slot.prompt.n_tokens() - n_tokens_cur; - const bool near_prompt_end = slot.task->n_tokens() < slot.prompt.n_tokens() + n_ubatch; + const bool near_prompt_end = n_prompt_in < slot.prompt.n_tokens() + n_ubatch; const bool is_user_start = spans.is_user_start(n_tokens_start); const bool is_last_user_message = n_tokens_start == last_user_pos; // entire prompt has been processed - if (slot.prompt.n_tokens() == slot.task->n_tokens()) { + if (slot.prompt.n_tokens() == n_prompt_in) { slot.state = SLOT_STATE_DONE_PROMPT; + if (slot.spec_prefill_active && slot.task) { + slot.stats.n_prompt_processed = slot.task->n_tokens() - slot.stats.n_prompt_cached; + } + GGML_ASSERT(batch.size() > 0); // extract the logits only for the last token