MTP: Avoid per step SSM copy (#1778)

* Avoid copying the per-step SSM state (CUDA)

* Avoid copying the per-step SSM state (CPU)

* Allocate only what is necessary for per-step SSM state

* Cleanup

ggml/include/ggml.h

@@ -2533,7 +2533,7 @@ extern "C" {
             struct ggml_tensor  * g,
             struct ggml_tensor  * beta,
             struct ggml_tensor  * state,
-            bool                  save_all_steps);
+            struct ggml_tensor  * saved_steps);
 
     // custom operators
 

ggml/src/ggml-cuda/delta-net.cu

@@ -36,13 +36,13 @@ __global__ void delta_net_recurrent_f32(
     const float * __restrict__ beta_in,   // [1, n_tokens, n_heads, n_seqs]
     const float * __restrict__ state_in,  // [HEAD_DIM, HEAD_DIM*n_heads, 1, n_seqs]
     float * __restrict__ dst,             // output + new_state(s) concatenated
+    float * __restrict__ saved_states,
     const int64_t n_heads,
     const int64_t gqa_ratio,
     const int repeat_type,
     const int64_t n_tokens,
     const int64_t n_seqs,
     const int64_t output_offset,          // offset where state starts in output
-    const int save_all_states,            // 1 = save per-step states, 0 = final only
     size_t vnb1, size_t vnb2, size_t vnb3) {
     constexpr int warps_per_head = HEAD_DIM/WARP_SIZE;
     const int batch_idx = blockIdx.x / (warps_per_head*n_heads);
@@ -160,8 +160,8 @@ __global__ void delta_net_recurrent_f32(
         }
 
         // Save per-step state if requested
-        if (save_all_states) {
-            float * state_step_dst = dst + output_offset + t * state_step_stride + batch_idx * state_batch_stride + state_head_offset;
+        if (saved_states && t < n_tokens - 1) {
+            float * state_step_dst = saved_states + batch_idx * state_batch_stride + state_head_offset + t * state_step_stride;
             for (int i = 0; i < HEAD_DIM/num_warps; ++i) {
                 int col = num_warps*i + col_idx_0;
                 state_step_dst[col*HEAD_DIM + row_out] = state_local[i];
@@ -176,11 +176,9 @@ __global__ void delta_net_recurrent_f32(
         __syncthreads();
     }
     // Copy the final state to its destination
-    if (!save_all_states) {
-        for (int i = 0; i < HEAD_DIM/num_warps; ++i) {
-            int col = num_warps*i + col_idx_0;
-            state_dst[col*HEAD_DIM + row_out] = state_local[i];
-        }
+    for (int i = 0; i < HEAD_DIM/num_warps; ++i) {
+        int col = num_warps*i + col_idx_0;
+        state_dst[col*HEAD_DIM + row_out] = state_local[i];
     }
 }
 
@@ -192,13 +190,13 @@ static void delta_net_f32_cuda(
     const float * beta,
     const float * state_in,
     float * dst,
+    float * saved_states,
     const int64_t head_dim,
     const int64_t n_tokens,
     const int64_t n_heads,
     const int64_t gqa_ratio,
     const int     repeat_type,
     const int64_t n_seqs,
-    const int     save_all_states,
     size_t vnb1, size_t vnb2, size_t vnb3,
     const int device_id,
     const int cc, // compute capability (e.g., 890 for SM 8.9, 1200 for SM 12.0)
@@ -220,19 +218,19 @@ static void delta_net_f32_cuda(
         constexpr int threads_per_block = 256;
         if (head_dim == 64) {
             delta_net_recurrent_f32<64, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
-                    q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, save_all_states, vnb1, vnb2, vnb3);
+                    q, k, v, g, beta, state_in, dst, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3);
         } else {
             delta_net_recurrent_f32<128, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
-                    q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, save_all_states, vnb1, vnb2, vnb3);
+                    q, k, v, g, beta, state_in, dst, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3);
         }
     } else {
         constexpr int threads_per_block = 128;
         if (head_dim == 64) {
             delta_net_recurrent_f32<64, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
-                    q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, save_all_states, vnb1, vnb2, vnb3);
+                    q, k, v, g, beta, state_in, dst, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3);
         } else {
             delta_net_recurrent_f32<128, threads_per_block><<<num_blocks, threads_per_block, smem_size, stream>>>(
-                    q, k, v, g, beta, state_in, dst, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, save_all_states, vnb1, vnb2, vnb3);
+                    q, k, v, g, beta, state_in, dst, saved_states, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, output_offset, vnb1, vnb2, vnb3);
         }
     }
 
