rtx-csm: Phase 13.8 — emotion2vec port, slice 2b (ProjectFeatures + Block)

Two more candle modules toward the full port.

ProjectFeatures — LayerNorm(512) → Linear(512→768). Sits between
LocalEncoder and the transformer. Pickle layout matches upstream:
project_features.1.* is the LayerNorm, project_features.2.* is the
Linear. Both have learnable affine params; the .1 LN is NOT just an
eps constant.

Block — pre-norm fused-QKV transformer block, the workhorse for both
ContextEncoder (4 instances) and MainEncoder (8 instances). Pickle
keys per block: norm1, attn.qkv (fused 768→2304), attn.proj,
norm2, mlp.fc1, mlp.fc2. GELU MLP activation. No positional encoding
inside the block — the conv-based positional bias lives at the encoder
boundary.

Attention uses the (B*H, T, D) 3D collapse-before-matmul Metal
workaround we shipped for Phase 8.8 Moonshine — candle's 4D batched
matmul still has the shape-mismatch bug.

2 new unit tests:
  - project_features_shape_check: (1, 50, 512) → (1, 50, 768)
  - block_residual_shape_check: random (2, 8, 768) → same shape AND
    all values finite (catches softmax NaN / attention overflow)

Lib suite 114/114 (was 112, +2 new).

Remaining within slice 2: relative_positional_encoder (5 Conv1d),
ContextEncoder (4 Blocks), MainEncoder (8 Blocks + LN), Classifier
(Linear 768→9), top-level Emotion2Vec + pickle .pt loader.

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

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

@@ -28,7 +28,9 @@
 
 use crate::error::Result as CsmResult;
 use candle_core::{Module, Tensor};
-use candle_nn::{conv1d_no_bias, layer_norm, Conv1d, Conv1dConfig, LayerNorm, VarBuilder};
+use candle_nn::{
+    conv1d_no_bias, layer_norm, linear, Conv1d, Conv1dConfig, LayerNorm, Linear, VarBuilder,
+};
 
 /// Top-level architectural hyperparameters from `config.yaml`. Only the
 /// fields needed by the inference path are surfaced; training-time fields
@@ -186,6 +188,163 @@ impl Module for LocalEncoder {
     }
 }
 
