model : add DSpark draft support for Spark-X2.6

conversion/spark2_5.py

@@ -1,19 +1,26 @@
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
+from copy import deepcopy
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf
+from .base import ModelBase, TextModel, gguf, torch
 
 
 @ModelBase.register("Spark2_5ForCausalLM")
 @ModelBase.example("XHToken/Spark-X2.5-1.7B")
 class Spark2_5Model(TextModel):
     model_arch = gguf.MODEL_ARCH.SPARK2_5
 
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if item[0].startswith("draft_model."):
+            return None
+        return super().filter_tensors(item)
+
     def set_gguf_parameters(self) -> None:
         super().set_gguf_parameters()
 
@@ -63,3 +70,62 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
                 )
 
         yield from super().modify_tensors(data_torch, name, bid)
+
+
+class Spark2_5DSparkModel(Spark2_5Model):
+    model_arch = gguf.MODEL_ARCH.DFLASH
+
+    def __init__(self, dir_model, *args, **kwargs):
+        hparams = deepcopy(kwargs.pop("hparams", None) or ModelBase.load_hparams(dir_model, False))
+        n_layers = hparams["num_draft_layers"]
+        hparams["num_hidden_layers"] = n_layers
+        hparams["layer_types"] = hparams.get("draft_layer_types", hparams["layer_types"][:n_layers])
+        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
+
+        if self.hparams.get("markov_head_type", "vanilla") != "vanilla":
+            raise ValueError("Spark DSpark requires a vanilla Markov head")
+        if self.hparams.get("draft_vocab_size", self.hparams["vocab_size"]) != self.hparams["vocab_size"]:
+            raise ValueError("Spark DSpark currently requires the full target vocabulary")
+        if not self.hparams.get("confidence_head_with_markov", True):
+            raise ValueError("Spark DSpark requires confidence_head_with_markov")
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("draft_model."):
+            return None
+        return TextModel.filter_tensors(("model." + name.removeprefix("draft_model."), gen))
+
+    def set_vocab(self) -> None:
+        super().set_vocab()
+        self.gguf_writer.add_mask_token_id(self.hparams["mask_token_id"])
+
+    def set_gguf_parameters(self) -> None:
+        super().set_gguf_parameters()
+        self.gguf_writer.add_block_size(self.hparams["dspark_block_size"])
+        # Target layer outputs are extracted as the next layer's inputs.
+        self.gguf_writer.add_target_layers([i + 1 for i in self.hparams["target_layer_ids"]])
+        self.gguf_writer.add_sample_from_anchor(self.hparams.get("sample_from_anchor", True))
+        self.gguf_writer.add_causal_attention(False)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith("model.fc_norm."):
+            if name == "model.fc_norm.0.weight":
+                norms = [self.model_tensors[f"model.fc_norm.{i}.weight"]().float()
+                         for i in range(len(self.hparams["target_layer_ids"]))]
+                yield self.format_tensor_name(gguf.MODEL_TENSOR.FC_NORM), torch.stack(norms)
+            return
+
+        if name.endswith(".self_attn.q_k_v_proj.weight"):
+            head_dim = self.hparams["head_dim"]
+            n_q = head_dim * self.hparams["num_attention_heads"]
+            n_kv = head_dim * self.hparams["num_key_value_heads"]
+            for tensor, value in zip(
+                (gguf.MODEL_TENSOR.ATTN_Q, gguf.MODEL_TENSOR.ATTN_K, gguf.MODEL_TENSOR.ATTN_V),
+                data_torch.split((n_q, n_kv, n_kv), dim=0),
+                strict=True,
+            ):
+                yield self.format_tensor_name(tensor, bid), value
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)

convert_hf_to_gguf.py

