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

ggml-metal: add chunked SSD MMA for Mamba-2 prefill optimization (#26647)

Browse Source 
* metal: WIP chunked SSD SSM_SCAN kernels for multi-token prefill

* metal: drop scalar SSD path; MMA + sequential tail

* drop WIP ssm scan test noise

* remove state_from_dst and rename CS and NSG constants

* remove unrelated  added whitespace padding

* added clarity to mma_tokens calculation

* added clarity to use_mma bool checks

* added comments to metal ssd op constants for clarity

* reserve K tokens for sequential kernel rollback snapshots

* reset concurrency between mma and seq tail

* remove print args no longer used

* fixed comment to no longer point to specific line

* add FC_SSM_SCAN so seq path skips token offlset unless it's mma tail

* added changes to new ssm.metal for rebase after ggml-metal.metal refactor

* specialize ssm_scan tail with a template instead of a function constant

---------

Co-authored-by: dpantaleoni <dominikpantaleoni@gmail.com>
Co-authored-by: forforever73 <690105611@qq.com>
This commit is contained in:
Dominik Pantaleoni
2026-08-26 11:57:07 +03:00
committed by GitHub
co-authored by dpantaleoni forforever73
parent 5d5cb4c3a4
commit 11cd988428
6 changed files with 279 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test-backend-ops.cpp
+10 -4
View File
@@ -4120,9 +4120,10 @@ struct test_ssm_scan : public test_case {
const int64_t n_seqs;
const bool xbc_overlap;
const int64_t K;
const bool weak_decay;
std::string vars() override {
return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K);
return VARS_TO_STR10(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K, weak_decay);
}
test_ssm_scan(ggml_type type = GGML_TYPE_F32,
@@ -4133,8 +4134,9 @@ struct test_ssm_scan : public test_case {
int64_t n_seq_tokens = 32,
int64_t n_seqs = 32,
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) {}
int64_t K = 1,
bool weak_decay = 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), K(K), weak_decay(weak_decay) {}
double max_nmse_err() override {
// SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
@@ -4187,7 +4189,7 @@ struct test_ssm_scan : public test_case {
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);
init_tensor_uniform(t, weak_decay ? -0.02f : -1.0f, weak_decay ? -0.005f : -0.5f);
} else {
init_tensor_uniform(t);
}
@@ -9111,6 +9113,10 @@ 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, 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_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 64, 4)); // Metal SSD one chunk MMA only, no seq tail
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 65, 2)); // SSD one chunk + 1-token sequential tail
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2)); // SSD multi-chunk, no tail (exercises the chunk-to-chunk state handoff)
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2, false, /*K=*/1, /*weak_decay=*/true)); // SSD multi-chunk, carried state not numerically negligible
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));