rtx-csm: Sprint 2 Phase A — activation steering API

Adds the apply hook for ActAdd-style activation steering on the Llama
backbone. Inspired by EmoSteer-TTS (arXiv 2508.03543), but adapted: the
paper is flow-matching-specific (DiT layers, 32 CFM steps, per-token
attribution search via mel synthesis), none of which apply to CSM's
autoregressive Llama-over-Mimi-tokens. What's portable is the
underlying difference-in-means construction with residual-stream
addition — the standard ActAdd / contrastive-steering pattern.

What lands:
- src/steering.rs: LayerSteering type, per-layer (1, embed_dim) tensors,
  global scale, safetensors load with keys `layer_<i>_steering`. Three
  unit tests covering empty/no-op, dimension validation, and apply math.
- src/csm_fork.rs LlamaModel: optional `steering: Option<LayerSteering>`
  field, applied after every layer's forward inside the for-loop. Adds
  ~3 LOC to the hot path; gated by the Option so unsteered generation
  has zero cost beyond a None check.
- src/csm_fork.rs Model::set_backbone_steering: installs steering only
  on the conditional backbone (cfg_backbone is intentionally left
  un-steered so CFG correctly subtracts an unsteered baseline).
- src/generator.rs Generator::set_steering: errors on quantized
  backend (only FP supported for now).
- examples/generate.rs: --steering-vec / --steering-scale flags.
- examples/steering_random.rs: smoke helper that writes random Gaussian
  vectors so the apply path can be exercised end-to-end before the
  real corpus extractor lands. Box-Muller via seeded rand to avoid an
  extra rand_distr dep.

Smoke test (16-layer random Gaussian, stddev=0.05, scale=0.5):
- baseline (no steering, same seed/text): 3.04 s @ RMS -19.5 dB
- steered (random vectors):              1.84 s @ RMS -16.2 dB,
                                          EOT triggered earlier
Output clearly differs — pathway is wired correctly. Random vectors
aren't musically meaningful; that's Phase B.

Phase B (next session): corpus extractor that runs forward passes over
emotion-labeled audio (we already have audio_to_manifest emitting
emotion_tag rows), captures per-layer post-residual activations, and
computes the difference-in-means between emotion_X and neutral pools.
Then A/B with quality_eval.

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

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

@@ -14,8 +14,8 @@
 //! `huggingface/candle/candle-transformers/src/models/csm.rs`.
 //! Apache-2.0 license preserved per upstream.
 
-use candle_core::{DType, Device, IndexOp, Module, Result, Tensor, D};
-use candle_nn::{embedding, linear_b, Embedding, Linear, RmsNorm, VarBuilder};
+use candle_core::{D, DType, Device, IndexOp, Module, Result, Tensor};
+use candle_nn::{Embedding, Linear, RmsNorm, VarBuilder, embedding, linear_b};
 use candle_transformers::generation::LogitsProcessor;
 use std::sync::Arc;
 
@@ -286,9 +286,10 @@ impl Attention {
             let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
             attn_weights.matmul(&value_states)?
         };
-        let pre_o = attn_output
-            .transpose(1, 2)?
-            .reshape((b_sz, q_len, self.num_heads * self.head_dim))?;
+        let pre_o =
+            attn_output
+                .transpose(1, 2)?
+                .reshape((b_sz, q_len, self.num_heads * self.head_dim))?;
         let out = self.o_proj.forward(&pre_o)?;
         let out = match &self.o_lora {
             Some(l) => (out + l.forward(&pre_o)?)?,
@@ -417,6 +418,10 @@ pub struct LlamaModel {
     norm: RmsNorm,
     pub(crate) device: Device,
     pub(crate) dtype: DType,
+    /// Optional per-layer activation-steering vectors. When set, the
+    /// `apply()` hook runs after each Layer's forward to add a steering
+    /// vector to the residual stream. See `crate::steering`.
+    steering: Option<crate::steering::LayerSteering>,
 }
 
 impl LlamaModel {
@@ -434,9 +439,21 @@ impl LlamaModel {
             norm,
             device: vb.device().clone(),
             dtype: vb.dtype(),
+            steering: None,
         })
     }
 
