Commit 998200491 for llama.cpp
commit 998200491327347eed85f007e4de7a1a7c871453
Author: sargentmajord <jasonlather13@gmail.com>
Date: Sun Nov 8 07:48:20 2026 -0600
ggml-cuda : add d_state=64 fast path to SSM_SCAN kernel (#31818)
Zamba2 and similar architectures use a state size of 64 in their
recurrent-scan step, which fell outside the GPU kernel's existing
fast paths (96, 128, 256). Every scan step was falling back to CPU,
bottlenecking generation to 6-7 tok/s against the 50-70 tok/s band
other architectures achieve on the same hardware.
Add a 4th template instantiation for d_state=64, following the
existing pattern; no change to the underlying math. Verified against
the existing CPU-reference correctness harness with 3 new test cases
built from Zamba2's real hyperparameters (28/28 passing).
Measured ~35% throughput improvement (11 -> 15.1 tok/s) on Zamba2-7B-
Instruct. Full parity with other architectures not yet reached - a
separate, unrelated fraction of scan calls (~28%) still falls back to
CPU for a cause not yet diagnosed.
Assisted-by: Claude Sonnet 5
Co-authored-by: Sargentmajord <331229907+sargentmajord@users.noreply.github.com>
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index a49c5247a..162a7d28d 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -6252,9 +6252,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
if (op->src[3]->ne[0] == 1) {
// Mamba2
- // (kernel only supports (d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
+ // (kernel only supports (d_state == 64 || d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
const int64_t d_state = op->src[0]->ne[0];
- return (d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
+ return (d_state == 64 || d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
} else {
if (K > 1) {
return false;
diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu
index dcc199f73..b88c9d6ea 100644
--- a/ggml/src/ggml-cuda/ssm-scan.cu
+++ b/ggml/src/ggml-cuda/ssm-scan.cu
@@ -256,7 +256,17 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
// NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
if (src3_nb1 == sizeof(float)) {
// Mamba-2
- if (d_state == 96) {
+ if (d_state == 64) { // Zamba2
+ constexpr int threads = 64;
+ constexpr int num_warps = threads/WARP_SIZE;
+
+ const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1);
+ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
+ ggml_cuda_kernel_launch(ssm_scan_f32_group<64/WARP_SIZE, 64>, launch_params,
+ src0, src1, src2, src3, src4, src5, src6, dst,
+ src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
+ src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_base, s_slot_bytes, n_head, head_dim, n_group, n_tok, K);
+ } else if (d_state == 96) {
constexpr int threads = 96;
constexpr int num_warps = threads/WARP_SIZE;
@@ -287,7 +297,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_base, s_slot_bytes, n_head, head_dim, n_group, n_tok, K);
} else {
- GGML_ABORT("doesn't support d_state!=(96, 128 or 256).");
+ GGML_ABORT("doesn't support d_state!=(64, 96, 128 or 256).");
}
} else {
// Mamba-1
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 852b62701..bfbd5322c 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11425,6 +11425,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 16, 1, 1024, 1, 32, 4)); // Mamba-1
+ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 64, 64, 112, 2, 32, 4)); // Zamba2 (real hparams: d_inner=7168, dt_rank=112, n_group=2 -> head_dim=7168/112=64)
+ test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 64, 64, 112, 2, 1, 1)); // Zamba2 decode (single-token, matches real per-token decode workload)
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 96, 64, 128, 8, 1, 1)); // Nemotron-3-Puzzle decode (scan path, unused warp in last block)
test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 96, 64, 16, 8, 300, 2, false, /*K=*/1, /*weak_decay=*/true)); // d_state=96 SSD multi-chunk (partial 2nd chunk, 2 seqs)
test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 96, 64, 16, 8, 8, 2, /*K=*/3)); // d_state=96 rollback snapshots match prefix states
@@ -11446,6 +11448,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 128, 64, 16, 2, 4, 1, 4));
test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 128, 64, 16, 2, 1, 1, 4)); // n_seq_tokens < K
test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 128, 64, 16, 2, 8, 1, 3)); // n_seq_tokens > K
+ test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 64, 64, 112, 2, 4, 1, 4)); // Zamba2
test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 96, 64, 16, 2, 4, 1, 4));
test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 256, 64, 8, 2, 4, 1, 4));
test_cases.emplace_back(new test_ssm_scan_cache_fusion(GGML_TYPE_F32, 128, 64, 16, 2, 1, 1, 1)); // K == 1, final state only