clear kv_cache

should fix big bug in llm daemon long-term running!
groovybits · Mar 23, 2024 · 3f80a03 · 3f80a03
1 parent 310755a
commit 3f80a03
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/scripts/twitch.sh b/scripts/twitch.sh
@@ -16,7 +16,7 @@ MODEL_ID=7b-it
 MAX_TOKENS=200
 ALIGNMENT=right
 TEMPERATURE=1.0
-CONTINUOUS=0
+CONTINUOUS=1
 POLL_INTERVAL=0
 PIPELINE_CONCURRENCY=3
 TWITCH_LLM_CONCURRENCY=3

diff --git a/src/candle_gemma.rs b/src/candle_gemma.rs
@@ -59,6 +59,7 @@ impl TextGeneration {
     async fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> {
         let verbose_prompt: bool = false;
 
+        self.model.clear_kv_cache();
         self.tokenizer.clear();
         let mut tokens = self
             .tokenizer

diff --git a/src/candle_mistral.rs b/src/candle_mistral.rs
@@ -67,6 +67,10 @@ impl TextGeneration {
 
     async fn run(&mut self, prompt: &str, sample_len: usize) -> Result<()> {
         let verbose_prompt: bool = false;
+        match &mut self.model {
+            Model::Mistral(m) => m.clear_kv_cache(),
+            Model::Quantized(m) => m.clear_kv_cache(),
+        };
         self.tokenizer.clear();
         let mut tokens = self
             .tokenizer
@@ -103,7 +107,6 @@ impl TextGeneration {
             let start_pos = tokens.len().saturating_sub(context_size);
             let ctxt = &tokens[start_pos..];
             let input = Tensor::new(ctxt, &self.device)?.unsqueeze(0)?;
-            //Model::Mistral7binstructV02(m) => m.forward(&input, start_pos)?,
             let logits = match &mut self.model {
                 Model::Mistral(m) => match m.forward(&input, start_pos) {
                     Ok(logits) => logits,