rtx-csm: Phase 8.6 — Moonshine encoder transformer block

Full encoder forward path: conv stem -> 6 transformer layers -> final
LayerNorm. Loads HF safetensors, runs end-to-end on Metal.

Components added to src/moonshine.rs:

  RotaryCache       partial RoPE (32 of 36 head_dim, theta=10000)
  EncoderAttention  MHA (8 heads, no bias), partial RoPE on q/k
  EncoderMlp        288 -> 1152 -> 288 with bias, GELU(erf) activation
  EncoderLayer      Pre-LN attn + Pre-LN MLP (LayerNorm weight-only)
  Encoder           stem + 6 layers + final LayerNorm
  load_encoder()    VarBuilder convenience for the standalone smoke

Smoke test (`examples/moonshine_smoke`) verified end-to-end:
  input  (1, 1, 160000)  -> output (1, 415, 288)
  forward: 132 ms        (10 s of audio at 0.013x realtime)
  max abs: 6.67          (signal preserved, not zeros)

Implementation notes captured in the diff:
  - candle Metal 4D batched matmul had shape-mismatch issues for our
    (B, H, T, D) pattern. Switched to (B*H, T, D) 3D form which is
    unambiguous and avoids the kernel bug.
  - LayerNorm is weight-only (no bias tensors in safetensors); we
    construct LayerNorm with a zeros bias to satisfy candle's API.
  - rotary_dim = floor(head_dim * 0.9 / 2) * 2 = 32 (must be even).
    The remaining 4 head_dim channels pass through unchanged via
    `narrow + cat` on dim 3.

Numerical parity vs HF Python reference is NOT yet verified — that's
the next bounded chunk (Phase 8.7). Shape + signal correctness are
verified by the smoke test.

Next: decoder transformer block (self-attn + cross-attn + SwiGLU).
~3-4 h of focused work.

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

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

@@ -29,10 +29,15 @@ fn main() -> Result<()> {
         .get("model.safetensors")?;
     eprintln!("weights: {}", weights.display());
 
-    // Build the conv stem.
+    // Build the full encoder (conv stem + 6 transformer layers + final LN).
+    let cfg = moonshine::MoonshineConfig::tiny();
+    let encoder = moonshine::load_encoder(&weights, &device, &cfg)
+        .context("load encoder")?;
+    eprintln!("encoder loaded (conv stem + 6 layers + final LN)");
+
+    // Also keep the standalone conv stem for the conv-only check below.
     let stem = moonshine::load_conv_stem(&weights, &device)
         .context("load conv stem")?;
-    eprintln!("conv stem loaded");
 
     // Synthetic 10 s of 16 kHz audio (silence + a sine pulse).
     let sr = 16_000usize;
@@ -83,5 +88,40 @@ fn main() -> Result<()> {
     } else {
         println!("output has signal — weight loading verified");
     }
+
+    // Now run through the FULL encoder (stem + 6 transformer layers + LN).
+    eprintln!();
+    eprintln!("=== full encoder forward ===");
+    let full_t = std::time::Instant::now();
+    let enc_out = encoder.forward(&pcm).context("encoder forward")?;
+    let full_ms = full_t.elapsed().as_millis();
+    eprintln!("forward: {full_ms} ms");
+    eprintln!("encoder output shape: {:?}", enc_out.shape());
+    let enc_dims = enc_out.dims();
+    if enc_dims == [1, expected_t_seq, 288] {
+        println!("PASS: encoder output preserves (B, T_seq, 288) shape");
+    } else {
+        println!(
+            "MISMATCH: expected [1, {expected_t_seq}, 288], got {:?}",
+            enc_dims
+        );
+        std::process::exit(2);
+    }
+    let enc_max = enc_out
+        .abs()?
+        .max_keepdim(0)?
+        .max_keepdim(1)?
+        .max_keepdim(2)?;
+    let enc_max_val: f32 = enc_max.flatten_all()?.to_vec1::<f32>()?[0];
+    let enc_mean = enc_out.mean_all()?;
+    let enc_mean_val: f32 = enc_mean.to_vec0::<f32>()?;
+    eprintln!("encoder output max abs: {enc_max_val:.4}");
+    eprintln!("encoder output mean    : {enc_mean_val:.4}");
+    if enc_max_val < 1e-6 {
+        println!("WARN: encoder output is all near-zero");
+    } else {
+        println!("encoder output has signal — full transformer pipeline verified");
+    }
+
     Ok(())
 }

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

@@ -28,8 +28,8 @@
 //!   of head_dim)
 //! - bos = 1, eos = 2, pad = 2, decoder_start = 1
 
