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, through kime-metal. macOS only, FP32 or FP16.
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    #[cfg(target_os = "macos")]
230    Metal(Box<Executor<kime_metal::MetalBackend>>),
231}
232
233impl Runner {
234    fn open(model: &Model, device: Device, precision: Precision) -> Result<Self, Error> {
235        let cpu = |threads: usize| {
236            let t = if threads == 0 { kime_cpu::par::available() } else { threads };
237            let backend = kime_cpu::CpuBackend::new(t).with_int8(precision == Precision::Int8);
238            Ok(Runner::Cpu(Box::new(kime_cpu::executor_with(model, backend)?)))
239        };
240        match device {
241            Device::Cpu { threads } => cpu(threads),
242            #[cfg(feature = "cuda")]
243            Device::Cuda(n) => {
244                Ok(Runner::Cuda(Box::new(kime_cuda::executor(model, n, cuda(precision)?)?)))
245            }
246            Device::Auto => Runner::auto(model, precision).map_or_else(|| cpu(0), Ok),
247            #[cfg(not(feature = "cuda"))]
248            Device::Cuda(_) => Err(Error::Unsupported("this build has no CUDA backend".into())),
249            #[cfg(target_os = "macos")]
250            Device::Metal => {
251                Ok(Runner::Metal(Box::new(kime_metal::executor(model, metal(precision)?)?)))
252            }
253            #[cfg(not(target_os = "macos"))]
254            Device::Metal => Err(Error::Unsupported("Metal needs macOS".into())),
255            Device::Ane => {
256                Err(Error::Unsupported("the Neural Engine backend arrives with M4".into()))
257            }
258        }
259    }
260
261    /// A GPU when there is one: CUDA, then on a Mac the Apple GPU, as Laya picks cuda, then mps.
262    /// `None` sends `Device::Auto` to the CPU, which is also where INT8 runs.
263    #[allow(unused_variables)]
264    fn auto(model: &Model, precision: Precision) -> Option<Self> {
265        #[cfg(feature = "cuda")]
266        if let Ok(e) = cuda(precision).and_then(|p| Ok(kime_cuda::executor(model, 0, p)?)) {
267            return Some(Runner::Cuda(Box::new(e)));
268        }
269        #[cfg(target_os = "macos")]
270        if let Ok(e) = metal(precision).and_then(|p| Ok(kime_metal::executor(model, p)?)) {
271            return Some(Runner::Metal(Box::new(e)));
272        }
273        None
274    }
275
276    fn add_graph(&mut self, g: kime_tensor::Graph) -> usize {
277        let b = Buckets::default();
278        match self {
279            Runner::Cpu(e) => e.add_graph(g, &b, EMBED_STAGE),
280            #[cfg(feature = "cuda")]
281            Runner::Cuda(e) => e.add_graph(g, &b, EMBED_STAGE),
282            #[cfg(target_os = "macos")]
283            Runner::Metal(e) => e.add_graph(g, &b, EMBED_STAGE),
284        }
285    }
286
287    fn prepare(&mut self, b: kime_tensor::Bucket) -> Result<(), Error> {
288        match self {
289            Runner::Cpu(e) => e.prepare(b)?,
290            #[cfg(feature = "cuda")]
291            Runner::Cuda(e) => e.prepare(b)?,
292            #[cfg(target_os = "macos")]
293            Runner::Metal(e) => e.prepare(b)?,
294        };
295        Ok(())
296    }
297
298    fn run(&mut self, buf: &BatchBuf, out: &mut Outputs) -> Result<(), Error> {
299        self.run_lane(0, buf, out)
300    }
301
302    fn run_lane(&mut self, lane: usize, buf: &BatchBuf, out: &mut Outputs) -> Result<(), Error> {
303        match self {
304            Runner::Cpu(e) => e.run_lane(lane, &buf.batch(), out)?,
305            #[cfg(feature = "cuda")]
306            Runner::Cuda(e) => e.run_lane(lane, &buf.batch(), out)?,
307            #[cfg(target_os = "macos")]
308            Runner::Metal(e) => e.run_lane(lane, &buf.batch(), out)?,
309        };
310        Ok(())
311    }
312
313    fn memory(&self) -> (usize, usize) {
314        match self {
315            Runner::Cpu(e) => e.memory(),
316            #[cfg(feature = "cuda")]
317            Runner::Cuda(e) => e.memory(),
318            #[cfg(target_os = "macos")]
319            Runner::Metal(e) => e.memory(),
320        }
321    }
322
323    fn int8(&self) -> bool {
324        match self {
325            Runner::Cpu(e) => e.backend().int8(),
326            #[cfg(feature = "cuda")]
327            Runner::Cuda(_) => false,
328            #[cfg(target_os = "macos")]
329            Runner::Metal(_) => false,
330        }
331    }
332
333    fn describe(&self) -> String {
334        match self {
335            Runner::Cpu(e) => {
336                let int8 = if e.backend().int8() { ", int8" } else { "" };
337                format!("cpu, {} threads{int8}", kime_tensor::Backend::caps(e.backend()).threads)
338            }
339            #[cfg(feature = "cuda")]
340            Runner::Cuda(e) => format!("cuda, {}", e.backend().name()),
341            #[cfg(target_os = "macos")]
342            Runner::Metal(e) => format!("metal, {}", e.backend().name()),
343        }
344    }
345}
346
347#[cfg(target_os = "macos")]
348fn metal(p: Precision) -> Result<kime_metal::Precision, Error> {
349    match p {
350        Precision::F16 => Ok(kime_metal::Precision::F16),
351        Precision::F32 => Ok(kime_metal::Precision::F32),
352        Precision::Int8 => Err(Error::Unsupported("INT8 runs on the CPU only for now".into())),
353    }
354}
355
356#[cfg(feature = "cuda")]
357fn cuda(p: Precision) -> Result<kime_cuda::Precision, Error> {
358    match p {
359        Precision::F16 => Ok(kime_cuda::Precision::F16),
360        Precision::F32 => Ok(kime_cuda::Precision::F32),
361        Precision::Int8 => Err(Error::Unsupported("INT8 runs on the CPU only for now".into())),
362    }
363}
364
365/// The buckets the pooled embedding runs in: sequences with no markers.
366const EMBED_STAGE: &str = "state";
367
368struct Session {
369    runner: Runner,
370    /// The lane of the pooled embedding graph on the runner.
371    embed: usize,
372    buf: BatchBuf,
373    out: Outputs,
374}
375
376struct Inner {
377    id: String,
378    tok: Tokenizer,
379    mask: String,
380    budget: CompatBudget,
381    temps: Temperatures,
382    /// The compat buckets, smallest first.
383    buckets: Vec<kime_tensor::Bucket>,
384    /// The buckets of the pooled embedding, smallest first.
385    embed_buckets: Vec<kime_tensor::Bucket>,
386    /// The width of the encoder, and of an embedding.
387    d: usize,
388    runner: Mutex<Session>,
389    /// [`Memory`], kept up to date after every forward pass so reading it needs no lock.
390    memory: [AtomicUsize; 2],
391    cache: Option<AnswerCache>,
392}
393
394/// A loaded model on a device. Clones share it, and it can be used from any thread.
395#[derive(Clone)]
396pub struct Kime {
397    inner: Arc<Inner>,
398}
399
400impl fmt::Debug for Kime {
401    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
402        f.debug_struct("Kime").field("model", &self.inner.id).finish_non_exhaustive()
403    }
404}
405
406/// Where the time of one [`Kime::decide_batch_timed`] call went.
407#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
408pub struct Timing {
409    /// Validating the requests and laying their questions out as token ids.
410    pub tokenize: Duration,
411    /// The device batches, from the first upload to the last result copied back.
412    pub device: Duration,
413    /// How many device batches the questions took.
414    pub batches: usize,
415    /// Questions whose state was cut to fit the model's sequence length.
416    pub truncated: usize,
417    /// State tokens left out of those questions.
418    pub cut_tokens: usize,
419    /// Questions answered from the answer cache, which took no device time.
420    pub cached: usize,
421}
422
423/// Bytes a model holds on its device.
424#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
425pub struct Memory {
426    /// The weights, in the device's layout.
427    pub weights: usize,
428    /// The plans built so far, one per bucket used: their arenas and staging buffers.
429    pub plans: usize,
430}
431
432/// One laid out question and where its answer goes.
433struct Item<'a> {
434    req: usize,
435    q: &'a Question,
436    seq: CompatSequence,
437    logits: Vec<f32>,
438    act: [f32; 2],
439}
440
441impl Item<'_> {
442    fn key(&self) -> cache::Key {
443        AnswerCache::key(&self.seq, self.q.qtype.index() as u8)
444    }
445
446    fn fill(&mut self, e: Entry) {
447        self.logits = e.logits.into_vec();
448        self.act = e.act;
449    }
450}
451
452/// The request's `kime.cache`, the default for anything but a mode's name.
453fn mode(req: &Request) -> CacheMode {
454    req.kime
455        .as_ref()
456        .and_then(|k| k.get("cache"))
457        .and_then(Value::as_str)
458        .and_then(CacheMode::parse)
459        .unwrap_or_default()
460}
461
462impl Kime {
463    /// Settings for a new engine.
464    #[must_use]
465    pub fn builder() -> Builder {
466        Builder::default()
467    }
468
469    /// The model's name, `laya` or `laya-multilingual` for the published checkpoints.
470    #[must_use]
471    pub fn model_id(&self) -> &str {
472        &self.inner.id
473    }
474
475    /// The device the model runs on, in words.
476    #[must_use]
477    pub fn device(&self) -> String {
478        self.lock().runner.describe()
479    }
480
481    /// The most tokens one question's row can hold: the state, the question and its options.
482    /// Longer states are cut to fit.
483    #[must_use]
484    pub fn max_row_tokens(&self) -> usize {
485        self.inner.budget.max_len
486    }
487
488    /// The bytes the model holds on its device, as of the last forward pass.
489    #[must_use]
490    pub fn memory(&self) -> Memory {
491        let m = &self.inner.memory;
492        Memory { weights: m[0].load(Ordering::Relaxed), plans: m[1].load(Ordering::Relaxed) }
493    }
494
495    /// The answer cache's counts, all 0 when it is off.
496    #[must_use]
497    pub fn cache_stats(&self) -> CacheStats {
498        self.inner.cache.as_ref().map(AnswerCache::stats).unwrap_or_default()
499    }
500
501    /// The answer to `req` when the answer cache holds every one of its questions, found without
502    /// the device lock, so the server can answer it without queueing it. `None` otherwise, and
503    /// then nothing is counted, since the request goes on to [`Kime::decide_batch`].
504    #[must_use]
505    pub fn cached(&self, req: &Request) -> Option<Response> {
506        let cache = self.inner.cache.as_ref()?;
507        if mode(req) != CacheMode::Use || req.questions.is_empty() {
508            return None;
509        }
510        let parsed = [parse(&req.to_json(), &Limits::LAYA).ok()?];
511        let mut items = self.lay_out(&parsed).ok()?;
512        let keys: Vec<_> = items.iter().map(Item::key).collect();
513        for (it, e) in items.iter_mut().zip(cache.all(&keys)?) {
514            it.fill(e);
515        }
516        self.respond(&parsed, &items).pop()
517    }
518
519    fn lock(&self) -> std::sync::MutexGuard<'_, Session> {
520        self.inner.runner.lock().unwrap_or_else(PoisonError::into_inner)
521    }
522
523    /// The tokens a request holds before anything is cut: its state once plus every question
524    /// with its options. The server refuses a request over its limit with this count, before it
525    /// is queued.
526    #[must_use]
527    pub fn count_tokens(&self, req: &Request) -> usize {
528        let inner = &*self.inner;
529        let mut n = inner.tok.encode(&compat_state(&req.state, &inner.mask)).len();
530        for q in &req.questions {
531            let text = compat_question(q, &inner.mask);
532            n += inner.tok.encode(&text.head).len();
533            n += text.options.iter().map(|o| inner.tok.encode(o).len()).sum::<usize>();
534        }
535        n
536    }
537
538    /// Answers one request.
539    ///
540    /// # Errors
541    ///
542    /// [`Error::Invalid`] for a request that does not validate, [`Error::TooLong`] for a question
543    /// whose options do not fit, and device errors.
544    pub fn decide(&self, req: &Request) -> Result<Response, Error> {
545        Ok(self.decide_batch(std::slice::from_ref(req))?.remove(0))
546    }
547
548    /// Answers many requests, packing all their questions into as few device batches as fit.
549    /// Each answer is the same bits it would be alone, whatever else is in the batch.
550    ///
551    /// # Errors
552    ///
553    /// As [`Kime::decide`]. One bad request fails the whole call.
554    pub fn decide_batch(&self, reqs: &[Request]) -> Result<Vec<Response>, Error> {
555        Ok(self.decide_batch_timed(reqs)?.0)
556    }
557
558    /// [`Kime::decide_batch`], and where the time went.
559    ///
560    /// # Errors
561    ///
562    /// As [`Kime::decide_batch`].
563    pub fn decide_batch_timed(&self, reqs: &[Request]) -> Result<(Vec<Response>, Timing), Error> {
564        let t0 = Instant::now();
565        let mut parsed = Vec::with_capacity(reqs.len());
566        for r in reqs {
567            parsed.push(parse(&r.to_json(), &Limits::LAYA).map_err(Error::Invalid)?);
568        }
569        let mut items = self.lay_out(&parsed)?;
570        let tokenize = t0.elapsed();
571        let t1 = Instant::now();
572        let (run, keys) = self.look_up(reqs, &mut items);
573        let batches = if run.is_empty() { 0 } else { self.run(&mut items, &run)? };
574        if let Some(c) = &self.inner.cache {
575            let keep = run.iter().filter(|&&i| mode(&reqs[items[i].req]) != CacheMode::Bypass);
576            c.insert(keep.map(|&i| {
577                (keys[i], Entry { logits: items[i].logits.clone().into(), act: items[i].act })
578            }));
579        }
580        let cached = items.len() - run.len();
581        let cut = items.iter().map(|it| it.seq.state_tokens - it.seq.state_tokens_used);
582        let timing = Timing {
583            tokenize,
584            device: t1.elapsed(),
585            batches,
586            truncated: cut.clone().filter(|&n| n > 0).count(),
587            cut_tokens: cut.sum(),
588            cached,
589        };
590        Ok((self.respond(&parsed, &items), timing))
591    }
592
593    /// Each question of `parsed` laid out as Laya lays it out, in request order.
594    fn lay_out<'a>(&self, parsed: &'a [Request]) -> Result<Vec<Item<'a>>, Error> {
595        let inner = &*self.inner;
596        let mut items = Vec::new();
597        for (i, r) in parsed.iter().enumerate() {
598            if r.questions.is_empty() {
599                continue;
600            }
601            // Laya keeps the end of a conversation and the start of anything else.
602            let cut = if matches!(r.state, Value::Array(_)) { Cut::Head } else { Cut::Tail };
603            let state = inner.tok.encode_state(&compat_state(&r.state, &inner.mask));
604            for q in &r.questions {
605                let text = compat_question(q, &inner.mask);
606                let seq =
607                    inner.tok.compat_sequence(&text.head, &text.options, &state, inner.budget, cut);
608                if seq.markers.len() != q.criteria.len() {
609                    return Err(Error::TooLong {
610                        question: q.id.clone(),
611                        options: q.criteria.len(),
612                        fit: seq.markers.len(),
613                    });
614                }
615                items.push(Item { req: i, q, seq, logits: Vec::new(), act: [0.0; 2] });
616            }
617        }
618        Ok(items)
619    }
620
621    /// Fills the items the cache holds, for the requests that read it. It gives the items left
622    /// to run, and every item's key, none when the cache is off.
623    fn look_up(&self, reqs: &[Request], items: &mut [Item<'_>]) -> (Vec<usize>, Vec<cache::Key>) {
624        let Some(c) = &self.inner.cache else { return ((0..items.len()).collect(), Vec::new()) };
625        let keys: Vec<_> = items.iter().map(Item::key).collect();
626        let look: Vec<usize> =
627            (0..items.len()).filter(|&i| mode(&reqs[items[i].req]) == CacheMode::Use).collect();
628        let found = c.get(&look.iter().map(|&i| keys[i]).collect::<Vec<_>>());
629        let mut hit = vec![false; items.len()];
630        for (&i, e) in look.iter().zip(found) {
631            if let Some(e) = e {
632                items[i].fill(e);
633                hit[i] = true;
634            }
635        }
636        ((0..items.len()).filter(|&i| !hit[i]).collect(), keys)
637    }
638
639    /// The responses, from items whose logits are in.
640    fn respond(&self, parsed: &[Request], items: &[Item<'_>]) -> Vec<Response> {
641        let inner = &*self.inner;
642        let mut out: Vec<Response> = parsed
643            .iter()
644            .map(|_| Response { model: LAYA_MODEL.into(), answers: Vec::new(), input_tokens: 0 })
645            .collect();
646        for it in items {
647            let res = &mut out[it.req];
648            res.input_tokens += it.seq.ids.len();
649            res.answers
650                .push((it.q.id.clone(), laya_answer(it.q, &it.logits, it.act, &inner.temps)));
651        }
652        out
653    }
654
655    /// Runs the items at `which` in the batches [`split::split`] picks, and says how many it
656    /// took.
657    fn run(&self, items: &mut [Item<'_>], which: &[usize]) -> Result<usize, Error> {
658        let sizes: Vec<(usize, usize)> =
659            which.iter().map(|&i| (items[i].seq.ids.len(), items[i].seq.markers.len())).collect();
660        let batches: Vec<Vec<usize>> = split::split(&self.inner.buckets, &sizes)
661            .into_iter()
662            .map(|b| b.into_iter().map(|j| which[j]).collect())
663            .collect();
664        let mut s = self.lock();
665        let Session { runner, buf, out, .. } = &mut *s;
666        for batch in &batches {
667            buf.clear();
668            for &i in batch {
669                let it = &items[i];
670                buf.push(&it.seq.ids, &it.seq.markers, it.q.qtype.index() as u8);
671            }
672            runner.run(buf, out)?;
673            let mut at = 0;
674            for (&i, a) in batch.iter().zip(&out.act) {
675                let it = &mut items[i];
676                let k = it.seq.markers.len();
677                it.logits = out.logits[at..at + k].to_vec();
678                it.act = *a;
679                at += k;
680            }
681        }
682        let (weights, plans) = runner.memory();
683        self.inner.memory[0].store(weights, Ordering::Relaxed);
684        self.inner.memory[1].store(plans, Ordering::Relaxed);
685        Ok(batches.len())
686    }
687
688    /// Each text's encoder output mean pooled over its tokens, as Laya's `embed_fn_from_agent`
689    /// computes it: `[CLS]`, the text cut to `max_length` tokens with the specials, `[SEP]`. It
690    /// runs no decision head. Each row is as wide as the encoder.
691    ///
692    /// # Errors
693    ///
694    /// Device errors, and [`Error::Unsupported`] on a backend without the pooled graph or in
695    /// INT8, which puts some rows far from Laya's.
696    pub fn embed(&self, texts: &[&str], max_length: usize) -> Result<Vec<Vec<f32>>, Error> {
697        let inner = &*self.inner;
698        if self.lock().runner.int8() {
699            // On 125 texts the worst row had a cosine of 0.53 to Laya's, too far to shortlist on.
700            return Err(Error::Unsupported("embeddings need FP32 or FP16, not INT8".into()));
701        }
702        let sp = inner.tok.specials();
703        let keep = max_length.saturating_sub(2);
704        let seqs: Vec<Vec<u32>> = texts
705            .iter()
706            .map(|t| {
707                let mut ids = Vec::with_capacity(keep.min(t.len()) + 2);
708                ids.push(sp.cls);
709                inner.tok.encode_into(t, &mut ids);
710                ids.truncate(keep + 1);
711                ids.push(sp.sep);
712                ids
713            })
714            .collect();
715        let sizes: Vec<(usize, usize)> = seqs.iter().map(|s| (s.len(), 0)).collect();
716        let batches = split::split(&inner.embed_buckets, &sizes);
717        let mut rows = vec![Vec::new(); texts.len()];
718        let mut s = self.lock();
719        let Session { runner, embed, buf, out } = &mut *s;
720        for batch in &batches {
721            buf.clear();
722            for &i in batch {
723                buf.push(&seqs[i], &[], 0);
724            }
725            runner.run_lane(*embed, buf, out)?;
726            for (&i, row) in batch.iter().zip(out.pooled.chunks_exact(inner.d)) {
727                rows[i] = row.to_vec();
728            }
729        }
730        let (weights, plans) = runner.memory();
731        inner.memory[0].store(weights, Ordering::Relaxed);
732        inner.memory[1].store(plans, Ordering::Relaxed);
733        Ok(rows)
734    }
735
736    /// [`Kime::decide`] on a thread of its own, for async callers. It works with any executor,
737    /// since it needs nothing from one but a waker.
738    #[must_use]
739    pub fn decide_async(&self, req: &Request) -> Decision {
740        let shared = Arc::new(Mutex::new((None, None::<Waker>)));
741        let (kime, req, done) = (self.clone(), req.clone(), shared.clone());
742        std::thread::spawn(move || {
743            let r = kime.decide(&req);
744            let mut g = done.lock().unwrap_or_else(PoisonError::into_inner);
745            g.0 = Some(r);
746            if let Some(w) = g.1.take() {
747                w.wake();
748            }
749        });
750        Decision { shared }
751    }
752}
753
754/// The future [`Kime::decide_async`] returns.
755#[derive(Debug)]
756pub struct Decision {
757    #[allow(clippy::type_complexity)]
758    shared: Arc<Mutex<(Option<Result<Response, Error>>, Option<Waker>)>>,
759}
760
761impl Future for Decision {
762    type Output = Result<Response, Error>;
763
764    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
765        let mut g = self.shared.lock().unwrap_or_else(PoisonError::into_inner);
766        match g.0.take() {
767            Some(r) => Poll::Ready(r),
768            None => {
769                g.1 = Some(cx.waker().clone());
770                Poll::Pending
771            }
772        }
773    }
774}
kime_engine/lib.rs

kime_engine/
lib.rs
