model: add DSpark support for Nemotron3.5

@@ -935,6 +935,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {     // dspark speculators     bool sample_from_anchor = true; +    // block-internal attention+    bool causal_attn = false;+     const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices     uint32_t        target_layer_ids_n = 0; @@ -972,12 +975,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {             if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {                 sample_from_anchor = std::strcmp(buf, "true") == 0;             }+            if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) {+                causal_attn = std::strcmp(buf, "true") == 0;+            }         }          selector_top_k = llama_model_dflash_selector_top_k(model_dft);         is_dflash2     = selector_top_k > 0;         mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); +        if (is_dspark && this->params.p_min > 0.0f) {+            char buf[16] = {};+            const bool has_conf =+                llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 ||+                std::strcmp(buf, "true") == 0;+            if (!has_conf) {+                throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0");+            }+        }+         LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());         LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);         LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__,@@ -1036,7 +1052,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {          // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.         llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);-        llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention+        llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise     }      ~common_speculative_impl_draft_dflash() override {

@@ -709,14 +709,20 @@ def set_gguf_parameters(self):             extract_layer_ids = [i + 1 for i in target_layer_ids]             self.gguf_writer.add_target_layers(extract_layer_ids) -        use_sliding_window = self.hparams.get("use_sliding_window", False)-        sliding_window = self.hparams.get("sliding_window")+        use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False)+        sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window")         layer_types = self.hparams.get("layer_types")         if use_sliding_window and sliding_window and layer_types:             is_swa = [lt == "sliding_attention" for lt in layer_types]             self.gguf_writer.add_sliding_window(sliding_window)             self.gguf_writer.add_sliding_window_pattern(is_swa) +        causal = self.hparams.get("is_causal")+        if causal is None:+            causal = dflash_config.get("causal")+        if causal is not None:+            self.gguf_writer.add_causal_attention(bool(causal))+         # M-RoPE target: the draft ropes on the temporal dim only, so write         # degenerate sections [n_rot/2, 0, 0, 0]         if self._target_uses_mrope():@@ -737,6 +743,8 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca         name, gen = item         if not name.startswith("model."):             name = "model." + name+        if "sink" in name and not name.endswith(".weight"):+            name += ".weight"         return super().filter_tensors((name, gen))      _ROPE_PERMUTE_SUFFIXES = (@@ -815,6 +823,10 @@ def set_gguf_parameters(self):         super().set_gguf_parameters()         self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) +        # confidence head is optional: vanilla-markov exports ship without it+        has_conf = any("confidence_head.proj" in name for name in self.model_tensors)+        self.gguf_writer.add_has_confidence_head(has_conf)+     @classmethod     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:         if item[0] == "t2d":  # not used at runtime@@ -833,7 +845,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter             self._d2t = data_torch             return -        if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")):+        if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"):             return          # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd

@@ -167,6 +167,7 @@ class LLM:         SELECTOR_RANK                     = "{arch}.selector_rank"         SELECTOR_TOP_K                    = "{arch}.selector_top_k"         SAMPLE_FROM_ANCHOR                = "{arch}.sample_from_anchor"+        HAS_CONFIDENCE_HEAD               = "{arch}.has_confidence_head"         NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"         NORM_BEFORE_FC                    = "{arch}.norm_before_fc" 

@@ -1008,6 +1008,9 @@ def add_selector_top_k(self, value: int) -> None:     def add_sample_from_anchor(self, value: bool) -> None:         self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) +    def add_has_confidence_head(self, value: bool) -> None:+        self.add_bool(Keys.LLM.HAS_CONFIDENCE_HEAD.format(arch=self.arch), value)+     def add_target_layers(self, value: Sequence[int]) -> None:         self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) 

@@ -672,6 +672,7 @@ struct llama_model {     // dspark     struct ggml_tensor * dspark_markov_w1   = nullptr;     struct ggml_tensor * dspark_markov_w2   = nullptr;+    struct ggml_tensor * dspark_markov_w2_s = nullptr;     struct ggml_tensor * dspark_conf_proj   = nullptr;     struct ggml_tensor * dspark_conf_proj_b = nullptr; 

@@ -115,10 +115,11 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {     if (markov_meta) {         const int64_t dspark_markov_rank = markov_meta->ne[0]; -        dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);-        dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);+        dspark_markov_w1   = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);+        dspark_markov_w2   = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);+        dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"),  { 1 }, TENSOR_NOT_REQUIRED); -        dspark_conf_proj   = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0);+        dspark_conf_proj   = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED);         dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"),   { 1 },             TENSOR_NOT_REQUIRED);          LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);@@ -219,6 +220,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {         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); +        // optional per-head attention sinks (e.g. Nemotron DSpark)+        layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), { n_head }, TENSOR_NOT_REQUIRED);+         layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0);         layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);         layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);@@ -290,7 +294,10 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &      ggml_tensor * w1 = model.dspark_markov_w1;     ggml_tensor * w2 = model.dspark_markov_w2;-    GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded");+    GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded");++    // confidence head is optional+    const bool has_conf = model.dspark_conf_proj != nullptr;      ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]     const int64_t n_vocab = base->ne[0];@@ -321,23 +328,22 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &     ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);     prev = ggml_cont_1d(ctx0, prev, n_blocks); -    // confidence head input: predicts per-position acceptance-    ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]-     ggml_tensor * cat      = nullptr;     ggml_tensor * cat_conf = nullptr;      if (!sample_from_anchor) {         // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column-        cat      = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));-        cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));+        cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));+        if (has_conf) {+            cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));+        }     }      // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final     //       token pick, not the Markov conditioning path     for (int64_t i = i_draft_beg; i < block_drafts; ++i) {-        ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev);   // [R, n_blocks]-        ggml_tensor * bias    = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks]+        ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev);                          // [R, n_blocks]+        ggml_tensor * bias    = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]         if (model.d2t) {             // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)             const int64_t n_draft_vocab = bias->ne[0];@@ -354,17 +360,21 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &          cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; -        // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b)  -- [1, n_blocks]-        ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,-                                                (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);-        ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);-        ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);-        if (model.dspark_conf_proj_b) {-            conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);-        }-        conf = ggml_sigmoid(ctx0, conf);+        if (has_conf) {+            // confidence head input: predicts per-position acceptance+            ggml_tensor * conf_inp   = res->t_embd; // [n_embd, n_tok]+            // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b)  -- [1, n_blocks]+            ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,+                                                    (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);+            ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);+            ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);+            if (model.dspark_conf_proj_b) {+                conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);+            }+            conf = ggml_sigmoid(ctx0, conf); -        cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;+            cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;+        }          if (i + 1 < block_drafts) {             prev = ggml_argmax(ctx0, col);@@ -376,7 +386,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &     out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]     out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); -    {+    if (has_conf) {         ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);         conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));         conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);@@ -707,8 +717,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra          // cache-aware, non-causal attention         ggml_tensor * cur = use_iswa-            ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)-            : build_attn(inp_attn,      layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);+            ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il)+            : build_attn(inp_attn,      layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il);          if (attn_dynamic) {             cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);