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

Rotating kv cache positions (#2901)

Browse Source 
* Retrieve the current positions for rotating KV caches.

* Add the function to the kv cache too.

* More testing.
This commit is contained in:
Laurent Mazare
2025-04-15 23:09:26 +02:00
committed by GitHub
parent 76e565c4ab
commit 7f0f83a7c1
2 changed files with 41 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -39,9 +39,16 @@ fn rotating_kv_cache() -> Result<()> {
assert_eq!(cache.current_seq_len(), 0);
let data = cache.current_data()?;
assert!(data.is_none());
assert_eq!(cache.positions(1), &[0]);
assert_eq!(cache.positions(2), &[0, 1]);
let t = Tensor::new(&[1., 2., 3.], &Device::Cpu)?;
let data = cache.append(&t)?;
assert_eq!(data.to_vec1::<f64>()?, [1., 2., 3.]);
assert_eq!(cache.positions(0), &[0, 1, 2]);
assert_eq!(cache.positions(1), &[0, 1, 2, 3]);
assert_eq!(cache.positions(2), &[0, 1, 2, 3, 4]);
assert_eq!(cache.positions(3), &[0, 1, 2, 3, 4, 5]);
assert_eq!(cache.positions(4), &[6, 1, 2, 3, 4, 5]);
let t = Tensor::new(&[4.], &Device::Cpu)?;
let data = cache.append(&t)?;
assert_eq!(data.to_vec1::<f64>()?, [1., 2., 3., 4.]);
@ -79,11 +86,17 @@ fn rotating_kv_cache() -> Result<()> {
mask.to_vec2::<u8>()?,
&[[0, 0, 1, 1, 0, 0], [0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 0]],
);
assert_eq!(cache.positions(0), &[12, 7, 8, 9, 10, 11]);
assert_eq!(cache.positions(2), &[12, 13, 14, 9, 10, 11]);
assert_eq!(cache.positions(3), &[12, 13, 14, 15, 10, 11]);
assert_eq!(cache.positions(8), &[13, 14, 15, 16, 17, 18, 19, 20]);
let t = Tensor::new(&[0., 1., 2., 3., 4., 5., 6., 7., 8.], &Device::Cpu)?;
let data = cache.append(&t)?;
assert_eq!(data.to_vec1::<f64>()?, [0., 1., 2., 3., 4., 5., 6., 7., 8.]);
assert_eq!(cache.current_seq_len(), 22);
assert_eq!(cache.offset(), 0);
assert_eq!(cache.positions(0), &[16, 17, 18, 19, 20, 21]);
assert_eq!(cache.positions(1), &[22, 17, 18, 19, 20, 21]);
let mask = cache.attn_mask(1, &Device::Cpu)?;
assert!(mask.is_none());