From 2dfe1b8dbec73c49e103351688cf439c0ee6055b Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Sat, 17 Oct 2026 14:37:50 +0200 Subject: [PATCH] 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 --- Cargo.lock | 1 + crates/sam2/Cargo.toml | 1 + crates/sam2/src/gpu/kernels.wgsl | 67 ++++++++++++++++++++++++++++++++ crates/sam2/src/gpu/mod.rs | 48 ++++++++++++++++++++--- crates/sam2/src/gpu/tests.rs | 33 ++++++++++++++++ 5 files changed, 144 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b4acc3a6c..b104f63cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1836,6 +1836,7 @@ dependencies = [ "candle-core", "candle-nn", "effectcraft-segment", + "half", "log", "pollster 0.4.0", "rayon", diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index ff8e64c04..e45cc8def 100644 --- a/crates/sam2/Cargo.toml +++ b/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" diff --git a/crates/sam2/src/gpu/kernels.wgsl b/crates/sam2/src/gpu/kernels.wgsl index cce4c2a05..faea20e76 100644 --- a/crates/sam2/src/gpu/kernels.wgsl +++ b/crates/sam2/src/gpu/kernels.wgsl @@ -96,3 +96,70 @@ fn matmul(@builtin(workgroup_id) wg: vec3, @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 w16: array; + +@compute @workgroup_size(16, 16) +fn matmul_w16(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_id) li: vec3) { + 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; + 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; + var br: array; + 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; + } + } + } +} diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index 5b091062c..152c67c50 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/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::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 = 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 { @@ -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)], ¶ms, (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 { 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 { 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, diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs index ed5b3528a..544df73e0 100644 --- a/crates/sam2/src/gpu/tests.rs +++ b/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"); +}