From 84e66d03218711dfcaa993e5f64a546abe82df9a Mon Sep 17 00:00:00 2001 From: Tarek Dakhran Date: Tue, 20 Oct 2026 23:17:13 +0200 Subject: [PATCH] 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 --------- Co-authored-by: Georgi Gerganov --- src/models/lfm-base.cpp | 114 ++++++++++++++++++++++++++++++++++++++++ src/models/lfm2.cpp | 112 +++------------------------------------ src/models/models.h | 22 +++++++- 3 files changed, 143 insertions(+), 105 deletions(-) create mode 100644 src/models/lfm-base.cpp 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 + +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(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(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 llama_model_lfm2::build_arch_graph(const llm_ template llama_model_lfm2::graph::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; using inp_attn_type = std::conditional_t; using mem_hybrid_ctx = std::conditional_t; - // 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(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::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(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(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::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(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::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 - struct graph : public llm_graph_context { + struct graph : public llm_build_lfm_base { graph(const llama_model & model, const llm_graph_params & params); };