rtx-csm: Phase 8.9 — Moonshine KV cache + profile binary

KV cache for the decoder turns greedy generation from O(T^2) into O(T)
total work. Per-token decode drops modestly on short transcripts
(7.1 -> 6.0 ms/token at 49 tokens) and compounds on longer ones.

New components in src/moonshine.rs:

  RotaryCache::apply_at(x, position, t)
      Apply RoPE for a window starting at `position`. Replaces
      `apply()` for cached step (which always called positions 0..T).

  DecoderSelfAttention::forward_step(xs, cache_k, cache_v, rope, position)
      Single-token cached self-attn. Appends new K/V to per-layer cache,
      attends across full accumulated history. No causal mask needed
      (cache only contains positions <= current).

  CrossAttention::precompute_kv(enc) -> (K, V)
      One-shot encoder K/V projection for cross-attn. Reused every step.

  CrossAttention::forward_step(xs, k, v)
      Cached cross-attn. Q computed from new token; K/V from precompute.

  DecoderCache { self_k: Vec<Option<Tensor>>, self_v, cross_k, cross_v, position }

  Decoder::precompute_cross_kv(enc) -> DecoderCache
  Decoder::step(token_id, &mut cache) -> logits (1, vocab)
  Decoder::generate_cached(enc, cfg, max_tokens) -> Vec<u32>
      Greedy loop using the cached step.

Profile (5 steady-state runs on /tmp/asr_test.flac, 10.42 s LibriSpeech):

  warm-up:                344 ms
  steady-state mean: 307 ms (p50 305, range 298-319)
  realtime factor: 0.0294x

Comparison across all STT in rtx-csm:

  Backend           RTF         Notes
  Kyutai STT 1B     1.01x       hardware-bound, 3 GB
  Whisper-tiny      0.020x      breaks CSM (in-process ggml conflict)
  Moonshine-tiny    0.0294x     pure candle, NO runtime conflict

Moonshine is the only fast STT path that integrates cleanly. ~34x
faster than realtime, ~17x faster than Kyutai 1B, no protobuf or
ggml linkage issues.

New `examples/moonshine_profile` mirrors `stt_profile` and
`whisper_profile` so all three STT backends report comparable numbers.

Phase 8.10 (next): wire as a third AsrEngine variant in converse_server
for English-only deploys. Replace the energy-VAD-gated Kyutai path
when --moonshine flag is set.

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

crates/models/rtx-csm/Cargo.toml

@@ -257,3 +257,7 @@ path = "examples/moonshine_smoke.rs"
 [[example]]
 name = "moonshine_transcribe"
 path = "examples/moonshine_transcribe.rs"
+
+[[example]]
+name = "moonshine_profile"
+path = "examples/moonshine_profile.rs"

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

