refactor: Make status and child states const in hybrid and iswa

Branch: HybridRecurrentCache

Signed-off-by: Gabe Goodhart <[email protected]>

Author: gabe-l-hart

src/llama-kv-cache-hybrid-recurrent.cpp

@@ -171,35 +171,32 @@ llama_kv_cache_recurrent * llama_kv_cache_hybrid_recurrent::get_kv_recurrent() c
 
 llama_kv_cache_hybrid_recurrent_state::llama_kv_cache_hybrid_recurrent_state(llama_memory_status status) : status(status) {}
 
-llama_kv_cache_hybrid_recurrent_state::llama_kv_cache_hybrid_recurrent_state(llama_kv_cache_hybrid_recurrent * kv)
-    : status(LLAMA_MEMORY_STATUS_SUCCESS) {
-    state_attn      = kv->get_kv_attn     ()->init_full();
-    state_recurrent = kv->get_kv_recurrent()->init_full();
-
-    status = llama_memory_status_combine(state_attn->get_status(), state_recurrent->get_status());
+llama_kv_cache_hybrid_recurrent_state::llama_kv_cache_hybrid_recurrent_state(llama_kv_cache_hybrid_recurrent * kv) :
+    state_attn     (kv->get_kv_attn     ()->init_full()),
+    state_recurrent(kv->get_kv_recurrent()->init_full()),
+    status(llama_memory_status_combine(state_attn->get_status(), state_recurrent->get_status())) {
 }
 
 llama_kv_cache_hybrid_recurrent_state::llama_kv_cache_hybrid_recurrent_state(
         llama_kv_cache_hybrid_recurrent * kv,
         llama_context * lctx,
-        bool optimize) : status(LLAMA_MEMORY_STATUS_SUCCESS) {
-    state_attn      = kv->get_kv_attn     ()->init_update(lctx, optimize);
-    state_recurrent = kv->get_kv_recurrent()->init_update(lctx, optimize);
-
-    status = llama_memory_status_combine(state_attn->get_status(), state_recurrent->get_status());
+        bool optimize) :
+    state_attn     (kv->get_kv_attn     ()->init_update(lctx, optimize)),
+    state_recurrent(kv->get_kv_recurrent()->init_update(lctx, optimize)),
+    status(llama_memory_status_combine(state_attn->get_status(), state_recurrent->get_status())) {
 }
 
 llama_kv_cache_hybrid_recurrent_state::llama_kv_cache_hybrid_recurrent_state(
         llama_kv_cache_hybrid_recurrent * kv,
                            llama_sbatch   sbatch,
                   std::vector<uint32_t>   heads_attn,
-              std::vector<llama_ubatch>   ubatches)
-    : status(LLAMA_MEMORY_STATUS_SUCCESS),
+              std::vector<llama_ubatch>   ubatches) :
     sbatch(std::move(sbatch)),
