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
65    /// Decodes `n` tokens at the cache tail and returns logits for every
66    /// position (`[1, n, vocab]`), not just the last row — the seam
67    /// multi-token verification needs. Architectures without it never take
68    /// the speculative path.
69    fn decode_all_logits(
70        &mut self,
71        _input: Tensor<B, 3>,
72        _cache: &mut dyn KVCache<B>,
73    ) -> crate::Result<Tensor<B, 3>> {
74        Err(crate::ModelError::Unsupported(
75            "this model does not expose per-position decode logits".to_string(),
76        ))
77    }
78
79    /// Whether [`GenerativeModel::decode_all_logits`] is implemented.
80    fn supports_decode_all_logits(&self) -> bool {
81        false
82    }
83
84    /// Runs (a chunk of) the prompt and returns the final-norm hidden
85    /// states for those positions, shape `[1, seq, hidden]` — the
86    /// embeddings path. Same cache/position contract as
87    /// [`GenerativeModel::prefill`]. Models that cannot expose hidden
88    /// states keep the default error.
89    fn prefill_hidden(
90        &mut self,
91        _input: Tensor<B, 3>,
92        _cache: &mut dyn KVCache<B>,
93        _pos: Range<u32>,
94    ) -> Result<Tensor<B, 3>> {
95        Err(crate::ModelError::Unsupported(
96            "this model does not expose hidden states for embeddings".to_string(),
97        ))
98    }
99
100    /// Whether [`GenerativeModel::prefill_hidden`] is implemented — the
101    /// capability flag `/v1/model/info` advertises as `embeddings`.
102    fn supports_hidden_states(&self) -> bool {
103        false
104    }
105
106    /// Runs (a chunk of) the prompt and returns logits for **every**
107    /// position, shape `[1, seq, vocab]` — the perplexity / speculative-
108    /// decode path. Same cache/position contract as
109    /// [`GenerativeModel::prefill`]. Memory scales with `seq × vocab`, so
110    /// callers chunk accordingly. Default: unsupported.
111    fn prefill_all_logits(
112        &mut self,
113        _input: Tensor<B, 3>,
114        _cache: &mut dyn KVCache<B>,
115        _pos: Range<u32>,
116    ) -> Result<Tensor<B, 3>> {
117        Err(crate::ModelError::Unsupported(
118            "this model does not expose per-position logits".to_string(),
119        ))
120    }
121}
122
123/// Speech-to-text models (Whisper-style encoder–decoder). A separate
124/// contract from [`GenerativeModel`]: the encoder runs once per audio
125/// window, then the decoder is stepped over token prefixes against the
126/// fixed encoder states.
127pub trait SpeechToTextModel<B: Backend>: Send {
128    /// Architecture + hyperparameter metadata.
129    fn metadata(&self) -> &ModelMetadata;
130
131    /// Mel bins the encoder expects (derived from its conv stem weights).
132    fn n_mels(&self) -> usize;
133
134    /// Encodes one `[1, n_mels, frames]` log-mel window into encoder
135    /// states `[1, frames/2, hidden]`.
136    fn encode_audio(&self, mel: Tensor<B, 3>) -> crate::Result<Tensor<B, 3>>;
137
138    /// Runs the decoder over the whole token prefix and returns the final
139    /// position's logits `[vocab]`.
140    fn decode_step(&self, tokens: &[u32], encoded: &Tensor<B, 3>)
141    -> crate::Result<Tensor<B, 1>>;
142}
combs_models/traits.rs

combs_models/
traits.rs
