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::atomic::{AtomicUsize, Ordering};
17use std::sync::{Arc, Mutex, PoisonError};
18use std::task::{Context, Poll, Waker};
19use std::time::{Duration, Instant};
20
21use kime_core::answer::{LAYA_MODEL, Response, Temperatures, laya_answer};
22use kime_core::render::{compat_question, compat_state};
23use kime_core::request::{Limits, Problem, Question, Request, parse};
24use kime_model::Model;
25use kime_tensor::{BatchBuf, Buckets, Executor, Outputs};
26use kime_tok::Tokenizer;
27use kime_tok::layout::{CompatBudget, CompatSequence, Cut};
28use serde_json::Value;
29
30pub mod hub;
31mod split;
32
33/// Where the model runs.
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
35pub enum Device {
36    /// The first CUDA GPU when there is one and this build has CUDA, the CPU otherwise.
37    #[default]
38    Auto,
39    /// The CPU, on `threads` threads, or every core for 0.
40    Cpu {
41        /// Worker threads.
42        threads: usize,
43    },
44    /// A CUDA GPU by ordinal.
45    Cuda(usize),
46    /// The Apple GPU. Arrives with M4.
47    Metal,
48    /// The Apple Neural Engine. Arrives with M4.
49    Ane,
50}
51
52/// The number format. The CPU computes in FP32 for both float settings.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
54pub enum Precision {
55    /// FP16 weights and GEMM inputs with FP32 accumulation, as Laya's autocast runs on a GPU.
56    #[default]
57    F16,
58    /// FP32 throughout, closest to Laya on the CPU.
59    F32,
60    /// INT8 weights and activations in the encoder and decision head GEMMs on the CPU, with the
61    /// scorer in FP32. Faster, and further from Laya than FP32: see spec/10-cpu.md for the gate a
62    /// checkpoint has to pass. Not on GPUs yet.
63    Int8,
64}
65
66/// What can go wrong.
67#[derive(Debug)]
68pub enum Error {
69    /// The request does not validate. Every problem is listed, in the wire format.
70    Invalid(Vec<Problem>),
71    /// The model name did not resolve.
72    NotFound(String),
73    /// The checkpoint did not load.
74    Model(kime_model::Error),
75    /// The checkpoint's tokenizer did not load.
76    Tokenizer(String),
77    /// The device failed, or the batch did not fit it.
78    Backend(kime_tensor::Error),
79    /// A question's head and options do not fit the checkpoint's budget, so some options have no
80    /// marker. Laya raises the same error.
81    TooLong {
82        /// The question id.
83        question: String,
84        /// Its options.
85        options: usize,
86        /// The options that fit.
87        fit: usize,
88    },
89    /// The device asked for is not in this build or not on this machine.
90    Unsupported(String),
91}
92
93impl fmt::Display for Error {
94    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95        match self {
96            Error::Invalid(p) => {
97                write!(f, "invalid request:")?;
98                for p in p {
99                    write!(f, " {};", p.to_json())?;
100                }
101                Ok(())
102            }
103            Error::NotFound(m) | Error::Tokenizer(m) | Error::Unsupported(m) => f.write_str(m),
104            Error::Model(e) => write!(f, "{e}"),
105            Error::Backend(e) => write!(f, "{e}"),
106            Error::TooLong { question, options, fit } => write!(
107                f,
108                "question {question:?} has {options} options but only {fit} fit the head budget"
109            ),
110        }
111    }
112}
113
114impl std::error::Error for Error {}
115
116impl From<kime_tensor::Error> for Error {
117    fn from(e: kime_tensor::Error) -> Self {
118        Error::Backend(e)
119    }
120}
121
122/// Settings for [`Kime`], from [`Kime::builder`].
123#[derive(Debug, Clone, Default)]
124pub struct Builder {
125    model: Option<String>,
126    device: Device,
127    precision: Precision,
128    preload: bool,
129}
130
131impl Builder {
132    /// The model: an alias such as `laya`, a checkpoint directory, a `.kime` file or an
133    /// `hf://org/repo[/subfolder]` reference. See [`hub`].
134    #[must_use]
135    pub fn model(mut self, name: impl Into<String>) -> Self {
136        self.model = Some(name.into());
137        self
138    }
139
140    /// Where to run.
141    #[must_use]
142    pub fn device(mut self, d: Device) -> Self {
143        self.device = d;
144        self
145    }
146
147    /// The number format on a GPU.
148    #[must_use]
149    pub fn precision(mut self, p: Precision) -> Self {
150        self.precision = p;
151        self
152    }
153
154    /// Builds the plans for the smallest batch shapes now, so the first requests do not pay for
155    /// them.
156    #[must_use]
157    pub fn preload(mut self, yes: bool) -> Self {
158        self.preload = yes;
159        self
160    }
161
162    /// Loads the model onto the device.
163    ///
164    /// # Errors
165    ///
166    /// When the model is not found or does not load, or the device cannot be opened.
167    pub fn build(self) -> Result<Kime, Error> {
168        let name = self.model.unwrap_or_else(|| "laya".into());
169        let path = hub::resolve(&name).map_err(Error::NotFound)?;
170        let model = Model::open(&path).map_err(Error::Model)?;
171        let tok_json = model
172            .file("tokenizer/tokenizer.json")
173            .ok_or_else(|| Error::Tokenizer(format!("{}: no tokenizer.json", path.display())))?;
174        let tok = Tokenizer::from_bytes(tok_json, model.file("tokenizer/tokenizer_config.json"))
175            .map_err(|e| Error::Tokenizer(e.to_string()))?;
176        let agent = &model.spec.agent;
177        let temps = Temperatures::new(agent.temperature, &agent.temperature_by_options);
178        let budget = CompatBudget { max_len: agent.max_len, head_max_len: agent.head_max_len };
179        let mut runner = Runner::open(&model, self.device, self.precision)?;
180        if self.preload {
181            for b in Buckets::default().stage("compat").iter().take(4) {
182                runner.prepare(*b)?;
183            }
184        }
185        let buckets = Buckets::default().stage("compat").to_vec();
186        let (weights, plans) = runner.memory();
187        Ok(Kime {
188            inner: Arc::new(Inner {
189                id: model.spec.id.clone(),
190                mask: tok.mask_text().to_string(),
191                tok,
192                budget,
193                temps,
194                buckets,
195                memory: [AtomicUsize::new(weights), AtomicUsize::new(plans)],
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 memory(&self) -> (usize, usize) {
261        match self {
262            Runner::Cpu(e) => e.memory(),
263            #[cfg(feature = "cuda")]
264            Runner::Cuda(e) => e.memory(),
265        }
266    }
267
268    fn describe(&self) -> String {
269        match self {
270            Runner::Cpu(e) => {
271                let int8 = if e.backend().int8() { ", int8" } else { "" };
272                format!("cpu, {} threads{int8}", kime_tensor::Backend::caps(e.backend()).threads)
273            }
274            #[cfg(feature = "cuda")]
275            Runner::Cuda(e) => format!("cuda, {}", e.backend().name()),
276        }
277    }
278}
279
280#[cfg(feature = "cuda")]
281fn cuda(p: Precision) -> Result<kime_cuda::Precision, Error> {
282    match p {
283        Precision::F16 => Ok(kime_cuda::Precision::F16),
284        Precision::F32 => Ok(kime_cuda::Precision::F32),
285        Precision::Int8 => Err(Error::Unsupported("INT8 runs on the CPU only for now".into())),
286    }
287}
288
289struct Session {
290    runner: Runner,
291    buf: BatchBuf,
292    out: Outputs,
293}
294
295struct Inner {
296    id: String,
297    tok: Tokenizer,
298    mask: String,
299    budget: CompatBudget,
300    temps: Temperatures,
301    /// The compat buckets, smallest first.
302    buckets: Vec<kime_tensor::Bucket>,
303    runner: Mutex<Session>,
304    /// [`Memory`], kept up to date after every forward pass so reading it needs no lock.
305    memory: [AtomicUsize; 2],
306}
307
308/// A loaded model on a device. Clones share it, and it can be used from any thread.
309#[derive(Clone)]
310pub struct Kime {
311    inner: Arc<Inner>,
312}
313
314impl fmt::Debug for Kime {
315    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
316        f.debug_struct("Kime").field("model", &self.inner.id).finish_non_exhaustive()
317    }
318}
319
320/// Where the time of one [`Kime::decide_batch_timed`] call went.
321#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
322pub struct Timing {
323    /// Validating the requests and laying their questions out as token ids.
324    pub tokenize: Duration,
325    /// The device batches, from the first upload to the last result copied back.
326    pub device: Duration,
327    /// How many device batches the questions took.
328    pub batches: usize,
329    /// Questions whose state was cut to fit the model's sequence length.
330    pub truncated: usize,
331    /// State tokens left out of those questions.
332    pub cut_tokens: usize,
333}
334
335/// Bytes a model holds on its device.
336#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
337pub struct Memory {
338    /// The weights, in the device's layout.
339    pub weights: usize,
340    /// The plans built so far, one per bucket used: their arenas and staging buffers.
341    pub plans: usize,
342}
343
344/// One laid out question and where its answer goes.
345struct Item<'a> {
346    req: usize,
347    q: &'a Question,
348    seq: CompatSequence,
349    logits: Vec<f32>,
350    act: [f32; 2],
351}
352
353impl Kime {
354    /// Settings for a new engine.
355    #[must_use]
356    pub fn builder() -> Builder {
357        Builder::default()
358    }
359
360    /// The model's name, `laya` or `laya-multilingual` for the published checkpoints.
361    #[must_use]
362    pub fn model_id(&self) -> &str {
363        &self.inner.id
364    }
365
366    /// The device the model runs on, in words.
367    #[must_use]
368    pub fn device(&self) -> String {
369        self.lock().runner.describe()
370    }
371
372    /// The most tokens one question's row can hold: the state, the question and its options.
373    /// Longer states are cut to fit.
374    #[must_use]
375    pub fn max_row_tokens(&self) -> usize {
376        self.inner.budget.max_len
377    }
378
379    /// The bytes the model holds on its device, as of the last forward pass.
380    #[must_use]
381    pub fn memory(&self) -> Memory {
382        let m = &self.inner.memory;
383        Memory { weights: m[0].load(Ordering::Relaxed), plans: m[1].load(Ordering::Relaxed) }
384    }
385
386    fn lock(&self) -> std::sync::MutexGuard<'_, Session> {
387        self.inner.runner.lock().unwrap_or_else(PoisonError::into_inner)
388    }
389
390    /// The tokens a request holds before anything is cut: its state once plus every question
391    /// with its options. The server refuses a request over its limit with this count, before it
392    /// is queued.
393    #[must_use]
394    pub fn count_tokens(&self, req: &Request) -> usize {
395        let inner = &*self.inner;
396        let mut n = inner.tok.encode(&compat_state(&req.state, &inner.mask)).len();
397        for q in &req.questions {
398            let text = compat_question(q, &inner.mask);
399            n += inner.tok.encode(&text.head).len();
400            n += text.options.iter().map(|o| inner.tok.encode(o).len()).sum::<usize>();
401        }
402        n
403    }
404
405    /// Answers one request.
406    ///
407    /// # Errors
408    ///
409    /// [`Error::Invalid`] for a request that does not validate, [`Error::TooLong`] for a question
410    /// whose options do not fit, and device errors.
411    pub fn decide(&self, req: &Request) -> Result<Response, Error> {
412        Ok(self.decide_batch(std::slice::from_ref(req))?.remove(0))
413    }
414
415    /// Answers many requests, packing all their questions into as few device batches as fit.
416    /// Each answer is the same bits it would be alone, whatever else is in the batch.
417    ///
418    /// # Errors
419    ///
420    /// As [`Kime::decide`]. One bad request fails the whole call.
421    pub fn decide_batch(&self, reqs: &[Request]) -> Result<Vec<Response>, Error> {
422        Ok(self.decide_batch_timed(reqs)?.0)
423    }
424
425    /// [`Kime::decide_batch`], and where the time went.
426    ///
427    /// # Errors
428    ///
429    /// As [`Kime::decide_batch`].
430    pub fn decide_batch_timed(&self, reqs: &[Request]) -> Result<(Vec<Response>, Timing), Error> {
431        let inner = &*self.inner;
432        let t0 = Instant::now();
433        let mut parsed = Vec::with_capacity(reqs.len());
434        for r in reqs {
435            parsed.push(parse(&r.to_json(), &Limits::LAYA).map_err(Error::Invalid)?);
436        }
437        let mut items = Vec::new();
438        for (i, r) in parsed.iter().enumerate() {
439            if r.questions.is_empty() {
440                continue;
441            }
442            // Laya keeps the end of a conversation and the start of anything else.
443            let cut = if matches!(r.state, Value::Array(_)) { Cut::Head } else { Cut::Tail };
444            let state = inner.tok.encode_state(&compat_state(&r.state, &inner.mask));
445            for q in &r.questions {
446                let text = compat_question(q, &inner.mask);
447                let seq =
448                    inner.tok.compat_sequence(&text.head, &text.options, &state, inner.budget, cut);
449                if seq.markers.len() != q.criteria.len() {
450                    return Err(Error::TooLong {
451                        question: q.id.clone(),
452                        options: q.criteria.len(),
453                        fit: seq.markers.len(),
454                    });
455                }
456                items.push(Item { req: i, q, seq, logits: Vec::new(), act: [0.0; 2] });
457            }
458        }
459        let tokenize = t0.elapsed();
460        let t1 = Instant::now();
461        let batches = self.run(&mut items)?;
462        let cut = items.iter().map(|it| it.seq.state_tokens - it.seq.state_tokens_used);
463        let timing = Timing {
464            tokenize,
465            device: t1.elapsed(),
466            batches,
467            truncated: cut.clone().filter(|&n| n > 0).count(),
468            cut_tokens: cut.sum(),
469        };
470        let mut out: Vec<Response> = parsed
471            .iter()
472            .map(|_| Response { model: LAYA_MODEL.into(), answers: Vec::new(), input_tokens: 0 })
473            .collect();
474        for it in &items {
475            let res = &mut out[it.req];
476            res.input_tokens += it.seq.ids.len();
477            res.answers
478                .push((it.q.id.clone(), laya_answer(it.q, &it.logits, it.act, &inner.temps)));
479        }
480        Ok((out, timing))
481    }
482
483    /// Runs every item in the batches [`split::split`] picks, and says how many it took.
484    fn run(&self, items: &mut [Item<'_>]) -> Result<usize, Error> {
485        let sizes: Vec<(usize, usize)> =
486            items.iter().map(|it| (it.seq.ids.len(), it.seq.markers.len())).collect();
487        let batches = split::split(&self.inner.buckets, &sizes);
488        let mut s = self.lock();
489        let Session { runner, buf, out } = &mut *s;
490        for batch in &batches {
491            buf.clear();
492            for &i in batch {
493                let it = &items[i];
494                buf.push(&it.seq.ids, &it.seq.markers, it.q.qtype.index() as u8);
495            }
496            runner.run(buf, out)?;
497            let mut at = 0;
498            for (&i, a) in batch.iter().zip(&out.act) {
499                let it = &mut items[i];
500                let k = it.seq.markers.len();
501                it.logits = out.logits[at..at + k].to_vec();
502                it.act = *a;
503                at += k;
504            }
505        }
506        let (weights, plans) = runner.memory();
507        self.inner.memory[0].store(weights, Ordering::Relaxed);
508        self.inner.memory[1].store(plans, Ordering::Relaxed);
509        Ok(batches.len())
510    }
511
512    /// [`Kime::decide`] on a thread of its own, for async callers. It works with any executor,
513    /// since it needs nothing from one but a waker.
514    #[must_use]
515    pub fn decide_async(&self, req: &Request) -> Decision {
516        let shared = Arc::new(Mutex::new((None, None::<Waker>)));
517        let (kime, req, done) = (self.clone(), req.clone(), shared.clone());
518        std::thread::spawn(move || {
519            let r = kime.decide(&req);
520            let mut g = done.lock().unwrap_or_else(PoisonError::into_inner);
521            g.0 = Some(r);
522            if let Some(w) = g.1.take() {
523                w.wake();
524            }
525        });
526        Decision { shared }
527    }
528}
529
530/// The future [`Kime::decide_async`] returns.
531#[derive(Debug)]
532pub struct Decision {
533    #[allow(clippy::type_complexity)]
534    shared: Arc<Mutex<(Option<Result<Response, Error>>, Option<Waker>)>>,
535}
536
537impl Future for Decision {
538    type Output = Result<Response, Error>;
539
540    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
541        let mut g = self.shared.lock().unwrap_or_else(PoisonError::into_inner);
542        match g.0.take() {
543            Some(r) => Poll::Ready(r),
544            None => {
545                g.1 = Some(cx.waker().clone());
546                Poll::Pending
547            }
548        }
549    }
550}
kime_engine/lib.rs

kime_engine/
lib.rs