-    ubatches(std::move(ubatches)) {
+    ubatches(std::move(ubatches)),
     // note: here we copy the ubatches. not sure if this is ideal
-    state_attn     .reset(new llama_kv_cache_unified_state  (kv->get_kv_attn(),      {}, std::move(heads_attn), this->ubatches));
-    state_recurrent.reset(new llama_kv_cache_recurrent_state(kv->get_kv_recurrent(), {},                        this->ubatches));
+    state_attn     (new llama_kv_cache_unified_state  (kv->get_kv_attn(),      {}, std::move(heads_attn), this->ubatches)),
+    state_recurrent(new llama_kv_cache_recurrent_state(kv->get_kv_recurrent(), {},                        this->ubatches)),
+    status(LLAMA_MEMORY_STATUS_SUCCESS) {
 }
 
 bool llama_kv_cache_hybrid_recurrent_state::next() {

src/llama-kv-cache-hybrid-recurrent.h

@@ -130,15 +130,15 @@ class llama_kv_cache_hybrid_recurrent_state : public llama_memory_state_i {
     const llama_kv_cache_recurrent_state * get_state_recurrent() const;
 
 private:
-    llama_memory_status status;
-
     llama_sbatch sbatch;
 
     // the index of the next ubatch to process
     size_t i_next = 0;
 
     std::vector<llama_ubatch> ubatches;
 
-    llama_memory_state_ptr state_attn;
-    llama_memory_state_ptr state_recurrent;
+    const llama_memory_state_ptr state_attn;
+    const llama_memory_state_ptr state_recurrent;
+
+    const llama_memory_status status;
 };

src/llama-kv-cache-unified-iswa.cpp

@@ -197,37 +197,33 @@ llama_kv_cache_unified * llama_kv_cache_unified_iswa::get_swa() const {
 llama_kv_cache_unified_iswa_state::llama_kv_cache_unified_iswa_state(llama_memory_status status) : status(status) {}
 
 llama_kv_cache_unified_iswa_state::llama_kv_cache_unified_iswa_state(
-        llama_kv_cache_unified_iswa * kv) : status(LLAMA_MEMORY_STATUS_SUCCESS) {
-    state_base = kv->get_base()->init_full();
-    state_swa  = kv->get_swa ()->init_full();
-
-    status = llama_memory_status_combine(state_base->get_status(), state_swa->get_status());
+        llama_kv_cache_unified_iswa * kv) :
+    state_base(kv->get_base()->init_full()),
+    state_swa (kv->get_swa ()->init_full()),
+    status(llama_memory_status_combine(state_base->get_status(), state_swa->get_status())) {
 }
 
 llama_kv_cache_unified_iswa_state::llama_kv_cache_unified_iswa_state(
         llama_kv_cache_unified_iswa * kv,
         llama_context * lctx,
-        bool optimize) : status(LLAMA_MEMORY_STATUS_SUCCESS) {
-    state_base = kv->get_base()->init_update(lctx, optimize);
-    state_swa  = kv->get_swa ()->init_update(lctx, optimize);
-
-    status = llama_memory_status_combine(state_base->get_status(), state_swa->get_status());
+        bool optimize) :
+    state_base(kv->get_base()->init_update(lctx, optimize)),
+    state_swa (kv->get_swa ()->init_update(lctx, optimize)),
+    status(llama_memory_status_combine(state_base->get_status(), state_swa->get_status())) {
 }
 
 llama_kv_cache_unified_iswa_state::llama_kv_cache_unified_iswa_state(
         llama_kv_cache_unified_iswa * kv,
         llama_sbatch sbatch,
         std::vector<uint32_t> heads_base,
         std::vector<uint32_t> heads_swa,
-        std::vector<llama_ubatch> ubatches)
-        : status(LLAMA_MEMORY_STATUS_SUCCESS),
-        sbatch(std::move(sbatch)),
-        ubatches(std::move(ubatches)) {
+        std::vector<llama_ubatch> ubatches) :
+    sbatch(std::move(sbatch)),
+    ubatches(std::move(ubatches)),
     // note: here we copy the ubatches. not sure if this is ideal
-    state_base.reset(new llama_kv_cache_unified_state(kv->get_base(), {}, std::move(heads_base), this->ubatches));
-    state_swa .reset(new llama_kv_cache_unified_state(kv->get_swa (), {}, std::move(heads_swa),  this->ubatches));
-
-    status = llama_memory_status_combine(state_base->get_status(), state_swa->get_status());
+    state_base(new llama_kv_cache_unified_state(kv->get_base(), {}, std::move(heads_base), this->ubatches)),
+    state_swa (new llama_kv_cache_unified_state(kv->get_swa (), {}, std::move(heads_swa),  this->ubatches)),
+    status(llama_memory_status_combine(state_base->get_status(), state_swa->get_status())) {
 }
 
 llama_kv_cache_unified_iswa_state:: ~llama_kv_cache_unified_iswa_state() = default;

src/llama-kv-cache-unified-iswa.h

@@ -117,8 +117,6 @@ class llama_kv_cache_unified_iswa_state : public llama_memory_state_i {
     const llama_kv_cache_unified_state * get_swa()  const;
 
 private:
-    llama_memory_status status;
-
     //llama_kv_cache_unified_iswa * kv;
 
     llama_sbatch sbatch;
@@ -128,6 +126,8 @@ class llama_kv_cache_unified_iswa_state : public llama_memory_state_i {
 
     std::vector<llama_ubatch> ubatches;
 
-    llama_memory_state_ptr state_base;
-    llama_memory_state_ptr state_swa;
+    const llama_memory_state_ptr state_base;
+    const llama_memory_state_ptr state_swa;
+
+    const llama_memory_status status;
 };