metrale_model_layers/layers/
vision_encoder.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The Qwen3-VL-shaped vision encoder: a ViT with optional DeepStack mergers.
4//!
5//! Owner: model-layers (vision).
6//! Invariants:
7//! - At most one scratch group is stored (`scratch` is a `OnceLock`); it is
8//!   set on the first image or by `VisionTower::ensure_scratch`.
9//!
10//! Pixels go through patch embedding, the ViT blocks and a final spatial
11//! merger. The first `Σmerged_p` rows of `buf_out` hold the final merger's
12//! output, `[Σmerged_p, out_hidden_size]` BF16, which the LLM embeds. Each
13//! DeepStack merger listed in `deepstack_indexes` writes its own region after
14//! that in the batched path; the LLM does not read those regions.
15
16use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
17
18pub(super) const IMAGE_PAD_TOKEN: u32 = 151_655;
19pub const IMAGE_PAD_TOKEN_ID: u32 = IMAGE_PAD_TOKEN;
20
21/// 2026-09-25: Fallback `<|video_pad|>` id, used when the checkpoint's config declares
22/// none (or 0). The Qwen3-VL, Qwen3.6 and Qwen3.8 checkpoints all put the
23/// video token directly after the image token. A declared id always wins, as
24/// for the image token.
25pub const VIDEO_PAD_TOKEN_ID: u32 = IMAGE_PAD_TOKEN + 1;
26
27/// 2026-09-25: The ViT's per-image scratch buffers, allocated as one group.
28///
29/// Sizes derive only from the encoder's geometry (`p_max` and the head/hidden
30/// dims), not from any image, so the group is built once, lazily, and reused
31/// for every image.
32pub struct VisionScratch {
33    pub buf_f32: DevicePtr,
34    pub buf_h1: DevicePtr,
35    pub buf_h2: DevicePtr,
36    pub buf_wide: DevicePtr,
37    pub buf_merge_in: DevicePtr,
38    pub buf_merge_fc1: DevicePtr,
39    pub buf_out: DevicePtr,
40    pub buf_pos_resampled: DevicePtr,
41    pub buf_rope_cos: DevicePtr,
42    pub buf_rope_sin: DevicePtr,
43    pub buf_qr: DevicePtr,
44    pub buf_kr: DevicePtr,
45    pub buf_vt: DevicePtr,
46    pub buf_scores: DevicePtr,
47    pub buf_probs: DevicePtr,
48    pub buf_o_stage: DevicePtr,
49}
50
51/// 2026-09-25: Flattened per-patch pixel dimension `C × temporal_patch_size × patch_size²`
52/// = 3 × 2 × 16 × 16 for this ViT. It is fixed in the encoder, not read from
53/// config: `buf_f32` is allocated at `p_max × PATCH_DIM × 4` and the
54/// patch-embed GEMM is issued with `K = PATCH_DIM`.
55///
56/// The host preprocessor (`vision_preprocess::preprocess_image`) sizes the
57/// pixel buffer from `vision_config`, so a checkpoint with a different
58/// `patch_size`/`temporal_patch_size` produces a buffer of a different
59/// length. `patch_embed`'s `check_pixel_len` refuses such a buffer.
60pub(crate) const PATCH_DIM: usize = 1536;
61
62pub struct ViTBlock {
63    pub norm1_w: DevicePtr,
64    pub norm1_b: DevicePtr,
65    pub qkv_w: DevicePtr,
66    pub qkv_b: DevicePtr,
67    pub proj_w: DevicePtr,
68    pub proj_b: DevicePtr,
69    pub norm2_w: DevicePtr,
70    pub norm2_b: DevicePtr,
71    pub fc1_w: DevicePtr,
72    pub fc1_b: DevicePtr,
73    pub fc2_w: DevicePtr,
74    pub fc2_b: DevicePtr,
75}
76
77pub struct MergerLayer {
78    pub norm_w: DevicePtr,
79    pub norm_b: DevicePtr,
80    pub fc1_w: DevicePtr,
81    pub fc1_b: DevicePtr,
82    pub fc2_w: DevicePtr,
83    pub fc2_b: DevicePtr,
84}
85
86pub struct VisionEncoder {
87    /// 2026-09-25: `[hidden_size, PATCH_DIM]` BF16.
88    pub patch_embed_w: DevicePtr,
89    /// 2026-09-25: `[hidden_size]` BF16.
90    pub patch_embed_b: DevicePtr,
91    /// 2026-09-25: `[num_position_embeddings, hidden_size]` BF16. The forward path
92    /// reads it directly only when `METRALE_VISION_POSINTERP=0`; otherwise
93    /// it uses the host copy `pos_embed_host_f32`.
94    pub pos_embed: DevicePtr,
95    pub blocks: Vec<ViTBlock>,
96    pub deepstack: Vec<MergerLayer>,
97    /// 2026-09-25: `deepstack[i]` runs after the block whose 1-based number is
98    /// `deepstack_indexes[i]` (`forward_batched` compares `block_idx + 1`).
99    pub deepstack_indexes: Vec<usize>,
100    /// 2026-09-25: The final merger, after the last block.
101    pub merger: MergerLayer,
102    k_gemm: KernelHandle,
103    // 2026-09-25: `k_gemm_pipelined`, `k_add_bias`, `k_rope_deint`,
104    // `k_softmax`, `k_scatter_head` and `k_gemm_f32` are bound with
105    // `try_kernel` and may be null. `vit_gemm_bias` uses `k_gemm` unless both
106    // `k_gemm_pipelined` and `k_add_bias` are set; the ViT attention uses
107    // `k_attn` when `k_rope_deint` is null.
108    k_gemm_pipelined: KernelHandle,
109    k_add_bias: KernelHandle,
110    k_norm: KernelHandle,
111    k_add: KernelHandle,
112    k_gelu: KernelHandle,
113    k_attn: KernelHandle,
114    k_rope_deint: KernelHandle,
115    k_softmax: KernelHandle,
116    k_scatter_head: KernelHandle,
117    k_gemm_f32: KernelHandle,
118    k_merge: KernelHandle,
119    k_f32_bf16: KernelHandle,
120    k_copy: KernelHandle,
121    pub hidden_size: usize,
122    pub num_heads: usize,
123    pub head_dim: usize,
124    pub spatial_merge_size: usize,
125    pub out_hidden_size: usize,
126    pub intermediate_size: usize,
127    /// 2026-09-25: Patch-row capacity of every scratch buffer (`derive_max_patches`).
128    pub p_max: usize,
129    /// 2026-09-25: `sqrt(num_position_embeddings)`; `new` refuses a non-square count.
130    pub num_grid_per_side: usize,
131    /// 2026-09-25: ViT scratch, allocated on the first image rather than at load, so
132    /// a serve that sees no image never allocates it.
133    ///
134    /// A `OnceLock`, because the forward path takes `&self`; if two callers
135    /// race, one group is kept.
136    scratch: std::sync::OnceLock<VisionScratch>,
137    /// 2026-09-25: `[num_position_embeddings × hidden_size]` row-major.
138    pos_embed_host_f32: Vec<f32>,
139    /// 2026-09-25: `[head_dim / 4]` RoPE inverse frequencies.
140    rope_inv_freq: Vec<f32>,
141}
142
143pub(crate) mod enc_impl;
metrale_model_layers/layers/vision_encoder.rs

metrale_model_layers/layers/
vision_encoder.rs