@@ -127,7 +127,7 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--dspark", action="store_true",
-        help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.",
+        help="Export only the DeepSeek-V4 or Spark2.5 DSpark draft tensors as a separate GGUF.",
     )
     parser.add_argument(
         "--mistral-format", action="store_true",
@@ -267,15 +267,21 @@ def main() -> None:
             sys.exit(1)
 
         if args.dspark:
-            if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
-                logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM")
+            if is_mistral_format:
+                logger.error("--dspark is not supported for Mistral format models")
                 sys.exit(1)
             if model_architecture == "DeepseekV41ForCausalLM":
                 from conversion.deepseek41 import DeepseekV41DSparkModel
                 model_class = DeepseekV41DSparkModel
-            else:
+            elif model_architecture == "DeepseekV4ForCausalLM":
                 from conversion.deepseek import DeepseekV4DSparkModel
                 model_class = DeepseekV4DSparkModel
+            elif model_architecture == "Spark2_5ForCausalLM":
+                from conversion.spark2_5 import Spark2_5DSparkModel
+                model_class = Spark2_5DSparkModel
+            else:
+                logger.error("--dspark is only supported for DeepseekV4ForCausalLM, DeepseekV41ForCausalLM, and Spark2_5ForCausalLM")
+                sys.exit(1)
 
         if args.mtp or args.no_mtp:
             if not model_class.supports_mtp_export:

gguf-py/gguf/constants.py

@@ -1328,6 +1328,7 @@ class MODEL_TENSOR(IntEnum):
     NEXTN_HC_HEAD_UP       = auto()
     # eagle3
     FC                     = auto()  # feature fusion layer
+    FC_NORM                = auto()
     D2T                    = auto()  # draft to target vocabulary mapping
     # dspark
     DSPARK_MARKOV_W1       = auto()  # markov head: prev-token embed
@@ -2184,6 +2185,7 @@ class MODEL_TENSOR(IntEnum):
     MODEL_TENSOR.NEXTN_HC_HEAD_DOWN:        "blk.{bid}.nextn.hc_head_down",
     MODEL_TENSOR.NEXTN_HC_HEAD_UP:          "blk.{bid}.nextn.hc_head_up",
     MODEL_TENSOR.FC:                        "fc",
+    MODEL_TENSOR.FC_NORM:                   "fc_norm",
     MODEL_TENSOR.DSPARK_MARKOV_W1:          "markov_w1",
     MODEL_TENSOR.DSPARK_MARKOV_W2:          "markov_w2",
     MODEL_TENSOR.DSPARK_CONF_PROJ:          "conf_proj",
@@ -5769,6 +5771,7 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.ATTN_K,
         MODEL_TENSOR.ATTN_V,
         MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_GATE,
         MODEL_TENSOR.ATTN_Q_NORM,
         MODEL_TENSOR.ATTN_K_NORM,
         MODEL_TENSOR.ATTN_SINKS,
@@ -5801,6 +5804,7 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.FFN_DOWN_SHEXP,
         MODEL_TENSOR.FFN_UP_SHEXP,
         MODEL_TENSOR.FC,
+        MODEL_TENSOR.FC_NORM,
         MODEL_TENSOR.ENC_OUTPUT_NORM,
         MODEL_TENSOR.D2T,
         # optional DSpark heads

src/llama-arch.cpp

@@ -792,6 +792,7 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
     { LLM_TENSOR_MASKED_EMBD_CENTROIDS,                  "masked_embd_centroids" },
     { LLM_TENSOR_MASKED_EMBD_ORDERING,                   "masked_embd_ordering" },
     { LLM_TENSOR_FC,                                     "fc" },
+    { LLM_TENSOR_FC_NORM,                                "fc_norm" },
     { LLM_TENSOR_D2T,                                    "d2t" },
     { LLM_TENSOR_DSPARK_MARKOV_W1,                       "markov_w1" },
     { LLM_TENSOR_DSPARK_MARKOV_W2,                       "markov_w2" },
@@ -1135,6 +1136,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_MASKED_EMBD_ORDERING,       {LLM_TENSOR_LAYER_INPUT,     GGML_OP_NONE}},
     // eagle3
     {LLM_TENSOR_FC,                         {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_FC_NORM,                    {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL}},
     {LLM_TENSOR_D2T,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_SET_ROWS}},
     // dspark
     {LLM_TENSOR_DSPARK_MARKOV_W1,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},

src/llama-arch.h

