Fix the attn mask generation.

candle-nn/src/kv_cache.rs

@@ -1,4 +1,4 @@
-use candle::{Result, Tensor};
+use candle::{Device, Result, Tensor};
 
 #[derive(Debug, Clone)]
 pub struct Cache {
@@ -255,6 +255,43 @@ impl RotatingCache {
  }
  }
  }
+
+ fn get_mask(&self, size1: usize, size2: usize, device: &Device) -> Result<Tensor> {
+ let context = self.max_seq_len;
+ let upd_offset = (self.offset + size1) % self.max_seq_len;
+ let mask: Vec<_> = (0..size1)
+ .flat_map(|pos_src| {
+ // The absolute position of the elements that will get added to the cache.
+ let pos_src = self.current_seq_len + pos_src;
+ (0..size2).map(move |pos_cache_rel| {
+ // The absolute position of the cache elements after the addition.
+ let pos_cache = self.current_seq_len + size1 + pos_cache_rel - upd_offset;
+ let pos_cache = if pos_cache_rel < upd_offset {
+ pos_cache
+ } else {
+ pos_cache - self.max_seq_len
+ };
+ u8::from(pos_cache > pos_src || pos_cache + context < pos_src)
+ })
+ })
+ .collect();
+ Tensor::from_slice(&mask, (size1, size2), device)
+ }
+
+ /// Returns the attn_mask to be applied *after* adding `seq_len` to the cache.
+ pub fn attn_mask(&self, seq_len: usize, device: &Device) -> Result<Option<Tensor>> {
+ let mask = if seq_len == 1 {
+ None
+ } else {
+ let cache_out_len = if seq_len < self.max_seq_len {
+ (self.current_seq_len + seq_len).min(self.max_seq_len)
+ } else {
+ seq_len
+ };
+ Some(self.get_mask(seq_len, cache_out_len, device)?)
+ };
+ Ok(mask)
+ }
 }
 
 #[derive(Debug, Clone)]
@@ -308,6 +345,10 @@ impl RotatingKvCache {
  self.k.current_seq_len()
  }
 
+ pub fn attn_mask(&self, seq_len: usize, device: &Device) -> Result<Option<Tensor>> {
+ self.k.attn_mask(seq_len, device)
+ }
+
  pub fn reset(&mut self) {
  self.k.reset();
  self.v.reset();

candle-transformers/src/models/mimi/transformer.rs

@@ -101,21 +101,6 @@ impl Module for LayerScale {
  }
 }
 
-pub(crate) fn get_mask(
- size1: usize,
- size2: usize,
- context: usize,
- device: &Device,
-) -> Result<Tensor> {
- let mask: Vec<_> = (0..size1)
- .flat_map(|i| {
- (0..size2)
- .map(move |j| u8::from(size1 + j > size2 + i || size1 + j + context < size2 + i))
- })
- .collect();
- Tensor::from_slice(&mask, (size1, size2), device)
-}
-
 #[derive(Debug, Clone)]
 pub struct StreamingMultiheadAttention {
  q_proj: Linear,
@@ -590,7 +575,6 @@ impl StreamingTransformerLayer {
 #[derive(Debug, Clone)]
 pub struct StreamingTransformer {
  layers: Vec<StreamingTransformerLayer>,
- context: usize,
  positional_embedding: PositionalEmbedding,
  max_period: usize,
 }
@@ -617,7 +601,6 @@ impl StreamingTransformer {
  }
  Ok(Self {
  layers,
- context: cfg.context,
  positional_embedding: cfg.positional_embedding,
  max_period: cfg.max_period,
  })
@@ -629,23 +612,11 @@ impl StreamingTransformer {
 
  pub fn forward_ca(&mut self, xs: &Tensor, ca_src: Option<&Tensor>) -> Result<Tensor> {
  let (_b, t, c) = xs.dims3()?;
- let pos = self.layers[0]
+ let pos = self.layers[0].self_attn.kv_cache.current_seq_len();
+ let mask = self.layers[0]
  .self_attn
  .kv_cache
- .k_cache()
- .current_seq_len();
- let mask = if t == 1 {
- None
- } else {
- let cache_out_len = if t < self.context {
- (pos + t).min(self.context)
- } else {
- t
- };
- // TODO: this is wrong, the mask depends on the kv-cache offset because of its rotating
- // nature.
- Some(get_mask(t, cache_out_len, self.context, xs.device())?)
- };
+ .attn_mask(t, xs.device())?;
  let mut xs = match self.positional_embedding {
  PositionalEmbedding::Rope | PositionalEmbedding::None => xs.clone(),
  PositionalEmbedding::Sin => {