model : support DSpark for LFM2 models (#27383)

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",

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):

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;

src/models/lfm2.cpp

@@ -2,6 +2,8 @@
 #include "../llama-memory-hybrid-iswa.h"
 #include "../llama-memory-hybrid.h"
 
+#include <algorithm>
+
 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<iswa>::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<int64_t>(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<iswa>::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<int>(hparams.n_layer_dense_lead);
 
         auto * prev_cur = cur;