Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml : recurrent state rollback for ggml_ssm_scan (#26623)

Browse Source 
* Initial changes for Recurrent state rollback for nemotron for cpu and cuda

* Removing CPU RS rollback. Will enable it in subsequent PRs

* addition of test case

* Removing assert and calling runtime API to check if op is supported

* removing extra API and updating the call sites for K

* replace static cuda detection to runtime fused_op api

* address review comments and fallback when SSM rollback not supprted

* Adding changes for supporting RS-rollback in CPU. Also added test-backend-ops for cpu and cuda

* removing memory manipulation as rs rollback is now supported in CPU

* removing the static probe which is not needed now

* correcting the format

* address review comments

* enabling test for all the backends, unsupported backends will fallback to CPU

* Apply suggestions from code review

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

* choose different graph based on the result of fused_ssm_op is supported or not and also handled memory->n_rs_seq >1 case incase of op is not supported

* Support K > 1 in ssm_scan for all backends

* Fix CI Issues

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Co-authored-by: Gaurav Garg <gaugarg@nvidia.com>
This commit is contained in:
lnigam
2026-08-14 17:20:40 +03:00
committed by GitHub
co-authored by Georgi Gerganov Gaurav Garg
parent 4c1a0af40d
commit 1692f9e50b
25 changed files with 291 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-arch.cpp
+2
View File
@@ -1001,6 +1001,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_DEEPSEEK4:
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
return true;
default:
return false;
src/llama-context.cpp
+1 -1
View File
@@ -103,7 +103,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)) {
LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model arch does not support recurrent partial rollback; clamping to 0\n",
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;
}
src/llama-model-loader.cpp
+1 -1
View File
@@ -1002,7 +1002,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
ggml_tensor * B = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
ggml_tensor * C = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids);
op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids, /*K=*/1);
} break;
case GGML_OP_RWKV_WKV6:
{
src/models/mamba-base.cpp
+32 -16
View File
@@ -2,6 +2,8 @@
#include "llama-memory-recurrent.h"
#include <algorithm>
llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {}
ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
@@ -118,7 +120,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
// Custom operator to optimize the parallel associative scan
// as described in the Annex D of the Mamba paper.
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
@@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
int il) const {
const auto * mctx_cur = inp->mctx;
const auto kv_head = mctx_cur->get_head();
const auto kv_head = mctx_cur->get_head();
const auto mem_size = mctx_cur->get_size();
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = hparams.ssm_d_inner;
@@ -164,6 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
GGML_ASSERT(n_seqs != 0);
GGML_ASSERT(ubatch.equal_seqs());
@@ -173,6 +177,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
const int64_t state_slots = ssm_states_all->ne[1];
ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
conv = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);
@@ -198,15 +203,19 @@ 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);
// copy last (d_conv - 1) columns back into the state cache
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 * (conv_x->nb[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);
ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
ggml_view_1d(ctx0, conv_states_all,
(d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs),
kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) *
ggml_element_size(conv_states_all))));
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)));
}
// 1D convolution
// The equivalent is to make a self-overlapping view of conv_x
@@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
// (this is necessary in order to properly use the states before they are overwritten,
// while avoiding to make unnecessary copies of the states)
auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {
ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size());
ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots);
// TODO: use semistructured matrices to implement state-space duality
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
// K > 1 asks the backend to return rollback snapshots in addition to the final state.
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
const int64_t D = d_state * d_inner;
const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
const size_t row_size = ggml_row_size(ssm_states_all->type, D);
const size_t y_row_size = ggml_row_size(y_ssm->type, D);
const size_t state_offset = ggml_nelements(x) * ggml_element_size(x);
// store last states
ggml_build_forward_expand(
gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]),
ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs,
kv_head * d_state * d_inner * ggml_element_size(ssm_states_all))));
gf, ggml_cpy(ctx0,
ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written,
y_row_size, y_row_size * n_seqs, state_offset),
ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written,
ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size)));
ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1],
n_seq_tokens * n_head * x->nb[1], 0);
src/models/plamo2.cpp
+1 -1
View File
@@ -382,7 +382,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
// Custom operator to optimize the parallel associative scan
// as described in the Annex D of the Mamba paper.
// => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
};
ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);