cont : fix

Author: ggerganov

src/llama-graph.cpp

@@ -235,6 +235,12 @@ void llm_graph_input_cls::set_input(const llama_ubatch * ubatch) {
     }
 }
 
+llm_graph_input_rs::llm_graph_input_rs(const llama_memory_recurrent_context * mctx) :
+    mctx(mctx),
+    head(mctx->get_head()),
+    rs_z(mctx->get_rs_z()) {
+}
+
 void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
     GGML_UNUSED(ubatch);
 
@@ -254,9 +260,6 @@ void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
 bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {
     const auto * mctx = static_cast<const llama_memory_recurrent_context *>(params.mctx);
 
-    const auto prev_head = this->mctx->get_head();
-    const auto prev_rs_z = this->mctx->get_rs_z();
-
     this->mctx = mctx;
 
     bool res = true;
@@ -266,8 +269,8 @@ bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {
     res &= s_copy_main->ne[0]  == params.ubatch.n_seqs;
     res &= s_copy_extra->ne[0] == mctx->get_n_rs() - params.ubatch.n_seqs;
 
-    res &= prev_head == mctx->get_head();
-    res &= prev_rs_z == mctx->get_rs_z();
+    res &= this->head == mctx->get_head();
+    res &= this->rs_z == mctx->get_rs_z();
 
     return res;
 }
@@ -478,9 +481,6 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
 bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
     const auto * mctx = static_cast<const llama_memory_hybrid_context *>(params.mctx);
 
-    const auto prev_head = this->mctx->get_recr()->get_head();
-    const auto prev_rs_z = this->mctx->get_recr()->get_rs_z();
-
     this->mctx = mctx;
 
     bool res = true;
@@ -496,8 +496,8 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
     res &= inp_rs->s_copy_main->ne[0]  == params.ubatch.n_seqs;
     res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs;
 
-    res &= prev_head == mctx->get_recr()->get_head();
-    res &= prev_rs_z == mctx->get_recr()->get_rs_z();
+    res &= inp_rs->head == mctx->get_recr()->get_head();
+    res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z();
 
     return res;
 }

src/llama-graph.h

@@ -219,7 +219,7 @@ class llm_graph_input_cls : public llm_graph_input_i {
 
 class llm_graph_input_rs : public llm_graph_input_i {
 public:
-    llm_graph_input_rs(const llama_memory_recurrent_context * mctx) : mctx(mctx) {}
+    llm_graph_input_rs(const llama_memory_recurrent_context * mctx);
     virtual ~llm_graph_input_rs() = default;
 
     void set_input(const llama_ubatch * ubatch) override;
@@ -234,6 +234,10 @@ class llm_graph_input_rs : public llm_graph_input_i {
     ggml_tensor * s_copy_extra;  // I32 [n_rs - n_seqs]
 
     const llama_memory_recurrent_context * mctx;
+
+    // need to match for valid graph reuse
+    const uint32_t head;
+    const  int32_t rs_z;
 };
 
 class llm_graph_input_cross_embd : public llm_graph_input_i {