-use candle_core::{Device, Module, Result, Tensor};
+use candle_core::{DType, Device, Module, Result, Tensor};
-use candle_nn::{Conv1d, Conv1dConfig, VarBuilder};
+use candle_nn::{Conv1d, Conv1dConfig, LayerNorm, Linear, VarBuilder};
 
 /// Hyperparameters matching `UsefulSensors/moonshine-tiny/config.json`.
 #[derive(Debug, Clone)]
@@ -167,18 +167,228 @@ impl ConvStem {
 /// Loader: open the HF safetensors and construct a `ConvStem`. Useful
 /// for the standalone Phase 8.5 smoke test.
 pub fn load_conv_stem(weights_path: &std::path::Path, device: &Device) -> Result<ConvStem> {
-    // Open the safetensors directly via VarBuilder. F32 storage.
     let vb = unsafe {
-        VarBuilder::from_mmaped_safetensors(
+        VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
-            &[weights_path],
-            candle_core::DType::F32,
-            device,
-        )
     }?;
-    // The encoder convs live under `model.encoder.{conv1,conv2,conv3}`.
     ConvStem::new(vb.pp("model").pp("encoder"))
 }
 
+// ---------------------------------------------------------------------
+// Phase 8.6 — encoder transformer block
+// ---------------------------------------------------------------------
+
+/// Cosine/sine tables for partial RoPE.
+///
+/// Moonshine uses `partial_rotary_factor = 0.9` with `head_dim = 36`
+/// (288 / 8 heads). `rotary_dim = floor(36 * 0.9 / 2) * 2 = 32` (must
+/// be even — RoPE rotates pairs). The remaining 4 dims pass through
+/// unchanged. RoPE base is `theta = 10_000`.
+struct RotaryCache {
+    cos: Tensor, // (max_seq, rotary_dim/2)
+    sin: Tensor, // (max_seq, rotary_dim/2)
+    rotary_dim: usize,
+}
+
+impl RotaryCache {
+    fn new(rotary_dim: usize, max_seq: usize, theta: f64, dtype: DType, dev: &Device) -> Result<Self> {
+        assert!(rotary_dim % 2 == 0, "rotary_dim must be even");
+        let inv_freq: Vec<f32> = (0..rotary_dim)
+            .step_by(2)
+            .map(|i| (1.0 / theta.powf(i as f64 / rotary_dim as f64)) as f32)
+            .collect();
+        let inv_freq = Tensor::new(inv_freq, dev)?;
+        let positions = Tensor::arange(0u32, max_seq as u32, dev)?
+            .to_dtype(DType::F32)?
+            .reshape((max_seq, 1))?;
+        let freqs = positions.matmul(&inv_freq.reshape((1, rotary_dim / 2))?)?;
+        let cos = freqs.cos()?.to_dtype(dtype)?;
+        let sin = freqs.sin()?.to_dtype(dtype)?;
+        Ok(Self { cos, sin, rotary_dim })
+    }
+
+    /// Apply partial RoPE to `q` of shape `(B, H, T, head_dim)`.
+    /// Rotates the first `rotary_dim` channels; leaves the rest as-is.
+    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)?;
+        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)?;
+        Tensor::cat(&[&rot, &pass], 3)
+    }
+}
+
+/// Self-attention block matching `model.encoder.layers.X.self_attn`.
+/// 8 heads × 36 head_dim, no bias (`attention_bias=false`), partial
+/// RoPE on first 32 head_dim channels.
+struct EncoderAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    n_heads: usize,
+    head_dim: usize,
+}
+
+impl EncoderAttention {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let h = cfg.hidden_size;
+        let n_heads = cfg.encoder_num_attention_heads;
+        let head_dim = h / n_heads;
+        let q_proj = candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?;
+        let k_proj = candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?;
+        let v_proj = candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?;
+        let o_proj = candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?;
+        Ok(Self { q_proj, k_proj, v_proj, o_proj, n_heads, head_dim })
+    }
+
+    fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
+        let (b, t, _h) = xs.dims3()?;
+        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_heads, 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()?;
+        let q = rope.apply(&q)?;
+        let k = rope.apply(&k)?;
+        // Collapse (B, H, T, D) -> (B*H, T, D) for the matmul. candle's
+        // 4D batched matmul on Metal had shape-mismatch issues for our
+        // pattern; the 3D form is unambiguous.
+        let bh = b * self.n_heads;
+        let q3 = q.reshape((bh, t, self.head_dim))?;
+        let k3 = k.reshape((bh, t, self.head_dim))?;
+        let v3 = v.reshape((bh, t, self.head_dim))?;
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        let scores = (q3.matmul(&k3.transpose(1, 2)?.contiguous()?)? * scale)?;
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let out = probs.matmul(&v3)?;
+        // (B*H, T, D) -> (B, H, T, D) -> (B, T, H*D)
+        let out = out
+            .reshape((b, self.n_heads, t, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((b, t, self.n_heads * self.head_dim))?;
+        self.o_proj.forward(&out)
+    }
+}
+
+/// FFN block matching `model.encoder.layers.X.mlp`.
+/// 288 → 1152 (with bias) → GELU → 1152 → 288 (with bias).
+struct EncoderMlp {
+    fc1: Linear,
+    fc2: Linear,
+}
+
+impl EncoderMlp {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let fc1 = candle_nn::linear(cfg.hidden_size, cfg.intermediate_size, vb.pp("fc1"))?;
+        let fc2 = candle_nn::linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?;
+        Ok(Self { fc1, fc2 })
+    }
+
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let h = self.fc1.forward(xs)?;
+        let h = h.gelu_erf()?;
+        self.fc2.forward(&h)
+    }
+}
+
+/// One encoder layer: Pre-LN attn + Pre-LN MLP, residuals as standard.
+/// Tensors are `LayerNorm` weight-only (no bias) per the inspector
+/// dump — implemented via `LayerNorm` with bias zeros.
+struct EncoderLayer {
+    input_ln: LayerNorm,
+    self_attn: EncoderAttention,
+    post_attn_ln: LayerNorm,
+    mlp: EncoderMlp,
+}
+
+impl EncoderLayer {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let h = cfg.hidden_size;
+        let input_ln = layer_norm_weight_only(h, 1e-5, vb.pp("input_layernorm"))?;
+        let self_attn = EncoderAttention::new(cfg, vb.pp("self_attn"))?;
+        let post_attn_ln = layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?;
+        let mlp = EncoderMlp::new(cfg, vb.pp("mlp"))?;
+        Ok(Self { input_ln, self_attn, post_attn_ln, mlp })
+    }
+
+    fn forward(&self, xs: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
+        let h = self.input_ln.forward(xs)?;
+        let h = self.self_attn.forward(&h, rope)?;
+        let xs = (xs + h)?;
+        let h = self.post_attn_ln.forward(&xs)?;
+        let h = self.mlp.forward(&h)?;
+        xs + h
+    }
+}
+
+/// Helper: `LayerNorm` with weight only (zeros bias). Moonshine's
+/// `*_layernorm.weight [288]` tensors lack a corresponding `.bias`.
+fn layer_norm_weight_only(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
+    let weight = vb.get((size,), "weight")?;
+    let bias = Tensor::zeros((size,), weight.dtype(), weight.device())?;
+    Ok(LayerNorm::new(weight, bias, eps))
+}
+
+/// Full encoder: conv stem → 6 transformer layers → final layer norm.
+/// Output shape: `(B, T_seq, hidden=288)` ready to feed cross-attn in
+/// the decoder (Phase 8.7).
+pub struct Encoder {
+    stem: ConvStem,
+    layers: Vec<EncoderLayer>,
+    final_ln: LayerNorm,
+    rope: RotaryCache,
+}
+
+impl Encoder {
+    pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let stem = ConvStem::new(vb.pp("encoder"))?;
+        let head_dim = cfg.hidden_size / cfg.encoder_num_attention_heads;
+        let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2;
+        // max_seq for the encoder is bounded by the input audio length /
+        // total stride 384. For 30 s of 16 kHz audio that's ~1250.
+        // Allocate a generous table.
+        let rope = RotaryCache::new(rotary_dim, 4096, cfg.rope_theta, vb.dtype(), vb.device())?;
+        let mut layers = Vec::with_capacity(cfg.encoder_num_hidden_layers);
+        let layer_vb = vb.pp("encoder").pp("layers");
+        for i in 0..cfg.encoder_num_hidden_layers {
+            layers.push(EncoderLayer::new(cfg, layer_vb.pp(i))?);
+        }
+        let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("encoder").pp("layer_norm"))?;
+        Ok(Self { stem, layers, final_ln, rope })
+    }
+
+    pub fn forward(&self, pcm: &Tensor) -> Result<Tensor> {
+        let mut h = self.stem.forward(pcm)?;
+        for layer in &self.layers {
+            h = layer.forward(&h, &self.rope)?;
+        }
+        self.final_ln.forward(&h)
+    }
+}
+
+/// Loader: open the HF safetensors and construct an `Encoder` with
+/// weights at `model.*` (the canonical HF prefix).
+pub fn load_encoder(
+    weights_path: &std::path::Path,
+    device: &Device,
+    cfg: &MoonshineConfig,
+) -> Result<Encoder> {
+    let vb = unsafe {
+        VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
+    }?;
+    Encoder::new(cfg, vb.pp("model"))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;