Skip to content

refactor : move LFM2 shortconv and FFN blocks into a shared LFM base #30874

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
CISC merged 2 commits into from
Oct 20, 2026
Merged

refactor : move LFM2 shortconv and FFN blocks into a shared LFM base #30874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
44c6aad
llama : move LFM2 shortconv and FFN blocks into a shared LFM base
tdakhran Oct 19, 2026
cf6914d
Update src/models/lfm-base.cpp
tdakhran Oct 20, 2026
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
114 changes: 114 additions & 0 deletions src/models/lfm-base.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
}
112 changes: 8 additions & 104 deletions src/models/lfm2.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand Down
22 changes: 21 additions & 1 deletion src/models/models.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

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

Expand Down
Loading