Skip to main content

GenerativeModel

Trait GenerativeModel 

Source
pub trait GenerativeModel<B: Backend>: Send {
    // Required methods
    fn metadata(&self) -> &ModelMetadata;
    fn load(source: &dyn ModelSource, device: &Device<B>) -> Result<Self>
       where Self: Sized;
    fn create_kv_cache(&self, config: &CacheConfig) -> Box<dyn KVCache<B>>;
    fn embed(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3>;
    fn prefill(
        &mut self,
        input: Tensor<B, 3>,
        cache: &mut dyn KVCache<B>,
        pos: Range<u32>,
    ) -> Tensor<B, 2>;
    fn decode(
        &mut self,
        input: Tensor<B, 3>,
        cache: &mut dyn KVCache<B>,
    ) -> Tensor<B, 2>;

    // Provided methods
    fn embed_multimodal(
        &self,
        tokens: Tensor<B, 2, Int>,
        images: &[Tensor<B, 4>],
    ) -> Result<Tensor<B, 3>> { ... }
    fn decode_all_logits(
        &mut self,
        _input: Tensor<B, 3>,
        _cache: &mut dyn KVCache<B>,
    ) -> Result<Tensor<B, 3>> { ... }
    fn supports_decode_all_logits(&self) -> bool { ... }
    fn prefill_hidden(
        &mut self,
        _input: Tensor<B, 3>,
        _cache: &mut dyn KVCache<B>,
        _pos: Range<u32>,
    ) -> Result<Tensor<B, 3>> { ... }
    fn supports_hidden_states(&self) -> bool { ... }
    fn prefill_all_logits(
        &mut self,
        _input: Tensor<B, 3>,
        _cache: &mut dyn KVCache<B>,
        _pos: Range<u32>,
    ) -> Result<Tensor<B, 3>> { ... }
}
Expand description

Fixed contract every generative architecture implements — the direct analog of MLC’s embed / prefill / decode / create_kv_cache function set. The runtime only ever talks to models through this trait.

Required Methods§

Source

fn metadata(&self) -> &ModelMetadata

Metadata this model was built from.

Source

fn load(source: &dyn ModelSource, device: &Device<B>) -> Result<Self>
where Self: Sized,

Loads all weights from a ModelSource onto device.

Source

fn create_kv_cache(&self, config: &CacheConfig) -> Box<dyn KVCache<B>>

Creates a fresh KV cache for a new generation session, sized and implemented according to config (paged arena vs contiguous baseline).

Source

fn embed(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3>

Embeds token ids: [batch, seq] -> [batch, seq, hidden].

Source

fn prefill( &mut self, input: Tensor<B, 3>, cache: &mut dyn KVCache<B>, pos: Range<u32>, ) -> Tensor<B, 2>

Runs (a chunk of) the prompt through the model, filling the KV cache for positions pos. pos.end - pos.start must equal the input sequence length, and pos.start must equal the cache’s current length (dense contiguous chunks). Returns the logits of the last position, shape [batch, vocab].

Source

fn decode( &mut self, input: Tensor<B, 3>, cache: &mut dyn KVCache<B>, ) -> Tensor<B, 2>

Runs one decode step (single new position at the end of the cache). Returns the logits of that position, shape [batch, vocab].

Provided Methods§

Source

fn embed_multimodal( &self, tokens: Tensor<B, 2, Int>, images: &[Tensor<B, 4>], ) -> Result<Tensor<B, 3>>

Embeds token ids, splicing vision-tower features into the image-token spans. images are preprocessed pixel batches [1, channels, H, W], one per image-token span, in order. Text-only models keep the default impl, which rejects non-empty media and otherwise defers to embed.

Source

fn decode_all_logits( &mut self, _input: Tensor<B, 3>, _cache: &mut dyn KVCache<B>, ) -> Result<Tensor<B, 3>>

Decodes n tokens at the cache tail and returns logits for every position ([1, n, vocab]), not just the last row — the seam multi-token verification needs. Architectures without it never take the speculative path.

Source

fn supports_decode_all_logits(&self) -> bool

Whether GenerativeModel::decode_all_logits is implemented.

Source

fn prefill_hidden( &mut self, _input: Tensor<B, 3>, _cache: &mut dyn KVCache<B>, _pos: Range<u32>, ) -> Result<Tensor<B, 3>>

Runs (a chunk of) the prompt and returns the final-norm hidden states for those positions, shape [1, seq, hidden] — the embeddings path. Same cache/position contract as GenerativeModel::prefill. Models that cannot expose hidden states keep the default error.

Source

fn supports_hidden_states(&self) -> bool

Whether GenerativeModel::prefill_hidden is implemented — the capability flag /v1/model/info advertises as embeddings.

Source

fn prefill_all_logits( &mut self, _input: Tensor<B, 3>, _cache: &mut dyn KVCache<B>, _pos: Range<u32>, ) -> Result<Tensor<B, 3>>

Runs (a chunk of) the prompt and returns logits for every position, shape [1, seq, vocab] — the perplexity / speculative- decode path. Same cache/position contract as GenerativeModel::prefill. Memory scales with seq × vocab, so callers chunk accordingly. Default: unsupported.

Dyn Compatibility§

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§