mlx-dspark 0.0.1__tar.gz

mlx_dspark-0.0.1/.gitignore

@@ -0,0 +1,17 @@
+.venv/
+__pycache__/
+*.pyc
+*.egg-info/
+build/
+dist/
+.ruff_cache/
+.pytest_cache/
+# local model/drafter weights
+checkpoints/
+*.safetensors
+scratch/
+# local Claude skill copy (not part of the package)
+.claude/
+# private dev/agent tracking notes (kept local, not published)
+CLAUDE.md
+NOTES.md

mlx_dspark-0.0.1/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/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.

mlx_dspark-0.0.1/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,87 @@
+# 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/benchmark.py

@@ -0,0 +1,52 @@
+"""Warm benchmark: greedy vs DSpark speculative (both warm), multi-trial.
+
+  python benchmark.py            # gemma4
+  python benchmark.py qwen3
+"""
+import sys, time, mlx.core as mx
+from mlx_dspark.load import load_pair
+from mlx_dspark.generate import (
+    speculative_generate, _make_target_cache, eos_token_ids, encode_prompt,
+)
+
+family = sys.argv[1] if len(sys.argv) > 1 else "gemma4"
+target, tok, drafter, cfg = load_pair(family)
+eos = eos_token_ids(tok)
+N = 100
+PROMPTS = [
+    "Explain how rainbows form.",
+    "Write a Python function to check if a string is a palindrome.",
+    "Give three tips for staying focused while working.",
+]
+
+def greedy(ids, n):
+    cache = _make_target_cache(target)
+    lg = target.plain(mx.array([ids]), cache); mx.eval(lg)
+    nx = int(mx.argmax(lg[0, -1]).item()); out = [nx]; t = time.time()
+    while len(out) < n and nx not in eos:
+        lg = target.plain(mx.array([[nx]]), cache); mx.eval(lg)
+        nx = int(mx.argmax(lg[0, -1]).item()); out.append(nx)
+    return len(out), time.time() - t
+
+print(f"[{family}] warming up (ramping clocks)...")
+for _ in range(2):
+    greedy(encode_prompt(tok, "Tell me about the sea.", True), 120)
+    speculative_generate(target, tok, drafter, "Tell me about the sea.",
+                         max_new_tokens=120, max_draft_tokens=3)
+
+print(f"\n{'prompt':<6} {'greedy':>9} | {'cap':>3} {'spec':>9} {'accept':>7} {'speedup':>8}")
+agg = {}
+for i, p in enumerate(PROMPTS):
+    ids = encode_prompt(tok, p, True)
+    gt, gs = greedy(ids, N); gtps = gt / gs
+    for cap in (2, 3, 4):
+        res = speculative_generate(target, tok, drafter, p, max_new_tokens=N,
+                                   max_draft_tokens=cap)
+        sp = gs / res.seconds
+        agg.setdefault(cap, []).append(sp)
+        tag = f"P{i}" if cap == 2 else ""
+        print(f"{tag:<6} {gtps:>7.1f}/s | {cap:>3} {res.tokens_per_sec:>7.1f}/s "
+              f"{res.mean_accept_len:>7.2f} {sp:>7.2f}x")
+print("\nmean speedup by cap:")
+for cap, v in agg.items():
+    print(f"  cap={cap}: {sum(v)/len(v):.2f}x")

mlx_dspark-0.0.1/pyproject.toml

@@ -0,0 +1,42 @@
+[project]
+name = "mlx-dspark"
+version = "0.0.1"
+description = "DSpark speculative decoding (semi-autoregressive drafting) for Apple Silicon via MLX"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [{ name = "erahim3", email = "erahim3@gmail.com" }]
+keywords = ["mlx", "apple-silicon", "speculative-decoding", "dspark", "deepspec", "llm-inference"]
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Operating System :: MacOS",
+    "Environment :: GPU",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "mlx>=0.31.2",
+    "mlx-lm>=0.31.3",
+    # mlx-vlm gives us Gemma-4 + the existing EAGLE3/DFlash spec-decode reference impl
+    "mlx-vlm>=0.6.3",
+    "numpy",
+    "huggingface-hub",
+]
+
+[project.optional-dependencies]
+dev = ["pytest", "ruff"]
+
+[project.urls]
+Homepage = "https://github.com/ARahim3/mlx-dspark"
+Repository = "https://github.com/ARahim3/mlx-dspark"
+Issues = "https://github.com/ARahim3/mlx-dspark/issues"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/mlx_dspark"]
+
+[tool.ruff]
+line-length = 100

mlx_dspark-0.0.1/src/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-0.0.1/src/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-0.0.1/src/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,
+        )