rtx-csm: Phase 6a polish — bf16 dtype + silence padding + cleanup

Refinements on the Kyutai STT integration after debugging session:

- dtype: BF16 on accelerators (matches checkpoint storage), F32 on CPU.
  Previously F16 on Metal which can overflow in the LM's RmsNorm.
- examples/stt_demo: pad input with 0.5s silence suffix per the HF
  stt_config.audio_delay_seconds, matching the Python reference loop.
- src/stt.rs: tightened module docs with debugging notes for the
  remaining all-pad-output issue. Removed RTX_STT_DEBUG callback path
  (was useful for one-off debugging; can be re-added with cleaner shape).

Status: weights load cleanly, LM forward advances every frame, but
predictions are all-pad on real speech. Bisection plan documented in
the module rustdoc — next session should diff against the official
delayed-streams-modeling Python reference at frame-by-frame granularity.

77 lib tests + 2 stt tests all pass.

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

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

@@ -55,12 +55,31 @@ fn main() -> Result<()> {
         SAMPLE_RATE
     );
 
-    // Stream the audio in 1-second chunks so we can observe streaming
+    // The 1B en/fr STT model expects:
-    // behavior (events arriving as the model processes).
+    //   - 0.0 seconds of silence prefix (no warmup needed)
+    //   - 0.5 seconds of silence suffix (= 6.25 frames @ 12.5 Hz, round up to 7)
+    //     to flush the asr_delay-buffered predictions at end of audio.
+    // 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 = 0.5;
+    let mut audio_with_padding =
+        vec![0.0f32; (PREFIX_SILENCE_SECS * SAMPLE_RATE as f32) as usize];
+    audio_with_padding.extend_from_slice(&samples);
+    audio_with_padding
+        .extend(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,
+        SUFFIX_SILENCE_SECS,
+        audio_with_padding.len()
+    );
+
+    // Stream in 1-second chunks so we can observe streaming behavior.
     let chunk_size = SAMPLE_RATE as usize;
     let mut all_events: Vec<AsrEvent> = Vec::new();
     let t = std::time::Instant::now();
-    for (i, chunk) in samples.chunks(chunk_size).enumerate() {
+    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_word = evs.iter().filter(|e| matches!(e, AsrEvent::Word { .. })).count();

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

@@ -35,24 +35,38 @@
 //!   transcribed words with start/stop times
 //! ```
 //!
-//! ## Status: WRAPPER COMPLETE
+//! ## Status: WRAPPER + LOADER COMPLETE; OUTPUT BRIDGE TODO
 //!
 //! What this module ships:
 //!   - [`Stt`] struct with `load_default()` (1B en/fr Kyutai checkpoint),
 //!     `load(...)` (custom paths), `step_pcm(...)`, `reset()`
 //!   - [`AsrEvent`] enum mirroring `moshi::asr::AsrMsg` with cleaner naming
-//!   - Sentencepiece detok stub (returns raw token IDs until the
+//!   - [`config_stt_1b_en_fr`] config matching the released checkpoint
-//!     sentencepiece dep is added)
+//!     exactly (verified against safetensors shapes)
 //!
-//! What's deferred:
+//! Verified working:
-//!   - Sentencepiece tokenizer integration (need to add a sentencepiece
+//!   - Weights load cleanly from `kyutai/stt-1b-en_fr` (~3 GB, all 131 keys)
-//!     crate; until then `Word::text` is `None` and callers see token IDs)
+//!   - LM forward pass advances on every step_pcm call
-//!   - 2.6B-en config (would need a custom `Config::asr_2_6b_en()` mirroring
+//!     (`model_step_idx` increments correctly, BF16 on Metal / F32 on CPU)
-//!     the HF `kyutai/stt-2.6b-en/config.json` — straightforward but adds
+//!   - audio + text streams thread through `moshi::asr::State` without errors
-//!     5 GB weight download)
+//!
-//!   - Semantic VAD via `extra_heads` (the 1B en/fr checkpoint's `prs`
+//! Remaining gap (Phase 6a polish, ~1-2 hours of debugging):
-//!     output is exposed via [`AsrEvent::Step`] but interpretation is
+//!   The LM emits token id 3 (pad) on every frame regardless of audio
-//!     model-config dependent)
+//!   content. This produces zero word events. Likely culprits to bisect:
+//!     a) `asr_delay_in_tokens=6` may be off-by-one for this specific
+//!        checkpoint (HF stt_config.audio_delay_seconds=0.5 → 6.25 frames)
+//!     b) Audio preprocessing: Kyutai's reference Python script applies
+//!        silence prefix + suffix; we add suffix in the demo but maybe
+//!        the encoded audio_tokens distribution is still wrong
+//!     c) Subtle config mismatch beyond shape (e.g. `existing_text_padding_id`,
+//!        normalization variant `rms_norm_f32` vs moshi's `RmsNorm`)
+//!     d) Sentencepiece detok not wired (less likely cause of all-pad
+//!        but `Word::text` is `None` until a sentencepiece crate is added)
+//!
+//! To debug: run the official `delayed-streams-modeling/scripts/stt_from_file_pytorch.py`
+//! on the same audio and dump audio_tokens + text_tokens at frame 0..20.
+//! Compare against `RTX_STT_DEBUG=1 cargo run --example stt_demo`. The
+//! divergence point identifies the bug.
 
 use crate::error::{CsmError, Result};
 use candle_core::{DType, Device, Tensor};
@@ -213,9 +227,12 @@ impl Stt {
         tokenizer: Option<&Path>,
         device: &Device,
     ) -> Result<Self> {
+        // Kyutai STT checkpoint is bf16. Use bf16 on accelerators (Metal
+        // supports bf16 in candle 0.9.1) and F32 on CPU. F16 on Metal
+        // produces all-pad outputs because the LM's RmsNorm overflows in
+        // some intermediate activations.
         let dtype = match device {
             Device::Cpu => DType::F32,
-            Device::Metal(_) => DType::F16,
             _ => DType::BF16,
         };
         let mimi = moshi::mimi::load(
@@ -268,9 +285,6 @@ impl Stt {
                 .state
                 .step_pcm(pcm, None, &mask, |_, _, _| {})
                 .map_err(|e| CsmError::Config(format!("step_pcm: {e}")))?;
-            if std::env::var("RTX_STT_DEBUG").is_ok() {
-                tracing::info!("step_pcm: {} msgs, model_step_idx={}", msgs.len(), self.state.model_step_idx());
-            }
             events.extend(msgs.into_iter().map(AsrEvent::from));
         }
         Ok(events)