ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration by BLSharda · Pull Request #22675 · ggml-org/llama.cpp

BLSharda · 2026-05-04T10:35:25Z

Overview

Adding SSD (State Space Duality) matmul path to the CUDA SSM scan for Mamba-2 prefill acceleration. For sequences >64 tokens, replaces the sequential O(T*N) scan with chunked parallel matmuls using cuBLAS batched GEMMs and a fused MMA tensor-core kernel that computes the causal decay mask on-the-fly in shared memory.

Benchmarks on RTX PRO 6000 Blackwell (CUDA 13.2):
-- mamba2-2.7B: +15-18% PP (pp128–pp32768)
-- Nemotron-Nano-4B (hybrid): +11-16% PP
-- Nemotron-Nano-9B (hybrid): +12-16% PP
-- Nemotron-Super-120B-A12B (MoE): +3-8% PP
-- TG unchanged (scan kernel still used for n_tokens ≤ 64). Perplexity verified identical on all models.

Nemotron-Nano-9B (hybrid) prompt processing on various GPUs

Hardware	ISL=1000	ISL=2000	ISL=4000	ISL=8000	ISL=16000	Max Δ%
RTX 5090	8.7%	8.2%	8.7%	10.9%	11.1%	11.1%
RTX 5080	2.3%	8.7%	10.2%	10.3%	10.4%	10.4%
RTX 5070	7.9%	9.1%	10.1%	10.6%	10.8%	10.8%
RTX 4090	7.4%	5.5%	12.6%	9.5%	9.7%	12.6%
RTX 2080 Ti	4.0%	2.1%	6.5%	5.9%	6.4%	6.5%
RTX 3090	4.6%	6.9%	13.0%	11.3%	9.5%	13.0%
RTX 5070 Ti	10.1%	9.6%	8.9%	10.8%	10.3%	10.8%

Additional information

Only ggml/src/ggml-cuda/ssm-scan.cu is modified. The existing sequential scan kernel is untouched and used as fallback for short sequences, non-NVIDIA GPUs, and pre-Turing architectures.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES - AI was used for kernel refactoring and optimizations from my initial design.

BLSharda · 2026-05-04T10:36:44Z

@ORippler and @gaugarg-nv for review.

BLSharda · 2026-05-12T05:36:27Z

The failure in CI / ubuntu-24-webgpu (pull_request) is unrelated to my change.

gaugarg-nv

I think for the perf data, you can keep the ISL constant to something like 16384, but change the ubatch and batch size from 32 till 8192. Something like-

llama-bench -m <model> -fa 1 -p 16384 -b 8192 -ub 32,64,128,256,512,1024,2048,4096,8192 -r 3

BLSharda · 2026-05-15T11:52:55Z

I think for the perf data, you can keep the ISL constant to something like 16384, but change the ubatch and batch size from 32 till 8192. Something like-

llama-bench -m <model> -fa 1 -p 16384 -b 8192 -ub 32,64,128,256,512,1024,2048,4096,8192 -r 3

Sure, below is the perf table for Nemotron 9b nano model

ubatch	Branch path	Master (t/s)	Branch SSD (t/s)	Δ
32	SCAN	2833.66	2826.70	−0.25%
64	SCAN	4071.67	4033.43	−0.94%
128	SSD	5170.13	5712.37	10.49%
256	SSD	6009.37	6828.75	13.64%
512	SSD	6539.02	7509.39	14.84%
1024	SSD	6701.96	7789.13	16.22%
2048	SSD	6601.55	7744.91	17.32%
4096	SSD	6357.18	7436.90	16.99%
8192	SSD	5951.29	6887.30	15.73%

Command used on PC with RTX 6000 Blackwell -
llama-bench.exe -m <model.gguf> -ngl 99 -fa 1 -p 16384 -b 8192 -ub 32,64,128,256,512,1024,2048,4096,8192 -r 3

BLSharda · 2026-05-20T12:11:20Z

I found some more optimization opportunities in CUDA kernel working on it and will update here soon. For now, converting to draft in the meantime.

BLSharda · 2026-06-09T14:03:30Z

My latest optimization improved performance, although the baseline appears to have shifted. The updated performance relative to baseline is shown below.

ubatch sweep - RTX 5090

