Speculative prefill

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

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<common_download_task> 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<common_download_task> tasks;
 
-    auto & plan      = handler.plan;
-    auto & plan_spec = handler.plan_spec;
+    auto & plan         = handler.plan;
+    auto & plan_spec    = handler.plan_spec;
+    auto & plan_prefill = handler.plan_prefill;
 
     auto opts = handler.opts; // copy
     opts.callback = callback;
@@ -478,6 +484,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
     handle_url(params.model);
     handle_url(params.mmproj);
     handle_url(params.speculative.draft.mparams);
+    handle_url(params.speculative.prefill.model);
 
     // optionally, if docker repo is set, resolve it
     if (!params.model.docker_repo.empty()) {
@@ -513,6 +520,13 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         tasks.push_back(task);
         had_spec_url = true;
     }
+    if (!params.speculative.prefill.model.url.empty()) {
+        common_download_task task;
+        task.url        = params.speculative.prefill.model.url;
+        task.local_path = params.speculative.prefill.model.path;
+        task.opts       = opts;
+        tasks.push_back(task);
+    }
 
     // handle hf_plan tasks
     auto add_tasks = [&opts, &tasks](const hf_cache::hf_files  & model_files,
@@ -626,6 +640,11 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         had_spec_url = true;
     }
 
+    // handle plan_prefill (e.g. --spec-prefill-hf)
+    if (!plan_prefill.model_files.empty()) {
+        add_tasks(plan_prefill.model_files, plan_prefill.primary, params.speculative.prefill.model);
+    }
+
     if (!plan.model_files.empty()) {
         add_tasks(plan.model_files, plan.primary, params.model);
     }
@@ -4245,6 +4264,111 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.speculative.draft.mparams.hf_file = value; // will be used if --spec-draft-hf is set
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL"));
+    add_opt(common_arg(
+        {"--spec-prefill", "--speculative-prefill"},
+        "enable speculative prefill using draft model to filter prompt tokens",
+        [](common_params & params) {
+            params.speculative.prefill.enabled = true;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL"));
+    add_opt(common_arg(
+        {"--spec-prefill-draft-model", "--spec-prefill-model", "-mpd", "--speculative-prefill-model", "--speculative-prefill-draft-model"}, "FNAME",
+        "draft model for speculative prefill (default: unused)",
+        [](common_params & params, const std::string & value) {
+            params.speculative.prefill.model.path    = value;
+            params.speculative.prefill.model.hf_file = value; // will be used if --spec-prefill-hf is set
+            params.speculative.prefill.enabled       = true;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_MODEL"));
+    add_opt(common_arg(
+        {"--spec-prefill-draft-hf", "--spec-prefill-hf", "-hfpd", "--speculative-prefill-hf", "--speculative-prefill-draft-hf"}, "<user>/<model>[:quant]",
+        "Hugging Face model repository for speculative prefill draft model (default: unused)",
+        [](common_params & params, const std::string & value) {
+            params.speculative.prefill.model.hf_repo = value;
+            params.speculative.prefill.enabled       = true;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_HF"));
+    add_opt(common_arg(
+        {"--spec-prefill-draft-ngl", "--spec-prefill-ngl", "-nglpd", "--speculative-prefill-ngl", "--speculative-prefill-draft-ngl"}, "N",
+        string_format("max. number of speculative prefill draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)",
+            params.speculative.prefill.n_gpu_layers == -1 ? "auto" : "all"),
+        [](common_params & params, const std::string & value) {
+            if (value == "auto") {
+                params.speculative.prefill.n_gpu_layers = -1;
+            } else if (value == "all") {
+                params.speculative.prefill.n_gpu_layers = -2;
+            } else {
+                params.speculative.prefill.n_gpu_layers = std::stoi(value);
+            }
+            if (!llama_supports_gpu_offload()) {
+                fprintf(stderr, "warning: no usable GPU found, --spec-prefill-ngl option will be ignored\n");
+                fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n");
+                fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n");
+            }
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_N_GPU_LAYERS"));
+    add_opt(common_arg(
+        {"--spec-prefill-draft-device", "--spec-prefill-device", "-devpd", "--speculative-prefill-device", "--speculative-prefill-draft-device"}, "<dev1,dev2,..>",
+        "comma-separated list of devices to use for offloading the speculative prefill draft model (none = don't offload)\n"
+        "use --list-devices to see a list of available devices",
+        [](common_params & params, const std::string & value) {
+            params.speculative.prefill.devices = parse_device_list(value);
+            params.speculative.prefill.enabled = true;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_DEVICE"));
+    add_opt(common_arg(
+        {"--spec-prefill-draft-ctx", "--spec-prefill-ctx", "--spec-prefill-ctx-size", "--spec-prefill-max-ctx", "-cpd", "--speculative-prefill-ctx", "--speculative-prefill-max-ctx"}, "N",
+        string_format("context size for speculative prefill draft model (default: %d, 0 = draft training limit or main context)", params.speculative.prefill.n_ctx),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("spec-prefill context size must be >= 0");
+            }
+            params.speculative.prefill.n_ctx   = value;
+            params.speculative.prefill.enabled = true;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_CTX_SIZE"));
+    add_opt(common_arg(
+        {"--spec-prefill-p", "--spec-prefill-percentage"}, "P",
+        string_format("fraction of prompt tokens to retain during speculative prefill (default: %.2f)", (double) params.speculative.prefill.percentage),
+        [](common_params & params, const std::string & value) {
+            const float val = std::stof(value);
+            if (val <= 0.0f || val > 1.0f) {
+                throw std::invalid_argument("spec-prefill percentage must be between 0.0 (exclusive) and 1.0 (inclusive)");
+            }
+            params.speculative.prefill.enabled    = true;
+            params.speculative.prefill.percentage = val;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_P"));
+    add_opt(common_arg(
+        {"--spec-prefill-chunk", "--spec-prefill-chunk-size"}, "N",
+        string_format("chunk grouping size for speculative prefill (default: %d, 0 to disable)", params.speculative.prefill.chunk_size),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("spec-prefill chunk size must be >= 0");
+            }
+            params.speculative.prefill.chunk_size = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--spec-prefill-lookahead", "--spec-prefill-lah"}, "N",
+        string_format("number of lookahead decode steps on draft model for attention estimation (default: %d)", params.speculative.prefill.look_ahead_cnt),
+        [](common_params & params, int value) {
+            if (value < 1) {
+                throw std::invalid_argument("spec-prefill lookahead count must be >= 1");
+            }
+            params.speculative.prefill.look_ahead_cnt = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--spec-prefill-pool-kernel"}, "N",
+        string_format("1D average pooling kernel size for attention smoothing (default: %d, 0 to disable)", params.speculative.prefill.pool_kernel_size),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("spec-prefill pool kernel size must be >= 0");
+            }
+            params.speculative.prefill.pool_kernel_size = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
     add_opt(common_arg(
         {"--spec-type"}, common_speculative_all_types_str(),
         string_format("comma-separated list of types of speculative decoding to use (default: %s)\n",

common/arg.h

@@ -138,6 +138,7 @@ void common_params_add_preset_options(std::vector<common_arg> & 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;
 };
 

common/common.h

@@ -367,6 +367,20 @@ struct common_params_speculative_ngram_cache {
     std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding
 };
 
+struct common_params_speculative_prefill {
+    bool                enabled          = false; // enable speculative prefill
+    common_params_model model;                    // draft model for speculative prefill
+    int32_t             n_ctx            = 0;     // context size for draft model (0 = default/target context size)
+    int32_t             n_gpu_layers     = -1;    // max draft model layers to store in VRAM (-1 - use default)
+    std::vector<ggml_backend_dev_t> devices;      // devices to use for offloading the draft model
+    float               percentage       = 0.3f;  // fraction of prompt tokens to retain (0.0 < p <= 1.0)
+    int32_t             chunk_size       = 32;    // chunk grouping size (0 to disable chunking)
+    int32_t             look_ahead_cnt   = 8;     // lookahead decode steps on draft model
+    int32_t             pool_kernel_size = 13;    // 1D average pooling kernel size for smoothing
+    bool                keep_bos         = true;  // preserve first token (BOS)
+    bool                keep_last        = true;  // preserve last token / tail chunk
+};
+
 struct common_params_speculative {
     std::vector<enum common_speculative_type> types = { COMMON_SPECULATIVE_TYPE_NONE };
 
@@ -383,6 +397,8 @@ struct common_params_speculative {
 
     common_params_speculative_ngram_cache ngram_cache;
 
+    common_params_speculative_prefill prefill;
+
     bool has_dft() const {
         return !draft.mparams.empty();
     }
@@ -391,6 +407,9 @@ struct common_params_speculative {
         return synth_len != -1.0 || !synth_rates.empty();
     }
 
+    bool has_prefill() const {
+        return prefill.enabled && !prefill.model.empty();
+    }
     uint32_t need_n_rs_seq() const {
         bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
             return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;