darthrubik/llama-cpp
1
0
Fork 0
Code Issues Pull Requests 1 Packages Projects Releases Wiki Activity

refactor : move LFM2 shortconv and FFN blocks into a shared LFM base (#30874)
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
Check Pre-Tokenizer Hashes / pre-tokenizer-hashes (push) Has been cancelled
Python check requirements.txt / check-requirements (push) Has been cancelled
Python Type-Check / python type-check (push) Has been cancelled
Update Operations Documentation / update-ops-docs (push) Has been cancelled

Browse Source 
* 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>
This commit is contained in:
Tarek Dakhran
2026-10-20 23:17:13 +02:00
committed by GitHub
parent 52ffa02745
commit 84e66d0321
3 changed files with 143 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/lfm-base.cpp
+114
View File
@@ -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);
}
src/models/lfm2.cpp
+8 -104
View File
@@ -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);
src/models/models.h
+21 -1
View File
@@ -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);
};