ggml-metal: add chunked SSD MMA for Mamba-2 prefill optimization by dpantaleoni · Pull Request #26647 · ggml-org/llama.cpp

dpantaleoni · 2026-08-05T18:48:12Z

Overview

Similar to BLSharda's PR, this PR adds a chunked SSD matmul path using Mamba-2 style prefill, but for Metal devices. For token sequences that fit 64 token chunks, parallel simdgroup matmuls are used in place of sequential, per-token scans. Any remaining tokens that do not fit in the 64 token chunks fallback to using sequential scan.

Performance gains are best at batched, larger prefills, see the benchmark results below:

Prompt tokens	Parallel prompts	Baseline (t/s)	MMA (t/s)	Δ
1024	1	660.03	693.80	+5.12%
1024	4	1010.01	1058.19	+4.77%
1024	8	1087.13	1137.92	+4.67%
4096	1	1319.01	1445.44	+9.59%
4096	4	1613.06	1784.91	+10.65%
4096	8	1658.10	1826.80	+10.17%
8192	1	1609.57	1797.35	+11.67%
8192	4	1790.94	2005.24	+11.97%
8192	8	1819.37	2029.98	+11.58%

Ran on M4 Pro:

./build/bin/llama-batched-bench -m <granite-4.0-h-1b-Q8_0.gguf> -c 131072 -b 2048 -ub 2048 -npp 1024,4096,8192 -ntg 128 -npl 1,4,8 -ngl 99

I found no change in Perplexity:

Metric	Baseline	SSD	Difference
Final PPL	9.1594	9.1593	-0.0001
PPL uncertainty	±0.06642	±0.06642	0

./build/bin/llama-perplexity -m <granite-4.0-h-1b-Q8_0.gguf> -f <wikitext-2-raw> -ngl 99 -c 8192 -fa 1

Additional information

Files changed:

ggml/src/ggml-metal/ggml-metal-device.cpp and ggml/src/ggml-metal/ggml-metal-device.h : added new pipeline struct to use the new ssd kernel
ggml/src/ggml-metal/ggml-metal-impl.h : added constants for new kernel op's Chunk Size and Number of Simdgroups; added token offset tracking args to kernel args struct
ggml/src/ggml-metal/ggml-metal-ops.cpp : added gate to the new ssd path so gpus without simdgroup architecture remain on sequential scan; added dispatch for the new ssd path with sequential scan fallback and token offset tail for remaining tokens
ggml/src/ggml-metal/ggml-metal.metal : added the new kernel and updated the existing sequential scan kernel to be used for leftover tokens via token offset tracking
tests/test-backend-ops.cpp : added one test_ssm_scan test case for a single 64 token chunk with no tail.

Sources:

Mamba-2 White Paper
@gabe-l-hart original mamba 1 metal implementation and work on mamba 2 implementation here
the CUDA equivalent mentioned earlier here

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: Yes, AI was used to help draft, refine and optimize the new kernel pipeline based on my initial designs I derived from the sources mentioned in Additional Information above.

gabe-l-hart

This looks really good! I've got a number of questions around places where things might be brittle by device or where some opaque acronyms could use comments. I haven't thoroughly reviewed the kernel math itself yet, but I think it's safe to take it out of Draft as is.

gabe-l-hart · 2026-08-07T17:07:38Z


-    GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+    const int64_t mma_tokens = n_seq_tokens / CHUNK * CHUNK;


Does this need parens? I always forget my exact order of operations in c++. Might be good to add them for clarity.

Added parenthesis for clarity and a comment as well

gabe-l-hart · 2026-08-07T17:11:12Z

+    const int64_t mma_tokens = n_seq_tokens / CHUNK * CHUNK;
+    const bool use_mma =
+        mma_tokens > 0 &&
+        ne30 == 1 &&


Do you know what this specific constraint is doing? I think this is the A matrix and I think this is saying that dim-0 must be unit, but I'm not sure if this is the batch dimension or what.

For reference, here's the primary place where ssm_scan is called: https://github.com/ggml-org/llama.cpp/blob/master/src/models/mamba-base.cpp#L121

Added comments for clarity! ne30 == 1 specifically checks that the first dimension of the A matrix is 1, since Mamba-2 is scalar decay per head, rather than Mamba-1's full state-space decay matrix

gabe-l-hart · 2026-08-07T17:15:03Z

@@ -155,6 +155,9 @@
 #define OP_SUM_ROWS_NUM_SUM_ROWS 10
 #define OP_SUM_ROWS_NUM_MEAN     11

+#define OP_SSM_SCAN_SSD_CS 64