n_ubatch	Baseline (scan)	Feature (SSD)	Delta	Path
32	3,514 ±76	3,547 ±57	+0.9%	scan
64	5,021 ±68	4,994 ±27	-0.5%	scan
128	5,981 ±146	6,563 ±35	+9.7%	scan
256	6,759 ±26	7,782 ±11	+15.1%	SSD
512	7,502 ±71	8,640 ±7	+15.2%	SSD
1024	7,736 ±41	8,967 ±4	+15.9%	SSD
2048	7,658 ±7	8,860 ±10	+15.7%	SSD
4096	7,444 ±16	8,634 ±5	+16.0%	SSD
8192	6,933 ±4	7,894 ±8	+13.9%	SSD

llama-bench.exe -m <nvidia_NVIDIA-Nemotron-Nano-9B-v2-Q4_K_M.gguf> -ngl 99 -fa 1 -p 16384 -b 8192 -ub 32,64,128,256,512,1024,2048,4096,8192 -r 3

BLSharda · 2026-06-11T17:54:57Z

ubatch sweep - DGX Spark

n_ubatch	Baseline (scan)	Feature (SSD)	Delta	Path
32	806.1 ±0.8	807.6 ±1.0	+0.2%	scan
64	1330.5 ±0.5	1335.2 ±0.7	+0.4%	scan
128	1822.6 ±2.1	1821.1 ±3.5	-0.1%	scan
256	1983.0 ±2.9	2040.1 ±3.2	+2.9%	SSD
512	2146.9 ±4.7	2215.4 ±2.6	+3.2%	SSD
1024	2055.4 ±1.3	2143.9 ±1.4	+4.3%	SSD
2048	2026.3 ±0.4	2131.0 ±2.2	+5.2%	SSD
4096	1906.2 ±1.2	2007.9 ±0.9	+5.3%	SSD
8192	1618.4 ±0.7	1694.4 ±1.0	+4.7%	SSD

llama-bench.exe -m <nvidia_NVIDIA-Nemotron-Nano-9B-v2-Q4_K_M.gguf> -ngl 99 -fa 1 -p 16384 -b 8192 -ub 32,64,128,256,512,1024,2048,4096,8192 -r 3

ORippler

Did we run any studies on TF32 vs. BF16 vs. F16 w.r.t numerical precision and speed?

At first glance, the changes required in test-backend-ops to get the new implementation to pass correctness are a bit concerning. Please elaborate on those seemingly magic bounds.

BLSharda · 2026-06-16T12:48:21Z

Did we run any studies on TF32 vs. BF16 vs. F16 w.r.t numerical precision and speed?

At first glance, the changes required in test-backend-ops to get the new implementation to pass correctness are a bit concerning. Please elaborate on those seemingly magic bounds.

I compared FP16 and BF16 for precision and BF16 shows significantly higher error than FP16.
Accuracy impact of baseline (scan) vs Feature (SSD) across Nemotron nano 4B, 9B and 30B all three models is negligible:

Model	PPL (Wikitext-2)	MMLU-Pro discriminative	MMLU-Pro CoT
4B	+0.012%	+0.108 pp	-0.25 pp
9B	+0.007%	+0.066 pp	+0.075 pp
30B	+0.091%	-0.033 pp	-0.125 pp

Also, I have now cleaned up test-backend-ops so that only one tensor has negative values, and I slightly adjusted the limit to better match FP16 matmul behavior.

gaugarg-nv · 2026-07-20T15:48:40Z

+    ggml_cuda_pool_alloc<float>    dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq);
+    ggml_cuda_pool_alloc<float>    cs_buf(ctx.pool(), n_tok * n_head * n_seq);
+    ggml_cuda_pool_alloc<float>    CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq);
+    ggml_cuda_pool_alloc<matmul_t> X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq);
+    ggml_cuda_pool_alloc<matmul_t> B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq);
+    ggml_cuda_pool_alloc<float>    C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq);


Do you have data on how much VRAM does this scratch space take for different ubatch sizes?

Below is the VRAM table for Nemotron Nano 4B, 9B, and Nemotron 3 Super models

ub 4B (n_head=96) 9B (n_head=128) 120B (n_head=128) path

256 35.9 MB 47.2 MB 46.2 MB SSD

512 36.1 MB 47.5 MB 46.5 MB SSD

1024 36.5 MB 48.0 MB 47.0 MB SSD

2048 37.2 MB 49.0 MB 48.0 MB SSD

4096 38.8 MB 51.0 MB 50.0 MB SSD

8192 41.8 MB 55.0 MB 54.0 MB SSD

Note: The scratch is per-layer but pool-allocated, so the peak VRAM footprint is one layer's worth at a time.

