Add env-gated prompt-lookup speculative decoding for greedy generation#396
Add env-gated prompt-lookup speculative decoding for greedy generation#396rwl4 wants to merge 1 commit into
Conversation
Env-gated (DS4_PROMPT_LOOKUP_DRAFT=1, greedy/temp-0 only): draft the continuation from the session's own token history -- the tokens that followed the most recent earlier occurrence of the current 4-gram -- and verify [anchor | drafts] in one batched target pass using the existing speculative verify and frontier-rollback machinery (metal_graph_verify_suffix_tops + spec_frontier_*), committing only target-agreed tokens. The spec shadow-state and spec_logits allocations now key on (MTP || prompt-lookup); no new GPU kernels, no second model, no MTP behavior changes. Proposals truncate at EOS; no-match cycles fall back to the ordinary one-token decode. Default draft depth is 7 so the fused verify pass (8 positions) stays on the small-batch mat-vec kernels (2..8 tokens, metal/dense.metal); batch 9+ falls onto the full prefill matmul path and roughly doubles the pass cost (measured 80 -> 167 ms), which erases the win. DS4_PROMPT_LOOKUP_MAX (1..15) tunes depth; the ablation in speed-bench/prompt-lookup-validation.md justifies the default. Measured on an M5 Max (Flash, full-resident, details and falsification pass in speed-bench/prompt-lookup-results.md and speed-bench/prompt-lookup-validation.md): copy-heavy ~2.2x, code edits ~1.7-1.9x, a realistic agent turn 1.47x; prose, long-context recall, and guaranteed-no-match prompts ~1.0x with microsecond-scale scan overhead. Output is byte-identical to the same decode loop across all validation workloads, including EOS-through-drafts, partial accepts, ambiguous n-grams, and context-wall stress; --long-context passes with the feature enabled. Shipped MTP (--mtp-draft 2) measures ~1.06x on the same probes. Known bounded worst case: adversarial short repeated-ambiguity can cost ~20% on tiny outputs (documented; miss-backoff is future work).
c93cfb5 to
359638f
Compare
|
Tested on Apple M3 Ultra (60-core GPU, 256 GB, Metal), macOS 26.5.
A/B:
Clean split: extract is faster on both columns, classify is slower on both. Prompt-lookup commit rate ~68% (extract-dominated). Win where output copies the input (extract); regression on short novel output (classify), where it mostly misses and also gives up MTP's speedup. Both correct (extract reconciled, 20/20 labels) — but extract output wasn't byte-identical (5,382 vs 5,401 tokens), unlike the MTP path. Useful for copy-heavy / extraction servers. On a mixed workload it's a net loss vs MTP draft-2, since it can't run alongside MTP. |
|
ROCm / gfx1151 (Strix Halo) validation — works, +23% on agent workloads, byte-identical. Tested this PR cherry-picked onto main (80ebbc3) together with #446/#498/#381/#385, built with Method: ds4-server on a test port, identical greedy request (3.1k-token "rewrite every handler in this module" code-edit prompt, think:false, temperature 0, 640 tokens), A/B with/without
Stats line: So the backend-shared verify path works fine on ROCm. Two notes: (1) ds4-bench doesn't exercise this path (only cli/server are wired), which cost me an hour of confusion — maybe worth a line in the PR description; (2) I've been running it in production since, no issues. |
OPS-NeoRetro
left a comment
OPS-NeoRetro
left a comment
There was a problem hiding this comment.
@rwl4 Well done! You used an AI agent to write most of the code and to also create writeups, but I think this is one of your first PRs. You should implement some of the changes that I recommended and also discuss more about your findings when reimplementing this for DeepSeek V4 and eventually, GLM-5.x.
| if (spec_state) { | ||
| g->spec_logits = ds4_gpu_tensor_alloc((uint64_t)16 * DS4_N_VOCAB * sizeof(float)); | ||
| } |
| 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(); |
| #ifdef DS4_NO_GPU | ||
| (void)eos_token; | ||
| snprintf(err, errlen, "GPU support is not compiled in"); | ||
| return -1; | ||
| #else |
There was a problem hiding this comment.
Please add a CPU fallback or else this is going to be a GPU-only or Metal-only PR
| /* 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 | ||
|
|
There was a problem hiding this comment.
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
| } else if (temperature <= 0.0f && | ||
| getenv("DS4_PROMPT_LOOKUP_DRAFT") != NULL) { | ||
| /* Prompt-lookup drafting covers greedy decode spans, which includes | ||
| * forced-greedy tool-call payloads on otherwise sampled requests. | ||
| * Committed tokens flow through the same per-token stop/stream | ||
| * handling below, like the MTP batch path. */ | ||
| ntok = ds4_session_eval_prompt_lookup_argmax(s->session, | ||
| token, | ||
| max_tokens - completion, | ||
| ds4_token_eos(s->engine), | ||
| toks, | ||
| (int)(sizeof(toks) / sizeof(toks[0])), | ||
| err, | ||
| sizeof(err)); | ||
| if (ntok < 0) { | ||
| finish = "error"; | ||
| break; | ||
| } |
There was a problem hiding this comment.
Good greedy decoding implementation
| } 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; | ||
| } |
There was a problem hiding this comment.
Same as the logic from ds4-server.c
| 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); |
|
Measurements for this PR on M3 Ultra 256 GB (60-core bin), Flash q4-imatrix, greedy, ctx 32768, merged onto the #555 branch. (Disclosure: AI-assisted — Claude driving the box; branch passes
Greedy output byte-identical to baseline in all three (same SHA-256). The no-match path is effectively free — the risk is only miss-heavy matching workloads burning verify passes. Two follow-ups on our branch (https://github.com/robotnursenyc/ds4/tree/dsv4-speculative-stack), offered if useful:
With #468-style per-position verify checkpointing also applied (see comment there), partial accepts drop from ~180 ms to ~112 ms and the extract case reaches 56.6 t/s. |
|
@rwl4 have you added the CPU fallback path? If you don't add a CPU fallback, or a non-Metal fallback sometime later before @antirez merges this, then please rename this PR to "Metal: Add env-gated prompt-lookup speculative decoding for greedy generation" if you only support Metal or "Metal/ROCm/CUDA: Add env-gated prompt-lookup speculative decoding for greedy generation" if you support Metal, CUDA and ROCm backends. |
5 participants
Add env-gated prompt-lookup speculative decoding for greedy generation
This adds a prompt-lookup draft source for greedy decoding. When enabled
(
DS4_PROMPT_LOOKUP_DRAFT=1), the CPU searches the existing token history for arepeat of the current 4-gram, proposes the tokens that followed it as a draft
(up to 7), and verifies
[anchor | drafts]in one batched pass with the existingtarget-model speculative verification path (
metal_graph_verify_suffix_tops+frontier rollback). Only target-agreed tokens commit; when no match is found the
path falls back to ordinary one-token decode. No new GPU kernels, no second
model, no MTP behavior changes (the speculative state buffers now allocate when
either MTP or prompt-lookup is active).
The feature is off by default and greedy/temp-0 only. It is workload-
dependent by design: it pays off when the output revisits text that exists in
context (code edits, file reproduction, quoting, structured/templated output —
i.e. coding-agent workloads), and stays out of the way otherwise.
Validation (M5 Max, Flash, full-resident; details in speed-bench/prompt-lookup-{results,validation}.md)
microsecond-scale scan overhead (e.g. 0.11 ms total over a 96-token run at 9K ctx)
validation workloads; correctness stress green: EOS arriving mid-draft, partial accepts,
ambiguous repeated n-grams, context-wall and sub-wall runs, no-match prompts
--long-contextpasses with the feature enabledthe small-batch mat-vec kernel ceiling (2..8 tokens); batch 9+ falls onto the
prefill matmul path and roughly doubles the pass cost.
DS4_PROMPT_LOOKUP_MAX(1..15) tunes it.
Known bounded worst case (documented in the validation note): an adversarial
short output over highly ambiguous repeated n-grams can lose ~20% to misfiring
verify passes; output remains byte-identical to the same greedy decode loop.
A miss-backoff is possible future
work, as are a hybrid with the MTP drafter for novel text and server/agent
wiring.
For comparison on the same probes,
--mtp --mtp-draft 2measures ~1.06×, whichmatches the project's own description of the MTP path as a slight speedup —
the two draft sources are complementary rather than competing (MTP drafts novel
text where lookup finds no match).