S2.14: sam2_bench example and first CPU speedups

Encode 5.0 s -> 3.4 s per 1024x1024 frame in release on 16 threads (fused
softmax, queries scaled before the matmul, flat 2-D linear layers); a click
on an encoded frame is about 120 ms. Reference tests unchanged.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

crates/sam2/examples/sam2_bench.rs

@@ -0,0 +1,51 @@
+//! Timing of the SAM 2 image path on this machine: load, encode one frame, decode a click.
+//!
+//! `cargo run -p effectcraft-sam2 --example sam2_bench -- [MODEL_DIR]` (default `sam2-tiny`).
+//! Prints a message and exits when the weights are not there.
+
+use std::time::Instant;
+
+use candle_core::{DType, Device, Tensor};
+use effectcraft_sam2::model::ImageModel;
+use effectcraft_sam2::prompt::Prompt;
+
+fn main() {
+    if let Err(e) = run() {
+        eprintln!("sam2_bench: {e}");
+        std::process::exit(1);
+    }
+}
+
+fn run() -> Result<(), Box<dyn std::error::Error>> {
+    let dir = std::env::args().nth(1).unwrap_or_else(|| "sam2-tiny".to_string());
+    let file = std::path::Path::new(&dir).join("model.safetensors");
+    if !file.exists() {
+        println!("no weights at {} (download facebook/sam2.1-hiera-tiny first)", file.display());
+        return Ok(());
+    }
+    let t = Instant::now();
+    let model = ImageModel::open(&file)?;
+    println!("load: {:.2} s", t.elapsed().as_secs_f64());
+    let pixels = Tensor::randn(0f32, 1.0, (1, 3, 1024, 1024), &Device::Cpu)?.to_dtype(DType::F32)?;
+    for run in 0..3 {
+        let t = Instant::now();
+        let enc = model.encode(&pixels)?;
+        let encode = t.elapsed().as_secs_f64();
+        let prompt = Prompt { points: vec![([512.0, 512.0], true)], bbox: None };
+        let t = Instant::now();
+        let out = model.segment(&enc, &prompt, true)?;
+        println!("run {run}: encode {encode:.2} s, click {:.0} ms, masks {:?}", t.elapsed().as_secs_f64() * 1e3, out.masks.dims());
+    }
+    // Where the encode time goes: patch embedding, then each Hiera block.
+    let hiera = &model.vision.hiera;
+    let t = Instant::now();
+    let mut x = hiera.patch_embed(&pixels)?;
+    println!("patch embed: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
+    for (i, b) in hiera.blocks.iter().enumerate() {
+        let t = Instant::now();
+        x = b.forward(&x)?;
+        println!("block {i}: {:.0} ms, out {:?}", t.elapsed().as_secs_f64() * 1e3, x.dims());
+    }
+    println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get()));
+    Ok(())
+}

crates/sam2/src/decoder.rs

@@ -6,9 +6,9 @@
 //! One object per call (batch 1, one point batch).
 
 use candle_core::{D, Tensor};
-use candle_nn::{LayerNorm, Linear, Module, ops::sigmoid, ops::softmax};
+use candle_nn::{Linear, ops::sigmoid, ops::softmax};
 
-use crate::nn::{layer_norm, linear};
+use crate::nn::{LayerNorm, apply, layer_norm, linear};
 use crate::weights::Weights;
 use crate::{Error, Result};
 
@@ -40,11 +40,11 @@ impl Attention {
     /// `q`: `[1, nq, 256]`, `k` and `v`: `[1, nk, 256]`.
     fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
         let split = |t: Tensor| -> Result<Tensor> { Ok(t.reshape((1, (), self.heads, t.dim(2)? / self.heads))?.transpose(1, 2)?.contiguous()?) };
-        let (q, k, v) = (split(self.q.forward(q)?)?, split(self.k.forward(k)?)?, split(self.v.forward(v)?)?);
+        let (q, k, v) = (split(apply(&self.q, q)?)?, split(apply(&self.k, k)?)?, split(apply(&self.v, v)?)?);
         let att = softmax(&(q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?, D::Minus1)?;
         let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?;
         let (n, h, d) = (out.dim(1)?, out.dim(2)?, out.dim(3)?);
-        Ok(self.o.forward(&out.reshape((1, n, h * d))?)?)
+        apply(&self.o, &out.reshape((1, n, h * d))?)
     }
 }
 
@@ -67,11 +67,11 @@ impl FeedForward {
     }
 
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        let mut x = self.proj_in.forward(x)?.relu()?;
+        let mut x = apply(&self.proj_in, x)?.relu()?;
         for l in &self.layers {
-            x = l.forward(&x)?.relu()?;
+            x = apply(l, &x)?.relu()?;
         }
-        let x = self.proj_out.forward(&x)?;
+        let x = apply(&self.proj_out, &x)?;
         Ok(if self.sigmoid_output { sigmoid(&x)? } else { x })
     }
 }

crates/sam2/src/hiera.rs

@@ -4,10 +4,10 @@
 //! The Hiera backbone: patch embedding, windowed positional embedding, and blocks with window or
 //! global attention and query pooling at the stage changes. Activations are NHWC like the reference.
 
-use candle_core::{D, Tensor};
-use candle_nn::{LayerNorm, Linear, Module, ops::softmax};
+use candle_core::Tensor;
+use candle_nn::{Linear, ops::softmax_last_dim};
 
