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-ops.cpp
+45 -15
View File
@@ -1677,6 +1677,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
ggml_metal_library_t lib = ctx->lib;
ggml_metal_encoder_t enc = ctx->enc;
const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev);
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
@@ -1722,6 +1723,8 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
/*.n_head =*/ n_head,
/*.n_group =*/ n_group,
/*.n_seq_tokens =*/ n_seq_tokens,
/*.n_seq_tokens_total =*/ n_seq_tokens,
/*.token_offset =*/ 0,
/*.n_seqs =*/ n_seqs,
/*.K =*/ K,
/*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float),
@@ -1751,26 +1754,53 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
/*.nb0 =*/ nb0,
};
auto pipeline = ggml_metal_library_get_pipeline_ssm_scan(lib, op);
constexpr int64_t CHUNK = OP_SSM_SCAN_SSD_CS;
GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
const int64_t snap_reserve = K > 1 ? K : 0; // tokens reserved for sequential kernel rollback snapshots
const int64_t mma_tokens = ((n_seq_tokens - snap_reserve) / CHUNK) * CHUNK; // largest multiple of CHUNK that leaves snap_reserve for the tail
const bool use_mma =
mma_tokens > 0 &&
ne30 == 1 && // checks that A tensor is set to scalar decay per head (A shape {1, n_head})
props_dev->has_simdgroup_mm && // hardware check for M1 or newer
d_state % 8 == 0 && // d_state must be multiple of 8 to align with simdgroup_float 8x8 tiles
d_inner == OP_SSM_SCAN_SSD_HD; // mma kernel is specialized for the Mamba-2 head dim; this checks it
const size_t smem = pipeline.smem;
const auto dispatch = [&](ggml_metal_pipeline_with_params pipeline, int64_t nth, int64_t n_tg_x) {
GGML_ASSERT(nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
GGML_ASSERT(pipeline.smem <= props_dev->max_theadgroup_memory_size);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 4);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 5);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), 6);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), 7);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 8);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 4);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 5);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), 6);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), 7);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 8);
ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0);
ggml_metal_encoder_dispatch_threadgroups(enc, n_tg_x, n_head, n_seqs, nth, 1, 1);
};
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
if (!use_mma) {
dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, false), d_state, d_inner);
return 1;
}
ggml_metal_encoder_dispatch_threadgroups(enc, d_inner, n_head, n_seqs, d_state, 1, 1);
args.n_seq_tokens = mma_tokens;
dispatch(
ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(lib, op),
OP_SSM_SCAN_SSD_NSG*32,
1);
if (mma_tokens < n_seq_tokens) {
ggml_metal_op_concurrency_reset(ctx);
args.n_seq_tokens = n_seq_tokens - mma_tokens;
args.token_offset = mma_tokens;
dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, true), d_state, d_inner);
}
return 1;
}