llama: R-SWA reference sliding window attention for Unlimited-OCR

conversion/deepseek.py

@@ -379,9 +379,9 @@ def set_gguf_parameters(self):
 
         self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"])
 
-        # Unlimited-OCR sliding window; written for metadata, the decoder ignores it (full MHA)
+        # Unlimited-OCR R-SWA sliding window; the deepseek2-ocr decoder reads it
         if is_ocr:
-            sliding_window = hparams.get("sliding_window_size") or hparams.get("sliding_window")
+            sliding_window = hparams.get("sliding_window_size", hparams.get("sliding_window"))
             if sliding_window:
                 self.gguf_writer.add_sliding_window(sliding_window)
 

include/llama.h

@@ -43,10 +43,10 @@
 #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'
 
 #define LLAMA_SESSION_MAGIC   LLAMA_FILE_MAGIC_GGSN
-#define LLAMA_SESSION_VERSION 11
+#define LLAMA_SESSION_VERSION 12
 
 #define LLAMA_STATE_SEQ_MAGIC   LLAMA_FILE_MAGIC_GGSQ
-#define LLAMA_STATE_SEQ_VERSION 4
+#define LLAMA_STATE_SEQ_VERSION 5
 
 #ifdef __cplusplus
 extern "C" {

src/llama-graph.cpp

@@ -449,6 +449,7 @@ static void print_mask(const T * data, int64_t n_tokens, int64_t n_kv, int64_t n
         case LLAMA_SWA_TYPE_STANDARD:  swa_type_str = "LLAMA_SWA_TYPE_STANDARD"; break;
         case LLAMA_SWA_TYPE_CHUNKED:   swa_type_str = "LLAMA_SWA_TYPE_CHUNKED"; break;
         case LLAMA_SWA_TYPE_SYMMETRIC: swa_type_str = "LLAMA_SWA_TYPE_SYMMETRIC"; break;
+        case LLAMA_SWA_TYPE_REFERENCE: swa_type_str = "LLAMA_SWA_TYPE_REFERENCE"; break;
     };
 
     LLAMA_LOG_DEBUG("%s: n_swa : %d, n_kv: %d, swa_type: %s\n", __func__, (int)n_swa, (int)n_kv, swa_type_str);
@@ -3209,7 +3210,9 @@ static std::unique_ptr<llm_graph_input_attn_kv> build_attn_inp_kv_impl(
     auto inp = std::make_unique<llm_graph_input_attn_kv>(hparams, cparams, mctx_cur);
 
     {
-        GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA");
+        // REFERENCE masks within this single cache; other SWA types need iswa
+        GGML_ASSERT((hparams.swa_type == LLAMA_SWA_TYPE_NONE ||
+                     hparams.swa_type == LLAMA_SWA_TYPE_REFERENCE) && "Use llama_kv_cache_iswa for SWA");
 
         inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch);
         inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch);

src/llama-hparams.h

@@ -26,6 +26,7 @@ enum llama_swa_type {
     LLAMA_SWA_TYPE_STANDARD  = 1,
     LLAMA_SWA_TYPE_CHUNKED   = 2,
     LLAMA_SWA_TYPE_SYMMETRIC = 3,
+    LLAMA_SWA_TYPE_REFERENCE = 4, // R-SWA: always-visible prefix + window over the rest
 };
 
 // how the non-causal mask should be constructed with llama_set_causal_attn(ctx, false)
@@ -529,7 +530,8 @@ struct llama_hparams {
     // note: inlined on purpose for performance reasons
     // TODO: think of a better place for this function
     // TODO: pack the SWA params in a struct?
-    static bool is_masked_swa(uint32_t n_swa, llama_swa_type swa_type, llama_pos p0, llama_pos p1) {
+    // n_ref = R-SWA prefix length L_m (always-visible positions); < 0 = unlatched, full causal
+    static bool is_masked_swa(uint32_t n_swa, llama_swa_type swa_type, llama_pos p0, llama_pos p1, llama_pos n_ref = -1) {
         assert(p0 >= 0 && p1 >= 0);
 
         switch (swa_type) {
@@ -560,6 +562,12 @@ struct llama_hparams {
                         return true;
                     }
                 } break;
+            case LLAMA_SWA_TYPE_REFERENCE:
+                {
+                    if (p1 - p0 >= (int32_t) n_swa && n_ref >= 0 && p0 >= n_ref) {
+                        return true;
+                    }
+                } break;
         }
 
         return false;

src/llama-kv-cache.cpp

@@ -87,6 +87,8 @@ llama_kv_cache::llama_kv_cache(
     v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
     v_cells(*v_cells_impl) {
 
+    n_ref.fill(-1);
+
     // shared cells view the source cache's K/V tensors, so the cell count
     // follows the source allocation: a fitted target can be smaller than the
     // draft default and oversized views would overflow the source tensors
@@ -375,6 +377,8 @@ llama_kv_cache::llama_kv_cache(
 }
 
 void llama_kv_cache::clear(bool data) {
+    n_ref.fill(-1);
+
     for (uint32_t s = 0; s < n_stream; ++s) {
         v_cells[s].reset();
         v_heads[s] = 0;
@@ -404,6 +408,20 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
         p1 = std::numeric_limits<llama_pos>::max();
     }
 
+    // removal at or into the prefix [0, n_ref) invalidates the latch; a re-extended
+    // prompt then re-latches at its new boundary on the next single-token decode
+    if (seq_id >= 0) {
+        if (p0 <= n_ref[seq_id]) {
+            n_ref[seq_id] = -1;
+        }
+    } else {
+        for (auto & r : n_ref) {
+            if (p0 <= r) {
+                r = -1;
+            }
+        }
+    }
+
     if (seq_id >= 0) {
         auto & cells = v_cells[seq_to_stream[seq_id]];
         auto & head  = v_heads[seq_to_stream[seq_id]];
@@ -465,6 +483,9 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll
     GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());
     GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());
 
+    // copy inherits the latched prefix
+    n_ref[seq_id_dst] = n_ref[seq_id_src];
+
     const auto s0 = seq_to_stream[seq_id_src];
     const auto s1 = seq_to_stream[seq_id_dst];
 
@@ -556,6 +577,13 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
 
     GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
 
+    // other seqs are purged -> drop their latched prefix
+    for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {
+        if ((llama_seq_id) s != seq_id) {
+            n_ref[s] = -1;
+        }
+    }
+
     auto & cells = v_cells[seq_to_stream[seq_id]];
     auto & head  = v_heads[seq_to_stream[seq_id]];
 
@@ -606,6 +634,12 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll
         return;
     }
 
+    // the prefix boundary is an absolute pos -> shift it with its cells; a shift below
+    // pos 0 empties the prefix rather than colliding with the -1 unlatched sentinel
+    if (n_ref[seq_id] >= 0 && n_ref[seq_id] >= p0 && n_ref[seq_id] < p1) {
+        n_ref[seq_id] = std::max(n_ref[seq_id] + shift, 0);
+    }
+
     for (uint32_t i = 0; i < cells.size(); ++i) {
         if (!cells.pos_in(i, p0, p1)) {
             continue;
@@ -653,6 +687,11 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in
         return;
     }
 
+    // the boundary is one-past-end -> ceil-divide so the last prefix cell (floored) stays inside
+    if (n_ref[seq_id] >= 0 && n_ref[seq_id] >= p0 && n_ref[seq_id] < p1) {
+        n_ref[seq_id] = (n_ref[seq_id] + d - 1) / d;
+    }
+
     for (uint32_t i = 0; i < cells.size(); ++i) {
         if (!cells.pos_in(i, p0, p1)) {
             continue;
@@ -755,6 +794,9 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
     // remember the old state of the cells so we can restore it in the end
     std::vector<state_t> states;
 
+    // apply_ubatch's R-SWA latch is a side effect that must not survive the dry run
+    const auto n_ref_old = n_ref;
+
     bool success = true;
 
     for (const auto & ubatch : ubatches) {
@@ -800,6 +842,8 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
         }
     }
 
+    n_ref = n_ref_old;
+
     if (!success) {
         return {};
     }
@@ -1050,6 +1094,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
                         const llama_seq_id seq_id_cell = cells.seq_get(idx);
 
                         // SWA mask
+                        // deliberately no n_ref arg: a REFERENCE cell must never be recycled (never-evict)
                         if (llama_hparams::is_masked_swa(n_swa, swa_type, pos_cell, cells.seq_pos_max(seq_id_cell) + 1)) {
                             can_use = true;
                         }
@@ -1096,6 +1141,32 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
         return;
     }
 
+    // latch L_m at the prefill->decode boundary: the first single-token append to a
+    // populated seq that requests output. prefill remainders (ubatch-split tails, mtmd
+    // separator chunks) carry no output flag and must not latch mid-prompt. assumes
+    // single-token decode; a prompt whose last ubatch is exactly 1 token latches one
+    // early, and multi-output decode (speculative) is excluded via ref_latch_enabled.
+    if (swa_type == LLAMA_SWA_TYPE_REFERENCE && ref_latch_enabled) {
+        uint32_t n_tok_seq[LLAMA_MAX_SEQ] = { 0 };
+        for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+            for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
+                n_tok_seq[ubatch.seq_id[i][s]]++;
+            }
+        }
+        for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+            if (!ubatch.output[i]) {
+                continue;
+            }
+            for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
+                const llama_seq_id seq_id = ubatch.seq_id[i][s];
+                if (n_ref[seq_id] < 0 && n_tok_seq[seq_id] == 1 &&
+                    v_cells[seq_to_stream[seq_id]].seq_pos_max(seq_id) >= 0) {
+                    n_ref[seq_id] = ubatch.pos[i];
+                }
+            }
+        }
+    }
+
     // keep track of the max sequence position that we would overwrite with this ubatch
     // for non-SWA cache, this would be always empty
     llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ];