@@ -247,6 +245,7 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
     const ggml_tensor * src3 = dst->src[3];  // g
     const ggml_tensor * src4 = dst->src[4];  // beta
     const ggml_tensor * src5 = dst->src[5];  // state
+    const ggml_tensor * src6 = dst->src[6];  // when not null, state for token 0...n_token-1
 
     GGML_ASSERT(src0->type == GGML_TYPE_F32);
     GGML_ASSERT(dst->type == GGML_TYPE_F32);
@@ -276,9 +275,12 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
     const int64_t state_size = head_dim * head_dim * n_heads * n_seqs;
 
     int repeat_type = dst->op_params[0];
-    int save_all_states = dst->op_params[1];
+    if (src6) {
+        GGML_ASSERT(src6->type == GGML_TYPE_F32);
+        GGML_ASSERT(src6->ne[0] >= (n_tokens - 1)*state_size);
+    }
 
-    const int64_t expected_size = save_all_states ? (output_size + n_tokens * state_size) : (output_size + state_size);
+    const int64_t expected_size = output_size + state_size;
     GGML_ASSERT(ggml_nelements(dst) == expected_size);
 
     GGML_ASSERT(head_dim <= 256);  // Reasonable limit for shared memory
@@ -295,8 +297,8 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
         (const float *)src4->data,
         (const float *)src5->data,
         (float *)dst->data,
+        src6 ? (float *)src6->data : nullptr,
         head_dim, n_tokens, n_heads, gqa_ratio, repeat_type, n_seqs,
-        save_all_states,
         src2->nb[1]/sizeof(float), src2->nb[2]/sizeof(float), src2->nb[3]/sizeof(float),
         device_id, cc,
         ctx.stream());

ggml/src/ggml.c

@@ -9975,7 +9975,7 @@ struct ggml_tensor * ggml_delta_net(
         struct ggml_tensor  * g,
         struct ggml_tensor  * beta,
         struct ggml_tensor  * state,
-        bool save_all_steps) {
+        struct ggml_tensor  * saved_steps) {
     GGML_ASSERT(ggml_is_contiguous(q));
     GGML_ASSERT(ggml_is_contiguous(k));
     GGML_ASSERT(ggml_is_contiguous(state));
@@ -10006,17 +10006,21 @@ struct ggml_tensor * ggml_delta_net(
     const int64_t output_size = S_v * H_v * n_tokens * n_seqs;
     const int64_t state_size  = S_v * S_v * H_v * n_seqs;
 
-    const int64_t state_slots = save_all_steps ? n_tokens : 1;
-    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_size + state_slots * state_size);
+    if (saved_steps) {
+        GGML_ASSERT(saved_steps->type == GGML_TYPE_F32);
+        GGML_ASSERT(saved_steps->ne[0] >= (n_tokens - 1)*state_size);
+    }
+
+    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_size + state_size);
 
     result->op     = GGML_OP_DELTA_NET;
-    result->op_params[1] = save_all_steps ? 1 : 0;
     result->src[0] = q;
     result->src[1] = k;
     result->src[2] = v;
     result->src[3] = g;
     result->src[4] = beta;
     result->src[5] = state;
+    result->src[6] = saved_steps;
 
     return result;
 }
