Add DSpark Markov draft-head support for Laguna

server/CMakeLists.txt

@@ -258,6 +258,7 @@ add_library(dflash_common STATIC
     src/laguna/laguna_dflash_target.cpp
     src/common/backend_ipc.cpp
     src/common/domino_head.cpp
+    src/common/dspark_head.cpp
     src/common/target_shard_ipc.cpp
     src/common/target_shard_ipc_daemon.cpp
     src/common/dflash_feature_ring.cpp

server/scripts/convert_dflash_to_gguf.py

@@ -244,6 +244,17 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray:
 }
 
 
+DSPARK_TENSOR_MAP = {
+    "dspark_markov_head.markov_w1.weight": ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16),
+    "dspark_markov_head.markov_w2.weight": ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16),
+}
+
+DSPARK_CONFIDENCE_TENSOR_MAP = {
+    "dspark_confidence_head.weight": ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16),
+    "dspark_confidence_head.bias": ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32),
+}
+
+
 def add_domino_aux_heads(writer, arch: str, aux_path: Path | None):
     if aux_path is None:
         return
@@ -265,6 +276,9 @@ def add_domino_aux_heads(writer, arch: str, aux_path: Path | None):
         print(f"[error] Domino aux heads file is not a tensor dict: {aux_path}", file=sys.stderr)
         sys.exit(1)
 
+    if not any(k in state for k in DOMINO_TENSOR_MAP):
+        return
+
     missing = [k for k in DOMINO_TENSOR_MAP if k not in state]
     if missing:
         print(f"[warn] incomplete Domino aux heads; missing {missing}; skipping Domino tensors")
@@ -294,6 +308,79 @@ def add_domino_aux_heads(writer, arch: str, aux_path: Path | None):
         print(f"[tensor] {gguf_name:50s} aux ->{raw_dtype.name:4s} {tuple(arr.shape)}")
 
 
