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

ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration (#22675)

Browse Source 
* ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration

* cuda: added SSD CICD fixes for CUDA / HIP / MUSA / MSVC.

* ggml-cuda: review comments fixed.

* ggml-cuda: Fuse M matrix materialization into pre_matmul kernel and enabled test.

* ggml-cuda: test updates and fixes

* ggml-cuda: test updates to remove hardcoding of tensor initialise data limits.

* ggml-cuda: ssd minor review comment fixed.

* ggml-cuda: ssd minor CICD fixed.

* CUDA SSD: Fixes correctness by promoting s0_stride_seq to int64_t, improves memory coalescing in ssm_ssd_prepare_dt_kernel, and boosts efficiency by merging B_weighted and C_scaled; also addresses prior review comments.

* cuda: fix sdata read-write race in prepare_dt fallback scan loop
This commit is contained in:
Bhavik Sharda
2026-07-28 17:33:42 +05:30
committed by GitHub
parent 84075273c8
commit b62b350981
2 changed files with 499 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test-backend-ops.cpp
+18 -3
View File
@@ -4000,7 +4000,7 @@ struct test_ssm_scan : public test_case {
test_ssm_scan(ggml_type type = GGML_TYPE_F32,
int64_t d_state = 32,
int64_t head_dim = 1, // non-zero for Mamba-2
int64_t head_dim = 1, // 1 = Mamba-1; > 1 = Mamba-2 (scalar A per head)
int64_t n_head = 32,
int64_t n_group = 1,
int64_t n_seq_tokens = 32,
@@ -4008,6 +4008,11 @@ struct test_ssm_scan : public test_case {
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) {}
double max_nmse_err() override {
// SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
return (head_dim > 1) ? 2e-7 : 1e-7;
}
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 * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs);
@@ -4034,14 +4039,14 @@ struct test_ssm_scan : public test_case {
return out;
}
// similar to test_mul_mat_id
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; }
// ids
// ids: permutation of [0..n_seqs)
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++) {
@@ -4050,6 +4055,11 @@ struct test_ssm_scan : public test_case {
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) {
// A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors)
init_tensor_uniform(t, -1.0f, -0.5f);
} else {
init_tensor_uniform(t);
}
@@ -8770,6 +8780,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, 64, 16, 2, 32, 4)); // Mamba-2
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap
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_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1));
test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1));
@@ -9990,6 +10003,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_ssm_conv_bias_silu(GGML_TYPE_F32, {4, 3328, 1, 1}, {4, 3328, 1, 1}, true)); // generate
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 512, 1)); // prefill
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 1, 1)); // generate
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B prefill
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 1, 1)); // Nemotron-9B generate
// acc
test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1));