Studio: drop the speculative drafter under Auto when only the model fits in VRAM - #8435
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1b3cbfbde
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # Same downgrade as the MLA branch below: ngram-mod costs no VRAM, and | ||
| # spec-off when the build lacks it. The drafter paths stay recorded, so | ||
| # the UI still names the kind and a repeat Apply still dedupes. | ||
| self._spec_fallback_reason = "drafter_no_vram" |
There was a problem hiding this comment.
When Auto resolved a DSpark sidecar and this fallback runs, _emit_ngram_mod() changes self._speculative_type to ngram-mod. The duplicate-load comparator identifies an Auto DSpark request only when that field is draft-dspark, so it subsequently compares the incoming MTP path (None) with the stored DSpark sidecar and rejects the match. As a result, every repeated Apply reloads the same healthy drafter-free server; use self._spec_drafter_kind or otherwise retain the resolved DSpark identity in that comparator.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The premise does not hold at head. _spec_drafter_kind is assigned at llama_cpp.py:13712, before any emit branch, from canonical_mode / dspark_draft_path / caps, and drafter_no_vram never touches it, so an Auto DSpark load still carries "dspark". The comparator consults it alongside _speculative_type (llama_cpp.py:4218), and self._mtp_draft_path still records the resolved sidecar, so the compare matches and no reload happens. test_auto_drops_the_drafter_when_only_the_target_fits asserts both.
for more information, see https://pre-commit.ci
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
for more information, see https://pre-commit.ci
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9daf6ba6b4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if _foot_w <= _budget_w: | ||
| _both_fit_somewhere = True |
There was a problem hiding this comment.
For an unsized drafter using the flat reserve, _probe_frac(True) lowers every GPU's usable budget, but the new per-device cap is evaluated only with _usable_wo from the target-only fraction. Fresh evidence in the current code is that _foot_w <= _budget_w immediately declares success without checking whether each device holds its pipeline and replicated compute reserve under the drafter-adjusted fraction. On a heterogeneous split the pooled total can pass while the weakest card fails the real planner's _every_gpu_holds_reserve check, so the drafter is kept and the context is still capped; repeat the per-device check using _probe_frac(True) before setting _both_fit_somewhere.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not taking this one, and I checked it properly rather than on the argument. In the sized case _probe_frac(True) is identical to _probe_frac(False), and the target-only per-device check has already run at this same _ctx_wo and capped if needed, so the extra check is provably a no-op. In the unsized case _mtp_bytes is 0, so _foot_w rises by only the 224 MiB draft-graph reserve while _budget_w drops by 5 percent of the subset's TOTAL VRAM: 1229 MiB on a 24 GB card, 4915 MiB on a 96 GB one. The pooled test therefore decides first in every configuration I could construct, and reaching the per-device check needs something like a 96 GB card with 24 GB free carrying a 17 GB compute buffer. I did write the check and then reverted it rather than ship a tightening I cannot demonstrate is reachable.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…ble charging a priced override
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
- Draft flags are last-wins, so a repo id followed by --model-draft leaves a path as the drafter. _extra_args_mtp_draft_source now returns the value and whether the flag that won carried a repo id, so a path is never priced as a Hugging Face repository. - split_listing_is_complete tracks distinct shard indices inside 1..total rather than counting matches, so 00001-of-00002 beside a stray 00003-of-00002 is no longer read as a whole set.
_cached_repo_gguf_bytes scanned huggingface_hub's import-time default, so a user who moved the cache had their cached drafter missed and priced at the flat reserve instead. Pass the active hub cache, as the rest of Studio's cache operations do. Also covers the underscore flag spelling llama.cpp accepts.
for more information, see https://pre-commit.ci
The reserve constant carried two overlapping headers, the extras-drafter note still described the local-file-only rule that was reverted once remote repos could be priced, and the zero-bound note no longer matched the code.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
- A remote split drafter the listing sized only in part is treated as unsized, so the guard measures the cache or falls back to the reserve instead of charging the shards it happened to know about. llama-server maps every shard. - A CPU-offloaded sidecar now clears the byte-accurate mtp_overhead_fn as well as the flat fraction. The callback is invoked whatever mtp_engaged says, so the fit went on cutting context for GPU bytes no drafter allocates. - The drafter probe walks both GPU rankings. An unsized drafter costs five points of pin fraction, which can reorder heterogeneous cards, and the placement loop sorts under the fraction the load ends up with.
|
Codex Review: Didn't find any major issues. Breezy! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
for more information, see https://pre-commit.ci
Comment-only pass over the two source files: 231 added comment lines down to 159, with the duplicated zero-bound note and the leftover reserve header gone. No code or docstring changes (comment_tools.py check passes).
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Speculative decoding is a speed option, but on a tight card Auto currently makes you pay for it with context length, and under
--fitit can make you pay for it with an OOM. This makes Auto give up the drafter first.The problem
The drafter is treated as non-negotiable in the VRAM fit, and the context is what gives way:
_mtp_will_engageis decided first, then_mtp_reserve_bytesis added to the budget alongside weights, KV and mmproj.Context auto-reduced: 32768 -> 8192), and failing that the load falls back to--fit onand offloads layers to host, where decode collapses roughly 3x (Studio: reserve CUDA context and mmproj/MTP soft overhead in the GGUF fit budget #6718).--fit, llama.cpp cannot size a separate drafter at all, so its weights are never reserved and the load can simply run out of VRAM.So on a card that holds the model but not the model plus its drafter, an Auto load either silently halves the context you asked for, or over-commits. In neither case did the user ask for speculative decoding: Auto chose it for them.
The change
Before the fit runs, price the load twice at the context the target alone would get: once without the drafter and once with it. If the target fits and target plus drafter does not, Auto drops the drafter and keeps the context.
--speculative-type, or via--spec-type/--model-draftin extra args is honored exactly as before, including today's context reduction.speculative_type, not_mtp_effective. By the time the fit runs, Auto has already been promoted to the drafter kind it resolved, so_mtp_effectivecan no longer tell a forceddsparkfrom an Auto one.drafter_no_vram, which takes the same downgrade the MLA branch already uses:ngram-mod, which costs no VRAM, or spec-off on a build that lacks it. Routing it through_build_speculative_flagsrather than blanking the drafter path matters for the embedded-MTP case, where there is no sidecar file to blank and emitting--spec-type draft-mtpagainst a reserve of zero is exactly the OOM this is meant to prevent.Applies to MTP, DSpark and DFlash alike: it lives in the shared fit path and keys on
_mtp_will_engage, so it covers the DFlash drafter in #8338 once that lands.What the user sees
A new
spec_fallback_reason,drafter_no_vram, and the amber notice in chat settings reads:Unlike
binary_no_mtp, it does not offer the llama.cpp update affordance, because an update does not fix a VRAM shortfall.Not covered
--fitloads are unchanged. llama.cpp cannot size the sidecar during fitting, so there is no reserve to compare against and nothing to decide from.routes/inference.pystill charges the sidecar under Auto, so with a training job running a chat load can still be refused for a drafter the loader would now shed. Matching it means replicating the fit math in the route, and loosening it risks evicting a training run, so it is left alone here.Tests
Four in
test_llama_cpp_placement.py, each driving a realload_modelagainst a stubbed 24 GB card and asserting on the composed argv:--model-draft, nodraft-dspark, and-c 8192survivesspec_fallback_reasonstaysNonedsparkoverrides the dropngram-modThe two new-behaviour tests were confirmed to fail without the change (they launch the drafter); the two guardrails pass either way, which is their point.
Codex follow-ups from #8338
Two items Codex raised on #8338 after it was approved, both landing on the guard this PR already touches:
--spec-type draft-dsparkand a local--model-draftin extra args,_build_speculative_flagsreturns before any mode branch, so Studio's own sidecar is neither fetched nor launched; only the named file becomes resident, and it is already charged as_extras_bytes. The guard was billing both and could return a spurious 409. Gated on both halves: owning the spec type while naming no drafter keeps the conservative charge, since a drafter can still arrive by a route the estimate cannot see and this guard protects a running training job._download_companion_ggufrefuses a split family whose encoded shard count is short, butdflash_budget_byteswas billing whatever the listing showed, so a repo caught mid-publication could be refused a load that fits. The DSpark path already filtered onsplit_listing_is_complete; the rule now lives in the budget helper where the shard-set logic already is, so the two cannot drift.One regression test each, both confirmed to fail without their fix.
Two further Codex items from #8338 are not in this PR, because each is a separate surface that deserves its own change and its own tests:
repo/Q4_K_M/model.gguf, discovery searchesrepo/for the sidecar butnative_gguf_companion_parent_allowedrejects it, so Auto silently loses DFlash on that layout. The fix has to thread the actual granted root through validation rather than the root derived from the weight path, or it widens the native-grant boundary.GatedRepoErrorsets no retry flag and DFlash is excluded from thedrafter_not_foundretry arm, so re-applying with a newly valid token dedupes against the drafterless server. Needs credential state in the dedup comparator, which is thrash-sensitive.