+    /// Install activation-steering vectors. Called once before generation
+    /// starts; subsequent `forward` calls apply the steering at every
+    /// step. Pass `None` to remove steering.
+    pub fn set_steering(&mut self, steering: Option<crate::steering::LayerSteering>) {
+        self.steering = steering;
+    }
+
+    pub fn steering(&self) -> Option<&crate::steering::LayerSteering> {
+        self.steering.as_ref()
+    }
+
     pub fn clear_kv_cache(&mut self) {
         for layer in self.layers.iter_mut() {
             layer.clear_kv_cache()
@@ -471,8 +488,11 @@ impl LlamaModel {
             Some(mask)
         };
         let mut xs = xs.clone();
-        for layer in self.layers.iter_mut() {
+        for (layer_idx, layer) in self.layers.iter_mut().enumerate() {
             xs = layer.forward(&xs, attention_mask.as_ref(), seqlen_offset)?;
+            if let Some(steering) = self.steering.as_ref() {
+                xs = steering.apply(layer_idx, &xs)?;
+            }
         }
         if std::env::var("CSM_LORA_DEBUG").is_ok() {
             eprintln!("LlamaModel: post-loop xs.track_op={}", xs.track_op());
@@ -590,9 +610,10 @@ impl Model {
         let dtype = self.backbone.dtype;
 
         // Closure to keep each per-target block readable and uniform.
-        let make = |name: &str, in_dim: usize, out_dim: usize|
-            -> candle_core::Result<crate::lora::LoraDelta>
-        {
+        let make = |name: &str,
+                    in_dim: usize,
+                    out_dim: usize|
+         -> candle_core::Result<crate::lora::LoraDelta> {
             crate::lora::LoraDelta::new(
                 cfg.rank,
                 cfg.alpha as f64,
@@ -657,7 +678,13 @@ impl Model {
         }
         tracing::info!(
             "LoRA injected: q={} k={} v={} o={} mlp_gate={} mlp_up={} mlp_down={}",
-            counts[0], counts[1], counts[2], counts[3], counts[4], counts[5], counts[6],
+            counts[0],
+            counts[1],
+            counts[2],
+            counts[3],
+            counts[4],
+            counts[5],
+            counts[6],
         );
         Ok(())
     }
@@ -669,22 +696,43 @@ impl Model {
     /// in place.
     pub fn refresh_lora(&mut self, vm: &candle_nn::VarMap) -> candle_core::Result<()> {
         for (i, layer) in self.backbone.layers.iter_mut().enumerate() {
-            let do_refresh = |slot: Option<&mut crate::lora::LoraDelta>, key: String|
-                -> candle_core::Result<()>
-            {
+            let do_refresh = |slot: Option<&mut crate::lora::LoraDelta>,
+                              key: String|
+             -> candle_core::Result<()> {
                 if let Some(d) = slot {
                     d.refresh_from(vm, &key)
                         .map_err(|e| candle_core::Error::Msg(e.to_string()))?;
                 }
                 Ok(())
             };
-            do_refresh(layer.attn.q_lora.as_mut(), format!("backbone.layers.{i}.attn.q_proj"))?;
-            do_refresh(layer.attn.k_lora.as_mut(), format!("backbone.layers.{i}.attn.k_proj"))?;
-            do_refresh(layer.attn.v_lora.as_mut(), format!("backbone.layers.{i}.attn.v_proj"))?;
-            do_refresh(layer.attn.o_lora.as_mut(), format!("backbone.layers.{i}.attn.output_proj"))?;
-            do_refresh(layer.mlp.gate_lora.as_mut(), format!("backbone.layers.{i}.mlp.w1"))?;
-            do_refresh(layer.mlp.up_lora.as_mut(),   format!("backbone.layers.{i}.mlp.w3"))?;
-            do_refresh(layer.mlp.down_lora.as_mut(), format!("backbone.layers.{i}.mlp.w2"))?;
+            do_refresh(
+                layer.attn.q_lora.as_mut(),
+                format!("backbone.layers.{i}.attn.q_proj"),
+            )?;
+            do_refresh(
+                layer.attn.k_lora.as_mut(),
+                format!("backbone.layers.{i}.attn.k_proj"),
+            )?;
+            do_refresh(
+                layer.attn.v_lora.as_mut(),
+                format!("backbone.layers.{i}.attn.v_proj"),
+            )?;
+            do_refresh(
+                layer.attn.o_lora.as_mut(),
+                format!("backbone.layers.{i}.attn.output_proj"),
+            )?;
+            do_refresh(
+                layer.mlp.gate_lora.as_mut(),
+                format!("backbone.layers.{i}.mlp.w1"),
+            )?;
+            do_refresh(
+                layer.mlp.up_lora.as_mut(),
+                format!("backbone.layers.{i}.mlp.w3"),
+            )?;
+            do_refresh(
+                layer.mlp.down_lora.as_mut(),
+                format!("backbone.layers.{i}.mlp.w2"),
+            )?;
         }
         Ok(())
     }
@@ -697,6 +745,22 @@ impl Model {
         }
     }
 
+    /// Install activation-steering vectors on the conditional backbone.
+    /// The unconditional CFG backbone (when present) is intentionally left
+    /// unmodified — steering should bias the conditional pathway, not the
+    /// unconditional baseline that CFG subtracts. Pass `None` to remove.
+    pub fn set_backbone_steering(&mut self, steering: Option<crate::steering::LayerSteering>) {
+        self.backbone.set_steering(steering);
+    }
+
+    pub fn backbone_steering(&self) -> Option<&crate::steering::LayerSteering> {
+        self.backbone.steering()
+    }
+
+    pub fn backbone_num_layers(&self) -> usize {
+        self.backbone.layers.len()
+    }
+
     /// Build the per-frame embedding tensor `(B, S, D)` from packed token slots.
     /// Shared by `generate_frame` and `generate_frame_cfg`.
     fn build_embeds(&self, tokens: &Tensor, tokens_mask: &Tensor) -> Result<Tensor> {
@@ -807,7 +871,11 @@ impl Model {
         let embeds = self.build_embeds(tokens, tokens_mask)?;
         let h = self.backbone.forward(&embeds, input_pos)?;
         if std::env::var("CSM_GRAD_DEBUG").is_ok() {
-            eprintln!("FL: embeds.track_op={}, h.track_op={}", embeds.track_op(), h.track_op());
+            eprintln!(
+                "FL: embeds.track_op={}, h.track_op={}",
+                embeds.track_op(),
+                h.track_op()
+            );
         }
 
         // c0 loss. cross_entropy expects F32 logits; cast if model runs in F16/BF16.
@@ -827,8 +895,7 @@ impl Model {
         }
 
         // Teacher-forced decoder: feed ground-truth previous tokens.
-        let c0_target_t =
-            Tensor::from_slice(&[target_codes[0]], (1, 1), &self.decoder.device)?;
+        let c0_target_t = Tensor::from_slice(&[target_codes[0]], (1, 1), &self.decoder.device)?;
         let c0_embed = self.audio_embeddings.forward(&c0_target_t)?;
         let mut curr_h = Tensor::cat(&[h, c0_embed], 1)?;
         self.decoder.clear_kv_cache();
@@ -839,8 +906,7 @@ impl Model {
             decoder_pos += curr_h.dim(1)?;
             let ci_logits = decoder_h.broadcast_matmul(&self.audio_head.get(i - 1)?)?;
             let ci_logits_2d = ci_logits.i((0, 0))?.unsqueeze(0)?.to_dtype(DType::F32)?;
-            let ci_target =
-                Tensor::from_slice(&[target_codes[i]], (1,), &self.decoder.device)?;
+            let ci_target = Tensor::from_slice(&[target_codes[i]], (1,), &self.decoder.device)?;
             let ci_loss = candle_nn::loss::cross_entropy(&ci_logits_2d, &ci_target)?;
             total_loss = (total_loss + ci_loss)?;
             // Teacher-force the next decoder input with the GT codebook id.