metrale_model_layers/layers/
glm_vit.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The GLM-5.3-Flash vision tower: ViT blocks with per-head
4//! QK-RMSNorm, 2D axial RoPE and clamped SwiGLU, then a 2×2 strided-conv
5//! downsample and a SwiGLU merger.
6//!
7//! A separate type from [`super::VisionEncoder`] (the Qwen3-VL tower), with
8//! its own kernel module (`glm_vit`); both use the `gemm` module's
9//! `dense_gemm_bf16_pipelined` and `dense_gemm_bf16_f32out`. Callers reach
10//! either through [`super::VisionTower`], which gives both the same
11//! `forward_batched` result and `out_row` contract.
12//!
13//! Owner: model-layers (vision).
14//! Invariants:
15//! - Assumes every weight pointer is a BF16 tensor. `GlmVit::new` does not
16//!   check; the loader (`load_glm5_next_vision`) refuses other dtypes.
17//! - The scratch is set at most once, for `p_max` patch rows.
18
19use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
20
21/// 2026-09-25: One ViT block's BF16 weights.
22pub struct GlmVitBlock {
23    pub norm1_w: DevicePtr,
24    pub qkv_w: DevicePtr,
25    pub qkv_b: DevicePtr,
26    /// 2026-09-25: Per-head RMSNorm weight for Q, `[head_dim]`. The kernel applies
27    /// it after the QKV split and before the rotary transform.
28    pub q_norm_w: DevicePtr,
29    pub k_norm_w: DevicePtr,
30    pub proj_w: DevicePtr,
31    pub proj_b: DevicePtr,
32    pub norm2_w: DevicePtr,
33    pub gate_w: DevicePtr,
34    pub gate_b: DevicePtr,
35    pub up_w: DevicePtr,
36    pub up_b: DevicePtr,
37    pub down_w: DevicePtr,
38    pub down_b: DevicePtr,
39}
40
41/// 2026-09-25: Everything after the last block: the post-norm, the 2×2 conv
42/// downsample and the merger MLP whose output fills `buf_out`.
43pub struct GlmVitMerger {
44    /// 2026-09-25: Weight-only RMSNorm over `hidden_size`, applied before the downsample.
45    pub post_layernorm_w: DevicePtr,
46    /// 2026-09-25: The 2×2 stride-2 conv weight read as `[out_hidden, hidden*4]`
47    /// in `(in_channel, kh, kw)` order, the order `glm_vit_im2col_2x2` writes.
48    pub downsample_w: DevicePtr,
49    pub downsample_b: DevicePtr,
50    pub proj_w: DevicePtr,
51    /// 2026-09-25: `post_projection_norm`, the tower's only LayerNorm (mean
52    /// subtracted, with bias); every other norm is weight-only RMSNorm.
53    pub norm_w: DevicePtr,
54    pub norm_b: DevicePtr,
55    pub gate_w: DevicePtr,
56    pub up_w: DevicePtr,
57    pub down_w: DevicePtr,
58}
59
60/// 2026-09-25: Device scratch, allocated by `scratch_init` on first use rather
61/// than at load. `buf_scores` and `buf_probs` hold `p_max²` elements.
62pub struct GlmVitScratch {
63    pub(crate) buf_f32: DevicePtr,
64    pub(crate) buf_h1: DevicePtr,
65    pub(crate) buf_h2: DevicePtr,
66    pub(crate) buf_gate: DevicePtr,
67    pub(crate) buf_up: DevicePtr,
68    pub(crate) buf_qr: DevicePtr,
69    pub(crate) buf_kr: DevicePtr,
70    pub(crate) buf_vt: DevicePtr,
71    pub(crate) buf_scores: DevicePtr,
72    pub(crate) buf_probs: DevicePtr,
73    pub(crate) buf_o_stage: DevicePtr,
74    pub(crate) buf_rope_cos: DevicePtr,
75    pub(crate) buf_rope_sin: DevicePtr,
76    pub(crate) buf_merge_a: DevicePtr,
77    pub(crate) buf_merge_b: DevicePtr,
78    pub(crate) buf_merge_g: DevicePtr,
79    pub(crate) buf_merge_u: DevicePtr,
80    pub buf_out: DevicePtr,
81}
82
83pub struct GlmVit {
84    pub(crate) patch_embed_w: DevicePtr,
85    pub(crate) patch_embed_b: DevicePtr,
86    pub(crate) blocks: Vec<GlmVitBlock>,
87    pub(crate) merger: GlmVitMerger,
88
89    pub(crate) k_gemm: KernelHandle,
90    pub(crate) k_gemm_f32: KernelHandle,
91    pub(crate) k_add_bias: KernelHandle,
92    pub(crate) k_rmsnorm: KernelHandle,
93    pub(crate) k_layernorm: KernelHandle,
94    pub(crate) k_gelu: KernelHandle,
95    pub(crate) k_swiglu: KernelHandle,
96    pub(crate) k_qknorm_rope: KernelHandle,
97    pub(crate) k_softmax: KernelHandle,
98    pub(crate) k_scatter_head: KernelHandle,
99    pub(crate) k_im2col: KernelHandle,
100    pub(crate) k_copy: KernelHandle,
101    pub(crate) k_add: KernelHandle,
102    pub(crate) k_f32_bf16: KernelHandle,
103
104    pub(crate) hidden_size: usize,
105    pub(crate) num_heads: usize,
106    pub(crate) head_dim: usize,
107    pub(crate) intermediate_size: usize,
108    pub(crate) spatial_merge_size: usize,
109    pub out_hidden_size: usize,
110    pub(crate) projection_intermediate_size: usize,
111    /// 2026-09-25: `3 × temporal_patch_size × patch_size²` floats per patch,
112    /// from the config; the Qwen encoder's `PATCH_DIM` is a constant 1536.
113    pub(crate) patch_dim: usize,
114    pub(crate) rms_norm_eps: f32,
115    pub(crate) swiglu_limit: f32,
116    /// 2026-09-25: Patch-row capacity of the scratch: the per-patch buffers hold
117    /// `p_max` rows, the merger buffers `mp_max()` merged rows.
118    pub(crate) p_max: usize,
119    /// 2026-09-25: `inv_freq[k] = theta^(-2k/spatial_dim)` for `k < spatial_dim/2`,
120    /// `spatial_dim = head_dim/2`. `rope::build_axial_rope_tables` uses every
121    /// entry once for the h axis and once for the w axis.
122    pub(crate) rope_inv_freq: Vec<f32>,
123
124    pub(crate) scratch: std::sync::OnceLock<GlmVitScratch>,
125}
126
127mod block;
128mod forward;
129pub mod init;
130mod merge;
131mod rope;
132
133pub use init::GlmVitGeometry;
134pub use rope::{block_major_position_ids, build_axial_rope_tables};
metrale_model_layers/layers/glm_vit.rs

metrale_model_layers/layers/
glm_vit.rs
