diff --git a/common/speculative.cpp b/common/speculative.cpp index f2ef82e5f42..9e4f7e22cbd 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -923,8 +923,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; - // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout - const bool is_dspark; + // draft-dspark: the draft carries a Markov head. Comes from the model, not the + // requested type. The DSpark path also truncates on confidence, and for + // sample_from_anchor models it starts one row earlier. + bool is_dspark = false; // dspark speculators bool sample_from_anchor = true; @@ -939,7 +941,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) : common_speculative_impl(type, n_seq) , params(params.draft) - , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; @@ -952,6 +953,19 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { target_layer_ids_n = llama_model_target_layer_ids_n(model_dft); GGML_ASSERT(target_layer_ids_n > 0 && "DFlash model has no target_layer_ids"); + // Both lineages declare general.architecture = dflash, so the requested type cannot + // pick the draft path. The Markov head is the on-disk marker and is already loaded. + is_dspark = llama_model_has_dspark_markov_head(model_dft); + const bool type_says_dspark = (type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK); + if (type_says_dspark != is_dspark) { + LOG_WRN("%s: draft model carries %s, but --spec-type requested %s. Using %s, which is " + "what the model needs. The wrong path drops confidence truncation, and for " + "sample_from_anchor models it also reads the drafts one row late.\n", __func__, + is_dspark ? "a DSpark Markov head" : "no DSpark Markov head", + common_speculative_type_to_str(type).c_str(), + is_dspark ? "DSpark" : "DFlash"); + } + n_embd_tgt = llama_model_n_embd(model_tgt); n_embd_dec = llama_model_n_embd(model_dft); n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; @@ -971,8 +985,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); - LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, - block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false"); + LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s, lineage=%s\n", __func__, + block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false", + is_dspark ? "dspark" : "dflash"); // DFlash input is [id_last, * (block_size-1)]: in-place denoising yields at most // block_size-1 draft tokens, anchor-first DSpark yields a full block_size draft tokens diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa..f1d9800a7b9 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -124,6 +124,9 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model); // returns the number of extracted layers from target model LLAMA_API uint32_t llama_model_target_layer_ids_n(const struct llama_model * model); +// returns true if the draft model carries a DSpark Markov head. Both lineages declare +// general.architecture = dflash, so this is how to tell them apart. +LLAMA_API bool llama_model_has_dspark_markov_head(const struct llama_model * model); // retrieves the whole token embedding matrix in F32 format (n_embd * n_vocab) // returns total number of elements or 0 on error diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c34700ff563..2882241abe6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -3123,6 +3123,10 @@ uint32_t llama_model_target_layer_ids_n(const struct llama_model * model) { return (uint32_t) model->target_layer_ids.size(); } +bool llama_model_has_dspark_markov_head(const struct llama_model * model) { + return model->dspark_markov_w1 != nullptr; +} + uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) { if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr) { return 0; diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 5837195006b..871e5abaf93 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -53,6 +53,32 @@ def test_with_and_without_draft(): assert tokens_no_draft == tokens_draft +def test_draft_acceptance_floor(): + # The other tests check that drafting happens, not that drafts are accepted. + # This drafter/target pair is mismatched on purpose, so accept is only 15-17%. + global server + server.start() + for prompt in [ + "I believe the meaning of life is", + "Once upon a time there was a little girl who", + ]: + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, + "temperature": 0.0, + "top_k": 1, + "n_predict": 64, + }) + assert res.status_code == 200 + draft_n = res.body["timings"]["draft_n"] + draft_n_accepted = res.body["timings"]["draft_n_accepted"] + assert draft_n > 0, f"nothing was drafted for {prompt!r}" + accept_rate = draft_n_accepted / draft_n + assert accept_rate > 0.05, ( + f"draft acceptance collapsed to {100 * accept_rate:.2f}% " + f"({draft_n_accepted}/{draft_n}) for {prompt!r}" + ) + + def test_different_draft_min_draft_max(): global server test_values = [