kime_tensor/
backend.rs

1//! The `Backend` trait and the executor that picks a bucket and replays its plan.
2
3use std::fmt;
4
5use crate::bucket::{Bucket, Buckets};
6use crate::dtype::DType;
7use crate::plan::Graph;
8
9/// A weight as it sits in the checkpoint, before a backend converts it.
10#[derive(Debug, Clone, Copy)]
11pub struct HostTensor<'a> {
12    /// Element type.
13    pub dtype: DType,
14    /// Shape.
15    pub shape: &'a [usize],
16    /// Little endian, row major.
17    pub bytes: &'a [u8],
18}
19
20/// A batch, flattened. Sequence `s` is tokens `cu[s]..cu[s + 1]` and markers
21/// `mcu[s]..mcu[s + 1]`, and each marker is a position within its own sequence.
22#[derive(Debug, Clone, Copy, Default)]
23pub struct Batch<'a> {
24    /// Token ids of every sequence, end to end.
25    pub ids: &'a [u32],
26    /// Sequence starts in `ids`, with the total at the end.
27    pub cu: &'a [u32],
28    /// Marker positions of every sequence, end to end.
29    pub markers: &'a [u32],
30    /// Sequence starts in `markers`, with the total at the end.
31    pub mcu: &'a [u32],
32    /// The question type of each sequence.
33    pub qtype: &'a [u8],
34}
35
36impl Batch<'_> {
37    /// Sequences.
38    #[must_use]
39    pub fn seqs(&self) -> usize {
40        self.cu.len().saturating_sub(1)
41    }
42
43    /// Checks that the pieces agree with each other, so a backend can index without checks.
44    ///
45    /// # Errors
46    ///
47    /// [`Error::Batch`] naming the first problem.
48    pub fn check(&self, vocab: usize, types: usize) -> Result<()> {
49        let bad = |m: String| Err(Error::Batch(m));
50        let s = self.seqs();
51        if self.cu.first() != Some(&0) || self.mcu.first() != Some(&0) {
52            return bad("cu and mcu must start at 0".into());
53        }
54        if self.mcu.len() != s + 1 || self.qtype.len() != s {
55            return bad(format!(
56                "{s} sequences but {} mcu and {} qtype",
57                self.mcu.len(),
58                self.qtype.len()
59            ));
60        }
61        if self.cu[s] as usize != self.ids.len() || self.mcu[s] as usize != self.markers.len() {
62            return bad("cu and mcu must end at the totals".into());
63        }
64        for i in 0..s {
65            let (a, b) = (self.cu[i], self.cu[i + 1]);
66            if b < a || self.mcu[i + 1] < self.mcu[i] {
67                return bad(format!("sequence {i} ends before it starts"));
68            }
69            if usize::from(self.qtype[i]) >= types {
70                return bad(format!("sequence {i} has type {} of {types}", self.qtype[i]));
71            }
72            for &m in &self.markers[self.mcu[i] as usize..self.mcu[i + 1] as usize] {
73                if m >= b - a {
74                    return bad(format!("marker {m} is past sequence {i} of {} tokens", b - a));
75                }
76            }
77        }
78        if let Some(&id) = self.ids.iter().find(|&&id| id as usize >= vocab) {
79            return bad(format!("token id {id} is past the vocabulary of {vocab}"));
80        }
81        Ok(())
82    }
83}
84
85/// What a batch produces. Reused across batches, so after the first few it does not allocate.
86#[derive(Debug, Clone, Default, PartialEq)]
87pub struct Outputs {
88    /// One logit per marker, in batch order.
89    pub logits: Vec<f32>,
90    /// The act head, one pair per sequence.
91    pub act: Vec<[f32; 2]>,
92    /// The pooled embedding, one row per sequence, for a graph that has one.
93    pub pooled: Vec<f32>,
94}
95
96/// What a backend can do.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct Caps {
99    /// `cpu`, `cuda`, `metal` or `ane`.
100    pub name: &'static str,
101    /// Worker threads, for the CPU.
102    pub threads: usize,
103    /// Whether plans are captured as device graphs, which makes padding to the bucket matter.
104    pub graphs: bool,
105    /// Whether host and device share memory.
106    pub unified_memory: bool,
107}
108
109/// Engine errors.
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub enum Error {
112    /// The batch is inconsistent.
113    Batch(String),
114    /// No bucket of the stage holds the batch.
115    NoBucket {
116        /// The stage.
117        stage: String,
118        /// Tokens in the batch.
119        tokens: usize,
120        /// Sequences.
121        seqs: usize,
122        /// Markers.
123        markers: usize,
124    },
125    /// The graph uses something the backend does not have, or weights of the wrong shape.
126    Unsupported(String),
127    /// The device or its driver failed.
128    Device(String),
129}
130
131impl fmt::Display for Error {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        match self {
134            Error::Batch(m) => write!(f, "bad batch: {m}"),
135            Error::NoBucket { stage, tokens, seqs, markers } => write!(
136                f,
137                "no {stage} bucket holds {tokens} tokens, {seqs} sequences and {markers} markers"
138            ),
139            Error::Unsupported(m) => write!(f, "unsupported: {m}"),
140            Error::Device(m) => write!(f, "device: {m}"),
141        }
142    }
143}
144
145impl std::error::Error for Error {}
146
147/// Engine results.
148pub type Result<T> = std::result::Result<T, Error>;
149
150/// A device that runs graphs. The engine is generic over it, so the path through it is
151/// monomorphized and there is no dynamic dispatch per op.
152pub trait Backend: Send + Sync {
153    /// Weights in the backend's own layout.
154    type Weights: Send + Sync;
155    /// A graph lowered for one bucket, with its arena.
156    type Plan: Send;
157
158    /// What the backend can do.
159    fn caps(&self) -> Caps;
160
161    /// Converts a checkpoint's tensors, indexed as the graph's weights are.
162    ///
163    /// # Errors
164    ///
165    /// When a tensor cannot be converted.
166    fn upload(&self, tensors: &[HostTensor<'_>], graph: &Graph) -> Result<Self::Weights>;
167
168    /// Lowers a graph for one bucket. This is where memory is allocated.
169    ///
170    /// # Errors
171    ///
172    /// [`Error::Unsupported`] for an op or a shape the backend cannot run.
173    fn lower(&self, w: &Self::Weights, graph: &Graph, bucket: Bucket) -> Result<Self::Plan>;
174
175    /// Runs a batch that fits the plan's bucket and has been checked.
176    ///
177    /// # Errors
178    ///
179    /// When the device fails.
180    fn run(&self, plan: &mut Self::Plan, batch: &Batch<'_>, out: &mut Outputs) -> Result<()>;
181
182    /// Bytes the weights take on the device.
183    fn weight_bytes(&self, _w: &Self::Weights) -> usize {
184        0
185    }
186
187    /// Bytes a plan takes on the device, its arena and staging buffers.
188    fn plan_bytes(&self, _p: &Self::Plan) -> usize {
189        0
190    }
191}
192
193/// A graph, its weights on one backend, and a plan per bucket built on first use. More graphs
194/// can share the weights, see [`Executor::add_graph`].
195#[derive(Debug)]
196pub struct Executor<B: Backend> {
197    backend: B,
198    weights: B::Weights,
199    /// The graph the weights were uploaded for first, then any added.
200    lanes: Vec<Lane<B>>,
201    vocab: usize,
202    types: usize,
203}
204
205/// One graph on an executor's weights, the buckets it runs in and its plans.
206struct Lane<B: Backend> {
207    graph: Graph,
208    stage: String,
209    buckets: Vec<Bucket>,
210    plans: Vec<(Bucket, B::Plan)>,
211}
212
213impl<B: Backend> fmt::Debug for Lane<B> {
214    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
215        let warm: Vec<Bucket> = self.plans.iter().map(|p| p.0).collect();
216        f.debug_struct("Lane")
217            .field("stage", &self.stage)
218            .field("warm", &warm)
219            .finish_non_exhaustive()
220    }
221}
222
223impl<B: Backend> Lane<B> {
224    fn new(graph: Graph, buckets: &Buckets, stage: &str) -> Self {
225        Self {
226            graph,
227            stage: stage.to_string(),
228            buckets: buckets.stage(stage).to_vec(),
229            plans: Vec::new(),
230        }
231    }
232}
233
234impl<B: Backend> Executor<B> {
235    /// Uploads the weights. `vocab` and `types` bound the token ids and question types a batch may
236    /// hold, and `stage` names the buckets to use from `buckets`.
237    ///
238    /// # Errors
239    ///
240    /// From [`Backend::upload`].
241    pub fn new(
242        backend: B,
243        tensors: &[HostTensor<'_>],
244        graph: Graph,
245        buckets: &Buckets,
246        stage: &str,
247        vocab: usize,
248        types: usize,
249    ) -> Result<Self> {
250        let weights = backend.upload(tensors, &graph)?;
251        Ok(Self { backend, weights, lanes: vec![Lane::new(graph, buckets, stage)], vocab, types })
252    }
253
254    /// Adds a graph that runs on the weights already uploaded, in the buckets of `stage`, and
255    /// returns the lane to pass to [`Executor::run_lane`]. The graph may only read weights the
256    /// first graph reads, the way the first graph reads them, such as the encoder alone.
257    pub fn add_graph(&mut self, graph: Graph, buckets: &Buckets, stage: &str) -> usize {
258        self.lanes.push(Lane::new(graph, buckets, stage));
259        self.lanes.len() - 1
260    }
261
262    /// The backend.
263    pub fn backend(&self) -> &B {
264        &self.backend
265    }
266
267    /// The buckets with a plan built.
268    pub fn warm(&self) -> impl Iterator<Item = Bucket> + '_ {
269        self.lanes[0].plans.iter().map(|p| p.0)
270    }
271
272    /// Bytes on the device: the weights, and the plans built so far.
273    pub fn memory(&self) -> (usize, usize) {
274        let plans =
275            self.lanes.iter().flat_map(|l| &l.plans).map(|p| self.backend.plan_bytes(&p.1)).sum();
276        (self.backend.weight_bytes(&self.weights), plans)
277    }
278
279    /// The plans built so far, to inspect or profile.
280    pub fn plans_mut(&mut self) -> impl Iterator<Item = &mut B::Plan> + '_ {
281        self.lanes[0].plans.iter_mut().map(|p| &mut p.1)
282    }
283
284    /// Builds the plan for `bucket` now rather than on first use.
285    ///
286    /// # Errors
287    ///
288    /// From [`Backend::lower`].
289    pub fn prepare(&mut self, bucket: Bucket) -> Result<usize> {
290        self.prepare_lane(0, bucket)
291    }
292
293    fn prepare_lane(&mut self, lane: usize, bucket: Bucket) -> Result<usize> {
294        let l = &mut self.lanes[lane];
295        if let Some(i) = l.plans.iter().position(|p| p.0 == bucket) {
296            return Ok(i);
297        }
298        let plan = self.backend.lower(&self.weights, &l.graph, bucket)?;
299        l.plans.push((bucket, plan));
300        Ok(l.plans.len() - 1)
301    }
302
303    /// Runs a batch in the smallest bucket that holds it. Once that bucket's plan exists and `out`
304    /// has grown to the batch size, this allocates nothing.
305    ///
306    /// # Errors
307    ///
308    /// [`Error::Batch`] for an inconsistent batch, [`Error::NoBucket`] for one too big, and
309    /// anything the backend reports.
310    pub fn run(&mut self, batch: &Batch<'_>, out: &mut Outputs) -> Result<Bucket> {
311        self.run_lane(0, batch, out)
312    }
313
314    /// [`Executor::run`] on the graph `lane` from [`Executor::add_graph`], 0 being the first.
315    ///
316    /// # Errors
317    ///
318    /// As [`Executor::run`].
319    pub fn run_lane(
320        &mut self,
321        lane: usize,
322        batch: &Batch<'_>,
323        out: &mut Outputs,
324    ) -> Result<Bucket> {
325        batch.check(self.vocab, self.types)?;
326        let (t, s, m) = (batch.ids.len(), batch.seqs(), batch.markers.len());
327        let l = &self.lanes[lane];
328        let bucket = l.buckets.iter().copied().find(|b| b.holds(t, s, m)).ok_or_else(|| {
329            Error::NoBucket { stage: l.stage.clone(), tokens: t, seqs: s, markers: m }
330        })?;
331        let i = self.prepare_lane(lane, bucket)?;
332        self.backend.run(&mut self.lanes[lane].plans[i].1, batch, out)?;
333        Ok(bucket)
334    }
335}
336
337/// Owned storage for a [`Batch`], reused from one batch to the next so that building one does
338/// not allocate once it has grown.
339#[derive(Debug, Clone, Default)]
340pub struct BatchBuf {
341    ids: Vec<u32>,
342    cu: Vec<u32>,
343    markers: Vec<u32>,
344    mcu: Vec<u32>,
345    qtype: Vec<u8>,
346}
347
348impl BatchBuf {
349    /// Empties it, keeping the memory.
350    pub fn clear(&mut self) {
351        for v in [&mut self.ids, &mut self.cu, &mut self.markers, &mut self.mcu] {
352            v.clear();
353        }
354        self.qtype.clear();
355    }
356
357    /// Appends a sequence.
358    ///
359    /// # Panics
360    ///
361    /// If the batch passes 2^32 tokens or markers.
362    pub fn push(&mut self, ids: &[u32], markers: &[u32], qtype: u8) {
363        if self.cu.is_empty() {
364            self.cu.push(0);
365            self.mcu.push(0);
366        }
367        self.ids.extend_from_slice(ids);
368        self.markers.extend_from_slice(markers);
369        self.cu.push(u32::try_from(self.ids.len()).expect("under 2^32 tokens"));
370        self.mcu.push(u32::try_from(self.markers.len()).expect("under 2^32 markers"));
371        self.qtype.push(qtype);
372    }
373
374    /// The batch.
375    #[must_use]
376    pub fn batch(&self) -> Batch<'_> {
377        const EMPTY: &[u32] = &[0];
378        let (cu, mcu) =
379            if self.cu.is_empty() { (EMPTY, EMPTY) } else { (&self.cu[..], &self.mcu[..]) };
380        Batch { ids: &self.ids, cu, markers: &self.markers, mcu, qtype: &self.qtype }
381    }
382}
kime_tensor/backend.rs

kime_tensor/
backend.rs
