rtx-csm: Stt::finish() — drain asr_delay buffer at end-of-stream

Phase 8.1.1 quality fix from the perf plan. Tokens emitted at LM step
`t` correspond to audio frame `t - ASR_DELAY_FRAMES` (6 frames /
0.48 s), so when a caller stops feeding audio without trailing
silence the last few words trail off — they're still inside the
delay pipeline.

finish() now steps ASR_DELAY_FRAMES additional silent frames after
handling any partial sub-frame buffer, giving the LM the chance to
emit those buffered tokens. Cost: 7 extra step_pcm calls per turn.

Verified end-to-end via stt_demo on a mid-utterance trim of the
LibriSpeech reference clip:
  pre-flush:  11 words ("...turnips and carrots and bruised")
  post-flush: 13 words ("...turnips and carrots and bruised potatoes and")

Also drops the now-redundant 2s silence suffix in stt_demo — the
flush replaces it. Affects converse_server's real-time end-of-turn
path where suffix padding wasn't possible.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/examples/stt_demo.rs

@@ -14,7 +14,10 @@
 
 use anyhow::Result;
 use clap::Parser;
-use rtx_csm::{audio_io, stt::{AsrEvent, Stt, SAMPLE_RATE}};
+use rtx_csm::{
+    audio_io,
+    stt::{AsrEvent, SAMPLE_RATE, Stt},
+};
 use std::path::PathBuf;
 
 #[derive(Debug, Parser)]
