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
tests/test-backend-ops.cpp
+117 -4
View File
@@ -4111,9 +4111,10 @@ struct test_ssm_scan : public test_case {
const int64_t n_seq_tokens;
const int64_t n_seqs;
const bool xbc_overlap;
const int64_t K;
std::string vars() override {
return VARS_TO_STR8(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap);
return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K);
}
test_ssm_scan(ggml_type type = GGML_TYPE_F32,
@@ -4123,8 +4124,9 @@ struct test_ssm_scan : public test_case {
int64_t n_group = 1,
int64_t n_seq_tokens = 32,
int64_t n_seqs = 32,
bool xbc_overlap = false)
: type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {}
bool xbc_overlap = false,
int64_t K = 1)
: type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K) {}
double max_nmse_err() override {
// SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
@@ -4153,7 +4155,7 @@ struct test_ssm_scan : public test_case {
C = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs);
}
ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids);
ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids, K);
return out;
}
@@ -4185,6 +4187,114 @@ struct test_ssm_scan : public test_case {
}
};
struct test_ssm_scan_rollback : public test_case {
const ggml_type type;
const int64_t d_state;
const int64_t head_dim;
const int64_t n_head;
const int64_t n_group;
const int64_t n_seq_tokens;
const int64_t n_seqs;
const int64_t K;
std::string vars() override {
return VARS_TO_STR8(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, K);
}
std::string op_desc(ggml_tensor * t) override {
GGML_UNUSED(t);
return "SSM_SCAN_ROLLBACK";
}
bool run_whole_graph() override {
return true;
}
double max_err() override {
return 1e-6;
}
double err(const float * a, const float * b, size_t n) override {
double result = 0.0;
for (size_t i = 0; i < n; ++i) {
result = std::max(result, (double) fabsf(a[i]));
result = std::max(result, (double) fabsf(b[i]));
}
return result;
}
test_ssm_scan_rollback(ggml_type type = GGML_TYPE_F32,
int64_t d_state = 32,
int64_t head_dim = 64,
int64_t n_head = 16,
int64_t n_group = 2,
int64_t n_seq_tokens = 8,
int64_t n_seqs = 2,
int64_t K = 3)
: type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group),
n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * s = ggml_new_tensor_4d(ctx, type, d_state, head_dim, n_head, n_seqs);
ggml_tensor * x = ggml_new_tensor_4d(ctx, type, head_dim, n_head, n_seq_tokens, n_seqs);
ggml_tensor * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs);
ggml_tensor * A = ggml_new_tensor_2d(ctx, type, 1, n_head);
ggml_tensor * B = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs);
ggml_tensor * C = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs);
ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
ggml_tensor * full = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids, K);
const int64_t y_elems = head_dim * n_head * n_seq_tokens * n_seqs;
const int64_t state_elems = d_state * head_dim * n_head * n_seqs;
ggml_tensor * out = nullptr;
for (int64_t slot = 0; slot < K; ++slot) {
const int64_t prefix_tokens = n_seq_tokens - slot;
ggml_tensor * x_prefix = ggml_cont(ctx, ggml_view_4d(ctx, x, head_dim, n_head, prefix_tokens, n_seqs, x->nb[1], x->nb[2], x->nb[3], 0));
ggml_tensor * dt_prefix = ggml_cont(ctx, ggml_view_3d(ctx, dt, n_head, prefix_tokens, n_seqs, dt->nb[1], dt->nb[2], 0));
ggml_tensor * B_prefix = ggml_cont(ctx, ggml_view_4d(ctx, B, d_state, n_group, prefix_tokens, n_seqs, B->nb[1], B->nb[2], B->nb[3], 0));
ggml_tensor * C_prefix = ggml_cont(ctx, ggml_view_4d(ctx, C, d_state, n_group, prefix_tokens, n_seqs, C->nb[1], C->nb[2], C->nb[3], 0));
ggml_tensor * prefix = ggml_ssm_scan(ctx, s, x_prefix, dt_prefix, A, B_prefix, C_prefix, ids, /*K=*/1);
ggml_tensor * full_state = ggml_view_1d(ctx, full, state_elems, (y_elems + slot*state_elems)*ggml_element_size(full));
ggml_tensor * prefix_state = ggml_view_1d(ctx, prefix, state_elems, (head_dim*n_head*prefix_tokens*n_seqs)*ggml_element_size(prefix));
ggml_tensor * diff = ggml_sum(ctx, ggml_sqr(ctx, ggml_sub(ctx, full_state, prefix_state)));
out = out == nullptr ? diff : ggml_add(ctx, out, diff);
}
return out;
}
void initialize_tensors(ggml_context * ctx) override {
std::random_device rd;
std::default_random_engine rng(rd());
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
if (t->type == GGML_TYPE_I32) {
if (ggml_is_view_op(t->op)) { continue; }
for (int64_t r = 0; r < ggml_nrows(t); r++) {
std::vector<int32_t> data(t->ne[0]);
for (int i = 0; i < t->ne[0]; i++) {
data[i] = i;
}
std::shuffle(data.begin(), data.end(), rng);
ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
}
} else if (ggml_is_view_op(t->op)) {
continue;
} else if (t->ne[1] == n_head && t->ne[2] == 1) {
init_tensor_uniform(t, -1.0f, -0.5f);
} else {
init_tensor_uniform(t);
}
}
}
};
// GGML_OP_RWKV_WKV6
struct test_rwkv_wkv6 : public test_case {
const ggml_type type;
@@ -8952,6 +9062,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks)
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs)
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 4, 2, false, /*K=*/4)); // Mamba-2 rollback snapshots
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, false, /*K=*/3)); // Mamba-2 rollback overflow
test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, /*K=*/3)); // rollback snapshots match prefix states
test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1));
test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1));