rustytorch/crates/models/rtx-csm/examples/stt_profile.rs

//! Profile Kyutai STT step_pcm latency. Loads the standard 1B en/fr
//! model, feeds PCM frames in fixed chunk sizes, and reports per-frame
//! latency stats so we can see where the ~50%-realtime cost comes from.
//!
//! Usage:
//! ```bash
//! cargo run -p rtx-csm --release --features metal --example stt_profile -- \
//!     --in /tmp/asr_test.flac --frame-batch 1
//! ```
//!
//! `--frame-batch 1` measures pure per-frame cost (one 80ms frame per
//! step_pcm call). `--frame-batch 8` simulates the converse_server's
//! receive pattern (200ms chunks ≈ 2.5 frames; 8 here is generous).

use anyhow::{Context, Result};
use clap::Parser;
use rtx_csm::{audio_io, stt::Stt};
use std::path::PathBuf;
use std::time::Instant;

#[derive(Debug, Parser)]
struct Cli {
    /// Input audio (any rate, mono OK; will be resampled to 24 kHz).
    #[arg(long = "in", default_value = "/tmp/asr_test.flac")]
    input: PathBuf,
    /// Number of 80ms frames (1920 samples each) per step_pcm call.
    /// 1 = pure per-frame cost. The server typically sees ~2-3.
    #[arg(long, default_value_t = 1)]
    frame_batch: usize,
    /// CPU only (skip Metal).
    #[arg(long)]
    cpu: bool,
}

const SAMPLES_PER_FRAME: usize = 1920;

fn main() -> Result<()> {
    tracing_subscriber::fmt().with_max_level(tracing::Level::WARN).init();
    let cli = Cli::parse();

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else if candle_core::utils::metal_is_available() {
        candle_core::Device::new_metal(0)?
    } else {
        candle_core::Device::Cpu
    };
    eprintln!("device: {device:?}");
    eprintln!("loading audio: {}", cli.input.display());
    let pcm = audio_io::load_mono_at_rate(&cli.input, 24_000).context("load audio")?;
    eprintln!("loaded {} samples ({:.2}s @ 24 kHz)", pcm.len(), pcm.len() as f32 / 24_000.0);

    eprintln!("loading Kyutai STT 1B en/fr (~3 GB)...");
    let load_t = Instant::now();
    let mut stt = Stt::load_default(&device).context("load stt")?;
    eprintln!("model loaded in {:.2}s", load_t.elapsed().as_secs_f32());

    let chunk_size = SAMPLES_PER_FRAME * cli.frame_batch.max(1);
    let n_chunks = pcm.len() / chunk_size;
    eprintln!(
        "feeding {} chunks of {} samples ({} frames each, {} ms)",
        n_chunks,
        chunk_size,
        cli.frame_batch,
        cli.frame_batch * 80,
    );

    // Warm up: first call pays Metal kernel compilation + first-frame KV
    // cache init. Don't include in stats.
    let warm_t = Instant::now();
    let _ = stt.step_pcm(&pcm[..chunk_size]).context("warm step_pcm")?;
    eprintln!("warm-up first chunk: {} ms", warm_t.elapsed().as_millis());

    // Steady-state: time each subsequent step_pcm call.
    let mut per_chunk_ms: Vec<f64> = Vec::with_capacity(n_chunks - 1);
    let mut total_events = 0usize;
    let total_t = Instant::now();
    for i in 1..n_chunks {
        let start = i * chunk_size;
        let end = start + chunk_size;
        let t = Instant::now();
        let events = stt
            .step_pcm(&pcm[start..end])
            .with_context(|| format!("step_pcm chunk {i}"))?;
        per_chunk_ms.push(t.elapsed().as_secs_f64() * 1000.0);
        total_events += events.len();
    }
    let total_wall_ms = total_t.elapsed().as_millis() as f64;

    // Stats.
    per_chunk_ms.sort_by(|a, b| a.partial_cmp(b).unwrap());
    let n = per_chunk_ms.len() as f64;
    let mean = per_chunk_ms.iter().sum::<f64>() / n;
    let p50 = per_chunk_ms[per_chunk_ms.len() / 2];
    let p95 = per_chunk_ms[((n * 0.95).min(n - 1.0)) as usize];
    let min = per_chunk_ms[0];
    let max = per_chunk_ms[per_chunk_ms.len() - 1];

    let chunk_audio_ms = (cli.frame_batch * 80) as f64;
    let realtime_factor = mean / chunk_audio_ms;

    println!();
    println!("=== Kyutai STT step_pcm profile ===");
    println!("frame_batch:  {} (= {} ms audio per call)", cli.frame_batch, cli.frame_batch * 80);
    println!("steady-state n: {}", per_chunk_ms.len());
    println!();
    println!("per-call latency:");
    println!("  mean = {mean:.1} ms");
    println!("  p50  = {p50:.1} ms");
    println!("  p95  = {p95:.1} ms");
    println!("  min  = {min:.1} ms");
    println!("  max  = {max:.1} ms");
    println!();
    println!("realtime factor: {realtime_factor:.2}x (mean step_pcm / audio duration)");
    println!("  < 1.0  = faster than real-time (good for streaming)");
    println!("  > 1.0  = slower than real-time (will lag behind incoming audio)");
    println!();
    println!("aggregate:");
    println!("  total wall   = {total_wall_ms:.0} ms");
    println!("  total audio  = {:.0} ms", n * chunk_audio_ms);
    println!("  total events = {total_events}");

    Ok(())
}