spark_model/layers/
vision_tower.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The vision tower a built model holds, whichever family bound it.
4//!
5//! Atlas serves two structurally different towers: the Qwen3-VL-shaped
6//! [`VisionEncoder`] that six loaders bind, and GLM-5.3's [`GlmVit`]. They
7//! agree on exactly three things, and those three are the whole interface the
8//! rest of the engine needs:
9//!
10//! * `forward_batched(images) -> per-image (post_h, post_w, merged_p)`,
11//! * a packed `[Σmerged_p, out_hidden_size]` BF16 output buffer, image-ordered,
12//! * `out_hidden_size` itself.
13//!
14//! So the enum lives here rather than a trait object: two variants, three
15//! methods, no dynamic dispatch on a path that already costs a GEMM, and —
16//! more to the point — no way for a third tower to be added without deciding
17//! out loud what its output contract is.
18
19use anyhow::Result;
20use spark_runtime::gpu::{DevicePtr, GpuBackend};
21
22use super::glm_vit::GlmVit;
23use super::vision_encoder::VisionEncoder;
24
25pub enum VisionTower {
26    /// Qwen3-VL / Qwen3.5 / Qwen3.6 / Qwen4 / LongCat / Mistral.
27    Qwen(Box<VisionEncoder>),
28    /// GLM-5.3-Flash (`glm5_next`).
29    Glm(Box<GlmVit>),
30}
31
32impl VisionTower {
33    pub fn qwen(encoder: VisionEncoder) -> Self {
34        Self::Qwen(Box::new(encoder))
35    }
36
37    pub fn glm(encoder: GlmVit) -> Self {
38        Self::Glm(Box::new(encoder))
39    }
40
41    /// Encode a batch of `(pixels, grid_h, grid_w)` temporal groups. Returns
42    /// `(post_merge_h, post_merge_w, merged_patches)` per group, in order.
43    pub fn forward_batched(
44        &self,
45        images: &[(&[f32], usize, usize)],
46        gpu: &dyn GpuBackend,
47        stream: u64,
48    ) -> Result<Vec<(usize, usize, usize)>> {
49        match self {
50            Self::Qwen(e) => e.forward_batched(images, gpu, stream),
51            Self::Glm(e) => e.forward_batched(images, gpu, stream),
52        }
53    }
54
55    /// Width of one merged embedding row — the LLM's hidden size.
56    pub fn out_hidden_size(&self) -> usize {
57        match self {
58            Self::Qwen(e) => e.out_hidden_size,
59            Self::Glm(e) => e.out_hidden_size,
60        }
61    }
62
63    /// Allocate the encoder's device scratch if it is not allocated yet.
64    ///
65    /// Normally the first image does this. A TENSOR-PARALLEL WORKER never sees
66    /// one — it receives the merged rows over NCCL instead — so it has to ask
67    /// for the buffers explicitly before `out_row` can name a destination.
68    pub fn ensure_scratch(&self, gpu: &dyn GpuBackend) -> Result<()> {
69        match self {
70            Self::Qwen(e) => e.scratch_init(gpu),
71            Self::Glm(e) => e.scratch_init(gpu),
72        }
73    }
74
75    /// Device pointer to merged row `row` of the packed output. Valid only
76    /// after `forward_batched` or `ensure_scratch` has allocated the scratch.
77    pub fn out_row(&self, row: usize) -> DevicePtr {
78        match self {
79            Self::Qwen(e) => e.scratch().buf_out.offset(row * e.out_hidden_size * 2),
80            Self::Glm(e) => e.out_row(row),
81        }
82    }
83}
spark_model/layers/vision_tower.rs

spark_model/layers/
vision_tower.rs
