mlx-dspark 0.0.1__py3-none-any.whl

mlx_dspark/__init__.py

@@ -0,0 +1,42 @@
+"""mlx-dspark: DSpark speculative decoding for Apple Silicon (MLX).
+
+DSpark (from DeepSeek's DeepSpec codebase) is a semi-autoregressive,
+EAGLE-family speculative-decoding drafter:
+
+  - a *parallel backbone* proposes base logits for all K draft positions at once,
+  - a tiny *sequential head* (low-rank, previous-token-conditioned) corrects
+    suffix decay,
+  - a *confidence head* scores how likely each drafted token survives
+    verification (used here for adaptive draft length instead of the
+    server-side load-aware scheduler, which is irrelevant single-user).
+
+This package targets single-user local inference on Apple Silicon.
+"""
+
+__version__ = "0.0.1"
+
+from .config import DSparkConfig
+from .load import (
+    DEFAULT_DRAFTER,
+    DEFAULT_TARGET,
+    PRESETS,
+    load_drafter,
+    load_pair,
+    load_target,
+)
+from .generate import GenResult, greedy_generate, speculative_generate
+from .target import Target
+
+__all__ = [
+    "DSparkConfig",
+    "Target",
+    "load_drafter",
+    "load_target",
+    "load_pair",
+    "speculative_generate",
+    "greedy_generate",
+    "GenResult",
+    "PRESETS",
+    "DEFAULT_TARGET",
+    "DEFAULT_DRAFTER",
+]

mlx_dspark/__main__.py

