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

kime_engine/
lib.rs
