feat: speculative decoding with paired drafter checkpoints - #5241
Conversation
Gemma 4 ships a small `*-it-assistant` drafter for every variant: it proposes several tokens per round that the target verifies in one pass, so decoding gets faster with identical output. Unlike DeepSeek-V3 style MTP the drafter is a separate checkpoint, so it has to be declared, downloaded and handed to the engine. Specs gain `draft_model_id` (optionally templated by `draft_quantizations`, which the MLX conversions publish up to eight of). `LLMCacheManager` grows a `use_draft_model` mode that reuses the existing per-hub download paths, so a drafter follows the same ModelScope/HuggingFace selection as its target and lands in a sibling `<cache_dir>-draft-<quantization>` directory. `create_llm_model_instance` resolves it when `enable_mtp` is passed and hands the local path to the engine; engines that cannot use one are rejected up front through `LLM.support_draft_model` rather than failing later. Per engine: - MLX loads the drafter on the same dedicated MLX thread as the target, lets mlx-vlm auto-detect the drafter kind from its `model_type`, and validates the pair at load time instead of on the first request. Measured 2.2x on decode for gemma-4 E4B, 8bit target with a bf16 drafter. - vLLM translates it into `speculative_config` with `method: mtp`, and rejects vllm<0.22.0, which treats the assistant as a generic draft model and fails to initialize against a multimodal target. - SGLang translates it into `--speculative-algorithm NEXTN` with matching step and draft-token counts. An explicitly provided `speculative_config` / `speculative_algorithm` always wins; the engine-neutral options are consumed either way so they never reach the engine. The Transformers engine is not supported: it runs its own continuous-batching loop rather than `generate()`, so there is nowhere to attach a drafter. Deleting a model's cache now removes its drafters too, instead of orphaning a directory per drafter quantization next to it. Also registers gemma-4 12B, the encoder-free Unified variant (`Gemma4UnifiedForConditionalGeneration`), across pytorch/gguf/mlx. It is the only variant whose drafter is published in more than one conversion, so it is what makes the drafter selection meaningful.
Advanced Configuration gains a Speculative Decoding section: an Enable MTP switch, a drafter quantization select when the spec publishes more than one conversion, and the tokens-per-round input. The whole section only appears for a format/size that actually declares a drafter, so it stays out of the way for the models that have none. The drafter options carry the same cached marker the model quantizations use, because turning the switch on downloads another checkpoint. Tooltips cover the parts that are not guessable: that the switch costs a download, why an unquantized drafter is recommended even for a quantized target, and that a value above the drafter's trained depth acts as a ceiling rather than a fixed block size.
Without NEXT_PUBLIC_API_URL the browser sends API requests to the Next dev server, whose rewrite proxy buffers the response body. Streaming endpoints then deliver the whole answer at once, which reads as the backend having lost its streaming support — the same UI served by the backend streams fine, because the static export is same-origin and has no proxy in front of it. Point the browser straight at the backend in dev through `.env.development`, which Next loads for `next dev` only, so the production static export keeps its same-origin relative URLs. The README said the variable was for non-default backends; it now explains what unsetting it costs.
llama.cpp learned the `gemma4-assistant` architecture in xllamacpp 2026.6.9713 (2026.6.9538 still fails to load it), which makes the tower drafter usable through its `draft-mtp` speculative implementation — the pre-existing one covers weights-internal NEXTN layers instead. Unlike the other engines the drafter is a single gguf published inside the target's own repository, so specs gain `draft_model_file_name_template` and the cache manager downloads that one file rather than a snapshot; `draft_model_id` falls back to the target's repo. Its quantizations are its own (`BF16`, `F16`, `Q8_0`), unrelated to the target's, so a Q4_K_M target pairs with an unquantized drafter by default.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for speculative decoding (MTP) with paired drafter checkpoints across multiple engines, including vLLM, SGLang, MLX, and llama.cpp, along with corresponding UI and documentation updates. The review feedback highlights two important issues: first, in xinference/core/worker.py, using os.path.realpath to resolve symlinks during cache deletion can lead to data loss by deleting shared files in the global HuggingFace or ModelScope cache; second, in xinference/model/llm/llama_cpp/core.py, glob.glob requires recursive=True to properly search nested directories for .gguf files.
The gguf specs declared only `draft_model_file_name_template` and leaned on the cache manager's fallback for the repository, but both the launch dialog and the cache-status probe key off `draft_model_id` — so the whole Speculative Decoding section never appeared for llama.cpp, and its drafters never reported whether they were downloaded. State the drafter repository explicitly and keep the fallback for user-registered specs. The per-round token count also read as if 4 were universal, which is the MLX drafter's trained depth. It is the engine that decides: 1 for vLLM, 6 for SGLang, 3 for llama.cpp. The field now names the default of the selected engine, and the tooltip spells out that only MLX treats a larger value as a ceiling rather than a fixed block size. llama.cpp gains the same three additional-parameter entries the other engines already offered, for a local drafter path.
Measured on this branch: gemma-4 31B 4bit goes 14.9 -> 31.0 tok/s with MTP, while 26B-A4B 4bit goes 73.2 -> 65.9. The switch is worth having off by default, and worth explaining why. The deciding ratio is the cost of a drafting step against a target decoding step. The drafter is a fixed 0.83 GB dense model, so a MoE target that reads only its activated slice per token shifts that ratio from ~5% to ~39%: the three drafting steps of a round then already cost more than one plain decoding step, before accounting for a lower acceptance rate. Also notes what is not the cause, since it is the intuitive suspect: verifying a block is cheap on both, a four-token forward costing ~40% more than a single-token one.
Deleting a model's cache resolved the drafter links and removed what they point at. That is right for a model file, which belongs to one cache entry, but wrong for a drafter: one download is shared by every quantization of its target, so all their `-draft-<quant>` entries link to the same thing. Removing it left the sibling quantizations reporting a cached drafter that no longer loads. Only our own cache entries are removed now; the hub cache keeps the blob for its own tooling to reclaim. `glob` with `**` and no `recursive=True` collapses to a single level, so a drafter sitting directly in the cache directory was not found — only the nested `MTP/` layout was, which is what the test happened to cover. Both are pinned now: one test asserts the shared download and the snapshot directory survive while the cache entries go, the other resolves a flat gguf. Both fail against the previous implementation.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements speculative decoding (MTP) support across multiple engines (vLLM, SGLang, MLX, and llama.cpp) in Xinference, complete with Web UI controls, documentation, and cache management. The feedback highlights a critical issue in LLMCacheManager where a GGUF draft model lacking a file name template could silently fall back to downloading the target model's files; storing use_draft_model as an instance variable and raising an explicit error in this scenario is recommended.
A gguf download is per file, so a drafter needs its own file name. With only `draft_model_id` set the file list fell through to the target's template and would have asked the drafter repository for the target's file name — a 404, or worse the target's own weights downloaded into the drafter cache directory. Not reachable through the built-in specs, which always declare both, but a user-registered spec can get there. The launch now fails in the cache manager's constructor, next to the other drafter validations, and points at `draft_model_path` for a local file. The file-list helper tracks draft mode so it can never silently fall back.
Gemma delimits string arguments with `<|"|>` instead of quoting them, so a
value is free to contain double quotes. The parser replaced the delimiters
with `"` textually and handed the result to json.loads, which then failed on
any value carrying a quote of its own:
reason:<|"|>The user said "你好" (Hello), a simple greeting.<|"|>
→ Expecting ',' delimiter: line 1 column 163
The whole block was returned as plain text, so a caller relying on tool calls
saw none — an agent asking the model to route through a decision tool reports
that the model returned no tool call, when the model did exactly what it was
asked. Quoting the user back is what most routing prompts produce, so this
fires constantly rather than at the margins.
Each string is now re-encoded with json.dumps instead, and bare keys are
quoted only in the structural text between strings. That also fixes a value
like `a,b:c`, whose inner `b:` was rewritten as a key, and values containing
backslashes or newlines.
rogercloud
left a comment
rogercloud
left a comment
There was a problem hiding this comment.
Summary
This PR adds speculative decoding for Gemma 4 via paired drafter checkpoints (a separate draft model, not inline MTP layers) across four engines: vLLM, SGLang, MLX, and llama.cpp. It extends the LLM spec schema and LLMCacheManager so drafter checkpoints can be declared, downloaded, and cache-managed alongside their target, adds a Web UI section gated on drafter availability, and registers gemma-4 12B as a second encoder-free architecture. Two unrelated fixes also ride along: a dev-proxy streaming fix and a gemma_tool_parser quoting fix. 30 files changed, +1679/-45, no linked issue.
Design verdict
Acceptable with reservations. The core shape is right: reusing LLMCacheManager by reparameterizing it instead of adding a second downloader is the correct call, LLM.support_draft_model fails fast before a multi-GB download, and the <cache_dir>-draft-<quantization> sibling directory scheme makes cascade delete a prefix scan with no new bookkeeping.
Design-level concerns:
- Drafter spec-shape validation lives in
LLMCacheManager.__init__rather than as a spec-level pydantic validator. This is the direct cause of the finding onworker.py/supervisor.pybelow: a malformed custom spec surfaces as a crash on a listing endpoint rather than at registration time. draft_model_id/draft_quantizations/draft_model_file_name_templateand theirvalidate_model_size_with_radixvalidator are duplicated near-verbatim across three spec classes inxinference/model/llm/llm_family.py(LlamaCppLLMSpecV2~88-94,PytorchLLMSpecV2~119-122,MLXLLMSpecV2~147-166). All three subclassBaseModeldirectly with identical types and defaults; a shared mixin would remove the drift risk. Maintainability, not correctness.- The drafter is downloaded in
create_llm_model_instance(xinference/model/llm/core.py:346-349) before any check of whether the user already supplied an explicit speculative config. The raw kwargs needed for that check are already in scope at that point, so hoisting the check earlier would avoid downloading and then discarding multiple GB of unused weights. enable_mtpreuses DeepSeek-MTP vocabulary for a mechanism this PR explicitly distinguishes from DeepSeek-style MTP. Naming nit only.- Three independent changes are bundled: the tool-parser quoting fix, the
.env.developmentdev-proxy streaming fix, and locale reflows. Each is fine standalone, but together they make the PR hard to revert atomically. Process note, not blocking.
Additional findings (not anchorable inline - pre-existing/unchanged lines)
- Minor (latent),
xinference/model/llm/cache_manager.py(cache_from_huggingface/cache_from_modelscope, ~line 203/271):self._llm_familyis never reassigned in theuse_draft_modelbranch, so the drafter'ssnapshot_downloadreuses the target'scache_config(e.g.allow_patterns/ignore_patterns), which is unchanged pre-existing code not touched by this PR. The same branch explicitly nulls_model_uriand_multimodal_projectorfor exactly this class of reason, butcache_configwas missed. No builtin family currently setscache_config, so this is dormant - but a future family combining a drafter with a target-scopedcache_configwould silently get an incomplete drafter download. Suggest resettingcache_configalongside the existing_model_uri/_multimodal_projectorreset. - Minor,
xinference/model/llm/llm_family.json(gemma-4 31Bfp4spec, ~line 30296): this spec has nodraft_*fields at all, while the 31Bpytorchspec at the same size does (this line is unchanged by the diff, so no drafter fields were added here either way). Perbackends.rstthe UI section simply stays hidden when a spec has no drafter, so this fails safe - but nothing states why fp4 lacks a drafter conversion. Worth a one-line comment confirming this is deliberate (no published NVFP4 drafter conversion) rather than an oversight. - Minor,
xinference/model/llm/mlx/core.py:1382(MLXVisionModel.stop(), unchanged by this diff): clears the executor but never clearsself._draft_model(norself._model, which has the same pre-existing gap). Not a regression introduced by this PR, but it extends an existing memory-retention pattern to a new, potentially large attribute. Worth a follow-up to release both references.
Simplification opportunities
xinference/model/llm/tool_parsers/gemma_tool_parser.pyL48-77:_arguments_to_jsonsplits on a literal string-token delimiter with a manual char-by-char loop usingstartswith/find.arg_block.split(token)yields the same alternating segments (even indices structural, odd indices literal string content), andlen(parts) % 2 == 0reproduces the unterminated-string check. Net: roughly -15 to -20 lines. Sits inside the unrelated tool-parser fix, not the speculative-decoding code.
Re-review: prior findings
All three prior findings from earlier review rounds are fixed, independently re-verified this round:
- Symlink
os.path.realpathcascade-delete data-loss risk (worker.py) - FIXED in7d49bbef. Norealpathremains in the drafter deletion block (worker.py:4114-4125); the drafter path is added topathsbefore anyislinkcheck, so only the symlink itself is unlinked, and directory drafters are walked with plainos.walk.test_list_deletable_models_keeps_shared_drafter_downloadsasserts a shared blob and shared snapshot survive. glob.globmissingrecursive=True(llama_cpp/core.py) - FIXED in7d49bbef.recursive=Trueis present atllama_cpp/core.py:166-167, withtest_cache_dir_resolves_a_flat_ggufcovering the flat case.- GGUF drafter silently falling back to target file names (
cache_manager.py) - FIXED in0a7cf5a9c._use_draft_modelis now a stored attribute (cache_manager.py:51) and the constructor raises aValueErrornaming the missingdraft_model_file_name_templatebefore any download begins. Covered bytest_gguf_drafter_without_a_file_template_is_rejectedandtest_gguf_drafter_downloads_only_its_own_file.
Test results
Ran the PR's new and changed unit tests in a fresh venv (editable install, without engine-specific optional deps): 96 passed, 2 skipped, 9 failed. All 9 failures are ModuleNotFoundError/ImportError for xllamacpp and mlx_lm, absent from the review environment - environmental, not code defects. Where a test could not execute past the missing-dependency gate (notably the glob-recursive fix and the gguf drafter template rejection), the underlying logic was confirmed by reading the code instead. This matches the author's stated gap: no actual vLLM/SGLang/llama.cpp run, for lack of hardware.
One malformed drafter spec took down model listing. Both cache-status probes construct a drafter cache manager per quantization, and its constructor rejects shapes no registration-time validator catches, so a single bad custom spec turned every listing call into a 500. Report that spec as not cached and log why instead. Both probes also only looked at `draft_model_id`, missing the gguf drafter identified by its file template alone; that guard and the Web UI now accept either, since the cache manager does. llama.cpp silently clobbered an explicitly configured drafter, unlike vLLM and SGLang which both bail with a log line. It now bails the same way when any `speculative*` passthrough key is present. The engine-neutral options are also pulled out before the dotted-key loop rather than inside `_apply_draft_model`: CommonParams has no such attributes, so the loop logged a failure for each of them on every speculative launch, teaching users to ignore real errors. SGLang overwrote an explicitly supplied `speculative_num_draft_tokens` while using `setdefault` for the two keys beside it; all three are consistent now, and the step count follows whichever value ends up in effect. The per-round defaults were described as "the engine's own default" in the docs, the frontend docblock and all four locales. They are not: only MLX reads the depth from the drafter, and the others are values from each engine's Gemma 4 guide. vLLM's is now a named constant recording that provenance, and the wording says it is a conservative starting point worth raising. There is no field in the checkpoint to resolve it from — Gemma 4's trained depth of 4 lives in mlx-vlm's own config class, not in the drafter's config.json. The Web UI also offered `enable_mtp` and `num_speculative_tokens` as free-form kwargs while both have dedicated controls; the kwargs list wins when merged, so adding one there silently overwrote the control. Dropped from every engine's list, `draft_model_path` kept since it has no control. And the Speculative Decoding section now also requires an engine that can run a drafter, exposed per engine through the engines endpoint, so Transformers no longer offers a switch whose launch is rejected. Finally, simplify the Gemma argument scanner to a split on the delimiter, which yields the same alternating segments in a third of the lines.
rogercloud
left a comment
rogercloud
left a comment
There was a problem hiding this comment.
Follow-up Review (commit bdff06a)
This round covers a single update commit (bdff06a, 18 files, +258/-79) in which the author replies to all 11 findings from the prior review at 44f2bac. The PR was re-read from a blank slate as a whole change, and every prior finding was re-verified independently against the code rather than accepted from the author's replies. The fix commit resolves most of what was raised, but it also introduces one regression that defeats its own UI fix by default, plus two new inconsistencies.
Design verdict (independent, whole-PR)
Verdict is unchanged: acceptable with reservations. This is an independent judgment of the change as it now stands, not a tally of "11/11 addressed" - several fixes are correct in mechanism while leaving the underlying shape untouched.
- Validation of malformed drafter specs still lives in
LLMCacheManager.__init__rather than in a spec-level validator. The fix commit wrapped the call site intry/exceptinstead of moving validation to where the spec is defined. Pragmatic and not wrong, but the layering gap is papered over rather than closed. No action requested this round; recording it as an open design note. - The "explicit speculative config wins" invariant now means three different things across three engines: vLLM keys on the presence of the
speculative_configobject, SGLang on thespeculative_algorithmselector, and llama.cpp (new in this commit) on any key under thespeculative/speculative.*namespace. The llama.cpp variant is the broadest and least consistent of the three - see the Major finding below. - The new per-engine
support_draft_modelcapability is exposed in only one of two structurally identical engine-listing code paths, and the untouched one is the default. See the Critical finding. - Four engines now ship four different hardcoded or undiscoverable "default speculation depth" values: vLLM 1, SGLang 6, MLX resolved from the drafter, and llama.cpp a documented "3" that does not exist in code. This is the same underlying gap flagged before - the drafter's trained depth is a property of the checkpoint, not of the engine, and the spec declares it nowhere.
What actually changed since the last review (independently verified, not taken from the author's replies)
| # | Finding | Status |
|---|---|---|
| 1 | Malformed drafter spec crashes model-listing (worker.py/supervisor.py) | FIXED - try/except added, test_malformed_drafter_spec_does_not_break_listing passes. (See new minor finding below: the catch is broader than necessary.) |
| 2 | draft_cache_status guard misses template-only gguf drafter |
FIXED - guard widened to draft_model_id or draft_model_file_name_template in both backend files and the frontend hasDrafter check. |
| 3 | llama.cpp missing explicit-config-wins guard | FIXED - guard added. (See new major finding below: the guard is broader than the sibling engines' equivalents.) |
| 4 | llama.cpp spurious ERROR logs on every speculative launch | FIXED - _pop_draft_options now runs before the setattr loop. (See new minor finding below: the pop is destructive and affects retries.) |
| 5 | vLLM default of 1 documented as "the engine's own default" | FIXED - docs/locales reworded to correctly describe it as a deliberately conservative Xinference-chosen starting point; the value itself (1) is unchanged, which the author explains is intentional given Gemma 4 drafters don't publish their trained depth. |
| 6 | SGLang overwrites explicit speculative_num_draft_tokens |
FIXED - all three keys now use .setdefault. (See new minor finding below: a residual 0/string edge case in the derived step count.) |
| 7 | Frontend UI gates on spec only, ignoring engine support | PARTIAL - the intended fix (exposing support_draft_model per engine and gating on it) is correctly implemented, but only wired into one of the two engine-listing code paths. See the new Critical finding below - this is effectively still open for any default-configuration deployment. |
| 8 | Frontend duplicate kwargs/dedicated-field mechanism | FIXED - enable_mtp/num_speculative_tokens removed from every engine's free-form kwargs list; draft_model_path correctly retained since it has no dedicated control. |
| 9 | Drafter download inherits target's cache_config (latent) |
NOT ADDRESSED - cache_manager.py is untouched by this commit; noted as dormant/latent in the original review and remains so, not a regression. |
| 10 | gemma-4 31B fp4 spec has no draft_* fields |
NOT ADDRESSED - llm_family.json is untouched by this commit; unchanged from before, likely intentional. |
| 11 | MLX stop() doesn't clear _draft_model |
NOT ADDRESSED - mlx/core.py is untouched by this commit; unchanged, not a regression. |
| - | gemma_tool_parser.py simplification (str.split) |
APPLIED CORRECTLY - verified the replacement preserves the even/odd segment alternation and the len(parts) % 2 == 0 unterminated-string check exactly. |
Tests
The PR's unit tests were re-run against the updated worktree: 100 passed, 2 skipped, 9 failed - up from 96 passing previously, reflecting the 4 new passing tests added by this commit. All 9 failures are the same pre-existing environmental ModuleNotFoundError/ImportError for xllamacpp/mlx_lm, absent from the review sandbox and unrelated to this change. No new test failures were introduced.
The previous commit taught only one of the two engine-discovery entry points to publish `support_draft_model`, and it taught the wrong one: virtualenv mode is on by default, and that path rebuilds engine entries from specs rather than from the registry, so the flag never reached a default deployment. Gating the Web UI on it therefore hid the Speculative Decoding section everywhere — worse than before that gating existed. Both paths now go through one helper for registry-built entries, and the spec-rebuilt lists are annotated from the engine's classes. A test asserts both entry points publish the flag, and fails if either regresses. Also from review: - Narrow llama.cpp's explicit-config guard to the two keys that actually select the mode. It matched any `speculative.*` key, so a tuning knob like `speculative.draft.n_min` silently disabled the drafter. - Read the neutral options instead of popping them, and skip them in the passthrough loop. Popping lost the drafter on the load retry, which reuses the instance and its config. - Stop claiming llama.cpp defaults to 3 tokens because this code says so: it does not set the value at all, and 3 is xllamacpp's own CommonParams default. The docs, docblock and four locales now name three provenances instead of one. - Coerce SGLang's `speculative_num_draft_tokens` in the config rather than only in the derived local, so a string from the Web UI does not reach the engine, and treat an explicit 0 as supplied. - Narrow the cache-status guards to `ValueError`, what the cache manager raises, so an unrelated bug there surfaces instead of reading as "not cached".
rogercloud
left a comment
rogercloud
left a comment
There was a problem hiding this comment.
Round 3 Re-Review - Speculative Decoding with Paired Drafter Checkpoints
This round covers the single new commit ebf04d8, which replies to all six remaining findings from round 2 (base bdff06a). Each prior finding was independently re-verified against the code rather than accepted from the reply text, and the fix commit itself was reviewed blank-slate for regressions. Result: three of six prior findings are genuinely closed, three are partial, and the partials share a single root cause worth naming explicitly.
Design assessment
The feature architecture remains sound: paired drafter checkpoints, per-engine translation of speculative options, and capability gating via support_draft_model are the right shape, and nothing in this round changes that verdict. The reservation is about the fix process, not the design.
This is now the third consecutive round in which a fix for one code path leaves an identical sibling path or sibling variable untouched. Round 1 added support_draft_model to one of two engine-listing functions. Round 2 found the second function missing it, and the round-2 fix patched that second function. Round 3 finds a third entry point - _force_virtualenv_engine_params - that rebuilds engine param lists from specs, bypassing both earlier paths, and which needed its own annotation logic added this round. Separately, the 0-vs-falsy coercion bug was fixed for speculative_num_draft_tokens but left in place for the sibling key num_speculative_tokens three lines above it - and that sibling pattern turns out to exist identically in all four engines, meaning it was never sglang-specific and was not addressed at all. And the except Exception narrowing was done by swapping the exception type without auditing what LLMCacheManager.__init__ actually raises across all its failure paths.
The recurring shape - fix N-1 of N duplicated things, or fix the exact reported key but not its sibling - suggests fixes are being applied by pattern-matching the specific finding text rather than by asking "where else does this exact shape occur." Three rounds of the same failure mode is a process signal rather than a code nit, and it is the main thing standing between this PR and a clean sign-off.
Prior findings - independently re-verified
Statuses below come from reading the code at ebf04d8, not from the author's replies.
| # | Round-2 finding | Severity | Status |
|---|---|---|---|
| 1 | support_draft_model missing from the default (virtualenv) discovery path |
Critical | Partial - presence fixed on all paths, but the two paths compute the value differently and can disagree (see Major finding below); currently latent for builtins only by coincidence |
| 2 | llama.cpp explicit-config-wins guard too broad (any speculative.* key) |
Major | Fixed - SPECULATIVE_SELECTOR_KEYS = ("speculative.types", "speculative.draft.mparams.path") correctly narrows the guard to the two mode selectors; verified in code |
| 3 | Docs/UI claimed a llama.cpp default of "3" not present in code | Major | Fixed - the value is confirmed accurate (xllamacpp CommonParams default), and docs plus all four locales now attribute it to xllamacpp, distinguishing each engine's actual provenance |
| 4 | except Exception too broad in worker.py/supervisor.py |
Minor | Partial - narrowed as claimed, but now too narrow and bypassable via KeyError/IndexError from .format() in the same constructor (see Major finding below) |
| 5 | sglang 0-vs-falsy and string coercion | Minor | Partial - the reported key is correctly fixed; the identical bug remains in the sibling num_speculative_tokens and is systemic across all four engines (see Major finding below) |
| 6 | llama.cpp destructive pop breaks retry |
Minor | Fixed for the reported scope - _pop_draft_options replaced by a non-destructive _draft_options reader using .get(), and the passthrough loop now skips DRAFT_OPTION_KEYS rather than relying on their absence; mechanism verified correct |
Aside on #6, out of scope and explicitly not a regression from this PR: load() is not safely retryable regardless, because it also calls self._llamacpp_model_config.pop("reasoning_content") with no default elsewhere in the function, which raises KeyError on a second load() for any llama.cpp model, speculative or not. This is pre-existing and unrelated, but it means the retry scenario motivating the fix does not work end-to-end today for a separate reason.
Tests
Re-ran the PR's unit tests: 103 passed, 2 skipped, 10 failed, up from 100 passed / 9 failed in round 2 - one additional test, which hits the same pre-existing environmental gate. All 10 failures are the same ModuleNotFoundError/ImportError for xllamacpp and mlx_lm, which are absent from the review sandbox. No new failures attributable to logic changes in this commit.
Closing
Not ready to sign off yet, and the reason is the pattern rather than any single finding. All three new findings this round are of the form "the reported instance was fixed, the sibling instance was not," and two of them are direct regressions of round-2 fixes rather than pre-existing issues. The remaining work is small in code volume - one shared coercion helper across four engines, one widened except clause, and either reusing the per-entry annotation logic or adding a parity test - but the right next step is a deliberate sibling sweep by the author before the next push: for each of the three fixes, grep for the same shape elsewhere and fix all occurrences at once, rather than the one named in the review.
A fourth pass is warranted, and it should be scoped narrowly to verifying that sweep plus the three fixes above, not a full re-read. If the sweep lands cleanly with the suggested tests (parity across discovery paths, a malformed-placeholder registration that does not break listing, and num_speculative_tokens=0 per engine), this should be mergeable at that point.
`num_speculative_tokens or DEFAULT` was written four times, once per engine, and all four silently replaced an explicit 0 with their default — the same bug class fixed for SGLang's engine-native key one commit ago, left in the shared one. There is now a single `parse_num_speculative_tokens`: None means unset so the engine's own default applies, anything else must be a positive integer, and a string from the Web UI is coerced. A zero, a negative or a non-integral float is rejected rather than quietly turned into something else. `_annotate_draft_support` reported whether *any* class registered under an engine name takes a drafter, which is wrong per entry: MLX registers a text class and a vision class and only the vision one does. It now resolves which classes claim an entry with the same `match_json` call the non-forced branch uses, falling back to the coarse answer only when nothing claims it — which is the normal case on the forced path, where matching cannot succeed yet. A test asserts both discovery paths report the same value for the same model. The cache manager's `.format()` calls could still raise KeyError or IndexError for a spec carrying an unexpected placeholder, bypassing the ValueError guard on the listing path and crashing the whole call again. They are normalized to ValueError naming the offending field.
rogercloud
left a comment
rogercloud
left a comment
There was a problem hiding this comment.
Round 4 Re-Review
This round covers a single new commit, b35cb69, which addresses all three major findings from round 3. The diff is narrow (8 files, ~200 lines): a new shared parse_num_speculative_tokens validator in xinference/model/llm/core.py, per-entry class resolution via match_json in _annotate_draft_support, and .format() error normalization at both cache_manager.py call sites. A blank-slate whole-PR pass was also run independently of the prior rounds' findings; its verdict and this round's new findings follow.
Design verdict and notes
Acceptable with reservations. Centralizing num_speculative_tokens validation into one function is real progress and is meaningfully different from what the previous three rounds delivered: instead of four independent per-engine patches (three of which still let 0 through silently), there is now one validator, one error-message format, and one place to get the semantics right. This is the first fix in four rounds that targets the root cause - duplicated logic - rather than the specific reported instance. That should be recognized as the correct instinct.
Two non-blocking design notes. First, _annotate_draft_support's family and specs parameters are typed Optional[Any] = None, but both call sites always pass concrete values; this is dead flexibility and the parameters should be made required with real types. Second, the match-checking block (the match_json lookup, the try/except, and the result normalization) now appears three times in xinference/model/utils.py - at both call sites in _force_virtualenv_engine_params and again inside _annotate_draft_support. Extracting a single _spec_matches_engine helper is worth doing here specifically because this exact shape of duplication is what produced the two-disagreeing-computations bug that round 3 reported.
MLX class resolution: confirmed correct
The match_json-based per-entry class resolution in _annotate_draft_support is verified to correctly distinguish MLXChatModel and MLXModel (support_draft_model=False) from MLXVisionModel (True) under the same engine name, which the previous engine-wide any() could not. The new parity test test_engine_discovery_paths_agree_on_draft_support genuinely exercises the mechanism rather than trivially asserting agreement: gemma-4 and qwen3 share the same global MLX_CLASSES list that the old any() iterated over, and qwen3 resolving to MLXChatModel rather than MLXVisionModel demonstrates class exclusion working independently of drafter presence. This is a solid fix. The one residual gap is in tests, not code: no test constructs a model that is both chat-only and has a drafter declared, so the single most direct form of the originally reported divergence is proven by code trace and by a structurally equivalent test, but not by one exact end-to-end scenario.
Prior findings - independently re-verified
| # | Round 3 finding | Status | Notes |
|---|---|---|---|
| 1 | support_draft_model computed two disagreeing ways across discovery paths |
Fixed | Per-entry match_json resolution now matches the strict path; the parity test proves the fix rather than asserting agreement. |
| 2 | except ValueError bypassable by KeyError/IndexError from cache_manager.py .format() |
Partial | Both call sites are wrapped and normalize to ValueError; the exact reported scenario is fixed and tested. The conditional gate on draft_model_id leaves an unguarded sibling case (see the Minor finding below), so "any malformed drafter template raises a clear ValueError" is not yet true. |
| 3 | 0-vs-falsy num_speculative_tokens across all four engines |
Partial | The neutral parameter is now correctly validated in all four engines via the centralized function - done well and architecturally right. The sglang-native speculative_num_draft_tokens was not migrated (see the Major finding below), so "an explicit speculation-depth value is never silently altered" does not yet hold everywhere. |
Test results
The PR's unit tests were re-run: 107 passed, 2 skipped, 10 failed, up from 103 passed / 10 failed in round 3, so four new tests pass. All 10 failures are the same pre-existing environmental ModuleNotFoundError/ImportError for xllamacpp and mlx_lm, which are absent from the review sandbox. No new failures are attributable to the logic changes in this commit.
Closing
The fix quality improved this round: centralizing the validator rather than patching four call sites individually is the right instinct, and the MLX class-resolution fix is both correct and well-tested. The remaining concern is a pattern rather than a single defect - this is the fourth consecutive round where the reported instance was fixed and an unmigrated sibling of the same shape remained, and this time the sibling is the neutral-versus-engine-native variant of round 3's own finding, sitting a few lines below the fixed code in the same function. Rather than a general appeal to look harder, the shape is now specific enough to grep for: before the next push, search the PR for any remaining direct int(...) casts, or DEFAULT fallbacks, or bare .format() calls applied to a speculation-depth or drafter-template value without going through parse_num_speculative_tokens or _format_draft_field. With the sglang speculative_num_draft_tokens fix and that sweep done, the remaining gaps are narrow and well-tested enough that this should be ready to merge.
rogercloud
left a comment
rogercloud
left a comment
There was a problem hiding this comment.
Approval
After four rounds of review (design pass, discovery, and independent verification of every prior finding against the code rather than against commit messages), the core feature - paired drafter checkpoints for speculative decoding across vLLM, SGLang, MLX, and llama.cpp - is sound, well-tested for the layers reachable without the optional engine dependencies, and the cache/deletion/UI-gating machinery has held up under repeated scrutiny. Approving on that basis.
Two items are still open from the round-4 pass (inline comments above). Flagging both again here so they are not lost, with emphasis on the first:
Major - not fully closed, recommend addressing before or immediately after merge
xinference/model/llm/sglang/core.py:333-337 - speculative_num_draft_tokens still bypasses the new parse_num_speculative_tokens validator. This is the fourth consecutive round where a fix closed the specific reported instance (the neutral num_speculative_tokens parameter, now correctly validated in all four engines) and left an identical sibling unfixed one function away. Concretely: an explicit speculative_num_draft_tokens=0 still passes through today, producing a self-contradictory speculative_num_draft_tokens=0, speculative_num_steps=1 configuration, and a non-integral float like 2.7 is silently truncated to 2 - exactly the class of silent substitution the new validator exists to prevent. It is not merge-blocking in the sense of breaking normal usage (it requires setting the SGLang-native key directly rather than the documented neutral option), but given the four-round pattern, it is worth closing in the very next commit rather than deferring indefinitely: route this key through parse_num_speculative_tokens too, and add a 0/float rejection test for it, matching what the neutral parameter already has.
Minor - low risk, fine to defer
xinference/model/llm/cache_manager.py:86 vs :99 - asymmetric placeholder validation. draft_model_id's new .format() guard is gated by if "{draft_quantization}" in draft_model_id:, while the sibling draft_model_file_name_template check has no such gate. A draft_model_id with some other unexpected placeholder that does not contain the literal {draft_quantization} substring skips validation and surfaces as a confusing download-time 404 instead of a clear registration-time ValueError. Low risk (custom-registration edge case only), safe to fold into the same follow-up.
Non-blocking design notes (already raised, restating for the record)
_annotate_draft_support'sfamily/specsparameters are typedOptional[Any] = Nonebut both call sites always pass concrete values - dead flexibility, make them required.- The match-checking block (match_json lookup, try/except, result normalization) is now triplicated across
xinference/model/utils.py- worth extracting into one_spec_matches_enginehelper, particularly because this exact duplication shape is what caused the discovery-path divergence bug fixed this round.
None of the above blocks merging as-is; recommend a fast follow-up PR or commit for the SGLang validator gap given the recurring pattern, so a fifth round isn't needed to catch a fifth sibling.
SGLang's engine-native `speculative_num_draft_tokens` kept a bare `int()` cast
while its neutral sibling moved to the shared validator, so an explicit 0
survived into a self-contradictory "propose 0 draft tokens, budget 1 step"
config and 2.7 truncated to 2. It goes through `parse_num_speculative_tokens`
now, like every other depth in the codebase.
`draft_model_id` only went through the placeholder normalization when it
contained `{draft_quantization}`, unlike the file template beside it. A field
carrying some other placeholder — `org/model-{quantization}`, or a bare `{}` —
skipped the check and was sent to the hub verbatim, turning a clear error at
launch into a 404 mid-download. Both fields are formatted unconditionally now,
and both require a quantization only when they actually carry the placeholder.
Tests cover 0, a negative, a non-integral float and a non-numeric string for
both SGLang keys, and the two placeholder shapes that used to slip through;
the latter fail against the previous gating.
|
ci is failing |
Use the shared Gemma 4 per-size speculation depth in SGLang and the launch UI, so E4B defaults to four tokens like vLLM. Upgrade the SGLang virtualenv stack to the first patched release that recognizes gemma4_assistant checkpoints.
3 participants
Gemma 4 ships a small
*-it-assistantdrafter for every variant: it proposes several tokens per round that the target verifies in one pass, so decoding gets faster with identical output. Unlike DeepSeek-V3 style MTP the drafter is a separate checkpoint that shares the target's KV cache, so it has to be declared, downloaded and handed to the engine.Launch with
--enable_mtp true, or flip the switch under Advanced Configuration → Speculative Decoding in the Web UI.Registry and cache
Specs gain
draft_model_id, optionally templated bydraft_quantizations— the MLX conversions of Gemma 4 12B publish eight, the gguf drafters three.LLMCacheManagergrows ause_draft_modelmode that reuses the existing per-hub download paths, so a drafter follows the same HuggingFace/ModelScope selection as its target and lands in a sibling<cache_dir>-draft-<quantization>directory. Every drafter repository referenced here was checked to exist on both hubs.create_llm_model_instanceresolves the drafter whenenable_mtpis passed and hands the local path to the engine. Engines that cannot use one are rejected up front viaLLM.support_draft_modelrather than failing somewhere inside engine startup.Deleting a model's cache now removes its drafters too, instead of orphaning a directory per drafter quantization next to it.
Per engine
vllm>=0.22.0speculative_configwithmethod: mtp--speculative-algorithm NEXTNplus matching step/draft-token countsmlx-vlm>=0.5.0(>=0.6.1for 12B)draft_modelongenerate_step, kind auto-detected from the drafter'smodel_typexllamacpp>=2026.6.9713draft-mtpspeculative implementationAn explicitly provided
speculative_config/speculative_algorithmalways wins; the engine-neutral options are consumed either way so they never reach the engine as unknown arguments. vLLM below 0.22.0 is rejected with a readable error because it treats the assistant as a generic draft model and fails to initialize against a multimodal target. Same for xllamacpp below 2026.6.9713, which cannot load thegemma4-assistantarchitecture at all.The Transformers engine is not supported: it runs its own continuous-batching loop rather than
generate(), so there is nowhere to attach a drafter.MLX loads the drafter on the same dedicated MLX thread as the target and validates the pair when the model loads, rather than on the first request.
gemma-4 12B
Also registers gemma-4 12B across pytorch/gguf/mlx. It is the encoder-free Unified variant (
Gemma4UnifiedForConditionalGeneration, draftergemma4_unified_assistant), which is why the family now lists two architectures. It is also the only variant whose drafter is published in more than one conversion, so it is what makes the drafter selection meaningful.Web UI
The Speculative Decoding section only appears for a format/size that actually declares a drafter. The drafter options carry the same cached marker the model quantizations use, because turning the switch on downloads another checkpoint. Tooltips cover the non-guessable parts: that the switch costs a download, why an unquantized drafter is recommended even for a quantized target, and that a value above the drafter's trained depth acts as a ceiling rather than a fixed block size.
One unrelated fix rides along, because it blocks anyone verifying this in the dev UI: without
NEXT_PUBLIC_API_URLthe browser sends API requests to the Next dev server, whose rewrite proxy buffers the body, so streaming endpoints deliver the whole answer at once..env.developmentpoints the browser at the backend directly; the production static export is unaffected (Next does not load that file for builds).Verification
Measured end to end on Apple silicon with gemma-4 E4B, 8bit target and a bf16 drafter: 2.2x on decode in isolation, ~15% on a full chat request. Driving
generate_stepdirectly shows the expected burst pattern — tokens arrive in groups the size of the draft block, and streaming is preserved.Unit tests cover the config translation for all four engines, the drafter selection and cache-directory rules, and the cascade delete. Not covered: an actual vLLM/SGLang/llama.cpp run, for lack of the hardware — those three are wired from their documented parameters and tested at the config-assembly level only.