@@ -22674,6 +22678,7 @@ static void ggml_compute_forward_delta_net_f32(
     const struct ggml_tensor * src3 = dst->src[3];
     const struct ggml_tensor * src4 = dst->src[4];
     const struct ggml_tensor * src5 = dst->src[5];
+    const struct ggml_tensor * src6 = dst->src[6];
 
     const int64_t head_dim = src0->ne[0];
     const int64_t n_tokens = src0->ne[1];
@@ -22691,19 +22696,24 @@ static void ggml_compute_forward_delta_net_f32(
     const float * beta_data = (const float *) src4->data;
     const float * state_in  = (const float *) src5->data;
     float * out_data  = (float *) dst->data;
+    float * saved_steps = src6 ? (float *)src6->data : NULL;
 
     const int ith = params->ith;
     const int nth = params->nth;
 
     int repeat_type = dst->op_params[0];
-    const int save_all_steps = dst->op_params[1];
     const int64_t state_step_stride = head_dim * head_dim * n_heads * n_seqs;
     float * state_working = out_data + output_size;
 
+    if (src6) {
+        GGML_ASSERT(src6->type == GGML_TYPE_F32);
+        GGML_ASSERT(src6->ne[0] >= (n_tokens - 1)*state_step_stride);
+    }
+
     if (iqk_fused_delta_net(head_dim, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
                 src2->nb[1]/sizeof(float), src2->nb[2]/sizeof(float), src2->nb[3]/sizeof(float),
                 q_data, k_data, v_data, g_data, beta_data, state_in,
-                out_data, state_working, save_all_steps, (int) state_step_stride, ith, nth)) {
+                out_data, state_working, saved_steps, (int) state_step_stride, ith, nth)) {
         return;
     }
 
@@ -22731,11 +22741,11 @@ static void ggml_compute_forward_delta_net_f32(
         const int64_t out_head_offset  = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim;
         const int64_t out_token_stride = head_dim * n_heads;
 
+        float * state = state_working + state_head_offset;
         for (int64_t i = 0; i < head_dim * head_dim; ++i) {
-            state_working[state_head_offset + i] = state_in[state_head_offset + i];
+            state[i] = state_in[state_head_offset + i];
         }
 
-        float * state = state_working + state_head_offset;
         const int64_t state_head_size = head_dim * head_dim;
 
         for (int64_t t = 0; t < n_tokens; ++t) {
@@ -22792,10 +22802,9 @@ static void ggml_compute_forward_delta_net_f32(
                 }
             }
 
-            if (save_all_steps && t + 1 < n_tokens) {
-                float * next_state = state_working + (t + 1) * state_step_stride + state_head_offset;
+            if (saved_steps && t + 1 < n_tokens) {
+                float * next_state = saved_steps + state_head_offset + t * state_step_stride;
                 memcpy(next_state, state, state_head_size * sizeof(float));
-                state = next_state;
             }
         }
     }

ggml/src/iqk/iqk_mul_mat.cpp

@@ -1438,7 +1438,7 @@ template <int head_dim>
 void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
         size_t vnb1, size_t vnb2, size_t vnb3,
         const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
-        const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth) {
+        const float * state_in, float * out_data, float * state_out, float * saved_steps, int state_step_stride, int ith, int nth) {
     const int total_heads = n_heads * n_seqs;
     const int heads_per_thread = (total_heads + nth - 1) / nth;
     const int h_start = ith * heads_per_thread;
@@ -1537,10 +1537,9 @@ void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type,
                 }
             }
 
-            if (save_all_steps && t + 1 < n_tokens) {
-                float * next_state = state_out + (t + 1) * state_step_stride + state_head_offset;
-                std::memcpy(next_state, state, head_dim * head_dim * sizeof(float));
-                state = next_state;
+            if (saved_steps && t + 1 < n_tokens) {
+                float * this_state = saved_steps + state_head_offset + t * state_step_stride;
+                std::memcpy(this_state, state, head_dim * head_dim * sizeof(float));
             }
         }
     }
@@ -1550,10 +1549,10 @@ template <int head_dim>
 void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
         size_t vnb1, size_t vnb2, size_t vnb3,
         const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
-        const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth) {
+        const float * state_in, float * out_data, float * state_out, float * saved_steps, int state_step_stride, int ith, int nth) {
 #ifdef __ARM_NEON
     iqk_fused_delta_net_neon_impl<head_dim>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3,
-            q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, save_all_steps, state_step_stride, ith, nth);
+            q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, saved_steps, state_step_stride, ith, nth);
     return;
 #endif
     const int total_heads = n_heads * n_seqs;
@@ -1711,10 +1710,9 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n
 #endif
 #endif
 
