Skip to content

superseded #66

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
khosravipasha wants to merge 1 commit into from
Closed

superseded #66

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 118 additions & 12 deletions examples/speculative-simple/speculative-simple.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@
#include "speculative.h"
#include "log.h"
#include "llama.h"
#include "gguf.h"
#include "../../src/llama-ext.h"

#include <algorithm>
#include <clocale>
#include <cstdio>
#include <cstring>
Expand All @@ -13,6 +16,43 @@
#include <vector>
#include <utility>

// 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<int32_t> 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<int32_t> 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");

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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();

Expand All @@ -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<int32_t> 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;
Expand All @@ -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);
}

Expand All @@ -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);
Expand All @@ -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()));
}

Expand All @@ -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);
Expand Down Expand Up @@ -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();
Expand Down
Loading
Loading