Struct VisionEncoder 

pub struct VisionEncoder {

    pub patch_embed_w: DevicePtr,
    pub patch_embed_b: DevicePtr,
    pub pos_embed: DevicePtr,
    pub blocks: Vec<ViTBlock>,
    pub deepstack: Vec<MergerLayer>,
    pub deepstack_indexes: Vec<usize>,
    pub merger: MergerLayer,
    pub hidden_size: usize,
    pub num_heads: usize,
    pub head_dim: usize,
    pub spatial_merge_size: usize,
    pub out_hidden_size: usize,
    pub intermediate_size: usize,
    pub p_max: usize,
    pub num_grid_per_side: usize,
    /* private fields */
}

Fields

patch_embed_w: DevicePtr

2026-09-25: [hidden_size, PATCH_DIM] BF16.

patch_embed_b: DevicePtr

2026-09-25: [hidden_size] BF16.

pos_embed: DevicePtr

2026-09-25: [num_position_embeddings, hidden_size] BF16. The forward path reads it directly only when METRALE_VISION_POSINTERP=0; otherwise it uses the host copy pos_embed_host_f32.

blocks: Vec<ViTBlock>

deepstack: Vec<MergerLayer>

deepstack_indexes: Vec<usize>

2026-09-25: deepstack[i] runs after the block whose 1-based number is deepstack_indexes[i] (forward_batched compares block_idx + 1).

merger: MergerLayer

2026-09-25: The final merger, after the last block.

hidden_size: usize

num_heads: usize

head_dim: usize

spatial_merge_size: usize

out_hidden_size: usize

intermediate_size: usize

p_max: usize

2026-09-25: Patch-row capacity of every scratch buffer (derive_max_patches).

num_grid_per_side: usize

2026-09-25: sqrt(num_position_embeddings); new refuses a non-square count.

Implementations

impl VisionEncoder

pub fn forward( &self, pixels: &[f32], grid_h: usize, grid_w: usize, gpu: &dyn GpuBackend, stream: u64, ) -> Result<usize>

2026-09-25: Single-image forward through forward_batched. pixels is [P, PATCH_DIM] f32. Returns (1 + deepstack_indexes.len()) * merged_p.

pub fn forward_batched( &self, images: &[(&[f32], usize, usize)], gpu: &dyn GpuBackend, stream: u64, ) -> Result<Vec<(usize, usize, usize)>>

2026-09-25: Batched forward over N images. Ops that do not depend on image geometry (patch_embed and every block’s GEMMs, norms, GELU and residuals) run once over M = Σpᵢ; host pos/rope prep, attention and the mergers loop per image.

buf_out layout, rows of out_hidden_size BF16:

• 0 .. Σmerged_p: the final merger, packed in image order (the rows the LLM reads);
• from (k+1)*Σmerged_p: DeepStack merger k, packed the same way.

When Σp <= p_max this path writes up to (1 + deepstack count) * Σmerged_p rows, which fits the p_max-row buf_out only while 1 + deepstack count <= spatial_merge_size²; nothing checks it. When Σp > p_max it encodes image by image without DeepStack rows.

Returns per-image (post_h, post_w, merged_p) in image order.

impl VisionEncoder

pub fn new( patch_embed_w: DevicePtr, patch_embed_b: DevicePtr, pos_embed: DevicePtr, num_position_embeddings: usize, blocks: Vec<ViTBlock>, deepstack: Vec<MergerLayer>, deepstack_indexes: Vec<usize>, merger: MergerLayer, hidden_size: usize, num_heads: usize, spatial_merge_size: usize, out_hidden_size: usize, intermediate_size: usize, patch_size: usize, max_pixels: Option<usize>, gpu: &dyn GpuBackend, ) -> Result<Self>