@@ -0,0 +1,114 @@
+//! Phase 8.9 — standalone profile of Moonshine STT, comparable to
+//! `examples/stt_profile` (Kyutai) and `examples/whisper_profile`
+//! (Whisper-tiny via whisper-rs). Runs N transcriptions of the same
+//! audio, reports per-call latency stats + realtime factor.
+//!
+//! Usage:
+//! ```bash
+//! cargo run -p rtx-csm --release --features metal --example moonshine_profile -- \
+//!     --in /tmp/asr_test.flac --repeat 5
+//! ```
+
+use anyhow::{Context, Result};
+use candle_core::{Device, Tensor};
+use clap::Parser;
+use hf_hub::api::sync::Api;
+use rtx_csm::{audio_io, moonshine};
+use std::path::PathBuf;
+use std::time::Instant;
+
+#[derive(Debug, Parser)]
+struct Cli {
+    #[arg(long = "in", default_value = "/tmp/asr_test.flac")]
+    input: PathBuf,
+    #[arg(long, default_value_t = 5)]
+    repeat: usize,
+    #[arg(long, default_value_t = 100)]
+    max_tokens: usize,
+    /// Force CPU device.
+    #[arg(long)]
+    cpu: bool,
+}
+
+fn main() -> Result<()> {
+    let cli = Cli::parse();
+    let device = if cli.cpu {
+        Device::Cpu
+    } else if candle_core::utils::metal_is_available() {
+        Device::new_metal(0)?
+    } else {
+        Device::Cpu
+    };
+    eprintln!("device: {device:?}");
+
+    let api = Api::new()?;
+    let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
+    let weights = repo.get("model.safetensors")?;
+    let tok_path = repo.get("tokenizer.json")?;
+
+    let pcm = audio_io::load_mono_at_rate(&cli.input, 16_000).context("load audio")?;
+    let audio_secs = pcm.len() as f32 / 16_000.0;
+    eprintln!("audio: {} samples ({audio_secs:.2}s @ 16 kHz)", pcm.len());
+
+    let cfg = moonshine::MoonshineConfig::tiny();
+    let load_t = Instant::now();
+    let (encoder, decoder) = moonshine::load_full(&weights, &device, &cfg)?;
+    let tok = moonshine::load_tokenizer(&tok_path)
+        .map_err(|e| anyhow::anyhow!("tok: {e}"))?;
+    eprintln!("load: {:.2}s", load_t.elapsed().as_secs_f32());
+
+    let pcm_tensor = Tensor::from_vec(pcm.clone(), (1, 1, pcm.len()), &device)?;
+
+    // Warm-up: first call pays JIT + cache init.
+    let warm_t = Instant::now();
+    let enc = encoder.forward(&pcm_tensor)?;
+    let _warm_tokens = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
+    eprintln!("warm-up: {} ms", warm_t.elapsed().as_millis());
+
+    // Steady-state runs.
+    let mut per_call_ms: Vec<f64> = Vec::with_capacity(cli.repeat);
+    let mut last_text = String::new();
+    let mut last_token_count = 0usize;
+    for i in 0..cli.repeat {
+        let t = Instant::now();
+        let enc = encoder.forward(&pcm_tensor)?;
+        let tokens = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
+        let ms = t.elapsed().as_secs_f64() * 1000.0;
+        per_call_ms.push(ms);
+        last_text = tok.decode(&tokens, true).map_err(|e| anyhow::anyhow!("detok: {e}"))?;
+        last_token_count = tokens.len();
+        eprintln!("  run {}: {ms:.0} ms ({} tokens)", i + 1, tokens.len());
+    }
+
+    per_call_ms.sort_by(|a, b| a.partial_cmp(b).unwrap());
+    let n = per_call_ms.len() as f64;
+    let mean = per_call_ms.iter().sum::<f64>() / n;
+    let p50 = per_call_ms[per_call_ms.len() / 2];
+    let realtime_factor = (mean / 1000.0) / audio_secs as f64;
+
+    println!();
+    println!("=== Moonshine-tiny profile ===");
+    println!("input: {} ({audio_secs:.2} s of audio)", cli.input.display());
+    println!("steady-state runs: {}", per_call_ms.len());
+    println!();
+    println!("per-call latency:");
+    println!("  mean = {mean:.0} ms");
+    println!("  p50  = {p50:.0} ms");
+    println!("  min  = {:.0} ms", per_call_ms[0]);
+    println!("  max  = {:.0} ms", per_call_ms[per_call_ms.len() - 1]);
+    println!();
+    println!("realtime factor: {realtime_factor:.4}x");
+    println!("  (mean / audio_duration; sub-1.0 = faster than realtime)");
+    println!();
+    println!("transcript ({last_token_count} tokens):");
+    println!("  {last_text}");
+
+    // For direct comparison with the other STT profiles in this crate:
+    println!();
+    println!("=== A/B against other STT backends in rtx-csm ===");
+    println!("  Kyutai STT 1B   ~1.01x realtime  (3 GB, hardware-bound)");
+    println!("  Whisper-tiny    ~0.020x          (in-process ggml, breaks CSM)");
+    println!("  Moonshine-tiny  {realtime_factor:.4}x          (pure candle, no conflict)");
+
+    Ok(())
+}

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

@@ -65,12 +65,12 @@ fn main() -> Result<()> {
         enc.shape()
     );
 
-    // Greedy decode.
+    // Greedy decode (cached path — O(T) per token instead of O(T²)).
     let dec_t = Instant::now();
-    let token_ids = decoder.generate(&enc, &cfg, cli.max_tokens)?;
+    let token_ids = decoder.generate_cached(&enc, &cfg, cli.max_tokens)?;
     let dec_ms = dec_t.elapsed().as_millis();
     eprintln!(
-        "decode:    {dec_ms} ms ({} tokens, {:.1} ms/token)",
+        "decode:    {dec_ms} ms ({} tokens, {:.1} ms/token, KV-cached)",
         token_ids.len(),
         dec_ms as f64 / token_ids.len().max(1) as f64
     );

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

