diff --git a/candle-transformers/src/models/mistral.rs b/candle-transformers/src/models/mistral.rs index fe3187b9..2a66515b 100644 --- a/candle-transformers/src/models/mistral.rs +++ b/candle-transformers/src/models/mistral.rs @@ -297,6 +297,10 @@ impl Attention { .reshape((b_sz, q_len, self.hidden_size))? .apply(&self.o_proj) } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } } #[derive(Debug, Clone)] @@ -340,6 +344,10 @@ impl DecoderLayer { let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; residual + xs } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } } #[derive(Debug, Clone)] @@ -423,4 +431,10 @@ impl Model { .apply(&self.norm)? .apply(&self.lm_head) } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } } diff --git a/candle-transformers/src/models/quantized_mistral.rs b/candle-transformers/src/models/quantized_mistral.rs index 9e306c67..f2cb3b27 100644 --- a/candle-transformers/src/models/quantized_mistral.rs +++ b/candle-transformers/src/models/quantized_mistral.rs @@ -198,6 +198,10 @@ impl Attention { .reshape((b_sz, q_len, self.hidden_size))? .apply(&self.o_proj) } + + fn clear_kv_cache(&mut self) { + self.kv_cache = None + } } #[derive(Debug, Clone)] @@ -241,6 +245,10 @@ impl DecoderLayer { let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?; residual + xs } + + fn clear_kv_cache(&mut self) { + self.self_attn.clear_kv_cache() + } } #[derive(Debug, Clone)] @@ -322,4 +330,10 @@ impl Model { .apply(&self.norm)? .apply(&self.lm_head) } + + pub fn clear_kv_cache(&mut self) { + for layer in self.layers.iter_mut() { + layer.clear_kv_cache() + } + } }