Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model: add DSpark support for Nemotron3.5 (#27804)

Browse Source 
* model: add DSpark support for Nemotron3.5

* Update src/models/dflash.cpp

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

---------

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
This commit is contained in:
Ruixiang Wang
2026-08-28 01:49:27 +02:00
committed by GitHub
co-authored by Sigbjørn Skjæret Xuan Son Nguyen
parent e70802a01f
commit ca3d5a3e10
6 changed files with 71 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-model.h
+1
View File
@@ -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;
src/models/dflash.cpp
+34 -24
View File
@@ -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);