@@ -798,6 +798,7 @@ enum llm_tensor {
     LLM_TENSOR_MASKED_EMBD_ORDERING,
     LLM_TENSOR_HRM_Z_L_INIT,
     LLM_TENSOR_FC,
+    LLM_TENSOR_FC_NORM,
     LLM_TENSOR_D2T,
     LLM_TENSOR_DSPARK_MARKOV_W1,
     LLM_TENSOR_DSPARK_MARKOV_W2,

src/llama-model.h

@@ -790,6 +790,7 @@ struct llama_model {
     struct ggml_tensor * fc      = nullptr;
     struct ggml_tensor * fc_s    = nullptr;
     struct ggml_tensor * fc_in_s = nullptr;
+    struct ggml_tensor * fc_norm = nullptr;
     struct ggml_tensor * d2t = nullptr;  // draft to target vocabulary mapping
 
     // dspark

src/models/dflash.cpp

@@ -113,12 +113,12 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
     }
 
     // optional interleaved sliding-window attention with per-layer pattern array.
-    // DFlash has a single rope, so the SWA rope == main rope.
     if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
         hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
         ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
         hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;
         hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+        ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
     }
 
     type = LLM_TYPE_UNKNOWN;
@@ -245,6 +245,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
     fc              = create_tensor(tn(LLM_TENSOR_FC,              "weight"), { n_embd_inp, n_embd }, 0);
     fc_s            = create_tensor(tn(LLM_TENSOR_FC,              "scale"),  { 1 }, TENSOR_NOT_REQUIRED);
     fc_in_s         = create_tensor(tn(LLM_TENSOR_FC,              "input_scale"), { 1 }, TENSOR_NOT_REQUIRED);
+    fc_norm         = create_tensor(tn(LLM_TENSOR_FC_NORM,         "weight"), { n_embd, n_embd_inp / n_embd }, TENSOR_NOT_REQUIRED);
     output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
     output_norm     = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM,    "weight"), { n_embd }, 0); // decoder final norm
 
@@ -305,6 +306,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
         return;
     }
 
+    const bool has_attn_gate = ml->get_tensor_meta("blk.0.attn_gate.weight") != nullptr;
     for (int i = 0; i < n_layer; ++i) {
         auto & layer = layers[i];
 
@@ -315,8 +317,12 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
         layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V,   "weight", i), { n_embd, n_embd_v_gqa }, TENSOR_NOT_REQUIRED);
         layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);
 
-        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
-        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
+        if (has_attn_gate) {
+            layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, n_head }, 0);
+        } else {
+            layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
+            layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
+        }
 
         layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED);
         layer.ffn_post_norm  = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM,  "weight", i), { n_embd }, TENSOR_NOT_REQUIRED);
@@ -364,6 +370,11 @@ template <>
 llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
     ggml_tensor * cur = build_inp_embd_enc();
 
+    if (model.fc_norm) {
+        cur = ggml_reshape_3d(ctx0, cur, n_embd, model.target_layer_ids.size(), n_tokens);
+        cur = build_norm(cur, model.fc_norm, nullptr, LLM_NORM_RMS, -1);
+        cur = ggml_reshape_2d(ctx0, cur, model.fc->ne[0], n_tokens);
+    }
     cur = build_lora_mm(model.fc, cur, model.fc_s);
     cb(cur, "fc_out", -1);
 
@@ -839,13 +850,16 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
     int sections[4];
     std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);
 
-    auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) {
+    auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos, int il) {
+        const int64_t n_rot_i = hparams.n_rot(il);
+        const float freq_base_i = model.get_rope_freq_base(cparams, il);
+        const float freq_scale_i = model.get_rope_freq_scale(cparams, il);
         return rope_type == GGML_ROPE_TYPE_MROPE
             ? ggml_rope_multi(ctx0, cur, pos, nullptr,
-                    n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
+                    n_rot_i, sections, rope_type, n_ctx_orig, freq_base_i, freq_scale_i,
                     ext_factor, attn_factor, beta_fast, beta_slow)
-            : ggml_rope_ext(ctx0, cur, pos, model.layers[0].rope_freqs,
-                    n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+            : ggml_rope_ext(ctx0, cur, pos, model.layers[il].rope_freqs,
+                    n_rot_i, rope_type, n_ctx_orig, freq_base_i, freq_scale_i,
                     ext_factor, attn_factor, beta_fast, beta_slow);
     };
 
@@ -863,6 +877,11 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
         res->add_input(std::move(inp));
 
         // fuse the target features through the encoder
+        if (model.fc_norm) {
+            inp_target = ggml_reshape_3d(ctx0, inp_target, n_embd, model.target_layer_ids.size(), n_tokens);
+            inp_target = build_norm(inp_target, model.fc_norm, nullptr, LLM_NORM_RMS, -1);
+            inp_target = ggml_reshape_2d(ctx0, inp_target, n_embd_inp, n_tokens);
+        }
         ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
         inp_g = build_norm(inp_g, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
         cb(inp_g, "inp_g_embeddings", -1);
@@ -877,11 +896,13 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
             Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
             Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
 
-            Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
+            if (layer.attn_k_norm) {
+                Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
+            }
             if (shared_kv) {
                 Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps);
             }
-            Kcur = build_rope(Kcur, inp_pos);
+            Kcur = build_rope(Kcur, inp_pos, il);
             cb(Kcur, "Kcur_injected", il);
             cb(Vcur, "Vcur_injected", il);
 
@@ -982,22 +1003,34 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
         Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
         Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
 
-        Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il);
-        Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
+        if (layer.attn_q_norm) {
+            Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il);
+            Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
+        }
         if (shared_kv) {
             Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps);
         }
 
