metrale_model_layers/layers/glm_vit/
init.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `GlmVit::new` and the device scratch allocated on first use.
4//!
5//! Owner: model-layers (vision).
6//! Invariants: `new` refuses a `hidden_size` not divisible by a nonzero
7//! `num_heads`, a `head_dim` not a multiple of 4, a `spatial_merge_size` other
8//! than 2, and a zero `patch_size` or `temporal_patch_size`.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12
13use super::{GlmVit, GlmVitBlock, GlmVitMerger, GlmVitScratch};
14use crate::layers::vision_encoder::enc_impl::init::derive_max_patches;
15
16/// 2026-09-25: Everything `GlmVit::new` needs that is not a device pointer.
17#[derive(Debug, Clone)]
18pub struct GlmVitGeometry {
19    pub hidden_size: usize,
20    pub num_heads: usize,
21    pub intermediate_size: usize,
22    pub spatial_merge_size: usize,
23    pub out_hidden_size: usize,
24    pub projection_intermediate_size: usize,
25    pub patch_size: usize,
26    pub temporal_patch_size: usize,
27    pub rms_norm_eps: f32,
28    pub swiglu_limit: f32,
29    pub rope_theta: f32,
30    pub max_pixels: Option<usize>,
31}
32
33impl GlmVit {
34    pub fn new(
35        patch_embed_w: DevicePtr,
36        patch_embed_b: DevicePtr,
37        blocks: Vec<GlmVitBlock>,
38        merger: GlmVitMerger,
39        geo: &GlmVitGeometry,
40        gpu: &dyn GpuBackend,
41    ) -> Result<Self> {
42        anyhow::ensure!(
43            geo.num_heads > 0 && geo.hidden_size.is_multiple_of(geo.num_heads),
44            "glm_vit: hidden_size {} is not divisible by num_heads {}",
45            geo.hidden_size,
46            geo.num_heads
47        );
48        let head_dim = geo.hidden_size / geo.num_heads;
49        anyhow::ensure!(
50            head_dim.is_multiple_of(4),
51            "glm_vit: head_dim {head_dim} must be a multiple of 4 — the axial RoPE splits it \
52             into an h half and a w half, each of which is then rotate-half'd"
53        );
54        anyhow::ensure!(
55            geo.spatial_merge_size == 2,
56            "glm_vit: spatial_merge_size {} — the 2x2 conv downsample and its im2col are \
57             written for 2 and only 2; a different merge needs a different kernel, not a \
58             different argument",
59            geo.spatial_merge_size
60        );
61        anyhow::ensure!(
62            geo.patch_size > 0 && geo.temporal_patch_size > 0,
63            "glm_vit: patch_size/temporal_patch_size must be non-zero"
64        );
65
66        let (p_max, asked_for) = derive_max_patches(geo.max_pixels, geo.patch_size);
67        if let Some(wanted) = asked_for {
68            tracing::warn!(
69                "GLM vision encoder capacity clamped to {p_max} patches; the resolved area \
70                 bound wanted {wanted}. The ViT materialises a full [seq, seq] score matrix, \
71                 so scratch is O(patches^2). Lower --vision-max-pixels to reclaim memory."
72            );
73        }
74
75        // 2026-09-25: `inv_freq[k] = theta^(-2k/spatial_dim)` for
76        // `k < spatial_dim/2`, i.e. `1 / theta ** (arange(0, spatial_dim, 2) / spatial_dim)`.
77        let spatial_dim = head_dim / 2;
78        let rope_inv_freq: Vec<f32> = (0..spatial_dim / 2)
79            .map(|k| 1.0 / geo.rope_theta.powf(2.0 * k as f32 / spatial_dim as f32))
80            .collect();
81
82        Ok(Self {
83            patch_embed_w,
84            patch_embed_b,
85            blocks,
86            merger,
87            k_gemm: gpu.kernel("gemm", "dense_gemm_bf16_pipelined")?,
88            k_gemm_f32: gpu.kernel("gemm", "dense_gemm_bf16_f32out")?,
89            k_add_bias: gpu.kernel("glm_vit", "glm_vit_add_bias")?,
90            k_rmsnorm: gpu.kernel("glm_vit", "glm_vit_rmsnorm")?,
91            k_layernorm: gpu.kernel("glm_vit", "glm_vit_layernorm")?,
92            k_gelu: gpu.kernel("glm_vit", "glm_vit_gelu_erf")?,
93            k_swiglu: gpu.kernel("glm_vit", "glm_vit_swiglu_clamp")?,
94            k_qknorm_rope: gpu.kernel("glm_vit", "glm_vit_qknorm_rope_deint")?,
95            k_softmax: gpu.kernel("glm_vit", "glm_vit_softmax_rows")?,
96            k_scatter_head: gpu.kernel("glm_vit", "glm_vit_scatter_head")?,
97            k_im2col: gpu.kernel("glm_vit", "glm_vit_im2col_2x2")?,
98            k_copy: gpu.kernel("glm_vit", "glm_vit_copy")?,
99            k_add: gpu.kernel("glm_vit", "glm_vit_add_inplace")?,
100            k_f32_bf16: gpu.kernel("glm_vit", "glm_vit_f32_to_bf16")?,
101            hidden_size: geo.hidden_size,
102            num_heads: geo.num_heads,
103            head_dim,
104            intermediate_size: geo.intermediate_size,
105            spatial_merge_size: geo.spatial_merge_size,
106            out_hidden_size: geo.out_hidden_size,
107            projection_intermediate_size: geo.projection_intermediate_size,
108            patch_dim: 3 * geo.temporal_patch_size * geo.patch_size * geo.patch_size,
109            rms_norm_eps: geo.rms_norm_eps,
110            swiglu_limit: geo.swiglu_limit,
111            p_max,
112            rope_inv_freq,
113            scratch: std::sync::OnceLock::new(),
114        })
115    }
116
117    /// 2026-09-25: Merged (post-downsample) row capacity of `buf_out`.
118    pub(crate) fn mp_max(&self) -> usize {
119        let sms2 = self.spatial_merge_size * self.spatial_merge_size;
120        self.p_max.div_ceil(sms2)
121    }
122
123    fn build_scratch(&self, gpu: &dyn GpuBackend) -> Result<GlmVitScratch> {
124        let p = self.p_max;
125        let mp = self.mp_max();
126        let h = self.hidden_size;
127        let inter = self.intermediate_size;
128        let out = self.out_hidden_size;
129        let pi = self.projection_intermediate_size;
130        // 2026-09-25: `buf_gate` holds the QKV output (`3 * hidden`), the MLP
131        // gate (`intermediate_size`) and the BF16 pixels (`patch_dim`), so it
132        // is sized for the widest.
133        let wide = inter.max(3 * h).max(self.patch_dim);
134        Ok(GlmVitScratch {
135            buf_f32: gpu.alloc(p * self.patch_dim * 4)?,
136            buf_h1: gpu.alloc(p * h * 2)?,
137            buf_h2: gpu.alloc(p * h * 2)?,
138            buf_gate: gpu.alloc(p * wide * 2)?,
139            buf_up: gpu.alloc(p * inter * 2)?,
140            buf_qr: gpu.alloc(p * h * 2)?,
141            buf_kr: gpu.alloc(p * h * 2)?,
142            buf_vt: gpu.alloc(p * h * 2)?,
143            // 2026-09-25: `p²` f32 scores and BF16 probs: 164 MB + 82 MB at
144            // 6400 patches.
145            buf_scores: gpu.alloc(p * p * 4)?,
146            buf_probs: gpu.alloc(p * p * 2)?,
147            buf_o_stage: gpu.alloc(p * self.head_dim * 2)?,
148            buf_rope_cos: gpu.alloc(p * self.head_dim * 2)?,
149            buf_rope_sin: gpu.alloc(p * self.head_dim * 2)?,
150            buf_merge_a: gpu.alloc(mp * out * 2)?,
151            buf_merge_b: gpu.alloc(mp * out * 2)?,
152            buf_merge_g: gpu.alloc(mp * pi * 2)?,
153            buf_merge_u: gpu.alloc(mp * pi * 2)?,
154            buf_out: gpu.alloc(mp * out * 2)?,
155        })
156    }
157
158    /// 2026-09-25: Allocate the scratch if it is not yet set. Called by
159    /// `forward_batched` and `VisionTower::ensure_scratch`.
160    pub(crate) fn scratch_init(&self, gpu: &dyn GpuBackend) -> Result<()> {
161        if self.scratch.get().is_none() {
162            let s = self.build_scratch(gpu)?;
163            let _ = self.scratch.set(s);
164            tracing::info!(
165                "GLM vision scratch allocated on first image: {} patches ({} merged rows)",
166                self.p_max,
167                self.mp_max()
168            );
169        }
170        Ok(())
171    }
172
173    pub(crate) fn scratch(&self) -> &GlmVitScratch {
174        self.scratch
175            .get()
176            .expect("glm_vit scratch: encode entry must call scratch_init(gpu) first")
177    }
178
179    /// 2026-09-25: Device pointer to row `row` of the packed merged output
180    /// (`buf_out`). Panics when the scratch is not allocated.
181    pub fn out_row(&self, row: usize) -> DevicePtr {
182        self.scratch()
183            .buf_out
184            .offset(row * self.out_hidden_size * 2)
185    }
186}
187
188#[cfg(test)]
189mod tests {
190
191    /// 2026-09-25: 3 × 2 × 14² = 1176 floats per patch, which is not the Qwen
192    /// encoder's `PATCH_DIM`.
193    #[test]
194    fn glm_patch_dim_is_a_runtime_value_of_1176() {
195        let patch_dim = 3 * 2 * 14 * 14;
196        assert_eq!(patch_dim, 1176);
197        assert_ne!(patch_dim, crate::layers::vision_encoder::PATCH_DIM);
198    }
199
200    /// 2026-09-25: The rope frequency table: 16 entries for head_dim 64, first
201    /// entry exactly 1, last `theta^(-30/32)`.
202    #[test]
203    fn rope_inv_freq_matches_the_reference_formula() {
204        let head_dim = 64usize;
205        let spatial_dim = head_dim / 2;
206        let theta = 10_000.0f32;
207        let f: Vec<f32> = (0..spatial_dim / 2)
208            .map(|k| 1.0 / theta.powf(2.0 * k as f32 / spatial_dim as f32))
209            .collect();
210        assert_eq!(f.len(), 16);
211        assert_eq!(f[0], 1.0);
212        let last = 1.0 / theta.powf(30.0 / 32.0);
213        assert!((f[15] - last).abs() < 1e-9);
214    }
215}
metrale_model_layers/layers/glm_vit/init.rs

metrale_model_layers/layers/glm_vit/
init.rs