@@ -1537,6 +1608,9 @@ struct args_set_input_kq_mask {
     uint32_t       n_swa;
     llama_swa_type swa_type;
 
+    // per-seq R-SWA prefix length L_m (-1 = unlatched), indexed by seq_id
+    const llama_pos * n_ref;
+
     int64_t n_kv;
     int64_t n_stream;
     int64_t n_tps;
@@ -1674,7 +1748,7 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
                 if (swa) {
                     // see llama_non_causal_type
                     const bool in_span = !causal && args.hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_SWA_FULL && p0 >= seq_pos_min[seq_id];
-                    if (!in_span && llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
+                    if (!in_span && llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1, args.n_ref[seq_id])) {
                         goto skip;
                     }
                 }
@@ -1763,6 +1837,7 @@ void llama_kv_cache::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * u
         /*.seq_to_stream    =*/ seq_to_stream,
         /*.n_swa            =*/ n_swa,
         /*.swa_type         =*/ swa_type,
+        /*.n_ref            =*/ n_ref.data(),
         /*.n_kv             =*/ n_kv,
         /*.n_stream         =*/ n_stream,
         /*.n_tps            =*/ n_tps,
@@ -2059,6 +2134,20 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
 
     io.write(&n_stream, sizeof(n_stream));
 
+    // persist n_ref; REFERENCE-guarded so other cache types' state format is unchanged.
+    // whole-cache case is count-prefixed to tolerate a different n_seq_max on restore.
+    if (swa_type == LLAMA_SWA_TYPE_REFERENCE) {
+        if (seq_id == -1) {
+            const uint32_t n_ref_count = n_seq_max;
+            io.write(&n_ref_count, sizeof(n_ref_count));
+            for (uint32_t i = 0; i < n_ref_count; ++i) {
+                io.write(&n_ref[i], sizeof(llama_pos));
+            }
+        } else {
+            io.write(&n_ref[seq_id], sizeof(llama_pos));
+        }
+    }
+
     for (uint32_t s = 0; s < n_stream; ++s) {
         cell_ranges_t cr { s, {} };
 
@@ -2077,6 +2166,7 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
             add_cell = add_cell && (seq_id == -1 || cells.seq_has(i, seq_id));
 
             // check the cell is not SWA-masked
+            // deliberately no n_ref arg: a REFERENCE seq must persist all its cells
             if (add_cell && seq_id != -1) {
                 const bool is_masked = llama_hparams::is_masked_swa(n_swa, swa_type, cells.pos_get(i), cells.seq_pos_max(seq_id));
 
@@ -2155,6 +2245,25 @@ const slot_info_vec_t *   sinfos_in) {
         throw std::runtime_error("n_stream mismatch");
     }
 
+    // read n_ref now but apply after the restore below; clear()/seq_rm() would reset it
+    std::array<llama_pos, LLAMA_MAX_SEQ> n_ref_restored;
+    n_ref_restored.fill(-1);
+    if (swa_type == LLAMA_SWA_TYPE_REFERENCE) {
+        if (seq_id == -1) {
+            uint32_t n_ref_count = 0;
+            io.read(&n_ref_count, sizeof(n_ref_count));
+            for (uint32_t i = 0; i < n_ref_count; ++i) {
+                llama_pos v;
+                io.read(&v, sizeof(v));
+                if (i < n_seq_max) {
+                    n_ref_restored[i] = v;
+                }
+            }
+        } else {
+            io.read(&n_ref_restored[seq_id], sizeof(llama_pos));
+        }
+    }
+
     // a whole-context restore replaces every stream, so the cache is emptied once here
     // clear() resets all streams at once, so doing it per stream below would keep only the last one
     if (seq_id == -1) {
@@ -2195,6 +2304,17 @@ const slot_info_vec_t *   sinfos_in) {
             (*sinfos_out)[s] = sinfo;
         }
     }
+
+    // cells restored -> reinstate n_ref
+    if (swa_type == LLAMA_SWA_TYPE_REFERENCE) {
+        if (seq_id == -1) {
+            for (uint32_t i = 0; i < n_seq_max; ++i) {
+                n_ref[i] = n_ref_restored[i];
+            }
+        } else {
+            n_ref[seq_id] = n_ref_restored[seq_id];
+        }
+    }
 }
 
 void llama_kv_cache::state_write_meta(llama_io_write_i & io, const cell_ranges_t & cr, llama_seq_id seq_id) const {

src/llama-kv-cache.h

@@ -133,6 +133,10 @@ class llama_kv_cache : public llama_memory_i {
 
     bool get_can_shift() const override;
 
+    void set_ref_latch_enabled(bool value) {
+        ref_latch_enabled = value;
+    }
+
     void clear(bool data) override;
 
     bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;
@@ -292,6 +296,14 @@ class llama_kv_cache : public llama_memory_i {
     // this is the SWA type of the cache - not to be confused with the model SWA type
     const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
 
+    // R-SWA per-seq prefix length L_m (-1 = unlatched -> full causal mask); latched at the
+    // prefill->decode boundary in apply_ubatch, read by set_input_kq_mask.
+    std::array<llama_pos, LLAMA_MAX_SEQ> n_ref;
+
+    // multi-output decode (speculative) breaks the single-token boundary inference,
+    // so the context disables the latch there -> full causal
+    bool ref_latch_enabled = true;
+
     // ggml contexts for the KV cache along with the allocated backend buffers:
     llama_memory_bufs ctxs_bufs;
 

src/llama-model.cpp

@@ -2859,6 +2859,37 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         filter_idx,
                         nullptr);
             } break;
+        case LLM_ARCH_DEEPSEEK2OCR:
+            {
+                // one full cache: the REFERENCE mask keeps the prefix visible, iswa would evict it.
+                // an F16 V cache truncates enough to garble dense-table OCR (HF attends in F32).
+                ggml_type type_v = params.type_v;
+                if (hparams.swa_type == LLAMA_SWA_TYPE_REFERENCE && type_v == GGML_TYPE_F16) {
+                    type_v = GGML_TYPE_F32;
+                    LLAMA_LOG_WARN("%s: promoting V cache type to f32 for R-SWA (pass a quantized -ctv to override)\n", __func__);
+                }
+                auto * kv = new llama_kv_cache(
+                        *this,
+                        params.no_alloc,
+                        hparams,
+                        params.type_k,
+                        type_v,
+                        !cparams.flash_attn,
+                        cparams.offload_kqv,
+                        cparams.kv_unified,
+                        cparams.n_ctx_seq,
+                        cparams.n_seq_max,
+                        hparams.n_swa,
+                        hparams.swa_type,
+                        nullptr,
+                        nullptr,
+                        nullptr,
+                        nullptr);
+                // the latch infers the prefill->decode boundary from single-token appends,
+                // which multi-output decode (speculative) breaks -> stay full causal there
+                kv->set_ref_latch_enabled(cparams.n_outputs_max_per_seq == 1);
+                res = kv;
+            } break;
         case LLM_ARCH_GLM_DSA:
         case LLM_ARCH_DEEPSEEK32:
             {
@@ -3489,6 +3520,11 @@ int32_t llama_model_n_swa(const llama_model * model) {
     if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) {
         return 0;
     }
+    // deepseek2-ocr R-SWA runs on a full cache that never evicts, so no SWA
+    // cache-invalidation applies
+    if (model->arch == LLM_ARCH_DEEPSEEK2OCR) {
+        return 0;
+    }
     return model->hparams.n_swa;
 }