-        Qcur = build_rope(Qcur, inp_pos);
-        Kcur = build_rope(Kcur, inp_pos);
+        Qcur = build_rope(Qcur, inp_pos, il);
+        Kcur = build_rope(Kcur, inp_pos, il);
         cb(Qcur, "Qcur", il);
         cb(Kcur, "Kcur", il);
         cb(Vcur, "Vcur", il);
 
         // cache-aware, non-causal attention
+        ggml_tensor * wo = layer.wqkv_gate ? nullptr : layer.wo;
         ggml_tensor * cur = use_iswa
-            ? build_attn(inp_attn_iswa, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il)
-            : build_attn(inp_attn,      layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il);
+            ? build_attn(inp_attn_iswa, wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il)
+            : build_attn(inp_attn,      wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il);
+
+        if (layer.wqkv_gate) {
+            ggml_tensor * gate = ggml_sigmoid(ctx0, build_lora_mm(layer.wqkv_gate, noise_norm));
+            cur = ggml_reshape_3d(ctx0, cur, n_embd_head, n_head, n_tokens);
+            gate = ggml_reshape_3d(ctx0, gate, 1, n_head, n_tokens);
+            cur = ggml_mul(ctx0, cur, gate);
+            cur = ggml_reshape_2d(ctx0, cur, n_embd_head * n_head, n_tokens);
+            cur = build_lora_mm(layer.wo, cur, layer.wo_s);
+        }
 
         if (hparams.f_attn_value_scale != 0.0f) {
             cur = ggml_scale(ctx0, cur, hparams.f_attn_value_scale);
@@ -1032,7 +1065,7 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
                 layer.ffn_gate, NULL, layer.ffn_gate_s,
                 layer.ffn_down, NULL, layer.ffn_down_s,
                 NULL,
-                hparams.llm_ffn_op, LLM_FFN_PAR, il);
+                layer.wqkv_gate ? LLM_FFN_GELU : hparams.llm_ffn_op, LLM_FFN_PAR, il);
         cb(cur, "ffn_out", il);
 
         if (ffn_dynamic) {

src/models/spark2-5.cpp

@@ -63,10 +63,13 @@ llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_pa
     ggml_tensor * inp_pos = build_inp_pos();
     auto * inp_attn = build_attn_inp_kv_iswa();
     ggml_tensor * inp_out_ids = build_inp_out_ids();
+    const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
 
     const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
         ggml_tensor * inpSA = inpL;
         ggml_tensor * cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
         cb(cur, "attn_norm", il);
@@ -108,7 +111,7 @@ llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_pa
         cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
         cb(cur, "attn_out_proj", il);
 
-        if (il == n_layer - 1 && inp_out_ids) {
+        if (il == n_layer - 1 && inp_out_ids && !extract_final_inp) {
             cur = ggml_get_rows(ctx0, cur, inp_out_ids);
             inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
         }
@@ -134,6 +137,10 @@ llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_pa
         inpL = cur;
     }
 
+    res->t_layer_inp[n_layer] = inpL;
+    if (extract_final_inp && inp_out_ids) {
+        inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+    }
     ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
     cb(cur, "result_norm", -1);
     res->t_embd = cur;