diff --git a/common/arg.cpp b/common/arg.cpp index 79480e06f9d2..9bda659a2621 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1624,6 +1624,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.checkpoint_min_step = value; } ).set_env("LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--rs-aligned"}, "N", + string_format("number of boundary-aligned deep-rollback state slots per sequence, for memories that support them (default: %d, 0 = disabled)", params.n_rs_aligned), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("rs-aligned must be non-negative"); + } + params.n_rs_aligned = value; + } + ).set_env("LLAMA_ARG_RS_ALIGNED").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-cram", "--cache-ram"}, "N", string_format("set the maximum cache size in MiB (default: %d, -1 - no limit, 0 - disable)" diff --git a/common/common.cpp b/common/common.cpp index ff27d392fb2e..b94464eaf17d 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1631,6 +1631,12 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.n_ctx = params.n_ctx; cparams.n_seq_max = params.n_parallel; cparams.n_rs_seq = params.speculative.need_n_rs_seq(); + cparams.n_rs_aligned = std::max(params.n_rs_aligned, 0); + if (cparams.n_rs_aligned > 0 && cparams.n_rs_seq == 0) { + // the aligned deep-rollback tier rides the per-token rollback machinery, + // and single-token re-eval removals below the tip need per-token depth 1 + cparams.n_rs_seq = 1; + } cparams.n_outputs_max = std::max(params.n_outputs_max, 0); cparams.n_batch = params.n_batch; cparams.n_ubatch = params.n_ubatch; diff --git a/common/common.h b/common/common.h index 919c0ea103a4..70a488762183 100644 --- a/common/common.h +++ b/common/common.h @@ -620,6 +620,7 @@ struct common_params { bool cache_idle_slots = true; // save and clear idle slots upon starting a new task int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints + int32_t n_rs_aligned = 0; // boundary-aligned deep-rollback slots per seq (DSV4-class memories, 0 = disabled) int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. std::string hostname = "127.0.0.1"; diff --git a/common/speculative.cpp b/common/speculative.cpp index 514fe852180f..c5bae8ba024a 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2325,8 +2325,9 @@ common_speculative_init_result::common_speculative_init_result( // note: for small models maybe we can set this to the maximum possible draft from all speculative types // the extra memory for small models is likely negligible? - cparams.n_rs_seq = 0; - cparams.ctx_other = ctx_tgt; + cparams.n_rs_seq = 0; + cparams.n_rs_aligned = 0; + cparams.ctx_other = ctx_tgt; std::string model_path; if (has_draft) { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 7f4e252dca39..ce6ab7286f3c 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -754,7 +754,7 @@ static bool ggml_is_view_op(enum ggml_op op) { #endif #ifndef GGML_SCHED_MAX_SPLIT_INPUTS -#define GGML_SCHED_MAX_SPLIT_INPUTS 30 +#define GGML_SCHED_MAX_SPLIT_INPUTS 48 #endif #ifndef GGML_SCHED_MAX_COPIES diff --git a/include/llama.h b/include/llama.h index 6e53e2297235..e13bd35e77f8 100644 --- a/include/llama.h +++ b/include/llama.h @@ -353,6 +353,7 @@ extern "C" { uint32_t n_ubatch; // physical maximum batch size uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] + uint32_t n_rs_aligned; // number of boundary-aligned deep-rollback slots per seq (0 = per-token tier only; requires n_rs_seq > 0) [EXPERIMENTAL] uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing @@ -557,6 +558,7 @@ extern "C" { LLAMA_API uint32_t llama_n_ubatch (const struct llama_context * ctx); LLAMA_API uint32_t llama_n_seq_max (const struct llama_context * ctx); LLAMA_API uint32_t llama_n_rs_seq (const struct llama_context * ctx); + LLAMA_API uint32_t llama_n_rs_aligned(const struct llama_context * ctx); DEPRECATED(LLAMA_API int32_t llama_n_ctx_train(const struct llama_model * model), "use llama_model_n_ctx_train instead"); DEPRECATED(LLAMA_API int32_t llama_n_embd (const struct llama_model * model), "use llama_model_n_embd instead"); @@ -791,6 +793,13 @@ extern "C" { // Check if the memory supports shifting LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); + // Position alignment of the memory's deep-rollback tier: in addition to the + // bounded per-token tier (see llama_n_rs_seq), partial sequence removal below + // the current tip may be requested at positions that are multiples of this + // value, subject to snapshot coverage (llama_memory_seq_rm reports acceptance). + // Returns 1 when the memory has no aligned deep-rollback tier. [EXPERIMENTAL] + LLAMA_API llama_pos llama_memory_seq_rm_align(llama_memory_t mem); + // // State / sessions // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 19cca7df1e9d..754e764d012d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -107,6 +107,14 @@ llama_context::llama_context( cparams.n_rs_seq = 0; } + // the aligned deep-rollback tier rides the per-token rollback machinery + cparams.n_rs_aligned = params.n_rs_aligned; + if (cparams.n_rs_aligned > 0 && cparams.n_rs_seq == 0) { + LLAMA_LOG_DEBUG("%s: n_rs_aligned=%u requested without n_rs_seq; clamping to 0\n", + __func__, cparams.n_rs_aligned); + cparams.n_rs_aligned = 0; + } + cparams.n_threads = params.n_threads; cparams.n_threads_batch = params.n_threads_batch; cparams.yarn_ext_factor = params.yarn_ext_factor >= 0.0f ? params.yarn_ext_factor : hparams.yarn_ext_factor; @@ -3487,6 +3495,7 @@ llama_context_params llama_context_default_params() { /*.n_ubatch =*/ 512, /*.n_seq_max =*/ 1, /*.n_rs_seq =*/ 0, + /*.n_rs_aligned =*/ 0, /*.n_outputs_max =*/ 0, /*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default /*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS, @@ -3649,6 +3658,10 @@ uint32_t llama_n_seq_max(const llama_context * ctx) { return ctx->n_seq_max(); } +uint32_t llama_n_rs_aligned(const llama_context * ctx) { + return ctx->get_cparams().n_rs_aligned; +} + uint32_t llama_n_rs_seq(const llama_context * ctx) { return ctx->get_cparams().n_rs_seq; } @@ -3968,6 +3981,14 @@ bool llama_memory_can_shift(llama_memory_t mem) { return mem->get_can_shift(); } +llama_pos llama_memory_seq_rm_align(llama_memory_t mem) { + if (!mem) { + return 1; + } + + return mem->get_seq_rm_align(); +} + // llama state API // deprecated diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed85e..0d7a29b76948 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -14,6 +14,7 @@ struct llama_cparams { uint32_t n_ubatch; uint32_t n_seq_max; uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback + uint32_t n_rs_aligned; // number of boundary-aligned deep-rollback slots per seq uint32_t n_outputs_max; // max outputs supported by the context int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 5caa05e8b07d..81a6bb4b16aa 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -432,6 +432,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( uint32_t kv_size, uint32_t n_stream, uint32_t n_rs_seq, + uint32_t n_aligned, const std::vector & rs_idx) { llama_kv_cache_dsv4_context::comp_plan plan; plan.n_visible.resize(ubatch.n_tokens); @@ -631,7 +632,10 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size); const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0; // Keep the restore graph fixed-width when no rollback is pending. - const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0; + // Markers beyond this state's planes (an aligned marker on a state + // without aligned slots) degrade to a no-op self-copy: such a state + // has nothing to restore, its rows are rewritten by the replay. + const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq + n_aligned ? (int64_t) rollback*state_rows : 0; for (uint32_t r = 0; r < state_size; ++r) { plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r)); plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r)); @@ -649,7 +653,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( } const uint32_t n_seq_tokens = (uint32_t) token_idxs.size(); - const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq); + const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq + n_aligned); for (uint32_t d = 1; d <= n_rs_seq; ++d) { const int64_t dst_plane = (int64_t) d*state_rows; @@ -674,6 +678,40 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r)); } } + + // Aligned-tier snapshots: when this ubatch commits an HCA boundary + // for this seq, capture the ring as of that boundary into the + // boundary's aligned plane slot. One write per slot (the latest + // boundary wins) so ggml_set_rows never sees duplicate rows. + if (n_aligned > 0) { + std::map slot_last_j; // slot -> index into token_idxs + for (uint32_t j = 0; j < n_seq_tokens; ++j) { + const llama_pos pos = ubatch.pos[token_idxs[j]]; + if (pos < 0 || (pos + 1) % DSV4_HCA_RATIO != 0) { + continue; + } + const uint32_t slot = (uint32_t) (((pos + 1)/DSV4_HCA_RATIO) % n_aligned); + slot_last_j[slot] = j; + } + + for (const auto & [slot, j_boundary] : slot_last_j) { + const int64_t dst_plane = (int64_t) (1 + n_rs_seq + slot)*state_rows; + const uint32_t prefix = j_boundary + 1; + + for (uint32_t r = 0; r < state_size; ++r) { + int32_t src = (int32_t) (stream_off + r); + for (uint32_t j = 0; j < prefix; ++j) { + const uint32_t i_tok = token_idxs[j]; + if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) { + src = (int32_t) (scratch_off + i_tok); + } + } + + plan.state_snapshot_src_idxs.push_back(src); + plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r)); + } + } + } } } @@ -700,12 +738,27 @@ static std::vector dsv4_build_comp_plans uint32_t kv_size, uint32_t n_stream, uint32_t n_rs_seq, + uint32_t n_aligned, const std::vector & rs_idx) { std::vector plans; plans.reserve(ubatches.size()); + // A pending restore belongs to the first ubatch of its seq only: replaying + // it in later ubatches would clobber ring rows the earlier ubatches have + // already persisted for their own positions. + std::vector rs_remaining = rs_idx; + for (const llama_ubatch & ubatch : ubatches) { - plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx)); + plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, n_aligned, rs_remaining)); + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { + const llama_seq_id seq_id = ubatch.seq_id[i][s]; + if (seq_id >= 0 && (size_t) seq_id < rs_remaining.size()) { + rs_remaining[seq_id] = 0; + } + } + } } return plans; @@ -793,7 +846,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( uint32_t state_size, uint32_t kv_size, uint32_t n_stream, - uint32_t n_rs_seq) { + uint32_t n_rs_seq, + uint32_t n_aligned) { llama_kv_cache_dsv4_context::comp_plan plan; plan.n_visible.resize(ubatch.n_tokens); plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream); @@ -812,7 +866,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( const uint64_t state_rows = (uint64_t) state_size*n_stream; const size_t n_persist = (size_t) std::min(ubatch.n_tokens, state_rows); const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max(1, ubatch.n_seqs_unq) : 0; - const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max(1, ubatch.n_seqs_unq); + const size_t n_aligned_writes = n_aligned > 0 ? (size_t) std::min(n_aligned, n_seq_tokens/128 + 1) : 0; + const size_t n_snapshot = (size_t) (n_rs_seq + n_aligned_writes)*state_size*std::max(1, ubatch.n_seqs_unq); plan.state_pos .resize(ubatch.n_tokens); plan.state_persist_src_idxs.resize(n_persist); @@ -847,13 +902,15 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( uint32_t state_size, uint32_t n_embd_state, uint32_t n_rs_seq, + uint32_t n_aligned, const char * name, const llama_memory_i::layer_filter_cb & filter) : ratio(ratio), state_size(state_size), n_embd_state(n_embd_state), n_stream(unified ? 1 : n_seq_max), - n_rs_seq(n_rs_seq) { + n_rs_seq(n_rs_seq), + n_aligned(n_aligned) { const llama_hparams & hparams = model.hparams; struct ggml_backend_buft_comparator { @@ -909,7 +966,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( throw std::runtime_error("failed to create ggml context for DSV4 compressor state"); } - const uint32_t n_planes = n_stream*(1 + n_rs_seq); + const uint32_t n_planes = n_stream*(1 + n_rs_seq + n_aligned); ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes); ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes); @@ -956,7 +1013,7 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) { GGML_ASSERT((uint32_t) seq_id < n_stream); for (const auto & layer : layers) { - for (uint32_t d = 0; d <= n_rs_seq; ++d) { + for (uint32_t d = 0; d <= n_rs_seq + n_aligned; ++d) { const uint32_t stream = d*n_stream + (uint32_t) seq_id; dsv4_clear_tensor_stream(layer.kv, stream); dsv4_clear_tensor_stream(layer.score, stream); @@ -1039,10 +1096,15 @@ void llama_dsv4_comp_state::state_write( std::vector stream_ids(ns); for (uint32_t s = 0; s < ns; ++s) { const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s; - if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) { + if (seq >= rs_idx.size()) { throw std::runtime_error("DSV4 recurrent state rollback index out of range"); } - stream_ids[s] = rs_idx[seq]*n_stream + s0 + s; + // Markers beyond this state's planes (an aligned marker on a state + // without aligned slots) degrade to the live plane - such a state has + // nothing snapshot-held to save, its rows are rewritten by the replay. + // Mirrors the restore graph's self-copy rule in dsv4_build_comp_plan. + const uint32_t plane = rs_idx[seq] <= n_rs_seq + n_aligned ? rs_idx[seq] : 0; + stream_ids[s] = plane*n_stream + s0 + s; } const uint32_t version = DSV4_COMP_STATE_VER; @@ -1109,14 +1171,18 @@ ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const int32_t ids = map_layer_ids.at(il); ggml_tensor * state = layers[ids].kv; - return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0); + return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq + n_aligned), state->nb[1], 0); +} + +uint32_t llama_dsv4_comp_state::get_n_aligned() const { + return n_aligned; } ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const { const int32_t ids = map_layer_ids.at(il); ggml_tensor * state = layers[ids].score; - return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0); + return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq + n_aligned), state->nb[1], 0); } ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const { @@ -1168,6 +1234,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( uint32_t n_ubatch, uint32_t n_pad, uint32_t n_rs_seq, + uint32_t n_rs_aligned, const layer_filter_cb & filter, const layer_reuse_cb & reuse) : hparams_raw(model.hparams), @@ -1176,7 +1243,9 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( hparams_lid(model.hparams), n_seq_max(n_seq_max), n_rs_seq(n_rs_seq), - rs_idx(n_seq_max, 0) { + n_rs_aligned(n_rs_aligned), + rs_idx(n_seq_max, 0), + aligned_pos(n_seq_max, std::vector(n_rs_aligned, -1)) { const layer_filter_cb filter_raw = [&](int32_t il) { if (filter && !filter(il)) { @@ -1262,19 +1331,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( csa_state = std::make_unique( model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, - 2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa); + 2*model.hparams.n_embd_head_k(), n_rs_seq, n_rs_aligned, "csa", filter_csa); LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__); hca_state = std::make_unique( model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO, - model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca); + model.hparams.n_embd_head_k(), n_rs_seq, 0, "hca", filter_hca); LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__); lid_state = std::make_unique( model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, - 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa); + 2*model.hparams.indexer_head_size, n_rs_seq, n_rs_aligned, "lid", filter_csa); // DSV4 attention reads compressed-K / compressor-state rows that the current // graph does not necessarily overwrite; uninitialized buffer contents would @@ -1391,6 +1460,11 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, std::move(lid_state->sc_info)); } +llama_pos llama_kv_cache_dsv4::get_seq_rm_align() const { + // the aligned deep-rollback tier snapshots at HCA block boundaries + return n_rs_aligned > 0 ? (llama_pos) DSV4_HCA_RATIO : 1; +} + bool llama_kv_cache_dsv4::get_can_shift() const { // Compressed row metadata uses block-derived positions. Keep shifting // disabled until DSV4 compressed-cache shift semantics are wired. @@ -1428,14 +1502,49 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 return false; } + // A pending un-consumed rollback must be replayed before another can + // be accepted: recomputing depth from the already-purged pos_max would + // silently restore the wrong plane. + if (rs_idx[seq_id] != 0) { + return false; + } + const llama_pos rollback = pos_max - (p0 - 1); - if (rollback < 1 || rollback > (llama_pos) n_rs_seq) { + if (rollback < 1) { + return false; + } + + uint32_t marker = 0; + if (rollback <= (llama_pos) n_rs_seq) { + marker = (uint32_t) rollback; + } else if (n_rs_aligned > 0 && p0 % (llama_pos) DSV4_HCA_RATIO == 0) { + // aligned deep tier: the target boundary must still hold a snapshot + const uint32_t slot = (uint32_t) ((p0/DSV4_HCA_RATIO) % n_rs_aligned); + if (aligned_pos[seq_id][slot] != p0) { + return false; + } + marker = n_rs_seq + 1 + slot; + } else { + return false; + } + + // Never apply a rollback the raw SWA window can no longer cover: the + // replay from p0 reads raw K back to p0 - n_swa, and reclaimed cells + // would silently produce wrong output. + const llama_pos swa_min = kv_raw->get_swa()->seq_pos_min(seq_id); + if (swa_min < 0 || swa_min > std::max(0, p0 - (llama_pos) hparams_raw.n_swa)) { return false; } const bool res = kv_raw->seq_rm(seq_id, p0, p1); if (res) { - rs_idx[seq_id] = (uint32_t) rollback; + rs_idx[seq_id] = marker; + // aligned slots holding rolled-back future boundaries are dead + for (llama_pos & pos : aligned_pos[seq_id]) { + if (pos > p0) { + pos = -1; + } + } } return res; @@ -1464,6 +1573,7 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds if (seq_id_src != seq_id_dst) { rs_idx[seq_id_dst] = 0; + std::fill(aligned_pos[seq_id_dst].begin(), aligned_pos[seq_id_dst].end(), -1); } } @@ -1610,8 +1720,12 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, if (seq_id >= 0) { GGML_ASSERT((uint32_t) seq_id < n_seq_max); rs_idx[seq_id] = 0; + std::fill(aligned_pos[seq_id].begin(), aligned_pos[seq_id].end(), -1); } else { std::fill(rs_idx.begin(), rs_idx.end(), 0); + for (auto & slots : aligned_pos) { + std::fill(slots.begin(), slots.end(), -1); + } } } @@ -1647,6 +1761,10 @@ uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const { return n_rs_seq; } +uint32_t llama_kv_cache_dsv4::get_n_rs_aligned() const { + return n_rs_aligned; +} + const std::vector & llama_kv_cache_dsv4::get_rs_idx() const { return rs_idx; } @@ -1668,6 +1786,29 @@ void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector & ubatches) { + if (n_rs_aligned == 0) { + return; + } + + for (const llama_ubatch & ubatch : ubatches) { + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + const llama_pos pos = ubatch.pos[i]; + if (pos < 0 || (pos + 1) % (llama_pos) DSV4_HCA_RATIO != 0) { + continue; + } + const llama_pos boundary = pos + 1; + const uint32_t slot = (uint32_t) ((boundary/DSV4_HCA_RATIO) % n_rs_aligned); + for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { + const llama_seq_id seq_id = ubatch.seq_id[i][s]; + if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) { + aligned_pos[seq_id][slot] = boundary; + } + } + } + } +} + void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { if (seq_id < 0) { kv_csa->clear(data); @@ -1697,8 +1838,12 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { if (seq_id >= 0) { rs_idx[seq_id] = 0; + std::fill(aligned_pos[seq_id].begin(), aligned_pos[seq_id].end(), -1); } else { std::fill(rs_idx.begin(), rs_idx.end(), 0); + for (auto & slots : aligned_pos) { + std::fill(slots.begin(), slots.end(), -1); + } } } @@ -1992,13 +2137,13 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( ubatches(std::move(ubatches)), plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true, kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(), - kv->get_n_rs_seq(), kv->get_rs_idx())), + kv->get_n_rs_seq(), kv->get_csa_state()->get_n_aligned(), kv->get_rs_idx())), plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false, kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(), - kv->get_n_rs_seq(), kv->get_rs_idx())), + kv->get_n_rs_seq(), kv->get_hca_state()->get_n_aligned(), kv->get_rs_idx())), plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true, kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(), - kv->get_n_rs_seq(), kv->get_rs_idx())), + kv->get_n_rs_seq(), kv->get_lid_state()->get_n_aligned(), kv->get_rs_idx())), ctx_raw(std::make_unique( kv->get_raw(), std::move(sinfos_raw_base_write), @@ -2025,6 +2170,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( hca_state(kv->get_hca_state()), lid_state(kv->get_lid_state()), status(ctx_raw->get_status()) { + kv->note_aligned_snapshots_for_ubatches(this->ubatches); kv->reset_rs_idx_for_ubatches(this->ubatches); } @@ -2161,7 +2307,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ reserve_plan_csa = dsv4_build_reserve_comp_plan( ubatch, DSV4_CSA_RATIO, true, - csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq()); + csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq(), csa_state->get_n_aligned()); return reserve_plan_csa; } @@ -2175,7 +2321,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ reserve_plan_hca = dsv4_build_reserve_comp_plan( ubatch, DSV4_HCA_RATIO, false, - hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq()); + hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq(), hca_state->get_n_aligned()); return reserve_plan_hca; } @@ -2189,7 +2335,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ reserve_plan_lid = dsv4_build_reserve_comp_plan( ubatch, DSV4_CSA_RATIO, true, - lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq()); + lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq(), lid_state->get_n_aligned()); return reserve_plan_lid; } diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index ce39867c0342..eb5dec2a9da8 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -23,6 +23,7 @@ class llama_dsv4_comp_state { uint32_t state_size, uint32_t n_embd_state, uint32_t n_rs_seq, + uint32_t n_aligned, const char * name, const llama_memory_i::layer_filter_cb & filter); @@ -34,6 +35,7 @@ class llama_dsv4_comp_state { uint32_t get_state_size() const; uint32_t get_n_stream() const; uint32_t get_n_rs_seq() const; + uint32_t get_n_aligned() const; uint32_t get_n_rows() const; std::map memory_breakdown() const; @@ -65,6 +67,7 @@ class llama_dsv4_comp_state { const uint32_t n_embd_state; const uint32_t n_stream; const uint32_t n_rs_seq; + const uint32_t n_aligned; std::vector> ctxs_bufs; @@ -100,6 +103,7 @@ class llama_kv_cache_dsv4 : public llama_memory_i { uint32_t n_ubatch, uint32_t n_pad, uint32_t n_rs_seq, + uint32_t n_rs_aligned, const layer_filter_cb & filter, const layer_reuse_cb & reuse); @@ -120,6 +124,8 @@ class llama_kv_cache_dsv4 : public llama_memory_i { bool get_can_shift() const override; + llama_pos get_seq_rm_align() const override; + void clear(bool data) override; bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; @@ -149,8 +155,10 @@ class llama_kv_cache_dsv4 : public llama_memory_i { llama_dsv4_comp_state * get_lid_state() const; uint32_t get_n_rs_seq() const; + uint32_t get_n_rs_aligned() const; const std::vector & get_rs_idx() const; void reset_rs_idx_for_ubatches(const std::vector & ubatches); + void note_aligned_snapshots_for_ubatches(const std::vector & ubatches); private: llama_hparams hparams_raw; @@ -160,9 +168,15 @@ class llama_kv_cache_dsv4 : public llama_memory_i { const uint32_t n_seq_max; const uint32_t n_rs_seq; + const uint32_t n_rs_aligned; + // rs_idx[seq] encodes the pending rollback restore: 0 = none, + // 1..n_rs_seq = per-token plane, n_rs_seq+1+slot = aligned plane slot std::vector rs_idx; + // boundary position held by each aligned slot per seq, -1 = empty + std::vector> aligned_pos; + std::unique_ptr kv_raw; std::unique_ptr kv_csa; std::unique_ptr kv_hca; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..cd0f0e1ee58b 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -100,6 +100,9 @@ struct llama_memory_i { // getters virtual bool get_can_shift() const = 0; + // position alignment of the memory's deep-rollback tier (1 = no aligned tier) + virtual llama_pos get_seq_rm_align() const { return 1; } + // // ops // diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 3812c594e795..c9c6df5a4f1a 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -214,8 +214,8 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_chexp); - add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp); - add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp); + add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, true); + add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, true); add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res); // add_kv(LLM_KV_TENSOR_DATA_LAYOUT, ???); add_kv(LLM_KV_EXPERT_COUNT, hparams.n_expert); @@ -285,6 +285,14 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true); add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true); + add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); + add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); + add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios); + add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base); + add_kv(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train; @@ -407,6 +415,9 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->cls_out); add_tensor(model->cls_out_b); add_tensor(model->cls_norm); + add_tensor(model->hc_head_fn); + add_tensor(model->hc_head_base); + add_tensor(model->hc_head_scale); for (const struct llama_layer & layer : model->layers) { for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e93641b63ddc..00331cc3834c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2177,6 +2177,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.n_ubatch, 1, cparams.n_rs_seq, + cparams.n_rs_aligned, nullptr, nullptr); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 881e55c75a1d..94ed8bc4765f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -199,11 +199,13 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") + # fixed seed: the recurrent-state rollback cases compare logits against + # references at a fixed epsilon, which must not vary run to run llama_test( test-llama-archs NAME test-generate-models LABEL main - ARGS -o "${MODEL_DIR}" + ARGS -o "${MODEL_DIR}" -s 1234 ) set_tests_properties(test-generate-models PROPERTIES FIXTURES_SETUP generate-models @@ -217,6 +219,18 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set_tests_properties(test-recurrent-state-rollback PROPERTIES FIXTURES_REQUIRED generate-models ) + + # CPU only: with n_rs_seq enabled the DSV4 graph's per-plane state inputs + # exceed GGML_SCHED_MAX_SPLIT_INPUTS when a single device takes the whole graph + llama_test( + test-recurrent-state-rollback + NAME test-recurrent-state-rollback-dsv4 + LABEL main + ARGS -m "${MODEL_DIR}/deepseek4-moe.gguf" --device none + ) + set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES + FIXTURES_REQUIRED generate-models + ) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 4336e4e13d4f..eb105a2eedcb 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -9,6 +9,7 @@ // TODO: replace with #include "llama-ext.h" in the future #include "../src/llama-arch.h" +#include "../src/llama-hparams.h" #include "../src/llama-model-saver.h" #include @@ -39,12 +40,18 @@ static double nmse(const std::vector & a, const std::vector & b) { return mse_a_b / mse_a_0; } +struct tensor_data_params { + size_t seed; + float sigma; +}; + static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { - size_t seed = *(const size_t *) userdata; + const tensor_data_params & tdp = *(const tensor_data_params *) userdata; + size_t seed = tdp.seed; std::hash hasher; seed ^= hasher(tensor->name); std::mt19937 gen(seed); - std::normal_distribution dis(0.0f, 1.0e-2f); + std::normal_distribution dis(0.0f, tdp.sigma); const int64_t ne = ggml_nelements(tensor); if (tensor->type == GGML_TYPE_F32) { @@ -109,6 +116,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_embd = 128; n_head = 1; n_ff = 192; + } else if (arch == LLM_ARCH_DEEPSEEK4) { + n_embd = 128; + n_head = 2; + n_ff = 192; + n_layer = 3; // one layer per compress ratio: raw, CSA, HCA } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { n_layer = 3; } else if (arch == LLM_ARCH_CHAMELEON) { @@ -168,6 +180,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); + } else if (arch == LLM_ARCH_DEEPSEEK4) { + // the graph requires key_length == value_length, K doubles as V in the k-only caches + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(64)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(64)); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(16)); } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -224,6 +241,23 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERTS_PER_GROUP, uint32_t(1)); } + if (arch == LLM_ARCH_DEEPSEEK4) { + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, uint32_t(1)); // MQA: a single shared KV vector per position + ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, n_ctx); // the raw window must cover positions the HCA has not yet compressed + ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, uint32_t(LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS)); // the loader rejects other scoring + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.5f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(2)); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(16)); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); // the graph builder asserts hc == 4 + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(20)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); + ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0)); // hash routing reads an I32 token->expert table, incompatible with random weights + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 4, 128})); + } + ms.add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH, n_embd); ms.add_kv(LLM_KV_POSNET_BLOCK_COUNT, n_layer); ms.add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, n_embd); @@ -261,9 +295,19 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) return true; } +// The saved DEEPSEEK4 fixture needs weights big enough that attention output +// is visible in the logits: the recurrent-state rollback test compares +// replayed logits against a reference, which is vacuous if the context +// contribution drowns in fp noise. Only the saved fixture gets the larger +// sigma - the backend NMSE comparison keeps the default, where cross-backend +// divergence stays within its threshold. +static float weight_sigma(const llm_arch arch) { + return arch == LLM_ARCH_DEEPSEEK4 ? 1.0e-1f : 1.0e-2f; +} + static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, float sigma = 1.0e-2f) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -280,7 +324,7 @@ static std::pair get_model_and_ctx( ctx_params.n_ubatch = 64; } - size_t tmp = seed; + tensor_data_params tmp = { seed, sigma }; llama_model_ptr model(gguf_ctx != nullptr ? llama_model_init_from_user(gguf_ctx, set_tensor_data, &tmp, model_params) : llama_model_load_from_file_ptr(file, model_params)); @@ -345,6 +389,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK: case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: @@ -424,9 +469,6 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } - if (arch == LLM_ARCH_DEEPSEEK4) { - return false; - } // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU @@ -481,7 +523,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml continue; } gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe); - auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}); + auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, weight_sigma(arch)); const std::string path = dir + "/" + llm_arch_name(arch) + (moe ? "-moe.gguf" : "-dense.gguf"); LOG_INF("%s: Saving %s model (%s) to %s...\n", __func__, llm_arch_name(arch), moe ? "MoE" : "dense", path.c_str()); llama_model_save_to_file(model_and_ctx.first.get(), path.c_str()); diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index 5d1f0140b623..180756abcd60 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include static llama_context * make_ctx(const common_params & params, llama_model * model) { @@ -35,6 +37,508 @@ static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos) { return ok; } +// +// deep-rollback cases (DSV4-architecture fixtures only) +// +// The per-token rollback tier above covers depths up to n_rs_seq. These cases +// pin the contract for depths beyond it and for the rollback lifecycle: +// a) a 128-aligned rollback deeper than n_rs_seq replays reference-equal +// b) a rollback the raw SWA window cannot cover is refused cleanly or +// falls back checkpoint-equivalent - never silently wrong output +// c) a second seq_rm before any decode is rejected or cumulative-correct +// d) a failed decode after a rollback is armed does not lose or corrupt the +// pending rollback state - the next decode replays reference-equal +// + +using pos_logits = std::map>; + +static llama_context * make_deep_ctx( + const common_params & params, llama_model * model, uint32_t n_rs_seq, uint32_t n_ubatch, uint32_t n_ctx) { + auto cparams = common_context_params_to_llama(params); + cparams.n_seq_max = 1; + cparams.n_rs_seq = n_rs_seq; + cparams.n_rs_aligned = 4; + cparams.n_ctx = n_ctx; + cparams.n_batch = std::max(n_ctx, n_rs_seq + 1); + cparams.n_ubatch = std::max(n_ubatch, n_rs_seq + 1); + return llama_init_from_model(model, cparams); +} + +static bool decode_singles( + llama_context * ctx, const std::vector & tokens, + llama_pos p0, llama_pos p1, pos_logits * logits_out, uint32_t n_vocab) { + for (llama_pos pos = p0; pos < p1; ++pos) { + if (!decode_one(ctx, tokens[pos], pos)) { + return false; + } + if (logits_out) { + const float * logits = llama_get_logits_ith(ctx, 0); + if (logits == nullptr) { + return false; + } + (*logits_out)[pos].assign(logits, logits + n_vocab); + } + } + return true; +} + +// every replayed position must match its reference row +static bool logits_equal(const pos_logits & ref, const pos_logits & got, const char * label) { + constexpr float eps = 1e-5f; + for (const auto & [pos, got_row] : got) { + const auto it = ref.find(pos); + if (it == ref.end()) { + fprintf(stderr, "%s: no reference logits at position %d\n", label, pos); + return false; + } + for (size_t tok = 0; tok < got_row.size(); ++tok) { + if (std::fabs(it->second[tok] - got_row[tok]) > eps) { + fprintf(stderr, "%s: logits mismatch at position %d, token %zu (%g != %g)\n", + label, pos, tok, (double) it->second[tok], (double) got_row[tok]); + return false; + } + } + } + return !got.empty(); +} + +// decode [p0, p1) as ONE batch (the context may split it into several +// ubatches), capturing logits for every position +static bool decode_range_batch( + llama_context * ctx, const std::vector & tokens, + llama_pos p0, llama_pos p1, pos_logits * logits_out, uint32_t n_vocab) { + const uint32_t count = (uint32_t) (p1 - p0); + llama_batch batch = llama_batch_init(count, 0, 1); + for (llama_pos pos = p0; pos < p1; ++pos) { + common_batch_add(batch, tokens[pos], pos, { 0 }, true); + } + const bool ok = llama_decode(ctx, batch) == 0; + if (ok && logits_out) { + for (uint32_t i = 0; i < count; ++i) { + const float * logits = llama_get_logits_ith(ctx, i); + if (logits == nullptr) { + llama_batch_free(batch); + return false; + } + (*logits_out)[p0 + (llama_pos) i].assign(logits, logits + n_vocab); + } + } + llama_batch_free(batch); + return ok; +} + +// reference logits: an untouched context decoding the same stream the same way +static bool build_reference( + const common_params & params, llama_model * model, const std::vector & tokens, + uint32_t n_rs_seq, uint32_t n_ubatch, uint32_t n_ctx, + llama_pos batch_end, llama_pos ref_from, llama_pos ref_to, pos_logits & out, uint32_t n_vocab) { + llama_context * ctx = make_deep_ctx(params, model, n_rs_seq, n_ubatch, n_ctx); + if (ctx == nullptr) { + return false; + } + bool ok = decode_tokens(ctx, tokens, batch_end); + ok = ok && decode_singles(ctx, tokens, batch_end, ref_from, nullptr, n_vocab); + ok = ok && decode_singles(ctx, tokens, ref_from, ref_to, &out, n_vocab); + llama_free(ctx); + return ok; +} + +// a) aligned deep rollback: from position 306 back to the 128-aligned 256, +// depth 50 > n_rs_seq, then replay must be reference-equal +static bool case_aligned_deep_rollback( + const common_params & params, llama_model * model, const std::vector & tokens, uint32_t n_vocab) { + const char * label = "case_aligned_deep_rollback"; + + pos_logits ref; + if (!build_reference(params, model, tokens, 8, 512, 1024, 256, 256, 300, ref, n_vocab)) { + fprintf(stderr, "%s: failed to build reference\n", label); + return false; + } + + llama_context * ctx = make_deep_ctx(params, model, 8, 512, 1024); + if (ctx == nullptr) { + fprintf(stderr, "%s: failed to create context\n", label); + return false; + } + + bool ok = decode_tokens(ctx, tokens, 256) && decode_singles(ctx, tokens, 256, 306, nullptr, n_vocab); + if (!ok) { + fprintf(stderr, "%s: initial decode failed\n", label); + llama_free(ctx); + return false; + } + + if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, 256, -1)) { + fprintf(stderr, "%s: aligned deep rollback to position 256 (depth 50) was refused\n", label); + llama_free(ctx); + return false; + } + + pos_logits got; + ok = decode_singles(ctx, tokens, 256, 300, &got, n_vocab); + if (!ok) { + fprintf(stderr, "%s: replay decode failed\n", label); + llama_free(ctx); + return false; + } + + ok = logits_equal(ref, got, label); + llama_free(ctx); + if (!ok) { + return false; + } + + // Multi-ubatch replay: the same deep rollback replayed as ONE batch that + // the context splits into several ubatches must also be reference-equal. + // The pending restore must run exactly once, in the first ubatch. + pos_logits ref_mu; + { + llama_context * ctx_r = make_deep_ctx(params, model, 8, 16, 1024); + ok = ctx_r != nullptr && + decode_tokens(ctx_r, tokens, 256) && + decode_range_batch(ctx_r, tokens, 256, 300, &ref_mu, n_vocab); + if (ctx_r != nullptr) { + llama_free(ctx_r); + } + if (!ok) { + fprintf(stderr, "%s: failed to build multi-ubatch reference\n", label); + return false; + } + } + + llama_context * ctx_mu = make_deep_ctx(params, model, 8, 16, 1024); + if (ctx_mu == nullptr) { + fprintf(stderr, "%s: failed to create multi-ubatch context\n", label); + return false; + } + + ok = decode_tokens(ctx_mu, tokens, 256) && decode_singles(ctx_mu, tokens, 256, 306, nullptr, n_vocab); + ok = ok && llama_memory_seq_rm(llama_get_memory(ctx_mu), 0, 256, -1); + if (!ok) { + fprintf(stderr, "%s: multi-ubatch rollback setup failed\n", label); + llama_free(ctx_mu); + return false; + } + + pos_logits got_mu; + ok = decode_range_batch(ctx_mu, tokens, 256, 300, &got_mu, n_vocab); + if (!ok) { + fprintf(stderr, "%s: multi-ubatch replay decode failed\n", label); + llama_free(ctx_mu); + return false; + } + + ok = logits_equal(ref_mu, got_mu, label); + if (!ok) { + fprintf(stderr, "%s: multi-ubatch replay diverged from reference\n", label); + } + llama_free(ctx_mu); + return ok; +} + +// b) uncovered depth: stacked rollbacks walk the sequence back beyond what the +// raw SWA window still holds; the memory must refuse at some point or +// replay reference-equal - silently wrong output fails the case +static bool case_uncovered_depth( + const common_params & params, llama_model * model, const std::vector & tokens, uint32_t n_vocab) { + const char * label = "case_uncovered_depth"; + + pos_logits ref; + if (!build_reference(params, model, tokens, 8, 16, 1024, 240, 263, 279, ref, n_vocab)) { + fprintf(stderr, "%s: failed to build reference\n", label); + return false; + } + + llama_context * ctx = make_deep_ctx(params, model, 8, 16, 1024); + if (ctx == nullptr) { + fprintf(stderr, "%s: failed to create context\n", label); + return false; + } + + bool ok = decode_tokens(ctx, tokens, 240) && decode_singles(ctx, tokens, 240, 900, nullptr, n_vocab); + if (!ok) { + fprintf(stderr, "%s: initial decode failed\n", label); + llama_free(ctx); + return false; + } + + // stack depth-7 rollbacks from 899 down to 263: the cumulative depth walks + // far behind anything the raw window ring still holds + bool refused = false; + for (llama_pos target = 893; target >= 263; target -= 7) { + if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, target, -1)) { + refused = true; + break; + } + } + + if (refused) { + fprintf(stderr, "%s: uncovered rollback cleanly refused\n", label); + llama_free(ctx); + return true; + } + + pos_logits got; + ok = decode_singles(ctx, tokens, 263, 279, &got, n_vocab); + if (!ok) { + fprintf(stderr, "%s: replay decode failed after unrefused uncovered rollback\n", label); + llama_free(ctx); + return false; + } + + + ok = logits_equal(ref, got, label); + if (!ok) { + fprintf(stderr, "%s: uncovered rollback was silently applied with wrong output\n", label); + } + llama_free(ctx); + return ok; +} + +// c) stacked rollback: a second seq_rm before any decode must be rejected or +// behave cumulative-correct (replay from the deeper target reference-equal) +static bool case_stacked_rollback( + const common_params & params, llama_model * model, const std::vector & tokens, uint32_t n_vocab) { + const char * label = "case_stacked_rollback"; + + // low positions on purpose: with only a handful of compressed blocks in + // view, per-position state rows corrupted by the stacked rollback dominate + // the attention output instead of vanishing into a long context + pos_logits ref; + if (!build_reference(params, model, tokens, 8, 512, 1024, 4, 10, 20, ref, n_vocab)) { + fprintf(stderr, "%s: failed to build reference\n", label); + return false; + } + + llama_context * ctx = make_deep_ctx(params, model, 8, 512, 1024); + if (ctx == nullptr) { + fprintf(stderr, "%s: failed to create context\n", label); + return false; + } + + bool ok = decode_tokens(ctx, tokens, 4) && decode_singles(ctx, tokens, 4, 20, nullptr, n_vocab); + if (!ok) { + fprintf(stderr, "%s: initial decode failed\n", label); + llama_free(ctx); + return false; + } + + // two stacked rollbacks with cumulative depth 10 > n_rs_seq, both targets + // mid-block so partial compressor state is live at each restore point + if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, 14, -1)) { + fprintf(stderr, "%s: first rollback (depth 6) was refused\n", label); + llama_free(ctx); + return false; + } + + const bool second = llama_memory_seq_rm(llama_get_memory(ctx), 0, 10, -1); + if (!second) { + fprintf(stderr, "%s: stacked rollback cleanly rejected\n", label); + // the pending first rollback must still replay correctly from 14 + pos_logits got; + ok = decode_singles(ctx, tokens, 14, 20, &got, n_vocab); + ok = ok && logits_equal(ref, got, label); + llama_free(ctx); + return ok; + } + + pos_logits got; + ok = decode_singles(ctx, tokens, 10, 20, &got, n_vocab); + if (!ok) { + fprintf(stderr, "%s: replay decode failed after stacked rollback\n", label); + llama_free(ctx); + return false; + } + + ok = logits_equal(ref, got, label); + if (!ok) { + fprintf(stderr, "%s: stacked rollback accepted but replay is not cumulative-correct\n", label); + } + llama_free(ctx); + return ok; +} + +// d) failed decode with a deep rollback armed: the failure must not lose or +// corrupt the pending rollback - the subsequent replay must be +// reference-equal +static bool case_failed_decode_lifecycle( + const common_params & params, llama_model * model, const std::vector & tokens, uint32_t n_vocab) { + const char * label = "case_failed_decode_lifecycle"; + + pos_logits ref; + if (!build_reference(params, model, tokens, 8, 512, 1024, 256, 256, 300, ref, n_vocab)) { + fprintf(stderr, "%s: failed to build reference\n", label); + return false; + } + + llama_context * ctx = make_deep_ctx(params, model, 8, 512, 1024); + if (ctx == nullptr) { + fprintf(stderr, "%s: failed to create context\n", label); + return false; + } + + bool ok = decode_tokens(ctx, tokens, 256) && decode_singles(ctx, tokens, 256, 306, nullptr, n_vocab); + if (!ok) { + fprintf(stderr, "%s: initial decode failed\n", label); + llama_free(ctx); + return false; + } + + // arm a deep aligned rollback, then fail a decode before the replay + if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, 256, -1)) { + fprintf(stderr, "%s: deep rollback to position 256 (depth 50) was refused\n", label); + llama_free(ctx); + return false; + } + + // a decode that must fail: position far beyond n_ctx + if (decode_one(ctx, tokens[0], 4000)) { + fprintf(stderr, "%s: decode at position 4000 unexpectedly succeeded, cannot exercise the hazard\n", label); + llama_free(ctx); + return false; + } + + pos_logits got; + ok = decode_singles(ctx, tokens, 256, 300, &got, n_vocab); + if (!ok) { + fprintf(stderr, "%s: replay decode failed after the failed decode\n", label); + llama_free(ctx); + return false; + } + + ok = logits_equal(ref, got, label); + if (!ok) { + fprintf(stderr, "%s: pending rollback state corrupted by the failed decode\n", label); + } + llama_free(ctx); + return ok; +} + +// e) state save with an aligned rollback pending: sequence state is sized and +// saved while the marker is armed (the server prompt-cache does this on +// every task start); states without aligned planes must degrade to their +// live plane instead of failing the save, and the saved state must restore +// reference-equal +static bool case_aligned_pending_state_save( + const common_params & params, llama_model * model, const std::vector & tokens, uint32_t n_vocab) { + const char * label = "case_aligned_pending_state_save"; + + pos_logits ref; + if (!build_reference(params, model, tokens, 8, 512, 1024, 256, 256, 300, ref, n_vocab)) { + fprintf(stderr, "%s: failed to build reference\n", label); + return false; + } + + llama_context * ctx_src = make_deep_ctx(params, model, 8, 512, 1024); + llama_context * ctx_dst = make_deep_ctx(params, model, 8, 512, 1024); + if (ctx_src == nullptr || ctx_dst == nullptr) { + fprintf(stderr, "%s: failed to create contexts\n", label); + return false; + } + + bool ok = decode_tokens(ctx_src, tokens, 256) && decode_singles(ctx_src, tokens, 256, 306, nullptr, n_vocab); + if (!ok) { + fprintf(stderr, "%s: initial decode failed\n", label); + llama_free(ctx_src); + llama_free(ctx_dst); + return false; + } + + if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, 256, -1)) { + fprintf(stderr, "%s: aligned deep rollback to position 256 was refused\n", label); + llama_free(ctx_src); + llama_free(ctx_dst); + return false; + } + + if (llama_state_seq_get_size(ctx_src, 0) == 0) { + fprintf(stderr, "%s: state size query failed with an aligned rollback pending\n", label); + llama_free(ctx_src); + llama_free(ctx_dst); + return false; + } + + common_prompt_checkpoint ckpt; + ckpt.update_tgt(ctx_src, 0, 0); + ckpt.load_tgt(ctx_dst, 0, 0); + llama_free(ctx_src); + + pos_logits got; + ok = decode_singles(ctx_dst, tokens, 256, 300, &got, n_vocab); + if (!ok) { + fprintf(stderr, "%s: replay decode failed after state restore\n", label); + llama_free(ctx_dst); + return false; + } + + // The save/load round trip repacks memory cells, which changes attention + // reduction order; the resulting drift is ~5e-5 at this context length and + // hits the pre-existing per-token save path equally (measured 2.4e-5 on the + // same prompt), so this case uses a round-trip tolerance instead of the + // in-place eps. Real state corruption measures ~1e-3 even on far shorter + // contexts, an order of magnitude above this tolerance. + constexpr float roundtrip_eps = 2e-4f; + ok = !got.empty(); + for (const auto & [pos, row] : got) { + const auto it = ref.find(pos); + if (it == ref.end()) { + fprintf(stderr, "%s: no reference logits at position %d\n", label, pos); + ok = false; + break; + } + for (size_t tok = 0; tok < row.size(); ++tok) { + if (std::fabs(it->second[tok] - row[tok]) > roundtrip_eps) { + fprintf(stderr, "%s: logits mismatch at position %d, token %zu (%g != %g)\n", + label, pos, tok, (double) it->second[tok], (double) row[tok]); + ok = false; + break; + } + } + if (!ok) { + break; + } + } + llama_free(ctx_dst); + return ok; +} + +static int run_deep_rollback_cases(const common_params & params, llama_model * model, uint32_t n_vocab) { + char arch[64] = {0}; + if (llama_model_meta_val_str(model, "general.architecture", arch, sizeof(arch)) < 0 || + strcmp(arch, "deepseek4") != 0) { + fprintf(stderr, "%s : skipping deep-rollback cases for non-DSV4 arch\n", __func__); + return 0; + } + + std::vector tokens; + tokens.reserve(1024); + for (uint32_t i = 0; i < 1024; ++i) { + tokens.push_back(1 + (llama_token) ((i*7) % (n_vocab - 1))); + } + + struct named_case { + const char * name; + bool (*fn)(const common_params &, llama_model *, const std::vector &, uint32_t); + }; + const named_case cases[] = { + { "aligned-deep-rollback", case_aligned_deep_rollback }, + { "stacked-rollback", case_stacked_rollback }, + { "failed-decode-lifecycle", case_failed_decode_lifecycle }, + { "uncovered-depth", case_uncovered_depth }, + { "aligned-pending-state-save", case_aligned_pending_state_save }, + }; + + int n_failed = 0; + for (const auto & c : cases) { + const bool ok = c.fn(params, model, tokens, n_vocab); + fprintf(stderr, "%s : deep-rollback case %-24s : %s\n", __func__, c.name, ok ? "PASS" : "FAIL"); + if (!ok) { + n_failed++; + } + } + return n_failed; +} + int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -220,5 +724,11 @@ int main(int argc, char ** argv) { llama_free(ctx_src); llama_free(ctx_dst); llama_free(ctx_dirty); + + const int n_deep_failed = run_deep_rollback_cases(params, model, n_vocab); + if (n_deep_failed > 0) { + fprintf(stderr, "%s : %d deep-rollback case(s) failed\n", __func__, n_deep_failed); + return 1; + } return 0; } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 4655b518e21f..f3a325cefa6a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -951,6 +951,10 @@ struct server_context_impl { // if swa_full is enabled, this is set to 0 to simulate a non-SWA model int32_t n_swa; + // set to llama_memory_seq_rm_align(memory) of the target context + // > 1 when the memory has an aligned deep-rollback tier (e.g. DSV4) + llama_pos n_seq_rm_align = 1; + // slots / clients std::vector slots; @@ -1134,7 +1138,8 @@ struct server_context_impl { if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; } - cparams_dft.n_rs_seq = 0; + cparams_dft.n_rs_seq = 0; + cparams_dft.n_rs_aligned = 0; std::vector devs; uint32_t hp_ngl = 0; @@ -1304,6 +1309,11 @@ struct server_context_impl { SRV_WRN("%s", "speculative decoding not supported by this context\n"); } + n_seq_rm_align = llama_memory_seq_rm_align(llama_get_memory(ctx_tgt)); + if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_seq_rm_align > 1) { + SRV_INF("deep-rollback tier active, seq_rm_align = %d\n", (int) n_seq_rm_align); + } + if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL) { SRV_TRC("%s", "speculative decoding will use checkpoints\n"); } @@ -3328,40 +3338,98 @@ struct server_context_impl { } if (pos_min >= pos_min_thold) { - // search for a context checkpoint - const auto it = std::find_if( - slot.prompt.checkpoints.rbegin(), - slot.prompt.checkpoints.rend(), - [&](const auto & cur) { - // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS] - SLT_TRC(slot, "checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold); - // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN] - if (cur.pos_max > pos_next) { - return false; + // deep-rollback-capable memories (e.g. DSV4 with the aligned tier) + // can serve the divergence themselves: the recurrent state rolls + // back to an aligned boundary and the removed suffix is re-decoded. + // try this before falling back to the checkpoint machinery + bool rs_deep = false; + + if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_seq_rm_align > 1 && + !slot.prompt.tokens.has_mtmd) { + const auto pos_max_seq = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id); + + if (pos_next > pos_max_seq) { + // nothing to remove below the tip - the live state simply + // continues; a full-prefix re-eval ([TAG_PROMPT_LOGITS]) + // removes at most one token, which the bounded per-token + // rollback tier covers + rs_deep = true; + } else { + // divergence below the tip: align the reuse point down to a + // boundary the deep tier can restore + llama_pos pos_rm = (pos_next/n_seq_rm_align)*n_seq_rm_align; + int n_past_new = (int) slot.prompt.tokens.size_up_to_pos(pos_rm); + + // keep at least one task token beyond the target so the + // [TAG_PROMPT_LOGITS] re-eval never needs a second removal + // below an armed rollback + if (n_past_new >= slot.task->n_tokens()) { + pos_rm -= n_seq_rm_align; + n_past_new = pos_rm > 0 ? (int) slot.prompt.tokens.size_up_to_pos(pos_rm) : 0; } - return cur.pos_min < pos_min_thold || cur.pos_min == 0; - } - ); - - bool do_reset = it == slot.prompt.checkpoints.rend(); - if (!do_reset) { - // restore the context checkpoint - it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - // restore the draft's speculative state - common_speculative_set_state(spec.get(), slot.id, it->data_spec); + // the memory is the authority on coverage - beyond it the + // rollback is refused and the checkpoint machinery below is + // the fallback, with n_past/pos_next untouched + if (pos_rm > 0 && llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot.id, pos_rm, -1)) { + SLT_INF(slot, "aligning n_past %d -> %d (seq_rm_align = %d)\n", n_past, n_past_new, (int) n_seq_rm_align); + SLT_INF(slot, "deep rollback engaged (pos_rm = %d, depth = %d)\n", pos_rm, (int) (pos_max_seq - pos_rm + 1)); - pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); - n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) it->n_tokens); - SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); + n_past = n_past_new; + pos_next = pos_rm; + rs_deep = true; + } + } } - if (do_reset) { - SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n", - "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); - pos_next = 0; - n_past = 0; + if (!rs_deep) { + // search for a context checkpoint + const auto it = std::find_if( + slot.prompt.checkpoints.rbegin(), + slot.prompt.checkpoints.rend(), + [&](const auto & cur) { + // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS] + SLT_TRC(slot, "checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold); + // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN] + // for aligned-removal memories the restore resumes at pos_max + 1, + // so a checkpoint with pos_max == pos_next would land the resume + // point past the divergence - require pos_max strictly below + if (cur.pos_max > pos_next || (n_seq_rm_align > 1 && cur.pos_max == pos_next)) { + return false; + } + return cur.pos_min < pos_min_thold || cur.pos_min == 0; + } + ); + + bool do_reset = it == slot.prompt.checkpoints.rend(); + + if (!do_reset) { + // restore the context checkpoint + it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + // restore the draft's speculative state + common_speculative_set_state(spec.get(), slot.id, it->data_spec); + + if (n_seq_rm_align > 1) { + // aligned-removal memories (e.g. DSV4) cannot remove at the restored + // tip: the restore invalidates the rollback history, so the only + // removal the memory still supports is a dead tail strictly past + // pos_max - resume at pos_max + 1 (the same shape as the speculative + // rollback below, which restores and then removes [pos_max + 1, end)) + pos_next = std::min(pos_next, it->pos_max + 1); + } else { + pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); + } + n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) it->n_tokens); + SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); + } + + if (do_reset) { + SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n", + "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); + pos_next = 0; + n_past = 0; + } } } }