Fix for gemma MQA.

huggingface · Apr 19, 2024 · adb1a23 · adb1a23
1 parent 9c532ae
commit adb1a23
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/candle-transformers/src/models/gemma.rs b/candle-transformers/src/models/gemma.rs
@@ -227,8 +227,9 @@ impl Attention {
  };
  self.kv_cache = Some((key_states.clone(), value_states.clone()));
 
- let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?;
- let value_states = crate::utils::repeat_kv(value_states, self.num_kv_groups)?;
+ let key_states = crate::utils::repeat_kv(key_states, self.num_kv_groups)?.contiguous()?;
+ let value_states =
+ crate::utils::repeat_kv(value_states, self.num_kv_groups)?.contiguous()?;
 
  let attn_output = {
  let scale = 1f64 / f64::sqrt(self.head_dim as f64);