Let's add comments on these explaining the acronyms and clarifying that they're Apple device constants (ideally with some citation)

gabe-l-hart · 2026-08-07T17:18:01Z

+
+        args.n_seq_tokens = n_seq_tokens - mma_tokens;
+        args.token_offset = mma_tokens;
+        dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op), d_state, d_inner);


Iiiinteresting. I think this is dispatching two underlying kernels for a single GGML op. That should be ok, but it's not what I've seen in most op implementations. I do see at least one more example in ggml_op_set, so just noting that this is a bit different.

gabe-l-hart · 2026-08-07T17:19:59Z

@@ -1717,6 +1718,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,


I don't think this added struct size should have any significant impact, but worth double checking since this is what traverses the device boundary on dispatch.

gabe-l-hart · 2026-08-07T17:22:48Z

+// MMA fast path for the common Mamba-2 layout (head_dim=64). One threadgroup owns a complete
+// (head, sequence) pair so the channel-independent C*B^T matrix is computed once per chunk and
+// reused for all eight 8-channel output tiles.
+kernel void kernel_ssm_scan_ssd_mma_f32(


One question that always comes up for me: Is there any benefit to making this a template over dtype and supporting bf16, f16, and potentially lower precision types as well? The downside is that extra template expanded kernels take up bloat in the JIT compile for the metal kernels, but it's possible that this could give an added perf boost for lower-precision models. Worth testing at least.

Looked into this and currently all ssm state Granite Hybrid models are allocated as f32 regardless of size, but it would be a great addition in the future.

gabe-l-hart · 2026-08-07T17:23:25Z

+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]]) {
+    constexpr short CS  = OP_SSM_SCAN_SSD_CS;
+    constexpr short TC  = 8;


Would be good to have a comment for each of these explaining the magic numbers and what they physically correspond to

added comments for these as well

gabe-l-hart · 2026-08-07T17:25:49Z

+    }
+
+    // acs/exp(acs)/state-decay vectors + dtX + SAM rows + two 8x8 tiles per simdgroup
+    res.smem = (3*OP_SSM_SCAN_SSD_CS +


This is some classic kernel math magic! I'm assuming that these numbers are tied to concrete device constraints. It would be good to figure out if this can be tested across a range of Apple Silicon devices (I know M5 made some big changes).

ggerganov · 2026-08-20T11:04:51Z

@dpantaleoni Thanks for the contribution. Could you rebase this on the latest master?

dpantaleoni · 2026-08-20T15:42:41Z

@ggerganov thank you for reviewing! I just rebased

dpantaleoni · 2026-08-20T15:59:13Z

Here's some additional benchmarking on a m4 air 24 gb:

Prompt tokens	Parallel prompts	Baseline (t/s)	MMA (t/s)	Δ
1024	1	668.83	673.09	+0.64%
1024	4	651.41	752.51	+15.52%
1024	8	642.02	755.21	+17.63%
4096	1	630.43	719.84	+14.18%
4096	4	624.16	716.51	+14.80%
4096	8	611.04	705.75	+15.50%
8192	1	591.15	691.77	+17.02%
8192	4	588.31	683.29	+16.15%
8192	8	575.49	652.93	+13.46%

./build/bin/llama-batched-bench -m <granite-4.0-h-1b-Q8_0.gguf> -c 131072 -b 2048 -ub 2048 -npp 1024,4096,8192 -ntg 128 -npl 1,4,8 -ngl 99

I believe the inconsistency in throughput is due to throttling (because of the lack of cooling) on the device.

On this MacBook Air, I also initially ran the noop experiment mentioned in PR #16982 to find the theoretical limit of prefill throughput by removing all compute cost in the ssm scan kernel.

Prompt tokens	Parallel prompts	Baseline (t/s)	MMA (t/s)	Noop ceiling (t/s)	Ceiling ratio	Captured
1024	1	668.83	673.09	979.93	1.47x	1%
1024	4	651.41	752.51	966.58	1.48x	32%
1024	8	642.02	755.21	944.28	1.47x	37%
4096	1	630.43	719.84	922.60	1.46x	31%
4096	4	624.16	716.51	914.41	1.46x	32%
4096	8	611.04	705.75	886.46	1.45x	34%
8192	1	591.15	691.77	854.09	1.44x	38%
8192	4	588.31	683.29	829.82	1.41x	39%
8192	8	575.49	652.93	797.71	1.39x	35%

So the new kernel was able to improve throughput over the baseline kernel by ~35% of the theoretical limit of improvement.

dpantaleoni · 2026-08-24T00:26:18Z

@forforever73 thank you for the thorough review and requested changes! I believe they are all complete now and I also rebased

…ctor

dpantaleoni · 2026-08-26T02:57:24Z

@forforever73 @ggerganov Just a check in to see if I could get a review after making the requested changes. I rebased again to accommodate the ggml-metal.metal refactor.

forforever73 · 2026-08-26T06:48:08Z

On the sequential path the function constant caused a regression, so I switched it to a template.

And test on my M4 Max

Prompt tokens	Parallel prompts	Baseline PP (t/s)	PR PP (t/s)	Delta PP	Baseline TG (t/s)	PR TG (t/s)	Delta TG
1024	1	2699.28	2805.40	+3.93%	120.91	121.48	+0.47%
1024	4	2766.71	2907.77	+5.10%	184.69	185.79	+0.60%
1024	8	2769.68	2925.28	+5.62%	201.46	201.48	+0.01%
4096	1	2690.70	2826.57	+5.05%	119.64	119.60	-0.03%
4096	4	2702.04	2838.20	+5.04%	182.77	183.20	+0.24%
4096	8	2702.52	2838.89	+5.05%	198.84	199.08	+0.12%
8192	1	2666.94	2798.30	+4.93%	119.03	118.84	-0.16%
8192	4	2670.34	2803.82	+5.00%	180.28	180.47	+0.11%
8192	8	2671.50	2804.47	+4.98%	195.89	196.01	+0.06%

…l-org#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>

github-actions Bot added testing Everything test related ggml changes relating to the ggml tensor library for machine learning Apple Metal https://en.wikipedia.org/wiki/Metal_(API) labels Aug 5, 2026

gabe-l-hart reviewed Aug 7, 2026

View reviewed changes

ggerganov self-assigned this Aug 7, 2026

dpantaleoni marked this pull request as ready for review August 7, 2026 22:02

dpantaleoni requested review from a team and ggerganov as code owners August 7, 2026 22:02

dpantaleoni force-pushed the ssd-metal branch from 2f4fec1 to 79ad37a Compare August 20, 2026 15:40

forforever73 requested changes Aug 22, 2026

View reviewed changes

Comment thread ggml/src/ggml-metal/ggml-metal-ops.cpp Outdated

Comment thread ggml/src/ggml-metal/ggml-metal.metal Outdated

Comment thread ggml/src/ggml-metal/ggml-metal-device.cpp Outdated

forforever73 reviewed Aug 22, 2026

View reviewed changes

dpantaleoni force-pushed the ssd-metal branch from 79ad37a to 3ab8dd4 Compare August 24, 2026 00:23

dpantaleoni and others added 14 commits August 25, 2026 19:14

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

c375a5a

metal: drop scalar SSD path; MMA + sequential tail

778f64c

drop WIP ssm scan test noise

6edd9bb

remove state_from_dst and rename CS and NSG constants

895c55e

remove unrelated added whitespace padding

b978c8f

added clarity to mma_tokens calculation

a4725e4

added clarity to use_mma bool checks

15668be

added comments to metal ssd op constants for clarity

5b2d152

reserve K tokens for sequential kernel rollback snapshots

4820ff8

reset concurrency between mma and seq tail

559a766

remove print args no longer used

61922f6

fixed comment to no longer point to specific line

0c4872d

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

0bc1e22

added changes to new ssm.metal for rebase after ggml-metal.metal refa…

f19941d

…ctor

dpantaleoni force-pushed the ssd-metal branch from 3ab8dd4 to f19941d Compare August 26, 2026 02:52

specialize ssm_scan tail with a template instead of a function constant

3954a13

forforever73 approved these changes Aug 26, 2026

View reviewed changes

ggerganov approved these changes Aug 26, 2026

View reviewed changes

ggerganov added the merge ready A maintainer can use this label to indicate that they consider the changes final and ready to merge. label Aug 26, 2026

ggerganov merged commit 11cd988 into ggml-org:master Aug 26, 2026
33 of 35 checks passed


		GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
		const int64_t mma_tokens = n_seq_tokens / CHUNK * CHUNK;

Uh oh!

Conversation

dpantaleoni commented Aug 5, 2026

Overview

Additional information

Requirements

Uh oh!

gabe-l-hart left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ggerganov commented Aug 20, 2026

Uh oh!

dpantaleoni commented Aug 20, 2026

Uh oh!

dpantaleoni commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dpantaleoni commented Aug 24, 2026

Uh oh!

dpantaleoni commented Aug 26, 2026

Uh oh!

forforever73 commented Aug 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dpantaleoni commented Aug 20, 2026 •

edited

Loading