+def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None):
+    if aux_path is None:
+        return
+    if not aux_path.exists():
+        return
+
+    try:
+        import torch
+    except ImportError as exc:
+        print(f"[error] --aux-heads requires torch: {exc}", file=sys.stderr)
+        sys.exit(1)
+
+    state = torch.load(aux_path, map_location="cpu")
+    if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict):
+        state = state["state_dict"]
+    if not isinstance(state, dict):
+        print(f"[error] DSpark aux heads file is not a tensor dict: {aux_path}", file=sys.stderr)
+        sys.exit(1)
+
+    missing = [k for k in DSPARK_TENSOR_MAP if k not in state]
+    if missing:
+        return
+
+    print(f"[info] reading DSpark aux heads from {aux_path}")
+    w1 = state["dspark_markov_head.markov_w1.weight"]
+    w2 = state["dspark_markov_head.markov_w2.weight"]
+    vocab = int(w1.shape[0])
+    rank = int(w1.shape[1])
+    if tuple(w2.shape) != (vocab, rank):
+        print(f"[error] DSpark markov_w2 shape {tuple(w2.shape)} != {(vocab, rank)}", file=sys.stderr)
+        sys.exit(1)
+
+    writer.add_uint32(f"{arch}.dflash.dspark.enabled", 1)
+    writer.add_uint32(f"{arch}.dflash.dspark.markov_rank", rank)
+    writer.add_uint32(f"{arch}.dflash.dspark.vocab_size", vocab)
+
+    for st_name, (gguf_name, raw_dtype) in DSPARK_TENSOR_MAP.items():
+        t = state[st_name]
+        if hasattr(t, "detach"):
+            t = t.detach().cpu()
+        arr = t.float().numpy().astype("<f2")
+        writer.add_tensor(gguf_name, arr, raw_dtype=raw_dtype)
+        print(f"[tensor] {gguf_name:50s} aux ->{raw_dtype.name:4s} {tuple(arr.shape)}")
+
+    conf_missing = [k for k in DSPARK_CONFIDENCE_TENSOR_MAP if k not in state]
+    if conf_missing:
+        print(f"[warn] incomplete DSpark confidence head; missing {conf_missing}; Markov head will still load")
+        return
+
+    conf_w = state["dspark_confidence_head.weight"]
+    conf_b = state["dspark_confidence_head.bias"]
+    confidence_dim = int(conf_w.shape[1])
+    if int(conf_w.shape[0]) != 1 or tuple(conf_b.shape) != (1,):
+        print(
+            f"[error] DSpark confidence shapes weight={tuple(conf_w.shape)} bias={tuple(conf_b.shape)}",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    writer.add_uint32(f"{arch}.dflash.dspark.confidence_dim", confidence_dim)
+    writer.add_uint32(f"{arch}.dflash.dspark.confidence.enabled", 1)
+    for st_name, (gguf_name, raw_dtype) in DSPARK_CONFIDENCE_TENSOR_MAP.items():
+        t = state[st_name]
+        if hasattr(t, "detach"):
+            t = t.detach().cpu()
+        arr = t.float().numpy()
+        if raw_dtype == gguf.GGMLQuantizationType.F16:
+            arr = arr.astype("<f2")
+        else:
+            arr = arr.astype("<f4")
+        writer.add_tensor(gguf_name, arr, raw_dtype=raw_dtype)
+        print(f"[tensor] {gguf_name:50s} aux ->{raw_dtype.name:4s} {tuple(arr.shape)}")
+
+
 # ──────────────────────────────────────────────────────────────────────
 # Main
 # ──────────────────────────────────────────────────────────────────────
@@ -400,6 +487,7 @@ def sort_key(t):
     if not args.no_aux_heads:
         aux_path = args.aux_heads if args.aux_heads is not None else args.safetensors.parent / "dflash_aux_heads.pt"
     add_domino_aux_heads(writer, ARCH, aux_path)
+    add_dspark_aux_heads(writer, ARCH, aux_path)
 
     print(f"[info] writing {args.out_gguf}")
     writer.write_header_to_file()

server/src/common/dspark_head.cpp

@@ -0,0 +1,174 @@
+#include "dspark_head.h"
+
+#include "ggml-alloc.h"
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+namespace dflash::common {
+
+namespace {
+
+bool dspark_step(const DraftWeights & dw,
+                 ggml_backend_t backend,
+                 int32_t prev_token,
+                 const float * draft_hidden,
+                 const float * base_logits,
+                 int vocab,
+                 int32_t & out_token,
+                 float * confidence_out) {
+    const int hidden = dw.n_embd;
+    const int rank = dw.dspark.markov_rank;
+    if (hidden <= 0 || rank <= 0 || vocab <= 0) return false;
+    if (!dw.dspark.markov_w1 || !dw.dspark.markov_w2) return false;
+
+    const bool want_conf =
+        confidence_out != nullptr &&
+        dw.dspark.confidence_w != nullptr &&
+        dw.dspark.confidence_b != nullptr &&
+        dw.dspark.confidence_dim > 0;
+
+    const size_t arena_size =
+        ggml_tensor_overhead() * 256 + ggml_graph_overhead() + 2 * 1024 * 1024;
+    static thread_local std::vector<uint8_t> g_arena;
+    if (g_arena.size() < arena_size) g_arena.resize(arena_size);
+
+    ggml_init_params ip{};
+    ip.mem_size = arena_size;
+    ip.mem_buffer = g_arena.data();
+    ip.no_alloc = true;
+    ggml_context * ctx = ggml_init(ip);
+    if (!ctx) return false;
+    ggml_cgraph * gf = ggml_new_graph_custom(ctx, 256, false);
+
+    ggml_tensor * inp_prev = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
+    ggml_tensor * inp_base = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, vocab, 1);
+    ggml_set_input(inp_prev);
+    ggml_set_input(inp_base);
+
+    ggml_tensor * prev_emb = ggml_get_rows(ctx, dw.dspark.markov_w1, inp_prev);
+    ggml_tensor * bias = ggml_mul_mat(ctx, dw.dspark.markov_w2, prev_emb);
+    ggml_tensor * corrected = ggml_add(ctx, inp_base, bias);
+    ggml_tensor * tok = ggml_argmax(ctx, corrected);
+    ggml_set_output(tok);
+    ggml_build_forward_expand(gf, tok);
+
+    ggml_tensor * conf = nullptr;
+    ggml_tensor * inp_hidden = nullptr;
+    if (want_conf) {
+        inp_hidden = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, 1);
+        ggml_set_input(inp_hidden);
+        ggml_tensor * conf_in = inp_hidden;
+        if (dw.dspark.confidence_dim == hidden + rank) {
+            conf_in = ggml_concat(ctx, inp_hidden, prev_emb, 0);
+        } else if (dw.dspark.confidence_dim != hidden) {
+            ggml_free(ctx);
+            return false;
+        }
+        conf = ggml_mul_mat(ctx, dw.dspark.confidence_w, conf_in);
+        conf = ggml_add(ctx, conf, ggml_reshape_2d(ctx, dw.dspark.confidence_b, 1, 1));
+        conf = ggml_sigmoid(ctx, conf);
+        ggml_set_output(conf);
+        ggml_build_forward_expand(gf, conf);
+    }
+
+    static thread_local ggml_gallocr_t galloc = nullptr;
+    if (!galloc) {
+        galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
+    }
+    if (!ggml_gallocr_alloc_graph(galloc, gf)) {
+        std::fprintf(stderr, "dspark_step: gallocr_alloc_graph failed\n");
+        ggml_free(ctx);
+        return false;
+    }
+
+    ggml_backend_tensor_set(inp_prev, &prev_token, 0, sizeof(prev_token));
+    ggml_backend_tensor_set(inp_base, base_logits, 0, sizeof(float) * (size_t)vocab);
+    if (want_conf) {
+        ggml_backend_tensor_set(inp_hidden, draft_hidden, 0,
+                                sizeof(float) * (size_t)hidden);
+    }
+
+    if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) {
+        std::fprintf(stderr, "dspark_step: graph_compute failed\n");
+        ggml_free(ctx);
+        return false;
+    }
+
+    ggml_backend_tensor_get(tok, &out_token, 0, sizeof(out_token));
+    if (want_conf) {
+        ggml_backend_tensor_get(conf, confidence_out, 0, sizeof(float));
+    }
+    ggml_free(ctx);
+    return true;
+}
+
+}  // namespace
+
+bool dspark_markov_correct_greedy_chain(const DraftWeights & dw,
+                                        ggml_backend_t backend,
+                                        DFlashTarget & target,
+                                        const float * local_hidden,
+                                        int q_len,
+                                        int32_t last_tok,
+                                        float confidence_threshold,
+                                        std::vector<int32_t> & draft_tok) {
+    if (!dw.dspark.enabled || q_len <= 1 || !local_hidden) return false;
+    const int hidden = dw.n_embd;
+    const int n_candidates = q_len - 1;
+    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);
+        return false;
+    }
+
+    draft_tok.clear();
+    draft_tok.reserve((size_t)q_len);
+    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;
+    }
+    return true;
+}
+
+}  // namespace dflash::common