gaugarg-nv · 2026-07-20T15:56:50Z

The PR looks good overall to me. Please rebase and update the PR description (which seems to be stale now).

@JohannesGaessler @am17an Can you please take a look as well?

ORippler · 2026-07-21T12:13:54Z

@BLSharda did we ever understand how #24561 managed to do FP16 GEMMs without relaxing numerical constraints in the corresponding test-backend-ops function? Was the baseline required precision lower?

BLSharda · 2026-07-21T14:35:07Z

@BLSharda did we ever understand how #24561 managed to do FP16 GEMMs without relaxing numerical constraints in the corresponding test-backend-ops function? Was the baseline required precision lower?

PR #24561 improves numerical stability by using log-space differences instead of dividing very small exponentials, which helps eliminate NaNs and underflow. It also ensures stability in tests by initializing g in the range [-20, -1e-4], so that 0 < exp(g) < 1. Without this constraint, positive g would give exp(g) > 1, causing the recurrent state to grow exponentially and potentially overflow.

For this Mamba-2 SSD PR, this change does not really apply since it already uses a stable formulation: decay = exp(A * (cs_out - cs_in)), with A < 0. This keeps values in the range and avoids overflow or 0/0 issues.

So, the implementation is already numerically stable. The minor remaining error (~2e-7) comes from FP16 quantization rather than instability and using log-space tricks is unlikely to reduce it further. Higher precision would be the more effective option, though the gap is already very small (~2e-7) and does not have a meaningful impact based on perplexity and KL divergence.

…nabled test.

…a limits.

…proves memory coalescing in ssm_ssd_prepare_dt_kernel, and boosts efficiency by merging B_weighted and C_scaled; also addresses prior review comments.

BLSharda · 2026-07-21T17:04:47Z

Perf sweep on latest changes

Nemotron-Nano-9B-v2

ub	base (scan)	branch (SSD)	speedup
128	5,404	5,351	−1% (both scan)
256	6,180	7,110	+15%
512	6,627	7,778	+17%
1k	6,814	8,152	+20%
2k	6,759	8,190	+21%
4k	6,660	8,118	+22%
8k	6,387	7,761	+22%

pp16384 tok/s (base=scan, branch=SSD at ub>128) tested on RTX 6000 Pro MaxQ.

BLSharda · 2026-07-27T12:19:35Z

The failing CI jobs appear to be pre-existing on master and are not caused by this PR.

The jobs ubuntu-22-hip-quality-check, gpu-webgpu-nvidia, gpu-webgpu-apple, and gpu-openvino-low-perf match failures already seen after the recent merged of PR #26046

The job windows (x64-vulkan, x64, Ninja Multi-Config, LLVM toolchain) is failing test test-thread-safety with exit code 0xc0000374. The same error also appears intermittently on master, including run 30240145058, so treating it as a pre-existing issue.

This PR modifies only the CUDA code path and does not affect Vulkan, HIP, WebGPU, OpenVINO, or the LLVM toolchain build configuration.

…gml-org#22675) * ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration * cuda: added SSD CICD fixes for CUDA / HIP / MUSA / MSVC. * ggml-cuda: review comments fixed. * ggml-cuda: Fuse M matrix materialization into pre_matmul kernel and enabled test. * ggml-cuda: test updates and fixes * ggml-cuda: test updates to remove hardcoding of tensor initialise data limits. * ggml-cuda: ssd minor review comment fixed. * ggml-cuda: ssd minor CICD fixed. * CUDA SSD: Fixes correctness by promoting s0_stride_seq to int64_t, improves memory coalescing in ssm_ssd_prepare_dt_kernel, and boosts efficiency by merging B_weighted and C_scaled; also addresses prior review comments. * cuda: fix sdata read-write race in prepare_dt fallback scan loop

BLSharda requested a review from a team as a code owner May 4, 2026 10:35

BLSharda mentioned this pull request May 4, 2026

ggml : add GPU support for Mamba models #6758

Open

github-actions Bot added Nvidia GPU Issues specific to Nvidia GPUs ggml changes relating to the ggml tensor library for machine learning labels May 4, 2026

gaugarg-nv reviewed May 12, 2026

View reviewed changes

Comment thread ggml/src/ggml-cuda/ssm-scan.cu Outdated

Comment thread ggml/src/ggml-cuda/ssm-scan.cu Outdated

