From 69e814705979356a2d28bb75d1b27bf65d6140fa Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:59:06 -0700 Subject: [PATCH] feat(mtp): accept llama.cpp dflash GGUFs as deepseek4 drafters, pull/validate treat drafter arches as companions --- CHANGELOG.md | 2 + docs/arch-coverage.md | 2 +- docs/cli.md | 2 +- gmlx/arch_table.py | 2 +- gmlx/discovery.py | 16 +++-- gmlx/manage.py | 25 +++++-- gmlx/mtp_load.py | 88 ++++++++++++++++++++++++- tests/test_dflash_normalize.py | 117 +++++++++++++++++++++++++++++++++ tests/test_discovery.py | 30 +++++++++ tests/test_manage.py | 25 +++++++ 10 files changed, 294 insertions(+), 15 deletions(-) create mode 100644 tests/test_dflash_normalize.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c6de72f..73fc3ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - Ctrl-T during a `chat` or `run` reply closes the model's open thinking block early (as if the thinking budget had just run out) so the answer starts now. +- DeepSeek-V4 speculative decoding now accepts llama.cpp dflash draft GGUFs + in addition to DwarfStar's. ### Changed diff --git a/docs/arch-coverage.md b/docs/arch-coverage.md index 65bddc7..65789f3 100644 --- a/docs/arch-coverage.md +++ b/docs/arch-coverage.md @@ -9,7 +9,7 @@ A GGUF loads iff gmlx maps its `general.architecture` to a `model_type` with a m | GGUF arch | model_type | model class | remap | synth | status | notes | |-----------|------------|:-----------:|-------|:-----:|--------|-------| | `deepseek2` | `deepseek_v3` | yes | yes | yes | loadable | DeepSeek-V3/R1 + GLM-4.x MLA conversions (GLM-4.7-Flash); MLA attention (absorbed embed_q/unembed_out via KQuantMultiLinear, native Q8_0), fine-grained sigmoid-gated MoE + shared expert + group routing, leading dense block. V2 (softmax gating) not yet supported | -| `deepseek4` | `deepseek_v4` | yes | yes | yes | loadable | DeepSeek V4 Flash (256x8.4B, dwarfstar 'deepseek4' arch, not a llama.cpp conversion; parity reference is the ds4 engine). MLA-lite attention (low-rank q, single shared 512-dim KV latent K=V, grouped low-rank output proj, per-head fp32 sinks, tail-64 NEOX rope with inverse-rope'd output) in three per-layer variants from compress_ratios: sliding-window(128), +compressed pool, +lightning-indexer top-512 sparse. Manifold-constrained hyper-connections (4-stream 4D hidden, Sinkhorn mixing). Every-layer 256-expert MoE, sqrt-softplus gating + selection-only correction bias, first 3 layers hash-routed (tid2eid), shared expert, clamped SwiGLU. QAT round-trips (fp8-E4M3 KV, Hadamard+fp4-E2M1 indexer) reproduced on-path for logit parity. Model class vendored from mlx-lm PR #1192; separate MTP drafter GGUF (--draft-gguf) for speculative decoding | +| `deepseek4` | `deepseek_v4` | yes | yes | yes | loadable | DeepSeek V4 Flash (256x8.4B, dwarfstar 'deepseek4' arch, not a llama.cpp conversion; parity reference is the ds4 engine). MLA-lite attention (low-rank q, single shared 512-dim KV latent K=V, grouped low-rank output proj, per-head fp32 sinks, tail-64 NEOX rope with inverse-rope'd output) in three per-layer variants from compress_ratios: sliding-window(128), +compressed pool, +lightning-indexer top-512 sparse. Manifold-constrained hyper-connections (4-stream 4D hidden, Sinkhorn mixing). Every-layer 256-expert MoE, sqrt-softplus gating + selection-only correction bias, first 3 layers hash-routed (tid2eid), shared expert, clamped SwiGLU. QAT round-trips (fp8-E4M3 KV, Hadamard+fp4-E2M1 indexer) reproduced on-path for logit parity. Model class vendored from mlx-lm PR #1192; separate MTP drafter GGUF (--draft-gguf, or autodetected next to the target) for speculative decoding - the DSpark sidecar in either container (gmlx deepseek4-dspark or llama.cpp dflash, e.g. the unsloth dspark release), or the legacy deepseek4_mtp_support nextn file | | `diffusion-gemma` | `diffusion_gemma` | yes | yes | yes | loadable | DiffusionGemma 26B-A4B; a non-autoregressive encoder-decoder block-diffusion model on the gemma-4 MoE backbone (each decoder layer runs a dense MLP + routed experts in parallel; encoder and decoder share weights bar a per-layer scalar). The decoder denoises a fixed-length canvas over reverse-diffusion steps with an entropy-bound sampler. Model class + denoising engine live in mlx-vlm (model_type diffusion_gemma); the backbone homes under model.decoder.* with the expert gate_up kept fused, plus the self-conditioning gated MLP. Single text-only GGUF (vision tower dropped at convert) | | `ernie4_5-moe` | `ernie4_5_moe` | yes | yes | yes | loadable | Baidu ERNIE-4.5-MoE (21B-A3B); fine-grained MoE (stacked routed experts -> SwitchGLU) + shared expert behind leading dense layers, softmax or aux-free sigmoid gating. NORM rope but mlx-lm uses traditional=True => Q/K pass through un-permuted (HF-native); the e_score_correction_bias is dropped on load (mlx-lm gates without it) | | `falcon-h1` | `falcon_h1` | yes | yes | yes | loadable | TII Falcon-H1 (0.5B-34B); parallel attention + Mamba2 in every layer (one input_layernorm feeds both, outputs summed), dense gated MLP under feed_forward.*, explicit head_dim from key_length. NEOX rope (no qk-permute). The muP multiplier zoo is folded into the wire weights at convert => synth pins every multiplier neutral; ffn_norm/ssm_a/ssm_d stored with no .weight suffix | diff --git a/docs/cli.md b/docs/cli.md index 71ee4c5..b172d34 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -144,7 +144,7 @@ with a warning); `--no-speculative`/`--no-mtp` forces it off. |------|---------| | `--speculative` / `--mtp` | Force MTP speculative decoding on. Native-head models (qwen3.5/3.6 `nextn`) need no companion; gemma4 needs `--draft-gguf`. Native heads are auto-enabled without this. Use it to force the path when a sampler flag would otherwise defer. | | `--no-speculative` / `--no-mtp` | Disable MTP. Overrides the native-head auto-enable and config `speculative: true`. | -| `--draft-gguf PATH` | Separate assistant-drafter GGUF (gemma4 two-GGUF MTP shape); implies `--speculative` (same as `serve`). | +| `--draft-gguf PATH` | Separate assistant-drafter GGUF (gemma4 two-GGUF MTP shape, or a deepseek4 DSpark/MTP sidecar - gmlx `deepseek4-dspark`, llama.cpp `dflash`, or legacy `deepseek4_mtp_support`); implies `--speculative` (same as `serve`). A sidecar in the target's directory is autodetected without the flag. | | `--draft-block-size N` | Override the MTP draft block size. | Speculative generation takes only `--temp`/`--top-p`/`--top-k`/`--min-p` plus a diff --git a/gmlx/arch_table.py b/gmlx/arch_table.py index d4dd161..adbe582 100644 --- a/gmlx/arch_table.py +++ b/gmlx/arch_table.py @@ -75,7 +75,7 @@ class ArchEntry: "mistral3": ("llama", "llama.cpp 'mistral3' = Ministral-3 / Mistral-Small-3.1, Llama layout"), "nemotron_h_moe": ("nemotron", "NVIDIA Nemotron-H MoE hybrid: Mamba2 SSM layers + sparse attention layers + MoE MLPs; layer_norm_epsilon (not rms_norm_eps)"), "deepseek2": ("deepseek", "DeepSeek-V3/R1 + GLM-4.x MLA conversions (GLM-4.7-Flash); MLA attention (absorbed embed_q/unembed_out via KQuantMultiLinear, native Q8_0), fine-grained sigmoid-gated MoE + shared expert + group routing, leading dense block. V2 (softmax gating) not yet supported"), - "deepseek4": ("deepseek", "DeepSeek V4 Flash (256x8.4B, dwarfstar 'deepseek4' arch, not a llama.cpp conversion; parity reference is the ds4 engine). MLA-lite attention (low-rank q, single shared 512-dim KV latent K=V, grouped low-rank output proj, per-head fp32 sinks, tail-64 NEOX rope with inverse-rope'd output) in three per-layer variants from compress_ratios: sliding-window(128), +compressed pool, +lightning-indexer top-512 sparse. Manifold-constrained hyper-connections (4-stream 4D hidden, Sinkhorn mixing). Every-layer 256-expert MoE, sqrt-softplus gating + selection-only correction bias, first 3 layers hash-routed (tid2eid), shared expert, clamped SwiGLU. QAT round-trips (fp8-E4M3 KV, Hadamard+fp4-E2M1 indexer) reproduced on-path for logit parity. Model class vendored from mlx-lm PR #1192; separate MTP drafter GGUF (--draft-gguf) for speculative decoding"), + "deepseek4": ("deepseek", "DeepSeek V4 Flash (256x8.4B, dwarfstar 'deepseek4' arch, not a llama.cpp conversion; parity reference is the ds4 engine). MLA-lite attention (low-rank q, single shared 512-dim KV latent K=V, grouped low-rank output proj, per-head fp32 sinks, tail-64 NEOX rope with inverse-rope'd output) in three per-layer variants from compress_ratios: sliding-window(128), +compressed pool, +lightning-indexer top-512 sparse. Manifold-constrained hyper-connections (4-stream 4D hidden, Sinkhorn mixing). Every-layer 256-expert MoE, sqrt-softplus gating + selection-only correction bias, first 3 layers hash-routed (tid2eid), shared expert, clamped SwiGLU. QAT round-trips (fp8-E4M3 KV, Hadamard+fp4-E2M1 indexer) reproduced on-path for logit parity. Model class vendored from mlx-lm PR #1192; separate MTP drafter GGUF (--draft-gguf, or autodetected next to the target) for speculative decoding - the DSpark sidecar in either container (gmlx deepseek4-dspark or llama.cpp dflash, e.g. the unsloth dspark release), or the legacy deepseek4_mtp_support nextn file"), "glm-dsa": ("deepseek", "GLM-5.2 (DeepSeek-V3.2): deepseek2 MLA + fine-grained sigmoid-gated MoE plus a per-layer DSA 'lightning indexer' (top-k sparse-attention key selection) + an MTP/nextn layer (dropped on load). Reuses the DEEPSEEK2 remap (indexer.* patterns appended) -> mlx-lm model_type glm_moe_dsa (subclasses deepseek_v32)"), "glm4moe": ("glm", "GLM-4.5 / 4.6 (incl. GLM-4.5-Air); standard MHA (qk-norm, partial rotary) + DeepSeek-V3-style fine-grained sigmoid-gated MoE (shared expert, group routing, leading dense block). NextN/MTP layer dropped on load"), "gpt-oss": ("gpt-oss", "OpenAI gpt-oss 20B/120B; MoE (no dense MLP) with per-head attention sinks, alternating sliding/full attention, YaRN rope, and native MXFP4 experts (packed-repacked in RAM, or zero-copy GGUF wire bytes for streaming/over-RAM - never dequantized). Attn/embed/output Q8_0; router/norms/biases F32"), diff --git a/gmlx/discovery.py b/gmlx/discovery.py index b986d81..f6fee1c 100644 --- a/gmlx/discovery.py +++ b/gmlx/discovery.py @@ -44,7 +44,8 @@ # hidden size it carries as a "backbone" field). The set is the fast path; the # backbone-field probe in `_looks_like_drafter` catches future naming. _DRAFTER_ARCHES = frozenset({"gemma4_assistant", "gemma4-assistant", "gemma4_mtp", - "deepseek4-dspark"}) + "deepseek4-dspark", + "dflash"}) # llama.cpp's arch name for DSpark _BACKBONE_FIELDS = ("backbone_embedding_length", "embedding_length_out", "n_embd_backbone") @@ -95,10 +96,16 @@ class ClassifiedGguf: # Classification -def _looks_like_drafter(meta, arch: str | None) -> bool: +def is_drafter_arch(arch: str | None) -> bool: + """Whether ``arch`` names a draft/assistant model - a speculative-decoding + companion for a target model, never a standalone servable model.""" if arch in _DRAFTER_ARCHES: return True - if arch and ("assistant" in arch or "_mtp" in arch): + return bool(arch) and ("assistant" in arch or "_mtp" in arch) + + +def _looks_like_drafter(meta, arch: str | None) -> bool: + if is_drafter_arch(arch): return True if arch: for suf in _BACKBONE_FIELDS: @@ -305,7 +312,8 @@ def header_sampling(path) -> dict: def find_mtp_companion( path: str, - drafter_arch: str | tuple = ("deepseek4-dspark", "deepseek4_mtp_support"), + drafter_arch: str | tuple = ("deepseek4-dspark", "dflash", + "deepseek4_mtp_support"), ) -> 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 diff --git a/gmlx/manage.py b/gmlx/manage.py index dca4da7..75839a4 100644 --- a/gmlx/manage.py +++ b/gmlx/manage.py @@ -165,7 +165,11 @@ def _verdict(ref: remote.Ref, report: remote.HeaderReport, *, # A LoRA adapter carries its base model's arch (general.type = "adapter"), # so like mmproj it's a valid companion file, never a standalone model. adapter = report.gguf_type == "adapter" - if mmproj or adapter: + # A draft model (MTP / DSpark / assistant) is likewise a companion: it + # rides a target model, so the standalone arch gate doesn't apply. + from .discovery import is_drafter_arch + drafter = not (mmproj or adapter) and is_drafter_arch(report.arch) + if mmproj or adapter or drafter: arch_ok, arch_err = False, None else: arch_ok, arch_err = _arch_status(report.arch, hf_source=hf_source) @@ -180,12 +184,14 @@ def _verdict(ref: remote.Ref, report: remote.HeaderReport, *, # unknown); companions (mmproj/adapter) ride a base model, so their # standalone fit is not judged. "fit": (classify_fit(size_bytes, ram) - if size_bytes and not (mmproj or adapter) else None), + if size_bytes and not (mmproj or adapter or drafter) + else None), "arch": report.arch, "arch_supported": arch_ok, "arch_error": arch_err, "mmproj": mmproj, "adapter": adapter, + "drafter": drafter, "n_shards": n_shards, "n_tensors": report.n_tensors, "codecs": dict(sorted(report.histogram.items())), @@ -195,7 +201,8 @@ def _verdict(ref: remote.Ref, report: remote.HeaderReport, *, # loadable = runs standalone; usable also admits a healthy companion # (mmproj / LoRA adapter). "usable": (report.loadable_codecs and arch_ok) - or ((mmproj or adapter) and report.loadable_codecs), + or ((mmproj or adapter or drafter) + and report.loadable_codecs), } @@ -207,6 +214,8 @@ def _print_report(v: dict) -> None: print(f" architecture: {arch} [mmproj companion]") elif v.get("adapter"): print(f" architecture: {arch} [LoRA adapter]") + elif v.get("drafter"): + print(f" architecture: {arch} [draft-model companion]") elif v["arch_supported"]: print(f" architecture: {arch} [supported]") else: @@ -242,13 +251,18 @@ def _print_report(v: dict) -> None: print(" => adapter companion: a LoRA adapter, not a standalone model -") print(" attach it to its base GGUF: --adapter on run/chat/serve, or") print(" `adapter:` per model in the server config.") + elif v.get("drafter") and v["codecs_loadable"]: + print(" => draft-model companion: a speculative-decoding drafter, not a") + print(" standalone model - place it in the same directory as its") + print(" target GGUF (autodetected) or pass --draft-gguf.") else: reasons = [] if not v["codecs_loadable"]: bad = ", ".join(f"{k}x{n}" for k, n in v["unsupported_codecs"].items()) reasons.append(f"unsupported codecs ({bad})") - if not v["arch_supported"] and not (v.get("mmproj") or v.get("adapter")): + if not v["arch_supported"] and not (v.get("mmproj") or v.get("adapter") + or v.get("drafter")): reasons.append("architecture not supported") print(f" => not loadable: {'; '.join(reasons)}") @@ -846,7 +860,8 @@ def cmd_pull(argv: list | None = None, prog: str = "gmlx pull") -> int: if not v["codecs_loadable"] and v["unsupported_codecs"]: reasons.append( "no kernel for codec(s): " + ", ".join(v["unsupported_codecs"])) - if not v["arch_supported"] and not (v.get("mmproj") or v.get("adapter")): + if not v["arch_supported"] and not (v.get("mmproj") or v.get("adapter") + or v.get("drafter")): reasons.append(f"unsupported arch: {v['arch'] or '?'}") why = f" - {'; '.join(reasons)}" if reasons else "" print(f"\nrefusing to download an unloadable GGUF{why}. Pass --force to " diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index 4fd5a0c..3ee531f 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -477,6 +477,11 @@ def _load_deepseek4_mtp_drafter( arrays, kquant_meta, d_arch, _meta, _shapes = load_gguf_wire_bytes( draft_gguf_path, zero_copy=zero_copy ) + if d_arch == "dflash": + arrays, kquant_meta, _meta = normalize_dflash_arrays( + arrays, kquant_meta, _meta + ) + d_arch = "deepseek4-dspark" if d_arch == "deepseek4-dspark": return _load_deepseek4_dspark_drafter( draft_gguf_path, @@ -489,9 +494,9 @@ def _load_deepseek4_mtp_drafter( ) if d_arch != "deepseek4_mtp_support": raise ValueError( - f"{draft_gguf_path}: expected a deepseek4_mtp_support or " - f"deepseek4-dspark drafter GGUF for a deepseek_v4 target, got " - f"arch {d_arch!r}" + f"{draft_gguf_path}: expected a deepseek4-dspark, dflash, or " + f"deepseek4_mtp_support drafter GGUF for a deepseek_v4 target, " + f"got arch {d_arch!r}" ) log( f"[mtp] drafter gguf ({d_arch}): {len(arrays)} arrays, " @@ -586,6 +591,83 @@ def _load_deepseek4_mtp_drafter( } +# llama.cpp packages the same DSpark drafter under arch "dflash" (its +# convert_hf_to_gguf --dspark output, e.g. the unsloth release): per-stage +# tensors under blk.{k}.* with the same leaf names, and drafter-level tensors +# renamed through its root map. normalize_dflash_arrays translates that +# container back to the deepseek4-dspark namespace so one remap serves both. +# Drafter-level base name -> dspark base name; {L} is the last stage index +# (stage placement only matters for the remap's per-stage bookkeeping: +# main_proj/main_norm are stage-0 entries, the head tensors final-stage). +_DFLASH_ROOT_MAP = { + "fc": "mtp.0.main_proj", + "enc.output_norm": "mtp.0.main_norm", + "output_norm": "mtp.{L}.norm", + "markov_w1": "mtp.{L}.markov_head.markov_w1", + "markov_w2": "mtp.{L}.markov_head.markov_w2", + "conf_proj": "mtp.{L}.confidence_head.proj", + "output_hc_fn": "mtp.{L}.hc_head_fn", + "output_hc_base": "mtp.{L}.hc_head_base", + "output_hc_scale": "mtp.{L}.hc_head_scale", +} +_DFLASH_SUFFIXES = (".weight", ".scales", ".biases", ".bias") + + +def _dflash_rename(name: str, last_stage: int) -> str: + """The deepseek4-dspark name for one dflash tensor entry.""" + if name.startswith("blk."): + return "mtp." + name[len("blk."):] + base, suffix = name, "" + for s in _DFLASH_SUFFIXES: + if name.endswith(s): + base, suffix = name[: -len(s)], s + break + mapped = _DFLASH_ROOT_MAP.get(base) + if mapped is None: + raise RuntimeError( + f"dflash normalize: unknown tensor {name!r} " + f"(the drafter tensor set is closed)" + ) + return mapped.replace("{L}", str(last_stage)) + suffix + + +def normalize_dflash_arrays(arrays: dict, kquant_meta: dict, meta: dict): + """Translate a llama.cpp ``dflash`` GGUF (tensor names and metadata) to + the ``deepseek4-dspark`` namespace. Returns ``(arrays, kquant_meta, + meta)`` ready for :func:`_load_deepseek4_dspark_drafter`.""" + stages = { + int(n.split(".")[1]) for n in arrays if n.startswith("blk.") + } + if not stages: + raise RuntimeError("dflash normalize: no blk.* stage tensors found") + last_stage = max(stages) + n_arrays = { + _dflash_rename(name, last_stage): arr for name, arr in arrays.items() + } + n_kquant = { + _dflash_rename(name, last_stage): codec + for name, codec in kquant_meta.items() + } + n_meta = dict(meta) + block_size = meta.get("dflash.block_size") + if block_size is not None: + n_meta["dspark.block_size"] = block_size + # llama.cpp's converter writes the capture layers shifted +1 (its layer 0 + # is the embedding) and carries the noise token as the tokenizer mask + # token; undo both. Verified against the source config ([40, 41, 42] for + # a 43-layer target) and conversion/deepseek.py's add_target_layers. + layers = meta.get("dflash.target_layers") + if layers is not None: + n_meta["dspark.target_layer_ids"] = [int(i) - 1 for i in layers] + mask = meta.get("tokenizer.ggml.mask_token_id") + if mask is not None: + n_meta["dspark.noise_token_id"] = int(mask) + w1 = n_arrays.get(f"mtp.{last_stage}.markov_head.markov_w1.weight") + if w1 is not None and "dspark.markov_rank" not in n_meta: + n_meta["dspark.markov_rank"] = int(min(w1.shape)) + return n_arrays, n_kquant, n_meta + + def remap_deepseek4_dspark_arrays( arrays: dict, kquant_meta: dict, diff --git a/tests/test_dflash_normalize.py b/tests/test_dflash_normalize.py new file mode 100644 index 0000000..6362e0d --- /dev/null +++ b/tests/test_dflash_normalize.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +"""dflash -> deepseek4-dspark normalization: llama.cpp's container for the +DSpark drafter (the unsloth release) translated to the gmlx-native namespace. +Name/metadata logic only - synthetic arrays, no GGUF files, no model load.""" +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core") + +from gmlx.mtp_load import ( # noqa: E402 + normalize_dflash_arrays, + remap_deepseek4_dspark_arrays, +) + +# Tensor skeleton of the unsloth dspark-DeepSeek-V4-Flash-0731-Q8_0.gguf: +# per-stage leaves match the native sidecar; drafter-level names go through +# llama.cpp's root map (fc / enc.output_norm / output_norm / markov_w1 / +# markov_w2 / conf_proj / output_hc_*). +_STAGE_LEAVES = ( + "attn_kv.weight", "attn_kv_a_norm.weight", "attn_norm.weight", + "attn_output_a.weight", "attn_output_b.weight", "attn_q_a.weight", + "attn_q_a_norm.weight", "attn_q_b.weight", "attn_sinks.weight", + "exp_probs_b.bias", "ffn_down_exps.weight", "ffn_down_shexp.weight", + "ffn_gate_exps.weight", "ffn_gate_inp.weight", "ffn_gate_shexp.weight", + "ffn_norm.weight", "ffn_up_exps.weight", "ffn_up_shexp.weight", + "hc_attn_base.weight", "hc_attn_fn.weight", "hc_attn_scale.weight", + "hc_ffn_base.weight", "hc_ffn_fn.weight", "hc_ffn_scale.weight", +) +_ROOT_LEAVES = ( + "fc.weight", "enc.output_norm.weight", "output_norm.weight", + "markov_w1.weight", "markov_w2.weight", "conf_proj.weight", + "output_hc_base.weight", "output_hc_fn.weight", "output_hc_scale.weight", +) +_META = { + "general.architecture": "dflash", + "dflash.block_size": 5, + "dflash.target_layers": [41, 42, 43], + "tokenizer.ggml.mask_token_id": 128799, +} + + +def _dflash_fixture(n_stages=3): + arrays = {} + for k in range(n_stages): + for leaf in _STAGE_LEAVES: + arrays[f"blk.{k}.{leaf}"] = mx.zeros((2, 2)) + for leaf in _ROOT_LEAVES: + arrays[leaf] = mx.zeros((2, 2)) + arrays["markov_w1.weight"] = mx.zeros((129280, 256)) + kquant = {"blk.0.attn_kv.weight": "q8_0", "fc.weight": "q8_0"} + arrays["blk.0.attn_kv.scales"] = mx.zeros((1,)) + arrays["fc.scales"] = mx.zeros((1,)) + return arrays, kquant, dict(_META) + + +def test_stage_tensors_land_under_mtp(): + arrays, kquant, meta = _dflash_fixture() + n_arrays, n_kquant, _ = normalize_dflash_arrays(arrays, kquant, meta) + assert "mtp.0.attn_kv.weight" in n_arrays + assert "mtp.2.ffn_gate_exps.weight" in n_arrays + assert n_kquant["mtp.0.attn_kv.weight"] == "q8_0" + assert "mtp.0.attn_kv.scales" in n_arrays # quant siblings ride along + assert not any(n.startswith("blk.") for n in n_arrays) + + +def test_root_tensors_get_dspark_names(): + arrays, kquant, meta = _dflash_fixture() + n_arrays, n_kquant, _ = normalize_dflash_arrays(arrays, kquant, meta) + assert "mtp.0.main_proj.weight" in n_arrays # fc + assert "mtp.0.main_proj.scales" in n_arrays + assert n_kquant["mtp.0.main_proj.weight"] == "q8_0" + assert "mtp.0.main_norm.weight" in n_arrays # enc.output_norm + assert "mtp.2.norm.weight" in n_arrays # output_norm -> final stage + assert "mtp.2.markov_head.markov_w1.weight" in n_arrays + assert "mtp.2.confidence_head.proj.weight" in n_arrays + assert "mtp.2.hc_head_fn.weight" in n_arrays # output_hc_fn + + +def test_meta_translation_unshifts_target_layers(): + """llama.cpp writes capture layers +1 (its layer 0 is the embedding) and + the noise token as the tokenizer mask token; both must be undone.""" + arrays, kquant, meta = _dflash_fixture() + _, _, n_meta = normalize_dflash_arrays(arrays, kquant, meta) + assert n_meta["dspark.target_layer_ids"] == [40, 41, 42] + assert n_meta["dspark.noise_token_id"] == 128799 + assert n_meta["dspark.block_size"] == 5 + assert n_meta["dspark.markov_rank"] == 256 # from markov_w1's shape + + +def test_unknown_tensor_is_hard_error(): + arrays, kquant, meta = _dflash_fixture() + arrays["mystery.weight"] = mx.zeros((2, 2)) + with pytest.raises(RuntimeError, match="unknown tensor"): + normalize_dflash_arrays(arrays, kquant, meta) + + +def test_normalized_set_survives_dspark_remap(): + """End to end through the closed-set dspark remap: every translated name + must be known and every drafter param must be produced.""" + arrays, kquant, meta = _dflash_fixture() + n_arrays, n_kquant, _ = normalize_dflash_arrays(arrays, kquant, meta) + # attn_output_a is reshaped (o_groups x o_lora_rank x -1) by the remap + for k in range(3): + n_arrays[f"mtp.{k}.attn_output_a.weight"] = mx.zeros((8 * 1024, 4)) + hf, hf_kq, stats = remap_deepseek4_dspark_arrays( + n_arrays, n_kquant, n_stages=3, o_groups=8, o_lora_rank=1024 + ) + for key in ("main_proj.weight", "main_norm.weight", "norm.weight", + "markov_w1.weight", "markov_w2.weight", + "confidence_proj.weight", "hc_head.fn", + "stages.0.block.attn.wkv.weight", + "stages.2.block.ffn.switch_mlp.gate_proj.weight"): + assert key in hf, key + assert stats["mapped"] == len( + [n for n in n_arrays if not n.endswith(".scales")] + ) diff --git a/tests/test_discovery.py b/tests/test_discovery.py index aaea259..bda4919 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -59,6 +59,36 @@ def test_mtp_support_arch_is_drafter(): assert c.kind == "drafter" +def test_dflash_arch_is_drafter(): + """llama.cpp packages the DSpark drafter under arch `dflash` (the unsloth + release); it must classify as a companion, not an unsupported model.""" + c = _classify({"general.architecture": "dflash"}, + "dspark-DeepSeek-V4-Flash-0731-Q8_0.gguf") + assert c.kind == "drafter" + + +def test_companion_prefers_native_dspark_over_dflash(tmp_path, monkeypatch): + """With both containers next to the target, the gmlx-native sidecar wins.""" + target = tmp_path / "model.gguf" + for name in ("a-dflash.gguf", "b-dspark.gguf", "model.gguf"): + (tmp_path / name).write_bytes(b"GGUF") + metas = {str(tmp_path / "a-dflash.gguf"): {"arch": "dflash"}, + str(tmp_path / "b-dspark.gguf"): {"arch": "deepseek4-dspark"}, + str(target): {"arch": "deepseek4"}} + monkeypatch.setattr(disc, "header_meta", lambda p: metas.get(str(p))) + assert disc.find_mtp_companion(str(target)) == str(tmp_path / "b-dspark.gguf") + + +def test_companion_accepts_dflash_alone(tmp_path, monkeypatch): + target = tmp_path / "model.gguf" + for name in ("dspark-q8.gguf", "model.gguf"): + (tmp_path / name).write_bytes(b"GGUF") + metas = {str(tmp_path / "dspark-q8.gguf"): {"arch": "dflash"}, + str(target): {"arch": "deepseek4"}} + monkeypatch.setattr(disc, "header_meta", lambda p: metas.get(str(p))) + assert disc.find_mtp_companion(str(target)) == str(tmp_path / "dspark-q8.gguf") + + def test_backbone_field_implies_drafter(): """A future/unknown drafter arch is still caught by its target-backbone field.""" c = _classify({"general.architecture": "gemma4-weird-draft", diff --git a/tests/test_manage.py b/tests/test_manage.py index a21d7d0..00e526b 100644 --- a/tests/test_manage.py +++ b/tests/test_manage.py @@ -884,6 +884,31 @@ def test_validate_mmproj_json(tmp_path, capsys): assert v["usable"] is True +# drafter companion files (MTP / DSpark sidecars, incl. llama.cpp's dflash) +@pytest.mark.parametrize("arch", ["deepseek4-dspark", "dflash", + "deepseek4_mtp_support"]) +def test_validate_local_drafter_companion(tmp_path, capsys, arch): + p = tmp_path / "drafter.gguf" + _mint(p, arch=arch, codec=GT.Q8_0) + rc = manage.cmd_validate([str(p)]) + out = capsys.readouterr().out + assert rc == 0 # valid for its purpose + assert "draft-model companion" in out + assert "--draft-gguf" in out # the message says how to use it + assert "not loadable" not in out and "[unsupported]" not in out + + +def test_validate_drafter_json(tmp_path, capsys): + p = tmp_path / "drafter.gguf" + _mint(p, arch="dflash", codec=GT.Q8_0) + rc = manage.cmd_validate([str(p), "--json"]) + v = json.loads(capsys.readouterr().out) + assert rc == 0 + assert v["drafter"] is True + assert v["loadable"] is False # not a standalone model + assert v["usable"] is True + + def test_pull_mmproj_no_force_needed(tmp_path, monkeypatch): _serve(monkeypatch, _mint_mmproj(tmp_path / "mmproj.gguf")) called = []