From 07822bddf80d73f1168e592c52e69caaff820f9c Mon Sep 17 00:00:00 2001 From: Tarek Dakhran Date: Thu, 20 Aug 2026 16:36:57 +0200 Subject: [PATCH] model : support DSpark for LFM2 models (#27383) --- conversion/__init__.py | 1 + conversion/qwen.py | 15 ++++++++++++++- src/llama-arch.cpp | 2 ++ src/models/lfm2.cpp | 25 +++++++++++++++++-------- 4 files changed, 34 insertions(+), 9 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index 5ae6ad819..4b8817ead 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -57,6 +57,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", + "Lfm2DSparkDraftModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DeepseekV4DSparkModel": "deepseek", "DistilBertForMaskedLM": "bert", diff --git a/conversion/qwen.py b/conversion/qwen.py index 26b10452b..355365763 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -709,7 +709,7 @@ class DFlashModel(Qwen3Model): yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator") +@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator", "Lfm2DSparkDraftModel") @ModelBase.example("satgeze/Qwen3.6-27B-DSpark") class DSparkModel(DFlashModel): # DSpark = DFlash + a semi-autoregressive Markov head. @@ -759,6 +759,13 @@ class DSparkModel(DFlashModel): return None return super().filter_tensors(item) + _ROPE_PERMUTE_SUFFIXES = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name == "model.d2t": self._d2t = data_torch @@ -767,6 +774,12 @@ class DSparkModel(DFlashModel): if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")): return + # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd + if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES): + head_dim = self.hparams["head_dim"] + shape = data_torch.shape + data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape) + yield from super().modify_tensors(data_torch, name, bid) def prepare_tensors(self): diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 955c2d796..408954401 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1032,6 +1032,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: + case LLM_ARCH_LFM2: + case LLM_ARCH_LFM2MOE: return true; default: return false; diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp index 70e837d6e..9a4295557 100644 --- a/src/models/lfm2.cpp +++ b/src/models/lfm2.cpp @@ -2,6 +2,8 @@ #include "../llama-memory-hybrid-iswa.h" #include "../llama-memory-hybrid.h" +#include + void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SHORTCONV_L_CACHE, hparams.n_shortconv_l_cache); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -202,15 +204,20 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ } GGML_ASSERT(bx->ne[0] > conv->ne[0]); - // last d_conv columns is a new conv state - auto * new_conv = ggml_view_3d(ctx0, bx, conv->ne[0], bx->ne[1], bx->ne[2], bx->nb[1], bx->nb[2], - (bx->ne[0] - conv->ne[0]) * ggml_element_size(bx)); - GGML_ASSERT(ggml_are_same_shape(conv, new_conv)); + // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback) + const int64_t K = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t n_written = std::min(n_seq_tokens, K); + const auto mem_size = mctx_cur->get_size(); + const size_t row_size = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd); - // write new conv conv state - ggml_build_forward_expand(gf, ggml_cpy(ctx0, new_conv, - ggml_view_1d(ctx0, conv_state, ggml_nelements(new_conv), - kv_head * d_conv * n_embd * ggml_element_size(new_conv)))); + for (int64_t slot = 0; slot < n_written; ++slot) { + auto * conv_snap = ggml_view_3d(ctx0, bx, d_conv, bx->ne[1], bx->ne[2], bx->nb[1], bx->nb[2], + (bx->ne[0] - d_conv - slot) * ggml_element_size(bx)); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap, + ggml_view_2d(ctx0, conv_state, (int64_t) d_conv * n_embd, n_seqs, + conv_state->nb[1], + ((size_t) slot * mem_size + kv_head) * row_size))); + } auto * conv_kernel = model.layers[il].shortconv.conv; auto * conv_out = ggml_ssm_conv(ctx0, bx, conv_kernel); @@ -242,6 +249,8 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * inp_out_ids = build_inp_out_ids(); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = cur; + const bool is_moe_layer = il >= static_cast(hparams.n_layer_dense_lead); auto * prev_cur = cur;