Repository navigation
Static-shape engine implementation for speculative decoding protocol, llm-runner flags - #386
Open
carinapeng wants to merge 3 commits into
Open
carinapeng wants to merge 3 commits into
carinapeng wants to merge 3 commits into
Conversation
StaticShapeEngine conforms to SpeculativeEngine, so the decoder can verify a draft block, and llm-runner exposes the decoder. - Verify tiles: a block is scored in the engine's fixed-width decode tiles. verifyTileWidth, oneForwardDraftBudget and maxVerifyBlockLength let the decoder shorten a proposal so one verify is one forward. - Per-position argmax and logits readback. LogitRowScan scans a row with explicit rules (first maximum wins, NaN counts as negative infinity, the final-logit soft cap is applied first). - Sliding-window rings rewind to the window stride: allowsRewind and the rewind alignment keep every key the restarted forward needs resident. reset(to:) throws for an index outside the committed range. - llm-runner: --prompt-lookup, --prompt-lookup-min-ngram, --draft-model, --num-draft-tokens, --adaptive-draft-count, --no-align-proposals and --seed (a seeded generator for speculative sampling). The options are validated before a model loads.
… prompt - processSpeculativeBlock slices batches from the engine's token history instead of a second buffer, advances processedTokenCount as the only cursor, and hoists the cap, the vocabulary width and the scratch row out of the batch loop. The capture mode is switched on once per batch, not per row. Logits are read in place only when the array is Float16 and contiguous; otherwise the read reports it, as the plain path does. - oneForwardDraftBudget uses verifyTileWidth. decodeBlockSize, which depended on the smallest decode graph, is removed. - The ring's maxVerifyBlockLength states where its bound comes from. - SpeculativeOptions is SpeculativeDecodingOptions. The sampling and image checks apply only when speculating. --video, --continuation and --batch-prototype are refused with speculative decoding, and the options it accepts but does not apply are named in a warning. - Plain and speculative image runs build the prompt with the same code, so --image-info applies to both and the USER/ASSISTANT fallback exists in both.
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
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.
Builds on the speculative decoding protocol PR (#378).
Type of Change
Motivation and Context
Runs the speculative decoder on the static-shape engine and exposes it in
llm-runner.StaticShapeEngine: SpeculativeEngine. A verify block is scored in the engine's fixed-width decode tiles, soverifyTileWidth,oneForwardDraftBudgetandmaxVerifyBlockLengthtell the decoder how much fits; it shortens a proposal so one verify is one forward instead of two.LogitRowScanselects from a row with explicit rules: the first maximum wins, NaN counts as negative infinity, and the final-logit soft cap is applied first (the cap's rounding can create ties, which resolve to the lowest index). It is used by the engine's verify scan only; plain greedy sampling is unchanged.window + stridekeys, and a restarted forward begins on a batch boundary, soallowsRewind(processed:to:alignment:)checks that the first query of the restart still finds every key it needs.reset(to:)throws for an index outside the committed range, and leaves the engine unchanged when the ring refuses.llm-runner:--prompt-lookup,--prompt-lookup-min-ngram,--draft-model,--num-draft-tokens,--adaptive-draft-count,--no-align-proposalsand--seed. Options are validated before a model loads. Speculation needs greedy decoding or--sampling-strategy temperature, and no--json-schema. Only the static-shape engine supports it; other engines report that.Known difference from plain decoding: plain greedy sampling uses
vDSP_maxvi, which returns index 0 when a row holds a NaN, while the verify scan treats a NaN as negative infinity. On a row with a NaN, speculative and plain greedy output can differ. On every non-NaN row tried, ties resolve to the lowest index in both.Testing
Tested locally
Added/updated tests
Added/updated docs
Unit tests for
LogitRowScan(ties, NaN, soft cap, top-k against a reference, and the engine's scan against plain decoding), the ring rewind rule (including a sweep of positions and targets), verify-block limits,reset(to:)validation, the runner option checks and the stop-token helpers. Full Swift suite on macOS passes except the Top-K latency test, an existing timing test that fails under the parallel run and passes alone.Gemma 4 E2B (static-shape engine, debug build, Mac): plain and
--prompt-lookup --num-draft-tokens 7give identical text on a repetition prompt; 95.1% of drafts accepted.