Use cat for faster MQA computation.

huggingface · Apr 11, 2024 · c40098d · c40098d
1 parent a0460cd
commit c40098d
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/candle-transformers/src/models/quantized_llama.rs b/candle-transformers/src/models/quantized_llama.rs
@@ -205,7 +205,7 @@ impl LayerWeights {
  };
  self.kv_cache = Some((k.clone(), v.clone()));
 
- // Support for MQA, useful for 70B models.
+ // Support for MQA, useful for 70B models and mistral.
  let k = self.repeat_kv(k)?;
  let v = self.repeat_kv(v)?;
 
@@ -231,11 +231,7 @@ impl LayerWeights {
  Ok(x)
  } else {
  let (b_sz, n_kv_head, seq_len, head_dim) = x.dims4()?;
- let x = x
- .unsqueeze(2)?
- .expand((b_sz, n_kv_head, n_rep, seq_len, head_dim))?
- .reshape((b_sz, n_kv_head * n_rep, seq_len, head_dim))?;
- Ok(x)
+ Tensor::cat(&vec![&x; n_rep], 2)?.reshape((b_sz, n_kv_head * n_rep, seq_len, head_dim))
  }
  }
 }