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
30mod cache;
31pub mod hub;
32mod split;
33
34use cache::{AnswerCache, Entry};
35pub use cache::{CacheMode, CacheStats};
36
37/// Where the model runs.
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
39pub enum Device {
40    /// The first CUDA GPU when there is one and this build has CUDA, the CPU otherwise.
41    #[default]
42    Auto,
43    /// The CPU, on `threads` threads, or every core for 0.
44    Cpu {
45        /// Worker threads.
46        threads: usize,
47    },
48    /// A CUDA GPU by ordinal.
49    Cuda(usize),
50    /// The Apple GPU. Arrives with M4.
51    Metal,
52    /// The Apple Neural Engine. Arrives with M4.
53    Ane,
54}
55
56/// The number format. The CPU computes in FP32 for both float settings.
57#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
58pub enum Precision {
59    /// FP16 weights and GEMM inputs with FP32 accumulation, as Laya's autocast runs on a GPU.
60    #[default]
61    F16,
62    /// FP32 throughout, closest to Laya on the CPU.
63    F32,
64    /// INT8 weights and activations in the encoder and decision head GEMMs on the CPU, with the
65    /// scorer in FP32. Faster, and further from Laya than FP32: see spec/10-cpu.md for the gate a
66    /// checkpoint has to pass. Not on GPUs yet.
67    Int8,
68}
69
70/// What can go wrong.
71#[derive(Debug)]
72pub enum Error {
73    /// The request does not validate. Every problem is listed, in the wire format.
74    Invalid(Vec<Problem>),
75    /// The model name did not resolve.
76    NotFound(String),
77    /// The checkpoint did not load.
78    Model(kime_model::Error),
79    /// The checkpoint's tokenizer did not load.
80    Tokenizer(String),
81    /// The device failed, or the batch did not fit it.
82    Backend(kime_tensor::Error),
83    /// A question's head and options do not fit the checkpoint's budget, so some options have no
84    /// marker. Laya raises the same error.
85    TooLong {
86        /// The question id.
87        question: String,
88        /// Its options.
89        options: usize,
90        /// The options that fit.
91        fit: usize,
92    },
93    /// The device asked for is not in this build or not on this machine.
94    Unsupported(String),
95}
96
97impl fmt::Display for Error {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        match self {
100            Error::Invalid(p) => {
101                write!(f, "invalid request:")?;
102                for p in p {
103                    write!(f, " {};", p.to_json())?;
104                }
105                Ok(())
106            }
107            Error::NotFound(m) | Error::Tokenizer(m) | Error::Unsupported(m) => f.write_str(m),
108            Error::Model(e) => write!(f, "{e}"),
109            Error::Backend(e) => write!(f, "{e}"),
110            Error::TooLong { question, options, fit } => write!(
111                f,
112                "question {question:?} has {options} options but only {fit} fit the head budget"
113            ),
114        }
115    }
116}
117
118impl std::error::Error for Error {}
119
120impl From<kime_tensor::Error> for Error {
121    fn from(e: kime_tensor::Error) -> Self {
122        Error::Backend(e)
123    }
124}
125
126/// Settings for [`Kime`], from [`Kime::builder`].
127#[derive(Debug, Clone, Default)]
128pub struct Builder {
129    model: Option<String>,
130    device: Device,
131    precision: Precision,
132    preload: bool,
133    answer_cache: usize,
134}
135
136impl Builder {
137    /// The model: an alias such as `laya`, a checkpoint directory, a `.kime` file or an
138    /// `hf://org/repo[/subfolder]` reference. See [`hub`].
139    #[must_use]
140    pub fn model(mut self, name: impl Into<String>) -> Self {
141        self.model = Some(name.into());
142        self
143    }
144
145    /// Where to run.
146    #[must_use]
147    pub fn device(mut self, d: Device) -> Self {
148        self.device = d;
149        self
150    }
151
152    /// The number format on a GPU.
153    #[must_use]
154    pub fn precision(mut self, p: Precision) -> Self {
155        self.precision = p;
156        self
157    }
158
159    /// Builds the plans for the smallest batch shapes now, so the first requests do not pay for
160    /// them.
161    #[must_use]
162    pub fn preload(mut self, yes: bool) -> Self {
163        self.preload = yes;
164        self
165    }
166
167    /// Keeps the answers to about `entries` questions, so the same question on the same state
168    /// is answered again without the device. 0, the default, keeps none.
169    #[must_use]
170    pub fn answer_cache(mut self, entries: usize) -> Self {
171        self.answer_cache = entries;
172        self
173    }
174
175    /// Loads the model onto the device.
176    ///
177    /// # Errors
178    ///
179    /// When the model is not found or does not load, or the device cannot be opened.
180    pub fn build(self) -> Result<Kime, Error> {
181        let name = self.model.unwrap_or_else(|| "laya".into());
182        let path = hub::resolve(&name).map_err(Error::NotFound)?;
183        let model = Model::open(&path).map_err(Error::Model)?;
184        let tok_json = model
185            .file("tokenizer/tokenizer.json")
186            .ok_or_else(|| Error::Tokenizer(format!("{}: no tokenizer.json", path.display())))?;
187        let tok = Tokenizer::from_bytes(tok_json, model.file("tokenizer/tokenizer_config.json"))
188            .map_err(|e| Error::Tokenizer(e.to_string()))?;
189        let agent = &model.spec.agent;
190        let temps = Temperatures::new(agent.temperature, &agent.temperature_by_options);
191        let budget = CompatBudget { max_len: agent.max_len, head_max_len: agent.head_max_len };
192        let mut runner = Runner::open(&model, self.device, self.precision)?;
193        let embed = runner.add_graph(model.graph.embed_plan(&model.spec));
194        if self.preload {
195            for b in Buckets::default().stage("compat").iter().take(4) {
196                runner.prepare(*b)?;
197            }
198        }
199        let buckets = Buckets::default().stage("compat").to_vec();
200        let embed_buckets = Buckets::default().stage(EMBED_STAGE).to_vec();
201        let (weights, plans) = runner.memory();
202        Ok(Kime {
203            inner: Arc::new(Inner {
204                id: model.spec.id.clone(),
205                mask: tok.mask_text().to_string(),
206                tok,
207                budget,
208                temps,
209                buckets,
210                embed_buckets,
211                d: model.spec.encoder.d,
212                memory: [AtomicUsize::new(weights), AtomicUsize::new(plans)],
213                cache: (self.answer_cache > 0).then(|| AnswerCache::new(self.answer_cache)),
214                runner: Mutex::new(Session {
215                    runner,
216                    embed,
217                    buf: BatchBuf::default(),
218                    out: Outputs::default(),
219                }),
220            }),
221        })
222    }
223}
224
225enum Runner {
226    Cpu(Box<Executor<kime_cpu::CpuBackend>>),
227    #[cfg(feature = "cuda")]
228    Cuda(Box<Executor<kime_cuda::CudaBackend>>),
229}
230
231impl Runner {
232    fn open(model: &Model, device: Device, precision: Precision) -> Result<Self, Error> {
233        let cpu = |threads: usize| {
234            let t = if threads == 0 { kime_cpu::par::available() } else { threads };
235            let backend = kime_cpu::CpuBackend::new(t).with_int8(precision == Precision::Int8);
236            Ok(Runner::Cpu(Box::new(kime_cpu::executor_with(model, backend)?)))
237        };
238        match device {
239            Device::Cpu { threads } => cpu(threads),
240            #[cfg(feature = "cuda")]
241            Device::Cuda(n) => {
242                Ok(Runner::Cuda(Box::new(kime_cuda::executor(model, n, cuda(precision)?)?)))
243            }
244            #[cfg(feature = "cuda")]
245            Device::Auto => {
246                match cuda(precision).and_then(|p| Ok(kime_cuda::executor(model, 0, p)?)) {
247                    Ok(e) => Ok(Runner::Cuda(Box::new(e))),
248                    Err(_) => cpu(0),
249                }
250            }
251            #[cfg(not(feature = "cuda"))]
252            Device::Auto => cpu(0),
253            #[cfg(not(feature = "cuda"))]
254            Device::Cuda(_) => Err(Error::Unsupported("this build has no CUDA backend".into())),
255            Device::Metal | Device::Ane => {
256                Err(Error::Unsupported("the Apple backends arrive with M4".into()))
257            }
258        }
259    }
260
261    fn add_graph(&mut self, g: kime_tensor::Graph) -> usize {
262        let b = Buckets::default();
263        match self {
264            Runner::Cpu(e) => e.add_graph(g, &b, EMBED_STAGE),
265            #[cfg(feature = "cuda")]
266            Runner::Cuda(e) => e.add_graph(g, &b, EMBED_STAGE),
267        }
268    }
269
270    fn prepare(&mut self, b: kime_tensor::Bucket) -> Result<(), Error> {
271        match self {
272            Runner::Cpu(e) => e.prepare(b)?,
273            #[cfg(feature = "cuda")]
274            Runner::Cuda(e) => e.prepare(b)?,
275        };
276        Ok(())
277    }
278
279    fn run(&mut self, buf: &BatchBuf, out: &mut Outputs) -> Result<(), Error> {
280        self.run_lane(0, buf, out)
281    }
282
283    fn run_lane(&mut self, lane: usize, buf: &BatchBuf, out: &mut Outputs) -> Result<(), Error> {
284        match self {
285            Runner::Cpu(e) => e.run_lane(lane, &buf.batch(), out)?,
286            #[cfg(feature = "cuda")]
287            Runner::Cuda(e) => e.run_lane(lane, &buf.batch(), out)?,
288        };
289        Ok(())
290    }
291
292    fn memory(&self) -> (usize, usize) {
293        match self {
294            Runner::Cpu(e) => e.memory(),
295            #[cfg(feature = "cuda")]
296            Runner::Cuda(e) => e.memory(),
297        }
298    }
299
300    fn int8(&self) -> bool {
301        match self {
302            Runner::Cpu(e) => e.backend().int8(),
303            #[cfg(feature = "cuda")]
304            Runner::Cuda(_) => false,
305        }
306    }
307
308    fn describe(&self) -> String {
309        match self {
310            Runner::Cpu(e) => {
311                let int8 = if e.backend().int8() { ", int8" } else { "" };
312                format!("cpu, {} threads{int8}", kime_tensor::Backend::caps(e.backend()).threads)
313            }
314            #[cfg(feature = "cuda")]
315            Runner::Cuda(e) => format!("cuda, {}", e.backend().name()),
316        }
317    }
318}
319
320#[cfg(feature = "cuda")]
321fn cuda(p: Precision) -> Result<kime_cuda::Precision, Error> {
322    match p {
323        Precision::F16 => Ok(kime_cuda::Precision::F16),
324        Precision::F32 => Ok(kime_cuda::Precision::F32),
325        Precision::Int8 => Err(Error::Unsupported("INT8 runs on the CPU only for now".into())),
326    }
327}
328
329/// The buckets the pooled embedding runs in: sequences with no markers.
330const EMBED_STAGE: &str = "state";
331
332struct Session {
333    runner: Runner,
334    /// The lane of the pooled embedding graph on the runner.
335    embed: usize,
336    buf: BatchBuf,
337    out: Outputs,
338}
339
340struct Inner {
341    id: String,
342    tok: Tokenizer,
343    mask: String,
344    budget: CompatBudget,
345    temps: Temperatures,
346    /// The compat buckets, smallest first.
347    buckets: Vec<kime_tensor::Bucket>,
348    /// The buckets of the pooled embedding, smallest first.
349    embed_buckets: Vec<kime_tensor::Bucket>,
350    /// The width of the encoder, and of an embedding.
351    d: usize,
352    runner: Mutex<Session>,
353    /// [`Memory`], kept up to date after every forward pass so reading it needs no lock.
354    memory: [AtomicUsize; 2],
355    cache: Option<AnswerCache>,
356}
357
358/// A loaded model on a device. Clones share it, and it can be used from any thread.
359#[derive(Clone)]
360pub struct Kime {
361    inner: Arc<Inner>,
362}
363
364impl fmt::Debug for Kime {
365    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
366        f.debug_struct("Kime").field("model", &self.inner.id).finish_non_exhaustive()
367    }
368}
369
370/// Where the time of one [`Kime::decide_batch_timed`] call went.
371#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
372pub struct Timing {
373    /// Validating the requests and laying their questions out as token ids.
374    pub tokenize: Duration,
375    /// The device batches, from the first upload to the last result copied back.
376    pub device: Duration,
377    /// How many device batches the questions took.
378    pub batches: usize,
379    /// Questions whose state was cut to fit the model's sequence length.
380    pub truncated: usize,
381    /// State tokens left out of those questions.
382    pub cut_tokens: usize,
383    /// Questions answered from the answer cache, which took no device time.
384    pub cached: usize,
385}
386
387/// Bytes a model holds on its device.
388#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
389pub struct Memory {
390    /// The weights, in the device's layout.
391    pub weights: usize,
392    /// The plans built so far, one per bucket used: their arenas and staging buffers.
393    pub plans: usize,
394}
395
396/// One laid out question and where its answer goes.
397struct Item<'a> {
398    req: usize,
399    q: &'a Question,
400    seq: CompatSequence,
401    logits: Vec<f32>,
402    act: [f32; 2],
403}
404
405impl Item<'_> {
406    fn key(&self) -> cache::Key {
407        AnswerCache::key(&self.seq, self.q.qtype.index() as u8)
408    }
409
410    fn fill(&mut self, e: Entry) {
411        self.logits = e.logits.into_vec();
412        self.act = e.act;
413    }
414}
415
416/// The request's `kime.cache`, the default for anything but a mode's name.
417fn mode(req: &Request) -> CacheMode {
418    req.kime
419        .as_ref()
420        .and_then(|k| k.get("cache"))
421        .and_then(Value::as_str)
422        .and_then(CacheMode::parse)
423        .unwrap_or_default()
424}
425
426impl Kime {
427    /// Settings for a new engine.
428    #[must_use]
429    pub fn builder() -> Builder {
430        Builder::default()
431    }
432
433    /// The model's name, `laya` or `laya-multilingual` for the published checkpoints.
434    #[must_use]
435    pub fn model_id(&self) -> &str {
436        &self.inner.id
437    }
438
439    /// The device the model runs on, in words.
440    #[must_use]
441    pub fn device(&self) -> String {
442        self.lock().runner.describe()
443    }
444
445    /// The most tokens one question's row can hold: the state, the question and its options.
446    /// Longer states are cut to fit.
447    #[must_use]
448    pub fn max_row_tokens(&self) -> usize {
449        self.inner.budget.max_len
450    }
451
452    /// The bytes the model holds on its device, as of the last forward pass.
453    #[must_use]
454    pub fn memory(&self) -> Memory {
455        let m = &self.inner.memory;
456        Memory { weights: m[0].load(Ordering::Relaxed), plans: m[1].load(Ordering::Relaxed) }
457    }
458
459    /// The answer cache's counts, all 0 when it is off.
460    #[must_use]
461    pub fn cache_stats(&self) -> CacheStats {
462        self.inner.cache.as_ref().map(AnswerCache::stats).unwrap_or_default()
463    }
464
465    /// The answer to `req` when the answer cache holds every one of its questions, found without
466    /// the device lock, so the server can answer it without queueing it. `None` otherwise, and
467    /// then nothing is counted, since the request goes on to [`Kime::decide_batch`].
468    #[must_use]
469    pub fn cached(&self, req: &Request) -> Option<Response> {
470        let cache = self.inner.cache.as_ref()?;
471        if mode(req) != CacheMode::Use || req.questions.is_empty() {
472            return None;
473        }
474        let parsed = [parse(&req.to_json(), &Limits::LAYA).ok()?];
475        let mut items = self.lay_out(&parsed).ok()?;
476        let keys: Vec<_> = items.iter().map(Item::key).collect();
477        for (it, e) in items.iter_mut().zip(cache.all(&keys)?) {
478            it.fill(e);
479        }
480        self.respond(&parsed, &items).pop()
481    }
482
483    fn lock(&self) -> std::sync::MutexGuard<'_, Session> {
484        self.inner.runner.lock().unwrap_or_else(PoisonError::into_inner)
485    }
486
487    /// The tokens a request holds before anything is cut: its state once plus every question
488    /// with its options. The server refuses a request over its limit with this count, before it
489    /// is queued.
490    #[must_use]
491    pub fn count_tokens(&self, req: &Request) -> usize {
492        let inner = &*self.inner;
493        let mut n = inner.tok.encode(&compat_state(&req.state, &inner.mask)).len();
494        for q in &req.questions {
495            let text = compat_question(q, &inner.mask);
496            n += inner.tok.encode(&text.head).len();
497            n += text.options.iter().map(|o| inner.tok.encode(o).len()).sum::<usize>();
498        }
499        n
500    }
501
502    /// Answers one request.
503    ///
504    /// # Errors
505    ///
506    /// [`Error::Invalid`] for a request that does not validate, [`Error::TooLong`] for a question
507    /// whose options do not fit, and device errors.
508    pub fn decide(&self, req: &Request) -> Result<Response, Error> {
509        Ok(self.decide_batch(std::slice::from_ref(req))?.remove(0))
510    }
511
512    /// Answers many requests, packing all their questions into as few device batches as fit.
513    /// Each answer is the same bits it would be alone, whatever else is in the batch.
514    ///
515    /// # Errors
516    ///
517    /// As [`Kime::decide`]. One bad request fails the whole call.
518    pub fn decide_batch(&self, reqs: &[Request]) -> Result<Vec<Response>, Error> {
519        Ok(self.decide_batch_timed(reqs)?.0)
520    }
521
522    /// [`Kime::decide_batch`], and where the time went.
523    ///
524    /// # Errors
525    ///
526    /// As [`Kime::decide_batch`].
527    pub fn decide_batch_timed(&self, reqs: &[Request]) -> Result<(Vec<Response>, Timing), Error> {
528        let t0 = Instant::now();
529        let mut parsed = Vec::with_capacity(reqs.len());
530        for r in reqs {
531            parsed.push(parse(&r.to_json(), &Limits::LAYA).map_err(Error::Invalid)?);
532        }
533        let mut items = self.lay_out(&parsed)?;
534        let tokenize = t0.elapsed();
535        let t1 = Instant::now();
536        let (run, keys) = self.look_up(reqs, &mut items);
537        let batches = if run.is_empty() { 0 } else { self.run(&mut items, &run)? };
538        if let Some(c) = &self.inner.cache {
539            let keep = run.iter().filter(|&&i| mode(&reqs[items[i].req]) != CacheMode::Bypass);
540            c.insert(keep.map(|&i| {
541                (keys[i], Entry { logits: items[i].logits.clone().into(), act: items[i].act })
542            }));
543        }
544        let cached = items.len() - run.len();
545        let cut = items.iter().map(|it| it.seq.state_tokens - it.seq.state_tokens_used);
546        let timing = Timing {
547            tokenize,
548            device: t1.elapsed(),
549            batches,
550            truncated: cut.clone().filter(|&n| n > 0).count(),
551            cut_tokens: cut.sum(),
552            cached,
553        };
554        Ok((self.respond(&parsed, &items), timing))
555    }
556
557    /// Each question of `parsed` laid out as Laya lays it out, in request order.
558    fn lay_out<'a>(&self, parsed: &'a [Request]) -> Result<Vec<Item<'a>>, Error> {
559        let inner = &*self.inner;
560        let mut items = Vec::new();
561        for (i, r) in parsed.iter().enumerate() {
562            if r.questions.is_empty() {
563                continue;
564            }
565            // Laya keeps the end of a conversation and the start of anything else.
566            let cut = if matches!(r.state, Value::Array(_)) { Cut::Head } else { Cut::Tail };
567            let state = inner.tok.encode_state(&compat_state(&r.state, &inner.mask));
568            for q in &r.questions {
569                let text = compat_question(q, &inner.mask);
570                let seq =
571                    inner.tok.compat_sequence(&text.head, &text.options, &state, inner.budget, cut);
572                if seq.markers.len() != q.criteria.len() {
573                    return Err(Error::TooLong {
574                        question: q.id.clone(),
575                        options: q.criteria.len(),
576                        fit: seq.markers.len(),
577                    });
578                }
579                items.push(Item { req: i, q, seq, logits: Vec::new(), act: [0.0; 2] });
580            }
581        }
582        Ok(items)
583    }
584
585    /// Fills the items the cache holds, for the requests that read it. It gives the items left
586    /// to run, and every item's key, none when the cache is off.
587    fn look_up(&self, reqs: &[Request], items: &mut [Item<'_>]) -> (Vec<usize>, Vec<cache::Key>) {
588        let Some(c) = &self.inner.cache else { return ((0..items.len()).collect(), Vec::new()) };
589        let keys: Vec<_> = items.iter().map(Item::key).collect();
590        let look: Vec<usize> =
591            (0..items.len()).filter(|&i| mode(&reqs[items[i].req]) == CacheMode::Use).collect();
592        let found = c.get(&look.iter().map(|&i| keys[i]).collect::<Vec<_>>());
593        let mut hit = vec![false; items.len()];
594        for (&i, e) in look.iter().zip(found) {
595            if let Some(e) = e {
596                items[i].fill(e);
597                hit[i] = true;
598            }
599        }
600        ((0..items.len()).filter(|&i| !hit[i]).collect(), keys)
601    }
602
603    /// The responses, from items whose logits are in.
604    fn respond(&self, parsed: &[Request], items: &[Item<'_>]) -> Vec<Response> {
605        let inner = &*self.inner;
606        let mut out: Vec<Response> = parsed
607            .iter()
608            .map(|_| Response { model: LAYA_MODEL.into(), answers: Vec::new(), input_tokens: 0 })
609            .collect();
610        for it in items {
611            let res = &mut out[it.req];
612            res.input_tokens += it.seq.ids.len();
613            res.answers
614                .push((it.q.id.clone(), laya_answer(it.q, &it.logits, it.act, &inner.temps)));
615        }
616        out
617    }
618
619    /// Runs the items at `which` in the batches [`split::split`] picks, and says how many it
620    /// took.
621    fn run(&self, items: &mut [Item<'_>], which: &[usize]) -> Result<usize, Error> {
622        let sizes: Vec<(usize, usize)> =
623            which.iter().map(|&i| (items[i].seq.ids.len(), items[i].seq.markers.len())).collect();
624        let batches: Vec<Vec<usize>> = split::split(&self.inner.buckets, &sizes)
625            .into_iter()
626            .map(|b| b.into_iter().map(|j| which[j]).collect())
627            .collect();
628        let mut s = self.lock();
629        let Session { runner, buf, out, .. } = &mut *s;
630        for batch in &batches {
631            buf.clear();
632            for &i in batch {
633                let it = &items[i];
634                buf.push(&it.seq.ids, &it.seq.markers, it.q.qtype.index() as u8);
635            }
636            runner.run(buf, out)?;
637            let mut at = 0;
638            for (&i, a) in batch.iter().zip(&out.act) {
639                let it = &mut items[i];
640                let k = it.seq.markers.len();
641                it.logits = out.logits[at..at + k].to_vec();
642                it.act = *a;
643                at += k;
644            }
645        }
646        let (weights, plans) = runner.memory();
647        self.inner.memory[0].store(weights, Ordering::Relaxed);
648        self.inner.memory[1].store(plans, Ordering::Relaxed);
649        Ok(batches.len())
650    }
651
652    /// Each text's encoder output mean pooled over its tokens, as Laya's `embed_fn_from_agent`
653    /// computes it: `[CLS]`, the text cut to `max_length` tokens with the specials, `[SEP]`. It
654    /// runs no decision head. Each row is as wide as the encoder.
655    ///
656    /// # Errors
657    ///
658    /// Device errors, and [`Error::Unsupported`] on a backend without the pooled graph or in
659    /// INT8, which puts some rows far from Laya's.
660    pub fn embed(&self, texts: &[&str], max_length: usize) -> Result<Vec<Vec<f32>>, Error> {
661        let inner = &*self.inner;
662        if self.lock().runner.int8() {
663            // On 125 texts the worst row had a cosine of 0.53 to Laya's, too far to shortlist on.
664            return Err(Error::Unsupported("embeddings need FP32 or FP16, not INT8".into()));
665        }
666        let sp = inner.tok.specials();
667        let keep = max_length.saturating_sub(2);
668        let seqs: Vec<Vec<u32>> = texts
669            .iter()
670            .map(|t| {
671                let mut ids = Vec::with_capacity(keep.min(t.len()) + 2);
672                ids.push(sp.cls);
673                inner.tok.encode_into(t, &mut ids);
674                ids.truncate(keep + 1);
675                ids.push(sp.sep);
676                ids
677            })
678            .collect();
679        let sizes: Vec<(usize, usize)> = seqs.iter().map(|s| (s.len(), 0)).collect();
680        let batches = split::split(&inner.embed_buckets, &sizes);
681        let mut rows = vec![Vec::new(); texts.len()];
682        let mut s = self.lock();
683        let Session { runner, embed, buf, out } = &mut *s;
684        for batch in &batches {
685            buf.clear();
686            for &i in batch {
687                buf.push(&seqs[i], &[], 0);
688            }
689            runner.run_lane(*embed, buf, out)?;
690            for (&i, row) in batch.iter().zip(out.pooled.chunks_exact(inner.d)) {
691                rows[i] = row.to_vec();
692            }
693        }
694        let (weights, plans) = runner.memory();
695        inner.memory[0].store(weights, Ordering::Relaxed);
696        inner.memory[1].store(plans, Ordering::Relaxed);
697        Ok(rows)
698    }
699
700    /// [`Kime::decide`] on a thread of its own, for async callers. It works with any executor,
701    /// since it needs nothing from one but a waker.
702    #[must_use]
703    pub fn decide_async(&self, req: &Request) -> Decision {
704        let shared = Arc::new(Mutex::new((None, None::<Waker>)));
705        let (kime, req, done) = (self.clone(), req.clone(), shared.clone());
706        std::thread::spawn(move || {
707            let r = kime.decide(&req);
708            let mut g = done.lock().unwrap_or_else(PoisonError::into_inner);
709            g.0 = Some(r);
710            if let Some(w) = g.1.take() {
711                w.wake();
712            }
713        });
714        Decision { shared }
715    }
716}
717
718/// The future [`Kime::decide_async`] returns.
719#[derive(Debug)]
720pub struct Decision {
721    #[allow(clippy::type_complexity)]
722    shared: Arc<Mutex<(Option<Result<Response, Error>>, Option<Waker>)>>,
723}
724
725impl Future for Decision {
726    type Output = Result<Response, Error>;
727
728    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
729        let mut g = self.shared.lock().unwrap_or_else(PoisonError::into_inner);
730        match g.0.take() {
731            Some(r) => Poll::Ready(r),
732            None => {
733                g.1 = Some(cx.waker().clone());
734                Poll::Pending
735            }
736        }
737    }
738}
kime_engine/lib.rs

kime_engine/
lib.rs
