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

Do not use the kv-cache on external key-value states. (#1054)

Browse Source
This commit is contained in:
Laurent Mazare
2023-10-07 22:37:19 +01:00
committed by GitHub
parent 823fe23f9b
commit 2e5fb0b251
3 changed files with 14 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
candle-examples/examples/quantized-t5/main.rs
View File

@ -153,7 +153,6 @@ fn main() -> Result<()> {
let args = Args::parse(); let args = Args::parse();
let _guard = if args.tracing { let _guard = if args.tracing {
println!("tracing...");
let (chrome_layer, guard) = ChromeLayerBuilder::new().build(); let (chrome_layer, guard) = ChromeLayerBuilder::new().build();
tracing_subscriber::registry().with(chrome_layer).init(); tracing_subscriber::registry().with(chrome_layer).init();
Some(guard) Some(guard)

14
candle-transformers/src/models/quantized_t5.rs
View File

@ -348,21 +348,21 @@ impl T5Attention {
.contiguous()?; .contiguous()?;
let mut k = k let mut k = k
.reshape((b_sz, kv_len, self.n_heads, self.d_kv))? .reshape((b_sz, kv_len, self.n_heads, self.d_kv))?
.transpose(1, 2)? .transpose(1, 2)?;
.contiguous()?;
let mut v = v let mut v = v
.reshape((b_sz, kv_len, self.n_heads, self.d_kv))? .reshape((b_sz, kv_len, self.n_heads, self.d_kv))?
.transpose(1, 2)? .transpose(1, 2)?;
.contiguous()?;
if self.use_cache { if self.use_cache && key_value_states.is_none() {
let _enter = self.span_cache.enter(); let _enter = self.span_cache.enter();
if let Some((kv_cache_k, kv_cache_v)) = &self.kv_cache { if let Some((kv_cache_k, kv_cache_v)) = &self.kv_cache {
k = Tensor::cat(&[kv_cache_k, &k], 2)?.contiguous()?; k = Tensor::cat(&[kv_cache_k, &k], 2)?;
v = Tensor::cat(&[kv_cache_v, &v], 2)?.contiguous()?; v = Tensor::cat(&[kv_cache_v, &v], 2)?;
}; };
self.kv_cache = Some((k.clone(), v.clone())); self.kv_cache = Some((k.clone(), v.clone()));
}; };
let k = k.contiguous()?;
let v = v.contiguous()?;
// TODO: Use flash_attn. // TODO: Use flash_attn.
let scores = { let scores = {
let _enter = self.span_mm.enter(); let _enter = self.span_mm.enter();

14
candle-transformers/src/models/t5.rs
View File

@ -348,21 +348,21 @@ impl T5Attention {
.contiguous()?; .contiguous()?;
let mut k = k let mut k = k
.reshape((b_sz, kv_len, self.n_heads, self.d_kv))? .reshape((b_sz, kv_len, self.n_heads, self.d_kv))?
.transpose(1, 2)? .transpose(1, 2)?;
.contiguous()?;
let mut v = v let mut v = v
.reshape((b_sz, kv_len, self.n_heads, self.d_kv))? .reshape((b_sz, kv_len, self.n_heads, self.d_kv))?
.transpose(1, 2)? .transpose(1, 2)?;
.contiguous()?;
if self.use_cache { if self.use_cache && key_value_states.is_none() {
let _enter = self.span_cache.enter(); let _enter = self.span_cache.enter();
if let Some((kv_cache_k, kv_cache_v)) = &self.kv_cache { if let Some((kv_cache_k, kv_cache_v)) = &self.kv_cache {
k = Tensor::cat(&[kv_cache_k, &k], 2)?.contiguous()?; k = Tensor::cat(&[kv_cache_k, &k], 2)?;
v = Tensor::cat(&[kv_cache_v, &v], 2)?.contiguous()?; v = Tensor::cat(&[kv_cache_v, &v], 2)?;
}; };
self.kv_cache = Some((k.clone(), v.clone())); self.kv_cache = Some((k.clone(), v.clone()));
}; };
let k = k.contiguous()?;
let v = v.contiguous()?;
// TODO: Use flash_attn. // TODO: Use flash_attn.
let scores = { let scores = {
let _enter = self.span_mm.enter(); let _enter = self.span_mm.enter();