Comment thread ggml/src/ggml-cuda/ssm-scan.cu Outdated

Comment thread ggml/src/ggml-cuda/ssm-scan.cu Outdated

BLSharda marked this pull request as draft May 20, 2026 12:11

github-actions Bot added the testing Everything test related label Jun 9, 2026

BLSharda marked this pull request as ready for review June 9, 2026 14:03

BLSharda requested a review from ggerganov as a code owner June 9, 2026 14:03

ORippler reviewed Jun 12, 2026

View reviewed changes

Comment thread ggml/src/ggml-cuda/ssm-scan.cu Outdated

Comment thread ggml/src/ggml-cuda/ssm-scan.cu

github-actions Bot added the CUDA Related to the CUDA backend label Jun 16, 2026

gaugarg-nv reviewed Jul 20, 2026

View reviewed changes

BLSharda added 8 commits July 21, 2026 21:44

ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration

8180cdb

cuda: added SSD CICD fixes for CUDA / HIP / MUSA / MSVC.

e9633ae

ggml-cuda: review comments fixed.

3f93219

ggml-cuda: Fuse M matrix materialization into pre_matmul kernel and e…

21f3745

…nabled test.

ggml-cuda: test updates and fixes

b0af8f4

ggml-cuda: test updates to remove hardcoding of tensor initialise dat…

f6ed777

…a limits.

ggml-cuda: ssd minor review comment fixed.

f35681f

ggml-cuda: ssd minor CICD fixed.

a942d1b

CUDA SSD: Fixes correctness by promoting s0_stride_seq to int64_t, im…

3c6938c

…proves memory coalescing in ssm_ssd_prepare_dt_kernel, and boosts efficiency by merging B_weighted and C_scaled; also addresses prior review comments.

BLSharda force-pushed the bhaviksharda/mamba2_SSD branch from 8e71a37 to 3c6938c Compare July 21, 2026 16:57

am17an approved these changes Jul 24, 2026

View reviewed changes

gaugarg-nv approved these changes Jul 27, 2026

View reviewed changes

gaugarg-nv requested changes Jul 28, 2026

View reviewed changes

Comment thread ggml/src/ggml-cuda/ssm-scan.cu

cuda: fix sdata read-write race in prepare_dt fallback scan loop

f96b0e5

BLSharda requested a review from gaugarg-nv July 28, 2026 11:32

gaugarg-nv approved these changes Jul 28, 2026

View reviewed changes

gaugarg-nv merged commit b62b350 into ggml-org:master Jul 28, 2026
26 of 30 checks passed

BLSharda deleted the bhaviksharda/mamba2_SSD branch August 5, 2026 08:48

dpantaleoni mentioned this pull request Aug 5, 2026

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

Merged

davetha mentioned this pull request Aug 9, 2026

Change sets 4-5: chunked SSD Mamba-2 prefill + stream-k tuning for CDNA2 (+23.5% pp16384) davetha/llama.cpp-mi210#1

Merged

ub	4B (n_head=96)	9B (n_head=128)	120B (n_head=128)	path
256	35.9 MB	47.2 MB	46.2 MB	SSD
512	36.1 MB	47.5 MB	46.5 MB	SSD
1024	36.5 MB	48.0 MB	47.0 MB	SSD
2048	37.2 MB	49.0 MB	48.0 MB	SSD
4096	38.8 MB	51.0 MB	50.0 MB	SSD
8192	41.8 MB	55.0 MB	54.0 MB	SSD

Uh oh!

Conversation

BLSharda commented May 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Additional information

Requirements

Uh oh!

BLSharda commented May 4, 2026

Uh oh!

BLSharda commented May 12, 2026

Uh oh!

gaugarg-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

BLSharda commented May 15, 2026

Uh oh!

BLSharda commented May 20, 2026

Uh oh!

BLSharda commented Jun 9, 2026

Uh oh!

BLSharda commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ORippler left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

BLSharda commented Jun 16, 2026

Uh oh!

Uh oh!

Uh oh!

gaugarg-nv Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

BLSharda Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gaugarg-nv commented Jul 20, 2026

Uh oh!

ORippler commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

BLSharda commented Jul 21, 2026

Uh oh!

BLSharda commented Jul 21, 2026

Uh oh!

BLSharda commented Jul 27, 2026

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

BLSharda commented May 4, 2026 •

edited

Loading

BLSharda commented Jun 11, 2026 •

edited

Loading

ORippler commented Jul 21, 2026 •

edited

Loading