spark_model/model/nllb/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Served NLLB-200 / M2M-100 encoder-decoder translation model.
4//!
5//! Atlas's production engine is decoder-only + GPU-only; NLLB is seq2seq
6//! (bidirectional encoder + decoder cross-attention + sinusoidal positions +
7//! ReLU FFN + biased LayerNorm). This module promotes the validated bf16 CUDA
8//! runtime (`examples/nllb_cuda_bf16`) into a first-class [`crate::traits::Model`]
9//! so `spark serve --model <nllb-dir>` translates through the SAME scheduler,
10//! OpenAI API, sampling and (later) LoRA path as every other model — no parallel
11//! server, no hardcoded paths (weights come from the standard `--model` store).
12//!
13//! The model owns ALL its KV (the scheduler's paged block cache is unused): a
14//! per-sequence decoder self-attn cache that grows one row per token plus a
15//! fixed cross-attn cache computed once from the encoder. See the `kv` module. Logits are
16//! bf16 — the scheduler's default sampling path — so no fp32 overrides are
17//! needed.
18
19use std::collections::HashMap;
20use std::sync::Mutex;
21use std::sync::atomic::{AtomicBool, Ordering};
22
23use anyhow::{Context, Result, ensure};
24use atlas_core::config::ModelConfig;
25use spark_runtime::gpu::{DevicePtr, GpuBackend};
26use spark_runtime::weights::{WeightDtype, WeightStore};
27
28mod beam;
29mod beam_compute;
30mod beam_multi;
31mod compute;
32mod kernels;
33mod kv;
34mod lang;
35mod lora;
36mod model_impl;
37mod util;
38
39pub use lang::NllbLang;
40
41use compute::DecScratch;
42use kernels::NllbKernels;
43use kv::NllbSeqKv;
44use lora::NllbLora;
45
46/// Decoder self-attn KV cache depth (max generated tokens per translation).
47/// NLLB `max_length` defaults to 200; 512 is a comfortable cap that keeps
48/// per-sequence KV small (`cache_rows·d·2·dec_layers·2` bytes).
49const DEFAULT_CACHE_ROWS: usize = 512;
50
51/// The served NLLB encoder-decoder model. `Send + Sync`: every field is either
52/// immutable after construction or behind a `Mutex`; `DevicePtr` is a `Copy`
53/// device handle and the GPU-side buffers are driven only through `&self`.
54pub struct NllbGpuModel {
55    gpu: Box<dyn GpuBackend>,
56    kernels: NllbKernels,
57    weights: HashMap<String, DevicePtr>,
58    embed_table: DevicePtr,
59    // dims
60    d: usize,
61    heads: usize,
62    head_dim: usize,
63    ffn: usize,
64    enc_layers: usize,
65    dec_layers: usize,
66    vocab: usize,
67    embed_scale: f32,
68    attn_scale: f32,
69    cache_rows: usize,
70    max_batch: usize,
71    lang: NllbLang,
72    // persistent device scratch (single-token decode) + outputs
73    dec: DecScratch,
74    /// bf16 decode logits `[max_batch, vocab]` — `decode_batch` writes CONTIGUOUS
75    /// rows `0..n` (batch position `i` ↔ `seqs[i]`), the scheduler's contract.
76    decode_logits: DevicePtr,
77    /// bf16 prefill logits `[max_batch, vocab]` — each concurrent prefill writes
78    /// its OWN row (indexed by `slot_idx`) so overlapping prefills don't collide.
79    prefill_logits: DevicePtr,
80    /// Decoder sinusoidal position table `[cache_rows, d]` bf16.
81    pos_table: DevicePtr,
82    // per-sequence KV + slot allocator
83    kv: Mutex<HashMap<usize, NllbSeqKv>>,
84    slots: Mutex<SlotAlloc>,
85    /// Optional PEFT LoRA adapter applied on every projection (runtime delta).
86    lora: Option<NllbLora>,
87    /// Per-request LoRA gate: set from `SequenceState.adapter_slot` before each
88    /// sequence's forward (`>=0` → apply the adapter, `-1` → base). Sound because
89    /// the scheduler drives prefill/decode serially on one thread.
90    lora_active: AtomicBool,
91}
92
93/// Trivial monotonic slot allocator with a free-list for reuse.
94#[derive(Default)]
95struct SlotAlloc {
96    next: usize,
97    free: Vec<usize>,
98}
99
100impl SlotAlloc {
101    fn claim(&mut self) -> usize {
102        self.free.pop().unwrap_or_else(|| {
103            let s = self.next;
104            self.next += 1;
105            s
106        })
107    }
108    fn release(&mut self, slot: usize) {
109        self.free.push(slot);
110    }
111}
112
113impl NllbGpuModel {
114    /// Build from the standard `--model` weight store + GPU backend. `lang`
115    /// carries the tokenizer-resolved source/target language ids (resolved
116    /// server-side, where the tokenizer lives). `max_seq_len` caps the decoder
117    /// KV depth.
118    pub fn new(
119        config: &ModelConfig,
120        store: &WeightStore,
121        gpu: Box<dyn GpuBackend>,
122        lang: NllbLang,
123        max_seq_len: usize,
124        max_batch: usize,
125        lora_dir: Option<&std::path::Path>,
126    ) -> Result<Self> {
127        let d = config.hidden_size;
128        let heads = config.num_attention_heads;
129        let head_dim = config.head_dim;
130        let ffn = config.intermediate_size;
131        let dec_layers = config.num_hidden_layers;
132        // NLLB / M2M-100 are architecturally symmetric: encoder and decoder
133        // share layer count, head count and FFN width.
134        let enc_layers = dec_layers;
135        let vocab = config.vocab_size;
136        // `scale_embedding` is `true` for the NLLB family (√d_model embed scale).
137        let embed_scale = (d as f32).sqrt();
138        let attn_scale = (head_dim as f32).powf(-0.5);
139        let cache_rows = DEFAULT_CACHE_ROWS.max(max_seq_len.min(2048));
140
141        ensure!(
142            store.get("model.shared.weight")?.dtype == WeightDtype::BF16,
143            "nllb serving requires a bf16 checkpoint; convert with \
144             scripts/convert-safetensors-to-bf16.py"
145        );
146        let weights: HashMap<String, DevicePtr> = store
147            .names()
148            .map(|n| Ok((n.to_string(), store.get(n)?.ptr)))
149            .collect::<Result<_>>()?;
150        let embed_table = *weights
151            .get("model.shared.weight")
152            .context("nllb: missing tied embedding model.shared.weight")?;
153
154        // Make the CUDA context current on the construction thread before any
155        // kernel resolve / device alloc.
156        gpu.bind_to_thread()?;
157        let kernels = NllbKernels::new(gpu.as_ref())?;
158        let dec = DecScratch::new(gpu.as_ref(), d, ffn, vocab)?;
159        let max_batch = max_batch.max(1);
160        let decode_logits = gpu.alloc(max_batch * vocab * 2)?;
161        let prefill_logits = gpu.alloc(max_batch * vocab * 2)?;
162        let pos_table = gpu.alloc(cache_rows * d * 2)?;
163        let pos_host = util::decoder_pos_table_bf16(cache_rows, d);
164        gpu.copy_h2d(util::bf16_bytes(&pos_host), pos_table)?;
165
166        let lora = match lora_dir {
167            Some(dir) => Some(NllbLora::load(dir, gpu.as_ref(), cache_rows)?),
168            None => None,
169        };
170
171        tracing::info!(
172            "NLLB served model ready: d={d} heads={heads} enc={enc_layers} dec={dec_layers} \
173             vocab={vocab} src_lang_id={} tgt_lang_id={} cache_rows={cache_rows}",
174            lang.src_lang_id,
175            lang.tgt_lang_id,
176        );
177
178        Ok(Self {
179            gpu,
180            kernels,
181            weights,
182            embed_table,
183            d,
184            heads,
185            head_dim,
186            ffn,
187            enc_layers,
188            dec_layers,
189            vocab,
190            embed_scale,
191            attn_scale,
192            cache_rows,
193            max_batch,
194            lang,
195            dec,
196            decode_logits,
197            prefill_logits,
198            pos_table,
199            kv: Mutex::new(HashMap::new()),
200            slots: Mutex::new(SlotAlloc::default()),
201            lora,
202            lora_active: AtomicBool::new(false),
203        })
204    }
205
206    /// Device pointer for weight `name` (panics if absent — construction
207    /// resolves the full store, so a miss is a checkpoint/format bug).
208    #[inline]
209    pub(super) fn w(&self, name: &str) -> DevicePtr {
210        self.weights[name]
211    }
212
213    /// Arm/disarm the LoRA delta for the sequence about to be forwarded
214    /// (`adapter_slot >= 0` → apply). No-op when no adapter is loaded.
215    #[inline]
216    pub(super) fn set_lora_active(&self, adapter_slot: i32) {
217        self.lora_active
218            .store(self.lora.is_some() && adapter_slot >= 0, Ordering::Relaxed);
219    }
220
221    #[inline]
222    pub(super) fn lora_is_active(&self) -> bool {
223        self.lora_active.load(Ordering::Relaxed)
224    }
225}
spark_model/model/nllb/mod.rs

spark_model/model/nllb/
mod.rs
