rtx-csm: Phase 8.7 — Moonshine decoder transformer (encoder+decoder)

Full encoder-decoder Moonshine v2 working end-to-end on candle 0.9 +
Metal. Loads HF safetensors, runs through every transformer block, and
produces real logits.

Components added to src/moonshine.rs:

  CrossAttention      MHA with K/V from encoder output (no causal mask)
  DecoderSelfAttention  MHA with causal mask, partial RoPE on q/k
  DecoderMlp          SwiGLU: fused fc1 [2304, 288] split gate+up,
                      silu(gate) * up, fc2 [288, 1152] back to hidden
  DecoderLayer        Pre-LN self-attn + Pre-LN cross-attn + Pre-LN MLP
  Decoder             token embed -> 6 layers -> final LN -> tied LM head
  load_full()         convenience: returns (Encoder, Decoder)

Smoke test verifies end-to-end:
  encoder forward   :   1 ms   (cached after warm-up)
  decoder forward   :  85 ms   (1 token, prefill mode)
  logits shape      :  (1, 1, 32768)
  logit max abs     :  30.66   (real signal, not zeros)
  argmax token_id   :  379     (non-trivial prediction; eos=2)

Implementation notes:
  - Same (B*H, T, D) 3D matmul pattern as encoder to dodge candle's 4D
    Metal matmul shape-mismatch bug.
  - LM head tied to decoder.embed_tokens.weight (cached on Decoder for
    fast forward; logits = hidden @ embed.T).
  - Causal mask is a (T, T) -inf upper-triangular added to scores
    before softmax.
  - Decoder final LN tensor is `decoder.norm.weight` (NOT
    `decoder.layer_norm.weight` — encoder uses the latter naming).
  - No KV cache yet: this is prefill mode. Phase 8.8 will add the
    streaming-generation loop with cache + tokenizer.

NOT yet verified: numerical parity vs HF Python reference. The token
predicted (id=379) looks plausible for silent-mostly audio, but a
parity check is still needed (Phase 8.9). Architecture appears
correct based on shape + signal sanity.

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

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

@@ -10,7 +10,7 @@
 //! ```
 
 use anyhow::{Context, Result};
-use candle_core::{Device, Tensor};
+use candle_core::{DType, Device, Tensor};
 use hf_hub::api::sync::Api;
 use rtx_csm::moonshine;
 
@@ -29,11 +29,11 @@ fn main() -> Result<()> {
         .get("model.safetensors")?;
     eprintln!("weights: {}", weights.display());
 
-    // Build the full encoder (conv stem + 6 transformer layers + final LN).
+    // Build the full encoder + decoder (Phase 8.5/8.6/8.7).
     let cfg = moonshine::MoonshineConfig::tiny();
-    let encoder = moonshine::load_encoder(&weights, &device, &cfg)
+    let (encoder, decoder) = moonshine::load_full(&weights, &device, &cfg)
-        .context("load encoder")?;
+        .context("load full")?;
-    eprintln!("encoder loaded (conv stem + 6 layers + final LN)");
+    eprintln!("encoder + decoder loaded");
 
     // Also keep the standalone conv stem for the conv-only check below.
     let stem = moonshine::load_conv_stem(&weights, &device)
@@ -123,5 +123,43 @@ fn main() -> Result<()> {
         println!("encoder output has signal — full transformer pipeline verified");
     }
 
+    // Phase 8.7: decoder forward step. Feed [bos] + encoder output,
+    // expect logits over the 32 768 vocab. We don't care about the
+    // exact predicted token yet (parity check is Phase 8.9) — just
+    // shape + signal.
+    eprintln!();
+    eprintln!("=== decoder forward step ===");
+    let bos = cfg.bos_token_id;
+    let tokens =
+        Tensor::from_vec(vec![bos], (1, 1), &device).context("bos tokens")?;
+    let dec_t = std::time::Instant::now();
+    let logits = decoder.forward(&tokens, &enc_out).context("decoder")?;
+    let dec_ms = dec_t.elapsed().as_millis();
+    eprintln!("forward: {dec_ms} ms");
+    eprintln!("decoder logits shape: {:?}", logits.shape());
+    let logit_dims = logits.dims();
+    if logit_dims == [1, 1, cfg.vocab_size] {
+        println!("PASS: decoder produces (B=1, T=1, vocab=32768) logits");
+    } else {
+        println!(
+            "MISMATCH: expected [1, 1, {}], got {:?}",
+            cfg.vocab_size, logit_dims
+        );
+        std::process::exit(2);
+    }
+    let logit_max = logits
+        .abs()?
+        .max_keepdim(0)?
+        .max_keepdim(1)?
+        .max_keepdim(2)?;
+    let logit_max_val: f32 = logit_max.flatten_all()?.to_vec1::<f32>()?[0];
+    eprintln!("logit max abs: {logit_max_val:.4}");
+    // Greedy argmax over the vocab dim — what's the model's first
+    // prediction given silence-mostly audio?
+    let last_step = logits.narrow(1, 0, 1)?.squeeze(1)?; // (1, vocab)
+    let argmax = last_step.argmax(1)?;
+    let token_id: u32 = argmax.to_dtype(DType::U32)?.to_vec1::<u32>()?[0];
+    eprintln!("argmax token_id: {token_id} (eos={})", cfg.eos_token_id);
+
     Ok(())
 }

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

