From 2115b73d8ebdbd659075cce66c609506863bc826 Mon Sep 17 00:00:00 2001 From: Tiwei Bie Date: Sat, 22 Aug 2026 17:19:48 +0800 Subject: [PATCH] model : support DSpark for bailingmoe3 (#27508) --- conversion/__init__.py | 1 + conversion/qwen.py | 8 +++++++- src/llama-arch.cpp | 1 + src/models/bailingmoe3.cpp | 41 +++++++++++++++++++++++--------------- 4 files changed, 34 insertions(+), 17 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index aaaebba97..8de97e959 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -58,6 +58,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", "Lfm2DSparkDraftModel": "qwen", + "LingDSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DeepseekV4DSparkModel": "deepseek", "DistilBertForMaskedLM": "bert", diff --git a/conversion/qwen.py b/conversion/qwen.py index 355365763..cdba8a63e 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -709,7 +709,13 @@ class DFlashModel(Qwen3Model): yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator", "Lfm2DSparkDraftModel") +@ModelBase.register( + "Qwen3DSparkModel", + "DSparkDraftModel", + "DSparkSpeculator", + "Lfm2DSparkDraftModel", + "LingDSparkModel", +) @ModelBase.example("satgeze/Qwen3.6-27B-DSpark") class DSparkModel(DFlashModel): # DSpark = DFlash + a semi-autoregressive Markov head. diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 60b8e461c..025f9fb54 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1038,6 +1038,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: + case LLM_ARCH_BAILINGMOE3: return true; default: return false; diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp index f5855696e..0637931cc 100644 --- a/src/models/bailingmoe3.cpp +++ b/src/models/bailingmoe3.cpp @@ -1,6 +1,8 @@ #include "models.h" #include "llama-memory-recurrent.h" +#include + void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); @@ -179,7 +181,9 @@ static ggml_tensor * bailingmoe3_causal_conv1d( int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, - int64_t cache_head) { + int64_t cache_head, + uint32_t mem_size, + uint32_t n_rs_seq) { const int64_t d_inner = head_dim * n_head; const int64_t conv_state_size = (d_conv - 1) * d_inner; const int64_t total_state_size = 3 * conv_state_size; @@ -193,13 +197,18 @@ static ggml_tensor * bailingmoe3_causal_conv1d( x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); - ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, - conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x, - ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, - (d_conv - 1) * ggml_element_size(conv_states_all), - total_state_size * ggml_element_size(conv_states_all), - (cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + const int64_t K = (int64_t) n_rs_seq + 1; + const int64_t n_written = std::min(n_seq_tokens, K); + + for (int64_t slot = 0; slot < n_written; ++slot) { + ggml_tensor * conv_snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + total_state_size * ggml_element_size(conv_states_all), + ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + } ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight); @@ -237,6 +246,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + const auto & layer = model.layers[il]; ggml_tensor * inpSA = inpL; ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il); @@ -245,18 +256,19 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph if (hparams.is_recr(il)) { const auto * mctx_cur = inp_rs->mctx; const auto cache_head = mctx_cur->get_head(); + const auto mem_size = mctx_cur->get_size(); ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); ggml_tensor * q = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); ggml_tensor * k = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); ggml_tensor * v = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); gate = ggml_add(ctx0, gate, layer.ssm_dt_b); @@ -276,11 +288,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs); state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); - auto result = build_delta_net(q, k, v, gate, beta, state, il); - ggml_tensor * out = ggml_cont(ctx0, result.first); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second, - ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs, - cache_head * hparams.n_embd_s() * ggml_element_size(states_all)))); + ggml_tensor * out = ggml_cont(ctx0, build_recurrent_attn( + inp_rs, states_all, q, k, v, gate, beta, state, il)); ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens);