From 5aa52c5ba3e016232c18bca2df87775bd74662fd Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 01:26:16 +0200 Subject: [PATCH] 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 --- crates/sam2/examples/sam2_bench.rs | 51 ++++++++++++++++++++++++++++++ crates/sam2/src/decoder.rs | 14 ++++---- crates/sam2/src/hiera.rs | 20 ++++++------ crates/sam2/src/nn.rs | 35 +++++++++++++++++--- crates/sam2/src/prompt.rs | 3 +- 5 files changed, 101 insertions(+), 22 deletions(-) create mode 100644 crates/sam2/examples/sam2_bench.rs diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs new file mode 100644 index 000000000..7bc678b7c --- /dev/null +++ b/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> { + 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(()) +} diff --git a/crates/sam2/src/decoder.rs b/crates/sam2/src/decoder.rs index e42a7acd9..d7a7a6dcf 100644 --- a/crates/sam2/src/decoder.rs +++ b/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 { let split = |t: Tensor| -> Result { 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 { - 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 }) } } diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs index 5091d331b..90833f7da 100644 --- a/crates/sam2/src/hiera.rs +++ b/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 { 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 { 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 { 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)?); diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs index ef334432a..f90958eb1 100644 --- a/crates/sam2/src/nn.rs +++ b/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 { 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 { + 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 { - 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 { - 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 { + 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 = lead.to_vec(); + out.push(y.dim(1)?); + Ok(y.reshape(out)?) +} diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 95834c3cb..55dd901a2 100644 --- a/crates/sam2/src/prompt.rs +++ b/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};