server/src/common/dspark_head.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "dflash_target.h"
+#include "internal.h"
+
+#include <cstdint>
+#include <vector>
+
+namespace dflash::common {
+
+bool dspark_markov_correct_greedy_chain(const DraftWeights & dw,
+                                        ggml_backend_t backend,
+                                        DFlashTarget & target,
+                                        const float * local_hidden,
+                                        int q_len,
+                                        int32_t last_tok,
+                                        float confidence_threshold,
+                                        std::vector<int32_t> & draft_tok);
+
+}  // namespace dflash::common

server/src/draft/draft_gguf_loader.cpp

@@ -479,6 +479,10 @@ bool load_draft_gguf(const std::string & path,
     const uint32_t domino_meta_gru     = read_u32("dflash.domino.gru_hidden_dim", 0);
     const uint32_t domino_meta_emb     = read_u32("dflash.domino.emb_dim", 0);
     const uint32_t domino_meta_vocab   = read_u32("dflash.domino.vocab_size", 0);
+    const uint32_t dspark_meta_enabled = read_u32("dflash.dspark.enabled", 0);
+    const uint32_t dspark_meta_rank    = read_u32("dflash.dspark.markov_rank", 0);
+    const uint32_t dspark_meta_vocab   = read_u32("dflash.dspark.vocab_size", 0);
+    const uint32_t dspark_meta_conf    = read_u32("dflash.dspark.confidence_dim", 0);
     // Explicit captured target-layer ids (data-driven). Lets any DFlash drafter
     // load without a hardcoded per-arch set; the array length also backstops
     // n_target_layers when the scalar KV is absent.
@@ -665,6 +669,61 @@ bool load_draft_gguf(const std::string & path,
                      out.domino.vocab_size);
     }
 
