memra_engine/
vision.rs

1//! Vision tower for Qwen3.8-27B multimodal input (lane/vision, 2026-08-15).
2//!
3//! The qwen3_5_vision ViT (depth 27, hidden 1152, heads 16, gelu_pytorch_tanh, patch 16,
4//! spatial_merge 2, temporal_patch 2, LEARNED pos embeddings on a 48x48 grid) lives in the
5//! official checkpoint's `outside.safetensors` (the unquantized shard) — the quantized
6//! trunks (ct-NVFP4 etc.) strip it. `MEMRA_VISION_DIR` points at any directory carrying
7//! that shard; the tower output is plain [n_tokens, 5120] embeddings, so vision requests
8//! serve on ANY trunk. Text side uses standard sequential rope (rope_scaling is null on
9//! this model — no M-RoPE), so spliced image tokens take ordinary positions.
10//!
11//! v1 posture: correctness-first — cuBLASLt f32 GEMMs (`Engine::linear` + bias epilogue),
12//! `sdpa_naive(causal=false)` for the bidirectional attention, host-side permutes between
13//! stages (the tower is a small fraction of a vision request; optimize later). Parity gate:
14//! merger-output cosine vs the HF reference per VISION-LANE.md.
15
16use crate::Engine;
17use cudarc::driver::CudaSlice;
18use memra_gguf::dequant::bf16_to_f32;
19use memra_gguf::safetensors::StShard;
20use std::path::Path;
21
22pub const V_HIDDEN: usize = 1152;
23pub const V_HEADS: usize = 16;
24pub const V_HEAD_DIM: usize = V_HIDDEN / V_HEADS; // 72
25pub const V_INTER: usize = 4304;
26pub const V_DEPTH: usize = 27;
27pub const V_PATCH: usize = 16;
28pub const V_MERGE: usize = 2;
29pub const V_TEMPORAL: usize = 2;
30pub const V_POS_GRID: usize = 48; // 2304 learned positions = 48x48
31pub const V_OUT: usize = 5120;
32pub const V_PATCH_IN: usize = 3 * V_TEMPORAL * V_PATCH * V_PATCH; // 1536
33pub const V_MERGED_IN: usize = V_HIDDEN * V_MERGE * V_MERGE; // 4608
34const LN_EPS: f32 = 1e-6;
35
36/// Mixed-embedding prime overlay: image embeddings that replace `<|image_pad|>` token
37/// embeddings at prompt-relative positions during `prime_cache_overlaid`. `rows` holds all
38/// images' merger outputs concatenated ([total_rows, n_embd]); each span is
39/// `(prompt_pos, row_off, n_rows)` — rows `[row_off, row_off+n_rows)` land at prompt
40/// positions `[prompt_pos, prompt_pos+n_rows)`. Spans must not overlap.
41pub struct EmbedOverlay {
42    pub rows: CudaSlice<f32>,
43    pub spans: Vec<(usize, usize, usize)>,
44}
45
46impl EmbedOverlay {
47    /// Sub-window for a prime call covering prompt-relative `[off, off+len)`: spans clipped
48    /// and rebased so the callee sees call-relative positions (the serve prefill tick primes
49    /// a prompt across multiple `prime_cache_overlaid` calls). `rows` is an Arc clone, not a
50    /// copy. None = no image rows in this window (caller may prime plain).
51    pub fn window(&self, off: usize, len: usize) -> Option<EmbedOverlay> {
52        let spans: Vec<(usize, usize, usize)> = self
53            .spans
54            .iter()
55            .filter_map(|&(pos, row_off, n_rows)| {
56                let lo = pos.max(off);
57                let hi = (pos + n_rows).min(off + len);
58                (lo < hi).then(|| (lo - off, row_off + (lo - pos), hi - lo))
59            })
60            .collect();
61        (!spans.is_empty()).then(|| EmbedOverlay {
62            rows: self.rows.clone(),
63            spans,
64        })
65    }
66}
67
68struct Lin {
69    w: CudaSlice<f32>,
70    b: CudaSlice<f32>,
71    in_f: usize,
72    out_f: usize,
73}
74
75struct VisBlock {
76    norm1_w: CudaSlice<f32>,
77    norm1_b: CudaSlice<f32>,
78    norm2_w: CudaSlice<f32>,
79    norm2_b: CudaSlice<f32>,
80    qkv: Lin,
81    proj: Lin,
82    fc1: Lin,
83    fc2: Lin,
84}
85
86pub struct VisionTower {
87    patch: Lin,
88    /// Host copy of the learned pos table [2304, 1152] — bilinear-interpolated per grid.
89    pos: Vec<f32>,
90    blocks: Vec<VisBlock>,
91    merger_norm_w: CudaSlice<f32>,
92    merger_norm_b: CudaSlice<f32>,
93    merger_fc1: Lin,
94    merger_fc2: Lin,
95}
96
97fn read_f32(sh: &StShard, name: &str) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
98    let (info, raw) = sh
99        .raw(name)
100        .ok_or_else(|| format!("vision tensor missing: {name}"))?;
101    match info.dtype.as_str() {
102        "BF16" => Ok(raw
103            .chunks_exact(2)
104            .map(|c| bf16_to_f32(u16::from_le_bytes([c[0], c[1]])))
105            .collect()),
106        "F32" => Ok(raw
107            .chunks_exact(4)
108            .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
109            .collect()),
110        other => Err(format!("vision tensor {name}: unsupported dtype {other}").into()),
111    }
112}
113
114fn load_lin(
115    e: &Engine,
116    sh: &StShard,
117    stem: &str,
118    in_f: usize,
119    out_f: usize,
120) -> Result<Lin, Box<dyn std::error::Error>> {
121    let w = read_f32(sh, &format!("{stem}.weight"))?;
122    let b = read_f32(sh, &format!("{stem}.bias"))?;
123    assert_eq!(w.len(), in_f * out_f, "{stem}.weight shape");
124    assert_eq!(b.len(), out_f, "{stem}.bias shape");
125    Ok(Lin {
126        w: e.htod(&w)?,
127        b: e.htod(&b)?,
128        in_f,
129        out_f,
130    })
131}
132
133impl VisionTower {
134    /// Load the tower from a directory containing `outside.safetensors`.
135    pub fn load(e: &Engine, dir: &Path) -> Result<Self, Box<dyn std::error::Error>> {
136        let sh = StShard::open(dir.join("outside.safetensors"))?;
137        let p = "model.visual";
138        let patch = {
139            // conv [1152, 3, 2, 16, 16] flattens to Linear 1536 -> 1152 (HF patchify order:
140            // channel-major within the (c, t, h, w) patch — the preprocessor emits the
141            // matching flat order).
142            let w = read_f32(&sh, &format!("{p}.patch_embed.proj.weight"))?;
143            let b = read_f32(&sh, &format!("{p}.patch_embed.proj.bias"))?;
144            assert_eq!(w.len(), V_HIDDEN * V_PATCH_IN);
145            Lin {
146                w: e.htod(&w)?,
147                b: e.htod(&b)?,
148                in_f: V_PATCH_IN,
149                out_f: V_HIDDEN,
150            }
151        };
152        let pos = read_f32(&sh, &format!("{p}.pos_embed.weight"))?;
153        assert_eq!(pos.len(), V_POS_GRID * V_POS_GRID * V_HIDDEN);
154        let mut blocks = Vec::with_capacity(V_DEPTH);
155        for il in 0..V_DEPTH {
156            let bp = format!("{p}.blocks.{il}");
157            blocks.push(VisBlock {
158                norm1_w: e.htod(&read_f32(&sh, &format!("{bp}.norm1.weight"))?)?,
159                norm1_b: e.htod(&read_f32(&sh, &format!("{bp}.norm1.bias"))?)?,
160                norm2_w: e.htod(&read_f32(&sh, &format!("{bp}.norm2.weight"))?)?,
161                norm2_b: e.htod(&read_f32(&sh, &format!("{bp}.norm2.bias"))?)?,
162                qkv: load_lin(e, &sh, &format!("{bp}.attn.qkv"), V_HIDDEN, 3 * V_HIDDEN)?,
163                proj: load_lin(e, &sh, &format!("{bp}.attn.proj"), V_HIDDEN, V_HIDDEN)?,
164                fc1: load_lin(e, &sh, &format!("{bp}.mlp.linear_fc1"), V_HIDDEN, V_INTER)?,
165                fc2: load_lin(e, &sh, &format!("{bp}.mlp.linear_fc2"), V_INTER, V_HIDDEN)?,
166            });
167        }
168        let merger_norm_w = e.htod(&read_f32(&sh, &format!("{p}.merger.norm.weight"))?)?;
169        let merger_norm_b = e.htod(&read_f32(&sh, &format!("{p}.merger.norm.bias"))?)?;
170        let merger_fc1 = load_lin(
171            e,
172            &sh,
173            &format!("{p}.merger.linear_fc1"),
174            V_MERGED_IN,
175            V_MERGED_IN,
176        )?;
177        let merger_fc2 = load_lin(
178            e,
179            &sh,
180            &format!("{p}.merger.linear_fc2"),
181            V_MERGED_IN,
182            V_OUT,
183        )?;
184        eprintln!(
185            "[vision] tower loaded from {} ({} blocks, f32-resident)",
186            dir.display(),
187            V_DEPTH
188        );
189        Ok(Self {
190            patch,
191            pos,
192            blocks,
193            merger_norm_w,
194            merger_norm_b,
195            merger_fc1,
196            merger_fc2,
197        })
198    }
199
200    fn linear_bias(
201        &self,
202        e: &Engine,
203        x: &CudaSlice<f32>,
204        l: &Lin,
205        m: usize,
206    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
207        let mut y = e.linear(x, &l.w, m, l.in_f, l.out_f)?;
208        // Row-broadcast bias via add_row_inplace (per-row launch; the tower is small and
209        // v1 is correctness-first — the cuBLASLt bias epilogue is the later optimization).
210        for r in 0..m {
211            e.add_row_inplace(&mut y, &l.b, l.out_f, r * l.out_f)?;
212        }
213        Ok(y)
214    }
215
216    /// Bilinear-interpolate the 48x48 learned pos table to [gh, gw] and return host
217    /// [gh*gw, 1152] (added to the patch embeddings).
218    fn pos_for_grid(&self, gh: usize, gw: usize) -> Vec<f32> {
219        let g = V_POS_GRID as f32;
220        let mut out = vec![0f32; gh * gw * V_HIDDEN];
221        for y in 0..gh {
222            for x in 0..gw {
223                // HF fast_pos_embed_interpolate: linspace(0, 47, g) == align_corners=TRUE
224                let sy = if gh > 1 {
225                    y as f32 * (g - 1.0) / (gh as f32 - 1.0)
226                } else {
227                    0.0
228                };
229                let sx = if gw > 1 {
230                    x as f32 * (g - 1.0) / (gw as f32 - 1.0)
231                } else {
232                    0.0
233                };
234                let (y0, x0) = (sy.floor() as usize, sx.floor() as usize);
235                let (y1, x1) = ((y0 + 1).min(V_POS_GRID - 1), (x0 + 1).min(V_POS_GRID - 1));
236                let (fy, fx) = (sy - y0 as f32, sx - x0 as f32);
237                let dst = &mut out[(y * gw + x) * V_HIDDEN..(y * gw + x + 1) * V_HIDDEN];
238                for c in 0..V_HIDDEN {
239                    let p00 = self.pos[(y0 * V_POS_GRID + x0) * V_HIDDEN + c];
240                    let p01 = self.pos[(y0 * V_POS_GRID + x1) * V_HIDDEN + c];
241                    let p10 = self.pos[(y1 * V_POS_GRID + x0) * V_HIDDEN + c];
242                    let p11 = self.pos[(y1 * V_POS_GRID + x1) * V_HIDDEN + c];
243                    dst[c] = p00 * (1.0 - fy) * (1.0 - fx)
244                        + p01 * (1.0 - fy) * fx
245                        + p10 * fy * (1.0 - fx)
246                        + p11 * fy * fx;
247                }
248            }
249        }
250        out
251    }
252
253    /// Forward one image's patches -> [gh*gw/4, 5120] merged embeddings (device).
254    /// `patches` is host [gh*gw, 1536] in the preprocessor's (c, t, ph, pw) flat order.
255    pub fn forward(
256        &self,
257        e: &Engine,
258        patches: &[f32],
259        gh: usize,
260        gw: usize,
261    ) -> Result<CudaSlice<f32>, Box<dyn std::error::Error>> {
262        let n = gh * gw;
263        assert_eq!(patches.len(), n * V_PATCH_IN, "patch buffer shape");
264        let xd = e.htod(patches)?;
265        let mut x = self.linear_bias(e, &xd, &self.patch, n)?;
266        // + interpolated pos embed
267        let pos = self.pos_for_grid(gh, gw);
268        let pos_d = e.htod(&pos)?;
269        let mut x2 = e.zeros(n * V_HIDDEN)?;
270        e.add(&x, &pos_d, &mut x2, n * V_HIDDEN)?;
271        x = x2;
272        // dev-only stage dumps for the HF parity bisect (row-major grid order, f32 LE)
273        let dbg = std::env::var("MEMRA_VISION_DEBUG").ok();
274        let dump = |tag: &str, buf: &[f32]| {
275            if let Some(dir) = dbg.as_deref() {
276                let raw: Vec<u8> = buf.iter().flat_map(|v| v.to_le_bytes()).collect();
277                let _ = std::fs::write(format!("{dir}/rust_{tag}.bin"), raw);
278            }
279        };
280        if dbg.is_some() {
281            dump("pre_blocks", &e.dtoh(&x)?);
282        }
283        let scale = 1.0 / (V_HEAD_DIM as f32).sqrt();
284        // 2D vision rope (Qwen3_5VisionRotaryEmbedding, theta 10000): per token (y, x) the
285        // head_dim/2 = 36 rotation angles are [y * inv_freq[0..18], x * inv_freq[0..18]],
286        // GPT-NeoX pairing (d, d+36). Same table for every block/head — precompute cos/sin.
287        let half = V_HEAD_DIM / 2; // 36
288        let quarter = half / 2; // 18
289        let inv_freq: Vec<f32> = (0..quarter)
290            .map(|i| 10000f32.powf(-(i as f32) / quarter as f32))
291            .collect();
292        let mut rope_cos = vec![0f32; n * half];
293        let mut rope_sin = vec![0f32; n * half];
294        for t in 0..n {
295            let (y, x) = (t / gw, t % gw);
296            for d in 0..half {
297                let f = if d < quarter {
298                    y as f32 * inv_freq[d]
299                } else {
300                    x as f32 * inv_freq[d - quarter]
301                };
302                rope_cos[t * half + d] = f.cos();
303                rope_sin[t * half + d] = f.sin();
304            }
305        }
306        for (ib, blk) in self.blocks.iter().enumerate() {
307            // attn: ln1 -> qkv -> sdpa(causal=false) -> proj -> +res
308            let mut h = e.zeros(n * V_HIDDEN)?;
309            e.layer_norm_bias(&x, &blk.norm1_w, &blk.norm1_b, &mut h, V_HIDDEN, n, LN_EPS)?;
310            let qkv = self.linear_bias(e, &h, &blk.qkv, n)?;
311            // sdpa_naive consumes token-major [T, n_head, head_dim] — exactly the qkv GEMM
312            // row layout, so q/k/v are column splits of each row (no permute). Host pass
313            // applies the vision rope to q/k on the way (v untouched).
314            let qkv_h = e.dtoh(&qkv)?;
315            let mut qh = vec![0f32; n * V_HIDDEN];
316            let mut kh = vec![0f32; n * V_HIDDEN];
317            let mut vh = vec![0f32; n * V_HIDDEN];
318            for t in 0..n {
319                let row = &qkv_h[t * 3 * V_HIDDEN..(t + 1) * 3 * V_HIDDEN];
320                let dst = t * V_HIDDEN;
321                vh[dst..dst + V_HIDDEN].copy_from_slice(&row[2 * V_HIDDEN..3 * V_HIDDEN]);
322                for hd in 0..V_HEADS {
323                    let o = hd * V_HEAD_DIM;
324                    // rotate-half pairs (d, d+36), angles shared across heads
325                    for d in 0..half {
326                        let (c, sn) = (rope_cos[t * half + d], rope_sin[t * half + d]);
327                        let (qa, qb) = (row[o + d], row[o + d + half]);
328                        qh[dst + o + d] = qa * c - qb * sn;
329                        qh[dst + o + d + half] = qb * c + qa * sn;
330                        let (ka, kb) = (row[V_HIDDEN + o + d], row[V_HIDDEN + o + d + half]);
331                        kh[dst + o + d] = ka * c - kb * sn;
332                        kh[dst + o + d + half] = kb * c + ka * sn;
333                    }
334                }
335            }
336            let (qd, kd, vd) = (e.htod(&qh)?, e.htod(&kh)?, e.htod(&vh)?);
337            let mut od = e.zeros(n * V_HIDDEN)?;
338            e.sdpa_naive(
339                &qd, &kd, &vd, &mut od, V_HEAD_DIM, V_HEADS, V_HEADS, n, n, scale, false,
340            )?;
341            let attn = self.linear_bias(e, &od, &blk.proj, n)?;
342            let mut xr = e.zeros(n * V_HIDDEN)?;
343            e.add(&x, &attn, &mut xr, n * V_HIDDEN)?;
344            // mlp: ln2 -> fc1 -> gelu_tanh -> fc2 -> +res
345            let mut h2 = e.zeros(n * V_HIDDEN)?;
346            e.layer_norm_bias(
347                &xr,
348                &blk.norm2_w,
349                &blk.norm2_b,
350                &mut h2,
351                V_HIDDEN,
352                n,
353                LN_EPS,
354            )?;
355            let f1 = self.linear_bias(e, &h2, &blk.fc1, n)?;
356            let mut g = e.zeros(n * V_INTER)?;
357            e.gelu_tanh(&f1, &mut g, n * V_INTER)?;
358            let f2 = self.linear_bias(e, &g, &blk.fc2, n)?;
359            let mut xn = e.zeros(n * V_HIDDEN)?;
360            e.add(&xr, &f2, &mut xn, n * V_HIDDEN)?;
361            x = xn;
362            if dbg.is_some() && ib == 0 {
363                dump("blk0", &e.dtoh(&x)?);
364            }
365        }
366        if dbg.is_some() {
367            dump("post_blocks", &e.dtoh(&x)?);
368        }
369        // merger: LN over [n, 1152], then 2x2 spatial concat -> [n/4, 4608] -> fc1 -> gelu -> fc2
370        let mut ln = e.zeros(n * V_HIDDEN)?;
371        e.layer_norm_bias(
372            &x,
373            &self.merger_norm_w,
374            &self.merger_norm_b,
375            &mut ln,
376            V_HIDDEN,
377            n,
378            LN_EPS,
379        )?;
380        let lh = e.dtoh(&ln)?;
381        let (mh, mw) = (gh / V_MERGE, gw / V_MERGE);
382        let nm = mh * mw;
383        let mut merged = vec![0f32; nm * V_MERGED_IN];
384        for my in 0..mh {
385            for mx in 0..mw {
386                let dst =
387                    &mut merged[(my * mw + mx) * V_MERGED_IN..(my * mw + mx + 1) * V_MERGED_IN];
388                for sy in 0..V_MERGE {
389                    for sx in 0..V_MERGE {
390                        let t = (my * V_MERGE + sy) * gw + (mx * V_MERGE + sx);
391                        let seg = (sy * V_MERGE + sx) * V_HIDDEN;
392                        dst[seg..seg + V_HIDDEN]
393                            .copy_from_slice(&lh[t * V_HIDDEN..(t + 1) * V_HIDDEN]);
394                    }
395                }
396            }
397        }
398        let md = e.htod(&merged)?;
399        let f1 = self.linear_bias(e, &md, &self.merger_fc1, nm)?;
400        let mut g = e.zeros(nm * V_MERGED_IN)?;
401        e.gelu_tanh(&f1, &mut g, nm * V_MERGED_IN)?;
402        self.linear_bias(e, &g, &self.merger_fc2, nm)
403    }
404}
memra_engine/vision.rs

memra_engine/
vision.rs
