spark_model/layers/
vision_encoder.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Qwen3-VL vision encoder: 27-block ViT + DeepStack mergers.
4//!
5//! Processes patch embeddings (BF16) through a ViT backbone, extracts
6//! intermediate hidden states at deepstack indices [8, 16, 24, 27], applies
7//! 2×2 spatial merges + 2-layer MLPs, and concatenates the four outputs.
8//! Result: [num_patches, out_hidden_size=2048] BF16 ready for LLM embedding.
9
10use spark_runtime::gpu::{DevicePtr, KernelHandle};
11
12pub(super) const IMAGE_PAD_TOKEN: u32 = 151_655;
13pub const IMAGE_PAD_TOKEN_ID: u32 = IMAGE_PAD_TOKEN;
14
15/// Fallback `<|video_pad|>` id, used when the checkpoint's config declares
16/// none. Qwen3-VL's video token sits directly after its image token, and the
17/// same holds for Qwen3.6/3.8 (248056 / 248057) — but a checkpoint that
18/// declares its own always wins, exactly as for the image token.
19pub const VIDEO_PAD_TOKEN_ID: u32 = IMAGE_PAD_TOKEN + 1;
20
21/// The ViT's per-image scratch buffers, allocated as one group.
22///
23/// Sizes derive only from the encoder's geometry (`p_max` and the head/hidden
24/// dims), so nothing here depends on the image itself — which is why it can be
25/// built once, lazily, and reused for every image after.
26pub struct VisionScratch {
27    pub buf_f32: DevicePtr,
28    pub buf_h1: DevicePtr,
29    pub buf_h2: DevicePtr,
30    pub buf_wide: DevicePtr,
31    pub buf_merge_in: DevicePtr,
32    pub buf_merge_fc1: DevicePtr,
33    pub buf_out: DevicePtr,
34    pub buf_pos_resampled: DevicePtr,
35    pub buf_rope_cos: DevicePtr,
36    pub buf_rope_sin: DevicePtr,
37    pub buf_qr: DevicePtr,
38    pub buf_kr: DevicePtr,
39    pub buf_vt: DevicePtr,
40    pub buf_scores: DevicePtr,
41    pub buf_probs: DevicePtr,
42    pub buf_o_stage: DevicePtr,
43}
44
45/// Flattened per-patch pixel dimension `C × temporal_patch_size × patch_size²`
46/// = 3 × 2 × 16 × 16 for this ViT. It is baked into the encoder, not read from
47/// config: `buf_f32` is allocated at `p_max × PATCH_DIM × 4` and the
48/// patch-embed GEMM is issued with `K = PATCH_DIM`.
49///
50/// The host side computes the same quantity from `vision_config`
51/// (`vision_preprocess::preprocess_image`), so a checkpoint declaring a
52/// different `patch_size`/`temporal_patch_size` produces a pixel buffer of a
53/// DIFFERENT length. Every use of this constant that touches a host slice must
54/// therefore check the length rather than assume it — see `patch_embed`.
55pub(crate) const PATCH_DIM: usize = 1536;
56
57pub struct ViTBlock {
58    pub norm1_w: DevicePtr,
59    pub norm1_b: DevicePtr,
60    pub qkv_w: DevicePtr,
61    pub qkv_b: DevicePtr,
62    pub proj_w: DevicePtr,
63    pub proj_b: DevicePtr,
64    pub norm2_w: DevicePtr,
65    pub norm2_b: DevicePtr,
66    pub fc1_w: DevicePtr,
67    pub fc1_b: DevicePtr,
68    pub fc2_w: DevicePtr,
69    pub fc2_b: DevicePtr,
70}
71
72pub struct MergerLayer {
73    pub norm_w: DevicePtr,
74    pub norm_b: DevicePtr,
75    pub fc1_w: DevicePtr,
76    pub fc1_b: DevicePtr,
77    pub fc2_w: DevicePtr,
78    pub fc2_b: DevicePtr,
79}
80
81pub struct VisionEncoder {
82    pub patch_embed_w: DevicePtr,      // [1152, 1536] BF16
83    pub patch_embed_b: DevicePtr,      // [1152] BF16
84    pub pos_embed: DevicePtr,          // [2304, 1152] BF16 (untouched, kept for reference)
85    pub blocks: Vec<ViTBlock>,         // 27 blocks
86    pub deepstack: Vec<MergerLayer>,   // 3 deepstack mergers
87    pub deepstack_indexes: Vec<usize>, // [8, 16, 24] (1-indexed, after Nth block)
88    pub merger: MergerLayer,           // final merger (after block 27)
89    // kernel handles
90    k_gemm: KernelHandle, // vision_gemm_bias: C[M,N] = A[M,K]@B[N,K]^T + bias
91    k_gemm_pipelined: KernelHandle, // dense_gemm_bf16_pipelined (tensor-core, ~40×; no bias)
92    k_add_bias: KernelHandle, // vision_add_bias: C += bias[n] (fuses bias for the TC path)
93    k_norm: KernelHandle, // vision_layer_norm (biased, in-place)
94    k_add: KernelHandle,  // vision_add_inplace
95    k_gelu: KernelHandle, // vision_gelu (in-place)
96    k_attn: KernelHandle, // vision_attention_rope (legacy SDPA — ATLAS_VISION_ATTN_LEGACY=1)
97    k_rope_deint: KernelHandle, // vit_rope_deinterleave (rope + head-contig Qr/Kr + V transpose)
98    k_softmax: KernelHandle, // vit_softmax_rows (parallel row softmax)
99    k_scatter_head: KernelHandle, // vit_scatter_head (contig → interleaved O slot)
100    k_gemm_f32: KernelHandle, // dense_gemm_bf16_f32out (raw QKᵀ scores, f32 out)
101    k_merge: KernelHandle, // vision_spatial_merge (2×2)
102    k_f32_bf16: KernelHandle, // vision_f32_to_bf16
103    k_copy: KernelHandle, // vision_bf16_copy
104    // config
105    pub hidden_size: usize,        // 1152
106    pub num_heads: usize,          // 16
107    pub head_dim: usize,           // 72
108    pub spatial_merge_size: usize, // 2
109    pub out_hidden_size: usize,    // 2048
110    pub intermediate_size: usize,  // 4304
111    pub p_max: usize,              // 6400 (80×80 patches for 1280×1280 image)
112    // num_grid_per_side = sqrt(num_position_embeddings) = 48 for Qwen3-VL/3.6.
113    pub num_grid_per_side: usize,
114    /// ViT scratch, allocated on the FIRST IMAGE rather than at load.
115    ///
116    /// ~2.2 GB at the 16384-patch rung on Qwen3.8-27B — the fourth-largest
117    /// consumer in the process — and a text-only serve never touches a byte of
118    /// it. Deferring hands that back to the KV budget on every text workload
119    /// while costing an image request one allocation it used to pay at boot.
120    ///
121    /// `OnceLock` rather than a flag: the encoder's forward path takes `&self`,
122    /// and the buffers must be filled exactly once even if two images race.
123    scratch: std::sync::OnceLock<VisionScratch>,
124    // host-side prep state
125    pos_embed_host_f32: Vec<f32>, // [num_position_embeddings × hidden_size] row-major
126    rope_inv_freq: Vec<f32>,      // [head_dim / 4] frequencies
127}
128
129mod enc_impl;
spark_model/layers/vision_encoder.rs

spark_model/layers/
vision_encoder.rs