-            if (save_all_steps && t + 1 < n_tokens) {
-                float * next_state = state_out + (t + 1) * state_step_stride + state_head_offset;
-                std::memcpy(next_state, state, head_dim * head_dim * sizeof(float));
-                state = next_state;
+            if (saved_steps && t + 1 < n_tokens) {
+                float * this_state = saved_steps + state_head_offset + t * state_step_stride;
+                std::memcpy(this_state, state, head_dim * head_dim * sizeof(float));
             }
         }
     }
@@ -1724,16 +1722,16 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n
 bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
         size_t vnb1, size_t vnb2, size_t vnb3,
         const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
-        const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth) {
+        const float * state_in, float * out_data, float * state_out, float * saved_steps, int state_step_stride, int ith, int nth) {
     if (head_dim != 64 && head_dim != 128) {
         return false;
     }
     if (head_dim == 64) {
         iqk_fused_delta_net_impl<64>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, q_data, k_data, v_data, g_data, beta_data, state_in,
-                out_data, state_out, save_all_steps, state_step_stride, ith, nth);
+                out_data, state_out, saved_steps, state_step_stride, ith, nth);
     } else {
         iqk_fused_delta_net_impl<128>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, q_data, k_data, v_data, g_data, beta_data, state_in,
-                out_data, state_out, save_all_steps, state_step_stride, ith, nth);
+                out_data, state_out, saved_steps, state_step_stride, ith, nth);
     }
     return true;
 }
