Module glm_vit

Module glm_vit 

Source
Expand description

GLM-5.3-Flash vision tower (Glm5NextVisionModel) — 24-block ViT with per-head QK-RMSNorm, pure 2D axial RoPE, clamped SwiGLU, a strided-conv 2×2 downsample and a 4-stage SwiGLU merger.

A SEPARATE encoder from super::VisionEncoder (the Qwen3-VL tower) rather than a set of flags on it. The two share no op: GLM norms are weight-only RMSNorm where Qwen’s are biased LayerNorm, GLM’s attention normalises Q and K per head where Qwen’s does not, GLM’s MLP is 3-matrix clamped SwiGLU where Qwen’s is 2-matrix GELU, and GLM’s spatial merge is a convolution where Qwen’s is a concatenation. Threading if glm through the Qwen path would have put a branch inside every kernel launch in it.

What IS shared: the tensor-core GEMM (gemm::dense_gemm_bf16_pipelined), the f32-out score GEMM, the packed buf_out contract the embedding splice reads, and the (post_h, post_w, merged_p) return shape — so the server side (msg_entry → VisionItem → embed_chunk splice) is reused unchanged.

Reference math and its citations: runs/ws2-vision/REFERENCE.md in the spark-bench repo, derived from transformers glm5_next / glm_ocr.

Re-exports§

pub use init::GlmVitGeometry;

Modules§

init
GlmVit::new and the lazily-allocated device scratch.

Structs§

GlmVit
GlmVitBlock
One ViT block’s weights. Every tensor is BF16 on disk and stays BF16 — the tower is NOT quantised in nvidia/GLM-5.3-Flash-NVFP4 (347/347 tensors are BF16; only the text/MoE stack is NVFP4).
GlmVitMerger
Everything after the last block: the post-norm, the 2×2 conv downsample and the merger MLP whose output is what gets spliced into the LM.
GlmVitScratch
Device scratch, allocated on the FIRST image rather than at load — a text-only GLM serve never pays a byte of it, which matters because the quadratic score matrix dominates the group.

Functions§

block_major_position_ids
Per-token (h_id, w_id) in the reference’s block-major order.
build_axial_rope_tables
(cos, sin) tables of [num_tokens, head_dim] f32, in the layout the rotate-half kernel reads.