ggml : recurrent state rollback for ggml_ssm_scan (#26623)
* 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
co-authored by
Georgi Gerganov
Gaurav Garg
Georgi Gerganov
Gaurav Garg
parent
4c1a0af40d
commit
1692f9e50b
@@ -217,6 +217,16 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
|
||||
set_tests_properties(test-recurrent-state-rollback PROPERTIES
|
||||
FIXTURES_REQUIRED generate-models
|
||||
)
|
||||
|
||||
llama_test(
|
||||
test-recurrent-state-rollback
|
||||
NAME test-recurrent-state-rollback-nemotron-h
|
||||
LABEL main
|
||||
ARGS -m "${MODEL_DIR}/nemotron_h-dense.gguf"
|
||||
)
|
||||
set_tests_properties(test-recurrent-state-rollback-nemotron-h PROPERTIES
|
||||
FIXTURES_REQUIRED generate-models
|
||||
)
|
||||
endif()
|
||||
|
||||
llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)
|
||||
|
||||
+117
-4
@@ -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));
|
||||
|
||||
Reference in New Issue
Block a user