+    out.dspark = DraftDSparkWeights{};
+    out.dspark.markov_w1    = g("dflash.dspark.markov.w1");
+    out.dspark.markov_w2    = g("dflash.dspark.markov.w2");
+    out.dspark.confidence_w = g("dflash.dspark.confidence.weight");
+    out.dspark.confidence_b = g("dflash.dspark.confidence.bias");
+
+    const bool dspark_any =
+        out.dspark.markov_w1 || out.dspark.markov_w2 ||
+        out.dspark.confidence_w || out.dspark.confidence_b ||
+        dspark_meta_enabled != 0;
+    if (dspark_any) {
+        if (!out.dspark.markov_w1 || !out.dspark.markov_w2) {
+            set_last_error("draft GGUF: incomplete DSpark Markov tensors");
+            gguf_free(gctx);
+            return false;
+        }
+        out.dspark.markov_rank =
+            dspark_meta_rank != 0 ? (int)dspark_meta_rank : (int)out.dspark.markov_w1->ne[0];
+        out.dspark.vocab_size =
+            dspark_meta_vocab != 0 ? (int)dspark_meta_vocab : (int)out.dspark.markov_w1->ne[1];
+
+        const int64_t R = out.dspark.markov_rank;
+        const int64_t V = out.dspark.vocab_size;
+        char shape_err[320];
+        if (!check_shape_2d(out.dspark.markov_w1, R, V, "dspark.markov.w1", shape_err, sizeof(shape_err)) ||
+            !check_shape_2d(out.dspark.markov_w2, R, V, "dspark.markov.w2", shape_err, sizeof(shape_err))) {
+            set_last_error(shape_err);
+            gguf_free(gctx);
+            return false;
+        }
+
+        const bool conf_any = out.dspark.confidence_w || out.dspark.confidence_b || dspark_meta_conf != 0;
+        if (conf_any) {
+            if (!out.dspark.confidence_w || !out.dspark.confidence_b) {
+                set_last_error("draft GGUF: incomplete DSpark confidence tensors");
+                gguf_free(gctx);
+                return false;
+            }
+            out.dspark.confidence_dim =
+                dspark_meta_conf != 0 ? (int)dspark_meta_conf : (int)out.dspark.confidence_w->ne[0];
+            const int64_t C = out.dspark.confidence_dim;
+            if (!check_shape_2d(out.dspark.confidence_w, C, 1, "dspark.confidence.weight", shape_err, sizeof(shape_err)) ||
+                !check_shape_1d(out.dspark.confidence_b, 1, "dspark.confidence.bias", shape_err, sizeof(shape_err))) {
+                set_last_error(shape_err);
+                gguf_free(gctx);
+                return false;
+            }
+        }
+
+        out.dspark.enabled = true;
+        std::fprintf(stderr, "[draft GGUF] DSpark Markov head enabled: rank=%d vocab=%d confidence_dim=%d\n",
+                     out.dspark.markov_rank, out.dspark.vocab_size,
+                     out.dspark.confidence_dim);
+    }
+
     // GGUF Qwen3.6 drafters carry SWA metadata emitted by the converter:
     //   dflash-draft.attention.sliding_window = 2048
     //   dflash-draft.attention.sliding_window_pattern = [true,true,true,true,false]