feat: Add layer filter to recurrent cache

Branch: HybridCache Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
ggml-org · gabe-l-hart · May 20, 2025 · May 20, 2025 · May 27, 2025 · May 23, 2025
commit 6625248043bc39c1e8a89c591eeb7a4ac4f5ac7d
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
@@ -1951,12 +1951,13 @@ class llama_kv_cache_recurrent_decode_state_t : public llama_memory_decode_state
 };
 
 llama_kv_cache_recurrent::llama_kv_cache_recurrent(
-        const llama_model & model,
-                ggml_type   type_k,
-                ggml_type   type_v,
-                     bool   offload,
-                 uint32_t   kv_size,
-                 uint32_t   n_seq_max) : hparams(model.hparams), n_seq_max(n_seq_max) {
+        const llama_model &  model,
+          layer_filter_cb && filter,
+                ggml_type    type_k,
+                ggml_type    type_v,
+                     bool    offload,
+                 uint32_t    kv_size,
+                 uint32_t    n_seq_max) : hparams(model.hparams), n_seq_max(n_seq_max) {
     const int32_t n_layer = hparams.n_layer;
 
     LLAMA_LOG_INFO("%s: kv_size = %u, n_seq_max = %u, type_k = '%s', type_v = '%s', n_layer = %d\n",
@@ -1998,6 +1999,11 @@ llama_kv_cache_recurrent::llama_kv_cache_recurrent(
     v_l.reserve(n_layer);
 
     for (int i = 0; i < n_layer; i++) {
+        if (filter && !filter(i)) {
+            LLAMA_LOG_DEBUG("%s: layer %3d: skipped\n", __func__, i);
+            continue;
+        }
+
         const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i) + hparams.n_embd_k_s();
         const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i) + hparams.n_embd_v_s();
 

diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
@@ -327,12 +327,13 @@ class llama_kv_cache_unified_iswa : public llama_kv_cache {
 class llama_kv_cache_recurrent : public llama_kv_cache {
 public:
     llama_kv_cache_recurrent(
-            const llama_model & model,
-                    ggml_type   type_k,
-                    ggml_type   type_v,
-                         bool   offload,
-                     uint32_t   kv_size,
-                     uint32_t   n_seq_max);
+            const llama_model &  model,
+              layer_filter_cb && filter,
+                    ggml_type    type_k,
+                    ggml_type    type_v,
+                         bool    offload,
+                     uint32_t    kv_size,
+                     uint32_t    n_seq_max);
 
     ~llama_kv_cache_recurrent() = default;
 

diff --git a/src/llama-model.cpp b/src/llama-model.cpp
@@ -13208,6 +13208,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
             {
                 res = new llama_kv_cache_recurrent(
                         *this,
+                        nullptr,
                         GGML_TYPE_F32,
                         GGML_TYPE_F32,
                         cparams.offload_kqv,

diff --git a/tests/test-memory.cpp b/tests/test-memory.cpp
@@ -209,6 +209,7 @@ static void test_llama_kv_cache_recurrent_constructor() {
     auto model = _make_model(LLM_ARCH_MAMBA);
     llama_kv_cache_recurrent cache(
         /* model     */ *model,
+        /* filter    */ nullptr,
         /* type_k    */ GGML_TYPE_F32,
         /* type_v    */ GGML_TYPE_F16,
         /* offload   */ false,