Commit 040ce199a for llama.cpp
commit 040ce199a156eea0da7618be1bd1ae9e0dca6b8d
Author: Pascal <admin@serveurperso.com>
Date: Thu Nov 5 09:28:03 2026 +0100
graph: store the MTP conv snapshots of every recurrent model in one set_rows (#31660)
* graph: store the conv state of every rollback arch in one set_rows
The rollback slots of the causal conv history were written with one
copy per slot in qwen4exp, glm5-next, kimi-k3, bailingmoe3, lfm and
mamba2, while delta_net already wrote them with a single set_rows.
build_conv_state_store now does this for all of them: one copy without
rollback, otherwise one set_rows of the overlapping windows into the
min(n, n_rs_seq + 1) slots a rollback can read.
The slot order {n_rs_seq, ..., 0} is computed once per graph and each
conv layer takes a view of its tail. Its size does not depend on the
ubatch, so the reserved graph covers every ubatch shape, and the graph
input that held it is gone.
A model that is not causal never rolls back, so the context clamps
n_rs_seq to 0 for it, like for an arch without rollback support.
* nits: compute the conv slot order in the graph and keep the n_head name of bailingmoe3
* graph: build the conv slot order in each conv layer
The slot order {K-1, ..., 0} of the conv state store was one tensor
shared by every conv layer of the graph, built on the device of the
first one, so each layer on another device copied the view of its
tail. Each conv layer now builds its own order next to its set_rows,
and llm_graph_input_rs no longer holds it.
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 63b203201..267908fa7 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -202,7 +202,7 @@ llama_context::llama_context(
}
cparams.n_rs_seq = params.n_rs_seq;
- if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
+ if (cparams.n_rs_seq > 0 && (!llm_arch_supports_rs_rollback(model.arch) || !hparams.causal_attn)) {
LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",
__func__, cparams.n_rs_seq);
cparams.n_rs_seq = 0;
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index edf1ca71f..4488b880d 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -481,17 +481,6 @@ void llm_graph_input_cls::set_input(const llama_ubatch * ubatch) {
}
}
-static void set_input_s_slot_rev(ggml_tensor * s_slot_rev) {
- if (s_slot_rev) {
- GGML_ASSERT(ggml_backend_buffer_is_host(s_slot_rev->buffer));
- int32_t * data = (int32_t *) s_slot_rev->data;
-
- for (int64_t i = 0; i < s_slot_rev->ne[0]; ++i) {
- data[i] = (int32_t) (s_slot_rev->ne[0] - 1 - i);
- }
- }
-}
-
void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
GGML_UNUSED(ubatch);
@@ -506,8 +495,6 @@ void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
data[i] = mctx->s_copy(i);
}
}
-
- set_input_s_slot_rev(s_slot_rev);
}
bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {
@@ -1222,8 +1209,6 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
data[i] = mctx->get_recr()->s_copy(i);
}
}
-
- set_input_s_slot_rev(inp_rs->s_slot_rev);
}
bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
@@ -1270,8 +1255,6 @@ void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) {
data[i] = mctx->get_recr()->s_copy(i);
}
}
-
- set_input_s_slot_rev(inp_rs->s_slot_rev);
}
bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) {
@@ -1329,8 +1312,6 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) {
data[i] = mctx->get_recr()->s_copy(i);
}
}
-
- set_input_s_slot_rev(inp_rs->s_slot_rev);
}
bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params) {
@@ -4013,6 +3994,67 @@ ggml_tensor * llm_graph_context::build_rs(
get_state_rows);
}
+void llm_graph_context::build_conv_state_store(
+ llm_graph_input_rs * inp,
+ ggml_tensor * s,
+ ggml_tensor * conv_x,
+ int64_t n_cols,
+ size_t offs) const {
+ const auto * mctx_cur = inp->mctx;
+
+ const int64_t n = conv_x->ne[0] - n_cols; // n_seq_tokens
+ const int64_t channels = conv_x->ne[1];
+ const int64_t n_seqs = conv_x->ne[2];
+
+ const size_t row_size = s->nb[1];
+ const size_t cell_offs = mctx_cur->get_head()*row_size + offs;
+
+ if (cparams.n_rs_seq == 0) {
+ ggml_tensor * last = ggml_view_3d(ctx0, conv_x,
+ n_cols, channels, n_seqs,
+ conv_x->nb[1], conv_x->nb[2],
+ ggml_row_size(conv_x->type, n));
+
+ ggml_tensor * cell = ggml_view_2d(ctx0, s,
+ n_cols*channels, n_seqs,
+ row_size, cell_offs);
+
+ ggml_build_forward_expand(gf, ggml_cpy(ctx0, last, cell));
+ return;
+ }
+
+ // [TAG_RECURRENT_ROLLBACK_SPLITS]
+ // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside
+ // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument
+ // a rollback of k tokens reads slot k, and the memory allows k < n only (see rs_max), so slots from n on are left unwritten
+ const int64_t K = (int64_t) cparams.n_rs_seq + 1;
+ const int64_t Kw = std::min(n, K);
+
+ // each conv layer builds its own slot order {K-1, ..., 0}, on the device of the layer, and takes its tail
+ // {Kw-1, ..., 0}; the size K never depends on the ubatch, so the reserved graph covers them all
+ ggml_tensor * s_slot_rev = ggml_arange(ctx0, 0.0f, (float) K, 1.0f);
+ s_slot_rev = ggml_scale_bias(ctx0, s_slot_rev, -1.0f, (float) (K - 1));
+ s_slot_rev = ggml_cast(ctx0, s_slot_rev, GGML_TYPE_I32);
+ cb(s_slot_rev, "s_slot_rev", -1);
+
+ ggml_tensor * slot_order = ggml_view_1d(ctx0, s_slot_rev, Kw, ggml_row_size(GGML_TYPE_I32, K - Kw));
+
+ // window w starts at column n - (Kw - 1) + w (ascending time) and is rollback slot Kw - 1 - w,
+ // consecutive windows start one element apart, so they overlap
+ ggml_tensor * windows = ggml_view_4d(ctx0, conv_x,
+ n_cols, Kw, channels, n_seqs,
+ ggml_element_size(conv_x), conv_x->nb[1], conv_x->nb[2],
+ ggml_row_size(conv_x->type, n - (Kw - 1)));
+
+ // [n_cols, slot, channel, seq] view of the cells of this ubatch
+ ggml_tensor * slots = ggml_view_4d(ctx0, s,
+ n_cols, K, channels, n_seqs,
+ (size_t) mctx_cur->get_size()*row_size, ggml_row_size(s->type, n_cols), row_size,
+ cell_offs);
+
+ ggml_build_forward_expand(gf, ggml_set_rows(ctx0, slots, windows, slot_order));
+}
+
ggml_tensor * llm_graph_context::build_rwkv_token_shift_load(
llm_graph_input_rs * inp,
const llama_ubatch & ubatch,
diff --git a/src/llama-graph.h b/src/llama-graph.h
index e1804fb3d..5305bdb96 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -309,9 +309,6 @@ public:
ggml_tensor * s_copy_main; // I32 [n_seqs]
ggml_tensor * s_copy_tail; // I32 [n_rs - 1]
- // I32 [n_slot] = {n_slot - 1, ..., 0}, rollback slot of each conv window, created by the graphs that write them
- ggml_tensor * s_slot_rev = nullptr;
-
const llama_memory_recurrent_context * mctx;
// used in view offsets, need to match for valid graph reuse
@@ -1436,6 +1433,15 @@ struct llm_graph_context {
int32_t n_seqs,
const llm_graph_get_rows_fn & get_state_rows = nullptr) const;
+ // store the conv history of the ubatch cells, the last n_cols columns of conv_x {n_cols + n_seq_tokens, channels, n_seqs},
+ // at byte offset offs of each row of s, and with n_rs_seq > 0 once per rollback slot of this ubatch
+ void build_conv_state_store(
+ llm_graph_input_rs * inp,
+ ggml_tensor * s,
+ ggml_tensor * conv_x,
+ int64_t n_cols,
+ size_t offs) const;
+
ggml_tensor * build_rwkv_token_shift_load(
llm_graph_input_rs * inp,
const llama_ubatch & ubatch,
diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp
index 49c0f7fea..c93fae895 100644
--- a/src/models/bailingmoe3.cpp
+++ b/src/models/bailingmoe3.cpp
@@ -160,24 +160,10 @@ std::unique_ptr<llm_graph_context> llama_model_bailingmoe3::build_arch_graph(con
return std::make_unique<graph>(*this, params);
}
-static ggml_tensor * bailingmoe3_causal_conv1d(
- ggml_cgraph * gf,
- ggml_context * ctx0,
- ggml_tensor * conv_states_all,
- ggml_tensor * conv_state_all,
- int64_t qkv,
- ggml_tensor * x,
- ggml_tensor * proj_w,
- ggml_tensor * conv_w,
- int64_t d_conv,
- int64_t head_dim,
- int64_t n_head,
- int64_t n_seq_tokens,
- int64_t n_seqs,
- int64_t n_tokens,
- int64_t cache_head,
- uint32_t mem_size,
- uint32_t n_rs_seq) {
+ggml_tensor * llama_model_bailingmoe3::graph::build_causal_conv1d(
+ llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head, int64_t n_seq_tokens, int64_t n_seqs) {
const int64_t d_inner = head_dim * n_head;
const int64_t conv_state_size = (d_conv - 1) * d_inner;
const int64_t total_state_size = 3 * conv_state_size;
@@ -191,18 +177,8 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
- const int64_t K = (int64_t) n_rs_seq + 1;
- const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
-
- for (int64_t slot = 0; slot < n_written; ++slot) {
- ggml_tensor * conv_snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
- conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]);
- ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
- ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
- (d_conv - 1) * ggml_element_size(conv_states_all),
- total_state_size * ggml_element_size(conv_states_all),
- ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
- }
+ build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,
+ ggml_row_size(conv_states_all->type, qkv * conv_state_size));
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight);
@@ -253,20 +229,18 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
if (hparams.is_recr(il)) {
const auto * mctx_cur = inp_rs->mctx;
- const auto cache_head = mctx_cur->get_head();
- const auto mem_size = mctx_cur->get_size();
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
- ggml_tensor * q = bailingmoe3_causal_conv1d(
- gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
- d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
- ggml_tensor * k = bailingmoe3_causal_conv1d(
- gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
- d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
- ggml_tensor * v = bailingmoe3_causal_conv1d(
- gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
- d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+ ggml_tensor * q = build_causal_conv1d(
+ inp_rs, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
+ d_conv, head_dim, n_head, n_seq_tokens, n_seqs);
+ ggml_tensor * k = build_causal_conv1d(
+ inp_rs, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
+ d_conv, head_dim, n_head, n_seq_tokens, n_seqs);
+ ggml_tensor * v = build_causal_conv1d(
+ inp_rs, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
+ d_conv, head_dim, n_head, n_seq_tokens, n_seqs);
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
gate = ggml_add(ctx0, gate, layer.ssm_dt_b);
diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp
index 0a44db66f..4c0c24d6c 100644
--- a/src/models/delta-net-base.cpp
+++ b/src/models/delta-net-base.cpp
@@ -450,11 +450,6 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
int64_t conv_kernel_size,
int64_t conv_channels,
int il) {
- const auto * mctx_cur = inp->mctx;
-
- const auto kv_head = mctx_cur->get_head();
- const auto mem_size = mctx_cur->get_size();
-
const int64_t n_seqs = ubatch.n_seqs;
ggml_tensor * conv_states = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
@@ -469,62 +464,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
ggml_tensor * conv_input = ggml_concat(ctx0, conv_states, qkv_mixed, 0);
cb(conv_input, "conv_input", il);
- const int64_t row_count = (conv_kernel_size - 1) * conv_channels;
-
- const size_t row_size = ggml_row_size(conv_states_all->type, row_count);
-
- if (cparams.n_rs_seq == 0) {
- const int64_t s_idx = conv_input->ne[0] - conv_states->ne[0];
- const int64_t s_slot = 0;
-
- ggml_tensor * conv_state_last =
- ggml_view_3d(ctx0, conv_input,
- conv_kernel_size - 1, conv_channels, n_seqs,
- conv_input->nb[1], conv_input->nb[2],
- ggml_row_size(conv_input->type, s_idx));
- cb(conv_state_last, "conv_state_last", il);
-
- ggml_tensor * conv_state_update =
- ggml_view_2d(ctx0, conv_states_all,
- row_count, n_seqs, conv_states_all->nb[1],
- (s_slot * mem_size + kv_head) * row_size);
- cb(conv_state_update, "conv_state_update", il);
-
- ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update));
- } else {
- // [TAG_RECURRENT_ROLLBACK_SPLITS]
- // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside
- // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument
-
- const int64_t K = (int64_t) cparams.n_rs_seq + 1;
- const int64_t n = conv_input->ne[0] - conv_states->ne[0]; // n_seq_tokens
- const int64_t Kw = std::min(n + 1, K); // windows that exist in this ubatch
-
- // {Kw-1, ..., 0}, filled on the host and shared by every layer of the graph
- if (inp->s_slot_rev == nullptr) {
- inp->s_slot_rev = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, Kw);
- ggml_set_input(inp->s_slot_rev);
- ggml_set_name(inp->s_slot_rev, "rs_s_slot_rev");
- }
-
- // window w starts at column n - (Kw - 1) + w (ascending time) and is rollback slot Kw - 1 - w
- // slots reaching back past this ubatch are left unwritten, like the ssm snapshots in build_recurrent_attn
- // consecutive windows start one element apart, so they overlap
- ggml_tensor * windows = ggml_view_4d(ctx0, conv_input,
- conv_kernel_size - 1, Kw, conv_channels, n_seqs,
- ggml_element_size(conv_input), conv_input->nb[1], conv_input->nb[2],
- ggml_row_size(conv_input->type, n - (Kw - 1)));
-
- // [d-1, slot, channel, seq] view of the cells of this ubatch
- ggml_tensor * slots = ggml_view_4d(ctx0, conv_states_all,
- conv_kernel_size - 1, K, conv_channels, n_seqs,
- (size_t) mem_size * row_size,
- ggml_row_size(conv_states_all->type, conv_kernel_size - 1),
- conv_states_all->nb[1],
- kv_head * row_size);
-
- ggml_build_forward_expand(gf, ggml_set_rows(ctx0, slots, windows, inp->s_slot_rev));
- }
+ build_conv_state_store(inp, conv_states_all, conv_input, conv_kernel_size - 1, 0);
return conv_input;
}
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 100af6d23..c37792339 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -191,13 +191,11 @@ std::unique_ptr<llm_graph_context> llama_model_glm5_next::build_arch_graph(const
}
// Causal conv1d over one of Q/K/V
-static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
- ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
- int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
- int64_t d_conv, int64_t head_dim, int64_t n_head,
- int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,
- int64_t mem_size, int64_t K_rs) {
- const int64_t d_inner = head_dim * n_head;
+ggml_tensor * llama_model_glm5_next::graph::build_kda_conv1d(
+ llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs) {
+ const int64_t d_inner = head_dim * n_head_kda;
const int64_t conv_state_size = (d_conv - 1) * d_inner;
const int64_t n_embd_r_total = 3 * conv_state_size;
@@ -210,26 +208,15 @@ static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
- // group s holds the conv window s tokens back.
- // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.
- for (int64_t s = 0; s < K_rs; ++s) {
- const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);
- ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
- conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);
- ggml_build_forward_expand(gf,
- ggml_cpy(ctx0, conv_x_s,
- ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
- (d_conv - 1) * ggml_element_size(conv_states_all),
- n_embd_r_total * ggml_element_size(conv_states_all),
- ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
- }
+ build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,
+ ggml_row_size(conv_states_all->type, qkv * conv_state_size));
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
Xcur = ggml_silu(ctx0, Xcur);
- return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+ return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head_kda, n_seq_tokens, n_seqs);
}
@@ -812,17 +799,13 @@ ggml_tensor * llama_model_glm5_next::graph::build_kda_layer(
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
const auto * mctx_cur = inp_rs->mctx;
- const auto kv_head = mctx_cur->get_head();
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
- const int64_t mem_size = mctx_cur->get_size();
- const int64_t K_rs = (int64_t) cparams.n_rs_seq + 1;
-
- ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
- ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
- ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+ ggml_tensor * Qcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs);
+ ggml_tensor * Kcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs);
+ ggml_tensor * Vcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs);
cb(Qcur, "kda_q_conv", il);
cb(Kcur, "kda_k_conv", il);
cb(Vcur, "kda_v_conv", il);
diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp
index 27ceffe92..3644fe77a 100644
--- a/src/models/kimi-k3.cpp
+++ b/src/models/kimi-k3.cpp
@@ -1,6 +1,5 @@
#include "models.h"
-#include <algorithm>
#include "llama-memory-recurrent.h"
//
@@ -353,13 +352,11 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par
//
// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use
-static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
- ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
- int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
- int64_t d_conv, int64_t head_dim, int64_t n_head,
- int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,
- int64_t mem_size, int64_t K_rs) {
- const int64_t d_inner = head_dim * n_head;
+ggml_tensor * llama_model_kimi_k3::graph::build_kda_conv1d(
+ llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs) {
+ const int64_t d_inner = head_dim * n_head_kda;
const int64_t conv_state_size = (d_conv - 1) * d_inner;
const int64_t n_embd_r_total = 3 * conv_state_size;
@@ -372,26 +369,15 @@ static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
- // group s holds the conv window s tokens back.
- // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.
- for (int64_t s = 0; s < K_rs; ++s) {
- const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);
- ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
- conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);
- ggml_build_forward_expand(gf,
- ggml_cpy(ctx0, conv_x_s,
- ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
- (d_conv - 1) * ggml_element_size(conv_states_all),
- n_embd_r_total * ggml_element_size(conv_states_all),
- ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
- }
+ build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,
+ ggml_row_size(conv_states_all->type, qkv * conv_state_size));
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
Xcur = ggml_silu(ctx0, Xcur);
- return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+ return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head_kda, n_seq_tokens, n_seqs);
}
ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
@@ -400,17 +386,13 @@ ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
const auto * mctx_cur = inp_rs->mctx;
- const auto kv_head = mctx_cur->get_head();
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
- const int64_t mem_size = mctx_cur->get_size();
- const int64_t K_rs = (int64_t) cparams.n_rs_seq + 1;
-
- ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
- ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
- ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+ ggml_tensor * Qcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs);
+ ggml_tensor * Kcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs);
+ ggml_tensor * Vcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs);
cb(Qcur, "kda_q_conv", il);
cb(Kcur, "kda_k_conv", il);
cb(Vcur, "kda_v_conv", il);
diff --git a/src/models/lfm-base.cpp b/src/models/lfm-base.cpp
index 74944c13c..f2a6ee461 100644
--- a/src/models/lfm-base.cpp
+++ b/src/models/lfm-base.cpp
@@ -2,8 +2,6 @@
#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,
@@ -11,7 +9,6 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
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);
@@ -55,20 +52,7 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
}
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)));
- }
+ build_conv_state_store(inp_recr, conv_state, bx, d_conv, 0);
auto * conv_kernel = model.layers[il].shortconv.conv;
auto * conv_out = ggml_ssm_conv(ctx0, bx, conv_kernel);
diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp
index 63af3113d..8d67d39e4 100644
--- a/src/models/mamba-base.cpp
+++ b/src/models/mamba-base.cpp
@@ -206,19 +206,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
// => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}
ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);
- const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);
- const size_t row_size = ggml_row_size(conv_states_all->type, row_count);
- const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
-
- for (int64_t slot = 0; slot < n_written; ++slot) {
- ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
- conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);
-
- ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
- ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,
- conv_states_all->nb[1],
- ((size_t) slot * mem_size + kv_head) * row_size)));
- }
+ build_conv_state_store(inp, conv_states_all, conv_x, d_conv - 1, 0);
// 1D convolution
// The equivalent is to make a self-overlapping view of conv_x
diff --git a/src/models/models.h b/src/models/models.h
index dd3b7b1ed..bc980dd4d 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -2176,6 +2176,11 @@ struct llama_model_bailingmoe3 : public llama_model_base {
graph(const llama_model & model, const llm_graph_params & params);
const llama_model & model;
+
+ // causal conv1d over one of Q/K/V, qkv selects which third of the conv state row is used
+ ggml_tensor * build_causal_conv1d(llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head, int64_t n_seq_tokens, int64_t n_seqs);
};
struct graph_mtp : public llm_graph_context {
@@ -3057,6 +3062,11 @@ struct llama_model_kimi_k3 : public llama_model_base {
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
+ // causal conv1d over one of Q/K/V, qkv selects which third of the conv state row is used
+ ggml_tensor * build_kda_conv1d(llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs);
+
ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_attn_k * inp_attn_k,
llm_graph_input_attn_kv * inp_attn_kv,
@@ -3129,6 +3139,11 @@ struct llama_model_glm5_next : public llama_model_base {
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
+ // causal conv1d over one of Q/K/V, qkv selects which third of the conv state row is used
+ ggml_tensor * build_kda_conv1d(llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs);
+
ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, const llama_layer & layer,
const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il);
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index e4848cb77..de6d21fb3 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -1402,10 +1402,6 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
int64_t state_cols,
int64_t channels,
int il) {
- const auto * mctx_cur = inp->mctx;
-
- const auto kv_head = mctx_cur->get_head();
-
const int64_t n_seqs = ubatch.n_seqs;
const int64_t row_total = conv_states_all->ne[0];
@@ -1423,28 +1419,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);
- // [TAG_RECURRENT_ROLLBACK_SPLITS] keep the last state_cols columns once per rollback slot,
- // slot s ending s tokens earlier so a rollback of s tokens reads a history that never saw them
- const size_t row_size = ggml_row_size(conv_states_all->type, row_total);
- const uint32_t mem_size = mctx_cur->get_size();
-
- const int64_t n_slots = (int64_t) cparams.n_rs_seq + 1;
-
- for (int64_t slot = 0; slot < n_slots; ++slot) {
- const int64_t s_idx = std::max<int64_t>(0, conv_input->ne[0] - state_cols - slot);
-
- ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
- state_cols, channels, n_seqs,
- conv_input->nb[1], conv_input->nb[2],
- ggml_row_size(conv_input->type, s_idx));
-
- ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
- state_cols * channels, n_seqs,
- conv_states_all->nb[1],
- (slot * mem_size + kv_head) * row_size);
-
- ggml_build_forward_expand(gf, ggml_cpy(ctx0, tail, dst));
- }
+ build_conv_state_store(inp, conv_states_all, conv_input, state_cols, 0);
return conv_input;
}