metrale_model_engine/traits/model/
vision.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelVision`, one of the supertraits `Model` is made of. Its methods, default
4//! bodies and docs are the ones `Model` declared before the split.
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10
11/// 2026-09-26: Vision-encoder input for prompts with images.
12pub trait ModelVision {
13    /// 2026-09-25: Encode images with the vision encoder and keep the embeddings for the next
14    /// prefill. Call it before `prefill_chunk` for a prompt with image pad tokens. Default:
15    /// `Ok(())`.
16    fn prepare_vision_embed(&self, _images: &[metrale_model_layers::VisionItem]) -> Result<()> {
17        Ok(())
18    }
19
20    /// 2026-09-25: Give every rank the same merged vision rows for this prompt. Otherwise only
21    /// rank 0 has the encoder's rows, the other ranks keep the pad-token embedding, and the
22    /// per-layer all-reduce mixes the two. The head calls it after the prompt broadcast and the
23    /// worker at the matching point, so the collective pairs. Default: `Ok(())`.
24    fn ep_sync_vision_embeds(&self, _tokens: &[u32]) -> Result<()> {
25        Ok(())
26    }
27
28    /// 2026-09-25: Encode the images of several requests in one encoder pass. Returns, per
29    /// request in order, `(patch_row_offset, grid_index_offset, num_images, patch_row_count)`
30    /// in the shared output buffer. Default: an empty `Vec`.
31    fn prepare_vision_embed_batched(
32        &self,
33        _per_request: &[Vec<metrale_model_layers::VisionItem>],
34    ) -> Result<Vec<(usize, usize, usize, usize)>> {
35        Ok(Vec::new())
36    }
37
38    /// 2026-09-25: Where the next `prefill_chunk` reads its vision rows in the batched encoder
39    /// output: row offset, grid index offset and image count; `(0, 0, 0)` resets. Default:
40    /// no-op.
41    fn set_vision_slice_base(&self, _row_base: usize, _grid_base: usize, _owned_images: usize) {}
42
43    /// 2026-09-25: Whether `tokens` holds a vision pad token, whose KV came from image
44    /// embeddings that re-prefilling the tokens cannot reproduce. Decode preemption then does
45    /// not pick the sequence for requeue-and-re-prefill unless it can be spilled. Default
46    /// `false`.
47    fn tokens_contain_vision_pad(&self, _tokens: &[u32]) -> bool {
48        false
49    }
50}
metrale_model_engine/traits/model/vision.rs

metrale_model_engine/traits/model/
vision.rs
