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/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 cos = self.cos.narrow(0, 0, t)?;
-        let sin = self.sin.narrow(0, 0, t)?;
+        let (_b, _h, t, _) = x.dims4()?;
+        self.apply_at(x, 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::*;