@@ -1774,7 +1772,7 @@ extern "C" IQK_API bool iqk_moe_fused_up_gate(long /*Nx*/, long /*Ny*/, long /*n
 bool iqk_fused_delta_net(int, int, int, int, int, int,
         size_t, size_t, size_t,
         const float *, const float *, const float *, const float *, const float *,
-        const float *, float *, float *, int, int, int, int) {
+        const float *, float *, float *, float *, int, int, int) {
     return false;
 }
 

ggml/src/iqk/iqk_mul_mat.h

@@ -76,7 +76,7 @@ IQK_API void iqk_topk_moe(int n_experts, int n_experts_used, int nrows, const fl
 IQK_API bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
         size_t vnb1, size_t vnb2, size_t vnb3,
         const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
-        const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth);
+        const float * state_in, float * out_data, float * state_out, float * saved_steps, int state_step_stride, int ith, int nth);
 
 #ifdef __cplusplus
 }

src/llama-delta-net.cpp

@@ -80,8 +80,7 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
         ggml_tensor * q, ggml_tensor * k, ggml_tensor * v,
         ggml_tensor * g, ggml_tensor * beta, ggml_tensor * state,
         int il, const llm_build_cb & cb, int repeat_type,
-        bool save_all_steps,
-        ggml_cgraph * gf, ggml_tensor * per_step_ckpt) {
+        ggml_tensor * per_step_ckpt) {
 
     const int64_t S_k      = q->ne[0];
     const int64_t H_k      = q->ne[2];
@@ -123,7 +122,7 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
     cb(beta,      "beta_fused", il);
     cb(state_flat,"state_fused", il);
 
-    ggml_tensor * fused_result = ggml_delta_net(ctx0, q, k, v, g, beta, state_flat, save_all_steps);
+    ggml_tensor * fused_result = ggml_delta_net(ctx0, q, k, v, g, beta, state_flat, per_step_ckpt);
     cb(fused_result, "delta_net_fused_raw", il);
     fused_result->op_params[0] = repeat_type;
 
@@ -137,34 +136,13 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
             ggml_row_size(fused_result->type, S_v * H_v * n_tokens), 0);
     //output_tokens = ggml_cont_4d(ctx0, output_tokens, S_v, H_v, n_tokens, n_seqs);
 
-    // per-step states are at [output_size, output_size + n_tokens*state_size)
-    const int64_t last_state_offset = save_all_steps
-        ? (output_size + (n_tokens - 1) * state_size)
-        : output_size;
-
     ggml_tensor * new_state_flat = ggml_view_1d(ctx0, fused_result, state_size,
-            last_state_offset * ggml_element_size(fused_result));
+            output_size * ggml_element_size(fused_result));
     ggml_tensor * new_state = ggml_reshape_4d(ctx0, new_state_flat, S_v, S_v, H_v, n_seqs);
 
     cb(output_tokens, "output_tokens", il);
     cb(new_state,     "new_state", il);
 
-    // Copy all per-step SSM states to persistent checkpoint tensor
-    if (save_all_steps && per_step_ckpt != nullptr && gf != nullptr && n_tokens > 1) {
-        const int64_t per_step_total = n_tokens * state_size;
-        if (per_step_total <= ggml_nelements(per_step_ckpt)) {
-            ggml_tensor * all_steps_src = ggml_view_1d(ctx0, fused_result, per_step_total,
-                    output_size * ggml_element_size(fused_result));
-            ggml_tensor * ckpt_dst = ggml_view_1d(ctx0, per_step_ckpt, per_step_total, 0);
-            auto ckpt_cpy = ggml_cpy(ctx0, all_steps_src, ckpt_dst);
-            cb(ckpt_cpy, "per_step_ckpt_cpy", il);
-            ggml_build_forward_expand(gf, ckpt_cpy);
-        } else {
-            LLAMA_LOG_WARN("%s: per-step checkpoint tensor too small for %lld tokens (need %lld, have %lld), skipping per-step save\n",
-                    __func__, (long long)n_tokens, (long long)per_step_total, (long long)ggml_nelements(per_step_ckpt));
-        }
-    }
-
     return {output_tokens, new_state};
 }
 
@@ -307,7 +285,7 @@ ggml_tensor * delta_net::build_qkv(ggml_context * ctx0, ggml_tensor * state_stor
         int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads, int64_t ssm_d_conv,
         int64_t state_seq_id_local, uint32_t qnext_state_slots, bool reset_state_local,
         float eps_norm, int repeat_type, int il, const llm_build_cb & cb, ggml_cgraph * gf,
-        bool save_per_step_states, ggml_tensor * per_step_ckpt) {
+        ggml_tensor * per_step_ckpt) {
     const int64_t key_dim        = head_k_dim * num_k_heads;
     const int64_t value_dim      = head_v_dim * num_v_heads;
     const int64_t conv_dim       = key_dim * 2 + value_dim;
@@ -393,7 +371,7 @@ ggml_tensor * delta_net::build_qkv(ggml_context * ctx0, ggml_tensor * state_stor
     cb(k_conv, "k_conv_normed", il);
 
     auto [output, new_state] = build_fused_delta_net(ctx0, q_conv, k_conv, v_conv, gate, beta, state, il, cb, repeat_type,
-            save_per_step_states, gf, per_step_ckpt);
+            per_step_ckpt);
 
     cb(output, "attn_output", il);
     cb(new_state, "new_state", il);
@@ -557,7 +535,7 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
             auto output = build_qkv(ctx0, split_s_l->splits[id], split_ssm_conv1d->splits[id], qkv_mixed, inp_s_seq_qnext, beta, gate,
                                head_k_dim, num_k_heads_id, head_v_dim, num_v_heads_id, hparams.ssm_d_conv,
                                state_seq_id_local, qnext_state_slots, reset_state_local, hparams.f_norm_rms_eps,
-                               l.ssm_beta_alpha ? 0 : 1, il, cb, gf, save_per_step_states, per_step_ckpt);
+                               l.ssm_beta_alpha ? 0 : 1, il, cb, gf, per_step_ckpt);
             split_norm = (ggml_split_tensor_t *)l.ssm_norm->extra;
             GGML_ASSERT(split_norm && split_norm->splits[id]);
             auto split_ssm_out = (ggml_split_tensor_t *)l.ssm_out->extra;
@@ -629,8 +607,7 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
         qkv_mixed, inp_s_seq_qnext, beta, gate,
         head_k_dim, num_k_heads, head_v_dim, num_v_heads, hparams.ssm_d_conv,
         state_seq_id_local, qnext_state_slots, reset_state_local, hparams.f_norm_rms_eps,
-        model.layers[il].ssm_beta_alpha ? 0 : 1, il, cb, gf,
-        save_per_step_states, per_step_ckpt);
+        model.layers[il].ssm_beta_alpha ? 0 : 1, il, cb, gf, per_step_ckpt);
 
     auto gated_output = build_gated_output(lctx, ctx0, model.layers[il].ssm_norm, model.layers[il].ssm_out, output, z, head_v_dim, num_v_heads, n_tok, il, cb);
     if (inp_out_ids) {

src/llama-delta-net.h

@@ -15,8 +15,7 @@ struct delta_net {
                       ggml_tensor * q, ggml_tensor * k, ggml_tensor * v,
                       ggml_tensor * g, ggml_tensor * beta, ggml_tensor * state,
                       int il, const llm_build_cb & cb, int repeat_type,
-                      bool save_all_steps = false,
-                      ggml_cgraph * gf = nullptr, ggml_tensor * per_step_ckpt = nullptr);
+                      ggml_tensor * per_step_ckpt = nullptr);
 
     ggml_tensor * build_layer_attn_linear_core(ggml_context * ctx0, ggml_cgraph * gf,
             ggml_tensor * cur, ggml_tensor * inp_s_seq_qnext, ggml_tensor * inp_out_ids,
@@ -52,7 +51,7 @@ private:
             int64_t head_k_dim, int64_t num_k_heads, int64_t head_v_dim, int64_t num_v_heads, int64_t ssm_d_conv,
             int64_t state_seq_id_local, uint32_t qnext_state_slots, bool reset_state_local,
             float eps_norm, int repeat_type, int il, const llm_build_cb & cb, ggml_cgraph * gf,
-            bool save_per_step_states = false, ggml_tensor * per_step_ckpt = nullptr);
+            ggml_tensor * per_step_ckpt = nullptr);
 
     static ggml_tensor * build_gated_output(llama_context & lctx, ggml_context * ctx0, ggml_tensor * ssm_norm, ggml_tensor * ssm_out,
             ggml_tensor * output, ggml_tensor * z, int64_t head_v_dim, int64_t num_v_heads, int64_t n_tok, int il, const llm_build_cb & cb);

src/llama.cpp

@@ -1546,11 +1546,13 @@ bool llama_kv_cache::per_step_alloc(const llama_model & model, int max_tokens) {
             auto [il, id] = p;
             // SSM state: max_tokens * ssm_state_dim
             if (id < 0) {
-                GGML_ASSERT(ckpt.per_step_ssm[il].empty());
-                GGML_ASSERT(ckpt.per_step_qkv[il].empty());
-                ggml_tensor * t_ssm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)max_tokens * ssm_state_dim);
-                ggml_format_name(t_ssm, "per_step_ssm_l%d", il);
-                ckpt.per_step_ssm[il].push_back(t_ssm);
+                if (max_tokens > 1) {
+                    GGML_ASSERT(ckpt.per_step_ssm[il].empty());
+                    GGML_ASSERT(ckpt.per_step_qkv[il].empty());
+                    ggml_tensor * t_ssm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)(max_tokens - 1) * ssm_state_dim);
+                    ggml_format_name(t_ssm, "per_step_ssm_l%d", il);
+                    ckpt.per_step_ssm[il].push_back(t_ssm);
+                }
 
                 // Conv features (qkv_mixed): max_tokens * conv_dim
                 ggml_tensor * t_qkv = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)max_tokens * conv_dim);
@@ -1574,9 +1576,11 @@ bool llama_kv_cache::per_step_alloc(const llama_model & model, int max_tokens) {
                 int nv = split->ne[0] / head_v_dim; // number of heads handled by this device
                 auto [this_conv_dim, this_ssm_dim] = model.hparams.n_embd_v_s_dims(nv);
 
-                auto t_ssm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)max_tokens * this_ssm_dim);
-                ggml_format_name(t_ssm, "per_step_ssm_l%d_%d", il, id);
-                ckpt.per_step_ssm[il][id] = t_ssm;
+                if (max_tokens > 1) {
+                    auto t_ssm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)(max_tokens - 1) * this_ssm_dim);
+                    ggml_format_name(t_ssm, "per_step_ssm_l%d_%d", il, id);
+                    ckpt.per_step_ssm[il][id] = t_ssm;
+                }
 
                 auto t_qkv = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)max_tokens * this_conv_dim);
                 ggml_format_name(t_qkv, "per_step_qkv_l%d_%d", il, id);