From a1f3ad0cd9ff4f44641040be22008875b98c9793 Mon Sep 17 00:00:00 2001 From: rockofox Date: Mon, 24 Aug 2026 16:22:15 +0200 Subject: [PATCH 01/19] speculative-prefill : implement token importance estimation and sparse prefill Assisted-by: Antigravity --- common/CMakeLists.txt | 2 + common/arg.cpp | 49 +++ common/common.h | 12 + common/speculative-prefill.cpp | 313 ++++++++++++++++++ common/speculative-prefill.h | 37 +++ examples/CMakeLists.txt | 1 + examples/speculative-prefill/CMakeLists.txt | 5 + .../speculative-prefill.cpp | 192 +++++++++++ include/llama.h | 3 + scripts/compare_spec_prefill.py | 305 +++++++++++++++++ scripts/eval_spec_prefill_quality.py | 168 ++++++++++ scripts/run_paper_benchmarks.py | 296 +++++++++++++++++ src/llama-context.cpp | 15 + src/llama-context.h | 1 + 14 files changed, 1399 insertions(+) create mode 100644 common/speculative-prefill.cpp create mode 100644 common/speculative-prefill.h create mode 100644 examples/speculative-prefill/CMakeLists.txt create mode 100644 examples/speculative-prefill/speculative-prefill.cpp create mode 100755 scripts/compare_spec_prefill.py create mode 100644 scripts/eval_spec_prefill_quality.py create mode 100644 scripts/run_paper_benchmarks.py 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 86f8610a56d..54822516c38 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4169,6 +4169,55 @@ 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"}, + "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-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/common.h b/common/common.h index de49dac9f63..1c039a80927 100644 --- a/common/common.h +++ b/common/common.h @@ -366,6 +366,16 @@ 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 + 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 }; @@ -379,6 +389,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(); } diff --git a/common/speculative-prefill.cpp b/common/speculative-prefill.cpp new file mode 100644 index 00000000000..d411bceeb88 --- /dev/null +++ b/common/speculative-prefill.cpp @@ -0,0 +1,313 @@ +#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) { + 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(struct ggml_tensor * t, bool ask, void * user_data) { + if (ask) { + return true; + } + + 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; + } + + 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 (n_past_total < n_prompt || n_heads <= 0) { + return true; + } + + std::vector raw_buf((size_t) ggml_nelements(t)); + ggml_backend_tensor_get(t, raw_buf.data(), 0, ggml_nbytes(t)); + + std::vector layer_data((size_t) n_heads * n_prompt); + + for (int32_t h = 0; h < n_heads; ++h) { + const float * src = raw_buf.data() + (size_t) h * n_past_total; + float * dst = layer_data.data() + (size_t) h * n_prompt; + std::copy(src, src + n_prompt, dst); + } + + data->layer_attns.push_back(std::move(layer_data)); + + return true; +} + +common_speculative_prefill_result common_speculative_prefill_execute( + llama_context * ctx_dft, + struct 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 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 + const int32_t lookahead = std::max(1, params.look_ahead_cnt); + 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) { + for (size_t i = 0; i < prompt.size(); ++i) { + total_importance[i] /= (float) actual_steps; + } + } else { + std::fill(total_importance.begin(), total_importance.end(), 1.0f); + } + + 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..09b33a0dcf1 --- /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, + struct 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..e386529617e --- /dev/null +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -0,0 +1,192 @@ +#include "arg.h" +#include "common.h" +#include "log.h" +#include "llama.h" +#include "sampling.h" +#include "speculative.h" +#include "speculative-prefill.h" + +#include +#include +#include +#include +#include +#include +#include + +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; + } + + if (params.speculative.draft.mparams.empty()) { + LOG_ERR("%s: draft model is required for speculative prefill (specify with -md or --spec-draft-model)\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(); + + // 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.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); + if (!init_dft) { + LOG_ERR("%s: failed to load draft model\n", __func__); + return 1; + } + + llama_model * model_dft = init_dft->model(); + llama_context * ctx_dft = init_dft->context(); + + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); + + // 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_sampler_ptr smpl_dft(common_sampler_init(model_dft, params.sampling)); + + // 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); + + 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\n", __func__, ttft_ms); + + // 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 = (int32_t) spec_res.kept_indices.size(); + 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 a04177f9f7d..267207fec0e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1012,6 +1012,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_spec_prefill.py b/scripts/compare_spec_prefill.py new file mode 100755 index 00000000000..2f307905e33 --- /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("-md", "--draft-model", help="Path to draft model GGUF") + parser.add_argument("-hfd", "--hf-repo-draft", "--draft-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("-ngld", "--n-gpu-layers-draft", 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_quality.py b/scripts/eval_spec_prefill_quality.py new file mode 100644 index 00000000000..ae6bac49677 --- /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("-md", "--draft-model", help="Path to draft model GGUF") + parser.add_argument("-ngl", "--n-gpu-layers", type=int, default=99, help="GPU layers for target model") + parser.add_argument("-ngld", "--n-gpu-layers-draft", 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/run_paper_benchmarks.py b/scripts/run_paper_benchmarks.py new file mode 100644 index 00000000000..e535ac218c9 --- /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("-md", "--draft-model", required=True, help="Draft model path") + 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("--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 0402044da6b..49bd4ae7dd1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1150,6 +1150,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); @@ -3716,6 +3727,10 @@ 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) { + 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); From 490ce1fb801b80f074238b2a59adc06c73331d7b Mon Sep 17 00:00:00 2001 From: rockofox Date: Mon, 24 Aug 2026 20:08:20 +0200 Subject: [PATCH 02/19] speculative-prefill : conform to contributing style guidelines Assisted-by: Antigravity --- common/speculative-prefill.cpp | 4 ++-- common/speculative-prefill.h | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/common/speculative-prefill.cpp b/common/speculative-prefill.cpp index d411bceeb88..ddcd8d696d2 100644 --- a/common/speculative-prefill.cpp +++ b/common/speculative-prefill.cpp @@ -146,7 +146,7 @@ struct cb_attn_collector_data { std::vector> layer_attns; }; -static bool cb_collect_attn(struct ggml_tensor * t, bool ask, void * user_data) { +static bool cb_collect_attn(ggml_tensor * t, bool ask, void * user_data) { if (ask) { return true; } @@ -186,7 +186,7 @@ static bool cb_collect_attn(struct ggml_tensor * t, bool ask, void * user_data) common_speculative_prefill_result common_speculative_prefill_execute( llama_context * ctx_dft, - struct common_sampler * smpl_dft, + common_sampler * smpl_dft, const std::vector & prompt, llama_seq_id seq_id, const common_params_speculative_prefill & params) { diff --git a/common/speculative-prefill.h b/common/speculative-prefill.h index 09b33a0dcf1..0752449133d 100644 --- a/common/speculative-prefill.h +++ b/common/speculative-prefill.h @@ -31,7 +31,7 @@ std::vector common_speculative_prefill_select_indices( // execute draft prefill, lookahead decoding, attention collection, and index selection common_speculative_prefill_result common_speculative_prefill_execute( llama_context * ctx_dft, - struct common_sampler * smpl_dft, + common_sampler * smpl_dft, const std::vector & prompt, llama_seq_id seq_id, const common_params_speculative_prefill & params); From 061a3b50a885d75ae073f6fa5bd6a4367a91fd81 Mon Sep 17 00:00:00 2001 From: rockofox Date: Mon, 24 Aug 2026 23:50:40 +0200 Subject: [PATCH 03/19] speculative-prefill : add dedicated draft model parameters Assisted-by: Antigravity --- common/arg.cpp | 63 +++++++++++++++++-- common/arg.h | 1 + common/common.h | 16 ++--- .../speculative-prefill.cpp | 13 +++- scripts/compare_spec_prefill.py | 6 +- scripts/eval_spec_prefill_quality.py | 4 +- scripts/run_paper_benchmarks.py | 4 +- tests/test-arg-parser.cpp | 6 ++ 8 files changed, 93 insertions(+), 20 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 54822516c38..5b26c5d16de 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); } @@ -4170,12 +4189,48 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } ).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"}, + {"--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-p", "--spec-prefill-percentage"}, "P", string_format("fraction of prompt tokens to retain during speculative prefill (default: %.2f)", (double) params.speculative.prefill.percentage), 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 1c039a80927..21b49f1c76f 100644 --- a/common/common.h +++ b/common/common.h @@ -367,13 +367,15 @@ struct common_params_speculative_ngram_cache { }; struct common_params_speculative_prefill { - bool enabled = false; // enable speculative prefill - 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 + bool enabled = false; // enable speculative prefill + common_params_model model; // draft model for speculative prefill + int32_t n_gpu_layers = -1; // max draft model layers to store in VRAM (-1 - use default) + 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 { diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index e386529617e..b1014aaf12e 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -25,8 +25,13 @@ int main(int argc, char ** argv) { return 1; } - if (params.speculative.draft.mparams.empty()) { - LOG_ERR("%s: draft model is required for speculative prefill (specify with -md or --spec-draft-model)\n", __func__); + 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; } @@ -50,6 +55,10 @@ int main(int argc, char ** argv) { // 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_gpu_layers != -1) { + params_dft.n_gpu_layers = params.speculative.prefill.n_gpu_layers; + } 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); diff --git a/scripts/compare_spec_prefill.py b/scripts/compare_spec_prefill.py index 2f307905e33..bf4a5d6dbe3 100755 --- a/scripts/compare_spec_prefill.py +++ b/scripts/compare_spec_prefill.py @@ -15,11 +15,11 @@ def parse_args(): 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("-md", "--draft-model", help="Path to draft model GGUF") - parser.add_argument("-hfd", "--hf-repo-draft", "--draft-hf", help="Hugging Face repo for draft model (/[:quant])") + 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("-ngld", "--n-gpu-layers-draft", type=int, default=99, help="Number of GPU layers for draft 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") diff --git a/scripts/eval_spec_prefill_quality.py b/scripts/eval_spec_prefill_quality.py index ae6bac49677..299bc0214cb 100644 --- a/scripts/eval_spec_prefill_quality.py +++ b/scripts/eval_spec_prefill_quality.py @@ -11,9 +11,9 @@ def parse_args(): 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("-md", "--draft-model", help="Path to draft 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("-ngld", "--n-gpu-layers-draft", type=int, default=99, help="GPU layers for draft 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)") diff --git a/scripts/run_paper_benchmarks.py b/scripts/run_paper_benchmarks.py index e535ac218c9..cda17a751ab 100644 --- a/scripts/run_paper_benchmarks.py +++ b/scripts/run_paper_benchmarks.py @@ -37,9 +37,9 @@ def parse_args(): 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("-md", "--draft-model", required=True, help="Draft 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("-ngld", "--n-gpu-layers-draft", 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) diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index ba58f852eb4..73d0e4c488c 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -197,6 +197,12 @@ 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"}; + 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.enabled == true); + 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); From ee7f301651cca7f93e8b6794952d903db5c0b2b6 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 00:21:15 +0200 Subject: [PATCH 04/19] llama-bench : add speculative prefill benchmarking support Assisted-by: Antigravity --- scripts/compare-llama-bench.py | 16 +- tools/llama-bench/llama-bench.cpp | 743 ++++++++++++++++++++++++------ 2 files changed, 619 insertions(+), 140 deletions(-) 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/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index a2da93b9a28..8d9f9ce3f01 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 @@ -319,12 +321,33 @@ 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_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; @@ -364,47 +387,55 @@ 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 }, - /* 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_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 }, + /* 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) { @@ -440,6 +471,22 @@ 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(" -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()); @@ -565,6 +612,81 @@ 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 == "-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; @@ -1097,10 +1219,50 @@ 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_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; } @@ -1188,6 +1350,12 @@ 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_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; @@ -1299,6 +1467,12 @@ 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 & 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) @@ -1328,33 +1502,39 @@ 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, - /* .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_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, + /* .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); } @@ -1364,33 +1544,39 @@ 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, - /* .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_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, + /* .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); } @@ -1400,33 +1586,39 @@ 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, - /* .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_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, + /* .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); } @@ -1468,6 +1660,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; @@ -1501,12 +1699,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; @@ -1565,6 +1769,8 @@ struct test { "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", "tensor_buft_overrides", "load_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" }; @@ -1578,17 +1784,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") { + if (field == "load_mode" || field == "spec_prefill_model") { return STRING; } return STRING; @@ -1663,6 +1871,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), @@ -1854,6 +2068,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); @@ -1906,6 +2130,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; } @@ -1987,6 +2229,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"); @@ -2041,6 +2308,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; @@ -2140,6 +2423,70 @@ 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); + + 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); + + 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); @@ -2247,8 +2594,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 @@ -2307,21 +2656,113 @@ 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; + } + } + 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(); + dft_cparams.n_ctx = inst.n_prompt + inst.spec_prefill_lookahead + 16; + 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; + smpl_dft.reset(common_sampler_init(lmodel_dft, sparams_dft)); + + spf_params.enabled = true; + 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) { @@ -2331,7 +2772,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); } @@ -2342,12 +2785,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) { @@ -2355,11 +2803,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); } @@ -2371,7 +2822,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); } @@ -2380,6 +2833,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; @@ -2401,7 +2858,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); } @@ -2425,10 +2884,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); } @@ -2441,7 +2904,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); } @@ -2463,11 +2928,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) { From 96be5de918db2154b9d09b3183839680fb7c43e0 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 13:08:24 +0200 Subject: [PATCH 05/19] scripts : add speculative prefill stress and evaluation scripts Assisted-by: Antigravity --- scripts/eval_spec_prefill_entropy.py | 125 ++++++++++ scripts/eval_spec_prefill_failure_modes.py | 228 ++++++++++++++++++ ...al_spec_prefill_longbench_nonparametric.py | 105 ++++++++ scripts/eval_spec_prefill_longbench_real.py | 143 +++++++++++ .../eval_spec_prefill_targeted_failures.py | 214 ++++++++++++++++ 5 files changed, 815 insertions(+) create mode 100644 scripts/eval_spec_prefill_entropy.py create mode 100644 scripts/eval_spec_prefill_failure_modes.py create mode 100644 scripts/eval_spec_prefill_longbench_nonparametric.py create mode 100644 scripts/eval_spec_prefill_longbench_real.py create mode 100644 scripts/eval_spec_prefill_targeted_failures.py 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_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() From b3173fc2b5dfb05407e1b9b9cc177485e5c1a883 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 13:16:36 +0200 Subject: [PATCH 06/19] speculative-prefill : add dedicated draft device parameter Assisted-by: Antigravity --- common/arg.cpp | 9 +++++++++ common/common.h | 1 + examples/speculative-prefill/speculative-prefill.cpp | 5 +++++ tests/test-arg-parser.cpp | 8 +++++++- 4 files changed, 22 insertions(+), 1 deletion(-) diff --git a/common/arg.cpp b/common/arg.cpp index 5b26c5d16de..b92030116c3 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4231,6 +4231,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).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-p", "--spec-prefill-percentage"}, "P", string_format("fraction of prompt tokens to retain during speculative prefill (default: %.2f)", (double) params.speculative.prefill.percentage), diff --git a/common/common.h b/common/common.h index 21b49f1c76f..aee3cc9176d 100644 --- a/common/common.h +++ b/common/common.h @@ -370,6 +370,7 @@ struct common_params_speculative_prefill { bool enabled = false; // enable speculative prefill common_params_model model; // draft model for speculative prefill 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 diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index b1014aaf12e..1d84cdb990d 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -59,6 +59,11 @@ int main(int argc, char ** argv) { 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); diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 73d0e4c488c..3837a7756e0 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -197,10 +197,16 @@ 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"}; + argv = {"binary_name", "-mpd", "prefill-draft.gguf", "-nglpd", "24", "-devpd", "none"}; 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.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); argv = {"binary_name", "-lm", "none"}; From 4db32fa915e4108b2f232abad4f740ecae219e08 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 13:48:43 +0200 Subject: [PATCH 07/19] speculative-prefill : wire sparse prefill into server and CLI --- tools/server/server-context.cpp | 155 ++++++++++++++++++++++++++++---- 1 file changed, 136 insertions(+), 19 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b..d2902c6e24d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -14,6 +14,7 @@ #include "log.h" #include "sampling.h" #include "speculative.h" +#include "speculative-prefill.h" #include "mtmd.h" #include "mtmd-helper.h" @@ -212,6 +213,9 @@ struct server_slot { common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + 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; @@ -358,6 +362,9 @@ struct server_slot { // clear multimodal state mbatch.reset(); + + spec_prefill_active = false; + spec_prefill_tokens.clear(); } void init_sampler() const { @@ -425,6 +432,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"); @@ -842,6 +857,11 @@ struct server_context_impl { common_speculative_init_result_ptr spec_init; + common_init_result_ptr llama_init_spf; + 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; @@ -891,6 +911,11 @@ struct server_context_impl { spec.reset(); spec_init.reset(); + smpl_spf.reset(); + llama_init_spf.reset(); + ctx_spf = nullptr; + model_spf = nullptr; + ctx_dft = nullptr; model_dft = nullptr; @@ -953,6 +978,51 @@ 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_mtmd) { + return; + } + + const llama_tokens & prompt = slot.task->tokens.get_tokens(); + if (prompt.empty()) { + return; + } + + llama_memory_clear(llama_get_memory(ctx_spf), false); + common_sampler_reset(smpl_spf.get()); + + const common_speculative_prefill_result res = common_speculative_prefill_execute( + ctx_spf, smpl_spf.get(), prompt, 0, params_base.speculative.prefill); + + 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, false); + 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) { @@ -1102,6 +1172,48 @@ struct server_context_impl { load_progress_callback(1.0f, &load_progress_spec); } + if (params_base.speculative.prefill.enabled) { + common_params_model spf_model = params_base.speculative.prefill.model; + if (spf_model.empty()) { + spf_model = params_base.speculative.draft.mparams; + } + + if (spf_model.empty()) { + SRV_WRN("%s", "speculative prefill enabled but no draft model was provided; disabling\n"); + params_base.speculative.prefill.enabled = false; + } else { + common_params params_spf = common_base_params_to_speculative(params_base); + params_spf.model = spf_model; + 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 = 1; + params_spf.speculative.prefill.enabled = false; + params_spf.speculative.draft.mparams = {}; + + 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_spf = llama_init_spf ? llama_init_spf->model() : nullptr; + ctx_spf = llama_init_spf ? llama_init_spf->context() : nullptr; + if (model_spf == nullptr || ctx_spf == nullptr) { + SRV_ERR("failed to load speculative prefill draft model, '%s'\n", params_spf.model.path.c_str()); + 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"}}); @@ -3017,7 +3129,12 @@ 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) { + 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(); @@ -3029,7 +3146,7 @@ struct server_context_impl { 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) { @@ -3066,33 +3183,33 @@ 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; @@ -3175,7 +3292,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)); @@ -3191,7 +3308,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; @@ -3216,7 +3333,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; @@ -3284,8 +3401,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); } @@ -3311,7 +3428,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; } } @@ -3362,7 +3479,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; @@ -3405,7 +3522,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) { @@ -3451,7 +3568,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; } @@ -3467,13 +3584,13 @@ 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; GGML_ASSERT(batch.size() > 0); From dd7c7e26dd9e6a583d64fb5497ec09bb00d72569 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 14:15:36 +0200 Subject: [PATCH 08/19] speculative-prefill : fix attention capture, RoPE positions, and draft reuse Assisted-by: Antigravity --- common/common.h | 4 + common/speculative-prefill.cpp | 42 +++++---- .../speculative-prefill.cpp | 30 ++++++- src/llama-context.cpp | 4 + tools/llama-bench/llama-bench.cpp | 20 ++++- tools/server/server-context.cpp | 88 +++++++++++++++---- 6 files changed, 152 insertions(+), 36 deletions(-) diff --git a/common/common.h b/common/common.h index aee3cc9176d..9a962911b86 100644 --- a/common/common.h +++ b/common/common.h @@ -398,6 +398,10 @@ struct common_params_speculative { return !draft.mparams.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 index ddcd8d696d2..7221e3204d1 100644 --- a/common/speculative-prefill.cpp +++ b/common/speculative-prefill.cpp @@ -132,7 +132,14 @@ std::vector common_speculative_prefill_select_indices( } if (params.keep_last) { - selected.push_back(n_prompt - 1); + 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()); @@ -148,7 +155,7 @@ struct cb_attn_collector_data { static bool cb_collect_attn(ggml_tensor * t, bool ask, void * user_data) { if (ask) { - return true; + return strncmp(t->name, "kq_soft_max", 11) == 0; } if (strncmp(t->name, "kq_soft_max", 11) != 0) { @@ -160,23 +167,23 @@ static bool cb_collect_attn(ggml_tensor * t, bool ask, void * user_data) { 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 (n_past_total < n_prompt || n_heads <= 0) { + if (t->nb[0] != sizeof(float) || n_heads <= 0 || n_past_total < n_prompt) { return true; } - std::vector raw_buf((size_t) ggml_nelements(t)); - ggml_backend_tensor_get(t, raw_buf.data(), 0, ggml_nbytes(t)); - std::vector layer_data((size_t) n_heads * n_prompt); for (int32_t h = 0; h < n_heads; ++h) { - const float * src = raw_buf.data() + (size_t) h * n_past_total; - float * dst = layer_data.data() + (size_t) h * n_prompt; - std::copy(src, src + n_prompt, dst); + 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)); @@ -295,12 +302,17 @@ common_speculative_prefill_result common_speculative_prefill_execute( // detach callback llama_set_eval_callback(ctx_dft, nullptr, nullptr); - if (actual_steps > 0) { - for (size_t i = 0; i < prompt.size(); ++i) { - total_importance[i] /= (float) actual_steps; - } - } else { - std::fill(total_importance.begin(), total_importance.end(), 1.0f); + 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(); diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index 1d84cdb990d..d47daa918df 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -14,6 +14,8 @@ #include #include +#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 + int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -76,6 +78,27 @@ int main(int argc, char ** argv) { llama_context * ctx_dft = init_dft->context(); 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); @@ -135,7 +158,7 @@ int main(int argc, char ** argv) { 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); + common_batch_add(batch_tgt, prompt_tokens[orig_idx], (llama_pos) orig_idx, { seq_id }, is_last); } ret = llama_decode(ctx_tgt, batch_tgt); @@ -147,6 +170,9 @@ int main(int argc, char ** argv) { } 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; @@ -160,7 +186,7 @@ int main(int argc, char ** argv) { LOG("\n--- Generation Start ---\n"); int32_t n_predict = params.n_predict > 0 ? params.n_predict : 32; - int32_t cur_pos = (int32_t) spec_res.kept_indices.size(); + int32_t cur_pos = spec_res.kept_indices.empty() ? 0 : spec_res.kept_indices.back() + 1; int32_t n_generated = 0; llama_batch batch_gen = llama_batch_init(1, 0, 1); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 49bd4ae7dd1..4861d0e3a0b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3728,6 +3728,10 @@ void llama_set_abort_callback(llama_context * ctx, bool (*abort_callback)(void * } 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); } diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 8d9f9ce3f01..83dc147e539 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -2444,6 +2444,11 @@ static bool test_speculative_prefill_prompt( 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); @@ -2471,7 +2476,7 @@ static bool test_speculative_prefill_prompt( 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); + common_batch_add(batch_tgt, prompt_tokens[orig_idx], (llama_pos) orig_idx, { seq_id }, is_last); } const int ret = llama_decode(ctx_tgt, batch_tgt); @@ -2709,6 +2714,18 @@ int llama_bench(int argc, char ** argv) { } } + 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()); @@ -2746,6 +2763,7 @@ int llama_bench(int argc, char ** argv) { } common_params_sampling sparams_dft; + sparams_dft.temp = 0.0f; smpl_dft.reset(common_sampler_init(lmodel_dft, sparams_dft)); spf_params.enabled = true; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d2902c6e24d..e23eac19208 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -215,6 +215,8 @@ struct server_slot { bool spec_prefill_active = false; server_tokens spec_prefill_tokens; + std::vector spec_prefill_pos; // parallel to spec_prefill_tokens, orig indices + llama_pos spec_prefill_gen_pos = 0; // first generation pos = last orig + 1 // 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 @@ -365,6 +367,8 @@ struct server_slot { spec_prefill_active = false; spec_prefill_tokens.clear(); + spec_prefill_pos.clear(); + spec_prefill_gen_pos = 0; } void init_sampler() const { @@ -477,10 +481,15 @@ struct server_slot { // no speculative decoding i_batch = batch.size(); + llama_pos pos = prompt.tokens.pos_next(); + if (spec_prefill_active) { + pos = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); + } + 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", @@ -496,7 +505,10 @@ 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(); + if (spec_prefill_active) { + pos0 = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); + } add_ok &= batch.add(id, sampled, pos0++, true, false); for (auto token : spec_draft) { @@ -858,6 +870,7 @@ 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; @@ -909,13 +922,15 @@ struct server_context_impl { void destroy() { spec.reset(); - spec_init.reset(); smpl_spf.reset(); + ctx_spf_own.reset(); llama_init_spf.reset(); ctx_spf = nullptr; model_spf = nullptr; + spec_init.reset(); + ctx_dft = nullptr; model_dft = nullptr; @@ -981,6 +996,8 @@ struct server_context_impl { void apply_spec_prefill(server_slot & slot) { slot.spec_prefill_active = false; slot.spec_prefill_tokens.clear(); + slot.spec_prefill_pos.clear(); + slot.spec_prefill_gen_pos = 0; if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { return; @@ -994,11 +1011,11 @@ struct server_context_impl { return; } - llama_memory_clear(llama_get_memory(ctx_spf), false); + 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, 0, params_base.speculative.prefill); + ctx_spf, smpl_spf.get(), prompt, slot.id, params_base.speculative.prefill); if (res.kept_indices.empty() || (int32_t) res.kept_indices.size() >= (int32_t) prompt.size()) { return; @@ -1006,16 +1023,20 @@ struct server_context_impl { llama_tokens kept; kept.reserve(res.kept_indices.size()); + slot.spec_prefill_pos.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]); + slot.spec_prefill_pos.push_back(idx); } } if (kept.size() < 2) { + slot.spec_prefill_pos.clear(); return; } slot.spec_prefill_tokens = server_tokens(kept, false); + slot.spec_prefill_gen_pos = res.kept_indices.back() + 1; slot.spec_prefill_active = true; SLT_INF(slot, "speculative prefill kept %d / %d tokens (%.1f%%)\n", @@ -1192,24 +1213,51 @@ struct server_context_impl { } params_spf.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; params_spf.fit_params = false; - params_spf.n_parallel = 1; + params_spf.n_parallel = params_base.n_parallel; params_spf.speculative.prefill.enabled = false; params_spf.speculative.draft.mparams = {}; - 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_spf = llama_init_spf ? llama_init_spf->model() : nullptr; - ctx_spf = llama_init_spf ? llama_init_spf->context() : nullptr; + const bool reuse_dft = model_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()); + llama_context_params cparams = common_context_params_to_llama(params_spf); + ctx_spf_own.reset(llama_init_from_model(model_dft, cparams)); + model_spf = model_dft; + ctx_spf = ctx_spf_own.get(); + } 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_spf = llama_init_spf ? llama_init_spf->model() : nullptr; + ctx_spf = llama_init_spf ? llama_init_spf->context() : nullptr; + } if (model_spf == nullptr || ctx_spf == nullptr) { SRV_ERR("failed to load speculative prefill draft model, '%s'\n", params_spf.model.path.c_str()); 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; + + 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; disabling\n", vocab_diff); + params_base.speculative.prefill.enabled = false; + smpl_spf.reset(); + ctx_spf_own.reset(); + llama_init_spf.reset(); + ctx_spf = nullptr; + model_spf = nullptr; + } else { + 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; + } } } } @@ -3540,9 +3588,13 @@ 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(); + if (slot.spec_prefill_active) { + pos = (llama_pos) slot.spec_prefill_pos[slot.prompt.n_tokens()]; + } 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); From 57b8386fdde2ba1b86d99f903e2492eacc874068 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 15:40:29 +0200 Subject: [PATCH 09/19] server : fix prompt throughput calculation for speculative prefill Assisted-by: Gemini --- examples/speculative-prefill/speculative-prefill.cpp | 3 ++- tools/server/server-context.cpp | 11 +++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index d47daa918df..5bee66e86e8 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -178,7 +178,8 @@ int main(int argc, char ** argv) { 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\n", __func__, ttft_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)); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e23eac19208..b9312f1ddfb 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -604,7 +604,7 @@ struct server_slot { } 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; 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); @@ -678,7 +678,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); @@ -3178,6 +3178,7 @@ 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) { if (slot.state == SLOT_STATE_STARTED) { + slot.stats.update_prompt_start(); apply_spec_prefill(slot); } @@ -3189,8 +3190,6 @@ struct server_context_impl { // 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", @@ -3645,6 +3644,10 @@ struct server_context_impl { 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 From c4f085a444f76b5c9de72415476d05d701781dce Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 16:14:54 +0200 Subject: [PATCH 10/19] server : allow speculative prefill on text prompts when mmproj is loaded Assisted-by: Gemini --- tools/server/server-common.cpp | 4 ++-- tools/server/server-common.h | 2 ++ tools/server/server-context.cpp | 4 ++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 4f5b8202aca..9b7a0256cd4 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 f8ea82ef4cf..992419aefa6 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -219,6 +219,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 b9312f1ddfb..24e2df8de86 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1002,7 +1002,7 @@ struct server_context_impl { if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { return; } - if (!slot.task || !slot.task->need_logits() || slot.task->tokens.has_mtmd) { + if (!slot.task || !slot.task->need_logits() || slot.task->tokens.has_media()) { return; } @@ -1035,7 +1035,7 @@ struct server_context_impl { return; } - slot.spec_prefill_tokens = server_tokens(kept, false); + slot.spec_prefill_tokens = server_tokens(kept, slot.task->tokens.has_mtmd); slot.spec_prefill_gen_pos = res.kept_indices.back() + 1; slot.spec_prefill_active = true; From e3c4aac7f30684905342af3f0922ff0e334f7668 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 16:37:19 +0200 Subject: [PATCH 11/19] server : use contiguous positions for speculative prefill tokens Assisted-by: Gemini --- tools/server/server-context.cpp | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 24e2df8de86..d1fb294f70f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -215,8 +215,6 @@ struct server_slot { bool spec_prefill_active = false; server_tokens spec_prefill_tokens; - std::vector spec_prefill_pos; // parallel to spec_prefill_tokens, orig indices - llama_pos spec_prefill_gen_pos = 0; // first generation pos = last orig + 1 // 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 @@ -367,8 +365,6 @@ struct server_slot { spec_prefill_active = false; spec_prefill_tokens.clear(); - spec_prefill_pos.clear(); - spec_prefill_gen_pos = 0; } void init_sampler() const { @@ -482,9 +478,6 @@ struct server_slot { i_batch = batch.size(); llama_pos pos = prompt.tokens.pos_next(); - if (spec_prefill_active) { - pos = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); - } if (!inp_embd.empty()) { add_ok &= batch.add(id, inp_embd, pos, true, false); @@ -506,9 +499,6 @@ struct server_slot { } llama_pos pos0 = prompt.tokens.pos_next(); - if (spec_prefill_active) { - pos0 = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); - } add_ok &= batch.add(id, sampled, pos0++, true, false); for (auto token : spec_draft) { @@ -996,8 +986,6 @@ struct server_context_impl { void apply_spec_prefill(server_slot & slot) { slot.spec_prefill_active = false; slot.spec_prefill_tokens.clear(); - slot.spec_prefill_pos.clear(); - slot.spec_prefill_gen_pos = 0; if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { return; @@ -1023,20 +1011,16 @@ struct server_context_impl { llama_tokens kept; kept.reserve(res.kept_indices.size()); - slot.spec_prefill_pos.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]); - slot.spec_prefill_pos.push_back(idx); } } if (kept.size() < 2) { - slot.spec_prefill_pos.clear(); return; } slot.spec_prefill_tokens = server_tokens(kept, slot.task->tokens.has_mtmd); - slot.spec_prefill_gen_pos = res.kept_indices.back() + 1; slot.spec_prefill_active = true; SLT_INF(slot, "speculative prefill kept %d / %d tokens (%.1f%%)\n", @@ -3588,9 +3572,6 @@ struct server_context_impl { // 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(); - if (slot.spec_prefill_active) { - pos = (llama_pos) slot.spec_prefill_pos[slot.prompt.n_tokens()]; - } add_ok &= batch.add(slot.id, cur_tok, /* pos = */ pos, From 027adeaaf3370f5ec7c449af4507816d9f1a1f30 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 16:47:59 +0200 Subject: [PATCH 12/19] server : scale intermediate prompt progress throughput for speculative prefill Assisted-by: Gemini --- tools/server/server-context.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d1fb294f70f..be0d4ed955e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -593,11 +593,18 @@ struct server_slot { return; } - const double n_prompt_second = stats.n_prompt_tps(); 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 { From 2973ac3d450e40e2047f6557c95f2937c82f2478 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 25 Aug 2026 23:56:58 +0200 Subject: [PATCH 13/19] speculative-prefill : align server RoPE positions and set greedy lookahead in CLI Assisted-by: Antigravity --- .../speculative-prefill.cpp | 4 +++- tools/server/server-context.cpp | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index 5bee66e86e8..36546466173 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -120,7 +120,9 @@ int main(int argc, char ** argv) { llama_seq_id seq_id = 0; // initialize draft sampler for lookahead steps - common_sampler_ptr smpl_dft(common_sampler_init(model_dft, params.sampling)); + 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(); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index be0d4ed955e..7ad0b79a62a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -215,6 +215,8 @@ struct server_slot { bool spec_prefill_active = false; server_tokens spec_prefill_tokens; + std::vector spec_prefill_pos; // parallel to spec_prefill_tokens, orig indices + llama_pos spec_prefill_gen_pos = 0; // first generation pos = last orig + 1 // 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 @@ -365,6 +367,8 @@ struct server_slot { spec_prefill_active = false; spec_prefill_tokens.clear(); + spec_prefill_pos.clear(); + spec_prefill_gen_pos = 0; } void init_sampler() const { @@ -478,6 +482,9 @@ struct server_slot { i_batch = batch.size(); llama_pos pos = prompt.tokens.pos_next(); + if (spec_prefill_active) { + pos = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); + } if (!inp_embd.empty()) { add_ok &= batch.add(id, inp_embd, pos, true, false); @@ -499,6 +506,9 @@ struct server_slot { } llama_pos pos0 = prompt.tokens.pos_next(); + if (spec_prefill_active) { + pos0 = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); + } add_ok &= batch.add(id, sampled, pos0++, true, false); for (auto token : spec_draft) { @@ -993,6 +1003,8 @@ struct server_context_impl { void apply_spec_prefill(server_slot & slot) { slot.spec_prefill_active = false; slot.spec_prefill_tokens.clear(); + slot.spec_prefill_pos.clear(); + slot.spec_prefill_gen_pos = 0; if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { return; @@ -1018,16 +1030,20 @@ struct server_context_impl { llama_tokens kept; kept.reserve(res.kept_indices.size()); + slot.spec_prefill_pos.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]); + slot.spec_prefill_pos.push_back(idx); } } if (kept.size() < 2) { + slot.spec_prefill_pos.clear(); return; } slot.spec_prefill_tokens = server_tokens(kept, slot.task->tokens.has_mtmd); + slot.spec_prefill_gen_pos = res.kept_indices.back() + 1; slot.spec_prefill_active = true; SLT_INF(slot, "speculative prefill kept %d / %d tokens (%.1f%%)\n", @@ -3579,6 +3595,9 @@ struct server_context_impl { // 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(); + if (slot.spec_prefill_active) { + pos = (llama_pos) slot.spec_prefill_pos[slot.prompt.n_tokens()]; + } add_ok &= batch.add(slot.id, cur_tok, /* pos = */ pos, From cc935057e824f4da76bdb370d27448ac3d112012 Mon Sep 17 00:00:00 2001 From: rockofox Date: Wed, 26 Aug 2026 10:58:19 +0200 Subject: [PATCH 14/19] speculative-prefill : use contiguous positions and disable for recurrent/hybrid and DFlash/DSpark Assisted-by: opencode --- .../speculative-prefill.cpp | 9 ++++- tools/llama-bench/llama-bench.cpp | 7 +++- tools/server/server-context.cpp | 37 +++++++++---------- 3 files changed, 31 insertions(+), 22 deletions(-) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index 36546466173..a1111f1188f 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -54,6 +54,11 @@ int main(int argc, char ** argv) { llama_model * model_tgt = init_tgt->model(); llama_context * ctx_tgt = init_tgt->context(); + if (llama_model_is_recurrent(model_tgt) || llama_model_is_hybrid(model_tgt)) { + LOG_ERR("%s: speculative prefill is not supported for recurrent or hybrid 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); @@ -160,7 +165,7 @@ int main(int argc, char ** argv) { 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) orig_idx, { seq_id }, is_last); + common_batch_add(batch_tgt, prompt_tokens[orig_idx], (llama_pos) k, { seq_id }, is_last); } ret = llama_decode(ctx_tgt, batch_tgt); @@ -189,7 +194,7 @@ int main(int argc, char ** argv) { LOG("\n--- Generation Start ---\n"); int32_t n_predict = params.n_predict > 0 ? params.n_predict : 32; - int32_t cur_pos = spec_res.kept_indices.empty() ? 0 : spec_res.kept_indices.back() + 1; + int32_t cur_pos = n_kept_total; int32_t n_generated = 0; llama_batch batch_gen = llama_batch_init(1, 0, 1); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 83dc147e539..4853b925941 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -2439,6 +2439,11 @@ static bool test_speculative_prefill_prompt( 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) || llama_model_is_hybrid(model_tgt)) { + fprintf(stderr, "%s: speculative prefill is not supported for recurrent or hybrid 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); @@ -2476,7 +2481,7 @@ static bool test_speculative_prefill_prompt( 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) orig_idx, { seq_id }, is_last); + 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); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7ad0b79a62a..aa3e3592262 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -215,8 +215,6 @@ struct server_slot { bool spec_prefill_active = false; server_tokens spec_prefill_tokens; - std::vector spec_prefill_pos; // parallel to spec_prefill_tokens, orig indices - llama_pos spec_prefill_gen_pos = 0; // first generation pos = last orig + 1 // 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 @@ -367,8 +365,6 @@ struct server_slot { spec_prefill_active = false; spec_prefill_tokens.clear(); - spec_prefill_pos.clear(); - spec_prefill_gen_pos = 0; } void init_sampler() const { @@ -482,9 +478,6 @@ struct server_slot { i_batch = batch.size(); llama_pos pos = prompt.tokens.pos_next(); - if (spec_prefill_active) { - pos = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); - } if (!inp_embd.empty()) { add_ok &= batch.add(id, inp_embd, pos, true, false); @@ -506,9 +499,6 @@ struct server_slot { } llama_pos pos0 = prompt.tokens.pos_next(); - if (spec_prefill_active) { - pos0 = spec_prefill_gen_pos + (prompt.n_tokens() - (int) spec_prefill_tokens.size()); - } add_ok &= batch.add(id, sampled, pos0++, true, false); for (auto token : spec_draft) { @@ -1003,8 +993,6 @@ struct server_context_impl { void apply_spec_prefill(server_slot & slot) { slot.spec_prefill_active = false; slot.spec_prefill_tokens.clear(); - slot.spec_prefill_pos.clear(); - slot.spec_prefill_gen_pos = 0; if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { return; @@ -1030,20 +1018,16 @@ struct server_context_impl { llama_tokens kept; kept.reserve(res.kept_indices.size()); - slot.spec_prefill_pos.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]); - slot.spec_prefill_pos.push_back(idx); } } if (kept.size() < 2) { - slot.spec_prefill_pos.clear(); return; } slot.spec_prefill_tokens = server_tokens(kept, slot.task->tokens.has_mtmd); - slot.spec_prefill_gen_pos = res.kept_indices.back() + 1; slot.spec_prefill_active = true; SLT_INF(slot, "speculative prefill kept %d / %d tokens (%.1f%%)\n", @@ -1200,6 +1184,24 @@ 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) || llama_model_is_hybrid(model_tgt)) { + SRV_WRN("%s", "speculative prefill is not supported for recurrent or hybrid models; disabling\n"); + params_base.speculative.prefill.enabled = false; + } else { + const bool spec_dflash = std::find(params_base.speculative.types.begin(), + params_base.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params_base.speculative.types.end(); + const bool spec_dspark = std::find(params_base.speculative.types.begin(), + params_base.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params_base.speculative.types.end(); + if (spec_dflash || spec_dspark) { + SRV_WRN("%s", "speculative prefill is not supported with DFlash or DSpark speculative decoding; disabling\n"); + params_base.speculative.prefill.enabled = false; + } + } + } + if (params_base.speculative.prefill.enabled) { common_params_model spf_model = params_base.speculative.prefill.model; if (spf_model.empty()) { @@ -3595,9 +3597,6 @@ struct server_context_impl { // 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(); - if (slot.spec_prefill_active) { - pos = (llama_pos) slot.spec_prefill_pos[slot.prompt.n_tokens()]; - } add_ok &= batch.add(slot.id, cur_tok, /* pos = */ pos, From 284a787edace0ef00dd420e51af413e84112967d Mon Sep 17 00:00:00 2001 From: rockofox Date: Wed, 26 Aug 2026 11:14:30 +0200 Subject: [PATCH 15/19] speculative-prefill : allow hybrid models, keep gate for recurrent only Assisted-by: opencode --- examples/speculative-prefill/speculative-prefill.cpp | 4 ++-- tools/llama-bench/llama-bench.cpp | 4 ++-- tools/server/server-context.cpp | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index a1111f1188f..3778771ac44 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -54,8 +54,8 @@ int main(int argc, char ** argv) { llama_model * model_tgt = init_tgt->model(); llama_context * ctx_tgt = init_tgt->context(); - if (llama_model_is_recurrent(model_tgt) || llama_model_is_hybrid(model_tgt)) { - LOG_ERR("%s: speculative prefill is not supported for recurrent or hybrid models\n", __func__); + if (llama_model_is_recurrent(model_tgt)) { + LOG_ERR("%s: speculative prefill is not supported for recurrent models\n", __func__); return 1; } diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 4853b925941..13c61364de7 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -2439,8 +2439,8 @@ static bool test_speculative_prefill_prompt( 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) || llama_model_is_hybrid(model_tgt)) { - fprintf(stderr, "%s: speculative prefill is not supported for recurrent or hybrid models, skipping instance\n", __func__); + if (llama_model_is_recurrent(model_tgt)) { + fprintf(stderr, "%s: speculative prefill is not supported for recurrent models, skipping instance\n", __func__); return false; } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index aa3e3592262..551e54c9846 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1185,8 +1185,8 @@ struct server_context_impl { } if (params_base.speculative.prefill.enabled) { - if (llama_model_is_recurrent(model_tgt) || llama_model_is_hybrid(model_tgt)) { - SRV_WRN("%s", "speculative prefill is not supported for recurrent or hybrid models; disabling\n"); + if (llama_model_is_recurrent(model_tgt)) { + SRV_WRN("%s", "speculative prefill is not supported for recurrent models; disabling\n"); params_base.speculative.prefill.enabled = false; } else { const bool spec_dflash = std::find(params_base.speculative.types.begin(), From c0edaa2e99851662d534d1529f90518fd75f2f2f Mon Sep 17 00:00:00 2001 From: rockofox Date: Fri, 28 Aug 2026 09:07:36 +0200 Subject: [PATCH 16/19] server : clear speculative prefill draft context KV on slot release Assisted-by: Antigravity --- tools/server/server-context.cpp | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e2fb01e786a..6246525879c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -242,6 +242,7 @@ struct server_slot { llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; + llama_context * ctx_spf = nullptr; common_memory mem; @@ -339,6 +340,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(); } @@ -568,6 +572,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(); @@ -749,6 +757,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; @@ -1060,6 +1072,8 @@ struct server_context_impl { 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; } @@ -1438,6 +1452,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; From ea8fc46d8a464508565738a7fa251caf5b932441 Mon Sep 17 00:00:00 2001 From: rockofox Date: Fri, 28 Aug 2026 09:54:32 +0200 Subject: [PATCH 17/19] server : fix speculative prefill with dflash2 and target-dependent draft models Assisted-by: Antigravity --- .../speculative-prefill.cpp | 11 ++ tools/server/server-context.cpp | 151 +++++++++--------- 2 files changed, 88 insertions(+), 74 deletions(-) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index 3778771ac44..0034654d6c8 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -5,6 +5,7 @@ #include "sampling.h" #include "speculative.h" #include "speculative-prefill.h" +#include "../../src/llama-ext.h" #include #include @@ -82,6 +83,16 @@ int main(int argc, char ** argv) { llama_model * model_dft = init_dft->model(); llama_context * ctx_dft = init_dft->context(); + 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; + } + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index bf116d064dd..37f36e1a3f3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -15,6 +15,7 @@ #include "sampling.h" #include "speculative.h" #include "speculative-prefill.h" +#include "src/llama-ext.h" #include "mtmd.h" #include "mtmd-helper.h" @@ -1248,88 +1249,90 @@ struct server_context_impl { if (params_base.speculative.prefill.enabled) { if (llama_model_is_recurrent(model_tgt)) { - SRV_WRN("%s", "speculative prefill is not supported for recurrent models; disabling\n"); - params_base.speculative.prefill.enabled = false; - } else { - const bool spec_dflash = std::find(params_base.speculative.types.begin(), - params_base.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params_base.speculative.types.end(); - const bool spec_dspark = std::find(params_base.speculative.types.begin(), - params_base.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params_base.speculative.types.end(); - if (spec_dflash || spec_dspark) { - SRV_WRN("%s", "speculative prefill is not supported with DFlash or DSpark speculative decoding; disabling\n"); - params_base.speculative.prefill.enabled = false; - } + SRV_ERR("%s", "speculative prefill is not supported for recurrent models\n"); + return false; } - } - if (params_base.speculative.prefill.enabled) { + 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()) { + if (spf_model.empty() && !has_target_dependent_dft) { spf_model = params_base.speculative.draft.mparams; } if (spf_model.empty()) { - SRV_WRN("%s", "speculative prefill enabled but no draft model was provided; disabling\n"); - params_base.speculative.prefill.enabled = false; + 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_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()); + llama_context_params cparams = common_context_params_to_llama(params_spf); + ctx_spf_own.reset(llama_init_from_model(model_dft, cparams)); + model_spf = model_dft; + ctx_spf = ctx_spf_own.get(); } else { - common_params params_spf = common_base_params_to_speculative(params_base); - params_spf.model = spf_model; - 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 && ( - 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()); - llama_context_params cparams = common_context_params_to_llama(params_spf); - ctx_spf_own.reset(llama_init_from_model(model_dft, cparams)); - model_spf = model_dft; - ctx_spf = ctx_spf_own.get(); - } 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_spf = llama_init_spf ? llama_init_spf->model() : nullptr; - ctx_spf = llama_init_spf ? llama_init_spf->context() : nullptr; - } - if (model_spf == nullptr || ctx_spf == nullptr) { - SRV_ERR("failed to load speculative prefill draft model, '%s'\n", params_spf.model.path.c_str()); - return false; - } + 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_spf = llama_init_spf ? llama_init_spf->model() : nullptr; + ctx_spf = llama_init_spf ? llama_init_spf->context() : nullptr; + } + if (model_spf == nullptr || ctx_spf == nullptr) { + SRV_ERR("failed to load speculative prefill draft model, '%s'\n", params_spf.model.path.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; disabling\n", vocab_diff); - params_base.speculative.prefill.enabled = false; - smpl_spf.reset(); - ctx_spf_own.reset(); - llama_init_spf.reset(); - ctx_spf = nullptr; - model_spf = nullptr; - } else { - 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 (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; } } @@ -3375,7 +3378,7 @@ struct server_context_impl { 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); From 47c24ae8e5ccffd7d37100cc085e699a80fd25c3 Mon Sep 17 00:00:00 2001 From: rockofox Date: Sun, 30 Aug 2026 00:09:28 +0200 Subject: [PATCH 18/19] speculative-prefill : add max context option for draft model Assisted-by: Antigravity --- common/arg.cpp | 11 ++++++++ common/common.h | 1 + common/speculative-prefill.cpp | 14 +++++++++- .../speculative-prefill.cpp | 3 +++ tests/test-arg-parser.cpp | 3 ++- tools/llama-bench/llama-bench.cpp | 27 ++++++++++++++++++- tools/server/server-context.cpp | 3 +++ 7 files changed, 59 insertions(+), 3 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index baf53632395..8226abba5d2 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4316,6 +4316,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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 = main context size)", 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), diff --git a/common/common.h b/common/common.h index cb5ea368b39..53785f82bb2 100644 --- a/common/common.h +++ b/common/common.h @@ -370,6 +370,7 @@ struct common_params_speculative_ngram_cache { 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) diff --git a/common/speculative-prefill.cpp b/common/speculative-prefill.cpp index 7221e3204d1..ab20ad08f1e 100644 --- a/common/speculative-prefill.cpp +++ b/common/speculative-prefill.cpp @@ -212,6 +212,19 @@ common_speculative_prefill_result common_speculative_prefill_execute( 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); @@ -250,7 +263,6 @@ common_speculative_prefill_result common_speculative_prefill_execute( res.t_draft_eval_us = t_prefill_end - t_start; // 2. lookahead decode steps with attention extraction - const int32_t lookahead = std::max(1, params.look_ahead_cnt); std::vector total_importance(prompt.size(), 0.0f); cb_attn_collector_data cb_data; diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index 0034654d6c8..f7c330be742 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -64,6 +64,9 @@ int main(int argc, char ** argv) { 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; } diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 0b7802f8c73..b823858405e 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -254,11 +254,12 @@ 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"}; + 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"}; diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 6d33fea63ef..090c560a47f 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -343,6 +343,7 @@ struct cmd_params { 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; @@ -395,6 +396,7 @@ static const cmd_params cmd_params_defaults = { /* 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 }, @@ -478,6 +480,8 @@ static void print_usage(int /* argc */, char ** argv) { 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"); @@ -635,6 +639,15 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } 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) { @@ -1248,6 +1261,12 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { 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; } @@ -1351,6 +1370,7 @@ 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; @@ -1468,6 +1488,7 @@ static std::vector get_cmd_params_instances(const cmd_param // 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)) @@ -1504,6 +1525,7 @@ static std::vector get_cmd_params_instances(const cmd_param cmd_params_instance instance = { /* .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, @@ -1546,6 +1568,7 @@ static std::vector get_cmd_params_instances(const cmd_param cmd_params_instance instance = { /* .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, @@ -1588,6 +1611,7 @@ static std::vector get_cmd_params_instances(const cmd_param cmd_params_instance instance = { /* .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, @@ -2747,7 +2771,7 @@ int llama_bench(int argc, char ** argv) { if (has_spec_prefill) { llama_context_params dft_cparams = llama_context_default_params(); - dft_cparams.n_ctx = inst.n_prompt + inst.spec_prefill_lookahead + 16; + dft_cparams.n_ctx = inst.spec_prefill_n_ctx > 0 ? inst.spec_prefill_n_ctx : (inst.n_prompt + inst.spec_prefill_lookahead + 16); dft_cparams.n_batch = inst.n_batch; dft_cparams.n_ubatch = inst.n_ubatch; dft_cparams.type_k = inst.type_k; @@ -2772,6 +2796,7 @@ int llama_bench(int argc, char ** argv) { 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; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 37f36e1a3f3..e87ea516b58 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1275,6 +1275,9 @@ struct server_context_impl { 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; } From 79acc585ce9a7b3d8a82f69c814d7705577d0f41 Mon Sep 17 00:00:00 2001 From: rockofox Date: Tue, 1 Sep 2026 11:35:37 +0200 Subject: [PATCH 19/19] speculative-prefill : clamp default draft context to training limit Assisted-by: Antigravity --- common/arg.cpp | 2 +- .../speculative-prefill.cpp | 20 ++++++++++++++++-- tools/llama-bench/llama-bench.cpp | 3 ++- tools/server/server-context.cpp | 21 +++++++++++++------ 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 52584a3f65f..6b95cadd809 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4318,7 +4318,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).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 = main context size)", params.speculative.prefill.n_ctx), + 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"); diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp index f7c330be742..1a7ae656046 100644 --- a/examples/speculative-prefill/speculative-prefill.cpp +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -77,14 +77,17 @@ int main(int argc, char ** argv) { } 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); + 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(); - llama_context * ctx_dft = init_dft->context(); + 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__); @@ -96,6 +99,19 @@ int main(int argc, char ** argv) { 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); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 6e434e929f0..11d1e0a353b 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -2830,7 +2830,8 @@ int llama_bench(int argc, char ** argv) { if (has_spec_prefill) { llama_context_params dft_cparams = llama_context_default_params(); - dft_cparams.n_ctx = inst.spec_prefill_n_ctx > 0 ? inst.spec_prefill_n_ctx : (inst.n_prompt + inst.spec_prefill_lookahead + 16); + 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; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e87ea516b58..f41650ecb65 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1296,21 +1296,30 @@ struct server_context_impl { if (reuse_dft) { SRV_INF("reusing draft model for speculative prefill '%s'\n", params_spf.model.get_name().c_str()); - llama_context_params cparams = common_context_params_to_llama(params_spf); - ctx_spf_own.reset(llama_init_from_model(model_dft, cparams)); model_spf = model_dft; - ctx_spf = ctx_spf_own.get(); } 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); + llama_init_spf = common_init_from_params(params_spf, /*model_only=*/true); model_spf = llama_init_spf ? llama_init_spf->model() : nullptr; - ctx_spf = llama_init_spf ? llama_init_spf->context() : nullptr; } - if (model_spf == nullptr || ctx_spf == 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;