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]>
This commit is contained in:
osobh
@@ -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::*;
|
||||
|
||||
Reference in New Issue
Block a user