From 4b251e643aaf41f278c18b55ba6b4fc8f338af2c Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Wed, 12 Aug 2026 20:12:40 -0700 Subject: [PATCH] feat(discover): pair a sibling drafter into its model as draft_gguf on init, sync-models, and pull --- CHANGELOG.md | 7 ++ docs/cli.md | 5 ++ docs/server-config.md | 7 +- gmlx/arch_table.py | 29 +++++++ gmlx/cli.py | 4 +- gmlx/discovery.py | 153 ++++++++++++++++++++++++++++++----- gmlx/mtp_load.py | 8 +- gmlx/server.py | 59 ++++++++++---- tests/test_discovery.py | 175 +++++++++++++++++++++++++++++++++++++++- tests/test_server.py | 39 +++++++++ 10 files changed, 446 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d312f7a..c0da328 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Changed + +- `init`, `sync-models`, and `pull` now pair a sibling drafter (gemma4 + assistant, DSpark, DFlash) into the model it serves as that model's + `draft_gguf`, which turns speculative decoding on. `sync-models` also adds + the key to a model the config already carries. Before, you added it by hand. + ## [0.3.1] - 2026-08-12 ### Added diff --git a/docs/cli.md b/docs/cli.md index c91a28c..5bfcad4 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -738,6 +738,11 @@ the first existing default config unless `--config` is given. Scanning recurses by default, because `pull` nests downloads under `/__/`. +A sibling drafter GGUF pairs into the model it serves as that model's +`draft_gguf`, reported as `update:`. This works on entries the config already +carries; an entry with its own `draft_gguf` keeps it. See +[`discover`](server-config.md#discover) for the rules. + With `--from-hf-cache` (or a config already carrying `server.hf_cache: true`) it also reconciles cache-resident GGUFs, adding new `hf:` entries and dropping ones that are no longer cached, and flips `server.hf_cache` on for them. diff --git a/docs/server-config.md b/docs/server-config.md index 71ac8a3..5822c2f 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -721,8 +721,11 @@ its target id (and profile, if any) must exist. Validated at load. Opt-in header-only directory scan (architecture + `nextn_predict_layers` only, zero tensor I/O). Native-head MTP models auto-enable speculative. -Sibling `mmproj*.gguf` pairs into the model it best matches. Assistant -drafters are reported but only wired when a model names one via `draft_gguf`. +Sibling `mmproj*.gguf` pairs into the model it best matches. A sibling +assistant drafter (gemma4 assistant, DSpark, DFlash) pairs in as that model's +`draft_gguf`, which turns speculative on; the arch, the hidden size, and the +filename must all agree. A streamed model gets no drafter, and +`speculative: false` stops the pairing. ```yaml discover: diff --git a/gmlx/arch_table.py b/gmlx/arch_table.py index 9e776fb..877240f 100644 --- a/gmlx/arch_table.py +++ b/gmlx/arch_table.py @@ -132,12 +132,41 @@ class UnsupportedArchError(Exception): }) +# model_types whose MTP drafter ships as a companion GGUF, and the arches that +# companion can carry, best container first. The loader reads a row to find the +# sidecar of a target; discovery reads it backwards to pair a drafter it found. +# Extend it together with the loader dispatch. +MTP_DRAFTER_ARCHES = { + # `dflash` is llama.cpp's DSpark container, `deepseek4_mtp_support` the + # legacy nextn one. + "deepseek_v4": ("deepseek4-dspark", "dflash", "deepseek4_mtp_support"), + "muse_glimmer": ("dflash",), + "gemma4_text": ("gemma4_assistant", "gemma4-assistant", "gemma4_mtp"), +} + + def mtp_wired(gguf_arch: str | None) -> bool: """True iff a native-head MTP GGUF of this arch has a wired target class.""" model_type = config_synth.GGUF_ARCH_TO_MODEL_TYPE.get(gguf_arch or "") return model_type in MTP_WIRED_MODEL_TYPES +def drafter_arches(model_type: str) -> tuple: + """The companion-GGUF arches that can draft for ``model_type``, best + container first. ``()`` if the type has no companion drafter.""" + return MTP_DRAFTER_ARCHES.get(model_type, ()) + + +def drafter_serves(drafter_arch: str | None, gguf_arch: str | None) -> bool | None: + """Whether a ``drafter_arch`` drafter can serve a ``gguf_arch`` target. + ``None`` when no row lists ``drafter_arch``: the table has no opinion.""" + rows = [t for t, arches in MTP_DRAFTER_ARCHES.items() + if drafter_arch in arches] + if not rows: + return None + return config_synth.GGUF_ARCH_TO_MODEL_TYPE.get(gguf_arch or "") in rows + + def has_synth(gguf_arch: str) -> bool: """True iff ``config_synth`` produces a complete config for this arch.""" return gguf_arch in config_synth.supported_arches() diff --git a/gmlx/cli.py b/gmlx/cli.py index a55c92e..d02d693 100644 --- a/gmlx/cli.py +++ b/gmlx/cli.py @@ -699,12 +699,14 @@ def _deepseek4_mtp_companion(gguf_path: str) -> str | None: it (auto enable; the loader re-resolves the same path when draft_gguf_path is not given). Header-cache peeks only.""" try: + from . import arch_table from .discovery import find_mtp_companion, header_meta meta = header_meta(gguf_path) if not meta or meta.get("arch") != "deepseek4": return None - return find_mtp_companion(gguf_path) + return find_mtp_companion(gguf_path, + arch_table.drafter_arches("deepseek_v4")) except Exception: return None diff --git a/gmlx/discovery.py b/gmlx/discovery.py index 41b8240..1403503 100644 --- a/gmlx/discovery.py +++ b/gmlx/discovery.py @@ -11,8 +11,9 @@ - ``general.architecture == "clip"`` (or a ``mmproj*`` filename) -> **mmproj**: a VLM companion, not a standalone model. - arch is an assistant shape (``gemma4_assistant`` / ``gemma4-assistant`` / - ``gemma4_mtp``, or carries a target-backbone field) -> **drafter**: only paired - when a model explicitly names it via ``draft_gguf``; never standalone. + ``gemma4_mtp``, or carries a target-backbone field) -> **drafter**: a + speculative-decoding companion, never a standalone model. A drafter pairs + into a sibling model as its ``draft_gguf`` (see :func:`_drafter_targets`). - ``.nextn_predict_layers > 0`` -> a **native-head MTP** model (the drafter lives inside the target GGUF; ``speculative: auto`` enables it). - otherwise -> a plain text **model**. @@ -94,6 +95,8 @@ class ClassifiedGguf: loadable: bool # arch builds a model with no hf override (model kind) moe: bool = False # routed experts present (model kind only) name: str | None = None # general.name, for family refinement (model kind) + # hidden size; on a drafter, its target's (see `_drafter_hidden_size`) + n_embd: int | None = None # Classification @@ -115,6 +118,23 @@ def _looks_like_drafter(meta, arch: str | None) -> bool: return False +def _hidden_size(meta, arch: str | None) -> int | None: + """The hidden size that ``arch`` declares, or ``None`` if it declares none.""" + return read_int(meta, f"{arch}.embedding_length") if arch else None + + +def _drafter_hidden_size(meta, arch: str | None) -> int | None: + """The hidden size of the target model for this drafter: its backbone + field, else its own size (a gemma4 assistant declares both, and they + differ; a DSpark sidecar declares neither).""" + if arch: + for suf in _BACKBONE_FIELDS: + v = read_int(meta, f"{arch}.{suf}") + if v: + return v + return _hidden_size(meta, arch) + + def _embedding_kind(meta, basename: str, arch: str | None) -> str | None: """``"reranker"`` / ``"embedding"`` if this GGUF is an encoder-style retrieval model rather than a generative chat model, else ``None``. Signals (any one): @@ -148,7 +168,8 @@ def _classify_meta(meta, *, basename: str, path: str) -> ClassifiedGguf: or read_string(meta, "adapter.type") is not None): return ClassifiedGguf(ap, "adapter", arch, False, quant, False) if _looks_like_drafter(meta, arch): - return ClassifiedGguf(ap, "drafter", arch, False, quant, False) + return ClassifiedGguf(ap, "drafter", arch, False, quant, False, + n_embd=_drafter_hidden_size(meta, arch)) emb = _embedding_kind(meta, basename, arch) if emb: return ClassifiedGguf(ap, emb, arch, False, quant, False) @@ -159,7 +180,8 @@ def _classify_meta(meta, *, basename: str, path: str) -> ClassifiedGguf: loadable = bool(arch) and arch in supported_arches() return ClassifiedGguf(ap, "model", arch, mtp, quant, loadable, moe=bool(experts and experts > 0), - name=read_string(meta, "general.name")) + name=read_string(meta, "general.name"), + n_embd=_hidden_size(meta, arch)) def classify_gguf(path: str) -> ClassifiedGguf | None: @@ -314,14 +336,19 @@ def header_sampling(path) -> dict: def find_mtp_companion( path: str, - drafter_arch: str | tuple = ("deepseek4-dspark", "dflash", - "deepseek4_mtp_support"), + drafter_arch: str | tuple | None = None, ) -> str | None: """Path of an MTP drafter GGUF (arch in ``drafter_arch``) sitting in the - same directory as ``path``, or ``None``. Header-only peeks through + same directory as ``path``, or ``None``. ``drafter_arch`` defaults to every + arch in :data:`arch_table.MTP_DRAFTER_ARCHES`; a caller that knows the + target model type passes that row instead (see + :func:`arch_table.drafter_arches`). Header-only peeks through :func:`header_meta`'s stat-validated cache, so a directory scan costs one stat per already-seen sibling. Earlier arches in the tuple win over later ones (dspark over legacy nextn); within an arch, lexically first wins.""" + if drafter_arch is None: + drafter_arch = tuple(dict.fromkeys( + a for row in _arch_table.MTP_DRAFTER_ARCHES.values() for a in row)) arches = (drafter_arch,) if isinstance(drafter_arch, str) else tuple(drafter_arch) ap = os.path.abspath(os.path.expanduser(path)) parent = os.path.dirname(ap) @@ -463,24 +490,34 @@ def scan_dirs( *, known_ids=frozenset(), known_paths=frozenset(), + known_models=None, progress=False, stats=None, ) -> list[ModelCfg]: """Discover servable models from ``specs`` (each a :class:`config.DiscoverSpec`). A spec with ``dir=None`` scans ``model_dirs``. Native-head MTP models get - ``speculative`` per the spec (``auto``/``True`` -> on for MTP; ``False`` -> off); - assistant drafters are reported but never auto-wired (they need an explicit - ``draft_gguf``). Sibling mmproj files pair into the model they best match when + ``speculative`` per the spec (``auto``/``True`` -> on for MTP; ``False`` -> off). + A sibling assistant drafter pairs into the models it can serve as their + ``draft_gguf`` (see :func:`_drafter_targets`), which also turns + ``speculative`` on. Sibling mmproj files pair into the model they best match when ``pair_mmproj``. Every id carries its quant codec (see :func:`_assign_ids`). ``known_ids`` / ``known_paths`` (from configured ``models:``) are skipped/deduped against, as are paths an earlier spec/root in this same call already emitted. + ``known_models`` maps a configured model's resolved path to its + :class:`config.ModelCfg`. A drafter pairs into one of those too; the model + stays out of the return list, and the caller writes the key from + ``stats["draft_pairs"]``. ``progress`` streams per-file scan feedback to stderr (used by ``init``). ``stats``, if given, is a dict that receives ``skipped``: the count of .gguf files seen but unreadable as GGUF (so callers can distinguish an - empty dir from a dir of truncated downloads).""" + empty dir from a dir of truncated downloads); and ``draft_pairs``: + ``{configured id: drafter path}``.""" known_paths = {os.path.abspath(os.path.expanduser(p)) for p in known_paths} + configured = {os.path.abspath(os.path.expanduser(p)): mc + for p, mc in (known_models or {}).items()} + draft_pairs: dict[str, str] = {} used_ids = set(known_ids) out: list[ModelCfg] = [] skipped = 0 @@ -502,9 +539,11 @@ def scan_dirs( paths.append(p) classified = [c for c in _classify_each(paths, progress=progress) if c] skipped += len(paths) - len(classified) - _emit_dir(classified, spec, used_ids, out) + _emit_dir(classified, spec, used_ids, out, + configured=configured, draft_pairs=draft_pairs) if stats is not None: stats["skipped"] = skipped + stats["draft_pairs"] = draft_pairs return out @@ -683,8 +722,15 @@ def _fit_in_memory(mc: ModelCfg, c: ClassifiedGguf) -> None: f"on its entry: {c.path}", file=sys.stderr) -def _emit_dir(classified, spec, used_ids, out): - """Build ModelCfgs for one scan, pairing mmproj/draft per directory.""" +def _emit_dir(classified, spec, used_ids, out, *, configured=None, + draft_pairs=None): + """Build ModelCfgs for one scan, pairing mmproj/draft per directory. + + A drafter can also pair into a ``configured`` model (resolved path -> + :class:`config.ModelCfg`); ``draft_pairs`` collects those as + ``{id: drafter path}``.""" + configured = configured or {} + draft_pairs = {} if draft_pairs is None else draft_pairs by_dir: dict[str, list[ClassifiedGguf]] = {} for c in classified: by_dir.setdefault(os.path.dirname(c.path), []).append(c) @@ -701,9 +747,6 @@ def _emit_dir(classified, spec, used_ids, out): if c.kind == "model" and not c.loadable: print(f"[discover] skip (unsupported arch {c.arch!r}): {c.path}", file=sys.stderr) - if c.kind == "drafter": - print(f"[discover] assistant drafter (configure via " - f"draft_gguf: on its model): {c.path}", file=sys.stderr) if c.kind in ("embedding", "reranker"): wire = ("server.embeddings:" if c.kind == "embedding" else "the /v1/rerank endpoint") @@ -737,6 +780,27 @@ def _emit_dir(classified, spec, used_ids, out): if tgt is not None and tgt.mmproj is None: tgt.mmproj = mm.path + drafters = [c for c in group if c.kind == "drafter"] + # The usual case: the target is configured already, so it is not in + # `made` and only this lookup can pair it. + old = _configured_in_dir(_dir, configured) if drafters else [] + old_ids = {mc.id for mc, _c in old} + for dr in drafters: + targets = ([] if spec.speculative is False + else _drafter_targets(dr, made + old)) + for mc in targets: + mc.draft_gguf = dr.path + mc.speculative = True + if mc.id in old_ids: + draft_pairs[mc.id] = dr.path + if targets: + ids = ", ".join(mc.id for mc in targets) + print(f"[discover] drafter paired (speculative on for {ids}): " + f"{dr.path}", file=sys.stderr) + else: + print(f"[discover] assistant drafter (configure via " + f"draft_gguf: on its model): {dr.path}", file=sys.stderr) + def _assign_ids(models, used_ids: set) -> dict[str, str]: """Map each model's path to a unique friendly id, quant tag always included. @@ -815,6 +879,49 @@ def _best_mmproj_target(mm: ClassifiedGguf, made): return None +def _configured_in_dir(dirpath: str, configured) -> list: + """``(ModelCfg, ClassifiedGguf)`` for each configured model in ``dirpath``. + + The scan skips these files, so only this lookup can offer them to a + drafter beside them. One cached header read each.""" + out = [] + for ap, mc in configured.items(): + if os.path.dirname(ap) != dirpath: + continue + c = classify_gguf(ap) + if c is not None and c.kind == "model" and c.loadable: + out.append((mc, c)) + return out + + +def _drafter_targets(dr: ClassifiedGguf, made) -> list: + """The models the sibling drafter ``dr`` can serve; empty leaves it unpaired. + + A candidate must have a wired MTP class, a resident placement (MTP needs a + resident base), an arch that :data:`arch_table.MTP_DRAFTER_ARCHES` lists + for this drafter, and an equal hidden size where both files declare one. + The size is the only test that separates the two ``dflash`` drafters, the + DSpark sidecar and the Muse Glimmer one. + + The filename then selects among the candidates, on the + :func:`_best_mmproj_target` rule; a lone candidate pairs without it. Two + quants of one target both pair on purpose.""" + cands = [(mc, c) for mc, c in made + if _arch_table.mtp_wired(c.arch) and not mc.stream + and mc.draft_gguf is None + and _arch_table.drafter_serves(dr.arch, c.arch) is not False + and not (dr.n_embd and c.n_embd and dr.n_embd != c.n_embd)] + if not cands: + return [] + core, _ = derive_id(os.path.basename(dr.path)) # markers (dflash/draft) stripped + hits = [mc for mc, _c in cands + if core and core != "model" + and _common_prefix_len(core, mc.id) >= max(6, int(0.7 * len(core)))] + if hits: + return hits + return [cands[0][0]] if len(cands) == 1 else [] + + def _common_prefix_len(a: str, b: str) -> int: n = 0 for x, y in zip(a, b): @@ -845,7 +952,11 @@ def model_to_entry(mc: ModelCfg, model_dirs) -> dict: entry: dict = {"path": _rel(mc.path, model_dirs)} if mc.mmproj: entry["mmproj"] = _rel(mc.mmproj, model_dirs) - if mc.speculative: + if mc.draft_gguf: + entry["draft_gguf"] = _rel(mc.draft_gguf, model_dirs) + # A `draft_gguf` turns speculative decoding on at config load, thus a + # second `speculative: true` key adds nothing. + if mc.speculative and not mc.draft_gguf: entry["speculative"] = True if mc.stream: entry["stream"] = mc.stream @@ -1208,7 +1319,11 @@ def _scaffold_models_block(models, dirs) -> list[str]: lines.append(f" profile: {mc.profile}") if mc.mmproj: lines.append(f" mmproj: {_rel(mc.mmproj, dirs)}") - if mc.speculative: + if mc.draft_gguf: + lines.append(" # a companion drafter sits next to the model; " + "this key turns speculative decoding on") + lines.append(f" draft_gguf: {_rel(mc.draft_gguf, dirs)}") + elif mc.speculative: lines.append(" # native-head MTP (drafter inside the " "target GGUF)") lines.append(" speculative: true") diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index f501db8..2e34280 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -1191,9 +1191,11 @@ def load_mtp_model( # deepseek4_mtp_support), never as in-GGUF nextn tensors; the # native-head extraction below is qwen-shaped and cannot serve it, # even though the V4 metadata advertises mtp_num_hidden_layers. + from . import arch_table from .discovery import find_mtp_companion - draft_gguf_path = find_mtp_companion(gguf_path) + draft_gguf_path = find_mtp_companion( + gguf_path, arch_table.drafter_arches("deepseek_v4")) if draft_gguf_path is None: raise ValueError( "deepseek_v4 MTP needs its companion drafter GGUF (arch " @@ -1206,9 +1208,11 @@ def load_mtp_model( if not assistant and config_dict.get("model_type") == "muse_glimmer": # Muse Glimmer's drafter is likewise a companion GGUF (arch dflash), # never an in-file nextn block. + from . import arch_table from .discovery import find_mtp_companion - draft_gguf_path = find_mtp_companion(gguf_path, ("dflash",)) + draft_gguf_path = find_mtp_companion( + gguf_path, arch_table.drafter_arches("muse_glimmer")) if draft_gguf_path is None: raise ValueError( "muse_glimmer MTP needs its companion DFlash drafter GGUF " diff --git a/gmlx/server.py b/gmlx/server.py index 9a91984..aab9a42 100644 --- a/gmlx/server.py +++ b/gmlx/server.py @@ -464,6 +464,7 @@ def _cmd_sync(argv: list, prog: str = "gmlx sync-models") -> int: missing_roots = [d for d in cfg.model_dirs if not os.path.isdir(os.path.expanduser(os.path.expandvars(d)))] kept, removed, unverified, known_paths = [], [], [], set() + known_models = {} # resolved path -> ModelCfg for mid, mc in cfg.models.items(): is_hf = str(mc.path).startswith("hf:") try: @@ -489,6 +490,7 @@ def _cmd_sync(argv: list, prog: str = "gmlx sync-models") -> int: kept.append(mid) if rp: known_paths.add(rp) + known_models[rp] = mc if missing_roots: print(f"warning: model_dirs root(s) not on disk right now: " f"{', '.join(missing_roots)} - their entries are kept unverified", @@ -498,12 +500,16 @@ def _cmd_sync(argv: list, prog: str = "gmlx sync-models") -> int: "are kept unverified", file=sys.stderr) # Discover, skipping anything already configured (by id or by resolved path). - discovered = [] + # `draft_pairs`: sibling drafters that pair into an entry already configured. + discovered, draft_pairs = [], {} if dirs: specs = [DiscoverSpec(dir=d, recursive=recursive) for d in dirs] + stats = {} discovered += discovery.scan_dirs( specs, dirs, - known_ids=set(cfg.models), known_paths=known_paths, progress=True) + known_ids=set(cfg.models), known_paths=known_paths, + known_models=known_models, progress=True, stats=stats) + draft_pairs = stats.get("draft_pairs") or {} if scan_cache: known_refs = {mc.path for mc in cfg.models.values() if str(mc.path).startswith("hf:")} @@ -519,7 +525,9 @@ def _cmd_sync(argv: list, prog: str = "gmlx sync-models") -> int: print(f" remove: {mid} ({cfg.models[mid].path} - gone)") for m in discovered: print(f" add: {m.id} ({discovery._rel(m.path, dirs)})") - if not removed and not discovered: + for mid, dpath in draft_pairs.items(): + print(f" update: {mid} (draft_gguf: {discovery._rel(dpath, dirs)})") + if not removed and not discovered and not draft_pairs: print(" (already in sync)") return 0 if a.dry_run: @@ -528,8 +536,11 @@ def _cmd_sync(argv: list, prog: str = "gmlx sync-models") -> int: new_roots = ([d for d in dirs if d not in cfg.model_dirs] if a.models_dir else []) - _apply_sync(path, removed, discovered, dirs, new_roots=new_roots) - print(f"\nupdated {path} (+{len(discovered)} / -{len(removed)})") + _apply_sync(path, removed, discovered, dirs, new_roots=new_roots, + draft_pairs=draft_pairs) + changed = (f"\nupdated {path} (+{len(discovered)} / -{len(removed)}" + + (f" / ~{len(draft_pairs)}" if draft_pairs else "") + ")") + print(changed) _reload_running(path, skip=a.no_reload) return 0 @@ -544,12 +555,15 @@ def _hf_cache_readable() -> bool: return False -def _apply_sync(path, removed, discovered, dirs, new_roots=()) -> None: +def _apply_sync(path, removed, discovered, dirs, new_roots=(), + draft_pairs=None) -> None: """Rewrite ``path``'s ``models:`` block in place, preserving comments and hand-edits (ruamel round-trip): delete ``removed`` ids, splice in the newly ``discovered`` models. Untouched entries keep their exact formatting. ``new_roots`` are --models-dir override roots absent from - ``server.model_dirs``; they're appended so the new entries resolve.""" + ``server.model_dirs``; they're appended so the new entries resolve. + ``draft_pairs`` (configured id -> drafter path) adds a ``draft_gguf`` key + to that entry; an entry that already has one keeps its value.""" from ruamel.yaml.comments import CommentedMap def mutate(doc): @@ -587,6 +601,10 @@ def mutate(doc): if note: models.yaml_set_comment_before_after_key( mc.id, before=note, indent=2) + for mid, dpath in (draft_pairs or {}).items(): + ent = models.get(mid) + if isinstance(ent, dict) and not ent.get("draft_gguf"): + ent["draft_gguf"] = discovery._rel(dpath, dirs) # Cache-resident entries need server.hf_cache to resolve from the cache. if any(str(mc.path).startswith("hf:") for mc in discovered): srv = doc.get("server") @@ -604,7 +622,8 @@ def register_downloads(paths: list, config_path=None) -> None: ``model_dirs`` root (an explicit ``--to`` elsewhere - the server could never discover it), or when the file is already configured. Otherwise the same machinery as sync-models end to end: id derivation, mmproj pairing, - speculative detection, comment-preserving splice, and a SIGHUP so a + drafter pairing, speculative detection, comment-preserving splice, and a + SIGHUP so a running server serves the new entries immediately. Best-effort by contract: the download already succeeded, so a registration problem warns and returns instead of failing ``pull``.""" @@ -625,24 +644,36 @@ def register_downloads(paths: list, config_path=None) -> None: wanted.add(ap) if not wanted: return - known_paths = set() + known_paths, known_models = set(), {} for mc in cfg.models.values(): try: - known_paths.add(resolve_path(mc.path, cfg.model_dirs)) + rp = resolve_path(mc.path, cfg.model_dirs) except ConfigError: - pass + continue + known_paths.add(rp) + if rp: + known_models[rp] = mc # Scan the parent dirs (mmproj pairing needs the siblings), then keep # only models whose file is one we just downloaded - a neighbouring # file the user left unregistered stays unregistered. parents = sorted({os.path.dirname(p) for p in wanted}) specs = [DiscoverSpec(dir=d, recursive=False) for d in parents] + stats = {} found = discovery.scan_dirs(specs, cfg.model_dirs, known_ids=set(cfg.models), - known_paths=known_paths) + known_paths=known_paths, + known_models=known_models, stats=stats) newly = [m for m in found if os.path.abspath(m.path) in wanted] - if not newly: + # Only a drafter from this download may change a configured entry: a + # file the user left beside a model stays unregistered. + pairs = {mid: dp for mid, dp in (stats.get("draft_pairs") or {}).items() + if os.path.abspath(dp) in wanted} + if not newly and not pairs: return # already configured (a re-pull) - quiet - _apply_sync(path, [], newly, cfg.model_dirs) + _apply_sync(path, [], newly, cfg.model_dirs, draft_pairs=pairs) + for mid, dp in pairs.items(): + print(f"registered {os.path.basename(dp)} as the drafter of " + f"{mid} in {path}") for m in newly: extras = [w for w, on in (("vlm", m.mmproj), ("mtp", m.speculative)) if on] diff --git a/tests/test_discovery.py b/tests/test_discovery.py index f1ba6f9..231c910 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -201,7 +201,12 @@ def _f(path): if name.startswith("mmproj"): meta = {"general.architecture": "clip"} elif "assistant" in name: - meta = {"general.architecture": "gemma4_assistant"} + # the drafter carries the hidden size of the target it drafts for + meta = {"general.architecture": "gemma4_assistant", + "gemma4_assistant.embedding_length_out": 5376} + elif "gemma" in name: + meta = {"general.architecture": "gemma4", + "gemma4.embedding_length": 5376} elif "nomtp" in name: meta = {"general.architecture": "qwen3"} elif "bad" in name: @@ -256,7 +261,147 @@ def test_scan_drafter_not_standalone(tmp_path, fake_classify): models = _scan(root) assert len(models) == 1 # the assistant is not a model assert models[0].id == "gemma-4-31b-it-q6" - assert models[0].draft_gguf is None # never auto-wired + # the assistant pairs into the model it serves, not into an entry of its own + assert models[0].draft_gguf.endswith("gemma-4-31B-it-assistant.Q8_0.gguf") + assert models[0].speculative is True + + +# drafter pairing - a sibling drafter becomes the `draft_gguf` of its target +def _drafter_classify(monkeypatch, *, target_arch="muse-glimmer", + target_embd=6656, draft_embd=6656): + """Classify a `dflash*` file as a drafter and every other file as a + target of ``target_arch``, with the hidden size each one declares.""" + def _f(path): + name = os.path.basename(path) + if name.lower().startswith("dflash"): + meta = {"general.architecture": "dflash", + "dflash.embedding_length": draft_embd} + else: + meta = {"general.architecture": target_arch, + f"{target_arch}.embedding_length": target_embd} + return disc._classify_meta(meta, basename=name, path=path) + + monkeypatch.setattr(disc, "classify_gguf", _f) + + +def test_scan_pairs_sibling_dflash_drafter(tmp_path, monkeypatch): + _drafter_classify(monkeypatch) + root = _write(tmp_path, "Muse-Glimmer-30B-KQuant-17GB-Q4_K_M.gguf", + "dflash-Muse-Glimmer-30B-Q4_K_M.gguf") + models = _scan(root) + assert len(models) == 1 + assert models[0].draft_gguf.endswith("dflash-Muse-Glimmer-30B-Q4_K_M.gguf") + assert models[0].speculative is True + + +def test_drafter_pairs_into_every_quant_of_its_target(tmp_path, monkeypatch): + # Two quants of one target share the drafter: each entry can serve. + _drafter_classify(monkeypatch) + root = _write(tmp_path, "Muse-Glimmer-30B-Q4_K_M.gguf", + "Muse-Glimmer-30B-Q6_K.gguf", + "dflash-Muse-Glimmer-30B-Q4_K_M.gguf") + models = _scan(root) + assert len(models) == 2 + assert all(m.draft_gguf and m.speculative for m in models) + + +def test_drafter_unpaired_on_hidden_size_mismatch(tmp_path, monkeypatch, capsys): + # llama.cpp writes the DeepSeek DSpark sidecar under the same `dflash` arch, + # so the arch alone must not pair it with a muse-glimmer target. + _drafter_classify(monkeypatch, target_embd=4096, draft_embd=6656) + root = _write(tmp_path, "Muse-Glimmer-30B-Q4_K_M.gguf", + "dflash-DeepSeek-V4-Q4_K_M.gguf") + models = _scan(root) + assert models[0].draft_gguf is None + assert models[0].speculative is False + assert "assistant drafter (configure via" in capsys.readouterr().err + + +def test_drafter_unpaired_when_speculative_false(tmp_path, monkeypatch): + _drafter_classify(monkeypatch) + root = _write(tmp_path, "Muse-Glimmer-30B-Q4_K_M.gguf", + "dflash-Muse-Glimmer-30B-Q4_K_M.gguf") + models = _scan(root, speculative=False) + assert models[0].draft_gguf is None + assert models[0].speculative is False + + +def test_drafter_unpaired_when_arch_has_no_mtp_class(tmp_path, monkeypatch): + # `llama` builds a model, but the loader has no MTP target class for it, + # so `speculative` would fail at build time. + _drafter_classify(monkeypatch, target_arch="llama", target_embd=6656) + root = _write(tmp_path, "Some-Llama-8B-Q4_K_M.gguf", + "dflash-Some-Llama-8B-Q4_K_M.gguf") + models = _scan(root) + assert models[0].draft_gguf is None + + +def test_generic_drafter_name_unpaired_when_several_candidates( + tmp_path, monkeypatch): + # A filename with no model name in it carries no signal, and two targets + # of equal hidden size are both plausible. + _drafter_classify(monkeypatch) + root = _write(tmp_path, "Muse-Glimmer-30B-Q4_K_M.gguf", + "Other-Muse-Model-Q4_K_M.gguf", "dflash-Q4_K_M.gguf") + models = _scan(root) + assert all(m.draft_gguf is None for m in models) + + +def test_drafter_pairs_into_a_model_already_in_the_config(tmp_path, monkeypatch): + # The scan skips a configured file, so the drafter would have no candidate + # without `known_models`. The pairing goes to stats for the caller to write. + _drafter_classify(monkeypatch) + root = _write(tmp_path, "Muse-Glimmer-30B-Q4_K_M.gguf", + "dflash-Muse-Glimmer-30B-Q4_K_M.gguf") + target = str(root / "Muse-Glimmer-30B-Q4_K_M.gguf") + mc = ModelCfg(id="muse-30b", path=target) + spec = DiscoverSpec(dir=str(root), recursive=False, pair_mmproj=True, + speculative="auto") + stats = {} + models = disc.scan_dirs([spec], [str(root)], known_paths={target}, + known_models={target: mc}, stats=stats) + assert models == [] # the target stays configured + assert stats["draft_pairs"] == { + "muse-30b": str(root / "dflash-Muse-Glimmer-30B-Q4_K_M.gguf")} + assert mc.draft_gguf and mc.speculative is True + + +def test_drafter_without_a_hidden_size_pairs_with_its_lone_target( + tmp_path, monkeypatch): + # A DSpark sidecar declares no hidden size, and its filename shares too + # little with the target quant to pass the name test. The arch table plus + # a single candidate still settle it (llama.cpp's own sidecar layout). + def _f(path): + name = os.path.basename(path) + meta = ({"general.architecture": "deepseek4-dspark"} + if "DSpark" in name else {"general.architecture": "deepseek4"}) + return disc._classify_meta(meta, basename=name, path=path) + + monkeypatch.setattr(disc, "classify_gguf", _f) + root = _write(tmp_path, "DeepSeek-V4-Flash-0731-UD-IQ3_XXS.gguf", + "DeepSeek-V4-Flash-0731-DSpark-MXFP4-Q8_0.gguf") + models = _scan(root) + assert len(models) == 1 + assert models[0].draft_gguf.endswith("DSpark-MXFP4-Q8_0.gguf") + + +def test_drafter_unpaired_across_model_families(tmp_path, fake_classify): + # A gemma4 assistant cannot draft for a qwen target: the arch table says + # which model types each drafter arch serves. + root = _write(tmp_path, "Qwen3.6-27B-Q4_K_S.gguf", + "gemma-4-31B-it-assistant.Q8_0.gguf") + by_id = {m.id: m for m in _scan(root)} + assert by_id["qwen3.6-27b-q4"].draft_gguf is None + + +def test_streamed_model_gets_no_drafter(): + # MTP needs a fully resident base, thus an over-RAM entry stays plain. + mc = ModelCfg(id="m-q4", path="/m/M-Q4_K_M.gguf", stream="experts") + target = disc.ClassifiedGguf("/m/M-Q4_K_M.gguf", "model", "muse-glimmer", + False, "Q4_K_M", True, n_embd=6656) + drafter = disc.ClassifiedGguf("/m/dflash-M-Q4_K_M.gguf", "drafter", "dflash", + False, "Q4_K_M", False, n_embd=6656) + assert disc._drafter_targets(drafter, [(mc, target)]) == [] def test_scan_adapter_not_standalone(tmp_path, monkeypatch, capsys): @@ -383,6 +528,21 @@ def test_scaffold_round_trips_through_build_config(): assert cfg.models["qwen3.6-27b"].path == "qwen3.6-27b/m-Q4_K_S.gguf" +def test_scaffold_writes_a_paired_drafter(): + import yaml + models = [ModelCfg(id="muse-glimmer-30b-q4", + path="/models/muse/m-Q4_K_M.gguf", + draft_gguf="/models/muse/dflash-m-Q4_K_M.gguf", + speculative=True)] + text = disc.scaffold_yaml(models, model_dirs=["/models"]) + cfg = build_config(yaml.safe_load(text)) + mc = cfg.models["muse-glimmer-30b-q4"] + assert mc.draft_gguf == "muse/dflash-m-Q4_K_M.gguf" + assert mc.speculative is False # the drafter key turns it on at load + # no live `speculative` key: the reference block keeps the only mention + assert not [ln for ln in text.splitlines() if ln.strip() == "speculative: true"] + + def test_scaffold_anchors_relative_model_dirs(tmp_path, monkeypatch): # A cwd-relative root would resolve against the SERVER's cwd later (launchd # runs at /), silently serving zero models - the scaffold anchors it. `~` @@ -552,6 +712,17 @@ def test_model_to_entry_includes_mmproj_and_speculative(): } +def test_model_to_entry_draft_gguf_implies_speculative(): + # `draft_gguf` turns speculative decoding on by itself, so the entry keeps + # no second `speculative` key. + mc = ModelCfg(id="g", path="/models/llm-Q6_K.gguf", + draft_gguf="/models/dflash-llm-Q4_K_M.gguf", speculative=True) + assert disc.model_to_entry(mc, ["/models"]) == { + "path": "llm-Q6_K.gguf", + "draft_gguf": "dflash-llm-Q4_K_M.gguf", + } + + def test_model_to_entry_absolute_when_outside_model_dirs(): mc = ModelCfg(id="x", path="/elsewhere/x-Q4_0.gguf") assert disc.model_to_entry(mc, ["/models"]) == {"path": "/elsewhere/x-Q4_0.gguf"} diff --git a/tests/test_server.py b/tests/test_server.py index c6bed3b..6c5f3ea 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -461,6 +461,45 @@ def test_sync_adds_new_and_removes_gone(monkeypatch, tmp_path): assert cfg.models["newbie"].path == "newbie.gguf" # relative to model_dirs +def test_sync_adds_a_drafter_to_an_entry_already_in_the_config(monkeypatch, + tmp_path): + # The user drops a drafter next to a model the config already carries: + # sync adds `draft_gguf` to that entry and leaves the rest of it alone. + from gmlx.config import load_config + cfg_path, lib = _sync_config( + tmp_path, + "models:\n muse:\n path: muse.gguf\n pin: true\n") + (lib / "muse.gguf").write_bytes(b"x") + drafter = str(lib / "dflash-muse.gguf") + + def fake_scan(specs, dirs, **kw): + kw["stats"]["draft_pairs"] = {"muse": drafter} + return [] + + monkeypatch.setattr(srv.discovery, "scan_dirs", fake_scan) + assert srv._cmd_sync(["--config", str(cfg_path)]) == 0 + mc = load_config(cfg_path).models["muse"] + assert mc.draft_gguf == "dflash-muse.gguf" # relative to model_dirs + assert mc.pin is True # the hand-edit survives + + +def test_sync_dry_run_keeps_a_drafter_pairing_unwritten(monkeypatch, tmp_path, + capsys): + from gmlx.config import load_config + cfg_path, lib = _sync_config( + tmp_path, "models:\n muse:\n path: muse.gguf\n") + (lib / "muse.gguf").write_bytes(b"x") + + def fake_scan(specs, dirs, **kw): + kw["stats"]["draft_pairs"] = {"muse": str(lib / "dflash-muse.gguf")} + return [] + + monkeypatch.setattr(srv.discovery, "scan_dirs", fake_scan) + assert srv._cmd_sync(["--config", str(cfg_path), "--dry-run"]) == 0 + assert "update: muse" in capsys.readouterr().out + assert load_config(cfg_path).models["muse"].draft_gguf is None + + def test_sync_never_drops_hf_entries_on_unreadable_cache(monkeypatch, tmp_path, capsys): # An unreadable hf cache means hf: entries CANNOT BE VERIFIED - removing