-use crate::nn::{Mlp, layer_norm, linear};
+use crate::nn::{LayerNorm, Mlp, apply, layer_norm, linear};
 use crate::weights::Weights;
 use crate::{Error, Result};
 
@@ -89,7 +89,7 @@ struct Attention {
 impl Attention {
     fn forward(&self, x: &Tensor) -> Result<Tensor> {
         let (b, h, w, _) = x.dims4()?;
-        let qkv = self.qkv.forward(x)?.reshape((b, h * w, 3, self.heads, ()))?;
+        let qkv = apply(&self.qkv, x)?.reshape((b, h * w, 3, self.heads, ()))?;
         let part = |i: usize| -> Result<Tensor> { Ok(qkv.narrow(2, i, 1)?.squeeze(2)?) };
         let (mut q, k, v) = (part(0)?, part(1)?, part(2)?);
         let (mut oh, mut ow) = (h, w);
@@ -101,10 +101,12 @@ impl Attention {
         let q = q.transpose(1, 2)?.contiguous()?;
         let k = k.transpose(1, 2)?.contiguous()?;
         let v = v.transpose(1, 2)?.contiguous()?;
-        let att = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?;
-        let att = softmax(&att, D::Minus1)?;
+        // Scaling the (small) queries first equals scaling the scores, and the fused softmax is
+        // much faster than the generic one on the 4096 x 4096 global-attention maps.
+        let att = (q * self.scale)?.matmul(&k.transpose(2, 3)?.contiguous()?)?;
+        let att = softmax_last_dim(&att)?;
         let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?.reshape((b, oh, ow, ()))?;
-        Ok(self.proj.forward(&out)?)
+        apply(&self.proj, &out)
     }
 }
 
@@ -149,8 +151,8 @@ impl Block {
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
         let h = self.norm1.forward(x)?;
         let residual = match (&self.proj, self.query_stride) {
-            (Some(p), Some(s)) => do_pool(&p.forward(&h)?, s)?,
-            (Some(p), None) => p.forward(&h)?,
+            (Some(p), Some(s)) => do_pool(&apply(p, &h)?, s)?,
+            (Some(p), None) => apply(p, &h)?,
             (None, _) => x.clone(),
         };
         let (hh, ww) = (h.dim(1)?, h.dim(2)?);

crates/sam2/src/nn.rs

@@ -2,17 +2,31 @@
 //! Linear, layer norm and a few tensor helpers on candle.
 
 use candle_core::Tensor;
-use candle_nn::{LayerNorm, Linear, Module};
+use candle_nn::{Linear, Module};
 
-use crate::Result;
 use crate::weights::Weights;
+use crate::{Error, Result};
 
 pub fn linear(w: &Weights, prefix: &str, input: usize, output: usize) -> Result<Linear> {
     Ok(Linear::new(w.get(&format!("{prefix}.weight"), &[output, input])?, Some(w.get(&format!("{prefix}.bias"), &[output])?)))
 }
 
+/// Layer norm over the last dimension using candle's fused kernel (much faster than the
+/// composed `candle_nn::LayerNorm` on large activations).
+pub struct LayerNorm {
+    weight: Tensor,
+    bias: Tensor,
+    eps: f32,
+}
+
+impl LayerNorm {
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        Ok(candle_nn::ops::layer_norm(&x.contiguous()?, &self.weight, &self.bias, self.eps)?)
+    }
+}
+
 pub fn layer_norm(w: &Weights, prefix: &str, dim: usize, eps: f64) -> Result<LayerNorm> {
-    Ok(LayerNorm::new(w.get(&format!("{prefix}.weight"), &[dim])?, w.get(&format!("{prefix}.bias"), &[dim])?, eps))
+    Ok(LayerNorm { weight: w.get(&format!("{prefix}.weight"), &[dim])?, bias: w.get(&format!("{prefix}.bias"), &[dim])?, eps: eps as f32 })
 }
 
 /// Two-layer MLP with exact (erf) GELU, as `Sam2FeedForward` with `num_layers = 2`.
@@ -27,6 +41,19 @@ impl Mlp {
     }
 
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        Ok(self.proj_out.forward(&self.proj_in.forward(x)?.gelu_erf()?)?)
+        apply(&self.proj_out, &apply(&self.proj_in, x)?.gelu_erf()?)
     }
 }
+
+/// Apply a linear layer to `[.., in]` through a flat 2-D matmul. Candle's `Linear::forward`
+/// broadcasts the weight over every leading dimension (copying it per window), which is slow
+/// on the 4-D activations of the backbone.
+pub fn apply(l: &Linear, x: &Tensor) -> Result<Tensor> {
+    let dims = x.dims();
+    let (&last, lead) = dims.split_last().ok_or_else(|| Error::Compute("linear on a scalar".into()))?;
+    let rows: usize = lead.iter().product();
+    let y = l.forward(&x.reshape((rows, last))?)?;
+    let mut out: Vec<usize> = lead.to_vec();
+    out.push(y.dim(1)?);
+    Ok(y.reshape(out)?)
+}

crates/sam2/src/prompt.rs

@@ -4,9 +4,8 @@
 //! The prompt encoder: clicks and boxes become sparse tokens, a mask or "no mask" becomes a dense map.
 
 use candle_core::{Device, Tensor};
-use candle_nn::{LayerNorm, Module};
 
-use crate::nn::layer_norm;
+use crate::nn::{LayerNorm, layer_norm};
 use crate::weights::Weights;
 use crate::{Error, Result};