llama : add opt-in boundary aligned deep rollback for DSV4 recurrent state

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)"

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;

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";

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) {

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

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
     //

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

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