dspark-mlx 0.1.0__py3-none-any.whl

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

dspark_mlx/arch/gemma4.py

@@ -0,0 +1,327 @@
+# Copyright 2026 popfido
+# Licensed under the Apache License, Version 2.0 - see LICENSE file
+# Based on DeepSeek DSpark (deepseek-ai/DeepSpec: dspark/gemma4/modeling.py)
+
+"""Gemma4 DSpark draft backbone (standalone bf16 ``layers.*`` checkpoint).
+
+Stock Gemma4 decoder layers with the DSpark context/noise K/V split. Gemma deltas vs Qwen3:
+K=V sharing (no v_proj; separate scaled k_norm + weightless v_norm), attention scale 1.0,
+partial (proportional) RoPE — only ``partial_rotary_factor`` of head_dim rotates, the rest
+pass through — four sandwich norms + a per-layer ``layer_scalar``, GeGLU (gelu-tanh) MLP, and
+final-logit softcapping.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import re as _re
+from dataclasses import dataclass
+from typing import Mapping, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..model.heads import DSparkConfidenceHead, DSparkMarkovHead
+from ..model.norm_rope import RMSNorm
+from ..recipe import draft_block_decode
+from .backbone import DraftArch
+from .qwen3 import _apply_rope  # shared NeoX rotate_half application
+
+_GEMMA4_LAYER_RE = _re.compile(r"layers\.(\d+)\.(.+)$")
+
+
+@dataclass
+class Gemma4DSparkArgs:
+    vocab_size: int = 262144
+    hidden_size: int = 3840
+    num_hidden_layers: int = 5
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 1          # global KV head count (k=v)
+    head_dim: int = 512                   # global_head_dim
+    intermediate_size: int = 15360
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1000000.0
+    partial_rotary_factor: float = 0.25
+    attention_k_eq_v: bool = True
+    final_logit_softcapping: float = 30.0
+    target_layer_ids: Tuple[int, ...] = (5, 17, 29, 41, 46)
+    num_target_layers: int = 48
+    block_size: int = 7
+    mask_token_id: int = 4
+    markov_rank: int = 256
+    temperature: float = 0.0
+    max_position_embeddings: int = 262144
+
+    @property
+    def fc_in(self) -> int:
+        return self.hidden_size * len(self.target_layer_ids)
+
+    @classmethod
+    def from_dict(cls, params: Mapping) -> "Gemma4DSparkArgs":
+        d = dict(params)
+        rope = (d.get("rope_parameters") or {}).get("full_attention") or {}
+        if rope.get("rope_theta"):
+            d["rope_theta"] = rope["rope_theta"]
+        if "partial_rotary_factor" in rope:
+            d["partial_rotary_factor"] = rope["partial_rotary_factor"]
+        if d.get("global_head_dim"):
+            d["head_dim"] = d["global_head_dim"]
+        if d.get("num_global_key_value_heads") is not None:
+            d["num_key_value_heads"] = d["num_global_key_value_heads"]
+        names = {f.name for f in dataclasses.fields(cls)}
+        kwargs = {k: v for k, v in d.items() if k in names}
+        if "target_layer_ids" in kwargs:
+            kwargs["target_layer_ids"] = tuple(kwargs["target_layer_ids"])
+        return cls(**kwargs)
+
+
+def rope_tables(position_ids: mx.array, head_dim: int, theta: float, partial: float) -> Tuple[mx.array, mx.array]:
+    """Proportional (partial) RoPE: first ``partial*head_dim`` dims rotate, rest are identity."""
+    rope_angles = int(partial * head_dim // 2)
+    inv_rot = 1.0 / (theta ** (mx.arange(0, 2 * rope_angles, 2).astype(mx.float32) / head_dim))
+    nope = head_dim // 2 - rope_angles
+    inv_freq = mx.concatenate([inv_rot, mx.zeros((nope,), dtype=mx.float32)]) if nope > 0 else inv_rot
+    freqs = position_ids.astype(mx.float32)[:, None] * inv_freq[None, :]
+    emb = mx.concatenate([freqs, freqs], axis=-1)
+    return mx.cos(emb), mx.sin(emb)
+
+
+class Gemma4DSparkAttention(nn.Module):
+    def __init__(self, args: Gemma4DSparkArgs):
+        super().__init__()
+        h, nh, nkv, hd = args.hidden_size, args.num_attention_heads, args.num_key_value_heads, args.head_dim
+        self.nh, self.nkv, self.hd = nh, nkv, hd
+        self.k_eq_v = args.attention_k_eq_v
+        self.q_proj = nn.Linear(h, nh * hd, bias=False)
+        self.k_proj = nn.Linear(h, nkv * hd, bias=False)
+        self.v_proj = None if self.k_eq_v else nn.Linear(h, nkv * hd, bias=False)
+        self.o_proj = nn.Linear(nh * hd, h, bias=False)
+        self.q_norm = RMSNorm(hd, args.rms_norm_eps)
+        self.k_norm = RMSNorm(hd, args.rms_norm_eps)
+        self.v_norm = RMSNorm(hd, args.rms_norm_eps, with_scale=False)
+
+    def __call__(self, hidden: mx.array, target_ctx: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+        b, q, _ = hidden.shape
+        ctx = target_ctx.shape[1]
+        qh = self.q_norm(self.q_proj(hidden).reshape(b, q, self.nh, self.hd))
+        k_ctx, k_noise = self.k_proj(target_ctx), self.k_proj(hidden)
+        v_ctx, v_noise = (k_ctx, k_noise) if self.k_eq_v else (self.v_proj(target_ctx), self.v_proj(hidden))
+        k = self.k_norm(mx.concatenate([k_ctx, k_noise], axis=1).reshape(b, ctx + q, self.nkv, self.hd))
+        v = self.v_norm(mx.concatenate([v_ctx, v_noise], axis=1).reshape(b, ctx + q, self.nkv, self.hd))
+        qh = _apply_rope(qh, cos[-q:], sin[-q:])
+        k = _apply_rope(k, cos, sin)  # v is not rotated
+        qh = qh.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.transpose(0, 2, 1, 3)
+        out = mx.fast.scaled_dot_product_attention(qh, k, v, scale=1.0, mask=None)  # Gemma4 scale==1
+        out = out.transpose(0, 2, 1, 3).reshape(b, q, self.nh * self.hd)
+        return self.o_proj(out)
+
+    # --- cached path (Phase 3b): context K/V precomputed once, reused every block ---
+    def context_kv(self, proj_ctx: mx.array, cos: mx.array, sin: mx.array):
+        b, n, _ = proj_ctx.shape
+        kc = self.k_proj(proj_ctx)
+        vc = kc if self.k_eq_v else self.v_proj(proj_ctx)
+        k = self.k_norm(kc.reshape(b, n, self.nkv, self.hd))
+        v = self.v_norm(vc.reshape(b, n, self.nkv, self.hd))
+        k = _apply_rope(k, cos, sin)  # v is not rotated
+        return k.transpose(0, 2, 1, 3), v.transpose(0, 2, 1, 3)
+
+    def attend_cached(self, noise, ctx_k, ctx_v, cos, sin):
+        b, q, _ = noise.shape
+        qh = self.q_norm(self.q_proj(noise).reshape(b, q, self.nh, self.hd))
+        kn = self.k_proj(noise)
+        vn = kn if self.k_eq_v else self.v_proj(noise)
+        nk = self.k_norm(kn.reshape(b, q, self.nkv, self.hd))
+        nv = self.v_norm(vn.reshape(b, q, self.nkv, self.hd))
+        qh = _apply_rope(qh, cos, sin).transpose(0, 2, 1, 3)
+        nk = _apply_rope(nk, cos, sin).transpose(0, 2, 1, 3)
+        nv = nv.transpose(0, 2, 1, 3)
+        k = nk if ctx_k is None else mx.concatenate([ctx_k, nk], axis=2)
+        v = nv if ctx_v is None else mx.concatenate([ctx_v, nv], axis=2)
+        out = mx.fast.scaled_dot_product_attention(qh, k, v, scale=1.0, mask=None)
+        out = out.transpose(0, 2, 1, 3).reshape(b, q, self.nh * self.hd)
+        return self.o_proj(out)
+
+
+class Gemma4MLP(nn.Module):
+    def __init__(self, args: Gemma4DSparkArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Gemma4DSparkLayer(nn.Module):
+    def __init__(self, args: Gemma4DSparkArgs):
+        super().__init__()
+        h, eps = args.hidden_size, args.rms_norm_eps
+        self.self_attn = Gemma4DSparkAttention(args)
+        self.mlp = Gemma4MLP(args)
+        self.input_layernorm = RMSNorm(h, eps)
+        self.post_attention_layernorm = RMSNorm(h, eps)
+        self.pre_feedforward_layernorm = RMSNorm(h, eps)
+        self.post_feedforward_layernorm = RMSNorm(h, eps)
+        self.layer_scalar = mx.ones((1,), dtype=mx.float32)
+
+    def __call__(self, hidden: mx.array, target_ctx: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+        h = self.post_attention_layernorm(self.self_attn(self.input_layernorm(hidden), target_ctx, cos, sin))
+        hidden = hidden + h
+        h = self.post_feedforward_layernorm(self.mlp(self.pre_feedforward_layernorm(hidden)))
+        hidden = hidden + h
+        return hidden * self.layer_scalar
+
+    def context_kv(self, proj_ctx: mx.array, cos: mx.array, sin: mx.array):
+        return self.self_attn.context_kv(proj_ctx, cos, sin)
+
+    def forward_cached(self, hidden, ctx_k, ctx_v, cos, sin) -> mx.array:
+        h = self.post_attention_layernorm(self.self_attn.attend_cached(self.input_layernorm(hidden), ctx_k, ctx_v, cos, sin))
+        hidden = hidden + h
+        h = self.post_feedforward_layernorm(self.mlp(self.pre_feedforward_layernorm(hidden)))
+        hidden = hidden + h
+        return hidden * self.layer_scalar
+
+
+class Gemma4Backbone(nn.Module):
+    def __init__(self, args: Gemma4DSparkArgs):
+        super().__init__()
+        self.fc = nn.Linear(args.fc_in, args.hidden_size, bias=False)
+        self.hidden_norm = RMSNorm(args.hidden_size, args.rms_norm_eps)
+        self.layers = [Gemma4DSparkLayer(args) for _ in range(args.num_hidden_layers)]
+        self.norm = RMSNorm(args.hidden_size, args.rms_norm_eps)
+
+    def project_context(self, target_hidden: mx.array) -> mx.array:
+        return self.hidden_norm(self.fc(target_hidden))
+
+    def __call__(self, noise_embed: mx.array, target_ctx: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+        h = noise_embed
+        for layer in self.layers:
+            h = layer(h, target_ctx, cos, sin)
+        return self.norm(h)
+
+
+class Gemma4DSparkDrafter(nn.Module):
+    """Gemma4 DSpark drafter (DraftBackbone). Same loop as Qwen3 + partial RoPE + softcap."""
+
+    def __init__(self, args: Gemma4DSparkArgs, max_seq_len: int = 8192):
+        super().__init__()
+        self.args = args
+        self.block_size = args.block_size
+        self.temperature = args.temperature
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.backbone = Gemma4Backbone(args)
+        self.markov_head = DSparkMarkovHead(args.vocab_size, args.markov_rank)
+        self.confidence_head = DSparkConfidenceHead(args.hidden_size + args.markov_rank, bias=True)
+        self.reset()
+
+    def reset(self) -> None:
+        self._ctx = None
+        self._next_pos = 0
+        self._ctx_k = None      # per-layer cached context K (eager path)
+        self._ctx_v = None
+        self._committed = 0
+
+    def _project_one(self, main_hidden: mx.array) -> mx.array:
+        return self.backbone.project_context(main_hidden.reshape(main_hidden.shape[0], 1, -1))
+
+    def _rope(self, length: int):
+        return rope_tables(mx.arange(length), self.args.head_dim, self.args.rope_theta, self.args.partial_rotary_factor)
+
+    def _softcap(self, logits: mx.array) -> mx.array:
+        sc = self.args.final_logit_softcapping
+        return mx.tanh(logits / sc) * sc if sc else logits
+
+    def _embed(self, ids: mx.array) -> mx.array:
+        """Gemma scales token embeddings by sqrt(hidden) (Gemma4TextScaledWordEmbedding)."""
+        e = self.embed_tokens(ids)
+        return e * mx.array(self.args.hidden_size ** 0.5, dtype=e.dtype)
+
+    # --- reference-matched eager interface (see qwen3.py for the rationale) ---
+
+    def extend_context(self, new_hiddens: mx.array) -> None:
+        n = new_hiddens.shape[1]
+        if n == 0:
+            return
+        proj = self.backbone.project_context(new_hiddens)
+        pos = mx.arange(self._committed, self._committed + n)
+        cos, sin = rope_tables(pos, self.args.head_dim, self.args.rope_theta, self.args.partial_rotary_factor)
+        layers = self.backbone.layers
+        if self._ctx_k is None:
+            self._ctx_k = [None] * len(layers)
+            self._ctx_v = [None] * len(layers)
+        for i, layer in enumerate(layers):
+            k, v = layer.context_kv(proj, cos, sin)
+            self._ctx_k[i] = k if self._ctx_k[i] is None else mx.concatenate([self._ctx_k[i], k], axis=2)
+            self._ctx_v[i] = v if self._ctx_v[i] is None else mx.concatenate([self._ctx_v[i], v], axis=2)
+        self._committed += n
+
+    def draft(self, anchor_token: mx.array):
+        b = anchor_token.shape[0]
+        start = self._committed
+        anchor = anchor_token.astype(mx.int32).reshape(b, 1)
+        noise = mx.full((b, self.block_size - 1), self.args.mask_token_id, dtype=mx.int32)
+        noise_embed = self._embed(mx.concatenate([anchor, noise], axis=1))
+        block_pos = mx.arange(start, start + self.block_size)
+        cos, sin = rope_tables(block_pos, self.args.head_dim, self.args.rope_theta, self.args.partial_rotary_factor)
+        h = noise_embed
+        ck = self._ctx_k or [None] * len(self.backbone.layers)
+        cv = self._ctx_v or [None] * len(self.backbone.layers)
+        for i, layer in enumerate(self.backbone.layers):
+            h = layer.forward_cached(h, ck[i], cv[i], cos, sin)
+        block_hidden = self.backbone.norm(h)
+        logits = self._softcap(self.lm_head(block_hidden.astype(mx.float32)))
+        return draft_block_decode(
+            logits, block_hidden, anchor_token, self.markov_head, self.confidence_head,
+            self.block_size, self.temperature,
+        )
+
+    def forward_spec(self, input_ids: mx.array, main_hidden: mx.array, start_pos: int = 0):
+        if start_pos == 0:
+            full = self.backbone.project_context(main_hidden)
+            self._ctx = full[:, :-1]
+            self._next_pos = full.shape[1] - 1
+            return None
+        self._ctx = mx.concatenate([self._ctx, self._project_one(main_hidden)], axis=1)
+        self._next_pos = start_pos + 1
+        b = input_ids.shape[0]
+        anchor = input_ids.astype(mx.int32).reshape(b, 1)
+        noise = mx.full((b, self.block_size - 1), self.args.mask_token_id, dtype=mx.int32)
+        noise_embed = self._embed(mx.concatenate([anchor, noise], axis=1))
+        cos, sin = self._rope(self._ctx.shape[1] + self.block_size)
+        block_hidden = self.backbone(noise_embed, self._ctx, cos, sin)
+        logits = self.lm_head(block_hidden.astype(mx.float32))
+        sc = self.args.final_logit_softcapping
+        if sc:
+            logits = mx.tanh(logits / sc) * sc
+        return draft_block_decode(
+            logits, block_hidden, input_ids, self.markov_head, self.confidence_head,
+            self.block_size, self.temperature,
+        )
+
+    def advance(self, main_hidden: mx.array, position: int) -> None:
+        self._ctx = mx.concatenate([self._ctx, self._project_one(main_hidden)], axis=1)
+        self._next_pos = position + 1
+
+
+def gemma4_key_map(key):
+    if key in ("embed_tokens.weight", "lm_head.weight"):
+        return key
+    if key in ("fc.weight", "hidden_norm.weight", "norm.weight"):
+        return f"backbone.{key}"
+    if key.startswith("markov_head.") or key.startswith("confidence_head."):
+        return key
+    m = _GEMMA4_LAYER_RE.match(key)
+    if m:
+        return f"backbone.layers.{m.group(1)}.{m.group(2)}"
+    return None
+
+
+def build(config, *, max_seq_len: int = 8192) -> Gemma4DSparkDrafter:
+    return Gemma4DSparkDrafter(Gemma4DSparkArgs.from_dict(config), max_seq_len=max_seq_len)
+
+
+GEMMA4 = DraftArch(name="gemma4", model_types=("gemma4", "gemma4_text"), build=build, key_map=gemma4_key_map)