Mimi streaming fixes.

candle-examples/examples/mimi/main.rs

@@ -126,7 +126,10 @@ fn main() -> Result<()> {
                     for chunk_start in (0..seq_len).step_by(chunk_size) {
                         let chunk_len = usize::min(chunk_size, seq_len - chunk_start);
                         let codes = codes.narrow(candle::D::Minus1, chunk_start, chunk_len)?;
-                        pcm_chunks.push(model.decode(&codes)?)
+                        let pcm = model.decode_step(&codes.into())?;
+                        if let Some(pcm) = pcm.as_option() {
+                            pcm_chunks.push(pcm.clone())
+                        }
                     }
                     Tensor::cat(&pcm_chunks, candle::D::Minus1)?
                 }

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

@@ -216,6 +216,16 @@ impl StreamingMultiheadAttention {
             let pre_ws = match mask {
                 None => pre_ws,
                 Some(mask) => {
+                    // This is a bit cumbersome and slightly incorrect: when providing a new slice
+                    // the kv cache will have a slice offset rather than offset + t. In the mimi
+                    // context of an offset of 250, this would not make much difference though.
+                    let mask_len = mask.dim(D::Minus1)?;
+                    let pre_ws_len = pre_ws.dim(D::Minus1)?;
+                    let mask = if pre_ws_len < mask_len {
+                        mask.narrow(D::Minus1, mask_len - pre_ws_len, pre_ws_len)?
+                    } else {
+                        mask.clone()
+                    };
                     let mask = mask.broadcast_left((b, self.num_heads))?;
                     let neg_inf = self.neg_inf.broadcast_as(pre_ws.shape())?;
                     mask.where_cond(&neg_inf, &pre_ws)?