Clear cache for mistral

EricLBuehler · Dec 1, 2023 · 753b887 · 753b887
1 parent a6c3ec3
commit 753b887
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/openai/models/mistral.rs b/src/openai/models/mistral.rs
@@ -278,6 +278,10 @@ impl Attention {
             .reshape((b_sz, q_len, self.hidden_size))?
             .apply(&self.o_proj)
     }
+
+    fn clear_cache(&mut self) {
+        self.kv_cache = None;
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -404,4 +408,10 @@ impl Model {
             .apply(&self.norm)?
             .apply(&self.lm_head)
     }
+
+    pub fn clear_cache(&mut self) {
+        for block in &mut self.layers {
+            block.self_attn.clear_cache();
+        }
+    }
 }
diff --git a/src/openai/pipelines/mistral.rs b/src/openai/pipelines/mistral.rs
@@ -478,6 +478,8 @@ impl<'s> ModulePipeline<'s> for Mistral7BPipeline {
             }
         }
 
+        self.mistral.clear_cache();
+
         Ok((
             Some(choices),
             ChatCompletionUsageResponse {