Skip to content

Add env-gated prompt-lookup speculative decoding for greedy generation #396

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

Open
rwl4 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add env-gated prompt-lookup speculative decoding for greedy generation #396

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
268 changes: 262 additions & 6 deletions ds4.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10828,6 +10828,19 @@ static bool metal_graph_ensure_batch_ffn_out(ds4_gpu_graph *g) {
* Metal Release Graph Allocation.
* ========================================================================= */

/* Prompt-lookup speculative drafting (DS4_PROMPT_LOOKUP_DRAFT=1): greedy decode
* drafts the continuation from the session's own token history and verifies it
* with the target model in one batched pass, reusing the MTP speculative
* verify/rollback machinery with the context as a zero-cost drafter. */
static bool prompt_lookup_draft_enabled(void) {
static int cached = -1;
if (cached < 0) {
const char *e = getenv("DS4_PROMPT_LOOKUP_DRAFT");
cached = (e && e[0] && e[0] != '0') ? 1 : 0;
}
return cached == 1;
}

/* Allocate the Metal graph state for a chosen raw-cache capacity. The model
* weights are not copied here; tensors reference the mapped GGUF. */
static bool metal_graph_alloc_raw_cap(
Expand All @@ -10840,6 +10853,9 @@ static bool metal_graph_alloc_raw_cap(
bool enable_mtp) {
memset(g, 0, sizeof(*g));
g->mtp_enabled = enable_mtp;
/* The speculative verify/rollback buffers serve both the MTP drafter and
* the prompt-lookup drafter; allocate them when either is active. */
const bool spec_state = enable_mtp || prompt_lookup_draft_enabled();

@OPS-NeoRetro OPS-NeoRetro Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reasonable

if (raw_cap == 0) raw_cap = 1;
if (ctx_size == 0) ctx_size = raw_cap;
if (prefill_cap == 0) prefill_cap = 1;
Expand Down Expand Up @@ -10945,7 +10961,7 @@ static bool metal_graph_alloc_raw_cap(
(DS4_GPU_ATTN_COMP_CACHE_F16 ? sizeof(uint16_t) : sizeof(float)));
g->layer_attn_state_kv[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
g->layer_attn_state_score[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
if (enable_mtp) {
if (spec_state) {
g->spec_attn_state_kv[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
g->spec_attn_state_score[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
g->spec_prefix1_attn_state_kv[il] = ds4_gpu_tensor_alloc(attn_width * attn_rows * sizeof(float));
Expand All @@ -10968,7 +10984,7 @@ static bool metal_graph_alloc_raw_cap(
(uint64_t)g->layer_comp_cap[il] * DS4_N_INDEXER_HEAD_DIM * sizeof(float));
g->layer_index_state_kv[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
g->layer_index_state_score[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
if (enable_mtp) {
if (spec_state) {
g->spec_index_state_kv[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
g->spec_index_state_score[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
g->spec_prefix1_index_state_kv[il] = ds4_gpu_tensor_alloc(index_width * index_rows * sizeof(float));
Expand Down Expand Up @@ -11042,9 +11058,11 @@ static bool metal_graph_alloc_raw_cap(
g->mtp_raw_cache = metal_graph_alloc_kv_cache_tensor(
managed_kv_cache,
(uint64_t)raw_cap * DS4_N_HEAD_DIM * sizeof(float));
g->spec_logits = ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float));
g->mtp_n_raw = 0;
}
if (spec_state) {
g->spec_logits = ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float));
}
Comment on lines +11063 to +11065

@OPS-NeoRetro OPS-NeoRetro Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is reasonable


g->prefill_tokens = ds4_gpu_tensor_alloc(pc * sizeof(int32_t));
g->batch_cur_hc = ds4_gpu_tensor_alloc(pc * hc_dim * sizeof(float));
Expand Down Expand Up @@ -11099,7 +11117,7 @@ static bool metal_graph_alloc_raw_cap(
layer_cache_ok = g->layer_attn_comp_cache[il] != NULL &&
g->layer_attn_state_kv[il] != NULL &&
g->layer_attn_state_score[il] != NULL &&
(!enable_mtp ||
(!spec_state ||
(g->spec_attn_state_kv[il] != NULL &&
g->spec_attn_state_score[il] != NULL &&
g->spec_prefix1_attn_state_kv[il] != NULL &&
Expand All @@ -11109,7 +11127,7 @@ static bool metal_graph_alloc_raw_cap(
layer_cache_ok = g->layer_index_comp_cache[il] != NULL &&
g->layer_index_state_kv[il] != NULL &&
g->layer_index_state_score[il] != NULL &&
(!enable_mtp ||
(!spec_state ||
(g->spec_index_state_kv[il] != NULL &&
g->spec_index_state_score[il] != NULL &&
g->spec_prefix1_index_state_kv[il] != NULL &&
Expand Down Expand Up @@ -11140,7 +11158,8 @@ static bool metal_graph_alloc_raw_cap(
(g->mtp_embed && g->mtp_enorm && g->mtp_eproj &&
g->mtp_eproj_hc && g->mtp_hnorm_hc && g->mtp_hproj_hc &&
g->mtp_input_hc && g->mtp_state_hc && g->mtp_next_hc &&
g->mtp_raw_cache && g->spec_logits)) &&
g->mtp_raw_cache)) &&
(!spec_state || g->spec_logits) &&
g->prefill_tokens &&
g->batch_cur_hc && g->batch_next_hc && g->batch_flat_hc &&
g->batch_hc_mix && g->batch_hc_split &&
Expand Down Expand Up @@ -22224,6 +22243,13 @@ struct ds4_session {
int ctx_size;
bool checkpoint_valid;
bool mtp_draft_valid;
/* Prompt-lookup drafting stats (DS4_PROMPT_LOOKUP_DRAFT). */
uint64_t pl_passes;
uint64_t pl_drafted;
uint64_t pl_committed;
uint64_t pl_no_match;
uint64_t pl_free_miss;
double pl_scan_sec;
};

/* =========================================================================
Expand Down Expand Up @@ -25025,6 +25051,22 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) {

void ds4_session_free(ds4_session *s) {
if (!s) return;
if (s->pl_passes || s->pl_no_match || s->pl_free_miss) {
const uint64_t attempts = s->pl_passes + s->pl_no_match + s->pl_free_miss;
fprintf(stderr,
"ds4: prompt-lookup summary: attempts=%llu no_match=%llu first_miss=%llu "
"verify_passes=%llu drafted=%llu committed=%llu accept=%.1f%% "
"tokens/pass=%.2f scan=%.2f ms total\n",
(unsigned long long)attempts,
(unsigned long long)s->pl_no_match,
(unsigned long long)s->pl_free_miss,
(unsigned long long)s->pl_passes,
(unsigned long long)s->pl_drafted,
(unsigned long long)s->pl_committed,
s->pl_drafted ? 100.0 * (double)s->pl_committed / (double)s->pl_drafted : 0.0,
s->pl_passes ? (double)s->pl_committed / (double)s->pl_passes : 0.0,
s->pl_scan_sec * 1000.0);
}
ds4_dist_session_free(s->distributed);
if (ds4_session_is_cpu(s)) {
kv_cache_free(&s->cpu_cache);
Expand Down Expand Up @@ -26598,6 +26640,220 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
#endif
}

/* Prompt-lookup drafter: propose the tokens that followed the most recent
* earlier occurrence of the checkpoint's last DS4_PL_NGRAM tokens. Host-only,
* a backward memcmp scan of the token history (sub-millisecond even at very
* long contexts; the CPU is idle during decode). */
#define DS4_PL_NGRAM 4
/* Array bound; the effective depth defaults to 7 so the fused verify pass
* (anchor + drafts = 8 positions) stays on the small-batch mat-vec kernels,
* which cover 2..8 tokens (metal/dense.metal); batch 9+ falls onto the full
* prefill matmul path whose fixed cost eats the speculation win. Tunable via
* DS4_PROMPT_LOOKUP_MAX (1..15; spec_logits holds 16 rows). */
#define DS4_PL_MAX_DRAFT 15

Comment on lines +26643 to +26654

@OPS-NeoRetro OPS-NeoRetro Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can move those #defines into ds4.h or the first lines of the file so people who have problems or want to tune the greedy decoder can easily tune it because it's a compile-time constant

static int prompt_lookup_max_draft(void) {
static int cached = -1;
if (cached < 0) {
cached = 7;
const char *e = getenv("DS4_PROMPT_LOOKUP_MAX");
if (e && e[0]) {
char *end = NULL;
long v = strtol(e, &end, 10);
if (end != e && v >= 1 && v <= DS4_PL_MAX_DRAFT) cached = (int)v;
}
}
return cached;
}

/* The matched n-gram ends with `pending` (the sampled-but-not-yet-evaluated
* anchor token), so drafting happens before the anchor's forward pass and the
* anchor can join the verify batch. */
static int prompt_lookup_propose(const token_vec *ckpt, int pending,
int draft_cap, int *drafts) {
const int n = DS4_PL_NGRAM;
const int len = ckpt->len;
if (draft_cap <= 0 || len < n) return 0;
if (draft_cap > DS4_PL_MAX_DRAFT) draft_cap = DS4_PL_MAX_DRAFT;
const int *v = ckpt->v;
int suffix[DS4_PL_NGRAM];
memcpy(suffix, v + len - (n - 1), (size_t)(n - 1) * sizeof(suffix[0]));
suffix[n - 1] = pending;
for (int j = len - n; j >= 0; j--) {
if (memcmp(v + j, suffix, (size_t)n * sizeof(v[0])) != 0) continue;
const int follow = j + n;
int count = len - follow;
if (count > draft_cap) count = draft_cap;
if (count <= 0) continue;
memcpy(drafts, v + follow, (size_t)count * sizeof(v[0]));
return count;
}
return 0;
}

/* Greedy prompt-lookup speculative eval. Same contract and acceptance rule as
* ds4_session_eval_speculative_argmax, with the session's own token history as
* the drafter instead of the MTP head: commit one normal target token, draft
* the continuation from the most recent matching context span, verify the
* suffix with the target model in one batched pass, and commit only the
* agreeing prefix (rolling back speculative Metal state otherwise). The
* batched verifier is the same one the MTP path uses, with the same non-quality
* caveat about batched reductions on nearly-tied logits. */
int ds4_session_eval_prompt_lookup_argmax(ds4_session *s, int first_token,
int max_tokens, int eos_token,
int *accepted, int accepted_cap,
char *err, size_t errlen) {
if (!s || !accepted || max_tokens <= 0 || accepted_cap <= 0) return 0;
if (s->distributed || ds4_session_is_cpu(s)) {
if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1;
accepted[0] = first_token;
return 1;
}
#ifdef DS4_NO_GPU
(void)eos_token;
snprintf(err, errlen, "GPU support is not compiled in");
return -1;
#else
Comment on lines +26712 to +26716

@OPS-NeoRetro OPS-NeoRetro Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a CPU fallback or else this is going to be a GPU-only or Metal-only PR

ds4_engine *e = s->engine;

/*
* Fused one-pass cycle: unlike the MTP drafter, prompt lookup needs only the
* token history, so it can draft BEFORE the anchor token's forward pass and
* verify [anchor | drafts] together in one batched target pass -- one weight
* stream per cycle instead of the MTP path's sequential-eval-plus-verify two.
* On no match (or any uncertainty) the cycle is exactly the baseline eval.
*/
int draft_cap = prompt_lookup_max_draft();
if (draft_cap > max_tokens - 1) draft_cap = max_tokens - 1;
if (draft_cap > accepted_cap - 1) draft_cap = accepted_cap - 1;
const int room = s->ctx_size - s->checkpoint.len;
if (draft_cap > room - 2) draft_cap = room - 2;
if (draft_cap > (int)s->graph.prefill_cap - 1) draft_cap = (int)s->graph.prefill_cap - 1;

int drafts[DS4_PL_MAX_DRAFT];
int draft_n = 0;
if (draft_cap > 0 && s->graph.spec_logits && first_token != eos_token) {
const double scan_t0 = now_sec();
draft_n = prompt_lookup_propose(&s->checkpoint, first_token, draft_cap, drafts);
s->pl_scan_sec += now_sec() - scan_t0;
/* Never propose past an EOS: tokens after it must not enter the
* checkpoint even if the target would agree with them. */
for (int i = 0; i < draft_n; i++) {
if (drafts[i] == eos_token) {
draft_n = i + 1;
break;
}
}
}
if (draft_n <= 0) {
s->pl_no_match++;
if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1;
accepted[0] = first_token;
return 1;
}

const bool pl_log = getenv("DS4_PROMPT_LOOKUP_LOG") != NULL;
int row_tops[DS4_PL_MAX_DRAFT];
float *row_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(row_logits[0]));
ds4_spec_frontier frontier;
memset(&frontier, 0, sizeof(frontier));
const int start = s->checkpoint.len;
const int n_verify = 1 + draft_n;
const double verify_t0 = pl_log ? now_sec() : 0.0;
const bool have_frontier = spec_frontier_snapshot(&frontier, s);
bool ok = have_frontier;
bool verifier_may_have_mutated = false;
int n_accept = 0;
if (ok) {
token_vec_push(&s->checkpoint, first_token);
for (int i = 0; i < draft_n; i++) token_vec_push(&s->checkpoint, drafts[i]);
verifier_may_have_mutated = true;
ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights,
&s->checkpoint, (uint32_t)start,
(uint32_t)n_verify, false, row_tops, NULL);
}
if (ok) {
/* The anchor is the sampled target token and is always committed;
* row_tops[i] is the target's next token after position i, so it
* verifies drafts[i]. */
int commit_total = 1;
for (int i = 0; i < draft_n; i++) {
if (row_tops[i] != drafts[i]) break;
commit_total++;
}
if (commit_total == n_verify) {
ok = metal_graph_read_spec_logits_row(&s->graph,
(uint32_t)(n_verify - 1),
row_logits);
} else {
/* Partial accept: rebuild exact state for just the accepted prefix.
* A lone anchor replays through the normal decode path (exact
* one-token numerics, as the MTP partial path does); longer
* prefixes recommit through the batched pass. */
s->checkpoint.len = start;
ok = spec_frontier_restore(&frontier, s);
if (ok && commit_total == 1) {
s->pl_free_miss++;
spec_frontier_free(&frontier);
free(row_logits);
if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1;
accepted[0] = first_token;
if (pl_log) {
fprintf(stderr, "ds4: prompt-lookup anchor-only (drafted=%d) verify=%.3f ms\n",
draft_n, (now_sec() - verify_t0) * 1000.0);
}
return 1;
}
if (ok) {
token_vec_push(&s->checkpoint, first_token);
for (int i = 0; i + 1 < commit_total; i++) token_vec_push(&s->checkpoint, drafts[i]);
ok = metal_graph_verify_suffix_tops(&s->graph, &e->model, &e->weights,
&s->checkpoint, (uint32_t)start,
(uint32_t)commit_total, false,
row_tops, NULL);
}
if (ok) ok = metal_graph_read_spec_logits_row(&s->graph,
(uint32_t)(commit_total - 1),
row_logits);
}
if (ok) {
memcpy(s->logits, row_logits, (size_t)DS4_N_VOCAB * sizeof(s->logits[0]));
accepted[n_accept++] = first_token;
for (int i = 0; i + 1 < commit_total && n_accept < accepted_cap; i++) {
accepted[n_accept++] = drafts[i];
if (drafts[i] == eos_token) break;
}
s->checkpoint_valid = true;
s->mtp_draft_valid = false;
s->pl_passes++;
s->pl_drafted += (uint64_t)draft_n;
s->pl_committed += (uint64_t)(commit_total - 1);
if (pl_log) {
fprintf(stderr,
"ds4: prompt-lookup fused drafted=%d committed=%d verify=%.3f ms\n",
draft_n, commit_total - 1, (now_sec() - verify_t0) * 1000.0);
}
spec_frontier_free(&frontier);
free(row_logits);
return n_accept;
}
}
s->checkpoint.len = start;
if (have_frontier) (void)spec_frontier_restore(&frontier, s);
spec_frontier_free(&frontier);
free(row_logits);
if (!verifier_may_have_mutated) {
/* Snapshot failed before any state change: baseline cycle. */
if (ds4_session_eval(s, first_token, err, errlen) != 0) return -1;
accepted[0] = first_token;
return 1;
}
snprintf(err, errlen, "prompt-lookup verifier failed");
s->checkpoint_valid = false;
return -1;
#endif
}

void ds4_session_invalidate(ds4_session *s) {
s->checkpoint_valid = false;
s->checkpoint.len = 0;
Expand Down
7 changes: 7 additions & 0 deletions ds4.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,13 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
int max_tokens, int eos_token,
int *accepted, int accepted_cap,
char *err, size_t errlen);
/* Greedy prompt-lookup speculative eval (DS4_PROMPT_LOOKUP_DRAFT=1): same
* contract as ds4_session_eval_speculative_argmax, drafting the continuation
* from the session's own token history instead of an MTP head. */
int ds4_session_eval_prompt_lookup_argmax(ds4_session *s, int first_token,
int max_tokens, int eos_token,
int *accepted, int accepted_cap,
char *err, size_t errlen);
void ds4_session_invalidate(ds4_session *s);
void ds4_session_rewind(ds4_session *s, int pos);
int ds4_session_pos(ds4_session *s);
Expand Down
20 changes: 19 additions & 1 deletion ds4_cli.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -497,6 +497,23 @@ static int run_sampled_generation(ds4_engine *engine, const cli_config *cfg, con
ds4_session_free(session);
return 1;
}
} else if (cfg->gen.temperature <= 0.0f &&
getenv("DS4_PROMPT_LOOKUP_DRAFT") != NULL) {
cli_dist_busy_set(cfg, true);
ntok = ds4_session_eval_prompt_lookup_argmax(session,
token,
max_tokens - generated,
ds4_token_eos(engine),
toks,
(int)(sizeof(toks) / sizeof(toks[0])),
err,
sizeof(err));
cli_dist_busy_set(cfg, false);
if (ntok < 0) {
fprintf(stderr, "ds4: decode failed: %s\n", err);
ds4_session_free(session);
return 1;
}
Comment on lines +500 to +516

@OPS-NeoRetro OPS-NeoRetro Jul 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as the logic from ds4-server.c

} else {
cli_dist_busy_set(cfg, true);
int eval_rc = ds4_session_eval(session, token, err, sizeof(err));
Expand Down Expand Up @@ -921,7 +938,8 @@ static int run_generation(ds4_engine *engine, const cli_config *cfg) {
}
} else if (cfg->engine.distributed.role == DS4_DISTRIBUTED_COORDINATOR ||
cfg->gen.temperature > 0.0f ||
ds4_engine_mtp_draft_tokens(engine) > 1) {
ds4_engine_mtp_draft_tokens(engine) > 1 ||
getenv("DS4_PROMPT_LOOKUP_DRAFT") != NULL) {
rc = run_sampled_generation(engine, cfg, &prompt);
} else {
token_printer printer = {
Expand Down
Loading