Commit 84e66d032 for llama.cpp
commit 84e66d03218711dfcaa993e5f64a546abe82df9a
Author: Tarek Dakhran <tarek@liquid.ai>
Date: Tue Oct 20 23:17:13 2026 +0200
refactor : move LFM2 shortconv and FFN blocks into a shared LFM base (#30874)
* llama : move LFM2 shortconv and FFN blocks into a shared LFM base
* Update src/models/lfm-base.cpp
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/src/models/lfm-base.cpp b/src/models/lfm-base.cpp
new file mode 100644
index 000000000..5e722eca8
--- /dev/null
+++ b/src/models/lfm-base.cpp
@@ -0,0 +1,114 @@
+#include "models.h"
+
+#include "llama-memory-recurrent.h"
+
+#include <algorithm>
+
+llm_build_lfm_base::llm_build_lfm_base(const llm_graph_params & params) : llm_graph_context(params) {}
+
+ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
+ const llama_memory_recurrent_context * mctx_recr,
+ llm_graph_input_rs * inp_recr,
+ ggml_tensor * cur,
+ int il) {
+ const uint32_t kv_head = mctx_recr->get_head();
+ const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+ const int64_t n_seqs = ubatch.n_seqs;
+ GGML_ASSERT(n_seqs != 0);
+ GGML_ASSERT(ubatch.equal_seqs());
+ GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
+
+ GGML_ASSERT(hparams.n_shortconv_l_cache > 1);
+ const uint32_t d_conv = hparams.n_shortconv_l_cache - 1;
+ const int64_t n_embd = hparams.n_embd;
+
+ // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
+ cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
+
+ auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur);
+ cb(bcx, "model.layers.{}.conv.in_proj", il);
+
+ constexpr auto n_chunks = 3;
+ GGML_ASSERT(bcx->ne[0] % n_chunks == 0);
+ const auto chunk_size = bcx->ne[0] / n_chunks;
+ auto * b = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
+ 0 * chunk_size * ggml_element_size(bcx));
+ auto * c = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
+ 1 * chunk_size * ggml_element_size(bcx));
+ auto * x = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
+ 2 * chunk_size * ggml_element_size(bcx));
+
+ auto * bx = ggml_transpose(ctx0, ggml_mul(ctx0, b, x));
+
+ // read conv state
+ auto * conv_state = mctx_recr->get_r_l(il);
+ auto * conv_rs = build_rs(inp_recr, conv_state, hparams.n_embd_r(), n_seqs);
+ auto * conv = ggml_reshape_3d(ctx0, conv_rs, d_conv, n_embd, n_seqs);
+
+ // causal prepends the state, non-causal pads symmetrically for a centered window
+ if (hparams.causal_attn) {
+ bx = ggml_concat(ctx0, conv, bx, 0);
+ } else {
+ const int64_t pad = (hparams.n_shortconv_l_cache - 1) / 2;
+ auto * left = ggml_cont(ctx0,
+ ggml_view_3d(ctx0, conv, pad, n_embd, n_seqs, conv->nb[1], conv->nb[2], (d_conv - pad) * conv->nb[0]));
+ bx = ggml_pad_ext(ctx0, ggml_concat(ctx0, left, bx, 0), 0, pad, 0, 0, 0, 0, 0, 0);
+ }
+ GGML_ASSERT(bx->ne[0] > conv->ne[0]);
+
+ // 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_recr->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);
+ cb(conv_out, "model.layers.{}.conv.conv", il);
+
+ auto * y = ggml_mul(ctx0, c, conv_out);
+ y = build_lora_mm(model.layers[il].shortconv.out_proj, y);
+ cb(y, "model.layers.{}.conv.out_proj", il);
+ // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
+ y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs);
+
+ return y;
+}
+
+ggml_tensor * llm_build_lfm_base::build_lfm_dense_ffn(const llama_model & model,
+ ggml_tensor * cur,
+ int il) {
+ GGML_ASSERT(!model.layers[il].ffn_up_b);
+ GGML_ASSERT(!model.layers[il].ffn_gate_b);
+ GGML_ASSERT(!model.layers[il].ffn_down_b);
+ return build_ffn(cur,
+ model.layers[il].ffn_up, NULL, NULL,
+ model.layers[il].ffn_gate, NULL, NULL,
+ model.layers[il].ffn_down, NULL, NULL,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+}
+
+ggml_tensor * llm_build_lfm_base::build_lfm_moe_ffn(const llama_model & model,
+ ggml_tensor * cur,
+ int il) {
+ return build_moe_ffn(cur,
+ model.layers[il].ffn_gate_inp,
+ model.layers[il].ffn_up_exps,
+ model.layers[il].ffn_gate_exps,
+ model.layers[il].ffn_down_exps,
+ model.layers[il].ffn_exp_probs_b,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, true,
+ hparams.expert_weights_scale,
+ static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),
+ il);
+}
diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp
index 942d8de90..dd61f509a 100644
--- a/src/models/lfm2.cpp
+++ b/src/models/lfm2.cpp
@@ -147,35 +147,11 @@ std::unique_ptr<llm_graph_context> llama_model_lfm2::build_arch_graph(const llm_
template <bool iswa>
llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_params & params) :
- llm_graph_context(params) {
+ llm_build_lfm_base(params) {
using inp_hybrid_type = std::conditional_t<iswa, llm_graph_input_mem_hybrid_iswa, llm_graph_input_mem_hybrid>;
using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;
using mem_hybrid_ctx = std::conditional_t<iswa, llama_memory_hybrid_iswa_context, llama_memory_hybrid_context>;
- // lambda helpers for readability
- auto build_dense_feed_forward = [&model, this](ggml_tensor * cur, int il) -> ggml_tensor * {
- GGML_ASSERT(!model.layers[il].ffn_up_b);
- GGML_ASSERT(!model.layers[il].ffn_gate_b);
- GGML_ASSERT(!model.layers[il].ffn_down_b);
- return build_ffn(cur,
- model.layers[il].ffn_up, NULL, NULL,
- model.layers[il].ffn_gate, NULL, NULL,
- model.layers[il].ffn_down, NULL, NULL,
- NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
- };
- auto build_moe_feed_forward = [&model, this](ggml_tensor * cur, int il) -> ggml_tensor * {
- return build_moe_ffn(cur,
- model.layers[il].ffn_gate_inp,
- model.layers[il].ffn_up_exps,
- model.layers[il].ffn_gate_exps,
- model.layers[il].ffn_down_exps,
- model.layers[il].ffn_exp_probs_b,
- n_expert, n_expert_used,
- LLM_FFN_SILU, true,
- hparams.expert_weights_scale,
- static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),
- il);
- };
auto build_attn_block = [&model, this](ggml_tensor * cur,
ggml_tensor * inp_pos,
inp_attn_type * inp_attn,
@@ -207,82 +183,6 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
return cur;
};
- auto build_shortconv_block = [&model, this](ggml_tensor * cur,
- llm_graph_input_rs * inp_recr,
- int il) -> ggml_tensor * {
- const auto * mctx_cur = static_cast<const mem_hybrid_ctx *>(mctx)->get_recr();
- const uint32_t kv_head = mctx_cur->get_head();
- const int64_t n_seq_tokens = ubatch.n_seq_tokens;
- const int64_t n_seqs = ubatch.n_seqs;
- GGML_ASSERT(n_seqs != 0);
- GGML_ASSERT(ubatch.equal_seqs());
- GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
-
- GGML_ASSERT(hparams.n_shortconv_l_cache > 1);
- const uint32_t d_conv = hparams.n_shortconv_l_cache - 1;
-
- // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}
- cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
-
- auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur);
- cb(bcx, "model.layers.{}.conv.in_proj", il);
-
- constexpr auto n_chunks = 3;
- GGML_ASSERT(bcx->ne[0] % n_chunks == 0);
- const auto chunk_size = bcx->ne[0] / n_chunks;
- auto * b = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
- 0 * chunk_size * ggml_element_size(bcx));
- auto * c = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
- 1 * chunk_size * ggml_element_size(bcx));
- auto * x = ggml_view_3d(ctx0, bcx, chunk_size, bcx->ne[1], bcx->ne[2], bcx->nb[1], bcx->nb[2],
- 2 * chunk_size * ggml_element_size(bcx));
-
- auto * bx = ggml_transpose(ctx0, ggml_mul(ctx0, b, x));
-
- // read conv state
- auto * conv_state = mctx_cur->get_r_l(il);
- auto * conv_rs = build_rs(inp_recr, conv_state, hparams.n_embd_r(), n_seqs);
- auto * conv = ggml_reshape_3d(ctx0, conv_rs, d_conv, hparams.n_embd, n_seqs);
-
- // causal prepends the state, non-causal pads symmetrically for a centered window
- if (hparams.causal_attn) {
- bx = ggml_concat(ctx0, conv, bx, 0);
- } else {
- const int64_t pad = (hparams.n_shortconv_l_cache - 1) / 2;
- auto * left = ggml_cont(ctx0,
- ggml_view_3d(ctx0, conv, pad, hparams.n_embd, n_seqs, conv->nb[1], conv->nb[2], (d_conv - pad) * conv->nb[0]));
- bx = ggml_pad_ext(ctx0, ggml_concat(ctx0, left, bx, 0), 0, pad, 0, 0, 0, 0, 0, 0);
- }
- GGML_ASSERT(bx->ne[0] > conv->ne[0]);
-
- // 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);
- cb(conv_out, "model.layers.{}.conv.conv", il);
-
- auto * y = ggml_mul(ctx0, c, conv_out);
- y = build_lora_mm(model.layers[il].shortconv.out_proj, y);
- cb(y, "model.layers.{}.conv.out_proj", il);
- // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}
- y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs);
-
- return y;
- };
-
// actual graph construction starts here
ggml_tensor * cur = build_inp_embd(model.tok_embd);
cb(cur, "model.embed_tokens", -1);
@@ -308,8 +208,12 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
cur = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "model.layers.{}.operator_norm", il);
- cur = hparams.is_recr(il) ? build_shortconv_block(cur, inp_hybrid->get_recr(), il) :
- build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
+ if (hparams.is_recr(il)) {
+ const auto * mctx_recr = static_cast<const mem_hybrid_ctx *>(mctx)->get_recr();
+ cur = build_lfm_shortconv(model, mctx_recr, inp_hybrid->get_recr(), cur, il);
+ } else {
+ cur = build_attn_block(cur, inp_pos, inp_hybrid->get_attn(), il);
+ }
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
@@ -322,7 +226,7 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
cb(ffn_norm_out, "model.layers.{}.ffn_norm", il);
ggml_tensor * ffn_out =
- is_moe_layer ? build_moe_feed_forward(ffn_norm_out, il) : build_dense_feed_forward(ffn_norm_out, il);
+ is_moe_layer ? build_lfm_moe_ffn(model, ffn_norm_out, il) : build_lfm_dense_ffn(model, ffn_norm_out, il);
cb(ffn_norm_out, "model.layers.{}.ffn_out", il);
cur = ggml_add(ctx0, cur, ffn_out);
diff --git a/src/models/models.h b/src/models/models.h
index e6cf9d213..9aa450dea 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -100,6 +100,26 @@ struct llm_build_delta_net_base : public llm_graph_context {
int il);
};
+struct llm_build_lfm_base : public llm_graph_context {
+ llm_build_lfm_base(const llm_graph_params & params);
+
+ virtual ~llm_build_lfm_base() = default;
+
+ ggml_tensor * build_lfm_shortconv(const llama_model & model,
+ const llama_memory_recurrent_context * mctx_recr,
+ llm_graph_input_rs * inp_recr,
+ ggml_tensor * cur,
+ int il);
+
+ ggml_tensor * build_lfm_dense_ffn(const llama_model & model,
+ ggml_tensor * cur,
+ int il);
+
+ ggml_tensor * build_lfm_moe_ffn(const llama_model & model,
+ ggml_tensor * cur,
+ int il);
+};
+
struct llm_build_rwkv6_base : public llm_graph_context {
const llama_model & model;
@@ -2264,7 +2284,7 @@ struct llama_model_lfm2 : public llama_model_base {
void load_arch_tensors(llama_model_loader & ml) override;
template <bool iswa>
- struct graph : public llm_graph_context {
+ struct graph : public llm_build_lfm_base {
graph(const llama_model & model, const llm_graph_params & params);
};