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

* 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
+2 -1
View File
@@ -129,7 +129,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tail);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op);