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

Allow for growing the default KV cache when needed. (#2810)

Browse Source
This commit is contained in:
Laurent Mazare
2025-03-16 17:30:25 +01:00
committed by GitHub
parent cbf5fc80c2
commit 3afb04925a
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -11,6 +11,7 @@ pub struct Cache {
all_data: Option<Tensor>,
dim: usize,
current_seq_len: usize,
grow_by: usize,
max_seq_len: usize,
}
@ -20,6 +21,7 @@ impl Cache {
all_data: None,
dim,
current_seq_len: 0,
grow_by: max_seq_len,
max_seq_len,
}
}
@ -65,11 +67,11 @@ impl Cache {
};
let ad = self.all_data.as_mut().unwrap();
if self.current_seq_len + seq_len > self.max_seq_len {
candle::bail!(
"kv-cache: above max-seq-len {}+{seq_len}>{}",
self.current_seq_len,
self.max_seq_len
)
let mut shape = src.dims().to_vec();
shape[self.dim] = self.grow_by;
let next_ad = Tensor::zeros(shape, src.dtype(), src.device())?;
*ad = Tensor::cat(&[&*ad, &next_ad], self.dim)?;
self.max_seq_len += self.grow_by;
}
ad.slice_set(src, self.dim, self.current_seq_len)?;
self.current_seq_len += seq_len;