@@ -208,15 +208,21 @@ impl RotaryCache {
 
     /// Apply partial RoPE to `q` of shape `(B, H, T, head_dim)`.
     /// Rotates the first `rotary_dim` channels; leaves the rest as-is.
+    /// Uses positions `0..T` (prefill mode).
     fn apply(&self, x: &Tensor) -> Result<Tensor> {
-        let (_b, _h, t, head_dim) = x.dims4()?;
+        let (_b, _h, t, _) = x.dims4()?;
-        let cos = self.cos.narrow(0, 0, t)?;
+        self.apply_at(x, 0, t)
-        let sin = self.sin.narrow(0, 0, t)?;
+    }
+
+    /// Apply RoPE for a window starting at `position`, length `t`.
+    /// Used by the cached single-token step (`t=1`, `position=cache_len`).
+    fn apply_at(&self, x: &Tensor, position: usize, t: usize) -> Result<Tensor> {
+        let (_b, _h, _t, head_dim) = x.dims4()?;
+        let cos = self.cos.narrow(0, position, t)?;
+        let sin = self.sin.narrow(0, position, t)?;
         if head_dim == self.rotary_dim {
-            // Full rotary — straightforward path.
             return candle_nn::rotary_emb::rope_i(x, &cos, &sin);
         }
-        // Partial rotary: split into rotary head and pass-through.
         let rot = x.narrow(3, 0, self.rotary_dim)?.contiguous()?;
         let pass = x.narrow(3, self.rotary_dim, head_dim - self.rotary_dim)?;
         let rot = candle_nn::rotary_emb::rope_i(&rot, &cos, &sin)?;
@@ -739,6 +745,217 @@ pub fn load_tokenizer(path: &std::path::Path) -> std::result::Result<tokenizers:
     tokenizers::Tokenizer::from_file(path)
 }
 
+// ---------------------------------------------------------------------
+// Phase 8.9 — KV cache for single-token autoregressive decoding
+// ---------------------------------------------------------------------
+
+/// Per-layer K/V cache for one decode session.
+/// - `self_k[i]` / `self_v[i]`: cumulative self-attn K/V for layer i,
+///   shape `(B*H, position, head_dim)`. Grows by 1 per `step()`.
+/// - `cross_k[i]` / `cross_v[i]`: encoder cross-attn K/V for layer i,
+///   computed once via `precompute_cross_kv` and reused every step.
+pub struct DecoderCache {
+    self_k: Vec<Option<Tensor>>,
+    self_v: Vec<Option<Tensor>>,
+    cross_k: Vec<Tensor>,
+    cross_v: Vec<Tensor>,
+    pub position: usize,
+}
+
+impl DecoderSelfAttention {
+    /// Single-token cached self-attention. Appends new K/V to the
+    /// cache and attends across the full accumulated history.
+    /// `xs`: `(B=1, T=1, H)`. Returns `(B=1, T=1, H)`.
+    fn forward_step(
+        &self,
+        xs: &Tensor,
+        cache_k: &mut Option<Tensor>,
+        cache_v: &mut Option<Tensor>,
+        rope: &RotaryCache,
+        position: usize,
+    ) -> Result<Tensor> {
+        let (b, t, _h) = xs.dims3()?;
+        debug_assert_eq!(t, 1, "self-attn step expects T=1");
+        let q = self.q_proj.forward(xs)?;
+        let k = self.k_proj.forward(xs)?;
+        let v = self.v_proj.forward(xs)?;
+        // (B, T, H) -> (B, H, T, head_dim)
+        let q = q
+            .reshape((b, t, self.n_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?;
+        let k = k
+            .reshape((b, t, self.n_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?;
+        let v = v
+            .reshape((b, t, self.n_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?;
+        // RoPE for q and the new k at this position.
+        let q = rope.apply_at(&q, position, 1)?;
+        let k_new = rope.apply_at(&k, position, 1)?;
+        // Collapse to (B*H, T, head_dim) for the matmul.
+        let bh = b * self.n_heads;
+        let q3 = q.reshape((bh, 1, self.head_dim))?;
+        let k_new3 = k_new.reshape((bh, 1, self.head_dim))?;
+        let v3 = v.reshape((bh, 1, self.head_dim))?;
+        // Append to cache (or initialize on first step).
+        let k_full = match cache_k.take() {
+            Some(prev) => Tensor::cat(&[&prev, &k_new3], 1)?,
+            None => k_new3,
+        };
+        let v_full = match cache_v.take() {
+            Some(prev) => Tensor::cat(&[&prev, &v3], 1)?,
+            None => v3,
+        };
+        // No causal mask needed: K/V only contains positions <= current.
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let scores = (q3.matmul(&k_full.transpose(1, 2)?.contiguous()?)? * scale)?;
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let out = probs.matmul(&v_full)?;
+        // Write cache back.
+        *cache_k = Some(k_full);
+        *cache_v = Some(v_full);
+        // (B*H, 1, head_dim) -> (B, 1, H)
+        let out = out
+            .reshape((b, self.n_heads, 1, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((b, 1, self.n_heads * self.head_dim))?;
+        self.o_proj.forward(&out)
+    }
+}
+
+impl CrossAttention {
+    /// Pre-compute and cache encoder K/V projections. Returns
+    /// `(K, V)` each shaped `(B*H, T_enc, head_dim)`.
+    fn precompute_kv(&self, enc: &Tensor) -> Result<(Tensor, Tensor)> {
+        let (b, t_enc, _) = enc.dims3()?;
+        let k = self.k_proj.forward(enc)?;
+        let v = self.v_proj.forward(enc)?;
+        let bh = b * self.n_heads;
+        let k = k
+            .reshape((b, t_enc, self.n_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((bh, t_enc, self.head_dim))?;
+        let v = v
+            .reshape((b, t_enc, self.n_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((bh, t_enc, self.head_dim))?;
+        Ok((k, v))
+    }
+
+    /// Cached cross-attention step. K/V come from `precompute_kv`.
+    /// `xs`: `(B=1, T=1, H)`. Returns `(B=1, T=1, H)`.
+    fn forward_step(&self, xs: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
+        let (b, t_dec, _) = xs.dims3()?;
+        debug_assert_eq!(t_dec, 1, "cross-attn step expects T=1");
+        let q = self.q_proj.forward(xs)?;
+        let bh = b * self.n_heads;
+        let q = q
+            .reshape((b, t_dec, self.n_heads, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((bh, 1, self.head_dim))?;
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let scores = (q.matmul(&k.transpose(1, 2)?.contiguous()?)? * scale)?;
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let out = probs.matmul(v)?;
+        let out = out
+            .reshape((b, self.n_heads, 1, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((b, 1, self.n_heads * self.head_dim))?;
+        self.o_proj.forward(&out)
+    }
+}
+
+impl Decoder {
+    /// Pre-compute the encoder cross-attn K/V for every layer. Returns
+    /// an empty `DecoderCache` ready for `step()` calls.
+    pub fn precompute_cross_kv(&self, encoder_output: &Tensor) -> Result<DecoderCache> {
+        let n = self.layers.len();
+        let mut cross_k = Vec::with_capacity(n);
+        let mut cross_v = Vec::with_capacity(n);
+        for layer in &self.layers {
+            let (k, v) = layer.cross_attn.precompute_kv(encoder_output)?;
+            cross_k.push(k);
+            cross_v.push(v);
+        }
+        Ok(DecoderCache {
+            self_k: (0..n).map(|_| None).collect(),
+            self_v: (0..n).map(|_| None).collect(),
+            cross_k,
+            cross_v,
+            position: 0,
+        })
+    }
+
+    /// Single-token cached forward. `token_id` is the *new* token to
+    /// process. Returns logits for that one position, shape
+    /// `(1, vocab)`. Updates `cache.position` and per-layer `self_k`/
+    /// `self_v` in place.
+    pub fn step(&self, token_id: u32, cache: &mut DecoderCache) -> Result<Tensor> {
+        let device = self.embed_weight.device();
+        let token = Tensor::from_vec(vec![token_id], (1, 1), device)?;
+        let mut h = self.embed.forward(&token)?;
+        let position = cache.position;
+        for (i, layer) in self.layers.iter().enumerate() {
+            let normed = layer.input_ln.forward(&h)?;
+            let attn_out = layer.self_attn.forward_step(
+                &normed,
+                &mut cache.self_k[i],
+                &mut cache.self_v[i],
+                &self.rope,
+                position,
+            )?;
+            h = (h + attn_out)?;
+            let normed = layer.post_attn_ln.forward(&h)?;
+            let cross_out = layer.cross_attn.forward_step(&normed, &cache.cross_k[i], &cache.cross_v[i])?;
+            h = (h + cross_out)?;
+            let normed = layer.final_ln.forward(&h)?;
+            let mlp_out = layer.mlp.forward(&normed)?;
+            h = (h + mlp_out)?;
+        }
+        h = self.final_ln.forward(&h)?;
+        cache.position += 1;
+        // Logits: (1, 1, H) -> (1, H) -> (1, vocab) via tied LM head.
+        let lm_w = self.embed_weight.transpose(0, 1)?.contiguous()?;
+        let logits_3d = h.broadcast_matmul(&lm_w)?;
+        logits_3d.squeeze(1)
+    }
+
+    /// Cached greedy decode. Equivalent semantics to
+    /// [`Self::generate`] but with O(T) total work instead of O(T²).
+    pub fn generate_cached(
+        &self,
+        encoder_output: &Tensor,
+        cfg: &MoonshineConfig,
+        max_tokens: usize,
+    ) -> Result<Vec<u32>> {
+        let mut cache = self.precompute_cross_kv(encoder_output)?;
+        let mut next = cfg.decoder_start_token_id;
+        let mut out = Vec::with_capacity(max_tokens);
+        for _ in 0..max_tokens {
+            let logits = self.step(next, &mut cache)?; // (1, vocab)
+            let argmax = logits.argmax(1)?;
+            let id: u32 = argmax.to_dtype(DType::U32)?.to_vec1::<u32>()?[0];
+            if id == cfg.eos_token_id {
+                break;
+            }
+            out.push(id);
+            next = id;
+            if cache.position >= cfg.max_position_embeddings {
+                break;
+            }
+        }
+        Ok(out)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;