huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 11:56:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

More kv-cache testing.

Browse Source
This commit is contained in:
Laurent
2024-09-22 12:51:07 +02:00
parent f9579f80be
commit 6547c4bfc3
2 changed files with 27 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-nn/src/kv_cache.rs
View File

@ -228,7 +228,7 @@ impl RotatingCache {
self.offset = 0;
} else {
let rem_len = self.max_seq_len - self.offset;
if rem_len <= seq_len {
if seq_len <= rem_len {
ad.slice_set(src, self.dim, self.offset)?;
self.offset = (self.offset + seq_len) % self.max_seq_len;
} else {

26
candle-nn/tests/kv_cache.rs
View File

@ -30,3 +30,29 @@ fn kv_cache() -> Result<()> {
}
Ok(())
}
#[test]
fn rotating_kv_cache() -> Result<()> {
let mut cache = candle_nn::kv_cache::RotatingCache::new(0, 6);
for _ in [0, 1] {
assert_eq!(cache.current_seq_len(), 0);
let data = cache.current_data()?;
assert!(data.is_none());
let t = Tensor::new(&[1f32, 2., 3.], &Device::Cpu)?;
cache.append(&t)?;
let data = cache.current_data()?.unwrap();
assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3.]);
let t = Tensor::new(&[4f32], &Device::Cpu)?;
cache.append(&t)?;
let data = cache.current_data()?.unwrap();
assert_eq!(data.to_vec1::<f32>()?, [1., 2., 3., 4.]);
let t = Tensor::new(&[0f32, 5., 6., 7.], &Device::Cpu)?;
cache.append(&t)?;
let data = cache.current_data()?.unwrap();
assert_eq!(data.to_vec1::<f32>()?, [6., 7., 3., 4., 0., 5.]);
assert_eq!(cache.current_seq_len(), 8);
assert_eq!(cache.offset(), 2);
cache.reset();
}
Ok(())
}