diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 35f0c44ec421..1b170e7796de 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2461,6 +2461,17 @@ extern "C" { struct ggml_tensor * C, struct ggml_tensor * ids); + GGML_API struct ggml_tensor * ggml_ssm_scan_ext( + struct ggml_context * ctx, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids, + int64_t K); + // partition into non-overlapping windows with padding if needed // example: // a: 768 64 64 1 diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 42ec809ce521..b4c9f273c3af 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 int32_t K_param = ggml_get_op_params_i32(dst, 0); + const int64_t K = K_param > 0 ? K_param : 1; // 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(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; @@ -9678,6 +9681,8 @@ static void ggml_compute_forward_ssm_scan_f32( const float * B = (const float *) ((const char *) src4->data + i2*(src4->nb[2]) + i3*(src4->nb[3])); // {d_state, ng, nt, ns} const float * C = (const float *) ((const char *) src5->data + i2*(src5->nb[2]) + i3*(src5->nb[3])); // {d_state, ng, nt, ns} float * y = ( float *) (( char *) dst->data + i2*(nh*nr*sizeof(float)) + i3*(nt*nh*nr*sizeof(float))); // {dim, nh, nt, ns} + const int64_t slot = nt - 1 - i2; + float * s_snapshot = slot < K ? (float *) ((char *) dst->data + s_off + (slot*ns + i3)*(src0->nb[3])) : nullptr; if (src3->ne[0] == 1) { // Mamba-2 has a scalar decay factor per head; dA can be outside the state-wise loop @@ -9721,6 +9726,9 @@ static void ggml_compute_forward_ssm_scan_f32( sum = GGML_F32_VEC_FMA(sum, t0, t2); GGML_F32_VEC_STORE(s + i + j*ggml_f32_epr + ii*nc, t0); + if (s_snapshot && s_snapshot != s) { + GGML_F32_VEC_STORE(s_snapshot + i + j*ggml_f32_epr + ii*nc, t0); + } } } @@ -9754,6 +9762,9 @@ static void ggml_compute_forward_ssm_scan_f32( sum[j] = GGML_F32_VEC_FMA(sum[j], ax[j], az[j]); GGML_F32_VEC_STORE(s + i + j*GGML_F32_EPR + ii*nc, ax[j]); + if (s_snapshot && s_snapshot != s) { + GGML_F32_VEC_STORE(s_snapshot + i + j*GGML_F32_EPR + ii*nc, ax[j]); + } } } @@ -9772,6 +9783,9 @@ static void ggml_compute_forward_ssm_scan_f32( // y = rowwise_dotprod(state, C) sumf += state * C[ig]; s[i] = state; + if (s_snapshot && s_snapshot != s) { + s_snapshot[i] = state; + } } y[ii] = sumf; } diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index f3418c2af83d..8ab92f959a03 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,15 @@ __global__ void __launch_bounds__(d_state, 1) if (lane == 0) { y_warp[i * stride_y] = state_sum; } + + const int64_t slot = n_tok - 1 - i; + if (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 +241,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 +254,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 +264,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 +779,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 +792,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 +827,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 +855,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.c b/ggml/src/ggml.c index 59191c663eb0..df98eb171681 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5589,6 +5589,21 @@ struct ggml_tensor * ggml_ssm_scan( struct ggml_tensor * B, struct ggml_tensor * C, struct ggml_tensor * ids) { + return ggml_ssm_scan_ext(ctx, s, x, dt, A, B, C, ids, 1); +} + +struct ggml_tensor * ggml_ssm_scan_ext( + struct ggml_context * ctx, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + 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 +5640,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 +5656,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; } diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index ea0ddd114c0c..17d17af0e113 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -992,6 +992,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_NEMOTRON_H: + case LLM_ARCH_NEMOTRON_H_MOE: return true; default: return false; diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp index fd3fe3f03230..97ee2fc0b1b2 100644 --- a/src/models/mamba-base.cpp +++ b/src/models/mamba-base.cpp @@ -2,6 +2,8 @@ #include "llama-memory-recurrent.h" +#include + llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {} ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp, @@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, int il) const { const auto * mctx_cur = inp->mctx; - const auto kv_head = mctx_cur->get_head(); + const auto kv_head = mctx_cur->get_head(); + const auto mem_size = mctx_cur->get_size(); const int64_t d_conv = hparams.ssm_d_conv; const int64_t d_inner = hparams.ssm_d_inner; @@ -164,6 +167,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, const int64_t n_seqs = ubatch.n_seqs; const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t state_slots = (int64_t) mem_size * K; GGML_ASSERT(n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); @@ -198,15 +203,19 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, // => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs} ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0); - // copy last (d_conv - 1) columns back into the state cache - ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs, - conv_x->nb[1], conv_x->nb[2], n_seq_tokens * (conv_x->nb[0])); + const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state); + const size_t row_size = ggml_row_size(conv_states_all->type, row_count); + const int64_t n_written = std::min(n_seq_tokens, K); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv, - ggml_view_1d(ctx0, conv_states_all, - (d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs), - kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) * - ggml_element_size(conv_states_all)))); + for (int64_t slot = 0; slot < n_written; ++slot) { + ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs, + conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv, + ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs, + conv_states_all->nb[1], + ((size_t) slot * mem_size + kv_head) * row_size))); + } // 1D convolution // The equivalent is to make a self-overlapping view of conv_x @@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, // (this is necessary in order to properly use the states before they are overwritten, // while avoiding to make unnecessary copies of the states) auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) { - ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size()); + ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots); // TODO: use semistructured matrices to implement state-space duality // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs} - return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids); + return ggml_ssm_scan_ext(ctx, ssm, x, dt, A, B, C, ids, K); }; ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows); - // store last states + const int64_t D = d_state * d_inner; + const int64_t n_written = std::min(n_seq_tokens, K); + const size_t row_size = ggml_row_size(ssm_states_all->type, D); + const size_t y_row_size = ggml_row_size(y_ssm->type, D); + const size_t state_offset = ggml_nelements(x) * ggml_element_size(x); + ggml_build_forward_expand( - gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]), - ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs, - kv_head * d_state * d_inner * ggml_element_size(ssm_states_all)))); + gf, ggml_cpy(ctx0, + ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written, + y_row_size, y_row_size * n_seqs, state_offset), + ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written, + ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size))); ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1], n_seq_tokens * n_head * x->nb[1], 0); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 419e1eba4c2c..08c6f5a47963 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -217,6 +217,16 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set_tests_properties(test-recurrent-state-rollback PROPERTIES FIXTURES_REQUIRED generate-models ) + + llama_test( + test-recurrent-state-rollback + NAME test-recurrent-state-rollback-nemotron-h + LABEL main + ARGS -m "${MODEL_DIR}/nemotron_h-dense.gguf" + ) + set_tests_properties(test-recurrent-state-rollback-nemotron-h PROPERTIES + FIXTURES_REQUIRED generate-models + ) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)