spark_model/layers/
glm_vit.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GLM-5.3-Flash vision tower (`Glm5NextVisionModel`) — 24-block ViT with
4//! per-head QK-RMSNorm, pure 2D axial RoPE, clamped SwiGLU, a strided-conv 2×2
5//! downsample and a 4-stage SwiGLU merger.
6//!
7//! A SEPARATE encoder from [`super::VisionEncoder`] (the Qwen3-VL tower) rather
8//! than a set of flags on it. The two share no op: GLM norms are weight-only
9//! RMSNorm where Qwen's are biased LayerNorm, GLM's attention normalises Q and
10//! K per head where Qwen's does not, GLM's MLP is 3-matrix clamped SwiGLU where
11//! Qwen's is 2-matrix GELU, and GLM's spatial merge is a convolution where
12//! Qwen's is a concatenation. Threading `if glm` through the Qwen path would
13//! have put a branch inside every kernel launch in it.
14//!
15//! What IS shared: the tensor-core GEMM (`gemm::dense_gemm_bf16_pipelined`),
16//! the f32-out score GEMM, the packed `buf_out` contract the embedding splice
17//! reads, and the `(post_h, post_w, merged_p)` return shape — so the server
18//! side (`msg_entry` → `VisionItem` → `embed_chunk` splice) is reused unchanged.
19//!
20//! Reference math and its citations: `runs/ws2-vision/REFERENCE.md` in the
21//! spark-bench repo, derived from `transformers` `glm5_next` / `glm_ocr`.
22
23use spark_runtime::gpu::{DevicePtr, KernelHandle};
24
25/// One ViT block's weights. Every tensor is BF16 on disk and stays BF16 — the
26/// tower is NOT quantised in `nvidia/GLM-5.3-Flash-NVFP4` (347/347 tensors are
27/// BF16; only the text/MoE stack is NVFP4).
28pub struct GlmVitBlock {
29    pub norm1_w: DevicePtr,
30    pub qkv_w: DevicePtr,
31    pub qkv_b: DevicePtr,
32    /// Per-head RMSNorm on Q, `[head_dim]`. Applied after the QKV split and
33    /// BEFORE the rotary transform — the order is the reference's.
34    pub q_norm_w: DevicePtr,
35    pub k_norm_w: DevicePtr,
36    pub proj_w: DevicePtr,
37    pub proj_b: DevicePtr,
38    pub norm2_w: DevicePtr,
39    pub gate_w: DevicePtr,
40    pub gate_b: DevicePtr,
41    pub up_w: DevicePtr,
42    pub up_b: DevicePtr,
43    pub down_w: DevicePtr,
44    pub down_b: DevicePtr,
45}
46
47/// Everything after the last block: the post-norm, the 2×2 conv downsample and
48/// the merger MLP whose output is what gets spliced into the LM.
49pub struct GlmVitMerger {
50    /// Weight-only RMSNorm over `hidden_size`, applied before the downsample.
51    pub post_layernorm_w: DevicePtr,
52    /// `Conv2d(hidden, out_hidden, k=2, s=2)` flattened to `[out_hidden,
53    /// hidden*4]` in `(in_channel, kh, kw)` order — the conv's own layout, which
54    /// is what makes the im2col below a plain reshape.
55    pub downsample_w: DevicePtr,
56    pub downsample_b: DevicePtr,
57    pub proj_w: DevicePtr,
58    /// `post_projection_norm` is a real mean-subtracting LayerNorm with a bias
59    /// — the ONLY such norm in this tower. Every other one is weight-only
60    /// RMSNorm. The checkpoint census settles it: this is the only norm that
61    /// ships a `.bias`.
62    pub norm_w: DevicePtr,
63    pub norm_b: DevicePtr,
64    pub gate_w: DevicePtr,
65    pub up_w: DevicePtr,
66    pub down_w: DevicePtr,
67}
68
69/// Device scratch, allocated on the FIRST image rather than at load — a
70/// text-only GLM serve never pays a byte of it, which matters because the
71/// quadratic score matrix dominates the group.
72pub struct GlmVitScratch {
73    pub(crate) buf_f32: DevicePtr,
74    pub(crate) buf_h1: DevicePtr,
75    pub(crate) buf_h2: DevicePtr,
76    pub(crate) buf_gate: DevicePtr,
77    pub(crate) buf_up: DevicePtr,
78    pub(crate) buf_qr: DevicePtr,
79    pub(crate) buf_kr: DevicePtr,
80    pub(crate) buf_vt: DevicePtr,
81    pub(crate) buf_scores: DevicePtr,
82    pub(crate) buf_probs: DevicePtr,
83    pub(crate) buf_o_stage: DevicePtr,
84    pub(crate) buf_rope_cos: DevicePtr,
85    pub(crate) buf_rope_sin: DevicePtr,
86    pub(crate) buf_merge_a: DevicePtr,
87    pub(crate) buf_merge_b: DevicePtr,
88    pub(crate) buf_merge_g: DevicePtr,
89    pub(crate) buf_merge_u: DevicePtr,
90    pub buf_out: DevicePtr,
91}
92
93pub struct GlmVit {
94    pub(crate) patch_embed_w: DevicePtr,
95    pub(crate) patch_embed_b: DevicePtr,
96    pub(crate) blocks: Vec<GlmVitBlock>,
97    pub(crate) merger: GlmVitMerger,
98
99    // Shared kernels (module `gemm`).
100    pub(crate) k_gemm: KernelHandle,
101    pub(crate) k_gemm_f32: KernelHandle,
102    // GLM-specific kernels (module `glm_vit`).
103    pub(crate) k_add_bias: KernelHandle,
104    pub(crate) k_rmsnorm: KernelHandle,
105    pub(crate) k_layernorm: KernelHandle,
106    pub(crate) k_gelu: KernelHandle,
107    pub(crate) k_swiglu: KernelHandle,
108    pub(crate) k_qknorm_rope: KernelHandle,
109    pub(crate) k_softmax: KernelHandle,
110    pub(crate) k_scatter_head: KernelHandle,
111    pub(crate) k_im2col: KernelHandle,
112    pub(crate) k_copy: KernelHandle,
113    pub(crate) k_add: KernelHandle,
114    pub(crate) k_f32_bf16: KernelHandle,
115
116    // Geometry.
117    pub(crate) hidden_size: usize,
118    pub(crate) num_heads: usize,
119    pub(crate) head_dim: usize,
120    pub(crate) intermediate_size: usize,
121    pub(crate) spatial_merge_size: usize,
122    pub out_hidden_size: usize,
123    pub(crate) projection_intermediate_size: usize,
124    /// `C × temporal_patch_size × patch_size²` — 1176 for GLM (3×2×14×14).
125    /// A RUNTIME value, not the compiled `PATCH_DIM` constant the Qwen encoder
126    /// carries: that constant is 1536 and is exactly why the Qwen tower cannot
127    /// bind this checkpoint.
128    pub(crate) patch_dim: usize,
129    pub(crate) rms_norm_eps: f32,
130    pub(crate) swiglu_limit: f32,
131    /// Patch rows every buffer is sized for.
132    pub(crate) p_max: usize,
133    /// `inv_freq[k] = theta^(-2k/spatial_dim)`, `spatial_dim = head_dim/2`.
134    /// Half the entries feed the h axis and half the w axis.
135    pub(crate) rope_inv_freq: Vec<f32>,
136
137    pub(crate) scratch: std::sync::OnceLock<GlmVitScratch>,
138}
139
140mod block;
141mod forward;
142pub mod init;
143mod merge;
144mod rope;
145
146pub use init::GlmVitGeometry;
147pub use rope::{block_major_position_ids, build_axial_rope_tables};
spark_model/layers/glm_vit.rs

spark_model/layers/
glm_vit.rs
