metrale_model_engine/model/nllb/
model_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `impl Model for NllbGpuModel`. `prefill` runs the encoder and seeds the
4//! decoder with `[decoder_start, forced_bos]`, so the forced target-language token never
5//! enters the sampled stream. The speculative, SSM and MTP methods either fail or do
6//! nothing.
7//!
8//! Owner: model-engine.
9//! Invariants: none beyond the types.
10
11use anyhow::{Context, Result, bail};
12use metrale_gpu_runtime::gpu::DevicePtr;
13
14use super::NllbGpuModel;
15use super::kv::NllbSeqKv;
16use crate::traits::{
17    Model, ModelAdapters, ModelDeviceFeed, ModelDraft, ModelEp, ModelForward, ModelLifecycle,
18    ModelLogits, ModelSsmState, ModelStreams, ModelVerify, ModelVision, SequenceState,
19};
20
21impl NllbGpuModel {
22    /// 2026-09-25: Encoder pass and decoder seed, shared by `prefill` and `prefill_chunk`.
23    fn prefill_impl(&self, tokens: &[u32], seq: &mut SequenceState) -> Result<DevicePtr> {
24        let slot = seq.slot_idx;
25        self.set_lora_active(seq.adapter_slot);
26        // 2026-09-25: Per-request source/target language (0 → deployment default).
27        let src_id = if seq.src_lang_id != 0 {
28            seq.src_lang_id
29        } else {
30            self.lang.src_lang_id
31        };
32        let tgt_id = if seq.tgt_lang_id != 0 {
33            seq.tgt_lang_id
34        } else {
35            self.lang.tgt_lang_id
36        };
37        let src = self.lang.encoder_input_with(src_id, tokens);
38        {
39            let mut map = self.kv.lock().unwrap();
40            let kv = map
41                .get_mut(&slot)
42                .context("nllb prefill: no KV state for slot (alloc_sequence not called?)")?;
43            self.run_encoder(&src, kv)?;
44            // 2026-09-25: Seed the decoder: step 0 consumes `decoder_start`
45            // (its logits are overwritten), step 1 consumes `forced_bos =
46            // tgt_lang` and yields the logits for the first translated token.
47            let row = self.prefill_logits_row(slot);
48            self.forward_one(self.lang.decoder_start_id, kv, row)?;
49            self.forward_one(tgt_id, kv, row)?;
50        }
51        self.sync()?;
52        // 2026-09-25: The sequence records the decoder tokens, not the source;
53        // the model's own KV holds the computed state.
54        seq.tokens = vec![self.lang.decoder_start_id, tgt_id];
55        seq.seq_len = seq.tokens.len();
56        seq.prompt_len = seq.tokens.len();
57        Ok(self.prefill_logits_row(slot))
58    }
59}
60
61impl Model for NllbGpuModel {}
62
63impl ModelLifecycle for NllbGpuModel {
64    fn bind_gpu_to_thread(&self) -> Result<()> {
65        self.gpu.bind_to_thread()
66    }
67    fn alloc_sequence(&self) -> Result<SequenceState> {
68        let slot = self.slots.lock().unwrap().claim();
69        let kv = NllbSeqKv::new(self.gpu.as_ref(), self.dec_layers, self.cache_rows, self.d)?;
70        self.kv.lock().unwrap().insert(slot, kv);
71        // 2026-09-25: Host-only sequence state; the KV lives in `self.kv`, keyed
72        // by `slot`.
73        Ok(SequenceState::host_only(slot))
74    }
75    fn free_sequence(&self, seq: &mut SequenceState) -> Result<()> {
76        let slot = seq.slot_idx;
77        if slot == usize::MAX {
78            return Ok(()); // 2026-09-25: detached by `detach_slot_for_reuse`
79        }
80        if let Some(kv) = self.kv.lock().unwrap().remove(&slot) {
81            kv.free(self.gpu.as_ref())?;
82        }
83        self.slots.lock().unwrap().release(slot);
84        Ok(())
85    }
86    fn compact_sequence(&self, seq: &mut SequenceState, new_slot: usize) -> Result<()> {
87        let old = seq.slot_idx;
88        let mut map = self.kv.lock().unwrap();
89        if let Some(kv) = map.remove(&old) {
90            map.insert(new_slot, kv);
91        }
92        seq.slot_idx = new_slot;
93        Ok(())
94    }
95    fn detach_slot_for_reuse(&self, seq: &mut SequenceState) {
96        seq.slot_idx = usize::MAX;
97    }
98    fn cache_sequence(&self, _seq: &SequenceState) {
99        // 2026-09-25: No prefix reuse for translation: nothing is cached.
100    }
101    fn num_free_blocks(&self) -> usize {
102        // 2026-09-25: The KV lives outside the paged block cache, so there are no
103        // blocks to count; report 2^20.
104        1 << 20
105    }
106    fn kv_block_bytes(&self) -> usize {
107        // 2026-10-01: No paged cache (see `num_free_blocks`), so no block has a size; 0, as
108        // `num_total_blocks` reports 0 blocks.
109        0
110    }
111}
112
113impl ModelForward for NllbGpuModel {
114    fn prefill(&self, tokens: &[u32], seq: &mut SequenceState, _stream: u64) -> Result<DevicePtr> {
115        self.prefill_impl(tokens, seq)
116    }
117    fn prefill_chunk(
118        &self,
119        tokens: &[u32],
120        seq: &mut SequenceState,
121        chunk_start: usize,
122        chunk_len: usize,
123        is_last_chunk: bool,
124        _stream: u64,
125    ) -> Result<DevicePtr> {
126        if chunk_start != 0 || !is_last_chunk || chunk_len != tokens.len() {
127            bail!("nllb: chunked prefill unsupported — the source must arrive in one chunk");
128        }
129        self.prefill_impl(tokens, seq)
130    }
131    fn decode(&self, token: u32, seq: &mut SequenceState, _stream: u64) -> Result<DevicePtr> {
132        let slot = seq.slot_idx;
133        self.set_lora_active(seq.adapter_slot);
134        let row = self.decode_logits_row(0);
135        {
136            let mut map = self.kv.lock().unwrap();
137            let kv = map
138                .get_mut(&slot)
139                .context("nllb decode: no KV state for slot")?;
140            self.forward_one(token, kv, row)?;
141        }
142        self.sync()?;
143        seq.tokens.push(token);
144        seq.seq_len += 1;
145        Ok(row)
146    }
147    /// 2026-09-25: Batched decode: one `forward_one` per sequence, serially on
148    /// the default stream, into contiguous logit rows `0..n` (batch position `i`
149    /// ↔ `seqs[i]`); returns the base of `[n, vocab]`. Each sequence's KV is
150    /// found by its `slot_idx`. Fails when `n` is 0, differs from
151    /// `tokens.len()`, or exceeds `max_batch`.
152    fn decode_batch(
153        &self,
154        tokens: &[u32],
155        seqs: &mut [&mut SequenceState],
156        _stream: u64,
157    ) -> Result<DevicePtr> {
158        let n = seqs.len();
159        if n == 0 || tokens.len() != n {
160            bail!(
161                "nllb decode_batch: tokens/seqs length mismatch ({}, {n})",
162                tokens.len()
163            );
164        }
165        if n > self.max_batch {
166            bail!(
167                "nllb decode_batch: n={n} exceeds max_batch={}",
168                self.max_batch
169            );
170        }
171        {
172            let mut map = self.kv.lock().unwrap();
173            for (i, seq) in seqs.iter().enumerate() {
174                // 2026-09-25: Per-request LoRA gate, re-armed for each sequence.
175                self.set_lora_active(seq.adapter_slot);
176                let kv = map
177                    .get_mut(&seq.slot_idx)
178                    .context("nllb decode_batch: no KV state for slot")?;
179                self.forward_one(tokens[i], kv, self.decode_logits_row(i))?;
180            }
181        }
182        self.sync()?;
183        for (i, seq) in seqs.iter_mut().enumerate() {
184            seq.tokens.push(tokens[i]);
185            seq.seq_len += 1;
186        }
187        Ok(self.decode_logits_row(0))
188    }
189    fn supports_beam(&self) -> bool {
190        true
191    }
192    fn generate_beam_batch(&self, reqs: &[crate::traits::BeamReq]) -> Result<Vec<Vec<u32>>> {
193        // 2026-09-25: The requests' beams decode as one fused batch per step when
194        // they share an adapter and `num_beams`; otherwise `beam_batched_multi`
195        // runs them one at a time.
196        self.beam_batched_multi(reqs)
197    }
198}
199
200impl ModelLogits for NllbGpuModel {
201    fn vocab_size(&self) -> usize {
202        self.vocab
203    }
204    fn copy_logits_to_host(&self, logits_ptr: DevicePtr, dst: &mut [u8]) -> Result<()> {
205        self.gpu.copy_d2h(logits_ptr, dst)
206    }
207    fn logits_buffer_ptr(&self) -> DevicePtr {
208        self.decode_logits_row(0)
209    }
210    fn argmax_on_device(&self, logits_ptr: DevicePtr, _stream: u64) -> Result<u32> {
211        self.argmax_of(logits_ptr)
212    }
213    fn argmax_batch(&self, logits_ptr: DevicePtr, n: usize, _stream: u64) -> Result<Vec<u32>> {
214        // 2026-09-25: `logits_ptr` is `[n, vocab]` bf16; argmax each contiguous row.
215        (0..n)
216            .map(|i| self.argmax_of(logits_ptr.offset(i * self.vocab * 2)))
217            .collect()
218    }
219    fn hidden_after_norm(&self) -> DevicePtr {
220        DevicePtr(0)
221    }
222}
223
224impl ModelAdapters for NllbGpuModel {}
225
226impl ModelSsmState for NllbGpuModel {
227    fn checkpoint_ssm_states(&self, _seq: &mut SequenceState) -> Result<()> {
228        Ok(())
229    }
230    fn rollback_ssm_states(&self, _seq: &mut SequenceState, _n: usize) -> Result<()> {
231        Ok(())
232    }
233}
234
235impl ModelVerify for NllbGpuModel {
236    fn decode_verify(&self, _t: &[u32], _s: &mut SequenceState, _st: u64) -> Result<Vec<u32>> {
237        bail!("nllb: speculative verify not supported")
238    }
239    fn decode_verify_graphed(
240        &self,
241        _t: &[u32; 2],
242        _s: &mut SequenceState,
243        _st: u64,
244    ) -> Result<[u32; 2]> {
245        bail!("nllb: graphed verify not supported")
246    }
247    fn decode_verify_graphed_k3(
248        &self,
249        _t: &[u32; 3],
250        _s: &mut SequenceState,
251        _st: u64,
252    ) -> Result<[u32; 3]> {
253        bail!("nllb: graphed verify not supported")
254    }
255    fn decode_verify_graphed_k4(
256        &self,
257        _t: &[u32; 4],
258        _s: &mut SequenceState,
259        _st: u64,
260    ) -> Result<[u32; 4]> {
261        bail!("nllb: graphed verify not supported")
262    }
263}
264
265impl ModelDraft for NllbGpuModel {
266    fn generate_speculative(
267        &self,
268        _tokens: &[u32],
269        _params: &metrale_sampling::SamplingParams,
270        _num_drafts: usize,
271    ) -> Result<crate::engine::GenerateResult> {
272        bail!("nllb: speculative decoding not supported")
273    }
274    fn has_proposer(&self) -> bool {
275        false
276    }
277    fn has_self_speculative(&self) -> bool {
278        false
279    }
280    fn decode_draft(
281        &self,
282        _token: u32,
283        _seq: &mut SequenceState,
284        _stream: u64,
285    ) -> Result<DevicePtr> {
286        bail!("nllb: self-speculative draft not supported")
287    }
288    fn save_hidden_for_mtp(&self, _token_idx: usize, _stream: u64) -> Result<()> {
289        Ok(())
290    }
291    fn run_mtp_propose(
292        &self,
293        _token: u32,
294        _position: usize,
295        _seq: &mut SequenceState,
296        _stream: u64,
297    ) -> Result<Option<u32>> {
298        Ok(None)
299    }
300    fn run_mtp_propose_multi(
301        &self,
302        _token: u32,
303        _position: usize,
304        _num_drafts: usize,
305        _seq: &mut SequenceState,
306        _stream: u64,
307        _grammar_bitmask: Option<&[i32]>,
308    ) -> Result<Vec<u32>> {
309        Ok(Vec::new())
310    }
311    fn trim_proposer_state(
312        &self,
313        _seq: &mut SequenceState,
314        _num_accepted: usize,
315        _stream: u64,
316    ) -> Result<()> {
317        Ok(())
318    }
319}
320
321impl ModelVision for NllbGpuModel {}
322
323impl ModelEp for NllbGpuModel {}
324impl crate::traits::ModelCircuit for NllbGpuModel {}
325
326impl ModelStreams for NllbGpuModel {}
327
328impl ModelDeviceFeed for NllbGpuModel {}
metrale_model_engine/model/nllb/model_impl.rs

metrale_model_engine/model/nllb/
model_impl.rs
