taconite_sam3/
text.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 Brishen Hawkins
2// SPDX-License-Identifier: Apache-2.0
3
4//! The CLIP text encoder (24 pre-LN layers, width 1024, causal attention)
5//! and SAM3's projection to the DETR width, on the host.
6//!
7//! Only the valid tokens are encoded: attention is causal and the padding
8//! follows them, so they never see it, and every consumer of the text
9//! features masks the padded positions out (the DETR cross-attentions, the
10//! mask decoder's, the scoring's mean pool). Their rows are left zero. A
11//! short prompt is a handful of tokens, so this is weight-bandwidth bound
12//! (the weights ship as bf16).
13
14use crate::bundle::Store;
15use crate::cpu::{self, Attn, Rows, W};
16use crate::{Config, Error, Sam3};
17
18/// An encoded prompt: `feats [L, 256]` (L = 32), `valid [L]`.
19pub struct Text {
20    pub feats: Vec<f32>,
21    pub valid: Vec<bool>,
22}
23
24impl Sam3 {
25    /// The CLIP text encoder over the tokens (`ids`, `mask` as the tokenizer
26    /// gives them) -> the prompt features `[L, 256]`.
27    pub fn text(&self, ids: &[u32], mask: &[u32]) -> Result<Text, Error> {
28        encode(&self.store, &self.cfg, ids, mask)
29    }
30}
31
32/// The text encoder over the bundle's weights alone (so it can run on a
33/// side thread while the ViT has the NPU).
34pub(crate) fn encode(st: &Store, c: &Config, ids: &[u32], mask: &[u32]) -> Result<Text, Error> {
35    {
36        let (dim, l) = (c.text_dim, c.text_len);
37        let n = mask.iter().take_while(|&&m| m != 0).count();
38        if n == 0 || mask[n..].iter().any(|&m| m != 0) || ids.len() != l {
39            return Err(Error::Input("the attention mask must be a non-empty prefix".into()));
40        }
41        let tok = st.bf16("t.tok_emb")?;
42        let pos = st.f32("t.pos_emb")?;
43        let mut x = vec![0f32; n * dim];
44        for (t, row) in x.chunks_mut(dim).enumerate() {
45            let e = &tok[ids[t] as usize * dim..(ids[t] as usize + 1) * dim];
46            for (j, v) in row.iter_mut().enumerate() {
47                *v = taconite::bf16_to_f32(e[j]) + pos[t * dim + j];
48            }
49        }
50        let eps = c.text_eps;
51        for i in 0..c.text_layers {
52            let p = |s: &str| format!("t.{i}.{s}");
53            let h = cpu::layer_norm(&x, dim, st.f32(&p("ln1.w"))?, st.f32(&p("ln1.b"))?, eps);
54            let qkv = cpu::linear(&h, dim, W::Bf16(st.bf16(&p("qkv.w"))?), Some(st.f32(&p("qkv.b"))?));
55            let q: Vec<f32> = qkv.chunks(3 * dim).flat_map(|r| r[..dim].iter().copied()).collect();
56            let a = Attn { heads: c.text_heads, causal: true, ..Default::default() };
57            let o = cpu::attention(
58                &q,
59                dim,
60                Rows::strided(&qkv, n, 3 * dim, dim),
61                Rows::strided(&qkv, n, 3 * dim, 2 * dim),
62                &a,
63            );
64            let o = cpu::linear(&o, dim, W::Bf16(st.bf16(&p("o.w"))?), Some(st.f32(&p("o.b"))?));
65            cpu::add_(&mut x, &o);
66            let h = cpu::layer_norm(&x, dim, st.f32(&p("ln2.w"))?, st.f32(&p("ln2.b"))?, eps);
67            let mut f = cpu::linear(&h, dim, W::Bf16(st.bf16(&p("fc1.w"))?), Some(st.f32(&p("fc1.b"))?));
68            for v in f.iter_mut() {
69                *v = cpu::gelu(*v);
70            }
71            let f = cpu::linear(&f, c.text_ffn, W::Bf16(st.bf16(&p("fc2.w"))?), Some(st.f32(&p("fc2.b"))?));
72            cpu::add_(&mut x, &f);
73        }
74        let x = cpu::layer_norm(&x, dim, st.f32("t.final_ln.w")?, st.f32("t.final_ln.b")?, eps);
75        let y = cpu::linear(&x, dim, W::F32(st.f32("t.proj.w")?), Some(st.f32("t.proj.b")?));
76        let d = c.d_model;
77        let mut feats = vec![0f32; l * d];
78        feats[..n * d].copy_from_slice(&y);
79        Ok(Text { feats, valid: mask.iter().map(|&m| m != 0).collect() })
80    }
81}
taconite_sam3/text.rs

taconite_sam3/
text.rs
