furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

recurrent model: implement recurrent kernel checkpoint

Browse Source
This commit is contained in:
SamuelOliveirads
2026-04-20 12:22:29 -03:00
parent 9c44b117f9
commit dfac8f19f6
10 changed files with 517 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/server/server-context.cpp
+83 -43
View File
@@ -23,7 +23,9 @@ static void log_text(const gpt_params & params_base, const std::string & text) {
void server_speculative_checkpoint::clear() {
valid = false;
per_step_enabled = false;
n_past = 0;
sampled = LLAMA_TOKEN_NULL;
data.clear();
if (sampler != nullptr) {
@@ -36,6 +38,7 @@ static void discard_speculative_checkpoint(server_slot & slot, llama_context * c
slot.spec_ckpt.clear();
if (gpu_ckpt) {
llama_kv_cache_set_per_step_save(ctx, false, 0);
llama_kv_cache_checkpoint_delete(ctx);
}
}
@@ -46,6 +49,16 @@ static bool save_speculative_checkpoint(server_slot & slot, llama_model * model,
slot.spec_ckpt.sampled = slot.sampled;
if (gpu_ckpt) {
// Enable per-step SSM state saves for the verification forward pass.
// The delta_net kernel saves SSM state after each token step, allowing us
// to skip the expensive re-decode on rejection.
// Returns false if per-step is unsupported (e.g., split tensors in graph split mode).
const int max_tokens = (int)slot.drafted.size() + 1;
slot.spec_ckpt.per_step_enabled = llama_kv_cache_set_per_step_save(ctx, true, max_tokens);
// Always save full checkpoint (conv + SSM state).
// On rejection with per-step: restore conv from here, SSM from per-step.
// On rejection without per-step: full restore + re-decode (legacy path).
slot.spec_ckpt.valid = llama_kv_cache_checkpoint_save(ctx);
} else {
const size_t ckpt_size = llama_state_seq_get_size(ctx, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
@@ -3661,59 +3674,86 @@ void server_context::speculative_decoding_accept() {
// for recurrent/hybrid models: if any drafts were rejected, restore recurrent state
const bool any_rejected = (ids.size() - 1) < n_draft;
const bool gpu_ckpt = llama_kv_cache_checkpoint_supported(ctx);
const bool has_per_step = slot.spec_ckpt.valid && slot.spec_ckpt.per_step_enabled;
if (any_rejected && slot.spec_ckpt.valid) {
if (gpu_ckpt) {
llama_kv_cache_checkpoint_restore(ctx);
} else {
llama_state_seq_set_data(ctx, slot.spec_ckpt.data.data(), slot.spec_ckpt.data.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
if (has_per_step) {
const int step = (int)ids.size() - 1;
llama_kv_cache_seq_rm(ctx, slot.id, slot.spec_ckpt.n_past, -1);
// Keep accepted attention KV from the verification pass and only fix the
// recurrent state plus remove rejected positions.
llama_kv_cache_per_step_restore(ctx, step);
// restore sampler state (RNG, grammar, prev tokens)
if (slot.spec_ckpt.sampler) {
common_sampler_clone(slot.spec_ckpt.sampler, slot.ctx_sampling);
}
const llama_pos accepted_pos = slot.spec_ckpt.n_past + step;
llama_kv_cache_set_cell_pos(ctx, slot.id, accepted_pos);
llama_kv_cache_seq_rm(ctx, slot.id, accepted_pos + 1, -1);
if (!ids.empty()) {
// Re-decode to restore the recurrent state after checkpoint restore.
const int n_re = (int)ids.size();
llama_batch re_batch = llama_batch_init(n_re, 0, 1);
common_batch_add(re_batch, slot.spec_ckpt.sampled, slot.spec_ckpt.n_past, { slot.id }, n_re == 1);
for (int j = 0; j < n_re - 1; j++) {
const bool is_last = (j == n_re - 2);
common_batch_add(re_batch, ids[j], slot.spec_ckpt.n_past + 1 + j, { slot.id }, is_last);
// Restore sampler state
if (slot.spec_ckpt.sampler) {
common_sampler_clone(slot.spec_ckpt.sampler, slot.ctx_sampling);
}
if (slot.has_mtp) {
llama_set_embeddings(ctx, true);
}
const int ret = llama_decode(ctx, re_batch);
if (ret != 0) {
SLT_ERR(slot, "failed to re-decode accepted tokens after checkpoint restore: %d\n", ret);
}
if (slot.has_mtp) {
llama_set_embeddings(ctx, false);
const int n_embd = llama_model_n_embd(llama_get_model(ctx));
const float * emb = llama_get_embeddings_ith(ctx, -1);
if (emb) {
slot.mtp_hidden_state.resize(n_embd);
memcpy(slot.mtp_hidden_state.data(), emb, n_embd * sizeof(float));
}
}
for (llama_token id : ids) {
common_sampler_accept(slot.ctx_sampling, ctx, id, true);
}
llama_batch_free(re_batch);
SLT_DBG(slot, "spec checkpoint restored (gpu=%s): re-decoded %d tokens (rejected %d drafts)\n",
gpu_ckpt ? "yes" : "no", n_re, (int)(n_draft - (ids.size() - 1)));
}
SLT_DBG(slot, "per-step restore: step=%d (rejected %d drafts)\n",
step, (int)(n_draft - (ids.size() - 1)));
discard_speculative_checkpoint(slot, ctx, gpu_ckpt);
discard_speculative_checkpoint(slot, ctx, gpu_ckpt);
} else {
// Legacy path: full checkpoint restore + re-decode
if (gpu_ckpt) {
llama_kv_cache_checkpoint_restore(ctx);
} else {
llama_state_seq_set_data(ctx, slot.spec_ckpt.data.data(), slot.spec_ckpt.data.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
}
llama_kv_cache_seq_rm(ctx, slot.id, slot.spec_ckpt.n_past, -1);
// restore sampler state (RNG, grammar, prev tokens)
if (slot.spec_ckpt.sampler) {
common_sampler_clone(slot.spec_ckpt.sampler, slot.ctx_sampling);
}
if (!ids.empty()) {
// Re-decode to restore the recurrent state after checkpoint restore.
const int n_re = (int)ids.size();
llama_batch re_batch = llama_batch_init(n_re, 0, 1);
common_batch_add(re_batch, slot.spec_ckpt.sampled, slot.spec_ckpt.n_past, { slot.id }, n_re == 1);
for (int j = 0; j < n_re - 1; j++) {
const bool is_last = (j == n_re - 2);
common_batch_add(re_batch, ids[j], slot.spec_ckpt.n_past + 1 + j, { slot.id }, is_last);
}
if (slot.has_mtp) {
llama_set_embeddings(ctx, true);
}
const int ret = llama_decode(ctx, re_batch);
if (ret != 0) {
SLT_ERR(slot, "failed to re-decode accepted tokens after checkpoint restore: %d\n", ret);
}
if (slot.has_mtp) {
llama_set_embeddings(ctx, false);
const int n_embd = llama_model_n_embd(llama_get_model(ctx));
const float * emb = llama_get_embeddings_ith(ctx, -1);
if (emb) {
slot.mtp_hidden_state.resize(n_embd);
memcpy(slot.mtp_hidden_state.data(), emb, n_embd * sizeof(float));
}
}
for (llama_token id : ids) {
common_sampler_accept(slot.ctx_sampling, ctx, id, true);
}
llama_batch_free(re_batch);
SLT_DBG(slot, "spec checkpoint restored (gpu=%s): re-decoded %d tokens (rejected %d drafts)\n",
gpu_ckpt ? "yes" : "no", n_re, (int)(n_draft - (ids.size() - 1)));
}
discard_speculative_checkpoint(slot, ctx, gpu_ckpt);
}
} else {
llama_kv_cache_seq_rm(ctx, slot.id, slot.n_past, -1);
discard_speculative_checkpoint(slot, ctx, gpu_ckpt);
examples/server/server-context.h
+1
View File
@@ -24,6 +24,7 @@ enum slot_command {
struct server_speculative_checkpoint {
bool valid = false;
bool per_step_enabled = false; // per-step SSM checkpoints
llama_pos n_past = 0;
llama_token sampled = LLAMA_TOKEN_NULL;
std::vector<uint8_t> data;
ggml/include/ggml.h
+2 -1
View File
@@ -2526,7 +2526,8 @@ extern "C" {
struct ggml_tensor * v,
struct ggml_tensor * g,
struct ggml_tensor * beta,
struct ggml_tensor * state);
struct ggml_tensor * state,
bool save_all_steps);
// custom operators
ggml/src/ggml-cuda/delta-net.cu
+31 -11
View File
@@ -35,13 +35,14 @@ __global__ void delta_net_recurrent_f32(
const float * __restrict__ g, // [n_tokens, 1, n_heads, n_seqs]
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 concatenated
float * __restrict__ dst, // output + new_state(s) concatenated
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);
@@ -69,6 +70,9 @@ __global__ void delta_net_recurrent_f32(
const int64_t state_head_offset = head_idx * HEAD_DIM * HEAD_DIM;
const int64_t state_batch_stride = HEAD_DIM * HEAD_DIM * n_heads;
// State step stride for save_all_states: HEAD_DIM^2 * n_heads * n_seqs
const int64_t state_step_stride = HEAD_DIM * HEAD_DIM * n_heads * n_seqs;
// Pointers for this batch/head
const float * q_ptr = q + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
const float * k_ptr = k + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
@@ -157,12 +161,23 @@ __global__ void delta_net_recurrent_f32(
state_local[i] = new_state_val;
}
// 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;
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];
}
}
}
__syncthreads();
// Copy the final state to its destination
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];
// 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];
}
}
}
@@ -180,9 +195,10 @@ static void delta_net_f32_cuda(
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)
const int cc,
cudaStream_t stream) {
GGML_UNUSED(device_id);
GGML_UNUSED(cc);
@@ -201,19 +217,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, vnb1, vnb2, vnb3);
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);
} 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, vnb1, vnb2, vnb3);
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);
}
} 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, vnb1, vnb2, vnb3);
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);
} 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, vnb1, vnb2, vnb3);
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);
}
}
@@ -255,9 +271,12 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
// Verify output tensor size
const int64_t output_size = head_dim * n_tokens * n_heads * n_seqs;
const int64_t state_size = head_dim * head_dim * n_heads * n_seqs;
GGML_ASSERT(ggml_nelements(dst) == output_size + state_size);
int repeat_type = dst->op_params[0];
int save_all_states = dst->op_params[1];
const int64_t expected_size = save_all_states ? (output_size + n_tokens * state_size) : (output_size + state_size);
GGML_ASSERT(ggml_nelements(dst) == expected_size);
GGML_ASSERT(head_dim <= 256); // Reasonable limit for shared memory
@@ -274,6 +293,7 @@ void ggml_cuda_op_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
(const float *)src5->data,
(float *)dst->data,
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
+11 -6
View File
@@ -9940,7 +9940,8 @@ struct ggml_tensor * ggml_delta_net(
struct ggml_tensor * v,
struct ggml_tensor * g,
struct ggml_tensor * beta,
struct ggml_tensor * state) {
struct ggml_tensor * state,
bool save_all_steps) {
GGML_ASSERT(ggml_is_contiguous(q));
GGML_ASSERT(ggml_is_contiguous(k));
GGML_ASSERT(ggml_is_contiguous(state));
@@ -9971,9 +9972,11 @@ 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;
struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_size + state_size);
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);
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;
@@ -22654,17 +22657,19 @@ 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 * state_out = out_data + output_size;
const int ith = params->ith;
const int nth = params->nth;
int repeat_type = dst->op_params[0];
// op_params[1] (save_all_steps) is handled by the CUDA backend only;
// the CPU path always writes to the single state slot after the output.
float * state_working = out_data + output_size;
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_out, ith, nth)) {
out_data, state_working, ith, nth)) {
return;
}
@@ -22694,10 +22699,10 @@ static void ggml_compute_forward_delta_net_f32(
const int64_t out_token_stride = head_dim * n_heads;
for (int64_t i = 0; i < head_dim * head_dim; ++i) {
state_out[state_head_offset + i] = state_in[state_head_offset + i];
state_working[state_head_offset + i] = state_in[state_head_offset + i];
}
float * state = state_out + state_head_offset;
float * state = state_working + state_head_offset;
for (int64_t t = 0; t < n_tokens; ++t) {
const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
include/llama.h
+11 -1
View File
@@ -798,12 +798,22 @@ extern "C" {
LLAMA_API void llama_kv_cache_clear(
struct llama_context * ctx);
// GPU-resident checkpoint for recurrent/hybrid speculative decoding
// Checkpoint for recurrent/hybrid speculative decoding
LLAMA_API bool llama_kv_cache_checkpoint_save(struct llama_context * ctx);
LLAMA_API bool llama_kv_cache_checkpoint_restore(struct llama_context * ctx);
LLAMA_API void llama_kv_cache_checkpoint_delete(struct llama_context * ctx);
LLAMA_API bool llama_kv_cache_checkpoint_supported(struct llama_context * ctx);
// Per-step SSM state checkpoints for speculative decoding
// Enables saving recurrent state after each token step during forward pass
LLAMA_API bool llama_kv_cache_set_per_step_save(struct llama_context * ctx, bool enable, int max_tokens);
// Restore SSM state from per-step checkpoint at given step index
LLAMA_API bool llama_kv_cache_per_step_restore(struct llama_context * ctx, int step);
// Update the position of a recurrent cell for a given sequence.
// Used by per-step restore to fix the cell's position to match the restored state.
LLAMA_API void llama_kv_cache_set_cell_pos(struct llama_context * ctx, llama_seq_id seq_id, llama_pos pos);
// Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
// Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails
// seq_id < 0 : match any sequence
src/llama-context.h
+34
View File
@@ -59,6 +59,9 @@ struct llama_kv_cache {
std::vector<struct ggml_tensor *> v_l;
std::vector<struct ggml_tensor *> s_l; // per layer recurrent state storage (Qwen3Next)
// When true, the delta_net graph builder will enable per-step SSM state saves
bool save_per_step_ssm = false;
std::vector<llama_split_tensor> split_k_l;
std::vector<llama_split_tensor> split_v_l;
std::vector<llama_split_tensor> split_s_l;
@@ -84,6 +87,27 @@ struct llama_kv_cache {
std::vector<std::vector<ggml_tensor *>> split_s_l_shadow;
// Per-step SSM state checkpoints for speculative decoding.
// One tensor per recurrent layer. Each stores up to max_spec_tokens
// copies of the SSM state (not conv state).
std::vector<ggml_tensor *> per_step_ssm;
// Per-step conv feature buffer: stores qkv_mixed features from the
// verification forward pass so conv state can be reconstructed at any step.
// One tensor per recurrent layer, each sized [conv_dim * max_tokens].
std::vector<ggml_tensor *> per_step_qkv;
int32_t per_step_n_tokens = 0; // actual n_tokens stored in current checkpoint
int32_t per_step_max_allocated = 0; // max tokens the per-step buffer can hold
int64_t per_step_ssm_state_size = 0; // SSM state size per step per layer
int64_t per_step_conv_state_dim = 0; // conv state dim (for partial restore offset)
int64_t per_step_conv_dim = 0; // conv_dim = key_dim*2 + value_dim (feature vector length)
int32_t per_step_d_conv = 0; // d_conv from hparams
// Separate storage for per-step allocations (so we can reallocate independently)
std::vector<struct ggml_context *> per_step_ctxs;
std::vector<ggml_backend_buffer_t> per_step_bufs;
std::vector<struct ggml_context *> shadow_ctxs;
std::vector<ggml_backend_buffer_t> shadow_bufs;
@@ -97,6 +121,12 @@ struct llama_kv_cache {
for (ggml_backend_buffer_t buf : shadow_bufs) {
ggml_backend_buffer_free(buf);
}
for (struct ggml_context * ctx : per_step_ctxs) {
ggml_free(ctx);
}
for (ggml_backend_buffer_t buf : per_step_bufs) {
ggml_backend_buffer_free(buf);
}
}
};
@@ -108,6 +138,10 @@ struct llama_kv_cache {
bool checkpoint_restore();
void checkpoint_delete();
// Per-step checkpoint: allocate, restore step k's full state (SSM + conv) to cache
bool per_step_alloc(int max_tokens);
bool per_step_restore(int step);
~llama_kv_cache() {
for (struct ggml_context * ctx : ctxs) {
ggml_free(ctx);
src/llama-delta-net.cpp
+50 -7
View File
@@ -70,6 +70,7 @@ delta_net::delta_net(llama_context & _lctx, const llama_batch & _batch) : lctx(_
GGML_ASSERT((uint32_t) s < qnext_state_slots);
}
save_per_step_states = lctx.kv_self.save_per_step_ssm && batch.n_tokens > 1;
}
delta_net::~delta_net() = default;
@@ -77,7 +78,9 @@ delta_net::~delta_net() = default;
std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_context * ctx0,
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) {
int il, const llm_build_cb & cb, int repeat_type,
bool save_all_steps,
ggml_cgraph * gf, ggml_tensor * per_step_ckpt) {
const int64_t S_k = q->ne[0];
const int64_t H_k = q->ne[2];
@@ -119,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);
ggml_tensor * fused_result = ggml_delta_net(ctx0, q, k, v, g, beta, state_flat, save_all_steps);
cb(fused_result, "delta_net_fused_raw", il);
fused_result->op_params[0] = repeat_type;
@@ -131,15 +134,35 @@ std::pair<ggml_tensor *, ggml_tensor *> delta_net::build_fused_delta_net(ggml_co
ggml_row_size(fused_result->type, S_v),
ggml_row_size(fused_result->type, S_v * H_v),
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,
output_size * ggml_element_size(fused_result));
last_state_offset * 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};
}
@@ -281,7 +304,8 @@ ggml_tensor * delta_net::build_qkv(ggml_context * ctx0, ggml_tensor * state_stor
ggml_tensor * qkv_mixed, ggml_tensor * inp_s_seq_qnext, ggml_tensor * beta, ggml_tensor * gate,
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) {
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) {
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;
@@ -366,7 +390,8 @@ ggml_tensor * delta_net::build_qkv(ggml_context * ctx0, ggml_tensor * state_stor
cb(q_conv, "q_conv_normed", il);
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);
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);
cb(output, "attn_output", il);
cb(new_state, "new_state", il);
@@ -566,11 +591,29 @@ ggml_tensor * delta_net::build_layer_attn_linear_core(ggml_context * ctx0, ggml_
auto [beta, gate] = build_beta_gate(lctx, ctx0, model.layers[il].ssm_beta_alpha, model.layers[il].ssm_beta, model.layers[il].ssm_alpha,
model.layers[il].ssm_dt, model.layers[il].ssm_a, num_k_heads, num_v_heads, n_seqs, cur, il, cb, gf);
// Get per-step checkpoint tensor if available
ggml_tensor * per_step_ckpt = nullptr;
if (save_per_step_states && il < (int)kv_self.ckpt.per_step_ssm.size()) {
per_step_ckpt = kv_self.ckpt.per_step_ssm[il];
}
// Save qkv_mixed features for per-step conv state reconstruction
if (save_per_step_states && il < (int)kv_self.ckpt.per_step_qkv.size() && kv_self.ckpt.per_step_qkv[il] != nullptr) {
const int64_t conv_dim = qkv_mixed->ne[0];
const int64_t n_tok_qkv = qkv_mixed->ne[1] * qkv_mixed->ne[2];
ggml_tensor * qkv_flat = ggml_reshape_2d(ctx0, qkv_mixed, conv_dim, n_tok_qkv);
ggml_tensor * qkv_dst = ggml_view_2d(ctx0, kv_self.ckpt.per_step_qkv[il],
conv_dim, n_tok_qkv, conv_dim * sizeof(float), 0);
auto qkv_cpy = ggml_cpy(ctx0, qkv_flat, qkv_dst);
ggml_build_forward_expand(gf, qkv_cpy);
}
auto output = build_qkv(ctx0, kv_self.s_l[il], model.layers[il].ssm_conv1d,
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);
model.layers[il].ssm_beta_alpha ? 0 : 1, il, cb, gf,
save_per_step_states, 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
+9 -2
View File
@@ -8,10 +8,16 @@ struct delta_net {
delta_net(llama_context & lctx, const llama_batch & batch);
~delta_net();
// When true, the delta_net kernel saves recurrent state after each token step.
// Used for speculative decoding to enable per-step state checkpoint restoration.
bool save_per_step_states = false;
static std::pair<ggml_tensor *, ggml_tensor *> build_fused_delta_net(ggml_context * ctx0,
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);
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 * 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,
@@ -46,7 +52,8 @@ private:
ggml_tensor * qkv_mixed, ggml_tensor * inp_s_seq_qnext, ggml_tensor * beta, ggml_tensor * gate,
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);
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);
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
+285 -37
View File
@@ -1227,7 +1227,10 @@ bool llama_kv_cache::checkpoint_alloc_shadows() {
int split_idx; // -1 for non-split
};
std::map<ggml_backend_buffer_type_t, std::vector<tensor_entry>> buft_entries;
const bool conv_only_shadow = save_per_step_ssm && ckpt.per_step_conv_state_dim > 0;
std::vector<tensor_entry> nonsplit_entries;
std::map<ggml_backend_buffer_type_t, std::vector<tensor_entry>> split_buft_entries;
uint32_t split_s_idx = 0;
for (uint32_t il = 0; il < n_layer; ++il) {
@@ -1237,25 +1240,21 @@ bool llama_kv_cache::checkpoint_alloc_shadows() {
auto * extra = s_l[il]->extra;
if (extra != nullptr) {
// Split tensor — shadow each per-device split tensor individually
auto * split_info = (const ggml_split_tensor_t *)extra;
for (int d = 0; d < split_info->n_device; ++d) {
if (split_info->splits[d] == nullptr) continue;
ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(split_info->splits[d]->buffer);
buft_entries[buft].push_back({split_info->splits[d], il, d});
split_buft_entries[buft].push_back({split_info->splits[d], il, d});
}
split_s_idx++;
} else {
// Non-split tensor
ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(s_l[il]->buffer);
buft_entries[buft].push_back({s_l[il], il, -1});
nonsplit_entries.push_back({s_l[il], il, -1});
}
}
// Allocate shadow tensors grouped by buffer type
for (auto & [buft, entries] : buft_entries) {
if (!nonsplit_entries.empty()) {
ggml_init_params params = {
/*.mem_size =*/ entries.size() * ggml_tensor_overhead(),
/*.mem_size =*/ nonsplit_entries.size() * ggml_tensor_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
@@ -1265,35 +1264,63 @@ bool llama_kv_cache::checkpoint_alloc_shadows() {
return false;
}
for (auto & entry : entries) {
ggml_tensor * shadow = ggml_dup_tensor(ctx, entry.primary);
ggml_format_name(shadow, "shadow_s_l%d_d%d", entry.il, entry.split_idx);
entry.primary = shadow; // temporarily store shadow in entry for buffer allocation
for (auto & entry : nonsplit_entries) {
// Only need the conv portion when per-step is active.
const int64_t nelems = conv_only_shadow
? ckpt.per_step_conv_state_dim
: (int64_t)ggml_nelements(entry.primary);
ggml_tensor * shadow = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelems);
ggml_format_name(shadow, "shadow_s_l%d", entry.il);
ckpt.s_l_shadow[entry.il] = shadow;
}
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft);
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_cpu_buffer_type());
if (!buf) {
LLAMA_LOG_ERROR("%s: failed to allocate buffer for shadow tensors\n", __func__);
LLAMA_LOG_ERROR("%s: failed to allocate CPU buffer for shadow tensors\n", __func__);
ggml_free(ctx);
return false;
}
ggml_backend_buffer_clear(buf, 0);
LLAMA_LOG_INFO("%s: %10s shadow buffer size = %8.2f MiB\n", __func__,
LLAMA_LOG_INFO("%s: CPU shadow buffer = %8.2f MiB (%s)\n", __func__,
ggml_backend_buffer_get_size(buf) / 1024.0 / 1024.0,
conv_only_shadow ? "conv-state only" : "full recurrent state");
ckpt.shadow_ctxs.push_back(ctx);
ckpt.shadow_bufs.push_back(buf);
}
// Allocate split shadows on their respective devices
for (auto & [buft, entries] : split_buft_entries) {
ggml_init_params params = {
/*.mem_size =*/ entries.size() * ggml_tensor_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
ggml_context * ctx = ggml_init(params);
if (!ctx) {
LLAMA_LOG_ERROR("%s: failed to create ggml context for split shadow tensors\n", __func__);
return false;
}
for (auto & entry : entries) {
ggml_tensor * shadow = ggml_dup_tensor(ctx, entry.primary);
ggml_format_name(shadow, "shadow_s_l%d_d%d", entry.il, entry.split_idx);
entry.primary = shadow;
}
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft);
if (!buf) {
LLAMA_LOG_ERROR("%s: failed to allocate buffer for split shadow tensors\n", __func__);
ggml_free(ctx);
return false;
}
ggml_backend_buffer_clear(buf, 0);
LLAMA_LOG_INFO("%s: %10s split shadow buffer = %8.2f MiB\n", __func__,
ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf) / 1024.0 / 1024.0);
ckpt.shadow_ctxs.push_back(ctx);
ckpt.shadow_bufs.push_back(buf);
}
// Build shadow lookup: for non-split layers, store directly in s_l_shadow
for (auto & [buft, entries] : buft_entries) {
for (auto & entry : entries) {
if (entry.split_idx == -1) {
ckpt.s_l_shadow[entry.il] = entry.primary;
}
}
}
// Build shadow split tensors
// Build split shadow lookup
ckpt.split_s_l_shadow.resize(split_s_l.size());
split_s_idx = 0;
for (uint32_t il = 0; il < n_layer; ++il) {
@@ -1307,9 +1334,8 @@ bool llama_kv_cache::checkpoint_alloc_shadows() {
for (int d = 0; d < split_info->n_device; ++d) {
if (split_info->splits[d] == nullptr) continue;
// Find corresponding shadow tensor from buft_entries
ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(split_info->splits[d]->buffer);
for (auto & entry : buft_entries[buft]) {
for (auto & entry : split_buft_entries[buft]) {
if (entry.il == il && entry.split_idx == d) {
shadow_split[d] = entry.primary;
break;
@@ -1330,12 +1356,10 @@ bool llama_kv_cache::checkpoint_save() {
const uint32_t n_layer = (uint32_t)s_l.size();
// 1. Snapshot cell metadata (small, host-side)
ckpt.cells_snapshot = cells;
ckpt.head_snapshot = head;
ckpt.used_snapshot = used;
// 2. Copy tensors: primary -> shadow (GPU-to-GPU when on same device)
uint32_t split_s_idx = 0;
for (uint32_t il = 0; il < n_layer; ++il) {
if (s_l[il] == nullptr) {
@@ -1343,7 +1367,6 @@ bool llama_kv_cache::checkpoint_save() {
}
if (s_l[il]->extra != nullptr) {
// Split tensor — copy per-device splits individually
auto * split_info = (const ggml_split_tensor_t *)s_l[il]->extra;
auto & shadow_split = ckpt.split_s_l_shadow[split_s_idx];
for (int d = 0; d < split_info->n_device; ++d) {
@@ -1353,8 +1376,8 @@ bool llama_kv_cache::checkpoint_save() {
}
split_s_idx++;
} else {
// Non-split tensor
ggml_backend_tensor_copy(s_l[il], ckpt.s_l_shadow[il]);
const size_t nbytes = ggml_nbytes(ckpt.s_l_shadow[il]);
ggml_backend_tensor_get(s_l[il], ckpt.s_l_shadow[il]->data, 0, nbytes);
}
}
@@ -1370,12 +1393,10 @@ bool llama_kv_cache::checkpoint_restore() {
const uint32_t n_layer = (uint32_t)s_l.size();
// 1. Restore cell metadata
cells = ckpt.cells_snapshot;
head = ckpt.head_snapshot;
used = ckpt.used_snapshot;
// 2. Copy tensors: shadow -> primary (GPU-to-GPU)
uint32_t split_s_idx = 0;
for (uint32_t il = 0; il < n_layer; ++il) {
if (s_l[il] == nullptr) {
@@ -1383,7 +1404,6 @@ bool llama_kv_cache::checkpoint_restore() {
}
if (s_l[il]->extra != nullptr) {
// Split tensor — copy per-device splits individually
auto * split_info = (const ggml_split_tensor_t *)s_l[il]->extra;
auto & shadow_split = ckpt.split_s_l_shadow[split_s_idx];
for (int d = 0; d < split_info->n_device; ++d) {
@@ -1393,7 +1413,7 @@ bool llama_kv_cache::checkpoint_restore() {
}
split_s_idx++;
} else {
// Non-split tensor
GGML_ASSERT(ggml_nbytes(ckpt.s_l_shadow[il]) == ggml_nbytes(s_l[il]));
ggml_backend_tensor_copy(ckpt.s_l_shadow[il], s_l[il]);
}
}
@@ -1405,6 +1425,159 @@ void llama_kv_cache::checkpoint_delete() {
ckpt.saved = false;
}
bool llama_kv_cache::per_step_alloc(int max_tokens) {
if (ckpt.per_step_max_allocated >= max_tokens) {
return true;
}
if (!ckpt.per_step_ssm.empty()) {
for (struct ggml_context * ctx : ckpt.per_step_ctxs) {
ggml_free(ctx);
}
for (ggml_backend_buffer_t buf : ckpt.per_step_bufs) {
ggml_backend_buffer_free(buf);
}
ckpt.per_step_ctxs.clear();
ckpt.per_step_bufs.clear();
ckpt.per_step_ssm.clear();
ckpt.per_step_qkv.clear();
ckpt.per_step_max_allocated = 0;
}
const uint32_t n_layer = (uint32_t)s_l.size();
ckpt.per_step_ssm.resize(n_layer, nullptr);
ckpt.per_step_qkv.resize(n_layer, nullptr);
const int64_t ssm_state_dim = ckpt.per_step_ssm_state_size;
const int64_t conv_dim = ckpt.per_step_conv_dim;
if (ssm_state_dim <= 0 || conv_dim <= 0) {
LLAMA_LOG_ERROR("%s: per_step dimensions not set (ssm=%lld, conv_dim=%lld)\n",
__func__, (long long)ssm_state_dim, (long long)conv_dim);
return false;
}
std::map<ggml_backend_buffer_type_t, std::vector<std::pair<uint32_t, ggml_backend_buffer_type_t>>> buft_layers;
for (uint32_t il = 0; il < n_layer; ++il) {
if (s_l[il] == nullptr) continue;
if (s_l[il]->extra != nullptr) continue; // skip split tensors
ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(s_l[il]->buffer);
buft_layers[buft].push_back({il, buft});
}
for (auto & [buft, layers] : buft_layers) {
// 2 tensors per layer: SSM states + qkv features
ggml_init_params params = {
/*.mem_size =*/ layers.size() * 2 * ggml_tensor_overhead(),
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ true,
};
ggml_context * ctx = ggml_init(params);
if (!ctx) {
LLAMA_LOG_ERROR("%s: failed to create ggml context for per-step checkpoints\n", __func__);
return false;
}
for (auto & [il, bt] : layers) {
// SSM state: max_tokens * ssm_state_dim
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] = 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);
ggml_format_name(t_qkv, "per_step_qkv_l%d", il);
ckpt.per_step_qkv[il] = t_qkv;
}
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft);
if (!buf) {
LLAMA_LOG_ERROR("%s: failed to allocate buffer for per-step checkpoints\n", __func__);
ggml_free(ctx);
return false;
}
ggml_backend_buffer_clear(buf, 0);
LLAMA_LOG_INFO("%s: %10s per-step buffer = %8.2f MiB (max_tokens=%d, ssm+qkv)\n", __func__,
ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf) / 1024.0 / 1024.0, max_tokens);
ckpt.per_step_ctxs.push_back(ctx);
ckpt.per_step_bufs.push_back(buf);
}
ckpt.per_step_max_allocated = max_tokens;
return true;
}
bool llama_kv_cache::per_step_restore(int step) {
if (ckpt.per_step_ssm.empty() || step < 0) {
return false;
}
const int64_t ssm_state_dim = ckpt.per_step_ssm_state_size;
const int64_t conv_state_dim = ckpt.per_step_conv_state_dim;
const int64_t conv_dim = ckpt.per_step_conv_dim;
const int32_t d_conv = ckpt.per_step_d_conv;
if (ssm_state_dim <= 0 || conv_dim <= 0 || d_conv <= 1) return false;
const int64_t ssm_bytes = ssm_state_dim * sizeof(float);
const int64_t conv_bytes = conv_state_dim * sizeof(float);
const int32_t d_conv_m1 = d_conv - 1; // number of columns in conv state
std::vector<float> ssm_buf(ssm_state_dim);
std::vector<float> conv_buf(conv_state_dim); // reconstructed conv state
std::vector<float> old_conv_buf(conv_state_dim); // pre-spec conv state from shadow
const int64_t qkv_needed = (int64_t)(step + 1) * conv_dim;
std::vector<float> qkv_buf(qkv_needed);
const uint32_t n_layer = (uint32_t)s_l.size();
int n_restored = 0;
for (uint32_t il = 0; il < n_layer; ++il) {
if (s_l[il] == nullptr || ckpt.per_step_ssm[il] == nullptr) continue;
if (s_l[il]->extra != nullptr) continue;
ggml_backend_tensor_get(ckpt.per_step_ssm[il], ssm_buf.data(),
(size_t)step * ssm_bytes, ssm_bytes);
if (ckpt.s_l_shadow[il] != nullptr) {
ggml_backend_tensor_get(ckpt.s_l_shadow[il], old_conv_buf.data(), 0, conv_bytes);
} else {
memset(old_conv_buf.data(), 0, conv_bytes);
}
if (ckpt.per_step_qkv[il] != nullptr) {
ggml_backend_tensor_get(ckpt.per_step_qkv[il], qkv_buf.data(), 0, qkv_needed * sizeof(float));
} else {
memset(qkv_buf.data(), 0, qkv_needed * sizeof(float));
}
for (int32_t col = 0; col < d_conv_m1; col++) {
int32_t src_token = step - (d_conv_m1 - 1) + col; // e.g., K-2, K-1, K for d_conv=4
if (src_token >= 0) {
for (int64_t d = 0; d < conv_dim; d++) {
conv_buf[col + d * d_conv_m1] = qkv_buf[d + (int64_t)src_token * conv_dim];
}
} else {
int32_t old_col = d_conv_m1 + src_token; // maps to 0, 1, ... for early steps
if (old_col >= 0 && old_col < d_conv_m1) {
for (int64_t d = 0; d < conv_dim; d++) {
conv_buf[col + d * d_conv_m1] = old_conv_buf[old_col + d * d_conv_m1];
}
} else {
for (int64_t d = 0; d < conv_dim; d++) {
conv_buf[col + d * d_conv_m1] = 0.0f;
}
}
}
}
ggml_backend_tensor_set(s_l[il], conv_buf.data(), 0, conv_bytes);
ggml_backend_tensor_set(s_l[il], ssm_buf.data(), conv_bytes, ssm_bytes);
n_restored++;
}
return true;
}
static void llama_kv_cache_clear(struct llama_kv_cache & cache) {
for (int32_t i = 0; i < (int32_t) cache.size; ++i) {
cache.cells[i].pos = -1;
@@ -6676,6 +6849,81 @@ bool llama_kv_cache_checkpoint_supported(struct llama_context * ctx) {
return ctx->kv_self.checkpoint_supported();
}
bool llama_kv_cache_set_per_step_save(struct llama_context * ctx, bool enable, int max_tokens) {
auto & kv = ctx->kv_self;
if (!enable) {
kv.save_per_step_ssm = false;
return false;
}
// Check if any recurrent layer has split state (graph split mode).
// Per-step restore can't handle split tensors — fall back to legacy path.
// Also require ALL recurrent layers to be on GPU (non-host backend),
// since per-step saves only work in the CUDA delta_net kernel.
// If any recurrent layer is on CPU, the per-step restore would leave
// CPU layers at the pre-spec position while GPU layers advance to step K.
bool has_gpu_layer = false;
bool has_cpu_layer = false;
for (const auto * sl : kv.s_l) {
if (sl == nullptr) continue;
if (sl->extra != nullptr) {
kv.save_per_step_ssm = false;
return false;
}
if (sl->buffer && !ggml_backend_buffer_is_host(sl->buffer)) {
has_gpu_layer = true;
} else if (sl->buffer) {
has_cpu_layer = true;
}
}
if (!has_gpu_layer || has_cpu_layer) {
// All recurrent layers on CPU or mixed CPU/GPU — per-step not supported.
if (has_cpu_layer && has_gpu_layer) {
LLAMA_LOG_INFO("%s: per-step disabled — mixed CPU/GPU recurrent layers\n", __func__);
}
kv.save_per_step_ssm = false;
return false;
}
// Initialize per-step checkpoint dimensions from model hparams (if not already set)
if (kv.ckpt.per_step_ssm_state_size <= 0) {
const auto & hparams = ctx->model.hparams;
const int64_t num_v_heads = hparams.ssm_dt_rank;
const int64_t head_v_dim = hparams.ssm_d_inner / num_v_heads;
const int64_t head_k_dim = hparams.ssm_d_state;
const int64_t num_k_heads = hparams.ssm_n_group;
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;
kv.ckpt.per_step_ssm_state_size = head_v_dim * head_v_dim * num_v_heads;
kv.ckpt.per_step_conv_state_dim = (hparams.ssm_d_conv - 1) * conv_dim;
kv.ckpt.per_step_conv_dim = conv_dim;
kv.ckpt.per_step_d_conv = hparams.ssm_d_conv;
}
// Allocate (or reallocate if max_tokens is larger than current)
if (!kv.per_step_alloc(max_tokens)) {
kv.save_per_step_ssm = false;
return false;
}
kv.save_per_step_ssm = true;
return true;
}
bool llama_kv_cache_per_step_restore(struct llama_context * ctx, int step) {
return ctx->kv_self.per_step_restore(step);
}
void llama_kv_cache_set_cell_pos(struct llama_context * ctx, llama_seq_id seq_id, llama_pos pos) {
auto & kv = ctx->kv_self;
if (seq_id >= 0 && (uint32_t)seq_id < kv.size) {
kv.cells[seq_id].pos = pos;
}
}
bool llama_kv_cache_seq_rm(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
return llama_kv_cache_seq_rm(ctx->kv_self, seq_id, p0, p1);
}