S2.67: optional half-precision GPU weights (EFFECTCRAFT_SAM2_F16=1)

Halves the weight memory (Large: about 900 MB to 450 MB). Off by default: the
video logit error grows from about 0.003 to 0.06-0.12 (masks unchanged) for a
speed gain of under 10 percent.

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

Cargo.lock

@@ -1836,6 +1836,7 @@ dependencies = [
  "candle-core",
  "candle-nn",
  "effectcraft-segment",
+ "half",
  "log",
  "pollster 0.4.0",
  "rayon",

crates/sam2/Cargo.toml

@@ -22,3 +22,4 @@ candle-nn = "=0.9.2"
 wgpu = { workspace = true }
 pollster = { workspace = true }
 bytemuck = { workspace = true }
+half = "2.7"

crates/sam2/src/gpu/kernels.wgsl

@@ -96,3 +96,70 @@ fn matmul(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li
         }
     }
 }
+
+// The same, with B (a weight matrix) stored as half precision, two values per word.
+@group(0) @binding(1) var<storage, read> w16: array<u32>;
+
+@compute @workgroup_size(16, 16)
+fn matmul_w16(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
+    let m = u(0u); let n = u(1u); let k = u(2u);
+    let a_base = u(4u) + wg.z * u(5u);
+    let b_base = u(8u) + wg.z * u(9u);
+    let a_rs = u(6u); let a_cs = u(7u); let b_rs = u(10u); let b_cs = u(11u);
+    let row0 = wg.y * 64u;
+    let col0 = wg.x * 64u;
+    let tid = li.y * 16u + li.x;
+    let a_k_contig = a_cs == 1u;
+    let b_n_contig = b_cs == 1u;
+    var acc: array<f32, 16>;
+    for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; }
+    for (var kt = 0u; kt < k; kt += 16u) {
+        for (var i = 0u; i < 4u; i++) {
+            let idx = tid + i * 256u;
+            var r: u32; var c: u32;
+            if (a_k_contig) { r = idx / 16u; c = idx % 16u; } else { r = idx % 64u; c = idx / 64u; }
+            let gr = row0 + r; let gc = kt + c;
+            var va = 0.0;
+            if (gr < m && gc < k) { va = in0[a_base + gr * a_rs + gc * a_cs]; }
+            sa[r * 16u + c] = va;
+            var r2: u32; var c2: u32;
+            if (b_n_contig) { r2 = idx / 64u; c2 = idx % 64u; } else { r2 = idx % 16u; c2 = idx / 16u; }
+            let gr2 = kt + r2; let gc2 = col0 + c2;
+            var vb = 0.0;
+            if (gr2 < k && gc2 < n) { let bi = b_base + gr2 * b_rs + gc2 * b_cs; vb = unpack2x16float(w16[bi >> 1u])[bi & 1u]; }
+            sb[r2 * 64u + c2] = vb;
+        }
+        workgroupBarrier();
+        for (var kk = 0u; kk < 16u; kk++) {
+            var ar: array<f32, 4>;
+            var br: array<f32, 4>;
+            for (var i = 0u; i < 4u; i++) {
+                ar[i] = sa[(li.y * 4u + i) * 16u + kk];
+                br[i] = sb[kk * 64u + li.x * 4u + i];
+            }
+            for (var i = 0u; i < 4u; i++) {
+                for (var j = 0u; j < 4u; j++) { acc[i * 4u + j] += ar[i] * br[j]; }
+            }
+        }
+        workgroupBarrier();
+    }
+    let flags = u(15u);
+    let alpha = f(16u);
+    let c_base = u(12u) + wg.z * u(13u);
+    let c_rs = u(14u);
+    let act = (flags >> 1u) & 3u;
+    for (var i = 0u; i < 4u; i++) {
+        for (var j = 0u; j < 4u; j++) {
+            let gr = row0 + li.y * 4u + i;
+            let gc = col0 + li.x * 4u + j;
+            if (gr < m && gc < n) {
+                var v = acc[i * 4u + j] * alpha;
+                if ((flags & 1u) != 0u) { v += in2[gc]; }
+                if (act == 1u) { v = gelu_(v); } else if (act == 2u) { v = max(v, 0.0); }
+                let idx = c_base + gr * c_rs + gc;
+                if ((flags & 8u) != 0u) { v += in3[idx]; }
+                out0[idx] = v;
+            }
+        }
+    }
+}

crates/sam2/src/gpu/mod.rs

@@ -23,6 +23,8 @@ pub struct Buf {
     pub(crate) buffer: wgpu::Buffer,
     /// Elements (f32).
     pub len: usize,
+    /// Two half-precision values per 32-bit word (a weight matrix only; `len` still counts values).
+    pub(crate) f16: bool,
 }
 
 /// One strided matrix operand: `(offset, batch stride, row stride, column stride)` in elements.
@@ -89,6 +91,11 @@ fn u32_of(v: usize, what: &str) -> Result<u32> {
     u32::try_from(v).map_err(|_| Error::Compute(format!("{what} {v} does not fit a GPU index")))
 }
 
