dspark-mlx 0.1.0__tar.gz

dspark_mlx-0.1.0/LICENSE

@@ -0,0 +1,182 @@
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and
+distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright
+owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities
+that control, are controlled by, or are under common control with that entity.
+For the purposes of this definition, "control" means (i) the power, direct or
+indirect, to cause the direction or management of such entity, whether by
+contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the
+outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising
+permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including
+but not limited to software source code, documentation source, and configuration
+files.
+
+"Object" form shall mean any form resulting from mechanical transformation or
+translation of a Source form, including but not limited to compiled object code,
+generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form,
+made available under the License, as indicated by a copyright notice that is
+included in or attached to the work (an example is provided in the Appendix
+below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that
+is based on (or derived from) the Work and for which the editorial revisions,
+annotations, elaborations, or other modifications represent, as a whole, an
+original work of authorship. For the purposes of this License, Derivative Works
+shall not include works that remain separable from, or merely link (or bind by
+name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version
+of the Work and any modifications or additions to that Work or Derivative Works
+thereof, that is intentionally submitted to Licensor for inclusion in the Work by
+the copyright owner or by an individual or Legal Entity authorized to submit on
+behalf of the copyright owner. For the purposes of this definition, "submitted"
+means any form of electronic, verbal, or written communication sent to the
+Licensor or its representatives, including but not limited to communication on
+electronic mailing lists, source code control systems, and issue tracking
+systems that are managed by, or on behalf of, the Licensor for the purpose of
+discussing and improving the Work, but excluding communication that is
+conspicuously marked or otherwise designated in writing by the copyright owner
+as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf
+of whom a Contribution has been received by Licensor and subsequently
+incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this
+License, each Contributor hereby grants to You a perpetual, worldwide,
+non-exclusive, no-charge, royalty-free, irrevocable copyright license to
+reproduce, prepare Derivative Works of, publicly display, publicly perform,
+sublicense, and distribute the Work and such Derivative Works in Source or
+Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this License,
+each Contributor hereby grants to You a perpetual, worldwide, non-exclusive,
+no-charge, royalty-free, irrevocable (except as stated in this section) patent
+license to make, have made, use, offer to sell, sell, import, and otherwise
+transfer the Work, where such license applies only to those patent claims
+licensable by such Contributor that are necessarily infringed by their
+Contribution(s) alone or by combination of their Contribution(s) with the Work to
+which such Contribution(s) was submitted. If You institute patent litigation
+against any entity (including a cross-claim or counterclaim in a lawsuit)
+alleging that the Work or a Contribution incorporated within the Work
+constitutes direct or contributory patent infringement, then any patent licenses
+granted to You under this License for that Work shall terminate as of the date
+such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the Work or
+Derivative Works thereof in any medium, with or without modifications, and in
+Source or Object form, provided that You meet the following conditions:
+
+(a) You must give any other recipients of the Work or Derivative Works a copy of
+this License; and
+
+(b) You must cause any modified files to carry prominent notices stating that You
+changed the files; and
+
+(c) You must retain, in the Source form of any Derivative Works that You
+distribute, all copyright, patent, trademark, and attribution notices from the
+Source form of the Work, excluding those notices that do not pertain to any part
+of the Derivative Works; and
+
+(d) If the Work includes a "NOTICE" text file as part of its distribution, then
+any Derivative Works that You distribute must include a readable copy of the
+attribution notices contained within such NOTICE file, excluding those notices
+that do not pertain to any part of the Derivative Works, in at least one of the
+following places: within a NOTICE text file distributed as part of the Derivative
+Works; within the Source form or documentation, if provided along with the
+Derivative Works; or, within a display generated by the Derivative Works, if and
+wherever such third-party notices normally appear. The contents of the NOTICE
+file are for informational purposes only and do not modify the License. You may
+add Your own attribution notices within Derivative Works that You distribute,
+alongside or as an addendum to the NOTICE text from the Work, provided that such
+additional attribution notices cannot be construed as modifying the License.
+
+You may add Your own copyright statement to Your modifications and may provide
+additional or different license terms and conditions for use, reproduction, or
+distribution of Your modifications, or for any such Derivative Works as a whole,
+provided Your use, reproduction, and distribution of the Work otherwise complies
+with the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise, any
+Contribution intentionally submitted for inclusion in the Work by You to the
+Licensor shall be under the terms and conditions of this License, without any
+additional terms or conditions. Notwithstanding the above, nothing herein shall
+supersede or modify the terms of any separate license agreement you may have
+executed with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade names,
+trademarks, service marks, or product names of the Licensor, except as required
+for reasonable and customary use in describing the origin of the Work and
+reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or agreed to in
+writing, Licensor provides the Work (and each Contributor provides its
+Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied, including, without limitation, any warranties or
+conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+PARTICULAR PURPOSE. You are solely responsible for determining the
+appropriateness of using or redistributing the Work and assume any risks
+associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory, whether in
+tort (including negligence), contract, or otherwise, unless required by
+applicable law (such as deliberate and grossly negligent acts) or agreed to in
+writing, shall any Contributor be liable to You for damages, including any
+direct, indirect, special, incidental, or consequential damages of any character
+arising as a result of this License or out of the use or inability to use the
+Work (including but not limited to damages for loss of goodwill, work stoppage,
+computer failure or malfunction, or any and all other commercial damages or
+losses), even if such Contributor has been advised of the possibility of such
+damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing the Work or
+Derivative Works thereof, You may choose to offer, and charge a fee for,
+acceptance of support, warranty, indemnity, or other liability obligations
+and/or rights consistent with this License. However, in accepting such
+obligations, You may act only on Your own behalf and on Your sole
+responsibility, not on behalf of any other Contributor, and only if You agree to
+indemnify, defend, and hold each Contributor harmless for any liability incurred
+by, or claims asserted against, such Contributor by reason of your accepting any
+such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+To apply the Apache License to your work, attach the following boilerplate
+notice, with the fields enclosed by brackets "[]" replaced with your own
+identifying information. (Don't include the brackets!) The text should be
+enclosed in the appropriate comment syntax for the file format. We also
+recommend that a file or class name and description of purpose be included on
+the same "printed page" as the copyright notice for easier identification within
+third-party archives.
+
+Copyright 2026 bstnxbt
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use
+this file except in compliance with the License. You may obtain a copy of the
+License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed
+under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.

dspark_mlx-0.1.0/PKG-INFO

@@ -0,0 +1,108 @@
+Metadata-Version: 2.4
+Name: dspark-mlx
+Version: 0.1.0
+Summary: DSpark self-speculative decoding for Apple Silicon (MLX) — lossless, target-agnostic
+Author-email: popfido <wanghailin317@gmail.com>
+Maintainer-email: popfido <wanghailin317@gmail.com>
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/popfido/dspark-mlx
+Project-URL: Issues, https://github.com/popfido/dspark-mlx/issues
+Keywords: mlx,apple-silicon,speculative-decoding,dspark,llm-inference,deepseek,qwen3,gemma
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Operating System :: MacOS
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlx>=0.31.0
+Requires-Dist: mlx-lm>=0.31.0
+Requires-Dist: mlx-vlm>=0.6.0
+Requires-Dist: numpy
+Provides-Extra: bench
+Requires-Dist: datasets; extra == "bench"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: torch; extra == "dev"
+Requires-Dist: numpy; extra == "dev"
+Requires-Dist: datasets; extra == "dev"
+Dynamic: license-file
+
+# dspark-mlx
+
+Target-agnostic MLX implementation of DeepSeek **DSpark** self-speculative decoding.
+
+DSpark drafts a block of tokens from a small EAGLE-style draft model (projected target
+hidden states + a low-rank Markov logit bias + a per-token confidence head), which the
+host base model then verifies **losslessly**. This package owns the DSpark draft stack and
+the verify/accept policy; the base model is supplied by the host through a small adapter
+(`dspark_mlx/adapter.py`). The emitted stream is identical to greedy decoding from the base
+model alone.
+
+## Install
+
+```bash
+pip install dspark-mlx          # Qwen3 (mlx-lm) and Gemma (mlx-vlm) both supported out of the box
+```
+
+```bash
+dspark generate --model qwen3-4b "Explain speculative decoding." --quant-draft 8 --no-think
+dspark bench    --model qwen3-4b --quant-draft 8 --no-think
+dspark eval     --model qwen3-4b --dataset gsm8k --n 20
+```
+
+The package itself is tiny; `dspark` downloads the DSpark draft + the deployed **instruct** base
+on first use (some bases — e.g. `google/gemma-4-12b-it` — are gated and large, ~24 GB).
+
+## Architectures
+
+One DSpark recipe, three base-model backbones — selected by `model_type` via the registry
+(`dspark_mlx/registry.py`), mirroring `dflash-mlx`'s `TARGET_BACKENDS`:
+
+| Backbone | Checkpoint | Draft layer body |
+|---|---|---|
+| `deepseek_v4` | `DeepSeek-V4-Flash-DSpark` (bundled fp8/fp4, `mtp.*`) | MLA + hash-MoE + Hyper-Connections + windowed sparse attn |
+| `qwen3` | `dspark_qwen3_{4b,8b,14b}_block7` (standalone bf16, `layers.*`) | Qwen3 GQA + QK-norm + SwiGLU |
+| `gemma4` | `dspark_gemma4_12b_block7` (standalone bf16, `layers.*`) | Gemma4 GQA (K=V) + sandwich norms + GeGLU + partial RoPE + softcap |
+
+Add an architecture: implement a `DraftArch` (build + key_map) in `dspark_mlx/arch/<name>.py`
+and append it to `ARCH_REGISTRY` — `generate`/`verify`/`adapter` are unchanged.
+
+```python
+from dspark_mlx import resolve_arch, load_drafter, generate
+
+arch = resolve_arch(config)                      # by config["model_type"]
+drafter = arch.build(config, max_seq_len=...)
+load_drafter(drafter, weights, key_map=arch.key_map)
+for event in generate(adapter, drafter, prompt_tokens, max_new_tokens):
+    ...
+```
+
+## Benchmarks
+
+Lossless speedup on Apple Silicon (eager loop, greedy) and average **accepted length**
+(τ = tokens per verify step) vs the DSpark paper — the hardware-independent metric the paper
+reports. Both backbones land in the paper's acceptance band. Full methodology, all precisions,
+and findings in **[BENCHMARK.md](https://github.com/popfido/dspark-mlx/blob/main/BENCHMARK.md)**.
+
+| model | accepted length (GSM8K) | acceptance rate | speedup |
+|---|---|---|---|
+| Qwen3-4B (bf16) | 3.86 / 6.27† | 41% / 75%† | 1.17× |
+| Qwen3-14B (bf16) | 3.79 / 6.29† | 40% / 76%† | **1.98×**† |
+| Gemma4-12B-it (bf16) | **5.84** | **69%** | **2.06×** |
+
+†thinking off (`--no-think`). Two knobs: **acceptance length** is set by the draft + task
+(≈ size-independent — 4B ≈ 14B), while **speedup** is set by how expensive the base is (a
+costlier base amortizes the draft better, so Qwen3-4B 1.17× → 14B 1.98× → Gemma-12B 2.06×).
+**Quantize the draft** (`--quant-draft 8`, acceptance-lossless, ½ the draft size) for a free
+~10–20% on top of any base — Qwen3-4B bf16 1.26× → **1.67×**, 8-bit base 1.34× → **1.62×**,
+Qwen3-14B 2.07× → **2.36×** (see BENCHMARK.md for the base×draft-precision table). The draft must
+target the deployed **instruct** model (the pretrained Gemma base gives ~3× lower acceptance);
+Qwen3's `<think>` traces roughly halve acceptance vs `--no-think`.
+
+Based on `deepseek-ai/DeepSeek-V4-Flash-DSpark` and the DeepSpec codebase (`dspark/*`).
+Repo structure mirrors `dflash-mlx`.
+
+Status: the full draft→verify→accept pipeline is verified against the reference for all
+three backbones (parity on tiny weights; lossless end-to-end against toy bases; real Qwen3
+/ Gemma4 checkpoints load and run). See the test suite.

dspark_mlx-0.1.0/README.md

@@ -0,0 +1,78 @@
+# dspark-mlx
+
+Target-agnostic MLX implementation of DeepSeek **DSpark** self-speculative decoding.
+
+DSpark drafts a block of tokens from a small EAGLE-style draft model (projected target
+hidden states + a low-rank Markov logit bias + a per-token confidence head), which the
+host base model then verifies **losslessly**. This package owns the DSpark draft stack and
+the verify/accept policy; the base model is supplied by the host through a small adapter
+(`dspark_mlx/adapter.py`). The emitted stream is identical to greedy decoding from the base
+model alone.
+
+## Install
+
+```bash
+pip install dspark-mlx          # Qwen3 (mlx-lm) and Gemma (mlx-vlm) both supported out of the box
+```
+
+```bash
+dspark generate --model qwen3-4b "Explain speculative decoding." --quant-draft 8 --no-think
+dspark bench    --model qwen3-4b --quant-draft 8 --no-think
+dspark eval     --model qwen3-4b --dataset gsm8k --n 20
+```
+
+The package itself is tiny; `dspark` downloads the DSpark draft + the deployed **instruct** base
+on first use (some bases — e.g. `google/gemma-4-12b-it` — are gated and large, ~24 GB).
+
+## Architectures
+
+One DSpark recipe, three base-model backbones — selected by `model_type` via the registry
+(`dspark_mlx/registry.py`), mirroring `dflash-mlx`'s `TARGET_BACKENDS`:
+
+| Backbone | Checkpoint | Draft layer body |
+|---|---|---|
+| `deepseek_v4` | `DeepSeek-V4-Flash-DSpark` (bundled fp8/fp4, `mtp.*`) | MLA + hash-MoE + Hyper-Connections + windowed sparse attn |
+| `qwen3` | `dspark_qwen3_{4b,8b,14b}_block7` (standalone bf16, `layers.*`) | Qwen3 GQA + QK-norm + SwiGLU |
+| `gemma4` | `dspark_gemma4_12b_block7` (standalone bf16, `layers.*`) | Gemma4 GQA (K=V) + sandwich norms + GeGLU + partial RoPE + softcap |
+
+Add an architecture: implement a `DraftArch` (build + key_map) in `dspark_mlx/arch/<name>.py`
+and append it to `ARCH_REGISTRY` — `generate`/`verify`/`adapter` are unchanged.
+
+```python
+from dspark_mlx import resolve_arch, load_drafter, generate
+
+arch = resolve_arch(config)                      # by config["model_type"]
+drafter = arch.build(config, max_seq_len=...)
+load_drafter(drafter, weights, key_map=arch.key_map)
+for event in generate(adapter, drafter, prompt_tokens, max_new_tokens):
+    ...
+```
+
+## Benchmarks
+
+Lossless speedup on Apple Silicon (eager loop, greedy) and average **accepted length**
+(τ = tokens per verify step) vs the DSpark paper — the hardware-independent metric the paper
+reports. Both backbones land in the paper's acceptance band. Full methodology, all precisions,
+and findings in **[BENCHMARK.md](https://github.com/popfido/dspark-mlx/blob/main/BENCHMARK.md)**.
+
+| model | accepted length (GSM8K) | acceptance rate | speedup |
+|---|---|---|---|
+| Qwen3-4B (bf16) | 3.86 / 6.27† | 41% / 75%† | 1.17× |
+| Qwen3-14B (bf16) | 3.79 / 6.29† | 40% / 76%† | **1.98×**† |
+| Gemma4-12B-it (bf16) | **5.84** | **69%** | **2.06×** |
+
+†thinking off (`--no-think`). Two knobs: **acceptance length** is set by the draft + task
+(≈ size-independent — 4B ≈ 14B), while **speedup** is set by how expensive the base is (a
+costlier base amortizes the draft better, so Qwen3-4B 1.17× → 14B 1.98× → Gemma-12B 2.06×).
+**Quantize the draft** (`--quant-draft 8`, acceptance-lossless, ½ the draft size) for a free
+~10–20% on top of any base — Qwen3-4B bf16 1.26× → **1.67×**, 8-bit base 1.34× → **1.62×**,
+Qwen3-14B 2.07× → **2.36×** (see BENCHMARK.md for the base×draft-precision table). The draft must
+target the deployed **instruct** model (the pretrained Gemma base gives ~3× lower acceptance);
+Qwen3's `<think>` traces roughly halve acceptance vs `--no-think`.
+
+Based on `deepseek-ai/DeepSeek-V4-Flash-DSpark` and the DeepSpec codebase (`dspark/*`).
+Repo structure mirrors `dflash-mlx`.
+
+Status: the full draft→verify→accept pipeline is verified against the reference for all
+three backbones (parity on tiny weights; lossless end-to-end against toy bases; real Qwen3
+/ Gemma4 checkpoints load and run). See the test suite.

dspark_mlx-0.1.0/dspark_mlx/__init__.py

@@ -0,0 +1,45 @@
+# Copyright 2026 popfido
+# Licensed under the Apache License, Version 2.0 - see LICENSE file
+# Based on DeepSeek DSpark (DeepSeek-V4-Flash-DSpark, deepseek-ai/DeepSpec)
+
+__version__ = "0.1.0"
+
+from .adapter import BaseModelAdapter, BlockOut, StepOut
+from .arch.backbone import DraftArch, DraftBackbone
+from .events import SummaryEvent, TokenEvent
+from .generate import generate
+from .loader import KNOWN_MODELS, load_draft, load_host, resolve_model
+from .loading import is_dspark_checkpoint, load_drafter, map_checkpoint_key
+from .loop import generate_eager
+from .model.config import DSparkArgs
+from .model.drafter import DSparkDrafter
+from .quant import quantize_drafter
+from .registry import ARCH_REGISTRY, resolve_arch
+from .verify import AcceptResult, greedy_accept, speculative_sample_accept
+
+__all__ = [
+    "BaseModelAdapter",
+    "BlockOut",
+    "StepOut",
+    "DSparkArgs",
+    "DSparkDrafter",
+    "DraftArch",
+    "DraftBackbone",
+    "resolve_arch",
+    "ARCH_REGISTRY",
+    "generate",
+    "generate_eager",
+    "load_draft",
+    "load_host",
+    "resolve_model",
+    "KNOWN_MODELS",
+    "greedy_accept",
+    "speculative_sample_accept",
+    "AcceptResult",
+    "TokenEvent",
+    "SummaryEvent",
+    "load_drafter",
+    "map_checkpoint_key",
+    "is_dspark_checkpoint",
+    "quantize_drafter",
+]

dspark_mlx-0.1.0/dspark_mlx/adapter.py

@@ -0,0 +1,78 @@
+# Copyright 2026 popfido
+# Licensed under the Apache License, Version 2.0 - see LICENSE file
+# Based on DeepSeek DSpark (DeepSeek-V4-Flash-DSpark, deepseek-ai/DeepSpec)
+
+"""The seam between dspark-mlx (drafter + verify/accept loop) and a host base model.
+
+dspark-mlx is target-agnostic: it owns the DSpark draft stack and the lossless
+accept policy, but never the base model. The host (e.g. omlx over its
+``patches/deepseek_v4`` model) implements :class:`BaseModelAdapter` so the drafter
+can (a) read the ``main_hidden`` it conditions on, (b) get the base distribution for
+each candidate token during verify, and (c) snapshot/roll back base KV when a block is
+only partially accepted.
+
+Logit conventions (one decode cycle):
+- ``prefill`` / ``decode_step`` return ``StepOut.logits`` = ``p_1``, the base
+  distribution for the *first* drafted token. It is free — already computed by the
+  step that produced the anchor — so the verify forward never recomputes it.
+- ``verify_forward`` runs ONE base forward over the K drafted tokens and returns the
+  K base distributions ``p_2 .. p_{K+1}`` (``p_{K+1}`` is the bonus position).
+- The generate loop concatenates ``[p_1] + [p_2..p_{K+1}]`` into the ``[K+1, V]`` block
+  the accept policy consumes (see :mod:`dspark_mlx.verify`).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol, Tuple, runtime_checkable
+
+import mlx.core as mx
+
+
+@dataclass
+class StepOut:
+    """Output of a single base forward at the anchor position."""
+
+    logits: mx.array        # [b, V] base distribution for the next (first drafted) token
+    main_hidden: mx.array   # [b, dim * len(target_layer_ids)] concat of target-layer hiddens
+
+
+@dataclass
+class BlockOut:
+    """Output of the base verify forward over a K-token draft block."""
+
+    per_pos_logits: mx.array       # [b, K, V] base distributions p_2 .. p_{K+1}
+    per_pos_main_hidden: mx.array  # [b, K, D] main hidden at each verified position
+    main_hidden_last: mx.array     # [b, D] convenience alias for the last verified position
+
+
+@runtime_checkable
+class BaseModelAdapter(Protocol):
+    """Host contract. Implementations own the base model and its KV cache."""
+
+    #: Main-model layer indices whose hidden states are concatenated into ``main_hidden``.
+    target_layer_ids: Tuple[int, ...]
+
+    def prefill(self, tokens: mx.array) -> StepOut:
+        """Process the prompt; return logits for the first generated token + main_hidden."""
+        ...
+
+    def decode_step(self, token: mx.array) -> StepOut:
+        """Advance one token; return its next-token logits + main_hidden."""
+        ...
+
+    def verify_forward(self, block_tokens: mx.array) -> BlockOut:
+        """Run one base forward over K draft tokens; return p_2..p_{K+1} + main_hidden_last.
+
+        Appends K entries to the base KV cache speculatively; the caller rolls back the
+        rejected tail via :meth:`kv_rollback`.
+        """
+        ...
+
+    def kv_snapshot(self) -> Any:
+        """Opaque handle capturing base KV state before a speculative block."""
+        ...
+
+    def kv_rollback(self, n_keep: int) -> None:
+        """Drop speculatively-appended KV beyond ``n_keep`` accepted tokens."""
+        ...

dspark_mlx-0.1.0/dspark_mlx/arch/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2026 popfido
+# Licensed under the Apache License, Version 2.0 - see LICENSE file
+
+from .backbone import DraftArch, DraftBackbone
+
+__all__ = ["DraftArch", "DraftBackbone"]

dspark_mlx-0.1.0/dspark_mlx/arch/backbone.py

@@ -0,0 +1,57 @@
+# Copyright 2026 popfido
+# Licensed under the Apache License, Version 2.0 - see LICENSE file
+# Based on DeepSeek DSpark (DeepSeek-V4-Flash-DSpark, deepseek-ai/DeepSpec)
+
+"""The per-architecture seam for DSpark drafters.
+
+DSpark ships one recipe (EAGLE-style context projection + Markov bias + confidence head +
+block drafting) realized over different base-model decoder layers — DeepSeek-V4 (windowed
+MLA + MoE + Hyper-Connections, bundled fp8/fp4 ``mtp.*`` checkpoint), Qwen3 and Gemma4
+(standalone bf16 ``layers.*`` checkpoints, full-context GQA). ``generate()`` drives any of
+them through the ``DraftBackbone`` interface; a ``DraftArch`` descriptor registers how to
+build and load each one (see :mod:`dspark_mlx.registry`).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Protocol, Tuple, runtime_checkable
+
+import mlx.core as mx
+
+
+@runtime_checkable
+class DraftBackbone(Protocol):
+    """A loaded DSpark drafter for one base architecture (what ``generate`` consumes)."""
+
+    block_size: int
+
+    def forward_spec(
+        self, input_ids: mx.array, main_hidden: mx.array, start_pos: int = 0
+    ) -> Optional[Tuple[mx.array, mx.array, mx.array]]:
+        """Prefill (start_pos==0) seeds context; decode drafts (ids, logits, confidence)."""
+        ...
+
+    def advance(self, main_hidden: mx.array, position: int) -> None:
+        """Slide the drafter's context over one committed token."""
+        ...
+
+
+@dataclass(frozen=True)
+class DraftArch:
+    """Registry entry: how to build + load a DSpark drafter for a base architecture."""
+
+    name: str
+    model_types: Tuple[str, ...]
+    build: Callable[..., DraftBackbone]  # (config: dict, *, max_seq_len) -> DraftBackbone
+    key_map: Callable[[str], Optional[str]]  # checkpoint key -> drafter param path (or None)
+
+    def supports(self, model_type: Optional[str]) -> bool:
+        return model_type in self.model_types
+
+
+def config_model_type(config: Any) -> Optional[str]:
+    """Read ``model_type`` from a dict-like or attribute-like config."""
+    if isinstance(config, dict):
+        return config.get("model_type")
+    return getattr(config, "model_type", None)

dspark_mlx-0.1.0/dspark_mlx/arch/deepseek_v4.py

@@ -0,0 +1,35 @@
+# Copyright 2026 popfido
+# Licensed under the Apache License, Version 2.0 - see LICENSE file
+# Based on DeepSeek DSpark (DeepSeek-V4-Flash-DSpark, deepseek-ai/DeepSpec)
+
+"""DeepSeek-V4-Flash-DSpark backbone descriptor.
+
+The windowed MLA + hash-MoE + Hyper-Connections realization, drafting from the ``mtp.*``
+namespace of the bundled fp8/fp4 checkpoint. The model code lives under ``dspark_mlx.model``
+(its parity tests pin it); this module just registers it as a DraftArch.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from ..loading import map_checkpoint_key
+from ..model.config import DSparkArgs
+from ..model.drafter import DSparkDrafter
+from .backbone import DraftArch, DraftBackbone
+
+
+def build(config: dict, *, max_seq_len: int = 8192) -> DraftBackbone:
+    return DSparkDrafter(DSparkArgs.from_dict(config), max_seq_len=max_seq_len)
+
+
+def key_map(key: str) -> Optional[str]:
+    return map_checkpoint_key(key)  # mtp.N.* -> blocks.N.*, embed/head pass through
+
+
+DEEPSEEK_V4 = DraftArch(
+    name="deepseek_v4",
+    model_types=("deepseek_v4",),
+    build=build,
+    key_map=key_map,
+)