From b7ea8b19bf387e1023ce8f86b9e090cd1faed92c Mon Sep 17 00:00:00 2001 From: lnigam Date: Fri, 14 Aug 2026 19:50:40 +0530 Subject: [PATCH] ggml : recurrent state rollback for ggml_ssm_scan (llama/26623) * Initial changes for Recurrent state rollback for nemotron for cpu and cuda * Removing CPU RS rollback. Will enable it in subsequent PRs * addition of test case * Removing assert and calling runtime API to check if op is supported * removing extra API and updating the call sites for K * replace static cuda detection to runtime fused_op api * address review comments and fallback when SSM rollback not supprted * Adding changes for supporting RS-rollback in CPU. Also added test-backend-ops for cpu and cuda * removing memory manipulation as rs rollback is now supported in CPU * removing the static probe which is not needed now * correcting the format * address review comments * enabling test for all the backends, unsupported backends will fallback to CPU * Apply suggestions from code review Co-authored-by: Georgi Gerganov * choose different graph based on the result of fused_ssm_op is supported or not and also handled memory->n_rs_seq >1 case incase of op is not supported * Support K > 1 in ssm_scan for all backends * Fix CI Issues --------- Co-authored-by: Georgi Gerganov Co-authored-by: Gaurav Garg --- ggml/include/ggml.h | 3 ++- ggml/src/ggml-cpu/ggml-cpu.cpp | 2 ++ ggml/src/ggml-cpu/ops.cpp | 12 ++++++++- ggml/src/ggml-cuda/ggml-cuda.cu | 6 +++++ ggml/src/ggml-cuda/ssm-scan.cu | 27 ++++++++++++++----- .../src/ggml-et/et-kernels/src/ssm_scan_f32.c | 15 +++++++++-- ggml/src/ggml-et/ggml-et-ops.cpp | 1 + ggml/src/ggml-et/ggml-et-ops.h | 3 ++- ggml/src/ggml-metal/ggml-metal-device.m | 3 ++- ggml/src/ggml-metal/ggml-metal-impl.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 5 ++++ ggml/src/ggml-metal/ggml-metal.metal | 8 ++++++ ggml/src/ggml-sycl/ssm_scan.cpp | 22 ++++++++++++--- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 7 +++-- .../ggml-vulkan/vulkan-shaders/ssm_scan.comp | 10 +++++++ ggml/src/ggml-webgpu/ggml-webgpu.cpp | 1 + .../ggml-webgpu/wgsl-shaders/ssm_scan.wgsl | 11 ++++++++ ggml/src/ggml.c | 10 +++++-- 18 files changed, 127 insertions(+), 20 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 5cb49d0ee..c2ccd9725 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2459,7 +2459,8 @@ extern "C" { struct ggml_tensor * A, struct ggml_tensor * B, struct ggml_tensor * C, - struct ggml_tensor * ids); + struct ggml_tensor * ids, + int64_t K); // partition into non-overlapping windows with padding if needed // example: diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index c0c9aa3cf..8cece71f1 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -472,6 +472,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_CONV_2D: return ggml_is_contiguous(op->src[0]); + case GGML_OP_SSM_SCAN: + return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; default: return true; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 25bb74383..001e1ae85 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9644,11 +9644,13 @@ static void ggml_compute_forward_ssm_scan_f32( const int64_t ng = src4->ne[1]; const int64_t nt = src1->ne[2]; // number of tokens per sequence const int64_t ns = src1->ne[3]; // number of sequences in the batch + const int64_t K = ggml_get_op_params_i32(dst, 0); // can't use ggml_nbytes because src1 is not necessarily contiguous const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1); - GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst)); + GGML_ASSERT(K >= 1); + GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*ns == ggml_nelements(dst)); GGML_ASSERT(src0->nb[0] == sizeof(float)); GGML_ASSERT(src1->nb[0] == sizeof(float)); GGML_ASSERT(src2->nb[0] == sizeof(float)); @@ -9657,6 +9659,7 @@ static void ggml_compute_forward_ssm_scan_f32( GGML_ASSERT(src5->nb[0] == sizeof(float)); GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); GGML_ASSERT(nh % ng == 0); + GGML_ASSERT(src3->ne[0] == 1 || K == 1); // heads per thread const int dh = (nh + nth - 1)/nth; @@ -9831,6 +9834,13 @@ static void ggml_compute_forward_ssm_scan_f32( } } } + const int64_t slot = nt - 1 - i2; + if (K > 1 && slot > 0 && slot < K) { + float * s_snapshot = (float *) ((char *) dst->data + s_off + (slot*ns + i3)*(src0->nb[3])); + for (int h = ih0; h < ih1; ++h) { + memcpy((char *) s_snapshot + h*src0->nb[2], (char *) s + h*src0->nb[2], src0->nb[2]); + } + } // use the output as the source when it's not the first token-wise iteration s0 = s; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index cb7e9330c..598f3228c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5189,11 +5189,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); case GGML_OP_SSM_SCAN: { + const int32_t K = ggml_get_op_params_i32(op, 0); + if (op->src[3]->ne[0] == 1) { // Mamba2 // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0) return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0; } else { + if (K > 1) { + return false; + } + // Mamba // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1) return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1; diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index f3418c2af..ef342f01f 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -149,7 +149,7 @@ __global__ void __launch_bounds__(d_state, 1) const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3, const int src2_nb1, const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, - const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) { + const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok, const int64_t K) { const float * GGML_CUDA_RESTRICT src0 = src0_ptr; const float * GGML_CUDA_RESTRICT src1 = src1_ptr; const float * GGML_CUDA_RESTRICT src2 = src2_ptr; @@ -217,6 +217,16 @@ __global__ void __launch_bounds__(d_state, 1) if (lane == 0) { y_warp[i * stride_y] = state_sum; } + + // Slot 0 is the final state written below; slots 1..K-1 are rollback snapshots. + const int64_t slot = n_tok - 1 - i; + if (K > 1 && slot > 0 && slot < K) { + float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * gridDim.y + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state); +#pragma unroll + for (int j = 0; j < c_factor; j++) { + s_snapshot_warp[WARP_SIZE * j + lane] = state[j]; + } + } } // write back the state @@ -232,7 +242,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim, const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq, - cudaStream_t stream) { + const int64_t K, cudaStream_t stream) { // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition! if (src3_nb1 == sizeof(float)) { // Mamba-2 @@ -245,7 +255,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa ggml_cuda_kernel_launch(ssm_scan_f32_group<128/WARP_SIZE, 128>, 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_off, n_head, head_dim, n_group, n_tok); + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K); } else if (d_state == 256) { // Falcon-H1 constexpr int threads = 256; constexpr int num_warps = threads/WARP_SIZE; @@ -255,12 +265,13 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa ggml_cuda_kernel_launch(ssm_scan_f32_group<256/WARP_SIZE, 256>, 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_off, n_head, head_dim, n_group, n_tok); + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K); } else { GGML_ABORT("doesn't support d_state!=(128 or 256)."); } } else { // Mamba-1 + GGML_ASSERT(K == 1); constexpr int threads = 128; GGML_ASSERT(n_head % threads == 0); GGML_ASSERT(head_dim == 1); @@ -769,10 +780,12 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t ng = src4->ne[1]; // n_group const int64_t n_t = src1->ne[2]; // number of tokens per sequence const int64_t n_s = src1->ne[3]; // number of sequences in the batch + const int32_t K_param = ggml_get_op_params_i32(dst, 0); + const int64_t K = K_param > 0 ? K_param : 1; const int64_t s_off = ggml_nelements(src1) * sizeof(float); - GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*n_s == ggml_nelements(dst)); + GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*n_s == ggml_nelements(dst)); GGML_ASSERT(src0->nb[0] == sizeof(float)); GGML_ASSERT(src1->nb[0] == sizeof(float)); GGML_ASSERT(src2->nb[0] == sizeof(float)); @@ -780,6 +793,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src4->nb[0] == sizeof(float)); GGML_ASSERT(src5->nb[0] == sizeof(float)); GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); + GGML_ASSERT(src3->ne[0] == 1 || K == 1); const float * src0_d = (const float *) src0->data; const float * src1_d = (const float *) src1->data; @@ -814,6 +828,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const bool is_mamba2 = (src3->nb[1] == sizeof(float)); const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && K == 1 && n_t <= SSM_SSD_MAX_TOKENS && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING @@ -841,5 +856,5 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], - s_off, nc, nr, nh, ng, n_t, n_s, stream); + s_off, nc, nr, nh, ng, n_t, n_s, K, stream); } diff --git a/ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c b/ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c index c114e9981..82ac4309c 100644 --- a/ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c +++ b/ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c @@ -12,7 +12,8 @@ struct ggml_et_ssm_scan_params { struct ggml_tensor src4; // B: [d_state, n_group, n_seq_tokens, n_seqs] struct ggml_tensor src5; // C: [d_state, n_group, n_seq_tokens, n_seqs] struct ggml_tensor src6; // ids: [n_seqs] i32 - struct ggml_tensor dst; // packed [y, final_state] + struct ggml_tensor dst; // packed [y, states] + int32_t K; }; static inline float softplus_f32(float x) { @@ -72,6 +73,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) { const int64_t n_seq_tokens = src1->ne[2]; const int64_t n_seqs = src1->ne[3]; const int64_t y_elems = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3]; + const int64_t K = params->K; if (src0->nb[0] != sizeof(float) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(float) || src3->nb[0] != sizeof(float) || src4->nb[0] != sizeof(float) || src5->nb[0] != sizeof(float) || @@ -79,7 +81,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) { return -1; } - if (n_group <= 0 || n_head % n_group != 0) { + if (K < 1 || n_group <= 0 || n_head % n_group != 0) { return -1; } @@ -260,6 +262,15 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) { sumf += st * C_row[state_idx]; } + const int64_t slot = n_seq_tokens - 1 - token_idx; + if (slot > 0 && slot < K) { + float * state_snapshot = + (float *) ((char *) state_dst + (size_t) slot * n_seqs * src0->nb[3]); + for (int64_t i = 0; i < d_state; ++i) { + state_snapshot[i] = state_dst[i]; + } + } + dst_data[seq_idx * (n_seq_tokens * n_head * head_dim) + token_idx * (n_head * head_dim) + head_idx * head_dim + dim_idx] = sumf; } diff --git a/ggml/src/ggml-et/ggml-et-ops.cpp b/ggml/src/ggml-et/ggml-et-ops.cpp index 6c80fe8ac..7871d5240 100644 --- a/ggml/src/ggml-et/ggml-et-ops.cpp +++ b/ggml/src/ggml-et/ggml-et-ops.cpp @@ -2064,6 +2064,7 @@ bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_te params.src5 = *node->src[5]; params.src6 = *node->src[6]; params.dst = *node; + params.K = ggml_get_op_params_i32(node, 0); bool kernel_result = ggml_et_launch_kernel(dev_ctx, "ssm_scan_f32", ¶ms, sizeof(params), 0xFFFFFFFF); diff --git a/ggml/src/ggml-et/ggml-et-ops.h b/ggml/src/ggml-et/ggml-et-ops.h index 2c7ca7ece..032f7a263 100644 --- a/ggml/src/ggml-et/ggml-et-ops.h +++ b/ggml/src/ggml-et/ggml-et-ops.h @@ -218,7 +218,8 @@ struct ggml_et_ssm_scan_params { ggml_tensor src4; // B: [d_state, n_group, n_seq_tokens, n_seqs] ggml_tensor src5; // C: [d_state, n_group, n_seq_tokens, n_seqs] ggml_tensor src6; // ids: [n_seqs] i32 - ggml_tensor dst; // [y, final_state] packed output from ggml_ssm_scan() + ggml_tensor dst; // [y, states] packed output from ggml_ssm_scan() + int32_t K; }; struct ggml_et_rwkv_wkv6_params { diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index b70816c32..312b00dc4 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1376,9 +1376,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op->src[2]) && ggml_is_contiguous_rows(op->src[3]); - case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; + case GGML_OP_SSM_CONV: + return has_simdgroup_reduction; case GGML_OP_RWKV_WKV6: case GGML_OP_RWKV_WKV7: return true; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index cf32c5c5b..1f6e8c48b 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -880,6 +880,7 @@ typedef struct { int64_t n_group; int64_t n_seq_tokens; int64_t n_seqs; + int64_t K; uint64_t s_off; uint64_t nb00; uint64_t nb01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 6d324056d..b7f9b2d0d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1710,6 +1710,10 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { const int64_t n_group = ne41; const int64_t n_seq_tokens = ne12; const int64_t n_seqs = ne13; + const int64_t K = ggml_get_op_params_i32(op, 0); + + GGML_ASSERT(K >= 1); + GGML_ASSERT(ggml_nelements(op->src[1]) + K*d_state*d_inner*n_head*n_seqs == ggml_nelements(op)); ggml_metal_kargs_ssm_scan args = { /*.d_state =*/ d_state, @@ -1718,6 +1722,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { /*.n_group =*/ n_group, /*.n_seq_tokens =*/ n_seq_tokens, /*.n_seqs =*/ n_seqs, + /*.K =*/ K, /*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float), /*.nb00 =*/ nb00, /*.nb01 =*/ nb01, diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index b38b23edc..243c997fc 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2429,6 +2429,8 @@ kernel void kernel_ssm_scan_f32( const int32_t nh = args.n_head; const int32_t ng = args.n_group; const int32_t n_t = args.n_seq_tokens; + const int32_t n_s = args.n_seqs; + const int32_t K = args.K; const int32_t s_off = args.s_off; @@ -2487,6 +2489,12 @@ kernel void kernel_ssm_scan_f32( // recurse s0 = s; + const int32_t slot = n_t - 1 - (i2 + t); + if (slot > 0 && slot < K) { + device float * s_snapshot = (device float *) ((device char *) s_buff + (int64_t) slot*n_s*args.nb03); + s_snapshot[i] = s; + } + B += args.ns42; C += args.ns52; } diff --git a/ggml/src/ggml-sycl/ssm_scan.cpp b/ggml/src/ggml-sycl/ssm_scan.cpp index ae6529813..7fceb85d2 100644 --- a/ggml/src/ggml-sycl/ssm_scan.cpp +++ b/ggml/src/ggml-sycl/ssm_scan.cpp @@ -10,6 +10,7 @@ static void ssm_scan_f32_group( const int src2_nb1, const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok, + const int64_t K, const sycl::nd_item<2> & item) { const int lane = item.get_local_id(1) % WARP_SIZE; @@ -64,6 +65,15 @@ static void ssm_scan_f32_group( if (lane == 0) { y_warp[i * stride_y] = state_sum; } + + const int64_t slot = n_tok - 1 - i; + if (K > 1 && slot > 0 && slot < K) { + float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * item.get_group_range(0) + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state); +#pragma unroll + for (int j = 0; j < c_factor; j++) { + s_snapshot_warp[WARP_SIZE * j + lane] = state[j]; + } + } } #pragma unroll @@ -79,6 +89,7 @@ static void ssm_scan_f32_sycl( const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim, const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq, + const int64_t K, dpct::queue_ptr stream) { // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition! @@ -94,7 +105,7 @@ static void ssm_scan_f32_sycl( ssm_scan_f32_group<128 / WARP_SIZE, 128>( 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_off, n_head, head_dim, n_group, n_tok, item); + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K, item); }); } else if (d_state == 256) { constexpr int threads = 256; @@ -107,7 +118,7 @@ static void ssm_scan_f32_sycl( ssm_scan_f32_group<256 / WARP_SIZE, 256>( 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_off, n_head, head_dim, n_group, n_tok, item); + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K, item); }); } else { GGML_ABORT("ssm_scan: unsupported d_state (must be 128 or 256)"); @@ -133,9 +144,12 @@ inline void ggml_sycl_op_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor * const int64_t ng = src4->ne[1]; const int64_t n_t = src1->ne[2]; const int64_t n_s = src1->ne[3]; + const int64_t K = ggml_get_op_params_i32(dst, 0); const int64_t s_off = ggml_nelements(src1) * sizeof(float); - GGML_ASSERT(ggml_nelements(src1) + nc * nr * nh * n_s == ggml_nelements(dst)); + GGML_ASSERT(K >= 1); + GGML_ASSERT(ggml_nelements(src1) + K * nc * nr * nh * n_s == ggml_nelements(dst)); + GGML_ASSERT(src3->ne[0] == 1 || K == 1); dpct::queue_ptr stream = ctx.stream(); SYCL_CHECK(ggml_sycl_set_device(ctx.device)); @@ -147,7 +161,7 @@ inline void ggml_sycl_op_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor * static_cast(src6->data), static_cast(dst->data), src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], - s_off, nc, nr, nh, ng, n_t, n_s, stream); + s_off, nc, nr, nh, ng, n_t, n_s, K, stream); } void ggml_sycl_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c815d4ff9..ff4a33904 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1861,6 +1861,7 @@ struct vk_op_ssm_scan_push_constants { uint32_t nb42, nb43, nb52, nb53; uint32_t s_off; uint32_t n_head, d_head, n_group, n_tok; + uint32_t n_seq, K; }; struct vk_op_ssm_conv_push_constants { uint32_t nb01, nb02; @@ -12731,7 +12732,8 @@ static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, (uint32_t)src4->nb[2], (uint32_t)src4->nb[3], (uint32_t)src5->nb[2], (uint32_t)src5->nb[3], (uint32_t)s_off, - n_head, head_dim, n_group, n_tok + n_head, head_dim, n_group, n_tok, + n_seq, (uint32_t) ggml_get_op_params_i32(dst, 0) }; vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); @@ -19417,8 +19419,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * } else if (tensor->op == GGML_OP_ADD_ID) { tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]); } else if (tensor->op == GGML_OP_SSM_SCAN) { + const int32_t K = ggml_get_op_params_i32(tensor, 0); tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], - src_clone[3], src_clone[4], src_clone[5], src_clone[6]); + src_clone[3], src_clone[4], src_clone[5], src_clone[6], K); } else if (tensor->op == GGML_OP_SSM_CONV) { tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]); } else if (tensor->op == GGML_OP_ROLL) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp b/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp index c7416206d..4fecb3aa5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp @@ -33,6 +33,8 @@ layout(push_constant) uniform PushConstants { uint d_head; uint n_group; uint n_tok; + uint n_seq; + uint K; }; float softplus(float x) { @@ -114,6 +116,14 @@ void main() { if (lane == 0) { d[y_base_idx + i * stride_y] = state_sum; } + + const uint slot = n_tok - 1u - i; + if (slot > 0u && slot < K) { + const uint snapshot_base_idx = s_base_idx + slot * n_seq * (nb03 / 4u); + [[unroll]] for (uint j = 0; j < c_factor; j++) { + d[snapshot_base_idx + SUBGROUP_SIZE * j + lane] = state[j]; + } + } } // write back the state diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 6741752b3..394aeeda2 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -1327,6 +1327,7 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx, (uint32_t) src4->ne[1], (uint32_t) src1->ne[2], (uint32_t) ggml_nelements(src1), + (uint32_t) ggml_get_op_params_i32(dst, 0), }; std::vector entries = { diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl index 2d4c4e5a0..57f012ad0 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl @@ -41,6 +41,7 @@ struct Params { n_seq_tokens: u32, y_elems: u32, + K: u32, }; @group(0) @binding(0) var s_in: array; @@ -123,6 +124,7 @@ fn main( let head_seq = wg_linear / params.d_inner; let ir = head_seq % params.n_head; let i3 = head_seq / params.n_head; + let n_seqs = params.y_elems / (params.n_seq_tokens * params.n_head * params.d_inner); let state_slot = read_state_slot(i3); let g = ir / (params.n_head / params.n_group); @@ -179,6 +181,15 @@ fn main( #endif s_prev = s; + let slot = params.n_seq_tokens - 1u - token; + if (slot > 0u && slot < params.K) { + let snapshot_idx = + params.offset_dst + params.y_elems + tid + i1 * params.d_state + + ir * (params.d_state * params.d_inner) + + (slot * n_seqs + i3) * (params.d_state * params.d_inner * params.n_head); + dst[snapshot_idx] = s; + } + #ifdef USE_SUBGROUP_REDUCTION #ifdef XBC_OVERLAP let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx)); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index da7f3a5f2..d0d369c41 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5588,7 +5588,10 @@ struct ggml_tensor * ggml_ssm_scan( struct ggml_tensor * A, struct ggml_tensor * B, struct ggml_tensor * C, - struct ggml_tensor * ids) { + struct ggml_tensor * ids, + int64_t K) { + GGML_ASSERT(K >= 1); + GGML_ASSERT(K <= INT32_MAX); GGML_ASSERT(ggml_is_contiguous(s)); GGML_ASSERT(ggml_is_contiguous(dt)); GGML_ASSERT(ggml_is_contiguous(A)); @@ -5625,11 +5628,12 @@ struct ggml_tensor * ggml_ssm_scan( if (A->ne[0] != 1) { // Mamba-1 has more granular decay factors GGML_ASSERT(A->ne[0] == d_state); + GGML_ASSERT(K == 1); } } // concatenated y + ssm_states - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + K*s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); result->op = GGML_OP_SSM_SCAN; result->src[0] = s; @@ -5640,6 +5644,8 @@ struct ggml_tensor * ggml_ssm_scan( result->src[5] = C; result->src[6] = ids; + ggml_set_op_params_i32(result, 0, (int32_t) K); + return result; }