diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index a5c73f957c0c..a21cb6996406 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -4,7 +4,10 @@ #include "speculative.h" #include "log.h" #include "llama.h" +#include "gguf.h" +#include "../../src/llama-ext.h" +#include #include #include #include @@ -13,6 +16,43 @@ #include #include +// dspark drafters carry the target layer-id list as an array-typed GGUF KV, +// which llama_model's string-KV cache skips -- read it from the drafter file +// directly (same as tools/server + tests/test-dspark-real-eval). +static std::vector read_dspark_target_layers(const std::string & drafter_path) { + struct gguf_init_params gp = { /* .no_alloc = */ true, /* .ctx = */ nullptr }; + gguf_context * gctx = gguf_init_from_file(drafter_path.c_str(), gp); + if (gctx == nullptr) { + return {}; + } + + std::vector out; + + const int64_t arch_kid = gguf_find_key(gctx, "general.architecture"); + if (arch_kid >= 0) { + const std::string key = std::string(gguf_get_val_str(gctx, arch_kid)) + ".dspark.target_layers"; + const int64_t kid = gguf_find_key(gctx, key.c_str()); + if (kid >= 0 && gguf_get_kv_type(gctx, kid) == GGUF_TYPE_ARRAY) { + const enum gguf_type arr_type = gguf_get_arr_type(gctx, kid); + const size_t n = gguf_get_arr_n(gctx, kid); + const void * data = gguf_get_arr_data(gctx, kid); + out.reserve(n); + for (size_t i = 0; i < n; i++) { + switch (arr_type) { + case GGUF_TYPE_INT32: out.push_back(((const int32_t *) data)[i]); break; + case GGUF_TYPE_UINT32: out.push_back((int32_t) ((const uint32_t *) data)[i]); break; + case GGUF_TYPE_INT64: out.push_back((int32_t) ((const int64_t *) data)[i]); break; + case GGUF_TYPE_UINT64: out.push_back((int32_t) ((const uint64_t *) data)[i]); break; + default: out.clear(); i = n; break; + } + } + } + } + + gguf_free(gctx); + return out; +} + int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -75,6 +115,18 @@ int main(int argc, char ** argv) { } auto cparams = common_context_params_to_llama(params_dft); + + // dspark stages all context rows since its cache position in one batch + // (worst case the full window), so its batch must cover the context -- + // otherwise every draft round is skipped and speculation degrades to AR. + const bool spec_dspark = std::find(params.speculative.types.begin(), + params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params.speculative.types.end(); + if (spec_dspark) { + cparams.n_batch = std::max(cparams.n_batch, cparams.n_ctx); + cparams.n_ubatch = std::max(cparams.n_ubatch, cparams.n_batch); + } + ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); params.speculative.draft.ctx_tgt = ctx_tgt; @@ -129,10 +181,6 @@ int main(int argc, char ** argv) { // target model sampling context common_sampler_ptr smpl(common_sampler_init(model_tgt, params.sampling)); - // eval the prompt - llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1)); - llama_decode(ctx_dft.get(), llama_batch_get_one(inp.data(), inp.size() - 1)); - // note: keep the last token separate! llama_token id_last = inp.back(); @@ -142,15 +190,53 @@ int main(int argc, char ** argv) { int n_past = inp.size() - 1; - // init the speculator + // init the speculator BEFORE the prompt is evaluated: capture-type drafters + // (dspark) stage their context features from the prompt decode itself, and + // their begin() clears the staging window -- so the order must be + // begin() -> prompt decode -> process(). const auto & params_spec = params.speculative; struct common_speculative * spec = common_speculative_init(params.speculative, 1); - common_speculative_begin(spec, seq_id, prompt_tgt); + const bool spec_capture = common_speculative_need_embd_capture(spec); + if (spec_capture) { + const std::vector capture_layers = read_dspark_target_layers(params.speculative.draft.mparams.path); + if (capture_layers.empty()) { + LOG_ERR("draft-dspark: failed to read dspark.target_layers from '%s'\n", params.speculative.draft.mparams.path.c_str()); + return 1; + } + // masked=false: capture stays dense (a row for every position) while + // batch.logits stays narrow -- no per-row full-vocab lm_head (fork #63). + llama_set_capture_layers(ctx_tgt, capture_layers.data(), capture_layers.size(), /* masked = */ false); + LOG_INF("draft-dspark: target tap capture engaged on %zu layers\n", capture_layers.size()); + } llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); + // eval the prompt + if (spec_capture) { + // begin() first (it clears the capture staging window), then an explicit + // batch (positions + seq ids) so the drafter's process() can stage a + // capture row for every prompt position; the drafter context is NOT + // decoded directly -- its cache is managed inside draft(). + common_speculative_begin(spec, seq_id, prompt_tgt); + + common_batch_clear(batch_tgt); + for (size_t i = 0; i < prompt_tgt.size(); ++i) { + common_batch_add(batch_tgt, prompt_tgt[i], (llama_pos) i, { seq_id }, /* logits = */ false); + } + llama_decode(ctx_tgt, batch_tgt); + if (!common_speculative_process(spec, batch_tgt)) { + LOG_ERR("draft-dspark: common_speculative_process (prefill) failed\n"); + return 1; + } + } else { + llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1)); + llama_decode(ctx_dft.get(), llama_batch_get_one(inp.data(), inp.size() - 1)); + + common_speculative_begin(spec, seq_id, prompt_tgt); + } + size_t n_draft = 0; llama_tokens draft; @@ -174,7 +260,7 @@ int main(int argc, char ** argv) { llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), seq_id), llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), seq_id)); - if (use_ckpt_dft) { + if (!spec_capture && use_ckpt_dft) { ckpt.update_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } @@ -194,13 +280,16 @@ int main(int argc, char ** argv) { // save a checkpoint of the target context before evaluating the draft // this allows us to restore the state if partial draft acceptance occurs - if (!draft.empty()) { + // (capture mode uses common_context_seq_rm rollback instead, like the + // dspark eval harness -- and must never touch ctx_dft's state, which is + // managed inside common_speculative_draft() itself) + if (!spec_capture && !draft.empty()) { if (use_ckpt_tgt) { ckpt.update_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } } - { + if (!spec_capture) { ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1); @@ -225,17 +314,27 @@ int main(int argc, char ** argv) { //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str()); llama_decode(ctx_tgt, batch_tgt); + + if (spec_capture) { + // stage the verify rows' capture features for the next draft round + if (!common_speculative_process(spec, batch_tgt)) { + LOG_ERR("draft-dspark: common_speculative_process (verify) failed\n"); + return 1; + } + } } // evaluate the same batch with the draft model - { + if (!spec_capture) { // TODO: extend to support MTP, Eagle, etc. See server code for reference + // (dspark must NOT decode the verify batch on ctx_dft -- its drafter + // cache is advanced inside common_speculative_draft() itself) llama_decode(ctx_dft.get(), batch_tgt); } // only save the sampler sampler state if we use checkpoints common_sampler_ptr smpl_save; - if (use_ckpt_tgt) { + if (!spec_capture && use_ckpt_tgt) { smpl_save.reset(common_sampler_clone(smpl.get())); } @@ -255,7 +354,7 @@ int main(int argc, char ** argv) { // check for partial draft acceptance: // if the context doesn't support partial sequence removal, restore the checkpoint // and make the accepted tokens the new partial draft for the next iteration - if (use_ckpt_tgt && ids.size() - 1 < draft.size()) { + if (!spec_capture && use_ckpt_tgt && ids.size() - 1 < draft.size()) { LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, draft.size()); draft = std::move(ids); @@ -284,6 +383,13 @@ int main(int argc, char ** argv) { // full acceptance: consume the draft and commit accepted tokens n_past += ids.size() - 1; + + if (spec_capture) { + // drop the rejected tail of this round's verify batch from the target + // cache (bounded partial rollback, also valid for hybrid GDN state); + // dspark's own drafter cache was already cropped inside draft(). + common_context_seq_rm(ctx_tgt, seq_id, n_past, -1); + } n_drafted += n_draft; // note: we ignore the discarded small drafts n_accept += ids.size() - 1; n_predict += ids.size(); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 07759f417084..81a763419e8a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -16,6 +16,7 @@ #include "mtmd-helper.h" #include "ggml-cpp.h" +#include "gguf.h" // TODO: tmp until the mtmd draft processing is refactored [TAG_MTMD_DRAFT_PROCESSING] #include "../../src/llama-ext.h" @@ -56,6 +57,68 @@ static uint32_t server_n_outputs_max(const common_params & params) { return std::max(1, std::min(n_batch, n_outputs)); } +// dspark drafters carry the target layer-id list to capture as an array-typed +// GGUF KV, which llama_model's string-KV cache skips (array KVs are not +// exposed via llama_model_meta_val_str) -- read it from the drafter file +// directly, same as tests/test-dspark-real-eval.cpp. Returns empty on failure. +static std::vector server_read_dspark_target_layers(const std::string & drafter_path) { + struct gguf_init_params gp = { /* .no_alloc = */ true, /* .ctx = */ nullptr }; + gguf_context * gctx = gguf_init_from_file(drafter_path.c_str(), gp); + if (gctx == nullptr) { + return {}; + } + + std::vector out; + + const int64_t arch_kid = gguf_find_key(gctx, "general.architecture"); + if (arch_kid >= 0) { + const std::string key = std::string(gguf_get_val_str(gctx, arch_kid)) + ".dspark.target_layers"; + const int64_t kid = gguf_find_key(gctx, key.c_str()); + if (kid >= 0 && gguf_get_kv_type(gctx, kid) == GGUF_TYPE_ARRAY) { + const enum gguf_type arr_type = gguf_get_arr_type(gctx, kid); + const size_t n = gguf_get_arr_n(gctx, kid); + const void * data = gguf_get_arr_data(gctx, kid); + out.reserve(n); + for (size_t i = 0; i < n; i++) { + switch (arr_type) { + case GGUF_TYPE_INT32: out.push_back(((const int32_t *) data)[i]); break; + case GGUF_TYPE_UINT32: out.push_back((int32_t) ((const uint32_t *) data)[i]); break; + case GGUF_TYPE_INT64: out.push_back((int32_t) ((const int64_t *) data)[i]); break; + case GGUF_TYPE_UINT64: out.push_back((int32_t) ((const uint64_t *) data)[i]); break; + default: out.clear(); i = n; break; + } + } + } + } + + gguf_free(gctx); + return out; +} + +// read the dspark drafter's block size (draft tokens per round) from its GGUF. +// Returns 0 on failure. +static uint32_t server_read_dspark_block_size(const std::string & drafter_path) { + struct gguf_init_params gp = { /* .no_alloc = */ true, /* .ctx = */ nullptr }; + gguf_context * gctx = gguf_init_from_file(drafter_path.c_str(), gp); + if (gctx == nullptr) { + return 0; + } + + uint32_t out = 0; + + const int64_t arch_kid = gguf_find_key(gctx, "general.architecture"); + if (arch_kid >= 0) { + const std::string key = std::string(gguf_get_val_str(gctx, arch_kid)) + ".dspark.block_size"; + const int64_t kid = gguf_find_key(gctx, key.c_str()); + if (kid >= 0 && gguf_get_kv_type(gctx, kid) == GGUF_TYPE_UINT32) { + out = gguf_get_val_u32(gctx, kid); + } + } + + gguf_free(gctx); + return out; +} + // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 enum slot_state { SLOT_STATE_IDLE, @@ -954,6 +1017,39 @@ struct server_context_impl { cparams.n_rs_seq = 0; cparams.ctx_other = ctx_tgt; + // dspark drafts a full block per round regardless of the configured + // draft n_max: its drafter batch requests [anchor + block_size] output + // rows per sequence, which can exceed the generic (1 + n_max) sizing. + 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_dspark) { + const uint32_t block_size = server_read_dspark_block_size(params_spec.mparams.path); + if (block_size > 0) { + const uint32_t n_out_dspark = params_base.n_parallel * (1 + block_size); + if (cparams.n_outputs_max < n_out_dspark) { + SRV_INF("draft-dspark: raising draft ctx n_outputs_max %u -> %u (block_size=%u)\n", + cparams.n_outputs_max, n_out_dspark, block_size); + cparams.n_outputs_max = n_out_dspark; + } + } + + // dspark stages all context rows since its cache position in ONE + // batch -- worst case the whole window right after begin() (e.g. + // a follow-up request with >n_batch tokens of history). If the + // drafter's batch is smaller than its context, every round is + // skipped ("round needs N tokens > n_batch") and speculation + // silently degrades to plain AR. + if (cparams.n_batch < cparams.n_ctx) { + SRV_INF("draft-dspark: raising draft ctx n_batch %u -> %u (full-context staging)\n", + cparams.n_batch, cparams.n_ctx); + cparams.n_batch = cparams.n_ctx; + } + if (cparams.n_ubatch < cparams.n_batch) { + cparams.n_ubatch = cparams.n_batch; + } + } + ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); params_base.speculative.draft.ctx_tgt = ctx_tgt; @@ -1068,6 +1164,23 @@ struct server_context_impl { ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft.get()); } + // dspark (draft-dspark) needs the target context to capture the drafter's + // tap layers on every decode -- without this the first draft round fails + // (capture rows come back null). masked=false keeps batch.logits narrow + // (no per-row full-vocab lm_head), see llama_set_capture_layers (#63). + if (spec && common_speculative_need_embd_capture(spec.get())) { + const std::string & drafter_path = params_base.speculative.draft.mparams.path; + + const std::vector capture_layers = server_read_dspark_target_layers(drafter_path); + if (capture_layers.empty()) { + SRV_ERR("draft-dspark: failed to read dspark.target_layers from '%s' -- disabling speculative decoding\n", drafter_path.c_str()); + spec.reset(); + } else { + llama_set_capture_layers(ctx_tgt, capture_layers.data(), capture_layers.size(), /* masked = */ false); + SRV_INF("draft-dspark: target tap capture engaged on %zu layers\n", capture_layers.size()); + } + } + if (spec) { SRV_INF("%s", "speculative decoding context initialized\n"); } else { @@ -2639,6 +2752,16 @@ struct server_context_impl { slot.state = SLOT_STATE_PROCESSING_PROMPT; + // capture-type drafters (dspark): begin() clears the per-seq staged + // feature window, so it must run BEFORE the prompt is decoded -- + // the prompt chunks' capture rows are staged by the + // common_speculative_process() call in the decode loop below. + // (for other spec types begin() stays after prompt eval, see + // SLOT_STATE_DONE_PROMPT.) + if (slot.can_speculate() && common_speculative_need_embd_capture(spec.get())) { + common_speculative_begin(spec.get(), slot.id, slot.task->tokens.get_text_tokens()); + } + 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()); @@ -2709,7 +2832,15 @@ struct server_context_impl { continue; } - if (slot.task->params.cache_prompt) { + // capture-type drafters (dspark) need a capture row staged for every + // prompt position; KV-cache prefix reuse would skip decoding (and thus + // capturing) the reused positions, so force a full reprocess. + const bool spec_needs_full_prompt = slot.can_speculate() && common_speculative_need_embd_capture(spec.get()); + if (spec_needs_full_prompt && slot.task->params.cache_prompt) { + SLT_DBG(slot, "%s", "draft-dspark: disabling prompt cache reuse (capture rows needed for every position)\n"); + } + + if (slot.task->params.cache_prompt && !spec_needs_full_prompt) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); @@ -3359,7 +3490,10 @@ struct server_context_impl { // prompt evaluated for next-token prediction slot.state = SLOT_STATE_GENERATING; - if (slot.can_speculate()) { + // capture-type drafters already ran begin() before prompt decode (see + // SLOT_STATE_STARTED) -- running it again here would wipe the prompt's + // staged capture rows. + if (slot.can_speculate() && !common_speculative_need_embd_capture(spec.get())) { common_speculative_begin(spec.get(), slot.id, slot.prompt.tokens.get_text_tokens()); } } else if (slot.state != SLOT_STATE_GENERATING) {