@@ -0,0 +1,84 @@
+"""CLI: run DSpark speculative decoding on Apple Silicon (streams tokens live).
+
+Side-by-side demo (record each, then stack the two screen captures):
+
+  # left panel — plain target, no drafter
+  python -m mlx_dspark --mode baseline --prompt "Explain how rainbows form." --max-new-tokens 220
+
+  # right panel — DSpark speculative decoding (same prompt, same output, faster)
+  python -m mlx_dspark --mode dspark   --prompt "Explain how rainbows form." --max-new-tokens 220
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+
+from .generate import greedy_generate, speculative_generate
+from .load import PRESETS, load_drafter, load_target
+
+
+def _emit(s: str) -> None:
+    sys.stdout.write(s)
+    sys.stdout.flush()
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(prog="mlx_dspark")
+    ap.add_argument("--mode", choices=["dspark", "baseline"], default="dspark",
+                    help="dspark = speculative decoding; baseline = plain greedy target")
+    ap.add_argument("--family", choices=["gemma4", "qwen3"], default="gemma4",
+                    help="model preset (target + drafter); overridden by --target/--drafter")
+    ap.add_argument("--prompt", default="Explain how rainbows form, in a few sentences.")
+    ap.add_argument("--target", default=None)
+    ap.add_argument("--drafter", default=None)
+    ap.add_argument("--max-new-tokens", type=int, default=220)
+    ap.add_argument("--max-draft", type=int, default=4)
+    ap.add_argument("--confidence-threshold", type=float, default=0.0)
+    ap.add_argument("--drafter-bits", type=int, default=4)
+    ap.add_argument("--no-chat-template", action="store_true")
+    ap.add_argument("--no-stream", action="store_true")
+    args = ap.parse_args()
+    target_repo = args.target or PRESETS[args.family]["target"]
+    drafter_repo = args.drafter or PRESETS[args.family]["drafter"]
+
+    label = "DSpark speculative" if args.mode == "dspark" else "Baseline (plain greedy)"
+    print(f"loading {args.mode}: target={target_repo}"
+          + (f", drafter={drafter_repo}" if args.mode == "dspark" else ""))
+    target, tok = load_target(target_repo)
+    drafter = None
+    if args.mode == "dspark":
+        drafter, _ = load_drafter(drafter_repo, quantize=args.drafter_bits > 0,
+                                  bits=max(args.drafter_bits, 2))
+
+    on_text = None if args.no_stream else _emit
+    print("\n" + "=" * 64)
+    print(f"  ▶  {label}   ·   {target_repo.split('/')[-1]}")
+    print("=" * 64)
+
+    if args.mode == "dspark":
+        res = speculative_generate(
+            target, tok, drafter, args.prompt,
+            max_new_tokens=args.max_new_tokens, max_draft_tokens=args.max_draft,
+            confidence_threshold=args.confidence_threshold,
+            apply_chat_template=not args.no_chat_template, on_text=on_text,
+        )
+        extra = f" · accept {res.mean_accept_len:.2f}/round · {res.target_forwards} target fwds"
+    else:
+        res = greedy_generate(
+            target, tok, args.prompt, max_new_tokens=args.max_new_tokens,
+            apply_chat_template=not args.no_chat_template, on_text=on_text,
+        )
+        extra = ""
+    if args.no_stream:
+        print(res.text)
+
+    print("\n" + "-" * 64)
+    print(f"  {res.num_tokens} tokens · {res.seconds:.2f}s · "
+          f"\033[1m{res.tokens_per_sec:.1f} tok/s\033[0m{extra}")
+    print("-" * 64)
+
+
+if __name__ == "__main__":
+    main()

mlx_dspark/config.py

@@ -0,0 +1,145 @@
+"""DSpark drafter config — loaded from the HF checkpoint's config.json.
+
+Supports two drafter families with a shared inference path:
+  - gemma4  (gemma4_text): k_eq_v attention, v_norm, partial/proportional rope,
+            sandwich norms + layer_scalar, gelu-tanh MLP, logit softcap.
+  - qwen3   (qwen3):       standard GQA (separate v_proj, no v_norm), default rope,
+            Llama-style 2-norm layer, silu MLP, no softcap.
+Only the fields the MLX inference path needs are pulled out.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class DSparkConfig:
+    family: str = "gemma4"             # "gemma4" | "qwen3"
+
+    # core dims
+    hidden_size: int = 3840
+    vocab_size: int = 262144
+    num_hidden_layers: int = 5
+    intermediate_size: int = 15360
+    rms_norm_eps: float = 1e-6
+
+    # attention
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    num_global_key_value_heads: int = 1
+    head_dim: int = 256
+    global_head_dim: int = 512
+    attention_k_eq_v: bool = True
+    attention_bias: bool = False
+
+    # rope
+    rope_theta: float = 1_000_000.0
+    partial_rotary_factor: float = 0.25
+    rope_type: str = "proportional"
+
+    # dspark specifics
+    block_size: int = 7
+    mask_token_id: int = 4
+    target_layer_ids: list[int] = field(default_factory=lambda: [5, 17, 29, 41, 46])
+    num_target_layers: int = 48
+
+    # markov + confidence
+    markov_rank: int = 256
+    markov_head_type: str = "vanilla"
+    enable_confidence_head: bool = True
+    confidence_head_with_markov: bool = True
+
+    # logits
+    final_logit_softcapping: float | None = 30.0
+    pad_token_id: int = 0
+
+    # ---- family-derived knobs (set in from_json) ----
+    mlp_activation: str = "gelu_tanh"   # "gelu_tanh" | "silu"
+    norm_style: str = "gemma"           # "gemma" (sandwich+scalar) | "qwen" (llama 2-norm)
+    use_v_norm: bool = True             # gemma: RMSNormNoScale v_norm; qwen: none
+    attention_scaling: float | None = None  # None -> 1/sqrt(attn_head_dim)
+
+    @property
+    def attn_head_dim(self) -> int:
+        """Head dim used by the drafter's own attention."""
+        return self.global_head_dim if self.family == "gemma4" else self.head_dim
+
+    @property
+    def n_kv_heads(self) -> int:
+        if self.family == "gemma4" and self.attention_k_eq_v:
+            return self.num_global_key_value_heads
+        return self.num_key_value_heads
+
+    @property
+    def scaling(self) -> float:
+        if self.attention_scaling is not None:
+            return self.attention_scaling
+        return self.attn_head_dim ** -0.5 if self.family == "qwen3" else 1.0
+
+    @property
+    def rope_parameters(self) -> dict:
+        return {"rope_type": self.rope_type, "partial_rotary_factor": self.partial_rotary_factor}
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> "DSparkConfig":
+        with open(path) as f:
+            c = json.load(f)
+        mt = c.get("model_type", "")
+        family = "qwen3" if "qwen3" in mt else "gemma4"
+
+        if family == "qwen3":
+            rp = c.get("rope_parameters") or {}
+            return cls(
+                family="qwen3",
+                hidden_size=c["hidden_size"], vocab_size=c["vocab_size"],
+                num_hidden_layers=c["num_hidden_layers"],
+                intermediate_size=c["intermediate_size"],
+                rms_norm_eps=c.get("rms_norm_eps", 1e-6),
+                num_attention_heads=c["num_attention_heads"],
+                num_key_value_heads=c.get("num_key_value_heads", 8),
+                head_dim=c.get("head_dim", c["hidden_size"] // c["num_attention_heads"]),
+                attention_k_eq_v=False, attention_bias=c.get("attention_bias", False),
+                rope_theta=rp.get("rope_theta", c.get("rope_theta", 1_000_000.0)),
+                rope_type="default",
+                block_size=c["block_size"], mask_token_id=c["mask_token_id"],
+                target_layer_ids=list(c["target_layer_ids"]),
+                num_target_layers=c.get("num_target_layers", 36),
+                markov_rank=c.get("markov_rank", 256),
+                markov_head_type=c.get("markov_head_type", "vanilla"),
+                enable_confidence_head=c.get("enable_confidence_head", True),
+                confidence_head_with_markov=c.get("confidence_head_with_markov", True),
+                final_logit_softcapping=c.get("final_logit_softcapping", None),
+                pad_token_id=c.get("pad_token_id") or 0,
+                mlp_activation="silu", norm_style="qwen", use_v_norm=False,
+            )
+
+        rope = (c.get("rope_parameters") or {}).get("full_attention", {}) or {}
+        return cls(
+            family="gemma4",
+            hidden_size=c["hidden_size"], vocab_size=c["vocab_size"],
+            num_hidden_layers=c["num_hidden_layers"],
+            intermediate_size=c["intermediate_size"],
+            rms_norm_eps=c.get("rms_norm_eps", 1e-6),
+            num_attention_heads=c["num_attention_heads"],
+            num_key_value_heads=c.get("num_key_value_heads", 8),
+            num_global_key_value_heads=c.get("num_global_key_value_heads", 1),
+            head_dim=c.get("head_dim", 256), global_head_dim=c.get("global_head_dim", 512),
+            attention_k_eq_v=c.get("attention_k_eq_v", True),
+            attention_bias=c.get("attention_bias", False),
+            rope_theta=rope.get("rope_theta", 1_000_000.0),
+            partial_rotary_factor=rope.get("partial_rotary_factor", 0.25),
+            rope_type=rope.get("rope_type", "proportional"),
+            block_size=c["block_size"], mask_token_id=c["mask_token_id"],
+            target_layer_ids=list(c["target_layer_ids"]),
+            num_target_layers=c.get("num_target_layers", 48),
+            markov_rank=c.get("markov_rank", 256),
+            markov_head_type=c.get("markov_head_type", "vanilla"),
+            enable_confidence_head=c.get("enable_confidence_head", True),
+            confidence_head_with_markov=c.get("confidence_head_with_markov", True),
+            final_logit_softcapping=c.get("final_logit_softcapping", 30.0),
+            pad_token_id=c.get("pad_token_id", 0),
+            mlp_activation="gelu_tanh", norm_style="gemma", use_v_norm=True,
+        )

mlx_dspark/generate.py

@@ -0,0 +1,289 @@
+"""DSpark speculative decoding loop (greedy, batch=1) for Apple Silicon.
+
+Per round:
+  1. draft a block of K tokens from the parallel backbone + Markov head,
+  2. verify them in one target forward,
+  3. accept the matching prefix + 1 bonus token (so >=1 token/round always),
+  4. trim the target KV cache and grow the fused-hidden context buffer.
+
+Because the target verifies every token, the *output is exactly greedy target
+decoding* regardless of drafter quality — drafter quality only shows up as the
+acceptance length (tokens committed per target forward).
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+import mlx.core as mx
+
+TAP = None  # set from drafter config at call time
+
+
+@dataclass
+class GenResult:
+    text: str
+    token_ids: list[int]
+    num_tokens: int
+    num_rounds: int
+    accept_lengths: list[int]
+    target_forwards: int
+    seconds: float
+
+    @property
+    def mean_accept_len(self) -> float:
+        return self.num_tokens / max(self.num_rounds, 1)
+
+    @property
+    def tokens_per_sec(self) -> float:
+        return self.num_tokens / max(self.seconds, 1e-9)
+
+
+def encode_prompt(tokenizer, prompt: str, use_chat: bool = True) -> list[int]:
+    """Token ids for a user prompt, using the model's chat template when present.
+
+    Gemma-4 uses `<|turn>` / `<channel|>` markers (NOT Gemma-3's `<start_of_turn>`),
+    so the template must be applied via the tokenizer — hand-formatting breaks the
+    instruct model. apply_chat_template may return list[int] or a BatchEncoding.
+    """
+    if use_chat and getattr(tokenizer, "chat_template", None):
+        r = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}], add_generation_prompt=True
+        )
+        if isinstance(r, (list, tuple)):
+            if r and isinstance(r[0], int):
+                return list(r)
+            if r and isinstance(r[0], (list, tuple)):
+                return list(r[0])
+        ii = None
+        if hasattr(r, "__contains__") and "input_ids" in r:
+            ii = r["input_ids"]
+        elif hasattr(r, "input_ids"):
+            ii = r.input_ids
+        if ii is not None:
+            ii = list(ii)
+            return list(ii[0]) if ii and isinstance(ii[0], (list, tuple)) else ii
+        if hasattr(r, "ids"):
+            return list(r.ids)
+    return list(tokenizer.encode(prompt))
+
+
+def eos_token_ids(tokenizer) -> set[int]:
+    """Collect stop-token ids: eos + Gemma turn-end markers (Gemma-4 uses <turn|>=106;
+    note <end_of_turn> is the UNK id in Gemma-4, so it must be filtered out)."""
+    ids: set[int] = set()
+    e = getattr(tokenizer, "eos_token_ids", None)
+    if isinstance(e, int):
+        ids.add(e)
+    elif e:
+        ids.update(int(x) for x in e)
+    e1 = getattr(tokenizer, "eos_token_id", None)
+    if isinstance(e1, int):
+        ids.add(e1)
+    unk = getattr(tokenizer, "unk_token_id", None)
+    # Gemma-4 (<turn|>), Gemma-3 (<end_of_turn>), Qwen (<|im_end|>), raw eos
+    for t in ("<turn|>", "<end_of_turn>", "<|im_end|>", "<|endoftext|>", "<eos>"):
+        try:
+            i = tokenizer.convert_tokens_to_ids(t)
+        except Exception:
+            continue
+        if isinstance(i, int) and i >= 0 and i != unk:
+            ids.add(i)
+    return ids
+
+
+def greedy_generate(
+    target_model,
+    tokenizer,
+    prompt: str,
+    *,
+    max_new_tokens: int = 128,
+    apply_chat_template: bool = True,
+    on_text=None,
+) -> GenResult:
+    """Plain greedy decoding of the target (no drafter, no hidden-state capture) —
+    the fair 'run the model normally' baseline. Streams via on_text."""
+    eos_ids = eos_token_ids(tokenizer)
+    ids = encode_prompt(tokenizer, prompt, use_chat=apply_chat_template)
+    cache = target_model.make_cache()
+
+    t0 = time.time()
+    logits = target_model.plain(mx.array([ids]), cache)
+    nxt = int(mx.argmax(logits[0, -1]).item())
+    out_ids = [nxt]
+    streamed = 0
+
+    def _stream():
+        nonlocal streamed
+        if on_text is None:
+            return
+        disp = [t for t in out_ids if t not in eos_ids]
+        full = tokenizer.decode(disp)
+        if len(full) > streamed:
+            on_text(full[streamed:])
+            streamed = len(full)
+
+    _stream()
+    while len(out_ids) < max_new_tokens and nxt not in eos_ids:
+        logits = target_model.plain(mx.array([[nxt]]), cache)
+        nxt = int(mx.argmax(logits[0, -1]).item())
+        out_ids.append(nxt)
+        _stream()
+
+    secs = time.time() - t0
+    disp = [t for t in out_ids if t not in eos_ids]
+    return GenResult(
+        text=tokenizer.decode(disp),
+        token_ids=out_ids,
+        num_tokens=len(out_ids),
+        num_rounds=len(out_ids),
+        accept_lengths=[1] * len(out_ids),
+        target_forwards=len(out_ids),
+        seconds=secs,
+    )
+
+
+def _make_target_cache(target):
+    return target.make_cache()
+
+
+def _run_target(target, ids: mx.array, cache, tap: list[int]):
+    """ids: [1, L]. Returns (logits[1,L,V], fused_hidden[1,L,n_tap*H])."""
+    return target.run(ids, cache, tap)
+
+
+def speculative_generate(
+    target_model,
+    tokenizer,
+    drafter,
+    prompt: str,
+    *,
+    max_new_tokens: int = 128,
+    confidence_threshold: float = 0.0,
+    max_draft_tokens: int | None = 4,
+    apply_chat_template: bool = True,
+    on_text=None,
+    verbose: bool = False,
+) -> GenResult:
+    """Greedy speculative decoding. Output is target-greedy by construction (up to
+    fp tie-breaking on near-ties). ``max_draft_tokens`` caps how many of the 7-token
+    block are verified per round; on Apple Silicon the target verify cost grows with
+    tokens, so the optimum is ~= acceptance length (default 4). ``None`` = full block
+    (faithful but slower on M-series). ``confidence_threshold`` > 0 instead truncates
+    the block adaptively using the drafter's confidence head."""
+    cfg = drafter.config
+    tap = list(cfg.target_layer_ids)
+    k = cfg.block_size
+    mask_id = cfg.mask_token_id
+    cap = k if max_draft_tokens is None else max(1, min(max_draft_tokens, k))
+
+    eos_ids = eos_token_ids(tokenizer)
+
+    # --- tokenize prompt ---
+    ids = encode_prompt(tokenizer, prompt, use_chat=apply_chat_template)
+    prompt_ids = mx.array([ids])
+
+    cache = _make_target_cache(target_model)
+    ctx_caches = drafter.make_ctx_cache()
+    t0 = time.time()
+
+    # --- prefill ---
+    logits, fused = _run_target(target_model, prompt_ids, cache, tap)
+    n_cached = prompt_ids.shape[1]
+    drafter.update_context(fused, ctx_offset=0, ctx_caches=ctx_caches)
+    pending = int(mx.argmax(logits[0, -1]).item())  # first committed token
+    mx.eval([c.k for c in ctx_caches])
+
+    out_ids: list[int] = [pending]
+    accept_lengths: list[int] = []
+    target_forwards = 1
+    streamed = 0
+
+    def _stream():
+        nonlocal streamed
+        if on_text is None:
+            return
+        disp = [t for t in out_ids if t not in eos_ids]
+        full = tokenizer.decode(disp)
+        if len(full) > streamed:
+            on_text(full[streamed:])
+            streamed = len(full)
+
+    _stream()
+    while len(out_ids) < max_new_tokens and pending not in eos_ids:
+        # ---- 1. draft a block ----
+        block_ids = [pending] + [mask_id] * (k - 1)
+        noise = drafter.embed(mx.array([block_ids]))            # [1, k, H]
+        block_hidden = drafter.backbone(noise, n_cached, ctx_caches)
+        base_logits = drafter.compute_logits(block_hidden)[0]   # [k, V]
+        draft = drafter.sample_block(base_logits, first_prev_token=pending)
+        mx.eval(draft)
+        draft = [int(x) for x in draft.tolist()]
+
+        # optional confidence-based truncation (adaptive block length)
+        if confidence_threshold > 0.0 and drafter.confidence_head is not None:
+            prev_tokens = mx.array([pending] + draft[:-1])
+            conf = mx.sigmoid(drafter.confidence_logits(block_hidden[0], prev_tokens))
+            mx.eval(conf)
+            below = [i for i, c in enumerate(conf.tolist()) if c < confidence_threshold]
+            if below:
+                draft = draft[: below[0]]
+        if cap < len(draft):
+            draft = draft[:cap]
+        if not draft:
+            draft = [int(mx.argmax(base_logits[0]).item())]  # always propose >=1
+
+        # ---- 2. verify with the target ----
+        verify_ids = mx.array([[pending] + draft])             # [1, 1+len(draft)]
+        v_logits, v_fused = _run_target(target_model, verify_ids, cache, tap)
+        mx.eval(v_logits, v_fused)
+        target_forwards += 1
+        tt = mx.argmax(v_logits[0], axis=-1)                   # [1+len(draft)]
+        tt = [int(x) for x in tt.tolist()]
+
+        # ---- 3. accept matching prefix + bonus ----
+        n = 0
+        while n < len(draft) and draft[n] == tt[n]:
+            n += 1
+        bonus = tt[n]                                          # correction / continuation
+        committed = draft[:n] + [bonus]
+        accept_lengths.append(len(committed))
+
+        # ---- 4. update caches/context ----
+        trim = len(draft) - n
+        if trim > 0:
+            for c in cache:
+                if c is not None and hasattr(c, "trim"):
+                    c.trim(trim)
+        # commit [pending, accepted drafts] (positions n_cached..n_cached+n) as context
+        drafter.update_context(
+            v_fused[:, : n + 1, :], ctx_offset=n_cached, ctx_caches=ctx_caches
+        )
+        n_cached = n_cached + n + 1
+        mx.eval([c.k for c in ctx_caches])
+
+        for tok in committed:
+            out_ids.append(tok)
+            if tok in eos_ids:
+                break
+        pending = committed[-1]
+        _stream()
+
+        if verbose:
+            print(f"  round {len(accept_lengths):3d}: drafted {len(draft)}, "
+                  f"accepted {n}, committed {len(committed)}")
+
+    secs = time.time() - t0
+    # strip trailing eos for display
+    disp = [t for t in out_ids if t not in eos_ids]
+    text = tokenizer.decode(disp)
+    return GenResult(
+        text=text,
+        token_ids=out_ids,
+        num_tokens=len(out_ids),
+        num_rounds=len(accept_lengths),
+        accept_lengths=accept_lengths,
+        target_forwards=target_forwards,
+        seconds=secs,
+    )

mlx_dspark/load.py

@@ -0,0 +1,117 @@
+"""Loaders for the target (Gemma-4 via mlx-vlm) and the DSpark drafter."""
+
+from __future__ import annotations
+
+import glob
+import os
+
+import mlx.core as mx
+import mlx.nn as nn
+from huggingface_hub import snapshot_download
+
+from .config import DSparkConfig
+from .model import DSparkDrafter
+from .target import Target
+
+# The drafter must be paired with the *instruct* target it was trained against, at decent
+# precision. Presets below; pick with load_pair("gemma4") or load_pair("qwen3").
+PRESETS = {
+    "gemma4": {
+        "target": "mlx-community/gemma-4-12B-it-8bit",
+        "drafter": "deepseek-ai/dspark_gemma4_12b_block7",
+    },
+    "qwen3": {
+        "target": "mlx-community/Qwen3-4B-8bit",
+        "drafter": "deepseek-ai/dspark_qwen3_4b_block7",
+    },
+}
+DEFAULT_TARGET = PRESETS["gemma4"]["target"]
+DEFAULT_DRAFTER = PRESETS["gemma4"]["drafter"]
+
+
+def _resolve(repo_or_path: str) -> str:
+    if os.path.isdir(repo_or_path):
+        return repo_or_path
+    return snapshot_download(repo_or_path)
+
+
+def load_drafter(
+    repo_or_path: str = DEFAULT_DRAFTER,
+    *,
+    quantize: bool = True,
+    bits: int = 4,
+    group_size: int = 64,
+):
+    """Return (drafter, config). Loads bf16 weights 1:1 by matching key names.
+
+    The drafter is ~6.86 GB in bf16 and runs every speculative round, so by
+    default it is quantized to 4-bit (~1.8 GB) — this is what makes spec
+    decoding a net speedup on Apple Silicon. Output correctness is unaffected
+    (the target verifies every token); only acceptance length may change.
+    """
+    path = _resolve(repo_or_path)
+    config = DSparkConfig.from_json(os.path.join(path, "config.json"))
+    drafter = DSparkDrafter(config)
+
+    weights: dict[str, mx.array] = {}
+    for st in glob.glob(os.path.join(path, "*.safetensors")):
+        weights.update(mx.load(st))
+
+    # Diagnose name mismatches before loading.
+    model_keys = {k for k, _ in _flatten_params(drafter)}
+    ckpt_keys = set(weights.keys())
+    missing = sorted(model_keys - ckpt_keys)
+    unexpected = sorted(ckpt_keys - model_keys)
+    if missing or unexpected:
+        print(f"[load_drafter] WARNING key mismatch:")
+        if missing:
+            print(f"  missing in checkpoint ({len(missing)}): {missing[:8]}")
+        if unexpected:
+            print(f"  unexpected in checkpoint ({len(unexpected)}): {unexpected[:8]}")
+
+    drafter.load_weights(list(weights.items()), strict=not (missing or unexpected))
+
+    if quantize:
+        # Quantize Linear/Embedding weights; norms/scalars stay full precision.
+        nn.quantize(drafter, group_size=group_size, bits=bits)
+
+    mx.eval(drafter.parameters())
+    return drafter, config
+
+
+def _flatten_params(module) -> list[tuple[str, mx.array]]:
+    from mlx.utils import tree_flatten
+
+    return tree_flatten(module.parameters())
+
+
+def load_target(repo_or_path: str = DEFAULT_TARGET):
+    """Return (Target, tokenizer). Routes text models (Qwen3) to mlx-lm and VLM/unified
+    models (Gemma-4) to mlx-vlm, then wraps in a family-aware Target (hidden-state tap)."""
+    import json
+
+    path = _resolve(repo_or_path)
+    model_type = ""
+    cfg_path = os.path.join(path, "config.json")
+    if os.path.exists(cfg_path):
+        with open(cfg_path) as f:
+            model_type = json.load(f).get("model_type", "")
+
+    if "qwen3" in model_type and "moe" not in model_type:
+        from mlx_lm import load as lm_load
+
+        model, tokenizer = lm_load(path)
+    else:
+        from mlx_vlm import load as vlm_load
+
+        model, processor = vlm_load(path)
+        tokenizer = getattr(processor, "tokenizer", processor)
+    return Target(model, tokenizer), tokenizer
+
+
+def load_pair(family: str = "gemma4", *, target_bits: str | None = None):
+    """Convenience: load (target, tokenizer, drafter, cfg) for a preset family."""
+    p = PRESETS[family]
+    target, tok = load_target(p["target"])
+    drafter, cfg = load_drafter(p["drafter"])
+    return target, tok, drafter, cfg

mlx_dspark/model.py

@@ -0,0 +1,277 @@
+"""DSpark drafter in MLX — Gemma-4 and Qwen3 families.
+
+Faithful port of the DeepSpec inference path. The EAGLE-style cross-attention is shared:
+Q comes from the draft block, K/V from concat([fused_target_context, block]), with the
+context K/V cached per layer (CtxCache). Family differences (norm layout, rope, MLP act,
+v handling, logit softcap) are config-driven. Module attribute names match the HF
+checkpoints so weights load 1:1.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_vlm.models.gemma4.rope_utils import initialize_rope
+
+from .config import DSparkConfig
+
+
+class RMSNormNoScale(nn.Module):
+    """RMSNorm with no learnable weight (Gemma-4 v_norm)."""
+
+    def __init__(self, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return mx.fast.rms_norm(x, None, self.eps)
+
+
+def _act(name: str):
+    return nn.silu if name == "silu" else nn.gelu_approx
+
+
+class MLP(nn.Module):
+    def __init__(self, config: DSparkConfig):
+        super().__init__()
+        h, i = config.hidden_size, config.intermediate_size
+        self.gate_proj = nn.Linear(h, i, bias=False)
+        self.up_proj = nn.Linear(h, i, bias=False)
+        self.down_proj = nn.Linear(i, h, bias=False)
+        self.act = _act(config.mlp_activation)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
+
+
+def _repeat_kv(x: mx.array, n_rep: int) -> mx.array:
+    if n_rep == 1:
+        return x
+    b, n_kv, s, d = x.shape
+    x = mx.expand_dims(x, 2)
+    x = mx.broadcast_to(x, (b, n_kv, n_rep, s, d))
+    return x.reshape(b, n_kv * n_rep, s, d)
+
+
+class CtxCache:
+    """Per-layer cache of the target context's projected K/V (roped K, normed/raw V)."""
+
+    __slots__ = ("k", "v")
+
+    def __init__(self):
+        self.k = None
+        self.v = None
+
+    def append(self, k: mx.array, v: mx.array) -> None:
+        if self.k is None:
+            self.k, self.v = k, v
+        else:
+            self.k = mx.concatenate([self.k, k], axis=2)
+            self.v = mx.concatenate([self.v, v], axis=2)
+
+
+class DSparkAttention(nn.Module):
+    """Cross-attention: Q from the draft block, K/V from [target_context, block]."""
+
+    def __init__(self, config: DSparkConfig):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.head_dim = config.attn_head_dim
+        self.k_eq_v = config.attention_k_eq_v
+        self.n_kv_heads = config.n_kv_heads
+        self.n_rep = self.n_heads // self.n_kv_heads
+        self.scale = config.scaling
+        self.use_v_norm = config.use_v_norm
+
+        h = config.hidden_size
+        b = config.attention_bias
+        self.q_proj = nn.Linear(h, self.n_heads * self.head_dim, bias=b)
+        self.k_proj = nn.Linear(h, self.n_kv_heads * self.head_dim, bias=b)
+        if not self.k_eq_v:
+            self.v_proj = nn.Linear(h, self.n_kv_heads * self.head_dim, bias=b)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, h, bias=b)
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        if self.use_v_norm:
+            self.v_norm = RMSNormNoScale(eps=config.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            dims=self.head_dim, base=config.rope_theta, traditional=False,
+            scaling_config=config.rope_parameters,
+        )
+
+    def _kv(self, x: mx.array):
+        """Project x -> (roped+normed K, V). k_eq_v shares k_proj for V."""
+        B, S, _ = x.shape
+        kp = self.k_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_norm(kp)
+        if self.k_eq_v:
+            v = self.v_norm(kp)
+        else:
+            v = self.v_proj(x).reshape(B, S, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+            if self.use_v_norm:
+                v = self.v_norm(v)
+        return k, v
+
+    def update_ctx(self, fused_new: mx.array, ctx_offset: int, cache: CtxCache) -> None:
+        k, v = self._kv(fused_new)
+        cache.append(self.rope(k, offset=ctx_offset), v)   # V is not roped
+
+    def attend(self, hidden: mx.array, block_offset: int, cache: CtxCache) -> mx.array:
+        B, q_len, _ = hidden.shape
+        q = self.q_proj(hidden).reshape(B, q_len, self.n_heads, self.head_dim)
+        q = self.rope(self.q_norm(q).transpose(0, 2, 1, 3), offset=block_offset)
+
+        k_blk, v_blk = self._kv(hidden)
+        k_blk = self.rope(k_blk, offset=block_offset)
+        k = mx.concatenate([cache.k, k_blk], axis=2)
+        v = mx.concatenate([cache.v, v_blk], axis=2)
+
+        k = _repeat_kv(k, self.n_rep)
+        v = _repeat_kv(v, self.n_rep)
+        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=None)
+        out = out.transpose(0, 2, 1, 3).reshape(B, q_len, -1)
+        return self.o_proj(out)
+
+
+class DSparkDecoderLayer(nn.Module):
+    def __init__(self, config: DSparkConfig):
+        super().__init__()
+        eps = config.rms_norm_eps
+        self.norm_style = config.norm_style
+        self.self_attn = DSparkAttention(config)
+        self.mlp = MLP(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=eps)
+        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=eps)
+        if self.norm_style == "gemma":
+            self.pre_feedforward_layernorm = nn.RMSNorm(config.hidden_size, eps=eps)
+            self.post_feedforward_layernorm = nn.RMSNorm(config.hidden_size, eps=eps)
+            self.layer_scalar = mx.ones((1,))
+
+    def __call__(self, hidden, block_offset, cache: CtxCache):
+        if self.norm_style == "gemma":
+            residual = hidden
+            h = self.input_layernorm(hidden)
+            h = self.self_attn.attend(h, block_offset, cache)
+            h = self.post_attention_layernorm(h)
+            h = residual + h
+            residual = h
+            h = self.pre_feedforward_layernorm(h)
+            h = self.mlp(h)
+            h = self.post_feedforward_layernorm(h)
+            h = residual + h
+            return h * self.layer_scalar
+        # qwen / llama 2-norm
+        residual = hidden
+        h = self.input_layernorm(hidden)
+        h = self.self_attn.attend(h, block_offset, cache)
+        h = residual + h
+        residual = h
+        h = self.post_attention_layernorm(h)
+        h = self.mlp(h)
+        return residual + h
+
+
+class VanillaMarkov(nn.Module):
+    """Rank-256 previous-token correction: logits += w2(w1[prev_token])."""
+
+    def __init__(self, config: DSparkConfig):
+        super().__init__()
+        self.markov_w1 = nn.Embedding(config.vocab_size, config.markov_rank)
+        self.markov_w2 = nn.Linear(config.markov_rank, config.vocab_size, bias=False)
+
+    def prev_embeddings(self, token_ids: mx.array) -> mx.array:
+        return self.markov_w1(token_ids)
+
+    def step_bias(self, token_ids: mx.array) -> mx.array:
+        return self.markov_w2(self.markov_w1(token_ids))
+
+
+class ConfidenceHead(nn.Module):
+    def __init__(self, input_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(input_dim, 1)
+
+    def __call__(self, features: mx.array) -> mx.array:
+        return self.proj(features).squeeze(-1)
+
+
+class DSparkDrafter(nn.Module):
+    def __init__(self, config: DSparkConfig):
+        super().__init__()
+        self.config = config
+        self.block_size = config.block_size
+        self.mask_token_id = config.mask_token_id
+        self.embed_scale = (float(config.hidden_size) ** 0.5) if config.family == "gemma4" else 1.0
+        self.softcap = config.final_logit_softcapping
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.fc = nn.Linear(
+            len(config.target_layer_ids) * config.hidden_size, config.hidden_size, bias=False
+        )
+        self.hidden_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.layers = [DSparkDecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.markov_head = VanillaMarkov(config) if config.markov_rank > 0 else None
+        self.confidence_head = None
+        if config.enable_confidence_head:
+            in_dim = config.hidden_size + (config.markov_rank if config.confidence_head_with_markov else 0)
+            self.confidence_head = ConfidenceHead(in_dim)
+
+    def embed(self, ids: mx.array) -> mx.array:
+        return self.embed_tokens(ids) * self.embed_scale
+
+    def fuse_target(self, target_hidden_cat: mx.array) -> mx.array:
+        return self.hidden_norm(self.fc(target_hidden_cat))
+
+    def make_ctx_cache(self) -> list[CtxCache]:
+        return [CtxCache() for _ in self.layers]
+
+    def update_context(self, target_hidden_cat, ctx_offset, ctx_caches) -> None:
+        fused = self.fuse_target(target_hidden_cat)
+        for layer, cache in zip(self.layers, ctx_caches):
+            layer.self_attn.update_ctx(fused, ctx_offset, cache)
+
+    def backbone(self, noise_embedding, block_offset, ctx_caches) -> mx.array:
+        h = noise_embedding
+        for layer, cache in zip(self.layers, ctx_caches):
+            h = layer(h, block_offset, cache)
+        return self.norm(h)
+
+    def compute_logits(self, hidden: mx.array) -> mx.array:
+        logits = self.lm_head(hidden)
+        if self.softcap is not None:
+            logits = mx.tanh(logits / self.softcap) * self.softcap
+        return logits
+
+    def sample_block(self, base_logits: mx.array, first_prev_token: int) -> mx.array:
+        k = base_logits.shape[0]
+        if self.markov_head is None:
+            return mx.argmax(base_logits, axis=-1)
+        tokens = []
+        prev = mx.array([first_prev_token])
+        for i in range(k):
+            step = base_logits[i] + self.markov_head.step_bias(prev)[0]
+            nxt = mx.argmax(step, axis=-1, keepdims=True)
+            tokens.append(nxt)
+            prev = nxt
+        return mx.concatenate(tokens)
+
+    def confidence_logits(self, block_hidden, prev_token_ids):
+        if self.confidence_head is None:
+            return None
+        if self.config.confidence_head_with_markov:
+            feats = mx.concatenate(
+                [block_hidden, self.markov_head.prev_embeddings(prev_token_ids)], axis=-1
+            )
+        else:
+            feats = block_hidden
+        return self.confidence_head(feats)
+
+
+# Backwards-compatible alias
+Gemma4DSparkDrafter = DSparkDrafter

mlx_dspark/target.py

@@ -0,0 +1,64 @@
+"""Family-aware target wrapper: KV cache + a hidden-state tap at given layers.
+
+- gemma4 (mlx-vlm): uses the built-in ``capture_layer_ids`` / ``hidden_sink`` hook.
+- qwen3  (mlx-lm):  no hook exists, so we replicate the model's forward loop and
+  capture the residual stream after the tapped layers.
+
+Both expose: make_cache(), run(ids, cache, tap)->(logits, fused_hidden), and
+plain(ids, cache)->logits (no capture, for the greedy baseline).
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+class Target:
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+        # mlx-vlm models expose .language_model; mlx-lm models expose .model + (lm_head|tied)
+        self.is_vlm = hasattr(model, "language_model")
+        self.family = "gemma4" if self.is_vlm else "qwen3"
+        if not self.is_vlm:
+            self._tied = bool(getattr(getattr(model, "args", None), "tie_word_embeddings", False))
+
+    # -- cache --
+    def make_cache(self):
+        if self.is_vlm:
+            return self.model.language_model.make_cache()
+        from mlx_lm.models.cache import make_prompt_cache
+        return make_prompt_cache(self.model)
+
+    # -- forward with hidden-state tap --
+    def run(self, ids: mx.array, cache, tap: list[int]):
+        """ids [1,L] -> (logits [1,L,V], fused_hidden [1,L,n_tap*H])."""
+        if self.is_vlm:
+            out = self.model.language_model(inputs=ids, cache=cache, capture_layer_ids=tap)
+            return out.logits, mx.concatenate(out.hidden_states, axis=-1)
+        return self._run_mlxlm(ids, cache, tap)
+
+    def _run_mlxlm(self, ids, cache, tap):
+        from mlx_lm.models.base import create_attention_mask
+
+        mm = self.model.model
+        tapset = set(tap)
+        h = mm.embed_tokens(ids)
+        mask = create_attention_mask(h, cache[0])
+        captured = []
+        for i, (layer, c) in enumerate(zip(mm.layers, cache)):
+            h = layer(h, mask, c)
+            if i in tapset:
+                captured.append(h)
+        hn = mm.norm(h)
+        if self._tied:
+            logits = mm.embed_tokens.as_linear(hn)
+        else:
+            logits = self.model.lm_head(hn)
+        return logits, mx.concatenate(captured, axis=-1)
+
+    # -- plain forward (no capture) for the greedy baseline --
+    def plain(self, ids: mx.array, cache):
+        if self.is_vlm:
+            return self.model.language_model(inputs=ids, cache=cache).logits
+        return self.model(ids, cache=cache)

mlx_dspark-0.0.1.dist-info/METADATA

@@ -0,0 +1,115 @@
+Metadata-Version: 2.4
+Name: mlx-dspark
+Version: 0.0.1
+Summary: DSpark speculative decoding (semi-autoregressive drafting) for Apple Silicon via MLX
+Project-URL: Homepage, https://github.com/ARahim3/mlx-dspark
+Project-URL: Repository, https://github.com/ARahim3/mlx-dspark
+Project-URL: Issues, https://github.com/ARahim3/mlx-dspark/issues
+Author-email: erahim3 <erahim3@gmail.com>
+License: MIT
+License-File: LICENSE
+License-File: NOTICE
+Keywords: apple-silicon,deepspec,dspark,llm-inference,mlx,speculative-decoding
+Classifier: Environment :: GPU
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: huggingface-hub
+Requires-Dist: mlx-lm>=0.31.3
+Requires-Dist: mlx-vlm>=0.6.3
+Requires-Dist: mlx>=0.31.2
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# mlx-dspark
+
+**DSpark speculative decoding for Apple Silicon**, built on [MLX](https://github.com/ml-explore/mlx).
+
+DSpark is DeepSeek's semi-autoregressive, EAGLE-family speculative-decoding drafter,
+open-sourced in the [DeepSpec](https://github.com/deepseek-ai/DeepSpec) codebase and used to
+accelerate DeepSeek-V4. This project ports the **inference path** to MLX so the published
+drafter checkpoints run natively on a Mac.
+
+**Supported families** (auto-detected from the drafter config):
+
+| family | target | drafter | RAM |
+|---|---|---|---|
+| `gemma4` | `gemma-4-12B-it-8bit` | `deepseek-ai/dspark_gemma4_12b_block7` | ~32 GB+ |
+| `qwen3`  | `Qwen3-4B-8bit`        | `deepseek-ai/dspark_qwen3_4b_block7`   | ~16 GB |
+
+## How it works
+
+- A **parallel backbone** (5 Gemma-4 layers) consumes the target model's hidden states
+  (tapped at layers `[5,17,29,41,46]`, EAGLE3-style) and proposes a 7-token block at once.
+- A **rank-256 Markov head** adds a previous-token correction, sampled sequentially — the only
+  sequential cost, which kills "suffix decay" cheaply.
+- A **confidence head** scores each draft position (optional adaptive block length).
+- The target **verifies** every token, so output is **greedy-correct by construction**
+  (identical to plain greedy decoding, up to floating-point tie-breaking on near-ties).
+
+The drafter is loaded 1:1 from the HF checkpoint and **quantized to 4-bit** by default
+(~1.8 GB) so it's cheap to run every round.
+
+## Install
+
+```bash
+uv venv --python 3.12
+source .venv/bin/activate
+uv pip install -e .
+```
+
+## Use
+
+```bash
+# CLI — pick a family (downloads drafter + instruct target on first run)
+python -m mlx_dspark --family qwen3  --prompt "Explain how rainbows form."
+python -m mlx_dspark --family gemma4 --prompt "Explain how rainbows form." --max-new-tokens 256
+
+# side-by-side demo: baseline (plain target) vs dspark (record each, stack)
+python -m mlx_dspark --family qwen3 --mode baseline --prompt "..." --max-new-tokens 400
+python -m mlx_dspark --family qwen3 --mode dspark   --prompt "..." --max-new-tokens 400
+```
+
+```python
+from mlx_dspark import load_pair, speculative_generate
+
+target, tok, drafter, cfg = load_pair("qwen3")   # or "gemma4"
+res = speculative_generate(target, tok, drafter, "Explain how rainbows form.")
+print(res.text, res.mean_accept_len, res.tokens_per_sec)
+```
+
+## Results (M4 Pro, 48 GB; 8-bit instruct target, 4-bit drafter; warm — `python benchmark.py`)
+
+| family | drafter `d_0` | accept len | greedy (baseline) | dspark (this project) | speedup |
+|---|---|---|---|---|---|
+| **Gemma-4 12B** | ~82% | 2.5–3.6 | ~17 tok/s | ~28 tok/s | **~1.5–1.6×** |
+| **Qwen3-4B**    | ~85% | 2.1–2.8 | ~49 tok/s | ~66 tok/s | **~1.3–1.4×** |
+
+"greedy" = the plain target model decoding one token per forward (no drafter); "dspark" =
+speculative decoding with the DSpark drafter. Both produce **identical** output — DSpark is just
+faster (it diverges from sequential greedy only at logit-margin≈0 ties). Smaller/faster targets (Qwen3-4B) have a lower
+per-token verify cost, so the optimal `--max-draft` is smaller (~2).
+
+### Target choice matters
+
+The drafter is trained against a specific **instruct** model in **bf16** — use the matching
+instruct target (`gemma-4-12B-it` / `Qwen3-4B`), not the base model (base gave `d_0` ~47% vs
+82%). Higher precision raises acceptance; 8-bit is the sweet spot. `-bf16` maximizes acceptance;
+4-bit verifies faster.
+
+### Tuning
+
+On Apple Silicon the target verify cost grows with the number of tokens verified, so the
+optimum is to verify ~= the acceptance length: `--max-draft 4` (default). `--max-draft <block>`
+(full 7) is faithful but *slower* on M-series. `--confidence-threshold 0.6` instead truncates
+the block adaptively via the drafter's confidence head.
+
+## License
+
+MIT — see [`LICENSE`](LICENSE). This is an independent MLX port of the inference path of
+DeepSeek's DSpark drafter; see [`NOTICE`](NOTICE) for attribution. No model weights are bundled.

mlx_dspark-0.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+mlx_dspark/__init__.py,sha256=fclV1VbqP6UzILCu35xa6hTDYV_BLzeH9Fr_32hbDwk,1156
+mlx_dspark/__main__.py,sha256=YEfk3teljcUJrc17eMMz7WIILwGFvnEB4mNGmQqnKpQ,3417
+mlx_dspark/config.py,sha256=Ey-tRJ3w6rAXSbe6LhuVHiyLdgbqNvZDGW8blq7WKS8,6257
+mlx_dspark/generate.py,sha256=sOdrRg2GA8so3JH9MBKAMaJ3q6BYuop7lxTmdTs7xkc,10146
+mlx_dspark/load.py,sha256=xLr0QLCFzQa7Qw0LkubT0Jpz0Mt58bvMtjdRBQxBAeQ,3946
+mlx_dspark/model.py,sha256=mnI6Y_ZWZbeohF4FGsLRcvPFamUoy16_0gVFlIkDZew,10607
+mlx_dspark/target.py,sha256=8SlmnMM4UO1qL36HZwaqGfniIEDoupU1T5G_wcKEpXw,2494
+mlx_dspark-0.0.1.dist-info/METADATA,sha256=Wea6DUzWjUOBKFkwXY7D31Ss6TqYBUam6mXXIXfUxZc,5047
+mlx_dspark-0.0.1.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+mlx_dspark-0.0.1.dist-info/licenses/LICENSE,sha256=7fTQ89BHhdRKOLqXpVQFZxTnwQUUtg8EeamGIrOz7eU,1064
+mlx_dspark-0.0.1.dist-info/licenses/NOTICE,sha256=4sSSawyNMWFGbhUevnc0ZKPgOS3auWOdm77K_BqLRcs,908
+mlx_dspark-0.0.1.dist-info/RECORD,,

mlx_dspark-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

mlx_dspark-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 erahim3
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_dspark-0.0.1.dist-info/licenses/NOTICE

@@ -0,0 +1,22 @@
+mlx-dspark
+==========
+
+This project is an independent MLX/Apple-Silicon port of the *inference path* of
+DSpark, a speculative-decoding drafter open-sourced by DeepSeek as part of the
+DeepSpec codebase:
+
+  - DeepSpec: https://github.com/deepseek-ai/DeepSpec  (MIT License)
+
+It loads the published DSpark drafter checkpoints (also released by DeepSeek under
+their respective licenses):
+
+  - deepseek-ai/dspark_gemma4_12b_block7
+  - deepseek-ai/dspark_qwen3_4b_block7
+
+The drafter architecture, training, and checkpoints are the work of DeepSeek and
+the DeepSpec authors. This repository reimplements only the forward/verification
+path for MLX and contains no DeepSpec source code. Target models (Gemma-4, Qwen3)
+are downloaded at runtime from their respective publishers and are subject to
+their own licenses (e.g. the Gemma Terms of Use and the Qwen license).
+
+No model weights are bundled with this package.