Repository navigation
Conversation
| Follow the repository's spec-decoding standard instead: compare fresh, | ||
| independently started spec-off and spec-on runs with the same inputs and | ||
| settings via `tools/check_parity.py --method suffix --top-k 5`, and accept | ||
| `EXACT`, or a first divergence with mutual `TOP_K_MATCH` membership, with |
There was a problem hiding this comment.
Please do not weaken the spec-decoding contract to TOP_K_MATCH here. Keep suffix decoding unmerged until fresh spec-off/spec-on greedy runs produce identical token IDs. Any repository-wide contract change needs a separate proposal and evidence.
There was a problem hiding this comment.
Thanks for calling out the distinction between validation and numerical behavior. I'd like to propose keeping this PR as a minimal integration of vLLM's default suffix behavior, with strict reproducibility handled as a separate, opt-in capability.
The standard greedy verifier and request-lifecycle checks remain unchanged. This does not introduce approximate draft acceptance. vLLM's speculative-decoding documentation explicitly describes floating-point/batch-size variations, and batch invariance is a separate mode with a performance trade-off. Our current local strict implementation also has a measured cost: in a same-base Qwen3-4B-4bit, 4K-input, concurrency-8 service comparison on M5 Max, its target-only throughput was about 5.6% below native main (three paired repetitions). That is a result for this implementation/workload, not a claim that strict reproducibility must always be slower.
The strict kernels and prompt-NAX optimization are kept out of this PR. check_parity.py --method suffix still fails on any complete-token difference; the fresh diagnostic result is 3/4 exact sequences at B=1 and 4/4 at B=2 (strict CLI exits 1 and 0); preemption and cold/warm comparisons also retain non-exact results. I am not treating TOP_K_MATCH as proof of equivalence. The previously reported 14% first-pass concurrency-8 regression remains disclosed with its original benchmark revision and is deferred to a focused follow-up.
Focused checks pass 311 tests (10 skipped). Eight fresh boundary cases pass their functional assertions, and all 1,077 observed committed tokens match the corresponding serving argmax. Matched sleep/no-sleep outputs are exact; this is distinct from cold/warm-tree equality. The PR now explicitly limits support to MLX, text-only full-attention targets, plain greedy sampling and synchronous scheduling. Hybrid GDN and broader suffix features remain out of scope. Marking this ready is a request to review this proposed default-numerics scope, not a statement that strict equality is fixed or that this policy change has already been accepted. I am leaving this discussion unresolved for maintainer agreement; a future strict mode would be proposed independently and disabled by default.
There was a problem hiding this comment.
I merged main (086e57c) into the branch (11871be), re-ran the checks and updated the description. The strict diagnostic is unchanged and still exits nonzero on any complete-token difference.
Strict identity is still not met at B=1: check_parity.py --method suffix gives 33/40 identical sequences (same in two runs). At each first difference, the suffix token is among the target's own top candidates (3 exact ties, 4 near-ties with a top-1/top-2 gap ≤ 0.125). I have not isolated the cause with a controlled same-prefix experiment. B=2, run in a separate process with a cold suffix tree, gives 40/40 in two runs. All 884 recorded committed-token decisions match the serving-logit argmax, but that does not replace complete-sequence equality. These results reproduce identically on vLLM 0.32.0+cpu (unit, parity and boundary checks); serving numbers are from 0.31.0.
Since your condition (identical spec-off/spec-on token IDs) is not met, could you say which you prefer: accept this PR as a minimal integration under default numerics, with the remaining differences disclosed, or keep it blocked until I submit a separate correctness-contract proposal (or an opt-in strict mode)?
| def _stop(self, req_id: str) -> None: | ||
| anchor = self._last_len.pop(req_id, None) | ||
| state = self._states.pop(req_id, None) | ||
| cache = getattr(self._vllm, "suffix_cache", None) |
There was a problem hiding this comment.
Use self._vllm.suffix_cache directly. The pinned upstream proposer guarantees this field, and _stop needs it for cleanup. This fallback only supports test fakes and can silently skip cleanup.
There was a problem hiding this comment.
Fixed and retained in the rebased branch (current HEAD 0bde3a0). _stop accesses self._vllm.suffix_cache directly; the fake cache now implements the real contract, and the regression asserts stopped request IDs and an empty active set. The focused suite passes 311 tests (10 skipped), and fresh lifecycle checks also exercise automatic cleanup and tail ingestion.
Wrap vLLM's SuffixDecodingProposer (Snowflake suffix decoding, arXiv 2411.04975) as a Metal proposer: model-free drafting from suffix trees over prompts and past responses, learning across requests — complementary to n-gram, whose matching never leaves one request's history. Like n-gram there is no draft model and no draft KV cache; the tree state is plain process memory in the arctic_inference package (pure C++/nanobind core, no CUDA; sdist-only on PyPI — the macOS build recipe is documented). - suffix_proposer.py (new): translate each ProposeContext into the InputBatch-shaped stub upstream's engine loop runs off. Two Metal-specific decisions: (a) every live decode request keeps a stub row, not just drafting ones — upstream archives an active tree the first step its id is missing from the batch, so a drafting request that cannot draft this step would lose its live tree (rows that never drafted never start a tree); (b) trees are fed committed-history deltas since the last feeding (anchored at prompt_len on first sight, clamped on preemption rollback) rather than raw sampled ids. - model_runner.py: dispatch method="suffix" alongside ngram; the unsupported-method error lists it. - Draft lengths are dynamic per request (bounded by num_speculative_tokens); verify is the established greedy span path shared with n-gram. Targets are full-attention paged-attention models; the hybrid GDN gate stays ngram-only until suffix is validated there (follow-up). Validation: 8 marshalling tests against a recording fake (no arctic install needed, CI-safe) plus a real-arctic integration test (skipif-gated; locally verified); e2e on Qwen3-4B: 63 verify rounds, 2.21 avg committed tokens/round, 0.66 draft acceptance; greedy output matched the no-spec reference on that run (not guaranteed in general: verify-shape numerics can flip near-tie logits, as seen with ngram on Qwen3-8B). CI-equivalent suite green except three pre-existing environment failures (test_memory_info, test_mps_attention, one order-dependent GQA routing test) reproduced on pristine main. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Pass every request scheduled this step (decode and prefill rows) to upstream's SuffixDecodingProposer, as the GPU runner does with its whole InputBatch. Upstream archives any active tree missing from the batch, so a request recomputing its prefill after preemption lost its live tree, and a fresh request got no draft on the step its prefill finished. The final prefill chunk now drafts. When upstream holds no active tree for an id (first sighting, or archived while absent) the wrapper feeds the whole response since the prompt boundary, so a resumed request's tree is rebuilt with its earlier output rather than only the latest token. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Upstream's propose asserts the configured num_speculative_tokens, so keep passing that, then trim each draft to ctx.num_speculative_tokens (and the per-request cap). With a batch-size-dependent width the wrapper used to return 4 drafts for a step that scheduled 2. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Stop the request in arctic's cache on finished ids and on release (finished, preempted, resumed), guarded by membership in its active set. Upstream only archives an absent request on its next propose call, so a finished request's tree stayed active, and out of the global cache, for as long as the server sat idle. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
…nobs - The arctic-missing error in our wrapper was unreachable: vLLM's SpeculativeConfig validation raises ImportError (pinning arctic-inference==0.1.1) before any platform hook runs, so an earlier check with our recipe isn't possible. Drop the dead branch; the docs now say which error users see, that Metal was tested on 0.3.0, and that the package is Apache-2.0. - Sleep mode: allow suffix alongside ngram. Neither keeps KV of its own; suffix trees are host memory keyed by token ids and survive sleep. - Correct the rationale for keeping all scheduled rows: rows that never drafted never start a tree, so the rows protect drafting requests that cannot draft this step from being archived. - Log all four suffix_decoding_* knobs at startup. - Docs: greedy outputs can diverge from a no-spec run (verify-shape numerics, also seen with ngram on Qwen3-8B); compare by first divergent position, not an exact-match claim. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Cap each suffix draft at max_tokens - generated_tokens - 1, as the block draft proposer does (vllm-project#1078): verification also emits a correction or bonus token, so longer drafts are verified only to be discarded. A non-positive limit now skips the request instead of slicing. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
…trip Fake-based (CI, no arctic): prefill rows in the stub with the final chunk drafting, configured width to upstream with truncation to the step's width, remaining-output cap, and an absent request re-fed from its prompt boundary. The fake now asserts upstream's width contract. Real arctic (skipped when not installed): a preempted-then-resumed request keeps its response in the tree, finish archives immediately, the final prefill chunk drafts, and a mixed drafting/non-drafting batch with variable-length drafts round-trips through build_decode_segments and verify_greedy before the next delta is fed. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Upstream reads token_ids_cpu for a row only to seed a new tree (the prompt) and to take the match pattern (the last max_tree_depth tokens). Copy just those for drafting rows, and nothing for rows without a delta, instead of zeroing the buffer and copying every row's full history each step. Synthetic 8 rows x ~6k-token histories: 0.74-0.80 -> 0.05 ms CPU per propose; drafts over a randomized 900-step lifecycle trace are byte-identical to the full-copy version. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
History ingestion followed drafting eligibility, so a request's last committed tokens never reached its tree: a row that could not draft (its output budget was exhausted) was handed an empty delta -- which upstream skips before add_active_response -- and a request that finished after its last drafting step archived without the tail. A response [20, 21, 22] capped at max_tokens=3 taught the cache only [20, 21], so later requests could not reuse the ending. Decouple the two: a non-drafting row with a live tree still feeds its delta through the stub batch, and the wrapper holds each anchored request's state so release/finish can feed whatever committed after the last step before stop_request archives the tree. The held reference is the finished request's own state -- request_states cannot be used there because a finished id may already be reused by a new request. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
The suffix docs claimed equivalence is checked by reporting the first divergent token position, but locating a divergence neither shows it is a near-tie nor proves the runs equivalent. State the repository's actual standard instead: independently started spec-off/spec-on runs accepted as EXACT or first-divergence mutual TOP_K_MATCH with zero FAIL, and wire tools/check_parity.py --method suffix so that comparison can actually be run (suffix takes no draft checkpoint; the schedule key applies unchanged). Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Keep the ingestion anchor unchanged when upstream skips a context-capped row, and flush pending tokens before a missing request is archived. Cover both explicit release paths and a batch gap after drafting is disabled with real suffix trees. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Convert nonempty checkpoint arguments to Path before inspecting local directories, while preserving raw repository IDs for Hub resolution. Exercise suffix without a draft checkpoint through the existing CLI worker and report tests for both local and Hub models. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
The docs stated sleep mode is supported because the trees are host memory independent of the KV cache, but no sleep-to-wake run has actually exercised the suffix proposer across a sleep cycle. State the expectation and the pending validation instead of claiming support; restore the claim once the test exists. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
sleep(level=1) -> wake_up() was exercised end to end under suffix spec decode: the trees survived the cycle (host memory, KV-independent), post-wake generation completed, the consumer's drafted/accepted stats matched the no-sleep control (96/95 vs 95/95), and greedy outputs were identical. State the support with that scope; level 2 is not available on Metal. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Document the worker-level sleep and wake behavior without turning one experiment into a general guarantee about draft counts or token equality. Matched level-1 and no-sleep checks preserve cached responses and produce identical 128-token continuations with 95 drafts verified and accepted in both arms. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Install the optional arctic 0.3.0 backend in CI and fail explicitly if its native cache cannot load. Document the matching interpreter-pinned macOS source build. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Fail during argument validation before resolving models or starting workers. Cover both the orchestrator and direct suffix-worker entry points with model-free CLI regressions. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
6581226 to
0bde3a0
Compare
Enable real suffix-backend tests by default without a second pytest process. Reuse a compatible compiled Arctic wheel across CI runs, while retaining hard backend checks and an explicit local opt-out. Signed-off-by: suntp <20786998+suntp@users.noreply.github.com>
WindChimeRan
added
ci/build
ChatGPT Codex Connector
Add suffix speculative decoding to reuse response patterns across requests, without a draft model or additional draft KV cache. On the repeated workload, serving throughput is 2.18× at concurrency 1 and 1.43× at concurrency 8; first-pass concurrency 8 is 0.88×. Performance depends on reuse and acceptance.
Implementation and scope
Reuse vLLM's suffix proposer and Arctic's CPU suffix trees on the existing Metal greedy verifier. Committed history is ingested independently of drafting; request release, preemption and batch removal flush pending history. Request-state ownership isolates reused IDs, and drafts respect context and output limits.
The initial scope is MLX, text-only full-attention targets, plain greedy sampling and synchronous scheduling. Hybrid GDN is unsupported by suffix and remains n-gram-only. The optional dependency is
arctic-inference==0.3.0, with macOS installation instructions. CI enables the real-tree tests by default, caches the compiled Arctic wheel, and runs them in the main pytest process. No attention or quantized-model kernel changes are included.Numerical behavior and validation
The greedy verifier is exact for fixed target logits. As in vLLM's speculative decoding, floating-point and batch-shape effects can change tokens between ordinary and speculative execution, so this PR does not promise identical complete output sequences and adds no batch-invariant mode.
check_parity.py --method suffixis the strict diagnostic: it records complete-token equality and exits nonzero on any difference. On Qwen3-0.6B-4bit it gives 33/40 prompts with identical tokens at B=1, with the same result in two runs, so strict identical-token output is not achieved. In each of the 7 B=1 differences the suffix token is within the target's own top candidates at the first differing position: 3 are exact ties and 4 are near-ties (top-1/top-2 logprob gap at most 0.125; in one case a three-way near-tie within 0.0625). This is consistent with numerical changes from verifying several positions in one forward pass (different execution shape in bf16), but I have not isolated that cause with a controlled same-prefix experiment, and a small logprob gap is not by itself an acceptable-error bound. At B=2, run in a fresh process with a cold suffix tree, it gives 40/40 in two runs (160 drafted / 43 accepted tokens, similar drafting to B=1's 162 / 44). Because B=2 runs in a different execution shape from B=1, this does not show that larger batches are more consistent; it only shows no differences on this prompt set at B=2.Eight boundary cases cover EOS/stops, request-ID reuse, preemption/resume, output/context limits and worker-level sleep/wake. All 884 recorded committed-token decisions match their serving-logit argmax. History-ingestion and cleanup assertions pass; matched sleep/no-sleep outputs are identical. Preemption/no-preemption and cold/warm-tree comparisons contain numerical divergences, reported separately; these checks do not establish complete-token equivalence across execution schedules. The boundary cases and the focused tests (317 passed / 10 skipped; the skips need Ray) were run on
11871be, which merges main086e57cinto the previous headae8804c. These were first measured with vLLM 0.31.0+cpu and MLX 0.32.4, then repeated with vLLM 0.32.0+cpu (the version this branch targets) and MLX 0.32.4: the focused tests, both parity modes and all eight boundary cases gave the same results, with byte-identical parity outputs and the same 7 differing B=1 prompts. The serving numbers below were measured with vLLM 0.31.0+cpu only.Serving performance
Qwen3-8B-4bit, M5 Max 48 GiB, MLX 0.32.4, vLLM 0.31.0+cpu, three draft tokens, prefix caching off. Real
vllm serve+vllm bench serve, three paired repetitions with alternating arm order. Each measured pass contains 16 requests × 128 output tokens: ten natural and six repetitive prompts, with 20–91 input tokens. Both arms receive disjoint, input-length-matched warmup requests.Measured on
11871be(main086e57cmerged intoae8804c), both arms.Values are medians; speedup is the median paired throughput ratio. Concurrency is the client/server request limit, not a fixed scheduler batch size. The first pass can already learn across requests; the repeated pass reuses the same workload. All measured requests complete, with zero swap-out pages during the measurement windows.
Repeated-pass acceptance is 81.1% / 75.2% at concurrency 1 / 8. These gains apply to this short, intentionally reused workload. First-pass concurrency 8 loses 12% throughput. Median TTFT at concurrency 1 rises from about 40–41 ms to 44–47 ms, and at concurrency 8 it falls from about 148–154 ms to 120–131 ms.
Cache memory and limitations
max_cached_requestslimits response count, not bytes. In three-process cache-only measurements at depth 24, incremental RSS spans 0.70–9.00 MiB for 128 retained responses and 11.80–127.66 MiB for 2,048, across shared/diverse 64/256-token responses. At cap 128, growth from requests 512 to 2,048 stays within 0–0.5 MiB. These figures include allocator/Python effects and do not represent full-server overhead.Low-reuse workloads can lose throughput; the concurrency-8 first-pass regression is left as a follow-up. Hybrid GDN, broader suffix feature coverage and an opt-in strict reproducibility mode are outside this integration.