Skip to content

Initial changes for Recurrent state rollback for nemotron for cpu and… #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
lnigam wants to merge 1 commit into from
Closed

Initial changes for Recurrent state rollback for nemotron for cpu and… #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
16 changes: 15 additions & 1 deletion ggml/src/ggml-cpu/ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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);
}
}
}

Expand Down Expand Up @@ -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]);
}
}
}

Expand All @@ -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;
}
Expand Down
26 changes: 20 additions & 6 deletions ggml/src/ggml-cuda/ssm-scan.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -769,17 +779,20 @@ 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));
GGML_ASSERT(src3->nb[0] == sizeof(float));
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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
20 changes: 19 additions & 1 deletion ggml/src/ggml.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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;
Expand All @@ -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;
}

Expand Down
2 changes: 2 additions & 0 deletions src/llama-arch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
46 changes: 31 additions & 15 deletions src/models/mamba-base.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

#include "llama-memory-recurrent.h"

#include <algorithm>

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,
Expand Down Expand Up @@ -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;
Expand All @@ -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());
Expand Down Expand Up @@ -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<int64_t>(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
Expand Down Expand Up @@ -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<int64_t>(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);
Expand Down
10 changes: 10 additions & 0 deletions tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading