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
ggml/src/ggml-metal/ggml-metal-device.cpp
+23 -2
View File
@@ -572,7 +572,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched
return res;
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op) {
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op, bool tail) {
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
char base[256];
@@ -580,7 +580,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me
const int nsg = (ne00 + 31)/32;
snprintf(base, 256, "kernel_ssm_scan_%s", ggml_type_name(op->src[0]->type));
snprintf(base, 256, "kernel_ssm_scan_%s%s", ggml_type_name(op->src[0]->type), tail ? "_tail" : "");
snprintf(name, 256, "%s_nsg=%d", base, nsg);
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
@@ -598,6 +598,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me
return res;
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(ggml_metal_library_t lib, const ggml_tensor * op) {
char base[256];
char name[256];
snprintf(base, 256, "kernel_ssm_scan_ssd_mma_%s", ggml_type_name(op->src[0]->type));
snprintf(name, 256, "%s", base);
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
if (!res.pipeline) {
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
}
// acs/exp(acs)/state-decay vectors + dtX + SAM rows + two 8x8 tiles per simdgroup
res.smem = (3*OP_SSM_SCAN_SSD_CS +
OP_SSM_SCAN_SSD_CS*OP_SSM_SCAN_SSD_HD +
OP_SSM_SCAN_SSD_NSG*8*OP_SSM_SCAN_SSD_CS +
OP_SSM_SCAN_SSD_NSG*2*8*8)*sizeof(float);
return res;
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) {
char base[256];
char name[256];