+/// Linear `feature_dim → embed_dim` (512 → 768 for `_base`) — projects the
+/// LocalEncoder's per-frame output into the transformer's working
+/// dimension. State-dict keys: `project_features.2.weight/bias` (path
+/// includes a LayerNorm at index 1 that we surface separately).
+///
+/// Per the inspector output the upstream stores:
+///   `project_features.1.weight/bias`  → LayerNorm(512)
+///   `project_features.2.weight/bias`  → Linear(512 → 768)
+///
+/// So the projection module is `LayerNorm → Linear`, both at the feature
+/// dim (channel-last after the LocalEncoder transposes back to (B, T, C)
+/// for the transformer).
+#[derive(Debug)]
+pub struct ProjectFeatures {
+    norm: LayerNorm,
+    proj: Linear,
+}
+
+impl ProjectFeatures {
+    /// Build from a VarBuilder rooted at
+    /// `d2v_model.modality_encoders.AUDIO.project_features`.
+    pub fn new(
+        in_dim: usize,
+        out_dim: usize,
+        norm_eps: f64,
+        vb: VarBuilder,
+    ) -> CsmResult<Self> {
+        let norm = layer_norm(in_dim, norm_eps, vb.pp("1"))
+            .map_err(|e| crate::CsmError::Config(format!("project_features norm: {e}")))?;
+        let proj = linear(in_dim, out_dim, vb.pp("2"))
+            .map_err(|e| crate::CsmError::Config(format!("project_features linear: {e}")))?;
+        Ok(Self { norm, proj })
+    }
+}
+
+impl Module for ProjectFeatures {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        // xs: (B, T, in_dim) — LocalEncoder output already transposed
+        // to channel-last. Both ops are channel-last.
+        let xs = self.norm.forward(xs)?;
+        self.proj.forward(&xs)
+    }
+}
+
+/// One transformer block — pre-norm, fused QKV attention, MLP residual.
+///
+/// State-dict layout (per `blocks.{i}` and `context_encoder.blocks.{i}`):
+///   `norm1.weight/bias`           LayerNorm(embed_dim)
+///   `attn.qkv.weight/bias`        Linear(embed_dim → 3 × embed_dim)
+///   `attn.proj.weight/bias`       Linear(embed_dim → embed_dim)
+///   `norm2.weight/bias`           LayerNorm(embed_dim)
+///   `mlp.fc1.weight/bias`         Linear(embed_dim → mlp_dim)
+///   `mlp.fc2.weight/bias`         Linear(mlp_dim → embed_dim)
+///
+/// MLP activation is GELU. No positional encoding is added inside the
+/// block — the upstream applies a conv-based positional bias at the
+/// encoder boundary (slice 2c).
+#[derive(Debug)]
+pub struct Block {
+    norm1: LayerNorm,
+    qkv: Linear,
+    attn_proj: Linear,
+    norm2: LayerNorm,
+    fc1: Linear,
+    fc2: Linear,
+    num_heads: usize,
+    head_dim: usize,
+    scale: f64,
+}
+
+impl Block {
+    /// Build from a VarBuilder rooted at e.g. `d2v_model.blocks.{i}` or
+    /// `d2v_model.modality_encoders.AUDIO.context_encoder.blocks.{i}`.
+    pub fn new(
+        embed_dim: usize,
+        num_heads: usize,
+        mlp_dim: usize,
+        norm_eps: f64,
+        vb: VarBuilder,
+    ) -> CsmResult<Self> {
+        let head_dim = embed_dim / num_heads;
+        let scale = 1.0 / (head_dim as f64).sqrt();
+        let norm1 = layer_norm(embed_dim, norm_eps, vb.pp("norm1"))
+            .map_err(|e| crate::CsmError::Config(format!("block norm1: {e}")))?;
+        let qkv = linear(embed_dim, embed_dim * 3, vb.pp("attn").pp("qkv"))
+            .map_err(|e| crate::CsmError::Config(format!("block qkv: {e}")))?;
+        let attn_proj = linear(embed_dim, embed_dim, vb.pp("attn").pp("proj"))
+            .map_err(|e| crate::CsmError::Config(format!("block attn proj: {e}")))?;
+        let norm2 = layer_norm(embed_dim, norm_eps, vb.pp("norm2"))
+            .map_err(|e| crate::CsmError::Config(format!("block norm2: {e}")))?;
+        let fc1 = linear(embed_dim, mlp_dim, vb.pp("mlp").pp("fc1"))
+            .map_err(|e| crate::CsmError::Config(format!("block fc1: {e}")))?;
+        let fc2 = linear(mlp_dim, embed_dim, vb.pp("mlp").pp("fc2"))
+            .map_err(|e| crate::CsmError::Config(format!("block fc2: {e}")))?;
+        Ok(Self {
+            norm1,
+            qkv,
+            attn_proj,
+            norm2,
+            fc1,
+            fc2,
+            num_heads,
+            head_dim,
+            scale,
+        })
+    }
+
+    fn attention(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        // xs: (B, T, D)
+        let (b, t, d) = xs.dims3()?;
+        let qkv = self.qkv.forward(xs)?; // (B, T, 3D)
+        // Reshape to (B, T, 3, H, head_dim), permute to (3, B, H, T, hd)
+        let qkv = qkv
+            .reshape((b, t, 3, self.num_heads, self.head_dim))?
+            .permute((2, 0, 3, 1, 4))?
+            .contiguous()?;
+        // Split into Q, K, V along the leading 3-axis.
+        let q = qkv.narrow(0, 0, 1)?.squeeze(0)?; // (B, H, T, hd)
+        let k = qkv.narrow(0, 1, 1)?.squeeze(0)?;
+        let v = qkv.narrow(0, 2, 1)?.squeeze(0)?;
+        // Attention scores. For Metal compatibility (the 4D matmul issue
+        // we hit on Moonshine, Phase 8.8), collapse (B, H) to a single
+        // batch dim before matmul.
+        let bh = b * self.num_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 scores = (q3.matmul(&k3.transpose(1, 2)?)? * self.scale)?;
+        let attn = candle_nn::ops::softmax_last_dim(&scores)?;
+        let out = attn.matmul(&v3)?; // (B*H, T, hd)
+        let out = out
+            .reshape((b, self.num_heads, t, self.head_dim))?
+            .transpose(1, 2)?
+            .reshape((b, t, d))?
+            .contiguous()?;
+        self.attn_proj.forward(&out)
+    }
+
+    fn mlp(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let xs = self.fc1.forward(xs)?;
+        let xs = xs.gelu()?;
+        self.fc2.forward(&xs)
+    }
+}
+
+impl Module for Block {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        // Pre-norm: x = x + attn(norm1(x)); x = x + mlp(norm2(x))
+        let h = self.norm1.forward(xs)?;
+        let h = self.attention(&h)?;
+        let xs = (xs + h)?;
+        let h = self.norm2.forward(&xs)?;
+        let h = self.mlp(&h)?;
+        xs + h
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -199,6 +358,46 @@ mod tests {
         assert_eq!(cfg.conv_layers.len(), 7);
     }
 
+    #[test]
+    fn project_features_shape_check() {
+        let dev = Device::Cpu;
+        let cfg = Emotion2VecConfig::plus_base();
+        let vm = VarMap::new();
+        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
+        let pf = ProjectFeatures::new(cfg.feature_dim, cfg.embed_dim, cfg.norm_eps, vb)
+            .expect("build");
+        // (B=1, T=50, in=512) → expect (1, 50, 768)
+        let xs = Tensor::zeros((1, 50, cfg.feature_dim), DType::F32, &dev).unwrap();
+        let ys = <ProjectFeatures as Module>::forward(&pf, &xs).expect("forward");
+        assert_eq!(ys.dims(), &[1, 50, cfg.embed_dim]);
+    }
+
+    #[test]
+    fn block_residual_shape_check() {
+        let dev = Device::Cpu;
+        let cfg = Emotion2VecConfig::plus_base();
+        let mlp_dim = (cfg.embed_dim as f32 * cfg.mlp_ratio) as usize;
+        let vm = VarMap::new();
+        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
+        let block = Block::new(
+            cfg.embed_dim,
+            cfg.num_heads,
+            mlp_dim,
+            cfg.norm_eps,
+            vb,
+        )
+        .expect("build");
+        // 8-frame input of zeros — residual should pass through, output
+        // shape unchanged, and the LayerNorm guarantees no NaNs even when
+        // input is all zeros.
+        let xs = Tensor::randn(0f32, 1.0, (2, 8, cfg.embed_dim), &dev).unwrap();
+        let ys = <Block as Module>::forward(&block, &xs).expect("forward");
+        assert_eq!(ys.dims(), &[2, 8, cfg.embed_dim]);
+        // Sanity: output is finite (not NaN/Inf).
+        let v = ys.flatten_all().unwrap().to_vec1::<f32>().unwrap();
+        assert!(v.iter().all(|x| x.is_finite()), "block produced non-finite output");
+    }
+
     #[test]
     fn local_encoder_random_init_shape_check() {
         // Build with a fresh VarMap → random Kaiming init weights, then