metrale_model_layers/layers/
vision_tower.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The vision tower a built model holds, whichever family bound it.
4//!
5//! Owner: model-layers (vision).
6//! Invariants: none beyond the types.
7//!
8//! Two towers exist: the Qwen3-VL-shaped [`VisionEncoder`] and GLM-5.3's
9//! [`GlmVit`]. The engine uses them only through this enum:
10//!
11//! * `forward_batched(images) -> per-image (post_h, post_w, merged_p)`,
12//! * a packed `[Σmerged_p, out_hidden_size]` BF16 output, image-ordered, read
13//!   through `out_row`,
14//! * `out_hidden_size`,
15//! * `ensure_scratch`, for a tensor-parallel worker that receives rows instead
16//!   of encoding.
17//!
18//! An enum rather than a trait object, so adding a tower means adding a
19//! variant to every match here.
20
21use anyhow::Result;
22use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
23
24use super::glm_vit::GlmVit;
25use super::vision_encoder::VisionEncoder;
26
27pub enum VisionTower {
28    /// 2026-09-25: Built by `weight_loader::qwen3_vl` and `weight_loader::qwen35`
29    /// (which the qwen35_dense and qwen4_exp loaders call).
30    Qwen(Box<VisionEncoder>),
31    /// 2026-09-25: GLM-5.3-Flash (`glm5_next`).
32    Glm(Box<GlmVit>),
33}
34
35impl VisionTower {
36    pub fn qwen(encoder: VisionEncoder) -> Self {
37        Self::Qwen(Box::new(encoder))
38    }
39
40    pub fn glm(encoder: GlmVit) -> Self {
41        Self::Glm(Box::new(encoder))
42    }
43
44    /// 2026-09-25: Encode a batch of `(pixels, grid_h, grid_w)` temporal groups. Returns
45    /// `(post_merge_h, post_merge_w, merged_patches)` per group, in order.
46    pub fn forward_batched(
47        &self,
48        images: &[(&[f32], usize, usize)],
49        gpu: &dyn GpuBackend,
50        stream: u64,
51    ) -> Result<Vec<(usize, usize, usize)>> {
52        match self {
53            Self::Qwen(e) => e.forward_batched(images, gpu, stream),
54            Self::Glm(e) => e.forward_batched(images, gpu, stream),
55        }
56    }
57
58    /// 2026-09-25: Width of one merged embedding row.
59    pub fn out_hidden_size(&self) -> usize {
60        match self {
61            Self::Qwen(e) => e.out_hidden_size,
62            Self::Glm(e) => e.out_hidden_size,
63        }
64    }
65
66    /// 2026-09-25: Allocate the encoder's device scratch if it is not allocated yet.
67    ///
68    /// Normally the first image does this. A tensor-parallel worker receives
69    /// the merged rows by broadcast instead of encoding, so it calls this
70    /// before `out_row` names a destination.
71    pub fn ensure_scratch(&self, gpu: &dyn GpuBackend) -> Result<()> {
72        match self {
73            Self::Qwen(e) => e.scratch_init(gpu),
74            Self::Glm(e) => e.scratch_init(gpu),
75        }
76    }
77
78    /// 2026-09-25: Device pointer to merged row `row` of the packed output. The
79    /// Qwen arm panics unless `forward_batched` or `ensure_scratch` has
80    /// allocated the scratch.
81    pub fn out_row(&self, row: usize) -> DevicePtr {
82        match self {
83            Self::Qwen(e) => e.scratch().buf_out.offset(row * e.out_hidden_size * 2),
84            Self::Glm(e) => e.out_row(row),
85        }
86    }
87}
metrale_model_layers/layers/vision_tower.rs

metrale_model_layers/layers/
vision_tower.rs