@@ -389,6 +389,303 @@ pub fn load_encoder(
     Encoder::new(cfg, vb.pp("model"))
 }
 
+// ---------------------------------------------------------------------
+// Phase 8.7 — decoder transformer block
+// ---------------------------------------------------------------------
+
+/// Cross-attention block matching `model.decoder.layers.X.encoder_attn`.
+/// Queries come from the decoder hidden state; keys and values come
+/// from the encoder output (computed once per generation step, cached
+/// across generation tokens — caching deferred to Phase 8.8).
+struct CrossAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    n_heads: usize,
+    head_dim: usize,
+}
+
+impl CrossAttention {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let h = cfg.hidden_size;
+        let n_heads = cfg.decoder_num_attention_heads;
+        let head_dim = h / n_heads;
+        Ok(Self {
+            q_proj: candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?,
+            k_proj: candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?,
+            v_proj: candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?,
+            o_proj: candle_nn::linear_no_bias(h, h, vb.pp("o_proj"))?,
+            n_heads,
+            head_dim,
+        })
+    }
+
+    /// `xs`: decoder hidden `(B, T_dec, H)`.
+    /// `enc`: encoder output `(B, T_enc, H)`.
+    /// Returns: `(B, T_dec, H)`.
+    fn forward(&self, xs: &Tensor, enc: &Tensor) -> Result<Tensor> {
+        let (b, t_dec, _h) = xs.dims3()?;
+        let (_, t_enc, _) = enc.dims3()?;
+        let q = self.q_proj.forward(xs)?;
+        let k = self.k_proj.forward(enc)?;
+        let v = self.v_proj.forward(enc)?;
+        // (B, T, H) -> (B*H_heads, T, head_dim) — same 3D matmul
+        // pattern as EncoderAttention to dodge the candle 4D Metal bug.
+        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, t_dec, self.head_dim))?;
+        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))?;
+        let scale = 1.0 / (self.head_dim as f64).sqrt();
+        // Cross-attention has NO causal mask: each decoder token can
+        // attend to every encoder position.
+        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, t_dec, self.head_dim))?
+            .transpose(1, 2)?
+            .contiguous()?
+            .reshape((b, t_dec, self.n_heads * self.head_dim))?;
+        self.o_proj.forward(&out)
+    }
+}
+
+/// Decoder self-attention. Same shape as encoder attention but applies
+/// a causal mask so token `i` only attends to tokens `0..=i`. Partial
+/// RoPE on q/k as in the encoder.
+struct DecoderSelfAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    n_heads: usize,
+    head_dim: usize,
+}
+
+impl DecoderSelfAttention {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let h = cfg.hidden_size;
+        let n_heads = cfg.decoder_num_attention_heads;
+        let head_dim = h / n_heads;
+        Ok(Self {
+            q_proj: candle_nn::linear_no_bias(h, h, vb.pp("q_proj"))?,
+            k_proj: candle_nn::linear_no_bias(h, h, vb.pp("k_proj"))?,
+            v_proj: candle_nn::linear_no_bias(h, h, vb.pp("v_proj"))?,
+            o_proj: candle_nn::linear_no_bias(h, h, vb.pp("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)?;
+        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)?;
+        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 mut scores = (q3.matmul(&k3.transpose(1, 2)?.contiguous()?)? * scale)?;
+        // Causal mask: build a (T, T) upper-triangular -inf mask.
+        let mask = causal_mask(t, scores.dtype(), scores.device())?;
+        scores = scores.broadcast_add(&mask)?;
+        let probs = candle_nn::ops::softmax_last_dim(&scores)?;
+        let out = probs.matmul(&v3)?;
+        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)
+    }
+}
+
+fn causal_mask(t: usize, dtype: DType, device: &Device) -> Result<Tensor> {
+    // Upper-triangular `-inf` mask: value at (i, j) is 0 if j <= i, else -inf.
+    let mut data = vec![0.0f32; t * t];
+    for i in 0..t {
+        for j in (i + 1)..t {
+            data[i * t + j] = f32::NEG_INFINITY;
+        }
+    }
+    Tensor::from_vec(data, (t, t), device)?.to_dtype(dtype)
+}
+
+/// SwiGLU MLP for the decoder. `fc1` is the fused gate+up projection:
+/// weight shape `[2 * intermediate_size, hidden_size] = [2304, 288]`.
+/// Output is split on dim -1 into `gate` and `up`, both `[..., 1152]`.
+/// Activation: `silu(gate) * up`. Then `fc2: [hidden, intermediate] =
+/// [288, 1152]` projects back to `hidden`.
+struct DecoderMlp {
+    fc1: Linear,
+    fc2: Linear,
+    intermediate: usize,
+}
+
+impl DecoderMlp {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        // fc1 has bias; fc2 has bias (consistent with encoder MLP — the
+        // inspector dump shows `.bias` on both).
+        let fc1 = candle_nn::linear(cfg.hidden_size, cfg.intermediate_size * 2, vb.pp("fc1"))?;
+        let fc2 = candle_nn::linear(cfg.intermediate_size, cfg.hidden_size, vb.pp("fc2"))?;
+        Ok(Self {
+            fc1,
+            fc2,
+            intermediate: cfg.intermediate_size,
+        })
+    }
+
+    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        // (B, T, H) -> (B, T, 2 * intermediate)
+        let h = self.fc1.forward(xs)?;
+        let dims = h.dims();
+        let last = dims.len() - 1;
+        // Split on the last dim into gate and up.
+        let gate = h.narrow(last, 0, self.intermediate)?;
+        let up = h.narrow(last, self.intermediate, self.intermediate)?;
+        // silu(gate) * up
+        let activated = candle_nn::ops::silu(&gate)?.mul(&up)?;
+        self.fc2.forward(&activated)
+    }
+}
+
+/// One decoder layer: Pre-LN self-attn, Pre-LN cross-attn, Pre-LN MLP.
+/// Three layer norms per layer (input_layernorm, post_attention_layernorm,
+/// final_layernorm), all weight-only.
+struct DecoderLayer {
+    input_ln: LayerNorm,
+    self_attn: DecoderSelfAttention,
+    post_attn_ln: LayerNorm,
+    cross_attn: CrossAttention,
+    final_ln: LayerNorm,
+    mlp: DecoderMlp,
+}
+
+impl DecoderLayer {
+    fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let h = cfg.hidden_size;
+        Ok(Self {
+            input_ln: layer_norm_weight_only(h, 1e-5, vb.pp("input_layernorm"))?,
+            self_attn: DecoderSelfAttention::new(cfg, vb.pp("self_attn"))?,
+            post_attn_ln: layer_norm_weight_only(h, 1e-5, vb.pp("post_attention_layernorm"))?,
+            cross_attn: CrossAttention::new(cfg, vb.pp("encoder_attn"))?,
+            final_ln: layer_norm_weight_only(h, 1e-5, vb.pp("final_layernorm"))?,
+            mlp: DecoderMlp::new(cfg, vb.pp("mlp"))?,
+        })
+    }
+
+    fn forward(&self, xs: &Tensor, enc: &Tensor, rope: &RotaryCache) -> Result<Tensor> {
+        // Self-attn (causal)
+        let h = self.input_ln.forward(xs)?;
+        let h = self.self_attn.forward(&h, rope)?;
+        let xs = (xs + h)?;
+        // Cross-attn
+        let h = self.post_attn_ln.forward(&xs)?;
+        let h = self.cross_attn.forward(&h, enc)?;
+        let xs = (xs + h)?;
+        // MLP
+        let h = self.final_ln.forward(&xs)?;
+        let h = self.mlp.forward(&h)?;
+        xs + h
+    }
+}
+
+/// Full decoder: token embedding → 6 layers → final LN → tied LM head.
+pub struct Decoder {
+    embed: candle_nn::Embedding,
+    layers: Vec<DecoderLayer>,
+    final_ln: LayerNorm,
+    rope: RotaryCache,
+    /// Cached transposed embedding for the tied LM head.
+    /// Computing logits: hidden @ embed.weight.T -> (B, T, vocab).
+    embed_weight: Tensor,
+}
+
+impl Decoder {
+    pub fn new(cfg: &MoonshineConfig, vb: VarBuilder) -> Result<Self> {
+        let embed = candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("decoder").pp("embed_tokens"))?;
+        let head_dim = cfg.hidden_size / cfg.decoder_num_attention_heads;
+        let rotary_dim = ((head_dim as f64 * cfg.partial_rotary_factor) as usize / 2) * 2;
+        let rope = RotaryCache::new(
+            rotary_dim,
+            cfg.max_position_embeddings.max(512),
+            cfg.rope_theta,
+            vb.dtype(),
+            vb.device(),
+        )?;
+        let mut layers = Vec::with_capacity(cfg.decoder_num_hidden_layers);
+        let layer_vb = vb.pp("decoder").pp("layers");
+        for i in 0..cfg.decoder_num_hidden_layers {
+            layers.push(DecoderLayer::new(cfg, layer_vb.pp(i))?);
+        }
+        // Decoder uses `decoder.norm.weight` (the encoder uses
+        // `encoder.layer_norm.weight`). Inspector dump confirmed.
+        let final_ln = layer_norm_weight_only(cfg.hidden_size, 1e-5, vb.pp("decoder").pp("norm"))?;
+        let embed_weight = embed.embeddings().clone();
+        Ok(Self { embed, layers, final_ln, rope, embed_weight })
+    }
+
+    /// Forward over `tokens` shape `(B, T)` with encoder output `enc`
+    /// shape `(B, T_enc, H)`. Returns logits `(B, T, vocab_size)`.
+    /// **No KV cache** — this is the prefill / single-step path.
+    /// The full streaming-generation loop with KV cache lands in
+    /// Phase 8.8.
+    pub fn forward(&self, tokens: &Tensor, enc: &Tensor) -> Result<Tensor> {
+        let mut h = self.embed.forward(tokens)?;
+        for layer in &self.layers {
+            h = layer.forward(&h, enc, &self.rope)?;
+        }
+        h = self.final_ln.forward(&h)?;
+        // LM head tied to embedding: logits = hidden @ embed.T.
+        // embed.weight shape: (vocab, hidden). transpose -> (hidden, vocab).
+        let lm_w = self.embed_weight.transpose(0, 1)?.contiguous()?;
+        h.broadcast_matmul(&lm_w)
+    }
+}
+
+/// Loader: open the HF safetensors and construct an encoder + decoder.
+pub fn load_full(
+    weights_path: &std::path::Path,
+    device: &Device,
+    cfg: &MoonshineConfig,
+) -> Result<(Encoder, Decoder)> {
+    let vb = unsafe {
+        VarBuilder::from_mmaped_safetensors(&[weights_path], DType::F32, device)
+    }?;
+    let encoder = Encoder::new(cfg, vb.pp("model"))?;
+    let decoder = Decoder::new(cfg, vb.pp("model"))?;
+    Ok((encoder, decoder))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;