@@ -43,7 +46,10 @@ fn main() -> Result<()> {
 
     let t = std::time::Instant::now();
     let mut stt = Stt::load_default(&device)?;
-    println!("loaded Kyutai STT 1B en/fr in {:.2}s", t.elapsed().as_secs_f32());
+    println!(
+        "loaded Kyutai STT 1B en/fr in {:.2}s",
+        t.elapsed().as_secs_f32()
+    );
 
     // Load WAV at 24 kHz mono (Mimi's expected input rate).
     let samples = audio_io::load_mono_at_rate(&cli.input, SAMPLE_RATE)?;
@@ -55,19 +61,17 @@ fn main() -> Result<()> {
         SAMPLE_RATE
     );
 
-    // The 1B en/fr STT model expects:
+    // The 1B en/fr STT model needs a few frames of silence after real audio
-    //   - 0.0 seconds of silence prefix (no warmup needed)
+    // to drain the asr_delay buffer (tokens at LM step `t` correspond to
-    //   - 0.5 seconds of silence suffix (= 6.25 frames @ 12.5 Hz, round up to 7)
+    // audio frame `t - 6`). `Stt::finish()` does that automatically — feed
-    //     to flush the asr_delay-buffered predictions at end of audio.
+    // raw audio and then call finish().
-    // Without the suffix the model produces only pad tokens. See HF
-    // config.json `stt_config` and the reference Python script.
     const PREFIX_SILENCE_SECS: f32 = 0.0;
-    const SUFFIX_SILENCE_SECS: f32 = 2.0;
+    const SUFFIX_SILENCE_SECS: f32 = 0.0;
-    let mut audio_with_padding =
+    let mut audio_with_padding = vec![0.0f32; (PREFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize];
-        vec![0.0f32; (PREFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize];
     audio_with_padding.extend_from_slice(&samples);
-    audio_with_padding
+    audio_with_padding.extend(
-        .extend(std::iter::repeat(0.0f32).take((SUFFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize));
+        std::iter::repeat(0.0f32).take((SUFFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize),
+    );
     println!(
         "padded with {:.1}s prefix + {:.1}s suffix silence -> {} samples",
         PREFIX_SILENCE_SECS,
@@ -81,9 +85,18 @@ fn main() -> Result<()> {
     let t = std::time::Instant::now();
     for (i, chunk) in audio_with_padding.chunks(chunk_size).enumerate() {
         let evs = stt.step_pcm(chunk)?;
-        let n_step = evs.iter().filter(|e| matches!(e, AsrEvent::Step { .. })).count();
+        let n_step = evs
-        let n_word = evs.iter().filter(|e| matches!(e, AsrEvent::Word { .. })).count();
+            .iter()
-        let n_end = evs.iter().filter(|e| matches!(e, AsrEvent::EndWord { .. })).count();
+            .filter(|e| matches!(e, AsrEvent::Step { .. }))
+            .count();
+        let n_word = evs
+            .iter()
+            .filter(|e| matches!(e, AsrEvent::Word { .. }))
+            .count();
+        let n_end = evs
+            .iter()
+            .filter(|e| matches!(e, AsrEvent::EndWord { .. }))
+            .count();
         println!("[chunk {i}] events: step={n_step} word={n_word} endword={n_end}");
         all_events.extend(evs);
     }
@@ -98,17 +111,13 @@ fn main() -> Result<()> {
     for ev in &all_events {
         match ev {
             AsrEvent::Word {
-                tokens,
+                tokens, start_time, ..
-                start_time,
-                ..
             } => {
                 pending = Some((tokens.clone(), *start_time));
             }
             AsrEvent::EndWord { stop_time, .. } => {
                 if let Some((tokens, start)) = pending.take() {
-                    let text = stt
+                    let text = stt.decode_word_text(&tokens).unwrap_or_default();
-                        .decode_word_text(&tokens)
-                        .unwrap_or_default();
                     println!("  ({:.2}s - {:.2}s) {}", start, stop_time, text);
                     if !full_text.is_empty() && !text.is_empty() {
                         full_text.push(' ');

crates/models/rtx-csm/src/stt.rs

@@ -62,10 +62,10 @@
 
 use crate::error::{CsmError, Result};
 use candle_core::{DType, Device, Tensor};
+use moshi::StreamMask;
 use moshi::asr::AsrMsg;
 use moshi::lm;
 use moshi::transformer;
-use moshi::StreamMask;
 use sentencepiece::SentencePieceProcessor;
 use std::path::{Path, PathBuf};
 
@@ -168,10 +168,7 @@ pub const ASR_DELAY_FRAMES: usize = 6;
 pub enum AsrEvent {
     /// Per-step probabilities from the `extra_heads` output (semantic VAD,
     /// turn-taking, etc.). Each inner `Vec<f32>` is one head's output.
-    Step {
+    Step { step_idx: usize, prs: Vec<Vec<f32>> },
-        step_idx: usize,
-        prs: Vec<Vec<f32>>,
-    },
     /// A complete word (sequence of subword tokens) with timing.
     /// `text` is `None` until sentencepiece detok is wired (Phase 6a polish).
     Word {
@@ -281,11 +278,7 @@ impl Stt {
             Device::Cpu => DType::F32,
             _ => DType::BF16,
         };
-        let mimi = moshi::mimi::load(
+        let mimi = moshi::mimi::load(mimi.to_string_lossy().as_ref(), Some(32), device)
-            mimi.to_string_lossy().as_ref(),
-            Some(32),
-            device,
-        )
             .map_err(|e| CsmError::Config(format!("moshi::mimi::load: {e}")))?;
         let cfg = if vad {
             config_stt_1b_en_fr_vad()
@@ -378,24 +371,42 @@ impl Stt {
         Ok(events)
     }
 
-
+    /// End-of-stream flush. Drains any partial sub-frame buffer (zero-padded
-    /// Drain the current pending buffer (zero-padded to one full frame)
+    /// to a full 1920-sample frame), then steps `ASR_DELAY_FRAMES` additional
-    /// and run a final step. Useful at end-of-stream to flush any audio
+    /// silent frames so words emitted at LM step `t` for audio at frame
-    /// that's shorter than one frame.
+    /// `t - ASR_DELAY_FRAMES` finally surface — without this the last ~480 ms
+    /// of speech (a couple of words on long utterances) trails off after the
+    /// caller stops feeding audio.
+    ///
+    /// Cost: `ASR_DELAY_FRAMES + 1` extra `step_pcm` calls (≈ 7 × per-frame
+    /// compute on Metal). Quality fix only; not a latency optimization.
     pub fn finish(&mut self) -> Result<Vec<AsrEvent>> {
-        if self.pending.is_empty() {
+        let mut events = Vec::new();
-            return Ok(Vec::new());
+        let mask = StreamMask::empty();
-        }
+
+        if !self.pending.is_empty() {
             self.pending.resize(SAMPLES_PER_FRAME, 0.0);
             let frame = std::mem::take(&mut self.pending);
             let pcm = Tensor::from_vec(frame, (1, 1, SAMPLES_PER_FRAME), &self.device)
                 .map_err(|e| CsmError::Config(format!("finish tensor: {e}")))?;
-        let mask = StreamMask::empty();
             let msgs = self
                 .state
                 .step_pcm(pcm, None, &mask, |_, _, _| {})
                 .map_err(|e| CsmError::Config(format!("finish step: {e}")))?;
-        Ok(msgs.into_iter().map(AsrEvent::from).collect())
+            events.extend(msgs.into_iter().map(AsrEvent::from));
+        }
+
+        let zeros = vec![0.0f32; SAMPLES_PER_FRAME];
+        for _ in 0..ASR_DELAY_FRAMES {
+            let pcm = Tensor::from_vec(zeros.clone(), (1, 1, SAMPLES_PER_FRAME), &self.device)
+                .map_err(|e| CsmError::Config(format!("flush tensor: {e}")))?;
+            let msgs = self
+                .state
+                .step_pcm(pcm, None, &mask, |_, _, _| {})
+                .map_err(|e| CsmError::Config(format!("flush step: {e}")))?;
+            events.extend(msgs.into_iter().map(AsrEvent::from));
+        }
+        Ok(events)
     }
 }