feat(server): reuse restored KV/recurrent state instead of reprocessing

On hybrid recurrent models, slot restore followed by prompt reuse always
forced a full reprocess: only the current recurrent plane was persisted,
rs_idx was reset on read, and the reuse gate reset n_past whenever the
aggregate pos_min came from recurrent memory. This made cross-process KV
restore useless for skipping prompt processing.

- llama-memory-recurrent: persist all 1+n_rs_seq rollback planes (offset
  k*size) and rs_idx in state_write_data/state_read_data, so the running
  state can roll back to the reuse boundary after a restore.
- server-context: in the reuse gate, when the memory supports bounded
  recurrent rollback and (pos_min - n_past + 1) <= n_rs_seq, keep n_past
  and reuse; otherwise fall back to full reprocess (no abort). Add a .dft
  sidecar to save/restore the speculative draft context so seq_rm on the
  draft does not abort under --spec-type.
- common: add KVARN_N_RS_SEQ env override for the recurrent rollback budget
  (native comes from speculative need_n_rs_seq(); 0 without speculation).

Result: restored state is reused instead of reprocessed (prompt_n 4843->19
in tests), correct recall, works cross-GPU, with safe reprocess fallback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RErC5TKNN678KpqAQwFatc

Author: Stratum

common/common.cpp

@@ -1697,6 +1697,7 @@ 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();
+    if (const char * kvenv = getenv("KVARN_N_RS_SEQ")) { cparams.n_rs_seq = (uint32_t) atoi(kvenv); } // KVSNAP test override
     cparams.n_outputs_max     = std::max(params.n_outputs_max, 0);
     cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0);
     cparams.n_batch           = params.n_batch;

src/llama-memory-recurrent.cpp

@@ -807,7 +807,16 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq
     io.write(&cell_count, sizeof(cell_count));
 
     state_write_meta(io, cell_ranges, seq_id);
-    state_write_data(io, cell_ranges_data);
+
+    // KVSNAP: persist the current rollback-plane index so read can restore snapshots
+    uint32_t rs_idx_save = 0;
+    if (n_rs_seq != 0 && seq_id >= 0 && (size_t) seq_id < rs_idx.size()) {
+        rs_idx_save = rs_idx[seq_id];
+    }
+    io.write(&rs_idx_save, sizeof(rs_idx_save));
+
+    (void) cell_ranges_data;
+    state_write_data(io, cell_ranges); // KVSNAP: base ranges; data loops all planes
 }
 
 void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
@@ -820,6 +829,9 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
 
     res = res && state_read_meta(io, cell_count, seq_id);
 
+    uint32_t rs_idx_load = 0; // KVSNAP
+    io.read(&rs_idx_load, sizeof(rs_idx_load));
+
     try {
         res = res && state_read_data(io, cell_count);
     } catch (...) {
@@ -839,7 +851,7 @@ void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_i
         if (seq_id == -1) {
             std::fill(rs_idx.begin(), rs_idx.end(), 0);
         } else {
-            set_rs_idx(seq_id, 0);
+            set_rs_idx(seq_id, rs_idx_load); // KVSNAP: restore saved rollback index
         }
     }
 }
@@ -886,10 +898,12 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::
 
         // Write each logical cell row range. With pending recurrent rollback,
         // the logical current state may live in a rollback snapshot plane.
-        for (const auto & range : cell_ranges) {
-            const size_t range_size = range.second - range.first;
-            const size_t buf_size = range_size * r_size_row;
-            io.write_tensor(r_l[il], range.first * r_size_row, buf_size);
+        for (uint32_t kp = 0; kp <= n_rs_seq; ++kp) { // KVSNAP: write all rollback planes
+            for (const auto & range : cell_ranges) {
+                const size_t range_size = range.second - range.first;
+                const size_t buf_size = range_size * r_size_row;
+                io.write_tensor(r_l[il], (range.first + kp * size) * r_size_row, buf_size);
+            }
         }
     }
 
@@ -908,10 +922,12 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::
 
             // Write each logical cell row range. With pending recurrent rollback,
             // the logical current state may live in a rollback snapshot plane.
