mybigday · a-ghorbani · Jul 13, 2026 · Jul 13, 2026 · Jul 13, 2026 · Jul 14, 2026
diff --git a/.gitignore b/.gitignore
@@ -102,3 +102,5 @@ local-plan-done
 
 tests/*.gguf
 tests/status_changes.log
+.claude/settings.local.json
+tests/android/build-hexagon/
diff --git a/cpp/jsi/RNLlamaJSI.cpp b/cpp/jsi/RNLlamaJSI.cpp
@@ -508,14 +508,21 @@ namespace rnllama_jsi {
                     }
                 }
 
+                int stateCacheBudgetMb =
+                    getPropertyAsInt(runtime, params, "state_cache_budget_mb", 160);
+                int stateCacheMaxCheckpoints =
+                    getPropertyAsInt(runtime, params, "state_cache_max_checkpoints", 8);
+
                 return createPromiseTask(runtime, callInvoker, [
                     contextId,
                     cparams,
                     skipGpuDevices,
                     requestedDevices,
                     devicesProvided,
                     useProgressCallback,
-                    progressData
+                    progressData,
+                    stateCacheBudgetMb,
+                    stateCacheMaxCheckpoints
                 ]() mutable -> PromiseResultGenerator {
                     if (isContextLimitReached()) {
                         throw std::runtime_error("Context limit reached");
@@ -575,6 +582,13 @@ namespace rnllama_jsi {
                     }
 
                     auto ctx = new rnllama::llama_rn_context();
+                    // Prompt state cache tuning (multi-turn KV reuse on
+                    // recurrent/hybrid/SWA models). Budget in MiB; 0 disables it.
+                    {
+                        ctx->state_cache_budget_bytes =
+                            stateCacheBudgetMb > 0 ? (size_t) stateCacheBudgetMb * 1024 * 1024 : 0;
+                        ctx->state_cache_max_checkpoints = stateCacheMaxCheckpoints;
+                    }
                     if (ctx->loadModel(cparams)) {
                          ctx->attachThreadpoolsIfAvailable();
 

diff --git a/cpp/rn-completion.cpp b/cpp/rn-completion.cpp
diff --git a/cpp/rn-completion.h b/cpp/rn-completion.h
@@ -22,6 +22,20 @@ static inline void llama_batch_clear(llama_batch *batch) {
 // Forward declarations
 struct llama_rn_context;
 
+// A snapshot of the non-rollbackable memory state (recurrent/SWA cells,
+// PARTIAL_ONLY) at a token boundary.
+struct rn_state_checkpoint {
+    // Exact token sequence the snapshot represents (memory positions [0, n)).
+    std::vector<llama_token> tokens;
+    // PARTIAL_ONLY per-sequence state blob for seq_id 0.
+    std::vector<uint8_t> data;
+
+    size_t n_tokens() const { return tokens.size(); }
+    size_t size_bytes() const {
+        return data.size() + tokens.size() * sizeof(llama_token);
+    }
+};
+
 // Types defined in rn-llama.h (needed here for compilation)
 enum stop_type
 {
@@ -73,6 +87,32 @@ struct llama_rn_context_completion {
     llama_pos n_past = 0;
     size_t n_remain = 0;
     std::vector<llama_token> embd;
+
+    // --- Prompt state cache (recurrent / hybrid / SWA prefix reuse) ----------
+    // Saved snapshots, oldest first; empty for pure-attention models (seq_rm
+    // already reuses the prefix for free).
+    std::vector<rn_state_checkpoint> state_checkpoints;
+    bool state_cache_enabled = false;   // set once, from the model architecture
+    bool state_cache_probed = false;    // whether we've inspected the model yet
+    bool prompt_checkpoint_pending = false; // prompt-region snapshots armed for this ingest
+    // embedding()/rerank() run throwaway prompts through this path; they must
+    // not capture into the chat's cache. Set per loadPrompt.
+    bool state_cache_capture_allowed = true;
+    // Bounds, copied from the host-set config at probeStateCache() time. The
+    // byte budget is the real limiter (snapshots are tens of KiB to ~22 MiB
+    // per model, fixed in the token count).
+    size_t state_cache_max_checkpoints = 8;
+    size_t state_cache_budget_bytes = (size_t) 160 * 1024 * 1024;
+    // Minimum spacing between message-boundary restore points on a COLD ingest:
+    // a boundary less than this far past the previous one saves less reprocess
+    // than a full-state snapshot costs (see computeMessageBoundaries).
+    llama_pos state_ckpt_min_gap = 64;
+    // Fallback snapshot interval during a COLD prompt ingest (0 disables); used
+    // only when the prompt exposes no message boundaries (base models, multimodal).
+    size_t state_ckpt_prefill_interval = 256;
+    // Message-boundary snapshot positions for the current prompt, ascending
+    // (see computeMessageBoundaries). Computed per loadPrompt.
+    std::vector<llama_pos> boundary_ckpts;
     bool incomplete = false;
     bool context_full = false;
     bool truncated = false;
@@ -99,6 +139,13 @@ struct llama_rn_context_completion {
     llama_pos spec_n_past = 0;
     llama_tokens spec_draft;
     std::deque<completion_token_output> spec_pending_tokens;
+    // Number of prompt tokens the last MTP prompt eval actually decoded (vs.
+    // reused from the cache). Instrumentation for the reuse tests.
+    size_t mtp_prompt_reprocessed = 0;
+    // Set once a mem-shared MTP draft (ctx_other == target) is detected: the
+    // checkpoint cache is disabled for it, since upstream state save/restore is
+    // a no-op on shared-cell caches (TAG_KV_CACHE_SHARE_CELLS). Test-visible.
+    bool mtp_draft_mem_shared = false;
 
     // Constructor
     llama_rn_context_completion(llama_rn_context* parent);
@@ -109,8 +156,29 @@ struct llama_rn_context_completion {
     // Completion processing methods
     void rewind();
     bool initSampling();
+
+    // Prompt state cache helpers (see rn_state_checkpoint).
+    void probeStateCache();                       // detect whether the model needs it
+    void captureStateCheckpoint();                // snapshot memory at current n_past (embd)
+    // Snapshot at position n, tagged with seq[0, n) (MTP path uses spec_prompt).
+    void captureStateCheckpoint(const std::vector<llama_token> &seq, size_t n);
+    int  findStateCheckpoint(const std::vector<llama_token> &target, size_t max_len) const;
+    bool restoreStateCheckpoint(size_t index);    // restore snapshot into seq 0
+    // Restore the longest snapshot prefixing `target` (length <= max_reuse,
+    // < total_tokens) and truncate the live prefix to it. Sets n_past_out.
+    bool recoverStateCheckpoint(const std::vector<llama_token> &target, size_t max_reuse,
+                                size_t total_tokens, llama_pos &n_past_out);
+    void evictStateCheckpoints();                 // enforce count / byte bounds
+    void clearStateCheckpoints();                 // drop all snapshots
+    void eraseStateCheckpointAt(size_t n_tokens); // drop the snapshot at a boundary
     void truncatePrompt(std::vector<llama_token> &prompt_tokens);
-    void loadPrompt(const std::vector<std::string> &media_paths);
+    void loadPrompt(const std::vector<std::string> &media_paths, bool allow_state_cache = true);
+    // Ascending message-boundary positions: the first content token after each
+    // run of chat-template delimiter tokens (CONTROL/USER_DEFINED). Boundaries
+    // closer than min_gap to the previous one are dropped (min_gap == 1 keeps
+    // all — used to recover the last boundary as the cold-ingest frontier).
+    std::vector<llama_pos> computeMessageBoundaries(const std::vector<llama_token> &tokens,
+                                                    llama_pos min_gap) const;
     void beginCompletion();
     void beginCompletion(int chat_format, common_reasoning_format reasoning_format, const std::string &generation_prompt = "", const std::string &chat_parser = "");
     void endCompletion();

diff --git a/cpp/rn-llama.cpp b/cpp/rn-llama.cpp
@@ -782,6 +782,7 @@ void llama_rn_context::clearCache(bool clear_data) {
     if (completion != nullptr) {
         completion->embd.clear();
         completion->n_past = 0;
+        completion->clearStateCheckpoints();
         LOG_INFO("Cache cleared and completion state reset (clear_data=%s)", clear_data ? "true" : "false");
     } else {
         LOG_INFO("Cache cleared (clear_data=%s)", clear_data ? "true" : "false");

diff --git a/cpp/rn-llama.h b/cpp/rn-llama.h
@@ -102,6 +102,11 @@ struct llama_rn_context {
     common_chat_templates_ptr templates;
     int n_ctx = 0;
 
+    // Prompt state cache tuning (see rn-completion.h). Budget 0 disables it;
+    // no-op on pure-attention models.
+    size_t state_cache_budget_bytes = (size_t) 160 * 1024 * 1024; // 0 = disabled
+    int32_t state_cache_max_checkpoints = 8;
+
     // Completion context (DEPRECATED: Use slot_manager for parallel decoding)
     llama_rn_context_completion *completion = nullptr;
 

diff --git a/cpp/rn-mtmd.hpp b/cpp/rn-mtmd.hpp
@@ -8,15 +8,29 @@
 #include <string>
 #include <vector>
 #include <cstdint>
+#include <functional>
 
 namespace rnllama {
 
+// Optional state-cache callbacks so the single-sequence completion path can
+// reuse the prefix across turns (as on the text path); null on the parallel
+// slot-manager path (full-clear).
+//   recover(target, max_reuse, total_tokens, n_past_out) -> restored?
+//   capture(seq, n) -> snapshot the memory at position n, tagged with seq[0,n)
+using mtmd_state_recover_fn =
+    std::function<bool(const std::vector<llama_token> &, size_t, size_t, llama_pos &)>;
+using mtmd_state_capture_fn =
+    std::function<void(const std::vector<llama_token> &, size_t)>;
+
 // MTMD context structure
 struct llama_rn_context_mtmd {
     mtmd_context *mtmd_ctx = nullptr;
 
     // State fields
     std::vector<std::string> bitmap_past_hashes;
+    // Number of prompt tokens reused from the cache on the last processMedia call
+    // (the position the chunk eval resumed from). Instrumentation for the tests.
+    llama_pos last_reused_n_past = 0;
 
     // Constructor - Initialize multimodal
     llama_rn_context_mtmd(
@@ -46,7 +60,9 @@ struct llama_rn_context_mtmd {
         bool &context_full,
         common_sampler *ctx_sampling,
         std::vector<std::string> &bitmap_past_hashes,  // Per-slot bitmap hashes
-        int32_t seq_id  // Sequence ID for parallel slots
+        int32_t seq_id,  // Sequence ID for parallel slots
+        mtmd_state_recover_fn recover = nullptr,
+        mtmd_state_capture_fn capture = nullptr
     );
 
     // Check if multimodal is enabled
@@ -393,7 +409,9 @@ inline void llama_rn_context_mtmd::processMedia(
     bool &context_full,
     common_sampler *ctx_sampling,
     std::vector<std::string> &bitmap_past_hashes_ref,  // Per-slot bitmap hashes
-    int32_t seq_id  // Sequence ID for parallel slots
+    int32_t seq_id,  // Sequence ID for parallel slots
+    mtmd_state_recover_fn recover,
+    mtmd_state_capture_fn capture
 ) {
     // Multimodal path
     std::string full_prompt = prompt;
@@ -478,15 +496,82 @@ inline void llama_rn_context_mtmd::processMedia(
 
     bool clear_result = llama_memory_seq_rm(kv, seq_id, n_past, -1);
     if (!clear_result) {
-        LOG_ERROR("[DEBUG] llama_memory_seq_rm failed (likely using a non-Transformer model)! Trying full clear...");
-        llama_memory_clear(kv, false);
-        n_past = 0;
-        new_n_past = n_past;
+        // Recurrent/hybrid: restore a checkpoint instead of re-encoding the
+        // images. Must be chunk-aligned — the eval loop below does whole chunks,
+        // so a mid-chunk restore (except a final plain-text tail) shifts chunks.
+        llama_pos recovered_n_past = 0;
+        bool recovered_ok = false;
+        size_t recover_cap = (size_t) n_past;
+        while (recover && recover(all_tokens, recover_cap, all_tokens.size(), recovered_n_past)) {
+            // Locate the chunk containing the recovered position.
+            size_t ci = 0;
+            for (size_t i = 0; i < chunk_pos.size(); i++) {
+                if (chunk_pos[i] <= (size_t) recovered_n_past) ci = i;
+                else break;
+            }
+            const bool aligned = chunk_pos[ci] == (size_t) recovered_n_past;
+            bool text_tail_of_last_chunk = false;
+            if (!aligned && ci + 1 == chunk_pos.size()) {
+                text_tail_of_last_chunk = true;
+                for (size_t j = (size_t) recovered_n_past; j < all_tokens.size(); j++) {
+                    if (all_tokens[j] == LLAMA_TOKEN_NULL) {
+                        text_tail_of_last_chunk = false;
+                        break;
+                    }
+                }
+            }
+            if (aligned || text_tail_of_last_chunk) {
+                recovered_ok = true;
+                break;
+            }
+            if (chunk_pos[ci] == 0) {
+                break; // no shorter aligned candidate can exist
+            }
+            // Mid-interior-chunk: retry capped at the start of the chunk that
+            // contains the recovered position (strictly decreasing -> terminates).
+            recover_cap = chunk_pos[ci];
+        }
+        if (recovered_ok) {
+            n_past = recovered_n_past;
+            new_n_past = n_past;
+            LOG_INFO("[DEBUG] Restored multimodal state checkpoint: reusing %d tokens", n_past);
+        } else {
+            LOG_ERROR("[DEBUG] llama_memory_seq_rm failed (likely using a non-Transformer model)! Trying full clear...");
+            llama_memory_clear(kv, false);
+            n_past = 0;
+            new_n_past = n_past;
+        }
     }
 
 
     LOG_INFO("[DEBUG] Evaluating chunks: n_past=%d, n_batch=%d", n_past, n_batch);
 
+    // Record the reused prefix (position the eval resumes from) for the tests.
+    last_reused_n_past = n_past;
+
+    // Frontier snapshot (like the text path, rn-completion.cpp:461): the reused
+    // state rests at n_past; snapshot it before reprocessing the tail. Advances
+    // the checkpoint chain each turn. Cold turns (n_past==0) skip it -- the
+    // media-boundary anchor below covers those.
+    if (capture && n_past > 0) {
+        capture(all_tokens, (size_t) n_past);
+    }
+
+    // Position right after the last image chunk (== L if the prompt ends in an
+    // image). Anchored below so a later divergence past the image reuses it
+    // rather than re-encoding.
+    size_t media_boundary = all_tokens.size();
+    if (!chunk_pos_media.empty()) {
+        const size_t last_media = chunk_pos_media.back();
+        for (size_t i = 0; i < chunk_pos.size(); i++) {
+            if (chunk_pos[i] == last_media) {
+                media_boundary = (i + 1 < chunk_pos.size()) ? chunk_pos[i + 1]
+                                                            : all_tokens.size();
+                break;
+            }
+        }
+    }
+
     size_t num_chunks = mtmd_input_chunks_size(chunks);
 
     for (size_t i = 0; i < chunk_pos.size(); i++) {
@@ -513,13 +598,19 @@ inline void llama_rn_context_mtmd::processMedia(
                 throw std::runtime_error("Failed to evaluate chunks");
             }
             n_past = new_n_past;
+
+            // Anchor the image before the trailing text is decoded (token-exact,
+            // no rollback needed on any arch).
+            if (capture && n_past > 0 && (size_t) n_past == media_boundary) {
+                capture(all_tokens, (size_t) n_past);
+            }
         }
     }
 
-    if (n_past == all_tokens.size() && n_past > 0 && all_tokens[n_past - 1] != LLAMA_TOKEN_NULL) {
-        // we have to evaluate at least 1 token to generate logits.
-        n_past--;
-    }
+    // Leave n_past at L and sample from the chunk eval's last-token logits
+    // (chunk_logits_last requested them) -- decode once, like mtmd-cli/server.
+    // Don't trim to L-1 and re-decode the last token: seq_rm can't roll back a
+    // recurrent/hybrid state (it no-ops), so the token would be consumed twice.
 
     // Update embd with all tokens (both text and media)
     embd = all_tokens;

diff --git a/scripts/build-hexagon-htp.sh b/scripts/build-hexagon-htp.sh
@@ -129,21 +129,22 @@ build_htp_version() {
     mkdir -p "$build_dir"
     cd "$build_dir"
 
-    # Determine generator based on environment
-    CMAKE_GENERATOR=""
+    # Determine generator based on environment (array: "Unix Makefiles" must
+    # stay one argument — unquoted expansion split it into "-G Unix" "Makefiles")
+    CMAKE_GENERATOR=()
     if command -v ninja &> /dev/null; then
-        CMAKE_GENERATOR="-GNinja"
+        CMAKE_GENERATOR=(-G Ninja)
         echo "Using Ninja build system"
     elif command -v make &> /dev/null; then
-        CMAKE_GENERATOR="-G Unix Makefiles"
+        CMAKE_GENERATOR=(-G "Unix Makefiles")
         echo "Using Make build system"
     else
         echo "Warning: Neither ninja nor make found, using default generator"
     fi
 
     # Configure with Hexagon toolchain
     cmake "${HTP_SOURCE_DIR}" \
-        $CMAKE_GENERATOR \
+        "${CMAKE_GENERATOR[@]}" \
         -DCMAKE_TOOLCHAIN_FILE="${HTP_SOURCE_DIR}/cmake-toolchain.cmake" \
         -DCMAKE_BUILD_TYPE=Release \
         -DCMAKE_INSTALL_LIBDIR="${HTP_OUTPUT_DIR}" \

diff --git a/src/types.ts b/src/types.ts
@@ -187,6 +187,18 @@ export type NativeContextParams = {
    */
   n_cpu_moe?: number
 
+  /**
+   * Memory budget (MiB) for the cross-turn KV prefix cache on recurrent/hybrid
+   * models. 0 disables it; no-op on pure-attention models. Default 160.
+   */
+  state_cache_budget_mb?: number
+
+  /**
+   * Max snapshots to keep (secondary cap; the byte budget is primary).
+   * 0 = no count cap. Default 8.
+   */
+  state_cache_max_checkpoints?: number
+
   // Embedding params
   embedding?: boolean
   embd_normalize?: number