From eb796e96570247a9645b814373c8225d7e3a333b Mon Sep 17 00:00:00 2001
From: mrciffa <49000955+davide221@users.noreply.github.com>
Date: Sat, 1 Aug 2026 14:08:17 +0200
Subject: [PATCH 1/3] feat(draft): add validated DSpark Q8 GGUF converter

---
 .github/workflows/ci.yml                      |   3 +
 server/scripts/convert_dspark_to_gguf.py      | 776 ++++++++++++++++++
 server/scripts/test_convert_dspark_to_gguf.py | 197 +++++
 3 files changed, 976 insertions(+)
 create mode 100644 server/scripts/convert_dspark_to_gguf.py
 create mode 100644 server/scripts/test_convert_dspark_to_gguf.py

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 28a0ca994..f73001b7a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -34,6 +34,9 @@ jobs:
       - name: Lint Python surfaces touched by lucebox tooling
         run: uv run --frozen --extra dev ruff check .
 
+      - name: Test DSpark GGUF converter
+        run: uv run --frozen --no-sync python -m unittest server/scripts/test_convert_dspark_to_gguf.py
+
   build:
     name: Build (cmake + uv sync --extra megakernel)
     runs-on: ubuntu-latest
diff --git a/server/scripts/convert_dspark_to_gguf.py b/server/scripts/convert_dspark_to_gguf.py
new file mode 100644
index 000000000..557e9dacc
--- /dev/null
+++ b/server/scripts/convert_dspark_to_gguf.py
@@ -0,0 +1,776 @@
+#!/usr/bin/env python3
+"""Convert a config-described DSpark/DFlash checkpoint to a Q8_0 GGUF.
+
+The converter intentionally knows the DFlash/DSpark tensor contract, not a
+specific target model.  Target geometry, capture layers, draft geometry, RoPE,
+and auxiliary-head dimensions come from ``config.json`` and are checked against
+the safetensors shapes before any output is committed.
+
+Large matrices are encoded as Q8_0.  Norms and the tiny confidence head remain
+F32.  Unknown tensors, malformed safetensors, inconsistent configs, hash
+mismatches, and unexpected Q8 alignment fail closed.  Output and the optional
+JSON report are written atomically.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import os
+import re
+import struct
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py"))
+
+import gguf
+from gguf.quants import dequantize, quantize
+
+ARCH = "dflash-draft"
+CONVERTER_VERSION = 1
+Q8_BLOCK_SIZE = 32
+MAX_HEADER_BYTES = 128 * 1024 * 1024
+SUPPORTED_DTYPES = {"BF16": 2, "F16": 2, "F32": 4}
+
+
+class ConversionError(RuntimeError):
+    """Raised when the source cannot safely satisfy the GGUF contract."""
+
+
+@dataclass(frozen=True)
+class TensorEntry:
+    name: str
+    dtype: str
+    shape: tuple[int, ...]
+    start: int
+    end: int
+
+    @property
+    def n_elements(self) -> int:
+        return math.prod(self.shape)
+
+    @property
+    def n_bytes(self) -> int:
+        return self.end - self.start
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    hidden: int
+    draft_layers: int
+    target_layers: int
+    heads: int
+    kv_heads: int
+    head_dim: int
+    intermediate: int
+    vocab: int
+    context_length: int
+    rms_eps: float
+    rope_theta: float
+    rope_type: str
+    rope_factor: float
+    rope_original_context: int
+    capture_layer_ids: tuple[int, ...]
+    block_size: int
+    mask_token_id: int
+    bos_token_id: int | None
+    eos_token_id: int | None
+    pad_token_id: int | None
+    markov_rank: int
+    markov_type: str
+    confidence_enabled: bool
+    confidence_with_markov: bool
+    confidence_dim: int
+    sliding_window: int
+    sliding_pattern: tuple[bool, ...]
+
+    @property
+    def capture_count(self) -> int:
+        return len(self.capture_layer_ids)
+
+
+@dataclass(frozen=True)
+class ConversionOptions:
+    model_dir: Path
+    output: Path
+    report: Path | None = None
+    name: str | None = None
+    source_repo: str | None = None
+    source_revision: str | None = None
+    target_repo: str | None = None
+    expected_sha256: str | None = None
+    max_relative_rmse: float = 0.01
+    sample_elements: int = 1_000_000
+    force: bool = False
+
+
+def _positive_int(config: dict[str, Any], key: str) -> int:
+    value = config.get(key)
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ConversionError(f"config.{key} must be a positive integer, got {value!r}")
+    return value
+
+
+def _optional_token_id(config: dict[str, Any], key: str) -> int | None:
+    value = config.get(key)
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+        raise ConversionError(f"config.{key} must be a non-negative integer, got {value!r}")
+    return value
+
+
+def load_model_spec(config: dict[str, Any]) -> ModelSpec:
+    architectures = config.get("architectures")
+    if not isinstance(architectures, list) or not any(
+        isinstance(value, str) and "DSparkDraftModel" in value for value in architectures
+    ):
+        raise ConversionError("config.architectures must identify a DSparkDraftModel")
+    if config.get("attention_bias", False):
+        raise ConversionError("attention_bias=true is not supported by the DFlash GGUF contract")
+    if config.get("hidden_act", "silu") != "silu":
+        raise ConversionError("only the SiLU DFlash MLP contract is currently supported")
+
+    hidden = _positive_int(config, "hidden_size")
+    draft_layers = _positive_int(config, "num_hidden_layers")
+    target_layers = _positive_int(config, "num_target_layers")
+    heads = _positive_int(config, "num_attention_heads")
+    kv_heads = _positive_int(config, "num_key_value_heads")
+    head_dim = _positive_int(config, "head_dim")
+    intermediate = _positive_int(config, "intermediate_size")
+    vocab = _positive_int(config, "vocab_size")
+    context_length = _positive_int(config, "max_position_embeddings")
+    block_size = _positive_int(config, "block_size")
+    markov_rank = _positive_int(config, "markov_rank")
+
+    if heads % kv_heads:
+        raise ConversionError(
+            f"num_attention_heads={heads} is not divisible by num_key_value_heads={kv_heads}"
+        )
+
+    rms_eps = config.get("rms_norm_eps")
+    if not isinstance(rms_eps, (int, float)) or isinstance(rms_eps, bool) or rms_eps <= 0:
+        raise ConversionError(f"config.rms_norm_eps must be positive, got {rms_eps!r}")
+
+    dflash = config.get("dflash_config")
+    if not isinstance(dflash, dict):
+        raise ConversionError("config.dflash_config must be an object")
+    capture_ids = dflash.get("target_layer_ids")
+    if (
+        not isinstance(capture_ids, list)
+        or not capture_ids
+        or any(isinstance(value, bool) or not isinstance(value, int) for value in capture_ids)
+    ):
+        raise ConversionError("config.dflash_config.target_layer_ids must be a non-empty int array")
+    capture_layer_ids = tuple(capture_ids)
+    if tuple(sorted(set(capture_layer_ids))) != capture_layer_ids:
+        raise ConversionError("target_layer_ids must be unique and strictly increasing")
+    if capture_layer_ids[0] < 0 or capture_layer_ids[-1] >= target_layers:
+        raise ConversionError(
+            f"target_layer_ids must fall inside target block range [0, {target_layers})"
+        )
+    mask_token_id = dflash.get("mask_token_id")
+    if isinstance(mask_token_id, bool) or not isinstance(mask_token_id, int):
+        raise ConversionError("config.dflash_config.mask_token_id must be an integer")
+    if not 0 <= mask_token_id < vocab:
+        raise ConversionError(f"mask_token_id={mask_token_id} is outside vocab_size={vocab}")
+
+    rope = config.get("rope_parameters") or {}
+    if not isinstance(rope, dict):
+        raise ConversionError("config.rope_parameters must be an object")
+    rope_theta = rope.get("rope_theta", config.get("rope_theta", 10_000.0))
+    rope_type = str(rope.get("rope_type", "none")).lower()
+    rope_factor = rope.get("factor", 1.0)
+    rope_original_context = rope.get("original_max_position_embeddings", context_length)
+    if not isinstance(rope_theta, (int, float)) or rope_theta <= 0:
+        raise ConversionError(f"RoPE theta must be positive, got {rope_theta!r}")
+    if not isinstance(rope_factor, (int, float)) or rope_factor <= 0:
+        raise ConversionError(f"RoPE factor must be positive, got {rope_factor!r}")
+    if (
+        isinstance(rope_original_context, bool)
+        or not isinstance(rope_original_context, int)
+        or rope_original_context <= 0
+    ):
+        raise ConversionError("original_max_position_embeddings must be a positive integer")
+    if rope_type not in {"none", "yarn"}:
+        raise ConversionError(f"unsupported rope_type={rope_type!r}; supported: none, yarn")
+
+    markov_type = str(config.get("markov_head_type", "vanilla")).lower()
+    if markov_type != "vanilla":
+        raise ConversionError(f"unsupported markov_head_type={markov_type!r}")
+    confidence_enabled = bool(config.get("enable_confidence_head", False))
+    confidence_with_markov = bool(config.get("confidence_head_with_markov", False))
+    if confidence_with_markov and not confidence_enabled:
+        raise ConversionError("confidence_head_with_markov requires enable_confidence_head")
+    confidence_dim = hidden + (markov_rank if confidence_with_markov else 0)
+
+    raw_layer_types = config.get("layer_types") or ["full_attention"] * draft_layers
+    if not isinstance(raw_layer_types, list) or len(raw_layer_types) != draft_layers:
+        raise ConversionError(f"config.layer_types must contain {draft_layers} entries")
+    unknown_layer_types = set(raw_layer_types) - {"full_attention", "sliding_attention"}
+    if unknown_layer_types:
+        raise ConversionError(f"unsupported layer_types: {sorted(unknown_layer_types)}")
+    sliding_pattern = tuple(value == "sliding_attention" for value in raw_layer_types)
+    raw_sliding_window = config.get("sliding_window")
+    if any(sliding_pattern):
+        if (
+            isinstance(raw_sliding_window, bool)
+            or not isinstance(raw_sliding_window, int)
+            or raw_sliding_window <= 0
+        ):
+            raise ConversionError("sliding_attention layers require a positive sliding_window")
+        sliding_window = raw_sliding_window
+    else:
+        sliding_window = 0
+
+    return ModelSpec(
+        hidden=hidden,
+        draft_layers=draft_layers,
+        target_layers=target_layers,
+        heads=heads,
+        kv_heads=kv_heads,
+        head_dim=head_dim,
+        intermediate=intermediate,
+        vocab=vocab,
+        context_length=context_length,
+        rms_eps=float(rms_eps),
+        rope_theta=float(rope_theta),
+        rope_type=rope_type,
+        rope_factor=float(rope_factor),
+        rope_original_context=rope_original_context,
+        capture_layer_ids=capture_layer_ids,
+        block_size=block_size,
+        mask_token_id=mask_token_id,
+        bos_token_id=_optional_token_id(config, "bos_token_id"),
+        eos_token_id=_optional_token_id(config, "eos_token_id"),
+        pad_token_id=_optional_token_id(config, "pad_token_id"),
+        markov_rank=markov_rank,
+        markov_type=markov_type,
+        confidence_enabled=confidence_enabled,
+        confidence_with_markov=confidence_with_markov,
+        confidence_dim=confidence_dim,
+        sliding_window=sliding_window,
+        sliding_pattern=sliding_pattern,
+    )
+
+
+def load_safetensors_header(path: Path) -> tuple[int, dict[str, TensorEntry]]:
+    file_size = path.stat().st_size
+    with path.open("rb") as handle:
+        encoded_size = handle.read(8)
+        if len(encoded_size) != 8:
+            raise ConversionError("safetensors file is too short to contain a header")
+        header_size = struct.unpack("<Q", encoded_size)[0]
+        if header_size == 0 or header_size > MAX_HEADER_BYTES:
+            raise ConversionError(f"invalid safetensors header size: {header_size}")
+        raw_header = handle.read(header_size)
+        if len(raw_header) != header_size:
+            raise ConversionError("truncated safetensors header")
+    try:
+        decoded = json.loads(raw_header)
+    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise ConversionError(f"invalid safetensors header JSON: {exc}") from exc
+    if not isinstance(decoded, dict):
+        raise ConversionError("safetensors header must be a JSON object")
+
+    data_size = file_size - 8 - header_size
+    if data_size < 0:
+        raise ConversionError("safetensors header extends beyond the file")
+    entries: dict[str, TensorEntry] = {}
+    intervals: list[tuple[int, int, str]] = []
+    for name, info in decoded.items():
+        if name == "__metadata__":
+            continue
+        if not isinstance(name, str) or not isinstance(info, dict):
+            raise ConversionError("invalid tensor entry in safetensors header")
+        dtype = info.get("dtype")
+        shape = info.get("shape")
+        offsets = info.get("data_offsets")
+        if dtype not in SUPPORTED_DTYPES:
+            raise ConversionError(f"{name}: unsupported source dtype {dtype!r}")
+        if (
+            not isinstance(shape, list)
+            or not shape
+            or any(isinstance(dim, bool) or not isinstance(dim, int) or dim <= 0 for dim in shape)
+        ):
+            raise ConversionError(f"{name}: shape must contain positive integers")
+        if (
+            not isinstance(offsets, list)
+            or len(offsets) != 2
+            or any(isinstance(value, bool) or not isinstance(value, int) for value in offsets)
+        ):
+            raise ConversionError(f"{name}: invalid data_offsets")
+        start, end = offsets
+        expected_bytes = math.prod(shape) * SUPPORTED_DTYPES[dtype]
+        if start < 0 or end <= start or end > data_size or end - start != expected_bytes:
+            raise ConversionError(
+                f"{name}: invalid byte range [{start}, {end}) for {dtype} shape={shape}"
+            )
+        entry = TensorEntry(name, dtype, tuple(shape), start, end)
+        entries[name] = entry
+        intervals.append((start, end, name))
+
+    if not entries:
+        raise ConversionError("safetensors contains no tensors")
+    intervals.sort()
+    if intervals[0][0] != 0 or intervals[-1][1] != data_size:
+        raise ConversionError("safetensors tensor data does not cover the payload exactly")
+    for (_, previous_end, previous_name), (start, _, name) in zip(
+        intervals, intervals[1:], strict=False
+    ):
+        if start != previous_end:
+            relation = "overlaps" if start < previous_end else "leaves a gap after"
+            raise ConversionError(f"{name} {relation} {previous_name} in the data payload")
+    return header_size, entries
+
+
+def map_tensor_name(name: str) -> str | None:
+    singleton_map = {
+        "fc.weight": "dflash.fc.weight",
+        "hidden_norm.weight": "dflash.hidden_norm.weight",
+        "norm.weight": "output_norm.weight",
+        "markov_head.markov_w1.weight": "dflash.dspark.markov.w1",
+        "markov_head.markov_w2.weight": "dflash.dspark.markov.w2",
+        "dspark_markov_head.markov_w1.weight": "dflash.dspark.markov.w1",
+        "dspark_markov_head.markov_w2.weight": "dflash.dspark.markov.w2",
+        "mtp.2.markov_head.markov_w1.weight": "dflash.dspark.markov.w1",
+        "mtp.2.markov_head.markov_w2.weight": "dflash.dspark.markov.w2",
+        "confidence_head.proj.weight": "dflash.dspark.confidence.weight",
+        "confidence_head.proj.bias": "dflash.dspark.confidence.bias",
+        "dspark_confidence_head.weight": "dflash.dspark.confidence.weight",
+        "dspark_confidence_head.bias": "dflash.dspark.confidence.bias",
+        "mtp.2.confidence_head.proj.weight": "dflash.dspark.confidence.weight",
+        "mtp.2.confidence_head.proj.bias": "dflash.dspark.confidence.bias",
+    }
+    if name in singleton_map:
+        return singleton_map[name]
+    match = re.fullmatch(r"layers\.(\d+)\.(.+)", name)
+    if not match:
+        return None
+    layer = int(match.group(1))
+    suffix_map = {
+        "input_layernorm.weight": "attn_norm.weight",
+        "post_attention_layernorm.weight": "ffn_norm.weight",
+        "self_attn.q_proj.weight": "attn_q.weight",
+        "self_attn.k_proj.weight": "attn_k.weight",
+        "self_attn.v_proj.weight": "attn_v.weight",
+        "self_attn.o_proj.weight": "attn_output.weight",
+        "self_attn.q_norm.weight": "attn_q_norm.weight",
+        "self_attn.k_norm.weight": "attn_k_norm.weight",
+        "mlp.gate_proj.weight": "ffn_gate.weight",
+        "mlp.up_proj.weight": "ffn_up.weight",
+        "mlp.down_proj.weight": "ffn_down.weight",
+    }
+    suffix = suffix_map.get(match.group(2))
+    return f"blk.{layer}.{suffix}" if suffix else None
+
+
+def expected_source_shapes(spec: ModelSpec) -> dict[str, tuple[int, ...]]:
+    shapes: dict[str, tuple[int, ...]] = {
+        "fc.weight": (spec.hidden, spec.capture_count * spec.hidden),
+        "hidden_norm.weight": (spec.hidden,),
+        "norm.weight": (spec.hidden,),
+        "markov_head.markov_w1.weight": (spec.vocab, spec.markov_rank),
+        "markov_head.markov_w2.weight": (spec.vocab, spec.markov_rank),
+    }
+    if spec.confidence_enabled:
+        shapes["confidence_head.proj.weight"] = (1, spec.confidence_dim)
+        shapes["confidence_head.proj.bias"] = (1,)
+    q_dim = spec.heads * spec.head_dim
+    kv_dim = spec.kv_heads * spec.head_dim
+    for layer in range(spec.draft_layers):
+        prefix = f"layers.{layer}."
+        shapes.update(
+            {
+                prefix + "input_layernorm.weight": (spec.hidden,),
+                prefix + "post_attention_layernorm.weight": (spec.hidden,),
+                prefix + "self_attn.q_proj.weight": (q_dim, spec.hidden),
+                prefix + "self_attn.k_proj.weight": (kv_dim, spec.hidden),
+                prefix + "self_attn.v_proj.weight": (kv_dim, spec.hidden),
+                prefix + "self_attn.o_proj.weight": (spec.hidden, q_dim),
+                prefix + "self_attn.q_norm.weight": (spec.head_dim,),
+                prefix + "self_attn.k_norm.weight": (spec.head_dim,),
+                prefix + "mlp.gate_proj.weight": (spec.intermediate, spec.hidden),
+                prefix + "mlp.up_proj.weight": (spec.intermediate, spec.hidden),
+                prefix + "mlp.down_proj.weight": (spec.hidden, spec.intermediate),
+            }
+        )
+    return shapes
+
+
+def validate_tensor_contract(entries: dict[str, TensorEntry], spec: ModelSpec) -> None:
+    expected = expected_source_shapes(spec)
+    # Accept historical aliases only after mapping them to the canonical contract.
+    actual_by_gguf: dict[str, TensorEntry] = {}
+    for source_name, entry in entries.items():
+        gguf_name = map_tensor_name(source_name)
+        if gguf_name is None:
+            raise ConversionError(f"unmapped source tensor: {source_name}")
+        if gguf_name in actual_by_gguf:
+            raise ConversionError(f"multiple source tensors map to {gguf_name}")
+        actual_by_gguf[gguf_name] = entry
+
+    expected_by_gguf = {map_tensor_name(name): (name, shape) for name, shape in expected.items()}
+    missing = sorted(set(expected_by_gguf) - set(actual_by_gguf))
+    extra = sorted(set(actual_by_gguf) - set(expected_by_gguf))
+    if missing or extra:
+        raise ConversionError(f"tensor contract mismatch: missing={missing}, extra={extra}")
+    for gguf_name, entry in actual_by_gguf.items():
+        _, expected_shape = expected_by_gguf[gguf_name]
+        if entry.shape != expected_shape:
+            raise ConversionError(
+                f"{entry.name}: shape {entry.shape} does not match expected {expected_shape}"
+            )
+
+
+def _read_tensor(path: Path, header_size: int, entry: TensorEntry) -> np.ndarray:
+    with path.open("rb") as handle:
+        handle.seek(8 + header_size + entry.start)
+        raw = handle.read(entry.n_bytes)
+    if len(raw) != entry.n_bytes:
+        raise ConversionError(f"short read for {entry.name}")
+    if entry.dtype == "BF16":
+        words = np.frombuffer(raw, dtype="<u2").reshape(entry.shape)
+        return (words.astype("<u4") << 16).view("<f4").reshape(entry.shape)
+    if entry.dtype == "F16":
+        return np.frombuffer(raw, dtype="<f2").reshape(entry.shape).astype("<f4")
+    return np.frombuffer(raw, dtype="<f4").reshape(entry.shape).copy()
+
+
+def _sha256(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _quantization_kind(gguf_name: str, shape: tuple[int, ...]) -> str:
+    if len(shape) == 1 or gguf_name.startswith("dflash.dspark.confidence."):
+        return "F32"
+    if shape[-1] % Q8_BLOCK_SIZE:
+        raise ConversionError(f"{gguf_name}: last dimension {shape[-1]} is not Q8_0 block-aligned")
+    return "Q8_0"
+
+
+def _sort_tensor(item: tuple[str, TensorEntry]) -> tuple[int, int, str]:
+    gguf_name, _ = item
+    if gguf_name.startswith("dflash."):
+        return (0, 0, gguf_name)
+    if gguf_name == "output_norm.weight":
+        return (1, 0, gguf_name)
+    match = re.match(r"blk\.(\d+)\.", gguf_name)
+    return (2, int(match.group(1)) if match else 0, gguf_name)
+
+
+def _sample_error(reference: np.ndarray, encoded: np.ndarray, limit: int) -> tuple[float, int]:
+    restored = dequantize(encoded, gguf.GGMLQuantizationType.Q8_0).reshape(reference.shape)
+    stride = max(1, math.ceil(reference.size / limit))
+    ref = reference.reshape(-1)[::stride][:limit].astype(np.float64)
+    got = restored.reshape(-1)[::stride][:limit].astype(np.float64)
+    error_energy = float(np.dot(got - ref, got - ref))
+    reference_energy = float(np.dot(ref, ref))
+    relative_rmse = math.sqrt(error_energy / max(reference_energy, np.finfo(np.float64).tiny))
+    return relative_rmse, ref.size
+
+
+def _add_metadata(
+    writer: gguf.GGUFWriter,
+    spec: ModelSpec,
+    config: dict[str, Any],
+    options: ConversionOptions,
+    source_hash: str,
+) -> None:
+    name = options.name or f"{options.model_dir.name}-Q8_0"
+    writer.add_name(name)
+    writer.add_type("model")
+    writer.add_description("Q8_0 DSpark/DFlash speculative drafter")
+    writer.add_quantized_by("Lucebox")
+    writer.add_file_type(gguf.LlamaFileType.MOSTLY_Q8_0)
+    writer.add_quantization_version(gguf.GGML_QUANT_VERSION)
+    if options.source_repo:
+        writer.add_source_repo_url(f"https://huggingface.co/{options.source_repo}")
+        writer.add_string("general.source.repository", options.source_repo)
+    if options.source_revision:
+        writer.add_string("general.source.revision", options.source_revision)
+    writer.add_string("general.source.file", "model.safetensors")
+    writer.add_string("general.source.sha256", source_hash)
+
+    writer.add_uint32(f"{ARCH}.context_length", spec.context_length)
+    writer.add_uint32(f"{ARCH}.embedding_length", spec.hidden)
+    writer.add_uint32(f"{ARCH}.block_count", spec.draft_layers)
+    writer.add_uint32(f"{ARCH}.feed_forward_length", spec.intermediate)
+    writer.add_uint32(f"{ARCH}.attention.head_count", spec.heads)
+    writer.add_uint32(f"{ARCH}.attention.head_count_kv", spec.kv_heads)
+    writer.add_uint32(f"{ARCH}.attention.key_length", spec.head_dim)
+    writer.add_uint32(f"{ARCH}.attention.value_length", spec.head_dim)
+    writer.add_uint32(f"{ARCH}.vocab_size", spec.vocab)
+    writer.add_float32(f"{ARCH}.attention.layer_norm_rms_epsilon", spec.rms_eps)
+    writer.add_uint32(f"{ARCH}.rope.dimension_count", spec.head_dim)
+    writer.add_float32(f"{ARCH}.rope.freq_base", spec.rope_theta)
+    if spec.rope_type == "yarn":
+        writer.add_string(f"{ARCH}.rope.scaling.type", "yarn")
+        writer.add_float32(f"{ARCH}.rope.scaling.factor", spec.rope_factor)
+        writer.add_uint32(
+            f"{ARCH}.rope.scaling.original_context_length", spec.rope_original_context
+        )
+    if spec.sliding_window:
+        writer.add_uint32(f"{ARCH}.attention.sliding_window", spec.sliding_window)
+        writer.add_array(f"{ARCH}.attention.sliding_window_pattern", spec.sliding_pattern)
+
+    # n_target_layers is the number of captured features consumed by fc, not
+    # the target network's block count.  They happen to be equal for some old
+    # DFlash checkpoints but are 5 and 93 respectively for Kimi K3.
+    writer.add_uint32(f"{ARCH}.dflash.n_target_layers", spec.capture_count)
+    writer.add_uint32(f"{ARCH}.dflash.n_target_features", spec.capture_count * spec.hidden)
+    writer.add_uint32(f"{ARCH}.dflash.target.block_count", spec.target_layers)
+    writer.add_uint32(f"{ARCH}.dflash.block_size", spec.block_size)
+    writer.add_uint32(f"{ARCH}.dflash.mask_token_id", spec.mask_token_id)
+    writer.add_array(f"{ARCH}.dflash.target_layer_ids", spec.capture_layer_ids)
+    if options.target_repo:
+        writer.add_string(f"{ARCH}.dflash.target.repository", options.target_repo)
+
+    writer.add_uint32(f"{ARCH}.dflash.dspark.enabled", 1)
+    writer.add_uint32(f"{ARCH}.dflash.dspark.markov_rank", spec.markov_rank)
+    writer.add_uint32(f"{ARCH}.dflash.dspark.vocab_size", spec.vocab)
+    writer.add_string(f"{ARCH}.dflash.dspark.markov_type", spec.markov_type)
+    writer.add_bool(f"{ARCH}.dflash.dspark.confidence.enabled", spec.confidence_enabled)
+    writer.add_bool(f"{ARCH}.dflash.dspark.confidence.with_markov", spec.confidence_with_markov)
+    if spec.confidence_enabled:
+        writer.add_uint32(f"{ARCH}.dflash.dspark.confidence_dim", spec.confidence_dim)
+
+    if spec.bos_token_id is not None:
+        writer.add_uint32("tokenizer.ggml.bos_token_id", spec.bos_token_id)
+    if spec.eos_token_id is not None:
+        writer.add_uint32("tokenizer.ggml.eos_token_id", spec.eos_token_id)
+    if spec.pad_token_id is not None:
+        writer.add_uint32("tokenizer.ggml.padding_token_id", spec.pad_token_id)
+    writer.add_string(f"{ARCH}.source.config_json", json.dumps(config, sort_keys=True))
+
+
+def _write_json_atomic(path: Path, value: dict[str, Any], force: bool) -> None:
+    if path.exists() and not force:
+        raise ConversionError(f"report already exists: {path} (pass --force to replace it)")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(f".{path.name}.partial.{os.getpid()}")
+    try:
+        with temporary.open("w", encoding="utf-8") as handle:
+            json.dump(value, handle, indent=2, sort_keys=True)
+            handle.write("\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary, path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def convert_model(options: ConversionOptions) -> dict[str, Any]:
+    if options.sample_elements <= 0:
+        raise ConversionError("sample_elements must be positive")
+    if not 0 < options.max_relative_rmse < 1:
+        raise ConversionError("max_relative_rmse must fall inside (0, 1)")
+    if bool(options.source_repo) != bool(options.source_revision):
+        raise ConversionError("source_repo and source_revision must be supplied together")
+    if options.output.exists() and not options.force:
+        raise ConversionError(
+            f"output already exists: {options.output} (pass --force to replace it)"
+        )
+    if options.report and options.report.exists() and not options.force:
+        raise ConversionError(
+            f"report already exists: {options.report} (pass --force to replace it)"
+        )
+
+    config_path = options.model_dir / "config.json"
+    source_path = options.model_dir / "model.safetensors"
+    if not config_path.is_file() or not source_path.is_file():
+        raise ConversionError(f"{options.model_dir} must contain config.json and model.safetensors")
+    protected_paths = {source_path.resolve(), config_path.resolve()}
+    if options.output.resolve() in protected_paths:
+        raise ConversionError("output must not overwrite a source model file")
+    if options.report and options.report.resolve() in protected_paths | {options.output.resolve()}:
+        raise ConversionError("report must be distinct from the source and GGUF paths")
+    try:
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as exc:
+        raise ConversionError(f"cannot read config.json: {exc}") from exc
+    if not isinstance(config, dict):
+        raise ConversionError("config.json must contain an object")
+
+    spec = load_model_spec(config)
+    header_size, entries = load_safetensors_header(source_path)
+    validate_tensor_contract(entries, spec)
+
+    source_hash = _sha256(source_path)
+    expected_hash = options.expected_sha256.lower() if options.expected_sha256 else None
+    if expected_hash and not re.fullmatch(r"[0-9a-f]{64}", expected_hash):
+        raise ConversionError("expected_sha256 must be 64 lowercase or uppercase hex characters")
+    if expected_hash and source_hash != expected_hash:
+        raise ConversionError(
+            f"source SHA256 mismatch: expected {expected_hash}, measured {source_hash}"
+        )
+
+    mapped = [(map_tensor_name(name), entry) for name, entry in entries.items()]
+    if any(name is None for name, _ in mapped):
+        raise AssertionError("validate_tensor_contract allowed an unmapped tensor")
+    ordered = sorted(((name, entry) for name, entry in mapped if name), key=_sort_tensor)
+
+    options.output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = options.output.with_name(f".{options.output.name}.partial.{os.getpid()}")
+    writer: gguf.GGUFWriter | None = None
+    quantized_bytes = 0
+    quantized_params = 0
+    q8_metrics: list[dict[str, Any]] = []
+    type_counts: dict[str, int] = {"Q8_0": 0, "F32": 0}
+    type_bytes: dict[str, int] = {"Q8_0": 0, "F32": 0}
+    try:
+        writer = gguf.GGUFWriter(temporary, ARCH)
+        _add_metadata(writer, spec, config, options, source_hash)
+        for gguf_name, entry in ordered:
+            values = _read_tensor(source_path, header_size, entry)
+            kind = _quantization_kind(gguf_name, entry.shape)
+            if kind == "Q8_0":
+                encoded = quantize(values, gguf.GGMLQuantizationType.Q8_0)
+                relative_rmse, samples = _sample_error(values, encoded, options.sample_elements)
+                if not math.isfinite(relative_rmse) or relative_rmse > options.max_relative_rmse:
+                    raise ConversionError(
+                        f"{gguf_name}: sampled relative RMSE {relative_rmse:.6g} exceeds "
+                        f"limit {options.max_relative_rmse:.6g}"
+                    )
+                writer.add_tensor(gguf_name, encoded, raw_dtype=gguf.GGMLQuantizationType.Q8_0)
+                q8_metrics.append(
+                    {
+                        "name": gguf_name,
+                        "relative_rmse": relative_rmse,
+                        "sample_elements": samples,
+                    }
+                )
+                quantized_params += entry.n_elements
+            else:
+                encoded = values.astype("<f4", copy=False)
+                writer.add_tensor(gguf_name, encoded, raw_dtype=gguf.GGMLQuantizationType.F32)
+            type_counts[kind] += 1
+            type_bytes[kind] += encoded.nbytes
+            quantized_bytes += encoded.nbytes
+            print(
+                f"[tensor] {gguf_name:48s} {entry.dtype:4s}->{kind:4s} "
+                f"shape={entry.shape} bytes={encoded.nbytes:,}"
+            )
+
+        writer.write_header_to_file()
+        writer.write_kv_data_to_file()
+        writer.write_tensors_to_file()
+        writer.close()
+        writer = None
+        with temporary.open("rb") as handle:
+            os.fsync(handle.fileno())
+        os.replace(temporary, options.output)
+    except Exception:
+        if writer is not None:
+            writer.close()
+        raise
+    finally:
+        temporary.unlink(missing_ok=True)
+
+    output_hash = _sha256(options.output)
+    source_bytes = source_path.stat().st_size
+    output_bytes = options.output.stat().st_size
+    weighted_error = math.sqrt(
+        sum(metric["relative_rmse"] ** 2 * metric["sample_elements"] for metric in q8_metrics)
+        / max(1, sum(metric["sample_elements"] for metric in q8_metrics))
+    )
+    report: dict[str, Any] = {
+        "schema_version": 1,
+        "converter": {"name": Path(__file__).name, "version": CONVERTER_VERSION},
+        "source": {
+            "repository": options.source_repo,
+            "revision": options.source_revision,
+            "file": source_path.name,
+            "sha256": source_hash,
+            "bytes": source_bytes,
+            "tensor_count": len(entries),
+            "parameter_count": sum(entry.n_elements for entry in entries.values()),
+        },
+        "target": {
+            "repository": options.target_repo,
+            "block_count": spec.target_layers,
+            "capture_layer_ids": list(spec.capture_layer_ids),
+        },
+        "output": {
+            "file": options.output.name,
+            "sha256": output_hash,
+            "bytes": output_bytes,
+            "architecture": ARCH,
+            "file_type": "MOSTLY_Q8_0",
+            "source_size_ratio": output_bytes / source_bytes,
+        },
+        "quantization": {
+            "tensor_counts": type_counts,
+            "payload_bytes": type_bytes,
+            "payload_total_bytes": quantized_bytes,
+            "q8_parameter_count": quantized_params,
+            "sampled_relative_rmse_rms": weighted_error,
+            "sampled_relative_rmse_max": max(
+                (metric["relative_rmse"] for metric in q8_metrics), default=0.0
+            ),
+            "max_allowed_relative_rmse": options.max_relative_rmse,
+            "tensors": q8_metrics,
+        },
+    }
+    if options.report:
+        _write_json_atomic(options.report, report, options.force)
+    print(
+        f"[done] {source_bytes / 1e9:.3f} GB -> {output_bytes / 1e9:.3f} GB "
+        f"({output_bytes / source_bytes:.1%}), sampled relative RMSE={weighted_error:.6g}"
+    )
+    print(f"[sha256] source={source_hash}")
+    print(f"[sha256] output={output_hash}")
+    return report
+
+
+def parse_args(argv: list[str] | None = None) -> ConversionOptions:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model_dir", type=Path)
+    parser.add_argument("output", type=Path)
+    parser.add_argument("--report", type=Path)
+    parser.add_argument("--name")
+    parser.add_argument("--source-repo")
+    parser.add_argument("--source-revision")
+    parser.add_argument("--target-repo")
+    parser.add_argument("--expected-sha256")
+    parser.add_argument("--max-relative-rmse", type=float, default=0.01)
+    parser.add_argument("--sample-elements", type=int, default=1_000_000)
+    parser.add_argument("--force", action="store_true")
+    args = parser.parse_args(argv)
+    return ConversionOptions(
+        model_dir=args.model_dir,
+        output=args.output,
+        report=args.report,
+        name=args.name,
+        source_repo=args.source_repo,
+        source_revision=args.source_revision,
+        target_repo=args.target_repo,
+        expected_sha256=args.expected_sha256,
+        max_relative_rmse=args.max_relative_rmse,
+        sample_elements=args.sample_elements,
+        force=args.force,
+    )
+
+
+def main() -> int:
+    try:
+        convert_model(parse_args())
+    except (ConversionError, OSError, ValueError) as exc:
+        print(f"[error] {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/server/scripts/test_convert_dspark_to_gguf.py b/server/scripts/test_convert_dspark_to_gguf.py
new file mode 100644
index 000000000..1313e9b93
--- /dev/null
+++ b/server/scripts/test_convert_dspark_to_gguf.py
@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import struct
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPT_DIR))
+
+import convert_dspark_to_gguf as converter
+import gguf
+
+
+def _to_bf16_bytes(values: np.ndarray) -> bytes:
+    words = values.astype("<f4").view("<u4")
+    return (words >> 16).astype("<u2").tobytes()
+
+
+def _write_safetensors(path: Path, tensors: dict[str, np.ndarray]) -> None:
+    header: dict[str, object] = {"__metadata__": {"format": "pt"}}
+    payload = bytearray()
+    for name, values in tensors.items():
+        raw = _to_bf16_bytes(values)
+        start = len(payload)
+        payload.extend(raw)
+        header[name] = {
+            "dtype": "BF16",
+            "shape": list(values.shape),
+            "data_offsets": [start, len(payload)],
+        }
+    encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    encoded += b" " * (-len(encoded) % 8)
+    path.write_bytes(struct.pack("<Q", len(encoded)) + encoded + payload)
+
+
+def _fixture_tensors(seed: int = 7) -> dict[str, np.ndarray]:
+    rng = np.random.default_rng(seed)
+
+    def values(*shape: int) -> np.ndarray:
+        return rng.normal(0, 0.1, shape).astype("<f4")
+
+    return {
+        "confidence_head.proj.bias": values(1),
+        "confidence_head.proj.weight": values(1, 64),
+        "fc.weight": values(32, 32),
+        "hidden_norm.weight": values(32),
+        "layers.0.input_layernorm.weight": values(32),
+        "layers.0.mlp.down_proj.weight": values(32, 64),
+        "layers.0.mlp.gate_proj.weight": values(64, 32),
+        "layers.0.mlp.up_proj.weight": values(64, 32),
+        "layers.0.post_attention_layernorm.weight": values(32),
+        "layers.0.self_attn.k_norm.weight": values(8),
+        "layers.0.self_attn.k_proj.weight": values(16, 32),
+        "layers.0.self_attn.o_proj.weight": values(32, 32),
+        "layers.0.self_attn.q_norm.weight": values(8),
+        "layers.0.self_attn.q_proj.weight": values(32, 32),
+        "layers.0.self_attn.v_proj.weight": values(16, 32),
+        "markov_head.markov_w1.weight": values(64, 32),
+        "markov_head.markov_w2.weight": values(64, 32),
+        "norm.weight": values(32),
+    }
+
+
+def _fixture_config() -> dict[str, object]:
+    return {
+        "architectures": ["DSparkDraftModel"],
+        "attention_bias": False,
+        "block_size": 2,
+        "bos_token_id": 1,
+        "confidence_head_with_markov": True,
+        "dflash_config": {"mask_token_id": 63, "target_layer_ids": [1]},
+        "enable_confidence_head": True,
+        "eos_token_id": 2,
+        "head_dim": 8,
+        "hidden_act": "silu",
+        "hidden_size": 32,
+        "intermediate_size": 64,
+        "layer_types": ["full_attention"],
+        "markov_head_type": "vanilla",
+        "markov_rank": 32,
+        "max_position_embeddings": 4096,
+        "num_attention_heads": 4,
+        "num_hidden_layers": 1,
+        "num_key_value_heads": 2,
+        "num_target_layers": 3,
+        "pad_token_id": 0,
+        "rms_norm_eps": 1e-5,
+        "rope_parameters": {
+            "factor": 2.0,
+            "original_max_position_embeddings": 2048,
+            "rope_theta": 10000.0,
+            "rope_type": "yarn",
+        },
+        "tie_word_embeddings": False,
+        "vocab_size": 64,
+    }
+
+
+def _field_value(reader: gguf.GGUFReader, name: str):
+    field = reader.fields[name]
+    part = field.parts[field.data[0]]
+    if field.types[0] == gguf.GGUFValueType.STRING:
+        return bytes(part).decode("utf-8")
+    return part.tolist()[0]
+
+
+class ConvertDSparkToGGUFTest(unittest.TestCase):
+    def setUp(self) -> None:
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.root = Path(self.temp_dir.name)
+        self.model_dir = self.root / "draft"
+        self.model_dir.mkdir()
+        (self.model_dir / "config.json").write_text(json.dumps(_fixture_config()))
+        _write_safetensors(self.model_dir / "model.safetensors", _fixture_tensors())
+
+    def tearDown(self) -> None:
+        self.temp_dir.cleanup()
+
+    def test_converts_q8_with_complete_metadata_and_report(self) -> None:
+        output = self.root / "draft-q8_0.gguf"
+        report_path = self.root / "conversion-report.json"
+        source_hash = hashlib.sha256(
+            (self.model_dir / "model.safetensors").read_bytes()
+        ).hexdigest()
+
+        report = converter.convert_model(
+            converter.ConversionOptions(
+                model_dir=self.model_dir,
+                output=output,
+                report=report_path,
+                source_repo="example/draft",
+                source_revision="a" * 40,
+                target_repo="example/target",
+                expected_sha256=source_hash,
+                sample_elements=256,
+            )
+        )
+
+        self.assertTrue(output.is_file())
+        self.assertTrue(report_path.is_file())
+        self.assertEqual(report["source"]["tensor_count"], 18)
+        self.assertEqual(report["quantization"]["tensor_counts"], {"Q8_0": 10, "F32": 8})
+        self.assertLess(report["quantization"]["sampled_relative_rmse_max"], 0.01)
+
+        reader = gguf.GGUFReader(output)
+        self.assertEqual(_field_value(reader, "general.architecture"), "dflash-draft")
+        self.assertEqual(_field_value(reader, "dflash-draft.dflash.n_target_layers"), 1)
+        self.assertEqual(_field_value(reader, "dflash-draft.dflash.target.block_count"), 3)
+        self.assertEqual(_field_value(reader, "dflash-draft.dflash.dspark.markov_rank"), 32)
+        self.assertEqual(_field_value(reader, "general.source.sha256"), source_hash)
+
+        tensors = {tensor.name: tensor for tensor in reader.tensors}
+        self.assertEqual(len(tensors), 18)
+        self.assertEqual(
+            tensors["dflash.fc.weight"].tensor_type,
+            gguf.GGMLQuantizationType.Q8_0,
+        )
+        self.assertEqual(
+            tensors["dflash.dspark.confidence.weight"].tensor_type,
+            gguf.GGMLQuantizationType.F32,
+        )
+
+    def test_unknown_tensor_fails_without_partial_output(self) -> None:
+        tensors = _fixture_tensors()
+        tensors["unexpected.weight"] = np.ones((32, 32), dtype="<f4")
+        _write_safetensors(self.model_dir / "model.safetensors", tensors)
+        output = self.root / "should-not-exist.gguf"
+
+        with self.assertRaisesRegex(converter.ConversionError, "unmapped source tensor"):
+            converter.convert_model(
+                converter.ConversionOptions(model_dir=self.model_dir, output=output)
+            )
+
+        self.assertFalse(output.exists())
+        self.assertEqual(list(self.root.glob("*.partial.*")), [])
+
+    def test_hash_mismatch_fails_before_output(self) -> None:
+        output = self.root / "should-not-exist.gguf"
+        with self.assertRaisesRegex(converter.ConversionError, "SHA256 mismatch"):
+            converter.convert_model(
+                converter.ConversionOptions(
+                    model_dir=self.model_dir,
+                    output=output,
+                    expected_sha256="0" * 64,
+                )
+            )
+        self.assertFalse(output.exists())
+
+
+if __name__ == "__main__":
+    unittest.main()

From 1ccd979b80cdb8e539cbcb78195fc06425ca9374 Mon Sep 17 00:00:00 2001
From: mrciffa <49000955+davide221@users.noreply.github.com>
Date: Sat, 1 Aug 2026 15:37:21 +0200
Subject: [PATCH 2/3] feat(kimi-k3): add shared DSpark speculative runtime

---
 server/CMakeLists.txt                        |   1 +
 server/src/common/adaptive_verify_width.h    |  15 +-
 server/src/common/backend_args.h             |   2 +-
 server/src/common/backend_factory.cpp        |   4 +
 server/src/common/dflash_feature_ring.cpp    |  21 +-
 server/src/common/dflash_spec_decode.cpp     | 117 +++-
 server/src/common/dflash_target.h            |  12 +
 server/src/common/model_capabilities.h       |   2 +-
 server/src/common/step_graph.h               |   2 +
 server/src/draft/draft_gguf_loader.cpp       |   2 +
 server/src/kimi_k3/kimi_k3_backend.cpp       | 182 +++++-
 server/src/kimi_k3/kimi_k3_backend.h         |  17 +-
 server/src/kimi_k3/kimi_k3_dflash_target.cpp | 277 ++++++++
 server/src/kimi_k3/kimi_k3_dflash_target.h   |  90 +++
 server/src/kimi_k3/kimi_k3_graph.cpp         | 643 ++++++++++++++-----
 server/src/kimi_k3/kimi_k3_internal.h        |  67 +-
 server/test/smoke_kimi_k3_forward.cpp        |   5 +-
 server/test/test_feature_gate.cpp            |   5 +-
 18 files changed, 1263 insertions(+), 201 deletions(-)
 create mode 100644 server/src/kimi_k3/kimi_k3_dflash_target.cpp
 create mode 100644 server/src/kimi_k3/kimi_k3_dflash_target.h

diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt
index 1401112e5..8c2cade0f 100644
--- a/server/CMakeLists.txt
+++ b/server/CMakeLists.txt
@@ -400,6 +400,7 @@ add_library(dflash_common STATIC
     # Kimi-K3 hybrid KDA/MLA + latent-MoE target arch
     src/kimi_k3/kimi_k3_loader.cpp
     src/kimi_k3/kimi_k3_graph.cpp
+    src/kimi_k3/kimi_k3_dflash_target.cpp
     src/kimi_k3/kimi_k3_backend.cpp
     src/flashprefill_q8.cpp
     src/kv_cache.cpp
diff --git a/server/src/common/adaptive_verify_width.h b/server/src/common/adaptive_verify_width.h
index b7a0cb954..bb9f785f1 100644
--- a/server/src/common/adaptive_verify_width.h
+++ b/server/src/common/adaptive_verify_width.h
@@ -13,7 +13,7 @@
 //
 // On by default (theta 0.20). Debug overrides:
 //   DFLASH_ADAPTIVE_WIDTH_THETA=<0..1>  0 disables (legacy fixed/EWMA width)
-//   DFLASH_ADAPTIVE_WIDTH_MIN=<n>       minimum kept rows incl. seed (default 4)
+//   DFLASH_ADAPTIVE_WIDTH_MIN=<n>       minimum kept rows incl. seed
 //
 // Model-agnostic: any family loop that has per-slot drafter top-1
 // probabilities (e.g. from ggml_backend_cuda_topk_rows over the draft-head
@@ -38,20 +38,21 @@ inline float adaptive_verify_width_theta() {
     return theta;
 }
 
-inline int adaptive_verify_width_min() {
-    static const int mn = []() {
+inline int adaptive_verify_width_min(int fallback = 4) {
+    static const int configured = []() {
         const char * e = std::getenv("DFLASH_ADAPTIVE_WIDTH_MIN");
-        if (!e) return 4;
+        if (!e) return 0;
         const int v = std::atoi(e);
         if (v <= 0) {
             std::fprintf(stderr, "[adaptive-width] ignoring "
                                  "DFLASH_ADAPTIVE_WIDTH_MIN=\"%s\" "
-                                 "(want a positive int); using 4\n", e);
-            return 4;
+                                 "(want a positive int); using target default\n",
+                         e);
+            return 0;
         }
         return v;
     }();
-    return mn;
+    return configured > 0 ? configured : (fallback > 0 ? fallback : 4);
 }
 
 // top1_probs[(j-1)*stride]: drafter top-1 probability of candidate slot j.
diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h
index c0573d2de..cc8e7c712 100644
--- a/server/src/common/backend_args.h
+++ b/server/src/common/backend_args.h
@@ -36,7 +36,7 @@ struct BackendArgs {
     // Required
     const char *    model_path   = nullptr;   // target .gguf
 
-    // Optional: speculative decode draft model (qwen35 only)
+    // Optional: architecture-compatible speculative decode draft model.
     const char *    draft_path   = nullptr;
 
     // Device placement
diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp
index a24fa88ef..f0fd8ed2a 100644
--- a/server/src/common/backend_factory.cpp
+++ b/server/src/common/backend_factory.cpp
@@ -476,7 +476,11 @@ std::unique_ptr<ModelBackend> create_backend(
     } else if (arch == "kimi-k3") {
         KimiK3BackendConfig cfg;
         cfg.model_path = args.model_path;
+        cfg.draft_path = args.draft_path;
         cfg.device = args.device;
+        cfg.draft_gpu = args.draft_device.gpu;
+        cfg.draft_ctx_max = args.draft_ctx_max;
+        cfg.fast_rollback = args.fast_rollback;
         cfg.stream_fd = args.stream_fd;
         cfg.moe_storage = plan.moe_storage_policy();
 
diff --git a/server/src/common/dflash_feature_ring.cpp b/server/src/common/dflash_feature_ring.cpp
index 076a9be12..fc6f821af 100644
--- a/server/src/common/dflash_feature_ring.cpp
+++ b/server/src/common/dflash_feature_ring.cpp
@@ -520,14 +520,29 @@ bool copy_host_capture_slice_to_draft_ring(
     const size_t expected = (size_t)n_tokens * (size_t)hidden;
     if (host_elems != expected) return false;
     const size_t dst_stride = feature_ring.target_feat->nb[1];
-    const size_t row_bytes = (size_t)hidden * sizeof(float);
+    const size_t row_bytes =
+        ggml_row_size(feature_ring.storage_type, hidden);
+    std::vector<uint8_t> converted;
+    if (feature_ring.storage_type != GGML_TYPE_F32) {
+        converted.resize(row_bytes);
+    }
     for (int i = 0; i < n_tokens; ++i) {
         const int slot = (start_pos + i) % feature_ring.cap;
         const float * src = host + (size_t)i * (size_t)hidden;
+        const void * row = src;
+        if (!converted.empty()) {
+            if (!host_f32_to_feature_row(
+                    feature_ring.storage_type, src,
+                    converted.data(), static_cast<size_t>(hidden))) {
+                return false;
+            }
+            row = converted.data();
+        }
         const size_t dst_offset =
             (size_t)slot * dst_stride +
-            (size_t)capture_idx * (size_t)hidden * sizeof(float);
-        ggml_backend_tensor_set(feature_ring.target_feat, src, dst_offset, row_bytes);
+            (size_t)capture_idx * row_bytes;
+        ggml_backend_tensor_set(
+            feature_ring.target_feat, row, dst_offset, row_bytes);
     }
     return true;
 }
diff --git a/server/src/common/dflash_spec_decode.cpp b/server/src/common/dflash_spec_decode.cpp
index 60e28c733..caab7c244 100644
--- a/server/src/common/dflash_spec_decode.cpp
+++ b/server/src/common/dflash_spec_decode.cpp
@@ -5,6 +5,8 @@
 #include "internal.h"        // DraftWeights
 #include "io_utils.h"
 #include "dflash_draft_graph.h"  // build_draft_step
+#include "dspark_head.h"
+#include "adaptive_verify_width.h"
 #include "step_graph.h"
 
 #include <algorithm>
@@ -64,17 +66,21 @@ bool run_dflash_spec_decode(
     const bool use_remote_draft = remote_draft && remote_draft->active();
     if (!use_remote_draft && !feature_ring.target_feat) return false;
 
-    const int hidden = draft_weights.n_embd;
-    const int q_len  = draft_weights.block_size;
+    const int hidden    = draft_weights.n_embd;
+    const int max_q_len = draft_weights.block_size;
+    if (hidden <= 0 || max_q_len <= 0) return false;
+    const float width_theta = adaptive_verify_width_theta();
+    const int target_width_min = target.default_adaptive_verify_min_rows();
+    const int width_min = adaptive_verify_width_min(target_width_min);
 
     StepGraph draft_sg;
     StepGraphGuard draft_sg_guard{draft_sg};
 
-    std::vector<float>   noise_embed((size_t)hidden * q_len);
-    std::vector<int32_t> noise_ids(q_len);
-    std::vector<int32_t> draft_tok(q_len);
-    std::vector<int32_t> target_tok(q_len);
-    std::vector<int32_t> pos_q(q_len);
+    std::vector<float>   noise_embed((size_t)hidden * max_q_len);
+    std::vector<int32_t> noise_ids(max_q_len);
+    std::vector<int32_t> draft_tok(max_q_len);
+    std::vector<int32_t> target_tok(max_q_len);
+    std::vector<int32_t> pos_q(max_q_len);
     std::vector<int32_t> pos_k;
     std::vector<float>   local_hidden;       // host buffer for local draft hidden states
     std::vector<float>   remote_hidden;      // host buffer for remote-draft hidden states
@@ -84,6 +90,7 @@ bool run_dflash_spec_decode(
     int n_generated     = 0;
     int n_draft_steps   = 0;
     int n_accept_sum    = 0;
+    int n_verify_rows   = 0;
     int n_hint_proposed = 0;
     int n_hint_accepted = 0;
     const ChainRollbackPolicy rollback_policy = resolve_chain_rollback_policy();
@@ -92,11 +99,15 @@ bool run_dflash_spec_decode(
     auto t_dec0 = std::chrono::steady_clock::now();
     while (n_generated < n_gen) {
         const int need_commit_budget = n_gen - n_generated;
+        int q_len = max_q_len;
 
         // ── Build noise input for draft ────────────────────────────────────
         noise_ids[0] = last_tok;
-        for (int i = 1; i < q_len; i++) noise_ids[i] = target.mask_token_id();
-        if (!target.embed_tokens(noise_ids.data(), q_len, noise_embed.data())) {
+        for (int i = 1; i < max_q_len; i++) {
+            noise_ids[i] = target.mask_token_id();
+        }
+        if (!target.embed_tokens(
+                noise_ids.data(), max_q_len, noise_embed.data())) {
             std::fprintf(stderr, "dflash-spec noise embed failed\n");
             return false;
         }
@@ -135,9 +146,9 @@ bool run_dflash_spec_decode(
             }
             ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0,
                                     sizeof(float) * noise_embed.size());
-            pos_k.resize((size_t)draft_ctx + q_len);
-            for (int i = 0; i < q_len; i++) pos_q[i] = draft_ctx + i;
-            for (int i = 0; i < draft_ctx + q_len; i++) pos_k[i] = i;
+            pos_k.resize((size_t)draft_ctx + max_q_len);
+            for (int i = 0; i < max_q_len; i++) pos_q[i] = draft_ctx + i;
+            for (int i = 0; i < draft_ctx + max_q_len; i++) pos_k[i] = i;
             ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0,
                                     sizeof(int32_t) * pos_q.size());
             ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0,
@@ -149,31 +160,83 @@ bool run_dflash_spec_decode(
             }
             // Read draft hidden states out to host so the target adapter can
             // project them through its own LM head (target-internal layout).
-            local_hidden.resize((size_t)hidden * q_len);
+            local_hidden.resize((size_t)hidden * max_q_len);
             ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0,
                                     sizeof(float) * local_hidden.size());
             draft_hidden_host = local_hidden.data();
         }
 
-        // ── Project draft hidden → token IDs via target's LM head ─────────
-        if (!target.project_hidden_to_tokens(draft_hidden_host, q_len, draft_tok)) {
+        // ── Project draft hidden → token IDs via the shared draft head ────
+        // DSpark is an auxiliary head on the universal DFlash checkpoint, not
+        // a target-model implementation.  Keep its Markov correction and
+        // confidence policy here so Kimi, Qwen, Laguna, and future adapters
+        // share exactly one implementation.
+        std::vector<float> confidence;
+        std::vector<float> candidate_probs;
+        std::vector<int32_t> candidate_ids;
+        bool projected = false;
+        if (draft_backend && draft_weights.dspark.enabled && max_q_len > 1) {
+            ggml_backend_t head_backend = target.fused_head_backend();
+            const bool same_device_head =
+                head_backend && target.lm_head_tensor() &&
+                ggml_backend_get_device(head_backend) ==
+                    ggml_backend_get_device(draft_backend);
+            if (same_device_head) {
+                projected = dspark_markov_correct_greedy_chain_fused(
+                    draft_weights, draft_backend, target.lm_head_tensor(),
+                    draft_hidden_host, max_q_len, last_tok, draft_tok,
+                    &confidence);
+            }
+            if (!projected) {
+                projected = dspark_markov_correct_greedy_chain(
+                    draft_weights, draft_backend, target,
+                    draft_hidden_host, max_q_len, last_tok,
+                    /*confidence_threshold=*/0.0f, draft_tok);
+            }
+        }
+        if (!projected) {
+            const int candidate_k = width_theta > 0.0f && max_q_len > 2 ? 2 : 0;
+            projected = target.project_hidden_to_tokens_topk(
+                draft_hidden_host, max_q_len, draft_tok, candidate_k,
+                candidate_k > 0 ? &candidate_probs : nullptr,
+                candidate_k > 0 ? &candidate_ids : nullptr);
+        }
+        if (!projected || draft_tok.size() < (size_t)max_q_len) {
             std::fprintf(stderr, "dflash-spec projection failed\n");
             return false;
         }
         draft_tok[0] = last_tok;
 
+        if (confidence.size() >= (size_t)(max_q_len - 1)) {
+            q_len = adaptive_verify_width(
+                confidence.data(), 1, max_q_len, width_theta, width_min);
+        } else if (candidate_probs.size() >=
+                   static_cast<size_t>((max_q_len - 1) * 2)) {
+            q_len = adaptive_verify_width(
+                candidate_probs.data(), 2, max_q_len,
+                width_theta, width_min);
+        }
+
         // ── Tool call hint injection ──────────────────────────────────────
         // Override draft tokens with pre-known hint tokens for near-100%
         // acceptance on predictable structural positions.
         int hint_filled = 0;
         if (hint_tokens && n_generated < (int)hint_tokens->size()) {
             const int hint_avail = (int)hint_tokens->size() - n_generated;
-            hint_filled = std::min(hint_avail, q_len - 1);
+            q_len = max_q_len;
+            hint_filled = std::min(hint_avail, max_q_len - 1);
             for (int i = 0; i < hint_filled; i++) {
                 draft_tok[1 + i] = (*hint_tokens)[n_generated + i];
             }
         }
 
+        // Never verify rows that cannot fit in the remaining generation
+        // budget.  The draft graph remains fixed-width for graph reuse; this
+        // only trims expensive target/MoE work.
+        q_len = std::max(1, std::min(q_len, need_commit_budget));
+        draft_tok.resize(static_cast<size_t>(q_len));
+        target_tok.resize(static_cast<size_t>(q_len));
+
         // Notify observer with draft tokens for this step.
         if (io.observer) {
             io.observer("draft", draft_tok);
@@ -220,12 +283,24 @@ bool run_dflash_spec_decode(
         rollback_diag.record_accept(accept_n);
         const bool use_fast_rollback =
             target.supports_fast_rollback() &&
-            (accept_n >= rollback_policy.fast_rollback_threshold);
+            (target.prefer_fast_rollback_over_replay() ||
+             accept_n >= rollback_policy.fast_rollback_threshold);
 
         std::vector<int32_t> replay_tok((size_t)commit_n);
         for (int i = 0; i < commit_n; i++) {
             replay_tok[i] = (i < accept_n) ? draft_tok[i] : bonus_tok;
         }
+        // Never commit hidden state past EOS when a verify batch happens to
+        // accept additional candidates after it.
+        for (int i = 0; i < commit_n; ++i) {
+            if (target.is_eos(replay_tok[(size_t)i])) {
+                commit_n = i + 1;
+                accept_n = std::min(accept_n, commit_n);
+                bonus_tok = -1;
+                replay_tok.resize((size_t)commit_n);
+                break;
+            }
+        }
 
         bool fast_rolled_back = false;
         if (use_fast_rollback) {
@@ -275,11 +350,15 @@ bool run_dflash_spec_decode(
             io.emit(replay_tok[i]);
             if (io.cancelled) break;
             ++emitted;
-            if (target.is_eos(replay_tok[i])) hit_eos = true;
+            if (target.is_eos(replay_tok[i])) {
+                hit_eos = true;
+                break;
+            }
         }
         committed   += emitted;
         n_generated += emitted;
         n_accept_sum += std::min(accept_n, emitted);
+        n_verify_rows += q_len;
         n_draft_steps++;
 
         // Notify observer with accepted tokens for this step.
@@ -293,7 +372,7 @@ bool run_dflash_spec_decode(
     if (!use_remote_draft && draft_backend) ggml_backend_synchronize(draft_backend);
     auto t_dec1 = std::chrono::steady_clock::now();
     const double decode_s = std::chrono::duration<double>(t_dec1 - t_dec0).count();
-    const int total_draft_pos = std::max(1, n_draft_steps * q_len);
+    const int total_draft_pos = std::max(1, n_verify_rows);
     const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos;
     if (accept_rate_out) {
         *accept_rate_out = total_draft_pos > 0
diff --git a/server/src/common/dflash_target.h b/server/src/common/dflash_target.h
index af7cd756f..f14f0a4c5 100644
--- a/server/src/common/dflash_target.h
+++ b/server/src/common/dflash_target.h
@@ -62,6 +62,11 @@ struct DFlashTarget {
     // When true, verify_batch captures intermediates and rollback_to() works.
     virtual bool supports_fast_rollback() const { return false; }
 
+    // Cost hint for targets where applying captured recurrent transitions is
+    // always cheaper than restoring and replaying the accepted tokens.  The
+    // default preserves the shared acceptance-threshold policy.
+    virtual bool prefer_fast_rollback_over_replay() const { return false; }
+
     // Roll back recurrent state to position `commit_n` within the last
     // verify batch (0-indexed). Uses SSM intermediate states captured during
     // verify. Also truncates KV to `base_pos + commit_n`. No replay needed.
@@ -184,6 +189,13 @@ struct DFlashTarget {
     // Which target layers to capture intermediate activations from.
     // The draft model's fc layer expects exactly this many feature slices.
     virtual const std::vector<int> & capture_layer_ids() const = 0;
+
+    // Optional cost hint for adaptive speculative verification.  Zero keeps
+    // the shared runtime default.  Targets whose extra verify rows trigger
+    // expensive external-memory traffic can return a smaller floor without
+    // putting model- or storage-specific policy in the common decode loop.
+    // DFLASH_ADAPTIVE_WIDTH_MIN remains an explicit runtime override.
+    virtual int default_adaptive_verify_min_rows() const { return 0; }
 };
 
 } // namespace dflash::common
diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h
index 7187e9921..a35af6c5d 100644
--- a/server/src/common/model_capabilities.h
+++ b/server/src/common/model_capabilities.h
@@ -70,7 +70,7 @@ inline constexpr ArchCapabilities kArchCapabilities[] = {
     {"qwen3",      false, false, true,  false,   kNever, kNever, kNever, kNever, kNever, kNever},
     {"gemma4",     true,  false, false, false,   kMono, kNever, kNever, kBoth, kNever, kNever},
     {"deepseek4",  true,  false, false, false,   kNever, kNever, kNever, kNever, kNever, kMono},
-    {"kimi-k3",    false, false, false, false,   kNever, kNever, kNever, kNever, kNever, kMono},
+    {"kimi-k3",    false, false, false, false,   kMono,  kNever, kNever, kNever, kNever, kMono},
 };
 
 inline constexpr std::size_t kArchCount =
diff --git a/server/src/common/step_graph.h b/server/src/common/step_graph.h
index b09a91341..8e364b046 100644
--- a/server/src/common/step_graph.h
+++ b/server/src/common/step_graph.h
@@ -32,6 +32,7 @@ struct StepGraph {
     int alloc_reserved_ctx = 0;
 
     // Named inputs
+    ggml_tensor *   token_ids = nullptr;     // embedding-only utility graphs
     ggml_tensor *   inp_embed = nullptr;
     ggml_tensor *   positions = nullptr;
     ggml_tensor *   attn_mask = nullptr;     // may be null
@@ -71,6 +72,7 @@ struct StepGraph {
 inline void step_graph_free(StepGraph & sg) {
     if (sg.ctx)   { ggml_free(sg.ctx); sg.ctx = nullptr; }
     sg.gf = nullptr;
+    sg.token_ids = nullptr;
     sg.inp_embed = sg.positions = sg.attn_mask = nullptr;
     sg.target_hidden_cat = sg.positions_k = nullptr;
     sg.pad_mask_full = nullptr;
diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp
index e5a04721c..05089db48 100644
--- a/server/src/draft/draft_gguf_loader.cpp
+++ b/server/src/draft/draft_gguf_loader.cpp
@@ -204,6 +204,8 @@ bool load_draft_gguf(const std::string & path,
     // Store GGUF-declared config into DraftWeights (replaces hardcoded defaults).
     out.block_size = (int)block_sz;
     out.n_target_layers = (int)n_tgt_lay;
+    out.mask_token_id = (int32_t)read_u32(
+        "dflash.mask_token_id", (uint32_t)out.mask_token_id);
 
     // Propagate target model properties if available.
     if (target) {
diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp
index c34bb7715..35e29245e 100644
--- a/server/src/kimi_k3/kimi_k3_backend.cpp
+++ b/server/src/kimi_k3/kimi_k3_backend.cpp
@@ -3,8 +3,10 @@
 #endif
 
 #include "kimi_k3_backend.h"
+#include "kimi_k3_dflash_target.h"
 
 #include "common/dynamic_backend.h"
+#include "common/dflash_spec_decode.h"
 #include "common/moe_expert_package.h"
 #include "common/moe_hybrid_placement.h"
 #include "common/moe_stream_cache_policy.h"
@@ -39,6 +41,8 @@
 namespace dflash::common {
 namespace {
 
+constexpr int kMaxDsparkBlockSize = 16;
+
 void close_file_descriptor(int fd) {
 #if defined(_WIN32)
     ::_close(fd);
@@ -701,6 +705,77 @@ bool KimiK3Backend::init_streaming() {
     return true;
 }
 
+bool KimiK3Backend::init_draft() {
+    if (!cfg_.draft_path || !*cfg_.draft_path) return true;
+    if (draft_backend_ || draft_weights_.ctx) return true;
+
+    draft_backend_ = ggml_backend_cuda_init(std::max(0, cfg_.draft_gpu));
+    if (!draft_backend_) {
+        std::fprintf(stderr,
+                     "[kimi-k3-dspark] draft backend init failed for device %d\n",
+                     cfg_.draft_gpu);
+        return false;
+    }
+    if (!load_draft_gguf(cfg_.draft_path, draft_backend_, draft_weights_)) {
+        std::fprintf(stderr,
+                     "[kimi-k3-dspark] draft load failed: %s\n",
+                     dflash27b_last_error());
+        free_drafter();
+        return false;
+    }
+
+    bool compatible =
+        draft_weights_.n_embd == weights_.n_embd &&
+        draft_weights_.block_size > 1 &&
+        draft_weights_.block_size <= kMaxDsparkBlockSize &&
+        draft_weights_.n_target_layers > 0 &&
+        draft_weights_.n_target_layers ==
+            static_cast<int>(draft_weights_.capture_layer_ids.size()) &&
+        draft_weights_.mask_token_id >= 0 &&
+        draft_weights_.mask_token_id < weights_.n_vocab;
+    for (int layer : draft_weights_.capture_layer_ids) {
+        compatible = compatible && layer >= 0 && layer < weights_.n_layer;
+    }
+    if (draft_weights_.dspark.enabled) {
+        compatible = compatible &&
+            draft_weights_.dspark.vocab_size == weights_.n_vocab;
+    }
+    if (!compatible) {
+        std::fprintf(stderr,
+            "[kimi-k3-dspark] incompatible checkpoint: target "
+            "hidden/vocab/layers=%d/%d/%d, draft hidden/block/captures/"
+            "mask/vocab=%d/%d/%zu/%d/%d\n",
+            weights_.n_embd, weights_.n_vocab, weights_.n_layer,
+            draft_weights_.n_embd, draft_weights_.block_size,
+            draft_weights_.capture_layer_ids.size(),
+            draft_weights_.mask_token_id,
+            draft_weights_.dspark.vocab_size);
+        free_drafter();
+        return false;
+    }
+
+    const int ring_cap = std::min(
+        std::max(1, cfg_.device.max_ctx),
+        std::max(2048, cfg_.draft_ctx_max));
+    if (!draft_feature_mirror_init(
+            feature_ring_, draft_backend_, std::max(0, cfg_.draft_gpu),
+            cfg_.device.primary_gpu(), ring_cap,
+            draft_weights_.n_target_layers, weights_.n_embd)) {
+        std::fprintf(stderr,
+                     "[kimi-k3-dspark] feature-ring allocation failed\n");
+        free_drafter();
+        return false;
+    }
+    std::fprintf(stderr,
+        "[kimi-k3-dspark] shared DFlash runtime enabled: block=%d "
+        "captures=%zu ring=%d draft_gpu=%d target_gpu=%d dspark=%d\n",
+        draft_weights_.block_size,
+        draft_weights_.capture_layer_ids.size(), ring_cap,
+        cfg_.draft_gpu, cfg_.device.primary_gpu(),
+        draft_weights_.dspark.enabled ? 1 : 0);
+    return true;
+}
+
 bool KimiK3Backend::init() {
     if (!cfg_.model_path) {
         std::fprintf(stderr, "[kimi-k3] model path is null\n");
@@ -721,8 +796,12 @@ bool KimiK3Backend::init() {
                      dflash27b_last_error());
         return false;
     }
+    if (!init_draft()) return false;
     const int max_ctx = std::max(1, cfg_.device.max_ctx);
-    if (!create_kimi_k3_cache(backend_, weights_, max_ctx, cache_)) {
+    const int max_verify_tokens = draft_weights_.ctx
+        ? draft_weights_.block_size : 0;
+    if (!create_kimi_k3_cache(
+            backend_, weights_, max_ctx, cache_, max_verify_tokens)) {
         std::fprintf(stderr, "[kimi-k3] cache allocation failed (max_ctx=%d)\n",
                      max_ctx);
         return false;
@@ -747,8 +826,13 @@ void KimiK3Backend::print_ready_banner() const {
 }
 
 bool KimiK3Backend::park(ParkTarget target) {
-    if (!park_target_includes_target_model(target)) return false;
-    if (!parked_) {
+    bool handled = false;
+    if (park_target_includes_draft_model(target) && draft_backend_) {
+        free_drafter();
+        handled = true;
+    }
+    if (park_target_includes_target_model(target) && !parked_) {
+        dflash_target_.reset();
         maybe_save_routing_stats();
         dual_stream_executor_.destroy();
         stream_engine_.destroy();
@@ -756,13 +840,14 @@ bool KimiK3Backend::park(ParkTarget target) {
         release_expert_backend();
         free_kimi_k3_weights(weights_);
         parked_ = true;
+        handled = true;
     }
-    return true;
+    return handled;
 }
 
 bool KimiK3Backend::unpark(ParkTarget target) {
-    if (!park_target_includes_target_model(target)) return false;
-    if (parked_) {
+    bool handled = false;
+    if (park_target_includes_target_model(target) && parked_) {
         if (!load_kimi_k3_gguf(
                 cfg_.model_path, backend_, weights_,
                 cfg_.moe_storage != MoeStoragePolicy::Resident) ||
@@ -770,8 +855,14 @@ bool KimiK3Backend::unpark(ParkTarget target) {
             return false;
         }
         parked_ = false;
+        handled = true;
     }
-    return true;
+    if (park_target_includes_draft_model(target) &&
+        cfg_.draft_path && *cfg_.draft_path && !draft_backend_) {
+        if (!init_draft()) return false;
+        handled = true;
+    }
+    return handled;
 }
 
 int32_t KimiK3Backend::choose_token(const std::vector<float> & logits,
@@ -785,6 +876,35 @@ int32_t KimiK3Backend::choose_token(const std::vector<float> & logits,
         std::max_element(logits.begin(), logits.end())));
 }
 
+bool KimiK3Backend::supports_dflash_spec_decode() const {
+    return draft_backend_ && draft_weights_.ctx && feature_ring_.target_feat;
+}
+
+DFlashTarget * KimiK3Backend::dflash_target() {
+    if (!supports_dflash_spec_decode()) return nullptr;
+    if (!dflash_target_) {
+        dflash_target_ = std::make_unique<KimiK3DFlashTarget>(
+            weights_, cache_, backend_, feature_ring_,
+            draft_weights_.capture_layer_ids,
+            draft_weights_.mask_token_id, cfg_.fast_rollback,
+            weights_.routed_experts_streamed ? &stream_engine_ : nullptr,
+            dual_stream_executor_.is_ready()
+                ? &dual_stream_executor_ : nullptr,
+            &stream_owner_policy_, routing_stats_.get());
+    }
+    return dflash_target_.get();
+}
+
+void KimiK3Backend::free_drafter() {
+    dflash_target_.reset();
+    draft_feature_mirror_free(feature_ring_);
+    if (draft_weights_.ctx) free_draft_weights(draft_weights_);
+    if (draft_backend_) {
+        ggml_backend_free(draft_backend_);
+        draft_backend_ = nullptr;
+    }
+}
+
 GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req,
                                             const DaemonIO & io) {
     GenerateResult result;
@@ -810,14 +930,19 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req,
 
     reset_kimi_k3_cache(cache_);
     std::vector<float> logits;
+    auto * spec_target = static_cast<KimiK3DFlashTarget *>(dflash_target());
     const auto prefill_begin = std::chrono::steady_clock::now();
     for (size_t i = 0; i < req.prompt.size(); ++i) {
-        if (!kimi_k3_step(
+        const bool ok = spec_target
+            ? spec_target->forward_token(
+                req.prompt[i], static_cast<int>(i), logits)
+            : kimi_k3_step(
                 backend_, weights_, cache_, req.prompt[i],
                 static_cast<int>(i), logits, &stream_engine_,
                 dual_stream_executor_.is_ready()
                     ? &dual_stream_executor_ : nullptr,
-                &stream_owner_policy_, routing_stats_.get())) {
+                &stream_owner_policy_, routing_stats_.get());
+        if (!ok) {
             result.fail(GenerateErrorCode::PrefillFailed,
                         dflash27b_last_error());
             out_io.emit(-1);
@@ -827,7 +952,45 @@ GenerateResult KimiK3Backend::generate_impl(const GenerateRequest & req,
     const auto prefill_end = std::chrono::steady_clock::now();
     result.prefill_s = std::chrono::duration<double>(prefill_end - prefill_begin).count();
 
+    if (req.n_gen <= 0 || out_io.cancelled) {
+        maybe_save_routing_stats();
+        out_io.emit(-1);
+        result.succeed();
+        return result;
+    }
+
     const auto decode_begin = std::chrono::steady_clock::now();
+    const bool can_spec = spec_target && !req.force_ar_decode &&
+        req.budget_hook.close_token_ids.empty() &&
+        !req.sampler.needs_logit_processing();
+    if (can_spec) {
+        const int32_t seed = choose_token(logits, req.sampler, result.tokens);
+        DaemonIO spec_io = out_io.with_token_callback(
+            [&](int32_t token) -> bool {
+                result.tokens.push_back(token);
+                return true;
+            });
+        double accept_rate = 0.0;
+        const bool ok = run_dflash_spec_decode(
+            *spec_target, draft_weights_, draft_backend_, feature_ring_,
+            req.prompt, req.n_gen, seed, /*out_path=*/nullptr,
+            cfg_.draft_ctx_max, spec_io, /*remote_draft=*/nullptr,
+            req.hint_tokens, /*base_pos=*/0, &accept_rate);
+        result.decode_s = std::chrono::duration<double>(
+            std::chrono::steady_clock::now() - decode_begin).count();
+        result.accept_rate = static_cast<float>(accept_rate);
+        result.spec_decode_ran = true;
+        maybe_save_routing_stats();
+        spec_io.emit(-1);
+        if (!ok) {
+            result.fail(GenerateErrorCode::DecodeFailed,
+                        dflash27b_last_error());
+            return result;
+        }
+        result.succeed();
+        return result;
+    }
+
     bool budget_close_started = false;
     size_t close_inject_pos = 0;
     for (int i = 0; i < req.n_gen; ++i) {
@@ -917,6 +1080,7 @@ bool KimiK3Backend::handle_compress(const std::string & line,
 }
 
 void KimiK3Backend::shutdown() {
+    free_drafter();
     maybe_save_routing_stats();
     dual_stream_executor_.destroy();
     stream_engine_.destroy();
diff --git a/server/src/kimi_k3/kimi_k3_backend.h b/server/src/kimi_k3/kimi_k3_backend.h
index f5472e992..0df923ccd 100644
--- a/server/src/kimi_k3/kimi_k3_backend.h
+++ b/server/src/kimi_k3/kimi_k3_backend.h
@@ -1,9 +1,11 @@
 #pragma once
 
 #include "common/model_backend.h"
+#include "common/dflash_feature_ring.h"
 #include "common/moe_hybrid_routing_stats.h"
 #include "common/moe_hybrid_stream.h"
 #include "common/moe_storage_policy.h"
+#include "internal.h"
 #include "kimi_k3_internal.h"
 #include "placement/placement_config.h"
 
@@ -11,11 +13,17 @@
 #include <memory>
 #include <string>
 
+struct ggml_backend;
+
 namespace dflash::common {
 
 struct KimiK3BackendConfig {
     const char * model_path = nullptr;
+    const char * draft_path = nullptr;
     DevicePlacement device;
+    int draft_gpu = 0;
+    int draft_ctx_max = 4096;
+    bool fast_rollback = true;
     int stream_fd = -1;
     // -1 resolves DFLASH_MOE_TP_GPU and otherwise keeps the primary device
     // index. DFLASH_MOE_TP_BACKEND may select a different in-process runtime
@@ -53,11 +61,14 @@ class KimiK3Backend final : public ModelBackend {
 
     bool handle_compress(const std::string & line,
                          const DaemonIO & io) override;
-    void free_drafter() override {}
+    void free_drafter() override;
+    bool supports_dflash_spec_decode() const override;
+    DFlashTarget * dflash_target() override;
     void shutdown() override;
 
 private:
     bool init_streaming();
+    bool init_draft();
     void release_expert_backend();
     void maybe_save_routing_stats();
 
@@ -67,11 +78,15 @@ class KimiK3Backend final : public ModelBackend {
 
     KimiK3BackendConfig cfg_;
     ggml_backend_t backend_ = nullptr;
+    ggml_backend_t draft_backend_ = nullptr;
     ggml_backend_t expert_backend_ = nullptr;
     PlacementBackend expert_backend_kind_ = PlacementBackend::Auto;
     int expert_gpu_ = -1;
     KimiK3Weights weights_;
     KimiK3Cache cache_;
+    DraftWeights draft_weights_;
+    DraftFeatureMirror feature_ring_;
+    std::unique_ptr<class KimiK3DFlashTarget> dflash_target_;
     MoeHybridStreamEngine stream_engine_;
     MoeHybridStreamEngine secondary_stream_engine_;
     MoeStreamDualOwnerExecutor dual_stream_executor_;
diff --git a/server/src/kimi_k3/kimi_k3_dflash_target.cpp b/server/src/kimi_k3/kimi_k3_dflash_target.cpp
new file mode 100644
index 000000000..eeb1f9a69
--- /dev/null
+++ b/server/src/kimi_k3/kimi_k3_dflash_target.cpp
@@ -0,0 +1,277 @@
+#include "kimi_k3_dflash_target.h"
+
+#include "common/dflash_feature_ring.h"
+
+#include "ggml-alloc.h"
+
+#include <cstdio>
+#include <utility>
+
+namespace dflash::common {
+
+KimiK3DFlashTarget::KimiK3DFlashTarget(
+        KimiK3Weights & weights,
+        KimiK3Cache & cache,
+        ggml_backend_t backend,
+        DraftFeatureMirror & feature_ring,
+        std::vector<int> capture_layer_ids,
+        int mask_token_id,
+        bool fast_rollback,
+        MoeHybridStreamEngine * stream_engine,
+        MoeStreamDualOwnerExecutor * dual_stream_executor,
+        const MoeStreamDualOwnerPolicy * stream_owner_policy,
+        MoeHybridRoutingStats * routing_stats)
+    : weights_(weights),
+      cache_(cache),
+      backend_(backend),
+      feature_ring_(feature_ring),
+      capture_layer_ids_(std::move(capture_layer_ids)),
+      mask_token_id_(mask_token_id),
+      fast_rollback_(fast_rollback),
+      stream_engine_(stream_engine),
+      dual_stream_executor_(dual_stream_executor),
+      stream_owner_policy_(stream_owner_policy),
+      routing_stats_(routing_stats) {}
+
+KimiK3DFlashTarget::~KimiK3DFlashTarget() {
+    step_graph_destroy(embedding_graph_);
+    step_graph_destroy(projection_graph_);
+}
+
+bool KimiK3DFlashTarget::sync_captures(
+        const KimiK3ForwardResult & result,
+        int base_pos,
+        int n_tokens) {
+    const size_t capture_values =
+        static_cast<size_t>(weights_.n_embd) * n_tokens;
+    if (result.captured_hidden.size() !=
+        capture_values * capture_layer_ids_.size()) {
+        std::fprintf(stderr,
+            "[kimi-k3-dspark] target capture shape mismatch: got=%zu expected=%zu\n",
+            result.captured_hidden.size(),
+            capture_values * capture_layer_ids_.size());
+        return false;
+    }
+    for (size_t i = 0; i < capture_layer_ids_.size(); ++i) {
+        if (!copy_host_capture_slice_to_draft_ring(
+                feature_ring_, static_cast<int>(i), base_pos, n_tokens,
+                result.captured_hidden.data() + i * capture_values,
+                capture_values)) {
+            std::fprintf(stderr,
+                "[kimi-k3-dspark] feature-ring copy failed at capture %zu\n", i);
+            return false;
+        }
+    }
+    return true;
+}
+
+bool KimiK3DFlashTarget::forward_token(
+        int32_t token, int position, std::vector<float> & logits) {
+    KimiK3ForwardOptions options;
+    options.capture_layer_ids = &capture_layer_ids_;
+    options.read_logits = true;
+    options.read_argmax = false;
+    KimiK3ForwardResult result;
+    if (!kimi_k3_forward(
+            backend_, weights_, cache_, std::vector<int32_t>{token}, position,
+            options, result, stream_engine_, dual_stream_executor_,
+            stream_owner_policy_, routing_stats_) ||
+        !sync_captures(result, position, 1)) {
+        return false;
+    }
+    logits = std::move(result.logits);
+    return true;
+}
+
+bool KimiK3DFlashTarget::verify_batch(
+        const std::vector<int32_t> & tokens,
+        int base_pos,
+        int & last_tok,
+        std::vector<int32_t> * all_argmax,
+        bool capture_ssm_intermediates) {
+    KimiK3ForwardOptions options;
+    options.capture_layer_ids = &capture_layer_ids_;
+    options.capture_replay = fast_rollback_ && capture_ssm_intermediates;
+    options.read_logits = false;
+    options.read_argmax = true;
+    KimiK3ForwardResult result;
+    if (!kimi_k3_forward(
+            backend_, weights_, cache_, tokens, base_pos,
+            options, result, stream_engine_, dual_stream_executor_,
+            stream_owner_policy_, routing_stats_) ||
+        result.argmax.size() != tokens.size() ||
+        !sync_captures(result, base_pos, static_cast<int>(tokens.size()))) {
+        return false;
+    }
+    last_tok = result.argmax.back();
+    if (all_argmax) *all_argmax = std::move(result.argmax);
+    return true;
+}
+
+bool KimiK3DFlashTarget::snapshot_kv() {
+    return kimi_k3_replay_snapshot(backend_, cache_);
+}
+
+bool KimiK3DFlashTarget::restore_kv() {
+    return kimi_k3_replay_restore(backend_, cache_);
+}
+
+bool KimiK3DFlashTarget::supports_fast_rollback() const {
+    return fast_rollback_ && cache_.max_verify_tokens > 0 &&
+           cache_.snapshot_valid && cache_.replay_valid;
+}
+
+bool KimiK3DFlashTarget::prefer_fast_rollback_over_replay() const {
+    // ReplaySSM recomputes only the recurrent KDA transitions; replaying a
+    // token would also reread every routed MoE layer from external storage.
+    return stream_engine_ != nullptr;
+}
+
+bool KimiK3DFlashTarget::rollback_to(int base_pos, int commit_n) {
+    return kimi_k3_replay_commit(
+        backend_, weights_, cache_, base_pos, commit_n);
+}
+
+bool KimiK3DFlashTarget::is_eos(int token) const {
+    return token == weights_.eos_token_id;
+}
+
+bool KimiK3DFlashTarget::build_embedding_graph(int n_tokens) const {
+    StepGraph & graph = embedding_graph_;
+    step_graph_free(graph);
+    ggml_init_params params{};
+    params.mem_size = 4ull * 1024ull * 1024ull;
+    params.no_alloc = true;
+    graph.ctx = ggml_init(params);
+    if (!graph.ctx) return false;
+    graph.gf = ggml_new_graph_custom(graph.ctx, 256, false);
+    graph.token_ids = ggml_new_tensor_1d(
+        graph.ctx, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(graph.token_ids);
+    graph.hidden_states = ggml_get_rows(
+        graph.ctx, weights_.tok_embd, graph.token_ids);
+    if (graph.hidden_states->type != GGML_TYPE_F32) {
+        graph.hidden_states = ggml_cast(
+            graph.ctx, graph.hidden_states, GGML_TYPE_F32);
+    }
+    ggml_set_output(graph.hidden_states);
+    ggml_build_forward_expand(graph.gf, graph.hidden_states);
+    if (!graph.alloc) {
+        graph.alloc = ggml_gallocr_new(
+            ggml_backend_get_default_buffer_type(backend_));
+    }
+    return graph.alloc && ggml_gallocr_alloc_graph(graph.alloc, graph.gf);
+}
+
+bool KimiK3DFlashTarget::embed_tokens(
+        const int32_t * tokens, int n, float * out) const {
+    if (!tokens || !out || n <= 0 || !build_embedding_graph(n)) return false;
+    ggml_backend_tensor_set(
+        embedding_graph_.token_ids, tokens, 0, sizeof(int32_t) * n);
+    if (ggml_backend_graph_compute(backend_, embedding_graph_.gf) !=
+        GGML_STATUS_SUCCESS) {
+        return false;
+    }
+    ggml_backend_tensor_get(
+        embedding_graph_.hidden_states, out, 0,
+        sizeof(float) * static_cast<size_t>(weights_.n_embd) * n);
+    return true;
+}
+
+bool KimiK3DFlashTarget::build_projection_graph(int n_tokens) {
+    StepGraph & graph = projection_graph_;
+    step_graph_free(graph);
+    ggml_init_params params{};
+    params.mem_size = 4ull * 1024ull * 1024ull;
+    params.no_alloc = true;
+    graph.ctx = ggml_init(params);
+    if (!graph.ctx) return false;
+    graph.gf = ggml_new_graph_custom(graph.ctx, 256, false);
+    graph.hidden_input = ggml_new_tensor_2d(
+        graph.ctx, GGML_TYPE_F32, weights_.n_embd, n_tokens);
+    ggml_set_input(graph.hidden_input);
+    graph.logits = ggml_mul_mat(
+        graph.ctx, weights_.output, graph.hidden_input);
+    graph.argmax_tokens = ggml_argmax(graph.ctx, graph.logits);
+    ggml_set_output(graph.logits);
+    ggml_set_output(graph.argmax_tokens);
+    ggml_build_forward_expand(graph.gf, graph.logits);
+    ggml_build_forward_expand(graph.gf, graph.argmax_tokens);
+    if (!graph.alloc) {
+        graph.alloc = ggml_gallocr_new(
+            ggml_backend_get_default_buffer_type(backend_));
+    }
+    return graph.alloc && ggml_gallocr_alloc_graph(graph.alloc, graph.gf);
+}
+
+bool KimiK3DFlashTarget::project_hidden_to_tokens(
+        const float * hidden,
+        int n_tokens,
+        std::vector<int32_t> & tokens_out) {
+    if (!hidden || n_tokens <= 0 || !build_projection_graph(n_tokens)) return false;
+    ggml_backend_tensor_set(
+        projection_graph_.hidden_input, hidden, 0,
+        sizeof(float) * static_cast<size_t>(weights_.n_embd) * n_tokens);
+    if (ggml_backend_graph_compute(backend_, projection_graph_.gf) !=
+        GGML_STATUS_SUCCESS) {
+        return false;
+    }
+    tokens_out.resize(static_cast<size_t>(n_tokens));
+    ggml_backend_tensor_get(
+        projection_graph_.argmax_tokens, tokens_out.data(), 0,
+        sizeof(int32_t) * tokens_out.size());
+    return true;
+}
+
+bool KimiK3DFlashTarget::project_hidden_to_logits(
+        const float * hidden,
+        int n_tokens,
+        std::vector<float> & logits_out) {
+    if (!hidden || n_tokens <= 0 || !build_projection_graph(n_tokens)) return false;
+    ggml_backend_tensor_set(
+        projection_graph_.hidden_input, hidden, 0,
+        sizeof(float) * static_cast<size_t>(weights_.n_embd) * n_tokens);
+    if (ggml_backend_graph_compute(backend_, projection_graph_.gf) !=
+        GGML_STATUS_SUCCESS) {
+        return false;
+    }
+    logits_out.resize(
+        static_cast<size_t>(weights_.n_vocab) * n_tokens);
+    ggml_backend_tensor_get(
+        projection_graph_.logits, logits_out.data(), 0,
+        sizeof(float) * logits_out.size());
+    return true;
+}
+
+ggml_tensor * KimiK3DFlashTarget::lm_head_tensor() {
+    return weights_.output;
+}
+
+ggml_tensor * KimiK3DFlashTarget::gpu_embd_table() {
+    return weights_.tok_embd;
+}
+
+ggml_backend_t KimiK3DFlashTarget::fused_head_backend() {
+    return backend_;
+}
+
+int KimiK3DFlashTarget::hidden_size() const {
+    return weights_.n_embd;
+}
+
+int KimiK3DFlashTarget::mask_token_id() const {
+    return mask_token_id_;
+}
+
+const std::vector<int> & KimiK3DFlashTarget::capture_layer_ids() const {
+    return capture_layer_ids_;
+}
+
+int KimiK3DFlashTarget::default_adaptive_verify_min_rows() const {
+    // Each additional row can route another set of file-backed experts.  The
+    // DSpark confidence head should therefore be allowed to stop after the
+    // seed row; resident targets retain the shared runtime's default floor.
+    return stream_engine_ ? 1 : 0;
+}
+
+} // namespace dflash::common
diff --git a/server/src/kimi_k3/kimi_k3_dflash_target.h b/server/src/kimi_k3/kimi_k3_dflash_target.h
new file mode 100644
index 000000000..1b63faa72
--- /dev/null
+++ b/server/src/kimi_k3/kimi_k3_dflash_target.h
@@ -0,0 +1,90 @@
+#pragma once
+
+#include "common/dflash_target.h"
+#include "kimi_k3_internal.h"
+#include "step_graph.h"
+
+#include <vector>
+
+namespace dflash::common {
+
+struct DraftFeatureMirror;
+
+// Kimi-specific state/capture adapter for the shared DFlash/DSpark runtime.
+// Draft execution, Markov correction, confidence gating, acceptance, and
+// scheduling deliberately remain in common/.  This class owns only the Kimi
+// forward and ReplaySSM boundary.
+class KimiK3DFlashTarget final : public DFlashTarget {
+public:
+    KimiK3DFlashTarget(
+        KimiK3Weights & weights,
+        KimiK3Cache & cache,
+        ggml_backend_t backend,
+        DraftFeatureMirror & feature_ring,
+        std::vector<int> capture_layer_ids,
+        int mask_token_id,
+        bool fast_rollback,
+        MoeHybridStreamEngine * stream_engine,
+        MoeStreamDualOwnerExecutor * dual_stream_executor,
+        const MoeStreamDualOwnerPolicy * stream_owner_policy,
+        MoeHybridRoutingStats * routing_stats);
+    ~KimiK3DFlashTarget() override;
+
+    KimiK3DFlashTarget(const KimiK3DFlashTarget &) = delete;
+    KimiK3DFlashTarget & operator=(const KimiK3DFlashTarget &) = delete;
+
+    // Sequential prefill/AR entry point that shares the same feature-capture
+    // code as speculative verification.
+    bool forward_token(int32_t token, int position, std::vector<float> & logits);
+
+    bool verify_batch(const std::vector<int32_t> & tokens,
+                      int base_pos,
+                      int & last_tok,
+                      std::vector<int32_t> * all_argmax = nullptr,
+                      bool capture_ssm_intermediates = false) override;
+    bool snapshot_kv() override;
+    bool restore_kv() override;
+    bool supports_fast_rollback() const override;
+    bool prefer_fast_rollback_over_replay() const override;
+    bool rollback_to(int base_pos, int commit_n) override;
+
+    bool is_eos(int token) const override;
+    bool embed_tokens(const int32_t * tokens, int n, float * out) const override;
+    bool project_hidden_to_tokens(const float * hidden,
+                                  int n_tokens,
+                                  std::vector<int32_t> & tokens_out) override;
+    bool project_hidden_to_logits(const float * hidden,
+                                  int n_tokens,
+                                  std::vector<float> & logits_out) override;
+
+    ggml_tensor * lm_head_tensor() override;
+    ggml_tensor * gpu_embd_table() override;
+    ggml_backend_t fused_head_backend() override;
+    int hidden_size() const override;
+    int mask_token_id() const override;
+    const std::vector<int> & capture_layer_ids() const override;
+    int default_adaptive_verify_min_rows() const override;
+
+private:
+    bool sync_captures(const KimiK3ForwardResult & result,
+                       int base_pos,
+                       int n_tokens);
+    bool build_embedding_graph(int n_tokens) const;
+    bool build_projection_graph(int n_tokens);
+
+    KimiK3Weights & weights_;
+    KimiK3Cache & cache_;
+    ggml_backend_t backend_ = nullptr;
+    DraftFeatureMirror & feature_ring_;
+    std::vector<int> capture_layer_ids_;
+    int mask_token_id_ = -1;
+    bool fast_rollback_ = false;
+    MoeHybridStreamEngine * stream_engine_ = nullptr;
+    MoeStreamDualOwnerExecutor * dual_stream_executor_ = nullptr;
+    const MoeStreamDualOwnerPolicy * stream_owner_policy_ = nullptr;
+    MoeHybridRoutingStats * routing_stats_ = nullptr;
+    mutable StepGraph embedding_graph_;
+    StepGraph projection_graph_;
+};
+
+} // namespace dflash::common
diff --git a/server/src/kimi_k3/kimi_k3_graph.cpp b/server/src/kimi_k3/kimi_k3_graph.cpp
index 38788a4b4..3826f5823 100644
--- a/server/src/kimi_k3/kimi_k3_graph.cpp
+++ b/server/src/kimi_k3/kimi_k3_graph.cpp
@@ -46,52 +46,61 @@ struct AttnResBank {
     ggml_context * ctx = nullptr;
     float eps = 1.0e-5f;
     int64_t n_embd = 0;
+    int64_t n_tokens = 1;
     std::vector<ggml_tensor *> checkpoints;
-    ggml_tensor * stack = nullptr;
-    size_t stack_size = 0;
 
     void push(ggml_tensor * cur) {
-        checkpoints.push_back(ggml_reshape_3d(ctx, cur, n_embd, 1, 1));
-    }
-
-    ggml_tensor * get_stack() {
-        if (stack && stack_size == checkpoints.size()) return stack;
-        stack = checkpoints.front();
-        for (size_t i = 1; i < checkpoints.size(); ++i) {
-            stack = ggml_concat(ctx, stack, checkpoints[i], 1);
-        }
-        stack_size = checkpoints.size();
-        return stack;
+        checkpoints.push_back(
+            ggml_reshape_3d(ctx, cur, n_embd, n_tokens, 1));
     }
 
     ggml_tensor * mix(ggml_tensor * cur, ggml_tensor * score_weight) {
         if (checkpoints.empty()) return cur;
         const int64_t n = static_cast<int64_t>(checkpoints.size());
-        ggml_tensor * src = get_stack(); // [hidden, n_checkpoint, 1]
-
-        ggml_tensor * score_src = rms_norm(ctx, src, score_weight, eps);
-        score_src = ggml_sum_rows(ctx, score_src);
-        score_src = ggml_reshape_2d(ctx, score_src, n, 1);
-
-        ggml_tensor * score_cur = rms_norm(ctx, cur, score_weight, eps);
-        score_cur = ggml_sum_rows(ctx, score_cur);
-
-        ggml_tensor * probs = ggml_soft_max(ctx,
-            ggml_concat(ctx, score_src, score_cur, 0));
-        ggml_tensor * p_src = ggml_cont(ctx,
-            ggml_view_2d(ctx, probs, n, 1, probs->nb[1], 0));
-        ggml_tensor * p_cur = ggml_cont(ctx,
-            ggml_view_2d(ctx, probs, 1, 1, probs->nb[1],
-                         probs->nb[0] * static_cast<size_t>(n)));
-
-        // Reduce checkpoint dimension with an ordinary matrix product. The
-        // newer upstream ggml has a dedicated dsv4_hc_pre op for this exact
-        // contraction; the Lucebox ggml snapshot predates that API, and this
-        // algebra is identical: [checkpoint, hidden] x [checkpoint, 1].
-        ggml_tensor * src_t = ggml_cont(ctx,
-            ggml_permute(ctx, src, 1, 0, 2, 3));
-        ggml_tensor * out = ggml_mul_mat(ctx, src_t, p_src);
-        return ggml_add(ctx, out, ggml_mul(ctx, cur, p_cur));
+        ggml_tensor * mixed = nullptr;
+
+        // AttnRes chooses a different checkpoint mixture for every token.
+        // Express the small contraction as independent token slices.  This is
+        // the exact one-token algebra used by the original Kimi path, avoids
+        // relying on fragile rank-3 broadcast rules, and is cheap at the
+        // bounded speculative widths (currently <= 16).
+        for (int64_t token = 0; token < n_tokens; ++token) {
+            ggml_tensor * src = nullptr; // [hidden, checkpoint]
+            for (ggml_tensor * checkpoint : checkpoints) {
+                ggml_tensor * checkpoint_token = ggml_view_2d(
+                    ctx, checkpoint, n_embd, 1, checkpoint->nb[1],
+                    static_cast<size_t>(token) * checkpoint->nb[1]);
+                src = src
+                    ? ggml_concat(ctx, src, checkpoint_token, 1)
+                    : checkpoint_token;
+            }
+            ggml_tensor * cur_token = ggml_view_2d(
+                ctx, cur, n_embd, 1, cur->nb[1],
+                static_cast<size_t>(token) * cur->nb[1]);
+
+            ggml_tensor * score_src = rms_norm(ctx, src, score_weight, eps);
+            score_src = ggml_reshape_2d(
+                ctx, ggml_sum_rows(ctx, score_src), n, 1);
+            ggml_tensor * score_cur = ggml_sum_rows(
+                ctx, rms_norm(ctx, cur_token, score_weight, eps));
+            ggml_tensor * probs = ggml_soft_max(
+                ctx, ggml_concat(ctx, score_src, score_cur, 0));
+            ggml_tensor * p_src = ggml_cont(ctx,
+                ggml_view_2d(ctx, probs, n, 1, probs->nb[1], 0));
+            ggml_tensor * p_cur = ggml_cont(ctx,
+                ggml_view_2d(ctx, probs, 1, 1, probs->nb[1],
+                             probs->nb[0] * static_cast<size_t>(n)));
+
+            ggml_tensor * src_t = ggml_cont(
+                ctx, ggml_permute(ctx, src, 1, 0, 2, 3));
+            ggml_tensor * out_token = ggml_add(
+                ctx, ggml_mul_mat(ctx, src_t, p_src),
+                ggml_mul(ctx, cur_token, p_cur));
+            mixed = mixed
+                ? ggml_concat(ctx, mixed, out_token, 1)
+                : out_token;
+        }
+        return mixed;
     }
 };
 
@@ -104,28 +113,32 @@ ggml_tensor * kda_conv1d(ggml_context * ctx,
                          ggml_tensor * conv_weight,
                          int d_conv,
                          int head_dim,
-                         int n_head) {
+                         int n_head,
+                         bool commit_state) {
     const int64_t d_inner = static_cast<int64_t>(head_dim) * n_head;
     const int64_t state_rows = d_conv - 1;
+    const int64_t n_tokens = x->ne[1];
     const size_t block_offset = static_cast<size_t>(qkv) * d_inner * all_state->nb[1];
     ggml_tensor * state = ggml_view_3d(ctx, all_state,
         state_rows, d_inner, 1, all_state->nb[1], all_state->nb[2],
         block_offset);
 
     ggml_tensor * projected = ggml_mul_mat(ctx, projection, x);
-    projected = ggml_reshape_3d(ctx, projected, d_inner, 1, 1);
+    projected = ggml_reshape_3d(ctx, projected, d_inner, n_tokens, 1);
     ggml_tensor * conv_input = ggml_concat(ctx, state,
                                            ggml_transpose(ctx, projected), 0);
 
     // Drop the oldest row and persist the newest d_conv-1 values.
-    ggml_tensor * newest = ggml_view_3d(ctx, conv_input,
-        state_rows, d_inner, 1, conv_input->nb[1], conv_input->nb[2],
-        conv_input->nb[0]);
-    ggml_build_forward_expand(graph, ggml_cpy(ctx, newest, state));
+    if (commit_state) {
+        ggml_tensor * newest = ggml_view_3d(ctx, conv_input,
+            state_rows, d_inner, 1, conv_input->nb[1], conv_input->nb[2],
+            static_cast<size_t>(n_tokens) * conv_input->nb[0]);
+        ggml_build_forward_expand(graph, ggml_cpy(ctx, newest, state));
+    }
 
     ggml_tensor * cw = ggml_reshape_2d(ctx, conv_weight, d_conv, d_inner);
     ggml_tensor * out = ggml_silu(ctx, ggml_ssm_conv(ctx, conv_input, cw));
-    out = ggml_reshape_4d(ctx, out, head_dim, n_head, 1, 1);
+    out = ggml_reshape_4d(ctx, out, head_dim, n_head, n_tokens, 1);
     return out;
 }
 
@@ -134,36 +147,52 @@ ggml_tensor * build_kda(ggml_context * ctx,
                         const KimiK3Weights & w,
                         const KimiK3Layer & layer,
                         KimiK3LayerCache & cache,
-                        ggml_tensor * cur) {
+                        ggml_tensor * cur,
+                        bool commit_state,
+                        bool capture_replay) {
     const int head_dim = w.kda_head_dim;
     const int n_head = w.n_head;
+    const int n_tokens = static_cast<int>(cur->ne[1]);
     const int64_t d_inner = static_cast<int64_t>(head_dim) * n_head;
 
+    if (capture_replay) {
+        GGML_ASSERT(cache.replay_input != nullptr);
+        GGML_ASSERT(n_tokens <= cache.replay_input->ne[1]);
+        ggml_tensor * replay_dst = ggml_view_2d(
+            ctx, cache.replay_input, w.n_embd, n_tokens,
+            cache.replay_input->nb[1], 0);
+        ggml_build_forward_expand(graph, ggml_cpy(ctx, cur, replay_dst));
+    }
+
     ggml_tensor * q = kda_conv1d(ctx, graph, cache.conv_state, 0, cur,
-        layer.wq, layer.ssm_q_conv, w.ssm_d_conv, head_dim, n_head);
+        layer.wq, layer.ssm_q_conv, w.ssm_d_conv, head_dim, n_head,
+        commit_state);
     ggml_tensor * k = kda_conv1d(ctx, graph, cache.conv_state, 1, cur,
-        layer.wk, layer.ssm_k_conv, w.ssm_d_conv, head_dim, n_head);
+        layer.wk, layer.ssm_k_conv, w.ssm_d_conv, head_dim, n_head,
+        commit_state);
     ggml_tensor * v = kda_conv1d(ctx, graph, cache.conv_state, 2, cur,
-        layer.wv, layer.ssm_v_conv, w.ssm_d_conv, head_dim, n_head);
+        layer.wv, layer.ssm_v_conv, w.ssm_d_conv, head_dim, n_head,
+        commit_state);
 
     ggml_tensor * decay = ggml_mul_mat(ctx, layer.ssm_f_a, cur);
     decay = ggml_mul_mat(ctx, layer.ssm_f_b, decay);
     decay = ggml_add(ctx, decay, layer.ssm_dt_b);
     ggml_tensor * A = ggml_reshape_3d(ctx, layer.ssm_a, 1, n_head, 1);
     if (std::isfinite(w.kda_gate_lower_bound)) {
-        decay = ggml_reshape_3d(ctx, decay, head_dim, n_head, 1);
+        decay = ggml_reshape_3d(ctx, decay, head_dim, n_head, n_tokens);
         decay = ggml_mul(ctx, decay, A);
         decay = ggml_sigmoid(ctx, ggml_scale(ctx, decay, -1.0f));
         decay = ggml_scale(ctx, decay, w.kda_gate_lower_bound);
     } else {
         decay = ggml_softplus(ctx, decay);
-        decay = ggml_reshape_3d(ctx, decay, head_dim, n_head, 1);
+        decay = ggml_reshape_3d(ctx, decay, head_dim, n_head, n_tokens);
         decay = ggml_mul(ctx, decay, A);
     }
-    decay = ggml_reshape_4d(ctx, decay, head_dim, n_head, 1, 1);
+    decay = ggml_reshape_4d(ctx, decay, head_dim, n_head, n_tokens, 1);
 
     ggml_tensor * beta = ggml_mul_mat(ctx, layer.ssm_beta, cur);
-    beta = ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta, 1, n_head, 1, 1));
+    beta = ggml_sigmoid(ctx,
+        ggml_reshape_4d(ctx, beta, 1, n_head, n_tokens, 1));
 
     q = ggml_l2_norm(ctx, q, w.rms_eps);
     k = ggml_l2_norm(ctx, k, w.rms_eps);
@@ -174,25 +203,27 @@ ggml_tensor * build_kda(ggml_context * ctx,
 
     const size_t elt = ggml_element_size(packed);
     ggml_tensor * output = ggml_view_4d(ctx, packed,
-        head_dim, n_head, 1, 1,
+        head_dim, n_head, n_tokens, 1,
         static_cast<size_t>(head_dim) * elt,
         static_cast<size_t>(head_dim) * n_head * elt,
-        static_cast<size_t>(head_dim) * n_head * elt, 0);
+        static_cast<size_t>(head_dim) * n_head * n_tokens * elt, 0);
     ggml_tensor * new_state = ggml_view_4d(ctx, packed,
         head_dim, head_dim, n_head, 1,
         static_cast<size_t>(head_dim) * elt,
         static_cast<size_t>(head_dim) * head_dim * elt,
         static_cast<size_t>(head_dim) * head_dim * n_head * elt,
-        static_cast<size_t>(head_dim) * n_head * elt);
-    ggml_build_forward_expand(graph,
-        ggml_cpy(ctx, new_state, cache.ssm_state));
+        static_cast<size_t>(head_dim) * n_head * n_tokens * elt);
+    if (commit_state) {
+        ggml_build_forward_expand(graph,
+            ggml_cpy(ctx, new_state, cache.ssm_state));
+    }
 
     ggml_tensor * gate = ggml_mul_mat(ctx, layer.ssm_g, cur);
-    gate = ggml_reshape_3d(ctx, gate, head_dim, n_head, 1);
-    output = ggml_reshape_3d(ctx, output, head_dim, n_head, 1);
+    gate = ggml_reshape_3d(ctx, gate, head_dim, n_head, n_tokens);
+    output = ggml_reshape_3d(ctx, output, head_dim, n_head, n_tokens);
     output = rms_norm(ctx, output, layer.ssm_o_norm, w.rms_eps);
     output = ggml_mul(ctx, output, ggml_sigmoid(ctx, gate));
-    output = ggml_cont_2d(ctx, output, d_inner, 1);
+    output = ggml_cont_2d(ctx, output, d_inner, n_tokens);
     return ggml_mul_mat(ctx, layer.wo, output);
 }
 
@@ -202,7 +233,8 @@ ggml_tensor * build_mla(ggml_context * ctx,
                         const KimiK3Layer & layer,
                         KimiK3LayerCache & cache,
                         ggml_tensor * cur,
-                        int position) {
+                        int position,
+                        ggml_tensor * attn_mask) {
     const int n_head = w.n_head;
     const int kv_rank = w.kv_lora_rank;
     const int key_dim = w.mla_k_head_dim;
@@ -210,7 +242,8 @@ ggml_tensor * build_mla(ggml_context * ctx,
     const int rope_dim = w.rope_dim;
     const int nope_dim = key_dim - rope_dim;
     const int compact_dim = kv_rank + rope_dim;
-    const int kv_len = position + 1;
+    const int n_tokens = static_cast<int>(cur->ne[1]);
+    const int kv_len = position + n_tokens;
 
     ggml_tensor * gate_input = cur;
     ggml_tensor * q_cur = nullptr;
@@ -223,18 +256,18 @@ ggml_tensor * build_mla(ggml_context * ctx,
     }
 
     ggml_tensor * compact_pe = ggml_mul_mat(ctx, layer.wkv_a_mqa, cur);
-    ggml_tensor * compact = ggml_view_2d(ctx, compact_pe, kv_rank, 1,
+    ggml_tensor * compact = ggml_view_2d(ctx, compact_pe, kv_rank, n_tokens,
         ggml_row_size(compact_pe->type, compact_dim), 0);
-    ggml_tensor * k_pe = ggml_view_3d(ctx, compact_pe, rope_dim, 1, 1,
-        ggml_row_size(compact_pe->type, compact_dim),
+    ggml_tensor * k_pe = ggml_view_3d(ctx, compact_pe, rope_dim, n_tokens, 1,
         ggml_row_size(compact_pe->type, compact_dim),
+        ggml_row_size(compact_pe->type, compact_dim) * n_tokens,
         ggml_row_size(compact_pe->type, kv_rank));
     compact = rms_norm(ctx, compact, layer.wkv_a_norm, w.rms_eps);
 
-    ggml_tensor * q_nope = ggml_view_3d(ctx, q_cur, nope_dim, n_head, 1,
+    ggml_tensor * q_nope = ggml_view_3d(ctx, q_cur, nope_dim, n_head, n_tokens,
         ggml_row_size(q_cur->type, key_dim),
         ggml_row_size(q_cur->type, key_dim) * n_head, 0);
-    ggml_tensor * q_pe = ggml_view_3d(ctx, q_cur, rope_dim, n_head, 1,
+    ggml_tensor * q_pe = ggml_view_3d(ctx, q_cur, rope_dim, n_head, n_tokens,
         ggml_row_size(q_cur->type, key_dim),
         ggml_row_size(q_cur->type, key_dim) * n_head,
         ggml_row_size(q_cur->type, nope_dim));
@@ -243,11 +276,12 @@ ggml_tensor * build_mla(ggml_context * ctx,
     q_nope = ggml_permute(ctx, q_nope, 0, 2, 1, 3);
     ggml_tensor * q = ggml_concat(ctx, q_nope, q_pe, 0);
 
-    ggml_tensor * compact_3d = ggml_reshape_3d(ctx, compact, kv_rank, 1, 1);
+    ggml_tensor * compact_3d =
+        ggml_reshape_3d(ctx, compact, kv_rank, n_tokens, 1);
     ggml_tensor * current_k = ggml_concat(ctx, compact_3d, k_pe, 0);
 
-    ggml_tensor * dst = ggml_view_3d(ctx, cache.mla_k,
-        compact_dim, 1, 1, cache.mla_k->nb[1], cache.mla_k->nb[2],
+    ggml_tensor * dst = ggml_view_2d(ctx, cache.mla_k,
+        compact_dim, n_tokens, cache.mla_k->nb[2],
         static_cast<size_t>(position) * cache.mla_k->nb[2]);
     ggml_build_forward_expand(graph, ggml_cpy(ctx, current_k, dst));
 
@@ -265,7 +299,7 @@ ggml_tensor * build_mla(ggml_context * ctx,
     v = ggml_permute(ctx, v, 0, 2, 1, 3);
     ggml_tensor * scores = ggml_mul_mat(ctx, k, q);
     ggml_mul_mat_set_prec(scores, GGML_PREC_F32);
-    scores = ggml_soft_max_ext(ctx, scores, nullptr,
+    scores = ggml_soft_max_ext(ctx, scores, attn_mask,
                                1.0f / std::sqrt(static_cast<float>(key_dim)),
                                0.0f);
     if (!v_trans) v = ggml_cont(ctx, ggml_transpose(ctx, v));
@@ -273,7 +307,7 @@ ggml_tensor * build_mla(ggml_context * ctx,
     out = ggml_mul_mat(ctx, layer.wv_b, out);
     out = ggml_permute(ctx, out, 0, 2, 1, 3);
     out = ggml_cont_2d(ctx, out,
-                       static_cast<int64_t>(value_dim) * n_head, 1);
+                       static_cast<int64_t>(value_dim) * n_head, n_tokens);
 
     if (layer.wqkv_gate) {
         ggml_tensor * output_gate = ggml_sigmoid(ctx,
@@ -288,18 +322,21 @@ TopKMoeRouterResult build_kimi_router(ggml_context * ctx,
                                       const KimiK3Weights & w,
                                       const KimiK3Layer & layer,
                                       ggml_tensor * cur) {
+    const int n_tokens = static_cast<int>(cur->ne[1]);
     ggml_tensor * logits = ggml_mul_mat(ctx, layer.ffn_gate_inp, cur);
     TopKMoeRouterResult router;
     if (w.expert_gating_func == 2) {
         router = build_sigmoid_topk_moe_router(ctx, graph, logits,
-            layer.ffn_exp_probs_b, w.n_expert, w.n_expert_used, 1,
+            layer.ffn_exp_probs_b, w.n_expert, w.n_expert_used, n_tokens,
             w.expert_weights_norm, w.expert_weights_scale, false);
     } else {
         ggml_tensor * probs = ggml_soft_max(ctx, logits);
         ggml_tensor * selected = ggml_argsort_top_k(ctx, probs, w.n_expert_used);
-        ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, 1);
+        ggml_tensor * probs_3d =
+            ggml_reshape_3d(ctx, probs, 1, w.n_expert, n_tokens);
         ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, selected);
-        weights = ggml_reshape_2d(ctx, weights, w.n_expert_used, 1);
+        weights = ggml_reshape_2d(
+            ctx, weights, w.n_expert_used, n_tokens);
         if (w.expert_weights_norm) {
             ggml_tensor * sum = ggml_clamp(ctx, ggml_sum_rows(ctx, weights),
                                            6.103515625e-5f, INFINITY);
@@ -310,7 +347,8 @@ TopKMoeRouterResult build_kimi_router(ggml_context * ctx,
         }
         router.selected = selected;
         router.weights_2d = weights;
-        router.weights_3d = ggml_reshape_3d(ctx, weights, 1, w.n_expert_used, 1);
+        router.weights_3d = ggml_reshape_3d(
+            ctx, weights, 1, w.n_expert_used, n_tokens);
     }
     return router;
 }
@@ -326,7 +364,8 @@ ggml_tensor * build_latent_moe(ggml_context * ctx,
         build_kimi_router(ctx, graph, w, layer, identity);
 
     ggml_tensor * routed_3d = ggml_reshape_3d(ctx, routed_in,
-                                               w.n_expert_latent, 1, 1);
+                                               w.n_expert_latent,
+                                               1, cur->ne[1]);
     ggml_tensor * gate = ggml_mul_mat_id(ctx, layer.ffn_gate_exps,
                                          routed_3d, router.selected);
     ggml_tensor * up = ggml_mul_mat_id(ctx, layer.ffn_up_exps,
@@ -337,9 +376,10 @@ ggml_tensor * build_latent_moe(ggml_context * ctx,
                                             activated, router.selected);
     experts = ggml_mul(ctx, experts, router.weights_3d);
     ggml_tensor * sum_shape = ggml_new_tensor_3d(ctx, GGML_TYPE_F32,
-                                                 w.n_expert_latent, 1, 1);
+                                                 w.n_expert_latent,
+                                                 1, cur->ne[1]);
     ggml_tensor * moe = ggml_repeat_back(ctx, experts, sum_shape);
-    moe = ggml_reshape_2d(ctx, moe, w.n_expert_latent, 1);
+    moe = ggml_reshape_2d(ctx, moe, w.n_expert_latent, cur->ne[1]);
     if (layer.ffn_routed_norm) {
         moe = rms_norm(ctx, moe, layer.ffn_routed_norm, w.rms_eps);
     }
@@ -418,15 +458,17 @@ bool run_host_boundary_graph(ggml_backend_t backend,
 void populate_attn_res_bank(
         ggml_context * ctx,
         const KimiK3Weights & w,
+        int n_tokens,
         const std::vector<std::vector<float>> & host_checkpoints,
         AttnResBank & bank,
         std::vector<GraphInput> & inputs) {
     bank.ctx = ctx;
     bank.eps = w.rms_eps;
     bank.n_embd = w.n_embd;
+    bank.n_tokens = n_tokens;
     for (const std::vector<float> & checkpoint : host_checkpoints) {
         ggml_tensor * tensor = ggml_new_tensor_2d(
-            ctx, GGML_TYPE_F32, w.n_embd, 1);
+            ctx, GGML_TYPE_F32, w.n_embd, n_tokens);
         ggml_set_input(tensor);
         inputs.push_back({
             tensor, checkpoint.data(),
@@ -442,18 +484,40 @@ ggml_context * new_kimi_step_context() {
     return ggml_init(params);
 }
 
-bool streamed_kimi_k3_step(
+bool streamed_kimi_k3_forward(
         ggml_backend_t backend,
         const KimiK3Weights & w,
         KimiK3Cache & cache,
-        int32_t token,
-        int position,
-        std::vector<float> & logits,
+        const std::vector<int32_t> & tokens,
+        int base_pos,
+        const KimiK3ForwardOptions & options,
+        KimiK3ForwardResult & result,
         MoeHybridStreamEngine & stream_engine,
         MoeStreamDualOwnerExecutor * dual_stream_executor,
         const MoeStreamDualOwnerPolicy * stream_owner_policy,
         MoeHybridRoutingStats * routing_stats) {
-    std::vector<float> hidden(static_cast<size_t>(w.n_embd));
+    const int n_tokens = static_cast<int>(tokens.size());
+    const size_t hidden_values =
+        static_cast<size_t>(w.n_embd) * static_cast<size_t>(n_tokens);
+    std::vector<float> hidden(hidden_values);
+
+    std::vector<int> capture_at_layer(static_cast<size_t>(w.n_layer), -1);
+    const int n_capture = options.capture_layer_ids
+        ? static_cast<int>(options.capture_layer_ids->size()) : 0;
+    for (int i = 0; i < n_capture; ++i) {
+        capture_at_layer[static_cast<size_t>((*options.capture_layer_ids)[i])] = i;
+    }
+    result.captured_hidden.assign(
+        static_cast<size_t>(n_capture) * hidden_values, 0.0f);
+
+    const int kv_len = base_pos + n_tokens;
+    std::vector<float> mla_mask(
+        static_cast<size_t>(kv_len) * n_tokens, -INFINITY);
+    for (int q = 0; q < n_tokens; ++q) {
+        for (int k = 0; k <= base_pos + q; ++k) {
+            mla_mask[static_cast<size_t>(q) * kv_len + k] = 0.0f;
+        }
+    }
 
     {
         ggml_context * ctx = new_kimi_step_context();
@@ -464,13 +528,13 @@ bool streamed_kimi_k3_step(
         ggml_cgraph * graph =
             ggml_new_graph_custom(ctx, 1024, false);
         ggml_tensor * ids =
-            ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
+            ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
         ggml_set_input(ids);
         ggml_tensor * embedding =
             ggml_get_rows(ctx, w.tok_embd, ids);
         const bool ok = run_host_boundary_graph(
             backend, ctx, graph,
-            {{ids, &token, sizeof(token)}},
+            {{ids, tokens.data(), sizeof(int32_t) * tokens.size()}},
             {{embedding, hidden.data(),
               hidden.size() * sizeof(float)}},
             "embedding");
@@ -502,7 +566,7 @@ bool streamed_kimi_k3_step(
             ggml_new_graph_custom(ctx, 32768, false);
         std::vector<GraphInput> inputs;
         ggml_tensor * hidden_in = ggml_new_tensor_2d(
-            ctx, GGML_TYPE_F32, w.n_embd, 1);
+            ctx, GGML_TYPE_F32, w.n_embd, n_tokens);
         ggml_set_input(hidden_in);
         inputs.push_back({
             hidden_in, hidden.data(),
@@ -510,7 +574,7 @@ bool streamed_kimi_k3_step(
 
         AttnResBank residuals;
         populate_attn_res_bank(
-            ctx, w, checkpoints, residuals, inputs);
+            ctx, w, n_tokens, checkpoints, residuals, inputs);
         ggml_tensor * prefix = hidden_in;
         ggml_tensor * cur =
             residuals.mix(prefix, layer.attn_res_score);
@@ -518,12 +582,21 @@ bool streamed_kimi_k3_step(
 
         cur = rms_norm(
             ctx, cur, layer.attn_norm, w.rms_eps);
-        cur = layer.recurrent
-            ? build_kda(
-                ctx, graph, w, layer, layer_cache, cur)
-            : build_mla(
+        if (layer.recurrent) {
+            cur = build_kda(
+                ctx, graph, w, layer, layer_cache, cur,
+                /*commit_state=*/!options.capture_replay,
+                options.capture_replay);
+        } else {
+            ggml_tensor * mask = ggml_new_tensor_2d(
+                ctx, GGML_TYPE_F32, kv_len, n_tokens);
+            ggml_set_input(mask);
+            inputs.push_back({
+                mask, mla_mask.data(), mla_mask.size() * sizeof(float)});
+            cur = build_mla(
                 ctx, graph, w, layer, layer_cache,
-                cur, position);
+                cur, base_pos, mask);
+        }
         prefix = banked
             ? cur : ggml_add(ctx, prefix, cur);
         cur = residuals.mix(prefix, layer.ffn_res_score);
@@ -542,8 +615,7 @@ bool streamed_kimi_k3_step(
                 ctx, layer.ffn_down, dense);
             ggml_tensor * hidden_out =
                 ggml_add(ctx, prefix, dense);
-            std::vector<float> next_hidden(
-                static_cast<size_t>(w.n_embd));
+            std::vector<float> next_hidden(hidden_values);
             const bool ok = run_host_boundary_graph(
                 backend, ctx, graph, inputs,
                 {{hidden_out, next_hidden.data(),
@@ -553,6 +625,13 @@ bool streamed_kimi_k3_step(
             if (!ok) return false;
             if (banked) checkpoints.push_back(checkpoint_value);
             hidden.swap(next_hidden);
+            const int capture_idx = capture_at_layer[static_cast<size_t>(il)];
+            if (capture_idx >= 0) {
+                std::memcpy(
+                    result.captured_hidden.data() +
+                        static_cast<size_t>(capture_idx) * hidden_values,
+                    hidden.data(), hidden_values * sizeof(float));
+            }
             continue;
         }
 
@@ -577,16 +656,14 @@ bool streamed_kimi_k3_step(
         shared = ggml_mul_mat(
             ctx, layer.ffn_down_shexp, shared);
 
-        std::vector<float> prefix_host(
-            static_cast<size_t>(w.n_embd));
+        std::vector<float> prefix_host(hidden_values);
         std::vector<float> routed_input_host(
-            static_cast<size_t>(w.n_expert_latent));
+            static_cast<size_t>(w.n_expert_latent) * n_tokens);
         std::vector<int32_t> selected(
-            static_cast<size_t>(w.n_expert_used));
+            static_cast<size_t>(w.n_expert_used) * n_tokens);
         std::vector<float> route_weights(
-            static_cast<size_t>(w.n_expert_used));
-        std::vector<float> shared_host(
-            static_cast<size_t>(w.n_embd));
+            static_cast<size_t>(w.n_expert_used) * n_tokens);
+        std::vector<float> shared_host(hidden_values);
         const bool prep_ok = run_host_boundary_graph(
             backend, ctx, graph, inputs,
             {
@@ -637,7 +714,7 @@ bool streamed_kimi_k3_step(
         route_batch.layer = il - w.n_dense_lead;
         route_batch.n_expert = w.n_expert;
         route_batch.top_k = w.n_expert_used;
-        route_batch.n_tokens = 1;
+        route_batch.n_tokens = n_tokens;
         route_batch.inputs = routed_input_host.data();
         route_batch.selected_ids = selected.data();
         route_batch.selected_weights = route_weights.data();
@@ -689,11 +766,11 @@ bool streamed_kimi_k3_step(
         }
         graph = ggml_new_graph_custom(ctx, 4096, false);
         ggml_tensor * prefix_in = ggml_new_tensor_2d(
-            ctx, GGML_TYPE_F32, w.n_embd, 1);
+            ctx, GGML_TYPE_F32, w.n_embd, n_tokens);
         ggml_tensor * routed_out_in = ggml_new_tensor_2d(
-            ctx, GGML_TYPE_F32, w.n_expert_latent, 1);
+            ctx, GGML_TYPE_F32, w.n_expert_latent, n_tokens);
         ggml_tensor * shared_in = ggml_new_tensor_2d(
-            ctx, GGML_TYPE_F32, w.n_embd, 1);
+            ctx, GGML_TYPE_F32, w.n_embd, n_tokens);
         ggml_set_input(prefix_in);
         ggml_set_input(routed_out_in);
         ggml_set_input(shared_in);
@@ -708,8 +785,7 @@ bool streamed_kimi_k3_step(
             ggml_add(ctx, routed, shared_in);
         ggml_tensor * hidden_out =
             ggml_add(ctx, prefix_in, moe_shared);
-        std::vector<float> next_hidden(
-            static_cast<size_t>(w.n_embd));
+        std::vector<float> next_hidden(hidden_values);
         const bool join_ok = run_host_boundary_graph(
             backend, ctx, graph,
             {
@@ -726,6 +802,13 @@ bool streamed_kimi_k3_step(
         ggml_free(ctx);
         if (!join_ok) return false;
         hidden.swap(next_hidden);
+        const int capture_idx = capture_at_layer[static_cast<size_t>(il)];
+        if (capture_idx >= 0) {
+            std::memcpy(
+                result.captured_hidden.data() +
+                    static_cast<size_t>(capture_idx) * hidden_values,
+                hidden.data(), hidden_values * sizeof(float));
+        }
     }
 
     ggml_context * ctx = new_kimi_step_context();
@@ -737,30 +820,40 @@ bool streamed_kimi_k3_step(
         ggml_new_graph_custom(ctx, 8192, false);
     std::vector<GraphInput> inputs;
     ggml_tensor * hidden_in = ggml_new_tensor_2d(
-        ctx, GGML_TYPE_F32, w.n_embd, 1);
+        ctx, GGML_TYPE_F32, w.n_embd, n_tokens);
     ggml_set_input(hidden_in);
     inputs.push_back({
         hidden_in, hidden.data(),
         hidden.size() * sizeof(float)});
     AttnResBank residuals;
     populate_attn_res_bank(
-        ctx, w, checkpoints, residuals, inputs);
+        ctx, w, n_tokens, checkpoints, residuals, inputs);
     ggml_tensor * output_hidden =
         residuals.mix(hidden_in, w.output_res_score);
     output_hidden = rms_norm(
         ctx, output_hidden, w.output_norm, w.rms_eps);
     ggml_tensor * output =
         ggml_mul_mat(ctx, w.output, output_hidden);
-    logits.resize(static_cast<size_t>(w.n_vocab));
+    ggml_tensor * argmax = ggml_argmax(ctx, output);
+    std::vector<GraphOutput> outputs;
+    if (options.read_logits) {
+        result.logits.resize(static_cast<size_t>(w.n_vocab) * n_tokens);
+        outputs.push_back({
+            output, result.logits.data(), result.logits.size() * sizeof(float)});
+    }
+    if (options.read_argmax) {
+        result.argmax.resize(static_cast<size_t>(n_tokens));
+        outputs.push_back({
+            argmax, result.argmax.data(), result.argmax.size() * sizeof(int32_t)});
+    }
     const bool output_ok = run_host_boundary_graph(
         backend, ctx, graph, inputs,
-        {{output, logits.data(),
-          logits.size() * sizeof(float)}},
+        outputs,
         "output");
     ggml_free(ctx);
     if (!output_ok) return false;
 
-    cache.cur_pos = position + 1;
+    cache.cur_pos = base_pos + n_tokens;
     return true;
 }
 
@@ -769,13 +862,14 @@ bool streamed_kimi_k3_step(
 bool create_kimi_k3_cache(ggml_backend_t backend,
                           const KimiK3Weights & w,
                           int max_ctx,
-                          KimiK3Cache & out) {
+                          KimiK3Cache & out,
+                          int max_verify_tokens) {
     free_kimi_k3_cache(out);
     if (!backend || max_ctx <= 0) return false;
 
     ggml_init_params params{};
     params.mem_size = ggml_tensor_overhead() *
-        static_cast<size_t>(w.n_layer * 3 + 16) + 16384;
+        static_cast<size_t>(w.n_layer * 6 + 16) + 16384;
     params.no_alloc = true;
     out.ctx = ggml_init(params);
     if (!out.ctx) return false;
@@ -795,6 +889,23 @@ bool create_kimi_k3_cache(ggml_backend_t backend,
             ggml_set_name(layer_cache.conv_state, name);
             std::snprintf(name, sizeof(name), "kimi_k3_ssm_state_%d", il);
             ggml_set_name(layer_cache.ssm_state, name);
+            if (max_verify_tokens > 0) {
+                layer_cache.conv_state_snap = ggml_dup_tensor(
+                    out.ctx, layer_cache.conv_state);
+                layer_cache.ssm_state_snap = ggml_dup_tensor(
+                    out.ctx, layer_cache.ssm_state);
+                layer_cache.replay_input = ggml_new_tensor_2d(
+                    out.ctx, GGML_TYPE_F32, w.n_embd, max_verify_tokens);
+                std::snprintf(
+                    name, sizeof(name), "kimi_k3_conv_state_snap_%d", il);
+                ggml_set_name(layer_cache.conv_state_snap, name);
+                std::snprintf(
+                    name, sizeof(name), "kimi_k3_ssm_state_snap_%d", il);
+                ggml_set_name(layer_cache.ssm_state_snap, name);
+                std::snprintf(
+                    name, sizeof(name), "kimi_k3_replay_input_%d", il);
+                ggml_set_name(layer_cache.replay_input, name);
+            }
         } else {
             layer_cache.mla_k = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F16,
                 compact_dim, 1, max_ctx);
@@ -809,6 +920,7 @@ bool create_kimi_k3_cache(ggml_backend_t backend,
         return false;
     }
     out.max_ctx = max_ctx;
+    out.max_verify_tokens = std::max(0, max_verify_tokens);
     reset_kimi_k3_cache(out);
     return true;
 }
@@ -816,6 +928,12 @@ bool create_kimi_k3_cache(ggml_backend_t backend,
 void reset_kimi_k3_cache(KimiK3Cache & cache) {
     if (cache.buf) ggml_backend_buffer_clear(cache.buf, 0);
     cache.cur_pos = 0;
+    cache.snapshot_pos = -1;
+    cache.replay_base_pos = -1;
+    cache.replay_n_tokens = 0;
+    cache.snapshot_valid = false;
+    cache.replay_valid = false;
+    cache.recurrent_state_pristine = false;
 }
 
 void free_kimi_k3_cache(KimiK3Cache & cache) {
@@ -824,39 +942,89 @@ void free_kimi_k3_cache(KimiK3Cache & cache) {
     cache = KimiK3Cache{};
 }
 
-bool kimi_k3_step(ggml_backend_t backend,
-                  const KimiK3Weights & w,
-                  KimiK3Cache & cache,
-                  int32_t token,
-                  int position,
-                  std::vector<float> & logits,
-                  MoeHybridStreamEngine * stream_engine,
-                  MoeStreamDualOwnerExecutor * dual_stream_executor,
-                  const MoeStreamDualOwnerPolicy * stream_owner_policy,
-                  MoeHybridRoutingStats * routing_stats) {
-    if (!backend || !w.ctx || !cache.ctx || position < 0 ||
-        position >= cache.max_ctx || position != cache.cur_pos ||
-        token < 0 || token >= w.n_vocab) {
-        set_last_error("Kimi-K3 step: invalid backend, cache position, or token");
+bool kimi_k3_forward(ggml_backend_t backend,
+                     const KimiK3Weights & w,
+                     KimiK3Cache & cache,
+                     const std::vector<int32_t> & tokens,
+                     int base_pos,
+                     const KimiK3ForwardOptions & options,
+                     KimiK3ForwardResult & result,
+                     MoeHybridStreamEngine * stream_engine,
+                     MoeStreamDualOwnerExecutor * dual_stream_executor,
+                     const MoeStreamDualOwnerPolicy * stream_owner_policy,
+                     MoeHybridRoutingStats * routing_stats) {
+    result = KimiK3ForwardResult{};
+    const int n_tokens = static_cast<int>(tokens.size());
+    if (!backend || !w.ctx || !cache.ctx || n_tokens <= 0 || base_pos < 0 ||
+        base_pos != cache.cur_pos || base_pos + n_tokens > cache.max_ctx ||
+        (!options.read_logits && !options.read_argmax)) {
+        set_last_error("Kimi-K3 forward: invalid backend, output, or cache span");
         return false;
     }
+    for (int32_t token : tokens) {
+        if (token < 0 || token >= w.n_vocab) {
+            set_last_error("Kimi-K3 forward: token is outside the vocabulary");
+            return false;
+        }
+    }
+
+    std::vector<int> capture_at_layer(static_cast<size_t>(w.n_layer), -1);
+    const int n_capture = options.capture_layer_ids
+        ? static_cast<int>(options.capture_layer_ids->size()) : 0;
+    for (int i = 0; i < n_capture; ++i) {
+        const int layer = (*options.capture_layer_ids)[static_cast<size_t>(i)];
+        if (layer < 0 || layer >= w.n_layer ||
+            capture_at_layer[static_cast<size_t>(layer)] >= 0) {
+            set_last_error("Kimi-K3 forward: invalid or duplicate capture layer");
+            return false;
+        }
+        capture_at_layer[static_cast<size_t>(layer)] = i;
+    }
+    if (options.capture_replay &&
+        (n_tokens > cache.max_verify_tokens || !cache.snapshot_valid ||
+         cache.snapshot_pos != base_pos)) {
+        set_last_error("Kimi-K3 forward: ReplaySSM capture has no matching snapshot");
+        return false;
+    }
+
     if (w.routed_experts_streamed) {
         if (!stream_engine || !stream_engine->is_bound()) {
             set_last_error(
-                "Kimi-K3 step: file-backed experts require a bound stream engine");
+                "Kimi-K3 forward: file-backed experts require a bound stream engine");
             return false;
         }
         if (dual_stream_executor &&
             (!dual_stream_executor->is_ready() || !stream_owner_policy)) {
             set_last_error(
-                "Kimi-K3 step: dual-owner streaming requires a ready "
+                "Kimi-K3 forward: dual-owner streaming requires a ready "
                 "executor and an ownership policy");
             return false;
         }
-        return streamed_kimi_k3_step(
-            backend, w, cache, token, position,
-            logits, *stream_engine, dual_stream_executor,
-            stream_owner_policy, routing_stats);
+        if (!streamed_kimi_k3_forward(
+                backend, w, cache, tokens, base_pos, options, result,
+                *stream_engine, dual_stream_executor,
+                stream_owner_policy, routing_stats)) {
+            return false;
+        }
+        if (options.capture_replay) {
+            cache.replay_base_pos = base_pos;
+            cache.replay_n_tokens = n_tokens;
+            cache.replay_valid = true;
+            cache.recurrent_state_pristine = true;
+        } else {
+            cache.replay_valid = false;
+            cache.recurrent_state_pristine = false;
+        }
+        return true;
+    }
+
+    const int kv_len = base_pos + n_tokens;
+    std::vector<float> mla_mask(
+        static_cast<size_t>(kv_len) * n_tokens, -INFINITY);
+    for (int q = 0; q < n_tokens; ++q) {
+        for (int k = 0; k <= base_pos + q; ++k) {
+            mla_mask[static_cast<size_t>(q) * kv_len + k] = 0.0f;
+        }
     }
 
     ggml_init_params params{};
@@ -864,20 +1032,27 @@ bool kimi_k3_step(ggml_backend_t backend,
     params.no_alloc = true;
     ggml_context * ctx = ggml_init(params);
     if (!ctx) {
-        set_last_error("Kimi-K3 step: graph context allocation failed");
+        set_last_error("Kimi-K3 forward: graph context allocation failed");
         return false;
     }
     ggml_cgraph * graph = ggml_new_graph_custom(ctx, 32768, false);
 
-    ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
-    ggml_set_name(ids, "token_id");
+    ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens);
+    ggml_set_name(ids, "token_ids");
     ggml_set_input(ids);
     ggml_tensor * hidden = ggml_get_rows(ctx, w.tok_embd, ids);
+    ggml_tensor * mask = ggml_new_tensor_2d(
+        ctx, GGML_TYPE_F32, kv_len, n_tokens);
+    ggml_set_name(mask, "kimi_k3_mla_causal_mask");
+    ggml_set_input(mask);
 
+    std::vector<ggml_tensor *> capture_tensors(
+        static_cast<size_t>(n_capture), nullptr);
     AttnResBank residuals;
     residuals.ctx = ctx;
     residuals.eps = w.rms_eps;
     residuals.n_embd = w.n_embd;
+    residuals.n_tokens = n_tokens;
     for (int il = 0; il < w.n_layer; ++il) {
         const KimiK3Layer & layer = w.layers[static_cast<size_t>(il)];
         KimiK3LayerCache & layer_cache = cache.layers[static_cast<size_t>(il)];
@@ -888,8 +1063,10 @@ bool kimi_k3_step(ggml_backend_t backend,
 
         cur = rms_norm(ctx, cur, layer.attn_norm, w.rms_eps);
         cur = layer.recurrent
-            ? build_kda(ctx, graph, w, layer, layer_cache, cur)
-            : build_mla(ctx, graph, w, layer, layer_cache, cur, position);
+            ? build_kda(ctx, graph, w, layer, layer_cache, cur,
+                        /*commit_state=*/!options.capture_replay,
+                        options.capture_replay)
+            : build_mla(ctx, graph, w, layer, layer_cache, cur, base_pos, mask);
         prefix = banked ? cur : ggml_add(ctx, prefix, cur);
 
         cur = residuals.mix(prefix, layer.ffn_res_score);
@@ -903,40 +1080,202 @@ bool kimi_k3_step(ggml_backend_t backend,
             cur = build_latent_moe(ctx, graph, w, layer, cur);
         }
         hidden = ggml_add(ctx, prefix, cur);
+        const int capture_idx = capture_at_layer[static_cast<size_t>(il)];
+        if (capture_idx >= 0) {
+            capture_tensors[static_cast<size_t>(capture_idx)] = hidden;
+            ggml_set_output(hidden);
+            ggml_build_forward_expand(graph, hidden);
+        }
     }
 
     hidden = residuals.mix(hidden, w.output_res_score);
     hidden = rms_norm(ctx, hidden, w.output_norm, w.rms_eps);
     ggml_tensor * output = ggml_mul_mat(ctx, w.output, hidden);
     ggml_set_name(output, "logits");
-    ggml_set_output(output);
-    ggml_build_forward_expand(graph, output);
+    ggml_tensor * argmax = ggml_argmax(ctx, output);
+    if (options.read_logits) {
+        ggml_set_output(output);
+        ggml_build_forward_expand(graph, output);
+    }
+    if (options.read_argmax) {
+        ggml_set_output(argmax);
+        ggml_build_forward_expand(graph, argmax);
+    }
 
     ggml_gallocr_t allocator = ggml_gallocr_new(
         ggml_backend_get_default_buffer_type(backend));
     if (!allocator || !ggml_gallocr_alloc_graph(allocator, graph)) {
-        set_last_error("Kimi-K3 step: graph allocation failed");
+        set_last_error("Kimi-K3 forward: graph allocation failed");
         if (allocator) ggml_gallocr_free(allocator);
         ggml_free(ctx);
         return false;
     }
-    ggml_backend_tensor_set(ids, &token, 0, sizeof(token));
+    ggml_backend_tensor_set(
+        ids, tokens.data(), 0, sizeof(int32_t) * tokens.size());
+    ggml_backend_tensor_set(
+        mask, mla_mask.data(), 0, sizeof(float) * mla_mask.size());
     const ggml_status status = ggml_backend_graph_compute(backend, graph);
     if (status != GGML_STATUS_SUCCESS) {
-        set_last_error("Kimi-K3 step: graph compute failed with status " +
+        set_last_error("Kimi-K3 forward: graph compute failed with status " +
                        std::to_string(static_cast<int>(status)));
         ggml_gallocr_free(allocator);
         ggml_free(ctx);
         return false;
     }
 
-    logits.resize(static_cast<size_t>(w.n_vocab));
-    ggml_backend_tensor_get(output, logits.data(), 0,
-                            logits.size() * sizeof(float));
-    cache.cur_pos = position + 1;
+    if (options.read_logits) {
+        result.logits.resize(static_cast<size_t>(w.n_vocab) * n_tokens);
+        ggml_backend_tensor_get(output, result.logits.data(), 0,
+                                result.logits.size() * sizeof(float));
+    }
+    if (options.read_argmax) {
+        result.argmax.resize(static_cast<size_t>(n_tokens));
+        ggml_backend_tensor_get(argmax, result.argmax.data(), 0,
+                                result.argmax.size() * sizeof(int32_t));
+    }
+    const size_t hidden_values =
+        static_cast<size_t>(w.n_embd) * static_cast<size_t>(n_tokens);
+    result.captured_hidden.resize(
+        static_cast<size_t>(n_capture) * hidden_values);
+    for (int i = 0; i < n_capture; ++i) {
+        ggml_backend_tensor_get(
+            capture_tensors[static_cast<size_t>(i)],
+            result.captured_hidden.data() + static_cast<size_t>(i) * hidden_values,
+            0, hidden_values * sizeof(float));
+    }
+
+    cache.cur_pos = base_pos + n_tokens;
+    if (options.capture_replay) {
+        cache.replay_base_pos = base_pos;
+        cache.replay_n_tokens = n_tokens;
+        cache.replay_valid = true;
+        cache.recurrent_state_pristine = true;
+    } else {
+        cache.replay_valid = false;
+        cache.recurrent_state_pristine = false;
+    }
     ggml_gallocr_free(allocator);
     ggml_free(ctx);
     return true;
 }
 
+bool kimi_k3_replay_snapshot(ggml_backend_t backend, KimiK3Cache & cache) {
+    if (!backend || cache.max_verify_tokens <= 0) return false;
+    for (KimiK3LayerCache & layer : cache.layers) {
+        if (!layer.ssm_state) continue;
+        if (!layer.ssm_state_snap || !layer.conv_state_snap ||
+            !layer.replay_input) {
+            return false;
+        }
+        ggml_backend_tensor_copy_async(
+            backend, backend, layer.ssm_state, layer.ssm_state_snap);
+        ggml_backend_tensor_copy_async(
+            backend, backend, layer.conv_state, layer.conv_state_snap);
+    }
+    ggml_backend_synchronize(backend);
+    cache.snapshot_pos = cache.cur_pos;
+    cache.snapshot_valid = true;
+    cache.replay_valid = false;
+    cache.recurrent_state_pristine = true;
+    return true;
+}
+
+bool kimi_k3_replay_restore(ggml_backend_t backend, KimiK3Cache & cache) {
+    if (!backend || !cache.snapshot_valid || cache.snapshot_pos < 0) return false;
+    if (!cache.recurrent_state_pristine) {
+        for (KimiK3LayerCache & layer : cache.layers) {
+            if (!layer.ssm_state) continue;
+            if (!layer.ssm_state_snap || !layer.conv_state_snap) return false;
+            ggml_backend_tensor_copy_async(
+                backend, backend, layer.ssm_state_snap, layer.ssm_state);
+            ggml_backend_tensor_copy_async(
+                backend, backend, layer.conv_state_snap, layer.conv_state);
+        }
+        ggml_backend_synchronize(backend);
+    }
+    cache.cur_pos = cache.snapshot_pos;
+    cache.replay_valid = false;
+    cache.recurrent_state_pristine = true;
+    return true;
+}
+
+bool kimi_k3_replay_commit(ggml_backend_t backend,
+                           const KimiK3Weights & w,
+                           KimiK3Cache & cache,
+                           int base_pos,
+                           int commit_n) {
+    if (!backend || !cache.snapshot_valid || !cache.replay_valid ||
+        !cache.recurrent_state_pristine || cache.snapshot_pos != base_pos ||
+        cache.replay_base_pos != base_pos || commit_n <= 0 ||
+        commit_n > cache.replay_n_tokens) {
+        return false;
+    }
+
+    ggml_init_params params{};
+    params.mem_size = 64ull * 1024ull * 1024ull;
+    params.no_alloc = true;
+    ggml_context * ctx = ggml_init(params);
+    if (!ctx) return false;
+    ggml_cgraph * graph = ggml_new_graph_custom(ctx, 32768, false);
+    for (int il = 0; il < w.n_layer; ++il) {
+        const KimiK3Layer & layer = w.layers[static_cast<size_t>(il)];
+        if (!layer.recurrent) continue;
+        KimiK3LayerCache & layer_cache = cache.layers[static_cast<size_t>(il)];
+        if (!layer_cache.replay_input) {
+            ggml_free(ctx);
+            return false;
+        }
+        ggml_tensor * replay = ggml_view_2d(
+            ctx, layer_cache.replay_input, w.n_embd, commit_n,
+            layer_cache.replay_input->nb[1], 0);
+        (void)build_kda(ctx, graph, w, layer, layer_cache, replay,
+                        /*commit_state=*/true,
+                        /*capture_replay=*/false);
+    }
+
+    ggml_gallocr_t allocator = ggml_gallocr_new(
+        ggml_backend_get_default_buffer_type(backend));
+    if (!allocator || !ggml_gallocr_alloc_graph(allocator, graph)) {
+        if (allocator) ggml_gallocr_free(allocator);
+        ggml_free(ctx);
+        return false;
+    }
+    cache.recurrent_state_pristine = false;
+    const ggml_status status = ggml_backend_graph_compute(backend, graph);
+    ggml_gallocr_free(allocator);
+    ggml_free(ctx);
+    if (status != GGML_STATUS_SUCCESS) {
+        (void)kimi_k3_replay_restore(backend, cache);
+        return false;
+    }
+    cache.cur_pos = base_pos + commit_n;
+    cache.snapshot_valid = false;
+    cache.replay_valid = false;
+    return true;
+}
+
+bool kimi_k3_step(ggml_backend_t backend,
+                  const KimiK3Weights & w,
+                  KimiK3Cache & cache,
+                  int32_t token,
+                  int position,
+                  std::vector<float> & logits,
+                  MoeHybridStreamEngine * stream_engine,
+                  MoeStreamDualOwnerExecutor * dual_stream_executor,
+                  const MoeStreamDualOwnerPolicy * stream_owner_policy,
+                  MoeHybridRoutingStats * routing_stats) {
+    KimiK3ForwardOptions options;
+    options.read_logits = true;
+    options.read_argmax = false;
+    KimiK3ForwardResult result;
+    if (!kimi_k3_forward(
+            backend, w, cache, std::vector<int32_t>{token}, position,
+            options, result, stream_engine, dual_stream_executor,
+            stream_owner_policy, routing_stats)) {
+        return false;
+    }
+    logits = std::move(result.logits);
+    return true;
+}
+
 } // namespace dflash::common
diff --git a/server/src/kimi_k3/kimi_k3_internal.h b/server/src/kimi_k3/kimi_k3_internal.h
index 42d9f6cf7..446354dc7 100644
--- a/server/src/kimi_k3/kimi_k3_internal.h
+++ b/server/src/kimi_k3/kimi_k3_internal.h
@@ -3,7 +3,7 @@
 // This is intentionally split into three model-neutral boundaries:
 //   * GGUF loading owns tensor metadata/storage only;
 //   * KimiK3Cache owns recurrent/attention state only; and
-//   * kimi_k3_step owns the architecture graph only.
+//   * kimi_k3_forward owns the architecture graph only.
 //
 // Routed-expert placement can therefore replace the resident expert tensors
 // with the common MoE stream engine without changing KDA, MLA, AttnRes, or the
@@ -143,6 +143,15 @@ struct KimiK3LayerCache {
     ggml_tensor * conv_state = nullptr; // [d_conv-1, 3*d_inner], F32
     ggml_tensor * ssm_state  = nullptr; // [head_dim, head_dim, n_head], F32
     ggml_tensor * mla_k      = nullptr; // [kv_rank+rope_dim, 1, max_ctx], F16
+
+    // Speculative-decode state.  ReplaySSM captures the much smaller
+    // pre-KDA activation for every verify row, then re-runs only the accepted
+    // recurrent transitions at commit time.  The snapshots are a failure-safe
+    // for the commit graph; ordinary rejected verification leaves the live
+    // recurrent tensors untouched and therefore needs no restore copy.
+    ggml_tensor * conv_state_snap = nullptr;
+    ggml_tensor * ssm_state_snap  = nullptr;
+    ggml_tensor * replay_input    = nullptr; // [hidden, max_verify_tokens], F32
 };
 
 struct KimiK3Cache {
@@ -151,6 +160,29 @@ struct KimiK3Cache {
     std::vector<KimiK3LayerCache> layers;
     int max_ctx = 0;
     int cur_pos = 0;
+    int max_verify_tokens = 0;
+    int snapshot_pos = -1;
+    int replay_base_pos = -1;
+    int replay_n_tokens = 0;
+    bool snapshot_valid = false;
+    bool replay_valid = false;
+    bool recurrent_state_pristine = false;
+};
+
+// Model-neutral forward result shape used by the Kimi DFlash adapter.  Capture
+// rows are capture-major, then token-major:
+//   [capture_layer][token][hidden].
+struct KimiK3ForwardOptions {
+    const std::vector<int> * capture_layer_ids = nullptr;
+    bool capture_replay = false;
+    bool read_logits = false;
+    bool read_argmax = true;
+};
+
+struct KimiK3ForwardResult {
+    std::vector<float> logits;
+    std::vector<int32_t> argmax;
+    std::vector<float> captured_hidden;
 };
 
 bool load_kimi_k3_gguf(const std::string & path,
@@ -162,14 +194,37 @@ void free_kimi_k3_weights(KimiK3Weights & w);
 bool create_kimi_k3_cache(ggml_backend_t backend,
                           const KimiK3Weights & w,
                           int max_ctx,
-                          KimiK3Cache & out);
+                          KimiK3Cache & out,
+                          int max_verify_tokens = 0);
 void reset_kimi_k3_cache(KimiK3Cache & cache);
 void free_kimi_k3_cache(KimiK3Cache & cache);
 
-// Executes exactly one token. Token-at-a-time is deliberate for the first
-// correctness path: it makes the recurrent state transition explicit and is
-// numerically equivalent to chunked prefill. Persistent/captured decode graphs
-// and chunked KDA are performance layers added above this contract.
+// Batch forward used for target verification.  With capture_replay=true the
+// recurrent state is read-only: KDA inputs are persisted for a later
+// kimi_k3_replay_commit(), while MLA writes remain position-indexed and become
+// invisible simply by restoring cur_pos.
+bool kimi_k3_forward(ggml_backend_t backend,
+                     const KimiK3Weights & w,
+                     KimiK3Cache & cache,
+                     const std::vector<int32_t> & tokens,
+                     int base_pos,
+                     const KimiK3ForwardOptions & options,
+                     KimiK3ForwardResult & result,
+                     MoeHybridStreamEngine * stream_engine = nullptr,
+                     MoeStreamDualOwnerExecutor * dual_stream_executor = nullptr,
+                     const MoeStreamDualOwnerPolicy * stream_owner_policy = nullptr,
+                     MoeHybridRoutingStats * routing_stats = nullptr);
+
+bool kimi_k3_replay_snapshot(ggml_backend_t backend, KimiK3Cache & cache);
+bool kimi_k3_replay_restore(ggml_backend_t backend, KimiK3Cache & cache);
+bool kimi_k3_replay_commit(ggml_backend_t backend,
+                           const KimiK3Weights & w,
+                           KimiK3Cache & cache,
+                           int base_pos,
+                           int commit_n);
+
+// Compatibility wrapper for the ordinary one-token AR path. Speculative
+// verification calls kimi_k3_forward directly with a bounded token batch.
 bool kimi_k3_step(ggml_backend_t backend,
                   const KimiK3Weights & w,
                   KimiK3Cache & cache,
diff --git a/server/test/smoke_kimi_k3_forward.cpp b/server/test/smoke_kimi_k3_forward.cpp
index c635a84c8..08faac497 100644
--- a/server/test/smoke_kimi_k3_forward.cpp
+++ b/server/test/smoke_kimi_k3_forward.cpp
@@ -11,7 +11,8 @@ int main(int argc, char ** argv) {
     if (argc < 2) {
         std::fprintf(stderr,
             "usage: %s <kimi-k3.gguf> [gpu=0] [n_gen=16] [prompt] "
-            "[stream_experts=1] [expert_gpu=-1]\n",
+            "[stream_experts=1] [expert_gpu=-1] [draft.gguf] "
+            "[draft_gpu=0]\n",
             argv[0]);
         return 2;
     }
@@ -41,6 +42,8 @@ int main(int argc, char ** argv) {
         ? MoeStoragePolicy::Ssd
         : MoeStoragePolicy::Resident;
     config.expert_gpu = argc > 6 ? std::atoi(argv[6]) : -1;
+    config.draft_path = argc > 7 && argv[7][0] != '\0' ? argv[7] : nullptr;
+    config.draft_gpu = argc > 8 ? std::atoi(argv[8]) : gpu;
     KimiK3Backend backend(config);
     if (!backend.init()) return 1;
 
diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp
index 9f83ee650..42ca62d52 100644
--- a/server/test/test_feature_gate.cpp
+++ b/server/test/test_feature_gate.cpp
@@ -456,7 +456,8 @@ static void test_feature_warnings_report_inert_draft() {
     // These native AR backends never forward a draft model.
     TEST_ASSERT(warns_about(warn_result(args, "qwen3"), "--draft"));
     TEST_ASSERT(warns_about(warn_result(args, "deepseek4"), "--draft"));
-    TEST_ASSERT(warns_about(warn_result(args, "kimi-k3"), "--draft"));
+    // Kimi-K3 uses the generic DFlash/DSpark runtime on monolithic placement.
+    TEST_ASSERT(!warns_about(warn_result(args, "kimi-k3"), "--draft"));
     // laguna and gemma4 forward it only when monolithic.
     TEST_ASSERT(!warns_about(warn_result(args, "laguna"), "--draft"));
     TEST_ASSERT(!warns_about(warn_result(args, "gemma4"), "--draft"));
@@ -547,6 +548,8 @@ static void test_model_capability_tables() {
     TEST_ASSERT(arch_supports_moe_ssd_storage("deepseek4", false));
     TEST_ASSERT(!arch_supports_moe_ssd_storage("deepseek4", true));
     TEST_ASSERT(arch_supports_moe_ssd_storage("kimi-k3", false));
+    TEST_ASSERT(arch_supports_decode_draft("kimi-k3", false));
+    TEST_ASSERT(!arch_supports_decode_draft("kimi-k3", true));
     TEST_ASSERT(!arch_supports_moe_ssd_storage("qwen35moe", false));
 }
 

From e9aa6bd13702dabc042292341b46fe9f06e46734 Mon Sep 17 00:00:00 2001
From: mrciffa <49000955+davide221@users.noreply.github.com>
Date: Sat, 1 Aug 2026 21:50:35 +0200
Subject: [PATCH 3/3] fix(dspark): align proposal and verification widths

---
 server/src/common/dflash_spec_decode.cpp |  90 +++++---
 server/src/common/dspark_head.cpp        | 264 +++++++++++++++--------
 server/src/common/dspark_head.h          |  23 ++
 server/src/draft/draft_gguf_loader.cpp   |  48 +++++
 server/src/draft/draft_graph.cpp         |  45 ++--
 server/src/internal.h                    |   9 +
 server/src/kimi_k3/kimi_k3_backend.cpp   |   2 +-
 server/tests/test_deepseek4_unit.cpp     |  78 +++++++
 8 files changed, 410 insertions(+), 149 deletions(-)

diff --git a/server/src/common/dflash_spec_decode.cpp b/server/src/common/dflash_spec_decode.cpp
index caab7c244..7461f9a80 100644
--- a/server/src/common/dflash_spec_decode.cpp
+++ b/server/src/common/dflash_spec_decode.cpp
@@ -66,9 +66,11 @@ bool run_dflash_spec_decode(
     const bool use_remote_draft = remote_draft && remote_draft->active();
     if (!use_remote_draft && !feature_ring.target_feat) return false;
 
-    const int hidden    = draft_weights.n_embd;
-    const int max_q_len = draft_weights.block_size;
-    if (hidden <= 0 || max_q_len <= 0) return false;
+    const int hidden = draft_weights.n_embd;
+    const int draft_block_size = draft_weights.block_size;
+    const bool dspark_block = draft_weights.dspark.enabled;
+    const int max_verify_width = draft_weights.max_chain_verify_tokens();
+    if (hidden <= 0 || draft_block_size <= 0 || max_verify_width <= 0) return false;
     const float width_theta = adaptive_verify_width_theta();
     const int target_width_min = target.default_adaptive_verify_min_rows();
     const int width_min = adaptive_verify_width_min(target_width_min);
@@ -76,11 +78,11 @@ bool run_dflash_spec_decode(
     StepGraph draft_sg;
     StepGraphGuard draft_sg_guard{draft_sg};
 
-    std::vector<float>   noise_embed((size_t)hidden * max_q_len);
-    std::vector<int32_t> noise_ids(max_q_len);
-    std::vector<int32_t> draft_tok(max_q_len);
-    std::vector<int32_t> target_tok(max_q_len);
-    std::vector<int32_t> pos_q(max_q_len);
+    std::vector<float>   noise_embed((size_t)hidden * draft_block_size);
+    std::vector<int32_t> noise_ids(draft_block_size);
+    std::vector<int32_t> draft_tok(max_verify_width);
+    std::vector<int32_t> target_tok(max_verify_width);
+    std::vector<int32_t> pos_q(draft_block_size);
     std::vector<int32_t> pos_k;
     std::vector<float>   local_hidden;       // host buffer for local draft hidden states
     std::vector<float>   remote_hidden;      // host buffer for remote-draft hidden states
@@ -99,15 +101,15 @@ bool run_dflash_spec_decode(
     auto t_dec0 = std::chrono::steady_clock::now();
     while (n_generated < n_gen) {
         const int need_commit_budget = n_gen - n_generated;
-        int q_len = max_q_len;
+        int q_len = max_verify_width;
 
         // ── Build noise input for draft ────────────────────────────────────
         noise_ids[0] = last_tok;
-        for (int i = 1; i < max_q_len; i++) {
+        for (int i = 1; i < draft_block_size; i++) {
             noise_ids[i] = target.mask_token_id();
         }
         if (!target.embed_tokens(
-                noise_ids.data(), max_q_len, noise_embed.data())) {
+                noise_ids.data(), draft_block_size, noise_embed.data())) {
             std::fprintf(stderr, "dflash-spec noise embed failed\n");
             return false;
         }
@@ -146,9 +148,9 @@ bool run_dflash_spec_decode(
             }
             ggml_backend_tensor_set(draft_sg.inp_embed, noise_embed.data(), 0,
                                     sizeof(float) * noise_embed.size());
-            pos_k.resize((size_t)draft_ctx + max_q_len);
-            for (int i = 0; i < max_q_len; i++) pos_q[i] = draft_ctx + i;
-            for (int i = 0; i < draft_ctx + max_q_len; i++) pos_k[i] = i;
+            pos_k.resize((size_t)draft_ctx + draft_block_size);
+            for (int i = 0; i < draft_block_size; i++) pos_q[i] = draft_ctx + i;
+            for (int i = 0; i < draft_ctx + draft_block_size; i++) pos_k[i] = i;
             ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0,
                                     sizeof(int32_t) * pos_q.size());
             ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0,
@@ -160,7 +162,7 @@ bool run_dflash_spec_decode(
             }
             // Read draft hidden states out to host so the target adapter can
             // project them through its own LM head (target-internal layout).
-            local_hidden.resize((size_t)hidden * max_q_len);
+            local_hidden.resize((size_t)hidden * draft_block_size);
             ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0,
                                     sizeof(float) * local_hidden.size());
             draft_hidden_host = local_hidden.data();
@@ -174,46 +176,70 @@ bool run_dflash_spec_decode(
         std::vector<float> confidence;
         std::vector<float> candidate_probs;
         std::vector<int32_t> candidate_ids;
+        std::vector<int32_t> proposals;
         bool projected = false;
-        if (draft_backend && draft_weights.dspark.enabled && max_q_len > 1) {
+        if (draft_backend && dspark_block) {
             ggml_backend_t head_backend = target.fused_head_backend();
             const bool same_device_head =
                 head_backend && target.lm_head_tensor() &&
                 ggml_backend_get_device(head_backend) ==
                     ggml_backend_get_device(draft_backend);
             if (same_device_head) {
-                projected = dspark_markov_correct_greedy_chain_fused(
+                projected = dspark_markov_propose_greedy_block_fused(
                     draft_weights, draft_backend, target.lm_head_tensor(),
-                    draft_hidden_host, max_q_len, last_tok, draft_tok,
+                    draft_hidden_host, draft_block_size, last_tok, proposals,
                     &confidence);
             }
             if (!projected) {
-                projected = dspark_markov_correct_greedy_chain(
+                projected = dspark_markov_propose_greedy_block(
                     draft_weights, draft_backend, target,
-                    draft_hidden_host, max_q_len, last_tok,
-                    /*confidence_threshold=*/0.0f, draft_tok);
+                    draft_hidden_host, draft_block_size, last_tok, proposals);
             }
         }
-        if (!projected) {
-            const int candidate_k = width_theta > 0.0f && max_q_len > 2 ? 2 : 0;
+        if (dspark_block && !projected) {
+            // Preserve the DSpark block contract even if the auxiliary head
+            // cannot execute: every hidden row still predicts a proposal.
+            projected = target.project_hidden_to_tokens(
+                draft_hidden_host, draft_block_size, proposals);
+        }
+        if (dspark_block && projected) {
+            if (proposals.size() < static_cast<size_t>(draft_block_size)) {
+                projected = false;
+            } else {
+                draft_tok.clear();
+                draft_tok.reserve(static_cast<size_t>(max_verify_width));
+                draft_tok.push_back(last_tok);
+                draft_tok.insert(
+                    draft_tok.end(), proposals.begin(),
+                    proposals.begin() + draft_block_size);
+            }
+        } else if (!dspark_block) {
+            const int candidate_k =
+                width_theta > 0.0f && draft_block_size > 2 ? 2 : 0;
             projected = target.project_hidden_to_tokens_topk(
-                draft_hidden_host, max_q_len, draft_tok, candidate_k,
+                draft_hidden_host, draft_block_size, draft_tok, candidate_k,
                 candidate_k > 0 ? &candidate_probs : nullptr,
                 candidate_k > 0 ? &candidate_ids : nullptr);
+            if (projected &&
+                draft_tok.size() >= static_cast<size_t>(draft_block_size)) {
+                draft_tok[0] = last_tok;
+            }
         }
-        if (!projected || draft_tok.size() < (size_t)max_q_len) {
+        if (!projected ||
+            draft_tok.size() < static_cast<size_t>(max_verify_width)) {
             std::fprintf(stderr, "dflash-spec projection failed\n");
             return false;
         }
-        draft_tok[0] = last_tok;
 
-        if (confidence.size() >= (size_t)(max_q_len - 1)) {
+        if (dspark_block &&
+            confidence.size() >= static_cast<size_t>(draft_block_size)) {
             q_len = adaptive_verify_width(
-                confidence.data(), 1, max_q_len, width_theta, width_min);
+                confidence.data(), 1, max_verify_width,
+                width_theta, width_min);
         } else if (candidate_probs.size() >=
-                   static_cast<size_t>((max_q_len - 1) * 2)) {
+                   static_cast<size_t>((max_verify_width - 1) * 2)) {
             q_len = adaptive_verify_width(
-                candidate_probs.data(), 2, max_q_len,
+                candidate_probs.data(), 2, max_verify_width,
                 width_theta, width_min);
         }
 
@@ -223,8 +249,8 @@ bool run_dflash_spec_decode(
         int hint_filled = 0;
         if (hint_tokens && n_generated < (int)hint_tokens->size()) {
             const int hint_avail = (int)hint_tokens->size() - n_generated;
-            q_len = max_q_len;
-            hint_filled = std::min(hint_avail, max_q_len - 1);
+            q_len = max_verify_width;
+            hint_filled = std::min(hint_avail, max_verify_width - 1);
             for (int i = 0; i < hint_filled; i++) {
                 draft_tok[1 + i] = (*hint_tokens)[n_generated + i];
             }
diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp
index f0df52c10..d63b277df 100644
--- a/server/src/common/dspark_head.cpp
+++ b/server/src/common/dspark_head.cpp
@@ -3,9 +3,11 @@
 #include "ggml-alloc.h"
 #include "ddtree.h"
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <utility>
 #include <vector>
 
 namespace dflash::common {
@@ -106,8 +108,77 @@ bool dspark_step(const DraftWeights & dw,
     return true;
 }
 
+bool dspark_markov_candidates(const DraftWeights & dw,
+                              ggml_backend_t backend,
+                              DFlashTarget & target,
+                              const float * candidate_hidden,
+                              int n_candidates,
+                              int32_t anchor_token,
+                              float confidence_threshold,
+                              std::vector<int32_t> & candidates_out) {
+    if (!dw.dspark.enabled || !candidate_hidden || n_candidates <= 0) return false;
+    const int hidden = dw.n_embd;
+    if (hidden <= 0) return false;
+    confidence_threshold = std::clamp(confidence_threshold, 0.0f, 1.0f);
+    const bool use_confidence_gate =
+        confidence_threshold > 0.0f &&
+        dw.dspark.confidence_w != nullptr &&
+        dw.dspark.confidence_b != nullptr &&
+        dw.dspark.confidence_dim > 0;
+
+    std::vector<float> base_logits;
+    if (!target.project_hidden_to_logits(
+            candidate_hidden, n_candidates, base_logits)) {
+        return false;
+    }
+    if (base_logits.size() % static_cast<size_t>(n_candidates) != 0) return false;
+    const int vocab = static_cast<int>(
+        base_logits.size() / static_cast<size_t>(n_candidates));
+    if (dw.dspark.vocab_size > 0 && vocab != dw.dspark.vocab_size) {
+        std::fprintf(stderr,
+            "dspark_markov_candidates: vocab mismatch target=%d dspark=%d\n",
+            vocab, dw.dspark.vocab_size);
+        return false;
+    }
+
+    candidates_out.clear();
+    candidates_out.reserve(static_cast<size_t>(n_candidates));
+    int32_t previous = anchor_token;
+    for (int i = 0; i < n_candidates; ++i) {
+        int32_t token = -1;
+        float confidence = 0.0f;
+        if (!dspark_step(
+                dw, backend, previous,
+                candidate_hidden + static_cast<size_t>(i) * hidden,
+                base_logits.data() + static_cast<size_t>(i) * vocab,
+                vocab, token,
+                use_confidence_gate ? &confidence : nullptr)) {
+            return false;
+        }
+        if (use_confidence_gate && confidence < confidence_threshold) break;
+        candidates_out.push_back(token);
+        previous = token;
+    }
+    // An empty proposal list is a valid confidence-gated result.  Legacy
+    // callers still return the anchor token, while block callers can fall
+    // back to the non-speculative path when the full block is unavailable.
+    return true;
+}
+
 }  // namespace
 
+bool dspark_markov_propose_greedy_block(const DraftWeights & dw,
+                                        ggml_backend_t backend,
+                                        DFlashTarget & target,
+                                        const float * draft_hidden,
+                                        int proposal_len,
+                                        int32_t anchor_token,
+                                        std::vector<int32_t> & proposals_out) {
+    return dspark_markov_candidates(
+        dw, backend, target, draft_hidden, proposal_len, anchor_token,
+        /*confidence_threshold=*/0.0f, proposals_out);
+}
+
 bool dspark_markov_correct_greedy_chain(const DraftWeights & dw,
                                         ggml_backend_t backend,
                                         DFlashTarget & target,
@@ -122,53 +193,17 @@ bool dspark_markov_correct_greedy_chain(const DraftWeights & dw,
     if (hidden <= 0 || n_candidates <= 0) return false;
     if (confidence_threshold < 0.0f) confidence_threshold = 0.0f;
     if (confidence_threshold > 1.0f) confidence_threshold = 1.0f;
-    const bool use_confidence_gate =
-        confidence_threshold > 0.0f &&
-        dw.dspark.confidence_w != nullptr &&
-        dw.dspark.confidence_b != nullptr &&
-        dw.dspark.confidence_dim > 0;
-
-    std::vector<float> candidate_hidden((size_t)n_candidates * (size_t)hidden);
-    for (int i = 0; i < n_candidates; ++i) {
-        const float * src = local_hidden + (size_t)(i + 1) * (size_t)hidden;
-        std::memcpy(candidate_hidden.data() + (size_t)i * (size_t)hidden,
-                    src, sizeof(float) * (size_t)hidden);
-    }
-
-    std::vector<float> base_logits;
-    if (!target.project_hidden_to_logits(candidate_hidden.data(), n_candidates, base_logits)) {
-        return false;
-    }
-    if (base_logits.size() % (size_t)n_candidates != 0) return false;
-    const int vocab = (int)(base_logits.size() / (size_t)n_candidates);
-    if (dw.dspark.vocab_size > 0 && vocab != dw.dspark.vocab_size) {
-        std::fprintf(stderr, "dspark_markov_correct_greedy_chain: vocab mismatch target=%d dspark=%d\n",
-                     vocab, dw.dspark.vocab_size);
+    std::vector<int32_t> candidates;
+    if (!dspark_markov_candidates(
+            dw, backend, target,
+            local_hidden + static_cast<size_t>(hidden), n_candidates,
+            last_tok, confidence_threshold, candidates)) {
         return false;
     }
-
     draft_tok.clear();
-    draft_tok.reserve((size_t)q_len);
+    draft_tok.reserve(candidates.size() + 1);
     draft_tok.push_back(last_tok);
-    int32_t prefix_tok = last_tok;
-    for (int i = 0; i < n_candidates; ++i) {
-        int32_t tok = -1;
-        float confidence = 0.0f;
-        float * confidence_ptr = use_confidence_gate ? &confidence : nullptr;
-        if (!dspark_step(dw, backend, prefix_tok,
-                         candidate_hidden.data() + (size_t)i * (size_t)hidden,
-                         base_logits.data() + (size_t)i * (size_t)vocab,
-                         vocab,
-                         tok,
-                         confidence_ptr)) {
-            return false;
-        }
-        if (use_confidence_gate && confidence < confidence_threshold) {
-            break;
-        }
-        draft_tok.push_back(tok);
-        prefix_tok = tok;
-    }
+    draft_tok.insert(draft_tok.end(), candidates.begin(), candidates.end());
     return true;
 }
 
@@ -305,80 +340,127 @@ bool build_markov_chain_graph(const DraftWeights & dw,
     return true;
 }
 
-}  // namespace
-
-bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw,
-                                              ggml_backend_t backend,
-                                              ggml_tensor * lm_head,
-                                              const float * local_hidden,
-                                              int q_len,
-                                              int32_t last_tok,
-                                              std::vector<int32_t> & draft_tok,
-                                              std::vector<float> * confidence_out,
-                                              const float * confidence_hidden) {
-    if (q_len <= 1) return false;
-    if (!dspark_fused_usable(dw, backend, lm_head, local_hidden, "dspark_fused")) return false;
-    const int hdim   = dw.n_embd;
-    const int n_cand = q_len - 1;
+bool dspark_markov_propose_fused_impl(
+        const DraftWeights & dw,
+        ggml_backend_t backend,
+        ggml_tensor * lm_head,
+        const float * proposal_hidden,
+        int proposal_len,
+        int32_t anchor_token,
+        std::vector<int32_t> & proposals_out,
+        std::vector<float> * confidence_out,
+        const float * confidence_hidden) {
+    if (proposal_len <= 0 ||
+        !dspark_fused_usable(
+            dw, backend, lm_head, proposal_hidden, "dspark_fused")) {
+        return false;
+    }
+    const int hidden = dw.n_embd;
 
     static thread_local std::vector<uint8_t> g_arena_chain;
-    MarkovChainGraph g;
+    MarkovChainGraph graph;
     const bool want_confidence = confidence_out != nullptr;
     if (confidence_out) confidence_out->clear();
-    if (!build_markov_chain_graph(dw, lm_head, n_cand, /*first_corrected=*/0,
-                                  /*corrected_are_outputs=*/false,
-                                  /*confidence_are_outputs=*/want_confidence,
-                                  g_arena_chain, g)) {
+    if (!build_markov_chain_graph(
+            dw, lm_head, proposal_len, /*first_corrected=*/0,
+            /*corrected_are_outputs=*/false,
+            /*confidence_are_outputs=*/want_confidence,
+            g_arena_chain, graph)) {
         return false;
     }
 
     static thread_local ggml_gallocr_t galloc_chain = nullptr;
     if (!galloc_chain) {
-        galloc_chain = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
+        galloc_chain = ggml_gallocr_new(
+            ggml_backend_get_default_buffer_type(backend));
     }
-    if (!ggml_gallocr_alloc_graph(galloc_chain, g.gf)) {
+    if (!ggml_gallocr_alloc_graph(galloc_chain, graph.gf)) {
         std::fprintf(stderr, "dspark_fused: gallocr_alloc_graph failed\n");
-        ggml_free(g.ctx);
+        ggml_free(graph.ctx);
         return false;
     }
 
-    // Candidate hidden states start at position 1 (position 0 is the seed).
-    ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0,
-                            sizeof(float) * (size_t)hdim * (size_t)n_cand);
-    if (want_confidence && g.inp_confidence_hidden) {
-        const float * conf_src = confidence_hidden ? confidence_hidden : local_hidden;
-        ggml_backend_tensor_set(g.inp_confidence_hidden, conf_src + (size_t)hdim, 0,
-                                sizeof(float) * (size_t)hdim * (size_t)n_cand);
+    ggml_backend_tensor_set(
+        graph.inp_hidden, proposal_hidden, 0,
+        sizeof(float) * static_cast<size_t>(hidden) * proposal_len);
+    if (want_confidence && graph.inp_confidence_hidden) {
+        const float * source = confidence_hidden
+            ? confidence_hidden : proposal_hidden;
+        ggml_backend_tensor_set(
+            graph.inp_confidence_hidden, source, 0,
+            sizeof(float) * static_cast<size_t>(hidden) * proposal_len);
     }
-    ggml_backend_tensor_set(g.inp_seed, &last_tok, 0, sizeof(int32_t));
+    ggml_backend_tensor_set(
+        graph.inp_seed, &anchor_token, 0, sizeof(anchor_token));
 
-    if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) {
+    if (ggml_backend_graph_compute(backend, graph.gf) != GGML_STATUS_SUCCESS) {
         std::fprintf(stderr, "dspark_fused: graph_compute failed\n");
-        ggml_free(g.ctx);
+        ggml_free(graph.ctx);
         return false;
     }
 
-    draft_tok.assign((size_t)q_len, 0);
-    draft_tok[0] = last_tok;
-    // One synchronize instead of n_cand blocking readbacks.
-    int32_t t_out[16];
-    float c_out[16] = {};
-    const int n_get = n_cand < 16 ? n_cand : 16;
-    for (int i = 0; i < n_get; ++i) {
-        ggml_backend_tensor_get_async(backend, g.toks[(size_t)i], &t_out[i], 0, sizeof(int32_t));
-        if (want_confidence && g.confidence[(size_t)i]) {
+    proposals_out.resize(static_cast<size_t>(proposal_len));
+    std::vector<float> confidence(static_cast<size_t>(proposal_len), 0.0f);
+    for (int i = 0; i < proposal_len; ++i) {
+        ggml_backend_tensor_get_async(
+            backend, graph.toks[static_cast<size_t>(i)],
+            &proposals_out[static_cast<size_t>(i)], 0, sizeof(int32_t));
+        if (want_confidence && graph.confidence[static_cast<size_t>(i)]) {
             ggml_backend_tensor_get_async(
-                backend, g.confidence[(size_t)i], &c_out[i], 0, sizeof(float));
+                backend, graph.confidence[static_cast<size_t>(i)],
+                &confidence[static_cast<size_t>(i)], 0, sizeof(float));
         }
     }
     ggml_backend_synchronize(backend);
-    for (int i = 0; i < n_get; ++i) {
-        draft_tok[(size_t)i + 1] = t_out[i];
+    if (want_confidence && !graph.confidence.empty() && graph.confidence[0]) {
+        *confidence_out = std::move(confidence);
     }
-    if (want_confidence && !g.confidence.empty() && g.confidence[0]) {
-        confidence_out->assign(c_out, c_out + n_get);
+    ggml_free(graph.ctx);
+    return true;
+}
+
+}  // namespace
+
+bool dspark_markov_propose_greedy_block_fused(
+        const DraftWeights & dw,
+        ggml_backend_t backend,
+        ggml_tensor * lm_head,
+        const float * draft_hidden,
+        int proposal_len,
+        int32_t anchor_token,
+        std::vector<int32_t> & proposals_out,
+        std::vector<float> * confidence_out,
+        const float * confidence_hidden) {
+    return dspark_markov_propose_fused_impl(
+        dw, backend, lm_head, draft_hidden, proposal_len, anchor_token,
+        proposals_out, confidence_out, confidence_hidden);
+}
+
+bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw,
+                                              ggml_backend_t backend,
+                                              ggml_tensor * lm_head,
+                                              const float * local_hidden,
+                                              int q_len,
+                                              int32_t last_tok,
+                                              std::vector<int32_t> & draft_tok,
+                                              std::vector<float> * confidence_out,
+                                              const float * confidence_hidden) {
+    if (q_len <= 1) return false;
+    const int hidden = dw.n_embd;
+    const int n_cand = q_len - 1;
+    std::vector<int32_t> candidates;
+    const float * candidate_confidence_hidden = confidence_hidden
+        ? confidence_hidden + static_cast<size_t>(hidden) : nullptr;
+    if (!dspark_markov_propose_fused_impl(
+            dw, backend, lm_head,
+            local_hidden + static_cast<size_t>(hidden), n_cand, last_tok,
+            candidates, confidence_out, candidate_confidence_hidden)) {
+        return false;
     }
-    ggml_free(g.ctx);
+    draft_tok.clear();
+    draft_tok.reserve(candidates.size() + 1);
+    draft_tok.push_back(last_tok);
+    draft_tok.insert(draft_tok.end(), candidates.begin(), candidates.end());
     return true;
 }
 
diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h
index 9b97b261d..91658d2db 100644
--- a/server/src/common/dspark_head.h
+++ b/server/src/common/dspark_head.h
@@ -8,6 +8,29 @@
 
 namespace dflash::common {
 
+// Reference DSpark contract: each of the `proposal_len` draft hidden rows
+// predicts one proposal. The known target token is only the Markov-chain
+// anchor; it is not part of proposals_out. This differs from the historical
+// Luce DFlash chain helper below, whose output starts with that anchor.
+bool dspark_markov_propose_greedy_block(const DraftWeights & dw,
+                                        ggml_backend_t backend,
+                                        DFlashTarget & target,
+                                        const float * draft_hidden,
+                                        int proposal_len,
+                                        int32_t anchor_token,
+                                        std::vector<int32_t> & proposals_out);
+
+bool dspark_markov_propose_greedy_block_fused(
+    const DraftWeights & dw,
+    ggml_backend_t backend,
+    ggml_tensor * lm_head,
+    const float * draft_hidden,
+    int proposal_len,
+    int32_t anchor_token,
+    std::vector<int32_t> & proposals_out,
+    std::vector<float> * confidence_out = nullptr,
+    const float * confidence_hidden = nullptr);
+
 bool dspark_markov_correct_greedy_chain(const DraftWeights & dw,
                                         ggml_backend_t backend,
                                         DFlashTarget & target,
diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp
index 05089db48..de885e323 100644
--- a/server/src/draft/draft_gguf_loader.cpp
+++ b/server/src/draft/draft_gguf_loader.cpp
@@ -31,6 +31,7 @@
 
 #include <algorithm>
 #include <cinttypes>
+#include <cmath>
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
@@ -149,6 +150,11 @@ bool load_draft_gguf(const std::string & path,
         if (id < 0) return fallback;
         return gguf_get_val_f32(gctx, id);
     };
+    auto read_string = [&](const char * suffix) -> std::string {
+        std::snprintf(key, sizeof(key), "%s.%s", A, suffix);
+        const int64_t id = gguf_find_key(gctx, key);
+        return id >= 0 ? std::string(gguf_get_val_str(gctx, id)) : std::string();
+    };
 
     const uint32_t n_embd    = read_u32("embedding_length",        0);
     const uint32_t n_layer   = read_u32("block_count",             0);
@@ -243,10 +249,52 @@ bool load_draft_gguf(const std::string & path,
     out.head_dim  = (int)head_dim;
     out.n_embd    = (int)n_embd;
     out.n_ff      = (int)n_ff;
+    out.rms_eps = read_f32(
+        "attention.layer_norm_rms_epsilon", DFLASH27B_RMS_EPS);
+    if (!std::isfinite(out.rms_eps) || out.rms_eps <= 0.0f) {
+        set_last_error("draft GGUF: invalid attention.layer_norm_rms_epsilon");
+        ggml_free(meta_ctx);
+        out.ctx = nullptr;
+        gguf_free(gctx);
+        return false;
+    }
     out.rope_theta = read_f32("rope.freq_base", 0.0f);
     if (out.rope_theta == 0.0f) {
         fprintf(stderr, "[draft-gguf] WARNING: rope.freq_base not found in GGUF, draft RoPE will be wrong\n");
     }
+    const std::string rope_scaling_type = read_string("rope.scaling.type");
+    if (!rope_scaling_type.empty() && rope_scaling_type != "none") {
+        if (rope_scaling_type != "yarn") {
+            set_last_error("draft GGUF: unsupported rope.scaling.type=" +
+                           rope_scaling_type);
+            ggml_free(meta_ctx);
+            out.ctx = nullptr;
+            gguf_free(gctx);
+            return false;
+        }
+        const float factor = read_f32("rope.scaling.factor", 0.0f);
+        const uint32_t original_context =
+            read_u32("rope.scaling.original_context_length", 0);
+        if (!std::isfinite(factor) || factor <= 0.0f || original_context == 0) {
+            set_last_error("draft GGUF: incomplete or invalid YaRN metadata");
+            ggml_free(meta_ctx);
+            out.ctx = nullptr;
+            gguf_free(gctx);
+            return false;
+        }
+        // Match Transformers' YaRN defaults when the checkpoint does not
+        // provide explicit attention/beta overrides.
+        out.rope_freq_scale = 1.0f / factor;
+        out.rope_ext_factor = 1.0f;
+        out.rope_attn_factor = 0.1f * std::log(factor) + 1.0f;
+        out.rope_beta_fast = 32.0f;
+        out.rope_beta_slow = 1.0f;
+        out.rope_n_ctx_orig = static_cast<int>(original_context);
+        std::fprintf(stderr,
+            "[draft-gguf] YaRN enabled: factor=%.3f original_ctx=%u "
+            "attn_factor=%.6f\n",
+            factor, original_context, out.rope_attn_factor);
+    }
     out.layers.assign((size_t)n_layer, DraftLayer{});
 
     auto g = [&](const char * name) -> ggml_tensor * {
diff --git a/server/src/draft/draft_graph.cpp b/server/src/draft/draft_graph.cpp
index 472c214c9..75fcf8185 100644
--- a/server/src/draft/draft_graph.cpp
+++ b/server/src/draft/draft_graph.cpp
@@ -50,7 +50,7 @@ static ggml_tensor * draft_fuse_features(
     ggml_tensor *        target_hidden_cat,
     int                  n_rows,
     bool                 disable_aux_hidden_norms) {
-    const float eps = DFLASH27B_RMS_EPS;
+    const float eps = w.rms_eps;
     ggml_tensor * thc = target_hidden_cat;
     if (!disable_aux_hidden_norms && !w.aux_hidden_norms.empty()) {
         ggml_tensor * aux_cat = nullptr;
@@ -72,6 +72,17 @@ static ggml_tensor * draft_fuse_features(
     return target_feat;
 }
 
+static ggml_tensor * draft_rope(ggml_context * ctx,
+                                const DraftWeights & w,
+                                ggml_tensor * input,
+                                ggml_tensor * positions) {
+    return ggml_rope_ext(
+        ctx, input, positions, /*freq_factors=*/nullptr,
+        w.head_dim, GGML_ROPE_TYPE_NEOX, w.rope_n_ctx_orig,
+        w.rope_theta, w.rope_freq_scale, w.rope_ext_factor,
+        w.rope_attn_factor, w.rope_beta_fast, w.rope_beta_slow);
+}
+
 DraftGraphOutputs build_draft_graph(
     ggml_context *            ctx,
     const DraftWeights &      w,
@@ -82,8 +93,7 @@ DraftGraphOutputs build_draft_graph(
     const int n_head   = w.n_head;
     const int n_kv     = w.n_head_kv;
     const int head_dim = w.head_dim;
-    const float eps    = DFLASH27B_RMS_EPS;
-    const float rope_base = w.rope_theta;
+    const float eps    = w.rms_eps;
 
     // ── 1. Feature fusion: target_feat = rms_norm(fc @ target_hidden_cat, hidden_norm)
     //    fc:                [5*hidden, hidden]  (ggml: ne[0]=5*hidden, ne[1]=hidden)
@@ -185,14 +195,8 @@ DraftGraphOutputs build_draft_graph(
                 pk = ggml_view_1d(ctx, in.positions_k, eff_total_k,
                                   ctx_offset * ggml_element_size(in.positions_k));
             }
-            Q = ggml_rope_ext(ctx, Q, in.positions_q, /*freq_factors=*/nullptr,
-                              head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0,
-                              rope_base, /*freq_scale=*/1.0f,
-                              /*ext_factor=*/0.0f, /*attn_factor=*/1.0f,
-                              /*beta_fast=*/0.0f, /*beta_slow=*/0.0f);
-            K = ggml_rope_ext(ctx, K, pk, nullptr,
-                              head_dim, GGML_ROPE_TYPE_NEOX, 0,
-                              rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+            Q = draft_rope(ctx, w, Q, in.positions_q);
+            K = draft_rope(ctx, w, K, pk);
 
             // ── 2e. Permute into the layout flash_attn_ext wants
             //   q: [n_embd_k=head_dim, n_batch=q_len, n_head,   ne3]
@@ -299,7 +303,7 @@ static void draft_ctx_kv_rows(
     int                  n,
     ggml_tensor **       k_rows_out,
     ggml_tensor **       v_rows_out) {
-    const float eps = DFLASH27B_RMS_EPS;
+    const float eps = w.rms_eps;
     ggml_tensor * tf_kv = target_feat;
     if (w.context_kv_layer_norm) {
         tf_kv = ggml_rms_norm(ctx, tf_kv, eps);
@@ -309,11 +313,7 @@ static void draft_ctx_kv_rows(
     K = ggml_reshape_3d(ctx, K, w.head_dim, w.n_head_kv, n);
     K = ggml_rms_norm(ctx, K, eps);
     K = ggml_mul     (ctx, K, L.k_norm);
-    K = ggml_rope_ext(ctx, K, positions, /*freq_factors=*/nullptr,
-                      w.head_dim, GGML_ROPE_TYPE_NEOX, /*n_ctx_orig=*/0,
-                      w.rope_theta, /*freq_scale=*/1.0f,
-                      /*ext_factor=*/0.0f, /*attn_factor=*/1.0f,
-                      /*beta_fast=*/0.0f, /*beta_slow=*/0.0f);
+    K = draft_rope(ctx, w, K, positions);
     // rope output is contiguous [head_dim, n_kv, n] → head-major rows view
     *k_rows_out = ggml_view_2d(ctx, K, (int64_t)w.head_dim * w.n_head_kv, n,
                                K->nb[2], 0);
@@ -355,8 +355,7 @@ DraftGraphOutputs build_draft_kv_step(
     const int n_head   = w.n_head;
     const int n_kv     = w.n_head_kv;
     const int head_dim = w.head_dim;
-    const float eps    = DFLASH27B_RMS_EPS;
-    const float rope_base = w.rope_theta;
+    const float eps    = w.rms_eps;
     const int kv_total = cache.kv_total;
 
     static const bool disable_attn_gate =
@@ -380,18 +379,14 @@ DraftGraphOutputs build_draft_kv_step(
         Q = ggml_reshape_3d(ctx, Q, head_dim, n_head, q_len);
         Q = ggml_rms_norm(ctx, Q, eps);
         Q = ggml_mul     (ctx, Q, L.q_norm);
-        Q = ggml_rope_ext(ctx, Q, in.positions_q, nullptr,
-                          head_dim, GGML_ROPE_TYPE_NEOX, 0,
-                          rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        Q = draft_rope(ctx, w, Q, in.positions_q);
 
         // ── noise K/V into the scratch cache slots
         ggml_tensor * Kn = ggml_mul_mat(ctx, L.wk, hn);
         Kn = ggml_reshape_3d(ctx, Kn, head_dim, n_kv, q_len);
         Kn = ggml_rms_norm(ctx, Kn, eps);
         Kn = ggml_mul     (ctx, Kn, L.k_norm);
-        Kn = ggml_rope_ext(ctx, Kn, in.positions_q, nullptr,
-                           head_dim, GGML_ROPE_TYPE_NEOX, 0,
-                           rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        Kn = draft_rope(ctx, w, Kn, in.positions_q);
         ggml_tensor * Kn_rows = ggml_view_2d(ctx, Kn,
             (int64_t)head_dim * n_kv, q_len, Kn->nb[2], 0);
         ggml_tensor * Vn_rows = ggml_mul_mat(ctx, L.wv, hn);  // [kv_dim, q_len]
diff --git a/server/src/internal.h b/server/src/internal.h
index ceadfc18b..3b4f23ad3 100644
--- a/server/src/internal.h
+++ b/server/src/internal.h
@@ -299,6 +299,7 @@ struct DraftWeights {
     int n_embd    = DFLASH27B_TARGET_HIDDEN;           // 5120
     int n_ff      = DFLASH27B_TARGET_INTERMEDIATE;     // 17408
     int swa_window = 0;  // sliding window size (0 = disabled)
+    float rms_eps = DFLASH27B_RMS_EPS;
     float rope_theta = 0.0f;  // RoPE frequency base (must come from GGUF)
 
     // YaRN rope scaling (populated by loader; 0 = disabled / plain RoPE).
@@ -323,6 +324,14 @@ struct DraftWeights {
     // Optional DSpark/DeepSpec-style Markov correction head. When present,
     // greedy chain decode adds a low-rank previous-token bias before argmax.
     DraftDSparkWeights dspark;
+
+    // DFlash's historical chain contract replaces draft row 0 with the known
+    // target seed, so block_size rows are verified.  DSpark instead predicts
+    // one token from every draft row: block_size proposals are verified after
+    // the seed, matching SGLang's gamma+1 target window.
+    int max_chain_verify_tokens() const {
+        return block_size + (dspark.enabled ? 1 : 0);
+    }
 };
 
 bool load_draft_safetensors(const std::string & path,
diff --git a/server/src/kimi_k3/kimi_k3_backend.cpp b/server/src/kimi_k3/kimi_k3_backend.cpp
index 35e29245e..763fd3ab0 100644
--- a/server/src/kimi_k3/kimi_k3_backend.cpp
+++ b/server/src/kimi_k3/kimi_k3_backend.cpp
@@ -799,7 +799,7 @@ bool KimiK3Backend::init() {
     if (!init_draft()) return false;
     const int max_ctx = std::max(1, cfg_.device.max_ctx);
     const int max_verify_tokens = draft_weights_.ctx
-        ? draft_weights_.block_size : 0;
+        ? draft_weights_.max_chain_verify_tokens() : 0;
     if (!create_kimi_k3_cache(
             backend_, weights_, max_ctx, cache_, max_verify_tokens)) {
         std::fprintf(stderr, "[kimi-k3] cache allocation failed (max_ctx=%d)\n",
diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp
index b0ba2fd8e..060c453a1 100644
--- a/server/tests/test_deepseek4_unit.cpp
+++ b/server/tests/test_deepseek4_unit.cpp
@@ -411,6 +411,83 @@ static void test_dspark_confidence_uses_separate_hidden(ggml_backend_t backend)
     std::fprintf(stderr, g_failures ? " done\n" : " ok\n");
 }
 
+static void test_dspark_proposal_block_uses_every_hidden_row(
+        ggml_backend_t backend) {
+    std::fprintf(stderr,
+                 "  test_dspark_proposal_block_uses_every_hidden_row ...");
+
+    constexpr int hidden = 2;
+    constexpr int rank = 1;
+    constexpr int vocab = 3;
+    constexpr int proposal_len = 2;
+
+    ggml_context * ctx = make_test_context();
+    TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed");
+    if (!ctx) {
+        std::fprintf(stderr, " FAIL\n");
+        return;
+    }
+
+    ggml_tensor * lm_head =
+        ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, vocab);
+    ggml_tensor * markov_w1 =
+        ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rank, vocab);
+    ggml_tensor * markov_w2 =
+        ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rank, vocab);
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
+    TEST_ASSERT_MSG(buf != nullptr, "weight allocation failed");
+    if (!buf) {
+        ggml_free(ctx);
+        std::fprintf(stderr, " FAIL\n");
+        return;
+    }
+
+    // Hidden row 0 selects token 1; hidden row 1 selects token 2. A runtime
+    // that incorrectly treats row 0 as the already-known seed returns only
+    // the second token and cannot satisfy this contract.
+    const std::vector<float> lm_values = {
+        0.0f, 0.0f,
+        1.0f, 0.0f,
+        0.0f, 1.0f,
+    };
+    const std::vector<float> markov_zeros((size_t)rank * vocab, 0.0f);
+    const std::vector<float> draft_hidden = {
+        3.0f, 0.0f,
+        0.0f, 3.0f,
+    };
+    ggml_backend_tensor_set(
+        lm_head, lm_values.data(), 0, lm_values.size() * sizeof(float));
+    ggml_backend_tensor_set(
+        markov_w1, markov_zeros.data(), 0,
+        markov_zeros.size() * sizeof(float));
+    ggml_backend_tensor_set(
+        markov_w2, markov_zeros.data(), 0,
+        markov_zeros.size() * sizeof(float));
+
+    DraftWeights dw{};
+    dw.n_embd = hidden;
+    dw.block_size = proposal_len;
+    dw.dspark.enabled = true;
+    dw.dspark.markov_rank = rank;
+    dw.dspark.vocab_size = vocab;
+    dw.dspark.markov_w1 = markov_w1;
+    dw.dspark.markov_w2 = markov_w2;
+
+    std::vector<int32_t> proposals;
+    const bool ok = dspark_markov_propose_greedy_block_fused(
+        dw, backend, lm_head, draft_hidden.data(), proposal_len,
+        /*anchor_token=*/0, proposals);
+    TEST_ASSERT_MSG(ok, "DSpark proposal block failed");
+    TEST_ASSERT(proposals == std::vector<int32_t>({1, 2}));
+    TEST_ASSERT(dw.max_chain_verify_tokens() == proposal_len + 1);
+    dw.dspark.enabled = false;
+    TEST_ASSERT(dw.max_chain_verify_tokens() == proposal_len);
+
+    ggml_backend_buffer_free(buf);
+    ggml_free(ctx);
+    std::fprintf(stderr, g_failures ? " done\n" : " ok\n");
+}
+
 static float softplus_stable(float x) {
     if (x > 20.0f) {
         return x;
@@ -3412,6 +3489,7 @@ int main() {
     test_loader_rejects_truncated_tensor_data(backend);
     test_dspark_loader_contract_and_bounds(backend);
     test_dspark_confidence_uses_separate_hidden(backend);
+    test_dspark_proposal_block_uses_every_hidden_row(backend);
     test_safe_compressor_batch_tokens();
     test_dspark_park_all_releases_drafter();
     test_dspark_raw_ring_rollback_after_wrap(backend);