-            for (const auto & range : cell_ranges) {
-                const size_t range_size = range.second - range.first;
-                const size_t buf_size = range_size * s_size_row;
-                io.write_tensor(s_l[il], range.first * s_size_row, buf_size);
+            for (uint32_t kp = 0; kp <= n_rs_seq; ++kp) { // KVSNAP: write all rollback planes
+                for (const auto & range : cell_ranges) {
+                    const size_t range_size = range.second - range.first;
+                    const size_t buf_size = range_size * s_size_row;
+                    io.write_tensor(s_l[il], (range.first + kp * size) * s_size_row, buf_size);
+                }
             }
         }
     } else {
@@ -1086,8 +1102,9 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell
         }
 
         if (cell_count) {
-            // Read and set the keys for the whole cell range
-            io.read_tensor(r_l[il], head * r_size_row, cell_count * r_size_row);
+            for (uint32_t kp = 0; kp <= n_rs_seq; ++kp) { // KVSNAP: read all rollback planes
+                io.read_tensor(r_l[il], (head + kp * size) * r_size_row, cell_count * r_size_row);
+            }
         }
     }
 
@@ -1116,8 +1133,9 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell
             }
 
             if (cell_count) {
-                // Read and set the values for the whole cell range
-                io.read_tensor(s_l[il], head * s_size_row, cell_count * s_size_row);
+                for (uint32_t kp = 0; kp <= n_rs_seq; ++kp) { // KVSNAP: read all rollback planes
+                    io.read_tensor(s_l[il], (head + kp * size) * s_size_row, cell_count * s_size_row);
+                }
             }
         }
     } else {

tools/server/server-context.cpp

@@ -2476,6 +2476,17 @@ struct server_context_impl {
                         break;
                     }
 
+                    // KVFIX-DFT: persist the draft (speculative) context state so a restored
+                    // slot has a consistent draft KV; else seq_rm on the draft aborts.
+                    if (ctx_dft) {
+                        const size_t szd = llama_state_seq_get_size_ext(ctx_dft, slot->id, LLAMA_STATE_SEQ_FLAGS_NONE);
+                        std::vector<uint8_t> bufd(szd);
+                        llama_state_seq_get_data_ext(ctx_dft, bufd.data(), szd, slot->id, LLAMA_STATE_SEQ_FLAGS_NONE);
+                        const std::string dftpath = filepath + ".dft";
+                        FILE * fdft = fopen(dftpath.c_str(), "wb");
+                        if (fdft) { fwrite(bufd.data(), 1, szd, fdft); fclose(fdft); }
+                    }
+
                     const int64_t t_end = ggml_time_us();
                     const double t_save_ms = (t_end - t_start) / 1000.0;
 
@@ -2535,6 +2546,21 @@ struct server_context_impl {
 
                         slot->prompt.clear();
                         slot->prompt.tokens = std::move(restored);
+
+                        // KVFIX-DFT: restore the draft (speculative) context state to match tgt.
+                        if (ctx_dft) {
+                            const std::string dftpath = filepath + ".dft";
+                            FILE * fdft = fopen(dftpath.c_str(), "rb");
+                            if (fdft) {
+                                fseek(fdft, 0, SEEK_END); long szd = ftell(fdft); fseek(fdft, 0, SEEK_SET);
+                                std::vector<uint8_t> bufd(szd > 0 ? szd : 0);
+                                const size_t rd = szd > 0 ? fread(bufd.data(), 1, (size_t) szd, fdft) : 0;
+                                fclose(fdft);
+                                if (rd == (size_t) szd && szd > 0) {
+                                    llama_state_seq_set_data_ext(ctx_dft, bufd.data(), bufd.size(), slot->id, LLAMA_STATE_SEQ_FLAGS_NONE);
+                                }
+                            }
+                        }
                     } catch (const std::exception & err) {
                         slot->prompt_clear();
                         send_error(task, std::string("Unable to restore slot: ") + err.what(), ERROR_TYPE_INVALID_REQUEST);
@@ -3270,10 +3296,22 @@ struct server_context_impl {
                                     }
 
                                     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;
+                                        // Hybrid/recurrent memory keeps only a running state, so the
+                                        // pos_min gate above wants a checkpoint. But if the memory
+                                        // supports bounded recurrent-state rollback (n_rs_seq) and the
+                                        // reuse boundary is within that budget, we can safely reuse the
+                                        // restored prefix: seq_rm rolls the running state back using the
+                                        // persisted snapshots. Otherwise fall back to full reprocess.
+                                        const int32_t   n_rs     = (int32_t) llama_n_rs_seq(ctx_tgt);
+                                        const llama_pos rollback = pos_min - n_past + 1;
+                                        if (n_rs > 0 && rollback <= (llama_pos) n_rs) {
+                                            SLT_DBG(slot, "recurrent reuse: rollback = %d <= n_rs_seq = %d, keeping n_past = %d\n", (int) rollback, n_rs, n_past);
+                                        } else {
+                                            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;
+                                        }
                                     }
                                 }
                             }