From a8667896034b2d7d99e033504d6da9820830741c Mon Sep 17 00:00:00 2001
From: Layl Bongers <3094382+LaylBongers@users.noreply.github.com>
Date: Fri, 19 Apr 2024 10:06:00 +0200
Subject: [PATCH] Add workaround for kv cache

---
 src/llama.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/llama.cpp b/src/llama.cpp
index ce2f87ef9..2c97a97f7 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -5799,6 +5799,12 @@ static void llm_load_hparams(
                     default: model.type = e_model::MODEL_UNKNOWN;
                 }
             } break;
+        case LLM_ARCH_RWKV: 
+            {
+                // TODO: Re-using mamba keys right now, but RWKV isn't state-space
+                ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);
+                ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);
+            } break;
         default: (void)0;
     }