+/// Whether weight matrices are kept on the GPU in half precision (`EFFECTCRAFT_SAM2_F16=1`).
+pub fn half_weights() -> bool {
+    std::env::var("EFFECTCRAFT_SAM2_F16").is_ok_and(|v| matches!(v.trim(), "1" | "on" | "true"))
+}
+
 impl Gpu {
     /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders
     /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go,
@@ -158,7 +165,36 @@ impl Gpu {
         if !data.is_empty() {
             self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(data));
         }
-        Buf { buffer, len: data.len() }
+        Buf { buffer, len: data.len(), f16: false }
+    }
+
+    /// A weight matrix: stored as half precision (half the memory) when `EFFECTCRAFT_SAM2_F16` asks for it and every
+    /// value fits, as f32 otherwise.
+    pub fn upload_weights(&self, data: &[f32]) -> Buf {
+        if !half_weights() || data.is_empty() || data.iter().any(|v| v.abs() > 60_000.0) {
+            return self.upload(data);
+        }
+        self.upload_f16(data)
+    }
+
+    /// `data` as half precision, two values per word.
+    pub(crate) fn upload_f16(&self, data: &[f32]) -> Buf {
+        let words: Vec<u32> = data
+            .chunks(2)
+            .map(|p| {
+                let lo = half::f16::from_f32(p.first().copied().unwrap_or(0.0)).to_bits();
+                let hi = half::f16::from_f32(p.get(1).copied().unwrap_or(0.0)).to_bits();
+                u32::from(lo) | (u32::from(hi) << 16)
+            })
+            .collect();
+        let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
+            label: None,
+            size: (words.len() as u64) * 4,
+            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+        self.queue.write_buffer(&buffer, 0, bytemuck::cast_slice(&words));
+        Buf { buffer, len: data.len(), f16: true }
     }
 
     /// A zeroed buffer of `len` elements.
@@ -173,7 +209,7 @@ impl Gpu {
             usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
             mapped_at_creation: false,
         });
-        Ok(Buf { buffer, len })
+        Ok(Buf { buffer, len, f16: false })
     }
 
     fn pipeline(&self, entry: &'static str) -> Arc<wgpu::ComputePipeline> {
@@ -315,7 +351,7 @@ impl Gpu {
             }
         };
         self.run(
-            "matmul",
+            if b.f16 { "matmul_w16" } else { "matmul" },
             &[(0, a), (1, b), (2, bias), (3, c), (4, res)],
             &params,
             (u32_of(op.n.div_ceil(64), "groups")?, u32_of(op.m.div_ceil(64), "groups")?, u32_of(op.batch, "batch")?),
@@ -351,7 +387,7 @@ impl GLinear {
     /// From `{name}.weight` (`[out, in]`) and `{name}.bias`.
     pub fn load(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result<GLinear> {
         Ok(GLinear {
-            w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?),
+            w: g.upload_weights(&flat(w.get(&format!("{name}.weight"), &[out, inp])?)?),
             b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?),
             inp,
             out,
@@ -373,13 +409,13 @@ impl GLinear {
                 *b *= s;
             }
         }
-        Ok(GLinear { w: g.upload(&wt), b: g.upload(&bias), inp, out })
+        Ok(GLinear { w: g.upload_weights(&wt), b: g.upload(&bias), inp, out })
     }
 
     /// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens.
     pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result<GLinear> {
         Ok(GLinear {
-            w: g.upload(&flat(w.get(&format!("{name}.weight"), &[out, inp, 1, 1])?)?),
+            w: g.upload_weights(&flat(w.get(&format!("{name}.weight"), &[out, inp, 1, 1])?)?),
             b: g.upload(&flat(w.get(&format!("{name}.bias"), &[out])?)?),
             inp,
             out,

crates/sam2/src/gpu/tests.rs

@@ -344,3 +344,36 @@ fn convolutions_and_the_layernorm_gelu_epilogue_match_candle() {
         .collect();
     assert_close(&run_and_read(&g, &gy), &want, 1e-3, "layernorm + gelu");
 }
+
+#[test]
+fn half_precision_weights_give_the_same_product_to_half_precision() {
+    let Some(g) = gpu() else { return };
+    // odd sizes: the last word holds one value; the weight is a linear layer's [n, k]
+    let (m, n, k) = (70, 91, 41);
+    let (av, wv) = (noise(m * k, 3), noise(n * k, 4));
+    let (a, w16, w32) = (g.upload(&av), g.upload_f16(&wv), g.upload(&wv));
+    assert!(w16.f16 && w16.len == n * k);
+    let (c16, c32) = (g.zeros(m * n).unwrap(), g.zeros(m * n).unwrap());
+    let sb = Strides { offset: 0, batch: 0, row: 1, col: k };
+    for (w, c) in [(&w16, &c16), (&w32, &c32)] {
+        g.matmul(&Matmul {
+            m,
+            n,
+            k,
+            batch: 1,
+            a: (&a, Strides::row_major(k)),
+            b: (w, sb),
+            bias: None,
+            c: (c, Strides::row_major(n)),
+            residual: None,
+            alpha: 1.0,
+            act: Act::None,
+        })
+        .unwrap();
+    }
+    g.flush().unwrap();
+    let (got, want) = (g.download(&c16).unwrap(), g.download(&c32).unwrap());
+    // f16 keeps 11 bits: each weight is off by at most 2^-11 relative
+    assert_close(&got, &want, 2e-3 * (k as f32).sqrt(), "f16 vs f32 weights");
+    assert!(got != want, "the weights really were rounded to half precision");
+}