metrale_model_engine/
engine.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: A greedy generate loop for one request over the [`Model`]
4//! trait: prefill, then decode, taking each token with `argmax_on_device`.
5//!
6//! Owner: metrale-model-engine.
7//! Invariants:
8//! - The engine holds no state: [`generate`] and [`generate_streaming`]
9//!   allocate a fresh sequence per call and call `free_sequence` on it whether
10//!   generation returned `Ok` or `Err`.
11
12use anyhow::Result;
13use metrale_sampling::SamplingParams;
14
15use crate::traits::Model;
16
17/// 2026-09-25: Result of a generate call.
18pub struct GenerateResult {
19    /// 2026-09-25: Output tokens, without the prompt tokens.
20    pub output_tokens: Vec<u32>,
21    /// 2026-09-25: Why generation stopped: `"stop"` (a token in
22    /// `params.stop_token_ids`) or `"length"`.
23    pub finish_reason: String,
24}
25
26/// 2026-09-25: Generate response tokens from a prompt.
27///
28/// Runs prefill, takes its token, then decodes until `params.max_tokens`
29/// tokens are out or a token in `params.stop_token_ids` is produced. The
30/// prefill token is always returned, so the output has at least one token.
31pub fn generate(
32    model: &dyn Model,
33    prompt_tokens: &[u32],
34    params: &SamplingParams,
35) -> Result<GenerateResult> {
36    let mut seq = model.alloc_sequence()?;
37    let stream = 0u64;
38
39    let result = generate_inner(model, prompt_tokens, params, &mut seq, stream);
40
41    // 2026-09-25: Free the sequence whether or not generation failed. An
42    // error from `free_sequence` replaces the result.
43    model.free_sequence(&mut seq)?;
44
45    result
46}
47
48fn generate_inner(
49    model: &dyn Model,
50    prompt_tokens: &[u32],
51    params: &SamplingParams,
52    seq: &mut crate::traits::SequenceState,
53    stream: u64,
54) -> Result<GenerateResult> {
55    let logits_ptr = model.prefill(prompt_tokens, seq, stream)?;
56    let first_token = model.argmax_on_device(logits_ptr, stream)?;
57
58    let mut output_tokens = Vec::with_capacity(params.max_tokens);
59    output_tokens.push(first_token);
60
61    if params.stop_token_ids.contains(&first_token) {
62        return Ok(GenerateResult {
63            output_tokens,
64            finish_reason: "stop".to_string(),
65        });
66    }
67
68    for _step in 1..params.max_tokens {
69        let last_token = *output_tokens.last().unwrap();
70        let logits_ptr = model.decode(last_token, seq, stream)?;
71        let token = model.argmax_on_device(logits_ptr, stream)?;
72
73        output_tokens.push(token);
74
75        if params.stop_token_ids.contains(&token) {
76            return Ok(GenerateResult {
77                output_tokens,
78                finish_reason: "stop".to_string(),
79            });
80        }
81    }
82
83    Ok(GenerateResult {
84        output_tokens,
85        finish_reason: "length".to_string(),
86    })
87}
88
89/// 2026-09-25: Same as [`generate`], but calls `on_token(token_id)`
90/// synchronously after each token, the prefill token included.
91pub fn generate_streaming<F>(
92    model: &dyn Model,
93    prompt_tokens: &[u32],
94    params: &SamplingParams,
95    mut on_token: F,
96) -> Result<GenerateResult>
97where
98    F: FnMut(u32),
99{
100    let mut seq = model.alloc_sequence()?;
101    let stream = 0u64;
102
103    let result = generate_streaming_inner(
104        model,
105        prompt_tokens,
106        params,
107        &mut on_token,
108        &mut seq,
109        stream,
110    );
111
112    model.free_sequence(&mut seq)?;
113
114    result
115}
116
117fn generate_streaming_inner<F>(
118    model: &dyn Model,
119    prompt_tokens: &[u32],
120    params: &SamplingParams,
121    on_token: &mut F,
122    seq: &mut crate::traits::SequenceState,
123    stream: u64,
124) -> Result<GenerateResult>
125where
126    F: FnMut(u32),
127{
128    let logits_ptr = model.prefill(prompt_tokens, seq, stream)?;
129    let first_token = model.argmax_on_device(logits_ptr, stream)?;
130
131    let mut output_tokens = Vec::with_capacity(params.max_tokens);
132    output_tokens.push(first_token);
133    on_token(first_token);
134
135    if params.stop_token_ids.contains(&first_token) {
136        return Ok(GenerateResult {
137            output_tokens,
138            finish_reason: "stop".to_string(),
139        });
140    }
141
142    for _step in 1..params.max_tokens {
143        let last_token = *output_tokens.last().unwrap();
144        let logits_ptr = model.decode(last_token, seq, stream)?;
145        let token = model.argmax_on_device(logits_ptr, stream)?;
146
147        output_tokens.push(token);
148        on_token(token);
149
150        if params.stop_token_ids.contains(&token) {
151            return Ok(GenerateResult {
152                output_tokens,
153                finish_reason: "stop".to_string(),
154            });
155        }
156    }
157
158    Ok(GenerateResult {
159        output_tokens,
160        finish_reason: "length".to_string(),
161    })
162}
163
164/// 2026-09-25: Generate with MTP speculative decoding, delegated to
165/// `Model::generate_speculative`.
166pub fn generate_speculative(
167    model: &dyn Model,
168    prompt_tokens: &[u32],
169    params: &SamplingParams,
170    num_drafts: usize,
171) -> Result<GenerateResult> {
172    model.generate_speculative(prompt_tokens, params, num_drafts)
173}
174
175#[cfg(test)]
176mod tests;
metrale_model_engine/engine.rs

metrale_model_engine/
engine.rs
