Skip to content

spec : support DSpark for LFM2 models #27383

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ruixiang63 merged 1 commit into from
Aug 20, 2026
+35 −10
Merged

spec : support DSpark for LFM2 models #27383

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions conversion/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
"Qwen3DSparkModel": "qwen",
"DSparkDraftModel": "qwen",
"DSparkSpeculator": "qwen",
"Lfm2DSparkDraftModel": "qwen",
"DeepseekV4ForCausalLM": "deepseek",
"DeepseekV4DSparkModel": "deepseek",
"DistilBertForMaskedLM": "bert",
Expand Down
15 changes: 14 additions & 1 deletion conversion/qwen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -709,7 +709,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
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.
Expand Down Expand Up @@ -759,6 +759,13 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca
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
Expand All @@ -767,6 +774,12 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
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):
Expand Down
2 changes: 2 additions & 0 deletions src/llama-arch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1029,6 +1029,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;
Expand Down
27 changes: 18 additions & 9 deletions src/models/lfm2.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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 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))));
// 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);

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);
Expand Down Expand Up @@ -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;
Expand Down
Loading