From 4412bc0dd0dda17202fc186b00521f9002cdac68 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 11:44:06 +0000 Subject: [PATCH 01/26] Studio: launch a DFlash drafter automatically Studio has recognised dflash-*.gguf since #7811, but only to hide it from the quant picker. Nothing ever launched it, so a model that ships a DFlash sidecar fell through to no speculative decoding at all. Add DFlash as the third launchable drafter kind beside MTP and DSpark: a _is_dflash_drafter_path predicate, local and Hub discovery, a supports_dflash capability parsed from llama-server --help, and the --model-draft / --spec-type draft-dflash emission. Unlike DSpark it is on under Auto, since the published sidecar is 1.52 GiB and ships in the model's own GGUF repo rather than being an ~11 GB opt-in fetch. DSpark keeps first refusal when a repo somehow ships both, matching llama.cpp's own downloader. Discovery confirms general.architecture = dflash in the header rather than pairing on the filename: the published sidecar is dflash-kquant.gguf, which names no model family, so the DSpark pairing rule would reject the one file this exists to find. The dflash/ directory is still not a drafter marker, and DFlash is still excluded from companion reclaim, both because the name doubles as a family a publisher puts on real weights. --- CHANGELOG.md | 3 + studio/backend/core/inference/llama_cpp.py | 362 ++++++++++-- studio/backend/hub/utils/gguf.py | 9 +- studio/backend/models/inference.py | 20 +- studio/backend/routes/inference.py | 85 ++- .../tests/test_mtp_drafter_companion.py | 521 +++++++++++++++++- .../tests/test_native_gguf_companion.py | 2 + studio/backend/tests/test_tensor_parallel.py | 4 +- studio/backend/utils/models/model_config.py | 123 +++++ .../utils/openai_auto_switch_settings.py | 10 +- .../src/features/chat/chat-settings-sheet.tsx | 18 +- .../chat/stores/chat-runtime-store.ts | 9 +- .../frontend/src/features/chat/types/api.ts | 21 +- .../components/model-config-page.tsx | 1 + .../model-config/per-model-config.ts | 3 + studio/frontend/src/lib/speculative-modes.ts | 4 +- 16 files changed, 1117 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1974a219f51..23f30248ae0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,9 @@ rename the heading at release time. ## Unreleased +- DFlash speculative decoding now starts automatically for models that ship a + DFlash sidecar, with no setting to change. + ## 2026.8.7 ### What's Changed diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3f92187a55c..468fe0977e4 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -150,6 +150,7 @@ class GgufLoadIntent: mmproj_path: Optional[str] = None mtp_draft_path: Optional[str] = None dspark_draft_path: Optional[str] = None + dflash_draft_path: Optional[str] = None hf_repo: Optional[str] = None hf_variant: Optional[str] = None hf_token: Optional[str] = None @@ -1144,6 +1145,9 @@ def _is_mtp_model_name(model_identifier: Optional[str], gguf_path: Optional[str] _DRAFTER_KINDS = ("mtp", "dspark", "dflash") _DRAFTER_DIR_KINDS = ("mtp", "dspark") +# Human label per resolved spec mode, for launch logging and the UI notice. +_DRAFTER_DISPLAY_LABELS = {"dspark": "DSpark", "dflash": "DFlash"} + def _drafter_path_kind(path: str) -> Optional[str]: """Drafter kind naming *path*: basename prefix, or exact parent dir for @@ -1199,6 +1203,22 @@ def _is_dspark_drafter_path(path: str) -> bool: return _drafter_path_kind(path) == "dspark" +def _is_dflash_drafter_path(path: str) -> bool: + """True for a DFlash sidecar, excluding the separate DSpark method. + + Prefix only: ``dflash`` doubles as a family name a publisher puts on real + weights, so ``Qwen3.6-35B-A3B-DFlash-Q4_K_M.gguf`` IS the model, and a + user's ``dflash/`` folder holds whatever they put there (``dflash`` is + absent from ``_DRAFTER_DIR_KINDS`` for that reason). Only a root-level + ``dflash-*.gguf`` names a sidecar. + + Broader than model_config.detect_dflash_file, which additionally confirms + ``general.architecture = dflash`` in the header. Only ever used against repo + listings and cache snapshots, where a header read is not available. + """ + return _drafter_path_kind(path) == "dflash" + + _BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE) _GGUF_KNOWN_QUANT_RE = re.compile( r"(UD-)?" @@ -2314,6 +2334,13 @@ def _extra_args_requests_dspark( return "draft-dspark" in _accumulated_spec_types(extra_args, env) +def _extra_args_requests_dflash( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if DFlash lands in llama.cpp's accumulated spec-type vector.""" + return "draft-dflash" in _accumulated_spec_types(extra_args, env) + + @functools.lru_cache(maxsize = 1) def _metal_device_is_paravirtual() -> bool: """True when Metal is a virtualised Apple GPU, whose offload can corrupt output. @@ -2594,7 +2621,8 @@ def _extra_args_requests_separate_draft( _extra_args_requests_mtp: a later --spec-default or --spec-type cannot clear an inherited draft-simple, so reading only the last under-reserves the model it loads.""" return bool( - _accumulated_spec_types(extra_args, env) & {"draft-simple", "draft-eagle3", "draft-dspark"} + _accumulated_spec_types(extra_args, env) + & {"draft-simple", "draft-eagle3", "draft-dspark", "draft-dflash"} ) @@ -2919,14 +2947,24 @@ def _build_ngram_mod_flags( # Canonical Speculative Decoding modes exposed by the Unsloth chat UI. -# Dropdown renders six (auto, mtp, dspark, ngram, mtp+ngram, off); the load API -# also accepts legacy values the original Switch and external callers emit -# (default, draft-mtp, ngram-mod, ngram-simple). -_CANONICAL_SPEC_MODES = {"auto", "mtp", "dspark", "ngram", "mtp+ngram", "off", "ngram-simple"} +# Dropdown renders seven (auto, mtp, dspark, dflash, ngram, mtp+ngram, off); the +# load API also accepts legacy values the original Switch and external callers +# emit (default, draft-mtp, ngram-mod, ngram-simple). +_CANONICAL_SPEC_MODES = { + "auto", + "mtp", + "dspark", + "dflash", + "ngram", + "mtp+ngram", + "off", + "ngram-simple", +} _LEGACY_SPEC_MODE_MAP = { "default": "auto", "draft-mtp": "mtp", "draft-dspark": "dspark", + "draft-dflash": "dflash", "ngram-mod": "ngram", } @@ -2934,8 +2972,8 @@ def _build_ngram_mod_flags( def _canonicalize_spec_mode(value): """Map any accepted ``speculative_type`` input onto a canonical mode. - Returns ``auto``, ``mtp``, ``dspark``, ``ngram``, ``mtp+ngram``, ``off``, - ``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``). + Returns ``auto``, ``mtp``, ``dspark``, ``dflash``, ``ngram``, ``mtp+ngram``, + ``off``, ``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``). Unknown strings collapse to ``auto`` so a stale UI value or typo falls back to the safe platform-aware path. """ @@ -3170,6 +3208,7 @@ def __init__(self): self._spec_fallback_reason: Optional[str] = None self._spec_drafter_kind: Optional[str] = None self._dspark_sidecar_absent: bool = False + self._dflash_sidecar_absent: bool = False # Set after an auto-Vulkan crash recovers with all devices disabled. self._cpu_fallback_reason: Optional[str] = None self._cpu_fallback_runtime: Optional[_CpuFallbackRuntime] = None @@ -3971,8 +4010,9 @@ def _norm(value): if ( intent.gguf_path is None and self._spec_fallback_reason == "drafter_not_found" - and speculative_type in ("auto", "mtp", "mtp+ngram", "dspark") + and speculative_type in ("auto", "mtp", "mtp+ngram", "dspark", "dflash") and not (self._spec_drafter_kind == "dspark" and self._dspark_sidecar_absent) + and not (self._spec_drafter_kind == "dflash" and self._dflash_sidecar_absent) and not spec_owned_by_extra_args ): return False @@ -3983,7 +4023,7 @@ def _norm(value): compared_draft_n_max = self._last_load_intent.spec_draft_n_max if ( ( - self._speculative_type in ("draft-mtp", "draft-dspark") + self._speculative_type in ("draft-mtp", "draft-dspark", "draft-dflash") or self._spec_fallback_reason == "runtime_error" ) and intent.spec_draft_n_max is not None @@ -4000,17 +4040,24 @@ def _norm(value): "mtp", "mtp+ngram", "dspark", + "dflash", ): try: - # Auto counts as dspark once it resolved that way: the launch stored - # the DSpark sidecar, so comparing the MTP field (None for these + # Auto counts as dspark/dflash once it resolved that way: the launch + # stored that sidecar, so comparing the MTP field (None for these # repos) against it would reload a healthy server on every Apply. _compare_dspark = speculative_type == "dspark" or ( speculative_type == "auto" and self._speculative_type == "draft-dspark" ) - intent_draft = ( - intent.dspark_draft_path if _compare_dspark else intent.mtp_draft_path + _compare_dflash = speculative_type == "dflash" or ( + speculative_type == "auto" and self._speculative_type == "draft-dflash" ) + if _compare_dspark: + intent_draft = intent.dspark_draft_path + elif _compare_dflash: + intent_draft = intent.dflash_draft_path + else: + intent_draft = intent.mtp_draft_path requested_draft = Path(intent_draft).resolve() if intent_draft else None # A drafter the last load dropped on purpose counts as launched here: # the file is still there, so comparing it against the stored None @@ -4081,7 +4128,7 @@ def speculative_type(self) -> Optional[str]: @property def spec_drafter_kind(self) -> Optional[str]: - """Which drafter the resolution was about, ``mtp`` or ``dspark``. + """Which drafter the resolution was about: ``mtp``, ``dspark`` or ``dflash``. Distinct from ``requested_spec_mode``, which stays ``auto`` when Auto resolves the kind itself, and from ``speculative_type``, which reads @@ -4268,8 +4315,9 @@ def _dspark_release_is_broken(cls, release_tag: Optional[str]) -> bool: @classmethod def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]: """Parse `llama-server --help` for feature flags. Returns - {found, mtp_token, supports_mtp, supports_dspark, ngram_mod_flavor, - supports_ngram_mod, spec_draft_n_max_flag, cache flag support}. + {found, mtp_token, supports_mtp, supports_dspark, supports_dflash, + ngram_mod_flavor, supports_ngram_mod, spec_draft_n_max_flag, cache flag + support}. ``ngram_mod_flavor``: ``"new"`` when the post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args; @@ -4290,6 +4338,7 @@ def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, ob "mtp_token": None, "supports_mtp": False, "supports_dspark": False, + "supports_dflash": False, "mtp_probe_inconclusive": True, "ngram_mod_flavor": None, "supports_ngram_mod": False, @@ -4317,6 +4366,7 @@ def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, ob mtp_token: Optional[str] = None supports_dspark = False + supports_dflash = False ngram_mod_flavor: Optional[str] = None spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False @@ -4406,6 +4456,15 @@ def _is_real(flag: str) -> bool: spec_help.lower(), ) ) + # Same word-boundary match as DSpark: draft-dflash only exists on + # builds that carry the arch, and emitting a --spec-type the binary + # does not know aborts the launch instead of falling back. + supports_dflash = bool( + re.search( + r"(? bool: "mtp_token": mtp_token, "supports_mtp": supports_mtp, "supports_dspark": bool(supports_dspark and saw_spec_type and probe_ok), + "supports_dflash": bool(supports_dflash and saw_spec_type and probe_ok), "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": ngram_mod_flavor, "supports_ngram_mod": ngram_mod_flavor is not None, @@ -8026,6 +8086,106 @@ def _pick_dspark(candidates: list[str]) -> Optional[str]: self._dspark_sidecar_absent = outcome.get("listed") is False return found + def _cached_repo_dflash_drafter( + self, + hf_repo: str, + *, + cache_dir: Optional[str] = None, + ) -> Optional[str]: + """The preferred already-cached DFlash sidecar for a repo, Q8_0 first + (dflash_preference_key), so an offline reuse picks the same file the + online download would have fetched.""" + try: + from utils.models.model_config import ( + _iter_hf_cache_snapshots, + dflash_preference_key, + ) + + snapshots = ( + _iter_hf_cache_snapshots(hf_repo) + if cache_dir is None + else _iter_hf_cache_snapshots(hf_repo, cache_dir) + ) + candidates: list[Path] = [] + for snap in snapshots: + candidates.extend( + snap / name + for name in _gguf_snapshot_files(snap) + if _is_dflash_drafter_path(name) + ) + for candidate in sorted(candidates, key = lambda p: dflash_preference_key(p.name)): + if candidate.is_file(): + return str(candidate) + except Exception as exc: + logger.debug("Cached DFlash drafter lookup failed for %s: %s", hf_repo, exc) + return None + + def _download_dflash( + self, + *, + hf_repo: str, + hf_token: Optional[str] = None, + near_path: Optional[str] = None, + binary: Optional[str] = None, + ) -> Optional[str]: + """Download the published DFlash sidecar, preferring its Q8_0 copy. + + Unlike the ~11 GB DSpark sidecar, which is why that one is opt-in, the + published DFlash sidecar is ~1.5 GiB and already ships in the model's own + GGUF repo, so Auto fetches it the same way it fetches the MTP drafter. + Still gated on ``supports_dflash``: a binary with no usable + ``--spec-type draft-dflash`` falls back, so the download would never be + opened. A raised probe still fetches, since launch re-probes and may yet + engage. + """ + + def _pick_dflash(candidates: list[str]) -> Optional[str]: + from utils.models.model_config import dflash_preference_key + files = sorted( + (name for name in candidates if _is_dflash_drafter_path(name)), + key = dflash_preference_key, + ) + return files[0] if files else None + + cached = _companion_snapshot_sibling(near_path, _pick_dflash) if near_path else None + if not cached and _hf_env_offline(): + cached = self._cached_repo_dflash_drafter( + hf_repo, + cache_dir = _hub_cache_dir_for_snapshot_path(near_path), + ) + try: + if not self.probe_server_capabilities(binary).get("supports_dflash"): + logger.warning( + "Skipping the DFlash sidecar download: llama-server has no usable " + "--spec-type draft-dflash, so this load falls back to no speculative " + "decoding. Run `unsloth studio update`, then reload." + ) + # A sidecar already on disk is still reported, for the same reason as + # the DSpark path: the route rediscovers it on every Apply, so + # answering None would make the reuse check compare it against a + # launched None and reload the same server each time. + # _build_speculative_flags re-checks the capability and still falls back. + return cached + except Exception as exc: + logger.debug("DFlash capability probe failed before the sidecar fetch: %s", exc) + + if cached: + logger.info("Reusing cached DFlash drafter: %s", cached) + return cached + outcome: dict = {} + found = self._download_companion_gguf( + hf_repo = hf_repo, + hf_token = hf_token, + pick = _pick_dflash, + label = "DFlash drafter", + near_path = near_path, + outcome = outcome, + ) + # Distinguishes a repo that ships no sidecar from a fetch that failed and + # could yet succeed. + self._dflash_sidecar_absent = outcome.get("listed") is False + return found + def _resolve_launch_mmproj_path( self, *, model_path: str, mmproj_path: Optional[str] ) -> Optional[str]: @@ -9291,6 +9451,7 @@ def load_model(self, intent: GgufLoadIntent) -> bool: mmproj_path = intent.mmproj_path mtp_draft_path = intent.mtp_draft_path dspark_draft_path = intent.dspark_draft_path + dflash_draft_path = intent.dflash_draft_path hf_repo = intent.hf_repo hf_variant = intent.hf_variant hf_token = intent.hf_token @@ -9606,6 +9767,23 @@ def load_model(self, intent: GgufLoadIntent) -> bool: near_path = model_path, binary = binary, ) + # DFlash: same shape, and Auto is included for the same + # reason, but it is not gated behind an opt-in the way the + # ~11 GB DSpark sidecar is. The published sidecar is ~1.5 GiB + # and ships in the model's own GGUF repo, so under Auto it + # costs about what the MTP drafter costs. Repos without one + # no-op after a single listing. + if ( + not dflash_draft_path + and _spec_canon in ("auto", "dflash") + and not _extra_args_set_spec_type(extra_args) + ): + dflash_draft_path = self._download_dflash( + hf_repo = hf_repo, + hf_token = hf_token, + near_path = model_path, + binary = binary, + ) elif gguf_path: if not Path(gguf_path).is_file(): raise FileNotFoundError(f"GGUF file not found: {gguf_path}") @@ -9634,6 +9812,22 @@ def load_model(self, intent: GgufLoadIntent) -> bool: except Exception as exc: logger.debug("DSpark capability probe failed during Auto: %s", exc) + # DFlash is the other Auto promotion, on the same capability gate. + # DSpark keeps first refusal (llama.cpp's own downloader ranks it + # ahead too), so a repo shipping both is unchanged; in practice a + # repo ships one kind or neither. + if ( + _spec_canon == "auto" + and dflash_draft_path + and not _extra_args_set_spec_type(extra_args) + ): + try: + if self.probe_server_capabilities(binary).get("supports_dflash"): + _spec_canon = "dflash" + logger.info("Auto: DFlash sidecar available, using draft-dflash.") + except Exception as exc: + logger.debug("DFlash capability probe failed during Auto: %s", exc) + # MTP and DSpark are mutually exclusive, and the drafter sizing and # lifecycle path below is architecture agnostic, so it carries the # one drafter the mode selected. From here mtp_draft_path means @@ -9641,6 +9835,8 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # _spec_canon is the only thing that says which. if _spec_canon == "dspark": mtp_draft_path = dspark_draft_path + elif _spec_canon == "dflash": + mtp_draft_path = dflash_draft_path # Read GGUF metadata (context_length, chat_template); header-only. self._read_gguf_metadata(model_path) @@ -10068,8 +10264,9 @@ def _pool_budget_mib(subset, frac): # engage; auto only on an MTP model >= 3B; ngram/off never. A # separate drafter (Gemma) counts as an MTP model. # _spec_canon, not the raw request: Auto has already resolved to - # dspark above when a sidecar is available, and the reserve below - # differs by kind (no duplicated target-KV copy, depth 3). + # dspark or dflash above when a sidecar is available, and the + # reserve below differs by kind (no duplicated target-KV copy, + # and DSpark uses depth 3). _mtp_effective = _spec_canon _mtp_size_for_fit = _extract_model_size_b(model_identifier) # Sub-3B drops MTP only for an embedded head; a separate @@ -10110,9 +10307,13 @@ def _pool_budget_mib(subset, frac): if not _user_mtp_via_extras: try: _fit_caps = self.probe_server_capabilities(binary) or {} + _sidecar_cap = { + "dspark": "supports_dspark", + "dflash": "supports_dflash", + }.get(_mtp_effective) _mtp_binary_ok = bool( - _fit_caps.get("supports_dspark") - if _mtp_effective == "dspark" + _fit_caps.get(_sidecar_cap) + if _sidecar_cap else _fit_caps.get("mtp_token") ) except Exception: @@ -10122,7 +10323,7 @@ def _pool_budget_mib(subset, frac): not _extra_args_set_spec_type(extra_args) and _mtp_model_for_fit and ( - _mtp_effective in ("mtp", "mtp+ngram", "dspark") + _mtp_effective in ("mtp", "mtp+ngram", "dspark", "dflash") or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) ) and ( @@ -10143,10 +10344,11 @@ def _pool_budget_mib(subset, frac): # only charge it when the engaged mode is truly MTP. Those modes # arrive by two routes, hence the two conditions: draft-simple # and draft-eagle3 via _user_draft_via_extras (already excluded, - # they never set _user_mtp_via_extras), and DSpark via Studio's - # own resolution, which does set _auto_studio_mtp. + # they never set _user_mtp_via_extras), and DSpark/DFlash via + # Studio's own resolution, which does set _auto_studio_mtp. _engaged_is_mtp = bool( - _user_mtp_via_extras or (_auto_studio_mtp and _mtp_effective != "dspark") + _user_mtp_via_extras + or (_auto_studio_mtp and _mtp_effective not in ("dspark", "dflash")) ) # Effective draft depth: extras win (last-wins at launch), else @@ -11244,7 +11446,7 @@ def _restore_after_tensor_downgrade(): _draft_device = ",".join(f"Vulkan{i}" for i in _vulkan_pin_ids) launch_mtp_draft_path = self._resolve_launch_mtp_path( mtp_draft_path = mtp_draft_path, - drafter_label = "DSpark" if _spec_canon == "dspark" else "MTP", + drafter_label = _DRAFTER_DISPLAY_LABELS.get(_spec_canon, "MTP"), ) _pv_suppressed_draft_path: Optional[str] = None _pv_suppressed_spec_extra_args: Optional[List[str]] = None @@ -11319,9 +11521,13 @@ def _restore_after_tensor_downgrade(): model_path = model_path, gpus = bool(_detected_gpus), binary = binary, - mtp_draft_path = (None if _spec_canon == "dspark" else launch_mtp_draft_path), + mtp_draft_path = ( + None if _spec_canon in ("dspark", "dflash") else launch_mtp_draft_path + ), dspark_draft_path = (launch_mtp_draft_path if _spec_canon == "dspark" else None), dspark_fit_sized = not use_fit, + dflash_draft_path = (launch_mtp_draft_path if _spec_canon == "dflash" else None), + dflash_fit_sized = not use_fit, draft_device = _draft_device, ) # _build_speculative_flags judged the stripped list, so a user @@ -12261,6 +12467,9 @@ def _try_auto_vulkan_cpu_fallback( _spec_requested_dspark = any( "draft-dspark" in str(t).lower() for t in spec_flags ) or _extra_args_requests_dspark(extra_args, env = _launch_spec_env) + _spec_requested_dflash = any( + "draft-dflash" in str(t).lower() for t in spec_flags + ) or _extra_args_requests_dflash(extra_args, env = _launch_spec_env) # Is the launched server actually running MTP+tensor? Gates the # probe/watchdog/recovery; cleared if the MTP-drop fallback wins. _mtp_active_for_launched_server = bool( @@ -12324,7 +12533,7 @@ def _try_auto_vulkan_cpu_fallback( _spec_cpu_replay_cmd: Optional[List[str]] = None if ( not healthy - and (_spec_requested_mtp or _spec_requested_dspark) + and (_spec_requested_mtp or _spec_requested_dspark or _spec_requested_dflash) and not self._cancel_event.is_set() ): _spec_cpu_replay_cmd = list(_last_spawn_cmd) @@ -12676,6 +12885,8 @@ def _build_speculative_flags( mtp_draft_path: Optional[str] = None, dspark_draft_path: Optional[str] = None, dspark_fit_sized: bool = True, + dflash_draft_path: Optional[str] = None, + dflash_fit_sized: bool = True, draft_device: Optional[str] = None, ) -> List[str]: """Return the llama-server flag list for the requested spec mode. @@ -12784,16 +12995,17 @@ def _build_speculative_flags( # fallback can erase the evidence: the UI names the recovery from this, # and under Auto neither the requested mode nor the resolved # _speculative_type ("default" after a fallback) still carries the kind. - self._spec_drafter_kind = ( - "dspark" - if canonical_mode == "dspark" - or ( - (canonical_mode or "auto") == "auto" - and dspark_draft_path - and caps.get("supports_dspark") - ) - else "mtp" - ) + _auto_mode = (canonical_mode or "auto") == "auto" + if canonical_mode == "dspark" or ( + _auto_mode and dspark_draft_path and caps.get("supports_dspark") + ): + self._spec_drafter_kind = "dspark" + elif canonical_mode == "dflash" or ( + _auto_mode and dflash_draft_path and caps.get("supports_dflash") + ): + self._spec_drafter_kind = "dflash" + else: + self._spec_drafter_kind = "mtp" def _resolved_draft_n_max() -> int: # User override wins; else platform default (the B200 / x86 @@ -12865,6 +13077,73 @@ def _emit_dspark() -> None: _emit_dspark() return flags + def _emit_dflash() -> None: + """Append --model-draft --spec-type draft-dflash + n-max. + + Callers have already established the sidecar and the capability. Auto + reaches this too, since DFlash is the default whenever the model + ships one: unlike the ~11 GB DSpark sidecar it is ~1.5 GiB and comes + from the model's own GGUF repo. + """ + if not dflash_fit_sized: + # Same shape as DSpark: the sidecar ships no token_embd/output and + # borrows the target's, so llama.cpp cannot build a standalone + # draft context to measure it (llama-context.cpp:153-160) and + # skips the reserve (server-context.cpp:1190-1193). The load still + # works; only the ~1.5 GB is missing from the fit budget, which is + # an order of magnitude less than the DSpark case. + logger.info( + "DFlash under --fit on: llama.cpp cannot size the sidecar during " + "fitting, so its ~%.1f GB is not reserved. Pin GPU Layers in " + "Manual mode if the load runs out of VRAM.", + (self._get_gguf_size_bytes(dflash_draft_path) or 0) / 1e9, + ) + draft_n_max = _resolved_draft_n_max() + n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max" + flags.extend( + [ + "--model-draft", + dflash_draft_path, + "--spec-type", + "draft-dflash", + str(n_max_flag), + str(draft_n_max), + ] + ) + if draft_device: + # A separate drafter, like the Gemma MTP one, so it takes the same + # device pin (_emit_mtp does this); a baked-in head has nowhere to + # put it and DSpark's sidecar predates the flag being threaded here. + flags.extend(["--spec-draft-device", draft_device]) + self._speculative_type = "draft-dflash" + logger.info("Spec decoding: draft-dflash using %s", dflash_draft_path) + + if effective_mode == "dflash": + # Capability first, for the same reason as DSpark: the fetch is gated + # on the same answer, so a missing sidecar here is usually the + # binary's fault, and blaming the file would tell the user to place + # one and reload on every Apply. + if not caps.get("supports_dflash"): + logger.warning( + "DFlash requested but llama-server lacks --spec-type " + "draft-dflash; loading without speculative decoding." + ) + flags.append("--spec-default") + self._speculative_type = "default" + self._spec_fallback_reason = "binary_no_mtp" + return flags + if not dflash_draft_path: + logger.warning( + "DFlash requested but no matching dflash-*.gguf sidecar was found; " + "loading without speculative decoding." + ) + flags.append("--spec-default") + self._speculative_type = "default" + self._spec_fallback_reason = "drafter_not_found" + return flags + _emit_dflash() + return flags + def _emit_mtp(*, chain_ngram: bool) -> bool: """Append --spec-type mtp[/draft-mtp][,ngram-mod] + n-max.""" mtp_token = caps.get("mtp_token") if caps else None @@ -12989,6 +13268,14 @@ def _fallback_drafter_not_found() -> None: # this architecture (1.84x on 4x B200, 1.91x on one). Without it these # models fall through to --spec-default, i.e. no drafter at all. _emit_dspark() + elif dflash_draft_path and caps.get("supports_dflash"): + # DFlash next, on the same reasoning as DSpark and with the same + # capability gate: load_model only hands a sidecar down once it has + # one this binary can launch. Unlike DSpark this needs no opt-in: the + # published sidecar is ~1.5 GiB and already sits in the model's own + # GGUF repo, so Auto pays roughly what the MTP drafter costs. DSpark + # keeps first refusal so a repo shipping both is unchanged. + _emit_dflash() elif _auto_mla_embedded_mtp: # MLA embedded-MTP (GLM-5.2 et al.): the MTP path regresses vs spec-off # on llama.cpp today, so Auto drops it and falls back to ngram-mod (or @@ -13232,6 +13519,7 @@ def unload_model(self) -> bool: self._spec_fallback_reason = None self._spec_drafter_kind = None self._dspark_sidecar_absent = False + self._dflash_sidecar_absent = False self._cpu_fallback_reason = None self._last_load_intent = None self._mtp_runtime_fallback_active = False diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index fc4886df4b7..c921e634896 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -124,8 +124,13 @@ def is_reclaimable_drafter_path(path: str) -> bool: """Drafters a repo's last-variant delete may reclaim: MTP, fetched with every variant, and DSpark, fetched on opt-in. Both are useless once no main GGUF is left, and companion filtering hides them from the variant menu, so leaving one - behind is an invisible allocation (DSpark is ~11 GB). DFlash is excluded: the - name doubles as a family a user picks for real weights.""" + behind is an invisible allocation (DSpark is ~11 GB). DFlash is excluded even + though Auto now launches it: the name doubles as a family a user picks for + real weights, whole repos publish nothing but root-level ``dflash-*.gguf`` + (Lucebox/Qwen3.6-27B-DFlash-GGUF), and the two outcomes are not symmetric. + Reclaiming wrongly destroys weights a user chose; not reclaiming leaves + ~1.5 GiB, an order of magnitude under the DSpark case this rule was written + for. Locked by test_deleting_the_last_variant_keeps_a_dflash_weight.""" p = path.replace("\\", "/").lower() if not p.endswith(".gguf"): return False diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2910fa38f4b..9eee28b77be 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -106,13 +106,15 @@ def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optiona description = ( "Speculative decoding mode for GGUF models. Canonical values: " "'auto' (platform-aware: DSpark when the model ships a sidecar, " - "else MTP on MTP GGUFs, ngram-mod fallback for sub-3B), " + "else DFlash when it ships one, else MTP on MTP GGUFs, ngram-mod " + "fallback for sub-3B), " "'mtp' (force draft-mtp only on both GPU and CPU), " "'dspark' (force a draft-dspark sidecar), " + "'dflash' (force a draft-dflash sidecar), " "'ngram' (force ngram-mod only), 'mtp+ngram' (force " "ngram-mod+draft-mtp chain on both platforms), 'off' (disabled). " "Legacy values 'default' (-> auto), 'draft-mtp' (-> mtp), " - "'draft-dspark' (-> dspark), " + "'draft-dspark' (-> dspark), 'draft-dflash' (-> dflash), " "'ngram-mod' (-> ngram), and 'ngram-simple' (kept as-is) are " "still accepted. Ignored for non-GGUF models." ), @@ -122,11 +124,12 @@ def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optiona ge = 1, le = 16, description = ( - "Max draft tokens per step for MTP or DSpark speculative decoding " - "(--spec-draft-n-max). Defaults to 2 on GPU and 3 on CPU/Mac " - "when unset (upstream-bench sweet spot for dense Qwen3.6 MTP " - "quants). Only applied when speculative_type resolves to " - "'mtp', 'mtp+ngram', or 'dspark'." + "Max draft tokens per step for MTP, DSpark or DFlash speculative " + "decoding (--spec-draft-n-max). Defaults to 2 on GPU and 3 on " + "CPU/Mac when unset (upstream-bench sweet spot for dense Qwen3.6 " + "MTP quants, and the measured sweet spot for DFlash too). Only " + "applied when speculative_type resolves to 'mtp', 'mtp+ngram', " + "'dspark' or 'dflash'." ), ) n_parallel: Optional[int] = Field( @@ -842,7 +845,8 @@ class InferenceStatusResponse(_InferenceRuntimeFields): spec_drafter_kind: Optional[str] = Field( None, description = ( - "Which drafter the resolution was about, 'mtp' or 'dspark'. Needed " + "Which drafter the resolution was about: 'mtp', 'dspark' or " + "'dflash'. Needed " "because Auto resolves the kind itself, so speculative_type still " "reads 'auto', and a fallback leaves the engaged type at 'default': " "neither still says which file the UI should tell the user to fix." diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index abed1a7d405..c2859645ce0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1141,6 +1141,7 @@ def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: from utils.models.model_config import ( _local_gguf_companion_search_root, colocated_split_shards, + detect_dflash_file, detect_dspark_file, detect_mtp_file, load_model_defaults, @@ -1193,6 +1194,7 @@ def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: from utils.models.model_config import ( _local_gguf_companion_search_root, colocated_split_shards, + detect_dflash_file, detect_dspark_file, detect_mtp_file, load_model_defaults, @@ -3828,6 +3830,10 @@ def _loaded_is_local_model( _DRAFTER_NATIVE_RULES = { "mtp": ("MTP drafter", "mtp"), "dspark": ("DSpark drafter", "dspark"), + # No companion subdirectory: dflash/ is a family name a user picks for real + # weights, so detect_dflash_file only ever offers a root-level sidecar and + # nothing outside the model's own directory is in bounds. + "dflash": ("DFlash drafter", None), } @@ -3857,7 +3863,7 @@ def _validate_native_mtp_drafter( str(shard), gguf_path, label, - allowed_subdirs = (subdir,), + allowed_subdirs = (subdir,) if subdir else (), mtp_search_root = mtp_search_root, ) @@ -4050,7 +4056,7 @@ def _drafter_for_path( """ if not gguf_path: return None - detect = detect_dspark_file if kind == "dspark" else detect_mtp_file + detect = {"dspark": detect_dspark_file, "dflash": detect_dflash_file}.get(kind, detect_mtp_file) root = _local_gguf_companion_search_root(gguf_path, gguf_path) rejected = False accept = None @@ -4103,6 +4109,20 @@ def _dspark_draft_for_path( ) +def _dflash_draft_for_path( + gguf_path: Optional[str], + native_grant_backed: bool, + *, + log_native_fallback: bool = False, +) -> Optional[str]: + return _drafter_for_path( + gguf_path, + native_grant_backed, + kind = "dflash", + log_native_fallback = log_native_fallback, + ) + + def _active_gguf_intent( request: LoadRequest, llama_backend: LlamaCppBackend, @@ -4163,6 +4183,7 @@ def _active_gguf_intent( ), mtp_draft_path = _mtp_draft_for_path(llama_backend.gguf_path, native_grant_backed), dspark_draft_path = _dspark_draft_for_path(llama_backend.gguf_path, native_grant_backed), + dflash_draft_path = _dflash_draft_for_path(llama_backend.gguf_path, native_grant_backed), compare_mtp_draft = True, extra_args_inherited = inherits_extras and not batch_overrides_inherit, ) @@ -5284,16 +5305,21 @@ def _remote_gguf_companion_bytes( include_mmproj: bool, include_mtp: bool = True, include_dspark: bool = False, + include_dflash: bool = False, ) -> int: """Bytes of companion GGUFs the requested launch downloads. 0 on error.""" try: - from core.inference.llama_cpp import _is_dspark_drafter_path + from core.inference.llama_cpp import ( + _is_dflash_drafter_path, + _is_dspark_drafter_path, + ) from huggingface_hub import model_info - from utils.models.model_config import dspark_preference_key + from utils.models.model_config import dflash_preference_key, dspark_preference_key info = model_info(repo, token = hf_token, files_metadata = True) total = 0 dspark_candidates: list[tuple[str, int]] = [] + dflash_candidates: list[tuple[str, int]] = [] for sibling in info.siblings or []: name = sibling.rfilename or "" base = Path(name).name.lower() @@ -5306,10 +5332,14 @@ def _remote_gguf_companion_bytes( total += getattr(sibling, "size", 0) or 0 if include_dspark and _is_dspark_drafter_path(name): dspark_candidates.append((name, getattr(sibling, "size", 0) or 0)) + if include_dflash and _is_dflash_drafter_path(name): + dflash_candidates.append((name, getattr(sibling, "size", 0) or 0)) if dspark_candidates: # Same preference order the download uses, so the budget sizes the # file the launch will actually fetch. total += min(dspark_candidates, key = lambda c: dspark_preference_key(c[0]))[1] + if dflash_candidates: + total += min(dflash_candidates, key = lambda c: dflash_preference_key(c[0]))[1] return total except Exception as e: logger.warning(f"Could not size GGUF companions for {repo}: {e}") @@ -5498,6 +5528,7 @@ def _estimate_gguf_required_gb( try: from core.inference.llama_cpp import ( _canonicalize_spec_mode, + _extra_args_requests_dflash, _extra_args_requests_dspark, ) @@ -5528,11 +5559,34 @@ def _estimate_gguf_required_gb( _dspark_capable and (_forced_dspark or (_auto_dspark and getattr(config, "gguf_dspark_file", None))) ) + # DFlash: same shape as DSpark above, and Auto sizes it for the same + # reason. The sidecar is ~1.5 GiB rather than ~11 GB, but a guard that + # protects a running training job still has to charge for it. + _forced_dflash = bool( + _spec_mode == "dflash" or _extra_args_requests_dflash(llama_extra_args, env = {}) + ) + _auto_dflash = _spec_mode == "auto" + _dflash_capable = True + if _forced_dflash or _auto_dflash: + try: + _dflash_capable = bool( + LlamaCppBackend.probe_server_capabilities().get("supports_dflash") + ) + except Exception: + pass + # DSpark keeps first refusal under Auto, mirroring the loader. + dflash_requested = bool( + _dflash_capable + and not dspark_requested + and (_forced_dflash or (_auto_dflash and getattr(config, "gguf_dflash_file", None))) + ) # Forced DSpark on a binary that cannot run it falls back to --spec-default, # which loads no drafter at all, so charging the MTP one would refuse a load # that fits. Auto is different: it falls through to the MTP branch, and keeps # its charge. - _charge_no_drafter = _forced_dspark and not _dspark_capable + _charge_no_drafter = (_forced_dspark and not _dspark_capable) or ( + _forced_dflash and not _dflash_capable + ) total_bytes = 0 main = getattr(config, "gguf_file", None) if main and Path(main).is_file(): @@ -5542,7 +5596,12 @@ def _estimate_gguf_required_gb( # for a load that never opens it. _sized_attrs = ["gguf_mmproj_file"] if not _charge_no_drafter: - _sized_attrs.append("gguf_dspark_file" if dspark_requested else "gguf_mtp_file") + if dspark_requested: + _sized_attrs.append("gguf_dspark_file") + elif dflash_requested: + _sized_attrs.append("gguf_dflash_file") + else: + _sized_attrs.append("gguf_mtp_file") for attr in _sized_attrs: f = getattr(config, attr, None) if f and Path(f).is_file(): @@ -5583,8 +5642,12 @@ def _estimate_gguf_required_gb( # listing. Under Auto size both: a repo has one kind or the other, # the absent one contributes 0, and over-estimating is the safe # direction for a guard that protects a running training job. - include_mtp = (not _charge_no_drafter and (_auto_dspark or not dspark_requested)), + include_mtp = ( + not _charge_no_drafter + and (_auto_dspark or not (dspark_requested or dflash_requested)) + ), include_dspark = (_dspark_capable and (_auto_dspark or dspark_requested)), + include_dflash = (_dflash_capable and (_auto_dflash or dflash_requested)), ) total_gb = (main_bytes + companions) / (1024**3) # remote dims are unreadable; only the kq mask, linear in ubatch x ctx, can be sized here @@ -6006,12 +6069,19 @@ def _resolve_gguf_load_intent( True, log_native_fallback = True, ) + if config.gguf_dflash_file: + config.gguf_dflash_file = _dflash_draft_for_path( + config.gguf_file, + True, + log_native_fallback = True, + ) source = GgufLoadIntent( model_identifier = config.identifier, gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, mtp_draft_path = config.gguf_mtp_file, dspark_draft_path = config.gguf_dspark_file, + dflash_draft_path = config.gguf_dflash_file, hf_variant = config.gguf_variant, ) @@ -7009,6 +7079,7 @@ def _resolve_config(): gguf_intent, mtp_draft_path = _mtp_draft_for_path(llama_backend.gguf_path, False), dspark_draft_path = _dspark_draft_for_path(llama_backend.gguf_path, False), + dflash_draft_path = _dflash_draft_for_path(llama_backend.gguf_path, False), compare_mtp_draft = True, ) _effective_tensor = _effective_tensor_parallel( diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index b066f945698..07de5903ae1 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Separate-file drafter contracts: MTP (Gemma 4) and DSpark (DeepSeek V4 Flash). +"""Separate-file drafter contracts: MTP (Gemma 4), DSpark and DFlash. Pins: the drafter-path predicate and its two layering mirrors, Gemma effective-size extraction, companion classification in variant plans @@ -36,6 +36,7 @@ _is_mtp_drafter, _local_gguf_companion_search_root, detect_gguf_model, + detect_dflash_file, detect_dspark_file, detect_mtp_file, extract_model_size_b, @@ -1565,3 +1566,521 @@ def test_deleting_the_last_variant_still_reclaims_mtp_and_mmproj(tmp_path): assert not (snap / "model-Q4_K_M.gguf").is_symlink() assert not (snap / "mtp-model.gguf").is_symlink() assert not (snap / "mmproj-F16.gguf").is_symlink() + + +# ── DFlash: predicate, discovery, capability gate and emission ─────── +# +# DFlash is the third separate-file drafter kind. It differs from DSpark in two +# ways that these tests pin: +# * it is ON under Auto rather than opt-in, because the published sidecar is +# ~1.5 GiB and ships in the model's own GGUF repo (DSpark's is ~11 GB); +# * its published filename (``dflash-kquant.gguf``) names no model family, so +# discovery confirms the header's ``general.architecture`` instead of +# pairing on the filename. + + +DFLASH_PREDICATE_CASES = [ + # The published sidecar, and the family-named scheme ggml-org uses. + ("dflash-kquant.gguf", True), + ("dflash-Qwen3.6-27B-BF16.gguf", True), + ("dflash-draft-3.6-q8_0.gguf", True), + ("DFLASH-Qwen3.6-27B-Q8_0.gguf", True), + # Adversarial: dflash is also a family a publisher puts on real weights. + ("Qwen3.6-35B-A3B-DFlash-Q4_K_M.gguf", False), + ("qwen35-4b-dflash-Q8_0.gguf", False), + ("laguna-s-2.1-dflash-Q4_K_M.gguf", False), + # A user's own dflash/ folder holds whatever they downloaded, so unlike + # dspark/ and MTP/ the DIRECTORY is not a drafter marker (_DRAFTER_DIR_KINDS). + ("dflash/Qwen3.6-35B-A3B-DFlash-Q4_K_M.gguf", False), + ("foo/dflash/bar.gguf", False), + # The other kinds must not leak into this one: each needs its own --spec-type. + ("dspark-DeepSeek-V4-Flash-0731-Q8_0.gguf", False), + ("mtp-gemma-4-12b-it.gguf", False), + ("dflash-notes.txt", False), +] + + +@pytest.mark.parametrize("path,expected", DFLASH_PREDICATE_CASES) +def test_is_dflash_drafter_path(path, expected): + from core.inference.llama_cpp import ( + _is_dflash_drafter_path, + _is_dspark_drafter_path, + _is_mtp_only_drafter_path, + ) + + assert _is_dflash_drafter_path(path) is expected + if expected: + # The three kinds partition: a DFlash sidecar launched as MTP or DSpark + # would get a --spec-type its architecture cannot serve. + assert _is_dspark_drafter_path(path) is False + assert _is_mtp_only_drafter_path(path) is False + + +@pytest.mark.parametrize( + "value,expected", + [ + ("dflash", "dflash"), + ("DFlash", "dflash"), + (" dflash ", "dflash"), + ("draft-dflash", "dflash"), + ("DRAFT-DFLASH", "dflash"), + ], +) +def test_canonicalize_spec_mode_accepts_dflash(value, expected): + from core.inference.llama_cpp import _canonicalize_spec_mode + + assert _canonicalize_spec_mode(value) == expected + + +# ── Capability probe ───────────────────────────────────────────────── + +_NEEDS_BASH = pytest.mark.skipif( + sys.platform == "win32", + reason = "fake llama-server is a bash stub; Windows has no direct executor", +) + + +def _fake_llama_server(path: Path, help_text: str) -> Path: + path.write_text(f"#!/usr/bin/env bash\ncat <<'EOF'\n{help_text}\nEOF\n") + path.chmod(0o755) + return path + + +@_NEEDS_BASH +@pytest.mark.parametrize( + "spec_line,expected", + [ + ("--spec-type none,draft-mtp,draft-dflash,draft-dspark,ngram-mod", True), + # A published prebuilt that predates the arch: emitting draft-dflash + # would abort the launch instead of falling back. + ("--spec-type none,draft-mtp,draft-dspark,ngram-mod", False), + # Word boundaries, so a longer token cannot be read as support. + ("--spec-type none,draft-dflash2,ngram-mod", False), + ("--spec-type none,xdraft-dflash,ngram-mod", False), + ], +) +def test_probe_server_capabilities_reports_dflash(tmp_path, spec_line, expected): + from core.inference.llama_cpp import LlamaCppBackend + + fake = _fake_llama_server(tmp_path / "llama-server", spec_line) + LlamaCppBackend._capability_cache.clear() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["supports_dflash"] is expected + # DSpark's answer is read from the same block and must not move. + assert caps["supports_dspark"] is ("draft-dspark" in spec_line) + + +def test_missing_binary_reports_no_dflash(): + """The not-found dict is returned before any parsing, so it has to carry the + key: a caller reading it with .get() would otherwise treat "absent" as False + only by luck, and _estimate_gguf_required_gb default-denies on it.""" + from core.inference.llama_cpp import LlamaCppBackend + + caps = LlamaCppBackend.probe_server_capabilities("/nonexistent/llama-server") + assert caps["found"] is False + assert caps["supports_dflash"] is False + + +# ── Emission ───────────────────────────────────────────────────────── + + +def _spec_backend(monkeypatch, *, supports_dflash = True, supports_dspark = True): + from core.inference.llama_cpp import LlamaCppBackend + + caps = { + "found": True, + "mtp_token": "draft-mtp", + "supports_mtp": True, + "supports_dspark": supports_dspark, + "supports_dflash": supports_dflash, + "mtp_probe_inconclusive": False, + "ngram_mod_flavor": "new", + "supports_ngram_mod": True, + "spec_draft_n_max_flag": "--spec-draft-n-max", + } + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: caps), + ) + backend = LlamaCppBackend() + backend._nextn_predict_layers = None + return backend + + +def _spec_flags(backend, **kwargs): + base = dict( + speculative_type = None, + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/Muse-Glimmer-30B-GGUF", + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + base.update(kwargs) + return backend._build_speculative_flags(**base) + + +def test_auto_launches_dflash_when_a_sidecar_is_present(monkeypatch): + """The headline behaviour: unlike DSpark, DFlash needs no opt-in. On a + B200 with the published 1.52 GiB sidecar this is 1.21x-1.36x decode over + spec-off at n_max=2, at 61-77% draft acceptance.""" + backend = _spec_backend(monkeypatch) + flags = _spec_flags(backend, speculative_type = "auto", dflash_draft_path = "/m/dflash-kquant.gguf") + + assert flags == [ + "--model-draft", + "/m/dflash-kquant.gguf", + "--spec-type", + "draft-dflash", + "--spec-draft-n-max", + "2", + ] + assert backend._speculative_type == "draft-dflash" + assert backend._spec_drafter_kind == "dflash" + assert backend._spec_fallback_reason is None + + +def test_auto_uses_the_cpu_draft_depth_off_gpu(monkeypatch): + backend = _spec_backend(monkeypatch) + flags = _spec_flags( + backend, speculative_type = "auto", dflash_draft_path = "/m/d.gguf", gpus = False + ) + assert flags[-2:] == ["--spec-draft-n-max", "3"] + + +def test_a_user_draft_depth_override_reaches_dflash(monkeypatch): + backend = _spec_backend(monkeypatch) + flags = _spec_flags( + backend, speculative_type = "dflash", dflash_draft_path = "/m/d.gguf", spec_draft_n_max = 6 + ) + assert flags[-2:] == ["--spec-draft-n-max", "6"] + assert backend._spec_draft_n_max == 6 + + +def test_auto_does_not_emit_dflash_when_the_binary_cannot_run_it(monkeypatch): + """A --spec-type the binary does not know aborts the launch, so the sidecar + being on disk is not enough. Published prebuilts predate the arch.""" + backend = _spec_backend(monkeypatch, supports_dflash = False) + flags = _spec_flags(backend, speculative_type = "auto", dflash_draft_path = "/m/dflash-kquant.gguf") + + assert "draft-dflash" not in flags + assert "--model-draft" not in flags + assert flags == ["--spec-default"] + assert backend._speculative_type == "default" + + +def test_forced_dflash_without_the_capability_falls_back(monkeypatch): + backend = _spec_backend(monkeypatch, supports_dflash = False) + flags = _spec_flags(backend, speculative_type = "dflash", dflash_draft_path = "/m/d.gguf") + + assert flags == ["--spec-default"] + assert backend._speculative_type == "default" + assert backend._spec_fallback_reason == "binary_no_mtp" + + +def test_forced_dflash_without_a_sidecar_falls_back(monkeypatch): + backend = _spec_backend(monkeypatch) + flags = _spec_flags(backend, speculative_type = "dflash", dflash_draft_path = None) + + assert flags == ["--spec-default"] + assert backend._spec_fallback_reason == "drafter_not_found" + assert backend._spec_drafter_kind == "dflash" + + +def test_dspark_keeps_first_refusal_when_a_repo_ships_both(monkeypatch): + """Mirrors llama.cpp's own downloader, which ranks dspark ahead of dflash. + In practice a repo ships one kind or neither; this pins that adding DFlash + did not reorder the existing choice.""" + backend = _spec_backend(monkeypatch) + flags = _spec_flags( + backend, + speculative_type = "auto", + dspark_draft_path = "/m/dspark-x-Q8_0.gguf", + dflash_draft_path = "/m/dflash-kquant.gguf", + ) + assert "draft-dspark" in flags + assert "draft-dflash" not in flags + assert backend._spec_drafter_kind == "dspark" + + +def test_dspark_emission_is_unchanged(monkeypatch): + """Regression guard on the behaviour this PR must not touch.""" + backend = _spec_backend(monkeypatch) + flags = _spec_flags( + backend, speculative_type = "dspark", dspark_draft_path = "/m/dspark-x-Q8_0.gguf" + ) + assert flags == [ + "--model-draft", + "/m/dspark-x-Q8_0.gguf", + "--spec-type", + "draft-dspark", + "--spec-draft-n-max", + "3", + ] + assert backend._speculative_type == "draft-dspark" + + +def test_a_dflash_sidecar_alone_does_not_change_the_mtp_or_off_paths(monkeypatch): + """Forced modes stay forced: a sidecar sitting on disk must not promote.""" + backend = _spec_backend(monkeypatch) + assert _spec_flags(backend, speculative_type = "off", dflash_draft_path = "/m/d.gguf") == [] + flags = _spec_flags(backend, speculative_type = "ngram", dflash_draft_path = "/m/d.gguf") + assert "draft-dflash" not in flags + assert flags[:2] == ["--spec-type", "ngram-mod"] + + +# ── Local discovery ────────────────────────────────────────────────── + + +def _write_gguf(path: Path, architecture: str) -> Path: + """A real GGUF header carrying one general.architecture string, which is + what detect_dflash_file confirms.""" + import struct + + key = b"general.architecture" + value = architecture.encode() + blob = struct.pack(" tuple[int, str]: return dspark_precision_rank(name), Path(name).name.lower() +# DFlash publishes the same precision vocabulary (and the published sidecar +# carries no precision token at all, which lands in the catch-all rank), so the +# ordering is shared rather than duplicated. +dflash_precision_rank = dspark_precision_rank + + +def dflash_preference_key(name: str) -> tuple[int, str]: + """Sort key picking the preferred DFlash sidecar by name alone.""" + return dflash_precision_rank(name), Path(name).name.lower() + + def detect_mtp_file( path: str, search_root: Optional[str] = None, @@ -2042,6 +2053,115 @@ def _rank(candidate: Path) -> tuple[int, int, int, str]: return None +def detect_dflash_file( + path: str, + search_root: Optional[str] = None, + accept: Optional[Callable[[str], bool]] = None, +) -> Optional[str]: + """Find a DFlash sidecar for a local GGUF model. + + Two things differ from detect_dspark_file, both forced by how DFlash is + published: + + 1. Root level only. ``dspark/`` is always a publisher's companion folder, so + that scan is safe; ``dflash/`` is a family name a user picks for real + weights (the reason llama_cpp._DRAFTER_DIR_KINDS leaves it out), so + reaching into it would launch a weight copy as --model-draft. + 2. No filename pairing. The published sidecar is ``dflash-kquant.gguf``, + which names no model family at all, so _drafter_matches_weight would + reject the one file this exists to find. The header is checked instead: + a DFlash sidecar declares ``general.architecture = dflash``, which no + real weight does, and that is a stronger signal than a filename. It also + settles the adversarial case on its own, since a model merely CALLED + DFlash (``Qwen3.6-35B-A3B-DFlash-Q4_K_M.gguf``) reports its own + architecture. + + A sidecar that does name a family (``dflash-Qwen3.6-27B-BF16.gguf``, the + scheme ggml-org uses) still wins over an unnamed one for the weight it + matches, so a multi-model folder attaches the specific sidecar first. + + ``accept`` filters candidates in preference order, so a caller with extra + rules (a native lease) keeps scanning instead of treating the first + rejection as no sidecar at all. + """ + + def _rank(candidate: Path) -> tuple[int, int, int, int, str]: + # A sidecar naming THIS weight's family first, then any unpaired one, + # then precision, then total size so a split copy cannot outrank a + # smaller single file, then name for a stable order. + paired = _drafter_matches_weight(candidate.name, weight_name, kind = "dflash") + return ( + 0 if paired else 1, + _drafter_stem_rank(candidate.name, kind = "dflash") if paired else 0, + dflash_precision_rank(candidate.name), + _drafter_total_size(candidate), + candidate.name.lower(), + ) + + p = Path(path) + weight_name = p.name if p.suffix.lower() == ".gguf" else None + start_dir = p.parent if p.is_file() else p + dirs = [start_dir] + if search_root is not None: + dirs.append(Path(search_root)) + + candidates: list[Path] = [] + seen: set[Path] = set() + # dict.fromkeys: search_root is the weight's own parent for a flat layout, + # and scanning it twice doubles the directory reads for nothing. + for root in dict.fromkeys(dirs): + try: + entries = list(root.iterdir()) + except OSError: + continue + for candidate in entries: + lower = candidate.name.lower() + if not lower.endswith(".gguf"): + continue + # Drop the shard suffix first: a split copy under the old scheme is + # -Q8_0-dflash-00001-of-00002.gguf, whose stem does not end + # in -dflash. + stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", Path(lower).stem) + if not (lower.startswith("dflash-") or stem.endswith("-dflash")): + continue + try: + # Collapse a split copy to shard 1 before ranking. + launch = _local_gguf_load_path(candidate) + # is_file() follows the link, so this also drops a dangling + # snapshot symlink and a directory named like a sidecar. Without + # it --model-draft gets a path llama-server cannot open, which + # fails the whole load rather than falling back to no + # speculation (detect_dspark_file guards the same way). + if not (launch.is_file() and _drafter_split_is_complete(launch)): + continue + resolved = launch.resolve() + except OSError: + continue + if resolved in seen: + continue + seen.add(resolved) + candidates.append(launch) + + for candidate in sorted(candidates, key = _rank): + meta = read_gguf_general_metadata(str(candidate)) or {} + if (meta.get("general.architecture") or "").strip().lower() != "dflash": + logger.info( + "detect_dflash_file: dropped %s (architecture %r is not dflash)", + candidate.name, + meta.get("general.architecture"), + ) + continue + try: + launch = _drafter_launch_path(candidate) + except OSError: + continue + if accept is not None and not accept(launch): + continue + logger.info("Detected DFlash drafter: %s", launch) + return launch + return None + + def _registered_custom_model_root(path: str) -> Optional[Path]: try: from storage.studio_db import list_scan_folders @@ -3530,6 +3650,7 @@ class ModelConfig: gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection) gguf_mtp_file: Optional[str] = None # Full path to the separate MTP drafter (local mode) gguf_dspark_file: Optional[str] = None # Full path to a DSpark sidecar (local mode) + gguf_dflash_file: Optional[str] = None # Full path to a DFlash sidecar (local mode) gguf_hf_repo: Optional[str] = ( None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF") ) @@ -3682,6 +3803,7 @@ def from_identifier( if mtp_file: logger.info(f"Detected MTP drafter: {mtp_file}") dspark_file = detect_dspark_file(gguf_file, search_root = companion_root) + dflash_file = detect_dflash_file(gguf_file, search_root = companion_root) return cls( identifier = identifier, @@ -3702,6 +3824,7 @@ def from_identifier( gguf_mmproj_file = mmproj_file, gguf_mtp_file = mtp_file, gguf_dspark_file = dspark_file, + gguf_dflash_file = dflash_file, ) else: # Does the HF repo contain GGUF files? diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py index 8f55ca86146..22b68c0aad9 100644 --- a/studio/backend/utils/openai_auto_switch_settings.py +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -287,18 +287,22 @@ def set_openai_auto_switch( "auto", "mtp", "dspark", + "dflash", "ngram", "mtp+ngram", "off", "default", "draft-mtp", "draft-dspark", + "draft-dflash", "ngram-mod", "ngram-simple", } ) # Only these consume spec_draft_n_max (mirrors DRAFT_N_MAX_SPEC_TYPES in the UI). -DRAFT_N_MAX_SPEC_TYPES = frozenset({"mtp", "mtp+ngram", "draft-mtp", "dspark", "draft-dspark"}) +DRAFT_N_MAX_SPEC_TYPES = frozenset( + {"mtp", "mtp+ngram", "draft-mtp", "dspark", "draft-dspark", "dflash", "draft-dflash"} +) VALID_GPU_MEMORY_MODES = frozenset({"auto", "manual"}) # Mirrors MLX_KV_BITS_CHOICES in core/inference/mlx_inference.py; a set, not a range. VALID_MLX_KV_BITS = frozenset({8, 6, 5, 4, 3, 2}) @@ -375,8 +379,8 @@ def normalize_model_override(payload: dict[str, Any]) -> dict[str, Any]: speculative_type = _clean_str(payload.get("speculative_type"), VALID_SPECULATIVE_TYPES) if speculative_type: entry["speculative_type"] = speculative_type - # Only the modes that launch a drafter with a configurable depth (MTP - # and DSpark); storing it otherwise shows an edit the loader ignores. + # Only the modes that launch a drafter with a configurable depth (MTP, + # DSpark and DFlash); storing it otherwise shows an edit the loader ignores. if speculative_type in DRAFT_N_MAX_SPEC_TYPES: spec_draft_n_max = _bounded_int(payload.get("spec_draft_n_max"), minimum = 1, maximum = 16) if spec_draft_n_max: diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 8361bf78b43..0385c76126e 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -418,7 +418,7 @@ function specFallbackMessage({ updateAvailable, }: { reason: string; - drafter: "MTP" | "DSpark"; + drafter: "MTP" | "DSpark" | "DFlash"; isLocalGguf: boolean; updateAvailable: boolean; }): string { @@ -433,6 +433,11 @@ function specFallbackMessage({ ? "No matching DSpark sidecar was found. Place its dspark-*.gguf beside the model or in its dspark folder, then reload the model." : "The DSpark sidecar could not be downloaded, so this model is running without speculative decoding. Check network or Hugging Face access, then reload it."; } + if (drafter === "DFlash") { + return isLocalGguf + ? "No matching DFlash sidecar was found. Place its dflash-*.gguf beside the model, then reload the model." + : "The DFlash sidecar could not be downloaded, so this model is running without speculative decoding. Check network or Hugging Face access, then reload it."; + } return isLocalGguf ? "This local model supports MTP, but no matching drafter file was found. Place its mtp-*.gguf beside the model or in its MTP folder, then reload the model." : "This model supports MTP, but its drafter file could not be downloaded, so MTP is off and it falls back to n-gram speculative decoding where the llama.cpp build supports it. Check your network connection or Hugging Face access, then reload the model to retry the drafter."; @@ -525,8 +530,12 @@ export function ChatSettingsPanel({ // The loaded model's own kind, not the pending control: the notice explains a // fallback that already happened, so a staged edit (or a preset applied without // a reload) must not re-label it and point at the wrong file. - const speculativeDrafterLabel = - (specDrafterKind ?? speculativeType) === "dspark" ? "DSpark" : "MTP"; + const speculativeDrafterLabel: "MTP" | "DSpark" | "DFlash" = + (specDrafterKind ?? speculativeType) === "dspark" + ? "DSpark" + : (specDrafterKind ?? speculativeType) === "dflash" + ? "DFlash" + : "MTP"; const mtpUpdatable = specFallbackReason === "binary_no_mtp" || specFallbackReason === "binary_outdated"; @@ -561,7 +570,8 @@ export function ChatSettingsPanel({ (speculativeType === "auto" || speculativeType === "mtp" || speculativeType === "mtp+ngram" || - speculativeType === "dspark"); + speculativeType === "dspark" || + speculativeType === "dflash"); const showContextVramWarning = !isExternalModel && isGguf && diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 34f62e3a938..13cc7f4148e 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -94,7 +94,7 @@ export const CHAT_SPECULATIVE_TYPE_KEY = "unsloth_chat_speculative_type"; export const CHAT_GPU_MEMORY_MODE_KEY = "unsloth_chat_gpu_memory_mode"; // Persist only the model-agnostic intents (auto/ngram/off). The model-specific -// drafter modes (mtp/mtp+ngram/dspark) and spec_draft_n_max stay session-only: +// drafter modes (mtp/mtp+ngram/dspark/dflash) and spec_draft_n_max stay session-only: // a persisted choice would silently no-op on a model with no MTP head or no // DSpark sidecar. Unknown -> auto. const PERSISTED_SPEC_MODES = new Set(["auto", "ngram", "off"]); @@ -493,7 +493,7 @@ function saveString(key: string, value: string): void { } // Canonicalises any backend value onto the Speculative Decoding dropdown's -// modes ("auto"/"mtp"/"ngram"/"mtp+ngram"/"off"/null). Backend-only +// modes ("auto"/"mtp"/"dspark"/"dflash"/"ngram"/"mtp+ngram"/"off"/null). Backend-only // legacy aliases map to their closest UI mode. export function normalizeSpeculativeType( v: string | null | undefined, @@ -505,6 +505,7 @@ export function normalizeSpeculativeType( if (s === "off") return "off"; if (s === "mtp" || s === "draft-mtp") return "mtp"; if (s === "dspark" || s === "draft-dspark") return "dspark"; + if (s === "dflash" || s === "draft-dflash") return "dflash"; if (s === "ngram" || s === "ngram-mod" || s === "ngram-simple") { return "ngram"; } @@ -1071,8 +1072,8 @@ type ChatRuntimeStore = { */ specFallbackReason: string | null; /** - * Which drafter the loaded model's speculative resolution was about, "mtp" or - * "dspark". Paired with specFallbackReason: the reason alone cannot name the + * Which drafter the loaded model's speculative resolution was about: "mtp", + * "dspark" or "dflash". Paired with specFallbackReason: the reason alone cannot name the * file to fix, since Auto resolves the kind server-side and the requested mode * still reads "auto". */ diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index cd84d8678b9..4ed959bc9f3 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -57,16 +57,17 @@ export interface LoadModelRequest { mlx_kv_bits?: number | null; /** * Speculative decoding mode for GGUF models. Canonical values: "auto" - * (platform-aware: MTP on MTP GGUFs, ngram-mod fallback for sub-3B), "mtp" - * (force draft-mtp), "dspark" (force draft-dspark with a sidecar), - * "ngram" (force ngram-mod), "mtp+ngram" (ngram-mod + - * draft-mtp chain), "off". Legacy "default"/"draft-mtp"/"ngram-mod"/ - * "ngram-simple" are still accepted by the backend. + * (platform-aware: DSpark or DFlash when the model ships that sidecar, else + * MTP on MTP GGUFs, ngram-mod fallback for sub-3B), "mtp" (force draft-mtp), + * "dspark" (force draft-dspark with a sidecar), "dflash" (force draft-dflash + * with a sidecar), "ngram" (force ngram-mod), "mtp+ngram" (ngram-mod + + * draft-mtp chain), "off". Legacy "default"/"draft-mtp"/"draft-dspark"/ + * "draft-dflash"/"ngram-mod"/"ngram-simple" are still accepted by the backend. */ speculative_type?: string | null; /** - * Override --spec-draft-n-max for MTP speculative decoding. Applied only - * when speculative_type resolves to "mtp", "mtp+ngram", or "dspark". + * Override --spec-draft-n-max for drafter speculative decoding. Applied only + * when speculative_type resolves to "mtp", "mtp+ngram", "dspark" or "dflash". */ spec_draft_n_max?: number | null; /** @@ -345,14 +346,14 @@ export interface InferenceStatusResponse { * Why a speculative drafter was disabled despite being requested. * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable * it; "runtime_error" -> the current build could not run it; - * "drafter_not_found" -> its MTP or DSpark sidecar was unavailable; + * "drafter_not_found" -> its MTP, DSpark or DFlash sidecar was unavailable; * "mla_mtp_disabled" -> an Auto-mode policy downgrade for MLA models * (GLM-5.2 et al.) whose llama.cpp MTP path is slower than no speculation * (updating won't help; choose MTP in Settings to force it). Null otherwise. */ /** - * Which drafter the resolution was about, "mtp" or "dspark". Auto resolves the - * kind itself, so speculative_type still reads "auto", and a fallback leaves + * Which drafter the resolution was about: "mtp", "dspark" or "dflash". Auto + * resolves the kind itself, so speculative_type still reads "auto", and a fallback leaves * the engaged type at "default": neither names the file to fix. */ spec_drafter_kind?: string | null; diff --git a/studio/frontend/src/features/model-picker/components/model-config-page.tsx b/studio/frontend/src/features/model-picker/components/model-config-page.tsx index 7ad168062fd..99e4a49224b 100644 --- a/studio/frontend/src/features/model-picker/components/model-config-page.tsx +++ b/studio/frontend/src/features/model-picker/components/model-config-page.tsx @@ -103,6 +103,7 @@ const SPECULATIVE_TYPE_LABELS: Record< auto: "Auto", mtp: "MTP", dspark: "DSpark", + dflash: "DFlash", ngram: "Ngram", "mtp+ngram": "MTP+Ngram", off: "Off", diff --git a/studio/frontend/src/features/model-picker/model-config/per-model-config.ts b/studio/frontend/src/features/model-picker/model-config/per-model-config.ts index c7165409403..7bea44f5cbd 100644 --- a/studio/frontend/src/features/model-picker/model-config/per-model-config.ts +++ b/studio/frontend/src/features/model-picker/model-config/per-model-config.ts @@ -231,6 +231,9 @@ function canonicalizeSpeculativeType(value: string): string | null { if (s === "dspark" || s === "draft-dspark") { return "dspark"; } + if (s === "dflash" || s === "draft-dflash") { + return "dflash"; + } if (s === "ngram" || s === "ngram-mod" || s === "ngram-simple") { return "ngram"; } diff --git a/studio/frontend/src/lib/speculative-modes.ts b/studio/frontend/src/lib/speculative-modes.ts index 75d780ba8cb..85ca49753f8 100644 --- a/studio/frontend/src/lib/speculative-modes.ts +++ b/studio/frontend/src/lib/speculative-modes.ts @@ -15,6 +15,7 @@ export const SPECULATIVE_TYPES = [ "auto", "mtp", "dspark", + "dflash", "ngram", "mtp+ngram", "off", @@ -23,11 +24,12 @@ export const SPECULATIVE_TYPES = [ /** * The modes that consume spec_draft_n_max, i.e. the ones that launch a drafter * with a configurable depth. Named for the setting rather than for MTP: DSpark - * is in here too. Mirrors DRAFT_N_MAX_SPEC_TYPES in + * and DFlash are in here too. Mirrors DRAFT_N_MAX_SPEC_TYPES in * studio/backend/utils/openai_auto_switch_settings.py. */ export const DRAFT_N_MAX_SPEC_TYPES: ReadonlySet = new Set([ "mtp", "mtp+ngram", "dspark", + "dflash", ]); From 6c5485ada3b569a440c33c7c791584bd5623d2f0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:06:58 +0000 Subject: [PATCH 02/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/test_mtp_drafter_companion.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 07de5903ae1..3f2f130c42a 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -1607,7 +1607,6 @@ def test_is_dflash_drafter_path(path, expected): _is_dspark_drafter_path, _is_mtp_only_drafter_path, ) - assert _is_dflash_drafter_path(path) is expected if expected: # The three kinds partition: a DFlash sidecar launched as MTP or DSpark @@ -1628,7 +1627,6 @@ def test_is_dflash_drafter_path(path, expected): ) def test_canonicalize_spec_mode_accepts_dflash(value, expected): from core.inference.llama_cpp import _canonicalize_spec_mode - assert _canonicalize_spec_mode(value) == expected @@ -1684,7 +1682,12 @@ def test_missing_binary_reports_no_dflash(): # ── Emission ───────────────────────────────────────────────────────── -def _spec_backend(monkeypatch, *, supports_dflash = True, supports_dspark = True): +def _spec_backend( + monkeypatch, + *, + supports_dflash = True, + supports_dspark = True, +): from core.inference.llama_cpp import LlamaCppBackend caps = { @@ -1744,9 +1747,7 @@ def test_auto_launches_dflash_when_a_sidecar_is_present(monkeypatch): def test_auto_uses_the_cpu_draft_depth_off_gpu(monkeypatch): backend = _spec_backend(monkeypatch) - flags = _spec_flags( - backend, speculative_type = "auto", dflash_draft_path = "/m/d.gguf", gpus = False - ) + flags = _spec_flags(backend, speculative_type = "auto", dflash_draft_path = "/m/d.gguf", gpus = False) assert flags[-2:] == ["--spec-draft-n-max", "3"] @@ -1928,7 +1929,12 @@ def test_model_config_reports_a_local_dflash_sidecar(tmp_path): # ── Download gating ────────────────────────────────────────────────── -def _dflash_download_probe(monkeypatch, *, supports_dflash, cached = None): +def _dflash_download_probe( + monkeypatch, + *, + supports_dflash, + cached = None, +): import core.inference.llama_cpp as llama_cpp_module from core.inference.llama_cpp import LlamaCppBackend From 16646037e992826957f5d3d50eca71e2167e3360 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 12:23:24 +0000 Subject: [PATCH 03/26] Harden the DFlash drafter fallback, pairing and dedupe Four fixes from review of the auto-launch path. Strip user-supplied DFlash args on the drafterless retry. The gate that enters the retry counts a DFlash request, but the cleanup only recognised MTP and DSpark. llama.cpp accumulates speculative types, so prepending --spec-default while the DFlash group survived relaunched the drafter that had just failed, and a main model that loads fine without it was lost instead of recovered. Skip a DFlash sidecar that names another weight in the same folder. _drafter_matches_weight is False both for a sidecar naming no family and for one naming a different family, so ranking put them in one bucket and precision could float the foreign one to the top: loading model B beside dflash-model-A-Q8_0.gguf and dflash-kquant.gguf launched model A's drafter. Both files carry a real dflash header, so the architecture check behind the ranking cannot catch it. The decision is made against the weights actually present in the folder rather than by guessing which stems are precision tokens, which keeps the published unpaired sidecar eligible. Stand the Auto DFlash fetch down once DSpark has resolved. DSpark takes first refusal in the promotion, so for a repo shipping both kinds the DFlash sidecar could never launch and the fetch spent bandwidth and cache on a file that would not be used. An explicit dflash request still fetches. Keep Auto deduplicated after a failed DFlash drafter. _speculative_type is reset to "default" by a successful drafterless retry while the launch still records the resolved sidecar, so the next Apply compared the intent's empty MTP path against it and reloaded a healthy server. _spec_drafter_kind survives the fallback and now decides the comparison. test_mtp_drafter_companion.py, test_native_gguf_companion.py, test_llama_cpp_mtp_detection.py and test_resolve_quant_gguf.py: 489 passed, including two new tests for the foreign-sidecar case and for the paired sidecar still winning. --- studio/backend/core/inference/llama_cpp.py | 33 ++++++++++++++++--- .../tests/test_mtp_drafter_companion.py | 29 ++++++++++++++++ studio/backend/utils/models/model_config.py | 32 ++++++++++++++++++ 3 files changed, 89 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 468fe0977e4..9795e8b4ca2 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4046,11 +4046,25 @@ def _norm(value): # Auto counts as dspark/dflash once it resolved that way: the launch # stored that sidecar, so comparing the MTP field (None for these # repos) against it would reload a healthy server on every Apply. + # _speculative_type is reset to "default" when a drafter fails to + # start and the drafterless retry succeeds, but the launch still + # recorded the sidecar it resolved. Keying the comparison off it + # would then compare the intent's (empty) MTP path against that + # recorded sidecar and reload a healthy server on every Apply. + # _spec_drafter_kind survives the fallback, so it decides here. _compare_dspark = speculative_type == "dspark" or ( - speculative_type == "auto" and self._speculative_type == "draft-dspark" + speculative_type == "auto" + and ( + self._speculative_type == "draft-dspark" + or self._spec_drafter_kind == "dspark" + ) ) _compare_dflash = speculative_type == "dflash" or ( - speculative_type == "auto" and self._speculative_type == "draft-dflash" + speculative_type == "auto" + and ( + self._speculative_type == "draft-dflash" + or self._spec_drafter_kind == "dflash" + ) ) if _compare_dspark: intent_draft = intent.dspark_draft_path @@ -9773,9 +9787,16 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # and ships in the model's own GGUF repo, so under Auto it # costs about what the MTP drafter costs. Repos without one # no-op after a single listing. + # Under Auto, DSpark takes first refusal in the promotion + # below, so a repo that ships both kinds would never launch + # the DFlash sidecar. Fetching it anyway costs ~1.5 GiB of + # bandwidth and cache for a file that cannot be used, so the + # Auto fetch stands down once DSpark has resolved. An + # explicit "dflash" request still fetches. if ( not dflash_draft_path and _spec_canon in ("auto", "dflash") + and not (_spec_canon == "auto" and dspark_draft_path) and not _extra_args_set_spec_type(extra_args) ): dflash_draft_path = self._download_dflash( @@ -12583,9 +12604,11 @@ def _try_auto_vulkan_cpu_fallback( # failed and lose a main model that loads fine without it. The # tail loses its spec group and the env goes with it, the same way # the crash replay does it. - if _extra_args_requests_mtp( - extra_args, env = _launch_spec_env - ) or _extra_args_requests_dspark(extra_args, env = _launch_spec_env): + if ( + _extra_args_requests_mtp(extra_args, env = _launch_spec_env) + or _extra_args_requests_dspark(extra_args, env = _launch_spec_env) + or _extra_args_requests_dflash(extra_args, env = _launch_spec_env) + ): _fb_tail = strip_shadowing_flags( _fb_tail, strip_context = False, diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 3f2f130c42a..ab142637924 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2090,3 +2090,32 @@ def test_dflash_stays_unreclaimable_even_though_auto_now_launches_it(tmp_path): assert not (snap / "model-Q4_K_M.gguf").is_symlink() assert (snap / "dflash-kquant.gguf").is_symlink() assert not (snap / "dspark-model-Q8_0.gguf").is_symlink() + + +def test_detect_dflash_file_skips_a_sidecar_named_for_another_weight(tmp_path): + """A multi-model folder must not attach a foreign drafter. + + _drafter_matches_weight is False both for a sidecar naming no family and for + one naming a DIFFERENT family, so ranking alone bucketed them together and + precision could float the foreign one to the top: loading model B beside + dflash-model-A-Q8_0.gguf and the generic dflash-kquant.gguf launched model + A's drafter for model B. Both files carry a real dflash header, so the + architecture check behind the ranking cannot catch this one. + """ + weight = _write_gguf(tmp_path / "Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "muse-glimmer") + _write_gguf(tmp_path / "Qwen3.6-27B-Q4_K_M.gguf", "qwen3") + foreign = _write_gguf(tmp_path / "dflash-Qwen3.6-27B-Q8_0.gguf", "dflash") + generic = _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") + + assert foreign.exists() + assert detect_dflash_file(str(weight)) == str(generic.resolve()) + + +def test_detect_dflash_file_still_prefers_a_sidecar_that_names_this_weight(tmp_path): + """The skip above must not cost the paired case: a sidecar naming THIS + weight's family still wins over the generic one.""" + weight = _write_gguf(tmp_path / "Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "muse-glimmer") + paired = _write_gguf(tmp_path / "dflash-Muse-Glimmer-30B-Q8_0.gguf", "dflash") + _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") + + assert detect_dflash_file(str(weight)) == str(paired.resolve()) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index b88e1977013..244738b5768 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2106,6 +2106,7 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: dirs.append(Path(search_root)) candidates: list[Path] = [] + other_weights: list[str] = [] seen: set[Path] = set() # dict.fromkeys: search_root is the weight's own parent for a flat layout, # and scanning it twice doubles the directory reads for nothing. @@ -2123,6 +2124,10 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: # in -dflash. stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", Path(lower).stem) if not (lower.startswith("dflash-") or stem.endswith("-dflash")): + # Every other GGUF in the folder is a weight some sidecar could + # be naming. Recorded so a sidecar belonging to a NEIGHBOUR can + # be told apart from one naming no family at all (below). + other_weights.append(candidate.name) continue try: # Collapse a split copy to shard 1 before ranking. @@ -2142,6 +2147,33 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: seen.add(resolved) candidates.append(launch) + # A sidecar naming a family that belongs to a NEIGHBOUR weight is that + # neighbour's drafter, not a generic one. _drafter_matches_weight is False + # both for it and for a sidecar naming no family (dflash-kquant.gguf), so + # ranking alone bucketed the two together and precision could float the + # foreign one to the top: loading model B beside dflash-model-A-Q8_0.gguf + # and dflash-kquant.gguf launched model A's drafter for model B. Both carry + # a real dflash header, so the architecture check behind the ranking cannot + # catch it. Deciding against the weights actually present keeps the + # published unpaired sidecar eligible (its stem, "kquant", names no file + # here) without hardcoding which stems are precision tokens. + if weight_name is not None and other_weights: + kept: list[Path] = [] + for candidate in candidates: + if not _drafter_matches_weight( + candidate.name, weight_name, kind = "dflash" + ) and any( + _drafter_matches_weight(candidate.name, other, kind = "dflash") + for other in other_weights + ): + logger.info( + "detect_dflash_file: dropped %s (names another weight in this folder)", + candidate.name, + ) + continue + kept.append(candidate) + candidates = kept + for candidate in sorted(candidates, key = _rank): meta = read_gguf_general_metadata(str(candidate)) or {} if (meta.get("general.architecture") or "").strip().lower() != "dflash": From 7d474e9d84cdc252032edbec00e36717f6bcd90e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:24:38 +0000 Subject: [PATCH 04/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/utils/models/model_config.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 244738b5768..20898663aa3 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2160,9 +2160,7 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: if weight_name is not None and other_weights: kept: list[Path] = [] for candidate in candidates: - if not _drafter_matches_weight( - candidate.name, weight_name, kind = "dflash" - ) and any( + if not _drafter_matches_weight(candidate.name, weight_name, kind = "dflash") and any( _drafter_matches_weight(candidate.name, other, kind = "dflash") for other in other_weights ): From 3984978976efd1538752205cc78a06cb4410d513 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 12:34:15 +0000 Subject: [PATCH 05/26] Keep DFlash discovery, the training guard and the hints in step Discovery now accepts the dflash- prefix only. The shared companion predicates recognise DFlash by that prefix, so a -dflash.gguf accepted by discovery was also a selectable Q8_0 main model in the quant picker, and choosing that variant handed llama-server the drafter as the target. Teaching the predicate the suffix instead would hide a real model whose name merely ends in DFlash, which is the case #7811 exists to protect, so detection gives the form up rather than the picker giving up a model. No published sidecar uses it; the shipped one is dflash-kquant.gguf. The same mismatch exists for MTP on main and is left alone here. The training VRAM guard now sizes a drafter named through llama_extra_args. Discovery never fills gguf_dflash_file for a file outside the model directory, but load_model still passes that path to llama-server, so a load could be admitted beside a training run while nothing was charged for the sidecar it makes resident. The Speculative Decoding hint said Auto picks DSpark or else MTP / ngram and that everything but DSpark leaves output unchanged. Auto now picks DFlash too, and like DSpark it is not bit-identical on quantized targets. The Draft Tokens hint gained the DFlash default, which shares the MTP branch at 2 on GPU and 3 on CPU/Mac. 514 passed across the drafter, companion, detection, quant-resolution and picker suites, including two new tests pinning the suffix form out of discovery and the prefix form still in. --- studio/backend/routes/inference.py | 11 +++++++ .../tests/test_mtp_drafter_companion.py | 30 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 16 ++++++---- .../components/model-config-page.tsx | 15 +++++----- 4 files changed, 60 insertions(+), 12 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c2859645ce0..41bb43f5491 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5528,6 +5528,7 @@ def _estimate_gguf_required_gb( try: from core.inference.llama_cpp import ( _canonicalize_spec_mode, + _extra_args_mtp_draft_path, _extra_args_requests_dflash, _extra_args_requests_dspark, ) @@ -5602,6 +5603,16 @@ def _estimate_gguf_required_gb( _sized_attrs.append("gguf_dflash_file") else: _sized_attrs.append("gguf_mtp_file") + # A caller that owns speculation through llama_extra_args names the + # drafter with --model-draft, and discovery never populates + # gguf_dflash_file / gguf_dspark_file for a file outside the model + # directory. load_model still hands that path to llama-server, so + # without this the guard admits a load beside a training run while + # charging nothing for the sidecar that load is about to make resident. + _extras_draft = _extra_args_mtp_draft_path(llama_extra_args, env = {}) + if _extras_draft and Path(_extras_draft).is_file(): + total_bytes += LlamaCppBackend._get_gguf_size_bytes(str(_extras_draft)) + for attr in _sized_attrs: f = getattr(config, attr, None) if f and Path(f).is_file(): diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index ab142637924..08dae40c863 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2119,3 +2119,33 @@ def test_detect_dflash_file_still_prefers_a_sidecar_that_names_this_weight(tmp_p _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") assert detect_dflash_file(str(weight)) == str(paired.resolve()) + + +def test_detect_dflash_file_ignores_the_suffix_form_the_picker_cannot_hide(tmp_path): + """Discovery and the quant picker have to agree on what a sidecar is. + + The shared companion predicates know DFlash by the dflash- prefix only, so a + -dflash.gguf accepted here would be a drafter for discovery and at the + same time a selectable Q8_0 main model in the picker, and choosing that + variant would hand llama-server the drafter as the target. Detection gives + the form up rather than teaching the predicate a suffix that would hide a + real model merely named DFlash. + """ + from core.inference.llama_cpp import _is_companion_gguf_path + + weight = _write_gguf(tmp_path / "Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "muse-glimmer") + suffix_form = tmp_path / "Muse-Glimmer-30B-Q8_0-dflash.gguf" + _write_gguf(suffix_form, "dflash") + + assert detect_dflash_file(str(weight)) is None + # The invariant behind the choice: what discovery accepts, the picker hides. + assert _is_companion_gguf_path(suffix_form.name) is False + assert _is_companion_gguf_path("dflash-kquant.gguf") is True + + +def test_dflash_prefix_form_is_still_found_beside_the_weight(tmp_path): + """Dropping the suffix form must not cost the published sidecar.""" + weight = _write_gguf(tmp_path / "Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "muse-glimmer") + sidecar = _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") + + assert detect_dflash_file(str(weight)) == str(sidecar.resolve()) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 20898663aa3..5e86f10bd08 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2119,11 +2119,17 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: lower = candidate.name.lower() if not lower.endswith(".gguf"): continue - # Drop the shard suffix first: a split copy under the old scheme is - # -Q8_0-dflash-00001-of-00002.gguf, whose stem does not end - # in -dflash. - stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", Path(lower).stem) - if not (lower.startswith("dflash-") or stem.endswith("-dflash")): + # Prefix form only, deliberately. The shared companion predicates + # (_drafter_path_kind, is_mtp_drafter_path) know DFlash by the + # dflash- prefix, so accepting -dflash.gguf here would let one + # file be a drafter for discovery AND a selectable Q8_0 main model in + # the quant picker, and choosing that variant would hand llama-server + # the drafter as the target. Teaching the predicate the suffix + # instead would hide a real model whose name merely ends in DFlash, + # which is the case #7811 exists to protect, so detection gives the + # form up rather than the picker giving up a model. No published + # DFlash sidecar uses it; the shipped one is dflash-kquant.gguf. + if not lower.startswith("dflash-"): # Every other GGUF in the folder is a weight some sidecar could # be naming. Recorded so a sidecar belonging to a NEIGHBOUR can # be told apart from one naming no family at all (below). diff --git a/studio/frontend/src/features/model-picker/components/model-config-page.tsx b/studio/frontend/src/features/model-picker/components/model-config-page.tsx index 99e4a49224b..d4ebcd1082b 100644 --- a/studio/frontend/src/features/model-picker/components/model-config-page.tsx +++ b/studio/frontend/src/features/model-picker/components/model-config-page.tsx @@ -704,11 +704,12 @@ function GgufAdvancedSettings({ Speculative Decoding Faster generation. Auto picks the best strategy for the model and - platform: DSpark when the model ships a drafter sidecar, otherwise - MTP / ngram. Pick a strategy to force it, or Off to disable. - DSpark downloads a sidecar of about 11 GB and trades VRAM for speed; - on quantized targets its greedy output can differ from a non - speculative run. MTP and ngram do not change output. + platform: DSpark or DFlash when the model ships a drafter sidecar, + otherwise MTP / ngram. Pick a strategy to force it, or Off to + disable. DSpark downloads a sidecar of about 11 GB and DFlash one of + about 1.5 GB, both trading VRAM for speed; on quantized targets + their greedy output can differ from a non speculative run. MTP and + ngram do not change output. Date: Mon, 10 Aug 2026 12:46:49 +0000 Subject: [PATCH 06/26] Pair the remote DFlash sidecars with the selected weight detect_dflash_file already refuses a sidecar named after a NEIGHBOURING weight, so a folder holding two families cannot attach a foreign drafter locally. The download picker and the offline cache reuse still ranked every dflash-*.gguf by precision and name alone, never comparing a candidate against the weight being loaded, so in a repo hosting more than one family dflash-model-A-Q8_0.gguf outranked the generic dflash-kquant.gguf and model B downloaded and launched model A's drafter. The pairing rule now lives in one place, dflash_repo_preference_key, built on the same _drafter_names_other_weight predicate the local scan uses: a sidecar naming this weight's family first (most specific stem first, as detect_mtp_file does), then one naming no weight present here, then one naming a neighbour. The last is demoted rather than dropped, so a repo whose only sidecar looks foreign still has a fallback. Deciding against the weights actually present is what keeps the published unpaired sidecar eligible: dflash-kquant.gguf has a precision token for a stem, not a family name, so "the stem is non-empty" cannot stand in for "this names another model". Nothing changes for a repo with one sidecar, and with no weight in hand the order is precision only, as before. Tests cover a multi-family repo picking the generic sidecar, the same repo picking the specific one for its own weight, the shipped Muse-Glimmer layout still resolving, and the cached path agreeing with the download path. --- studio/backend/core/inference/llama_cpp.py | 53 ++++-- .../tests/test_mtp_drafter_companion.py | 163 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 70 +++++++- 3 files changed, 268 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 9795e8b4ca2..507718e2928 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8105,14 +8105,23 @@ def _cached_repo_dflash_drafter( hf_repo: str, *, cache_dir: Optional[str] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """The preferred already-cached DFlash sidecar for a repo, Q8_0 first - (dflash_preference_key), so an offline reuse picks the same file the - online download would have fetched.""" + (dflash_repo_preference_key), so an offline reuse picks the same file + the online download would have fetched. + + ``near_path`` is the weight this sidecar would draft for. A repo hosting + more than one family ships more than one sidecar, and precision alone + would hand model B the sidecar named after model A, so the candidates + are ranked against the weight actually being loaded (same rule as + detect_dflash_file). Without it the ordering is precision only, as + before. + """ try: from utils.models.model_config import ( _iter_hf_cache_snapshots, - dflash_preference_key, + dflash_repo_preference_key, ) snapshots = ( @@ -8120,14 +8129,22 @@ def _cached_repo_dflash_drafter( if cache_dir is None else _iter_hf_cache_snapshots(hf_repo, cache_dir) ) - candidates: list[Path] = [] + weight_name = Path(near_path).name if near_path else None + ranked: list[tuple[tuple[int, int, int, str], Path]] = [] for snap in snapshots: - candidates.extend( - snap / name - for name in _gguf_snapshot_files(snap) + names = _gguf_snapshot_files(snap) + # Every non-sidecar GGUF in the snapshot is a weight some sidecar + # could be naming; that is what tells a neighbour's sidecar apart + # from one naming no family at all. + others = [ + Path(name).name for name in names if not _is_dflash_drafter_path(name) + ] + ranked.extend( + (dflash_repo_preference_key(name, weight_name, others), snap / name) + for name in names if _is_dflash_drafter_path(name) ) - for candidate in sorted(candidates, key = lambda p: dflash_preference_key(p.name)): + for _, candidate in sorted(ranked, key = lambda entry: entry[0]): if candidate.is_file(): return str(candidate) except Exception as exc: @@ -8153,11 +8170,26 @@ def _download_dflash( engage. """ + weight_name = Path(near_path).name if near_path else None + def _pick_dflash(candidates: list[str]) -> Optional[str]: - from utils.models.model_config import dflash_preference_key + # Ranked against the weight being loaded, not by precision alone: a + # repo hosting two families ships a sidecar per family, and + # dflash-model-A-Q8_0.gguf beats the generic dflash-kquant.gguf on + # precision, so model B would download and launch model A's drafter. + # The rest of the listing supplies the neighbouring weights that make + # a foreign sidecar recognisable (a sidecar naming no family at all + # stays eligible, which is what the published one does). + from utils.models.model_config import dflash_repo_preference_key + + others = [ + Path(name).name + for name in candidates + if name.lower().endswith(".gguf") and not _is_dflash_drafter_path(name) + ] files = sorted( (name for name in candidates if _is_dflash_drafter_path(name)), - key = dflash_preference_key, + key = lambda name: dflash_repo_preference_key(name, weight_name, others), ) return files[0] if files else None @@ -8166,6 +8198,7 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: cached = self._cached_repo_dflash_drafter( hf_repo, cache_dir = _hub_cache_dir_for_snapshot_path(near_path), + near_path = near_path, ) try: if not self.probe_server_capabilities(binary).get("supports_dflash"): diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 08dae40c863..623eab1ea10 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2065,6 +2065,169 @@ def test_a_cached_dflash_drafter_is_never_launched_as_an_mtp_drafter(tmp_path, m assert b._cached_repo_dflash_drafter("org/repo") == str(snap / "dflash-kquant.gguf") +# ── Remote sidecars pair with the selected weight ──────────────────── +# +# detect_dflash_file already refuses a sidecar named after a NEIGHBOURING +# weight, so a multi-family folder cannot attach a foreign drafter locally. The +# download and the offline cache reuse ranked by precision and name alone, so +# dflash-model-A-Q8_0.gguf beat the generic dflash-kquant.gguf and model B was +# launched with model A's drafter. All three paths now share +# dflash_repo_preference_key. + +_MULTI_FAMILY_LISTING = [ + "model-A-Q4_K_M.gguf", + "model-B-Q4_K_M.gguf", + "dflash-model-A-Q8_0.gguf", + "dflash-kquant.gguf", +] + + +def _dflash_download_pick(monkeypatch, *, listing, near_path): + """The sidecar _download_dflash's picker chooses out of a repo listing.""" + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: {"supports_dflash": True}), + ) + monkeypatch.setattr( + llama_cpp_module, "_companion_snapshot_sibling", lambda near_path, pick: None + ) + picked = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + near_path = None, + outcome = None, + ): + picked["name"] = pick(listing) + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + b._download_dflash( + hf_repo = "org/repo", + near_path = near_path, + binary = "/fake/llama-server", + ) + return picked.get("name") + + +def test_download_dflash_skips_a_sidecar_named_after_another_weight(monkeypatch): + """Model B must get the generic sidecar, not the higher-precision one that + names model A.""" + assert ( + _dflash_download_pick( + monkeypatch, + listing = _MULTI_FAMILY_LISTING, + near_path = "/cache/snap/model-B-Q4_K_M.gguf", + ) + == "dflash-kquant.gguf" + ) + + +def test_download_dflash_takes_the_sidecar_naming_this_weight(monkeypatch): + """The other direction: model A's own sidecar still wins over the generic + one, as it does in detect_dflash_file.""" + assert ( + _dflash_download_pick( + monkeypatch, + listing = _MULTI_FAMILY_LISTING, + near_path = "/cache/snap/model-A-Q4_K_M.gguf", + ) + == "dflash-model-A-Q8_0.gguf" + ) + + +@pytest.mark.parametrize("sidecar", ["dflash-kquant.gguf", "dflash-bf16.gguf"]) +def test_download_dflash_keeps_the_single_published_sidecar(monkeypatch, sidecar): + """The shipped unsloth/Muse-Glimmer-30B-GGUF layout. Its sidecar's stem is a + precision token, not a family, so nothing may treat "names no weight here" + as a rejection.""" + assert ( + _dflash_download_pick( + monkeypatch, + listing = [ + "Muse-Glimmer-30B-UD-Q4_K_XL.gguf", + "mmproj-Muse-Glimmer-30B-Q8_0.gguf", + sidecar, + ], + near_path = "/cache/snap/Muse-Glimmer-30B-UD-Q4_K_XL.gguf", + ) + == sidecar + ) + + +def test_cached_dflash_lookup_pairs_with_the_selected_weight(tmp_path, monkeypatch): + """The offline reuse must reach the same file the download would have + fetched, or a reload swaps drafters as soon as the cache is warm.""" + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "snapshots" / "abc" + snap.mkdir(parents = True) + for name in _MULTI_FAMILY_LISTING: + (snap / name).write_bytes(b"x") + monkeypatch.setattr( + "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] + ) + + b = LlamaCppBackend() + for weight, expected in ( + ("model-B-Q4_K_M.gguf", "dflash-kquant.gguf"), + ("model-A-Q4_K_M.gguf", "dflash-model-A-Q8_0.gguf"), + ): + assert b._cached_repo_dflash_drafter( + "org/repo", near_path = str(snap / weight) + ) == str(snap / expected) + # No weight in hand: precision order, exactly as before. + assert b._cached_repo_dflash_drafter("org/repo") == str(snap / "dflash-model-A-Q8_0.gguf") + + +def test_cached_dflash_lookup_keeps_the_single_published_sidecar(tmp_path, monkeypatch): + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "snapshots" / "abc" + snap.mkdir(parents = True) + for name in ("Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "dflash-kquant.gguf"): + (snap / name).write_bytes(b"x") + monkeypatch.setattr( + "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] + ) + + b = LlamaCppBackend() + assert b._cached_repo_dflash_drafter( + "org/repo", near_path = str(snap / "Muse-Glimmer-30B-UD-Q4_K_XL.gguf") + ) == str(snap / "dflash-kquant.gguf") + + +def test_local_and_remote_dflash_pairing_agree(tmp_path): + """One rule, three call sites: the local scan, the download picker and the + cache lookup all go through dflash_repo_preference_key / + _drafter_names_other_weight.""" + from utils.models.model_config import dflash_repo_preference_key + + others = ["model-A-Q4_K_M.gguf", "model-B-Q4_K_M.gguf"] + ranked = sorted( + ("dflash-model-A-Q8_0.gguf", "dflash-kquant.gguf"), + key = lambda name: dflash_repo_preference_key(name, "model-B-Q4_K_M.gguf", others), + ) + assert ranked[0] == "dflash-kquant.gguf" + + for name in _MULTI_FAMILY_LISTING: + _write_gguf(tmp_path / name, "dflash" if name.startswith("dflash-") else "llama") + assert detect_dflash_file(str(tmp_path / "model-B-Q4_K_M.gguf")) == str( + (tmp_path / "dflash-kquant.gguf").resolve() + ) + + # ── Reclaim: deliberately unchanged ────────────────────────────────── diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5e86f10bd08..2d5009db9fc 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -30,7 +30,7 @@ import subprocess import sys from pathlib import Path -from typing import Callable, List, Tuple, Union +from typing import Callable, Iterable, List, Tuple, Union import hashlib import json import threading @@ -1776,6 +1776,61 @@ def dflash_preference_key(name: str) -> tuple[int, str]: return dflash_precision_rank(name), Path(name).name.lower() +def _drafter_names_other_weight( + candidate_name: str, + weight_name: Optional[str], + other_weight_names: Iterable[str], + *, + kind: str = "dflash", +) -> bool: + """Whether a sidecar names a DIFFERENT weight sitting beside it. + + A sidecar that names no family at all (the published ``dflash-kquant.gguf``, + whose stem is a precision token) has to stay eligible, so "does it name a + family" cannot be answered from the sidecar name alone. It is answered + against the weights actually present instead: only a stem that pairs with + some OTHER weight in the same repo/folder is evidence the sidecar belongs to + that neighbour rather than to the weight being loaded. + """ + if weight_name is None: + return False + if _drafter_matches_weight(candidate_name, weight_name, kind = kind): + return False + return any( + _drafter_matches_weight(candidate_name, other, kind = kind) + for other in other_weight_names + ) + + +def dflash_repo_preference_key( + name: str, + weight_name: Optional[str] = None, + other_weight_names: Iterable[str] = (), +) -> tuple[int, int, int, str]: + """Order DFlash sidecars in a repo listing / cache snapshot against the + weight actually being loaded. + + dflash_preference_key ranks by precision and name alone, which is all a + single-model repo needs. A repo hosting more than one family also has to be + told which weight each sidecar belongs to, or ``dflash-model-A-Q8_0.gguf`` + outranks the generic ``dflash-kquant.gguf`` on precision and model B is + launched with model A's drafter. Same rule the local scan applies in + detect_dflash_file, kept in one place so the download, the snapshot reuse + and the offline cache all pick the same file. + + Three buckets: a sidecar naming this weight's family (most specific stem + first, as detect_mtp_file does), then one naming no weight present here, + then one naming a neighbour. The last is demoted rather than dropped, so a + repo whose only sidecar looks foreign still gets a fallback and today's + single-sidecar behaviour is unchanged. + """ + precision, sort_name = dflash_preference_key(name) + if weight_name is not None and _drafter_matches_weight(name, weight_name, kind = "dflash"): + return 0, _drafter_stem_rank(name, kind = "dflash"), precision, sort_name + foreign = _drafter_names_other_weight(name, weight_name, other_weight_names) + return 2 if foreign else 1, 0, precision, sort_name + + def detect_mtp_file( path: str, search_root: Optional[str] = None, @@ -2160,16 +2215,15 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: # foreign one to the top: loading model B beside dflash-model-A-Q8_0.gguf # and dflash-kquant.gguf launched model A's drafter for model B. Both carry # a real dflash header, so the architecture check behind the ranking cannot - # catch it. Deciding against the weights actually present keeps the - # published unpaired sidecar eligible (its stem, "kquant", names no file - # here) without hardcoding which stems are precision tokens. + # catch it. _drafter_names_other_weight decides against the weights actually + # present, which keeps the published unpaired sidecar eligible (its stem, + # "kquant", names no file here) without hardcoding which stems are precision + # tokens. Shared with the remote paths through dflash_repo_preference_key, + # so a download and a local scan agree on which sidecar belongs here. if weight_name is not None and other_weights: kept: list[Path] = [] for candidate in candidates: - if not _drafter_matches_weight(candidate.name, weight_name, kind = "dflash") and any( - _drafter_matches_weight(candidate.name, other, kind = "dflash") - for other in other_weights - ): + if _drafter_names_other_weight(candidate.name, weight_name, other_weights): logger.info( "detect_dflash_file: dropped %s (names another weight in this folder)", candidate.name, From d4933b7baaca85a6ad657ee9fdf9c3b947581448 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:47:54 +0000 Subject: [PATCH 07/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 4 +--- studio/backend/tests/test_mtp_drafter_companion.py | 6 +++--- studio/backend/utils/models/model_config.py | 3 +-- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 507718e2928..fefa7ef09eb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8136,9 +8136,7 @@ def _cached_repo_dflash_drafter( # Every non-sidecar GGUF in the snapshot is a weight some sidecar # could be naming; that is what tells a neighbour's sidecar apart # from one naming no family at all. - others = [ - Path(name).name for name in names if not _is_dflash_drafter_path(name) - ] + others = [Path(name).name for name in names if not _is_dflash_drafter_path(name)] ranked.extend( (dflash_repo_preference_key(name, weight_name, others), snap / name) for name in names diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 623eab1ea10..5903d772aa4 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2184,9 +2184,9 @@ def test_cached_dflash_lookup_pairs_with_the_selected_weight(tmp_path, monkeypat ("model-B-Q4_K_M.gguf", "dflash-kquant.gguf"), ("model-A-Q4_K_M.gguf", "dflash-model-A-Q8_0.gguf"), ): - assert b._cached_repo_dflash_drafter( - "org/repo", near_path = str(snap / weight) - ) == str(snap / expected) + assert b._cached_repo_dflash_drafter("org/repo", near_path = str(snap / weight)) == str( + snap / expected + ) # No weight in hand: precision order, exactly as before. assert b._cached_repo_dflash_drafter("org/repo") == str(snap / "dflash-model-A-Q8_0.gguf") diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 2d5009db9fc..6fc775d7a0f 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1797,8 +1797,7 @@ def _drafter_names_other_weight( if _drafter_matches_weight(candidate_name, weight_name, kind = kind): return False return any( - _drafter_matches_weight(candidate_name, other, kind = kind) - for other in other_weight_names + _drafter_matches_weight(candidate_name, other, kind = kind) for other in other_weight_names ) From a10b30ad40b55e2eee8daeacce1ca6a1866a26f5 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 12:59:07 +0000 Subject: [PATCH 08/26] Validate DFlash candidates and size the extras drafter once Three fixes found in review of the DFlash drafter work. detect_dflash_file read a candidate's GGUF header before asking the caller's accept callback about it, so a dflash-*.gguf symlink in a directory reached through a native grant had its out-of-lease target opened before the grant check ran, and no later rejection takes a read back. The loop now resolves the launch path, runs accept, and only then parses the header and applies the architecture check. accept still receives the resolved launch path, and callers that pass no accept see the same candidates in the same order as before. The training admission guard charged the llama_extra_args --model-draft sidecar on top of the local one discovery had already found, so a 1.5 GiB drafter was billed as 3 GiB and the guard could refuse an inference load that fits. The effective draft path is now sized exactly once, with identity taken from the resolved path so a symlink or another spelling of the same file dedupes too. That same charge also satisfied the local-weights early return on its own. Loading a remote GGUF repo has no local main weight, so a local --model-draft made the guard return the drafter alone and skip the listing that prices the target model, which could admit a load that then exhausts VRAM next to a running training job. The local branch now fires only when a local weight is actually present, and the drafter is added to whichever branch produces the estimate, including the remote one. Regression tests for all three. --- studio/backend/routes/inference.py | 50 +++++++-- .../tests/test_chat_load_during_training.py | 102 ++++++++++++++++++ .../tests/test_mtp_drafter_companion.py | 46 ++++++++ studio/backend/utils/models/model_config.py | 24 +++-- 4 files changed, 204 insertions(+), 18 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 41bb43f5491..459363f9ced 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5588,10 +5588,22 @@ def _estimate_gguf_required_gb( _charge_no_drafter = (_forced_dspark and not _dspark_capable) or ( _forced_dflash and not _dflash_capable ) + def _same_file_key(p: str) -> str: + # Identity by resolved path, so a symlinked or differently spelled + # copy of one file is still one file. + try: + return os.path.realpath(p) + except OSError: + return str(p) + total_bytes = 0 + # Only the files already charged above, so the extras drafter below can + # tell "another sidecar" from "the one discovery already found". + _sized_keys: set[str] = set() main = getattr(config, "gguf_file", None) if main and Path(main).is_file(): total_bytes += LlamaCppBackend._get_gguf_size_bytes(str(main)) + _sized_keys.add(_same_file_key(str(main))) # Only the drafter the launch will load: the modes are exclusive, and a # 10 GB DSpark sidecar merely sitting on disk must not inflate the guard # for a load that never opens it. @@ -5603,15 +5615,6 @@ def _estimate_gguf_required_gb( _sized_attrs.append("gguf_dflash_file") else: _sized_attrs.append("gguf_mtp_file") - # A caller that owns speculation through llama_extra_args names the - # drafter with --model-draft, and discovery never populates - # gguf_dflash_file / gguf_dspark_file for a file outside the model - # directory. load_model still hands that path to llama-server, so - # without this the guard admits a load beside a training run while - # charging nothing for the sidecar that load is about to make resident. - _extras_draft = _extra_args_mtp_draft_path(llama_extra_args, env = {}) - if _extras_draft and Path(_extras_draft).is_file(): - total_bytes += LlamaCppBackend._get_gguf_size_bytes(str(_extras_draft)) for attr in _sized_attrs: f = getattr(config, attr, None) @@ -5620,8 +5623,31 @@ def _estimate_gguf_required_gb( # so stat() alone would size a split drafter at one shard and let the # guard admit a load that evicts the training run it protects. total_bytes += LlamaCppBackend._get_gguf_size_bytes(str(f)) + _sized_keys.add(_same_file_key(str(f))) + + # A caller that owns speculation through llama_extra_args names the + # drafter with --model-draft. load_model hands that path to llama-server, + # so it has to be charged, but it is charged exactly once: the same file + # is often the local sidecar discovery already put in gguf_dflash_file / + # gguf_dspark_file / gguf_mtp_file, and adding it twice billed a 1.5 GiB + # drafter as 3 GiB and refused loads that fit. When the drafter really is + # outside the model directory nothing above named it and the charge lands + # here. It is a companion either way, never evidence of a local main + # weight, so it does not decide which branch below produces the estimate: + # a remote repo with a local --model-draft still has to price its weights + # through the listing, and returning the drafter alone under-estimated a + # load by the whole target model. + _extras_bytes = 0 + _extras_draft = _extra_args_mtp_draft_path(llama_extra_args, env = {}) + if ( + _extras_draft + and Path(_extras_draft).is_file() + and _same_file_key(str(_extras_draft)) not in _sized_keys + ): + _extras_bytes = LlamaCppBackend._get_gguf_size_bytes(str(_extras_draft)) + if total_bytes > 0: - return total_bytes / (1024**3) + _estimate_gguf_kv_gb( + return (total_bytes + _extras_bytes) / (1024**3) + _estimate_gguf_kv_gb( main, max_seq_length, llama_extra_args, @@ -5660,7 +5686,9 @@ def _estimate_gguf_required_gb( include_dspark = (_dspark_capable and (_auto_dspark or dspark_requested)), include_dflash = (_dflash_capable and (_auto_dflash or dflash_requested)), ) - total_gb = (main_bytes + companions) / (1024**3) + # Plus the local --model-draft, if the caller named one: the repo + # listing cannot see it, and it is resident next to these weights. + total_gb = (main_bytes + companions + _extras_bytes) / (1024**3) # remote dims are unreadable; only the kq mask, linear in ubatch x ctx, can be sized here from core.inference.llama_server_args import parse_ctx_override diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 186a684dc79..0069b4c71e6 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1326,6 +1326,108 @@ def test_split_dspark_sidecar_counts_every_shard(self): gb = self.route._estimate_gguf_required_gb(cfg, speculative_type = "dspark") self.assertAlmostEqual(gb, 9000 / (1024**3), places = 9) # 2000 + 3000 + 4000 + @staticmethod + def _dflash_capable(supported = True): + """Same shape as _dspark_capable: the DFlash sizing gate asks the binary + whether it can run draft-dflash.""" + from core.inference.llama_cpp import LlamaCppBackend + return patch.object( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: {"supports_dflash": supported}), + ) + + def test_extra_args_drafter_is_charged_once_when_it_is_the_local_sidecar(self): + """--model-draft usually names the very sidecar discovery already found, + and charging it on both paths billed a 1.5 GiB drafter as 3 GiB, so the + guard refused an inference load that fits. Identity is the resolved path, + so a symlink or another spelling of the same file dedupes too, while a + genuinely separate drafter outside the model directory is still charged. + """ + import os + import tempfile + + with tempfile.TemporaryDirectory() as d: + p = Path(d) + target = p / "model.gguf" + sidecar = p / "dflash-kquant.gguf" + target.write_bytes(b"x" * 2000) + sidecar.write_bytes(b"y" * 3000) + link = p / "linked-dflash.gguf" + os.symlink(sidecar, link) + elsewhere = p / "other" / "dflash-elsewhere.gguf" + elsewhere.parent.mkdir() + elsewhere.write_bytes(b"z" * 4000) + cfg = SimpleNamespace( + gguf_file = str(target), + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = str(sidecar), + gguf_hf_repo = None, + gguf_variant = None, + ) + with ( + patch.object(self.route, "_estimate_gguf_kv_gb", return_value = 0.0), + self._dflash_capable(), + ): + plain = self.route._estimate_gguf_required_gb(cfg, speculative_type = "dflash") + same = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "dflash", + llama_extra_args = ["--model-draft", str(sidecar)], + ) + through_link = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "dflash", + llama_extra_args = ["--model-draft", str(link)], + ) + separate = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "dflash", + llama_extra_args = ["--model-draft", str(elsewhere)], + ) + self.assertAlmostEqual(plain, 5000 / (1024**3), places = 9) + self.assertAlmostEqual(same, 5000 / (1024**3), places = 9) # not 8000 + self.assertAlmostEqual(through_link, 5000 / (1024**3), places = 9) + self.assertAlmostEqual(separate, 9000 / (1024**3), places = 9) # 2000+3000+4000 + + def test_remote_weights_stay_in_the_estimate_beside_a_local_extra_args_drafter(self): + """A remote repo has no local main weight, so a local --model-draft was + the only thing making the local branch fire: it returned ~1.5 GiB and + skipped the listing that prices the target model entirely. The drafter is + a companion, not evidence of local weights, so it is added to whichever + branch produces the estimate.""" + import tempfile + + import utils.models.model_config as mc + + cfg = SimpleNamespace( + gguf_file = None, + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = None, + gguf_hf_repo = "org/repo", + gguf_variant = "Q4_K_M", + ) + variant = SimpleNamespace(quant = "Q4_K_M", size_bytes = 10 * 1024**3) + with tempfile.TemporaryDirectory() as d: + drafter = Path(d) / "dflash-kquant.gguf" + drafter.write_bytes(b"y" * 3000) + with ( + patch.object(mc, "list_gguf_variants", return_value = ([variant], False)), + patch.object(self.route, "_remote_gguf_companion_bytes", return_value = 0), + self._dflash_capable(), + ): + gb = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "dflash", + llama_extra_args = ["--model-draft", str(drafter)], + ) + # The 10 GB target weights, not just the drafter beside them. + self.assertAlmostEqual(gb, 10.0 + 3000 / (1024**3), places = 9) + def test_remote_threads_token_and_adds_companions(self): import utils.models.model_config as mc diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 5903d772aa4..2a29ad8b976 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2312,3 +2312,49 @@ def test_dflash_prefix_form_is_still_found_beside_the_weight(tmp_path): sidecar = _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") assert detect_dflash_file(str(weight)) == str(sidecar.resolve()) + + +def test_detect_dflash_file_validates_a_candidate_before_reading_its_header(tmp_path, monkeypatch): + """A native grant answers through ``accept``, and its answer has to arrive + before the file is opened. + + A dflash-*.gguf inside a leased directory can be a symlink whose target sits + outside the lease. Parsing the header first opened that target, and no later + rejection takes a read back, so the order is: resolve, ask accept, then read. + """ + import os + + import utils.models.model_config as mc + + leased = tmp_path / "leased" + leased.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + weight = _write_gguf(leased / "Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "muse-glimmer") + target = _write_gguf(outside / "dflash-kquant.gguf", "dflash") + os.symlink(target, leased / "dflash-kquant.gguf") + + reads: list[str] = [] + real_read = mc.read_gguf_general_metadata + + def _recording_read(path, *args, **kwargs): + reads.append(str(path)) + return real_read(path, *args, **kwargs) + + monkeypatch.setattr(mc, "read_gguf_general_metadata", _recording_read) + + def _inside_the_lease(launch: str) -> bool: + # accept is handed the resolved launch path, not the candidate. + return leased in Path(launch).parents + + assert detect_dflash_file(str(weight), accept = _inside_the_lease) is None + assert reads == [] # the out-of-grant target was never opened + + +def test_detect_dflash_file_still_checks_the_header_of_an_accepted_candidate(tmp_path): + """The reorder must not cost the architecture check every other caller relies + on: an accepted candidate that is not a dflash model is still dropped.""" + weight = _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + _write_gguf(tmp_path / "dflash-something-Q8_0.gguf", "llama") + + assert detect_dflash_file(str(weight), accept = lambda launch: True) is None diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 6fc775d7a0f..ccd6e6df767 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2232,7 +2232,23 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: candidates = kept for candidate in sorted(candidates, key = _rank): - meta = read_gguf_general_metadata(str(candidate)) or {} + # Resolve and validate before opening anything. A dflash-*.gguf in a + # directory reached through a native grant can be a symlink whose target + # sits outside the lease, and ``accept`` is what decides that; reading the + # header first opened the target before the answer arrived, which a later + # rejection cannot undo. Callers without a grant pass accept = None and + # see the same candidates in the same order as before. + try: + launch = _drafter_launch_path(candidate) + except OSError: + continue + if accept is not None and not accept(launch): + logger.info( + "detect_dflash_file: dropped %s (outside the granted directory)", + candidate.name, + ) + continue + meta = read_gguf_general_metadata(launch) or {} if (meta.get("general.architecture") or "").strip().lower() != "dflash": logger.info( "detect_dflash_file: dropped %s (architecture %r is not dflash)", @@ -2240,12 +2256,6 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: meta.get("general.architecture"), ) continue - try: - launch = _drafter_launch_path(candidate) - except OSError: - continue - if accept is not None and not accept(launch): - continue logger.info("Detected DFlash drafter: %s", launch) return launch return None From 2ca75d8594ab2c27bd10d05c6937ee515c108ff4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:23:17 +0000 Subject: [PATCH 09/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 459363f9ced..bed386e90c9 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5588,6 +5588,7 @@ def _estimate_gguf_required_gb( _charge_no_drafter = (_forced_dspark and not _dspark_capable) or ( _forced_dflash and not _dflash_capable ) + def _same_file_key(p: str) -> str: # Identity by resolved path, so a symlinked or differently spelled # copy of one file is still one file. From 0b9525635780c2ba6db2d3bbe428fa61260410fe Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 13:49:59 +0000 Subject: [PATCH 10/26] Validate remote DFlash files by header and stop charging unused DFlash bytes Two fixes to the DFlash sidecar paths. Remote and cached DFlash candidates are now confirmed by their GGUF header, not by their filename. _pick_dflash and _cached_repo_dflash_drafter selected with _is_dflash_drafter_path, a dflash- prefix test, while the local scan in detect_dflash_file also required general.architecture == dflash. A remote repo holding an ordinary weight whose basename starts with dflash- therefore had that full weight downloaded and handed to llama-server as --model-draft, which falls back at startup after the bytes are already spent. The architecture rule moves into is_dflash_architecture in model_config, beside the naming rules and shared by every path, the way dflash_repo_preference_key already is. The header is only readable once the file is on disk, so the download validates after the fetch and falls through to the next candidate instead of returning None; the prefix-only naming rule is unchanged. The training coexistence guard no longer charges DFlash bytes a load under Auto will never fetch. _remote_gguf_companion_bytes added the preferred DSpark and the preferred DFlash sidecar whenever the repo listed both, but the loader stands down on the DFlash fetch once DSpark resolves under Auto, so those bytes are never resident and the guard could 409 a load that fits. The new dspark_first flag mirrors that selection. Where the choice is genuinely unknown the deliberate over-estimate stands, and an explicitly forced DFlash still pays for its sidecar. Regression tests cover the fetch falling through an impostor to the real sidecar, an all-impostor repo recording a permanent absence, the snapshot reuse and offline cache lookups applying the same rule, and the Auto guard charging DSpark only when a repo publishes both kinds. --- studio/backend/core/inference/llama_cpp.py | 102 +++++++-- studio/backend/routes/inference.py | 24 ++- .../tests/test_chat_load_during_training.py | 62 ++++++ .../tests/test_mtp_drafter_companion.py | 197 ++++++++++++++++-- studio/backend/utils/models/model_config.py | 27 ++- 5 files changed, 381 insertions(+), 31 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fefa7ef09eb..6a0959c270d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8122,6 +8122,7 @@ def _cached_repo_dflash_drafter( from utils.models.model_config import ( _iter_hf_cache_snapshots, dflash_repo_preference_key, + is_dflash_architecture, ) snapshots = ( @@ -8143,8 +8144,21 @@ def _cached_repo_dflash_drafter( if _is_dflash_drafter_path(name) ) for _, candidate in sorted(ranked, key = lambda entry: entry[0]): - if candidate.is_file(): - return str(candidate) + if not candidate.is_file(): + continue + # The ranking above works off names, and dflash- is only a naming + # convention: a cached weight whose basename happens to start with + # it would be launched as --model-draft and refused at startup. + # Same header rule, same helper, as the local scan, so an offline + # reuse and a local scan cannot disagree about what a sidecar is. + # Skipped rather than fatal, since the snapshot may hold a real one. + if not is_dflash_architecture(str(candidate)): + logger.info( + "Ignoring cached DFlash candidate %s: general.architecture is not dflash", + candidate, + ) + continue + return str(candidate) except Exception as exc: logger.debug("Cached DFlash drafter lookup failed for %s: %s", hf_repo, exc) return None @@ -8169,6 +8183,11 @@ def _download_dflash( """ weight_name = Path(near_path).name if near_path else None + # Basenames whose header turned out not to say dflash. A name lands here + # only once the file is readable on disk, and _pick_dflash then skips it, + # so a repo whose best-ranked candidate is an impostor still reaches the + # real sidecar behind it instead of the search giving up on the repo. + rejected: set[str] = set() def _pick_dflash(candidates: list[str]) -> Optional[str]: # Ranked against the weight being loaded, not by precision alone: a @@ -8186,12 +8205,53 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: if name.lower().endswith(".gguf") and not _is_dflash_drafter_path(name) ] files = sorted( - (name for name in candidates if _is_dflash_drafter_path(name)), + ( + name + for name in candidates + if _is_dflash_drafter_path(name) and Path(name).name not in rejected + ), key = lambda name: dflash_repo_preference_key(name, weight_name, others), ) return files[0] if files else None - cached = _companion_snapshot_sibling(near_path, _pick_dflash) if near_path else None + def _validated(path: Optional[str]) -> Optional[str]: + """``path`` back iff its header really says ``dflash``. + + _is_dflash_drafter_path is a FILENAME test, deliberately (the prefix + form is the only one accepted, so a sidecar cannot double as a + selectable main model in the quant picker), which means a remote repo + holding an ordinary weight called dflash-*.gguf passes it. The local + scan settles that with the architecture in the header; the remote + paths have to as well, or the load spends gigabytes of download on a + file llama-server then refuses as --model-draft. Same helper as + detect_dflash_file so the two rules cannot drift. + """ + from utils.models.model_config import is_dflash_architecture + + if not path: + return None + if is_dflash_architecture(path): + return path + rejected.add(Path(path).name) + logger.warning("Ignoring DFlash candidate %s: general.architecture is not dflash", path) + return None + + # Retried rather than abandoned: every rejection removes one name from the + # pool, so the next pick returns a different file and the scan ends at the + # last candidate (or at the first real sidecar). Answers already seen end + # it too, so a scan that ignores the pool cannot spin here. + cached: Optional[str] = None + seen_siblings: set[str] = set() + while near_path: + sibling = _companion_snapshot_sibling(near_path, _pick_dflash) + if sibling is None or sibling in seen_siblings: + break + seen_siblings.add(sibling) + cached = _validated(sibling) + if cached: + break + # _cached_repo_dflash_drafter applies the same header check to its own + # candidates, so anything it hands back is already validated. if not cached and _hf_env_offline(): cached = self._cached_repo_dflash_drafter( hf_repo, @@ -8218,16 +8278,32 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: logger.info("Reusing cached DFlash drafter: %s", cached) return cached outcome: dict = {} - found = self._download_companion_gguf( - hf_repo = hf_repo, - hf_token = hf_token, - pick = _pick_dflash, - label = "DFlash drafter", - near_path = near_path, - outcome = outcome, - ) + found: Optional[str] = None + # The header can only be read once the bytes are here, so validation is a + # post-fetch step and a rejection has to be able to fall through to the + # next candidate rather than end the search. Bounded the same way as the + # reuse scan above: a fetch that answers with a file already rejected has + # nothing further to offer. + fetched: set[str] = set() + while True: + candidate = self._download_companion_gguf( + hf_repo = hf_repo, + hf_token = hf_token, + pick = _pick_dflash, + label = "DFlash drafter", + near_path = near_path, + outcome = outcome, + ) + if candidate is None or candidate in fetched: + break + fetched.add(candidate) + found = _validated(candidate) + if found: + break # Distinguishes a repo that ships no sidecar from a fetch that failed and - # could yet succeed. + # could yet succeed. A repo whose only dflash-*.gguf is an ordinary weight + # exits the loop through a listing that picked nothing, which records the + # same permanent absence: there is nothing usable here to fetch next time. self._dflash_sidecar_absent = outcome.get("listed") is False return found diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index bed386e90c9..61c61c9b114 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5306,8 +5306,14 @@ def _remote_gguf_companion_bytes( include_mtp: bool = True, include_dspark: bool = False, include_dflash: bool = False, + dspark_first: bool = False, ) -> int: - """Bytes of companion GGUFs the requested launch downloads. 0 on error.""" + """Bytes of companion GGUFs the requested launch downloads. 0 on error. + + ``dspark_first`` mirrors the loader's Auto rule: the DFlash fetch stands down + once DSpark has resolved, so a repo publishing both kinds only ever pays for + the DSpark sidecar. + """ try: from core.inference.llama_cpp import ( _is_dflash_drafter_path, @@ -5338,7 +5344,16 @@ def _remote_gguf_companion_bytes( # Same preference order the download uses, so the budget sizes the # file the launch will actually fetch. total += min(dspark_candidates, key = lambda c: dspark_preference_key(c[0]))[1] - if dflash_candidates: + # Under Auto the loader stands down on the DFlash fetch as soon as DSpark + # resolves (DSpark takes first refusal in the promotion), so a repo that + # publishes both kinds never has the DFlash bytes resident. The caller + # asks for both because which kind a repo ships is unknown before the + # listing, and over-estimating is the safe direction for a guard + # protecting a running training job -- but the listing has answered by + # here, so with both present the outcome is known rather than unknown, and + # charging the unused ~1.5 GiB only makes the guard 409 a load that fits. + # An explicitly forced DFlash is not the Auto race and still pays. + if dflash_candidates and not (dspark_first and dspark_candidates): total += min(dflash_candidates, key = lambda c: dflash_preference_key(c[0]))[1] return total except Exception as e: @@ -5686,6 +5701,11 @@ def _same_file_key(p: str) -> str: ), include_dspark = (_dspark_capable and (_auto_dspark or dspark_requested)), include_dflash = (_dflash_capable and (_auto_dflash or dflash_requested)), + # ... except where the listing settles it: a repo shipping BOTH + # kinds only loads the DSpark one under Auto, so charging the + # DFlash sidecar too is not caution, it is a refusal for bytes + # that never land. + dspark_first = _auto_dspark, ) # Plus the local --model-draft, if the caller named one: the repo # listing cannot see it, and it is resident next to these weights. diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 0069b4c71e6..81720f533ba 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1494,6 +1494,68 @@ def test_remote_companions_choose_preferred_dspark_sidecar(self): ) self.assertEqual(dspark_only, 200) + def test_auto_charges_only_dspark_when_a_repo_publishes_both_sidecars(self): + """The loader stands down on the DFlash fetch once DSpark has resolved + under Auto, so those bytes are never resident. Charging both is not the + safe over-estimate it is for an unlisted repo -- the listing has answered + by then -- it is a 409 for a load that fits.""" + both = [ + SimpleNamespace(rfilename = "dspark/dspark-model-Q8_0.gguf", size = 200), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 400), + ] + + def _companion_bytes(siblings, **kwargs): + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + return self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + **kwargs, + ) + + self.assertEqual( + _companion_bytes(both, include_dspark = True, include_dflash = True, dspark_first = True), + 200, + ) + # Only one kind published: Auto still charges whichever the repo has. + self.assertEqual( + _companion_bytes( + [both[1]], include_dspark = True, include_dflash = True, dspark_first = True + ), + 400, + ) + # An explicit DFlash request is not the Auto race and still pays for it. + self.assertEqual(_companion_bytes(both, include_dflash = True), 400) + + def test_auto_tells_the_companion_sizing_that_dspark_comes_first(self): + """The remote branch is where both kinds can be asked for at once, so it + is the caller that has to pass the loader's Auto rule down.""" + import utils.models.model_config as mc + + cfg = SimpleNamespace( + gguf_file = None, + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = None, + gguf_hf_repo = "org/repo", + gguf_variant = "Q4_K_M", + ) + variant = SimpleNamespace(quant = "Q4_K_M", size_bytes = 1024**3) + with ( + patch.object(mc, "list_gguf_variants", lambda repo, hf_token = None: ([variant], False)), + patch.object(self.route, "_remote_gguf_companion_bytes", return_value = 0) as comp, + self._dspark_capable(), + ): + self.route._estimate_gguf_required_gb(cfg, speculative_type = "auto") + self.assertTrue(comp.call_args.kwargs["dspark_first"]) + self.route._estimate_gguf_required_gb(cfg, speculative_type = "dflash") + self.assertFalse(comp.call_args.kwargs["dspark_first"]) + def test_remote_unknown_variant_returns_none(self): import utils.models.model_config as mc cfg = SimpleNamespace( diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 2a29ad8b976..c46bdd1fcfd 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -1930,6 +1930,7 @@ def test_model_config_reports_a_local_dflash_sidecar(tmp_path): def _dflash_download_probe( + tmp_path, monkeypatch, *, supports_dflash, @@ -1969,7 +1970,8 @@ def _fake_companion( ) if outcome is not None: outcome["listed"] = True - return "/cache/dflash-kquant.gguf" + # A real file: the fetch is only accepted once its header says dflash. + return str(_write_gguf(tmp_path / "dflash-kquant.gguf", "dflash")) b = LlamaCppBackend() b._download_companion_gguf = _fake_companion @@ -1981,27 +1983,32 @@ def _fake_companion( return got, reached -def test_download_dflash_fetches_when_the_binary_supports_it(monkeypatch): - got, reached = _dflash_download_probe(monkeypatch, supports_dflash = True) - assert got == "/cache/dflash-kquant.gguf" +def test_download_dflash_fetches_when_the_binary_supports_it(tmp_path, monkeypatch): + got, reached = _dflash_download_probe(tmp_path, monkeypatch, supports_dflash = True) + assert got == str(tmp_path / "dflash-kquant.gguf") assert reached["hit"] is True # The picker must select the sidecar, not the weight or the projector. assert reached["picked"] == "dflash-kquant.gguf" -def test_download_dflash_skips_the_fetch_when_the_binary_cannot_run_it(monkeypatch): +def test_download_dflash_skips_the_fetch_when_the_binary_cannot_run_it(tmp_path, monkeypatch): """Same gate as DSpark: _build_speculative_flags drops DFlash outright on a binary without --spec-type draft-dflash, so the file would never be opened.""" - got, reached = _dflash_download_probe(monkeypatch, supports_dflash = False) + got, reached = _dflash_download_probe(tmp_path, monkeypatch, supports_dflash = False) assert got is None assert reached.get("hit", False) is False -def test_download_dflash_still_reports_a_cached_sidecar_it_cannot_run(monkeypatch): +def test_download_dflash_still_reports_a_cached_sidecar_it_cannot_run(tmp_path, monkeypatch): """The route rediscovers it on every Apply, so answering None would compare - it against a launched None and reload the same server each time.""" - cached = "/cache/snap/dflash-kquant.gguf" - got, reached = _dflash_download_probe(monkeypatch, supports_dflash = False, cached = cached) + it against a launched None and reload the same server each time. + + A real file on disk, since the reuse now confirms the header says dflash + before handing the path back.""" + cached = str(_write_gguf(tmp_path / "dflash-kquant.gguf", "dflash")) + got, reached = _dflash_download_probe( + tmp_path, monkeypatch, supports_dflash = False, cached = cached + ) assert got == cached assert reached.get("hit", False) is False @@ -2053,8 +2060,10 @@ def test_a_cached_dflash_drafter_is_never_launched_as_an_mtp_drafter(tmp_path, m snap = tmp_path / "snapshots" / "abc" snap.mkdir(parents = True) + # Real headers: the cached lookup confirms general.architecture before it + # hands a path to --model-draft. for name in ("dflash-kquant.gguf", "model-Q4_K_M.gguf"): - (snap / name).write_bytes(b"x") + _write_gguf(snap / name, "dflash" if name.startswith("dflash-") else "llama") monkeypatch.setattr( "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] ) @@ -2174,7 +2183,7 @@ def test_cached_dflash_lookup_pairs_with_the_selected_weight(tmp_path, monkeypat snap = tmp_path / "snapshots" / "abc" snap.mkdir(parents = True) for name in _MULTI_FAMILY_LISTING: - (snap / name).write_bytes(b"x") + _write_gguf(snap / name, "dflash" if name.startswith("dflash-") else "llama") monkeypatch.setattr( "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] ) @@ -2197,7 +2206,7 @@ def test_cached_dflash_lookup_keeps_the_single_published_sidecar(tmp_path, monke snap = tmp_path / "snapshots" / "abc" snap.mkdir(parents = True) for name in ("Muse-Glimmer-30B-UD-Q4_K_XL.gguf", "dflash-kquant.gguf"): - (snap / name).write_bytes(b"x") + _write_gguf(snap / name, "dflash" if name.startswith("dflash-") else "muse-glimmer") monkeypatch.setattr( "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] ) @@ -2358,3 +2367,165 @@ def test_detect_dflash_file_still_checks_the_header_of_an_accepted_candidate(tmp _write_gguf(tmp_path / "dflash-something-Q8_0.gguf", "llama") assert detect_dflash_file(str(weight), accept = lambda launch: True) is None + + +# ── Remote candidates are validated by header, not by name ─────────── +# +# _is_dflash_drafter_path is a dflash- FILENAME test, and deliberately only the +# prefix form (widening it would let one file be both a drafter and a selectable +# main model in the quant picker). detect_dflash_file backs that name test with +# the architecture in the GGUF header; the download and the cache reuse did not, +# so a repo holding an ordinary weight called dflash-*.gguf had it downloaded in +# full and handed to llama-server as --model-draft, which falls back at startup +# after the bytes are already spent. Same helper on every path. + + +def _dflash_repo_download( + tmp_path, + monkeypatch, + *, + listing, + sibling = None, +): + """Drive _download_dflash over a repo listing whose files exist in tmp_path.""" + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: {"supports_dflash": True}), + ) + monkeypatch.setattr( + llama_cpp_module, "_companion_snapshot_sibling", lambda near_path, pick: sibling + ) + fetched: list[str] = [] + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + near_path = None, + outcome = None, + ): + target = pick(listing) + if outcome is not None: + outcome["listed"] = target is not None + if target is None: + return None + fetched.append(target) + return str(tmp_path / target) + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + got = b._download_dflash( + hf_repo = "org/repo", + near_path = str(tmp_path / "model-Q4_K_M.gguf"), + binary = "/fake/llama-server", + ) + return b, got, fetched + + +def test_download_dflash_falls_through_a_candidate_that_is_not_a_dflash_model( + tmp_path, monkeypatch +): + """The impostor outranks the real sidecar on both name rules (it pairs with + this weight, and Q8_0 beats an unmarked precision), so the fetch reaches it + first. Its header is what disqualifies it, and only after the fetch, so the + search has to move on to the next candidate instead of returning None.""" + _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + _write_gguf(tmp_path / "dflash-model-Q8_0.gguf", "llama") + sidecar = _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") + + b, got, fetched = _dflash_repo_download( + tmp_path, + monkeypatch, + listing = ["model-Q4_K_M.gguf", "dflash-model-Q8_0.gguf", "dflash-kquant.gguf"], + ) + + assert got == str(sidecar) + assert fetched == ["dflash-model-Q8_0.gguf", "dflash-kquant.gguf"] + assert b._dflash_sidecar_absent is False + + +def test_download_dflash_reports_no_sidecar_when_every_candidate_is_a_weight(tmp_path, monkeypatch): + """A repo whose only dflash-*.gguf is an ordinary model publishes no sidecar, + so the absence is recorded and the next Apply does not re-list forever.""" + _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + _write_gguf(tmp_path / "dflash-model-Q8_0.gguf", "llama") + + b, got, fetched = _dflash_repo_download( + tmp_path, + monkeypatch, + listing = ["model-Q4_K_M.gguf", "dflash-model-Q8_0.gguf"], + ) + + assert got is None + assert fetched == ["dflash-model-Q8_0.gguf"] # tried once, never re-picked + assert b._dflash_sidecar_absent is True + + +def test_download_dflash_validates_the_snapshot_sibling_it_reuses(tmp_path, monkeypatch): + """The reuse path never downloads anything, but it hands the same file to + --model-draft, so it applies the same header rule and keeps scanning.""" + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "models--org--repo" / "snapshots" / "abc" + snap.mkdir(parents = True) + _write_gguf(snap / "model-Q4_K_M.gguf", "llama") + _write_gguf(snap / "dflash-model-Q8_0.gguf", "llama") + sidecar = _write_gguf(snap / "dflash-kquant.gguf", "dflash") + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: {"supports_dflash": True}), + ) + + b = LlamaCppBackend() + b._download_companion_gguf = lambda **kwargs: pytest.fail("the reuse must not download") + assert b._download_dflash( + hf_repo = "org/repo", + near_path = str(snap / "model-Q4_K_M.gguf"), + binary = "/fake/llama-server", + ) == str(sidecar) + + +def test_cached_dflash_lookup_skips_a_prefixed_file_of_another_architecture(tmp_path, monkeypatch): + """The offline cache lookup ranks by name too, so a cached weight named like + a sidecar would be launched as the drafter with nothing left to catch it.""" + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "snapshots" / "abc" + snap.mkdir(parents = True) + _write_gguf(snap / "model-Q4_K_M.gguf", "llama") + _write_gguf(snap / "dflash-model-Q8_0.gguf", "llama") + _write_gguf(snap / "dflash-kquant.gguf", "dflash") + monkeypatch.setattr( + "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] + ) + + b = LlamaCppBackend() + assert b._cached_repo_dflash_drafter( + "org/repo", near_path = str(snap / "model-Q4_K_M.gguf") + ) == str(snap / "dflash-kquant.gguf") + + +def test_local_and_remote_dflash_architecture_checks_agree(tmp_path): + """One rule, one place: detect_dflash_file and the remote paths both ask + is_dflash_architecture, so neither can start trusting the name alone.""" + from utils.models.model_config import is_dflash_architecture + + weight = _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + impostor = _write_gguf(tmp_path / "dflash-model-Q8_0.gguf", "llama") + sidecar = _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") + + assert is_dflash_architecture(str(impostor)) is False + assert is_dflash_architecture(str(sidecar)) is True + assert is_dflash_architecture(str(tmp_path / "missing.gguf")) is False + assert detect_dflash_file(str(weight)) == str(sidecar.resolve()) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index ccd6e6df767..f31aae4d7cb 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1830,6 +1830,25 @@ def dflash_repo_preference_key( return 2 if foreign else 1, 0, precision, sort_name +def is_dflash_architecture(path: str) -> bool: + """Whether a GGUF really is a DFlash sidecar, decided by its header. + + ``dflash-`` is a filename convention an ordinary weight can satisfy, by + accident or otherwise, and llama-server only discovers that at startup: it + refuses the file as ``--model-draft`` and the load falls back to no + speculation, after the bytes were already fetched. A DFlash sidecar declares + ``general.architecture = dflash``, which no real weight does, so that is what + settles it. + + Kept here, beside the naming rules, because the local scan + (detect_dflash_file) and the download / cache reuse in llama_cpp all have to + apply it -- a remote path that trusted the prefix alone would download + gigabytes the launch then cannot use. + """ + meta = read_gguf_general_metadata(str(path)) or {} + return (meta.get("general.architecture") or "").strip().lower() == "dflash" + + def detect_mtp_file( path: str, search_root: Optional[str] = None, @@ -2248,12 +2267,14 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: candidate.name, ) continue - meta = read_gguf_general_metadata(launch) or {} - if (meta.get("general.architecture") or "").strip().lower() != "dflash": + if not is_dflash_architecture(launch): logger.info( "detect_dflash_file: dropped %s (architecture %r is not dflash)", candidate.name, - meta.get("general.architecture"), + # Re-read only on the reject path, and header reads are cached by + # (path, mtime, size), so naming the offending architecture in the + # log costs nothing. + (read_gguf_general_metadata(launch) or {}).get("general.architecture"), ) continue logger.info("Detected DFlash drafter: %s", launch) From 812c88902205d6da507e3de56544912097db024b Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 14:01:28 +0000 Subject: [PATCH 11/26] Gate the DFlash stand-down and the guard's sizing on what the load actually does The Auto DFlash fetch stood down whenever _download_dspark answered with a path, but that call deliberately reports an already-cached DSpark sidecar even on a binary with no usable --spec-type draft-dspark (so the route's reuse check does not reload the same server on every Apply). The promotion refuses such a path, so on a DFlash-capable binary a repo shipping both companions suppressed the DFlash fetch for a sidecar that can never launch and the load ended up with no drafter at all. The capability gate now lives in _dspark_wins_auto, shared by the fetch and the promotion so the two cannot disagree. _remote_gguf_companion_bytes still ranked DFlash candidates with the name-only dflash_preference_key while the loader moved to the family-aware dflash_repo_preference_key, so in a multi-family repo the guard could price a different, smaller sidecar than the one that lands. The selected weight name is threaded down and the guard now sorts with the downloader's key over the neighbouring weights from the same listing. --- studio/backend/core/inference/llama_cpp.py | 73 ++++++++--- studio/backend/routes/inference.py | 40 +++++- .../tests/test_chat_load_during_training.py | 63 ++++++++++ .../tests/test_mtp_drafter_companion.py | 116 ++++++++++++++++++ 4 files changed, 270 insertions(+), 22 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 6a0959c270d..3358dfcee60 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8307,6 +8307,39 @@ def _validated(path: Optional[str]) -> Optional[str]: self._dflash_sidecar_absent = outcome.get("listed") is False return found + def _dspark_wins_auto( + self, + *, + binary: Optional[str], + dspark_draft_path: Optional[str], + spec_canon: str, + extra_args: Optional[list[str]], + ) -> bool: + """Whether Auto will actually resolve this load to DSpark. + + A DSpark path on its own does not answer that. _download_dspark reports + an already-cached sidecar even when the binary has no usable + ``--spec-type draft-dspark`` (deliberately: the route rediscovers it on + every Apply, and answering None there would reload the same server each + time), so the capability has to be applied by everything that reads the + path as "DSpark is what Auto picks". Without it a binary that supports + DFlash but not DSpark stands down on the DFlash fetch for a sidecar it + can never launch and the load ends up with no drafter at all. + + A raised probe answers False, which keeps Auto on its other options + exactly as the promotion does. probe_server_capabilities caches on + (binary path, mtime), so asking here and at the promotion is one probe. + """ + if spec_canon != "auto" or not dspark_draft_path: + return False + if _extra_args_set_spec_type(extra_args): + return False + try: + return bool(self.probe_server_capabilities(binary).get("supports_dspark")) + except Exception as exc: + logger.debug("DSpark capability probe failed during Auto: %s", exc) + return False + def _resolve_launch_mmproj_path( self, *, model_path: str, mmproj_path: Optional[str] ) -> Optional[str]: @@ -9898,12 +9931,21 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # below, so a repo that ships both kinds would never launch # the DFlash sidecar. Fetching it anyway costs ~1.5 GiB of # bandwidth and cache for a file that cannot be used, so the - # Auto fetch stands down once DSpark has resolved. An - # explicit "dflash" request still fetches. + # Auto fetch stands down once DSpark has resolved. It has to + # be the same "resolved" the promotion means, capability and + # all: a DSpark path this binary cannot launch loses the + # promotion, so standing down on it would leave a + # DFlash-capable binary with no drafter at all. An explicit + # "dflash" request still fetches. if ( not dflash_draft_path and _spec_canon in ("auto", "dflash") - and not (_spec_canon == "auto" and dspark_draft_path) + and not self._dspark_wins_auto( + binary = binary, + dspark_draft_path = dspark_draft_path, + spec_canon = _spec_canon, + extra_args = extra_args, + ) and not _extra_args_set_spec_type(extra_args) ): dflash_draft_path = self._download_dflash( @@ -9924,21 +9966,18 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # Auto resolves to DSpark whenever a sidecar is available and this # binary can run it: 1.84x decode on 4x B200 and 1.91x on one, against - # the ngram-mod fallback this architecture would otherwise get. Gated - # on the capability because _download_dspark also reports a cached - # sidecar an incapable binary cannot launch, and promoting there would - # turn Auto's fallback into no speculative decoding at all. - if ( - _spec_canon == "auto" - and dspark_draft_path - and not _extra_args_set_spec_type(extra_args) + # the ngram-mod fallback this architecture would otherwise get. The + # capability gate lives in _dspark_wins_auto, shared with the DFlash + # fetch above so the fetch and the promotion cannot disagree about + # whether Auto is going to end up on DSpark. + if self._dspark_wins_auto( + binary = binary, + dspark_draft_path = dspark_draft_path, + spec_canon = _spec_canon, + extra_args = extra_args, ): - try: - if self.probe_server_capabilities(binary).get("supports_dspark"): - _spec_canon = "dspark" - logger.info("Auto: DSpark sidecar available, using draft-dspark.") - except Exception as exc: - logger.debug("DSpark capability probe failed during Auto: %s", exc) + _spec_canon = "dspark" + logger.info("Auto: DSpark sidecar available, using draft-dspark.") # DFlash is the other Auto promotion, on the same capability gate. # DSpark keeps first refusal (llama.cpp's own downloader ranks it diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 61c61c9b114..d84a3ccc2e8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5307,12 +5307,20 @@ def _remote_gguf_companion_bytes( include_dspark: bool = False, include_dflash: bool = False, dspark_first: bool = False, + weight_name: Optional[str] = None, ) -> int: """Bytes of companion GGUFs the requested launch downloads. 0 on error. ``dspark_first`` mirrors the loader's Auto rule: the DFlash fetch stands down once DSpark has resolved, so a repo publishing both kinds only ever pays for the DSpark sidecar. + + ``weight_name`` is the basename of the main GGUF this load selects. A repo + hosting more than one family ships a DFlash sidecar per family, and the + loader pairs them against that weight, so the guard has to be told which + weight it is pricing or it can charge a different (possibly smaller) sidecar + than the one the load will fetch, and wave through a load that then exhausts + VRAM beside a running training job. """ try: from core.inference.llama_cpp import ( @@ -5320,17 +5328,24 @@ def _remote_gguf_companion_bytes( _is_dspark_drafter_path, ) from huggingface_hub import model_info - from utils.models.model_config import dflash_preference_key, dspark_preference_key + from utils.models.model_config import dflash_repo_preference_key, dspark_preference_key info = model_info(repo, token = hf_token, files_metadata = True) total = 0 dspark_candidates: list[tuple[str, int]] = [] dflash_candidates: list[tuple[str, int]] = [] + # The weights a DFlash sidecar could be naming instead of this one, which + # is what tells a neighbour's sidecar apart from one naming no family at + # all. Derived from the listing exactly as _download_dflash derives it, + # so the guard ranks the candidates off the same evidence. + other_weight_names: list[str] = [] for sibling in info.siblings or []: name = sibling.rfilename or "" base = Path(name).name.lower() if not base.endswith(".gguf"): continue + if not _is_dflash_drafter_path(name): + other_weight_names.append(Path(name).name) # Root-level mtp- only: -hf auto-fetches the repo-root drafter, not # the MTP/ subdir copies (which now share the mtp- prefix too). is_root_mtp = "/" not in name and base.startswith("mtp-") @@ -5354,7 +5369,13 @@ def _remote_gguf_companion_bytes( # charging the unused ~1.5 GiB only makes the guard 409 a load that fits. # An explicitly forced DFlash is not the Auto race and still pays. if dflash_candidates and not (dspark_first and dspark_candidates): - total += min(dflash_candidates, key = lambda c: dflash_preference_key(c[0]))[1] + # dflash_repo_preference_key, not the name-only key: it is the key the + # downloader sorts with, and in a multi-family repo the two disagree + # about which sidecar this weight gets. + total += min( + dflash_candidates, + key = lambda c: dflash_repo_preference_key(c[0], weight_name, other_weight_names), + )[1] return total except Exception as e: logger.warning(f"Could not size GGUF companions for {repo}: {e}") @@ -5682,11 +5703,15 @@ def _same_file_key(p: str) -> str: from utils.models.model_config import list_gguf_variants variants, has_vision = list_gguf_variants(repo, hf_token = hf_token) - main_bytes = next( - (v.size_bytes for v in variants if v.quant.lower() == variant.lower()), None - ) + selected = next((v for v in variants if v.quant.lower() == variant.lower()), None) + main_bytes = selected.size_bytes if selected is not None else None if main_bytes is None: return None + # The variant record names the file this load opens, which is what the + # DFlash sizing needs to price the sidecar the loader will pair with + # it. A lister that reported no name leaves it None and the ranking + # falls back to precision alone, exactly as before. + selected_weight = Path(getattr(selected, "filename", "") or "").name or None companions = _remote_gguf_companion_bytes( repo, hf_token = hf_token, @@ -5706,6 +5731,11 @@ def _same_file_key(p: str) -> str: # DFlash sidecar too is not caution, it is a refusal for bytes # that never land. dspark_first = _auto_dspark, + # The weight this load actually opens. A multi-family repo ships a + # DFlash sidecar per family and the loader pairs them by name, so + # without it the guard can price a foreign (and smaller) sidecar + # than the one that lands. + weight_name = selected_weight, ) # Plus the local --model-draft, if the caller named one: the repo # listing cannot see it, and it is resident next to these weights. diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 81720f533ba..e7c46e49b7b 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1556,6 +1556,69 @@ def test_auto_tells_the_companion_sizing_that_dspark_comes_first(self): self.route._estimate_gguf_required_gb(cfg, speculative_type = "dflash") self.assertFalse(comp.call_args.kwargs["dspark_first"]) + _MULTI_FAMILY_SIBLINGS = [ + SimpleNamespace(rfilename = "model-A-Q4_K_M.gguf", size = 10 * 1024**3), + SimpleNamespace(rfilename = "model-B-Q4_K_M.gguf", size = 10 * 1024**3), + # Named after model A and higher precision, so the name-only key ranks it + # first for every weight in the repo. + SimpleNamespace(rfilename = "dflash-model-A-Q8_0.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 4 * 1024**3), + ] + + def test_remote_dflash_sizing_prices_the_sidecar_this_weight_pairs_with(self): + """The loader ranks DFlash candidates against the weight being loaded + (dflash_repo_preference_key), so a multi-family repo hands model B the + generic sidecar. Sizing by the name-only key priced model A's smaller + one instead, and the guard admitted a load that then exhausts VRAM + beside a running training job.""" + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = self._MULTI_FAMILY_SIBLINGS), + ): + for weight, expected_gib in ( + ("model-B-Q4_K_M.gguf", 4), + ("model-A-Q4_K_M.gguf", 1), + ): + total = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dflash = True, + weight_name = weight, + ) + self.assertEqual(total, expected_gib * 1024**3, weight) + + def test_remote_estimate_passes_the_selected_weight_to_the_dflash_sizing(self): + """End to end: the guard's own estimate has to carry the selected + filename down, or the sizing has nothing to pair the sidecar against.""" + import utils.models.model_config as mc + + cfg = SimpleNamespace( + gguf_file = None, + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = None, + gguf_hf_repo = "org/repo", + gguf_variant = "Q4_K_M", + ) + variant = SimpleNamespace( + filename = "model-B-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10 * 1024**3 + ) + with ( + patch.object(mc, "list_gguf_variants", lambda repo, hf_token = None: ([variant], False)), + patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = self._MULTI_FAMILY_SIBLINGS), + ), + self._dflash_capable(), + ): + gb = self.route._estimate_gguf_required_gb(cfg, speculative_type = "dflash") + # 10 GiB of weights plus the 4 GiB generic sidecar model B actually gets, + # not the 1 GiB one named after model A. + self.assertAlmostEqual(gb, 14.0, places = 6) + def test_remote_unknown_variant_returns_none(self): import utils.models.model_config as mc cfg = SimpleNamespace( diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index c46bdd1fcfd..a3079c6c999 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2529,3 +2529,119 @@ def test_local_and_remote_dflash_architecture_checks_agree(tmp_path): assert is_dflash_architecture(str(sidecar)) is True assert is_dflash_architecture(str(tmp_path / "missing.gguf")) is False assert detect_dflash_file(str(weight)) == str(sidecar.resolve()) + + +# ── Auto only stands down on DFlash for a DSpark it can launch ─────── +# +# _download_dspark reports an already-cached sidecar even when the binary has no +# usable --spec-type draft-dspark (so the route's reuse check does not reload the +# same server on every Apply), and the promotion refuses that path. The DFlash +# fetch read the bare path as "DSpark won" and stood down, so a repo shipping +# both companions left a DFlash-capable binary with NO drafter at all. + + +class _StopAfterDownloads(Exception): + """Ends the load once Phase 2 is done, which is all these tests observe.""" + + +def _dflash_fetch_during_auto_load(monkeypatch, *, supports_dspark, supports_dflash, dspark_cached): + """Whether an Auto load fetches the DFlash sidecar, and what it resolves to. + + Drives the real load path: the suppression lives inline in load_model's + download phase, so nothing short of running it can pin the interaction. + """ + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import GgufLoadIntent, LlamaCppBackend + + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod( + lambda cls, binary = None: { + "found": True, + "supports_dspark": supports_dspark, + "supports_dflash": supports_dflash, + } + ), + ) + monkeypatch.setattr(llama_cpp_module, "_resolve_repo_id_casing", lambda repo: repo) + monkeypatch.setattr( + llama_cpp_module, + "_hf_offline_if_unreachable", + lambda: __import__("contextlib").nullcontext(), + ) + + backend = LlamaCppBackend() + seen: dict = {"dflash_fetched": False} + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: False) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 4096, 8192)]) + monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: False) + monkeypatch.setattr(backend, "_kill_process", lambda: None) + monkeypatch.setattr( + backend, "_download_gguf", lambda **_kwargs: "/cache/snap/model-Q4_K_M.gguf" + ) + monkeypatch.setattr(backend, "_download_mtp", lambda **_kwargs: None) + # Exactly what _download_dspark does for a cached sidecar on a binary that + # cannot run it: the path comes back regardless of the capability. + monkeypatch.setattr(backend, "_download_dspark", lambda **_kwargs: dspark_cached) + + def _fetch_dflash(**_kwargs): + seen["dflash_fetched"] = True + return "/cache/snap/dflash-kquant.gguf" + + monkeypatch.setattr(backend, "_download_dflash", _fetch_dflash) + + def _stop(*_args, **_kwargs): + raise _StopAfterDownloads + + # The first call past the download phase; the resolved drafter is already + # settled by then. + monkeypatch.setattr(backend, "_read_gguf_metadata", _stop) + + with pytest.raises(_StopAfterDownloads): + backend.load_model( + GgufLoadIntent( + hf_repo = "org/repo", + hf_variant = "Q4_K_M", + model_identifier = "org/repo", + speculative_type = "auto", + ) + ) + return seen + + +def test_auto_still_fetches_dflash_when_the_binary_cannot_run_dspark(monkeypatch): + """The regression: a cached DSpark sidecar this binary cannot launch is not + a reason to skip the drafter it CAN launch.""" + seen = _dflash_fetch_during_auto_load( + monkeypatch, + supports_dspark = False, + supports_dflash = True, + dspark_cached = "/cache/snap/dspark-model-Q8_0.gguf", + ) + assert seen["dflash_fetched"] is True + + +def test_auto_stands_down_on_dflash_for_a_dspark_it_can_launch(monkeypatch): + """Unchanged where the stand-down was right: DSpark takes first refusal, so + the ~1.5 GiB DFlash fetch would buy a file the load never opens.""" + seen = _dflash_fetch_during_auto_load( + monkeypatch, + supports_dspark = True, + supports_dflash = True, + dspark_cached = "/cache/snap/dspark-model-Q8_0.gguf", + ) + assert seen["dflash_fetched"] is False + + +def test_auto_fetches_dflash_when_the_repo_ships_no_dspark_sidecar(monkeypatch): + """Positive control: nothing about the DSpark capability gates a repo that + publishes only the DFlash companion.""" + seen = _dflash_fetch_during_auto_load( + monkeypatch, + supports_dspark = True, + supports_dflash = True, + dspark_cached = None, + ) + assert seen["dflash_fetched"] is True From 87f23735d2ed2b2f792e504e8666a9ef9ee3c4a2 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 15:09:05 +0000 Subject: [PATCH 12/26] Apply the load's boundaries to drafter discovery, and size Auto's one drafter ModelConfig.from_identifier ran the local companion scan with no way for the caller to say what was in bounds, so a native-grant load read the header of a dflash-*.gguf symlinked out of the granted directory. The validated rescan on the load route rejected it afterwards, which does not take a read back. The boundary now travels into the scan, for all three drafter kinds, so the two passes cannot disagree about what is in bounds. Remote DFlash discovery matched the basename in any nested directory, but the local contract is root level only: a quants/dflash-*.gguf is an ordinary weight detect_dflash_file would never offer, and the header can only be read once the bytes are here, so the whole weight downloaded before the rejection. Checked through a separate predicate so the prefix-only naming rule the other callers share stays exactly as it is. A split companion is only usable as a whole set, since llama-server resolves the sibling shards from the first one's directory. Fetching just the picked shard left a drafter whose header reads fine and which the server cannot open, so the load fell back to no speculation with nothing to show for the download. The companion download now resolves its shards with the same helper the main-model download uses, and neither reuse path reports a half set as a cache hit. The remote sizing charged the first-ranked DFlash candidate, but a rejected candidate falls through to the next name in the ranking, which can be a larger file; headers are unreadable from a listing, so the bound now covers every candidate the fallback can reach. And under Auto the guard charged the MTP drafter on top of the DFlash sidecar that replaces it. Auto launches exactly one drafter, in a fixed order, so dspark_first now expresses the whole promotion: DSpark alone when the repo publishes one, otherwise the larger of the DFlash bound and the MTP drafter, since every DFlash candidate can still be turned away on its header and the load then keeps the MTP one. --- studio/backend/core/inference/llama_cpp.py | 101 +++++- studio/backend/routes/inference.py | 152 ++++---- .../tests/test_chat_load_during_training.py | 197 +++++++++-- .../tests/test_mtp_drafter_companion.py | 327 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 45 ++- 5 files changed, 719 insertions(+), 103 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3358dfcee60..c4c8547c746 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1219,6 +1219,25 @@ def _is_dflash_drafter_path(path: str) -> bool: return _drafter_path_kind(path) == "dflash" +def _is_root_dflash_drafter_path(path: str) -> bool: + """A DFlash sidecar as a repo listing or cache snapshot offers one: the + ``dflash-`` prefix AND the repository root. + + _is_dflash_drafter_path tests the basename, which is all its other callers + have (a bare filename carries no parent to look at). Remote discovery does + have the repository-relative path, and needs it: the local contract is root + level only, so ``quants/dflash-model-Q8_0.gguf`` is an ordinary weight that + detect_dflash_file would never offer. Going by basename there made it a + candidate, and since the header can only be read once the bytes are here, + the whole weight was downloaded before the rejection. + + Checked here rather than inside the basename predicate so the prefix-only + naming rule stays exactly as it is for the callers that share it. + """ + normalized = path.replace("\\", "/") + return "/" not in normalized and _is_dflash_drafter_path(normalized) + + _BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE) _GGUF_KNOWN_QUANT_RE = re.compile( r"(UD-)?" @@ -1560,7 +1579,17 @@ def _companion_snapshot_sibling( if not sibling: return None candidate = snap / sibling - return str(candidate) if candidate.is_file() else None + if not candidate.is_file(): + return None + # A split companion is usable only as a whole set (llama-server resolves the + # siblings from this path's directory), so a snapshot holding shard 1 alone + # is not a reuse -- reporting it would skip the download that completes it and + # leave the load with a drafter the server cannot open. Same rule the local + # scan applies. Non-split names are a complete one-file set, so this is a + # no-op for every companion published today. + from utils.models.model_config import _drafter_split_is_complete + + return str(candidate) if _drafter_split_is_complete(candidate) else None def _pick_mmproj(candidates: list[str]) -> Optional[str]: @@ -7779,11 +7808,28 @@ def _download_companion_gguf( Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. ``cancel_event`` overrides ``self._cancel_event`` (defaults to it). + + Split-aware, like the main-model download: a companion published as a + split GGUF is only usable as a complete set, since llama-server resolves + the sibling shards from the first one's directory. Fetching just the + picked shard left a drafter whose header reads fine and which the server + then cannot open, so the load fell back to no speculation with nothing to + show for the download -- and it disagreed with the local scan, which + accepts a split drafter only when every shard is present. """ cancel_event = cancel_event if cancel_event is not None else self._cancel_event if cancel_event.is_set(): return None + # The listing that produced the pick, kept so the shard siblings can be + # resolved from the same names the pick chose among. + available: list[str] = [] + + def _pick_from(names: list[str]) -> Optional[str]: + nonlocal available + available = list(names) + return pick(available) + # Keep companion files in the main GGUF's snapshot. if near_path: cached = _companion_snapshot_sibling(near_path, pick) @@ -7814,7 +7860,7 @@ def _download_companion_gguf( if cancel_event.is_set(): return None try: - target = pick(list_repo_files(hf_repo, token = hf_token)) + target = _pick_from(list_repo_files(hf_repo, token = hf_token)) listing_answered = True break except Exception as e: @@ -7838,7 +7884,7 @@ def _download_companion_gguf( from utils.models.model_config import _iter_hf_cache_snapshots for snap in _iter_hf_cache_snapshots(hf_repo, companion_cache_dir): rel_files = _gguf_snapshot_files(snap) - target = pick(rel_files) + target = _pick_from(rel_files) if target is not None: logger.info("Resolved %s %s from local HF cache", label, target) break @@ -7871,19 +7917,43 @@ def _download_companion_gguf( cache_dir = companion_cache_dir, ) if cached: - logger.info("Resolved %s from local HF cache: %s", label, cached) - return cached - + from utils.models.model_config import _drafter_split_is_complete + + # Same whole-set rule as the snapshot reuse above: half a split + # companion is not a companion, and offline there is no fetch to + # complete it, so answering None leaves the load without a + # drafter instead of with one llama-server cannot open. + if _drafter_split_is_complete(Path(cached)): + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + + # A split companion is one file to pick and N files to fetch. Same helper + # the main-model download resolves its shards with, so the two cannot + # disagree about what belongs to a set; empty for the single-file case, + # which is every mmproj and every published sidecar so far. + extra_shards = _gguf_extra_shards(available, target) try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). - return hf_hub_download_with_xet_fallback( + local_path = hf_hub_download_with_xet_fallback( hf_repo, target, hf_token, cancel_event = cancel_event, cache_dir = companion_cache_dir, ) + for shard in extra_shards: + if cancel_event.is_set(): + return None + logger.info(f"Downloading {label} shard: {hf_repo}/{shard}") + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = cancel_event, + cache_dir = companion_cache_dir, + ) + return local_path except Exception as e: logger.warning(f"Could not download {label}: {e}") return None @@ -8136,12 +8206,16 @@ def _cached_repo_dflash_drafter( names = _gguf_snapshot_files(snap) # Every non-sidecar GGUF in the snapshot is a weight some sidecar # could be naming; that is what tells a neighbour's sidecar apart - # from one naming no family at all. - others = [Path(name).name for name in names if not _is_dflash_drafter_path(name)] + # from one naming no family at all. A nested dflash-*.gguf counts + # as one of those weights, since only a root-level file is a + # sidecar here. + others = [ + Path(name).name for name in names if not _is_root_dflash_drafter_path(name) + ] ranked.extend( (dflash_repo_preference_key(name, weight_name, others), snap / name) for name in names - if _is_dflash_drafter_path(name) + if _is_root_dflash_drafter_path(name) ) for _, candidate in sorted(ranked, key = lambda entry: entry[0]): if not candidate.is_file(): @@ -8199,16 +8273,19 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: # stays eligible, which is what the published one does). from utils.models.model_config import dflash_repo_preference_key + # Root level only, as the local scan is: a nested dflash-*.gguf is an + # ordinary weight, and offering it here spends its entire download + # before the header check can turn it away. others = [ Path(name).name for name in candidates - if name.lower().endswith(".gguf") and not _is_dflash_drafter_path(name) + if name.lower().endswith(".gguf") and not _is_root_dflash_drafter_path(name) ] files = sorted( ( name for name in candidates - if _is_dflash_drafter_path(name) and Path(name).name not in rejected + if _is_root_dflash_drafter_path(name) and Path(name).name not in rejected ), key = lambda name: dflash_repo_preference_key(name, weight_name, others), ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d84a3ccc2e8..92417042e6c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4084,6 +4084,26 @@ def accept(candidate): return detected +def _native_drafter_accept(candidate: str, gguf_path: str, kind: str, search_root: str) -> bool: + """The native lease rule, in the shape ModelConfig.from_identifier takes. + + Discovery inside from_identifier runs before this route ever sees a path, and + the DFlash scan opens a candidate's header to confirm the architecture. A + native grant covers one directory, so handing the boundary down is what keeps + a sidecar symlinked out of the lease from being read at all -- rejecting it + afterwards, which _resolve_gguf_load_intent still does, cannot undo a read. + Same predicate the rescan uses, so the two passes cannot disagree about what + is in bounds. + """ + return _native_gguf_companion_usable( + candidate, + gguf_path, + kind = kind, + mtp_search_root = search_root, + log_rejection = True, + ) + + def _mtp_draft_for_path( gguf_path: Optional[str], native_grant_backed: bool, @@ -5307,76 +5327,84 @@ def _remote_gguf_companion_bytes( include_dspark: bool = False, include_dflash: bool = False, dspark_first: bool = False, - weight_name: Optional[str] = None, ) -> int: - """Bytes of companion GGUFs the requested launch downloads. 0 on error. - - ``dspark_first`` mirrors the loader's Auto rule: the DFlash fetch stands down - once DSpark has resolved, so a repo publishing both kinds only ever pays for - the DSpark sidecar. - - ``weight_name`` is the basename of the main GGUF this load selects. A repo - hosting more than one family ships a DFlash sidecar per family, and the - loader pairs them against that weight, so the guard has to be told which - weight it is pricing or it can charge a different (possibly smaller) sidecar - than the one the load will fetch, and wave through a load that then exhausts - VRAM beside a running training job. + """Bytes of companion GGUFs the requested launch keeps resident. 0 on error. + + ``dspark_first`` says this is an Auto load, which is the only caller that can + ask for several drafter kinds at once, so the loader's promotion order gets + to say which single one is charged: DSpark, else DFlash, else the MTP + drafter. The loader replaces mtp_draft_path with whichever kind wins the + promotion, so at most one drafter is ever launched; a sidecar that is fetched + and then not opened costs disk, not VRAM, and this guard sizes VRAM. Off, the + caller has already narrowed the request to one kind and the sum is that kind. """ try: from core.inference.llama_cpp import ( - _is_dflash_drafter_path, _is_dspark_drafter_path, + _is_root_dflash_drafter_path, ) from huggingface_hub import model_info - from utils.models.model_config import dflash_repo_preference_key, dspark_preference_key + from utils.models.model_config import dspark_preference_key info = model_info(repo, token = hf_token, files_metadata = True) total = 0 + mtp_bytes = 0 dspark_candidates: list[tuple[str, int]] = [] - dflash_candidates: list[tuple[str, int]] = [] - # The weights a DFlash sidecar could be naming instead of this one, which - # is what tells a neighbour's sidecar apart from one naming no family at - # all. Derived from the listing exactly as _download_dflash derives it, - # so the guard ranks the candidates off the same evidence. - other_weight_names: list[str] = [] + dflash_sizes: list[int] = [] for sibling in info.siblings or []: name = sibling.rfilename or "" base = Path(name).name.lower() if not base.endswith(".gguf"): continue - if not _is_dflash_drafter_path(name): - other_weight_names.append(Path(name).name) + size = getattr(sibling, "size", 0) or 0 # Root-level mtp- only: -hf auto-fetches the repo-root drafter, not # the MTP/ subdir copies (which now share the mtp- prefix too). is_root_mtp = "/" not in name and base.startswith("mtp-") - if (include_mtp and is_root_mtp) or (include_mmproj and "mmproj" in base): - total += getattr(sibling, "size", 0) or 0 + if include_mtp and is_root_mtp: + mtp_bytes += size + elif include_mmproj and "mmproj" in base: + total += size if include_dspark and _is_dspark_drafter_path(name): - dspark_candidates.append((name, getattr(sibling, "size", 0) or 0)) - if include_dflash and _is_dflash_drafter_path(name): - dflash_candidates.append((name, getattr(sibling, "size", 0) or 0)) + dspark_candidates.append((name, size)) + # Root level only, exactly as _download_dflash's picker is: a nested + # dflash-*.gguf is an ordinary weight there and never a candidate, so + # counting it here would price a file the load cannot fetch. + if include_dflash and _is_root_dflash_drafter_path(name): + dflash_sizes.append(size) + # Same preference order the download uses, so the budget sizes the file + # the launch will actually fetch. DSpark has no post-fetch rejection, so + # the best-ranked candidate is the one that lands. + dspark_bytes = ( + min(dspark_candidates, key = lambda c: dspark_preference_key(c[0]))[1] + if dspark_candidates + else 0 + ) + # The largest candidate the fetch could end up on, not the best-ranked + # one. _download_dflash can only read a candidate's header once it has + # paid for the bytes, and a rejection falls through to the next name in + # the ranking, so any candidate can be the file that lands -- and the + # whole point of the fallback is the case where it is a different, bigger + # one. Headers are unreadable from a listing, so the ranking cannot + # narrow that down here, and over-estimating is the established safe + # direction for a guard protecting a running training job. + dflash_bytes = max(dflash_sizes, default = 0) + if not dspark_first: + return total + mtp_bytes + dspark_bytes + dflash_bytes if dspark_candidates: - # Same preference order the download uses, so the budget sizes the - # file the launch will actually fetch. - total += min(dspark_candidates, key = lambda c: dspark_preference_key(c[0]))[1] - # Under Auto the loader stands down on the DFlash fetch as soon as DSpark - # resolves (DSpark takes first refusal in the promotion), so a repo that - # publishes both kinds never has the DFlash bytes resident. The caller - # asks for both because which kind a repo ships is unknown before the - # listing, and over-estimating is the safe direction for a guard - # protecting a running training job -- but the listing has answered by - # here, so with both present the outcome is known rather than unknown, and - # charging the unused ~1.5 GiB only makes the guard 409 a load that fits. - # An explicitly forced DFlash is not the Auto race and still pays. - if dflash_candidates and not (dspark_first and dspark_candidates): - # dflash_repo_preference_key, not the name-only key: it is the key the - # downloader sorts with, and in a multi-family repo the two disagree - # about which sidecar this weight gets. - total += min( - dflash_candidates, - key = lambda c: dflash_repo_preference_key(c[0], weight_name, other_weight_names), - )[1] - return total + # DSpark takes first refusal in the Auto promotion, so a listed + # sidecar settles the load: the DFlash fetch stands down and + # mtp_draft_path is replaced by the DSpark one. Charging the other + # two is not the safe over-estimate it is for a repo whose listing + # has not answered yet, it is a 409 for a load that fits. + return total + dspark_bytes + if dflash_sizes: + # DFlash is the other Auto promotion and replaces mtp_draft_path the + # same way, so the two are never resident together. Which of them it + # is stays genuinely unknown here: every DFlash candidate can still be + # turned away on its header, and the load then keeps the MTP drafter + # it has already fetched. The larger of the two covers both outcomes. + return total + max(dflash_bytes, mtp_bytes) + return total + mtp_bytes except Exception as e: logger.warning(f"Could not size GGUF companions for {repo}: {e}") return 0 @@ -5707,11 +5735,6 @@ def _same_file_key(p: str) -> str: main_bytes = selected.size_bytes if selected is not None else None if main_bytes is None: return None - # The variant record names the file this load opens, which is what the - # DFlash sizing needs to price the sidecar the loader will pair with - # it. A lister that reported no name leaves it None and the ranking - # falls back to precision alone, exactly as before. - selected_weight = Path(getattr(selected, "filename", "") or "").name or None companions = _remote_gguf_companion_bytes( repo, hf_token = hf_token, @@ -5726,16 +5749,11 @@ def _same_file_key(p: str) -> str: ), include_dspark = (_dspark_capable and (_auto_dspark or dspark_requested)), include_dflash = (_dflash_capable and (_auto_dflash or dflash_requested)), - # ... except where the listing settles it: a repo shipping BOTH - # kinds only loads the DSpark one under Auto, so charging the - # DFlash sidecar too is not caution, it is a refusal for bytes - # that never land. + # ... except where the listing settles it. Auto launches exactly + # one drafter, in a fixed order, so once the listing says which + # kinds the repo has, charging the losers is not caution, it is a + # refusal for bytes that never become resident. dspark_first = _auto_dspark, - # The weight this load actually opens. A multi-family repo ships a - # DFlash sidecar per family and the loader pairs them by name, so - # without it the guard can price a foreign (and smaller) sidecar - # than the one that lands. - weight_name = selected_weight, ) # Plus the local --model-draft, if the caller named one: the repo # listing cannot see it, and it is resident next to these weights. @@ -7128,6 +7146,10 @@ def _resolve_config(): model_id = model_identifier, hf_token = request.hf_token, gguf_variant = request.gguf_variant, + # A native grant covers one directory, and this is the first + # pass that touches a drafter candidate, so the boundary has + # to travel with it rather than being applied afterwards. + drafter_accept = _native_drafter_accept if native_grant_backed else None, ) # Guard and call go to the worker together: from_identifier can import transformers @@ -7838,6 +7860,10 @@ def _resolve_config(): model_id = model_identifier, hf_token = request.hf_token, gguf_variant = request.gguf_variant, + # A native grant covers one directory, and this is the first + # pass that touches a drafter candidate, so the boundary has + # to travel with it rather than being applied afterwards. + drafter_accept = _native_drafter_accept if native_grant_backed else None, ) config = await asyncio.to_thread(_resolve_config) diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index e7c46e49b7b..e8b47b4970d 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1565,33 +1565,54 @@ def test_auto_tells_the_companion_sizing_that_dspark_comes_first(self): SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 4 * 1024**3), ] - def test_remote_dflash_sizing_prices_the_sidecar_this_weight_pairs_with(self): - """The loader ranks DFlash candidates against the weight being loaded - (dflash_repo_preference_key), so a multi-family repo hands model B the - generic sidecar. Sizing by the name-only key priced model A's smaller - one instead, and the guard admitted a load that then exhausts VRAM - beside a running training job.""" + def test_remote_dflash_sizing_bounds_every_candidate_the_fallback_can_reach(self): + """_download_dflash reads a candidate's header only after paying for the + bytes, and a rejection falls through to the next name in the ranking, so + the file that lands can be any candidate -- including one LARGER than the + best-ranked pick. Sizing the first-ranked entry alone under-charged model + A by 3 GiB and admitted a load that then exhausts VRAM beside a running + training job. Headers are unreadable from a listing, so the bound has to + cover the whole reachable set.""" with patch( "huggingface_hub.model_info", return_value = SimpleNamespace(siblings = self._MULTI_FAMILY_SIBLINGS), ): - for weight, expected_gib in ( - ("model-B-Q4_K_M.gguf", 4), - ("model-A-Q4_K_M.gguf", 1), - ): - total = self.route._remote_gguf_companion_bytes( - "org/repo", - hf_token = None, - include_mmproj = False, - include_mtp = False, - include_dflash = True, - weight_name = weight, - ) - self.assertEqual(total, expected_gib * 1024**3, weight) + total = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dflash = True, + ) + # 4 GiB, the largest reachable candidate, for either weight in the repo: + # model A's own 1 GiB sidecar is merely the one tried FIRST. + self.assertEqual(total, 4 * 1024**3) + + def test_remote_dflash_sizing_ignores_a_nested_dflash_named_weight(self): + """The picker is root level only, so a quants/dflash-*.gguf is an ordinary + weight there and can never be fetched as the drafter. Charging it made the + bound track a file the load cannot reach.""" + siblings = [ + SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10 * 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 1024**3), + SimpleNamespace(rfilename = "quants/dflash-model-Q8_0.gguf", size = 9 * 1024**3), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + total = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dflash = True, + ) + self.assertEqual(total, 1024**3) - def test_remote_estimate_passes_the_selected_weight_to_the_dflash_sizing(self): - """End to end: the guard's own estimate has to carry the selected - filename down, or the sizing has nothing to pair the sidecar against.""" + def test_remote_estimate_bounds_the_dflash_fallback_end_to_end(self): + """End to end: the guard's own estimate has to carry the same bound, or + the multi-family repo above is under-charged by the whole difference.""" import utils.models.model_config as mc cfg = SimpleNamespace( @@ -1604,7 +1625,7 @@ def test_remote_estimate_passes_the_selected_weight_to_the_dflash_sizing(self): gguf_variant = "Q4_K_M", ) variant = SimpleNamespace( - filename = "model-B-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10 * 1024**3 + filename = "model-A-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10 * 1024**3 ) with ( patch.object(mc, "list_gguf_variants", lambda repo, hf_token = None: ([variant], False)), @@ -1615,10 +1636,136 @@ def test_remote_estimate_passes_the_selected_weight_to_the_dflash_sizing(self): self._dflash_capable(), ): gb = self.route._estimate_gguf_required_gb(cfg, speculative_type = "dflash") - # 10 GiB of weights plus the 4 GiB generic sidecar model B actually gets, - # not the 1 GiB one named after model A. + # 10 GiB of weights plus the 4 GiB the fallback can still land on, not the + # 1 GiB candidate that merely goes first. self.assertAlmostEqual(gb, 14.0, places = 6) + # ── Auto charges ONE drafter, the one the promotion leaves resident ── + + def _auto_companion_bytes(self, siblings): + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + return self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = True, + include_dspark = True, + include_dflash = True, + dspark_first = True, + ) + + def test_auto_does_not_charge_the_mtp_drafter_dflash_replaces(self): + """Under Auto the caller asks for MTP and DFlash together, but the loader + promotes DFlash and overwrites mtp_draft_path with it, so the two are + never resident at once. Charging the sum was a 409 for a load that fits. + + The DFlash sidecar is the larger of the two here, so the bound is its + size alone -- the MTP bytes are not added on top.""" + siblings = [ + SimpleNamespace(rfilename = "mtp-model.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 3 * 1024**3), + ] + self.assertEqual(self._auto_companion_bytes(siblings), 3 * 1024**3) + + def test_auto_keeps_the_mtp_charge_when_the_dflash_candidates_may_all_fail(self): + """The other half of the same rule: every DFlash candidate can still be + turned away on its header, and the load then keeps the MTP drafter it has + already fetched. That outcome is genuinely unknown from a listing, so the + larger of the two is charged -- here the MTP one.""" + siblings = [ + SimpleNamespace(rfilename = "mtp-model.gguf", size = 5 * 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 1024**3), + ] + self.assertEqual(self._auto_companion_bytes(siblings), 5 * 1024**3) + + def test_auto_charges_the_largest_reachable_dflash_against_the_mtp_drafter(self): + """Items 2 and 5 together, which is the only way they are coherent: the + DFlash side of the comparison is the whole reachable candidate set (4 + GiB), not the first-ranked pick (1 GiB), and it is compared against the + MTP drafter rather than added to it. Fixing only one of the two lands on + the wrong number from either side: summing the first-ranked pick charges + 3 GiB, and comparing against the first-ranked pick charges 2 GiB.""" + siblings = [ + SimpleNamespace(rfilename = "mtp-model.gguf", size = 2 * 1024**3), + *self._MULTI_FAMILY_SIBLINGS, + ] + self.assertEqual(self._auto_companion_bytes(siblings), 4 * 1024**3) + + def test_auto_charges_dspark_alone_over_both_of_the_others(self): + """DSpark takes first refusal in the promotion and has no post-fetch + rejection, so a listed sidecar settles the load: the DFlash fetch stands + down and mtp_draft_path is replaced. Neither of the other two is + resident.""" + siblings = [ + SimpleNamespace(rfilename = "mtp-model.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dspark/dspark-model-Q8_0.gguf", size = 2 * 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 3 * 1024**3), + ] + self.assertEqual(self._auto_companion_bytes(siblings), 2 * 1024**3) + + def test_auto_still_charges_the_mtp_drafter_when_the_repo_ships_no_sidecar(self): + """Positive control: with nothing to promote, Auto launches the MTP + drafter and it keeps its charge.""" + siblings = [ + SimpleNamespace(rfilename = "mtp-model.gguf", size = 1024**3), + SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10 * 1024**3), + ] + self.assertEqual(self._auto_companion_bytes(siblings), 1024**3) + + def test_an_explicit_request_is_not_the_auto_race(self): + """dspark_first off means the caller already narrowed the kinds to the one + it asked for, so nothing here may drop a charge it passed in.""" + siblings = [ + SimpleNamespace(rfilename = "mtp-model.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 3 * 1024**3), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + total = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = True, + include_dflash = True, + ) + self.assertEqual(total, 4 * 1024**3) + + def test_native_drafter_accept_applies_the_lease_before_the_scan_reads(self): + """The load route's boundary, in the shape ModelConfig.from_identifier + takes. Discovery runs inside from_identifier and opens a DFlash + candidate's header, so a dflash-*.gguf symlinked out of the granted + directory was read before the validated rescan could reject it, and no + later rejection takes a read back.""" + import os + import tempfile + + with tempfile.TemporaryDirectory() as d: + leased = Path(d) / "leased" + leased.mkdir() + outside = Path(d) / "outside" + outside.mkdir() + weight = leased / "model-Q4_K_M.gguf" + weight.write_bytes(b"x") + inside = leased / "dflash-kquant.gguf" + inside.write_bytes(b"y") + target = outside / "dflash-escape.gguf" + target.write_bytes(b"z") + escape = leased / "dflash-escape.gguf" + os.symlink(target, escape) + + accept = self.route._native_drafter_accept + self.assertTrue( + accept(str(inside), str(weight), "dflash", str(leased)) + ) + self.assertFalse( + accept(str(target.resolve()), str(weight), "dflash", str(leased)) + ) + def test_remote_unknown_variant_returns_none(self): import utils.models.model_config as mc cfg = SimpleNamespace( diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index a3079c6c999..98afe857cd6 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2645,3 +2645,330 @@ def test_auto_fetches_dflash_when_the_repo_ships_no_dspark_sidecar(monkeypatch): dspark_cached = None, ) assert seen["dflash_fetched"] is True + + +# ── The caller's boundary reaches discovery, not just the rescan ────── +# +# ModelConfig.from_identifier runs the local companion scan, and the DFlash scan +# opens a candidate's header to confirm the architecture. A native grant covers +# one directory, so a dflash-*.gguf inside it can be a symlink whose target sits +# outside the lease. The load route rejects that afterwards, which cannot undo a +# read, so the boundary has to travel INTO the scan. + + +def test_from_identifier_hands_the_boundary_to_every_drafter_kind(tmp_path, monkeypatch): + """All three kinds, not only the one that reads a header: they are the same + discovery, and a kind that skipped the check would hand the load route a + sidecar it has to reject a second time.""" + import utils.models.model_config as mc + + weight = _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + seen: dict[str, tuple] = {} + + def _recorder(kind): + def _detect(path, search_root = None, accept = None, **kwargs): + seen[kind] = (path, search_root, accept) + return None + + return _detect + + for kind, name in ( + ("mtp", "detect_mtp_file"), + ("dspark", "detect_dspark_file"), + ("dflash", "detect_dflash_file"), + ): + monkeypatch.setattr(mc, name, _recorder(kind)) + + calls: list[tuple[str, str, str, str]] = [] + + def _accept(candidate, gguf_file, kind, search_root): + calls.append((candidate, gguf_file, kind, search_root)) + return False + + config = ModelConfig.from_identifier(str(weight), drafter_accept = _accept) + + assert config is not None + assert set(seen) == {"mtp", "dspark", "dflash"} + for kind, (path, search_root, accept) in seen.items(): + assert path == str(weight) + assert accept is not None, kind + # Bound to this load's file, this kind and this search root, so the three + # closures cannot be swapped for one another. + assert accept("/candidate.gguf") is False + assert calls[-1] == ("/candidate.gguf", str(weight), kind, search_root) + + +def test_from_identifier_without_a_boundary_scans_exactly_as_before(tmp_path, monkeypatch): + """Every caller that has no lease to impose passes nothing, and must see the + same candidates in the same order.""" + import utils.models.model_config as mc + + weight = _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + accepts: list = [] + + def _detect(path, search_root = None, accept = None, **kwargs): + accepts.append(accept) + return None + + for name in ("detect_mtp_file", "detect_dspark_file", "detect_dflash_file"): + monkeypatch.setattr(mc, name, _detect) + + ModelConfig.from_identifier(str(weight)) + assert accepts == [None, None, None] + + +def test_from_identifier_never_reads_a_sidecar_outside_the_boundary(tmp_path, monkeypatch): + """End to end: the escaping symlink's target is never opened, and the config + reports no DFlash sidecar rather than one the load route would reject.""" + import os + + import utils.models.model_config as mc + + leased = tmp_path / "leased" + leased.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + weight = _write_gguf(leased / "model-Q4_K_M.gguf", "llama") + target = _write_gguf(outside / "dflash-kquant.gguf", "dflash") + os.symlink(target, leased / "dflash-kquant.gguf") + + reads: list[str] = [] + real_check = mc.is_dflash_architecture + + def _recording_check(path, *args, **kwargs): + reads.append(str(path)) + return real_check(path, *args, **kwargs) + + monkeypatch.setattr(mc, "is_dflash_architecture", _recording_check) + + def _inside_the_lease(candidate, gguf_file, kind, search_root): + return Path(search_root) in Path(candidate).parents + + config = ModelConfig.from_identifier(str(weight), drafter_accept = _inside_the_lease) + + assert config is not None + assert config.gguf_dflash_file is None + assert reads == [] # the out-of-lease target's header was never read + + +# ── Remote DFlash discovery is root level only, like the local scan ─── +# +# The local contract is a root-level dflash- file (detect_dflash_file never +# offers a nested one, since dflash/ is a family name a user picks for real +# weights). The remote paths matched the basename in any nested directory, so an +# ordinary quants/dflash-*.gguf weight became a candidate -- and the header can +# only be read once the bytes are here, so the whole weight downloaded before the +# rejection. + + +@pytest.mark.parametrize( + "path,expected", + [ + ("dflash-kquant.gguf", True), + ("dflash-model-Q8_0.gguf", True), + ("quants/dflash-kquant.gguf", False), + ("dflash/dflash-kquant.gguf", False), + (r"quants\dflash-kquant.gguf", False), + ("model-Q4_K_M.gguf", False), + ("model-dflash-Q8_0.gguf", False), # prefix-only naming rule, unchanged + ], +) +def test_is_root_dflash_drafter_path(path, expected): + from core.inference.llama_cpp import _is_root_dflash_drafter_path + + assert _is_root_dflash_drafter_path(path) is expected + + +def test_the_basename_predicate_keeps_its_own_semantics(): + """The root check is a separate predicate on purpose: _is_dflash_drafter_path + is shared with callers that only ever have a bare filename.""" + from core.inference.llama_cpp import _is_dflash_drafter_path + + assert _is_dflash_drafter_path("quants/dflash-kquant.gguf") is True + assert _is_dflash_drafter_path("dflash-kquant.gguf") is True + + +def test_download_dflash_never_fetches_a_nested_dflash_named_weight(tmp_path, monkeypatch): + """The regression: the nested file is an ordinary weight the local scan would + never offer, and picking it spent its entire download before the header check + could turn it away.""" + _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + + b, got, fetched = _dflash_repo_download( + tmp_path, + monkeypatch, + listing = ["model-Q4_K_M.gguf", "quants/dflash-model-Q8_0.gguf"], + ) + + assert got is None + assert fetched == [] # nothing was paid for + assert b._dflash_sidecar_absent is True + + +def test_download_dflash_still_takes_the_root_sidecar_beside_a_nested_one(tmp_path, monkeypatch): + """Positive control: the nested name is skipped, not the whole repo.""" + _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") + sidecar = _write_gguf(tmp_path / "dflash-kquant.gguf", "dflash") + + b, got, fetched = _dflash_repo_download( + tmp_path, + monkeypatch, + listing = ["model-Q4_K_M.gguf", "quants/dflash-model-Q8_0.gguf", "dflash-kquant.gguf"], + ) + + assert got == str(sidecar) + assert fetched == ["dflash-kquant.gguf"] + + +def test_cached_dflash_lookup_ignores_a_nested_dflash_named_weight(tmp_path, monkeypatch): + """Same rule on the offline cache scan, which hands its answer straight to + --model-draft with no download to be rejected first.""" + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "snapshots" / "abc" + (snap / "quants").mkdir(parents = True) + _write_gguf(snap / "model-Q4_K_M.gguf", "llama") + _write_gguf(snap / "quants" / "dflash-model-Q8_0.gguf", "dflash") + monkeypatch.setattr( + "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] + ) + + b = LlamaCppBackend() + assert b._cached_repo_dflash_drafter( + "org/repo", near_path = str(snap / "model-Q4_K_M.gguf") + ) is None + + +# ── A split companion is fetched as a whole set ────────────────────── +# +# llama-server resolves a split drafter's sibling shards from the first shard's +# directory, so fetching only the picked shard left a drafter whose header reads +# fine and which the server then cannot open: the load fell back to no +# speculation with nothing to show for the download. The main-model downloader +# already resolves its shards with _gguf_extra_shards; the companion path reuses +# it rather than growing a second rule. + + +def _split_companion_download(tmp_path, monkeypatch, listing): + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr( + "huggingface_hub.list_repo_files", lambda repo, token = None: list(listing) + ) + monkeypatch.setattr(llama_cpp_module, "_hub_download_in_flight", lambda hf_repo: False) + fetched: list[str] = [] + + def _fake_download(repo, filename, token, *, cancel_event = None, cache_dir = None): + fetched.append(filename) + path = tmp_path / filename + path.parent.mkdir(parents = True, exist_ok = True) + path.write_bytes(b"x") + return str(path) + + monkeypatch.setattr(llama_cpp_module, "hf_hub_download_with_xet_fallback", _fake_download) + + def _pick(names): + return next((n for n in sorted(names) if Path(n).name.startswith("dflash-")), None) + + got = LlamaCppBackend()._download_companion_gguf( + hf_repo = "org/repo", + hf_token = None, + pick = _pick, + label = "DFlash drafter", + ) + return got, fetched + + +def test_download_companion_gguf_fetches_every_shard_of_a_split_sidecar(tmp_path, monkeypatch): + got, fetched = _split_companion_download( + tmp_path, + monkeypatch, + [ + "model-Q4_K_M.gguf", + "dflash-kquant-00001-of-00002.gguf", + "dflash-kquant-00002-of-00002.gguf", + ], + ) + + # The launch path is still shard 1, which is what llama-server is given. + assert got == str(tmp_path / "dflash-kquant-00001-of-00002.gguf") + assert fetched == [ + "dflash-kquant-00001-of-00002.gguf", + "dflash-kquant-00002-of-00002.gguf", + ] + + +def test_download_companion_gguf_leaves_a_single_file_sidecar_alone(tmp_path, monkeypatch): + """Every companion published today is one file, so the split handling must be + a no-op for them: exactly one fetch, same path back.""" + got, fetched = _split_companion_download( + tmp_path, monkeypatch, ["model-Q4_K_M.gguf", "dflash-kquant.gguf"] + ) + + assert got == str(tmp_path / "dflash-kquant.gguf") + assert fetched == ["dflash-kquant.gguf"] + + +def test_companion_snapshot_reuse_skips_an_incomplete_split_sidecar(tmp_path): + """The reuse scan is the other half: a snapshot holding shard 1 alone is not + a reuse. Reporting it would skip the download that completes the set and + leave the load with a drafter llama-server cannot open.""" + from core.inference.llama_cpp import _companion_snapshot_sibling + + snap = tmp_path / "models--org--repo" / "snapshots" / "abc" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"x") + first = snap / "dflash-kquant-00001-of-00002.gguf" + first.write_bytes(b"y") + + def _pick(names): + return next((n for n in sorted(names) if Path(n).name.startswith("dflash-")), None) + + near = str(snap / "model-Q4_K_M.gguf") + assert _companion_snapshot_sibling(near, _pick) is None + + (snap / "dflash-kquant-00002-of-00002.gguf").write_bytes(b"z") + assert _companion_snapshot_sibling(near, _pick) == str(first) + + +def test_offline_companion_cache_hit_skips_an_incomplete_split(tmp_path, monkeypatch): + """The offline cache lookup is the third way a shard can reach --model-draft, + and offline there is no fetch left to complete the set.""" + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr(llama_cpp_module, "_hub_download_in_flight", lambda hf_repo: False) + monkeypatch.setattr( + "huggingface_hub.list_repo_files", + lambda repo, token = None: ["dflash-kquant-00001-of-00002.gguf"], + ) + first = tmp_path / "dflash-kquant-00001-of-00002.gguf" + first.write_bytes(b"x") + monkeypatch.setattr( + llama_cpp_module, "_cached_hf_snapshot_file", lambda *a, **k: str(first) + ) + def _offline_fetch(*_args, **_kwargs): + # What the Hub raises offline, which the caller swallows to None. + raise RuntimeError("offline mode is enabled") + + monkeypatch.setattr( + llama_cpp_module, "hf_hub_download_with_xet_fallback", _offline_fetch + ) + + def _pick(names): + return next((n for n in sorted(names) if Path(n).name.startswith("dflash-")), None) + + b = LlamaCppBackend() + # The half set is not reported as a cache hit, so the load ends with no + # drafter rather than one llama-server cannot open. + assert b._download_companion_gguf( + hf_repo = "org/repo", hf_token = None, pick = _pick, label = "DFlash drafter" + ) is None + + (tmp_path / "dflash-kquant-00002-of-00002.gguf").write_bytes(b"y") + assert b._download_companion_gguf( + hf_repo = "org/repo", hf_token = None, pick = _pick, label = "DFlash drafter" + ) == str(first) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f31aae4d7cb..bfca5623eef 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -3833,6 +3833,7 @@ def from_identifier( hf_token: Optional[str] = None, is_lora: bool = False, gguf_variant: Optional[str] = None, + drafter_accept: Optional[Callable[[str, str, str, str], bool]] = None, ) -> Optional["ModelConfig"]: """Create ModelConfig from a clean model identifier (HF repo or local path), for FastAPI routes that send sanitized paths. @@ -3843,6 +3844,16 @@ def from_identifier( is_lora: Whether this is a LoRA adapter gguf_variant: Optional GGUF quant variant (e.g. "Q4_K_M") to load via -hf for remote repos; None auto-selects via _pick_best_gguf(). + drafter_accept: ``(candidate, gguf_file, kind, search_root) -> bool``, + the caller's extra admission rule for a discovered drafter. A + native-grant load passes the lease boundary here so it is applied + BEFORE this scan inspects a candidate: detect_dflash_file reads + the header of the file it is about to accept, and a + ``dflash-*.gguf`` symlink in a granted directory can point at a + target outside the lease, which the validated rescan on the load + route rejects only after the read already happened. Left None by + every caller that has no boundary to impose, which sees the same + candidates in the same order as before. Returns: ModelConfig or None if it cannot be created. @@ -3911,6 +3922,18 @@ def from_identifier( # Direct file selections may point into a quant subdir while mmproj-*.gguf sits at the root. companion_root = _local_gguf_companion_search_root(path, gguf_file) + + # One accept per drafter kind, bound to the file this load opens. + # Each kind admits a different companion directory, so they cannot + # share one closure or an MTP load would take a sidecar out of + # dspark/. + def _drafter_accept_for(kind: str) -> Optional[Callable[[str], bool]]: + if drafter_accept is None: + return None + return lambda candidate: drafter_accept( + candidate, gguf_file, kind, companion_root + ) + mmproj_file = detect_mmproj_file(gguf_file, search_root = companion_root) if mmproj_file: gguf_is_vision = True @@ -3919,11 +3942,27 @@ def from_identifier( logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}") # Separate MTP drafter sibling (Gemma 4), mirroring mmproj. - mtp_file = detect_mtp_file(gguf_file, search_root = companion_root) + mtp_file = detect_mtp_file( + gguf_file, + search_root = companion_root, + accept = _drafter_accept_for("mtp"), + ) if mtp_file: logger.info(f"Detected MTP drafter: {mtp_file}") - dspark_file = detect_dspark_file(gguf_file, search_root = companion_root) - dflash_file = detect_dflash_file(gguf_file, search_root = companion_root) + # DSpark and DFlash take the boundary for the same reason, even + # though only the DFlash scan opens a candidate: all three are the + # same discovery, and a kind that skipped the check would hand the + # load route a sidecar it has to reject a second time. + dspark_file = detect_dspark_file( + gguf_file, + search_root = companion_root, + accept = _drafter_accept_for("dspark"), + ) + dflash_file = detect_dflash_file( + gguf_file, + search_root = companion_root, + accept = _drafter_accept_for("dflash"), + ) return cls( identifier = identifier, From a2e02bec68edbb57eb5e38cad3f055288c9d38d1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:12:42 +0000 Subject: [PATCH 13/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/test_chat_load_during_training.py | 8 +-- .../tests/test_mtp_drafter_companion.py | 57 +++++++++++-------- 2 files changed, 36 insertions(+), 29 deletions(-) diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index e8b47b4970d..4fba33f9ddc 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1759,12 +1759,8 @@ def test_native_drafter_accept_applies_the_lease_before_the_scan_reads(self): os.symlink(target, escape) accept = self.route._native_drafter_accept - self.assertTrue( - accept(str(inside), str(weight), "dflash", str(leased)) - ) - self.assertFalse( - accept(str(target.resolve()), str(weight), "dflash", str(leased)) - ) + self.assertTrue(accept(str(inside), str(weight), "dflash", str(leased))) + self.assertFalse(accept(str(target.resolve()), str(weight), "dflash", str(leased))) def test_remote_unknown_variant_returns_none(self): import utils.models.model_config as mc diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 98afe857cd6..2df31dc0708 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2666,7 +2666,12 @@ def test_from_identifier_hands_the_boundary_to_every_drafter_kind(tmp_path, monk seen: dict[str, tuple] = {} def _recorder(kind): - def _detect(path, search_root = None, accept = None, **kwargs): + def _detect( + path, + search_root = None, + accept = None, + **kwargs, + ): seen[kind] = (path, search_root, accept) return None @@ -2706,7 +2711,12 @@ def test_from_identifier_without_a_boundary_scans_exactly_as_before(tmp_path, mo weight = _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") accepts: list = [] - def _detect(path, search_root = None, accept = None, **kwargs): + def _detect( + path, + search_root = None, + accept = None, + **kwargs, + ): accepts.append(accept) return None @@ -2775,7 +2785,6 @@ def _inside_the_lease(candidate, gguf_file, kind, search_root): ) def test_is_root_dflash_drafter_path(path, expected): from core.inference.llama_cpp import _is_root_dflash_drafter_path - assert _is_root_dflash_drafter_path(path) is expected @@ -2834,9 +2843,9 @@ def test_cached_dflash_lookup_ignores_a_nested_dflash_named_weight(tmp_path, mon ) b = LlamaCppBackend() - assert b._cached_repo_dflash_drafter( - "org/repo", near_path = str(snap / "model-Q4_K_M.gguf") - ) is None + assert ( + b._cached_repo_dflash_drafter("org/repo", near_path = str(snap / "model-Q4_K_M.gguf")) is None + ) # ── A split companion is fetched as a whole set ────────────────────── @@ -2854,13 +2863,18 @@ def _split_companion_download(tmp_path, monkeypatch, listing): from core.inference.llama_cpp import LlamaCppBackend monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) - monkeypatch.setattr( - "huggingface_hub.list_repo_files", lambda repo, token = None: list(listing) - ) + monkeypatch.setattr("huggingface_hub.list_repo_files", lambda repo, token = None: list(listing)) monkeypatch.setattr(llama_cpp_module, "_hub_download_in_flight", lambda hf_repo: False) fetched: list[str] = [] - def _fake_download(repo, filename, token, *, cancel_event = None, cache_dir = None): + def _fake_download( + repo, + filename, + token, + *, + cancel_event = None, + cache_dir = None, + ): fetched.append(filename) path = tmp_path / filename path.parent.mkdir(parents = True, exist_ok = True) @@ -2894,10 +2908,7 @@ def test_download_companion_gguf_fetches_every_shard_of_a_split_sidecar(tmp_path # The launch path is still shard 1, which is what llama-server is given. assert got == str(tmp_path / "dflash-kquant-00001-of-00002.gguf") - assert fetched == [ - "dflash-kquant-00001-of-00002.gguf", - "dflash-kquant-00002-of-00002.gguf", - ] + assert fetched == ["dflash-kquant-00001-of-00002.gguf", "dflash-kquant-00002-of-00002.gguf"] def test_download_companion_gguf_leaves_a_single_file_sidecar_alone(tmp_path, monkeypatch): @@ -2947,16 +2958,13 @@ def test_offline_companion_cache_hit_skips_an_incomplete_split(tmp_path, monkeyp ) first = tmp_path / "dflash-kquant-00001-of-00002.gguf" first.write_bytes(b"x") - monkeypatch.setattr( - llama_cpp_module, "_cached_hf_snapshot_file", lambda *a, **k: str(first) - ) + monkeypatch.setattr(llama_cpp_module, "_cached_hf_snapshot_file", lambda *a, **k: str(first)) + def _offline_fetch(*_args, **_kwargs): # What the Hub raises offline, which the caller swallows to None. raise RuntimeError("offline mode is enabled") - monkeypatch.setattr( - llama_cpp_module, "hf_hub_download_with_xet_fallback", _offline_fetch - ) + monkeypatch.setattr(llama_cpp_module, "hf_hub_download_with_xet_fallback", _offline_fetch) def _pick(names): return next((n for n in sorted(names) if Path(n).name.startswith("dflash-")), None) @@ -2964,9 +2972,12 @@ def _pick(names): b = LlamaCppBackend() # The half set is not reported as a cache hit, so the load ends with no # drafter rather than one llama-server cannot open. - assert b._download_companion_gguf( - hf_repo = "org/repo", hf_token = None, pick = _pick, label = "DFlash drafter" - ) is None + assert ( + b._download_companion_gguf( + hf_repo = "org/repo", hf_token = None, pick = _pick, label = "DFlash drafter" + ) + is None + ) (tmp_path / "dflash-kquant-00002-of-00002.gguf").write_bytes(b"y") assert b._download_companion_gguf( From f0033198d7d9713f662cc1a3b74d1a51e45b16eb Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 15:30:27 +0000 Subject: [PATCH 14/26] Budget a split DFlash sidecar as a set, and reject half a cached one The remote sizing bounded the DFlash fetch with the largest candidate the post-fetch fallback could land on, but each entry is one shard, while _download_companion_gguf fetches the whole shard set the picked file belongs to and llama-server keeps every shard resident. A sidecar published as two 1 GiB shards was budgeted at 1 GiB, and under-charging is the direction that waves a load through and then exhausts VRAM beside a running training job. The candidates are grouped into their sets with _gguf_extra_shards, the same helper the download resolves shards with, and the bound is the largest set total. _cached_repo_dflash_drafter's offline fallback accepted a candidate on is_file plus its header, so a snapshot holding shard 1 alone was handed back as the drafter with no fetch left to complete the set. The header reads fine, then llama-server cannot open the siblings it resolves from that directory and the load falls back to no speculation. Same _drafter_split_is_complete rule the snapshot reuse already applies, and skipped rather than fatal like the header check, since another snapshot may hold the complete copy. --- studio/backend/core/inference/llama_cpp.py | 14 ++++++ studio/backend/routes/inference.py | 19 +++++-- .../tests/test_chat_load_during_training.py | 49 +++++++++++++++++++ .../tests/test_mtp_drafter_companion.py | 45 +++++++++++++++++ 4 files changed, 124 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c4c8547c746..a1215788d4e 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8190,6 +8190,7 @@ def _cached_repo_dflash_drafter( """ try: from utils.models.model_config import ( + _drafter_split_is_complete, _iter_hf_cache_snapshots, dflash_repo_preference_key, is_dflash_architecture, @@ -8232,6 +8233,19 @@ def _cached_repo_dflash_drafter( candidate, ) continue + # Same whole-set rule _companion_snapshot_sibling applies, for the + # same reason: llama-server resolves the sibling shards from this + # one's directory, and what this lookup returns is handed straight + # back as the drafter with no fetch left to complete it. Half a + # set reads as a valid header and then cannot be opened, so the + # load quietly drops speculation. Skipped rather than fatal, since + # another snapshot may hold the complete copy. + if not _drafter_split_is_complete(candidate): + logger.info( + "Ignoring cached DFlash candidate %s: split shard set is incomplete", + candidate, + ) + continue return str(candidate) except Exception as exc: logger.debug("Cached DFlash drafter lookup failed for %s: %s", hf_repo, exc) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 92417042e6c..60c14d80e14 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5340,6 +5340,7 @@ def _remote_gguf_companion_bytes( """ try: from core.inference.llama_cpp import ( + _gguf_extra_shards, _is_dspark_drafter_path, _is_root_dflash_drafter_path, ) @@ -5350,7 +5351,7 @@ def _remote_gguf_companion_bytes( total = 0 mtp_bytes = 0 dspark_candidates: list[tuple[str, int]] = [] - dflash_sizes: list[int] = [] + dflash_sizes: dict[str, int] = {} for sibling in info.siblings or []: name = sibling.rfilename or "" base = Path(name).name.lower() @@ -5370,7 +5371,7 @@ def _remote_gguf_companion_bytes( # dflash-*.gguf is an ordinary weight there and never a candidate, so # counting it here would price a file the load cannot fetch. if include_dflash and _is_root_dflash_drafter_path(name): - dflash_sizes.append(size) + dflash_sizes[name] = size # Same preference order the download uses, so the budget sizes the file # the launch will actually fetch. DSpark has no post-fetch rejection, so # the best-ranked candidate is the one that lands. @@ -5387,7 +5388,19 @@ def _remote_gguf_companion_bytes( # one. Headers are unreadable from a listing, so the ranking cannot # narrow that down here, and over-estimating is the established safe # direction for a guard protecting a running training job. - dflash_bytes = max(dflash_sizes, default = 0) + # Each entry summed is a whole shard SET, not one file: a split sidecar + # is picked as its first shard and _download_companion_gguf then fetches + # every sibling, all of which llama-server keeps resident. Sizing one + # shard would halve a two-shard sidecar, and under-estimating is the + # direction that waves a load through and then exhausts VRAM. + dflash_bytes = max( + ( + size + + sum(dflash_sizes.get(shard, 0) for shard in _gguf_extra_shards(dflash_sizes, name)) + for name, size in dflash_sizes.items() + ), + default = 0, + ) if not dspark_first: return total + mtp_bytes + dspark_bytes + dflash_bytes if dspark_candidates: diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 4fba33f9ddc..4e37d63ca5e 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1588,6 +1588,55 @@ def test_remote_dflash_sizing_bounds_every_candidate_the_fallback_can_reach(self # model A's own 1 GiB sidecar is merely the one tried FIRST. self.assertEqual(total, 4 * 1024**3) + def test_remote_dflash_sizing_totals_every_shard_of_a_split_sidecar(self): + """A split sidecar is picked as its first shard, and the download then + fetches every sibling; llama-server keeps the whole set resident. Sizing + one shard budgeted a two-shard 2 GiB sidecar at 1 GiB and let it lose the + comparison to a smaller single-file candidate, which is the direction that + admits a load and then exhausts VRAM beside a running training job.""" + siblings = [ + SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10 * 1024**3), + SimpleNamespace(rfilename = "dflash-split-00001-of-00002.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dflash-split-00002-of-00002.gguf", size = 1024**3), + # Bigger than either shard, smaller than the set they form. + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 3 * 1024**3 // 2), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + total = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dflash = True, + ) + # The set totals 2 GiB and is the largest thing the fallback can land on. + self.assertEqual(total, 2 * 1024**3) + + def test_remote_dflash_sizing_charges_a_split_set_once(self): + """The other half: every shard is a listed dflash- name, so a rule that + totalled the candidates rather than taking the safe maximum across shard + SETS would double-charge this repo and 409 a load that fits.""" + siblings = [ + SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10 * 1024**3), + SimpleNamespace(rfilename = "dflash-split-00001-of-00002.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dflash-split-00002-of-00002.gguf", size = 1024**3), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + total = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dflash = True, + ) + self.assertEqual(total, 2 * 1024**3) + def test_remote_dflash_sizing_ignores_a_nested_dflash_named_weight(self): """The picker is root level only, so a quants/dflash-*.gguf is an ordinary weight there and can never be fetched as the drafter. Charging it made the diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 2df31dc0708..41b7212780d 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -2983,3 +2983,48 @@ def _pick(names): assert b._download_companion_gguf( hf_repo = "org/repo", hf_token = None, pick = _pick, label = "DFlash drafter" ) == str(first) + + +def test_cached_dflash_lookup_skips_an_incomplete_split_set(tmp_path, monkeypatch): + """The fourth way a shard can reach --model-draft: _download_dflash's offline + fallback hands this lookup's answer back as the drafter with no fetch left to + complete the set, so a lone shard reads as a valid header and llama-server + then cannot open the siblings it resolves from that directory. The load drops + speculation silently.""" + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "snapshots" / "abc" + snap.mkdir(parents = True) + weight = _write_gguf(snap / "model-Q4_K_M.gguf", "llama") + first = _write_gguf(snap / "dflash-kquant-00001-of-00002.gguf", "dflash") + monkeypatch.setattr( + "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] + ) + + b = LlamaCppBackend() + assert b._cached_repo_dflash_drafter("org/repo", near_path = str(weight)) is None + + _write_gguf(snap / "dflash-kquant-00002-of-00002.gguf", "dflash") + assert b._cached_repo_dflash_drafter("org/repo", near_path = str(weight)) == str(first) + + +def test_cached_dflash_lookup_falls_through_from_a_half_split_to_a_whole_one( + tmp_path, monkeypatch +): + """Skipped, not fatal, exactly as the header check is: the half set merely + ranks first, and the snapshot still holds a sidecar that can be launched.""" + from core.inference.llama_cpp import LlamaCppBackend + + snap = tmp_path / "snapshots" / "abc" + snap.mkdir(parents = True) + weight = _write_gguf(snap / "model-Q4_K_M.gguf", "llama") + # Q8_0 outranks BF16, so the incomplete set is the candidate tried first. + _write_gguf(snap / "dflash-kquant-Q8_0-00001-of-00002.gguf", "dflash") + whole = _write_gguf(snap / "dflash-kquant-BF16.gguf", "dflash") + monkeypatch.setattr( + "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] + ) + + assert LlamaCppBackend()._cached_repo_dflash_drafter( + "org/repo", near_path = str(weight) + ) == str(whole) From 3f22e03d6e492b26c61451d688a8343ad86b9c33 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:31:47 +0000 Subject: [PATCH 15/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 4 +++- studio/backend/tests/test_mtp_drafter_companion.py | 10 ++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 60c14d80e14..2a594a57e46 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5396,7 +5396,9 @@ def _remote_gguf_companion_bytes( dflash_bytes = max( ( size - + sum(dflash_sizes.get(shard, 0) for shard in _gguf_extra_shards(dflash_sizes, name)) + + sum( + dflash_sizes.get(shard, 0) for shard in _gguf_extra_shards(dflash_sizes, name) + ) for name, size in dflash_sizes.items() ), default = 0, diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 41b7212780d..77bb37c3818 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -3008,9 +3008,7 @@ def test_cached_dflash_lookup_skips_an_incomplete_split_set(tmp_path, monkeypatc assert b._cached_repo_dflash_drafter("org/repo", near_path = str(weight)) == str(first) -def test_cached_dflash_lookup_falls_through_from_a_half_split_to_a_whole_one( - tmp_path, monkeypatch -): +def test_cached_dflash_lookup_falls_through_from_a_half_split_to_a_whole_one(tmp_path, monkeypatch): """Skipped, not fatal, exactly as the header check is: the half set merely ranks first, and the snapshot still holds a sidecar that can be launched.""" from core.inference.llama_cpp import LlamaCppBackend @@ -3025,6 +3023,6 @@ def test_cached_dflash_lookup_falls_through_from_a_half_split_to_a_whole_one( "utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: [snap] ) - assert LlamaCppBackend()._cached_repo_dflash_drafter( - "org/repo", near_path = str(weight) - ) == str(whole) + assert LlamaCppBackend()._cached_repo_dflash_drafter("org/repo", near_path = str(weight)) == str( + whole + ) From e4c6a6a7c32cc8a6f873cabd2517d748bba8f04a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 16:14:08 +0000 Subject: [PATCH 16/26] Move drafter naming, ranking and DFlash discovery into a drafters package Pure structural move, no behaviour change. The shared primitives, the ranking keys and the DFlash detector were spread through model_config alongside unrelated model handling, and the same rules are reached from four different paths, so they now live in one package. model_config re-exports every moved name, so callers and tests that import them from there keep working. The package deliberately does not import model_config at module import time. The GGUF split and quant naming helpers stay where they are, since non-drafter code shares them, and are imported per call instead. --- .../backend/utils/models/drafters/__init__.py | 47 +++ .../backend/utils/models/drafters/common.py | 148 +++++++ .../backend/utils/models/drafters/dflash.py | 216 ++++++++++ .../utils/models/drafters/preference.py | 80 ++++ studio/backend/utils/models/model_config.py | 369 +----------------- 5 files changed, 512 insertions(+), 348 deletions(-) create mode 100644 studio/backend/utils/models/drafters/__init__.py create mode 100644 studio/backend/utils/models/drafters/common.py create mode 100644 studio/backend/utils/models/drafters/dflash.py create mode 100644 studio/backend/utils/models/drafters/preference.py diff --git a/studio/backend/utils/models/drafters/__init__.py b/studio/backend/utils/models/drafters/__init__.py new file mode 100644 index 00000000000..fc6d8d16abb --- /dev/null +++ b/studio/backend/utils/models/drafters/__init__.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Speculative drafter discovery, ranking and budgeting. + +One home for the rules a drafter sidecar obeys. Import from this facade rather +than from the submodules, so the internal split can change without touching +every call site. +""" + +from utils.models.drafters.common import ( + _drafter_launch_path, + _drafter_matches_weight, + _drafter_names_other_weight, + _drafter_pairing_stem, + _drafter_split_is_complete, + _drafter_stem_rank, + _drafter_total_size, +) +from utils.models.drafters.preference import ( + dflash_precision_rank, + dflash_preference_key, + dflash_repo_preference_key, + dspark_precision_rank, + dspark_preference_key, +) +from utils.models.drafters.dflash import ( + detect_dflash_file, + is_dflash_architecture, +) + +__all__ = [ + "_drafter_launch_path", + "_drafter_matches_weight", + "_drafter_names_other_weight", + "_drafter_pairing_stem", + "_drafter_split_is_complete", + "_drafter_stem_rank", + "_drafter_total_size", + "detect_dflash_file", + "dflash_precision_rank", + "dflash_preference_key", + "dflash_repo_preference_key", + "dspark_precision_rank", + "dspark_preference_key", + "is_dflash_architecture", +] diff --git a/studio/backend/utils/models/drafters/common.py b/studio/backend/utils/models/drafters/common.py new file mode 100644 index 00000000000..9fc251a3eae --- /dev/null +++ b/studio/backend/utils/models/drafters/common.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Primitives shared by every speculative drafter kind. + +MTP, DSpark and DFlash all pair a sidecar with a main weight by name, resolve a +launch path through a split set, and answer whether that set is complete. Those +rules live here, once, so a change to the pairing rule cannot reach one kind and +miss another. Nothing here is DFlash specific. +""" + +import os +import re +import sys +from pathlib import Path +from typing import Iterable, Optional + +# model_config imports this module, so the split and quant naming helpers below +# are pulled in per call rather than at module import time. They are constants +# and pure functions shared with non-drafter code (gguf_variants, the auto +# download paths), which is why they stay where they are rather than moving +# here: this package must not become a second home for the GGUF naming rules. + + +def _drafter_pairing_stem(name: str, *, kind: str) -> str: + """The model family a drafter filename names, stripped of its own markers. + + Both published schemes are handled: ``-`` and the older + ``-``. The shard suffix sits outside the quant token, so it + goes first or the anchored quant strip below cannot match. Full quant + vocabulary, not a subset: K/IQ/UD/MXFP drafters pair too, and the optional + bpw modifier goes with it, as _extract_quant_label does. + """ + stem = Path(name).stem.lower() + if stem.startswith(f"{kind}-"): + stem = stem[len(kind) + 1 :] + stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", stem) + if stem.endswith(f"-{kind}"): + stem = stem[: -(len(kind) + 1)] + from utils.models.model_config import _GGUF_KNOWN_QUANT_RE + + return re.sub( + rf"-(?:{_GGUF_KNOWN_QUANT_RE.pattern})(?:-[0-9]+(?:\.[0-9]+)?bpw)?$", + "", + stem, + flags = re.IGNORECASE, + ) + + +def _drafter_matches_weight(candidate_name: str, weight_name: Optional[str], *, kind: str) -> bool: + """Whether a drafter pairs with the weight, by name. + + A multi-model folder must not attach a foreign drafter, so the family the + drafter names has to PREFIX the weight filename at a non-alphanumeric + boundary. That blocks one direction of a ``DeepSeek-V4-Flash-Lite`` / + ``DeepSeek-V4-Flash`` pair but not the other: the shorter family name is a + prefix of the longer weight, so a base-family sidecar still matches a + longer-named sibling's weights. Exact equality cannot replace the prefix + rule -- ``mtp-gemma-4-12B-it.gguf`` really does ship beside + ``gemma-4-12B-it-qat-*.gguf`` -- so the remaining direction is settled by + ranking: callers prefer the longest matching stem (see _drafter_stem_rank). + """ + if weight_name is None: + return True + stem = _drafter_pairing_stem(candidate_name, kind = kind) + weight = weight_name.lower() + return ( + bool(stem) + and weight.startswith(stem) + and (len(weight) == len(stem) or not weight[len(stem)].isalnum()) + ) + + +def _drafter_stem_rank(candidate_name: str, *, kind: str) -> int: + """Sort key placing the most specific family first (longest stem wins). + + Both ``mtp-DeepSeek-V4-Flash-BF16.gguf`` and + ``mtp-DeepSeek-V4-Flash-0731-BF16.gguf`` prefix-match a 0731 weight, and + only the second is really its drafter. + """ + return -len(_drafter_pairing_stem(candidate_name, kind = kind) or "") + + +def _drafter_launch_path(candidate: Path) -> str: + """The path llama-server should receive for *candidate*. + + llama-server takes shard 1 as the model path, and a split copy must stay on + its snapshot path: the blob target has no sibling shard names. Single-file + drafters still resolve, as callers expect. + """ + from utils.models.model_config import _GGUF_SPLIT_FILE_RE, _local_gguf_load_path + + loadable = _local_gguf_load_path(candidate) + if _GGUF_SPLIT_FILE_RE.match(loadable.name): + return str(loadable) + return str(loadable.resolve()) + + +def _drafter_split_is_complete(candidate: Path) -> bool: + """False for a partial split set, which would fail llama-server's draft + startup and disable speculation entirely; skip it so a complete copy wins.""" + from utils.models.model_config import colocated_split_shards + + try: + _, complete = colocated_split_shards(candidate) + except OSError: + return False + return complete + + +def _drafter_total_size(candidate: Path) -> int: + """Bytes across every shard. Candidates are collapsed to shard 1, so a split + copy must be summed or it would outrank a smaller single file.""" + from utils.models.model_config import colocated_split_shards + + try: + shards, _ = colocated_split_shards(candidate) + return sum(shard.stat().st_size for shard in shards) + except OSError: + return sys.maxsize + + + +def _drafter_names_other_weight( + candidate_name: str, + weight_name: Optional[str], + other_weight_names: Iterable[str], + *, + kind: str = "dflash", +) -> bool: + """Whether a sidecar names a DIFFERENT weight sitting beside it. + + A sidecar that names no family at all (the published ``dflash-kquant.gguf``, + whose stem is a precision token) has to stay eligible, so "does it name a + family" cannot be answered from the sidecar name alone. It is answered + against the weights actually present instead: only a stem that pairs with + some OTHER weight in the same repo/folder is evidence the sidecar belongs to + that neighbour rather than to the weight being loaded. + """ + if weight_name is None: + return False + if _drafter_matches_weight(candidate_name, weight_name, kind = kind): + return False + return any( + _drafter_matches_weight(candidate_name, other, kind = kind) for other in other_weight_names + ) + + diff --git a/studio/backend/utils/models/drafters/dflash.py b/studio/backend/utils/models/drafters/dflash.py new file mode 100644 index 00000000000..45095c0971c --- /dev/null +++ b/studio/backend/utils/models/drafters/dflash.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""DFlash sidecar discovery for a local GGUF model. + +DFlash is published as a ``dflash-`` prefixed sibling of the weights it drafts +for, so discovery is a naming question first and a header question second. The +order of those two matters and is deliberate: a caller with a directory lease +(a native grant) hands in ``accept``, and that has to answer before anything is +opened, because reading the header of a symlink pointing out of the lease is +the very thing the lease exists to prevent, and no later rejection takes a read +back. +""" + +import logging +import os +from pathlib import Path +from typing import Callable, Optional + +from utils.models.gguf_metadata import read_gguf_general_metadata +from utils.models.drafters.common import ( + _drafter_launch_path, + _drafter_matches_weight, + _drafter_names_other_weight, + _drafter_split_is_complete, + _drafter_stem_rank, + _drafter_total_size, +) +from utils.models.drafters.preference import ( + dflash_precision_rank, + dflash_preference_key, + dflash_repo_preference_key, +) + +logger = logging.getLogger(__name__) + + +def is_dflash_architecture(path: str) -> bool: + """Whether a GGUF really is a DFlash sidecar, decided by its header. + + ``dflash-`` is a filename convention an ordinary weight can satisfy, by + accident or otherwise, and llama-server only discovers that at startup: it + refuses the file as ``--model-draft`` and the load falls back to no + speculation, after the bytes were already fetched. A DFlash sidecar declares + ``general.architecture = dflash``, which no real weight does, so that is what + settles it. + + Kept here, beside the naming rules, because the local scan + (detect_dflash_file) and the download / cache reuse in llama_cpp all have to + apply it -- a remote path that trusted the prefix alone would download + gigabytes the launch then cannot use. + """ + meta = read_gguf_general_metadata(str(path)) or {} + return (meta.get("general.architecture") or "").strip().lower() == "dflash" + + + +def detect_dflash_file( + path: str, + search_root: Optional[str] = None, + accept: Optional[Callable[[str], bool]] = None, +) -> Optional[str]: + """Find a DFlash sidecar for a local GGUF model. + + Two things differ from detect_dspark_file, both forced by how DFlash is + published: + + 1. Root level only. ``dspark/`` is always a publisher's companion folder, so + that scan is safe; ``dflash/`` is a family name a user picks for real + weights (the reason llama_cpp._DRAFTER_DIR_KINDS leaves it out), so + reaching into it would launch a weight copy as --model-draft. + 2. No filename pairing. The published sidecar is ``dflash-kquant.gguf``, + which names no model family at all, so _drafter_matches_weight would + reject the one file this exists to find. The header is checked instead: + a DFlash sidecar declares ``general.architecture = dflash``, which no + real weight does, and that is a stronger signal than a filename. It also + settles the adversarial case on its own, since a model merely CALLED + DFlash (``Qwen3.6-35B-A3B-DFlash-Q4_K_M.gguf``) reports its own + architecture. + + A sidecar that does name a family (``dflash-Qwen3.6-27B-BF16.gguf``, the + scheme ggml-org uses) still wins over an unnamed one for the weight it + matches, so a multi-model folder attaches the specific sidecar first. + + ``accept`` filters candidates in preference order, so a caller with extra + rules (a native lease) keeps scanning instead of treating the first + rejection as no sidecar at all. + """ + + # Imported per call: model_config imports this module. + from utils.models.model_config import _local_gguf_load_path + + def _rank(candidate: Path) -> tuple[int, int, int, int, str]: + # A sidecar naming THIS weight's family first, then any unpaired one, + # then precision, then total size so a split copy cannot outrank a + # smaller single file, then name for a stable order. + paired = _drafter_matches_weight(candidate.name, weight_name, kind = "dflash") + return ( + 0 if paired else 1, + _drafter_stem_rank(candidate.name, kind = "dflash") if paired else 0, + dflash_precision_rank(candidate.name), + _drafter_total_size(candidate), + candidate.name.lower(), + ) + + p = Path(path) + weight_name = p.name if p.suffix.lower() == ".gguf" else None + start_dir = p.parent if p.is_file() else p + dirs = [start_dir] + if search_root is not None: + dirs.append(Path(search_root)) + + candidates: list[Path] = [] + other_weights: list[str] = [] + seen: set[Path] = set() + # dict.fromkeys: search_root is the weight's own parent for a flat layout, + # and scanning it twice doubles the directory reads for nothing. + for root in dict.fromkeys(dirs): + try: + entries = list(root.iterdir()) + except OSError: + continue + for candidate in entries: + lower = candidate.name.lower() + if not lower.endswith(".gguf"): + continue + # Prefix form only, deliberately. The shared companion predicates + # (_drafter_path_kind, is_mtp_drafter_path) know DFlash by the + # dflash- prefix, so accepting -dflash.gguf here would let one + # file be a drafter for discovery AND a selectable Q8_0 main model in + # the quant picker, and choosing that variant would hand llama-server + # the drafter as the target. Teaching the predicate the suffix + # instead would hide a real model whose name merely ends in DFlash, + # which is the case #7811 exists to protect, so detection gives the + # form up rather than the picker giving up a model. No published + # DFlash sidecar uses it; the shipped one is dflash-kquant.gguf. + if not lower.startswith("dflash-"): + # Every other GGUF in the folder is a weight some sidecar could + # be naming. Recorded so a sidecar belonging to a NEIGHBOUR can + # be told apart from one naming no family at all (below). + other_weights.append(candidate.name) + continue + try: + # Collapse a split copy to shard 1 before ranking. + launch = _local_gguf_load_path(candidate) + # is_file() follows the link, so this also drops a dangling + # snapshot symlink and a directory named like a sidecar. Without + # it --model-draft gets a path llama-server cannot open, which + # fails the whole load rather than falling back to no + # speculation (detect_dspark_file guards the same way). + if not (launch.is_file() and _drafter_split_is_complete(launch)): + continue + resolved = launch.resolve() + except OSError: + continue + if resolved in seen: + continue + seen.add(resolved) + candidates.append(launch) + + # A sidecar naming a family that belongs to a NEIGHBOUR weight is that + # neighbour's drafter, not a generic one. _drafter_matches_weight is False + # both for it and for a sidecar naming no family (dflash-kquant.gguf), so + # ranking alone bucketed the two together and precision could float the + # foreign one to the top: loading model B beside dflash-model-A-Q8_0.gguf + # and dflash-kquant.gguf launched model A's drafter for model B. Both carry + # a real dflash header, so the architecture check behind the ranking cannot + # catch it. _drafter_names_other_weight decides against the weights actually + # present, which keeps the published unpaired sidecar eligible (its stem, + # "kquant", names no file here) without hardcoding which stems are precision + # tokens. Shared with the remote paths through dflash_repo_preference_key, + # so a download and a local scan agree on which sidecar belongs here. + if weight_name is not None and other_weights: + kept: list[Path] = [] + for candidate in candidates: + if _drafter_names_other_weight(candidate.name, weight_name, other_weights): + logger.info( + "detect_dflash_file: dropped %s (names another weight in this folder)", + candidate.name, + ) + continue + kept.append(candidate) + candidates = kept + + for candidate in sorted(candidates, key = _rank): + # Resolve and validate before opening anything. A dflash-*.gguf in a + # directory reached through a native grant can be a symlink whose target + # sits outside the lease, and ``accept`` is what decides that; reading the + # header first opened the target before the answer arrived, which a later + # rejection cannot undo. Callers without a grant pass accept = None and + # see the same candidates in the same order as before. + try: + launch = _drafter_launch_path(candidate) + except OSError: + continue + if accept is not None and not accept(launch): + logger.info( + "detect_dflash_file: dropped %s (outside the granted directory)", + candidate.name, + ) + continue + if not is_dflash_architecture(launch): + logger.info( + "detect_dflash_file: dropped %s (architecture %r is not dflash)", + candidate.name, + # Re-read only on the reject path, and header reads are cached by + # (path, mtime, size), so naming the offending architecture in the + # log costs nothing. + (read_gguf_general_metadata(launch) or {}).get("general.architecture"), + ) + continue + logger.info("Detected DFlash drafter: %s", launch) + return launch + return None + + diff --git a/studio/backend/utils/models/drafters/preference.py b/studio/backend/utils/models/drafters/preference.py new file mode 100644 index 00000000000..db15e00c75d --- /dev/null +++ b/studio/backend/utils/models/drafters/preference.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Ranking keys that decide which sidecar a load prefers. + +The download, the snapshot reuse, the offline cache lookup and the local scan +all order candidates with these, so a repo resolves to the same file whichever +path reaches it first. +""" + +from pathlib import Path +from typing import Iterable, Optional + +from utils.models.drafters.common import ( + _drafter_matches_weight, + _drafter_names_other_weight, + _drafter_stem_rank, +) + + +def dspark_precision_rank(name: str) -> int: + """Sidecar precision preference: Q8_0 first, the precision the DSpark model + card recommends. Shared with the hub download and VRAM-sizing paths so the + file Studio budgets for is the file it fetches and launches.""" + base = Path(name).name.lower() + if "-q8_0" in base: + return 0 + if "-q4_0" in base: + return 1 + if "-bf16" in base or "-f16" in base: + return 2 + return 3 + + +def dspark_preference_key(name: str) -> tuple[int, str]: + """Sort key picking the preferred sidecar by name alone (no filesystem).""" + return dspark_precision_rank(name), Path(name).name.lower() + + +# DFlash publishes the same precision vocabulary (and the published sidecar +# carries no precision token at all, which lands in the catch-all rank), so the +# ordering is shared rather than duplicated. +dflash_precision_rank = dspark_precision_rank + + +def dflash_preference_key(name: str) -> tuple[int, str]: + """Sort key picking the preferred DFlash sidecar by name alone.""" + return dflash_precision_rank(name), Path(name).name.lower() + + + +def dflash_repo_preference_key( + name: str, + weight_name: Optional[str] = None, + other_weight_names: Iterable[str] = (), +) -> tuple[int, int, int, str]: + """Order DFlash sidecars in a repo listing / cache snapshot against the + weight actually being loaded. + + dflash_preference_key ranks by precision and name alone, which is all a + single-model repo needs. A repo hosting more than one family also has to be + told which weight each sidecar belongs to, or ``dflash-model-A-Q8_0.gguf`` + outranks the generic ``dflash-kquant.gguf`` on precision and model B is + launched with model A's drafter. Same rule the local scan applies in + detect_dflash_file, kept in one place so the download, the snapshot reuse + and the offline cache all pick the same file. + + Three buckets: a sidecar naming this weight's family (most specific stem + first, as detect_mtp_file does), then one naming no weight present here, + then one naming a neighbour. The last is demoted rather than dropped, so a + repo whose only sidecar looks foreign still gets a fallback and today's + single-sidecar behaviour is unchanged. + """ + precision, sort_name = dflash_preference_key(name) + if weight_name is not None and _drafter_matches_weight(name, weight_name, kind = "dflash"): + return 0, _drafter_stem_rank(name, kind = "dflash"), precision, sort_name + foreign = _drafter_names_other_weight(name, weight_name, other_weight_names) + return 2 if foreign else 1, 0, precision, sort_name + + diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index bfca5623eef..4c9354ba7e6 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1656,198 +1656,27 @@ def _add(d: Path) -> None: return str(best[1]) -def _drafter_pairing_stem(name: str, *, kind: str) -> str: - """The model family a drafter filename names, stripped of its own markers. - - Both published schemes are handled: ``-`` and the older - ``-``. The shard suffix sits outside the quant token, so it - goes first or the anchored quant strip below cannot match. Full quant - vocabulary, not a subset: K/IQ/UD/MXFP drafters pair too, and the optional - bpw modifier goes with it, as _extract_quant_label does. - """ - stem = Path(name).stem.lower() - if stem.startswith(f"{kind}-"): - stem = stem[len(kind) + 1 :] - stem = re.sub(r"-[0-9]{5}-of-[0-9]{5}$", "", stem) - if stem.endswith(f"-{kind}"): - stem = stem[: -(len(kind) + 1)] - return re.sub( - rf"-(?:{_GGUF_KNOWN_QUANT_RE.pattern})(?:-[0-9]+(?:\.[0-9]+)?bpw)?$", - "", - stem, - flags = re.IGNORECASE, - ) - - -def _drafter_matches_weight(candidate_name: str, weight_name: Optional[str], *, kind: str) -> bool: - """Whether a drafter pairs with the weight, by name. - - A multi-model folder must not attach a foreign drafter, so the family the - drafter names has to PREFIX the weight filename at a non-alphanumeric - boundary. That blocks one direction of a ``DeepSeek-V4-Flash-Lite`` / - ``DeepSeek-V4-Flash`` pair but not the other: the shorter family name is a - prefix of the longer weight, so a base-family sidecar still matches a - longer-named sibling's weights. Exact equality cannot replace the prefix - rule -- ``mtp-gemma-4-12B-it.gguf`` really does ship beside - ``gemma-4-12B-it-qat-*.gguf`` -- so the remaining direction is settled by - ranking: callers prefer the longest matching stem (see _drafter_stem_rank). - """ - if weight_name is None: - return True - stem = _drafter_pairing_stem(candidate_name, kind = kind) - weight = weight_name.lower() - return ( - bool(stem) - and weight.startswith(stem) - and (len(weight) == len(stem) or not weight[len(stem)].isalnum()) - ) - - -def _drafter_stem_rank(candidate_name: str, *, kind: str) -> int: - """Sort key placing the most specific family first (longest stem wins). - - Both ``mtp-DeepSeek-V4-Flash-BF16.gguf`` and - ``mtp-DeepSeek-V4-Flash-0731-BF16.gguf`` prefix-match a 0731 weight, and - only the second is really its drafter. - """ - return -len(_drafter_pairing_stem(candidate_name, kind = kind) or "") - - -def _drafter_launch_path(candidate: Path) -> str: - """The path llama-server should receive for *candidate*. - - llama-server takes shard 1 as the model path, and a split copy must stay on - its snapshot path: the blob target has no sibling shard names. Single-file - drafters still resolve, as callers expect. - """ - loadable = _local_gguf_load_path(candidate) - if _GGUF_SPLIT_FILE_RE.match(loadable.name): - return str(loadable) - return str(loadable.resolve()) - - -def _drafter_split_is_complete(candidate: Path) -> bool: - """False for a partial split set, which would fail llama-server's draft - startup and disable speculation entirely; skip it so a complete copy wins.""" - try: - _, complete = colocated_split_shards(candidate) - except OSError: - return False - return complete - - -def _drafter_total_size(candidate: Path) -> int: - """Bytes across every shard. Candidates are collapsed to shard 1, so a split - copy must be summed or it would outrank a smaller single file.""" - try: - shards, _ = colocated_split_shards(candidate) - return sum(shard.stat().st_size for shard in shards) - except OSError: - return sys.maxsize - - -def dspark_precision_rank(name: str) -> int: - """Sidecar precision preference: Q8_0 first, the precision the DSpark model - card recommends. Shared with the hub download and VRAM-sizing paths so the - file Studio budgets for is the file it fetches and launches.""" - base = Path(name).name.lower() - if "-q8_0" in base: - return 0 - if "-q4_0" in base: - return 1 - if "-bf16" in base or "-f16" in base: - return 2 - return 3 - - -def dspark_preference_key(name: str) -> tuple[int, str]: - """Sort key picking the preferred sidecar by name alone (no filesystem).""" - return dspark_precision_rank(name), Path(name).name.lower() - - -# DFlash publishes the same precision vocabulary (and the published sidecar -# carries no precision token at all, which lands in the catch-all rank), so the -# ordering is shared rather than duplicated. -dflash_precision_rank = dspark_precision_rank - - -def dflash_preference_key(name: str) -> tuple[int, str]: - """Sort key picking the preferred DFlash sidecar by name alone.""" - return dflash_precision_rank(name), Path(name).name.lower() - - -def _drafter_names_other_weight( - candidate_name: str, - weight_name: Optional[str], - other_weight_names: Iterable[str], - *, - kind: str = "dflash", -) -> bool: - """Whether a sidecar names a DIFFERENT weight sitting beside it. - - A sidecar that names no family at all (the published ``dflash-kquant.gguf``, - whose stem is a precision token) has to stay eligible, so "does it name a - family" cannot be answered from the sidecar name alone. It is answered - against the weights actually present instead: only a stem that pairs with - some OTHER weight in the same repo/folder is evidence the sidecar belongs to - that neighbour rather than to the weight being loaded. - """ - if weight_name is None: - return False - if _drafter_matches_weight(candidate_name, weight_name, kind = kind): - return False - return any( - _drafter_matches_weight(candidate_name, other, kind = kind) for other in other_weight_names - ) - - -def dflash_repo_preference_key( - name: str, - weight_name: Optional[str] = None, - other_weight_names: Iterable[str] = (), -) -> tuple[int, int, int, str]: - """Order DFlash sidecars in a repo listing / cache snapshot against the - weight actually being loaded. - - dflash_preference_key ranks by precision and name alone, which is all a - single-model repo needs. A repo hosting more than one family also has to be - told which weight each sidecar belongs to, or ``dflash-model-A-Q8_0.gguf`` - outranks the generic ``dflash-kquant.gguf`` on precision and model B is - launched with model A's drafter. Same rule the local scan applies in - detect_dflash_file, kept in one place so the download, the snapshot reuse - and the offline cache all pick the same file. - - Three buckets: a sidecar naming this weight's family (most specific stem - first, as detect_mtp_file does), then one naming no weight present here, - then one naming a neighbour. The last is demoted rather than dropped, so a - repo whose only sidecar looks foreign still gets a fallback and today's - single-sidecar behaviour is unchanged. - """ - precision, sort_name = dflash_preference_key(name) - if weight_name is not None and _drafter_matches_weight(name, weight_name, kind = "dflash"): - return 0, _drafter_stem_rank(name, kind = "dflash"), precision, sort_name - foreign = _drafter_names_other_weight(name, weight_name, other_weight_names) - return 2 if foreign else 1, 0, precision, sort_name - - -def is_dflash_architecture(path: str) -> bool: - """Whether a GGUF really is a DFlash sidecar, decided by its header. - - ``dflash-`` is a filename convention an ordinary weight can satisfy, by - accident or otherwise, and llama-server only discovers that at startup: it - refuses the file as ``--model-draft`` and the load falls back to no - speculation, after the bytes were already fetched. A DFlash sidecar declares - ``general.architecture = dflash``, which no real weight does, so that is what - settles it. - - Kept here, beside the naming rules, because the local scan - (detect_dflash_file) and the download / cache reuse in llama_cpp all have to - apply it -- a remote path that trusted the prefix alone would download - gigabytes the launch then cannot use. - """ - meta = read_gguf_general_metadata(str(path)) or {} - return (meta.get("general.architecture") or "").strip().lower() == "dflash" - +# Drafter naming, ranking and DFlash discovery live in utils.models.drafters, so +# one rule serves the local scan, the download, the snapshot reuse and the +# offline cache alike. Re-exported here deliberately: these names were part of +# this module's surface and are imported from it across the backend and the +# tests, so the move stays source compatible. +from utils.models.drafters import ( # noqa: E402 + _drafter_launch_path, + _drafter_matches_weight, + _drafter_names_other_weight, + _drafter_pairing_stem, + _drafter_split_is_complete, + _drafter_stem_rank, + _drafter_total_size, + detect_dflash_file, + dflash_precision_rank, + dflash_preference_key, + dflash_repo_preference_key, + dspark_precision_rank, + dspark_preference_key, + is_dflash_architecture, +) def detect_mtp_file( path: str, @@ -2126,162 +1955,6 @@ def _rank(candidate: Path) -> tuple[int, int, int, str]: return None -def detect_dflash_file( - path: str, - search_root: Optional[str] = None, - accept: Optional[Callable[[str], bool]] = None, -) -> Optional[str]: - """Find a DFlash sidecar for a local GGUF model. - - Two things differ from detect_dspark_file, both forced by how DFlash is - published: - - 1. Root level only. ``dspark/`` is always a publisher's companion folder, so - that scan is safe; ``dflash/`` is a family name a user picks for real - weights (the reason llama_cpp._DRAFTER_DIR_KINDS leaves it out), so - reaching into it would launch a weight copy as --model-draft. - 2. No filename pairing. The published sidecar is ``dflash-kquant.gguf``, - which names no model family at all, so _drafter_matches_weight would - reject the one file this exists to find. The header is checked instead: - a DFlash sidecar declares ``general.architecture = dflash``, which no - real weight does, and that is a stronger signal than a filename. It also - settles the adversarial case on its own, since a model merely CALLED - DFlash (``Qwen3.6-35B-A3B-DFlash-Q4_K_M.gguf``) reports its own - architecture. - - A sidecar that does name a family (``dflash-Qwen3.6-27B-BF16.gguf``, the - scheme ggml-org uses) still wins over an unnamed one for the weight it - matches, so a multi-model folder attaches the specific sidecar first. - - ``accept`` filters candidates in preference order, so a caller with extra - rules (a native lease) keeps scanning instead of treating the first - rejection as no sidecar at all. - """ - - def _rank(candidate: Path) -> tuple[int, int, int, int, str]: - # A sidecar naming THIS weight's family first, then any unpaired one, - # then precision, then total size so a split copy cannot outrank a - # smaller single file, then name for a stable order. - paired = _drafter_matches_weight(candidate.name, weight_name, kind = "dflash") - return ( - 0 if paired else 1, - _drafter_stem_rank(candidate.name, kind = "dflash") if paired else 0, - dflash_precision_rank(candidate.name), - _drafter_total_size(candidate), - candidate.name.lower(), - ) - - p = Path(path) - weight_name = p.name if p.suffix.lower() == ".gguf" else None - start_dir = p.parent if p.is_file() else p - dirs = [start_dir] - if search_root is not None: - dirs.append(Path(search_root)) - - candidates: list[Path] = [] - other_weights: list[str] = [] - seen: set[Path] = set() - # dict.fromkeys: search_root is the weight's own parent for a flat layout, - # and scanning it twice doubles the directory reads for nothing. - for root in dict.fromkeys(dirs): - try: - entries = list(root.iterdir()) - except OSError: - continue - for candidate in entries: - lower = candidate.name.lower() - if not lower.endswith(".gguf"): - continue - # Prefix form only, deliberately. The shared companion predicates - # (_drafter_path_kind, is_mtp_drafter_path) know DFlash by the - # dflash- prefix, so accepting -dflash.gguf here would let one - # file be a drafter for discovery AND a selectable Q8_0 main model in - # the quant picker, and choosing that variant would hand llama-server - # the drafter as the target. Teaching the predicate the suffix - # instead would hide a real model whose name merely ends in DFlash, - # which is the case #7811 exists to protect, so detection gives the - # form up rather than the picker giving up a model. No published - # DFlash sidecar uses it; the shipped one is dflash-kquant.gguf. - if not lower.startswith("dflash-"): - # Every other GGUF in the folder is a weight some sidecar could - # be naming. Recorded so a sidecar belonging to a NEIGHBOUR can - # be told apart from one naming no family at all (below). - other_weights.append(candidate.name) - continue - try: - # Collapse a split copy to shard 1 before ranking. - launch = _local_gguf_load_path(candidate) - # is_file() follows the link, so this also drops a dangling - # snapshot symlink and a directory named like a sidecar. Without - # it --model-draft gets a path llama-server cannot open, which - # fails the whole load rather than falling back to no - # speculation (detect_dspark_file guards the same way). - if not (launch.is_file() and _drafter_split_is_complete(launch)): - continue - resolved = launch.resolve() - except OSError: - continue - if resolved in seen: - continue - seen.add(resolved) - candidates.append(launch) - - # A sidecar naming a family that belongs to a NEIGHBOUR weight is that - # neighbour's drafter, not a generic one. _drafter_matches_weight is False - # both for it and for a sidecar naming no family (dflash-kquant.gguf), so - # ranking alone bucketed the two together and precision could float the - # foreign one to the top: loading model B beside dflash-model-A-Q8_0.gguf - # and dflash-kquant.gguf launched model A's drafter for model B. Both carry - # a real dflash header, so the architecture check behind the ranking cannot - # catch it. _drafter_names_other_weight decides against the weights actually - # present, which keeps the published unpaired sidecar eligible (its stem, - # "kquant", names no file here) without hardcoding which stems are precision - # tokens. Shared with the remote paths through dflash_repo_preference_key, - # so a download and a local scan agree on which sidecar belongs here. - if weight_name is not None and other_weights: - kept: list[Path] = [] - for candidate in candidates: - if _drafter_names_other_weight(candidate.name, weight_name, other_weights): - logger.info( - "detect_dflash_file: dropped %s (names another weight in this folder)", - candidate.name, - ) - continue - kept.append(candidate) - candidates = kept - - for candidate in sorted(candidates, key = _rank): - # Resolve and validate before opening anything. A dflash-*.gguf in a - # directory reached through a native grant can be a symlink whose target - # sits outside the lease, and ``accept`` is what decides that; reading the - # header first opened the target before the answer arrived, which a later - # rejection cannot undo. Callers without a grant pass accept = None and - # see the same candidates in the same order as before. - try: - launch = _drafter_launch_path(candidate) - except OSError: - continue - if accept is not None and not accept(launch): - logger.info( - "detect_dflash_file: dropped %s (outside the granted directory)", - candidate.name, - ) - continue - if not is_dflash_architecture(launch): - logger.info( - "detect_dflash_file: dropped %s (architecture %r is not dflash)", - candidate.name, - # Re-read only on the reject path, and header reads are cached by - # (path, mtime, size), so naming the offending architecture in the - # log costs nothing. - (read_gguf_general_metadata(launch) or {}).get("general.architecture"), - ) - continue - logger.info("Detected DFlash drafter: %s", launch) - return launch - return None - - def _registered_custom_model_root(path: str) -> Optional[Path]: try: from storage.studio_db import list_scan_folders From ebd0ef29321ab98be55666ecf748fe8c5892b730 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 10 Aug 2026 16:16:35 +0000 Subject: [PATCH 17/26] Give the guard's DFlash bound a name and a home The bound was fifteen lines of generator plus the comment explaining why it is a max over shard sets rather than the best-ranked candidate, inline in the middle of a function that also sizes mmproj, MTP and DSpark. It is pure arithmetic over a listing, so it moves to drafters.budget with the reasoning attached to it. --- studio/backend/routes/inference.py | 27 ++---------- .../backend/utils/models/drafters/__init__.py | 2 + .../backend/utils/models/drafters/budget.py | 43 +++++++++++++++++++ 3 files changed, 49 insertions(+), 23 deletions(-) create mode 100644 studio/backend/utils/models/drafters/budget.py diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2a594a57e46..87c4863cdbe 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5345,6 +5345,7 @@ def _remote_gguf_companion_bytes( _is_root_dflash_drafter_path, ) from huggingface_hub import model_info + from utils.models.drafters import dflash_budget_bytes from utils.models.model_config import dspark_preference_key info = model_info(repo, token = hf_token, files_metadata = True) @@ -5380,29 +5381,9 @@ def _remote_gguf_companion_bytes( if dspark_candidates else 0 ) - # The largest candidate the fetch could end up on, not the best-ranked - # one. _download_dflash can only read a candidate's header once it has - # paid for the bytes, and a rejection falls through to the next name in - # the ranking, so any candidate can be the file that lands -- and the - # whole point of the fallback is the case where it is a different, bigger - # one. Headers are unreadable from a listing, so the ranking cannot - # narrow that down here, and over-estimating is the established safe - # direction for a guard protecting a running training job. - # Each entry summed is a whole shard SET, not one file: a split sidecar - # is picked as its first shard and _download_companion_gguf then fetches - # every sibling, all of which llama-server keeps resident. Sizing one - # shard would halve a two-shard sidecar, and under-estimating is the - # direction that waves a load through and then exhausts VRAM. - dflash_bytes = max( - ( - size - + sum( - dflash_sizes.get(shard, 0) for shard in _gguf_extra_shards(dflash_sizes, name) - ) - for name, size in dflash_sizes.items() - ), - default = 0, - ) + # Bounded rather than picked: see dflash_budget_bytes for why the max + # over whole shard sets is the answer a listing can give. + dflash_bytes = dflash_budget_bytes(dflash_sizes, _gguf_extra_shards) if not dspark_first: return total + mtp_bytes + dspark_bytes + dflash_bytes if dspark_candidates: diff --git a/studio/backend/utils/models/drafters/__init__.py b/studio/backend/utils/models/drafters/__init__.py index fc6d8d16abb..31eac8ed860 100644 --- a/studio/backend/utils/models/drafters/__init__.py +++ b/studio/backend/utils/models/drafters/__init__.py @@ -24,6 +24,7 @@ dspark_precision_rank, dspark_preference_key, ) +from utils.models.drafters.budget import dflash_budget_bytes from utils.models.drafters.dflash import ( detect_dflash_file, is_dflash_architecture, @@ -38,6 +39,7 @@ "_drafter_stem_rank", "_drafter_total_size", "detect_dflash_file", + "dflash_budget_bytes", "dflash_precision_rank", "dflash_preference_key", "dflash_repo_preference_key", diff --git a/studio/backend/utils/models/drafters/budget.py b/studio/backend/utils/models/drafters/budget.py new file mode 100644 index 00000000000..aee4ab5a108 --- /dev/null +++ b/studio/backend/utils/models/drafters/budget.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""What a drafter costs the training coexistence guard. + +The guard admits an inference load only if it fits beside a running training +job, so it has to price the drafter that load will actually make resident. It +sees a repository listing, never a header, which is what makes this its own +problem rather than a detail of discovery: the rules that decide WHICH sidecar +lands cannot all be evaluated here, so the budget bounds them instead. +""" + +from typing import Callable, Mapping + + +def dflash_budget_bytes( + sizes: Mapping[str, int], + extra_shards: Callable[[Mapping[str, int], str], list], +) -> int: + """A safe bound on the DFlash sidecar a load may end up resident on. + + The largest candidate the fetch could end up on, not the best-ranked one. + The download can only read a candidate's header once it has paid for the + bytes, and a rejection falls through to the next name in the ranking, so any + candidate can be the file that lands, and the whole point of the fallback is + the case where it is a different, bigger one. Headers are unreadable from a + listing, so the ranking cannot narrow that down here, and over-estimating is + the established safe direction for a guard protecting a running training + job. + + Each entry summed is a whole shard SET, not one file: a split sidecar is + picked as its first shard and the companion download then fetches every + sibling, all of which llama-server keeps resident. Sizing one shard would + halve a two-shard sidecar, and under-estimating is the direction that waves + a load through and then exhausts VRAM. + """ + return max( + ( + size + sum(sizes.get(shard, 0) for shard in extra_shards(sizes, name)) + for name, size in sizes.items() + ), + default = 0, + ) From 852e642f0be67a5969a46e9fe7a03769fad78e36 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:20:00 +0000 Subject: [PATCH 18/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/utils/models/drafters/budget.py | 3 +-- studio/backend/utils/models/drafters/common.py | 4 ---- studio/backend/utils/models/drafters/dflash.py | 3 --- studio/backend/utils/models/drafters/preference.py | 3 --- studio/backend/utils/models/model_config.py | 1 + 5 files changed, 2 insertions(+), 12 deletions(-) diff --git a/studio/backend/utils/models/drafters/budget.py b/studio/backend/utils/models/drafters/budget.py index aee4ab5a108..44aa1a03388 100644 --- a/studio/backend/utils/models/drafters/budget.py +++ b/studio/backend/utils/models/drafters/budget.py @@ -14,8 +14,7 @@ def dflash_budget_bytes( - sizes: Mapping[str, int], - extra_shards: Callable[[Mapping[str, int], str], list], + sizes: Mapping[str, int], extra_shards: Callable[[Mapping[str, int], str], list] ) -> int: """A safe bound on the DFlash sidecar a load may end up resident on. diff --git a/studio/backend/utils/models/drafters/common.py b/studio/backend/utils/models/drafters/common.py index 9fc251a3eae..45a5eda436c 100644 --- a/studio/backend/utils/models/drafters/common.py +++ b/studio/backend/utils/models/drafters/common.py @@ -112,7 +112,6 @@ def _drafter_total_size(candidate: Path) -> int: """Bytes across every shard. Candidates are collapsed to shard 1, so a split copy must be summed or it would outrank a smaller single file.""" from utils.models.model_config import colocated_split_shards - try: shards, _ = colocated_split_shards(candidate) return sum(shard.stat().st_size for shard in shards) @@ -120,7 +119,6 @@ def _drafter_total_size(candidate: Path) -> int: return sys.maxsize - def _drafter_names_other_weight( candidate_name: str, weight_name: Optional[str], @@ -144,5 +142,3 @@ def _drafter_names_other_weight( return any( _drafter_matches_weight(candidate_name, other, kind = kind) for other in other_weight_names ) - - diff --git a/studio/backend/utils/models/drafters/dflash.py b/studio/backend/utils/models/drafters/dflash.py index 45095c0971c..c9dd1fdcb88 100644 --- a/studio/backend/utils/models/drafters/dflash.py +++ b/studio/backend/utils/models/drafters/dflash.py @@ -54,7 +54,6 @@ def is_dflash_architecture(path: str) -> bool: return (meta.get("general.architecture") or "").strip().lower() == "dflash" - def detect_dflash_file( path: str, search_root: Optional[str] = None, @@ -212,5 +211,3 @@ def _rank(candidate: Path) -> tuple[int, int, int, int, str]: logger.info("Detected DFlash drafter: %s", launch) return launch return None - - diff --git a/studio/backend/utils/models/drafters/preference.py b/studio/backend/utils/models/drafters/preference.py index db15e00c75d..9fc3a840307 100644 --- a/studio/backend/utils/models/drafters/preference.py +++ b/studio/backend/utils/models/drafters/preference.py @@ -48,7 +48,6 @@ def dflash_preference_key(name: str) -> tuple[int, str]: return dflash_precision_rank(name), Path(name).name.lower() - def dflash_repo_preference_key( name: str, weight_name: Optional[str] = None, @@ -76,5 +75,3 @@ def dflash_repo_preference_key( return 0, _drafter_stem_rank(name, kind = "dflash"), precision, sort_name foreign = _drafter_names_other_weight(name, weight_name, other_weight_names) return 2 if foreign else 1, 0, precision, sort_name - - diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 4c9354ba7e6..ce8ce03c5ec 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1678,6 +1678,7 @@ def _add(d: Path) -> None: is_dflash_architecture, ) + def detect_mtp_file( path: str, search_root: Optional[str] = None, From 32d1227ac2c6a98f5c078bbdb7df35b270ae05d9 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 08:42:45 +0000 Subject: [PATCH 19/26] Studio: fix the DFlash lint gate, and carry over what #6747 got right Five changes on top of the DFlash drafter work. Lint gate. The compatibility shim re-exporting the moved drafter helpers from model_config tripped scripts/verify_import_hoist.py, whose __all__ exemption is scoped to package __init__.py and which ships a reexport_in_ordinary_module_is_still_blocked self-test. The shim is gone: the module imports only what it still calls, and every other call site imports from utils.models.drafters directly. dspark_preference_key stays reachable from model_config as a delegating def, because repointing routes/inference.py's pre-existing function-local import is the verifier's TARGET-CHANGED case and its relocation exemption only covers module-level imports. Download plan. preferred_dflash_sibling in hub/utils/gguf_plan.py, and the sidecar as an expected file on every GgufVariantPlan. The sidecar was fetched but the hub manifest never knew about it, so download progress under-counted by ~1.5 GiB. Ranked with dflash_repo_preference_key, so the plan and the loader cannot disagree, and per variant, so a multi-family repo does not hand variant B the drafter named after variant A. Capability-regained retry. A load that stood down because llama-server could not run the drafter told the user to update, then deduped the reload the update was meant to repair. spec_binary_fallback_can_retry re-reads the binary, asking about the capability the drafter kind actually needs rather than the reason code, since every kind records the same binary_no_mtp. Transient fetch retry. _download_companion_gguf gained on_transient_failure, so a listing that never answered or a download that dropped is worth one more Apply. Permanent Hub errors, a full or unwritable cache, offline mode and cancellation are unaffected, and a header rejection still falls through to the next candidate rather than counting as transient. The probe cache key moved from (path, int(mtime)) to (path, st_mtime_ns, st_size), so an update landing in the same second as the probe is not answered with the old build's capabilities. CLI. unsloth_cli/_inference.py passes gguf_dflash_file into the GGUF load, so the managed CLI path engages DFlash instead of silently running without it. No vision gate, now measured rather than argued. Muse-Glimmer-30B UD-Q4_K_XL with mmproj-kquant and dflash-kquant, llama.cpp b10342, one B200, n_max=2, greedy, on a prompt carrying ~545 image tokens: 92.1 to 114.2 tok/s at 0.646 acceptance, greedy output byte-identical to the drafter-free run, no load failure. The comment at the Auto promotion site cited llama.cpp #22673, which is an MTP result, for a DFlash decision; it now cites the measurement. Also fixes test_from_identifier_never_reads_a_sidecar_outside_the_boundary, which patched is_dflash_architecture on the re-exporting module rather than the one detect_dflash_file resolves it in, so its reads == [] assertion held whether or not the lease boundary worked. --- studio/backend/core/inference/llama_cpp.py | 240 ++++++++++++++++-- studio/backend/hub/utils/gguf_plan.py | 63 ++++- .../tests/test_llama_cpp_mtp_detection.py | 216 ++++++++++++++++ .../tests/test_mtp_drafter_companion.py | 196 +++++++++++++- studio/backend/utils/models/model_config.py | 33 ++- unsloth_cli/_inference.py | 1 + unsloth_cli/tests/test_inference_chat.py | 91 +++++-- 7 files changed, 775 insertions(+), 65 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index a1215788d4e..e7f8f582409 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10,6 +10,7 @@ import ast import atexit import contextlib +import errno import functools import json import logging @@ -218,6 +219,39 @@ class _CpuFallbackRuntime(NamedTuple): "LLAMA_ARG_HFD_REPO", ) +# Hub failures that answer the question for good: the repo, the revision or the file is +# not there, or it is there and this token may not have it, or the Hub was never going to +# be asked. Retrying any of them, on this load or the next Apply, buys nothing. +_PERMANENT_HUB_ERRORS = ( + "RepositoryNotFoundError", + "GatedRepoError", + "RevisionNotFoundError", + "EntryNotFoundError", + "OfflineModeIsEnabled", +) + +# The other kind of settled failure a download can hit: this machine, not the Hub. Told +# apart by errno rather than by exception type, because the Hub client raises OSError +# subclasses for network trouble too (requests' ConnectionError is one of them) and +# those DO deserve another attempt. +_UNRECOVERABLE_LOCAL_ERRNOS = frozenset( + { + errno.EACCES, # cache dir not writable + errno.EPERM, + errno.ENOSPC, # disk full + errno.EROFS, + } +) + +# What llama-server has to advertise for each drafter kind. Every kind records the same +# "binary_no_mtp" when it stands down, so the reason alone cannot say which capability +# to ask about again. +_SPEC_KIND_CAPABILITY: dict[str, str] = { + "dspark": "supports_dspark", + "dflash": "supports_dflash", + "mtp": "supports_mtp", +} + _PARAVIRTUAL_DIFFUSION_NO_NGL_ERROR = ( "This Mac's Metal device is virtualised, where offloaded layers can return " "corrupt output, and the installed unsloth_zoo diffusion shim has no --ngl, " @@ -3238,6 +3272,16 @@ def __init__(self): self._spec_drafter_kind: Optional[str] = None self._dspark_sidecar_absent: bool = False self._dflash_sidecar_absent: bool = False + # Set when the DFlash sidecar could not be fetched for a reason that says + # nothing about the repo (a listing blip, a download that dropped). Unlike + # _dflash_sidecar_absent, which is the permanent answer for almost every repo, + # this one is worth one more Apply: under Auto the fetch failing silently + # leaves the load with no drafter and nothing else recording why. + self._dflash_retry_needed: bool = False + # Which llama-server file the last load resolved. `unsloth studio update` + # replaces it in place, so a load that blamed the binary can tell an update + # from the same build still being installed. + self._launch_binary_revision: tuple = () # Set after an auto-Vulkan crash recovers with all devices disabled. self._cpu_fallback_reason: Optional[str] = None self._cpu_fallback_runtime: Optional[_CpuFallbackRuntime] = None @@ -3494,6 +3538,50 @@ def spec_fallback_reason(self) -> Optional[str]: """Why MTP was disabled on the last MTP-requesting load, else None.""" return self._spec_fallback_reason + def _binary_changed_since_launch(self) -> bool: + """Whether a different llama-server file is installed than the live server was + launched from. False when either side is unreadable, so the install's own window + is not mistaken for a finished install: the next Apply asks again.""" + current = self._binary_revision(self._find_llama_server_binary()) + if not current or not self._launch_binary_revision: + return False + return current != self._launch_binary_revision + + def spec_binary_fallback_can_retry(self) -> bool: + """Whether the binary has since gained what the last load stood down for. + + A load that drops speculative decoding because llama-server cannot run it tells + the user to run `unsloth studio update`. Doing so changes nothing about the + request, so the duplicate-load comparators see an identical intent and skip the + reload: the model keeps serving without a drafter until something unrelated + forces a relaunch, and the hint reads as a lie. Re-reading the binary here is + what turns it into a working instruction. + + Which question to ask depends on what the binary was blamed for. A missing + ``--spec-type`` is answered by the capability the drafter kind needs, since every + kind records the same "binary_no_mtp" and asking about MTP after a DSpark + stand-down would say yes on a build that never gained draft-dspark. An + architecture the build did not know is advertised by no flag at all, so that one + can only compare the file. + """ + if self._spec_fallback_reason == "binary_outdated": + return self._binary_changed_since_launch() + if self._spec_fallback_reason != "binary_no_mtp": + return False + if self._launch_binary_revision and not self._binary_changed_since_launch(): + # An untouched file cannot advertise anything new, and this runs on every + # duplicate-load check, where the probe below is a subprocess on a cold + # cache -- which is precisely the state an update leaves it in. + return False + capability = _SPEC_KIND_CAPABILITY.get(self._spec_drafter_kind or "") + if capability is None: + return False + try: + return bool(self.probe_server_capabilities().get(capability)) + except Exception as exc: + logger.debug("Could not recheck the speculative capability: %s", exc) + return False + @property def extra_args(self) -> Optional[List[str]]: """Extra llama-server flags from the last load (a copy). None = @@ -4045,6 +4133,21 @@ def _norm(value): and not spec_owned_by_extra_args ): return False + # The other recoverable stand-down, and the one the UI actively asks the user to + # fix: the drafter was there and llama-server could not run it. Updating the + # binary leaves the request identical, so without this the load the update was + # meant to repair is exactly the load that never happens again. + if ( + speculative_type in ("auto", "mtp", "mtp+ngram", "dspark", "dflash") + and self.spec_binary_fallback_can_retry() + ): + return False + # A DFlash fetch that failed on the way out leaves no other trace under Auto: + # the promotion never ran, so no fallback reason was recorded and the branch + # above cannot see it. A repo that simply publishes no sidecar is + # _dflash_sidecar_absent's business and is never retried here. + if self._dflash_retry_needed and speculative_type in ("auto", "dflash"): + return False compared_draft_n_max = self._spec_draft_n_max if self._spec_fallback_reason == "runtime_error" and self._last_load_intent is not None: # The MTP-free recovery clears the runtime value but retains the @@ -4352,8 +4455,12 @@ def _dspark_release_is_broken(cls, release_tag: Optional[str]) -> bool: match = re.match(r"b(\d+)", str(release_tag or "")) return bool(match) and int(match.group(1)) in cls._BROKEN_DSPARK_BUILDS - # Cached on (path, mtime); `unsloth studio update` bumps mtime. - _capability_cache: dict[tuple[str, int], dict[str, object]] = {} + # Cached on the revision of the file, not on the path: `unsloth studio update` + # replaces the binary in place. Nanoseconds and size rather than int(st_mtime), + # since an update landing in the same second as the probe kept the key identical and + # the new build was answered with the old one's capabilities -- which is exactly the + # moment a capability the user just installed has to become visible. + _capability_cache: dict[tuple[str, int, int], dict[str, object]] = {} @classmethod def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]: @@ -4399,10 +4506,13 @@ def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, ob "spec_draft_ngl_flag": None, } try: - mtime = int(Path(bin_path).stat().st_mtime) + binary_stat = Path(bin_path).stat() + mtime_ns = binary_stat.st_mtime_ns + size = binary_stat.st_size except OSError: - mtime = 0 - cache_key = (bin_path, mtime) + mtime_ns = 0 + size = 0 + cache_key = (bin_path, mtime_ns, size) cached = cls._capability_cache.get(cache_key) if cached is not None: return cached @@ -7799,6 +7909,7 @@ def _download_companion_gguf( cancel_event: Optional[threading.Event] = None, near_path: Optional[str] = None, outcome: Optional[dict] = None, + on_transient_failure: Optional[Callable[[], None]] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. @@ -7816,6 +7927,13 @@ def _download_companion_gguf( then cannot open, so the load fell back to no speculation with nothing to show for the download -- and it disagreed with the local scan, which accepts a split drafter only when every shard is present. + + ``on_transient_failure`` fires when the companion was lost to something that + says nothing about the repo: a listing that never completed, or a download that + dropped. That is the one None worth another attempt, and it is a callback rather + than a second ``outcome`` key because the DFlash caller runs this in a loop and + needs the answer as each attempt ends. Permanent errors, offline mode and a + cancelled load never fire it. """ cancel_event = cancel_event if cancel_event is not None else self._cancel_event if cancel_event.is_set(): @@ -7852,6 +7970,10 @@ def _pick_from(names: list[str]) -> Optional[str]: # completed (offline, transient Hub failure) leaves target None for a # reason that says nothing about the repo's contents. listing_answered = False + # Whether the last listing attempt died of something that could yet succeed. + # Distinct from listing_answered: a permanent error also leaves the listing + # unanswered, and retrying that one buys nothing. + listing_failed = False from huggingface_hub import list_repo_files # Retry a transient listing blip; permanent repo/auth errors and offline @@ -7862,17 +7984,13 @@ def _pick_from(names: list[str]) -> Optional[str]: try: target = _pick_from(list_repo_files(hf_repo, token = hf_token)) listing_answered = True + listing_failed = False break except Exception as e: - if type(e).__name__ in ( - "RepositoryNotFoundError", - "GatedRepoError", - "RevisionNotFoundError", - "EntryNotFoundError", - "OfflineModeIsEnabled", - ): + if type(e).__name__ in _PERMANENT_HUB_ERRORS: logger.debug(f"Could not list repo files for {label}: {e}") break + listing_failed = True logger.debug( f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}" ) @@ -7903,6 +8021,10 @@ def _pick_from(names: list[str]) -> Optional[str]: ): outcome["listed"] = target is not None if target is None or cancel_event.is_set(): + # The listing is the only step that can fail this far in, and it failing + # transiently is the one None the caller may want to come back for. + if target is None and listing_failed and not cancel_event.is_set(): + self._report_transient_companion_failure(on_transient_failure, label) return None # Offline, resolve the companion straight from the cache snapshot that @@ -7956,8 +8078,36 @@ def _pick_from(names: list[str]) -> Optional[str]: return local_path except Exception as e: logger.warning(f"Could not download {label}: {e}") + # The listing already named the file, so this is the fetch itself dropping: + # worth another attempt, unless the Hub answered for good, this machine did + # (a full disk stays full), the load was cancelled, or there was never going + # to be a download (offline). + if ( + not cancel_event.is_set() + and type(e).__name__ not in _PERMANENT_HUB_ERRORS + and getattr(e, "errno", None) not in _UNRECOVERABLE_LOCAL_ERRNOS + and not _hf_env_offline() + ): + self._report_transient_companion_failure(on_transient_failure, label) return None + @staticmethod + def _report_transient_companion_failure( + on_transient_failure: Optional[Callable[[], None]], label: str + ) -> None: + """Tell the caller its companion was lost to something retryable. + + Best-effort in both directions: no caller cares (the callback is optional), and + a caller whose callback raises must not lose the load over bookkeeping for a + companion that is already best-effort. + """ + if on_transient_failure is None: + return + try: + on_transient_failure() + except Exception as exc: + logger.debug("Could not record the transient %s failure: %s", label, exc) + def _download_mmproj( self, *, @@ -8189,12 +8339,12 @@ def _cached_repo_dflash_drafter( before. """ try: - from utils.models.model_config import ( + from utils.models.drafters import ( _drafter_split_is_complete, - _iter_hf_cache_snapshots, dflash_repo_preference_key, is_dflash_architecture, ) + from utils.models.model_config import _iter_hf_cache_snapshots snapshots = ( _iter_hf_cache_snapshots(hf_repo) @@ -8271,6 +8421,12 @@ def _download_dflash( """ weight_name = Path(near_path).name if near_path else None + # Whatever a previous load concluded, this one is asking again. + self._dflash_retry_needed = False + + def _mark_retry_needed() -> None: + self._dflash_retry_needed = True + # Basenames whose header turned out not to say dflash. A name lands here # only once the file is readable on disk, and _pick_dflash then skips it, # so a repo whose best-ranked candidate is an impostor still reaches the @@ -8285,7 +8441,7 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: # The rest of the listing supplies the neighbouring weights that make # a foreign sidecar recognisable (a sidecar naming no family at all # stays eligible, which is what the published one does). - from utils.models.model_config import dflash_repo_preference_key + from utils.models.drafters import dflash_repo_preference_key # Root level only, as the local scan is: a nested dflash-*.gguf is an # ordinary weight, and offering it here spends its entire download @@ -8317,7 +8473,7 @@ def _validated(path: Optional[str]) -> Optional[str]: file llama-server then refuses as --model-draft. Same helper as detect_dflash_file so the two rules cannot drift. """ - from utils.models.model_config import is_dflash_architecture + from utils.models.drafters import is_dflash_architecture if not path: return None @@ -8384,6 +8540,10 @@ def _validated(path: Optional[str]) -> Optional[str]: label = "DFlash drafter", near_path = near_path, outcome = outcome, + # A header rejection is not this: that candidate is settled, and the + # loop moves on to the next one. Only the Hub dropping out from under + # the fetch leaves a sidecar this repo really publishes unfetched. + on_transient_failure = _mark_retry_needed, ) if candidate is None or candidate in fetched: break @@ -8418,8 +8578,8 @@ def _dspark_wins_auto( can never launch and the load ends up with no drafter at all. A raised probe answers False, which keeps Auto on its other options - exactly as the promotion does. probe_server_capabilities caches on - (binary path, mtime), so asking here and at the promotion is one probe. + exactly as the promotion does. probe_server_capabilities caches on the + binary's revision, so asking here and at the promotion is one probe. """ if spec_canon != "auto" or not dspark_draft_path: return False @@ -9562,6 +9722,21 @@ def _binary_stamp(path: Path) -> tuple: return () return (stat.st_ino, stat.st_size, stat.st_mtime_ns) + @staticmethod + def _binary_revision(binary: Optional[str]) -> tuple: + """Which build of llama-server a path currently holds, () when unreadable. + + The path alone cannot answer that: an update swaps the file in place, which is + why the CPU staging above stamps its source too. Empty on an unreadable file, so + a caller comparing two revisions treats "cannot tell" as "unchanged" -- the + install window itself is unreadable, and reading that as a finished update would + tear a healthy server down for a binary that is not there yet. + """ + if not binary: + return () + stamp = LlamaCppBackend._binary_stamp(Path(binary)) + return (str(binary),) + stamp if stamp else () + def _cleanup_cpu_fallback_runtime(self) -> None: runtime = getattr(self, "_cpu_fallback_runtime", None) self._cpu_fallback_runtime = None @@ -9826,6 +10001,14 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() + # Which build this load is about to run, kept so a stand-down blamed on the + # binary can tell a later `unsloth studio update` from the same file still + # being installed. + self._launch_binary_revision = self._binary_revision(binary) + # Cleared per load, not only where the fetch runs: a local-file load never + # reaches the DFlash fetch, and a verdict left over from the previous model + # would make every Apply for this one reload. + self._dflash_retry_needed = False is_vulkan_backend = self._is_vulkan_backend(binary) _vulkan_ordinal_pin = ( is_vulkan_backend and bool(gpu_ids) and gpu_ids_are_vulkan_ordinals is not False @@ -10119,6 +10302,14 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # reloads a healthy diffusion server on every Apply. self._mtp_draft_path = None self._mtp_draft_suppressed_path = None + # And the verdict on why the last load had no drafter, for the same + # reason once more: _build_speculative_flags clears it at the top of + # every load that reaches it, and this path never does. Both retry + # rules in the dedupe read it -- the drafter one and the binary one -- + # and a stale answer would relaunch this server on every Apply for a + # drafter it was never going to carry. + self._spec_fallback_reason = None + self._spec_drafter_kind = None with self._lock: if self._cancel_event.is_set(): logger.info("Load cancelled before diffusion server start") @@ -13520,8 +13711,15 @@ def _fallback_drafter_not_found() -> None: _emit_mtp(chain_ngram = chain_ngram) return flags - # effective_mode == "auto": the promotion path. llama.cpp #22673: - # MTP is compatible with mmproj, so there's no vision gate. + # effective_mode == "auto": the promotion path. No vision gate on any + # drafter kind. MTP's mmproj compatibility is llama.cpp #22673; DFlash + # is a different --spec-type and #22105 says nothing either way, so it + # was measured rather than assumed. Muse-Glimmer-30B UD-Q4_K_XL with + # mmproj-kquant and dflash-kquant, b10342, one B200, n_max=2, greedy, + # on a prompt carrying ~545 image tokens: 92.1 -> 114.2 tok/s at 0.646 + # acceptance, greedy output byte-identical to the drafter-free run. A + # vision gate would cost the flagship model 1.24x on the workload it + # ships for. if dspark_draft_path and caps.get("supports_dspark"): # DSpark first: load_model only hands a sidecar down once it has one # this binary can launch, and it beats every other Auto outcome for @@ -13780,6 +13978,8 @@ def unload_model(self) -> bool: self._spec_drafter_kind = None self._dspark_sidecar_absent = False self._dflash_sidecar_absent = False + self._dflash_retry_needed = False + self._launch_binary_revision = () self._cpu_fallback_reason = None self._last_load_intent = None self._mtp_runtime_fallback_active = False diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index b05ed430d13..21c65c7c6f4 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -131,6 +131,42 @@ def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: return candidates[0] if candidates else None +def preferred_dflash_sibling( + siblings: Sequence, + weight_name: Optional[str] = None, + other_weight_names: Sequence[str] = (), +) -> Optional[object]: + """The DFlash sidecar to fetch alongside ``weight_name``. + + Root level only, like preferred_mtp_sibling above and for the same reason + the loader's remote discovery is: the local contract in detect_dflash_file + never offers a nested ``quants/dflash-*.gguf``, and a listing cannot read a + header, so going by basename would put a whole ordinary weight in the plan + before anything could reject it. + + Ordered by dflash_repo_preference_key, the same key the download, the + snapshot reuse and the offline cache use, so the file the manifest promises + is the file the loader ends up launching. + """ + from utils.models.drafters import dflash_repo_preference_key + + candidates = [ + s + for s in siblings + if (name := _gguf_rfilename(s)) + and "/" not in name + and name.lower().startswith("dflash-") + ] + if not candidates: + return None + return min( + candidates, + key = lambda s: dflash_repo_preference_key( + getattr(s, "rfilename"), weight_name, other_weight_names + ), + ) + + def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: main: dict[str, list] = {} all_mmproj = mmproj_siblings(siblings) @@ -166,13 +202,38 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: main.setdefault(quant, []).append(sibling) plans: dict[str, GgufVariantPlan] = {} + # Every weight in the listing, so the DFlash ranking can tell a sidecar that + # names a neighbouring family from one that names this variant's. + all_weight_names = [ + name.rsplit("/", 1)[-1] + for quant_siblings in main.values() + for sibling in quant_siblings + if (name := _gguf_rfilename(sibling)) + ] for quant, target_main_siblings in main.items(): main_expected = tuple( file for sibling in target_main_siblings if (file := expected_file_from_sibling(sibling)) is not None ) - expected_files = (*main_expected, *companions_expected) + # The DFlash sidecar is per variant, unlike mmproj and the MTP drafter: + # its ranking is relative to the weight being fetched, so a multi-family + # repo does not hand variant B the drafter that names variant A. + target_weight = _gguf_rfilename(target_main_siblings[0]) + target_weight_name = target_weight.rsplit("/", 1)[-1] if target_weight else None + dflash_sibling = preferred_dflash_sibling( + siblings, + target_weight_name, + [n for n in all_weight_names if n != target_weight_name], + ) + dflash_expected = ( + expected_file_from_sibling(dflash_sibling) if dflash_sibling is not None else None + ) + expected_files = ( + *main_expected, + *companions_expected, + *((dflash_expected,) if dflash_expected is not None else ()), + ) plans[quant] = plan_from_expected_files( quant, expected_files, diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 9b23441c391..b4f982a0993 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -770,6 +770,19 @@ def _make_fake_llama_server(path: Path, help_text: str) -> Path: return path +# One fixed wall-clock second, so two revisions of a file differ only below the +# resolution a whole-second mtime can see. +_FIXED_MTIME_SECOND = 1_700_000_000 + + +def _pin_mtime(path: Path, *, nanos: int) -> Path: + """Pin a file's mtime `nanos` into one fixed second, so a later rewrite of it + differs only below the resolution a whole-second mtime can see.""" + stamp = _FIXED_MTIME_SECOND * 1_000_000_000 + nanos + os.utime(path, ns = (stamp, stamp)) + return path + + _NEEDS_BASH = pytest.mark.skipif( sys.platform == "win32", reason = "fake llama-server is a bash stub; Windows has no direct executor", @@ -806,6 +819,31 @@ def test_probe_server_capabilities_detects_dspark(tmp_path): assert caps["supports_dspark"] is True +_DFLASH_SPEC_HELP = "--spec-type none,draft-mtp,draft-dflash,ngram-mod" +# Padded to the same length as the DFlash one, so the two builds are the same size on +# disk and only the sub-second mtime tells them apart. +_PRE_DFLASH_SPEC_HELP = "--spec-type none,draft-mtp,ngram-mod".ljust(len(_DFLASH_SPEC_HELP)) + + +@_NEEDS_BASH +def test_probe_server_capabilities_rereads_a_binary_replaced_in_the_same_second(tmp_path): + """`unsloth studio update` overwrites llama-server in place. Keyed on whole + seconds, an update landing in the second the old build was probed in kept the key + identical and the new build was answered with the old one's capabilities -- and + "the user just installed the missing capability" is exactly the moment the cache + has to notice.""" + binary = _make_fake_llama_server(tmp_path / "llama-server", _PRE_DFLASH_SPEC_HELP) + _pin_mtime(binary, nanos = 100_000) + _clear_caps_cache() + assert LlamaCppBackend.probe_server_capabilities(str(binary))["supports_dflash"] is False + + before = binary.stat().st_size + _make_fake_llama_server(binary, _DFLASH_SPEC_HELP) + _pin_mtime(binary, nanos = 900_000) + assert binary.stat().st_size == before + assert LlamaCppBackend.probe_server_capabilities(str(binary))["supports_dflash"] is True + + @_NEEDS_BASH def test_probe_server_capabilities_gates_known_broken_dspark_prebuilt(tmp_path, monkeypatch): fake = _make_fake_llama_server( @@ -2562,6 +2600,184 @@ def test_already_in_target_state_retries_after_hf_drafter_not_found(): assert _matches(ok, **_drafter_not_found_kwargs()) is True +# ── A binary that has since gained the drafter ─────────────────────── +# +# Standing down on speculative decoding because llama-server cannot run it tells the +# user to run `unsloth studio update`. The update changes nothing about the request, so +# the comparators see the same intent and skip the reload: the one load the update +# exists to fix is the one that never happens again. + + +def _binary_fallback_kwargs(): + """An Auto request for the model the fallen-back server is already running.""" + return dict( + model_identifier = "unsloth/Muse-Glimmer-30B-GGUF", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gguf_path = None, + ) + + +def _stood_down_backend(**overrides): + """Live server that dropped its drafter because the binary could not run it.""" + state = dict( + _model_identifier = "unsloth/Muse-Glimmer-30B-GGUF", + _speculative_type = "default", + _spec_fallback_reason = "binary_no_mtp", + _gguf_path = None, + ) + state.update(overrides) + return _mtp_backend(**state) + + +def _fake_caps(monkeypatch, **capabilities): + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: dict(capabilities)), + ) + + +@pytest.mark.parametrize( + ("kind", "capability"), + [("dflash", "supports_dflash"), ("dspark", "supports_dspark"), ("mtp", "supports_mtp")], +) +def test_already_in_target_state_reloads_once_the_binary_can_run_the_drafter( + monkeypatch, kind, capability +): + _fake_caps(monkeypatch, **{capability: True}) + backend = _stood_down_backend(_spec_drafter_kind = kind) + assert _matches(backend, **_binary_fallback_kwargs()) is False + + +@pytest.mark.parametrize("kind", ["dflash", "dspark", "mtp"]) +def test_already_in_target_state_keeps_deduping_while_the_binary_still_cannot(monkeypatch, kind): + """The half that stops this becoming a reload loop: nothing has changed, so the + healthy drafterless server has to be left alone.""" + _fake_caps( + monkeypatch, + supports_dflash = False, + supports_dspark = False, + supports_mtp = False, + ) + backend = _stood_down_backend(_spec_drafter_kind = kind) + assert _matches(backend, **_binary_fallback_kwargs()) is True + + +def test_already_in_target_state_asks_about_the_drafter_that_actually_stood_down(monkeypatch): + """Every kind records the same "binary_no_mtp", so a check keyed on the reason + alone would read a DSpark stand-down as answered by any build carrying MTP -- and + tear down a healthy server on every Apply for a capability it never gained.""" + _fake_caps(monkeypatch, supports_mtp = True, supports_dspark = False, supports_dflash = False) + backend = _stood_down_backend(_spec_drafter_kind = "dspark") + assert _matches(backend, **_binary_fallback_kwargs()) is True + + +def test_already_in_target_state_never_reprobes_a_binary_nothing_has_touched(tmp_path, monkeypatch): + """The steady state, and the reason it has to be cheap: this runs on every Apply, + while the probe behind it spawns `llama-server --help` on a cold cache -- and the + cache is cold exactly when the binary was just replaced.""" + binary = tmp_path / "llama-server" + binary.write_bytes(b"unchanged build") + monkeypatch.setattr( + LlamaCppBackend, + "_find_llama_server_binary", + staticmethod(lambda **_kwargs: str(binary)), + ) + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: pytest.fail("an untouched binary was reprobed")), + ) + backend = _stood_down_backend(_spec_drafter_kind = "dflash") + backend._launch_binary_revision = LlamaCppBackend._binary_revision(str(binary)) + + assert _matches(backend, **_binary_fallback_kwargs()) is True + + +def test_already_in_target_state_sits_out_an_install_still_in_flight(tmp_path, monkeypatch): + """An update is not atomic, and mid-install the binary is unreadable. Reading that + as "a different build is installed" would tear the server down for a file that is + not there yet, and the reload would kill the process before finding that out.""" + binary = tmp_path / "llama-server" + binary.write_bytes(b"old build") + monkeypatch.setattr( + LlamaCppBackend, + "_find_llama_server_binary", + staticmethod(lambda **_kwargs: str(binary)), + ) + backend = _stood_down_backend( + _spec_fallback_reason = "binary_outdated", + _spec_drafter_kind = "mtp", + ) + backend._launch_binary_revision = LlamaCppBackend._binary_revision(str(binary)) + + binary.unlink() + assert LlamaCppBackend._binary_revision(str(binary)) == () + assert _matches(backend, **_binary_fallback_kwargs()) is True + + +def test_already_in_target_state_reloads_when_the_crashed_binary_was_replaced( + tmp_path, monkeypatch +): + """A binary_outdated stand-down comes from a launch that died on an architecture + the build did not know, and no --help flag advertises those, so this one has to + compare the file itself.""" + binary = tmp_path / "llama-server" + binary.write_bytes(b"old build") + _pin_mtime(binary, nanos = 100_000) + monkeypatch.setattr( + LlamaCppBackend, + "_find_llama_server_binary", + staticmethod(lambda **_kwargs: str(binary)), + ) + backend = _stood_down_backend( + _spec_fallback_reason = "binary_outdated", + _spec_drafter_kind = "mtp", + ) + backend._launch_binary_revision = LlamaCppBackend._binary_revision(str(binary)) + assert _matches(backend, **_binary_fallback_kwargs()) is True + + # Same path, same size, same second: an update landing right after the crash. + binary.write_bytes(b"new build") + _pin_mtime(binary, nanos = 900_000) + assert _matches(backend, **_binary_fallback_kwargs()) is False + + +def test_diffusion_load_clears_the_previous_models_spec_fallback(): + """The diffusion early-return skips _build_speculative_flags, which is what clears + the stand-down on every other load, and only /unload clears it otherwise. Both + retry rules in the dedupe read it, so an MTP model's verdict left behind by a + switch to DiffusionGemma relaunches the diffusion server on every Apply -- forever, + since the relaunch takes this same path and leaves the verdict exactly as it was.""" + src = inspect.getsource(LlamaCppBackend.load_model) + diffusion = src.find("if self._is_diffusion:") + assert diffusion != -1 + start = src.find("return self._start_diffusion_server", diffusion) + assert start != -1 + assert "self._spec_fallback_reason = None" in src[diffusion:start] + assert "self._spec_drafter_kind = None" in src[diffusion:start] + + +def test_already_in_target_state_reloads_after_a_dflash_fetch_that_dropped(): + """Under Auto a lost sidecar leaves no fallback reason at all -- the promotion + never ran -- so the flag is the only thing that can ask for one more attempt.""" + backend = _mtp_backend( + _model_identifier = "unsloth/Muse-Glimmer-30B-GGUF", + _speculative_type = "default", + _gguf_path = None, + ) + assert _matches(backend, **_binary_fallback_kwargs()) is True + + backend._dflash_retry_needed = True + assert _matches(backend, **_binary_fallback_kwargs()) is False + + _MODERN_DRAFT_NGL_HELP = """usage: llama-server [options] --spec-draft-ngl N layers to offload for the draft model diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 77bb37c3818..41ed40140ad 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -1959,6 +1959,7 @@ def _fake_companion( cancel_event = None, near_path = None, outcome = None, + on_transient_failure = None, ): reached["hit"] = True reached["picked"] = pick( @@ -2039,6 +2040,7 @@ def _fake( cancel_event = None, near_path = None, outcome = None, + on_transient_failure = None, ): if outcome is not None: outcome["listed"] = listed @@ -2116,6 +2118,7 @@ def _fake_companion( cancel_event = None, near_path = None, outcome = None, + on_transient_failure = None, ): picked["name"] = pick(listing) return None @@ -2221,7 +2224,7 @@ def test_local_and_remote_dflash_pairing_agree(tmp_path): """One rule, three call sites: the local scan, the download picker and the cache lookup all go through dflash_repo_preference_key / _drafter_names_other_weight.""" - from utils.models.model_config import dflash_repo_preference_key + from utils.models.drafters import dflash_repo_preference_key others = ["model-A-Q4_K_M.gguf", "model-B-Q4_K_M.gguf"] ranked = sorted( @@ -2411,6 +2414,7 @@ def _fake_companion( cancel_event = None, near_path = None, outcome = None, + on_transient_failure = None, ): target = pick(listing) if outcome is not None: @@ -2519,7 +2523,7 @@ def test_cached_dflash_lookup_skips_a_prefixed_file_of_another_architecture(tmp_ def test_local_and_remote_dflash_architecture_checks_agree(tmp_path): """One rule, one place: detect_dflash_file and the remote paths both ask is_dflash_architecture, so neither can start trusting the name alone.""" - from utils.models.model_config import is_dflash_architecture + from utils.models.drafters import is_dflash_architecture weight = _write_gguf(tmp_path / "model-Q4_K_M.gguf", "llama") impostor = _write_gguf(tmp_path / "dflash-model-Q8_0.gguf", "llama") @@ -2732,7 +2736,10 @@ def test_from_identifier_never_reads_a_sidecar_outside_the_boundary(tmp_path, mo reports no DFlash sidecar rather than one the load route would reject.""" import os - import utils.models.model_config as mc + # Patch the module detect_dflash_file resolves the name in, not the one that + # used to re-export it: a patch on the re-exporting module never intercepts, + # so `reads == []` below would hold whether or not the boundary works. + import utils.models.drafters.dflash as dflash_mod leased = tmp_path / "leased" leased.mkdir() @@ -2743,13 +2750,13 @@ def test_from_identifier_never_reads_a_sidecar_outside_the_boundary(tmp_path, mo os.symlink(target, leased / "dflash-kquant.gguf") reads: list[str] = [] - real_check = mc.is_dflash_architecture + real_check = dflash_mod.is_dflash_architecture def _recording_check(path, *args, **kwargs): reads.append(str(path)) return real_check(path, *args, **kwargs) - monkeypatch.setattr(mc, "is_dflash_architecture", _recording_check) + monkeypatch.setattr(dflash_mod, "is_dflash_architecture", _recording_check) def _inside_the_lease(candidate, gguf_file, kind, search_root): return Path(search_root) in Path(candidate).parents @@ -3026,3 +3033,182 @@ def test_cached_dflash_lookup_falls_through_from_a_half_split_to_a_whole_one(tmp assert LlamaCppBackend()._cached_repo_dflash_drafter("org/repo", near_path = str(weight)) == str( whole ) + + +# ── A fetch that dropped is worth one more Apply ───────────────────── +# +# _dflash_sidecar_absent answers "this repo publishes none", which is permanent and +# must never be retried. The other None -- the Hub going away mid-fetch -- is invisible +# under Auto: no promotion happened, so no fallback reason was recorded, and the next +# Apply reuses a server that has no drafter for a repo that does publish one. + + +def _dflash_hub_download(tmp_path, monkeypatch, *, listing, fetch): + """Drive _download_dflash through the real _download_companion_gguf, with only + the two Hub calls stubbed out.""" + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: {"supports_dflash": True}), + ) + monkeypatch.setattr(llama_cpp_module, "_hub_download_in_flight", lambda hf_repo: False) + monkeypatch.setattr("huggingface_hub.list_repo_files", listing) + monkeypatch.setattr(llama_cpp_module, "hf_hub_download_with_xet_fallback", fetch) + # The local cache is not part of what is under test, and an unstubbed scan would + # answer from whatever this machine happens to have downloaded. + monkeypatch.setattr("utils.models.model_config._iter_hf_cache_snapshots", lambda *a, **k: []) + + b = LlamaCppBackend() + got = b._download_dflash(hf_repo = "org/repo", binary = "/fake/llama-server") + return b, got + + +def _never_fetched(*_args, **_kwargs): + raise AssertionError("nothing should be downloaded") + + +def test_download_dflash_asks_again_after_a_listing_that_never_answered(tmp_path, monkeypatch): + """An unreachable Hub says nothing about the repo, so recording it as "publishes + none" would strand the model without the sidecar it does publish.""" + + def _listing(repo, token = None): + raise ConnectionError("hub unreachable") + + b, got = _dflash_hub_download(tmp_path, monkeypatch, listing = _listing, fetch = _never_fetched) + + assert got is None + assert b._dflash_sidecar_absent is False + assert b._dflash_retry_needed is True + + +def test_download_dflash_asks_again_after_a_download_that_dropped(tmp_path, monkeypatch): + """The listing named the file, so this repo definitely publishes one: the bytes + are all that is missing.""" + + def _fetch( + repo, + filename, + token, + *, + cancel_event = None, + cache_dir = None, + ): + raise ConnectionError("connection reset") + + b, got = _dflash_hub_download( + tmp_path, + monkeypatch, + listing = lambda repo, token = None: ["model-Q4_K_M.gguf", "dflash-kquant.gguf"], + fetch = _fetch, + ) + + assert got is None + assert b._dflash_sidecar_absent is False + assert b._dflash_retry_needed is True + + +def test_download_dflash_does_not_ask_again_after_a_permanent_hub_error(tmp_path, monkeypatch): + """A repo that is gone, gated, or being asked about offline is answered for good; + retrying it on every Apply would relaunch an identical server forever.""" + + class RepositoryNotFoundError(Exception): + pass + + def _listing(repo, token = None): + raise RepositoryNotFoundError("404") + + b, got = _dflash_hub_download(tmp_path, monkeypatch, listing = _listing, fetch = _never_fetched) + + assert got is None + assert b._dflash_retry_needed is False + + +def test_download_dflash_does_not_ask_again_when_this_machine_is_the_problem(tmp_path, monkeypatch): + """A full disk or an unwritable cache stays that way, and the retry costs a full + unload plus another ~1.5 GiB attempt. Classified on errno, since the Hub client + raises OSError subclasses for network trouble too.""" + import errno + + def _fetch( + repo, + filename, + token, + *, + cancel_event = None, + cache_dir = None, + ): + raise OSError(errno.ENOSPC, "No space left on device") + + b, got = _dflash_hub_download( + tmp_path, + monkeypatch, + listing = lambda repo, token = None: ["model-Q4_K_M.gguf", "dflash-kquant.gguf"], + fetch = _fetch, + ) + + assert got is None + assert b._dflash_retry_needed is False + + +def test_download_dflash_treats_a_header_rejection_as_settled(tmp_path, monkeypatch): + """A candidate whose header does not say dflash is permanently not a sidecar: the + search falls through to the next one, and if that was the last one the repo + publishes none. Reading it as a dropped fetch would reload on every Apply for a + file that can never be launched.""" + + def _fetch( + repo, + filename, + token, + *, + cancel_event = None, + cache_dir = None, + ): + # An ordinary weight that merely carries the sidecar's naming. + return str(_write_gguf(tmp_path / filename, "llama")) + + b, got = _dflash_hub_download( + tmp_path, + monkeypatch, + listing = lambda repo, token = None: ["model-Q4_K_M.gguf", "dflash-model-Q8_0.gguf"], + fetch = _fetch, + ) + + assert got is None + assert b._dflash_sidecar_absent is True + assert b._dflash_retry_needed is False + + +def test_download_dflash_reaches_the_real_sidecar_behind_an_impostor(tmp_path, monkeypatch): + """The other half of the same rule: the rejection only removes that one name from + the pool, so the sidecar ranked behind it is still fetched and nothing is flagged + for a retry.""" + + def _fetch( + repo, + filename, + token, + *, + cancel_event = None, + cache_dir = None, + ): + arch = "dflash" if filename == "dflash-kquant.gguf" else "llama" + return str(_write_gguf(tmp_path / filename, arch)) + + b, got = _dflash_hub_download( + tmp_path, + monkeypatch, + listing = lambda repo, token = None: [ + "model-Q4_K_M.gguf", + "dflash-model-Q8_0.gguf", + "dflash-kquant.gguf", + ], + fetch = _fetch, + ) + + assert got == str(tmp_path / "dflash-kquant.gguf") + assert b._dflash_retry_needed is False diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index ce8ce03c5ec..91400262e96 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -30,7 +30,7 @@ import subprocess import sys from pathlib import Path -from typing import Callable, Iterable, List, Tuple, Union +from typing import Callable, List, Tuple, Union import hashlib import json import threading @@ -1658,25 +1658,36 @@ def _add(d: Path) -> None: # Drafter naming, ranking and DFlash discovery live in utils.models.drafters, so # one rule serves the local scan, the download, the snapshot reuse and the -# offline cache alike. Re-exported here deliberately: these names were part of -# this module's surface and are imported from it across the backend and the -# tests, so the move stays source compatible. +# offline cache alike. Only the names this module still calls are imported here; +# everything else is imported from utils.models.drafters directly at its use +# site. A re-export shim would be the friendlier move, but scripts/ +# verify_import_hoist.py scopes its __all__ exemption to package __init__.py and +# blocks unused imports in an ordinary module on purpose -- see its +# reexport_in_ordinary_module_is_still_blocked self-test. from utils.models.drafters import ( # noqa: E402 _drafter_launch_path, _drafter_matches_weight, - _drafter_names_other_weight, - _drafter_pairing_stem, _drafter_split_is_complete, _drafter_stem_rank, _drafter_total_size, detect_dflash_file, - dflash_precision_rank, - dflash_preference_key, - dflash_repo_preference_key, dspark_precision_rank, - dspark_preference_key, - is_dflash_architecture, ) +from utils.models.drafters import ( # noqa: E402 + dspark_preference_key as _drafters_dspark_preference_key, +) + + +def dspark_preference_key(name: str) -> Tuple[int, str]: + """Sort key picking the preferred DSpark sidecar by name alone. + + Delegates rather than re-exports: routes/inference.py has imported this from + model_config since #7968, and repointing that call site would trip + scripts/verify_import_hoist.py's TARGET-CHANGED rule, whose relocation + exemption only covers module-level imports. One implementation still, in + utils.models.drafters. + """ + return _drafters_dspark_preference_key(name) def detect_mtp_file( diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index a49f2b8edfe..5e925ed9c5a 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -474,6 +474,7 @@ def _load_gguf_backend( mmproj_path = model_config.gguf_mmproj_file, mtp_draft_path = model_config.gguf_mtp_file, dspark_draft_path = model_config.gguf_dspark_file, + dflash_draft_path = model_config.gguf_dflash_file, ) if speculative_type is not None: intent_fields["speculative_type"] = speculative_type diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py index d32bf46d016..648ade6af88 100644 --- a/unsloth_cli/tests/test_inference_chat.py +++ b/unsloth_cli/tests/test_inference_chat.py @@ -677,34 +677,9 @@ def test_deferred_error_helper_defaults_a_missing_status(): assert excinfo.value.code == 500 -@pytest.mark.parametrize( - ("source", "expected_source"), - [ - ( - {"gguf_hf_repo": "org/model-GGUF"}, - {"hf_repo": "org/model-GGUF", "hf_token": "hf_x"}, - ), - ( - { - "gguf_hf_repo": None, - "gguf_file": "/models/model.gguf", - "gguf_mmproj_file": "/models/mmproj.gguf", - "gguf_mtp_file": "/models/mtp.gguf", - "gguf_dspark_file": "/models/dspark-model.gguf", - }, - { - "gguf_path": "/models/model.gguf", - "mmproj_path": "/models/mmproj.gguf", - "mtp_draft_path": "/models/mtp.gguf", - "dspark_draft_path": "/models/dspark-model.gguf", - }, - ), - ], - ids = ("hugging-face", "local"), -) -def test_load_gguf_backend_forwards_source_and_runtime_options( - monkeypatch, source, expected_source -): +def _stub_studio_gguf_load(monkeypatch): + """Stand in for the studio backend `_load_gguf_backend` imports in-venv, and + return the list the intents it builds land in.""" import unsloth_cli._inference as inference calls = [] @@ -739,6 +714,42 @@ async def _passthrough( monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args) monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback) monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None) + return calls + + +@pytest.mark.parametrize( + ("source", "expected_source"), + [ + ( + {"gguf_hf_repo": "org/model-GGUF"}, + {"hf_repo": "org/model-GGUF", "hf_token": "hf_x"}, + ), + ( + { + "gguf_hf_repo": None, + "gguf_file": "/models/model.gguf", + "gguf_mmproj_file": "/models/mmproj.gguf", + "gguf_mtp_file": "/models/mtp.gguf", + "gguf_dspark_file": "/models/dspark-model.gguf", + "gguf_dflash_file": "/models/dflash-kquant.gguf", + }, + { + "gguf_path": "/models/model.gguf", + "mmproj_path": "/models/mmproj.gguf", + "mtp_draft_path": "/models/mtp.gguf", + "dspark_draft_path": "/models/dspark-model.gguf", + "dflash_draft_path": "/models/dflash-kquant.gguf", + }, + ), + ], + ids = ("hugging-face", "local"), +) +def test_load_gguf_backend_forwards_source_and_runtime_options( + monkeypatch, source, expected_source +): + import unsloth_cli._inference as inference + + calls = _stub_studio_gguf_load(monkeypatch) config = SimpleNamespace( gguf_variant = "Q4_K_M", @@ -773,6 +784,30 @@ async def _passthrough( ] +def test_load_gguf_backend_hands_a_local_dflash_sidecar_to_the_load(monkeypatch): + """The managed CLI resolves the sidecar next to a local weight exactly as Studio + does, and dropping it here is silent: the load simply comes up with no drafter and + nothing says the sidecar sitting beside the model was ever found.""" + import unsloth_cli._inference as inference + + calls = _stub_studio_gguf_load(monkeypatch) + config = SimpleNamespace( + gguf_variant = "Q4_K_M", + identifier = "org/model-GGUF", + is_vision = False, + gguf_hf_repo = None, + gguf_file = "/models/model.gguf", + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = "/models/dflash-kquant.gguf", + ) + + inference._load_gguf_backend(config, hf_token = None, max_seq_length = 8192) + + assert [intent.dflash_draft_path for intent in calls] == ["/models/dflash-kquant.gguf"] + + def test_load_gguf_backend_exits_cleanly_on_invalid_extra_args(monkeypatch): import unsloth_cli._inference as inference From 06d8f9538d5e98bdb9ccff52f5f4eb53e3e5261c Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 09:18:58 +0000 Subject: [PATCH 20/26] Tighten the DFlash comments for PR #8338 --- studio/backend/core/inference/llama_cpp.py | 144 ++++++++------------ studio/backend/hub/utils/gguf_plan.py | 24 ++-- studio/backend/utils/models/model_config.py | 22 ++- 3 files changed, 75 insertions(+), 115 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e7f8f582409..5bacc469ec3 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -219,9 +219,7 @@ class _CpuFallbackRuntime(NamedTuple): "LLAMA_ARG_HFD_REPO", ) -# Hub failures that answer the question for good: the repo, the revision or the file is -# not there, or it is there and this token may not have it, or the Hub was never going to -# be asked. Retrying any of them, on this load or the next Apply, buys nothing. +# Hub failures that are settled: retrying them buys nothing. _PERMANENT_HUB_ERRORS = ( "RepositoryNotFoundError", "GatedRepoError", @@ -230,10 +228,8 @@ class _CpuFallbackRuntime(NamedTuple): "OfflineModeIsEnabled", ) -# The other kind of settled failure a download can hit: this machine, not the Hub. Told -# apart by errno rather than by exception type, because the Hub client raises OSError -# subclasses for network trouble too (requests' ConnectionError is one of them) and -# those DO deserve another attempt. +# Settled failures of this machine, not the Hub. By errno, not exception type: the Hub +# client raises OSError subclasses for network trouble too, and those are retryable. _UNRECOVERABLE_LOCAL_ERRNOS = frozenset( { errno.EACCES, # cache dir not writable @@ -243,9 +239,8 @@ class _CpuFallbackRuntime(NamedTuple): } ) -# What llama-server has to advertise for each drafter kind. Every kind records the same -# "binary_no_mtp" when it stands down, so the reason alone cannot say which capability -# to ask about again. +# Every kind records the same "binary_no_mtp", so the reason alone cannot say which +# capability to recheck. _SPEC_KIND_CAPABILITY: dict[str, str] = { "dspark": "supports_dspark", "dflash": "supports_dflash", @@ -3272,15 +3267,12 @@ def __init__(self): self._spec_drafter_kind: Optional[str] = None self._dspark_sidecar_absent: bool = False self._dflash_sidecar_absent: bool = False - # Set when the DFlash sidecar could not be fetched for a reason that says - # nothing about the repo (a listing blip, a download that dropped). Unlike - # _dflash_sidecar_absent, which is the permanent answer for almost every repo, - # this one is worth one more Apply: under Auto the fetch failing silently - # leaves the load with no drafter and nothing else recording why. + # Sidecar fetch lost to a blip, not to the repo lacking one + # (_dflash_sidecar_absent). Worth one more Apply: under Auto nothing else + # records that the load ended up with no drafter. self._dflash_retry_needed: bool = False - # Which llama-server file the last load resolved. `unsloth studio update` - # replaces it in place, so a load that blamed the binary can tell an update - # from the same build still being installed. + # Which llama-server file this load ran. `unsloth studio update` replaces it + # in place, so a stand-down blamed on the binary can spot a real update. self._launch_binary_revision: tuple = () # Set after an auto-Vulkan crash recovers with all devices disabled. self._cpu_fallback_reason: Optional[str] = None @@ -3539,9 +3531,9 @@ def spec_fallback_reason(self) -> Optional[str]: return self._spec_fallback_reason def _binary_changed_since_launch(self) -> bool: - """Whether a different llama-server file is installed than the live server was - launched from. False when either side is unreadable, so the install's own window - is not mistaken for a finished install: the next Apply asks again.""" + """Whether a different llama-server is installed than the live one was launched + from. False when either side is unreadable, so an install still in flight is not + read as a finished one; the next Apply asks again.""" current = self._binary_revision(self._find_llama_server_binary()) if not current or not self._launch_binary_revision: return False @@ -3550,28 +3542,22 @@ def _binary_changed_since_launch(self) -> bool: def spec_binary_fallback_can_retry(self) -> bool: """Whether the binary has since gained what the last load stood down for. - A load that drops speculative decoding because llama-server cannot run it tells - the user to run `unsloth studio update`. Doing so changes nothing about the - request, so the duplicate-load comparators see an identical intent and skip the - reload: the model keeps serving without a drafter until something unrelated - forces a relaunch, and the hint reads as a lie. Re-reading the binary here is - what turns it into a working instruction. - - Which question to ask depends on what the binary was blamed for. A missing - ``--spec-type`` is answered by the capability the drafter kind needs, since every - kind records the same "binary_no_mtp" and asking about MTP after a DSpark - stand-down would say yes on a build that never gained draft-dspark. An - architecture the build did not know is advertised by no flag at all, so that one - can only compare the file. + The UI tells the user to run `unsloth studio update`, which leaves the request + identical, so the duplicate-load comparators would otherwise dedup the very + reload that update was meant to enable. + + Ask about the capability the drafter kind needs, not the reason: every kind + records "binary_no_mtp", so asking about MTP after a DSpark stand-down says yes + on a build that never gained draft-dspark. "binary_outdated" is an unknown + architecture, which no flag advertises, so that one compares the file. """ if self._spec_fallback_reason == "binary_outdated": return self._binary_changed_since_launch() if self._spec_fallback_reason != "binary_no_mtp": return False if self._launch_binary_revision and not self._binary_changed_since_launch(): - # An untouched file cannot advertise anything new, and this runs on every - # duplicate-load check, where the probe below is a subprocess on a cold - # cache -- which is precisely the state an update leaves it in. + # An untouched file advertises nothing new, and this runs on every + # duplicate-load check where the probe below is a subprocess. return False capability = _SPEC_KIND_CAPABILITY.get(self._spec_drafter_kind or "") if capability is None: @@ -4133,19 +4119,15 @@ def _norm(value): and not spec_owned_by_extra_args ): return False - # The other recoverable stand-down, and the one the UI actively asks the user to - # fix: the drafter was there and llama-server could not run it. Updating the - # binary leaves the request identical, so without this the load the update was - # meant to repair is exactly the load that never happens again. + # The stand-down the UI asks the user to fix by updating llama.cpp. That leaves + # the request identical, so without this the repaired load never happens. if ( speculative_type in ("auto", "mtp", "mtp+ngram", "dspark", "dflash") and self.spec_binary_fallback_can_retry() ): return False - # A DFlash fetch that failed on the way out leaves no other trace under Auto: - # the promotion never ran, so no fallback reason was recorded and the branch - # above cannot see it. A repo that simply publishes no sidecar is - # _dflash_sidecar_absent's business and is never retried here. + # Under Auto a failed fetch records no fallback reason, so the branch above + # cannot see it. A repo with no sidecar is _dflash_sidecar_absent's business. if self._dflash_retry_needed and speculative_type in ("auto", "dflash"): return False compared_draft_n_max = self._spec_draft_n_max @@ -4455,11 +4437,9 @@ def _dspark_release_is_broken(cls, release_tag: Optional[str]) -> bool: match = re.match(r"b(\d+)", str(release_tag or "")) return bool(match) and int(match.group(1)) in cls._BROKEN_DSPARK_BUILDS - # Cached on the revision of the file, not on the path: `unsloth studio update` - # replaces the binary in place. Nanoseconds and size rather than int(st_mtime), - # since an update landing in the same second as the probe kept the key identical and - # the new build was answered with the old one's capabilities -- which is exactly the - # moment a capability the user just installed has to become visible. + # Keyed on the file's revision, since `unsloth studio update` replaces it in place. + # Nanoseconds and size, not int(st_mtime): an update landing in the same second as + # the probe kept the key identical and got the old build's capabilities. _capability_cache: dict[tuple[str, int, int], dict[str, object]] = {} @classmethod @@ -7928,12 +7908,11 @@ def _download_companion_gguf( show for the download -- and it disagreed with the local scan, which accepts a split drafter only when every shard is present. - ``on_transient_failure`` fires when the companion was lost to something that - says nothing about the repo: a listing that never completed, or a download that - dropped. That is the one None worth another attempt, and it is a callback rather - than a second ``outcome`` key because the DFlash caller runs this in a loop and - needs the answer as each attempt ends. Permanent errors, offline mode and a - cancelled load never fire it. + ``on_transient_failure`` fires when the companion was lost to a listing that + never completed or a download that dropped, the one None worth another attempt. + A callback rather than an ``outcome`` key because the DFlash caller loops and + needs the answer per attempt. Permanent errors, offline and cancellation do not + fire it. """ cancel_event = cancel_event if cancel_event is not None else self._cancel_event if cancel_event.is_set(): @@ -7970,9 +7949,8 @@ def _pick_from(names: list[str]) -> Optional[str]: # completed (offline, transient Hub failure) leaves target None for a # reason that says nothing about the repo's contents. listing_answered = False - # Whether the last listing attempt died of something that could yet succeed. - # Distinct from listing_answered: a permanent error also leaves the listing - # unanswered, and retrying that one buys nothing. + # Whether the listing died of something retryable. Distinct from + # listing_answered: a permanent error also leaves the listing unanswered. listing_failed = False from huggingface_hub import list_repo_files @@ -8021,8 +7999,7 @@ def _pick_from(names: list[str]) -> Optional[str]: ): outcome["listed"] = target is not None if target is None or cancel_event.is_set(): - # The listing is the only step that can fail this far in, and it failing - # transiently is the one None the caller may want to come back for. + # The listing is the only step that can fail this far in. if target is None and listing_failed and not cancel_event.is_set(): self._report_transient_companion_failure(on_transient_failure, label) return None @@ -8078,10 +8055,8 @@ def _pick_from(names: list[str]) -> Optional[str]: return local_path except Exception as e: logger.warning(f"Could not download {label}: {e}") - # The listing already named the file, so this is the fetch itself dropping: - # worth another attempt, unless the Hub answered for good, this machine did - # (a full disk stays full), the load was cancelled, or there was never going - # to be a download (offline). + # The listing named the file, so this is the fetch dropping: retryable + # unless the Hub or this machine answered for good, or we were cancelled. if ( not cancel_event.is_set() and type(e).__name__ not in _PERMANENT_HUB_ERRORS @@ -8097,9 +8072,8 @@ def _report_transient_companion_failure( ) -> None: """Tell the caller its companion was lost to something retryable. - Best-effort in both directions: no caller cares (the callback is optional), and - a caller whose callback raises must not lose the load over bookkeeping for a - companion that is already best-effort. + Best-effort both ways: the callback is optional, and one that raises must not + lose the load over bookkeeping for an already best-effort companion. """ if on_transient_failure is None: return @@ -8540,9 +8514,8 @@ def _validated(path: Optional[str]) -> Optional[str]: label = "DFlash drafter", near_path = near_path, outcome = outcome, - # A header rejection is not this: that candidate is settled, and the - # loop moves on to the next one. Only the Hub dropping out from under - # the fetch leaves a sidecar this repo really publishes unfetched. + # Not a header rejection: that candidate is settled and the loop moves + # on. Only the Hub dropping out leaves a real sidecar unfetched. on_transient_failure = _mark_retry_needed, ) if candidate is None or candidate in fetched: @@ -9726,11 +9699,10 @@ def _binary_stamp(path: Path) -> tuple: def _binary_revision(binary: Optional[str]) -> tuple: """Which build of llama-server a path currently holds, () when unreadable. - The path alone cannot answer that: an update swaps the file in place, which is - why the CPU staging above stamps its source too. Empty on an unreadable file, so - a caller comparing two revisions treats "cannot tell" as "unchanged" -- the - install window itself is unreadable, and reading that as a finished update would - tear a healthy server down for a binary that is not there yet. + The path alone cannot say, since an update swaps the file in place. Empty on an + unreadable file so callers read "cannot tell" as "unchanged": the install window + is itself unreadable, and treating it as a finished update would tear a healthy + server down for a binary that is not there yet. """ if not binary: return () @@ -10001,9 +9973,8 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() - # Which build this load is about to run, kept so a stand-down blamed on the - # binary can tell a later `unsloth studio update` from the same file still - # being installed. + # Kept so a stand-down blamed on the binary can later tell a real update + # from the same file still being installed. self._launch_binary_revision = self._binary_revision(binary) # Cleared per load, not only where the fetch runs: a local-file load never # reaches the DFlash fetch, and a verdict left over from the previous model @@ -13711,15 +13682,12 @@ def _fallback_drafter_not_found() -> None: _emit_mtp(chain_ngram = chain_ngram) return flags - # effective_mode == "auto": the promotion path. No vision gate on any - # drafter kind. MTP's mmproj compatibility is llama.cpp #22673; DFlash - # is a different --spec-type and #22105 says nothing either way, so it - # was measured rather than assumed. Muse-Glimmer-30B UD-Q4_K_XL with - # mmproj-kquant and dflash-kquant, b10342, one B200, n_max=2, greedy, - # on a prompt carrying ~545 image tokens: 92.1 -> 114.2 tok/s at 0.646 - # acceptance, greedy output byte-identical to the drafter-free run. A - # vision gate would cost the flagship model 1.24x on the workload it - # ships for. + # effective_mode == "auto": the promotion path. No vision gate on any drafter + # kind. MTP's mmproj compatibility is llama.cpp #22673; DFlash is a different + # --spec-type and #22105 says nothing, so it was measured: Muse-Glimmer-30B + # UD-Q4_K_XL + mmproj-kquant + dflash-kquant, b10342, B200, n_max=2, greedy, + # ~545 image tokens gave 92.1 -> 114.2 tok/s at 0.646 acceptance with output + # byte-identical to the drafter-free run. if dspark_draft_path and caps.get("supports_dspark"): # DSpark first: load_model only hands a sidecar down once it has one # this binary can launch, and it beats every other Auto outcome for diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 21c65c7c6f4..8eb501edb73 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -138,15 +138,12 @@ def preferred_dflash_sibling( ) -> Optional[object]: """The DFlash sidecar to fetch alongside ``weight_name``. - Root level only, like preferred_mtp_sibling above and for the same reason - the loader's remote discovery is: the local contract in detect_dflash_file - never offers a nested ``quants/dflash-*.gguf``, and a listing cannot read a - header, so going by basename would put a whole ordinary weight in the plan - before anything could reject it. - - Ordered by dflash_repo_preference_key, the same key the download, the - snapshot reuse and the offline cache use, so the file the manifest promises - is the file the loader ends up launching. + Root level only, like preferred_mtp_sibling: detect_dflash_file never offers a + nested ``quants/dflash-*.gguf``, and a listing cannot read a header, so matching + the basename would plan a whole ordinary weight nothing could reject in time. + + Ordered by dflash_repo_preference_key, as the download, snapshot reuse and offline + cache are, so the manifest promises the file the loader launches. """ from utils.models.drafters import dflash_repo_preference_key @@ -202,8 +199,8 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: main.setdefault(quant, []).append(sibling) plans: dict[str, GgufVariantPlan] = {} - # Every weight in the listing, so the DFlash ranking can tell a sidecar that - # names a neighbouring family from one that names this variant's. + # Every weight in the listing, so the ranking can tell a sidecar naming a + # neighbouring family from one naming this variant's. all_weight_names = [ name.rsplit("/", 1)[-1] for quant_siblings in main.values() @@ -216,9 +213,8 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: for sibling in target_main_siblings if (file := expected_file_from_sibling(sibling)) is not None ) - # The DFlash sidecar is per variant, unlike mmproj and the MTP drafter: - # its ranking is relative to the weight being fetched, so a multi-family - # repo does not hand variant B the drafter that names variant A. + # Per variant, unlike mmproj and the MTP drafter: ranked against the weight + # being fetched, so a multi-family repo does not hand B the drafter naming A. target_weight = _gguf_rfilename(target_main_siblings[0]) target_weight_name = target_weight.rsplit("/", 1)[-1] if target_weight else None dflash_sibling = preferred_dflash_sibling( diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 91400262e96..735fe9223a0 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1656,14 +1656,12 @@ def _add(d: Path) -> None: return str(best[1]) -# Drafter naming, ranking and DFlash discovery live in utils.models.drafters, so -# one rule serves the local scan, the download, the snapshot reuse and the -# offline cache alike. Only the names this module still calls are imported here; -# everything else is imported from utils.models.drafters directly at its use -# site. A re-export shim would be the friendlier move, but scripts/ -# verify_import_hoist.py scopes its __all__ exemption to package __init__.py and -# blocks unused imports in an ordinary module on purpose -- see its -# reexport_in_ordinary_module_is_still_blocked self-test. +# Drafter naming, ranking and DFlash discovery live in utils.models.drafters, so one +# rule serves the local scan, the download, the snapshot reuse and the offline cache. +# Only the names this module calls are imported here; the rest import from drafters at +# their use site. A re-export shim would read better, but verify_import_hoist.py scopes +# its __all__ exemption to package __init__.py on purpose (see its +# reexport_in_ordinary_module_is_still_blocked self-test). from utils.models.drafters import ( # noqa: E402 _drafter_launch_path, _drafter_matches_weight, @@ -1681,11 +1679,9 @@ def _add(d: Path) -> None: def dspark_preference_key(name: str) -> Tuple[int, str]: """Sort key picking the preferred DSpark sidecar by name alone. - Delegates rather than re-exports: routes/inference.py has imported this from - model_config since #7968, and repointing that call site would trip - scripts/verify_import_hoist.py's TARGET-CHANGED rule, whose relocation - exemption only covers module-level imports. One implementation still, in - utils.models.drafters. + Delegates rather than re-exports: routes/inference.py has imported this from here + since #7968, and repointing that function-local import trips verify_import_hoist.py's + TARGET-CHANGED rule, whose relocation exemption only covers module-level imports. """ return _drafters_dspark_preference_key(name) From 80c5af0c48a5e26cab8262de0ba64a20c58e1e05 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 09:20:01 +0000 Subject: [PATCH 21/26] Apply ruff-format kwarg spacing for PR #8338 --- studio/backend/hub/utils/gguf_plan.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 8eb501edb73..4b2b089e53c 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -150,9 +150,7 @@ def preferred_dflash_sibling( candidates = [ s for s in siblings - if (name := _gguf_rfilename(s)) - and "/" not in name - and name.lower().startswith("dflash-") + if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("dflash-") ] if not candidates: return None From 2534936d79128b99886c60d848e36810c3a9397e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 09:42:31 +0000 Subject: [PATCH 22/26] Fix the DFlash download plan and two stale-state reloads for PR #8338 Five review items, all reproduced first. The download plan promised the wrong files. A split sidecar contributed only its first shard, so the variant read complete while the loader's completeness check then refused the companion; it now carries the whole shard family. The pairing weight came from the listing's first sibling while plan_from_expected_files keeps the lexicographically first family, so a two-family variant key planned the discarded family's sidecar; both now use the kept family. And a root-level dflash- prefix is one real weights carry, which a listing cannot tell apart from a drafter, so a 54 GB model was planned as a companion to a 15 GB variant; a candidate is now bounded by the weights it would draft for, since a drafter is a few layers of its target and cannot outweigh it. The training coexistence guard charged the Auto DFlash sidecar even when extra args owned --spec-type, which stops the loader's promotion, so a chat load could be refused with 409 for bytes nothing would open. Extra args asking for draft-dflash keep the charge. The diffusion early-return cleared the speculative fallback state but not the DFlash retry flag, and discovery runs before the metadata read that classifies the model, so a transient sidecar failure tore down a healthy diffusion server on every Apply. Each fix has a regression test that fails without it. --- studio/backend/core/inference/llama_cpp.py | 4 ++ studio/backend/hub/utils/gguf_plan.py | 57 ++++++++++++--- studio/backend/routes/inference.py | 6 +- .../tests/test_chat_load_during_training.py | 40 +++++++++++ .../tests/test_llama_cpp_mtp_detection.py | 4 ++ .../tests/test_mtp_drafter_companion.py | 71 +++++++++++++++++++ 6 files changed, 170 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 206ee8a3108..d8a0d8c3fd6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10286,6 +10286,10 @@ def load_model(self, intent: GgufLoadIntent) -> bool: # drafter it was never going to carry. self._spec_fallback_reason = None self._spec_drafter_kind = None + # DFlash discovery runs before the metadata read that classifies this as + # diffusion, so a transient sidecar failure can have set the retry flag + # for a server that carries no drafter at all. + self._dflash_retry_needed = False with self._lock: if self._cancel_event.is_set(): logger.info("Load cancelled before diffusion server start") diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index fafb7a2dc18..7904d4b13cf 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -162,6 +162,43 @@ def preferred_dflash_sibling( ) +def dflash_plan_files( + siblings: Sequence, + weight_name: Optional[str] = None, + other_weight_names: Sequence[str] = (), + *, + max_bytes: int = 0, +) -> tuple[ExpectedFile, ...]: + """Every shard of the DFlash sidecar to plan alongside ``weight_name``, or (). + + Whole shard family, not just the ranked file: the loader refuses a companion whose + split set is incomplete, so planning shard 1 alone reports the variant complete and + then loses DFlash on the load. + + Bounded by ``max_bytes``, the variant's own weights. ``dflash-`` is a prefix real + weights carry too (Lucebox/Qwen3.6-27B-DFlash-GGUF), and a listing cannot read the + ``general.architecture`` the loader rejects them by. A drafter is a few layers of + its target, so one at least as large as the target cannot be drafting for it, and + an unknown size on either side stays out rather than risk planning a whole model. + """ + best = preferred_dflash_sibling(siblings, weight_name, other_weight_names) + if best is None: + return () + family = gguf_variant_family(getattr(best, "rfilename")) + shards = tuple( + file + for sibling in siblings + if (name := _gguf_rfilename(sibling)) + and "/" not in name + and gguf_variant_family(name) == family + and (file := expected_file_from_sibling(sibling)) is not None + ) + total = sum(max(0, int(file.size or 0)) for file in shards) + if not total or max_bytes <= 0 or total >= max_bytes: + return () + return tuple(sorted(shards, key = lambda file: file.path)) + + def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: main: dict[str, list] = {} all_mmproj = mmproj_siblings(siblings) @@ -213,21 +250,19 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: ) # Per variant, unlike mmproj and the MTP drafter: ranked against the weight # being fetched, so a multi-family repo does not hand B the drafter naming A. - target_weight = _gguf_rfilename(target_main_siblings[0]) - target_weight_name = target_weight.rsplit("/", 1)[-1] if target_weight else None - dflash_sibling = preferred_dflash_sibling( + # Ranked against the family plan_from_expected_files will KEEP, not the first + # in the listing, or a two-family variant key pairs the discarded one's sidecar. + kept_main = _one_shard_family(main_expected) + target_weight_name = ( + min(file.path for file in kept_main).rsplit("/", 1)[-1] if kept_main else None + ) + dflash_expected = dflash_plan_files( siblings, target_weight_name, [n for n in all_weight_names if n != target_weight_name], + max_bytes = sum(max(0, int(file.size or 0)) for file in kept_main), ) - dflash_expected = ( - expected_file_from_sibling(dflash_sibling) if dflash_sibling is not None else None - ) - expected_files = ( - *main_expected, - *companions_expected, - *((dflash_expected,) if dflash_expected is not None else ()), - ) + expected_files = (*main_expected, *companions_expected, *dflash_expected) plans[quant] = plan_from_expected_files( quant, expected_files, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9b4c4b6e2ff..3a6df3106f8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5694,6 +5694,7 @@ def _estimate_gguf_required_gb( _extra_args_mtp_draft_path, _extra_args_requests_dflash, _extra_args_requests_dspark, + _extra_args_set_spec_type, ) _spec_mode = _canonicalize_spec_mode(speculative_type) or "auto" @@ -5729,7 +5730,10 @@ def _estimate_gguf_required_gb( _forced_dflash = bool( _spec_mode == "dflash" or _extra_args_requests_dflash(llama_extra_args, env = {}) ) - _auto_dflash = _spec_mode == "auto" + # Extra args owning --spec-type stop the loader's Auto promotion, so charging + # the sidecar here would refuse a load for ~1.5 GiB nothing will open. Extra + # args asking for draft-dflash are _forced_dflash above and keep the charge. + _auto_dflash = _spec_mode == "auto" and not _extra_args_set_spec_type(llama_extra_args) _dflash_capable = True if _forced_dflash or _auto_dflash: try: diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 4e37d63ca5e..5b597de86c7 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1689,6 +1689,46 @@ def test_remote_estimate_bounds_the_dflash_fallback_end_to_end(self): # 1 GiB candidate that merely goes first. self.assertAlmostEqual(gb, 14.0, places = 6) + def test_auto_does_not_charge_dflash_when_extra_args_own_speculation(self): + """Extra args setting --spec-type stop the loader's Auto promotion, so the + sidecar is never opened. Charging it anyway refused a chat load with 409 for + ~1.5 GiB nothing would load. Extra args asking for draft-dflash still pay.""" + import utils.models.model_config as mc + + cfg = SimpleNamespace( + gguf_file = None, + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = None, + gguf_hf_repo = "org/repo", + gguf_variant = "Q4_K_M", + ) + variant = SimpleNamespace( + filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10 * 1024**3 + ) + siblings = [SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 2 * 1024**3)] + with ( + patch.object(mc, "list_gguf_variants", lambda repo, hf_token = None: ([variant], False)), + patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ), + self._dflash_capable(), + ): + owned = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "auto", + llama_extra_args = ["--spec-type", "ngram-mod"], + ) + asked = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "auto", + llama_extra_args = ["--spec-type", "draft-dflash"], + ) + self.assertAlmostEqual(owned, 10.0, places = 6) + self.assertAlmostEqual(asked, 12.0, places = 6) + # ── Auto charges ONE drafter, the one the promotion leaves resident ── def _auto_companion_bytes(self, siblings): diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index b4f982a0993..0993ff82ffa 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -2762,6 +2762,10 @@ def test_diffusion_load_clears_the_previous_models_spec_fallback(): assert start != -1 assert "self._spec_fallback_reason = None" in src[diffusion:start] assert "self._spec_drafter_kind = None" in src[diffusion:start] + # And the DFlash retry flag: discovery runs before the metadata read that + # classifies this as diffusion, so a transient sidecar failure can set it for a + # server that will never carry a drafter, and the dedupe reads it too. + assert "self._dflash_retry_needed = False" in src[diffusion:start] def test_already_in_target_state_reloads_after_a_dflash_fetch_that_dropped(): diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 41ed40140ad..8d8ec20bbf9 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -3212,3 +3212,74 @@ def _fetch( assert got == str(tmp_path / "dflash-kquant.gguf") assert b._dflash_retry_needed is False + + +# ── The DFlash sidecar in the download plan ────────────────────────── +# +# The plan has to promise exactly what the loader will open: every shard of it, +# paired with the weight family the plan keeps, and never a whole model that +# merely carries the prefix. + + +def test_variant_plans_carry_every_shard_of_a_split_dflash_sidecar(): + """The loader refuses a companion whose split set is incomplete, so planning + shard 1 alone reports the variant complete and then loses DFlash.""" + plans = build_gguf_variant_plans( + [ + _sib("model-Q4_K_M.gguf", 15_000, "main"), + _sib("dflash-kquant-00001-of-00002.gguf", 800, "d1"), + _sib("dflash-kquant-00002-of-00002.gguf", 800, "d2"), + ] + ) + targets = set(plans["q4_k_m"].target_filenames) + assert "dflash-kquant-00001-of-00002.gguf" in targets + assert "dflash-kquant-00002-of-00002.gguf" in targets + + +def test_variant_plans_pair_dflash_with_the_weight_family_they_keep(): + """_one_shard_family keeps the lexicographically first family, so ranking the + sidecar against the listing's first weight pairs the discarded one.""" + plans = build_gguf_variant_plans( + [ + # Listing order puts the discarded family first. + _sib("QwQ-32B.BF16-00001-of-00002.gguf", 30_000, "b1"), + _sib("QwQ-32B.BF16-00002-of-00002.gguf", 30_000, "b2"), + _sib("QwQ-32B-BF16-00001-of-00002.gguf", 30_000, "a1"), + _sib("QwQ-32B-BF16-00002-of-00002.gguf", 30_000, "a2"), + _sib("dflash-QwQ-32B-BF16-Q8_0.gguf", 2_000, "da"), + _sib("dflash-QwQ-32B.BF16-Q8_0.gguf", 2_000, "db"), + ] + ) + plan = plans["bf16"] + assert plan.main_filenames == frozenset( + {"QwQ-32B-BF16-00001-of-00002.gguf", "QwQ-32B-BF16-00002-of-00002.gguf"} + ) + assert "dflash-QwQ-32B-BF16-Q8_0.gguf" in plan.target_filenames + assert "dflash-QwQ-32B.BF16-Q8_0.gguf" not in plan.target_filenames + + +def test_variant_plans_skip_a_dflash_prefixed_file_too_big_to_be_a_drafter(): + """dflash- is a prefix real weights carry (Lucebox/Qwen3.6-27B-DFlash-GGUF) and a + listing cannot read the architecture, so size is the only bound available: a + drafter is a few layers of its target and cannot outweigh it.""" + plans = build_gguf_variant_plans( + [ + _sib("Qwen3.6-27B-Q4_K_M.gguf", 15_000, "main"), + _sib("dflash-Qwen3.6-27B-BF16.gguf", 54_000, "impostor"), + ] + ) + plan = plans["q4_k_m"] + assert "dflash-Qwen3.6-27B-BF16.gguf" not in plan.target_filenames + assert plan.download_size_bytes == 15_000 + + +def test_variant_plans_still_carry_the_published_dflash_sidecar(): + """The Muse-Glimmer shape the feature ships for stays planned.""" + plans = build_gguf_variant_plans( + [ + _sib("Muse-Glimmer-30B-UD-Q4_K_XL.gguf", 15_878, "main"), + _sib("mmproj-kquant.gguf", 1_400, "mmproj"), + _sib("dflash-kquant.gguf", 1_631, "dflash"), + ] + ) + assert "dflash-kquant.gguf" in plans["ud-q4_k_xl"].target_filenames From 0460d2c52bd45ab11860320bc25ae9184cc71580 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 10:03:48 +0000 Subject: [PATCH 23/26] Carry the DFlash plan bounds into the runtime paths for PR #8338 Four review items from the second round, each reproduced first. The budget still charged a forced dflash mode when extra args owned --spec-type. _build_speculative_flags returns before any mode branch in that case, so neither the forced mode nor the Auto promotion reaches the sidecar; only extra args asking for draft-dflash themselves still pay. The runtime picker had no size bound, so a root-level ordinary weight carrying the dflash- prefix downloaded in full before its header could be read, which is exactly what the download plan now refuses. It applies the same bound, sized from the repo listing, and an unavailable size leaves the candidate eligible as before. A permanent listing error records no answer at all, so _dflash_sidecar_absent stayed False and the drafter_not_found arm relaunched a healthy drafter-free server on every Apply. DFlash asks through _dflash_retry_needed instead, which is set only for the failures worth another attempt. A listing holding part of a split companion returned its first shard as usable, contradicting the complete-set checks on snapshot and cache reuse and handing llama-server a set it cannot open. The filename carries the set size, so the listing is now checked before the download. Each fix has a regression test that fails without it. --- studio/backend/core/inference/llama_cpp.py | 53 +++++++++++++++- studio/backend/routes/inference.py | 13 ++-- .../tests/test_chat_load_during_training.py | 8 +++ .../tests/test_llama_cpp_mtp_detection.py | 17 +++++ .../tests/test_mtp_drafter_companion.py | 63 +++++++++++++++++++ 5 files changed, 147 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d8a0d8c3fd6..d35c4f2a34a 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4120,7 +4120,12 @@ def _norm(value): and self._spec_fallback_reason == "drafter_not_found" and speculative_type in ("auto", "mtp", "mtp+ngram", "dspark", "dflash") and not (self._spec_drafter_kind == "dspark" and self._dspark_sidecar_absent) - and not (self._spec_drafter_kind == "dflash" and self._dflash_sidecar_absent) + # DFlash asks through _dflash_retry_needed below instead. A permanent + # listing error (gated repo, offline) records no answer at all, so + # _dflash_sidecar_absent stays False and this arm relaunched a healthy + # server on every Apply; the flag is set only for the failures worth + # another attempt. + and self._spec_drafter_kind != "dflash" and not spec_owned_by_extra_args ): return False @@ -8036,6 +8041,19 @@ def _pick_from(names: list[str]) -> Optional[str]: # disagree about what belongs to a set; empty for the single-file case, # which is every mmproj and every published sidecar so far. extra_shards = _gguf_extra_shards(available, target) + # The filename carries the set size, so a listing missing part of it is + # answerable before the download rather than after: half a split companion is + # not a companion, the same rule the snapshot and cache reuse above apply. + _shard_total = _SHARD_FULL_RE.match(target) + if _shard_total and len(extra_shards) + 1 != int(_shard_total.group(3)): + logger.info( + "Skipping %s: %s lists %d of %s shards.", + label, + hf_repo, + len(extra_shards) + 1, + _shard_total.group(3), + ) + return None try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -8430,11 +8448,22 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: for name in candidates if name.lower().endswith(".gguf") and not _is_root_dflash_drafter_path(name) ] + # Same bound dflash_plan_files applies, for the same reason: dflash- is a + # prefix real weights carry, the header only reads once the bytes are here, + # and a drafter is a few layers of its target so it cannot outweigh it. + # Sizes the listing does not carry leave the candidate in, as before. + sizes = self._remote_root_gguf_sizes(hf_repo, hf_token) + try: + target_bytes = (self._get_gguf_size_bytes(near_path) or 0) if near_path else 0 + except OSError: + target_bytes = 0 files = sorted( ( name for name in candidates - if _is_root_dflash_drafter_path(name) and Path(name).name not in rejected + if _is_root_dflash_drafter_path(name) + and Path(name).name not in rejected + and not (target_bytes and sizes.get(Path(name).name, 0) >= target_bytes) ), key = lambda name: dflash_repo_preference_key(name, weight_name, others), ) @@ -8536,6 +8565,26 @@ def _validated(path: Optional[str]) -> Optional[str]: self._dflash_sidecar_absent = outcome.get("listed") is False return found + @staticmethod + def _remote_root_gguf_sizes(hf_repo: str, hf_token: Optional[str] = None) -> dict[str, int]: + """Root-level GGUF basenames to their listed byte size, {} when unavailable. + + list_repo_files answers names only, so the size bound needs the metadata + listing. Best effort: an empty map leaves every candidate eligible, which is + the behaviour before the bound existed. + """ + try: + from huggingface_hub import model_info + info = model_info(hf_repo, token = hf_token, files_metadata = True) + return { + name: int(getattr(sibling, "size", 0) or 0) + for sibling in (getattr(info, "siblings", None) or []) + if isinstance(name := getattr(sibling, "rfilename", None), str) and "/" not in name + } + except Exception as exc: + logger.debug("Could not size the DFlash candidates for %s: %s", hf_repo, exc) + return {} + def _dspark_wins_auto( self, *, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3a6df3106f8..2b877989ee7 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5727,13 +5727,16 @@ def _estimate_gguf_required_gb( # DFlash: same shape as DSpark above, and Auto sizes it for the same # reason. The sidecar is ~1.5 GiB rather than ~11 GB, but a guard that # protects a running training job still has to charge for it. + # Extra args owning --spec-type end _build_speculative_flags before any mode + # branch runs, so neither the forced mode nor the Auto promotion reaches the + # sidecar and charging it refuses a load for bytes nothing will open. Extra + # args asking for draft-dflash themselves are the one case that still pays. + _extra_args_own_spec = _extra_args_set_spec_type(llama_extra_args) _forced_dflash = bool( - _spec_mode == "dflash" or _extra_args_requests_dflash(llama_extra_args, env = {}) + _extra_args_requests_dflash(llama_extra_args, env = {}) + or (_spec_mode == "dflash" and not _extra_args_own_spec) ) - # Extra args owning --spec-type stop the loader's Auto promotion, so charging - # the sidecar here would refuse a load for ~1.5 GiB nothing will open. Extra - # args asking for draft-dflash are _forced_dflash above and keep the charge. - _auto_dflash = _spec_mode == "auto" and not _extra_args_set_spec_type(llama_extra_args) + _auto_dflash = _spec_mode == "auto" and not _extra_args_own_spec _dflash_capable = True if _forced_dflash or _auto_dflash: try: diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 5b597de86c7..61cb2b901ae 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1726,8 +1726,16 @@ def test_auto_does_not_charge_dflash_when_extra_args_own_speculation(self): speculative_type = "auto", llama_extra_args = ["--spec-type", "draft-dflash"], ) + # Same for the forced mode: _build_speculative_flags returns before any + # mode branch when extra args own --spec-type, so dflash never emits. + forced = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "dflash", + llama_extra_args = ["--spec-type", "ngram-mod"], + ) self.assertAlmostEqual(owned, 10.0, places = 6) self.assertAlmostEqual(asked, 12.0, places = 6) + self.assertAlmostEqual(forced, 10.0, places = 6) # ── Auto charges ONE drafter, the one the promotion leaves resident ── diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 0993ff82ffa..25dc680ddec 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -2768,6 +2768,23 @@ def test_diffusion_load_clears_the_previous_models_spec_fallback(): assert "self._dflash_retry_needed = False" in src[diffusion:start] +def test_already_in_target_state_settles_a_dflash_listing_that_never_answered(): + """A permanent listing error (gated repo, offline) records no answer, so + _dflash_sidecar_absent stays False. The drafter_not_found arm read that as "worth + another go" and relaunched a healthy drafter-free server on every Apply. DFlash + asks through _dflash_retry_needed instead, which a permanent error never sets.""" + backend = _mtp_backend( + _model_identifier = "unsloth/Muse-Glimmer-30B-GGUF", + _speculative_type = "default", + _gguf_path = None, + _spec_fallback_reason = "drafter_not_found", + _spec_drafter_kind = "dflash", + _dflash_sidecar_absent = False, + _dflash_retry_needed = False, + ) + assert _matches(backend, **_binary_fallback_kwargs()) is True + + def test_already_in_target_state_reloads_after_a_dflash_fetch_that_dropped(): """Under Auto a lost sidecar leaves no fallback reason at all -- the promotion never ran -- so the flag is the only thing that can ask for one more attempt.""" diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 8d8ec20bbf9..a349deb20f3 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -3283,3 +3283,66 @@ def test_variant_plans_still_carry_the_published_dflash_sidecar(): ] ) assert "dflash-kquant.gguf" in plans["ud-q4_k_xl"].target_filenames + + +def test_download_dflash_skips_a_root_weight_too_big_to_be_a_drafter(monkeypatch, tmp_path): + """The runtime picker needs the bound the plan has: a root dflash-*.gguf that is + an ordinary weight passes the filename test, and the header can only answer once + the whole object is on disk.""" + from core.inference.llama_cpp import LlamaCppBackend + + weight = tmp_path / "Qwen3.6-27B-Q4_K_M.gguf" + weight.write_bytes(b"x" * 4_000) + monkeypatch.setattr( + LlamaCppBackend, + "_remote_root_gguf_sizes", + staticmethod( + lambda repo, token = None: { + "dflash-Qwen3.6-27B-BF16.gguf": 40_000, + "dflash-kquant.gguf": 500, + } + ), + ) + picked = _dflash_download_pick( + monkeypatch, + listing = ["Qwen3.6-27B-Q4_K_M.gguf", "dflash-Qwen3.6-27B-BF16.gguf", "dflash-kquant.gguf"], + near_path = str(weight), + ) + assert picked == "dflash-kquant.gguf" + + +def test_download_companion_refuses_a_listing_missing_part_of_a_split_set(monkeypatch): + """The snapshot and cache paths both refuse half a split companion; the download + path returned shard 1 and handed llama-server a set it cannot open.""" + import core.inference.llama_cpp as llama_cpp_module + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr( + llama_cpp_module, "_companion_snapshot_sibling", lambda near_path, pick: None + ) + # Patched at the source: _download_companion_gguf imports list_repo_files inside + # its own body, so a module attribute on llama_cpp is never consulted. + import huggingface_hub + + monkeypatch.setattr( + huggingface_hub, + "list_repo_files", + lambda repo, token = None: ["dflash-kquant-00001-of-00002.gguf"], + ) + downloads: list = [] + monkeypatch.setattr( + llama_cpp_module, + "hf_hub_download_with_xet_fallback", + lambda *a, **k: downloads.append(a) or "/tmp/x.gguf", + raising = False, + ) + b = LlamaCppBackend() + got = b._download_companion_gguf( + hf_repo = "org/repo", + hf_token = None, + pick = lambda files: next(iter(files), None), + label = "DFlash drafter", + ) + assert got is None + assert downloads == [] From c05c2c249ed1f98a65dd34da8f99608004e0e1e7 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 10:27:55 +0000 Subject: [PATCH 24/26] Make the DFlash size and split rules agree across plan, fetch and guard for PR #8338 Six review items from the third round, each reproduced first. Five are places the previous round's rules had not reached. dflash_plan_files now filters candidate families before ranking rather than after, so a half-published split set or an oversized ordinary weight at the top of the order steps aside for a usable sidecar behind it instead of taking the plan down with it. It also applies the split-completeness rule the runtime got last round, since planning a set the listing only half carries reports the download complete and then loses DFlash. The runtime size bound compared the picked shard rather than its whole set, so a split ordinary weight whose halves each sit under the target still downloaded in full. It sums the family now, through a shared helper. The training coexistence guard took the maximum over every root candidate with no size bound at all, charging gigabytes for files the fetch itself refuses. dflash_budget_bytes takes the target size and drops them. The incomplete-split rejection added last round lands after outcome["listed"] is set, so DSpark read a settled answer as retryable and relaunched a healthy server on every Apply. It records absence explicitly. SpeculativeType omitted dflash, so Typer rejected --speculative-type dflash before any of the new loading code ran and the mode was reachable only through Auto. Each fix has a regression test that fails without it. --- studio/backend/core/inference/llama_cpp.py | 18 +++- studio/backend/hub/utils/gguf_plan.py | 58 ++++++++---- studio/backend/routes/inference.py | 11 ++- .../tests/test_chat_load_during_training.py | 22 +++++ .../tests/test_mtp_drafter_companion.py | 91 +++++++++++++++++++ .../backend/utils/models/drafters/budget.py | 18 +++- unsloth_cli/_inference.py | 4 +- .../tests/test_studio_run_parallel_flag.py | 11 +++ 8 files changed, 208 insertions(+), 25 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 731cd2877bc..2b7cb7e2d04 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1788,6 +1788,13 @@ def wrapped(self, intent: GgufLoadIntent): return wrapped +def _drafter_set_bytes(sizes: Mapping[str, int], name: str) -> int: + """Total listed bytes of ``name``'s whole shard set (its own size when single).""" + return (sizes.get(name, 0) or 0) + sum( + sizes.get(shard, 0) or 0 for shard in _gguf_extra_shards(sizes, name) + ) + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -8115,6 +8122,11 @@ def _pick_from(names: list[str]) -> Optional[str]: len(extra_shards) + 1, _shard_total.group(3), ) + # Settled, not retryable: the listing answered and this is what it said, so + # the caller records absence rather than reloading on every Apply hoping the + # same listing says something else. + if outcome is not None: + outcome["listed"] = False return None try: logger.info(f"Downloading {label}: {hf_repo}/{target}") @@ -8532,7 +8544,11 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: for name in candidates if _is_root_dflash_drafter_path(name) and Path(name).name not in rejected - and not (target_bytes and sizes.get(Path(name).name, 0) >= target_bytes) + # Whole shard set, not the picked shard: each half of a split + # ordinary weight can sit under the target while the set does not. + and not ( + target_bytes and _drafter_set_bytes(sizes, Path(name).name) >= target_bytes + ) ), key = lambda name: dflash_repo_preference_key(name, weight_name, others), ) diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 7904d4b13cf..ade768699fc 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -3,6 +3,7 @@ from __future__ import annotations +import re from dataclasses import dataclass from typing import Optional, Sequence @@ -131,6 +132,12 @@ def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: return candidates[0] if candidates else None +def _shard_set_size(path: str) -> Optional[int]: + """How many shards ``path``'s name says its set has, or None when it is single.""" + match = re.search(r"-\d{5}-of-(\d{5})\.gguf$", path, re.IGNORECASE) + return int(match.group(1)) if match else None + + def preferred_dflash_sibling( siblings: Sequence, weight_name: Optional[str] = None, @@ -173,30 +180,49 @@ def dflash_plan_files( Whole shard family, not just the ranked file: the loader refuses a companion whose split set is incomplete, so planning shard 1 alone reports the variant complete and - then loses DFlash on the load. + then loses DFlash on the load. A family the listing only half publishes is dropped + for the same reason, before it can be ranked. Bounded by ``max_bytes``, the variant's own weights. ``dflash-`` is a prefix real - weights carry too (Lucebox/Qwen3.6-27B-DFlash-GGUF), and a listing cannot read the + weights carry (Lucebox/Qwen3.6-27B-DFlash-GGUF), and a listing cannot read the ``general.architecture`` the loader rejects them by. A drafter is a few layers of its target, so one at least as large as the target cannot be drafting for it, and an unknown size on either side stays out rather than risk planning a whole model. + + Both rules filter families BEFORE the ranking, so an oversized or half-published + name at the top of the order steps aside for a usable sidecar behind it instead of + taking the plan down with it. """ - best = preferred_dflash_sibling(siblings, weight_name, other_weight_names) - if best is None: + from utils.models.drafters import dflash_repo_preference_key + + families: dict[str, list[ExpectedFile]] = {} + for sibling in siblings: + name = _gguf_rfilename(sibling) + if not name or "/" in name or not name.lower().startswith("dflash-"): + continue + file = expected_file_from_sibling(sibling) + if file is not None: + families.setdefault(gguf_variant_family(name), []).append(file) + + eligible: dict[str, tuple[ExpectedFile, ...]] = {} + for family, files in families.items(): + shards = tuple(sorted(files, key = lambda file: file.path)) + listed = _shard_set_size(shards[0].path) + if listed is not None and len(shards) != listed: + continue + total = sum(max(0, int(file.size or 0)) for file in shards) + if not total or max_bytes <= 0 or total >= max_bytes: + continue + eligible[family] = shards + if not eligible: return () - family = gguf_variant_family(getattr(best, "rfilename")) - shards = tuple( - file - for sibling in siblings - if (name := _gguf_rfilename(sibling)) - and "/" not in name - and gguf_variant_family(name) == family - and (file := expected_file_from_sibling(sibling)) is not None + best = min( + eligible, + key = lambda family: dflash_repo_preference_key( + eligible[family][0].path, weight_name, other_weight_names + ), ) - total = sum(max(0, int(file.size or 0)) for file in shards) - if not total or max_bytes <= 0 or total >= max_bytes: - return () - return tuple(sorted(shards, key = lambda file: file.path)) + return eligible[best] def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2b877989ee7..1e0be8d90df 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5430,6 +5430,7 @@ def _remote_gguf_companion_bytes( include_dspark: bool = False, include_dflash: bool = False, dspark_first: bool = False, + weight_bytes: int = 0, ) -> int: """Bytes of companion GGUFs the requested launch keeps resident. 0 on error. @@ -5485,8 +5486,10 @@ def _remote_gguf_companion_bytes( else 0 ) # Bounded rather than picked: see dflash_budget_bytes for why the max - # over whole shard sets is the answer a listing can give. - dflash_bytes = dflash_budget_bytes(dflash_sizes, _gguf_extra_shards) + # over whole shard sets is the answer a listing can give. Bounded by the + # target too, so the guard stops charging for the oversized candidates the + # fetch itself now refuses. + dflash_bytes = dflash_budget_bytes(dflash_sizes, _gguf_extra_shards, weight_bytes) if not dspark_first: return total + mtp_bytes + dspark_bytes + dflash_bytes if dspark_candidates: @@ -5855,6 +5858,10 @@ def _same_file_key(p: str) -> str: ), include_dspark = (_dspark_capable and (_auto_dspark or dspark_requested)), include_dflash = (_dflash_capable and (_auto_dflash or dflash_requested)), + # The size the DFlash bound measures candidates against, so the guard + # stops charging for weights the fetch would refuse as too big to be + # a drafter. + weight_bytes = int(main_bytes or 0), # ... except where the listing settles it. Auto launches exactly # one drafter, in a fixed order, so once the listing says which # kinds the repo has, charging the losers is not caution, it is a diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 61cb2b901ae..522b897e88e 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1737,6 +1737,28 @@ def test_auto_does_not_charge_dflash_when_extra_args_own_speculation(self): self.assertAlmostEqual(asked, 12.0, places = 6) self.assertAlmostEqual(forced, 10.0, places = 6) + def test_remote_dflash_sizing_drops_a_candidate_too_big_to_be_a_drafter(self): + """The fetch refuses an oversized root dflash-*.gguf, so charging for it is a + 409 for bytes that will never be resident.""" + siblings = [ + SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10 * 1024**3), + SimpleNamespace(rfilename = "dflash-model-BF16.gguf", size = 40 * 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 1024**3), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + charged = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dflash = True, + weight_bytes = 10 * 1024**3, + ) + self.assertEqual(charged, 1024**3) + # ── Auto charges ONE drafter, the one the promotion leaves resident ── def _auto_companion_bytes(self, siblings): diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index a349deb20f3..ae9c808cb9b 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -3346,3 +3346,94 @@ def test_download_companion_refuses_a_listing_missing_part_of_a_split_set(monkey ) assert got is None assert downloads == [] + + +def test_variant_plans_skip_a_half_published_split_sidecar(): + """Planning a set the listing only half carries reports the download complete and + then loses DFlash, since the loader refuses the partial set.""" + plans = build_gguf_variant_plans( + [ + _sib("model-Q4_K_M.gguf", 15_000, "main"), + _sib("dflash-kquant-00001-of-00002.gguf", 800, "d1"), + ] + ) + assert not [f for f in plans["q4_k_m"].target_filenames if f.startswith("dflash-")] + + +def test_variant_plans_fall_through_to_a_usable_sidecar_behind_an_oversized_one(): + """Both plan rules filter before the ranking, so the impostor at the top of the + order steps aside instead of taking the whole plan down with it.""" + plans = build_gguf_variant_plans( + [ + _sib("model-B-Q4_K_M.gguf", 15_000, "main"), + # Ranks first (names this weight) but is a whole model. + _sib("dflash-model-B-BF16.gguf", 54_000, "impostor"), + _sib("dflash-kquant.gguf", 900, "real"), + ] + ) + targets = plans["q4_k_m"].target_filenames + assert "dflash-kquant.gguf" in targets + assert "dflash-model-B-BF16.gguf" not in targets + + +def test_download_dflash_sums_a_split_family_before_the_size_bound(monkeypatch, tmp_path): + """Each shard of a split ordinary weight can sit under the target while the set + does not, so bounding the picked shard alone still fetched the whole thing.""" + from core.inference.llama_cpp import LlamaCppBackend + + weight = tmp_path / "model-Q4_K_M.gguf" + weight.write_bytes(b"x" * 10_000) + monkeypatch.setattr( + LlamaCppBackend, + "_remote_root_gguf_sizes", + staticmethod( + lambda repo, token = None: { + "dflash-big-00001-of-00002.gguf": 6_000, + "dflash-big-00002-of-00002.gguf": 6_000, + "dflash-kquant.gguf": 500, + } + ), + ) + picked = _dflash_download_pick( + monkeypatch, + listing = [ + "model-Q4_K_M.gguf", + "dflash-big-00001-of-00002.gguf", + "dflash-big-00002-of-00002.gguf", + "dflash-kquant.gguf", + ], + near_path = str(weight), + ) + assert picked == "dflash-kquant.gguf" + + +def test_download_companion_records_an_incomplete_listing_as_settled(): + """The completeness rejection lands after outcome["listed"] was set true, so + without this the caller reads a settled answer as one worth retrying forever.""" + import core.inference.llama_cpp as llama_cpp_module + import huggingface_hub + from core.inference.llama_cpp import LlamaCppBackend + + import pytest as _pytest + + mp = _pytest.MonkeyPatch() + try: + mp.delenv("HF_HUB_OFFLINE", raising = False) + mp.setattr(llama_cpp_module, "_companion_snapshot_sibling", lambda near_path, pick: None) + mp.setattr( + huggingface_hub, + "list_repo_files", + lambda repo, token = None: ["dspark-kquant-00001-of-00002.gguf"], + ) + outcome: dict = {} + got = LlamaCppBackend()._download_companion_gguf( + hf_repo = "org/repo", + hf_token = None, + pick = lambda files: next(iter(files), None), + label = "DSpark drafter", + outcome = outcome, + ) + finally: + mp.undo() + assert got is None + assert outcome.get("listed") is False diff --git a/studio/backend/utils/models/drafters/budget.py b/studio/backend/utils/models/drafters/budget.py index 44aa1a03388..d8065b32f31 100644 --- a/studio/backend/utils/models/drafters/budget.py +++ b/studio/backend/utils/models/drafters/budget.py @@ -14,7 +14,9 @@ def dflash_budget_bytes( - sizes: Mapping[str, int], extra_shards: Callable[[Mapping[str, int], str], list] + sizes: Mapping[str, int], + extra_shards: Callable[[Mapping[str, int], str], list], + target_bytes: int = 0, ) -> int: """A safe bound on the DFlash sidecar a load may end up resident on. @@ -32,11 +34,17 @@ def dflash_budget_bytes( sibling, all of which llama-server keeps resident. Sizing one shard would halve a two-shard sidecar, and under-estimating is the direction that waves a load through and then exhausts VRAM. + + ``target_bytes`` drops the candidates the fetch itself now refuses: a drafter is + a few layers of its target, so a set at least that large is an ordinary weight + wearing the prefix and is never made resident. Zero means unknown, which keeps + every candidate, as before. """ + totals = ( + size + sum(sizes.get(shard, 0) for shard in extra_shards(sizes, name)) + for name, size in sizes.items() + ) return max( - ( - size + sum(sizes.get(shard, 0) for shard in extra_shards(sizes, name)) - for name, size in sizes.items() - ), + (total for total in totals if not target_bytes or total < target_bytes), default = 0, ) diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index 5e925ed9c5a..8a2a80d3ea6 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -18,7 +18,9 @@ # _CANONICAL_SPEC_MODES. Named once so the CLI's option annotations, the HTTP # payload builders and the in-process loader cannot drift apart when a mode is # added; typer reads it at runtime to validate --speculative-type. -SpeculativeType = Literal["auto", "mtp", "dspark", "ngram", "mtp+ngram", "off", "ngram-simple"] +SpeculativeType = Literal[ + "auto", "mtp", "dspark", "dflash", "ngram", "mtp+ngram", "off", "ngram-simple" +] _THINK_OPEN = "" _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py index d52c6cfbcf5..33c0ff35486 100644 --- a/unsloth_cli/tests/test_studio_run_parallel_flag.py +++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py @@ -444,6 +444,17 @@ def test_reexec_forwards_speculative_options(monkeypatch): assert _value_after(argv, "--spec-draft-n-max") == "3", argv +def test_reexec_forwards_dflash_speculative_type(monkeypatch): + """SpeculativeType is what Typer validates the option against, so a mode missing + from the literal is rejected before any loading code runs.""" + result, captured = _invoke_run( + monkeypatch, + _BASE + ["--speculative-type", "dflash"], + ) + assert len(captured) == 1, result.output + assert _value_after(captured[0]["argv"], "--speculative-type") == "dflash" + + def test_reexec_omits_unset_speculative_options(monkeypatch): result, captured = _invoke_run(monkeypatch, _BASE) assert len(captured) == 1, result.output From f3f7201563315cb76a10fd01a51e581fadde56aa Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 10:50:29 +0000 Subject: [PATCH 25/26] Price DSpark by shard set and share one split-listing rule for PR #8338 Four review items from the fourth round, two of them under-charges that could admit a load beside a running training job and then exhaust VRAM. The guard priced a remote DSpark sidecar as the single file the ranking picked, while llama-server maps every shard of a split set, so a two-shard sidecar was budgeted at roughly half its resident weight. DSpark candidates are grouped into shard families now and the selected family's total is charged, matching what DFlash already did. Auto granted DSpark first refusal on the strength of the listing alone. Since the fetch now refuses an incomplete split set, the load falls through to DFlash, which can be the larger of the two, and the guard had already returned the DSpark figure. Only a complete set settles it. The runtime DFlash picker filtered incomplete families after ranking rather than before, so a half-published set at the top of the order returned a shard, _download_companion_gguf refused it, and the loop ended instead of reaching the complete sidecar behind it. Extras owning --spec-type with their own --model-draft charged the discovered sidecar as well, though _build_speculative_flags returns before that one is emitted. Only the drafter that launches is charged now. Extras without --spec-type still charge both, since Studio emits its own and which lands is genuinely unknown. The listing completeness rule was about to have three copies, so it moved into utils.models.drafters as split_listing_is_complete and the plan, the fetch and the guard all call it. Each fix has a regression test that fails without it. --- studio/backend/core/inference/llama_cpp.py | 9 +- studio/backend/hub/utils/gguf_plan.py | 12 +-- studio/backend/routes/inference.py | 49 +++++++++-- .../tests/test_chat_load_during_training.py | 84 +++++++++++++++++++ .../tests/test_mtp_drafter_companion.py | 25 ++++++ .../backend/utils/models/drafters/__init__.py | 2 + .../backend/utils/models/drafters/common.py | 22 +++++ 7 files changed, 183 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2b7cb7e2d04..02d9d8ff4d4 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8519,7 +8519,10 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: # The rest of the listing supplies the neighbouring weights that make # a foreign sidecar recognisable (a sidecar naming no family at all # stays eligible, which is what the published one does). - from utils.models.drafters import dflash_repo_preference_key + from utils.models.drafters import ( + dflash_repo_preference_key, + split_listing_is_complete, + ) # Root level only, as the local scan is: a nested dflash-*.gguf is an # ordinary weight, and offering it here spends its entire download @@ -8544,6 +8547,10 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: for name in candidates if _is_root_dflash_drafter_path(name) and Path(name).name not in rejected + # Before the ranking, as dflash_plan_files is: a set the listing + # only half carries is refused by the fetch, and returning it + # ends the loop instead of reaching the complete one behind it. + and split_listing_is_complete(candidates, name) # Whole shard set, not the picked shard: each half of a split # ordinary weight can sit under the target while the set does not. and not ( diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index ade768699fc..fb26f9e81cb 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -3,7 +3,6 @@ from __future__ import annotations -import re from dataclasses import dataclass from typing import Optional, Sequence @@ -132,12 +131,6 @@ def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: return candidates[0] if candidates else None -def _shard_set_size(path: str) -> Optional[int]: - """How many shards ``path``'s name says its set has, or None when it is single.""" - match = re.search(r"-\d{5}-of-(\d{5})\.gguf$", path, re.IGNORECASE) - return int(match.group(1)) if match else None - - def preferred_dflash_sibling( siblings: Sequence, weight_name: Optional[str] = None, @@ -193,7 +186,7 @@ def dflash_plan_files( name at the top of the order steps aside for a usable sidecar behind it instead of taking the plan down with it. """ - from utils.models.drafters import dflash_repo_preference_key + from utils.models.drafters import dflash_repo_preference_key, split_listing_is_complete families: dict[str, list[ExpectedFile]] = {} for sibling in siblings: @@ -207,8 +200,7 @@ def dflash_plan_files( eligible: dict[str, tuple[ExpectedFile, ...]] = {} for family, files in families.items(): shards = tuple(sorted(files, key = lambda file: file.path)) - listed = _shard_set_size(shards[0].path) - if listed is not None and len(shards) != listed: + if not split_listing_is_complete([f.path for f in shards], shards[0].path): continue total = sum(max(0, int(file.size or 0)) for file in shards) if not total or max_bytes <= 0 or total >= max_bytes: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1e0be8d90df..b2ca7596a66 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5449,7 +5449,7 @@ def _remote_gguf_companion_bytes( _is_root_dflash_drafter_path, ) from huggingface_hub import model_info - from utils.models.drafters import dflash_budget_bytes + from utils.models.drafters import dflash_budget_bytes, split_listing_is_complete from utils.models.model_config import dspark_preference_key info = model_info(repo, token = hf_token, files_metadata = True) @@ -5477,12 +5477,25 @@ def _remote_gguf_companion_bytes( # counting it here would price a file the load cannot fetch. if include_dflash and _is_root_dflash_drafter_path(name): dflash_sizes[name] = size - # Same preference order the download uses, so the budget sizes the file - # the launch will actually fetch. DSpark has no post-fetch rejection, so - # the best-ranked candidate is the one that lands. + # Same preference order the download uses, so the budget sizes the file the + # launch will actually fetch, and by whole shard SET: llama-server maps every + # shard, so pricing the one the ranking picked halved a two-shard sidecar and + # let the guard admit a load that evicts the training run it protects. + # Incomplete sets are dropped rather than priced, because the fetch now refuses + # them: a listing missing a shard is not a sidecar this load can end up on. + _dspark_sizes = dict(dspark_candidates) + dspark_families = [ + ( + name, + size + + sum(_dspark_sizes.get(s, 0) for s in _gguf_extra_shards(_dspark_sizes, name)), + ) + for name, size in dspark_candidates + if split_listing_is_complete(_dspark_sizes, name) + ] dspark_bytes = ( - min(dspark_candidates, key = lambda c: dspark_preference_key(c[0]))[1] - if dspark_candidates + min(dspark_families, key = lambda c: dspark_preference_key(c[0]))[1] + if dspark_families else 0 ) # Bounded rather than picked: see dflash_budget_bytes for why the max @@ -5492,12 +5505,14 @@ def _remote_gguf_companion_bytes( dflash_bytes = dflash_budget_bytes(dflash_sizes, _gguf_extra_shards, weight_bytes) if not dspark_first: return total + mtp_bytes + dspark_bytes + dflash_bytes - if dspark_candidates: + if dspark_families: # DSpark takes first refusal in the Auto promotion, so a listed # sidecar settles the load: the DFlash fetch stands down and # mtp_draft_path is replaced by the DSpark one. Charging the other # two is not the safe over-estimate it is for a repo whose listing - # has not answered yet, it is a 409 for a load that fits. + # has not answered yet, it is a 409 for a load that fits. Only a + # COMPLETE set settles it, since the fetch falls through to DFlash on + # one it has to reject, and that one can be the larger of the two. return total + dspark_bytes if dflash_sizes: # DFlash is the other Auto promotion and replaces mtp_draft_path the @@ -5786,7 +5801,23 @@ def _same_file_key(p: str) -> str: if dspark_requested: _sized_attrs.append("gguf_dspark_file") elif dflash_requested: - _sized_attrs.append("gguf_dflash_file") + # Only when the extras own --spec-type: _build_speculative_flags then + # returns before discovery's sidecar is ever emitted, so llama-server + # opens the extras' --model-draft alone and charging both billed two + # drafters for one. Without --spec-type Studio still emits its own, and + # which of the two lands is genuinely unknown, so both stay charged. + _manual_draft = ( + _extra_args_mtp_draft_path(llama_extra_args, env = {}) + if _extra_args_own_spec + else None + ) + _configured = getattr(config, "gguf_dflash_file", None) + if not ( + _manual_draft + and _configured + and _same_file_key(str(_manual_draft)) != _same_file_key(str(_configured)) + ): + _sized_attrs.append("gguf_dflash_file") else: _sized_attrs.append("gguf_mtp_file") diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 522b897e88e..278ff6000b9 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -1392,6 +1392,46 @@ def test_extra_args_drafter_is_charged_once_when_it_is_the_local_sidecar(self): self.assertAlmostEqual(through_link, 5000 / (1024**3), places = 9) self.assertAlmostEqual(separate, 9000 / (1024**3), places = 9) # 2000+3000+4000 + def test_extras_owning_spec_type_charge_only_their_own_drafter(self): + """--spec-type in the extras ends _build_speculative_flags before discovery's + sidecar is emitted, so llama-server opens the extras' --model-draft alone and + charging the configured one too billed two drafters for the one that loads.""" + import tempfile + + with tempfile.TemporaryDirectory() as d: + p = Path(d) + target = p / "model.gguf" + sidecar = p / "dflash-kquant.gguf" + custom = p / "custom-dflash.gguf" + target.write_bytes(b"x" * 2000) + sidecar.write_bytes(b"y" * 3000) + custom.write_bytes(b"z" * 4000) + cfg = SimpleNamespace( + gguf_file = str(target), + gguf_mmproj_file = None, + gguf_mtp_file = None, + gguf_dspark_file = None, + gguf_dflash_file = str(sidecar), + gguf_hf_repo = None, + gguf_variant = None, + ) + with ( + patch.object(self.route, "_estimate_gguf_kv_gb", return_value = 0.0), + self._dflash_capable(), + ): + owned = self.route._estimate_gguf_required_gb( + cfg, + speculative_type = "auto", + llama_extra_args = [ + "--spec-type", + "draft-dflash", + "--model-draft", + str(custom), + ], + ) + # 2000 weights + 4000 for the drafter that actually launches, not 9000. + self.assertAlmostEqual(owned, 6000 / (1024**3), places = 9) + def test_remote_weights_stay_in_the_estimate_beside_a_local_extra_args_drafter(self): """A remote repo has no local main weight, so a local --model-draft was the only thing making the local branch fire: it returned ~1.5 GiB and @@ -1759,6 +1799,50 @@ def test_remote_dflash_sizing_drops_a_candidate_too_big_to_be_a_drafter(self): ) self.assertEqual(charged, 1024**3) + def test_remote_dspark_sizing_totals_every_shard_of_a_split_sidecar(self): + """llama-server maps every shard, so pricing the one the ranking picked + halved a two-shard sidecar and let the guard admit a load that evicts + the training run it protects.""" + siblings = [ + SimpleNamespace(rfilename = "dspark/dspark-00001-of-00002.gguf", size = 5 * 1024**3), + SimpleNamespace(rfilename = "dspark/dspark-00002-of-00002.gguf", size = 5 * 1024**3), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + charged = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dspark = True, + ) + self.assertEqual(charged, 10 * 1024**3) + + def test_auto_budgets_dflash_when_the_dspark_set_is_incomplete(self): + """A listing missing a DSpark shard is not a load this can end up on: the + fetch refuses it and falls through to DFlash, which can be the larger of + the two, so granting first refusal on the listing under-charged.""" + siblings = [ + SimpleNamespace(rfilename = "dspark/dspark-00001-of-00002.gguf", size = 1024**3), + SimpleNamespace(rfilename = "dflash-kquant.gguf", size = 4 * 1024**3), + ] + with patch( + "huggingface_hub.model_info", + return_value = SimpleNamespace(siblings = siblings), + ): + charged = self.route._remote_gguf_companion_bytes( + "org/repo", + hf_token = None, + include_mmproj = False, + include_mtp = False, + include_dspark = True, + include_dflash = True, + dspark_first = True, + ) + self.assertEqual(charged, 4 * 1024**3) + # ── Auto charges ONE drafter, the one the promotion leaves resident ── def _auto_companion_bytes(self, siblings): diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index ae9c808cb9b..d6317447e79 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -3437,3 +3437,28 @@ def test_download_companion_records_an_incomplete_listing_as_settled(): mp.undo() assert got is None assert outcome.get("listed") is False + + +def test_download_dflash_reaches_the_complete_family_behind_an_incomplete_one( + monkeypatch, tmp_path +): + """A shard from a half-published set makes _download_companion_gguf answer None, + which ends the loop, so the complete sidecar behind it was never reached.""" + from core.inference.llama_cpp import LlamaCppBackend + + weight = tmp_path / "model-B-Q4_K_M.gguf" + weight.write_bytes(b"x" * 10_000) + monkeypatch.setattr( + LlamaCppBackend, "_remote_root_gguf_sizes", staticmethod(lambda repo, token = None: {}) + ) + picked = _dflash_download_pick( + monkeypatch, + listing = [ + "model-B-Q4_K_M.gguf", + # Ranks first by naming this weight, but the set is missing shard 2. + "dflash-model-B-00001-of-00002.gguf", + "dflash-kquant.gguf", + ], + near_path = str(weight), + ) + assert picked == "dflash-kquant.gguf" diff --git a/studio/backend/utils/models/drafters/__init__.py b/studio/backend/utils/models/drafters/__init__.py index 31eac8ed860..3ca1840aae2 100644 --- a/studio/backend/utils/models/drafters/__init__.py +++ b/studio/backend/utils/models/drafters/__init__.py @@ -14,6 +14,7 @@ _drafter_names_other_weight, _drafter_pairing_stem, _drafter_split_is_complete, + split_listing_is_complete, _drafter_stem_rank, _drafter_total_size, ) @@ -36,6 +37,7 @@ "_drafter_names_other_weight", "_drafter_pairing_stem", "_drafter_split_is_complete", + "split_listing_is_complete", "_drafter_stem_rank", "_drafter_total_size", "detect_dflash_file", diff --git a/studio/backend/utils/models/drafters/common.py b/studio/backend/utils/models/drafters/common.py index 45a5eda436c..2d5b9df8381 100644 --- a/studio/backend/utils/models/drafters/common.py +++ b/studio/backend/utils/models/drafters/common.py @@ -142,3 +142,25 @@ def _drafter_names_other_weight( return any( _drafter_matches_weight(candidate_name, other, kind = kind) for other in other_weight_names ) + + +_LISTED_SHARD_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", re.IGNORECASE) + + +def split_listing_is_complete(names: Iterable[str], name: str) -> bool: + """Whether ``names`` carries every shard of the set ``name`` belongs to. + + The listing counterpart of _drafter_split_is_complete, which needs the files on + disk. A repo mid-upload lists part of a set, and the fetch refuses that, so the + plan and the budget have to agree with it: a half-published sidecar is not one + this load can end up on. True for a single-file name, which encodes no set. + """ + match = _LISTED_SHARD_RE.match(Path(name).name) + if not match: + return True + stem, total = match.group(1), match.group(3) + sibling = re.compile( + r"^" + re.escape(stem) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$", + re.IGNORECASE, + ) + return sum(1 for other in names if sibling.match(Path(other).name)) == int(total) From 471ef9e3783516af58ce8db3537f2603c73dd167 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 11 Aug 2026 11:12:13 +0000 Subject: [PATCH 26/26] Tighten the DFlash review-round comments for PR #8338 Comments and docstrings only, no code change: verified with comment_tools.py check and the prepush gate's comment-only mode. --- studio/backend/core/inference/llama_cpp.py | 39 ++++++++----------- studio/backend/hub/utils/gguf_plan.py | 25 ++++++------ studio/backend/routes/inference.py | 29 ++++++-------- .../backend/utils/models/drafters/budget.py | 7 ++-- .../backend/utils/models/drafters/common.py | 7 ++-- 5 files changed, 45 insertions(+), 62 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 02d9d8ff4d4..0230cfaf065 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4134,11 +4134,9 @@ def _norm(value): and self._spec_fallback_reason == "drafter_not_found" and speculative_type in ("auto", "mtp", "mtp+ngram", "dspark", "dflash") and not (self._spec_drafter_kind == "dspark" and self._dspark_sidecar_absent) - # DFlash asks through _dflash_retry_needed below instead. A permanent - # listing error (gated repo, offline) records no answer at all, so - # _dflash_sidecar_absent stays False and this arm relaunched a healthy - # server on every Apply; the flag is set only for the failures worth - # another attempt. + # DFlash asks through _dflash_retry_needed below, which is set only for + # retryable failures. A permanent listing error records no answer, so + # _dflash_sidecar_absent stays False and this arm relaunched forever. and self._spec_drafter_kind != "dflash" and not spec_owned_by_extra_args ): @@ -8110,9 +8108,8 @@ def _pick_from(names: list[str]) -> Optional[str]: # disagree about what belongs to a set; empty for the single-file case, # which is every mmproj and every published sidecar so far. extra_shards = _gguf_extra_shards(available, target) - # The filename carries the set size, so a listing missing part of it is - # answerable before the download rather than after: half a split companion is - # not a companion, the same rule the snapshot and cache reuse above apply. + # The filename carries the set size, so a short listing is answerable before + # the download: half a split companion is not one, as the reuse paths above say. _shard_total = _SHARD_FULL_RE.match(target) if _shard_total and len(extra_shards) + 1 != int(_shard_total.group(3)): logger.info( @@ -8122,9 +8119,8 @@ def _pick_from(names: list[str]) -> Optional[str]: len(extra_shards) + 1, _shard_total.group(3), ) - # Settled, not retryable: the listing answered and this is what it said, so - # the caller records absence rather than reloading on every Apply hoping the - # same listing says something else. + # Settled, not retryable: the listing answered, so the caller records + # absence rather than reloading on every Apply. if outcome is not None: outcome["listed"] = False return None @@ -8532,10 +8528,9 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: for name in candidates if name.lower().endswith(".gguf") and not _is_root_dflash_drafter_path(name) ] - # Same bound dflash_plan_files applies, for the same reason: dflash- is a - # prefix real weights carry, the header only reads once the bytes are here, - # and a drafter is a few layers of its target so it cannot outweigh it. - # Sizes the listing does not carry leave the candidate in, as before. + # Same bound as dflash_plan_files: dflash- is a prefix real weights carry, + # the header only reads once the bytes are here, and a drafter is a few + # layers of its target. An unlisted size leaves the candidate in, as before. sizes = self._remote_root_gguf_sizes(hf_repo, hf_token) try: target_bytes = (self._get_gguf_size_bytes(near_path) or 0) if near_path else 0 @@ -8547,12 +8542,11 @@ def _pick_dflash(candidates: list[str]) -> Optional[str]: for name in candidates if _is_root_dflash_drafter_path(name) and Path(name).name not in rejected - # Before the ranking, as dflash_plan_files is: a set the listing - # only half carries is refused by the fetch, and returning it - # ends the loop instead of reaching the complete one behind it. + # Before the ranking, as dflash_plan_files is: the fetch refuses a + # half-listed set, and returning it ends the loop. and split_listing_is_complete(candidates, name) - # Whole shard set, not the picked shard: each half of a split - # ordinary weight can sit under the target while the set does not. + # Whole set, not the picked shard: each half of a split weight can + # sit under the target while the set does not. and not ( target_bytes and _drafter_set_bytes(sizes, Path(name).name) >= target_bytes ) @@ -10444,9 +10438,8 @@ def _launch_caps(bin_path): # drafter it was never going to carry. self._spec_fallback_reason = None self._spec_drafter_kind = None - # DFlash discovery runs before the metadata read that classifies this as - # diffusion, so a transient sidecar failure can have set the retry flag - # for a server that carries no drafter at all. + # Discovery runs before the metadata read that says diffusion, so a + # transient sidecar failure can set this for a drafterless server. self._dflash_retry_needed = False with self._lock: if self._cancel_event.is_set(): diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index fb26f9e81cb..ce563ea86b2 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -171,20 +171,17 @@ def dflash_plan_files( ) -> tuple[ExpectedFile, ...]: """Every shard of the DFlash sidecar to plan alongside ``weight_name``, or (). - Whole shard family, not just the ranked file: the loader refuses a companion whose - split set is incomplete, so planning shard 1 alone reports the variant complete and - then loses DFlash on the load. A family the listing only half publishes is dropped - for the same reason, before it can be ranked. + Whole shard family, not the ranked file alone: the loader refuses an incomplete + split set, so planning shard 1 reports the variant complete and then loses DFlash. + A half-published family is dropped for the same reason. Bounded by ``max_bytes``, the variant's own weights. ``dflash-`` is a prefix real - weights carry (Lucebox/Qwen3.6-27B-DFlash-GGUF), and a listing cannot read the - ``general.architecture`` the loader rejects them by. A drafter is a few layers of - its target, so one at least as large as the target cannot be drafting for it, and - an unknown size on either side stays out rather than risk planning a whole model. - - Both rules filter families BEFORE the ranking, so an oversized or half-published - name at the top of the order steps aside for a usable sidecar behind it instead of - taking the plan down with it. + weights carry (Lucebox/Qwen3.6-27B-DFlash-GGUF) and a listing cannot read the + ``general.architecture`` the loader rejects them by, but a drafter is a few layers + of its target and cannot outweigh it. An unknown size stays out. + + Both rules filter BEFORE the ranking, so an oversized or half-published name at the + top steps aside for a usable sidecar behind it. """ from utils.models.drafters import dflash_repo_preference_key, split_listing_is_complete @@ -268,8 +265,8 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: ) # Per variant, unlike mmproj and the MTP drafter: ranked against the weight # being fetched, so a multi-family repo does not hand B the drafter naming A. - # Ranked against the family plan_from_expected_files will KEEP, not the first - # in the listing, or a two-family variant key pairs the discarded one's sidecar. + # Against the family plan_from_expected_files KEEPS, not the listing's first, + # or a two-family variant key pairs the discarded one's sidecar. kept_main = _one_shard_family(main_expected) target_weight_name = ( min(file.path for file in kept_main).rsplit("/", 1)[-1] if kept_main else None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b2ca7596a66..ca84127210c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5477,12 +5477,10 @@ def _remote_gguf_companion_bytes( # counting it here would price a file the load cannot fetch. if include_dflash and _is_root_dflash_drafter_path(name): dflash_sizes[name] = size - # Same preference order the download uses, so the budget sizes the file the - # launch will actually fetch, and by whole shard SET: llama-server maps every - # shard, so pricing the one the ranking picked halved a two-shard sidecar and - # let the guard admit a load that evicts the training run it protects. - # Incomplete sets are dropped rather than priced, because the fetch now refuses - # them: a listing missing a shard is not a sidecar this load can end up on. + # The download's preference order, by whole shard SET: llama-server maps every + # shard, so pricing the picked one halved a two-shard sidecar and let the guard + # admit a load that evicts the training run. Incomplete sets are dropped, not + # priced, since the fetch refuses them. _dspark_sizes = dict(dspark_candidates) dspark_families = [ ( @@ -5746,9 +5744,8 @@ def _estimate_gguf_required_gb( # reason. The sidecar is ~1.5 GiB rather than ~11 GB, but a guard that # protects a running training job still has to charge for it. # Extra args owning --spec-type end _build_speculative_flags before any mode - # branch runs, so neither the forced mode nor the Auto promotion reaches the - # sidecar and charging it refuses a load for bytes nothing will open. Extra - # args asking for draft-dflash themselves are the one case that still pays. + # branch, so neither forced nor Auto reaches the sidecar and charging it refuses + # a load for nothing. Extras asking for draft-dflash themselves still pay. _extra_args_own_spec = _extra_args_set_spec_type(llama_extra_args) _forced_dflash = bool( _extra_args_requests_dflash(llama_extra_args, env = {}) @@ -5801,11 +5798,10 @@ def _same_file_key(p: str) -> str: if dspark_requested: _sized_attrs.append("gguf_dspark_file") elif dflash_requested: - # Only when the extras own --spec-type: _build_speculative_flags then - # returns before discovery's sidecar is ever emitted, so llama-server - # opens the extras' --model-draft alone and charging both billed two - # drafters for one. Without --spec-type Studio still emits its own, and - # which of the two lands is genuinely unknown, so both stay charged. + # Only when extras own --spec-type: _build_speculative_flags then + # returns before discovery's sidecar is emitted, so llama-server opens + # theirs alone. Without it Studio emits its own too and which lands is + # unknown, so both stay charged. _manual_draft = ( _extra_args_mtp_draft_path(llama_extra_args, env = {}) if _extra_args_own_spec @@ -5889,9 +5885,8 @@ def _same_file_key(p: str) -> str: ), include_dspark = (_dspark_capable and (_auto_dspark or dspark_requested)), include_dflash = (_dflash_capable and (_auto_dflash or dflash_requested)), - # The size the DFlash bound measures candidates against, so the guard - # stops charging for weights the fetch would refuse as too big to be - # a drafter. + # What the DFlash bound measures candidates against, so the guard stops + # charging for weights the fetch refuses as too big to be a drafter. weight_bytes = int(main_bytes or 0), # ... except where the listing settles it. Auto launches exactly # one drafter, in a fixed order, so once the listing says which diff --git a/studio/backend/utils/models/drafters/budget.py b/studio/backend/utils/models/drafters/budget.py index d8065b32f31..f89d4e4f53e 100644 --- a/studio/backend/utils/models/drafters/budget.py +++ b/studio/backend/utils/models/drafters/budget.py @@ -35,10 +35,9 @@ def dflash_budget_bytes( halve a two-shard sidecar, and under-estimating is the direction that waves a load through and then exhausts VRAM. - ``target_bytes`` drops the candidates the fetch itself now refuses: a drafter is - a few layers of its target, so a set at least that large is an ordinary weight - wearing the prefix and is never made resident. Zero means unknown, which keeps - every candidate, as before. + ``target_bytes`` drops what the fetch itself refuses: a drafter is a few layers of + its target, so a set at least that large is an ordinary weight wearing the prefix. + Zero means unknown and keeps every candidate. """ totals = ( size + sum(sizes.get(shard, 0) for shard in extra_shards(sizes, name)) diff --git a/studio/backend/utils/models/drafters/common.py b/studio/backend/utils/models/drafters/common.py index 2d5b9df8381..1502448aaaa 100644 --- a/studio/backend/utils/models/drafters/common.py +++ b/studio/backend/utils/models/drafters/common.py @@ -150,10 +150,9 @@ def _drafter_names_other_weight( def split_listing_is_complete(names: Iterable[str], name: str) -> bool: """Whether ``names`` carries every shard of the set ``name`` belongs to. - The listing counterpart of _drafter_split_is_complete, which needs the files on - disk. A repo mid-upload lists part of a set, and the fetch refuses that, so the - plan and the budget have to agree with it: a half-published sidecar is not one - this load can end up on. True for a single-file name, which encodes no set. + The listing counterpart of _drafter_split_is_complete, which needs files on disk. + A repo mid-upload lists part of a set and the fetch refuses that, so the plan and + the budget must agree. True for a single-file name, which encodes no set. """ match = _LISTED_SHARD_RE.match(Path(name).name) if not match: