combs_models/
traits.rs

1//! The model-agnostic generation contract.
2
3use std::ops::Range;
4
5use burn::tensor::{Int, Tensor, backend::Backend, Device};
6use combs_formats::{ModelMetadata, ModelSource};
7
8use crate::Result;
9use crate::kv::{CacheConfig, KVCache};
10
11/// Fixed contract every generative architecture implements — the direct
12/// analog of MLC's `embed / prefill / decode / create_kv_cache` function set.
13/// The runtime only ever talks to models through this trait.
14pub trait GenerativeModel<B: Backend>: Send {
15    /// Metadata this model was built from.
16    fn metadata(&self) -> &ModelMetadata;
17
18    /// Loads all weights from a [`ModelSource`] onto `device`.
19    fn load(source: &dyn ModelSource, device: &Device<B>) -> Result<Self>
20    where
21        Self: Sized;
22
23    /// Creates a fresh KV cache for a new generation session, sized and
24    /// implemented according to `config` (paged arena vs contiguous
25    /// baseline).
26    fn create_kv_cache(&self, config: &CacheConfig) -> Box<dyn KVCache<B>>;
27
28    /// Embeds token ids: `[batch, seq] -> [batch, seq, hidden]`.
29    fn embed(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3>;
30
31    /// Embeds token ids, splicing vision-tower features into the image-token
32    /// spans. `images` are preprocessed pixel batches `[1, channels, H, W]`,
33    /// one per image-token span, in order. Text-only models keep the default
34    /// impl, which rejects non-empty media and otherwise defers to `embed`.
35    fn embed_multimodal(
36        &self,
37        tokens: Tensor<B, 2, Int>,
38        images: &[Tensor<B, 4>],
39    ) -> Result<Tensor<B, 3>> {
40        if !images.is_empty() {
41            return Err(crate::ModelError::UnsupportedMedia(format!(
42                "{} image(s) passed to a text-only model",
43                images.len()
44            )));
45        }
46        Ok(self.embed(tokens))
47    }
48
49    /// Runs (a chunk of) the prompt through the model, filling the KV cache
50    /// for positions `pos`. `pos.end - pos.start` must equal the input
51    /// sequence length, and `pos.start` must equal the cache's current
52    /// length (dense contiguous chunks). Returns the logits of the **last**
53    /// position, shape `[batch, vocab]`.
54    fn prefill(
55        &mut self,
56        input: Tensor<B, 3>,
57        cache: &mut dyn KVCache<B>,
58        pos: Range<u32>,
59    ) -> Tensor<B, 2>;
60
61    /// Runs one decode step (single new position at the end of the cache).
62    /// Returns the logits of that position, shape `[batch, vocab]`.
63    fn decode(&mut self, input: Tensor<B, 3>, cache: &mut dyn KVCache<B>) -> Tensor<B, 2>;
64}
combs_models/traits.rs

combs_models/
traits.rs
