Add backend-agnostic speculative decoding to llm-runner (Part 1) - #284
Open
carinapeng wants to merge 3 commits into
Open
carinapeng wants to merge 3 commits into
carinapeng wants to merge 3 commits into
Conversation
Add advance, forwardWithPerPositionLogits, and forwardPerPositionArgmax: prefill-without-logits plus single-pass per-position logits/argmax. These are the propose/verify primitives a speculative decoder drives the target and draft engines through; greedy uses the argmax fast path to avoid a per-position vocab copy.
Introduce a pluggable speculative-decoding core: - SpeculativeEngine: the propose/verify/accept/rollback surface a decoder needs from the draft and target engines, so they may run on different backends and the loop is testable against an in-memory double. CoreAISequentialEngine conforms. - DrafterEngine: the seam behind which drafting mechanisms live, with an AutoregressiveDrafter (separate draft model) and a model-free PromptLookupDrafter (n-gram). - SpeculativeDecoder: proposes k tokens and verifies them in one target pass; greedy acceptance is token-for-token identical to the target, sampling uses the Chen et al. (2023) accept/reject scheme. - DraftProposal / DrafterFeatures / VerificationPolicy value types and SpeculativeMath selection primitives.
Add --draft-model, --num-draft-tokens, and --prompt-lookup. Draft and target are built via EngineFactory and driven by SpeculativeDecoder; the engine must conform to SpeculativeEngine (else a clear error). Report acceptance rate, tokens per target pass, and generation-only throughput.
carinapeng
force-pushed
the
carina/spec-decode-generic
branch
from
September 22, 2026 16:29
0edbe27 to
211cd0f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a pluggable, backend-agnostic speculative-decoding path to
llm-runner, plus the framework it runs on. Draft and target engines sit behind a small protocol so different drafting mechanisms — a separate draft model, a model-free n-gram prompt-lookup drafter, or a future attached head / tree drafter — share one propose → verify → accept/rollback loop.What's included
SpeculativeEngine— the propose/verify/accept/rollback surface a decoder needs from the draft and target engines. Draft and target may run on different backends, and the loop is testable against an in-memory double with no model assets.CoreAISequentialEngineconforms via three added primitives (advance,forwardWithPerPositionLogits,forwardPerPositionArgmax).DrafterEngine— the seam behind which drafting mechanisms live:AutoregressiveDrafter(a separate draft model on anySpeculativeEngine) and a model-freePromptLookupDrafter(n-gram from the running context).SpeculativeDecoder— proposesktokens and verifies them in a single target pass. Greedy acceptance is token-for-token identical to the target; sampling uses the Chen et al. (2023) accept/reject scheme.llm-runnerflags —--draft-model,--num-draft-tokens,--prompt-lookup. Draft and target are built viaEngineFactory; the runner reports acceptance rate, tokens per target pass, and generation-only throughput.Design
This PR intentionally general:
DrafterFeaturesandDraftProposal.parentsalready model attached-head (feature-injecting) and tree drafting, so a fused-target or attached-head drafter plugs into the same decoder without changing the verify/accept path.Scope / follow-ups
SpeculativeEngineby implementing the same primitives.