huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expose a function to clear the KV cache on mixformers. (#964)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-26 05:41:07 +01:00
committed by GitHub
parent a084f65f9a
commit 1fcac4afed
2 changed files with 24 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
candle-transformers/src/models/mixformer.rs
View File

@ -287,6 +287,10 @@ impl MHA {
.flatten_from(D::Minus2)?;
attn_output.apply(&self.out_proj)
}
fn clear_kv_cache(&mut self) {
self.kv_cache = None
}
}
#[derive(Debug)]
@ -318,6 +322,10 @@ impl ParallelBlock {
let feed_forward_hidden_states = self.mlp.forward(&xs)?;
attn_outputs + feed_forward_hidden_states + residual
}
fn clear_kv_cache(&mut self) {
self.mixer.clear_kv_cache()
}
}
#[derive(Debug)]
@ -360,4 +368,8 @@ impl MixFormerSequentialForCausalLM {
}
xs.narrow(1, seq_len - 1, 1)?.apply(&self.head)?.squeeze(1)
}
pub fn clear_kv_cache(&mut self) {
self.blocks.iter_mut().for_each(|b| b.clear_kv_cache())
}
}

12
candle-transformers/src/models/quantized_mixformer.rs
View File

@ -268,6 +268,10 @@ impl MHA {
.flatten_from(D::Minus2)?;
attn_output.apply(&self.out_proj)
}
fn clear_kv_cache(&mut self) {
self.kv_cache = None
}
}
#[derive(Debug)]
@ -299,6 +303,10 @@ impl ParallelBlock {
let feed_forward_hidden_states = self.mlp.forward(&xs)?;
attn_outputs + feed_forward_hidden_states + residual
}
fn clear_kv_cache(&mut self) {
self.mixer.clear_kv_cache()
}
}
#[derive(Debug)]
@ -341,4 +349,8 @@ impl MixFormerSequentialForCausalLM {
}
xs.narrow(1, seq_len - 1, 1)?.apply(&self.head)?.squeeze(1)
}
pub fn clear_kv_cache(&mut self) {
self.blocks.iter_mut().for_each(|b| b.clear_kv_cache())
}
}