kime_engine/
lib.rs

1//! Session, scheduler, batching, the state memory cache, the answer cache and the tokenization cache. This is the in process API that the server, the CLI and the language bindings all sit on. See spec/07-engine.md and spec/11-serving.md.
2//!
3//! What is here today is the compat path end to end: open a Laya checkpoint, lay each question out
4//! as Laya does, run the questions of one or many requests in shared batches on the CPU or a CUDA
5//! GPU, and build Laya's answers from the logits. The scheduler that merges requests from many
6//! callers arrives with the server, so for now a call runs on the caller's thread and callers
7//! share the device through a lock.
8//!
9//! The `kime` crate re-exports all of it, and its docs hold a full example.
10
11#![forbid(unsafe_code)]
12
13use std::fmt;
14use std::future::Future;
15use std::pin::Pin;
16use std::sync::{Arc, Mutex, PoisonError};
17use std::task::{Context, Poll, Waker};
18
19use kime_core::answer::{LAYA_MODEL, Response, Temperatures, laya_answer};
20use kime_core::render::{compat_question, compat_state};
21use kime_core::request::{Limits, Problem, Question, Request, parse};
22use kime_model::Model;
23use kime_tensor::{BatchBuf, Buckets, Executor, Outputs};
24use kime_tok::Tokenizer;
25use kime_tok::layout::{CompatBudget, CompatSequence, Cut};
26use serde_json::Value;
27
28pub mod hub;
29
30/// Where the model runs.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
32pub enum Device {
33    /// The first CUDA GPU when there is one and this build has CUDA, the CPU otherwise.
34    #[default]
35    Auto,
36    /// The CPU, on `threads` threads, or every core for 0.
37    Cpu {
38        /// Worker threads.
39        threads: usize,
40    },
41    /// A CUDA GPU by ordinal.
42    Cuda(usize),
43    /// The Apple GPU. Arrives with M4.
44    Metal,
45    /// The Apple Neural Engine. Arrives with M4.
46    Ane,
47}
48
49/// The number format. The CPU computes in FP32 for both float settings.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
51pub enum Precision {
52    /// FP16 weights and GEMM inputs with FP32 accumulation, as Laya's autocast runs on a GPU.
53    #[default]
54    F16,
55    /// FP32 throughout, closest to Laya on the CPU.
56    F32,
57    /// INT8 weights and activations in the encoder and decision head GEMMs on the CPU, with the
58    /// scorer in FP32. Faster, and further from Laya than FP32: see spec/10-cpu.md for the gate a
59    /// checkpoint has to pass. Not on GPUs yet.
60    Int8,
61}
62
63/// What can go wrong.
64#[derive(Debug)]
65pub enum Error {
66    /// The request does not validate. Every problem is listed, in the wire format.
67    Invalid(Vec<Problem>),
68    /// The model name did not resolve.
69    NotFound(String),
70    /// The checkpoint did not load.
71    Model(kime_model::Error),
72    /// The checkpoint's tokenizer did not load.
73    Tokenizer(String),
74    /// The device failed, or the batch did not fit it.
75    Backend(kime_tensor::Error),
76    /// A question's head and options do not fit the checkpoint's budget, so some options have no
77    /// marker. Laya raises the same error.
78    TooLong {
79        /// The question id.
80        question: String,
81        /// Its options.
82        options: usize,
83        /// The options that fit.
84        fit: usize,
85    },
86    /// The device asked for is not in this build or not on this machine.
87    Unsupported(String),
88}
89
90impl fmt::Display for Error {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        match self {
93            Error::Invalid(p) => {
94                write!(f, "invalid request:")?;
95                for p in p {
96                    write!(f, " {};", p.to_json())?;
97                }
98                Ok(())
99            }
100            Error::NotFound(m) | Error::Tokenizer(m) | Error::Unsupported(m) => f.write_str(m),
101            Error::Model(e) => write!(f, "{e}"),
102            Error::Backend(e) => write!(f, "{e}"),
103            Error::TooLong { question, options, fit } => write!(
104                f,
105                "question {question:?} has {options} options but only {fit} fit the head budget"
106            ),
107        }
108    }
109}
110
111impl std::error::Error for Error {}
112
113impl From<kime_tensor::Error> for Error {
114    fn from(e: kime_tensor::Error) -> Self {
115        Error::Backend(e)
116    }
117}
118
119/// Settings for [`Kime`], from [`Kime::builder`].
120#[derive(Debug, Clone, Default)]
121pub struct Builder {
122    model: Option<String>,
123    device: Device,
124    precision: Precision,
125    preload: bool,
126}
127
128impl Builder {
129    /// The model: an alias such as `laya`, a checkpoint directory, a `.kime` file or an
130    /// `hf://org/repo[/subfolder]` reference. See [`hub`].
131    #[must_use]
132    pub fn model(mut self, name: impl Into<String>) -> Self {
133        self.model = Some(name.into());
134        self
135    }
136
137    /// Where to run.
138    #[must_use]
139    pub fn device(mut self, d: Device) -> Self {
140        self.device = d;
141        self
142    }
143
144    /// The number format on a GPU.
145    #[must_use]
146    pub fn precision(mut self, p: Precision) -> Self {
147        self.precision = p;
148        self
149    }
150
151    /// Builds the plans for the smallest batch shapes now, so the first requests do not pay for
152    /// them.
153    #[must_use]
154    pub fn preload(mut self, yes: bool) -> Self {
155        self.preload = yes;
156        self
157    }
158
159    /// Loads the model onto the device.
160    ///
161    /// # Errors
162    ///
163    /// When the model is not found or does not load, or the device cannot be opened.
164    pub fn build(self) -> Result<Kime, Error> {
165        let name = self.model.unwrap_or_else(|| "laya".into());
166        let path = hub::resolve(&name).map_err(Error::NotFound)?;
167        let model = Model::open(&path).map_err(Error::Model)?;
168        let tok_json = model
169            .file("tokenizer/tokenizer.json")
170            .ok_or_else(|| Error::Tokenizer(format!("{}: no tokenizer.json", path.display())))?;
171        let tok = Tokenizer::from_bytes(tok_json, model.file("tokenizer/tokenizer_config.json"))
172            .map_err(|e| Error::Tokenizer(e.to_string()))?;
173        let agent = &model.spec.agent;
174        let temps = Temperatures::new(agent.temperature, &agent.temperature_by_options);
175        let budget = CompatBudget { max_len: agent.max_len, head_max_len: agent.head_max_len };
176        let mut runner = Runner::open(&model, self.device, self.precision)?;
177        if self.preload {
178            for b in Buckets::default().stage("compat").iter().take(4) {
179                runner.prepare(*b)?;
180            }
181        }
182        let buckets = Buckets::default();
183        let largest = buckets.stage("compat").last().copied().unwrap_or(kime_tensor::Bucket {
184            tokens: 16384,
185            seqs: 1024,
186            markers: 8192,
187        });
188        Ok(Kime {
189            inner: Arc::new(Inner {
190                id: model.spec.id.clone(),
191                mask: tok.mask_text().to_string(),
192                tok,
193                budget,
194                temps,
195                limits: (largest.tokens, largest.seqs, largest.markers),
196                runner: Mutex::new(Session {
197                    runner,
198                    buf: BatchBuf::default(),
199                    out: Outputs::default(),
200                }),
201            }),
202        })
203    }
204}
205
206enum Runner {
207    Cpu(Box<Executor<kime_cpu::CpuBackend>>),
208    #[cfg(feature = "cuda")]
209    Cuda(Box<Executor<kime_cuda::CudaBackend>>),
210}
211
212impl Runner {
213    fn open(model: &Model, device: Device, precision: Precision) -> Result<Self, Error> {
214        let cpu = |threads: usize| {
215            let t = if threads == 0 { kime_cpu::par::available() } else { threads };
216            let backend = kime_cpu::CpuBackend::new(t).with_int8(precision == Precision::Int8);
217            Ok(Runner::Cpu(Box::new(kime_cpu::executor_with(model, backend)?)))
218        };
219        match device {
220            Device::Cpu { threads } => cpu(threads),
221            #[cfg(feature = "cuda")]
222            Device::Cuda(n) => {
223                Ok(Runner::Cuda(Box::new(kime_cuda::executor(model, n, cuda(precision)?)?)))
224            }
225            #[cfg(feature = "cuda")]
226            Device::Auto => {
227                match cuda(precision).and_then(|p| Ok(kime_cuda::executor(model, 0, p)?)) {
228                    Ok(e) => Ok(Runner::Cuda(Box::new(e))),
229                    Err(_) => cpu(0),
230                }
231            }
232            #[cfg(not(feature = "cuda"))]
233            Device::Auto => cpu(0),
234            #[cfg(not(feature = "cuda"))]
235            Device::Cuda(_) => Err(Error::Unsupported("this build has no CUDA backend".into())),
236            Device::Metal | Device::Ane => {
237                Err(Error::Unsupported("the Apple backends arrive with M4".into()))
238            }
239        }
240    }
241
242    fn prepare(&mut self, b: kime_tensor::Bucket) -> Result<(), Error> {
243        match self {
244            Runner::Cpu(e) => e.prepare(b)?,
245            #[cfg(feature = "cuda")]
246            Runner::Cuda(e) => e.prepare(b)?,
247        };
248        Ok(())
249    }
250
251    fn run(&mut self, buf: &BatchBuf, out: &mut Outputs) -> Result<(), Error> {
252        match self {
253            Runner::Cpu(e) => e.run(&buf.batch(), out)?,
254            #[cfg(feature = "cuda")]
255            Runner::Cuda(e) => e.run(&buf.batch(), out)?,
256        };
257        Ok(())
258    }
259
260    fn describe(&self) -> String {
261        match self {
262            Runner::Cpu(e) => {
263                let int8 = if e.backend().int8() { ", int8" } else { "" };
264                format!("cpu, {} threads{int8}", kime_tensor::Backend::caps(e.backend()).threads)
265            }
266            #[cfg(feature = "cuda")]
267            Runner::Cuda(e) => format!("cuda, {}", e.backend().name()),
268        }
269    }
270}
271
272#[cfg(feature = "cuda")]
273fn cuda(p: Precision) -> Result<kime_cuda::Precision, Error> {
274    match p {
275        Precision::F16 => Ok(kime_cuda::Precision::F16),
276        Precision::F32 => Ok(kime_cuda::Precision::F32),
277        Precision::Int8 => Err(Error::Unsupported("INT8 runs on the CPU only for now".into())),
278    }
279}
280
281struct Session {
282    runner: Runner,
283    buf: BatchBuf,
284    out: Outputs,
285}
286
287struct Inner {
288    id: String,
289    tok: Tokenizer,
290    mask: String,
291    budget: CompatBudget,
292    temps: Temperatures,
293    /// The largest bucket: tokens, sequences and markers one batch can hold.
294    limits: (usize, usize, usize),
295    runner: Mutex<Session>,
296}
297
298/// A loaded model on a device. Clones share it, and it can be used from any thread.
299#[derive(Clone)]
300pub struct Kime {
301    inner: Arc<Inner>,
302}
303
304impl fmt::Debug for Kime {
305    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
306        f.debug_struct("Kime").field("model", &self.inner.id).finish_non_exhaustive()
307    }
308}
309
310/// One laid out question and where its answer goes.
311struct Item<'a> {
312    req: usize,
313    q: &'a Question,
314    seq: CompatSequence,
315    logits: Vec<f32>,
316    act: [f32; 2],
317}
318
319impl Kime {
320    /// Settings for a new engine.
321    #[must_use]
322    pub fn builder() -> Builder {
323        Builder::default()
324    }
325
326    /// The model's name, `laya` or `laya-multilingual` for the published checkpoints.
327    #[must_use]
328    pub fn model_id(&self) -> &str {
329        &self.inner.id
330    }
331
332    /// The device the model runs on, in words.
333    #[must_use]
334    pub fn device(&self) -> String {
335        self.lock().runner.describe()
336    }
337
338    fn lock(&self) -> std::sync::MutexGuard<'_, Session> {
339        self.inner.runner.lock().unwrap_or_else(PoisonError::into_inner)
340    }
341
342    /// Answers one request.
343    ///
344    /// # Errors
345    ///
346    /// [`Error::Invalid`] for a request that does not validate, [`Error::TooLong`] for a question
347    /// whose options do not fit, and device errors.
348    pub fn decide(&self, req: &Request) -> Result<Response, Error> {
349        Ok(self.decide_batch(std::slice::from_ref(req))?.remove(0))
350    }
351
352    /// Answers many requests, packing all their questions into as few device batches as fit.
353    /// Each answer is the same bits it would be alone, whatever else is in the batch.
354    ///
355    /// # Errors
356    ///
357    /// As [`Kime::decide`]. One bad request fails the whole call.
358    pub fn decide_batch(&self, reqs: &[Request]) -> Result<Vec<Response>, Error> {
359        let inner = &*self.inner;
360        let mut parsed = Vec::with_capacity(reqs.len());
361        for r in reqs {
362            parsed.push(parse(&r.to_json(), &Limits::LAYA).map_err(Error::Invalid)?);
363        }
364        let mut items = Vec::new();
365        for (i, r) in parsed.iter().enumerate() {
366            if r.questions.is_empty() {
367                continue;
368            }
369            // Laya keeps the end of a conversation and the start of anything else.
370            let cut = if matches!(r.state, Value::Array(_)) { Cut::Head } else { Cut::Tail };
371            let state = inner.tok.encode_state(&compat_state(&r.state, &inner.mask));
372            for q in &r.questions {
373                let text = compat_question(q, &inner.mask);
374                let seq =
375                    inner.tok.compat_sequence(&text.head, &text.options, &state, inner.budget, cut);
376                if seq.markers.len() != q.criteria.len() {
377                    return Err(Error::TooLong {
378                        question: q.id.clone(),
379                        options: q.criteria.len(),
380                        fit: seq.markers.len(),
381                    });
382                }
383                items.push(Item { req: i, q, seq, logits: Vec::new(), act: [0.0; 2] });
384            }
385        }
386        self.run(&mut items)?;
387        let mut out: Vec<Response> = parsed
388            .iter()
389            .map(|_| Response { model: LAYA_MODEL.into(), answers: Vec::new(), input_tokens: 0 })
390            .collect();
391        for it in &items {
392            let res = &mut out[it.req];
393            res.input_tokens += it.seq.ids.len();
394            res.answers
395                .push((it.q.id.clone(), laya_answer(it.q, &it.logits, it.act, &inner.temps)));
396        }
397        Ok(out)
398    }
399
400    /// Runs every item in as few batches as the largest bucket allows, in order.
401    fn run(&self, items: &mut [Item<'_>]) -> Result<(), Error> {
402        let (max_t, max_s, max_m) = self.inner.limits;
403        let mut s = self.lock();
404        let Session { runner, buf, out } = &mut *s;
405        let mut start = 0;
406        while start < items.len() {
407            let (mut t, mut m, mut end) = (0, 0, start);
408            while end < items.len() && end - start < max_s {
409                let it = &items[end];
410                if end > start && (t + it.seq.ids.len() > max_t || m + it.seq.markers.len() > max_m)
411                {
412                    break;
413                }
414                t += it.seq.ids.len();
415                m += it.seq.markers.len();
416                end += 1;
417            }
418            buf.clear();
419            for it in &items[start..end] {
420                buf.push(&it.seq.ids, &it.seq.markers, it.q.qtype.index() as u8);
421            }
422            runner.run(buf, out)?;
423            let mut at = 0;
424            for (it, a) in items[start..end].iter_mut().zip(&out.act) {
425                let k = it.seq.markers.len();
426                it.logits = out.logits[at..at + k].to_vec();
427                it.act = *a;
428                at += k;
429            }
430            start = end;
431        }
432        Ok(())
433    }
434
435    /// [`Kime::decide`] on a thread of its own, for async callers. It works with any executor,
436    /// since it needs nothing from one but a waker.
437    #[must_use]
438    pub fn decide_async(&self, req: &Request) -> Decision {
439        let shared = Arc::new(Mutex::new((None, None::<Waker>)));
440        let (kime, req, done) = (self.clone(), req.clone(), shared.clone());
441        std::thread::spawn(move || {
442            let r = kime.decide(&req);
443            let mut g = done.lock().unwrap_or_else(PoisonError::into_inner);
444            g.0 = Some(r);
445            if let Some(w) = g.1.take() {
446                w.wake();
447            }
448        });
449        Decision { shared }
450    }
451}
452
453/// The future [`Kime::decide_async`] returns.
454#[derive(Debug)]
455pub struct Decision {
456    #[allow(clippy::type_complexity)]
457    shared: Arc<Mutex<(Option<Result<Response, Error>>, Option<Waker>)>>,
458}
459
460impl Future for Decision {
461    type Output = Result<Response, Error>;
462
463    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
464        let mut g = self.shared.lock().unwrap_or_else(PoisonError::into_inner);
465        match g.0.take() {
466            Some(r) => Poll::Ready(r),
467            None => {
468                g.1 = Some(cx.waker().clone());
469                Poll::Pending
470            }
471        }
472    }
473}
kime_engine/lib.rs

kime_engine/
lib.rs
